High-performance SpMM kernel implementation method based on dense tensor core

Abstract

The invention relates to the technical field of sparse matrix multiplication, and discloses a high-performance SpMM kernel implementation method based on a dense tensor core. The method comprises the following steps: storing a sparse matrix in a BlockNM sparse format, and pre-fetching metadata in the sparse matrix into a shared memory of a GPU (Graphics Processing Unit); asynchronously copying the sparse matrix and the dense matrix to a shared memory, setting a four-stage pipeline scheduling strategy, and controlling SpMM operations of different iteration processes; loading the sparse matrix and the dense matrix in the shared memory into a register, setting an offset for each thread accessing the shared memory, and enabling each thread to access different banks in the shared memory based on the offset; the matrix multiplication and addition operation is executed in the register to obtain a calculation result, the calculation result is temporarily stored in the shared memory from the register, then the calculation result in the shared memory is written back to the global memory, and the efficiency of sparse matrix multiplication is improved.

Description

Technical Field

[0001] This application relates to the technical field of sparse matrix multiplication, and in particular to a high-performance SpMM kernel implementation method based on a dense tensor core. Background Art

[0002] Sparse matrix-matrix multiplication (SpMM) is a core computational operation in many scientific and engineering fields, including machine learning, scientific simulation, data analysis, and graphics processing. The main feature of SpMM is that one of its operands is a sparse matrix, that is, most elements in the matrix are zero, and only a small number of non-zero elements. Due to the characteristics of sparse matrices, traditional matrix multiplication methods are less efficient when dealing with SpMM because they cannot effectively utilize sparsity to reduce the amount of computation and memory access.

[0003] For example, the existing Chinese patent application with the publication number CN117454946A proposes a tensor core architecture system that supports unstructured sparse matrix calculations, including a matrix loading unit, a decoding unit, an array of processing units, and a matrix storage unit. Among them: after the MLU loads the data of the dense matrix and the sparse matrix from the HBM to the chip, it outputs the corresponding data to the DU. The DU selects the corresponding elements, that is, the valid elements, in the dense matrix block according to the index of the decoded sparse matrix block and outputs them to the PU array. After each PE array in the PU array completes all calculations, it stores the intermediate results in the MSU, and the MSU accumulates the intermediate results of the PU array to the corresponding on-chip buffer.

[0004] Another example is the Chinese patent application with the publication number CN117609677A, which discloses a sparse matrix multiplication acceleration method, an FPGA, a computing system, and a storage medium. This method is used to calculate the product of a sparse matrix A and a matrix B. The sparse matrix A is stored in an off-chip memory in units of sparse blocks. The method includes: configuring the parallelism parameter of the computing unit according to the available computing resources on the FPGA and the sparse block size of the sparse matrix A; determining the storage structure of the off-chip memory and the storage form of the data in the on-chip buffer according to the parallelism parameter and the data bit width of the off-chip memory; transferring the data in the on-chip memory to the on-chip buffer, and the sparse core computing unit on the FPGA calculates the data stored in the on-chip buffer to obtain the product of the sparse matrix A and the matrix B.

[0005] However, the above two documents have relatively low calculation efficiency for the product of the sparse matrix A and the matrix B. Therefore, a high-performance SpMM kernel implementation method based on a dense tensor core is needed to solve the above problems. Summary of the Invention

[0006] To solve the above technical problems, the present application provides a method for implementing a high-performance SpMM kernel based on a dense tensor core, which realizes efficient calculation of sparse matrix multiplication through block alignment, pipeline overlap, and hardware adaptation.

[0007] In a first aspect, the present application provides a method for implementing a high-performance SpMM kernel based on a dense tensor core, the method comprising: Step S1: Divide the sparse matrix into matrix blocks of a preset size, store the non-zero elements in each matrix block in the BlockNM sparse format, and obtain the column index data of the non-zero elements, defined as metadata, and prefetch the metadata to the shared memory of the GPU through a first instruction; Step S2: Read the dense matrix from the global memory, and asynchronously copy the sparse matrix and the dense matrix to the shared memory through a second instruction. In the shared memory, the dense matrix is stored in the row-major order format. After asynchronously copying to the shared memory, set a four-stage pipeline scheduling strategy, and control the matrix multiply-add operations in different iteration processes based on the four-stage pipeline scheduling strategy; Step S3: Load the sparse matrix and the dense matrix in the shared memory into the register in batches. During the loading process, set an offset for each thread accessing the shared memory, and based on the offset, each thread accesses different banks in the shared memory; Step S4: Execute matrix multiply-add operations in the register based on a third instruction to obtain a calculation result, and temporarily store the calculation result from the register to the shared memory, and write the calculation result in the shared memory back to the global memory based on a fourth instruction.

