Low latency matrix multiply unit
Methods, systems, and apparatus for a matrix multiply unit implemented as a systolic array of cells are disclosed. Each cell of the matrix multiply includes: a weight matrix register configured to receive a weight input from either a transposed or a non-transposed weight shift register; a transposed weight shift register configured to receive a weight input from a horizontal direction to be stored in the weight matrix register; a non-transposed weight shift register configured to receive a weight input from a vertical direction to be stored in the weight matrix register; and a multiply unit that is coupled to the weight matrix register and configured to multiply the weight input of the weight matrix register with a vector data input in order to obtain a multiplication result.
1 . A processor comprising:
a processing core having an array of cells;
a first plurality of shift chains corresponding to rows of the array of cells;
a second plurality of shift chains corresponding to columns of the array of cells,
wherein the processor is configured to execute instructions that cause the processor to perform operations comprising:
selectively loading weight values using the first plurality of shift chains corresponding to the rows of the array of cells or using the second plurality of shift chains corresponding to the columns of the array of cells, and
performing one or more matrix multiplication operations using the loaded weight values and a vector input.
2 . The processor of claim 1 , further comprising alternately loading weight values using the first plurality of shift chains and the second plurality of shift chains.
3 . The processor of claim 2 , wherein alternately loading the weight values corresponds to different steps of a neural network training process.
4 . The processor of claim 1 , wherein the first plurality of shift chains corresponds to transposed weights and the second plurality of shift chains corresponds to non-transposed weights or wherein the first plurality of shift chains corresponds to non-transposed weights and the second plurality of shift chains corresponds to transposed weights.
5 . The processor of claim 1 , wherein the operations further comprises loading the vector input using the first plurality of shift chains or the second plurality of shift chains.
6 . The processor of claim 1 , further comprising a plurality of weight registers coupled to the first plurality of shift chains and the second plurality of shift chains.
7 . The processor of claim 6 , further comprising a plurality of multiplexors configured to select between loading the weight registers with data from the first plurality of shift chains or the second plurality of shift chains.
8 . The processor of claim 1 , wherein the processing core further includes one or more of a scalar memory, a vector memory, a scalar processing unit, one or more vector registers, a transposed unit, or a reduction and permutation unit.
9 . A method performed by a processor, wherein the processor comprises a processing core having an array of cells, a first plurality of shift chains, and a second plurality of shift chains, the method comprising:
selectively loading weight values using the first plurality of shift chains corresponding to the rows of the array of cells or using the second plurality of shift chains corresponding to the columns of the array of cells; and
performing one or more matrix multiplication operations using the loaded weight values and a vector input.
10 . The method of claim 9 , further comprising alternately loading weight values using the first plurality of shift chains and the second plurality of shift chains.
11 . The method of claim 10 , wherein alternately loading the weight values corresponds to different steps of a neural network training process.
12 . The method of claim 9 , wherein the first plurality of shift chains corresponds to transposed weights and the second plurality of shift chains corresponds to non-transposed weights or wherein the first plurality of shift chains corresponds to non-transposed weights and the second plurality of shift chains corresponds to transposed weights.
13 . The method of claim 9 , wherein the operations further comprises loading the vector input using the first plurality of shift chains or the second plurality of shift chains.
14 . The method of claim 9 , further comprising a plurality of weight registers coupled to the first plurality of shift chains and the second plurality of shift chains.
15 . The method of claim 14 , further comprising a plurality of multiplexors configured to select between loading the weight registers with data from the first plurality of shift chains or the second plurality of shift chains.
16 . The method of claim 9 , wherein the processing core further includes one or more of a scalar memory, a vector memory, a scalar processing unit, one or more vector registers, a transposed unit, or a reduction and permutation unit.
17 . A non-transitory computer program product storing instructions that, when executed by a programmable processor comprising a processing core having an array of cells, a first plurality of shift chains, and a second plurality of shift chains, cause the processor to perform operations comprising:
selectively loading weight values using the first plurality of shift chains corresponding to the rows of the array of cells or using the second plurality of shift chains corresponding to the columns of the array of cells; and
performing one or more matrix multiplication operations using the loaded weight values and a vector input.
18 . The non-transitory computer program of claim 17 , further comprising alternately loading weight values using the first plurality of shift chains and the second plurality of shift chains.
19 . The non-transitory computer program of claim 18 , wherein alternately loading the weight values corresponds to different steps of a neural network training process.