LOOP OPTIMIZATION-BASED APPROACH FOR TASK SCHEDULING IN GRAPH MACHINE LEARNING MODELS

A processor-implemented method for representing computation in a machine learning (ML) model as loops includes receiving the ML model. The ML model is initially represented as a graph having multiple nodes coupled by edges. A value number is assigned to each node in the graph based on a similarity in characteristics of the multiple nodes. A loop for computations in the graph is reconstructed based on the value number. The loop bounds are determined using an affine scalar evolution analysis technique.

FIELD OF THE DISCLOSURE

Aspects of the present disclosure generally relate to machine learning, and more particularly to scheduling tasks in graph machine learning models using loop optimizations.

BACKGROUND

Artificial neural networks may comprise interconnected groups of artificial neurons (e.g., neuron models). The artificial neural network (ANN) may be a computational device or be represented as a method to be performed by a computational device. Convolutional neural networks (CNNs) are a type of feed-forward ANN. Convolutional neural networks may include collections of neurons that each have a receptive field and that collectively tile an input space. Convolutional neural networks, such as deep convolutional neural networks (DCNs), have numerous applications. In particular, these neural network architectures are used in various technologies, such as image recognition, speech recognition, acoustic scene classification, keyword spotting, autonomous driving, and other classification tasks.

The tasks of a machine learning (ML) model may be scheduled, and data allocated by the ML optimization framework. Optimal graph scheduling under resource constraints may be a non-deterministic polynomial (NP)-complete problem, which may result in poor tractability of the ML optimizers as the width of the models increase. As a result, the time it takes for these optimizers to optimize an ML model may be unacceptable to a user.

SUMMARY

In some aspects of the present disclosure, a processor-implemented method includes receiving a machine learning (ML) model. The ML model is represented as a graph having multiple nodes coupled by edges. The processor-implemented method also includes assigning a value number to each node in the graph based on a similarity in characteristics of the multiple nodes. The processor-implemented method additionally includes reconstructing a loop for computations in the graph based on the value number to generate a reconstructed loop. The processor-implemented method further includes determining loop bounds of the reconstructed loop using an affine scalar evolution analysis technique.

Various aspects of the present disclosure are directed to an apparatus including means for receiving a machine learning (ML) model. The ML model is represented as a graph having multiple nodes coupled by edges. The apparatus also includes means for assigning a value number to each node in the graph based on a similarity in characteristics of the multiple nodes. The apparatus additionally includes means for reconstructing a loop for computations in the graph based on the value number to generate a reconstructed loop. The apparatus further includes means for determining loop bounds of the reconstructed loop using an affine scalar evolution analysis technique.

Some aspects of the present disclosure are directed to an apparatus having at least one memory and one or more processors coupled to the at least one memory. The processor(s) is configured to receive a machine learning (ML) model. The ML model is represented as a graph having multiple nodes coupled by edges. The processor(s) is also configured to assign a value number to each node in the graph based on a similarity in characteristics of the multiple nodes. The processor(s) is additionally configured to reconstruct a loop for computations in the graph based on the value number to generate a reconstructed loop. The processor(s) is further configured to determine loop bounds of the reconstructed loop using an affine scalar evolution analysis technique.

DETAILED DESCRIPTION

The word “exemplary” is used to mean “serving as an example, instance, or illustration.” Any aspect described as “exemplary” is not necessarily to be construed as preferred or advantageous over other aspects.

Many machine learning (ML) optimization tools (e.g., compilers) represent the computations of the ML model as a graph of tasks, in which vertices represent computations, and edges represent data dependencies among the computations. The compilers may utilize an unrolled representation of the ML model to perform such optimizations. In some aspects, the unrolled representation may be unknowingly produced. That is, an ML author may not realize that the process performed by the ML model corresponds to unrolling. For example, a generative pre-trained transformer (GPT) model may include a sequence of transformers, which may produce unrolled loops.

The graph may comprise a directed acyclic graph of layers, for instance. The layers may be tiled and may represent a task. The optimizations may include scheduling execution of tasks and/or memory allocation, for instance. However, the size of the graph increases as the width of the model increases. The size of the inputs may have direct impact on the number of nodes in the graph. This is because tile sizes may be fixed (e.g., constant). The tensors of the ML model may be two-dimensional (2D) or higher, and thus, the size of the graph may increase at a rate of O(n2).

