Method, device and equipment for realizing matrix operation on TPU (Thermoplastic Polyurethane) and medium

Abstract

The invention provides a method, a device and equipment for realizing matrix operation on TPU (thermoplastic polyurethane) and a medium, and belongs to the technical field of electric digital data processing. The method comprises the following steps: dividing an input matrix according to the number of tensor calculation cores in TPU to obtain to-be-processed sub-matrixes corresponding to the tensor calculation cores; obtaining the matrix dimension of the to-be-processed sub-matrix and the capacity of a vector memory; according to the matrix dimension and the vector memory capacity, matching a target data processing template from the multiple pre-compiled data processing templates; according to the target data processing template, loading corresponding matrix data in the to-be-processed sub-matrix into a vector memory, so that a tensor calculation core performs matrix operation on the matrix data to obtain a sub-matrix operation result; and splicing the sub-matrix operation results to obtain a final matrix operation result. According to the method, the performance advantages of the TPU can be exerted, and efficient matrix operation on the TPU is achieved.

Description

Technical Field

[0001] The present invention relates to the technical field of electronic digital data processing, and in particular to a method, device, equipment and medium for implementing matrix operations on a TPU. Background Art

[0002] In the computing world, different computing architectures have their own strengths and weaknesses. The Central Processing Unit (CPU), as the traditional computing core, offers strong versatility. However, due to its serial processing nature, it faces significant performance bottlenecks and low computational efficiency when faced with the large number of complex matrix operations required in deep learning and high-performance computing. While the Graphics Processing Unit (GPU) has achieved breakthroughs in parallel computing and is commonly used to accelerate graphics processing and some computing tasks, its architecture is not specifically designed for the matrix operations required for deep learning. Consequently, it consumes high energy when handling large-scale, high-intensity matrix operations, and it struggles to achieve optimal computing resource utilization.

[0003] In contrast, the Tensor Processing Unit (TPU) is purpose-built for matrix operations in deep learning and high-performance computing. The TPU's systolic array architecture significantly enhances the parallel processing capabilities of matrix operations, enabling it to complete massive matrix operations within a single clock cycle, achieving computational efficiency far exceeding that of CPUs and GPUs. Furthermore, the TPU is highly adaptable to low-precision computations. When using low-precision data types such as 16-bit floating point (BF16) and 8-bit signed integers (INT8), it not only maintains the accuracy requirements of deep learning tasks but also significantly accelerates computation and reduces energy consumption, demonstrating an exceptional energy-efficiency ratio. Furthermore, the TPU's hardware design is optimized for deep learning algorithms. Data transmission within the chip is shortened, reducing transmission energy consumption. When processing the same computational task, the energy consumption is significantly lower than that of CPUs and GPUs. Furthermore, the TPU's simple and efficient architectural design results in lower costs, offering significant price-performance advantages for large-scale deployment of deep learning applications.

[0004] In the GPU field, stencil technology has been widely used to optimize matrix multiplication. By precompiling kernel code for different block sizes, optimized code is generated for different input sizes. This approach optimizes code for specific data volumes at compile time, reducing the overhead of dynamic adjustments at runtime.

[0005] When implementing matrix operations on TPUs, related technologies mostly apply template technology from the GPU field. However, the performance structure of TPUs is significantly different from that of GPUs. Directly transplanting the template solution to the GPU cannot fully utilize the performance advantages of TPUs and there is an efficiency bottleneck. Summary of the Invention

[0006] The embodiments of the present invention provide a method, apparatus, device and medium for implementing matrix operations on a TPU to solve the problem that existing template solutions cannot fully utilize the performance advantages of the TPU and have efficiency bottlenecks.

[0007] In a first aspect, an embodiment of the present invention provides a method for implementing matrix operations on a TPU, including: Divide the input matrix in the high-bandwidth memory according to the number of tensor computing cores in the TPU to obtain the sub-matrix to be processed corresponding to each tensor computing core; For each tensor computing core, obtain the matrix dimension of the sub-matrix to be processed corresponding to the tensor computing core and the vector memory capacity corresponding to the tensor computing core; matching at least one target data processing template from a plurality of precompiled data processing templates according to the matrix dimension and the vector memory capacity; According to the at least one target data processing template, the matrix data corresponding to the sub-matrix to be processed is loaded into the vector memory corresponding to the tensor computing core, so that the tensor computing core performs a matrix operation on the matrix data to obtain a sub-matrix operation result; 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.

[0008] In a possible implementation, matching at least one target data processing template from a plurality of precompiled data processing templates according to the matrix dimension and the vector memory capacity includes: According to the matrix dimension, determining at least one optimal data processing template matching the matrix dimension from a plurality of pre-compiled data processing templates; For each optimal data processing template, a target data processing template corresponding to the optimal data processing template is determined according to a relationship between the amount of data corresponding to the optimal data processing template and the capacity of the vector memory.