[0008] In combination with the first aspect, in the first implementation manner of the first aspect of the present application, storing in the BlockNM sparse format includes: The BlockNM sparse format sets the width of each matrix block in the horizontal direction to 64, and sets a value array and a column index array. The non-zero elements in the sparse matrix are stored in the value array in the order of the rows where they are located. The value array includes a numerical bit and a padding bit, and the padding bit is filled with zero elements. Generate a column index based on the column where each non-zero element is located, and store the column index in the column index array. The column index array includes an index bit and a flag bit, and the flag bit is used to mark the padding bit in the value array.

[0009] In combination with the first aspect, in the second implementation manner of the first aspect of the present application, asynchronously copying to the shared memory through a second instruction includes: The second instruction is an instruction to load 128-bit data from the global memory each time. During the copying process, the sparse matrix is divided into multiple BMxBK blocks along the K dimension, and the dense matrix is divided into multiple BKxBN blocks, where BM = 64, BK = 32, BN = 64, and K = 32.

[0010] Combined with the first aspect, in the third implementation manner of the first aspect of the present application, loading the sparse matrix and the dense matrix in the shared memory into the register in batches includes: The BMxBK block is sliced into multiple first sub-blocks with a size of 16x16, and each first sub-block is loaded into the register in row-major order. The BKxBN block is sliced into multiple second sub-blocks with a size of 16x8. The data in each second sub-block is transposed, and the second sub-block is converted from row-major order to column-major order and loaded into the register in the column-major order.

[0011] Combined with the first aspect, in the fourth implementation manner of the first aspect of the present application, setting a four-stage pipeline scheduling strategy includes: The four-stage pipeline scheduling strategy includes a first stage, a second stage, a third stage, and a fourth stage. The first stage is to prefetch the metadata in the (n + 1)-th iteration process. The second stage and the third stage are to wait for the sparse matrix and the dense matrix in the n-th iteration to be copied to the shared memory to complete. The fourth stage is to call the dense tensor core to execute the third instruction to complete the matrix multiplication and addition calculation.

[0012] Combined with the first aspect, in the fifth implementation manner of the first aspect of the present application, setting an offset for each thread accessing the shared memory includes: The shared memory is divided into multiple banks. An ID is assigned to each thread accessing the shared memory. When multiple threads access the banks in the shared memory, an offset is set for each thread, and the new address of each thread is calculated based on the first formula , and the first formula is: , where is the original memory address of each thread without an offset, , represents thread x, is the number of non-zero elements in each bank, is a multiple of the length of 32 banks, is the offset of each thread.

[0013] In combination with the first aspect, in the sixth implementation manner of the first aspect of the present application, performing a matrix multiply-accumulate operation in the register based on the third instruction includes: The third instruction is an MMA instruction. The MMA instruction is used to calculate the matrix multiplication of each first sub-block in the sparse matrix and each second sub-block in the dense matrix to obtain a result block. The size of the result block is 16x8, and the result block is accumulated into the result matrix, and the result matrix is a dense matrix pre-stored in the register.

[0014] In combination with the first aspect, in the sixth implementation manner of the first aspect of the present application, temporarily storing the calculation result from the register to the shared memory includes: During the process of writing the calculation result to the shared memory, an offset is set for each thread, and consecutive threads are scattered into different banks based on the offset.

[0015] In combination with the first aspect, in the seventh implementation manner of the first aspect of the present application, writing the calculation result in the shared memory back to the global memory based on the fourth instruction includes: The fourth instruction is a vectorized store instruction, which is used to write 128-bit data in the shared memory back to the global memory at one time. During the write-back process, it is set that each thread writes back 16 bytes at one time and aligns with the global memory address.

