Quantized neural network training and inference

1 . A computer-implemented method for processing a neural network, the method comprising:

receiving a first quantized two-dimensional square matrix that corresponds to a portion of a multi-dimensional input tensor and has been quantized based on a first scale factor, wherein the multi-dimensional input tensor has more than two dimensions, and wherein a first two-dimensional square matrix sub-divided from the multi-dimensional input tensor is quantized using the first scale factor to generate the first quantized two-dimensional square matrix, and wherein the first two-dimensional square matrix has a number of rows and a number of columns corresponding to a size of a vector multiply-accumulate (MAC) unit that performs computations on the first quantized two-dimensional square matrix; and

performing one or more computational operations using the first quantized two-dimensional square matrix and the first scale factor to generate one or more data values that correspond to a first portion of a multi-dimensional output tensor, wherein the one or more computational operations include one or more memory read operations.

2 . The computer-implemented method of claim 1 , further comprising:

receiving a second quantized two-dimensional square matrix that corresponds to a portion of a second multi-dimensional input tensor and has been quantized based on a second scale factor;

wherein performing the one or more computational operations further comprises using the second quantized two-dimensional square matrix and the second scale factor to generate the one or more data values.

3 . The computer-implemented method of claim 2 , wherein performing the one or more computational operations comprises computing one or more dot-product values using the first quantized two-dimensional square matrix and the second quantized two-dimensional square matrix.

4 . The computer-implemented method of claim 2 , wherein performing the one or more computational operations comprises scaling each dot-product value included in one or more dot-product values using the first scale factor and the second scale factor.

5 . The computer-implemented method of claim 1 , wherein the first quantized two-dimensional square matrix comprises a transposed matrix.

6 . The computer-implemented method of claim 1 , wherein the first quantized two-dimensional square matrix is included in a plurality of quantized two-dimensional square matrices, wherein each two-dimensional quantized square matrix included in the plurality of quantized two-dimensional square matrices corresponds to a different portion of the multi-dimensional input tensor and has been quantized using a different scale factor.

7 . The computer-implemented method of claim 6 , wherein performing the one or more computational operations further comprises, for each quantized two-dimensional square matrix included in the plurality of quantized two-dimensional square matrices, using the quantized two-dimensional square matrix and an associated scale factor to generate one or more associated data values that correspond to a different portion of the multi-dimensional output tensor.

8 . The computer-implemented method of claim 1 , further comprising scaling the multi-dimensional output tensor based on a second scale factor associated with the multi-dimensional output tensor.

9 . The computer-implemented method of claim 1 , further comprising quantizing the multi-dimensional output tensor by subdividing the multi-dimensional output tensor into a plurality of two-dimensional square matrices and quantizing each two-dimensional square matrix included in the plurality of two-dimensional square matrices.

10 . The computer-implemented method of claim 1 , further comprising quantizing the multi-dimensional output tensor by:

computing, for each two-dimensional square matrix included in a plurality of two-dimensional square matrices that comprise the multi-dimensional output tensor, a corresponding scale factor based on one or more data values included in the two-dimensional square matrix; and

quantizing each two-dimensional square matrix included in the plurality of two-dimensional square matrices based on the corresponding scale factor.

11 . One or more non-transitory computer-readable media storing instructions that, when executed by one or more processors, cause the one or more processors to perform the steps of:

receiving a first quantized two-dimensional square matrix that corresponds to a portion of a multi-dimensional input tensor and has been quantized based on a first scale factor, wherein the multi-dimensional input tensor has more than two dimensions, and wherein a first two-dimensional square matrix sub-divided from the multi-dimensional input tensor is quantized using the first scale factor to generate the first quantized two-dimensional square matrix, and wherein the first two-dimensional square matrix has a number of rows and a number of columns corresponding to a size of a vector multiply-accumulate (MAC) unit that performs computations on the first quantized two-dimensional square matrix; and