[0009] In a possible implementation, determining, according to the matrix dimension, at least one optimal data processing template that matches the matrix dimension from a plurality of precompiled data processing templates includes: obtaining template sizes of the plurality of precompiled data processing templates; If, among the multiple precompiled data processing templates, there is a data processing template whose template size is greater than or equal to the matrix dimension, selecting the data processing template with the smallest template size from the data processing templates whose template size is greater than or equal to the matrix dimension as the optimal data processing template; If the template sizes of the multiple pre-compiled data processing templates are all smaller than the matrix dimension, then multiple data processing templates whose spliced ​​template sizes are greater than or equal to the matrix dimension and whose spliced ​​template sizes are the smallest are selected as the optimal data processing templates.

[0010] In a possible implementation, determining a target data processing template corresponding to the optimal data processing template according to a relationship between the amount of data corresponding to the optimal data processing template and the capacity of the vector memory includes: Detecting whether the amount of data corresponding to the optimal data processing template is greater than the capacity of the vector memory; If the data volume corresponding to the optimal data processing template is greater than the capacity of the vector memory, a plurality of suboptimal data processing templates having data volumes smaller than the capacity of the vector memory 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 the suboptimal data processing templates 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 volume corresponding to the optimal data processing template is less than or equal to the capacity of the vector memory, the optimal data processing template is determined as the target data processing template.

[0011] In a possible implementation, intermediate operation results of the tensor computing core in the process of performing matrix operations on matrix data are stored in the corresponding vector memory.

[0012] In one possible implementation, the input matrix in the high bandwidth memory and the matrix data in the vector memory are both stored on 512-byte boundaries.

[0013] In one possible implementation, the tensor computing core uses tensor instructions of the TPU to process matrix data in batches during the process of performing matrix operations on matrix data.

[0014] In a second aspect, an embodiment of the present invention provides a device for implementing matrix operations on a TPU, including: A partitioning module is used to partition the input matrix in the high-bandwidth memory according to the number of tensor computing cores in the TPU, and obtain the sub-matrix to be processed corresponding to each tensor computing core; Matching modules for: For each tensor computing core, obtain the matrix dimension of the sub-matrix to be processed corresponding to the tensor computing core and the vector memory capacity corresponding to the tensor computing core; matching at least one target data processing template from a plurality of precompiled data processing templates according to the matrix dimension and the vector memory capacity; Compute module for: According to the at least one target data processing template, the matrix data corresponding to the sub-matrix to be processed is loaded into the vector memory corresponding to the tensor computing core, so that the tensor computing core performs a matrix operation on the matrix data to obtain a sub-matrix operation result; 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.

[0015] In a third aspect, an embodiment of the present invention provides an electronic device comprising a memory and a processor, wherein the memory stores a computer program, and when the processor executes the computer program, the method in the first aspect or any possible implementation of the first aspect is implemented.

[0016] In a fourth aspect, an embodiment of the present invention provides a computer-readable storage medium storing a computer program, which, when executed by a processor, implements the method in the first aspect or any possible implementation of the first aspect.

[0017] In a fifth aspect, an embodiment of the present invention provides a computer program product, including a computer program, which, when executed by a processor, implements the method in the first aspect or any possible implementation of the first aspect.

[0018] The embodiment of the present invention takes into account both the matrix dimension and the vector memory capacity to dynamically match the target data processing template from multiple pre-compiled data processing templates, and loads the data from the high-bandwidth memory into the vector memory according to the target data processing template to perform matrix operations. This method can achieve deep adaptation of the target data processing template and the TPU hardware memory level, which can not only give full play to the advantages of the vector memory unique to the TPU, reduce the number of accesses to the high-latency high-bandwidth memory, reduce memory access latency, and comprehensively improve the overall computing performance of the system, but also increase the flexibility of matrix operations on the TPU. In addition, the embodiment of the present invention also splits the input matrix in the high-bandwidth memory according to the hardware characteristics of the TPU, so that each tensor computing core can perform parallel collaborative calculations, which can reduce the computing pressure of each tensor computing core and improve computing efficiency. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 11 is a hardware architecture diagram of a TPU provided in an embodiment of the present invention; Figure 2 This is a flowchart of a method for implementing matrix operations on a TPU provided by an embodiment of the present invention; Figure 3 Schematic diagram of the division of the input matrix provided by an embodiment of the present invention; Figure 4 This is a schematic diagram of data filling when the number of target data processing templates provided by an embodiment of the present invention is 1; Figure 5 This is a schematic diagram of data partitioning when the number of target data processing templates provided by an embodiment of the present invention is 2; Figure 6 1 is a schematic structural diagram of a method and apparatus for implementing matrix operations on a TPU provided by an embodiment of the present invention; Figure 7 is a schematic diagram of an electronic device provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0020] The embodiments of the present invention will be described in detail below with reference to the accompanying drawings.

