Lazy compilation and kernel fusion in dynamic computation graphs

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for lazy compilation and kernel fusion in dynamic computation graphs. One of the operations is performed by generating an input graph based on translation of user code into an expression graph. The expression graph represents control flow dependencies of operations of the generated input graph. Optimization of the input graph is then performed by iterative application of optimization rules to the input graph. An optimized version of the input graph results from the application of the optimization rules. A transformation graph then is generated by comparing changes made from the original input graph to the final optimized version of the input graph. The transformation graph provides a blueprint such that the system may recreate the optimization of a similarly structured later generated input graph without having to reapply the optimization rules. During the optimization of the input graph, the system may also generate fused blocks of just-in-time operations which the system may later optimize for parallel processing on one or more graphic processing units, and/or hardware accelerated computation units.

BACKGROUND

Deep learning frameworks rely on two paradigms that are based on either static computation graphs or dynamic computation graphs. In many situations, a computation graph needs to be optimized such that the computation graph performs more efficiently and executes faster on a given hardware platform. With static computation graphs, operations to be performed by the computation graph are known in advance, and thus the operations can be easily revised and optimized prior to execution of the computation graph thereby improving performance.

However, with dynamic computation graphs, operations to be performed by underlying written code are not known ahead of time. The operations are determined dynamically as the processing of the dynamic computation graph occurs. In other words, the particular operations of a dynamic computation graph are not decided until execution of the computation graph with particular operations to be performed often depending on the contents or data of the dynamic computation graph. Current techniques and systems do not provide for effective optimization of dynamic computation graphs.

SUMMARY

Described herein is an exemplary system, method and a computer program product for lazy compilation and kernel fusion in dynamic computation graphs. The system, method and computer program product provide a lazy computation framework where code is initially written as to a desired computation graph. Graph operations are queued and then performed at the last possible moment when a result is needed. When that moment is called, then the queued operations are converted into a temporary input graph that may be optimized by the system to find opportunities to combine, reduce and/or delete certain operations thereby increasing the processing speed of the computation graph.

The system improves over existing frameworks regardless of whether the framework is dynamic or static. The system may perform optimizations as needed (i.e., “on the fly”), and can find groups of operations that can be compiled into unique kernels that would be used just this once for a computation in the form of what is called JITting. The system's JITting process is able to intelligently combine many sub-operations into a larger one (referred to herein as a JIT block) in a manner that has not been seen in the prior art. The system considers aspects of how operations may be combined together or simplified and determines the processing impact of combining those operations.

In general, one innovative aspect of the subject described in this specification can be embodied in systems, computer readable media, and methods that include operations for lazy compilation and kernel fusion in dynamic computation graphs. One of the operations is performed by generating an input graph based on translation of user code into an expression graph. The expression graph represents control flow dependencies of operations of the generated input graph. Optimization of the input graph is then performed by iterative application of optimization rules to the input graph. An optimized version of the input graph results from the application of the optimization rules. A transformation graph then is generated by comparing changes made from the original input graph to the final optimized version of the input graph. The transformation graph provides a blueprint such that the system may recreate the optimization of a similarly structured later generated input graph without having to reapply the optimization rules. During the optimization of the input graph, the system may also generate fused blocks of just-in-time operations (e.g., JIT blocks) which may later be optimized for parallel processing on one or more graphic processing units (GPUs), and/or hardware accelerated computation units (GPUs, Tensor Processing Units, Cerebras, Graphcore devices, etc.).

Further areas of applicability of the present disclosure will become apparent from the detailed description, the claims and the drawings. The detailed description and specific examples are intended for illustration only and are not intended to limit the scope of the disclosure.

DETAILED DESCRIPTION

FIG. 1illustrates a block diagram of an example system100utilized in performing lazy compilation and kernel fusion in dynamic computation graphs. The system100includes an input graph module104, an optimization module106, a transformation graph module108, an A* search module110, a JIT block module112, and application program interface (API) module114.

The input graph module104receives user code102as an input and translates the user code102into an input graph representing the computational operations of the user code102. The optimization module106performs an iterative optimization process against the input graph by applying a set of optimization rules. The system100generates a final optimized input graph representing the application of the set of optimization rules. The transformation graph module108compares an original state of the input graph to a final optimized state of the input graph and generates a transformation graph. The A* search module performs an optimization process to determine an optimal approach to parallel processing of JIT blocks operations. The API module114allows the system100to receive inputs and commands from other systems and/or programs and provide data in response to the received inputs and commands. Other systems, programs, processes, and/or software may bridge into the system100where the system100provides an external library of functionality that may perform real-time dynamic optimizations of computation graphs.

The system100includes one or more electronic data stores120,122,124,126,128that may be locally or remotely accessed. These data stores (e.g., libraries, databases, memory cache, file systems, etc.), include, but are not limited to, an input graph database120which may be used to store information about an input graph and revised iterations of the input graph, an optimization rules library122which may include optimization rules for optimizing the input graph; a transformation graph library124which may include data describing a transformation graph and associated hash values, a JIT Kernels database126which may store code for execution of JIT kernels, and an operation registry128which may store operations and information about a referenced library from which the operations may be executed. While the data stores120,122,124,126and128are displayed separately, the data and information maintained in a data store may be combined together or further separated in a manner the promotes retrieval and storage efficiency and/or data security. Moreover, the described modules and data stores may be configured in a distributed architecture where processing of the modules may occur on one or more user devices (such as a desktop, cell phone, tablet device), one or more separate servers or computers, cloud-based processing services and/or a combination thereof.

Process Overview

FIG. 2illustrates a flow chart of an example process for performing lazy compilation and kernel fusion in dynamic computation graphs. The system100receives user code and generates an input graph based on evaluation and translation of the received user code (block210). The user code is of the type used for computation graph processing and may be written in various programming languages (e.g., Python, C++, Ruby, JavaScript, etc.). The system100stores the generated input graph into an electronic data store.

