Method and device for implementing matrix operation on tpu, equipment and medium

Abstract

The application provides an implementation method, device and equipment of matrix operation on TPU and a medium, and belongs to the technical field of electric digital data processing. The method comprises the following steps: dividing input matrices according to the number of tensor calculation cores in the TPU to obtain to-be-processed sub-matrices corresponding to each tensor calculation core; obtaining the matrix dimension and vector memory capacity of the to-be-processed sub-matrices; matching a target data processing template from a plurality of pre-compiled data processing templates according to the matrix dimension and the vector memory capacity; loading corresponding matrix data in the to-be-processed sub-matrices into the vector memory according to the target data processing template, so that the tensor calculation core performs matrix operation on the matrix data to obtain sub-matrix operation results; and splicing the sub-matrix operation results to obtain a final matrix operation result. The application can exert the performance advantage of the TPU and realize efficient matrix operation on the TPU.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of digital data processing, and in particular to a method and device for implementing matrix operation on TPU, equipment and medium. BACKGROUND

[0002] In the field of computing, different computing architectures have their own advantages. As a traditional computing core, the central processing unit (CPU) has strong versatility, but when facing a large number of complex matrix operations in deep learning and high-performance computing, the performance bottleneck is obvious due to its serial processing characteristics, and the computing efficiency is relatively low. Although the graphics processing unit (GPU) has made breakthroughs in parallel computing, it is commonly used to accelerate graphics processing and part of the computing task, but its architecture is not specifically designed for matrix operations in deep learning. When processing large-scale and high-intensity matrix operations, the energy consumption is high and the utilization rate of computing resources is difficult to reach the optimum.

[0003] In contrast, the tensor processing unit (TPU) is born for matrix operations in deep learning and high-performance computing. The systolic array architecture adopted by TPU greatly improves the parallel processing capability of matrix operations, and can complete a large number of matrix operations in one clock cycle, with a computing efficiency far exceeding CPU and GPU. At the same time, TPU is highly adaptable to low-precision computing. When using low-precision data types such as 16-bit floating point (Brain Floating point 16-bit, BF16) and 8-bit signed integer (8-bit Integer, INT8), it not only guarantees the accuracy requirement of deep learning tasks, but also significantly speeds up the computing speed and reduces the energy consumption, showing excellent energy consumption ratio. In addition, the hardware design of TPU is optimized for deep learning algorithms, with short data transmission distance inside the chip, reducing transmission energy consumption, and much lower energy consumption than CPU and GPU when processing the same computing task. Moreover, the architecture design of TPU is simple and efficient, with lower cost, and obvious cost-effective advantage in large-scale deployment of deep learning applications.

[0004] In the field of GPU, template technology has been widely used in the optimization of matrix multiplication. By pre-compiling kernel code of different block sizes, specific optimization code is generated for different input sizes. This method can complete code optimization for specific data volume at the compilation stage, thereby reducing the overhead of dynamic adjustment at runtime.

[0005] The related art mostly applies the template technology in the GPU field when implementing matrix operation on the TPU. However, the performance structure of the TPU is significantly different from that of the GPU, and directly transplanting the template scheme of the GPU cannot fully exert the performance advantage of the TPU, and there is an efficiency bottleneck. SUMMARY

[0006] Embodiments of the present application provide a method, device, equipment and medium for implementing matrix operation on a TPU, to solve the problem that the existing template scheme cannot fully exert the performance advantage of the TPU and there is an efficiency bottleneck.

[0007] In a first aspect, embodiments of the present application provide a method for implementing matrix operation on a TPU, comprising:

[0008] According to the number of tensor calculation cores in the TPU, the input matrix in the high-bandwidth memory is divided to obtain a to-be-processed sub-matrix corresponding to each tensor calculation core;

[0009] For each tensor calculation core, the matrix dimension of the to-be-processed sub-matrix corresponding to the tensor calculation core and the vector memory capacity corresponding to the tensor calculation core are obtained;

[0010] According to the matrix dimension and the vector memory capacity, at least one target data processing template is matched from a plurality of pre-compiled data processing templates;

[0011] According to the at least one target data processing template, the corresponding matrix data in the to-be-processed sub-matrix is loaded into the vector memory corresponding to the tensor calculation core, so that the tensor calculation core performs matrix operation on the matrix data to obtain a sub-matrix operation result;

[0012] The sub-matrix operation results of each tensor calculation core are stored in the high-bandwidth memory, and each sub-matrix operation result is spliced to obtain a final matrix operation result.

[0013] In a possible implementation manner, according to the matrix dimension and the vector memory capacity, at least one target data processing template is matched from a plurality of pre-compiled data processing templates, comprising:

[0014] According to the matrix dimension, at least one optimal data processing template matching the matrix dimension is determined from a plurality of pre-compiled data processing templates;

[0015] For each optimal data processing template, according to the size relationship between the data amount corresponding to the optimal data processing template and the vector memory capacity, a target data processing template corresponding to the optimal data processing template is determined.

[0016] In a possible implementation, the determining, according to the matrix dimension, at least one optimal data processing template matching the matrix dimension from a plurality of precompiled data processing templates, comprises the following steps.

[0017] Obtaining template sizes of the plurality of precompiled data processing templates;

[0018] If there is a data processing template with a template size greater than or equal to the matrix dimension in the plurality of precompiled data processing templates, selecting, as the optimal data processing template, a data processing template with the smallest template size from the data processing templates with the template size greater than or equal to the matrix dimension;