[0021] Template technology is an optimization method that dynamically adapts input parameters to predefined code structures (such as module size and data type). Its core goal is to reduce the overhead of dynamic adjustments during matrix operations and improve computational efficiency. Most related technologies apply template technology used in GPUs to TPUs to perform matrix operations. However, the hardware architecture of TPUs (such as vector instruction sets and memory hierarchies) differs significantly from that of GPUs. Directly applying GPU template solutions to TPUs fails to fully exploit the performance advantages of TPUs and creates efficiency bottlenecks.

[0022] To achieve efficient matrix operations on the TPU and fully leverage its performance advantages, the present invention dynamically matches a target data processing template from multiple pre-compiled data processing templates, taking into account both matrix dimensions and vector memory capacity. Data is then loaded from high-bandwidth memory into vector memory according to the target data processing template to perform matrix operations. This approach allows for deep adaptation of the target data processing template and the TPU hardware memory hierarchy, fully leveraging the advantages of the TPU's unique vector memory, reducing the number of accesses to high-latency, high-bandwidth memory, lowering memory access latency, and comprehensively improving the overall computing performance of the system.

[0023] First, let's briefly introduce the hardware architecture of TPU. Figure 1The TPU includes at least two tensor cores, which can be used to perform matrix operations. Each tensor core is equipped with a corresponding vector memory (VMEM), and at the same time, all tensor cores share the same high bandwidth memory (HBM).

[0024] Considering that matrix multiplication is the core operation in deep learning and high-performance computing, the present invention takes matrix multiplication as an example and introduces the implementation method of matrix operations on the TPU provided by the present invention based on the above TPU architecture. Figure 2 , which shows a flowchart of the implementation method of matrix operations on the TPU provided by an embodiment of the present invention, as detailed below: Step 201 : Divide the input matrix in the high bandwidth memory according to the number of tensor computing cores in the TPU to obtain a sub-matrix to be processed corresponding to each tensor computing core.

[0025] Each tensor core in the TPU can be used to perform matrix multiplication. To fully utilize the advantages of the TPU's tensor cores, the present invention divides the input matrix in the HBM according to the number of tensor cores within the TPU. This allows each tensor core to perform matrix operations on the sub-matrices to be processed in parallel, reducing the computing pressure on a single core.

[0026] Here, the input matrix in the HBM is the matrix input to the TPU for performing matrix multiplication. It can be understood that there are two matrices involved in the matrix multiplication, and accordingly, the number of input matrices in the HBM is two.

[0027] In order to facilitate the parallel computing of each tensor computing core and avoid data interaction between each tensor computing core during the matrix operation process, the embodiment of the present invention can divide the right matrix in the input matrix based on the principle of "splitting the number of columns of the right matrix in the matrix multiplication according to the number of tensor computing cores" to obtain multiple sub-right matrices, and respectively determine the left matrix and the divided sub-right matrices as the sub-matrices to be processed corresponding to each tensor computing core.

[0028] For example, see Figure 3 , the number of tensor computing cores is 2. When the input matrix is ​​an 8×128 left matrix A and a 128×128 right matrix B, when calculating the product of the left matrix A and the right matrix B, the columns of the right matrix B are divided into two 128×64 sub-right matrices B1 and B2. The left matrix A and the divided sub-matrix B1 are used as the sub-matrix to be processed corresponding to one tensor computing core, and the left matrix A and the divided sub-right matrix B2 are used as the sub-matrix to be processed corresponding to another tensor computing core.

[0029] Step 202: For each tensor computing core, obtain the matrix dimension of the sub-matrix to be processed corresponding to the tensor computing core and the vector memory capacity corresponding to the tensor computing core.

[0030] In this embodiment of the present invention, the matrix dimensions (ah, aw, bw), data type, and storage layout of the submatrix to be processed can be obtained through tensor metadata. Ah represents the number of rows in the left matrix, aw represents the number of columns in the left matrix (i.e., the number of rows in the right submatrix), and bw represents the number of columns in the right submatrix.

[0031] Data types include, but are not limited to, 32-bit floating point numbers (FP32), 16-bit brain floating point numbers (BF16), or 8-bit signed integers (INT8). Storage layouts include row-major and column-major.

[0032] Step 203 : Match at least one target data processing template from a plurality of pre-compiled data processing templates according to the matrix dimension and the vector memory capacity.

[0033] In the embodiment of the present invention, data processing templates of various template sizes can be pre-compiled to cover common matrix dimensions, such as 128×256×128, 256×256×256, and 512×512×256.

[0034] The embodiment of the present invention can match at least one target data processing template from a variety of pre-compiled data processing templates according to the matrix dimension to perform matrix multiplication. At the same time, the embodiment of the present invention also takes into account the VMEM advantage unique to the TPU, and can store matrix data in VMEM when each tensor computing core performs matrix operations, so that the tensor computing core can load the matrix data in real time. To this end, when matching the target data processing template, the embodiment of the present invention also considers the size of the vector memory capacity to avoid the situation where VMEM cannot load the matrix data in the target data processing template at one time and needs to access HBM midway.