[0016] The first instruction is an instruction for loading 64-bit data from the global memory each time.

[0017] Compared with the prior art, the beneficial effects of the present invention are at least as follows: The present invention receives a sparse matrix A and a dense matrix B as inputs through SpMM and outputs a dense matrix C. In FastVector, the sparse matrix is stored in the proposed BlockNM format, and the metadata is represented by the low bits. By prefetching frequently accessed metadata in batches into the shared memory, the global memory access is minimized, thereby reducing the data access latency; in order to eliminate the bank conflicts generated when reading from matrix B and writing data back to matrix C to the shared memory, an offset is applied to the memory address, and multiple memory accesses are interleaved into different banks to avoid bank conflicts; in addition, the present invention adopts an efficient pipeline strategy to overlap memory access and calculation; the GPU parallelism is further improved by selecting the optimal tile size, thereby improving the calculation throughput. Description of the Drawings

[0018] To more clearly illustrate the technical solutions of the embodiments of the present invention, the following will briefly introduce the drawings required for the description of the embodiments. Obviously, the drawings in the following description are some embodiments of the present invention. For those of ordinary skill in the art, without creative efforts, other drawings can also be obtained based on these drawings.

[0019] Figure 1 It is a schematic diagram of an embodiment of a high-performance SpMM kernel implementation method based on a dense tensor core in an embodiment of the present application; Figure 2 It is a schematic diagram of various sparse formats in an embodiment of the present application; Figure 3 It is a schematic diagram of the FastVector method design in an embodiment of the present application; Figure 4 It is a schematic diagram of the BlockNM sparse format in an embodiment of the present application; Figure 5 It is a schematic diagram of a four-stage pipeline scheduling strategy in an embodiment of the present application; Figure 6 It is a schematic diagram of the access modes for reading a dense matrix from shared memory with and without offsets in an embodiment of the present application; Figure 7 It is a schematic diagram of the memory access position remapping of threads and shared banks in an embodiment of the present application; Figure 8 The write-back process of the shared memory in an embodiment of the present application; Figure 9 It is a relationship diagram of the mapping between C tiles and the global matrix C in an embodiment of the present application. Detailed implementation manners

[0020] The embodiments of the present application provide a high-performance SpMM kernel implementation method based on a dense tensor core. The terms "first", "second", "third", "fourth", etc. (if any) in the specification, claims and the above drawings of the present application are used to distinguish similar objects and do not necessarily describe a specific order or sequence. It should be understood that such data can be interchanged under appropriate circumstances so that the embodiments described herein can be implemented in an order different from that shown or described herein. In addition, the term "comprising" or "having" and any variation thereof are intended to cover non-exclusive inclusion. For example, a process, method, system, product or device that includes a series of steps or units does not necessarily limit to those clearly listed steps or units, but may include other steps or units not clearly listed or inherent to these processes, methods, products or devices.

[0021] For ease of understanding, the specific process of the embodiments of the present application will be described below. Please refer to Figure 1 An embodiment of the method for implementing a high-performance SpMM kernel based on a dense tensor core in the embodiments of the present application includes: Step S1: Divide the sparse matrix into matrix blocks of a preset size. The non-zero elements in each matrix block are stored in the BlockNM sparse format, and the column index data of the non-zero elements is obtained and defined as metadata. The metadata is prefetched into the shared memory of the GPU through a first instruction.

[0022] Specifically, as the basis of the SpMM kernel, the sparse format is crucial for reducing storage and preprocessing overhead. Structured sparsity can be divided into four categories: element-wise (EW), vector-wise (VW), block-wise (BW), and N:M sparsity. As shown in Figure 2, it is a schematic diagram of various sparse formats. Among them, EW sparsity involves individual weight pruning, which has the least impact on model accuracy. However, due to irregular computational patterns and scattered memory access, the capabilities of the GPU are usually not fully utilized. While VW and BW sparsity prune weight groups at a coarser granularity, which makes GPU execution more efficient but may lead to a decrease in model accuracy. In contrast, N:M sparsity can achieve both high model accuracy and inference efficiency. It enforces a structured distribution of non-zero weights, where every M elements must contain N non-zero elements (EW-N:M sparsity). The present invention proposes a new sparse format, namely the BlockNM format. Among them, this format provides an efficient storage format for accelerating the application of general N:M sparsity on the TC. When performing SpMM calculations, the input sparse matrix A is converted into the BlockNM format, and the column index data of the sparse matrix A is prefetched into the shared memory through a first instruction. The first instruction is the LDG.64 instruction, that is, an instruction to load 64-bit data from the global memory.