[0019] If the template sizes of the plurality of precompiled data processing templates are all smaller than the matrix dimension, selecting, as the optimal data processing template, a plurality of data processing templates with the template sizes greater than or equal to the matrix dimension after splicing and with the smallest template sizes after splicing.

[0020] In a possible implementation, the determining, according to a size relationship between the data amount corresponding to the optimal data processing template and the vector memory capacity, of a target data processing template corresponding to the optimal data processing template, comprises the following steps.

[0021] Detecting whether the data amount corresponding to the optimal data processing template is greater than the vector memory capacity;

[0022] If the data amount corresponding to the optimal data processing template is greater than the vector memory capacity, selecting, as the target data processing template, a plurality of suboptimal data processing templates with data amounts smaller than the vector memory capacity from the plurality of precompiled data processing templates; wherein each suboptimal data processing template has a template size greater than or equal to a template size of the optimal data processing template after splicing and has the smallest template size after splicing;

[0023] If the data amount corresponding to the optimal data processing template is smaller than or equal to the vector memory capacity, determining the optimal data processing template as the target data processing template.

[0024] In a possible implementation, intermediate operation results in a process of performing matrix operation on matrix data by the tensor calculation core are all stored in the corresponding vector memory.

[0025] In a possible implementation, the input matrix in the high-bandwidth memory and the matrix data in the vector memory are both stored according to a 512-byte boundary.

[0026] In a possible implementation, in the process of performing the matrix operation on the matrix data by the tensor calculation core, the matrix data is batch processed by using a tensor instruction of the TPU.

[0027] In a second aspect, an embodiment of the present application provides an implementation device for matrix operation on a TPU, including:

[0028] The division module is configured to divide the input matrix in the high-bandwidth memory according to the number of the tensor calculation cores in the TPU, to obtain a to-be-processed sub-matrix corresponding to each tensor calculation core;

[0029] The matching module is configured to:

[0030] For each tensor calculation core, the matrix dimension of the to-be-processed sub-matrix corresponding to the tensor calculation core and the vector memory capacity corresponding to the tensor calculation core are obtained;

[0031] At least one target data processing template is matched from a plurality of pre-compiled data processing templates according to the matrix dimension and the vector memory capacity;

[0032] The calculation module is configured to:

[0033] According to the at least one target data processing template, the corresponding matrix data in the to-be-processed sub-matrix is loaded into the vector memory corresponding to the tensor calculation core, so that the tensor calculation core performs the matrix operation on the matrix data to obtain a sub-matrix operation result;

[0034] The sub-matrix operation results of the tensor calculation cores are all stored in the high-bandwidth memory, and the sub-matrix operation results are spliced to obtain a final matrix operation result.

[0035] In a third aspect, an embodiment of the present application provides an electronic device including a memory and a processor, the memory stores a computer program, and the processor implements the method in the first aspect or any possible implementation manner of the first aspect when executing the computer program.

[0036] In a fourth aspect, an embodiment of the present application provides a computer readable storage medium, which stores a computer program, and the computer program is executed by a processor to implement the method in the first aspect or any possible implementation manner of the first aspect.

[0037] In a fifth aspect, an embodiment of the present application provides a computer program product, including a computer program, and the computer program is executed by a processor to implement the method in the first aspect or any possible implementation manner of the first aspect.

[0038] The embodiment of the present application dynamically matches a target data processing template from a plurality of pre-compiled data processing templates from the aspects of matrix dimension and vector memory capacity, loads data from a high bandwidth memory to a vector memory according to the target data processing template, and performs matrix operation. This way can realize deep adaptation of the target data processing template and the memory level of the TPU hardware, can not only give full play to the advantages of the vector memory unique to the TPU, reduce the access times to the high-delay high-bandwidth memory, reduce the memory access delay, comprehensively improve the overall computing performance of the system, but also increase the flexibility of matrix operation on the TPU. In addition, the embodiment of the present application splits the input matrix in the high-bandwidth memory according to the hardware characteristics of the TPU, so that each tensor calculation core can perform parallel and collaborative calculation, which can reduce the calculation pressure of each tensor calculation core and improve the calculation efficiency. BRIEF DESCRIPTION OF DRAWINGS

[0039] Figure 1 is a hardware architecture diagram of the TPU provided by the embodiment of the present application;

[0040] Figure 2 is an implementation flowchart of the implementation method of the matrix operation on the TPU provided by the embodiment of the present application;

[0041] Figure 3 is a division schematic diagram of the input matrix provided by the embodiment of the present application;

[0042] Figure 4 is a data filling schematic diagram when the number of target data processing templates is 1 provided by the embodiment of the present application;

[0043] Figure 5 is a data division schematic diagram when the number of target data processing templates is 2 provided by the embodiment of the present application;

[0044] Figure 6 is a structure schematic diagram of the implementation method and device of the matrix operation on the TPU provided by the embodiment of the present application;

[0045] Figure 7 is a schematic diagram of an electronic device provided by the embodiment of the present application. DETAILED DESCRIPTION

[0046] The embodiment of the present application will be described in detail below with reference to the accompanying drawings.

[0047] Template technology is an optimization method by predefining code structure (such as module size, data type) and dynamically adapting input parameters. The core goal is to reduce the overhead of dynamic adjustment during matrix operation and improve computing efficiency. In related technologies, the template technology applied to the GPU is applied to the TPU to implement matrix operation. However, the hardware architecture (such as vector instruction set, memory hierarchy) of the TPU is significantly different from that of the GPU. Directly applying the template scheme of the GPU on the TPU cannot fully exert the performance advantage of the TPU, and there is an efficiency bottleneck.