[0035] By comprehensively considering the matrix dimension and vector memory capacity, the embodiment of the present invention can dynamically match a target data processing template suitable for the input matrix and TPU from a variety of pre-compiled data processing templates.

[0036] It should be noted here that all template sizes are integer multiples of 128, which is aligned with the 512-byte vector instruction width of the TPU, thereby dynamically adapting to the TPU instruction set.

[0037] The related technology only supports fixed template sizes when applying GPU template technology, and requires pre-generating multiple codes for different input matrices, resulting in code expansion. In addition, there are differences between the GPU instruction set and the TPU's 512-byte vector instructions. When applying GPU template technology, the TPU needs to simulate GPU instructions through complex logic, which increases delay calculation.

[0038] On the one hand, the embodiments of the present invention dynamically match target data processing templates from multiple data processing templates by comprehensively considering matrix dimensions and vector memory capacity, without causing code bloat and increasing the flexibility of template matching. On the other hand, the data processing templates in the embodiments of the present invention can dynamically adapt to the TPU instruction set, eliminating the need to emulate GPU instructions, effectively reducing computational latency.

[0039] In addition, the data processing template also encapsulates different data types and storage layouts to adapt to the data type and storage layout of the input matrix. This eliminates the need for developers to manually handle data type conversion and layout adaptation, reduces the development complexity of cross-operator collaboration, and improves the programming efficiency of the matrix multiplication module in the framework.

[0040] This embodiment of the present invention generates specialized kernel code at compile time for different data types to adapt to specific matrix multiplication scenarios and other possible computing needs. This ensures that the code maintains good flexibility and high performance under different precision requirements, meeting the needs of diverse application scenarios.

[0041] The data processing template in the embodiment of the present invention can be compatible with different storage layouts of row priority and column priority by generating load instruction sequences of different storage formats during compilation, and adapt to the TPU's Direct Memory Access (DMA) transfer unit for efficient reading of 512-byte aligned data.

[0042] The TPU's 512-byte vector instructions (such as load with offset and load with Vmask) and vector memory VMEM provide the hardware foundation for matrix multiplication operations. By designing a data processing template that is compatible with the TPU instruction set and VMEM capacity, the embodiments of the present invention can give full play to the hardware advantages of the TPU and solve the computing bottleneck problem caused by applying GPU template technology.

[0043] Step 204 , according to at least one target data processing template, load the matrix data corresponding to the sub-matrix to be processed into the vector memory corresponding to the tensor computing core, so that the tensor computing core can perform matrix operations on the matrix data to obtain sub-matrix operation results.

[0044] After the target data processing template is determined, the matrix data corresponding to the sub-matrix to be processed may be loaded into the vector memory according to the template size of the target data processing template.

[0045] For example, see 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 matrix data in the sub-matrix to be processed are loaded into VMEM and filled in VMEM according to the template size. After filling, the data is strictly aligned with 512 bytes to adapt to the 512-byte vector instructions of TPU, ensuring that the vector instruction loads a complete row of data in a single time.

[0046] See also 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 sub-matrix to be processed will be split according to the template size, and the split matrix data will be loaded into VMEM in sequence for the tensor computing core to perform matrix operations.

[0047] This embodiment of the present invention dynamically matches the most appropriate target data processing template from a variety of precompiled data processing templates based on the input matrix's dimensions, data type, and storage layout. If the matrix dimensions don't exactly match the template size, a padding or splicing strategy is employed to ensure efficient adaptation while meeting the kernel parameter range, improving computational efficiency.

[0048] When the number of target data processing templates is 1, the tensor computing core directly performs matrix operations on the matrix data corresponding to the target data processing template to obtain sub-matrix operation results.

[0049] When the number of target data processing templates is greater than one, the tensor computing core performs matrix operations on the matrix data corresponding to each target data processing template in turn, obtaining intermediate operation results. These intermediate operation results are then further operated on to obtain the final sub-matrix operation result. The matrix operation process here follows conventional mathematical logic and will not be further explained here.

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

[0051] Here, by storing the matrix data used to perform matrix operations and the intermediate calculation results in the matrix operation process in the vector memory for real-time access and call by the tensor computing core, the advantages of the vector memory cache in the TPU can be fully utilized to minimize the number of accesses to high-latency and high-bandwidth memory, effectively reduce memory access latency, and comprehensively improve the overall computing performance of the system.

[0052] In some embodiments, the input matrix in the high bandwidth memory and the matrix data in the vector memory are stored on 512-byte boundaries.

[0053] Here, you can use macros such as ALIGN128 and ALIGN256 to align matrix data or intermediate calculation results to ensure that the data is stored on 512-byte boundaries, thereby matching the load unit of the TPU vector instruction. For example, you can use the macro instruction head_dim_aligned=ALIGN128(head_dim) to enable vector instructions such as v_f32_ld_tnsr_st_msk to load efficiently and avoid additional processing of non-aligned data.