Each node may represent a task in the graph. The tasks may be scheduled, and data allocated by the ML optimization framework. Optimal graph scheduling under resource constraints may be a non-deterministic polynomial time (NP)-complete (NP-complete) problem, which may result in poor tractability of the ML optimizers as the width of the models increase. As a result, the time it takes for these optimizers to optimize an ML model may be unacceptable to a user.

To address these and other problems, scheduling of larger task graphs may leverage the highly repetitive nature of such graphs to convert the graphs to a loop representation. That is, the graph may be reconstructed in the form of loops. Each task becomes an iteration of a loop-based program. The loops may access data elements, which may be described as a function (e.g., linear or non-linear) of the loop counter, so the tasks and the data consumed and produced may be indexed and represented as a loop counter (then, loop optimization techniques may be applied to determine a schedule for the tasks).

Particular aspects of the subject matter described in this disclosure can be implemented to realize one or more of the following potential advantages. In some examples, the described techniques (e.g., assigning value numbers to similar substructures of the graph, reconstructing loops based on the value numbers, and performing scalar evolution analysis) may beneficially reduce latency in computing a schedule for executing nodes of the ML model. This is because the time to compute a schedule is constant as a function of the model width.

In addition, by representing the ML computations using loops, various loop optimization techniques may be applied to improve scheduling performance and ML model processing efficiency. For instance, such loop optimizations may include loop fusion, loop tiling, placement across multiple cores, direct memory access (DMA) generation and coarsening, multi-buffering (sometimes referred to as high-level software pipelining), and other loop optimizations.

Moreover, performance debugging may also be simplified because the behavior of many tasks (of the same loop) may be analyzed at once, as opposed to analyzing the sum of their individual behavior, as is the case in some conventional approaches.

Furthermore, optimization decisions may be performed in bulk over the original tasks, making for a more regular behavior. Regularity may enable allocation of the same local memory offsets for a certain piece of data on all the cores. A multicast transmission to the same address everywhere may be more efficient than allocation of the same address to a different address on each core.

FIG. 1 illustrates an example implementation of a system-on-a-chip (SOC) 100, which may include a central processing unit (CPU) 102 or a multi-core CPU configured for scheduling tasks in graph machine learning models using loop optimizations. Variables (e.g., neural signals and synaptic weights), system parameters associated with a computational device (e.g., neural network with weights), delays, frequency bin information, and task information may be stored in a memory block associated with a neural processing unit (NPU) 108, in a memory block associated with a CPU 102, in a memory block associated with a graphics processing unit (GPU) 104, in a memory block associated with a digital signal processor (DSP) 106, in a memory block 118, or may be distributed across multiple blocks. Instructions executed at the CPU 102 may be loaded from a program memory associated with the CPU 102 or may be loaded from a memory block 118.

The SOC 100 may also include additional processing blocks tailored to specific functions, such as a GPU 104, a DSP 106, a connectivity block 110, which may include fifth generation (5G) connectivity, fourth generation long term evolution (4G LTE) connectivity, Wi-Fi connectivity, USB connectivity, Bluetooth connectivity, and the like, and a multimedia processor 112 that may, for example, detect and recognize gestures. In one implementation, the NPU 108 is implemented in the CPU 102, DSP 106, and/or GPU 104. The SOC 100 may also include a sensor processor 114, image signal processors (ISPs) 116, and/or navigation module 120, which may include a global positioning system.

The SOC 100 may be based on an ARM, RISC-V (RISC-five), or any reduced instruction set computing (RISC) architecture. In aspects of the present disclosure, the instructions loaded into the general-purpose processor 102 may include code to receive a machine learning (ML) model. The ML model is represented as a graph having multiple nodes coupled by edges. The instructions loaded into the general-purpose processor 102 may also include code to assign a value number to each node in the graph based on a similarity in characteristics of the multiple nodes. The instructions loaded into the general-purpose processor 102 may additionally include code to reconstruct a loop for computations in the graph based on the value number to generate a reconstructed loop. The instructions loaded into the general-purpose processor 102 may further include code to determine loop bounds of the reconstructed loop using an affine scalar evolution analysis technique.