[0023] Step S2: Read the dense matrix from the global memory, and asynchronously copy the sparse matrix and the dense matrix into the shared memory through a second instruction. The dense matrix is stored in the row-major order format in the shared memory. After asynchronously copying into the shared memory, set a four-stage pipeline scheduling strategy, and control the matrix multiplication and addition operations of different iteration processes based on the four-stage pipeline scheduling strategy.

[0024] Specifically, read the data of the dense matrix B from the global memory. The global memory is a large-capacity but slow-access memory in the GPU, which is used to store data that needs to be accessed frequently. Copy the non-zero value data in the sparse matrix A and the dense matrix B into the shared memory through a second instruction. The second instruction is the LDG.128 instruction, that is, an instruction to load 128-bit data from the global memory.

[0025] During the process of moving the sparse matrix A and the dense matrix B to the shared memory, the tensor core (TC) cannot obtain the data required for calculation in a timely manner, resulting in insufficient utilization of the TC pipeline. On the other hand, the positioning of the sparse matrix B depends on the metadata of matrix A, which means that B must wait for the data movement of A to complete. This data dependency introduces additional synchronization overhead. Therefore, simply using a two-stage pipeline to load the matrices and perform TC calculations is inefficient for kernel execution. Therefore, the present invention designs a four-stage pipeline strategy combined with metadata prefetching. By using vector access to prefetch the metadata of multiple iterations, after copying the tiles of the sparse matrix A and the dense matrix B of the nth iteration to the shared memory, the thread immediately starts to asynchronously copy the tiles of the sparse matrix A and the dense matrix B of the (n + 1)th iteration, thus releasing the thread. Then, the thread will synchronously perform the TC calculation of the nth stage, enabling the data movement of the next stage to overlap with the current calculation. Reading 8-bit metadata using a larger data type (such as int) based on the four-stage pipeline strategy increases the length of the requested data, thereby improving the utilization of memory transactions. The dense matrix B no longer depends on the completion of the sparse matrix A reading, allowing the asynchronous copying of A and B to start simultaneously, reducing the synchronization overhead, and promoting the overlap of calculation and data movement, ultimately improving the utilization of the TC pipeline. The specific four-stage pipeline strategy will be elaborated later.

[0026] Step S3: Load the sparse matrix and the dense matrix in the shared memory into the register in batches. During the loading process, set an offset for each thread accessing the shared memory, and based on the offset, each thread accesses different banks in the shared memory.

[0027] Specifically, to make full use of the parallel computing power and memory hierarchy of the GPU, the present invention proposes a calculation method optimized for sparse matrix-matrix multiplication (SpMM), namely the FastVector method. As Figure 3 shown, it is a schematic diagram of the FastVector method design. Among them, SpMM receives a sparse matrix A and a dense matrix B as inputs and outputs a dense matrix C. When loading the sparse matrix A and the dense matrix B into the register, it is necessary to perform tiling processing on the sparse matrix A and the dense matrix B. The specific processing method will be elaborated later. In FastVector, the sparse matrix is stored in the proposed BlockNM format, and the metadata is represented by the lower bits. To eliminate the bank conflicts generated when loading from the dense matrix B and the sparse matrix into the register, set an offset for each thread accessing the shared memory, and stagger multiple memory accesses to different banks to avoid bank conflicts.

[0028] Step S4: Perform a matrix multiply-accumulate operation in the register based on the third instruction, obtain the calculation result, temporarily store the calculation result from the register in the shared memory, and write the calculation result in the shared memory back to the global memory based on the fourth instruction.

