High-Performance SpMM Kernel Implementation Method Based on Dense Tensor Cores

Abstract

This application 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 includes: storing the sparse matrix in the BlockNM sparse format, and prefetching the metadata in the sparse matrix into the shared memory of the GPU; asynchronously copying the sparse matrix and the dense matrix into the shared memory, setting a four-stage pipeline scheduling strategy to control the SpMM operations in different iteration processes; loading the sparse matrix and the dense matrix in the shared memory into the registers, setting an offset for each thread accessing the shared memory, and making each thread access different banks in the shared memory based on the offset; performing matrix multiply-accumulate operations in the registers to obtain the calculation result, temporarily storing the calculation result from the registers to the shared memory, and then writing the calculation result in the shared memory back to the global memory. This application improves the efficiency of sparse matrix multiplication.

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 few 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 onto 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, the intermediate results are stored in the MSU, and the MSU accumulates the intermediate results of the PU array into 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 high-performance SpMM kernel implementation method 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 high-performance SpMM kernel implementation method based on a dense tensor core, and the method includes:

[0008] 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, which is defined as metadata. The metadata is prefetched into the shared memory of the GPU through a first instruction.

[0009] 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. In the shared memory, the dense matrix is stored in the row-major order format. After being asynchronously copied into the shared memory, set a four-stage pipeline scheduling strategy, and control the matrix multiply-accumulate operations in different iteration processes based on the four-stage pipeline scheduling strategy.

[0010] Step S3: Load the sparse matrix and the dense matrix in the shared memory into the registers 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.

[0011] Step S4: Execute the matrix multiply-accumulate operation in the registers based on a third instruction to obtain a calculation result, and temporarily store the calculation result from the registers into the shared memory. Write the calculation result in the shared memory back to the global memory based on a fourth instruction.

[0012] Combined with the first aspect, in the first implementation manner of the first aspect of the present application, storing in the BlockNM sparse format includes:

[0013] 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, 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.

[0014] In combination with the first aspect, in the second implementation manner of the first aspect of the present application, asynchronously copying the second instruction into the shared memory includes:

[0015] The second instruction is an instruction for loading 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.

[0016] In combination 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:

[0017] 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, 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.

[0018] In combination 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:

[0019] 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 multiply-add calculation.

[0020] In combination 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:

[0021] 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 offset, , represents thread x, is the number of non-zero elements in each bank, It is a multiple of the length of 32 banks, which is the offset for each thread.

[0022] Combined with the first aspect, in the sixth implementation manner of the first aspect of the present application, performing a matrix multiply-accumulate operation on the register based on the third instruction includes:

[0023] The third instruction is an MMA instruction, and 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.

[0024] Combined 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:

[0025] During the process of writing the calculation result to the shared memory, an offset is set for each thread, and consecutive threads are scattered to different banks based on the offset.

[0026] Combined 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:

[0027] 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, each thread is set to write back 16 bytes at one time and is aligned with the global memory address.

[0028] The first instruction is an instruction that loads 64-bit data from the global memory each time.

[0029] Compared with the prior art, the beneficial effects of the present invention are at least as follows:

[0030] 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 lower 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; by selecting the optimal tile size, the GPU parallelism is further improved, thereby increasing the calculation throughput. Brief Description of the Drawings

[0031] 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.

[0032] Figure 1 Schematic diagram of an embodiment of the high-performance SpMM kernel implementation method based on a dense tensor core in the embodiments of the present application;

[0033] Figure 2 Schematic diagrams of various sparse formats in the embodiments of the present application;

[0034] Figure 3 Schematic diagram of the FastVector method design in the embodiments of the present application;

[0035] Figure 4 Schematic diagram of the BlockNM sparse format in the embodiments of the present application;

[0036] Figure 5 Schematic diagram of the four-stage pipeline scheduling strategy in the embodiments of the present application;

[0037] Figure 6 Schematic diagram of the access modes for reading a dense matrix from shared memory with and without offsets in the embodiments of the present application;

[0038] Figure 7 Schematic diagram of the memory access location remapping of threads and shared banks in the embodiments of the present application;

[0039] Figure 8 Write-back process of the shared memory in the embodiments of the present application;

[0040] Figure 9 Relationship diagram of the mapping between C tiles and the global matrix C in the embodiments of the present application. Detailed Embodiments

[0041] The embodiments of the present application provide a method for implementing a high-performance SpMM kernel based on a dense tensor core. The terms "first", "second", "third", "fourth", etc. (if any) in the specification, claims and the above-mentioned 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 the data used in this way can be interchanged under appropriate circumstances so that the embodiments described here can be implemented in an order different from that illustrated or described here. In addition, the term "including" or "having" and any variations 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 have to be limited to those steps or units clearly listed, but may include other steps or units not clearly listed or inherent to these processes, methods, products or devices.