[0048] In order to implement efficient matrix operation on the TPU and fully exert the performance advantage of the TPU, the embodiment of the present application considers matrix dimension and vector memory capacity from two aspects to dynamically match the target data processing template from a plurality of pre-compiled data processing templates, and loads data from the high bandwidth memory to the vector memory according to the target data processing template to execute matrix operation. This way can realize the deep adaptation of the target data processing template and the TPU hardware memory hierarchy, can fully exert the advantage of the vector memory unique to the TPU, reduce the access times to the high latency high bandwidth memory, reduce the memory access delay, and comprehensively improve the overall computing performance of the system.

[0049] First, the hardware architecture of the TPU is briefly introduced. Referring to Figure 1 , the TPU includes at least two tensor computing cores that can be used to perform matrix operation work. Each tensor computing core is equipped with a corresponding vector memory (Vector Memory, VMEM), and each tensor computing core shares the same high bandwidth memory (High Bandwidth Memory, HBM).

[0050] Considering that matrix multiplication in matrix operation is a core operation in deep learning and high-performance computing, the embodiment of the present application takes matrix multiplication as an example and introduces the implementation method of the matrix operation on the TPU provided by the embodiment of the present application based on the above TPU architecture. Referring to Figure 2 , which shows the implementation flowchart of the implementation method of the matrix operation on the TPU provided by the embodiment of the present application, which is described in detail as follows:

[0051] Step 201, according to the number of tensor computing cores in the TPU, the input matrix in the high bandwidth memory is divided to obtain the to-be-processed sub-matrix corresponding to each tensor computing core.

[0052] Each tensor computing core in the TPU can be used to perform matrix multiplication. To fully utilize the advantages of the TPU tensor computing core, the input matrix in the HBM is divided according to the number of the TPU internal tensor computing core, so that each tensor computing core can process the matrix operation of the to-be-processed sub-matrix in parallel, and the single-core computing pressure is reduced.

[0053] Here, the input matrix in the HBM is the matrix used to perform matrix multiplication in the TPU. It can be understood that the number of matrices participating in matrix multiplication is 2, and accordingly, the number of input matrices in the HBM is 2.

[0054] To facilitate parallel computing of each tensor computing core and avoid data interaction of each tensor computing core in the matrix operation process, the present embodiment can divide the right matrix in the input matrix according to the principle of “splitting the column number of the right matrix in the matrix multiplication according to the number of the tensor computing core”, to obtain a plurality of sub-right matrices, and determine the left matrix and each divided sub-right matrix as the to-be-processed sub-matrix corresponding to each tensor computing core.

[0055] Exemplarily, referring to Figure 3 , when the number of tensor computing cores is 2, the input matrix is an 8x128 left matrix A and a 128x128 right matrix B, and the product of the left matrix A and the right matrix B is calculated, the columns of the right matrix B are bisected, and the right matrix B is split into two 128x64 sub-right matrices B1 and B2, the left matrix A and the divided sub-right matrix B1 are taken as a to-be-processed sub-matrix corresponding to one tensor computing core, and the left matrix A and the divided sub-right matrix B2 are taken as a to-be-processed sub-matrix corresponding to another tensor computing core.

[0056] In step 202, for each tensor computing core, the matrix dimension of the to-be-processed sub-matrix corresponding to the tensor computing core and the vector memory capacity corresponding to the tensor computing core are obtained.

[0057] In the present embodiment, the matrix dimension (ah, aw, bw), data type and storage layout of the to-be-processed sub-matrix can be obtained through tensor metadata. Wherein, ah represents the number of rows of the left matrix, aw represents the number of columns of the left matrix, that is, the number of rows of the sub-right matrix, and bw represents the number of columns of the sub-right matrix.

[0058] Here, the data type includes but is not limited to: 32-bit floating point number (Float Point 32, FP32), 16-bit floating point number (Brain Floating point 16-bit, BF16) or 8-bit signed integer (8-bit Integer, INT8), etc. The storage layout includes: row-major or column-major.

[0059] In step 203, at least one target data processing template is matched from a plurality of pre-compiled data processing templates according to the matrix dimension and the vector memory capacity.

[0060] In the embodiment of the present application, a plurality of data processing templates of different template sizes can be pre-compiled to cover common matrix dimensions. For example, 128x256x128, 256x256x256, 512x512x256, etc.

[0061] In the embodiment of the present application, at least one target data processing template can be matched from a plurality of pre-compiled data processing templates according to the matrix dimension, which is used to perform matrix multiplication. Meanwhile, the embodiment of the present application also takes into account the VMEM advantage of TPU, and can store matrix data in VMEM when performing matrix operation on each tensor calculation core, so that the tensor calculation core can load matrix data in real time. Therefore, when matching the target data processing template, the embodiment of the present application also considers the size of the vector memory capacity, so as to avoid the situation that the VMEM cannot load the matrix data in the target data processing template at one time and needs to access the HBM in the middle.

[0062] In the embodiment of the present application, by comprehensively considering the matrix dimension and the vector memory capacity, a target data processing template suitable for the input matrix and TPU can be dynamically matched from a plurality of pre-compiled data processing templates.

[0063] Here, it should be noted that all template sizes are integer multiples of 128, which are aligned with the 512-byte vector instruction width of TPU, so as to dynamically adapt to the TPU instruction set.

[0064] In the related art, when the template technology of GPU is applied, only fixed template sizes are supported, and multiple codes need to be pre-generated for different input matrices, which leads to code inflation. Moreover, there is a difference between the GPU instruction set and the 512-byte vector instruction of TPU, and TPU needs to simulate the GPU instruction through complex logic when applying the template technology of GPU, which increases the delay calculation.