[0029] Specifically, call the Dense Tensor Core (DETC) to execute the third instruction (mma.m16n8k16) to complete the matrix multiply-accumulate calculation, temporarily store the calculation result from the register in the shared memory buffer, apply the address offset strategy to eliminate bank conflicts, and write the calculation result in the shared memory back to the global memory by the fourth instruction (STS.128 instruction) in consecutive addresses to ensure a fully merged transaction.

[0030] In a specific embodiment, the storage in the BlockNM sparse format specifically includes the following steps: The BlockNM sparse format sets the width of each matrix block in the horizontal direction to 64, and sets a value array and a column index array. The non-zero elements in the sparse matrix are stored in the value array in the order of their rows. The value array includes a numerical bit and a padding bit, and the padding bit is filled with zero elements. Generate a column index based on the column where each non-zero element is located, store the column index in the column index array, and the column index array includes an index bit and a flag bit, and the flag bit is used to mark the padding bit in the value array.

[0031] Specifically, as Figure 4 shown, it is a schematic diagram of the BlockNM sparse format. By dividing the sparse matrix A into multiple blocks (block), the width of each block in the horizontal direction is fixed (for example, 64). By fixing the block width, the locality of memory access is improved, which is convenient for parallel processing. Each non-zero element is stored in the value array in the order of its row. The value array includes a numerical bit and a padding bit, and the padding bit is filled with zero elements to ensure that the length of each block is the same. For example, for a 64x64 block, if there are 3 non-zero elements in the first row and 2 non-zero elements in the second row, then store the 3 non-zero elements in the first row first, and then store the 2 non-zero elements in the second row. For each non-zero element, record its column index. For example, if the non-zero element is in the third column, its column index is 3. Store the column index in the column index array, and the flag bit in the column index array is used to mark the padding bit in the value array, and the flag bit is represented by -1 to distinguish valid elements and padding elements when reading data.

[0032] Compared with the existing format, the proposed BlockNM format of the present invention has two main advantages. First, it minimizes the storage overhead. When using the BlockNM format, only H×K×(1 - sparsity) + 2 elements are required to represent the sparse matrix A, and only an 8-bit array is used to store the indices of non-zero elements. Second, it can directly meet the data distribution requirements of DETCs. This favorable feature enables the non-zero elements of the sparse matrix to be directly loaded into the register in sequence, eliminating the need for additional preprocessing steps and thus minimizing the preprocessing overhead.

[0033] In a specific embodiment, the asynchronous copying to the shared memory through the second instruction specifically includes the following steps: The second instruction is an instruction that loads 128 bits of data from the global memory each time. During the copying process, the sparse matrix is divided into multiple BMxBK blocks along the K dimension, and the dense matrix is divided into multiple BKxBN blocks, where BM = 64, BK = 32, BN = 64, and K = 32.

[0034] Specifically, to ensure effective conflict elimination, the sparse matrix A is divided into multiple BMxBK blocks along the K dimension, where BM = 64 and BK = 32. For example, for a 128x128 sparse matrix A, it can be divided into 2 64x32 blocks. The dense matrix B is divided into multiple BKxBN blocks, where BK = 32 and BN = 64. For a 128x128 dense matrix B, it can be divided into 4 32x64 blocks. For a 64x32 block of the sparse matrix A, 128 bits of data are loaded each time, and the data of the entire block is gradually loaded. Asynchronous copying allows the GPU to perform other computing tasks while the data is being transferred. For example, when loading a block of the sparse matrix A into the shared memory, a block of the dense matrix B can be loaded into the shared memory at the same time.

[0035] In a specific embodiment, the step of loading the sparse matrix and the dense matrix in the shared memory into the register in batches specifically includes the following steps: The BMxBK block is divided into multiple first sub-blocks with a size of 16x16, and each first sub-block is loaded into the register in row-major order. The BKxBN block is divided into multiple second sub-blocks with a size of 16x8. The data in each second sub-block is transposed, converting the second sub-block from row-major order to column-major order, and the second sub-block is loaded into the register in column-major order.