Deep learning architectures may perform an object recognition task by learning to represent inputs at successively higher levels of abstraction in each layer, thereby building up a useful feature representation of the input data. In this way, deep learning addresses a major bottleneck of traditional machine learning. Prior to the advent of deep learning, a machine learning approach to an object recognition problem may have relied heavily on human engineered features, perhaps in combination with a shallow classifier. A shallow classifier may be a two-class linear classifier, for example, in which a weighted sum of the feature vector components may be compared with a threshold to predict to which class the input belongs. Human engineered features may be templates or kernels tailored to a specific problem domain by engineers with domain expertise. Deep learning architectures, in contrast, may learn to represent features that are similar to what a human engineer might design, but through training. Furthermore, a deep network may learn to represent and recognize new types of features that a human might not have considered.

The connections between layers of a neural network may be fully connected or locally connected. FIG. 2A illustrates an example of a fully connected neural network 202. In a fully connected neural network 202, a neuron in a first layer may communicate its output to every neuron in a second layer, so that each neuron in the second layer will receive input from every neuron in the first layer. FIG. 2B illustrates an example of a locally connected neural network 204. In a locally connected neural network 204, a neuron in a first layer may be connected to a limited number of neurons in the second layer. More generally, a locally connected layer of the locally connected neural network 204 may be configured so that each neuron in a layer will have the same or a similar connectivity pattern, but with connections strengths that may have different values (e.g., 210, 212, 214, and 216). The locally connected connectivity pattern may give rise to spatially distinct receptive fields in a higher layer because the higher layer neurons in a given region may receive inputs that are tuned through training to the properties of a restricted portion of the total input to the network.

One example of a locally connected neural network is a convolutional neural network. FIG. 2C illustrates an example of a convolutional neural network 206. The convolutional neural network 206 may be configured such that the connection strengths associated with the inputs for each neuron in the second layer are shared (e.g., 208). Convolutional neural networks may be well suited to problems in which the spatial location of inputs is meaningful.

One type of convolutional neural network is a deep convolutional network (DCN). FIG. 2D illustrates a detailed example of a DCN 200 designed to recognize visual features from an image 226 input from an image capturing device 230, such as a car-mounted camera. The DCN 200 of the current example may be trained to identify traffic signs and a number provided on the traffic sign. Of course, the DCN 200 may be trained for other tasks, such as identifying lane markings or identifying traffic lights.

The DCN 200 may be trained with supervised learning. During training, the DCN 200 may be presented with an image, such as the image 226 of a speed limit sign, and a forward pass may then be computed to produce an output 222. The DCN 200 may include a feature extraction section and a classification section. Upon receiving the image 226, a convolutional layer 232 may apply convolutional kernels (not shown) to the image 226 to generate a first set of feature maps 218. As an example, the convolutional kernel for the convolutional layer 232 may be a 5×5 kernel that generates 28×28 feature maps. In the present example, because four different feature maps are generated in the first set of feature maps 218, four different convolutional kernels were applied to the image 226 at the convolutional layer 232. The convolutional kernels may also be referred to as filters or convolutional filters.

The first set of feature maps 218 may be subsampled by a max pooling layer (not shown) to generate a second set of feature maps 220. The max pooling layer reduces the size of the first set of feature maps 218. That is, a size of the second set of feature maps 220, such as 14×14, is less than the size of the first set of feature maps 218, such as 28×28. The reduced size provides similar information to a subsequent layer while reducing memory consumption. The second set of feature maps 220 may be further convolved via one or more subsequent convolutional layers (not shown) to generate one or more subsequent sets of feature maps (not shown).

In the example of FIG. 2D, the second set of feature maps 220 is convolved to generate a first feature vector 224. Furthermore, the first feature vector 224 is further convolved to generate a second feature vector 228. Each feature of the second feature vector 228 may include a number that corresponds to a possible feature of the image 226, such as “sign,” “60,” and “100.” A softmax function (not shown) may convert the numbers in the second feature vector 228 to a probability. As such, an output 222 of the DCN 200 may be a probability of the image 226 including one or more features.

In the present example, the probabilities in the output 222 for “sign” and “60” are higher than the probabilities of the others of the output 222, such as “30,” “40,” “50,” “70,” “80,” “90,” and “100”. Before training, the output 222 produced by the DCN 200 may likely be incorrect. Thus, an error may be calculated between the output 222 and a target output. The target output is the ground truth of the image 226 (e.g., “sign” and “60”). The weights of the DCN 200 may then be adjusted so the output 222 of the DCN 200 is more closely aligned with the target output.

