Arithmetic-intensity based load cloning

1. A computer-implemented method comprising:

receiving a neural network model that includes memory load operations;

estimating a maximal local memory need for implementing the neural network model;

determining that the estimated maximal local memory need is greater than a size of a local memory of a processing engine for implementing the neural network model;

determining a respective arithmetic intensity factor (AIF) for each memory load operation of the memory load operations;

modifying the neural network model by performing, for each memory load operation with the respective AIF greater than a threshold value:

grouping computation operations that use weights loaded into the local memory by the memory load operation into one or more clusters of computation operations; and

providing the memory load operation before each cluster of the one or more clusters of computation operations, wherein the memory load operation provided before each cluster of the one or more clusters of computation operations has a respective AIF greater than the threshold value; and

generating, based on the modified neural network model, instructions for performing the memory load operations and the computation operations.

2. The computer-implemented method of claim 1 , wherein the threshold value is an arithmetic intensity factor of a ridge point of a roofline model of the processing engine, and wherein the processing engine includes a processing element array.

3. The computer-implemented method of claim 1 , wherein the neural network model includes a batch-N convolutional neural network model, where N is greater than 1.

4. The computer-implemented method of claim 1 , wherein grouping the computation operations that use the weights loaded into the local memory by the memory load operation into the one or more clusters of computation operations comprises:

clustering the computation operations based on memory location objects the computation operations write to, wherein each set of computation operations that write to a same memory location object maps to a respective node in an aggregated cluster; and

splitting, based on distances between nodes in the aggregated cluster, the nodes in the aggregated cluster into one or more clusters of nodes representing the one or more clusters of computation operations, wherein a distance between a first node and a second node in the aggregated cluster is equal to a number of computation operations in the neural network model that are between the set of computation operations mapped to the first node and the set of computation operations mapped to the second node.

5. A computer-implemented method comprising:

receiving a neural network model that includes memory load operations;

estimating a maximal local memory need for implementing the neural network model;

determining that the estimated maximal local memory need is greater than a size of a local memory of a processing engine for implementing the neural network model;

selecting, from the memory load operations and based on arithmetic intensity factors (AIFs) of the memory load operations, a memory load operation having an AIF greater than a threshold value;

grouping computation operations associated with data loaded by the selected memory load operation into two or more clusters of computation operations; and

modifying the neural network model by providing the selected memory load operation before each cluster of the two or more clusters of computation operations.

6. The computer-implemented method of claim 5 , wherein the threshold value is determined based on an arithmetic intensity factor of a ridge point of a roofline model of a processing engine for implementing the neural network model.

7. The computer-implemented method of claim 5 , wherein the neural network model includes a batch-N convolutional neural network model, where N is greater than 1.

8. The computer-implemented method of claim 5 , further comprising determining the AIF for the selected memory load operation based on a ratio between a number of floating point operations in the computation operations associated with the data loaded by the selected memory load operation and a number of weights in the data loaded by the selected memory load operation.

9. The computer-implemented method of claim 5 , wherein the data loaded by the selected memory load operation includes weights or filter tensors loaded into a local memory of a processing engine for implementing the neural network model.

10. The computer-implemented method of claim 5 , wherein grouping the computation operations associated with the data loaded by the selected memory load operation into the two or more clusters of computation operations comprises:

clustering the computation operations based on memory location objects that the computation operations write to, wherein each set of computation operations that write to a same memory location object maps to a respective node in an aggregated cluster; and

splitting, based on distances between nodes in the aggregated cluster, the nodes in the aggregated cluster into two or more clusters of nodes that represent the two or more clusters of computation operations, respectively, wherein a distance between a first node and a second node in the aggregated cluster is equal to a number of operations between the set of computation operations mapped to the first node and the set of computation operations mapped to the second node in the neural network model.

11. The computer-implemented method of claim 10 , wherein splitting the nodes in the aggregated cluster into the two or more clusters of nodes comprises:

splitting an intermediate cluster of nodes in the aggregated cluster into a first candidate cluster of nodes and a second candidate cluster of nodes;

determining that both an AIF associated with the first candidate cluster of nodes and an AIF associated with the second candidate cluster of nodes are greater than the threshold value;

determining a quality of clustering after splitting the intermediate cluster of nodes; and

including the first candidate cluster of nodes and the second candidate cluster of nodes in the two or more clusters of nodes based on determining that the quality of clustering after splitting the intermediate cluster of nodes is higher than a quality of clustering before splitting the intermediate cluster of nodes.

12. The computer-implemented method of claim 10 , wherein splitting the nodes in the aggregated cluster into the two or more clusters of nodes comprises:

splitting an intermediate cluster of nodes in the aggregated cluster into a first candidate cluster of nodes and a second candidate cluster of nodes;

determining that an AIF associated with the first candidate cluster of nodes or an AIF associated with the second candidate cluster of nodes is lower than the threshold value; and

rejecting the splitting of the intermediate cluster of nodes into the first candidate cluster of nodes and the second candidate cluster of nodes.

13. The computer-implemented method of claim 5 , further comprising generating, based on the modified neural network model that includes the selected memory load operation before each cluster of the two or more clusters of computation operations, instructions for performing the selected memory load operations and the computation operations.

14. A non-transitory computer readable medium having stored therein instructions that, when executed by one or more processors, cause the one or more processors to execute a compiler, the compiler performing operations including:

receiving a neural network model that includes memory load operations;

estimating a maximal local memory need for implementing the neural network model;

determining that the estimated maximal local memory need is greater than a size of a local memory of a processing engine for implementing the neural network model;

selecting, from the memory load operations and based on arithmetic intensity factors (AIFs) of the memory load operations, a memory load operation having an AIF greater than a threshold value;

grouping computation operations associated with data loaded by the selected memory load operation into two or more clusters of computation operations; and

modifying the neural network model by providing the selected memory load operation before each cluster of the two or more clusters of computation operations.

15. The non-transitory computer readable medium of claim 14 , wherein the threshold value is determined based on an arithmetic intensity factor of a ridge point of a roofline model of a processing engine for implementing the neural network model.

16. The non-transitory computer readable medium of claim 14 , wherein the neural network model includes a batch-N convolutional neural network model, where N is greater than 1.

17. The non-transitory computer readable medium of claim 14 , wherein the operations further comprise determining the AIF for the selected memory load operation based on a ratio between a number of floating point operations in the computation operations associated with the data loaded by the selected memory load operation and a number of weights in the data loaded by the selected memory load operation.

18. The non-transitory computer readable medium of claim 14 , wherein grouping the computation operations associated with the data loaded by the selected memory load operation into the two or more clusters of computation operations comprises:

clustering the computation operations based on memory location objects the computation operations write to, wherein each set of computation operations that write to a same memory location object maps to a respective node in an aggregated cluster; and

splitting, based on distances between nodes in the aggregated cluster, the nodes in the aggregated cluster into two or more clusters of nodes that represent the two or more clusters of computation operations, respectively, wherein a distance between a first node and a second node in the aggregated cluster is equal to a number of operations between the set of computation operations mapped to the first node and the set of computation operations mapped to the second node in the neural network model.

19. The non-transitory computer readable medium of claim 14 , wherein the operations further comprise generating, based on the modified neural network model that includes the selected memory load operation before each cluster of the two or more clusters of computation operations, instructions for performing the memory load operations and the computation operations.