[0065] On the one hand, the embodiment of the present application dynamically matches the target data processing template from a plurality of data processing templates by comprehensively considering the matrix dimension and the vector memory capacity, which does not cause code inflation and can increase the flexibility of template matching. On the other hand, the data processing template in the embodiment of the present application can dynamically adapt to the TPU instruction set, without the need to simulate the GPU instruction, which can effectively reduce the calculation delay.

[0066] In addition, different data types and storage layouts are also encapsulated in the data processing template, which is used to adapt to the data type and storage layout of the input matrix, without the need for the developer to manually process the data type conversion and layout adaptation, thereby reducing the development complexity of cross-operator cooperation and improving the programming efficiency of the matrix multiplication module in the framework.

[0067] The embodiment of the present application generates special kernel code at the compilation stage to adapt to specific matrix multiplication scenarios and other possible computing requirements for different data types. The code ensures good flexibility and high performance under different precision requirements, meeting the needs of diversified application scenarios.

[0068] The data processing template in the embodiment of the present application can be compatible with different storage layouts of row priority and column priority by generating loading instruction sequences of different storage formats at the compilation stage, and adapt to the efficient reading of 512-byte aligned data by the direct memory access (DMA) transfer unit of TPU.

[0069] The 512-byte vector instructions (such as load with offset, load with Vmask) and vector memory VMEM of TPU provide a hardware basis for matrix multiplication operations. The embodiment of the present application can fully utilize the hardware advantages of TPU by designing a data processing template that is adapted to the instruction set and VMEM capacity of TPU, and solve the computing bottleneck problem caused by applying GPU template technology.

[0070] In step 204, according to at least one target data processing template, the corresponding matrix data in the to-be-processed sub-matrix is loaded into the vector memory corresponding to the tensor calculation core, so that the tensor calculation core performs matrix operation on the matrix data to obtain the sub-matrix operation result.

[0071] After determining the target data processing template, the corresponding matrix data in the to-be-processed sub-matrix can be loaded into the vector memory according to the template size of the target data processing template.

[0072] For example, referring to Figure 4 When the number of target data processing templates is 1, the template size is 128*256*128, and the matrix dimension is 128*200*128, all the matrix data in the to-be-processed sub-matrix is loaded into the VMEM, and the VMEM is filled completely according to the template size. After filling, the data is strictly 512-byte aligned to adapt to the 512-byte vector instructions of TPU, and ensure that the vector instruction loads complete row data at a time.

[0073] For example, referring to Figure 5 When the number of target data processing templates is 2, the template size is 128*256*128, and the matrix dimension is 256*256*128, the matrix data in the to-be-processed sub-matrix is split according to the template size, and the split matrix data is loaded into the VMEM in sequence for matrix operation by the tensor calculation core.

[0074] According to the matrix dimension, data type and storage layout of the input matrix, the embodiments of the present application dynamically match the most suitable target data processing template from the pre-compiled multiple data processing templates. If the matrix dimension does not completely match the template size, padding or splicing strategies are adopted to ensure efficient adaptation and improve computing efficiency under the condition of meeting the kernel parameter range.

[0075] When the number of target data processing templates is 1, the tensor computing core directly performs matrix operation on the matrix data corresponding to the target data processing template, and obtains the sub-matrix operation result.

[0076] When the number of target data processing templates is greater than 1, the tensor computing core sequentially performs matrix operation on the matrix data corresponding to each target data processing template to obtain intermediate operation results, and performs operation on the intermediate operation results again to obtain the final sub-matrix operation result. Here, the process of matrix operation is the conventional mathematical logic, which will not be introduced here.

[0077] In some embodiments, the intermediate operation results in the process of matrix operation of the tensor computing core on the matrix data are all stored in the corresponding vector memory.

[0078] Here, by storing the matrix data used to perform matrix operation and the intermediate operation results in the process of matrix operation in the vector memory for real-time access and calling by the tensor computing core, the advantages of vector memory cache in TPU can be fully utilized, the number of accesses to high-delay high-bandwidth memory is maximally reduced, the memory access delay is effectively reduced, and the overall computing performance of the system is comprehensively improved.

[0079] In some embodiments, the input matrix in the high-bandwidth memory and the matrix data in the vector memory are all stored according to 512-byte boundary.

[0080] Here, ALIGN128, ALIGN256 and other macros can be used to align the matrix data or intermediate operation results to ensure that the data is stored according to 512-byte boundary, so as to match the loading unit of TPU vector instruction, for example, the macro instruction head_dim_aligned=ALIGN128(head_dim) can be used to make the vector instruction v_f32_ld_tnsr_st_msk efficiently loaded, and the additional processing of unaligned data is avoided.

[0081] In some embodiments, in the process of matrix operation of the tensor computing core on the matrix data, the tensor instruction of TPU is used to batch process the matrix data.

[0082] In the matrix multiplication operation, the tensor instructions (such as load with offset, m_matmul_single, etc.) using the TPU architecture can realize data processing optimization. These tensor instructions have the ability to process data in batches, unlike the scattered 32-bit floating-point operations in other architectures, the TPU tensor instructions can load multiple data elements at a time according to the 512-byte boundary (for example, the load with offset instruction can load 1024 FP32 elements at a time), relying on the parallel hardware design of the TPU vector processing unit, so that multiple data elements can be processed at a time. This batch processing mechanism greatly improves the parallelism of data processing, reduces the redundancy of instruction execution, significantly speeds up the operation process compared to the traditional operation mode, and fully releases the high-performance potential of TPU in the matrix multiplication scenario.