In practice, the error gradient of weights may be calculated over a small number of examples, so that the calculated gradient approximates the true error gradient. This approximation method may be referred to as stochastic gradient descent. Stochastic gradient descent may be repeated until the achievable error rate of the entire system has stopped decreasing or until the error rate has reached a target level. After learning, the DCN 200 may be presented with new images (e.g., the speed limit sign of the image 226) and a forward pass through the DCN 200 may yield an output 222 that may be considered an inference or a prediction of the DCN 200.

FIG. 3 is a block diagram illustrating a DCN 350. The DCN 350 may include multiple different types of layers based on connectivity and weight sharing. As shown in FIG. 3, the DCN 350 includes the convolution blocks 354A, 354B. Each of the convolution blocks 354A, 354B may be configured with a convolution layer (CONV) 356, a normalization layer (LNorm) 358, and a max pooling layer (MAX POOL) 360.

Although only two of the convolution blocks 354A, 354B are shown, the present disclosure is not so limiting, and instead, any number of the convolution blocks 354A, 354B may be included in the DCN 350 according to design preference.

The convolution layers 356 may include one or more convolutional filters, which may be applied to the input data to generate a feature map. The normalization layer 358 may normalize the output of the convolution filters. For example, the normalization layer 358 may provide whitening or lateral inhibition. The max pooling layer 360 may provide down sampling aggregation over space for local invariance and dimensionality reduction.

The parallel filter banks, for example, of a deep convolutional network may be loaded on a CPU 102 or GPU 104 of an SOC 100 (e.g., FIG. 1) to achieve high performance and low power consumption. In alternative embodiments, the parallel filter banks may be loaded on the DSP 106 or an ISP 116 of an SOC 100. In addition, the DCN 350 may access other processing blocks that may be present on the SOC 100, such as sensor processor 114 and navigation module 120, dedicated, respectively, to sensors and navigation.

The DCN 350 may also include one or more fully connected layers 362 (FC1 and FC2). The DCN 350 may further include a logistic regression (LR) layer 364. Between each layer 356, 358, 360, 362, 364 of the DCN 350 are weights (not shown) that are to be updated. The output of each of the layers (e.g., 356, 358, 360, 362, 364) may serve as an input of a succeeding one of the layers (e.g., 356, 358, 360, 362, 364) in the DCN 350 to learn hierarchical feature representations from input data 352 (e.g., images, audio, video, sensor data and/or other input data) supplied at the first of the convolution blocks 354A. The output of the DCN 350 is a classification score 366 for the input data 352. The classification score 366 may be a set of probabilities, where each probability is the probability of the input data including a feature from a set of features.

FIG. 4 is a block diagram illustrating an exemplary software architecture 400 that may modularize artificial intelligence (AI) functions. Using the architecture 400, applications may be designed that may cause various processing blocks of an SOC 420 (for example a CPU 422, a DSP 424, a GPU 426 and/or an NPU 428) (which may be similar to SOC 100 of FIG. 1) to support loop optimization-based task scheduling for an AI application 402, according to aspects of the present disclosure. The architecture 400 may, for example, be included in a computational device, such as a smartphone.

The AI application 402 may be configured to call functions defined in a user space 404 that may, for example, provide for the detection and recognition of a scene indicative of the location at which the computational device including the architecture 400 currently operates. The AI application 402 may, for example, configure a microphone and a camera differently depending on whether the recognized scene is an office, a lecture hall, a restaurant, or an outdoor setting such as a lake. The AI application 402 may make a request to compiled program code associated with a library defined in an AI function application programming interface (API) 406. This request may ultimately rely on the output of a deep neural network configured to provide an inference response based on video and positioning data, for example.