[0036] Specifically, the sparse matrix A has been divided into multiple 64x32 blocks, and the 64x32 blocks are further divided into 8 first sub-blocks, each with a size of 16x16. The dense matrix B has been divided into multiple 32x64 blocks, and the 32x64 blocks are divided into 16 second sub-blocks, with the size of the second sub-blocks being 16x8. The elements of the 16x16 first sub-block are read in row order from the shared memory, and the read elements are loaded into the registers of the GPU. Each thread is responsible for loading a part of the data. In the registers, the data in the first sub-block remains in row-major order, which is convenient for subsequent matrix multiplication operations. The elements of the 16x8 second sub-block are read in row order from the shared memory, the second sub-block is transposed to obtain an 8x16 sub-block, and the transposed 8x16 matrix is loaded into the registers in column-major order. Through reasonable matrix partitioning and transposition operations, the memory access pattern is optimized, thereby improving the computational efficiency of SpMM.

[0037] In a specific embodiment, setting the four-stage pipeline scheduling strategy specifically includes the following steps: The four-stage pipeline scheduling strategy includes a first stage, a second stage, a third stage, and a fourth stage. The first stage is to prefetch the metadata in the (n + 1)-th iteration process. The second and third stages are to wait for the sparse matrix and the dense matrix in the n-th iteration to be copied to the shared memory to complete. The fourth stage is to call the dense tensor core to execute the third instruction to complete the matrix multiply-add calculation.

[0038] Specifically, as Figure 5 shown, it is a schematic diagram of the four-stage pipeline scheduling strategy. Vector access is used to prefetch the metadata of multiple iterations. After the sparse matrix A and the dense matrix B tiles in the n-th iteration are copied to the shared memory, buffer 1 in the shared memory is used, and at the same time, the thread immediately starts to asynchronously copy the A and B tiles in the (n + 1)-th iteration to buffer 2 in the shared memory. Meanwhile, the thread will synchronously perform the TC calculation in the n-th stage, so that the data movement in the next stage overlaps with the current calculation. This method brings two benefits: First, by using a larger data type (such as int) to read 8-bit metadata, the utilization rate of memory transactions is improved by increasing the length of the requested data. Second, B no longer depends on the completion of the reading of A, allowing the asynchronous copying of A and B to start simultaneously. This reduces the synchronization overhead and promotes the overlap of calculation and data movement, ultimately improving the utilization rate of the TC pipeline.

[0039] In a specific embodiment, setting the offset for each thread accessing the shared memory specifically includes the following steps: The shared memory is divided into multiple banks, and an ID is assigned to each thread accessing the shared memory. When multiple threads access the banks in the shared memory, an offset is set for each thread, and the new address of each thread is calculated based on the first formula , the first formula is: , where is the original memory address of each thread without offset, , represents thread x, is the number of non-zero elements in each bank, is a multiple of the length of 32 banks, is the offset of each thread.

[0040] Specifically, as Figure 6 shown, it is a schematic diagram of the access patterns for reading a dense matrix from shared memory with and without offsets. The consecutive column indices of the sparse matrix exhibit a scattered distribution, which means that the elements involved in the calculation of each column of matrix B will also be scattered. Therefore, if B is stored in column-major order, the addresses accessed by different threads will be non-consecutive, resulting in low memory transaction efficiency. In contrast, storing B in row-major order can alleviate this problem because it allows a consecutive memory access pattern. Therefore, we adopt a row-major storage format to store matrix B and use vector instructions to load it into shared memory. However, reading matrix B from shared memory into registers may cause significant bank conflicts.

[0041] As Figure 6 shown, 32 threads in one warp access the same four banks. To solve these bank conflicts, the original consecutive memory addresses are interleaved across different banks by applying a memory address offset. Specifically, during each loop iteration of the ldmatrix.trans.x4 instruction execution, an offset is introduced for each thread to ensure that each thread accesses a unique address. The new address for each thread can calculate the offset through the first formula, and the memory offset ensures that the banks accessed by two consecutive threads in a quarter warp differ by exactly four banks, Figure 7 shows the remapping of the memory access positions of threads and shared banks, demonstrating that each loop iteration allows a quarter warp to access different banks, thus effectively achieving bank-conflict-free access.