[0083] The embodiment of the application fully utilizes the 512-byte vector instruction of TPU to process matrix data, and more data can be processed by a single instruction, thereby improving parallel efficiency.

[0084] In addition, considering the loop iteration of multiple judgments and jumps in matrix operations, a judgment delay is generated, and therefore, the embodiment of the application can also use the #pragma unroll instruction in TPU to optimize the loop structure, and convert the loop iteration of multiple judgments and jumps into a continuous sequential instruction sequence.

[0085] Exemplarily, assuming that a loop needs to perform 10 operations, and each iteration needs to judge "whether it is completed" and jump, which will cause the processor pipeline to frequently stop. Through instruction pipeline optimization, the 10 iterations can be unfolded into 8 continuous operations (dynamically adjusted according to hardware capability), and the intermediate conditional judgment link is eliminated. The processor can therefore continuously receive the instruction stream, like a production line running at full speed without interruption, so that the different types of processing units of the processor are always in full load working state, and the continuity and parallelism of instruction execution are significantly improved.

[0086] In step 205, the sub-matrix operation results of each tensor calculation core are stored in the high-bandwidth memory, and the sub-matrix operation results are spliced to obtain the final matrix operation result.

[0087] Still taking Figure 3 for example, the tensor calculation core calculates the left matrix A and the sub-right matrix B1 to obtain the sub-matrix operation result 8x64; another tensor calculation core calculates the matrix A and the sub-right matrix B2 to obtain the sub-matrix operation result 8x64. The two sub-matrix operation results are spliced to obtain the final matrix operation result 8x128.

[0088] Compared with the prior art, the embodiment of the application considers both the matrix dimension and the vector memory capacity to dynamically match a target data processing template from a plurality of pre-compiled data processing templates, and loads data from the high-bandwidth memory to the vector memory according to the target data processing template to perform matrix operation. In this way, the depth adaptation of the target data processing template and the memory hierarchy of the TPU hardware can be achieved, not only the advantages of the vector memory unique to the TPU can be fully utilized, the number of accesses to the high-latency high-bandwidth memory can be reduced, the memory access delay can be reduced, and the overall computing performance of the system can be improved, but also the flexibility of matrix operation on the TPU can be increased. In addition, according to the hardware characteristics of the TPU, the input matrix in the high-bandwidth memory is split to enable parallel and collaborative computation of each tensor calculation core, so that the computing pressure of each tensor calculation core can be reduced and the computing efficiency can be improved.

[0089] The specific implementation of matching and determining the target data processing template from the plurality of pre-compiled data processing templates will be described below.

[0090] In some embodiments, at least one optimal data processing template matching the matrix dimension can be determined from the plurality of pre-compiled data processing templates according to the matrix dimension first; then, for each optimal data processing template, a target data processing template corresponding to the optimal data processing template can be determined according to the size relationship between the data amount corresponding to the optimal data processing template and the vector memory capacity.

[0091] For each dimension (ah / aw / bw), the minimum template size not less than the dimension is selected (for example, when ah=100, 128 is selected, and when aw=200, 256 is selected), and a data processing template that simultaneously satisfies the boundaries of all dimensions is preferentially matched (for example, for a matrix dimension of 100x200x200, a 128x256x256 template size is matched).

[0092] Specifically, the template sizes of the plurality of pre-compiled data processing templates are obtained; if there is a data processing template with a template size greater than or equal to the matrix dimension in the plurality of pre-compiled data processing templates, a data processing template with the smallest template size is selected from the data processing templates with the template size greater than or equal to the matrix dimension as the optimal data processing template; if the template sizes of the plurality of pre-compiled data processing templates are all less than the matrix dimension, a plurality of data processing templates with the template size greater than or equal to the matrix dimension after splicing and with the smallest template size after splicing are selected as the optimal data processing template.

[0093] When the data processing templates with the template size greater than or equal to the matrix dimension exist, in order to ensure that the matrix data can be loaded by using one data processing template at one time, reduce the data filling overhead, and avoid wasting the computing resources, an embodiment of the present application selects the data processing template with the smallest template size from the data processing templates with the template size greater than or equal to the matrix dimension as the optimal data processing template.

[0094] When the template sizes of the plurality of pre-compiled templates are all less than the matrix dimension, a plurality of data processing templates are selected from the plurality of pre-compiled data processing templates by using the template splicing manner, the template size of the plurality of data processing templates after splicing is greater than or equal to the matrix dimension, and the template size after splicing is the smallest, so that the matrix data can be completely loaded while reducing the data splicing overhead and occupying the least computing resources.

[0095] On the basis of the above-mentioned consideration of the matrix dimension for determining the optimal data processing template, an embodiment of the present application further considers the vector memory capacity, determines the target data processing template, divides the matrix data according to the target data processing template, and performs the subsequent matrix operation work.

[0096] Specifically, for each optimal data processing template, it is detected whether the data amount corresponding to the optimal data processing template is greater than the vector memory capacity; if the data amount corresponding to the optimal data processing template is less than or equal to the vector memory capacity, the optimal data processing template is determined as the target data processing template.

[0097] When the data amount corresponding to the optimal data processing template is less than or equal to the vector memory capacity, the TPU can load the corresponding matrix data into the vector memory at one time according to the optimal data processing template, so as to be accessed and called by the tensor calculation core at any time, at this time, the optimal data processing template can be directly determined as the target data processing template.

