Automatic low level operator loop generation, parallelization and vectorization for tensor computations
A method is provided for transforming a high-level language representation of a tensor computation graph into a low level language. The method includes assigning a tensor shape and a loop primitive. The method also includes generating, from the tensor computation graph and the assigned loop primitives, an initial loop structure. The method further includes positioning the layers of the tensor computation graph within a nested loop structure to provide a final loop structure, collapsing loops in the final loop structure, and mapping the collapsed loops to hardware components configured to execute the collapsed loops. The method can be applied to artificial intelligence (AI) and machine learning (ML) use cases for improved optimization of neural networks including compilation optimization for improving performance of simulations such as medical simulations, healthcare simulations, weather simulations, and/or simulations related to other complex systems, which can also support decision making.
1 . A computer-implemented method for transforming a high-level language representation of a tensor computation graph into a low level language, the computer-implemented method comprising:
assigning, for each input edge of each node in the tensor computation graph, a tensor shape;
assigning, for each dimension of the input and output of each layer of the tensor computation graph, a loop primitive;
generating, from the tensor computation graph and the assigned loop primitives, an initial loop structure;
positioning the layers of the tensor computation graph within a nested loop structure to provide a final loop structure;
collapsing loops in the final loop structure;
mapping the collapsed loops to hardware components configured to execute the collapsed loops; and
executing the collapsed loops using the hardware components based on the mapping of the collapsed loops.
2 . The computer-implemented method according to claim 1 , further comprising defining, for the tensor computation graph, an available set of operators and loop primitives.
3 . The computer-implemented method according to claim 1 , wherein the generating the initial loop structure comprises:
generating an initial schedule for executing layers of the tensor computation graph;
ranking, based on the initial schedule, the layers of the tensor computation graph;
initializing linear loops according to loop primitives assigned to the dimensions of the input and output of each layer of the tensor computation graph;
initializing non-linear loops based on the initialized linear loops; and
inheriting linear loops to direct neighbor layers in the tensor computation graph.
4 . The computer-implemented method according to claim 3 , wherein Einstein-Notation dimension mapping is used for initializing the linear loops.
5 . The computer-implemented method according to claim 4 , wherein initializing linear loops according to the loop primitives assigned to the dimensions of the input and output of each layer of the tensor computation graph is performed based on the ranking of the layers of the tensor computation graph.
6 . The computer-implemented method according to claim 1 , wherein the positioning the layers of the tensor computation graph within a nested loop structure to provide a final loop structure comprises:
generating a hierarchical tree structure that includes loops and all layers of the tensor computation graph;
determining, based on the hierarchical tree structure, dependencies of the loops; and
classifying the loops as one of a write-loop or a read-loop.
7 . The computer-implemented method according to claim 1 , wherein mapping the collapsed loops to hardware components configured to execute the collapsed loops is performed by an auto-vectorizing compiler.
8 . The computer-implemented method according to claim 1 , wherein the high level language is a tensor-based programming model.
9 . The computer-implemented method according to claim 1 , wherein the low level language is machine code.
10 . The computer-implemented method according to claim 9 , wherein the machine code is code for a central processing unit (CPU) or a graphics processing unit (GPU).
11 . The computer-implemented method according to claim 1 , wherein the low level language is one of C/C++, very high-speed integrated circuits program (VHSIC) hardware description language (VHDL), or Verilog.
12 . The computer-implemented method of claim 1 , wherein the hardware components comprise one or more central processing units (CPUs), one or more vector processors, or one or more graphics processing units (GPUs), wherein mapping the collapsed loops to the hardware components comprises mapping the collapsed loops to the one or more CPUs, the one or more vector processors, or the one or more GPUs, and wherein executing the collapsed loops comprises executing the collapsed loops based on the mapping of the collapsed loops to the one or more CPUs, the one or more vector processors, or the one or more GPUs.
13 . The computer-implemented method of claim 1 , wherein the hardware components comprise at least one central processing unit (CPU) comprising a plurality of processor cores, wherein mapping the collapsed loops to the hardware components comprises:
mapping a first collapsed loop of the collapsed loops to a first processor core of the plurality of processor cores; and
mapping a second collapsed loop of the collapsed loops to a second processor core of the plurality of processor cores, and
wherein executing the collapsed loops comprises executing the first collapsed loop using the first processor core and executing the second collapsed loop using the second processor core.
14 . The computer-implemented method of claim 1 , wherein the hardware components comprise a plurality of processor cores, wherein the plurality of processor cores comprise scalar units and/or vector units, wherein each of the vector units comprises a plurality of Compute Unified Device Architecture (CUDA) cores, and wherein mapping the collapsed loops to the hardware components comprises mapping the collapsed loops to the scalar units and/or the vector units of the plurality of processor cores.
15 . The computer-implemented method of claim 14 , wherein the collapsed loops comprise loops that access an inner-most dimension of tensors, wherein the mapping the collapsed loops to the vector units is based on the loops that access the inner-most dimension of tensors, and wherein executing the collapsed loops comprises executing the loops that access the inner-most dimension of tensors using the vector units.
16 . The computer-implemented method of claim 14 , wherein the collapsed loops comprise inner loops and outer loops, and wherein mapping the collapsed loops to the scalar units is based on whether a collapsed loop, from the collapsed loops, is one of the inner loops or one of the outer loops.
17 . A non-transitory computer readable medium having stored thereon instructions for carrying out a method for transforming a high-level language representation of a tensor computation graph into a low level language, the method comprising:
assigning, for each input edge of each node in the tensor computation graph, a tensor shape;
assigning, for each dimension of the input and output of each layer of the tensor computation graph, a loop primitive;
generating, from the tensor computation graph and the assigned loop primitives, an initial loop structure;
positioning the layers of the tensor computation graph within a nested loop structure to provide a final loop structure;
collapsing loops in the final loop structure;
mapping the collapsed loops to hardware components configured to execute the collapsed loops; and
executing the collapsed loops using the hardware components based on the mapping of the collapsed loops.
18 . The non-transitory computer readable medium according to claim 17 , wherein the generating the initial loop structure comprises:
generating an initial schedule for executing layers of the tensor computation graph;
ranking, based on the initial schedule, the layers of the tensor computation graph;
initializing linear loops according to loop primitives assigned to the dimensions of the input and output of each layer of the tensor computation graph;
initializing non-linear loops based on the initialized linear loops; and
inheriting linear loops to direct neighbor layers in the tensor computation graph.
19 . The non-transitory computer readable medium according to claim 17 , wherein the positioning the layers of the tensor computation graph within a nested loop structure to provide a final loop structure comprises:
generating a hierarchical tree structure that includes loops and all layers of the tensor computation graph;
determining, based on the hierarchical tree structure, dependencies of the loops; and
classifying the loops as one of a write-loop or a read-loop.
20 . A system for transforming a high-level language representation of a tensor computation graph into a low level language, the system comprising:
processing circuitry configured to:
assign, for each input edge of each node in the tensor computation graph, a tensor shape;
assign, for each dimension of the input and output of each layer of the tensor computation graph, a loop primitive;
generate, from the tensor computation graph and the assigned loop primitives, an initial loop structure;
position the layers of the tensor computation graph within a nested loop structure to provide a final loop structure;
collapse loops in the final loop structure;
map the collapsed loops to hardware components configured to execute the collapsed loops; and
execute the collapsed loops using the hardware components based on the mapping of the collapsed loops.