The system100computes a hash value of the input graph and stores the hash value for later lookups of other generated input graphs to determine a graph structure match. The system100optionally looks up the hash value of a new input graph to determine whether a similarly structured input graph had been generated before. If the system100determines a hash value match exists to a previously stored hash value, then the system100may forgo the optimization processing of the new input graph and use a stored transformation graph to reconstruct an optimized version of the new input graph. A hash match indicates that the structure of a previous input graph has already been optimized and a transformation graph was generated.

The system100performs optimization of the input graph by applying a set of optimization rules to the generated input graph (block230). The result of the application of the optimization rules is a final optimized input graph. The optimization rules make changes to the original structure of the generated input graph. Before applying the optimization rules, the system100optionally stores in an electronic data store, a copy of the structure of the generated input graph to save a copy of its original state (e.g., the structure). The system100creates and stores another copy of the structure of the input graph which is then used for in-situ modification where changes are made based on the requirements of particular optimization rules.

Some of the optimization rules may generate JIT blocks including portions or segments of the input graph nodes and/or operations (block240). An optimization rule may evaluate the input graph to identify nodes and/or operations of the input graph that can be combined to form a JIT block that will later be evaluated by the system100for optimized GPU processing.

The system100also may determine implementation blocks for segments of the input graph that are non-JIT blocks (block250). The system100may evaluate nodes of the optimized input graph and connect operations to vendor or reference libraries from an operation registry so that the system100may perform a particular operation of the input graph when the operation is to be executed. An implementation block includes the operations of the optimized input graph and a reference to a respective referenced library. For example, the system100may determine that a matrix multiply operation on a particular device may be executed using a specific library's function, such as using a matrix multiply function from a referenced library, such as an OpenBlas or CuBlas library.

The operation registry may be a database or other data structure where the system may reference a particular operation by name and retrieve an associated library function from which the operation may be performed. For example, the system100may include a registry of different types of operations and one or more libraries available to the system100for performing a specific operation. The system100may choose one library over another library based on the particular hardware configuration of the system100performing the operation. The system100may include an API where particular operations and a referenced library may be defined. This API allows the system100to define the registry with the identified operations and the one or more libraries that may be used perform the operation.

After the system100applies the optimization rules to the input graph, a final input graph will have been created. The system100derives a transformation graph by comparison of the structure of the original stored copy of the generated input graph to a structure of the final optimized version of the input graph (block260). The transformation graph is stored in a data store and provides information to the system100about how to recreate the final optimized version of the input graph without having to reapply the optimization rules.

The system100then performs a device-specific optimization process of the JIT blocks based on the target device (e.g., CPU, GPU, TPU, etc.) upon which the kernel will be performed (block270). In one embodiment, the JIT blocks include operations of the original input graph that the system100has combined together for GPU processing. The system100uses an A* search-based model to determine a processing approach that relies on a basic cost model of the GPU's memory and parallelism. The system100determines in what manner to parallelize operations of the JIT block (e.g., processing by blocks, threads, blocks+threads, or none (run sequentially)). The system generates GPU code based on the determined parallel processing approach. The system100generates GPU and/or CPU code which is then compiled to create one or more kernels to perform the operations of the JIT block, thereby creating a fused JIT kernel. The code for the fused JIT kernel (in machine or string form) may be stored in memory or on a storage device for later access and retrieval when the system100needs to perform the operations of the fused JIT kernel.

While general GPU processing should be understood by one skilled in the art, the following paragraph may aid one's understanding of the inventive subject matter described herein. A thread is an execution unit that runs kernels on a GPU. Kernels are functions that are executed in parallel on a GPU (e.g., a C language written function or program). Each thread has its own identifier and uses an index to access elements in an array. A collection of threads may cooperatively execute the same kernel for example to process a data set. A block (also referred to as a thread block) is a group of threads that can be executed serially or in parallel. All of the threads in a single thread block can communicate with one another. Multiple thread blocks may be combined to form a grid with each of the blocks of the grid containing the same number of threads. With a grid, a kernel is launched as a collection of thread blocks.

The system100then performs the optimized input graph code using optimized JIT blocks and the operations for the determined implementation block(s) (block280). When a result is requested for an input graph, the system100performs the optimized version of the input graph. In some instances, the optimized input graph may not include any JIT blocks. In this instance, the operations of the optimized input graph are performed by the system100, and a result of the optimized input graph is provided to a calling system, process, program or other code. The system100performs the operations of the non-JIT blocks using the referenced libraries of the implementation block(s). In other instances, the optimized input graph may include JIT blocks. In this case, the system100will access and execute the stored fused JIT kernel associated with the JIT block of the optimized input graph. In one embodiment of the system100, some of the non-JIT block operations are performed by a CPU and the fused JIT kernels are performed by a GPU. In other embodiments, the system100may generate code and compile kernels for hardware accelerated processors such as Tensor Processing Units, Cerebras, Graphcore device and memristors.

Input Graph Generation

FIG. 3illustrates an example of an input graph in the form of an expression graph generated from user code. The system100translates received user code300into an input graph as represented by the expression graph330. The system100translates the user code300into an input graph where ordering of operations becomes a directed graph. The input graph allows the system100to perform the operations with some knowledge of what is needed to be completed before the input graph can return the results of what the user code had otherwise expected.

While translating the user code300, the system100iteratively appends new nodes and edges to create the expression graph330in a manner that logically represents the intended operations of the user code300. The system100translates the user code300into dynamic control flow objects that have control dependencies with other objects where the other objects must be executed or performed before a respective control flow object may continue processing. Operations that modify memory replace the affected object by the associated transaction operation (e.g., assignment, read, etc.) last applied to it, thereby ensuring the user code300imperative order translates to the same partial order in the generated input graph. Operations are queued by the system100in a way that respects the original intent of the user code300. Operations of the input graph generated by translation of the user code300are queued up until the moment when the operations are needed to be executed.