[0098] However, when the data amount corresponding to the optimal data processing template is greater than the vector memory capacity, the vector memory cannot completely load all the matrix data corresponding to the optimal data processing template, which causes the tensor calculation core to access the high-delay HBM to obtain the matrix data not loaded into the VMEM when performing the matrix operation according to the optimal data processing template, and further causes the calculation access delay problem.

[0099] In order to avoid the above-mentioned calculation access delay problem, the target data processing template needs to be further determined according to the optimal data processing template, the matrix data needs to be divided according to the target data processing template, and the subsequent matrix operation work needs to be performed.

[0100] Specifically, if the data amount corresponding to the optimal data processing template is greater than the vector memory capacity, a plurality of suboptimal data processing templates with a data amount less than the vector memory capacity are selected from a plurality of precompiled data processing templates, and the plurality of suboptimal data processing templates are determined as the target data processing templates. The template size after splicing of each suboptimal data processing template is greater than or equal to the template size of the optimal data processing template, and the template size after splicing is the smallest.

[0101] If the data amount of the optimal data processing template is greater than the vector memory capacity, a plurality of suboptimal data processing templates are reselected from a plurality of precompiled data processing templates for splicing to replace the optimal data processing template. Here, the data amount of each suboptimal data processing template is less than the vector memory capacity, and the template size after splicing of each suboptimal data processing template is greater than or equal to the minimum template size of the optimal data processing template.

[0102] Exemplarily, if the VMEM capacity is 16M and the data type of the matrix data is FP32 (4 bytes / element). Here, to ensure the storage performance of the VMEM, its available capacity can be reserved by 75%, and the available capacity is: 16MBx75%=12MB=12582912 bytes. The maximum data amount of the suboptimal data processing template is: 12582912÷4=3145728 elements. That is, the template size of the suboptimal data processing template satisfies xxyxz≤3145728. Here, x, y, and z respectively represent the template size of each dimension.

[0103] The embodiment of the application can dynamically select a target data processing template for calculation according to the matrix dimension of the to-be-processed submatrix, so as to reduce the data padding and splicing overhead caused by the mismatch of the template size and improve the calculation efficiency.

[0104] To verify the implementation method of the matrix operation on the TPU provided by the embodiment of the application, the embodiment of the application respectively performs statistical testing and performance testing.

[0105] Statistical testing:

[0106] The matrix multiplication results of the TPU and the GPU are compared to ensure numerical consistency. The matrix multiplication result calculated by the TPU is compared with the result calculated by the GPU. The GPU is widely used in matrix calculation, and its calculation result has high reliability. When a large number of different scale matrix multiplication tests are performed, the floating point error of the TPU and the GPU result is accurately calculated. The floating point error qualified standard is set, and after multiple rounds of testing, the TPU calculation result meets the standard. The matrix multiplication implemented on the TPU of the application still has high numerical accuracy and can be reliably applied to various computing tasks.

[0107] Performance testing:

[0108] The efficiency of the TPU templating implementation and the general GPU scheme is compared under a typical template size (1024x4096x4096). Under the same calculation task and test environment, the matrix multiplication program is run multiple times, and the number of calculation loop structures is recorded. The results show that, by using the implementation method of the matrix operation on the TPU provided in the embodiment of the application, the calculation time of the TPU templating implementation is still significantly shortened compared with the general GPU scheme, and the efficiency is improved by more than 20%. This fully proves that the application can effectively improve the calculation performance of the matrix multiplication on the TPU, and has good practicability and superiority.

[0109] It should be understood that the size of the serial number of each step in the above embodiment does not mean the order of execution, and the execution order of each process should be determined according to its function and inherent logic, and should not constitute any limitation on the implementation process of the embodiment of the application.

[0110] The following is a device embodiment of the application, and for details not described in detail, reference can be made to the corresponding method embodiments described above.

[0111] Figure 6 The structure schematic diagram of the implementation device of the matrix operation on the TPU provided in the embodiment of the application is shown, only the part related to the embodiment of the application is shown for the convenience of description, and the details are as follows:

[0112] As shown in Figure 6 The implementation device 6 of the matrix operation on the TPU includes a division module 61, a matching module 62 and a calculation module 63.

[0113] The division module 61 is configured to divide the input matrix in the high bandwidth memory according to the number of tensor calculation cores in the TPU, to obtain a to-be-processed sub-matrix corresponding to each tensor calculation core;

[0114] The matching module 62 is configured to:

[0115] For each tensor calculation core, the matrix dimension of the to-be-processed sub-matrix corresponding to the tensor calculation core and the vector memory capacity corresponding to the tensor calculation core are obtained;

[0116] At least one target data processing template is matched from a plurality of pre-compiled data processing templates according to the matrix dimension and the vector memory capacity;

[0117] The calculation module 63 is configured to:

[0118] According to the at least one target data processing template, the corresponding matrix data in the to-be-processed sub-matrix is loaded into the vector memory corresponding to the tensor calculation core, so that the tensor calculation core performs matrix operation on the matrix data to obtain a sub-matrix operation result;

[0119] The sub-matrix operation results of each tensor calculation core are stored in the high-bandwidth memory, and the sub-matrix operation results are spliced to obtain a final matrix operation result.

[0120] In a possible implementation, the matching module 62 is specifically configured to:

[0121] According to the matrix dimension, at least one optimal data processing template matching the matrix dimension is determined from a plurality of pre-compiled data processing templates;

