Allocation of machine learning tasks into a shared cache
The subject technology receives code corresponding to a neural network (NN) model, the code including particular operations that are performed by the NN model. The subject technology determines, among the particular operations, a set of operations that are to be allocated to a cache of the electronic device that is to execute the NN model. The subject technology generates a set of cache indicators corresponding to the determined set of operations. The subject technology compiles the code and the generated set of cache indicators to provide a compiled binary for the NN model to execute on a target device.
1. A method comprising:
receiving code corresponding to a neural network (NN) model, the code including particular operations that are performed by the NN model, wherein at least some of the particular operations include respective data that is to be stored in memory of an electronic device during execution of the NN model;
determining, among the particular operations, a set of operations that are to be allocated to a cache of the electronic device that is to execute the NN model;
generating a set of cache indicators corresponding to the determined set of operations, wherein the set of cache indicators includes information indicating whether to request an allocation of memory in the cache; and
compiling the code and the generated set of cache indicators to provide a compiled binary for the NN model to execute on a target device.
2. The method of claim 1 , wherein the particular operations are performed by at least one of a neural processor, a GPU, or a CPU, and each of the particular operations corresponds to at least a machine learning operation performed by the NN model, and the cache is shared between the neural processor, the GPU, and the CPU.
3. The method of claim 2 , wherein a respective quota of memory is assigned to at least one of the neural processor, the GPU, or the CPU based at least in part on a predetermined amount of memory that the particular operations uses when the NN model is executed by the target device.
4. The method of claim 3 , wherein the respective quota of memory is constrained based at least in part on a size of cache memory provided by the target device, and
the respective quota of memory is dynamic such that, during execution of the NN model by the target device, a particular processor of the target device in enabled to request an allocation of memory based at least in part on the respective quota of memory.
5. The method of claim 1 , wherein the set of operations includes only one operation.
6. The method of claim 1 , wherein generating the set of cache indicators corresponding to the determined set of operations further comprises generating further information indicating the particular operation uses data only a single time and that the data is to be stored in a second memory slower than the cache.
7. The method of claim 1 , wherein generating the set of cache indicators corresponding to the determined set of operations further comprises generating further information indicating the particular operation uses data multiple times and that the data is to be stored in the cache.
8. The method of claim 1 , wherein generating the set of cache indicators corresponding to the determined set of operations comprises generating further information indicating a cache drop operation to invalidate a portion of the cache corresponding to data that is no longer utilized by the determined set of operations.
9. The method of claim 1 , wherein determining the set of operations is based at least in part on whether a particular operation uses data that is accessed more than a single time during execution of the particular operation.
10. The method of claim 1 , wherein the set of operations that are to be allocated to the cache is based at least in part on a set of priorities, the set of priorities indicating particular data is given priority over other data for placement in the cache based on performance requirements or energy requirements.
11. A system comprising;
a processor;
a memory device containing instructions, which when executed by the processor cause the processor to:
receive code corresponding to a neural network (NN) model, the code including particular operations that are performed by the NN model, wherein at least some of the particular operations include respective data that is to be stored in memory of an electronic device during execution of the NN model;
determine, among the particular operations, a set of operations that are to be allocated to a cache of the electronic device that is to execute the NN model;
generate a set of cache indicators corresponding to the determined set of operations, wherein the set of cache indicators includes information indicating whether to request an allocation of memory in the cache; and
compile the code and the generated set of cache indicators to provide a compiled binary for the NN model to execute on a target device.
12. The system of claim 11 , wherein the particular operations are performed by at least one of a neural processor, a GPU, or a CPU, and each of the particular operations corresponds to at least a machine learning operation performed by the NN model, and the cache is shared between the neural processor, the GPU, and the CPU.
13. The system of claim 12 , wherein a respective quota of memory is assigned to at least one of the neural processor, the GPU, or the CPU based at least in part on a predetermined amount of memory that the particular operations uses when the NN model is executed by the target device.
14. The system of claim 13 , wherein the respective quota of memory is constrained based at least in part on a size of cache memory provided by the target device, and
the respective quota of memory is dynamic such that, during execution of the NN model by the target device, a particular processor of the target device in enabled to request an allocation of memory based at least in part on the respective quota of memory.
15. The system of claim 14 , wherein the set of operations includes only one operation.
16. The system of claim 11 , wherein to generate the set of cache indicators corresponding to the determined set of operations further causes the processor to generate further information indicating the particular operation uses data only a single time and that the data is to be stored in a second memory slower than the cache.
17. The system of claim 11 , wherein to generate the set of cache indicators corresponding to the determined set of operations further causes the processor to generate further information indicating the particular operation uses data multiple times and that the data is to be stored in the cache.
18. The system of claim 11 , wherein to generate the set of cache indicators corresponding to the determined set of operations further causes the processor to generate further information indicating a cache drop operation to invalidate a portion of the cache corresponding to data that is no longer utilized by the determined set of operations.
19. The system of claim 11 , wherein to determine the set of operations is based at least in part on whether a particular operation uses data that is accessed more than a single time during execution of the particular operation.
20. A non-transitory computer-readable medium comprising instructions, which when executed by a computing device, cause the computing device to perform operations comprising:
receiving a request for performing an operation by a neural network model, the request including a cache indicator with information indicating whether the operation is to include an allocation of memory in a cache provided by the computing device;
determining that a request for the allocation of memory in the cache is to be made based at least in part on the cache indicator and the operation; and
sending the request for the allocation of memory to a cache engine to complete the allocation of memory in the cache.