[0042] 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:

[0043] 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.

[0044] 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, it usually cannot fully utilize the capabilities of the GPU. In contrast, 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 comparison, 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.

[0045] 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 the second instruction. In the shared memory, the dense matrix is stored in row-major order. After asynchronously copying to the shared memory, set a four-stage pipeline scheduling strategy, and control the matrix multiply-accumulate operations in different iteration processes based on the four-stage pipeline scheduling strategy.

[0046] 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. Asynchronously copy the non-zero value data in the sparse matrix A and the dense matrix B to the shared memory through the second instruction. The second instruction is the LDG.128 instruction, that is, the instruction to load 128-bit data from the global memory.

[0047] 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 time, resulting in insufficient utilization of the TC pipeline. On the other hand, the positioning of the sparse matrix B depends on the metadata of the matrix A, which means that B must wait for the data movement of A to complete. This data dependence introduces additional synchronization overhead. Therefore, simply using a two-stage pipeline to load the matrix and perform TC calculations is inefficient for kernel execution. Therefore, the present invention designs a four-stage pipeline strategy combined with metadata prefetching. Use 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, thereby releasing the thread. Then, the thread will synchronously perform the TC calculation of the nth stage, so that the data movement in the next stage overlaps with the current calculation. Read 8-bit metadata using a larger data type (such as int) based on the four-stage pipeline strategy, increasing 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 reading of the sparse matrix A, 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 described in detail later.

[0048] 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 a different bank in the shared memory.

[0049] Specifically, to make full use of the parallel computing power and memory hierarchy of the GPU, the present invention proposes a computational method optimized for sparse matrix-matrix multiplication (SpMM), namely the FastVector method. As Figure 3 shown, it is a schematic diagram of the design of the FastVector method. 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 tile the sparse matrix A and the dense matrix B. The specific processing method will be described 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, an offset is set for each thread accessing the shared memory, and multiple memory accesses are interleaved into different banks to avoid bank conflicts.

[0050] Step S4: Execute the matrix multiply-add operation in the register based on the third instruction to obtain the calculation result, and temporarily store the calculation result from the register to the shared memory. Write the calculation result in the shared memory back to the global memory based on the fourth instruction.

[0051] Specifically, call the dense tensor core (DETC) to execute the third instruction (mma.m16n8k16) to complete the matrix multiply-add calculation, temporarily store the calculation result from the register to 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 at consecutive addresses through the fourth instruction (STS.128 instruction) to ensure a fully merged transaction.

[0052] In a specific embodiment, the storage using the BlockNM sparse format specifically includes the following steps:

[0053] 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 numerical bits and padding bits, and the padding bits are filled with zero elements. Column indices are generated based on the columns where each non-zero element is located, and the column indices are stored in the column index array. The column index array includes index bits and flag bits, and the flag bits are used to mark the padding bits in the value array.

[0054] Specifically, as Figure 4As shown, it is a schematic diagram of the BlockNM sparse format. By dividing the sparse matrix A into multiple blocks, 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, facilitating parallel processing. Each non-zero element is stored in the value array in the order of its row. The value array includes numerical bits and padding bits, and the padding bits are 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 the 3 non-zero elements in the first row are stored first, and then the 2 non-zero elements in the second row are stored. For each non-zero element, its column index is recorded. For example, if the non-zero element is in the 3rd column, its column index is 3. The column indices are stored in the column index array, and the marker bits in the column index array are used to mark the padding bits in the value array, with the marker bits represented by -1 to distinguish valid elements and padding elements when reading data.

[0055] Compared with the existing format, the BlockNM format proposed by 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 needed to represent the sparse matrix A, and only one 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 registers in order, eliminating the need for additional preprocessing steps, thus minimizing the preprocessing overhead.

[0056] In a specific embodiment, the asynchronous copying to the shared memory by the second instruction specifically includes the following steps:

[0057] The second instruction is an instruction that loads 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.

[0058] Specifically, to ensure the effective implementation of 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-bit data is loaded each time, and the data of the entire block is loaded step by step. Asynchronous copying allows the GPU to perform other computing tasks while 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.

[0059] In a specific embodiment, loading the sparse matrix and the dense matrix in the shared memory into the register in batches specifically includes the following steps:

[0060] 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, and the data in each second sub-block is transposed, converting the second sub-block from row-major order to column-major order, and loading the second sub-block into the register in column-major order.

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

[0062] In a specific embodiment, setting a four-stage pipeline scheduling strategy specifically includes the following steps:

[0063] 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 prefetch the metadata in the (n + 1)-th iteration. The second and third stages are to wait for the sparse matrix and dense matrix of 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.