[0054] In some embodiments, during the process of the tensor computing core performing matrix operations on matrix data, the tensor instructions of the TPU are used to batch process the matrix data.

[0055] In matrix multiplication operations, the TPU architecture's tensor instructions (such as load with offset and m_matmul_single) can be used to optimize data processing. These tensor instructions are capable of batch processing. Unlike the distributed 32-bit floating-point operations used in other architectures, TPU tensor instructions can load multiple data elements at once on 512-byte boundaries (for example, the load with offset instruction can load 1024 FP32 elements at a time). Leveraging the parallel hardware design of the TPU's vector processing unit, a single operation can process multiple data elements simultaneously. This batch processing mechanism significantly increases data processing parallelism and reduces instruction execution redundancy. Compared to traditional operation modes, it significantly accelerates the computational process and fully unleashes the TPU's high-performance potential in matrix multiplication scenarios.

[0056] The embodiment of the present invention fully utilizes the 512-byte vector instructions of the TPU to process matrix data, and can use a single instruction to process more data, thereby improving parallel efficiency.

[0057] In addition, considering that the loop iterations in matrix operations require multiple judgments and jumps, which results in judgment delays, the embodiments of the present invention can also use the #pragma unroll instruction in the TPU to optimize the loop structure, converting the loop iterations that originally required multiple judgments and jumps into a continuous sequential instruction sequence.

[0058] For example, suppose a loop needs to execute 10 operations. Each iteration needs to determine "completeness" and jump, which will cause frequent stalls in the processor pipeline. Through instruction pipeline optimization, the 10 iterations can be expanded into 8 consecutive operations (dynamically adjusted according to hardware capabilities), eliminating the intermediate conditional judgment link. The processor can therefore continuously receive instruction streams, just like a production line running at full speed without interruption, so that the processor's different processing units are always working at full capacity, significantly improving the continuity and parallelism of instruction execution.

[0059] 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 concatenated to obtain the final matrix operation result.

[0060] Still Figure 3 For example, a tensor core computes the left matrix A and the right sub-matrix B1, obtaining a sub-matrix result of 8×64. Another tensor core computes the matrix A and the right sub-matrix B2, obtaining a sub-matrix result of 8×64. The two sub-matrix results are concatenated to obtain the final matrix result of 8×128.

[0061] Compared with the prior art, the embodiment of the present invention takes into account both the matrix dimension and the vector memory capacity to dynamically match the target data processing template from multiple pre-compiled data processing templates, and loads the data from the high-bandwidth memory into the vector memory according to the target data processing template to perform matrix operations. This method can achieve deep adaptation of the target data processing template and the TPU hardware memory level, which can not only give full play to the advantages of the vector memory unique to the TPU, reduce the number of accesses to the high-latency high-bandwidth memory, reduce memory access latency, and comprehensively improve the overall computing performance of the system, but also increase the flexibility of matrix operations on the TPU. In addition, the embodiment of the present invention also splits the input matrix in the high-bandwidth memory according to the hardware characteristics of the TPU, so that each tensor computing core can perform parallel collaborative calculations, which can reduce the computing pressure of each tensor computing core and improve computing efficiency.

[0062] The following describes a specific implementation method for matching and determining a target data processing template from a plurality of pre-compiled data processing templates.

[0063] In some embodiments, at least one optimal data processing template that matches the matrix dimensions can be first determined from a plurality of pre-compiled data processing templates according to the matrix dimensions; then, for each optimal data processing template, a target data processing template corresponding to the optimal data processing template is determined based on the relationship between the amount of data corresponding to the optimal data processing template and the capacity of the vector memory.

[0064] For each dimension (ah / aw / bw), select the minimum template size that is not less than the dimension (for example, when ah=100, select 128, when aw=200, select 256), and prioritize matching data processing templates that meet the boundaries of all dimensions at the same time (for example, if the matrix dimension is 100×200×200, match the template size of 128×256×256).

[0065] Specifically, the template sizes of multiple precompiled data processing templates are obtained; if among the multiple precompiled data processing templates, there is a data processing template whose template size is greater than or equal to the matrix dimension, then 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 sizes of multiple precompiled data processing templates are all smaller than the matrix dimension, then multiple data processing templates whose spliced ​​template sizes are greater than or equal to the matrix dimension and whose spliced ​​template sizes are the smallest are selected as the optimal data processing templates.

[0066] When there are data processing templates with template sizes greater than or equal to the matrix dimensions, in order to ensure that the matrix data can be loaded in one go using one data processing template while reducing data filling overhead and avoiding wasting computing resources, an embodiment of the present invention selects the data processing template with the smallest template size from the data processing templates with template sizes greater than or equal to the matrix dimensions as the optimal data processing template.

[0067] When the sizes of multiple precompiled templates are all smaller than the matrix dimension, template splicing is used to select multiple data processing templates from multiple precompiled data processing templates, and the template size after splicing these multiple data processing templates is greater than or equal to the matrix dimension. At the same time, the size of the spliced ​​template is the smallest, thereby ensuring the complete loading of matrix data while reducing data splicing overhead and occupying the least computing resources.