[0122] For each optimal data processing template, a target data processing template corresponding to the optimal data processing template is determined according to the size relationship between the data amount corresponding to the optimal data processing template and the vector memory capacity.

[0123] In a possible implementation, the matching module 62 is specifically configured to:

[0124] Template sizes of the plurality of pre-compiled data processing templates are obtained;

[0125] If there is a data processing template with a template size greater than or equal to the matrix dimension in the plurality of pre-compiled data processing templates, a data processing template with the smallest template size is selected from the data processing templates with the template size greater than or equal to the matrix dimension as the optimal data processing template;

[0126] If the template sizes of the plurality of pre-compiled data processing templates are all smaller than the matrix dimension, a plurality of data processing templates with the spliced template sizes greater than or equal to the matrix dimension and the smallest spliced template sizes are selected as the optimal data processing templates.

[0127] In a possible implementation, the matching module 62 is specifically configured to:

[0128] It is detected whether the data amount corresponding to the optimal data processing template is greater than the vector memory capacity;

[0129] If the data amount corresponding to the optimal data processing template is greater than the vector memory capacity, a plurality of suboptimal data processing templates with the data amount smaller than the vector memory capacity are selected from the plurality of pre-compiled data processing templates, and the plurality of suboptimal data processing templates are determined as the target data processing templates; wherein the spliced template sizes of the suboptimal data processing templates are greater than or equal to the template size of the optimal data processing template, and the spliced template sizes are the smallest;

[0130] If the data amount corresponding to the optimal data processing template is smaller than or equal to the vector memory capacity, the optimal data processing template is determined as the target data processing template.

[0131] In a possible implementation, the intermediate operation results in the process of the matrix operation of the matrix data by the tensor calculation core are all stored into the corresponding vector memories.

[0132] In a possible implementation, the input matrix in the high-bandwidth memory and the matrix data in the vector memory are both stored according to a 512-byte boundary.

[0133] In a possible implementation, in the process of the matrix operation of the matrix data by the tensor calculation core, the matrix data are processed in batches by using the tensor instruction of the TPU.

[0134] The device embodiment can be used to implement the method embodiments, and has the same technical principles and implementation effects as the method embodiments, which will not be repeated here.

[0135] Figure 7 is a schematic diagram of an electronic device provided by an embodiment of the present application. As shown in Figure 7 The electronic device 7 of this embodiment includes a processor 70 and a memory 71. The memory 71 stores a computer program 72. The processor 70 implements the steps in each of the method embodiments when executing the computer program 72. Alternatively, the processor 70 implements the functions of each module / unit in each of the device embodiments when executing the computer program 72.

[0136] For example, the computer program 72 can be divided into one or more modules / units, which are stored in the memory 71 and executed by the processor 70 to complete the present application. The one or more modules / units can be a series of computer program instruction segments that can complete a specific function, which are used to describe the execution process of the computer program 72 in the electronic device 7.

[0137] The electronic device 7 can include, but is not limited to, the processor 70 and the memory 71. Those skilled in the art can understand that Figure 7 The electronic device 7 is only an example and does not constitute a limitation on the electronic device 7, and can include more or fewer components than those shown, or combine certain components, or different components, for example, the electronic device 7 can also include an input / output device, a network access device, a bus, etc.

[0138] The processor 70 can be a central processing unit (CPU), and can also be other general-purpose processors, a digital signal processor (DSP), an application specific integrated circuit (ASIC), a field-programmable gate array (FPGA) or other programmable logic device, discrete gate or transistor logic, discrete hardware components, or the like. The general-purpose processor can be a microprocessor or the processor can also be any conventional processor.

[0139] The memory 71 can be an internal storage unit of the electronic device 7, such as a hard disk or a memory of the electronic device 7. The memory 71 can also be an external storage device of the electronic device 7, such as a plug-in hard disk, a smart media card (SMC), a secure digital (SD) card, a flash card, or the like equipped on the electronic device 7. Further, the memory 71 can also include both the internal storage unit and the external storage device of the electronic device 7. The memory 71 is used to store the computer program 72 and other programs and data required by the electronic device 7. The memory 71 can also be used to temporarily store data that has been output or will be output.

[0140] For the convenience and brevity of description, only the above-mentioned division of the functional modules / units is exemplified, and in actual application, the above-mentioned functions can be completed by different functional modules / units according to needs. The above-mentioned modules / units can be realized in the form of hardware, software, or a combination of hardware and software.

[0141] The embodiments of the present application also provide a computer readable storage medium storing a computer program. When the computer program is executed by a processor, the method in any of the above method embodiments is implemented.

[0142] The embodiments of the present application also provide a computer program product comprising a computer program. When the computer program is executed by a processor, the method in any of the above method embodiments is implemented.

[0143] The computer program includes computer program code, which can be in the form of source code, object code, executable files or some intermediate forms, etc. The computer readable medium can include any entity or device capable of carrying computer program code, recording medium, U disk, mobile hard disk, magnetic disk, optical disk, computer memory, read-only memory (ROM), random access memory (RAM), electrical carrier signal, telecommunication signal and software distribution medium, etc.

[0144] In the above embodiments, the description of each embodiment has its own focus, and the parts not described or recorded in a certain embodiment can be referred to the relevant description of other embodiments. If there is no special description and logical conflict, the terms and / or descriptions of different embodiments are consistent and can be mutually referred to, and the technical features in different embodiments can be combined to form new embodiments according to their inherent logical relationship.

[0145] The above-described embodiments are only used to illustrate the technical solutions of the present application, rather than limit them; although the present application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that the technical solutions recorded in the foregoing embodiments can still be modified, or some technical features can be replaced by equivalents; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of the embodiments of the present application, and should be included in the protection scope of the present application.

