Scratchpad tensor allocation
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.
1 . 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.
2 . The computer-implemented method of claim 1 , wherein performing the address assignments includes:
generating a directed interference graph and a set of undecided edges for the tensors that have not been selected for spilling; and
providing constraints from the directed interference graph and the set of undecided edges to the SMT solver.
3 . The computer-implemented method of claim 2 , wherein performing the address assignments further includes identifying an additional tensor from the remaining tensors to spill based on a constraint that the SMT solver fails to satisfy.
4 . The computer-implemented method of claim 1 , wherein the one or more heuristics include a tensor size and a loop-level live interval of each tensor candidate.
5 . The computer-implemented method of claim 1 , wherein each loop-level live interval is represented using an affine schedule for the corresponding tensor.
6 . 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.
7 . The computer-implemented method of claim 6 , wherein the set of directed edges includes edges in the interference graph that are each between a pair of tensors in which a live range of a first tensor in the pair completely overlaps a live range of a second tensor in the pair.
8 . The computer-implemented method of claim 6 , wherein the set of undecided edges includes edges in the interference graph that are each between a pair of tensors in which a live range of a first tensor in the pair partially overlaps a live range of a second tensor in the pair.
9 . The computer-implemented method of claim 6 , wherein forming the directed interference graph includes:
sorting the set of directed edges into an order based on a set of one or more heuristics; and
performing operations on each of the directed edges in the sorted order, the operations including:
selecting a direction for a directed edge; and
adding the directed edge to the directed interference graph.
10 . The computer-implemented method of claim 9 , further comprising:
verifying a longest path in the directed interference graph is less than a size of the buffer memory.
11 . The computer-implemented method of claim 9 , further comprising:
verifying that the directed interference graph corresponds to a directed acyclic graph.
12 . The computer-implemented method of claim 6 , wherein assigning the addresses results in an unsatisfiable core, and the method further comprises performing additional iterations of address assignment by:
updating the directed interference graph by removing one or more directed edges from the directed interference graph, and updating the set of undecided edges to include the removed one or more directed edges; and
providing updated constraints from the updated directed interference graph and the updated set of undecided edges to the constraints solver.
13 . The computer-implemented method of claim 12 , wherein when a constraint remains unsatisfiable after removing all removable directed edges from the directed interference graph, the method further includes:
identifying a spill candidate associated with the constraint; and
updating the interference graph by removing the spill candidate associated with the constraint from the interference graph.
14 . The computer-implemented method of claim 6 , wherein the interference graph is obtained after determining an initial set of tensors to spill at timestamps having a memory usage estimation exceeding a size of the buffer memory.
15 . The computer-implemented method of claim 14 , wherein the memory usage estimation is determined based on loop-level live intervals of the tensors represented as affine schedules.
16 . The computer-implemented method of claim 15 , wherein the loop-level live intervals are approximated by replacing iteration variables in the affine schedules with numeric values.
17 . 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.
18 . The non-transitory computer readable medium of claim 17 , wherein the operations include:
determining an initial set of tensors to spill at timestamps having a memory usage estimation exceeding a size of the buffer memory.
19 . The non-transitory computer readable medium of claim 17 , wherein the operations include:
updating the directed interference graph by removing one or more directed edges from the directed interference graph, and updating the set of undecided edges to include the removed one or more directed edges; and
providing updated constraints from the updated directed interference graph and the updated set of undecided edges to the constraints solver.
20 . The non-transitory computer readable medium of claim 17 , wherein the operations include:
identifying a spill candidate associated with an unsatisfiable constraint; and
updating the interference graph by removing the spill candidate associated with the unsatisfiable constraint from the interference graph.