[0064] 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 dense matrix B tiles of the n-th iteration are copied to the shared memory, buffer 1 in the shared memory is used. At the same time, the thread immediately starts to asynchronously copy the A and B tiles of the (n + 1)-th iteration to buffer 2 in the shared memory. Meanwhile, the thread will synchronously perform the TC calculation of the n-th stage, making the data movement in the next stage overlap 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 A's read, 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.

[0065] In a specific embodiment, setting an offset for each thread accessing the shared memory specifically includes the following steps:

[0066] 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 , 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.

[0067] Specifically, as Figure 6As 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 are also scattered. Therefore, if B is stored in column-major order, the addresses accessed by different threads will be non-consecutive, resulting in inefficient memory transactions. In contrast, storing B in row-major order can alleviate this problem because it allows for 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.

[0068] 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, 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. 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.

[0069] In a specific embodiment, performing a matrix multiply-accumulate operation in registers based on a third instruction specifically includes the following steps:

[0070] 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, which is a dense matrix pre-stored in the registers.

[0071] 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 registers. The dense matrix B is divided into multiple 16x8 sub-blocks, and the data of each sub-block has been stored in the registers. 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 into the pre-existing 16x8 result matrix C. Repeat the above operation for all sub-blocks until the entire matrix multiply-accumulate operation is completed.

[0072] In a specific embodiment, temporarily storing the calculation result from the register to the shared memory specifically includes the following steps:

[0073] During the process of writing the calculation result into the shared memory, an offset is set for each thread, and based on the offset, consecutive threads are dispersed into different banks.

[0074] 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, an 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.

[0075] In a specific embodiment, writing the calculation result in the shared memory back to the global memory based on the fourth instruction specifically includes the following steps:

[0076] 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, each thread is set to write back 16 bytes at one time and aligns with the global memory address.

[0077] Specifically, when writing the result of matrix C back to the global memory, the data in the registers between adjacent threads may not be completely continuous. Writing directly back to the global memory may prevent the realization 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 realizing the 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 technology to prevent these conflicts from occurring.

[0078] 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.

[0079] When 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 this 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 may be a personal computer, a server, or a network device, etc.) to execute all or part of the steps of the methods of various embodiments of this 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.

[0080] In the above, the above embodiments are only used to illustrate the technical solutions of this application, rather than to limit it; although this 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 make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of various embodiments of this application.

Claims

1. A high-performance SpMM kernel implementation method based on dense tensor cores, characterized in that The method 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. 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. In the shared memory, the dense matrix is stored in the row-major order format. After being asynchronously copied into 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 the matrix multiply-add operation in the register based on a third instruction to obtain a calculation result, and temporarily store the calculation result from the register into the shared memory. Write the calculation result in the shared memory back to the global memory based on a fourth instruction.

2. The method according to claim 1, characterized in that Storing in the BlockNM sparse format 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, 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 bits in the value array.

3. The method according to claim 1, characterized in that Asynchronously copying into the shared memory through a second instruction includes the following steps: The second instruction is an instruction that loads 128-bit data from the global memory each time. During the copying process, divide the sparse matrix into multiple BMxBK blocks along the K dimension, and divide the dense matrix into multiple BKxBN blocks, where BM = 64, BK = 32, BN is an integer multiple of the number of banks, and K = 32.

4. The method according to claim 3, wherein Loading the sparse matrix and the dense matrix in the shared memory into the register in batches includes the following steps: Divide the BMxBK block into multiple first sub-blocks with a size of 16x16, and load each first sub-block into the register in the row-major order format. Divide the BKxBN block into multiple second sub-blocks with a size of 16x8, perform a transpose operation on the data in each second sub-block, convert the second sub-block from the row-major order format to the column-major order format, and load the second sub-block into the register in the column-major order format.

5. The method according to claim 3, characterized in that Setting the four-stage pipeline scheduling strategy 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. 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 tensor core to execute the third instruction to complete the matrix multiply-accumulate calculation.

6. The method according to claim 1, wherein Setting an offset for each thread accessing the shared memory includes the following steps: 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 a new address for each thread is calculated based on the first formula , and the first formula is: , where is the original memory address of each thread without offset, , is the number of elements of the basic data type that can be accommodated in each bank, is an integer multiple of the number of banks, is the offset of each thread.

7. The method according to claim 4, characterized in that, Performing matrix multiply-accumulate operations on the registers based on the third instruction 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, which is a dense matrix pre-stored in the register.

8. The method according to claim 5, characterized in that, Temporarily storing the calculation result from the register to the shared memory 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.

9. The method according to claim 1, wherein 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 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, set each thread to write back 16 bytes at one time and align with the global memory address.

10. The method according to claim 1, wherein The first instruction is an instruction to load 64-bit data from the global memory each time.

Images