Efficient matrix multiply and add with a group of warps
This specification describes techniques for implementing matrix multiply and add (MMA) operations in graphics processing units (GPU)s and other processors. The implementations provide for a plurality of warps of threads to collaborate in generating the result matrix by enabling each thread to share its respective register files to be accessed by the datapaths associated with other threads in the group of warps. A state machine circuit controls a MMA execution among the warps executing on asynchronous computation units. A group MMA (GMMA) instruction provides for a descriptor to be provided as parameter where the descriptor may include information regarding size and formats of input data to be loaded into shared memory and/or the datapath.
1 . A processor comprising:
a plurality of cores, each core comprising a datapath;
a work distributor circuit configured to distribute a plurality of warps to the plurality of cores, the plurality of warps including a plurality of threads cooperating to execute a matrix multiply and accumulate (MMA) instruction to determine a result matrix based on input matrices;
a state machine circuitry configured to coordinate execution of the MMA instruction to generate the result matrix;
a register file;
a shared memory; and
an interconnection switch,
wherein each core executes the MMA instruction to generate one or more elements of the result matrix at an output of the datapath of said each core, and wherein the state machine circuitry is further configured to copy operands of said input matrices from at least one of a portion of the shared memory or a portion of the register file of any thread of any said plurality of warps to an input of the datapath of any of said plurality of cores.
2 . The processor according to claim 1 , further comprising a register file and an interconnection switch configured to copy operands of said input matrices from a portion of the register file of any thread of any said plurality of warps to an input of the datapath of any of said plurality of cores.
3 . The processor according to claim 1 , wherein the processor further comprises a shared memory and an interface to an external memory, and the interconnection switch is further configured to copy operands of said input matrices from a portion of the shared memory to an input of the datapath of any of said plurality of cores.
4 . The processor according to claim 1 , wherein the state machine circuitry is further configured to initiate said copying of operands of the input matrices responsive to all of said plurality of warps satisfying a synchronization barrier.
5 . The processor according to claim 4 , wherein the state machine circuitry is further configured to signal to all of said plurality of warps an end of the MMA instruction.
6 . The processor according to claim 5 , wherein the state machine circuitry is further configured to update a second synchronization barrier to signal the end of the MMA instruction to all of said plurality of warps.
7 . The processor according to claim 5 , wherein the state machine circuitry is further configured to, between said initiating of the copying of operands of the input matrices and the signaling of the end of the MMA instruction to the plurality of warps, update one or more additional synchronization barriers to signal completion of calculating respective parts of the result matrix to the plurality of the warps.
8 . The processor according to claim 4 , wherein an instruction stream executed by the plurality of warps includes the MMA instruction and a MMA synchronization instruction that is arranged after the MMA instruction in the instruction stream, the MMA synchronization instruction indicating a completion of one or more previous MMA instructions to the plurality of warps.
9 . The processor according to claim 1 , comprising a plurality of streaming multiprocessors, wherein a streaming multiprocessor of the plurality of streaming multiprocessors comprises the plurality of cores and the work distributor circuit.
10 . The processor according to claim 1 , wherein the datapath includes a tensor core.
11 . A processor comprising:
a plurality of cores, each core comprising a datapath;
a work distributor circuit configured to distribute a plurality of warps to the plurality of cores, the plurality of warps including a plurality of threads cooperating to execute a matrix multiply and accumulate (MMA) instruction to determine a result matrix based on input matrices, wherein each core executes the MMA instruction to generate one or more elements of the result matrix at an output of the datapath of said each core, and wherein, for calculating elements of a particular row of the result matrix, a same portion of an input matrix is provided to one or more cores in the plurality of cores by providing data of the same portion obtained from shared memory with each read operation to said one or more cores.
12 . The processor according to claim 11 , wherein the calculating uses a read first input matrix and read portions of a second input matrix from the respective shared memories, and a third input matrix in respective register files.
13 . The processor according to claim 12 , wherein respective portions of the result matrix is written to the respective register files of each thread.
14 . The processor according to claim 2 , wherein the read first input matrix and the read portions of the second input matrix are subsequently loaded to inputs of the datapaths of the plurality of cores without being stored in register files of the plurality of warps.
15 . The processor according to claim 12 , wherein an instruction to read the first input matrix or the second input matrix includes a descriptor identifying a location and/or layout of data in the memory and wherein the processor further comprises a state machine circuitry configured to obtain the first input matrix or the second input matrix in accordance with a descriptor.
16 . The processor according to claim 12 , wherein the processor further includes transpose circuitry configured to transpose the data read from the memory before the data is written to the shared memory and/or the register files.
17 . The processor according to claim 12 , wherein a sequence of instructions further causes, at least one of the first input matrix to be read or respectively different portions of a second input matrix to be read to be multicast to the respective cores.
18 . A method comprising receive an instruction stream including a matrix multiply and accumulate (MMA) instruction to determine a result matrix based on input matrices;
distribute a plurality of warps to a plurality of cores, each core comprising a datapath, each warp including a plurality of threads executing instructions from the instruction stream;
execute, by each core of the plurality of cores, the MMA instruction to generate one or more elements of the result matrix at an output of the datapath of said each core;
coordinating, by a state machine circuitry, the execution of the MMA instruction to generate the result matrix; and
copying, by the state machine circuitry, operands of said input matrices from at least one of a portion of a shared memory of any thread of any said plurality of warps or a portion of a register file of any thread of any said plurality of warps to an input of the datapath of any of said plurality of cores.