[0042] In a specific embodiment, performing a matrix multiply-add operation on a register based on a third instruction specifically includes the following steps: The third instruction is an MMA instruction. The MMA instruction is used to calculate the matrix multiplication of each first sub-block in the sparse matrix and each second sub-block in the dense matrix to obtain a result block. The size of the result block is 16x8, and the result block is accumulated into the result matrix, where the result matrix is a dense matrix pre-stored in the register.

[0043] Specifically, assume that the sparse matrix A is divided into multiple 16x16 first sub-blocks, and the data of each first sub-block has been stored in the register. The dense matrix B is divided into multiple 16x8 sub-blocks, and the data of each sub-block has been stored in the register. Use the MMA instruction to calculate the matrix multiplication of each 16x16 sub-block of the sparse matrix A and each 16x8 sub-block of the dense matrix B, and accumulate the resulting 16x8 result matrix C from the matrix multiplication into the pre-existing 16x8 result matrix C. Repeat the above operations for all sub-blocks until the entire matrix multiply-add operation is completed.

[0044] In a specific embodiment, temporarily storing the calculation result from the register to the shared memory specifically includes the following steps: During the process of writing the calculation result to the shared memory, set an offset for each thread, and disperse consecutive threads to different banks based on the offset.

[0045] Specifically, as Figure 8 shown, it is the write-back process of the shared memory, where each fragment of the C matrix is written back to two rows of the C tile, and 32 threads access different banks. Figure 9 It shows the mapping between the C tile and the global matrix C, indicating that each thread writes back data from four banks to the global memory. Whether reading from or writing to the shared memory, the offset calculated based on the first formula is used to generate a new address. This ensures that during the read and write operations, each thread in a quarter of the warp accesses a different bank, thus effectively preventing bank conflicts.

[0046] In a specific embodiment, writing back the calculation result in the shared memory to the global memory based on the fourth instruction specifically includes the following steps: The fourth instruction is a vectorized store instruction, which is used to write back 128-bit data in the shared memory to the global memory at one time. During the write-back process, set each thread to write back 16 bytes at one time and align it with the global memory address.

[0047] Specifically, when writing back the result of matrix C to the global memory, the data in the registers between adjacent threads may not be completely continuous. Writing back directly to the global memory may prevent the implementation of a fully merged transaction and make it complex to utilize vector instructions (such as STG.128) to enhance throughput. To solve this problem, we use the shared memory as a buffer, store the data from different registers in a continuous manner, and then write it back to the global memory through vector instructions. This method ensures that all data in the transaction is continuous and valid in the global memory, thus fully implementing a merged transaction. However, this write-back process involves read and write operations in the shared memory, which may cause bank conflicts. To solve this problem, we adopt a similar memory address offset technique to prevent these conflicts from occurring.

[0048] Those skilled in the art can clearly understand that for the convenience and conciseness of description, the specific working processes of the above-described systems, systems, and units can refer to the corresponding processes in the foregoing method embodiments, and will not be elaborated herein.

[0049] If the integrated unit is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present application, in essence, or the part that contributes to the prior art, or all or part of this technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions for causing a computer device (which can be a personal computer, a server, or a network device, etc.) to execute all or part of the steps of the methods in various embodiments of the present application. The foregoing storage medium includes: various media such as USB flash drives, mobile hard disks, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical discs that can store program codes.

[0050] The above embodiments are only used to illustrate the technical solutions of the present application and are not intended to limit them; although the present application has been described in detail with reference to the foregoing embodiments, those of ordinary skill in the art should understand that they can still modify the technical solutions recorded in the foregoing embodiments, or perform equivalent replacements for some of the technical features; and these modifications or replacements do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of various embodiments of the present application.

Claims

1. A high-performance SpMM kernel implementation method based on dense tensor cores, characterized in that: The method comprises: Step S1: Divide the sparse matrix into matrix blocks of a preset size, store the non-zero elements in each matrix block in BlockNM sparse format, and obtain column index data of the non-zero elements, which is defined as metadata, and pre-fetch the metadata into the shared memory of the GPU through a first instruction; Step S2: reading a dense matrix from a global memory, and asynchronously copying the sparse matrix and the dense matrix to the shared memory through a second instruction, wherein the dense matrix is ​​stored in a row-major order format in the shared memory, and after being asynchronously copied to the shared memory, setting a four-stage pipeline scheduling strategy, and controlling matrix multiplication and addition operations of different iterative processes based on the four-stage pipeline scheduling strategy; Step S3: loading the sparse matrix and the dense matrix in the shared memory into registers in batches, and during the loading process, setting an offset for each thread accessing the shared memory, so that each thread accesses a different bank in the shared memory based on the offset; Step S4: Based on the third instruction, perform matrix multiplication and addition operations in the register to obtain a calculation result, and temporarily store the calculation result from the register to the shared memory, and based on the fourth instruction, write the calculation result in the shared memory back to the global memory.

