Method and system for memory management within machine learning inference engine
A method includes in a first iteration receiving a machine learning (ML) network model comprising a plurality of ML operations in high-level code; generating an internal representation (IR) for the ML network model, the IR mapped to one or more components in a multi-processing tile device; generating primitive functions based on the IR; generating an allocation list based on the primitive functions; determining when a tensor data within the allocation list is no longer needed; and inserting a deallocation function associated with the tensor data to the primitive functions to form an updated primitive functions, the inserting frees up a memory space associated with the tensor data when the tensor is no longer needed; and in a second iteration generating a compilation of the updated primitive functions to map the IR to the one or more components in the multi-processing device, wherein the compilation is a low-level instructions.
1 . A computer-implemented method comprising:
in a first iteration of code generation:
receiving a machine learning (ML) network model comprising a plurality of ML operations in high-level code;
partitioning the ML network model into a plurality of sub-graphs;
generating an internal representation (IR) for each sub-graph of the plurality of sub-graphs, wherein the IR is mapped to one or more components in a multi-processing tile device;
generating one or more primitive functions based on the IR;
generating an allocation list based on the one or more primitive functions;
determining one or more memory space within the multi-processing tile device associated with a tensor data that can be deallocated and timing associated therewith, wherein the determining is based on the allocation list; and
inserting a deallocation function to the one or more primitive functions to deallocated the one or more memory space, wherein location of the deallocation function within the primitive functions is based on the timing, and wherein an updated primitive functions is formed; and
in a second iteration of the code generation:
generating a compilation of the updated primitive functions to map the IR to the one or more components in a multi-processing device.
2 . The compiler implemented method of claim 1 , wherein the first iteration and the second iteration occur within an ahead of time (AOT) compiler.
3 . The compiler implemented method of claim 1 further comprising inserting a code after the deallocation function is inserted, wherein the code is configured to check whether de-allocation of the one or more memory space has occurred.
4 . The compiler implemented method of claim 1 , wherein the allocation list includes one or more of model input tensor data, model output tensor data, model constant tensor data, an input activation tensor data, an output activation tensor data, and an intermediate tensor data.
5 . The compiler implemented method of claim 1 , wherein the deallocation function frees up space in the one or more memory space.
6 . The compiler implemented method of claim 5 , wherein the one or more memory space is an on-chip memory (OCM) of a processing tile within the multi-processing device.
7 . The compiler implemented method of claim 6 , wherein the multi-processing device comprises a plurality of processing tiles, wherein each processing tile of the plurality of processing tiles includes a plurality of processing elements and its dedicated OCM.
8 . The compiler implemented method of claim 1 , wherein the compilation of the updated primitive functions that map the IR to the one or more components in the multi-processing device generates low-level instructions, wherein the method further comprises transmitting the low-level instructions to the multi-processing tile device.
9 . The compiler implemented method of claim 1 , wherein the deallocation function removes data associated with the tensor data to free up memory space within the multi-processing tile device.
10 . A method comprising:
in a first iteration of code generation:
receiving a machine learning (ML) network model comprising a plurality of ML operations in high-level code;
generating an internal representation (IR) for the ML network model, wherein the IR is mapped to one or more components in a multi-processing tile device;
generating one or more primitive functions based on the IR;
generating an allocation list based on the one or more primitive functions;
determining when a tensor data within the allocation list is no longer needed; and
inserting a deallocation function associated with the tensor data to the one or more primitive functions to form an updated primitive functions, wherein the inserting frees up a memory space associated with the tensor data when the tensor is no longer needed; and
in a second iteration:
generating a compilation of the updated primitive functions to map the IR to the one or more components in the multi-processing device.
11 . The method of claim 10 , wherein the first iteration and the second iteration occur within an ahead of time (AOT) compiler.
12 . The method of claim 10 further comprising inserting a code after the deallocation function is inserted, wherein the code is configured to check whether de-allocation of the one or more memory space has occurred.
13 . The method of claim 10 , wherein the allocation list includes one or more of model input tensor data, model output tensor data, model constant tensor data, an input activation tensor data, an output activation tensor data, and an intermediate tensor data.
14 . The method of claim 10 , wherein the memory space is an on-chip memory (OCM) of a processing tile within the multi-processing device.
15 . The method of claim 14 , wherein the multi-processing device comprises a plurality of processing tiles, wherein each processing tile of the plurality of processing tiles includes a plurality of processing elements and its dedicated OCM.
16 . The method of claim 10 , wherein the compilation of the updated primitive functions to map the IR to the one or more components in the multi-processing device generates low-level instruction, and wherein the method further comprises transmitting the low-level instructions to the multi-processing tile device.
17 . The method of claim 10 , wherein the deallocation function removes data associated with the tensor data that is no longer needed.
18 . A system comprising:
a compiler comprising:
a code generator sequence engine configured to:
receive a machine learning (ML) network model comprising a plurality of ML operations in high-level code; and
generate an internal representation (IR) for the ML network model, wherein the IR is mapped to one or more components in a multi-processing tile device;
a primitive code generator in a first iteration configured to generate one or more primitive functions based on the IR and wherein the primitive code generator is further configured to generate an allocation list based on the one or more primitive functions;
a memory engine in the first iteration configured to receive the allocation list from the primitive code generator and determine when a tensor data within the allocation list is no longer needed; and
a library module in the first iteration configured to receive the allocation list from the memory engine and wherein the library module is configured to generate one or more instruction set architecture (ISA) based on the allocation list,
wherein the memory engine in the first iteration is further configured to cause the library module to generate an ISA to deallocate the tensor data when the tensor data is no longer needed,
wherein the memory engine in the first iteration is further configured to cause the primitive code generator to insert the generate ISA to the one or more primitive functions to form an updated primitive functions, wherein the inserting deallocate a memory space associated with the tensor data when the tensor data is no longer needed,
wherein the primitive code generator in a second iteration is configured to generate a compilation of the updated primitive functions to map the IR to the one or more components in the multi-processing device.
19 . The system of claim 18 , wherein the compiler is an ahead of time (AOT) compiler.
20 . The system of claim 18 , wherein the memory engine is further configured to cause the library module to generate another ISA and for the primitive code generator to insert the another ISA after the generated ISA to deallocate the memory space to check whether de-allocation of the one or more memory space has occurred.
21 . The system of claim 18 , wherein the allocation list includes one or more of model input tensor data, model output tensor data, model constant tensor data, an input activation tensor data, an output activation tensor data, and an intermediate tensor data.
22 . The system of claim 18 , wherein the memory space is an on-chip memory (OCM) of a processing tile within the multi-processing device.
23 . The system of claim 22 , wherein the multi-processing device comprises a plurality of processing tiles, wherein each processing tile of the plurality of processing tiles includes a plurality of processing elements and its dedicated OCM.
24 . The system of claim 18 , wherein the primitive code generator generates low level instructions and wherein the primitive code generator is further configured to transmit the low-level instructions to the multi-processing tile device.
25 . The system of claim 18 , wherein the insertion of the ISA to deallocate the memory space removes data associated with the tensor data that is no longer needed.
26 . A method comprising:
in a first iteration of code generation:
receiving a machine learning (ML) network model comprising a plurality of ML operations in high-level code;
generating an internal representation (IR) for the ML network model, wherein the IR is mapped to one or more components in a multi-processing tile device;
generating one or more primitive functions based on the IR;
generating a memory footprint of the ML network model;
generating an allocation list based on the one or more primitive functions; and
determining sizes of memory regions associated with the ML network mode; and
in a second iteration:
predicting a memory utilization for a given memory size within a given processing tile based on the memory footprint and further based on the sizes of memory regions;
changing a sequence associated with the allocation list based on the predicting to reduce a number of re-allocation; and
generating a compilation of the updated primitive functions to map the IR to the one or more components in the multi-processing device based on the changing.
27 . The method of claim 26 , wherein the first iteration and the second iteration occur within an ahead of time (AOT) compiler.
28 . The method of claim 26 , wherein the allocation list includes one or more of model input tensor data, model output tensor data, model constant tensor data, an input activation tensor data, an output activation tensor data, and an intermediate tensor data.
29 . The method of claim 26 , wherein the memory space is an on-chip memory (OCM) of a processing tile within the multi-processing device.
30 . The method of claim 29 , wherein the multi-processing device comprises a plurality of processing tiles, wherein each processing tile of the plurality of processing tiles includes a plurality of processing elements and its dedicated OCM.
31 . The method of claim 26 , wherein the compilation of the updated primitive functions to map the IR to the one or more components in the multi-processing device generates low-level instructions, wherein the method further comprises transmitting the low-level instructions to the multi-processing tile device.