[0068] Based on the above consideration of the matrix dimension to determine the optimal data processing template, an embodiment of the present invention 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 subsequent matrix operations.

[0069] Specifically, for each optimal data processing template, it is detected whether the data volume corresponding to the optimal data processing template is greater than the vector memory capacity; if the data volume 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.

[0070] When the amount of data 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 that the tensor computing core can access and call it at any time. At this time, the optimal data processing template can be directly determined as the target data processing template.

[0071] However, if the amount of data corresponding to the optimal data processing template is greater than the capacity of the vector memory, the vector memory cannot completely load all the matrix data corresponding to the optimal data processing template. This means that when the tensor computing core performs matrix operations according to the optimal data processing template, it needs to access the high-latency HBM to obtain the matrix data that has not been loaded into VMEM, thereby calculating the access delay problem.

[0072] In order to avoid the above-mentioned computational access delay problem, it is necessary to further determine the target data processing template based on the optimal data processing template, divide the matrix data according to the target data processing template, and perform subsequent matrix operations.

[0073] Specifically, if the data size corresponding to the optimal data processing template is greater than the vector memory capacity, multiple suboptimal data processing templates with data sizes smaller than the vector memory capacity are selected from the plurality of pre-compiled data processing templates, and the multiple suboptimal data processing templates are determined as target data processing templates. The template size of the combined suboptimal data processing templates is greater than or equal to the template size of the optimal data processing template, and the combined template size is the smallest.

[0074] If the data size of the optimal data processing template exceeds the capacity of the vector memory, multiple suboptimal data processing templates are selected from the multiple pre-compiled data processing templates and spliced ​​together to replace the optimal data processing template. Here, the data size of each suboptimal data processing template is less than the capacity of the vector memory, and the template size after splicing the suboptimal data processing templates is greater than or equal to the minimum template size of the optimal data processing template.

[0075] For example, if the VMEM capacity is 16MB and the matrix data type is FP32 (4 bytes / element), to ensure VMEM storage performance, 75% of its available capacity can be reserved. The available capacity is: 16MB x 75% = 12MB = 12,582,912 bytes. The maximum data size of the suboptimal data processing template is: 12,582,912 ÷ 4 = 3,145,728 elements. In other words, the template size in the suboptimal data processing template satisfies x x y x z ≤ 3,145,728. Here, x, y, and z represent the template size in each dimension, respectively.

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

[0077] To verify the implementation method of matrix operations on the TPU provided by the embodiment of the present invention, the embodiment of the present invention conducted statistical tests and performance tests.

[0078] Statistical tests: Compare the matrix multiplication results of the TPU and GPU to ensure numerical consistency. Compare the matrix multiplication results calculated by the TPU with those calculated by the GPU. GPUs are widely used in matrix calculations, and their calculation results are highly reliable. When performing a large number of matrix multiplication tests of different scales, accurately calculate the floating-point error of the TPU and GPU results. Set a floating-point error qualification standard, and after multiple rounds of testing, the TPU calculation results all meet this standard. The matrix multiplication implemented by the present invention on the TPU still has a high degree of numerical accuracy and can be reliably applied to various computing tasks.

[0079] Performance testing: The efficiency of a TPU templated implementation was compared with a general-purpose GPU solution at a typical template size (1024×4096×4096). Under the same computational task and test environment, a matrix multiplication program was run multiple times, and the number of computational loop structures was recorded. The results showed that the TPU templated implementation, using the matrix operations implementation method provided by the present invention, significantly reduced the computation time compared to a general-purpose GPU solution, improving efficiency by over 20%. This fully demonstrates that the present invention can effectively improve the computational performance of matrix multiplication on the TPU, demonstrating its excellent practicality and superiority.

[0080] It should be understood that the size of the serial numbers of the steps in the above embodiments does not mean the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiments of the present invention.

[0081] The following are device embodiments of the present invention. For details not fully described therein, reference may be made to the corresponding method embodiments described above.

[0082] Figure 6 The following is a schematic diagram of the structure of a device for implementing matrix operations on a TPU provided by an embodiment of the present invention. For ease of explanation, only the parts related to the embodiment of the present invention are shown, which are detailed as follows: like Figure 6 As shown, the device 6 for implementing matrix operations on the TPU includes: a division module 61, a matching module 62 and a calculation module 63.

[0083] A partitioning module 61 is configured to partition the input matrix in the high bandwidth memory according to the number of tensor computing cores in the TPU, and obtain a sub-matrix to be processed corresponding to each tensor computing core; The matching module 62 is configured to: For each tensor computing core, obtain the matrix dimension of the sub-matrix to be processed corresponding to the tensor computing core and the vector memory capacity corresponding to the tensor computing core; matching at least one target data processing template from a plurality of precompiled data processing templates according to matrix dimensions and vector memory capacity; The calculation module 63 is used to: According to at least one target data processing template, the matrix data corresponding to the sub-matrix to be processed is loaded into the vector memory corresponding to the tensor computing core, so that the tensor computing core performs matrix operations on the matrix data to obtain sub-matrix operation results; 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.