2. The method according to claim 1, characterized in that: Using BlockNM sparse format for storage includes the following steps: The BlockNM sparse format sets the horizontal width of each matrix block to 64, and sets a value array and a column index array. The non-zero elements in the sparse matrix are stored in the value array in the order of their rows. The value array includes numerical bits and padding bits. The padding bits are filled with zero elements. A column index is generated based on the column where each non-zero element is located, and the column index is stored in the column index array. The column index array includes index bits and mark bits. The mark bits are used to mark the padding bits in the value array.

3. The method according to claim 1, characterized in that Asynchronously copying to the shared memory by the second instruction includes the following steps: The second instruction is an instruction to load 128 bits of data from the global memory each time. During the copying process, the sparse matrix is ​​divided into multiple BMxBK blocks along the K dimension, and the dense matrix is ​​divided into multiple BKxBN blocks, where BM=64, BK=32, BN=64, and K=32.

4. The method according to claim 3, characterized in that Loading the sparse matrix and the dense matrix in the shared memory into registers in batches comprises the following steps: The BMxBK block is divided into a plurality of first sub-blocks, the size of the first sub-blocks is 16x16, and each of the first sub-blocks is loaded into the register in a row-major format, the BKxBN block is divided into a plurality of second sub-blocks, the size of the second sub-blocks is 16x8, a transposition operation is performed on the data in each second sub-block, the second sub-blocks are converted from a row-major format to a column-major format, and the second sub-blocks are loaded into the register in the column-major format.

5. The method according to claim 3, characterized in that: Setting up a four-stage pipeline scheduling strategy involves the following steps: The four-stage pipeline scheduling strategy includes the first stage, the second stage, the third stage and the fourth stage. The first stage is to pre-fetch metadata during the n+1th iteration. The second stage and the third stage are to wait for the sparse matrix of the nth iteration and the dense matrix to be copied to the shared memory. The fourth stage is to call the dense tensor core to execute the third instruction to complete the matrix multiplication and addition calculation.

6. The method according to claim 1, characterized in that Setting the offset for each thread accessing the shared memory comprises the following steps: The shared memory is divided into multiple banks, and an ID is assigned to each thread accessing the shared memory. When multiple threads access a bank in the shared memory, an offset is set for each thread, and a new address of each thread is calculated based on the first formula. , the first formula is: ,in, is the original memory address of each thread without the offset, , represents thread x, is the number of non-zero elements in each bank, A multiple of 32 bank lengths, The offset for each thread.

7. The method according to claim 4, characterized in that Performing a matrix multiplication and addition operation in the register based on the third instruction comprises the following steps: The third instruction is an MMA instruction, which is used to calculate the matrix multiplication of each first sub-block in the sparse matrix and each second sub-block in the dense matrix to obtain a result block, the size of which is 16x8, and accumulate the result block into a result matrix, which is a dense matrix pre-stored in a register.

8. The method according to claim 5, characterized in that Temporarily storing the calculation result from the register to the shared memory comprises the following steps: In the process of writing the calculation result into the shared memory, an offset is set for each thread, and continuous threads are dispersed into different banks based on the offset.

9. The method according to claim 1, characterized in that: Writing the calculation result in the shared memory back to the global memory based on the fourth instruction includes the following steps: The fourth instruction is a vectorized storage instruction, which is used to write the 128-bit data in the shared memory back to the global memory at one time. During the writing back process, each thread is set to write back 16 bytes at one time and align with the global memory address.

10. The method according to claim 1, characterized in that The first instruction is an instruction for loading 64 bits of data each time from the global memory.

Images