Native tensor processor, and partitioning of tensor contractions
A native tensor processor calculates tensor contractions using a sum of outer products. In one implementation, the native tensor processor preferably is implemented as a single integrated circuit and includes an input buffer and a contraction engine. The input buffer buffers tensor elements retrieved from off-chip and transmits the elements to the contraction engine as needed. The contraction engine calculates the tensor contraction by executing calculations from an equivalent matrix multiplications, as if the tensors were unfolded into matrices, but avoiding the overhead of expressly unfolding the tensors. The contraction engine includes a plurality of outer product units that calculate matrix mutiplications by a sum of outer products. By using outer products, the equivalent matrix multiplications can be partitioned into smaller matrix multiplications, each of which is localized with respect to which tensor elements are required.
1. A method for calculating a contraction of tensors TX and TY using a native tensor processor; wherein the contraction of tensors TX and TY is characterized by an equivalent matrix multiplication X(l,n)×Y(n,m)=Z(l,m), where X is an unfolded matrix for tensor TX and Y is an unfolded matrix for tensor TY, and l=1 . . . L and m=1 . . . M are free indices and n=1 . . . N is a contraction index; and the native tensor processor comprises a distribution section, a processing section and a collection section, the processing section comprising a plurality of atomic processing elements (APEs) that (a) calculate outer products α×β where α is an ι×1 column vector and β is a 1×φ row vector and (b) accumulate the outer products α×β over the contraction index n; the method comprising:
partitioning the matrix multiplication X×Y into a set of component matrix multiplications A×B based on a sum of outer products, the partitioning implemented by a combination of (a) hardware parallelism in the distribution and collection sections, (b) time division multiplexing (TDM), and (c) accumulation of the outer products α×β by the APEs;
calculating the component matrix multiplications A×B=C using the APEs; and
combining the component products C into the product Z=X×Y.
2. The method of claim 1 wherein partitioning implemented by hardware parallelism in the distribution and collection sections comprises:
scattering with respect to the contraction index n in the distribution section, and a corresponding reduction in the collection section.
3. The method of claim 1 wherein partitioning implemented by hardware parallelism in the distribution and collection sections comprises:
scattering with respect to the free indices l,m in the distribution section, and a corresponding gathering in the collection section.
4. The method of claim 1 wherein the hardware parallelism in the distribution and collection sections is configurable, and partitioning the matrix multiplication X×Y comprises:
configuring the hardware parallelism in the distribution and collection sections according to the sizes L,M,N.
5. The method of claim 1 wherein partitioning the matrix multiplication X×Y with respect to the contraction index n is implemented entirely by (a) hardware parallelism in the distribution and collection sections, and (c) accumulation of the outer products α×β by the APEs.
6. The method of claim 1 wherein partitioning the matrix multiplication X×Y with respect to the contraction index n comprises implementation by TDM in the distribution and collection sections.
7. The method of claim 1 wherein hardware parallelism with respect to the contraction index n is fixed, and accumulation of outer products α×β over the contraction index n by the APEs is determined by the size N of the contraction index.
8. The method of claim 1 wherein partitioning the matrix multiplication X×Y with respect to the free indices l,m is implemented using both (a) hardware parallelism in the distribution and collection sections, and (b) time division multiplexing (TDM).
9. The method of claim 1 wherein hardware parallelism with respect to the free indices l,m is fixed, and time division multiplexing (TDM) is determined by the size L×M of the free indices.
10. The method of claim 1 wherein the native tensor processor further comprises a plurality of outer product units (OPUs) that calculate matrix mutiplications by a sum of outer products, and partitioning the matrix multiplication X×Y comprises:
partitioning the X×Y matrix multiplication with respect to the contraction index k into a plurality of Xk×Yk outer products and directing the Xk×Yk outer products to the OPUs.
11. The method of claim 10 wherein the OPUs calculate the component matrix multiplications A(i,k)×B(k,j)=C(i,j) by a sum of outer products, where i=1 . . . I and j=1 . . . J are free indices and k=1 . . . K is a contraction index; and partitioning the matrix multiplication X×Y comprises:
partitioning the X×Y matrix multiplication into a set of A×B matrix multiplications executable by the OPUs.
12. The method of claim 11 wherein I, J, K are configurable, and partitioning the X×Y matrix multiplication comprises determining values of I, J, K based on the sizes L, M, N.
13. The method of claim 11 wherein I, J, K are configurable, and partitioning the X×Y matrix multiplication comprises determining values of I, J, K based on the sizes L, M, N, provided that the product I×J is constant.
14. The method of claim 11 wherein accumulation of the outer products α×β by the APEs depends on K.
15. The method of claim 10 wherein all partitioning of the X×Y matrix multiplication with respect to the contraction index n is implemented entirely by (a) hardware parallelism in the distribution section that directs the Xk×Yk outer products to the OPUs, and (c) accumulation of the outer products α×β by the APEs.
16. The method of claim 10 wherein all partitioning of the X×Y matrix multiplication with respect to the free indices l,m is implemented entirely by (a) hardware parallelism in the distribution section between the OPUs and the APEs, and (b) time division multiplexing.
17. The method of claim 10 wherein the native tensor processor further comprises a plurality of intermediate processing elements (IPEs) coupling the OPUs to the APEs, and partitioning the X×Y matrix multiplication comprises:
determining partitioning from the OPUs to the IPEs; and
determining partitioning from the IPEs to the APEs.
18. The method of claim 17 wherein determining partitioning from the IPEs to the APEs maximizes broadcasting of one of the vectors α and β to the APEs coupled to each IPE.
19. The method of claim 1 wherein the APEs operate in parallel to calculate the component matrix multiplications A×B.
20. The method of claim 1 wherein calculating the component matrix multiplications A×B comprises:
retrieving elements of the tensors TX and TY required to calculate the component matrix multiplications A×B without expressly unfolding the tensors TX and TY into the equivalent matrices X and Y.