Claims

1. A method for implementing matrix operation on TPU, characterized in that, The method comprises the following steps: According to the number of tensor calculation cores in the TPU, the input matrix in the high-bandwidth memory is divided to obtain the to-be-processed sub-matrix corresponding to each tensor calculation core; wherein, the TPU comprises at least two tensor calculation cores, each tensor calculation core is equipped with a corresponding vector memory, and each tensor calculation core shares the same high-bandwidth memory; For each tensor calculation core, the matrix dimension of the to-be-processed sub-matrix corresponding to the tensor calculation core is obtained, and the vector memory capacity corresponding to the tensor calculation core is obtained; According to the matrix dimension and the vector memory capacity, at least one target data processing template is matched from a plurality of pre-compiled data processing templates; According to the at least one target data processing template, the corresponding matrix data in the to-be-processed sub-matrix is loaded into the vector memory corresponding to the tensor calculation core, so that the tensor calculation core performs matrix operation on the matrix data to obtain the sub-matrix operation result; the intermediate operation result in the process of the tensor calculation core performing matrix operation on the matrix data is stored in the corresponding vector memory; The sub-matrix operation results of each tensor calculation core are stored in the high-bandwidth memory, and each sub-matrix operation result is spliced to obtain the final matrix operation result.

2. The method for implementing matrix operation on TPU according to claim 1, characterized in that, According to the matrix dimension and the vector memory capacity, at least one target data processing template is matched from a plurality of pre-compiled data processing templates, which comprises: According to the matrix dimension, at least one optimal data processing template matching the matrix dimension is determined from a plurality of pre-compiled data processing templates; For each optimal data processing template, according to the size relationship between the data amount corresponding to the optimal data processing template and the vector memory capacity, the target data processing template corresponding to the optimal data processing template is determined.

3. The method for implementing matrix operation on TPU according to claim 2, characterized in that, According to the matrix dimension, at least one optimal data processing template matching the matrix dimension is determined from a plurality of pre-compiled data processing templates, which comprises: The template size of the plurality of pre-compiled data processing templates is obtained; If there is a data processing template with a template size greater than or equal to the matrix dimension in the plurality of pre-compiled data processing templates, the data processing template with the smallest template size is selected from the data processing templates with a template size greater than or equal to the matrix dimension as the optimal data processing template; If the template size of the plurality of pre-compiled data processing templates is smaller than the matrix dimension, a plurality of data processing templates with a spliced template size greater than or equal to the matrix dimension and the smallest spliced template size are selected as the optimal data processing template.

4. The method for implementing matrix operation on TPU according to claim 2, characterized in that, According to the size relationship between the data amount corresponding to the optimal data processing template and the vector memory capacity, the target data processing template corresponding to the optimal data processing template is determined, which comprises: It is detected whether the data amount corresponding to the optimal data processing template is greater than the vector memory capacity; If the data amount corresponding to the optimal data processing template is greater than the vector memory capacity, a plurality of suboptimal data processing templates with data amounts less than the vector memory capacity are selected from the plurality of precompiled data processing templates, and the plurality of suboptimal data processing templates are determined as the target data processing templates; wherein the template size after splicing of each suboptimal data processing template is greater than or equal to the template size of the optimal data processing template, and the template size after splicing is the smallest; If the data amount corresponding to the optimal data processing template is less than or equal to the vector memory capacity, the optimal data processing template is determined as the target data processing template.

5. The method for implementing matrix operation on TPU according to any one of claims 1-4, characterized in that, The input matrix in the high-bandwidth memory and the matrix data in the vector memory are stored according to a 512-byte boundary.

6. The method for implementing matrix operation on TPU according to any one of claims 1-4, characterized in that, In the process of performing matrix operation on the matrix data by the tensor calculation core, the tensor instruction of the TPU is used to batch process the matrix data.

7. An apparatus for implementing matrix operation on TPU, characterized in that, Comprise: The division module is configured to divide the input matrix in the high-bandwidth memory according to the number of tensor calculation cores in the TPU to obtain a to-be-processed sub-matrix corresponding to each tensor calculation core; wherein the TPU comprises at least two tensor calculation cores, each tensor calculation core is equipped with a corresponding vector memory, and each tensor calculation core shares the same high-bandwidth memory; The matching module is configured to: For each tensor calculation core, obtain the matrix dimension of the to-be-processed sub-matrix corresponding to the tensor calculation core and the vector memory capacity corresponding to the tensor calculation core; According to the matrix dimension and the vector memory capacity, at least one target data processing template is matched from a plurality of precompiled data processing templates; The calculation module is configured to: According to the at least one target data processing template, load the corresponding matrix data in the to-be-processed sub-matrix into the vector memory corresponding to the tensor calculation core, so that the tensor calculation core performs matrix operation on the matrix data to obtain a sub-matrix operation result; the intermediate operation result in the process of performing matrix operation on the matrix data by the tensor calculation core is stored in the corresponding vector memory; The sub-matrix operation results of each tensor calculation core are stored in the high-bandwidth memory, and each sub-matrix operation result is spliced to obtain a final matrix operation result.

8. An electronic device, comprising: The computer readable storage medium stores a computer program, and the processor executes the computer program to realize the method in any one of claims 1 to 6.

9. A computer-readable storage medium, characterized in that, The computer readable storage medium stores a computer program, and the processor executes the computer program to realize the method in any one of claims 1 to 6.

Images