The run-time engine 408, which may be compiled code of a runtime framework, may be further accessible to the AI application 402. The AI application 402 may cause the run-time engine 408, for example, to request an inference at a particular time interval or triggered by an event detected by the user interface of the AI application 402. When caused to provide an inference response, the run-time engine 408 may in turn send a signal to an operating system in an operating system (OS) space 410, such as a Kernel 412, running on the SOC 420. In some examples, the Kernel 412 may be a LINUX Kernel. The operating system, in turn, may cause a continuous relaxation of quantization to be performed on the CPU 422, the DSP 424, the GPU 426, the NPU 428, or some combination thereof. The CPU 422 may be accessed directly by the operating system, and other processing blocks may be accessed through a driver, such as a driver 414, 416, or 418 for, respectively, the DSP 424, the GPU 426, or the NPU 428. In the exemplary example, the deep neural network may be configured to run on a combination of processing blocks, such as the CPU 422, the DSP 424, and the GPU 426, or may be run on the NPU 428.

FIG. 5 is a block diagram illustrating an example value numbering process 500, in accordance with various aspects of the present disclosure. Referring to FIG. 5, the example value numbering process 500 may receive a graph 502 as input. The graph 502 may represent a ML model. The ML model may be (but is not limited to) a transformer model, for example. The graph 502 may be (but is not limited to) a directed acyclic graph (DAG), for example. The graph 502 includes multiple nodes connected by edges. The graph 502 includes input nodes (e.g., input1 and input2), a weight node, a bias node, and convolution nodes (e.g., Conv1 and Conv2). The number and types of nodes shown are merely an example for ease of illustration and understanding, and not limiting. Rather, graph 502 may include any number of nodes and different types of nodes.

The value numbering process 500 may conduct a value numbering pass 506. In the value numbering pass 506, the value numbering process 500 may assign a value number to each of the inputs, nodes, and constants. The value numbering process 500 may use similar characteristics (e.g., single static assignments (SSA) and intermediate representation (IR)) as well as control flow to group the nodes. In various aspects, input nodes may be assigned constant value numbers by hashing according to an operation name or a group identifier, for instance. For example, the input nodes (e.g., input1 and input2) may be grouped because of the similarity in node type. The input nodes may be assigned value number 1. The other nodes may be assigned by hashing on operation name, number of operands, for example.

Accordingly, as shown in FIG. 5, the weight node and bias node have a different node type than other nodes in the graph 502, and may be respectively assigned value numbers 2 and 3. Finally, the convolution nodes (Conv1 and Conv2), have a same node type and both consume an input (e.g., input1 or input2) as well as the weight node and the bias node. Thus, the convolution nodes (e.g., Conv1 and Conv2) may be grouped and assigned a value number 4. In various aspects, each of the groups (may also be referred to as buckets) may represent identical computation expressions. An analysis of the node similarity may be complexity bounded by a number of edges.

FIG. 6 is a block diagram illustrating an example process 600 for loop re-rolling, in accordance with various aspects of the present disclosure. Re-rolling may refer to reconstructing loops and loop access based on edges in the graph (e.g., 502 of FIG. 5). The example process 600 for loop re-rolling may conduct loop re-rolling 602 to reconstruct loops and loop accesses. In various aspects, the identical computation expressions may indicate an individual iteration of an unrolled loop. For example,

may be expressed as:

A new DAG 604 may be generated. The nodes of the DAG 604 may be identified using valued numbers. Each valued number may represent an operation in a loop body. For example, node 1 may represent the input nodes (e.g., input1 and input2 from graph 502 of FIG. 5), nodes 2 and 3 may respectively represent the weight node and the bias node, and node 4 may represent the convolution nodes (e.g., Conv1 and Conv2).

Information about loop bounds, access functions, and number of loop variables/loop nests within each loop body, however, may be unknown. For instance, the start iteration, final iteration, and stride may be unknown, as follows:

To determine these loop characteristics (e.g., start iteration, final iteration, or stride), affine expressions that define each induction variable/derived induction variables in the loop may be determined, for example, using scalar evolution analysis.

FIG. 7 is a block diagram illustrating an example scalar evolution analysis 700, in accordance with various aspects of the present disclosure. Scalar evolution analysis refers to a compiler technique that involves an analysis for deriving the stride and upper and lower bounds of a loop. A split history of each operation in a graph may provide the offset of each node (e.g., operation) value across the tiled dimensions. The offsets for each node in each dimension may be mapped to an iteration sequence. The stride of the loop may be derived by taking the difference between neighboring offsets.

Referring to FIG. 7, a split history offset 704 corresponding to input nodes (e.g., input1, input2) may be determined. A difference array 706 may store the result of a difference between the split history offsets (e.g., (2)−(1)=+1). Accordingly, the stride of the loop may be +1.