For example inFIG. 3, the user code300begins by creating an empty 3×3 matrix for the variable dest with the code “auto dest=Array({3, 3});”. The expression graph330represents the matrix by the object called dest332that points to a 3×3 matrix234.

In the next step of the user code300“for (int I=0; I<2; i++)”, a loop from the number zero to the number1is performed. In this loop, the user code300selects some row and column of the matrix dest and will increment the sub-value of the matrix by 2. In terms of computation, whenever a process accesses this matrix again, the process expects the value of this specific location to be incremented by 2, and that result should be available in the matrix dest.

The expression graph330represents the array object236as a submatrix of the matrix334. The expression graph represents a new operation object “Assign+=”338and a new scalar object340which indicates to increment the scalar object “Scalar(2)”340by a value of two. The assign operation object238describes what needs to be performed before accessing the matrix334or the object dest332.

Dashed lines364,366indicate variable dest332redirected to reflect a memory transaction. Below dashed line364, the expression graph330represents that a control dependency342exists whenever the dest object is accessed again. To access the dest object, the assign operation object338needs to be executed before processing can continue. The control dependency342creates a locking mechanism or transaction by connecting the operation object338to the dest object via the control flow object344. The dest object uses the control flow object344as its new reference.

The dashed line366near the bottom of the expression graph330is a similar control flow requirement repeated again. The expression graph330represents objects in which the processing is going to increment the control flow object346by a value of two. The expression graph330represents the same operation348as operation338and a new scalar object350which represents the increment by 2 as indicated by the scalar object “Scalar(2)”350.

The expression graph330represents another control flow object352. The second control flow object352is now relying on two events to be completed before processing can proceed. The control flow object352has two control dependencies354,356. Before being able to return its contents to any other object, the control flow object352waits for the first increment by two operation338and the second increment by two operation348.