[0084] In a possible implementation, the matching module 62 is specifically configured to: According to the matrix dimension, determining at least one optimal data processing template matching the matrix dimension from a plurality of pre-compiled data processing templates; For each optimal data processing template, a target data processing template corresponding to the optimal data processing template is determined according to the relationship between the amount of data corresponding to the optimal data processing template and the capacity of the vector memory.

[0085] In a possible implementation, the matching module 62 is specifically configured to: Get template sizes for various precompiled data processing templates; If there is a data processing template whose template size is greater than or equal to the matrix dimension among the multiple pre-compiled data processing templates, then the data processing template with the smallest template size is selected from the data processing templates whose template size is greater than or equal to the matrix dimension as the optimal data processing template; If the template sizes of multiple pre-compiled data processing templates are all smaller than the matrix dimension, multiple data processing templates whose spliced ​​template sizes are greater than or equal to the matrix dimension and whose spliced ​​template sizes are smallest are selected as the optimal data processing templates.

[0086] In a possible implementation, the matching module 62 is specifically configured to: Detecting whether the data volume corresponding to the optimal data processing template is greater than the vector memory capacity; If the data volume corresponding to the optimal data processing template is larger than the capacity of the vector memory, multiple suboptimal data processing templates with data volumes smaller than the capacity of the vector memory are selected from the multiple pre-compiled data processing templates, and the multiple suboptimal data processing templates are determined as target data processing templates; wherein the template size after splicing the suboptimal data processing templates is larger 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 volume 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.

[0087] In one possible implementation, intermediate operation results of the tensor computing core in the process of performing matrix operations on matrix data are stored in the corresponding vector memory.

[0088] In one possible implementation, the input matrix in the high bandwidth memory and the matrix data in the vector memory are both stored on 512-byte boundaries.

[0089] In one possible implementation, the tensor core performs matrix operations on matrix data using tensor instructions of the TPU to process matrix data in batches.

[0090] This device embodiment can be used to implement the above method embodiment. Its technical principles and implementation effects are the same as those of the above method embodiment, and will not be repeated here.

[0091] Figure 7 Schematic diagram of an electronic device provided by an embodiment of the present invention. Figure 7 As shown, the electronic device 7 of this embodiment includes a processor 70 and a memory 71. The memory 71 stores a computer program 72. When the processor 70 executes the computer program 72, the steps of the above-described method embodiments are implemented. Alternatively, when the processor 70 executes the computer program 72, the functions of the modules / units in the above-described device embodiments are implemented.

[0092] Exemplarily, the computer program 72 may be divided into one or more modules / units, which are stored in the memory 71 and executed by the processor 70 to implement the present invention. The one or more modules / units may be a series of computer program instruction segments capable of implementing specific functions, and the instruction segments are used to describe the execution process of the computer program 72 in the electronic device 7.

[0093] The electronic device 7 may include, but is not limited to, a processor 70 and a memory 71. Those skilled in the art will appreciate that Figure 7It is only an example of the electronic device 7 and does not constitute a limitation of the electronic device 7. It may include more or fewer components than shown in the figure, or a combination of certain components, or different components. For example, the electronic device 7 may also include input and output devices, network access devices, buses, etc.

[0094] The processor 70 may be a central processing unit (CPU), or other general-purpose processors, digital signal processors (DSP), application-specific integrated circuits (ASIC), field-programmable gate arrays (FPGA), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. The general-purpose processor may be a microprocessor or any conventional processor.

[0095] The memory 71 can be an internal storage unit of the electronic device 7, such as the hard drive or 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 drive, a Smart Media Card (SMC), a Secure Digital (SD) card, a flash memory card, etc. equipped on the electronic device 7. Furthermore, the memory 71 can include both the internal storage unit of the electronic device 7 and an external storage device. 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 is about to be output.

[0096] For the sake of convenience and brevity, the division of the above functional modules / units is only used as an example. In actual applications, the above functions can be assigned to different functional modules / units as needed. The above modules / units can be implemented in the form of hardware, software, or a combination of hardware and software.

[0097] An embodiment of the present invention further provides a computer-readable storage medium storing a computer program that, when executed by a processor, implements the methods in the above-mentioned method embodiments.

[0098] An embodiment of the present invention further provides a computer program product, including a computer program, which, when executed by a processor, implements the methods in the above-mentioned method embodiments.

[0099] The term "computer program" includes computer program code, which may be in source code form, object code form, executable file, or some intermediate form. Computer-readable media may include any entity or device capable of carrying computer program code, recording media, USB flash drives, removable hard drives, magnetic disks, optical disks, computer memory, read-only memory (ROM), random access memory (RAM), electrical carrier signals, telecommunications signals, and software distribution media.