A first element in the array (of the split history offset 704) may represent the start value and the last element may represent the final value. Accordingly, the start value may be determined to be 1 and the final value may be determined to be 2.

The affine scalar evolution (SCEV) may be derived for the operands and the result of the loop using their respective split history (e.g., 704). The SCEV of all components may aid in deriving the following: (a) loop variables based on a number of dimensions, (b) an iteration length across each dimension (final-start), (c) consistency across all SCEV expressions to define the loop control construct, and (d) the SCEV expression, which may also define an access function of each node in the loop. Accordingly, the re-rolled loop expression may be expressed as follows:

FIG. 8 is a flow diagram illustrating a processor-implemented method 800 for scheduling tasks in graph machine learning models using loop optimizations, in accordance with various aspects of the present disclosure. The processor-implemented method 800 may be performed by one or more processors such as the CPU (e.g., 102, 422), GPU (e.g., 104, 426), and/or other processing unit (e.g., DSP 424, NPU 428), for example. In some aspects, the processor-implemented method may be performed in a compiler, for instance.

At block 802, the at least one processor receives a machine learning (ML) model, the ML model being represented as a graph having multiple nodes coupled by edges. As described, for example, with reference to FIG. 5, the example value numbering process 500 may receive a graph 502 as input. The graph 502 may represent a ML model. The ML model may be (but is not limited to) a transformer model, for example. The graph 502 may be (but is not limited to) a directed acyclic graph (DAG), for example. The graph 502 may include multiple nodes connected by edges. The graph 502 may include input nodes (e.g., input1 and input2), a weight node, a bias node, and convolution nodes (e.g., Conv1 and Conv2).

At block 804, the at least one processor assigns a value number to each node in the graph based on a similarity in characteristics of the multiple nodes. For instance, as described with reference to FIG. 5, the value numbering process 500 may conduct a value numbering pass 506. In the value numbering pass 506, the value numbering process 500 may assign a value number to each of the inputs, nodes, and constants. The value numbering process 500 may use similar characteristics (e.g., single static assignments (SSA) and intermediate representation (IR)) as well as control flow to group the nodes. In various aspects, input nodes may be assigned constant value numbers by hashing according to an operation name or a group identifier, for instance. For example, the input nodes (e.g., input1 and input2) may be grouped because of the similarity in node type. The input nodes may be assigned value number 1. The other nodes may be assigned by hashing on operation name, number of operands, for example.

At block 806, the at least one processor reconstructs a loop for computations in the graph based on the value number to generate a reconstructed loop. For example, as described with reference to FIG. 6, loops may be reconstructed (may also be referred to as re-rolled) based on edges in the graph (e.g., 502 of FIG. 5). The example process 600 for loop re-rolling may conduct loop re-rolling 602 to reconstruct loops and loop accesses. In various aspects, the identical computation expressions may indicate an individual iteration of an unrolled loop.

At block 808, the at least one processor determine loop bounds of the reconstructed loop using an affine scalar evolution analysis technique. As described, for instance, with reference to FIG. 7, scalar evolution analysis refers to a compiler technique that involves an analysis for deriving the stride and upper and lower bounds of a loop. The affine scalar evolution (SCEV) may be derived for the operands and the result of the loop using the respective split history (e.g., 704) of the operands. A split history of each operation in a graph may provide the offset of each node (e.g., operation) value across the tiled dimensions. The offsets for each node in each dimension may be mapped to an iteration sequence. The stride of the loop may be derived by taking the difference between neighboring offsets.

The SCEV of all components may aid in deriving the following: (a) loop variables based on a number of dimensions, (b) an iteration length across each dimension (final-start), (c) consistency across all SCEV expressions to define the loop control construct, and (d) the SCEV expression, which may also define an access function of each node in the loop.

EXAMPLE ASPECTS

In one aspect, the receiving means, assigning means, reconstructing means and/or determining means may be the CPU 102, program memory associated with the CPU 102, NPU 108, the dedicated memory block 118, fully connected layers 362, NPU 428 and/or the routing connection processing unit 216 configured to perform the functions recited. In another configuration, the aforementioned means may be any module or any apparatus configured to perform the functions recited by the aforementioned means.