Now the user code300“auto res=op::matmul (dest, op::uniform(−1, 1, {3, 3});”, would compute a matrix multiply between dest and another matrix of uniform random noise sampled from the values of −1 to 1. The expression graph represents an object358“MatMul” for the matrix multiply operation, and an object360“Uniform” for the matrix of uniform random noise. The matrix multiply object358performs a matrix multiply operation using two objects, the control flow object352and the uniform matrix object360. The result of the matrix multiply is represented by the object362“res”.

By generation of the expression graph, the system100now has created a framework of the input graph and has determined a processing control flow for which nodes and operations of the computation graph must be completed and in which order. The system100may then optimize the input graph via application of optimization rules.

Input Graph Optimization

FIG. 4illustrates an example of the system100optimizing an input graph410to create a final optimized input graph440and generating a transformation graph470from the final optimized input graph440. The system100receives as an input the generated input graph (i.e., the expression graph) and optimizes the input graph using optimization rules. The input graph410represents the generated expression graph300as described above. While not shown in the input graph410, execution of the input graph410uses the generated control flow dependencies for execution of the input graph410and/or the final optimized input graph440.

Optimization Rules

The system100optimizes the input graph410by applying a series of optimization rules and determining whether or not a modification may be made to a node or portion of the input graph410. If an optimization rule is applicable to the node or portion of the input graph410, then the system100applies the rule and modifies the particular node or portion of the input graph410. After all of the optimization rules have been applied to the input graph410, a final optimized input graph is generated440. As further discussed below, the system100compares the original input graph to the final optimized input graph340to derive a transformation graph470.

Optimization rules may be defined by a test function which receives a specific node in the input graph410and returns whether the optimization rule can be applied at the particular location of the node. An optimization rule may also include a transformation function that makes desired edits to the input graph410. The optimization rules may be stored, locally or remotely, as a compiled library, database, file, listing or collection of rules in an electronic form retrievable by the system100.

The optimization rules may define a test for whether or not certain operations of the input graph310should be run in their current form, or whether the operations can be rewritten and/or combined into a more efficient form for processing. For example, the system100may apply an optimization rule to determine whether a combination of a log and an exponential operation exists in the input graph410. If the system100identifies the occurrence of this combination of operations in the input graph410, the system100may delete the log and exponential operations and replace them with an identity expression. In another example, the system100may apply an optimization rule, such as a mathematical simplification rule, to evaluate the occurrence of certain operations performed by the input graph410and simplify those operations into a basic set of operations. A non-exhaustive example listing of optimization rules is provided inFIG. 5.

Multiple Optimization Passes

The system100uses an iterative, multiple optimization pass approach to optimize the input graph410. Optimization rules may be registered or defined with a priority, ranking or execution order which indicates the order each pass of the optimization rules will be applied to the input graph410. In addition to the registered optimization rules being applied, a simplification pass may be run by the system100on the input graph410to unify duplicated calls and assignments.

To ensure that transformations of input graph410nodes run in isolation of one another, the system100may run a particular optimization rule over the entire input graph410before running the next optimization rule. In other words, the system100traverses through each of the nodes of the input graph and applies the same optimization rule to each node of the input graph. For example, the system100would apply a first optimization rule to each node of the input graph410starting from the leaf nodes and moving to the other nodes of the input graph410, and then after completing processing of the first optimization rule, the system100would apply a second optimization rule to each node of the input graph410.

During the optimization processing of the input graph410, the system100may apply a particular type of optimization rule to determine JIT-able blocks of operations. A just-in-time (JIT) block, also referred to as a JIT block, is a collection of operations that may be fused together into a kernel for improved GPU processing (as discussed further below). The system100may determine whether nodes or portions of nodes of the input graph410may be combined together to form a JIT block. The system100may evaluate for nodes that are connected, or could have been connected, to form a JIT block. For example, the system100may apply an optimization rule that determines that a sigmoid and scalar operation may be combined into a JIT block of operations. The system100may apply heuristics that include a set of measurements to determine whether two or more operations of an input graph310should combined into a larger JIT block.

In some cases, the system100may recognize the structure of a JIT-able block of operations. For example, the system100may have already fused together a similar set of JIT block operations and created a kernel for those operations (such as softmax, attention or crossentropy operations). Rather than having to optimize the JIT block, the system100may instead forgo the optimization process, and invoke a previously stored JIT kernel for execution of the JIT block.

In-Place Modification of Input Graph

The system100in-place modifies a data structure of the input graph410in a memory cache when applying the optimization rules. The system100creates a data structure of the input graph410in the memory cache and modifies the data structure as applicable optimization rules are processed and make modifications to the input graph410. For example, once the system100has applied a first optimization rule to the created data structure of the input graph410, the data structure is revised with changes to the data structure indicating graph modification requirements of the first optimization rule. The system100may apply subsequent optimization rules which may also make revisions to the data structure. The resulting data structure would reflect the additional changes being made by subsequent optimization rules.

The system100applies a series of optimization rules and makes sequenced in-situ changes to the data structure of the input graph410in the memory cache. A final result of the series of optimization rules being applied to the first input graph410is the generation of a final optimized input graph440(i.e., an optimized final data structure of the input graph). The system100may then take a snapshot or make a copy of the final optimized input graph440data structure and save the structure in a memory cache or to a data storage device.

By having made in-situ changes to the data structure of the input graph440, the system100may then compare the final optimized data structure to a previously stored version of the data structure of the original generated input graph. The system100may determine which nodes of the original input graph have changed, and which of the nodes have not. The system100may compare the data structure of the original input graph410to the final data structure of the optimized input graph440. As further described below, the system100may evaluate the differences of the original input graph410and the final optimized input graph440to generate the transformation graph470.

While the system100does not need to store or create transactions or logs about what modifications are being made to the input graph410during the optimization process, the system100may, nevertheless, be configured to generate transactional logs of the changes being made to the input graph as the optimization rules are being applied. The transactional logs may assist in understanding the step-by-step application of the optimization rules to the input graph.

Optimization Example

Referring again toFIG. 4, a more detailed discussion of the optimization of an input graph ensues. As noted previously, the input graph410logically represents a computation graph of written user code, and what is expected to happen when the user code is run. The input graph410represents the sequence of elements the system100needs to obtain in order to determine the result of the addition operation422. In order to obtain the result of the addition operation422, the system100needs to determine the result of the sigmoid operation418and scalar operation402. The sigmoid operation418needs to compute the result of the matmul operation416. The matmul operation416needs to compute the result of the uniform operation414, and so forth.

Input graph410includes two possible inputs for a first array412that is dimensioned100by200, and second array414of uniform random numbers that is dimensioned200by50. A matrix multiply operation416is performed on the two arrays412,414. A sigmoid nonlinearity function418is applied to the output of the matrix multiply416operation generating a result. Then some scalar number420is added to that result.

Graph440exemplifies a final optimized input graph after optimization rules have been applied to the input graph410. The system100saves every operation of the input graph410to some memory location. Each operation needs to have some output and is represented by the assign objects452,454and456. The assign object452indicates that the uniform random numbers444needs to be stored in memory and will be stored in an array458dimensioned200by50. The assign object454indicates that the result of the matrix multiply operation446needs to be stored in memory and will be stored in an array460dimensioned100by50. The assign object456indicates that addition of the output of the sigmoid function and the scalar value needs to be store in memory and will be stored in an array462dimensioned100by50.

As noted previously, the system100may evaluate operations to combine them into JIT blocks. In the example optimized input graph440, the system100has combined the sigmoid operation418, scalar operation420, and add operation of the original input graph410into a single JIT block464that did not previously exist. The new JIT block464receives as input the output of the matrix multiply operation446. The system100saves the JIT block464to an assigned location, e.g., in a memory cache or data store, and executes the JIT block when it is ready to be performed. As discussed below, the system100may perform further processing on the JIT block to determine an optimized fused kernel approach.

The system100may perform the operations of the final optimized input graph440when the result of the final optimized input graph400is needed (e.g., when the result of the assign object456is needed). For instance, the system100may receive a request from another control source, program, software, process, function or program code asking or calling for the result. Until such time that the result is needed or requested, the system100may have already queued the operations of the final optimized input graph440in memory, and then perform the operations of the graph440when the result is needed. Alternatively, the system100when receiving a request for the result, may in response to the request, in real-time (i.e., dynamically) generate the input graph, optimize the input graph and then perform the operations of the optimized input graph.

When processing the operations of the final optimized input graph440, the system100determines the required inputs for other operations of the graph in a hierarchical or precedential order. For example, the system100would first determine the assign operation452, determine the assign operation454and then lastly determine the assign operation456(i.e., the result of the input graph). The system100stages and computes the three assign operations452,454,456in order, executes the respect assign operation before handing off its results to another node or function.

For each of the operations to be performed by the final optimized input graph440, the system100may use a library of pre-written and/or pre-compiled code to perform the operation of a respective node of the graph. In other words, for a respective node of the final optimized input graph440, the system100may make a call to a library function to perform a particular operation required by the node. For example, for the matmul operation446, the system100may access a library of functions, where one of the library functions is able to take two input matrices and compute a matrix multiply. The library of functions may be a database, file, program, process and/or an API accessible to the system100.

The system100may generate a mapping from the nodes of the final optimized input graph440to a set of functions from the library that should be run by the system100when a particular node is executed. In some cases, common functions for a set of computation graph operations may exist, such as a matrix multiply operation or operations for image convolutions. In other cases, a referenced function may not exist for the operation of the particular node being executed. This may occur, for example, when the system100executes a JIT block464for the first time. The system100would need to first compile and/or optimize the JIT block (as described herein), and then the system100may then perform the fused JIT block. The system100may optionally store the compiled and/or optimized code for later retrieval should the system100need to execute the JIT block again as to another generated input graph.

Generation of Transformation Graph

Having completed generation of the optimized input graph440, the system100constructs a transformation graph470. The system100may create a data structure in a memory cache to represent the transformation graph479. The transformation graph470represents short-cuts describing the operations performed by the final optimized input graph440. The system100builds a skeleton (e.g. a data structure) of the output of the final optimized input graph440.

To recover the transformation that takes the input graph410to the final optimized input graph440, the system100traces the effect of each transformation of the optimization rules from the optimization passes. The system100constructs in a memory cache, a placeholder target graph which has the same graph structure as the final optimized input graph440, but with all nodes now containing reconstruction information. Each node of the placeholder target graph contains one of two types of reconstruction information:

1. Retrieval from the input graph410: if a node in the optimized graph440was left unchanged from its state in the input graph410, then the system100stores a path from the root of the input graph410to the nodes486,488,492,494,496,498marked “Copy” in the transformation graph470. Any suitable path reference identification construct may be used to locate and reference the nodes of the graph. In one embodiment, each expression contains zero or more arguments, so a path can be generated by storing a list of indices indicating which argument of a node that must be traversed to arrive at the node desired to be retrieved. In other embodiment, the path may be a single index indicating which of the input nodes the system100may reference (e.g., traversing particular nodes of the graph in some pre-determined order.)

2. A function that constructs the node anew along with some recovery information. For example, if temporary storage is created during the optimization pass to store intermediary results, then the system100saves instructions480,481,484stating “New Array” in the transformation graph470, with a shape given by retrieving the strides and dimensions from the input graph410

In places within the optimized input graph440where some object was copied wholesale from the original input graph410, the transformation graph470indicates to copy the object over unchanged. The system100identifies the object's location and where to find it in memory.

In places where some new object that did not exist before has been created in the optimized input graph440, the system100generates instructions about how to make those changes again. For example, in the transformation graph470, the new assign object482indicates the system100to assign data into a new array484that has dimensions that are copied over from the input graph410as indicated by copy objects486,488. One of the inputs to the new assign482operation is going to be the random uniform box444from the optimized input graph440as indicated by the copy object490. The new assign objects476,478,482of the transformation respectively correspond to the assign objects452,454,456of the of the optimized graph440.

The copy object490provides instructions of how to find the copied object from a memory location. The system100may evaluate the differences from the data structures of the original input graph410and the data structure of final optimized input graph440to determine what are the copy operations for the transformation graph470, and which operations are new.

From the point of view of the transformation graph470regarding data or graph contents, the system100essentially forgets about the edges of the input graph410. The system100treats the input graph410as a placeholder in which wholesale nodes may be copied, and each of the edges get recreated in the output. The edges in the transformation graph470are the same as the edges in the final optimized input graph440. The transformation graph470has a similar skeleton of the final optimized input graph440. Both graphs440,470have a similar backbone of edges and nodes. However, the contents of the nodes in what the system100is going to modify by evaluating a new input graph410would change.

Evaluation of New Input Graphs

Via the transformation graph470, the system100caches the transformations that were applied to an earlier generated input graph and reapplies them to new or subsequently generated input graphs that have the same structure and operations of the earlier generated input graph, but which may have different and/or identical contents or data. Applying a previously stored transformation graph470to a newly generated input graph allows the system100to forgo applying the optimization rules to the newly generated input graph thereby increasing overall processing efficiency and performance.

After the transformation graph470has been generated, the system100may store the transformation graph in a memory cache and/or on a storage device for subsequent access. Also, the system100generates a hash value of the original generated non-optimized version of the input graph410, and stores the hash value in a memory cache and/or on a storage device, and associates the hash value of the original generated non-optimized version of the input graph with the transformation graph470. The hash value represents a unique identifier of the structure of the original non-optimized input graph410.

During processing of a newly generated input graph, the system100may calculate a hash value for the newly generated input graph. The system100compares the hash value of the new input graph to hash values that had been previously stored in a lookup table in the memory cache. If a new input graph hashes to what the system100has identified before (e.g., a hash value match), the system100may skip the optimization process for the newly generated input graph. In other words, the system100does not need to apply the optimization rules to the newly generated input graph. By determining a hash value match, the system100recognizes that the structure of the newly generated input graph has a similar structure of a previously optimized input graph. The system100may then retrieve the stored transformation graph information from stored memory based on the hash value. The system100then uses the retrieved transformation graph470information and processes the retrieved transformation graph470against the newly generated input graph to derive an optimized version of the newly generated input graph.

A hash match or other suitable identifier of the structure of input graph may be used to identify an input graph. In one embodiment, the system100uses a sequence of shifts and XOR operations on integers and booleans that are indicative of states of different elements of the input graph. During the hashing process, the system100follows a deterministic order in which the system traverses an input graph1004.

The system100evaluates the input graph in a top-down manner where the top is considered the root of the input graph, and down direction are the leaves of the input graph. The system100builds up a hash identity for each local node using its type. For example, the hash identity may include a string for the name of the node's operation (such as an addition, multiplication, control flow, matrix multiply, etc.). The system100uses the hashing of the string using a shift in XOR of the character integers thereby creating an identify for the name. Additional information for the dimensionalities for the node (e.g., 1-d, 2-d, 3-d, n-d) may be included. Each of the node of the input graph has associated with it a dimensionality to describe its shape. This information can be used to create a more unique identifier associated with the node. Each node of the input graph also may have an API that allows the node to specify additional meta-data when its being hashed. This additional information allows the system100to hash on this information as well to create a unique identifier for the node. For example, if the node handles memory storage, the node may have a boolean to identify whether the node is using strided or contiguous memory. In another example, if the node handles casting to a different data type, such as converting numbers from a float to an integer representation. The system100may add to the hash values related to the precision of the integer representation or the type of casting that is going to be by the node where the node specifies the type of truncation (e.g., values 1, 2 or 3 to specify the type of truncation). The hashing process thereby creates a local hashed identify for each of the nodes of the input graph. The system100may store the collective local hashed identify on a data storage device or accessible memory for later retrieval.

To determine a hash match, the system100may compare hash values for the first hashed input graph to a second hashed input graph. The system100may take the local identity of each of the nodes in the manner in which they are stored, and compare the root node and respective child nodes to determine whether each of the local identities of nodes of the first input graph match the local identities of nodes of the second input graph. If the local identities of the nodes of both graphs match, then the system100may confirm that the input graphs are similar, and use a transformation graph as discussed below to optimize the second input graph.

Basically, the transformation graph470lays out a blueprint of what happened as a result of the system100processing a non-optimized input graph to create a final optimized input graph. The transformation graph470is not a data structure that will be executed per se to obtain the results of the computation of an input graph. Rather, the transformation graph470is a set of information that the system100may use to generate an optimized version of a newly generated input graph. During optimization processing, there were nodes of the final optimized input graph that were copied over from the original input graph, and there were new nodes created as part of the final optimized input graph.

In transforming a new input graph to an optimized input graph, the system100creates a data structure in memory to build out an optimized version of the new input graph. The system may start at node478, and at that location the system100understands to create a new assignment node that takes two inputs. The system100create a new assignment object in the data structure which represents the assignment operation478. The system100recursively determines the inputs to the assignment operation478. For example, the system determines how to recreate the assignment operation482and the assignment operation476. The system100walks through the structure of transformation graph478and creates new objects in the data structure in memory based on the “new” designation for a particular node of the transformation graph478(e.g., nodes481,476,480,482and484).

When the system reaches a node in the transformation graph478indicated as “copy” (e.g., nodes486,488,490,492,494,496,498) the system100may copy the node and/or contents of the node into the new data structure in memory. For example, when evaluating node490of the transformation graph, the system100, instead of creating a new object, copies the object of the new generated input graph that is located at location 0,0,1 (reading node levels from bottom to top) of the structure of the generated input graph. As indicated before, the newly generated input graph has a similar structure as the original input graph.

As a result of transforming a newly generated input graph via the transformation graph470, the system100create a new input graph structure similar to that of the final optimized input graph440, yet the new input graph structure may have different data contents than the contents of the original input graph. The system100uses the new input graph structure to perform computation graph operations when a result is requested for the newly generated input graph.

As benefit of using the cached transformation graph470, the system100effectively has removed processing for the x number of optimization rules that otherwise would have been applied to the newly generated input graph. By omitting the optimization process for the newly generated input graph, the system100may save a significant amount of processing time.

A* Fused CUDA Kernel Search

JIT Block Fusion

As discussed above, an optimization rule may determine that various nodes of an input graph may be combined to form one or more JIT blocks. Some nodes of an input graph are determined to be JITable, while others are not.FIG. 6represents a non-comprehensive table of operations that the system100may fuse and JIT-compile. Some JIT operations are specialized using functors. Unary specializations of elementwise operations are shown inFIG. 7, and binary specializations show inFIG. 8.

To construct a JIT block, the system100may consider connected subtrees of the full input graph that include JIT-able expressions (e.g., Sigmoid418+Scalar420inFIG. 4). If the parent of a JIT-able expression is also JIT-able, the system100may opt to combine those two operations into a single computation unit through fusion to form a JIT block. The combined set of operations may now be acted upon and treated by the system100as a single operation. This process of connecting JIT-able expressions is repeated by the system100until identified JIT-able expressions are fused to their neighbors.

The system100may determine that smaller JIT blocks can be combined into a larger JIT block. The system100may also determine that a larger JIT block may be divided into smaller JIT blocks, even though the nodes could have otherwise been combined into a single JIT block. For example, the system100may set a maximum number of operations that may be combined together to form a JIT block. The system100combines the JITable nodes of an input graph until reaching a predetermined number of operations to be included in the JIT block. Once the threshold number of operations is reached, the system100may combine subsequent JITable nodes into another JIT block. Some input graphs would then have multiple JIT block of various JITable operations.

The system100may also determine whether or not combined JITable operations would increase the number of input arrays needed for the JIT block. If the system100determines that the number of input arrays needed for the JIT block meets or exceeds some threshold value of maximum input arrays, then the system100may truncate the JIT block at the location of the last operation and create a new JIT block for subsequent remaining operations.

Referring back toFIG. 4, a very simple JIT block464was created in the optimized input graph440and is also reflected as a new JIT block493in the transformation graph470. The system100dynamically generates one or more segments of code by evaluating the JIT blocks of the final optimized input graph440and determines an optimized GPU and/or CPU processing approach. For example, the operations of JIT block464, the sigmoid, scalar and addition operation would be fused together and run as a segment of system100generated code to be executed on one or more graphical processing units (GPUs) and/or computer processing units (CPUs).

The system100evaluates the JIT block440operations (e.g., the sigmoid, scalar, and addition operations), and determines what dimensional space (e.g., a multi-dimensional array or matrix) is needed for performing the JIT block, and the shape of the inputs to the JIT block440. The system100uses this information to generate an approach to run the JIT block operations in parallel on one or more GPUs and/or CPUs. For example, the system100may combine JIT block operations into code strings that are later used to compile JIT kernels (such as CUDA C or CPU C code). CUDA is a parallel computing platform and programming model used to program code for GPU processing.

The system100may determine whether JIT block operations use some high dimensional space. The system100may determine that in some cases GPU processing of certain operations may be optimally performed using block level parallelism for some dimensions, and that certain operations may be optimally performed using thread level parallelism for other dimensions. And for other operations the system may determine that such operation should be performed sequentially in a non-parallel manner. The determination of optimal block or thread level parallelism is performed by an A* search process.

A* Search Process

The system100performs an optimization process to determine an approach to process the JIT blocks on the one or more GPUs and/or CPUs. The system100uses a search-based approach that relies on a basic cost model of the GPU's memory and parallelism. The search approach uses a heuristic to maximize parallelism which does not require compilation to estimate performance. The heuristic weighs the total number of loops of the final optimized input graph against the number of unique parallelism types exposed while penalizing the number of loops that have to be run sequentially. The heuristic sets unassigned loop to be maximally parallel. This provides a valid upper bound on the remaining parallelism opportunities.

Once the system100has determined all of the loops and data-consistency constraints related to the input graph, the system100may apply graph search over the assignment of parallelism types to loops while avoiding incorrect assignments (e.g., two loops that are nested cannot both rely on blocks to do their iteration, since the sub-expression is only evaluated when the parent is active, the inner loop's block iteration will never fully execute).

Each JIT-able expression provides information about the for-loops and access pattern over its inputs. This information enables the system100to determine an appropriate way of combining multiple reductions and elementwise operations into a single kernel while remaining data consistent. Iterators belonging to a loop can be annotated to state they will access data non-sequentially, indicating a possible violation of data-consistency if the read data was generated in a different block. If this iterator is used to access an intermediary result in the kernel, the system100may disallow block-level parallelism over the loop that created this iterator. Loops of JIT block operations are annotated with the available ways to parallelize them (blocks, threads, blocks+threads, or none (run sequentially)). The system100looks for repeated sub-expressions or broadcasts within a JIT subgraph and adds temporary assignments within the generated code to avoid recomputing the results.

The system100uses an A* search algorithm which utilizes a couple of pieces of information. The system100evaluates what is the desired goal state, and what is the transition cost incurred every time a temporary solution is modified. The system100may query a cost function before the system100performs a change to the temporary solution. The system100uses a heuristic function that tells the system100the expected cost to go forward with the solution. The system100may use an informative heuristic to perform the A* search to perform fast search that avoids large parts of the search base that is otherwise not useful.

Using the heuristic and a cost estimate, the system100may use the A* search algorithm to search through the space of valid solutions and find an optimal solution. The system100may perform redundancy detection checks to determine if two loops are descendants of each other, have the same symbolic size and/or whether they have equal access to using threads or blocks.

To evaluate the quality of a solution, the system100may use the following cost estimate g, where N is the total number of loops to parallelize, Nblocks, Nthreads, NblocksThreadsare the numbers of loops using blocks, threads, or both, while eleafparallel is the event when a for-loop affecting the lowest stride (innermost dimension) uses threads or blocks:
NparallelNblocks+Nthreads+NblocksThreads,
Tparallel=1.0+1Nblocks+NblocksThreads>0·8·0+
1Nthreads+NblocksThreads>0·8·0+
1eleaf parallel·0·5+
1Nblocks+blocksThreads+Nthreads>0+
Nparallel·0·2;
g=(N−Nparallel)·16+max(0,Nparallel·16−Tparallel):

To obtain a heuristic h the system100may compute g over the assigned loop variables, and for the remainder assume that parallelism was achieved using both threads and blocks, thereby providing an admissible heuristic that always underestimates the cost estimate of the kernel.

With the combination of these elements, the system100may identify areas of the optimized final input graph in which something needs to be filled in, and a way of identifying every time a change was made to a temporary solution of whether or not the change reduced or increased the processing cost. The system100may determine an estimate given all of the operations that have been performed so far, and what the cost will be of performing the remaining operations.

This system100begins with a long list that starts out empty, and for every element in the list of operations, the system100determines what the options are available to the system100for processing an item in the list. The system100processes the list in sequence, and every time there is an option in the list for which the system100has not yet made a decision, the system100generates all possible decisions that could have made, and evaluates the heuristic cost for the solution assuming that the system100has made that choice.

Example Fused Kernel Solution

FIG. 9. represents a solution900for a fused kernel that the system100has determined based on the A* search process. At the beginning of the A* search process the system100would not have any information specifically about what the entire solution is going to look like. Initially, all of the blocks (e.g., blocks912,914,916) for an assignment group910,930,950,970that are currently identified inFIG. 9as “blocks”, “thread” or “null” would be empty as the system100has not yet determined the type of parallelism operation that may be used.

In this example case for a softmax operation, the system100may want to parallelize the max over threads operation920. There is an assignment operation910that is going to receive the max operation and needs to save it to a three dimensional object. For every one of the boxes (blocks912, blocks914and null916) in the assignment operation910, the system100considers one of those approaches to use for parallelism.

There are different places where the system100needs to determine what kind of parallelism the system100should use. For example, the max operation920indicates “Threads” which means the max operation920only has the ability to be run sequentially or with threads. As such there are only two parallelism options about how the max operation920can be executed. The “Threads” requirement limits choices of the type of parallel processing may occur in the area of max operation920, and this will have implications for the rest of the graph. For example, an illustration of this limiting effect is that box912is one loop of execution that is going to call into box914multiple times, which will call box916multiple times which will the call box920. If box920and any of the boxes above it are using the same kind of parallelism, then the way the system100may iterate through one of the higher boxes is going to essentially prevent the lower one from using the same kind of parallelism. Thus, conflicts that can arise if in the parent-child relationship, the parent and one of the child use the same manner of parallelism.

The cost for the solution at the beginning is zero and the heuristic for the rest of the search is going to be an assumption that every possible decision that the system100is going to make is going to benefit from the maximum amount of parallelism available for a set of operations.

Referring toFIG. 10, the diagram illustrates a search tree for determining parallelism processing options. The system100performs an iterative search to determine what parallelism processing options should be for each of the assignments910,930,950,970as shown inFIG. 9. So for the first assignment910, assuming no real constraints of what can be done by the system100, the system100may evaluate using four different parallelism processing options: (1) blocks1010, (2) threads1020, (3) sequential1030, or (4) using a mix of blocks and thread1040. The system100may consider all four of these options. However, one of the options may be too expensive to run, cheap to run, or equally cheap to run as another option that is run sequentially but incurs some cost otherwise because it is a non-parallel loop.

The system100may then place the particular considered solution at the bottom of a queue of running solutions. However, the heuristic for all of the solutions left may indicate that all of these options, still leave most every other variable in this graph free to be allocated to a variety of options. Initially, few options have been ruled out. However, if, for instance the system100picked threads as an option, the ability to use threads in the max operation920that follows below, then the system920might have made the heuristic cost for the future go up a little. So essentially, the system100determines all solutions that are a descendant of the current one, might be a lot more expensive to run. At this point, the system100may want to delay or abort this particular branch of the search. The system100runs through the A* search process and determines the most suitable type of parallelism for the operations of the JIT block.

The system100completing the A* Search process then would have generated an optimized solution with one or more assignment groups for the respective operation of the JIT block. The system100may store a data structure and information representing the solution of the optimized JIT block. In the example, the solution is composed of four assignments910,930,950,970and has two reductions. Each reduction and assignment can be parallelized in a variety of ways. The final assignment of the loop parallelism for each operation is shown in the assignment groups910,930,950,970. For example, in assignment group910the loop parallelism is shown as blocks912, blocks914, and null916.

GPU Code Generation

The system100may generate GPU code (e.g., C or C++ code) by evaluating the data structure of the solution900generated by the A* Search process. To generate GPU code for solution of the optimized JIT block, the system100generates a series of code strings associated with each of the assign operations910,930,950,970and the operations of the sub-graphs associated with the respective assign operation. For example, for the max over threads object920, the system100may generate a string that describes how to run “max” on a GPU on an input array of some dimension.

The system100compiles the comprehensive set of code strings to generate machine code that may be performed by one or more GPUs, and/or CPUs. While reference is made to optimization of GPU and/or CPU processors, the system100may also generate code for memristor-based processors, ARM mali processors, other variants of custom processors and hardware accelerated processors such as Tensor Processing Units, Cerebras, and Graphcore device. The system100may reference the compiled GPU code whenever the system100needs to perform the JIT block of operations and determine a result.

Certain nodes of a sub-graph (e.g., the expand_dim918, max920, exp940, subtract942, broadcast944, etc.) have two type of strings associated with them. One type of string is used for definitions, and the other type of string is used to describe how to call the particular node. Both types of strings are added to a system100generated file (stored in memory or on a storage device) in some specified order. The resulting file with the system100generated code may then be dynamically compiled to generate a kernel of code (e.g., a CUDA kernel) that may perform the operations of the JIT block on one or more GPUs.

The first type of string associated with a particular node may describe some header file for a library that needs to be included in the generated GPU code so that a particular operation for the node may be performed. The first type of string includes code that does not have to be run such as definitions and/or header references to include libraries to other code or functions. For example, a library may be needed during compilation to perform specific mathematical operations of the JIT block.

The second type of string associated with a particular node may describe how the respective node may be found or called by other nodes. The system100when creating the GPU code may recursively iterate through the operations of the JIT block. For example, the system100may create code for the exponential operation940that may “call functor of exponential of (and ask subtract942how to call subtract942)”. The system100then recursively continues through the sub-graph to the broadcast operation944asking broadcast how to call itself. The system100continues until it has produced a piece of code that describes building the relevant sub-graph. In this manner, the system100ensures that the GPU code includes the right headers and libraries to perform operations, and that the GPU code identifies the right order and context for performing the specific operations of the sub-tree of the JIT block (such as broadcast944, then subtract942and the exponential940).

In further detail, the system100performs a for loop that uses the dimensions given by the assignment side and the sub-graph side. For example as to assignment group930, the system100determines for i up to the first dimension of exponent940, j up to the second dimension, etc. The variable that is represented by the string of940, uses a method where if it is called using a specific accessor (can be found by its value located at i, j (the ith, jth value of940). The system100would compute that individual item and return a specific value (e.g., an actual value such as float, integer, double, etc.). And similarly for the assignment of temp2946the system generates code representing the temporary variable, and wrapper code around reference to memory indicating there is going to be some method for accessing the memory at location i, j of temp2946. And for assignment group930the system100is going to generate a string that identifies to obtain the ith, jth value of940, and save it the ith, jth location of temp2946, and then increment i, and/or increment j, and so on and so forth.

The system100also generates GPU code indicating the left side of the assignment (e.g., assignment groups910,930,950,970) identifying the memory location where the result of the assignment will be saved to. Also, the system100may create strings to generate variables for the right side of the assignment where results for operations can be saved. For example, for the exponential940, and subtract942, the system100may generate GPU code to save the output of the respective operations to two temporary variables.

The GPU code is considered device code which is executed on a GPU with device variables residing in GPU memory. In some embodiments, the system100may solely generate device code for performing the operations of a JIT block. In other embodiments, the system100may generate a combination of both device code and host code. Host code is executed on a CPU with host variables residing in CPU memory. For examples, certain sequential operations may be performed by host code and parallel operations may be performed by the device code.

The example computer system1100includes a processing device1102, a main memory1104(e.g., read-only memory (ROM), flash memory, dynamic random access memory (DRAM) such as synchronous DRAM (SDRAM) or Rambus DRAM (RDRAM), etc.), a static memory1106(e.g., flash memory, static random access memory (SRAM), etc.), and a data storage device1118, which communicate with each other via a bus1130.

The computer system1100may further include a network interface device1008to communicate over the network1120. The computer system1000also may include a video display unit1110(e.g., a liquid crystal display (LCD) or a cathode ray tube (CRT)), an alphanumeric input device1112(e.g., a keyboard), a cursor control device1144(e.g., a mouse), a graphics processing unit1122, a signal generation device1116(e.g., a speaker), graphics processing unit1122, video processing unit1128, and audio processing unit1132.

The data storage device1118may include a machine-readable storage medium1124(also known as a computer-readable medium) on which is stored one or more sets of instructions or software1126embodying any one or more of the methodologies or functions described herein. The instructions1126may also reside, completely or at least partially, within the main memory1104and/or within the processing device1102during execution thereof by the computer system1100, the main memory1104and the processing device1002also constituting machine-readable storage media.