Template-based High-performance GPU Tensor Contraction Method

Abstract

The present invention discloses a template-based high-performance GPU tensor contraction method. First, the user defines tensor contraction, classifies its indices and dimensions to obtain four index sequences and four dimension sequences; then performs dimensionality reduction on it to obtain a memory access function and implicit dimensions, defines placeholders, and writes a calculation template according to the BLAS library implementation and optimization method. Then, template dispatch during compilation is performed, substituting the memory access function and implicit dimensions into the placeholders of the selected calculation template to generate CUDA C / C++ code, which is compiled into a reusable executable program. Finally, the data of each tensor and the specific values of each dimension are input to complete the calculation. The present invention can support any tensor contraction with a finite set of templates; while maintaining ease of use, the system performance is still close to or even exceeds manually optimized computing libraries such as cuBLAS and cuDNN.

Description

Technical Field

[0001] The present invention belongs to the field of high-performance computing, and particularly relates to a template-based GPU high-performance tensor contraction method. Background Art

[0002] The concept of tensor is a generalization of vectors, and scalars, vectors, and matrices can be understood as 0, 1, and 2-order tensors respectively. It is widely used in fields such as machine learning, data analysis, and scientific computing, providing a concise and clear abstract representation for data and algorithms. However, tensor calculations in system implementation are very complex: even though memory access-intensive operators such as normalization and downsampling are relatively easy to optimize, compute-intensive operators such as matrix multiplication and convolution are difficult to approach the peak performance of the hardware and often require experts familiar with the target platform architecture for meticulous manual optimization. At the same time, users often highly customize algorithms according to real-world requirements, and the calculations involved often exceed the expression range of common operators, leading to the so-called "long tail" problem, that is, the data layout and calculation mode of tensors are not supported by fixed-function computing libraries such as cuBLAS and cuDNN, and writing high-performance implementations brings huge time and mental burdens to users.

[0003] Tensor calculations generally involve some common patterns. For example, tensors are often stored as multi-dimensional arrays, and the calculation process often consists of multiple nested loops. Therefore, system implementers can design domain-specific languages (DSLs) to specifically express and optimize this type of operation.

[0004] For example, the machine learning software stack PlaidML (Intel.Plaidml.https: / / github.com / plaidml / plaidml.Accessed:2021-10-01.) includes an eDSL for users to customize long-tail operators, and Facebook designed Tensor Comprehensions (TensorComp; Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary Devito, William S. Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated gpukernels, automatically. ACM Trans. Archit. Code Optim., 16(4), October 2019.) for the same purpose. The syntax and semantics of both are based on the Einstein summation convention and can support arbitrary tensor contractions, i.e., the generalization of matrix multiplication in tensors: some dimensions of the input tensors are retained in the output tensors, while the elements on other dimensions are pairwise multiplied and accumulated. The PlaidML eDSL and TensorComp are easy to use and perform well in memory-intensive operators, but they have not considered the actual architectural characteristics of the backend platform, so there are performance defects in computationally intensive scenarios.

[0005] According to the peak computing performance and maximum communication bandwidth of the target platform, tensor computations in fields such as machine learning, data analysis, and scientific computing can be roughly divided into two categories. Operators such as downsampling, reduction, and regularization generate relatively small amounts of computation, and the performance bottleneck lies mainly in the bandwidth. Therefore, the underlying implementation of such operations does not require particularly fine optimization. As long as it can adapt to the storage structure of the target platform, eliminate invalid memory accesses, and maintain sufficient parallelism, it is sufficient to maximize the bandwidth usage, and at this time, factors such as block size, parallel granularity, and instruction arrangement will not have a significant impact. In contrast, operators such as matrix multiplication and convolution exhibit computationally intensive characteristics, and highly optimized implementations can approach the peak arithmetic performance of the system, but achieving this goal requires a considerable amount of effort: the implementer needs to reasonably partition the data onto multi-level parallel units, balance parallelism and data reuse, and carefully arrange different types of instructions to reduce pipeline stalls, etc.

[0006] Tensor contraction is an arbitrary generalization of matrix multiplication in the high-dimensional case, which can be represented by the Einstein summation convention. With a slight extension, it can also cover convolution operations and their various variants. That is to say, common compute-intensive operators are within the scope of tensor contraction, which can serve as the basic unit of tensor computation. Additionally, tensor contraction can also support many memory-intensive operators. Operations such as pooling and reduction are equivalent to special forms of tensor contraction, and more complex operators such as regularization and softmax can also be expressed as combinations of multiple tensor contractions. Therefore, the high-performance implementation of any tensor contraction is taken as the target problem.