[0100] In the above embodiments, the descriptions of each embodiment have their own focus. For parts not described or recorded in detail in one embodiment, please refer to the relevant descriptions of other embodiments. Unless otherwise specified or there is a logical conflict, the terms and / or descriptions between different embodiments are consistent and can be referenced to each other. The technical features of different embodiments can be combined to form new embodiments based on their inherent logical relationships.

[0101] The embodiments described above are only used to illustrate the technical solutions of the present invention, rather than to limit the same. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. These modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the various embodiments of the present invention, and should all be included in the scope of protection of the present invention.

Claims

1. A method for implementing matrix operations on a TPU, characterized in that: include: Divide the input matrix in the high-bandwidth memory according to the number of tensor computing cores in the TPU to obtain the sub-matrix to be processed corresponding to each tensor computing core; For each tensor computing core, obtain the matrix dimension of the sub-matrix to be processed corresponding to the tensor computing core and the vector memory capacity corresponding to the tensor computing core; matching at least one target data processing template from a plurality of precompiled data processing templates according to the matrix dimension and the vector memory capacity; According to the at least one target data processing template, the matrix data corresponding to the sub-matrix to be processed is loaded into the vector memory corresponding to the tensor computing core, so that the tensor computing core performs a matrix operation on the matrix data to obtain a sub-matrix operation result; 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.

2. The method for implementing matrix operations on a TPU according to claim 1, wherein: Matching at least one target data processing template from a plurality of precompiled data processing templates according to the matrix dimension and the vector memory capacity includes: According to the matrix dimension, determining at least one optimal data processing template matching the matrix dimension from a plurality of pre-compiled data processing templates; For each optimal data processing template, a target data processing template corresponding to the optimal data processing template is determined according to a relationship between the amount of data corresponding to the optimal data processing template and the capacity of the vector memory.

3. The method for implementing matrix operations on a TPU according to claim 2, wherein: Determining at least one optimal data processing template matching the matrix dimension from a plurality of pre-compiled data processing templates according to the matrix dimension includes: obtaining template sizes of the plurality of precompiled data processing templates; If, among the multiple precompiled data processing templates, there is a data processing template whose template size is greater than or equal to the matrix dimension, selecting the data processing template with the smallest template size from the data processing templates whose template size is greater than or equal to the matrix dimension as the optimal data processing template; If the template sizes of the multiple pre-compiled data processing templates are all smaller than the matrix dimension, then multiple data processing templates whose spliced ​​template sizes are greater than or equal to the matrix dimension and whose spliced ​​template sizes are the smallest are selected as the optimal data processing templates.

4. The method for implementing matrix operations on a TPU according to claim 2, wherein: The determining, based on the relationship between the amount of data corresponding to the optimal data processing template and the capacity of the vector memory, a target data processing template corresponding to the optimal data processing template includes: Detecting whether the amount of data corresponding to the optimal data processing template is greater than the capacity of the vector memory; If the data volume corresponding to the optimal data processing template is greater than the capacity of the vector memory, a plurality of suboptimal data processing templates having data volumes smaller than the capacity of the vector memory 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 the suboptimal data processing templates 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 volume corresponding to the optimal data processing template is less than or equal to the capacity of the vector memory, the optimal data processing template is determined as the target data processing template.

5. The method for implementing matrix operations on a TPU according to any one of claims 1 to 4, characterized in that: The intermediate calculation results of the tensor calculation core in the process of performing matrix operations on matrix data are all stored in the corresponding vector memory.

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

7. The method for implementing matrix operations on a TPU according to any one of claims 1 to 4, characterized in that: During the process of performing matrix operations on matrix data by the tensor computing core, the matrix data is processed in batches using the tensor instructions of the TPU.

8. A device for implementing matrix operations on a TPU, characterized in that: include: A partitioning module is used to partition the input matrix in the high-bandwidth memory according to the number of tensor computing cores in the TPU, and obtain the sub-matrix to be processed corresponding to each tensor computing core; Matching modules for: For each tensor computing core, obtain the matrix dimension of the sub-matrix to be processed corresponding to the tensor computing core and the vector memory capacity corresponding to the tensor computing core; matching at least one target data processing template from a plurality of precompiled data processing templates according to the matrix dimension and the vector memory capacity; Compute module for: According to the at least one target data processing template, the matrix data corresponding to the sub-matrix to be processed is loaded into the vector memory corresponding to the tensor computing core, so that the tensor computing core performs a matrix operation on the matrix data to obtain a sub-matrix operation result; 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.

9. An electronic device, characterized in that: The method comprises a memory and a processor, wherein the memory stores a computer program, and when the processor executes the computer program, the method according to any one of claims 1 to 7 is implemented.

10. A computer-readable storage medium, characterized in that The computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, the method according to any one of claims 1 to 7 is implemented.

Images