Tile assignment for matrix multiplication packing
A technique for arranging matrix multiplications for concurrent execution in an integrated circuit device may include obtaining a representation of a data dependency graph of a neural network model. The data dependency graph may include having an accumulation group (AG) pack of accumulation groups (AGs), in which each of the AGs has one or more matrix multipartition instructions. A representation of a memory location base partition constraint graph of the AG pack can be generated, and an AG row group constraint graph can be generated based on the memory location base partition constraint graph. The AGs of the AG pack can then be assigned to tiles in an integrated circuit device based on the AG row group constraint graph.
1 . A computer-implemented method for compiling a neural network model, the method comprising:
obtaining a description of the neural network model by a compiler;
generating, by compiling the description by the compiler, a data dependency graph of the neural network model;
reordering, by the compiler, the data dependency graph to form an accumulation group (AG) pack of adjacent accumulation groups (AGs), wherein each AG includes one or more matrix multiplication instructions, and wherein the AG pack represents a consecutive sequence of matrix multiplication instructions;
generating, by the compiler, a memory location base partition constraint graph (mlBPCG) of the AG pack according to a base partition constraint of tensors of the AG pack, the mlBPCG generated for memory locations of the tensors;
generating, by the compiler, an AG row group constraint graph (agRGCG) based on the memory location base partition constraint graph and a starting row offset of tensors of the AGs, the agRGCG including row group constraints of the AGs for adjacency representation of the AGs;
assigning, by the compiler, the AGs of the AG pack to tiles of a processing engine array based on the AG row group constraint graph; and
generating, by the compiler, machine instructions for the neural network model based on the reordered data dependency graph and the assignment of the AGs to the tiles of the processing engine array.
2 . The computer-implemented method of claim 1 , wherein assigning the AGs of the AG pack to the tiles of the processing engine array includes assigning the AGs in descending order of tensor size.
3 . The computer-implemented method of claim 1 , further comprising assigning an unpacked AG based on tensor sharing between the unpacked AG and the AG pack.
4 . The computer-implemented method of claim 1 , wherein the processing engine array includes at least 4 row groups in a tile arrangement.
5 . The computer-implemented method of claim 1 , wherein generating the memory location base partition constraint graph of the AG pack includes: determining a first base partition constraint of tensors of a matrix multiplication instruction in an AG; and assigning the first base partition constraint to an edge between the tensors of the matrix multiplication instruction.
6 . The computer-implemented method of claim 1 , wherein generating the memory location base partition constraint graph of the AG pack includes: determining a transitive base partition constraint between tensors of different AGs; and assigning the transitive base partition constraint to an edge between the tensors of the different AGs.
7 . The computer-implemented method of claim 1 , wherein generating the AG row group constraint graph includes: determining a row constraint between the AGs of the AG pack based on edges of the memory location base partition constraint graph and the starting row offset between the tensors of the AGs; and assigning the row constraint to an edge between the AGs.
8 . The computer-implemented method of claim 7 , wherein assigning the AGs of the AG pack to the tiles of the processing engine array includes minimizing an objective function based on the row constraint in the AG row group constraint graph.
9 . The computer-implemented method of claim 1 , further comprising associating an AG pack identifier to the one or more matrix multiplication instructions of the AG pack.
10 . The computer-implemented method of claim 1 , further comprising inserting an assigned tile position to a matrix multiplication instruction of the AG pack.
11 . 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 a description of a neural network model;
generating, by compiling the descriptio,n a data dependency graph of the neural network model;
reordering the data dependency graph to form an accumulation group (AG) pack of adjacent accumulation groups (AGs), wherein each AG includes one or more matrix multiplication instructions, and wherein the AG pack represents a consecutive sequence of matrix multiplication instructions;
generating a memory location base partition constraint graph (mlBPCG) of the AG pack according to a base partition constraint of tensors of the AG pack, the mlBPCG generated for memory locations of the tensors;
generating an AG row group constraint graph (agRGCG) based on the memory location base partition constraint graph and a starting row offset of tensors of the AGs, the agRGCG including row group constraints of the AGs for adjacency representation of the AGs;
assigning the AGs of the AG pack to tiles of a processing engine array based on the AG row group constraint graph; and
generating machine instructions for the neural network model based on the reordered data dependency graph and the assignment of the AGs to the tiles of the processing engine array.
12 . The non-transitory computer-readable medium of claim 11 , wherein the operations further include inserting an assigned tile position to a matrix multiplication instruction of the AG pack.
13 . The non-transitory computer-readable medium of claim 12 , wherein the operations further include scheduling the matrix multiplication instruction with the assigned tile position.
14 . The non-transitory computer-readable medium of claim 11 , wherein generating the memory location base partition constraint graph of the AG pack includes: determining a first base partition constraint of tensors of a matrix multiplication instruction in an AG; assigning the first base partition constraint to an edge between the tensors of the matrix multiplication instruction; determining a transitive base partition constraint between tensors of different AGs; and assigning the transitive base partition constraint to an edge between the tensors of the different AGs.
15 . The non-transitory computer-readable medium of claim 11 , wherein assigning the AGs of the AG pack to the tiles of the processing engine array includes minimizing an objective function based on the row group constraints in the AG row group constraint graph.
16 . The non-transitory computer-readable medium of claim 11 , wherein assigning the AGs of the AG pack to tiles of processing engine array includes assigning the AGs in descending order of tensor size.
17 . The non-transitory computer-readable medium of claim 11 , further comprising assigning an unpacked AG based on tensor sharing between the unpacked AG and the AG pack.