[0007] PlaidML eDSL and TensorComp introduced the polyhedral model to attempt to solve this problem. They transform the front-end user input into multi-level nested loops, perform platform-independent general loop transformations based on the polyhedral model in the middle end, and sink the intermediate representation to machine code for a specific backend. This approach performs well in memory-intensive tasks but ignores the architectural characteristics of the target platform. Therefore, its performance is far lower than the theoretical peak in tensor contraction, which is mostly compute-intensive. This problem is even more serious on heterogeneous computing platforms. For example, the performance of TensorComp in large-scale matrix multiplication can be as low as 20% of the native linear algebra library cuBLAS on NVIDIA GPUs. However, due to its excellent parallel processing ability, general-purpose GPUs are widely used in tensor computation and have become de facto infrastructure in fields such as deep learning. Therefore, high-performance tensor contraction on GPUs is a topic of great practical significance: The solution based on the native computing library is trapped by the long-tail problem and has a narrow application range, while the existing tensor computation DSLs are flexible and easy to use, but their performance still needs to be improved.

[0008] Tensor contraction is suitable for computational acceleration on general GPUs because it can be recursively divided into many independent sub-problems, and the control flow of the computational process is static and has no dependence on the input data. The difficulty in implementing tensor contraction on this multi-core architecture lies in that the implementer must comprehensively consider the computational and storage characteristics of the target platform, such as the number of parallel processing units, the latency, bandwidth, and capacity of each level of cache, and divide the tensor into appropriate sizes in each dimension to improve data reuse in the blocks as much as possible and avoid bandwidth bottlenecks. At the same time, it is necessary to alternately arrange memory access and computational instructions to hide communication latency, reduce blocking, and approach peak computational performance. In addition, the SIMT (Single Instruction, Multiple Threads) model of general GPU computing (Erik Lindholm, John Nickolls, Stuart Oberman, and John Montrym. Nvidia tesla: A unified graphics and computing architecture. IEEE Micro, 28(2): 39-55, 2008.) introduces many specific problems such as memory access transaction merging and bank conflict, which pose more subtle requirements for the algorithm implementation of tensor contraction.

[0009] In addition, if arbitrary tensor contractions are to be supported, system implementers will face a huge problem space because the dimensions of higher-order tensors can be arranged in any order, resulting in factorial-level complexity. In the BLAS numerical computing libraries of various platforms, matrix multiplication generally only handles four possible cases, namely the 2×2 combinations generated by each input tensor being column-major or row-major; but the cases faced by tensor contraction are often much more than this. Take the tensor contraction in the form of C a,b,d,e = A a,b,c × B c,d,e as an example. When contracting one dimension c between two third-order tensors, there are a total of 4!×3!×3! = 846 cases for all possible data layouts of this operation; if contracting the two dimensions b and c in the form of C a,d = A a,b,c × B b,c,d , there are also 2!×3!×3! = 72 cases. It is obviously infeasible to provide implementations for each case separately, which requires the system to have good applicability and scalability and be able to quickly solve a class of similar computations. Summary of the Invention

[0010] The purpose of the present invention is to provide a template-based high-performance GPU tensor contraction computing technology in view of the deficiencies of the prior art. The present invention can compute any tensor contraction on a general GPU and performs particularly well in terms of performance.

[0011] The object of the present invention is achieved by the following technical solutions: A GPU high-performance tensor contraction method based on templates, comprising the following steps:

[0012] (1) The user inputs the definition of tensor contraction, classifies its indices and dimensions, and obtains four index sequences p, x, y, r and four dimension sequences P, Y, X, R;

[0013] (2) Reduce the dimension of the index sequences p, x, y, r and the dimension sequences P, Y, X, R to obtain a memory access function and implicit dimensions;

[0014] (3) Define placeholders to represent the content related to the memory access function and implicit dimensions in the implicit batch matrix multiplication;

[0015] (4) Write a calculation template according to the BLAS library implementation and optimization method, and reserve the placeholders described in step (3);

[0016] (5) Perform template dispatch during compilation, substitute the memory access function and implicit dimensions into the placeholders of the selected calculation template, and generate CUDA C / C++ code;

[0017] (6) Compile the code generated in step (5) into a reusable executable program;

[0018] (7) Input the data of each tensor and the specific values of each dimension, and use the executable program compiled in step (6) to complete the calculation.

[0019] Further, the step (1) includes the following sub-steps:

[0020] (1.1) Notation convention: For an n-dimensional tensor T with dimensions from high to low being D = D n , …, D 1 , it is denoted as T D ; use T[i] = T[i n ,..., i 1 to represent the scalar at the index sequence i = (i n ,..., i 1 ); the function ext(i k ) = D k , k = 1,..., n maps a single index to the corresponding dimension, and Ext(i) = (ext(i n ),..., ext(i 1 )) = (D n ,..., D 1 ) = D maps the index sequence to the corresponding dimension sequence; any tensor contraction C = A × B input by the user is represented in the Einstein summation convention:

[0021] C[ic] = A[ia] × B[ib] (1)

[0022] Among them, ic, ia, and ib are the index sequences of tensors C, A, and B respectively.

[0023] (1.2) Define the batch index sequence: Define p = ic ∩ ia ∩ ib, that is, the sequence composed of all batch indexes, and each element is used to index all tensors. The order in which each element appears in ic, ia, and ib may be inconsistent. Any order can be adopted here, and the same is true for the following several sequences;

[0024] (1.3) Define the free index sequences y and x of the input tensors A and B respectively; Define y = (ic ∩ ia)\p, that is, the sequence composed of all free indexes of A, which is only used to index C and A; Define x = (ic ∩ ib)\p as the sequence composed of all free indexes of B, which is only used to index C and B.

[0025] (1.4) Define the reduction index sequence: Define r = (ia ∩ ib)\p as the reduction index sequence. The elements of tensors A and B on Ext(r) are multiplied pairwise and reduced to a scalar.

[0026] (1.5) Define the permutation function: Perform an equivalent transformation on formula (1), substitute p, y, x, r for ic, ia, ib, and introduce three permutation functions shfl C (·), shfl A (·) and shfl B (·), so as to rewrite the original formula as:

[0027] C[shfl C (p, y, x)] = A[shfl A (p, y, r)] × B[shfl B (p, r, x)] (2)

[0028] (1.6) Define the dimension sequences corresponding to each index sequence: Define the batch dimension sequence P = Ext(p), where each element is the dimension corresponding to each index in the batch index sequence; Define the free dimension sequence of A as Y = Ext(y), the free dimension sequence of B as X = Ext(x), and the reduction dimension sequence as R = Ext(r).

[0029] Furthermore, the step (2) includes the following sub-steps:

[0030] (2.1) Define the dimension reduction auxiliary function: First, define the mapping i = gath from the index sequence to the scalar index d(i), where d is an arbitrary sequence of dimensions, but must be the same length as the index sequence i. This function maps the offsets i in each dimension of d to the linear address i:

[0031]

[0032] Define the inverse function of the above function i = scat d (i) = gath d -1 (i). It decomposes the linear address i back into the offsets i in each dimension of d:

[0033]

[0034] With the help of the above two auxiliary functions, the high-order tensor in Equation (2) is implicitly transformed into a matrix in a specific storage format. (2.2) Reduce the dimensions of the indices so that elements in the tensor can be accessed using scalar indices

[0035] Convert the index sequences p, y, x, r into the components of the scalar indices in the corresponding dimensions respectively. Define p = gath P (p), where p is a scalar index with a value range of [0, Π P ) and represent the index sequence in reverse as p = gath P -1 (p) = scat P (p); represent the other index sequences as y = scat Y (y), x = scat X (x), r = scat R (r), where the value ranges of the newly defined scalar indices are y ∈ [0, ΠY), x ∈ [0, ΠX), r ∈ [0, ΠR); rewrite Equation (2) as:

[0036]

[0037] (2.3) Reduce the dimensions of the dimensions and represent a tensor of any order and any dimension arrangement as a special matrix

[0038] The actual calculation that occurs in Equation (7) is:

[0039]

[0040] Denote the access to the output tensor on the left side of Equation (8) as C(p, y, x) = C[shfl C (scat P (p), scat Y (y), scat X (x))], and denote the access to the two input tensors on the right side as A(p, y, r) = A[shflA (scat P (p),scat Y (y),scat R (r))]、B(p,r,x) = B[shfl B (scat P (p),scat R (r),scat X (x))], rewrite Equation (8) as:

[0041]

[0042] where it only contains 4 scalar dimensions, i.e., P = ΠP, Y = ΠY, X = ΠX, R = ΠR.

[0043] (2.4) Implicitly convert tensor contraction into matrix multiplication: Denote C(p, y, x), A(p, x, r), and B(p, r, x) as the memory access functions of the output and input tensors, which abstract the tensors into batch matrices with specific data layouts. Call the four scalars P, Y, X, and R implicit dimensions, which are the batch dimension, the free dimension of A, the free dimension of B, and the product of the reduction dimensions respectively. Equation (9) hides the data reading and writing process with the memory access functions and expresses the traversal and iteration ranges with the implicit dimensions, thus converting tensor contraction into batch matrix multiplication. With the aid of Einstein summation convention, Equation (9) can be expressed as:

[0044] C(p, y, x) = A(p, y, r) × B(p, r, x) (10)

[0045] Furthermore, step (3) includes the following sub - steps:

[0046] (3.1) Define implicit dimension placeholders: Equation (9) contains four implicit dimensions P, Y, X, R, which represent the scales of the implicit batch matrices and are the products of the corresponding dimension sequences in the original tensors respectively. Therefore, define four corresponding implicit dimension placeholders ph_P, ph_Y, ph_X, ph_R, corresponding to the above - mentioned implicit dimensions respectively.

[0047] (3.2) Define memory access function placeholders: In Equation (10), the reading and writing of the implicit matrices are abstracted by three memory access functions A(p, y, r), B(p, r, x), and C(p, y, x) respectively; define three memory access function placeholders ph_A(p, y, k), ph_B(p, k, x), and ph_C(p, y, x) to represent the elements at the corresponding positions.

[0048] Further, the step (4) is specifically as follows: write the calculation template code for linear algebra operations according to the BLAS library implementation and optimization method; each time the matrix dimension is referenced, write it as the placeholder for the four implicit dimensions of ph_P, ph_Y, ph_X, and ph_R; each time any element of the matrix is accessed, write it as the placeholder for the three memory access functions of ph_A(p, y, k), ph_B(p, k, x), and ph_C(p, y, x).

[0049] Further, the step (5) is specifically as follows: analyze the content of the four dimension sequences of P, Y, X, and R from the symbolic representation formula (9) of tensor contraction, and select a specific linear algebra calculation template at compile time according to whether they are empty:

[0050]

[0051] After selecting the template, substitute the placeholder into the calculation template to generate legal CUDA C / C++ code.

[0052] The beneficial effects of the present invention are as follows:

[0053] (1) The present invention designs a general tensor calculation framework. Through concepts such as memory access functions and implicit dimensions, any tensor contraction is transformed into basic linear algebra operations represented by GEMM. Through the dispatch in two stages of compile time and run time, the input tensor contraction is mapped to a specific calculation template, so as to support any tensor contraction with a finite set of templates;

[0054] (2) The present invention designs a DSL based on Einstein summation convention as the front-end representation of the general framework and implements it on NVIDIA GPUs. On the premise of maintaining ease of use, the system performance is still close to or even exceeds the manually optimized computing libraries such as cuBLAS and cuDNN. Description of the Drawings

[0055] Figure 1 is the flowchart of the method of the present invention. Detailed Embodiments

[0056] The present invention designs a tensor contraction method based on a calculation template. Without occupying additional space, tensors with arbitrary orders and arbitrary dimensional arrangements are implicitly transformed into matrices, so as to transform any tensor contraction into a set of finite linear algebra operations, and the implementation of these linear algebra operations is supported by the calculation template.

[0057] The calculation template uses the same syntax as the native code of the target platform, that is, CUDA C / C++. A set of placeholders to be completed are reserved, and specific content is filled in according to the definition of tensor contraction, and executable code is generated by compilation. Finally, the tensor data input of the user is accepted, and the calculation result of tensor contraction is obtained.

[0058] The performance provided by the computing template is sufficient to rival highly optimized numerical computing libraries on the target platform, such as cuBLAS, cuDNN, etc. on NVIDIA GPUs, while retaining a certain code generation ability to support a class of similar computations.

[0059] The method includes the following steps:

[0060] 1. The user inputs the definition of tensor contraction, classifies its indices and dimensions, and obtains four index sequences p, x, y, r and four dimension sequences P, Y, X, R.

[0061] 1.1 Notation convention

[0062] First, the notations used in the present invention are agreed upon. The dimensions from high to low are D = D n , …, D 1 of the n-dimensional tensor T, denoted as T D . Taking specific offsets on each of the n dimensions can uniquely determine an element in T, and T[i] = T[i n ,..., i 1 represents the scalar at the index sequence i = (i n ,..., i 1 ). The function ext(i k ) = D k , k = 1,..., n maps a single index to the corresponding dimension, and Ext(i) = (ext(i n ),..., ext(i 1 )) = (D n ,..., D 1 ) = D maps the index sequence to the corresponding dimension sequence.

[0063] Any tensor contraction C = A × B input by the user is expressed in the Einstein summation convention:

[0064] C[ic] = A[ia] × B[ib] (1)

[0065] where ic, ia, and ib are the index sequences of tensors C, A, and B respectively.

[0066] 1.2 Define the batch index sequence

[0067] Define p = ic ∩ ia ∩ ib, that is, the sequence composed of all batch indices, and each element is used to index all tensors. The order of the elements in p may be inconsistent in ic, ia, and ib, and any order can be adopted here, and the same applies to the following several sequences.

[0068] 1.3 Define the free index sequences y and x of the input tensors A and B respectively

[0069] Define y = (ic ∩ ia)\p, that is, the sequence composed of all free indices of A, which is only used to index C and A.

[0070] Similarly, define x = (ic ∩ ib)\p as the sequence composed of all free indices of B, which is only used to index C and B.

[0071] 1.4 Define the reduction index sequence

[0072] Define r = (ia ∩ ib)\p as the reduction index sequence. The elements of tensors A and B on Ext(r) are multiplied pairwise and reduced to a scalar.

[0073] 1.5 Define the permutation function

[0074] Perform an equivalent transformation on Equation (1), substitute p, y, x, r for ic, ia, ib, and introduce three permutation functions shfl C (·), shfl A (·) and shfl B (·), so as to rewrite the original formula as

[0075] C[shfl C (p, y, x)] = A[shfl A (p, y, r)] × B[shfl B (p, r, x)] (2)

[0076] The reason for introducing three permutation functions shfl C (·) and setting ic = shfl C (p, x, y) is as follows: All indices in p, y, x may appear alternately in any order and serve as the offsets on each dimension of tensor C; shfl C (·) abstracts all possible permutations of the indices on tensor C. Correspondingly, it also represents all possible dimension permutations of tensor C. Similarly, the corresponding permutation functions shfl A (·) and shfl B (·) are introduced for the two input tensors, representing all possible dimension permutations of tensors A and B.

[0077] 1.6 Define the dimension sequences corresponding to each index sequence

[0078] Define the batch dimension sequence P = Ext(p), where each element is the dimension corresponding to each index in the batch index sequence. Similarly, define the free dimension sequence of A as Y = Ext(y), the free dimension sequence of B as X = Ext(x), and the reduction dimension sequence as R = Ext(r).

[0079] Example 1: Batch matrix multiplication C T,M,N = A T,M,K × B T,K,N , which can be expressed using Einstein summation convention as

[0080] C[t, m, n] = A[t, m, k] × B[t, k, n]. (3)

[0081] t appears in all tensors and is the batch index, T = ext(t) is the corresponding batch dimension. m appears in the output tensor C and the input tensor A and is the free index of A, M = ext(m) is called the free dimension of A; similarly, n is the free index of the input tensor B, and N = ext(n) is the free dimension of B. k that only appears in the two input tensors is the reduction index, K = ext(k) is the reduction dimension, and the matrix multiplication multiplies the elements of the input tensors pairwise in this dimension and accumulates (reduces) them to a scalar.

[0082] Therefore, in the batch matrix multiplication shown in Equation (3), the index sequences are p = (t), y = (m), x = (n), r = (k), and the dimension sequences are P = (T), Y = (M), X = (N), R = (K).

[0083] Example 2: Equation (4) represents that the fourth-order tensors A and B are contracted in the dimensions where the reduction indices a and b are located. c and d are the free indices of A, e and f are the free indices of B, and there is no batch index.

[0084] C[d, f, c, e] = A[a, b, c, d] × B[a, e, b, f] (4)

[0085] Therefore, in the tensor contraction shown in Equation (4), the index sequences are p = (), y = (c, d), x = (e, f), r = (a, b), and the permutation functions are shfl C ((), (c, d), (e, f)) = (d, f, c, e), shfl A ((), (c, d), (a, b)) = (a, b, c, d), shfl B ((), (a, b), (e, f)) = (a, e, b, f).

[0086] 2 Dimensionality reduction is performed on the 2 pairs of index sequences p, x, y, r and dimension sequences P, Y, X, R to obtain the memory access function and implicit dimensions

[0087] The tensor is represented as a matrix in a special storage format, enabling the reuse of the same calculation process for tensor contraction and matrix multiplication (and their respective degenerate cases).

[0088] 2.1 Define the dimensionality reduction auxiliary function

[0089] The present invention establishes a mapping between the dimensions of tensors and matrices, thereby transforming the multi-linear indices on high-order tensors into scalar indices on matrix batches, rows, and columns.

[0090] The index sequence in tensor contraction corresponds to the scalar index in matrix multiplication. Therefore, first define the mapping of the index sequence to the scalar index i = gath d (i), where d can be any sequence of dimensions but must be the same length as the index sequence i. This function maps the offset i on each dimension in d to the linear address i:

[0091]

[0092] Correspondingly, define the inverse function of the above function i = scat d (i) = gath d -1 (i). It re-decomposes the linear address i into the offset i on each dimension in d:

[0093]

[0094] With the help of the above two auxiliary functions, the high-order tensor in Equation (2) can be implicitly transformed into a matrix with a specific storage format.

[0095] 2.2 Reduce the dimensions of the indices to access the elements in the tensor with scalar indices

[0096] Convert the previously defined index sequences p, y, x, r into the components of the scalar index on the corresponding dimensions respectively. Define p = gath P (p). It is easy to know that p is a scalar index with a value range of [0, ΠP). Therefore, the index sequence can be expressed conversely as p = gath P -1 (p) = scat P (p). Similarly, express the other index sequences as y = scat Y (y), x = scat X (x), r = scat R (r), where the value ranges of the newly defined scalar indices are y ∈ [0, ΠY), x ∈ [0, ΠX), r ∈ [0, ΠR).

[0097] Use the scalar indices p, y, x, r to substitute the original index sequence, and rewrite Equation (2) as:

[0098]

[0099] At this time, any element in each tensor can be uniquely determined by three scalar indices.

[0100] 2.3 Reduce the dimension, and represent a tensor arranged in any order and any dimension as a special matrix

[0101] The actual calculation in Equation (7) is as follows:

[0102]

[0103] The above equation uses index mapping functions such as scat P (·), permutation functions such as shfl C (·), which contain the specific information of tensor contraction, including the data layout of the tensor and the calculations occurring between them. Once a certain tensor contraction is given, these functions are fixed, while the four scalar indices p, y, x, r can be regarded as variables to traverse the output and input tensors and perform calculations. Therefore, the access to the output tensor on the left side of Equation (8) is denoted as C(p, y, x) = C[shfl C (scat P (p), scat Y (y), scat X (x))], the access to the two input tensors on the right side is denoted as A(p, y, r) = A[shfl A (scat P (p), scat Y (y), scat R (r))] and B(p, r, x) = B[shfl B (scat P (p), scat R (r), scat X (x))], and rewrite Equation (8) as:

[0104]

[0105] which only contains 4 scalar dimensions, i.e., P = ΠP, Y = ΠY, X = ΠX, R = ΠR.

[0106] 2.4 Implicitly convert tensor contraction into matrix multiplication

[0107] Comparing the tensor contraction formula (9) with the batch matrix multiplication (3), it can be seen that their forms are exactly the same. Regarding C(p, y, x), A(p, x, r), and B(p, r, x) as the memory access functions of the output and input tensors, they abstract the tensors into batch matrices with specific data layouts. Four scalars P, Y, X, and R are called implicit dimensions, which are the batch dimension, the free dimension of A, the free dimension of B, and the product of the reduction dimensions respectively. Equation (9) hides the data reading and writing process with the memory access function and expresses the traversal and iteration range with the implicit dimensions, thus transforming the tensor contraction into batch matrix multiplication. With the aid of the Einstein summation convention, Equation (9) can be expressed as:

[0108] C(p,y,x) = A(p,y,r) × B(p,r,x) (10)

[0109] So far, any tensor contraction has been implicitly transformed into batch matrix multiplication.

[0110] 3 Define placeholders representing the content related to the memory access function and implicit dimensions in the implicit batch matrix multiplication

[0111] 3.1 The present invention completes the calculation of tensor contraction through the code implementation of batch matrix multiplication. To achieve this goal, placeholders need to be reserved in the batch matrix multiplication and specific content is filled in according to the definition of tensor contraction.

[0112] 3.2 Define the implicit dimension placeholder

[0113] Equation (9) contains four implicit dimensions P, Y, X, and R, which represent the scale of the implicit batch matrix and are respectively the products of the corresponding dimension sequences in the original tensors. Therefore, four corresponding placeholders ph_P, ph_Y, ph_X, and ph_R are introduced, corresponding to the above implicit dimensions respectively.

[0114] Example 1: For the tensor contraction C M,N = A K1,M,K2 × B K1,N,K2 as an example, it is very close to matrix multiplication, but there are reductions in two dimensions K1 and K2. Then at this time, the implicit reduction dimension ph_R = R = K 1 K 2 .

[0115] 3.3 Define the memory access function placeholder

[0116] In batch matrix multiplication, three scalar indices can uniquely determine any element in the batch matrix. In Equation (10), the reading and writing of the implicit matrix are abstracted by three memory access functions A(p, y, r), B(p, r, x), and C(p, y, x) respectively. Correspondingly, the present invention introduces three placeholders ph_A(p, y, k), ph_B(p, k, x), and ph_C(p, y, x) to represent the elements at the corresponding positions.

[0117] Embodiment 2: Each dimension of the tensor is a symbolic representation known at compile time, but its specific value is a variable that can only be determined at runtime. Therefore, the logically high-order tensor must be flattened into a one-dimensional array during implementation. For example, in single-precision batch matrix multiplication, the output tensor C T,M,N may be "float*C;" in implementation, and its shape is defined by the dimension variables "int T, M, N;". At this time, the result obtained by instantiating the placeholder ph_C(p, y, x) of the output tensor is "O[p*M*N + y*N + x]".

[0118] Embodiment 3: The implicit dimension and the memory access function together reflect the mutual conversion between the tensor and the implicit batch matrix. Table 1 further illustrates the meanings of common placeholders and gives examples of the application of these two types of placeholders, namely the implicit dimension and the memory access function, in a two-dimensional depth convolution O with N feature maps, K groups of convolution kernels, and a channel number of C N,K,C,Ho,Wo = I N,C,Hi,Wi × W K,C,R,S Among them, represents truncated integer division a % b represents the modulo operation a mod b.

[0119] Table 1: Explanation and Examples of Placeholders

[0120]

[0121]

[0122] 4 Write a calculation template according to the implementation and optimization method of the BLAS library, where the placeholders described in step (3) are reserved

[0123] 4.1 The solution provided by the present invention is that the computation template is still written in the native language of the target platform, namely CUDA C / C++. The computation template is legal CUDA code at the syntax and semantic levels. The only exception is that a series of special symbol representations are defined, which are called placeholders (abbreviated as ph in pseudocode). The batch matrices implicitly converted from tensors are abstracted by placeholders, and the placeholders define all relevant information such as memory access functions, implicit dimensions, data types, and the names of tensors and dimensions. Through the information contained in the placeholders, the computation template reads tensor data into the cache and implements subsequent high-performance batch matrix multiplication independent of memory access. The present invention embeds the data representation of tensors into the positions where the placeholders are located through mechanisms such as macros, inline functions, and template parameters, thereby generating legal and readable CUDA code.

[0124] 4.2 Using the existing BLAS library implementation and optimization methods, write a computation template for matrix multiplication.

[0125] There are three possible degenerate cases for matrix multiplication (GEMM), namely inner product (GEDOT), outer product (GER), and matrix-vector multiplication (GEMV). Write computation templates for these four cases respectively.

[0126] According to the existing BLAS library implementation and optimization methods, write the computation template code for the above linear algebra operations. However, each time the matrix dimensions are referenced, they are written as placeholders for the four implicit dimensions ph_P, ph_Y, ph_X, ph_R; each time any element of the matrix is accessed, it is written as placeholders for the three memory access functions ph_A(p, y, k), ph_B(p, k, x), and ph_C(p, y, x).

[0127] 5 According to the symbolic representation of tensor contraction, perform template dispatch during compilation to generate CUDA C / C++ code

[0128] In the previous step, tensor contraction has been implicitly represented as batch matrix multiplication. Among them, the four dimension sequences P, Y, X, R in tensor contraction logically correspond to the four dimensions of batch matrix multiplication respectively. Therefore, when a certain dimension sequence in P, Y, X, R is empty, the corresponding dimension of batch matrix multiplication is also equal to 1, thus degenerating into a simpler linear algebra operation. We design specific computation templates for different degenerate cases to maximize computational efficiency.

[0129] The content of the four dimension sequences P, Y, X, R has been analyzed from the symbolic representation formula (9) of tensor contraction. According to whether they are empty, select a specific linear algebra computation template during compilation:

[0130]

[0131] After selecting a suitable template, substitute the placeholder into the calculation template to generate legal CUDA C / C++ code.

[0132] 6 For the code generated by instantiating the calculation template, generate a dynamic link library through JIT (Just-In-Time) compilation technology and load it into the address space of the user program. In this way, a reusable executable program is obtained.

[0133] 7 The user inputs the data of each tensor and the specific values of each dimension, and makes the executable program compiled in the previous step complete the calculation.

Claims

1. A template-based high-performance tensor contraction method for GPU, characterized in that, it includes the following steps: (1) The user inputs the definition of tensor contraction, classifies its indices and dimensions, and obtains four index sequences p, x, y, r and four dimension sequences P, Y, X, R; (1.1) Notation convention: For an n-dimensional tensor T with dimensions from highest to lowest being D = D n , …, D 1 , it is denoted as T D ; T[i] = T[i n , …, i 1 represents the scalar at the index sequence i = (i n , …, i 1 ); the function ext(i k ) = D k , k = 1, …, n, maps a single index to the corresponding dimension, and Ext(i) = (ext(i n ), …, ext(i 1 )) = (D n , …, D 1 ) = D maps the index sequence to the corresponding dimension sequence; any tensor contraction C = A × B input by the user is expressed in Einstein summation convention: C[ic] = A[ia] × B[ib] (1) where ic, ia, ib are the index sequences of tensors C, A, B respectively; (2) Reduce the dimensions of the index sequences p, x, y, r and the dimension sequences P, Y, X, R to obtain the memory access function and implicit dimensions; (3) Define placeholders to represent the content related to the memory access function and implicit dimensions in the implicit batch matrix multiplication; (4) Write a calculation template according to the BLAS library implementation and optimization method, and reserve the placeholders described in step (3); (5) Perform template dispatch during compilation, substitute the memory access function and implicit dimensions into the placeholders of the selected calculation template to generate CUDA C / C++ code; according to whether the content of the four dimension sequences P, Y, X, R is empty, select a specific linear algebra calculation template during compilation: After selecting the template, substitute the placeholder into the calculation template to generate legal CUDA C / C++ code; where GEDOT represents the inner product, GEMV represents the matrix-vector multiplication, GER represents the outer product, and GEMM represents the matrix multiplication; (6) Compile the code generated in step (5) into a reusable executable program; (7) Input the data of each tensor and the specific values of each dimension, and use the executable program compiled in step (6) to complete the calculation.

2. The template-based high-performance tensor contraction method for GPU according to claim 1, characterized in that, the step (1) further includes the following sub-steps: (1.2) Define the batch index sequence: Define p = ic ∩ ia ∩ ib, that is, the sequence composed of all batch indices, and each element is used to index all tensors; the order of the elements in p in ic, ia, ib can be the same or different, and any order is adopted here, and the same applies to the following sequences; (1.3) Define the free index sequences y, x of the input tensors A, B respectively; Define y = (ic ∩ ia)\p, that is, the sequence composed of all free indices of A, which is only used to index C and A; Define x = (ic ∩ ib)\p as the sequence composed of all free indices of B, which is only used to index C and B; (1.4) Define the reduction index sequence: Define r = (ia ∩ ib)\p as the reduction index sequence, and the elements of tensors A and B on Ext(r) are multiplied pairwise and reduced to a scalar; (1.5) Define the permutation functions: By equivalently transforming Equation (1) and substituting \(p\), \(y\), \(x\), \(r\) for \(ic\), \(ia\), \(ib\), three permutation functions \(shfl_{ C}(·)\), \(shfl_{ A}(·)\) and \(shfl_{ B}(·)\) are introduced, and thus the original formula is rewritten as: C (·)、shfl A (·) and shfl B (·), thereby rewriting the original formula as: C[shfl C (p, y, x)] = A[shfl A (p, y, r)] × B[shfl B (p, r, x)] (2) (1.6) Define the dimension sequences corresponding to each index sequence: Define the batch dimension sequence P = Ext(p), where each element is the dimension corresponding to each index in the batch index sequence; Define the free dimension sequence of A as Y = Ext(y), the free dimension sequence of B as X = Ext(x), and the reduction dimension sequence as R = Ext(r).

3. The template-based high-performance tensor contraction method for GPU according to claim 2, characterized in that, the step (2) includes the following sub-steps: (2.1) Define the dimensionality reduction auxiliary function: First, define the mapping i = gath d (i) from the index sequence to the scalar index, where d is an arbitrary dimensionality sequence but must be the same length as the index sequence i; this function maps the offset i on each dimension in d to the linear address i: Define the inverse function of the above function i = scat d (i) = gath d -1 (i); it re - decomposes the linear address i into the offsets i on each dimension in d: With the help of the above two auxiliary functions, the high-order tensors in Equation (2) are implicitly transformed into matrices in a specific storage format; (2.2) Reduce the dimensions of the indices, so as to access the elements in the tensor with scalar indices: Convert the index sequences p, y, x, r into the components of the scalar indices in the corresponding dimensions; define p = gath P (p), where p is a scalar index with a value range of [0, ∏P), and represent the index sequence in reverse as p = gath P -1 (p) = scat P (p); represent the other index sequences as y = scat Y (y), x = scat X (x), r = scat R (r), where the value ranges of the newly defined scalar indices are y ∈ [0, ∏Y), x ∈ [0, ∏X), r ∈ [0, ∏R); rewrite Equation (2) as: (2.3) Reduce the dimensions of the dimensions, and represent tensors with any order and any dimension arrangement as special matrices: The actual calculation that occurs in Equation (7) is: Denote the access to the output tensor on the left side of Equation (8) as C(p, y, x) = C[shfl C (scat P (p), scat Y (y), scat X (x))], denote the access to the two input tensors on the right side as A(p, y, r) = A[shfl A (scat P (p), scat Y (y), scat R (r))] and B(p, r, x) = B[shfl B (scat P (p), scat R (r), scat X (x))], and rewrite Equation (8) as: which only contains 4 scalar dimensions, that is, P = ΠP, Y = ∏Y, X = ∏X, R = ∏R; (2.4) Implicitly transform the tensor contraction into matrix multiplication: Let C(p, y, x), A(p, x, r), and B(p, r, x) be the memory access functions of the output and input tensors. They abstract the tensors into batch matrices with a specific data layout; Let the four scalars P, Y, X, and R be the implicit dimensions, which are the batch dimension, the free dimension of A, the free dimension of B, and the product of the reduction dimensions respectively; Equation (9) hides the data reading and writing process with the memory access function and expresses the range of traversal and iteration with the implicit dimensions, thus transforming the tensor contraction into batch matrix multiplication; With the help of Einstein summation convention, Equation (9) is expressed as: C(p, y, x) = A(p, y, r) × B(p, r, x) (10).

4. The template-based GPU high-performance tensor contraction method according to claim 3, characterized in that the step (3) includes the following sub-steps: (3.1) Define implicit dimension placeholders: Equation (9) contains four implicit dimensions P, Y, X, R, which represent the scales of the implicit batch matrices and are the products of the corresponding dimension sequences in the original tensors respectively; Therefore, define four corresponding implicit dimension placeholders ph_P, ph_Y, ph_X, ph_R, corresponding to the above implicit dimensions respectively; (3.2) Define memory access function placeholders: In Equation (10), the reading and writing of the implicit matrix are abstracted by three memory access functions A(p, y, r), B(p, r, x), and C(p, y, x) respectively; Define three memory access function placeholders ph_A(p, y, k), ph_B(p, k, x), and ph_C(p, y, x) to represent the elements in the corresponding positions.

5. The template-based GPU high-performance tensor contraction method according to claim 4, characterized in that the step (4) is specifically: Write the calculation template code for linear algebra operations according to the BLAS library implementation and optimization method; Each time the matrix dimension is referenced, write it as the placeholders of the four implicit dimensions ph_P, ph_Y, ph_X, ph_R; Each time any element of the matrix is accessed, write it as the placeholders of the three memory access functions ph_A(p, y, k), ph_B(p, k, x), and ph_C(p, y, x).

6. The template-based GPU high-performance tensor contraction method according to claim 5, characterized in that the step (5) further includes: Analyze and obtain the content of the four dimension sequences of P, Y, X, and R through Equation (9).

Images