performing one or more computational operations using the first quantized two-dimensional square matrix and the first scale factor to generate one or more data values that correspond to a first portion of a multi-dimensional output tensor, wherein the one or more computational operations include one or more memory read operations.

12 . The one or more non-transitory computer-readable media of claim 11 , further comprising:

receiving a second quantized two-dimensional square matrix that corresponds to a portion of a second multi-dimensional input tensor and has been quantized based on a second scale factor;

wherein performing the one or more computational operations further comprises using the second quantized two-dimensional square matrix and the second scale factor to generate the one or more data values.

13 . The one or more non-transitory computer-readable media of claim 12 , wherein performing the one or more computational operations comprises computing a first dot-product value between a first vector included in the first quantized two-dimensional square matrix and a second vector included in the second quantized two-dimensional square matrix.

14 . The one or more non-transitory computer-readable media of claim 12 , wherein performing the one or more computational operations comprises scaling a first dot-product value based on at least one of the first scale factor or the second scale factor.

15 . The one or more non-transitory computer-readable media of claim 11 , further comprising receiving the first scale factor and storing the first quantized two-dimensional square matrix and the first scale factor in a data buffer.

16 . The one or more non-transitory computer-readable media of claim 15 , wherein performing the one or more computational operations comprises reading a first set of input values included in the first quantized two-dimensional square matrix from the data buffer; and reading the first scale factor from the data buffer, wherein the first set of input values also is used to perform the one or more computational operations.

17 . The one or more non-transitory computer-readable media of claim 15 , wherein the data buffer comprises a plurality of rows, and wherein storing the first quantized two-dimensional square matrix and the first scale factor in the data buffer comprises storing the first quantized two-dimensional square matrix and the first scale factor in a first row included in the plurality of rows.

18 . The one or more non-transitory computer-readable media of claim 15 , wherein the data buffer comprises a plurality of memory banks, wherein storing the first quantized two-dimensional square matrix in the data buffer comprises for each column included in the first quantized two-dimensional square matrix, storing the column in a different memory bank included in the plurality of memory banks.

19 . The one or more non-transitory computer-readable media of claim 11 , wherein the first quantized two-dimensional square matrix is included in a plurality of quantized two-dimensional square matrices, wherein each quantized two-dimensional square matrix included in the plurality of quantized two-dimensional square matrices corresponds to a different portion of the multi-dimensional input tensor and has been quantized using a different scale factor, and further comprising for each quantized two-dimensional square matrix included in the plurality of quantized two-dimensional square matrices, receiving the quantized two-dimensional square matrix and an associated scale factor, and storing the quantized two-dimensional square matrix and the associated scale factor in a different row of a data buffer.

20 . A system comprising:

one or more memories storing instructions; and

one or more processors that are coupled to the one or more memories and, when executing the instructions, perform the steps of:

receiving a first quantized two-dimensional square matrix that corresponds to a portion of a multi-dimensional input tensor and has been quantized based on a first scale factor, wherein the multi-dimensional input tensor has more than two dimensions, and wherein a first two-dimensional square matrix sub-divided from the multi-dimensional input tensor is quantized using the first scale factor to generate the first quantized two-dimensional square matrix, and wherein the first two-dimensional square matrix has a number of rows and a number of columns corresponding to a size of a vector multiply-accumulate (MAC) unit that performs computations on the first quantized two-dimensional square matrix; and

performing one or more computational operations using the first quantized two-dimensional square matrix and the first scale factor to generate one or more data values that correspond to a first portion of a multi-dimensional output tensor, wherein the one or more computational operations include one or more memory read operations.

21 . The computer-implemented method of claim 1 , wherein the number of rows and the number of columns of the first two-dimensional square matrix are selected to match a number of input pairs that the vector MAC unit receives as input, and wherein the vector MAC unit performs computations on an entire row or an entire column of the first two-dimensional square matrix.