Scratchpad tensor allocation
AssigneeAmazon Technologies, Inc.
InventorsHongbin Zheng, Parivallal Kannan, Yunxuan Yu
Life of this patent
Estimating the term…
Abstract
A technique to perform memory allocation can include obtaining an interference graph of tensors of a neural network model to store in a buffer memory, and determining a set of directed edges and a set of undecided edges in the interference graph. A directed interference graph can be formed from the set of directed edges, and constraints based on the directed interference graph and the set of undecided edges are provided to a constraints solver to assign addresses in the buffer memory for storing the tensors of the neural network model.
Claims20
A computer-implemented method for compiling a neural network model, the method comprising:
- obtaining a description of the neural network model;
- determining loop-level live intervals of tensors in the neural network model;
- identifying a timestamp having a largest memory usage that exceeds a size of a scratchpad memory based on the loop-level live intervals;
- identifying tensor candidates for spilling to external memory to reduce memory usage at the identified timestamp to be below the size of the scratchpad memory;
- selecting one or more tensor candidates to spill based on one or more heuristics;
- performing, using a satisfiability modulo theories (SMT) solver, address assignments for the scratchpad memory to store remaining tensors that have not been selected for spilling;
- and
- generating machine instructions based on the address assignments and spill operations for the tensors of the neural network model.
A computer-implemented method comprising:
- obtaining an interference graph having nodes and edges connecting node pairs, wherein each node represents a tensor of a neural network model to store in a buffer memory, and each edge between a node pair indicates that tensors represented by the nodes of the node pair have overlapping live ranges;
- determining a set of directed edges and a set of undecided edges in the interference graph;
- forming a directed interference graph from the set of directed edges;
- and
- providing constraints based on the directed interference graph and the set of undecided edges to a constraints solver to assign addresses in the buffer memory for storing the tensors of the neural network model.
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:
- obtaining an interference graph having nodes and edges connecting node pairs, wherein each node represents a tensor of a neural network model to store in a buffer memory, and each edge between a node pair indicates that tensors represented by the nodes of the node pair have overlapping live ranges;
- determining a set of directed edges and a set of undecided edges in the interference graph;
- forming a directed interference graph from the set of directed edges;
- and
- providing constraints based on the directed interference graph and the set of undecided edges to a constraints solver to assign addresses in the buffer memory for storing the tensors of the neural network model.
Dependencies are inferred from the claim text.