Sparse matrix multiplication acceleration method, system and equipment based on tensor processor

Abstract

The invention provides a sparse matrix multiplication acceleration method based on a tensor processor, which relates to the field of high-performance computing, and comprises the following steps: carrying out block processing on a sparse matrix to obtain sparse matrix blocks matched with the register granularity of the tensor processor; performing block processing on the dense matrix to obtain dense matrix blocks; generating a matrix multiplication task set based on the sparse matrix block and the dense matrix block; and distributing the task set to a plurality of threads, controlling the plurality of threads to respectively call the tensor processors to execute matrix multiplication operation, and aggregating calculation results of the threads to obtain an output matrix. Compared with the prior art, the sparse matrix blocks are processed into the sparse matrix blocks matched with the register granularity of the tensor processor, the index coding and decoding overhead, the reconstruction overhead and the communication overhead of the sparse matrix are reduced, and therefore the operation performance of sparse matrix multiplication based on the tensor processor is improved. The system has the same beneficial effects.

Description

Technical Field

[0001] This application relates to the field of high-performance computing, and in particular to a method, system, and device for accelerating sparse matrix multiplication based on a tensor processor. Background Technology

[0002] Sparse matrix-matrix multiplication (SpMM) is a crucial operation for multiplying a sparse matrix A with a dense matrix B, and it is widely used in machine learning (such as neural networks), scientific computing, and engineering simulation. With the evolution of CPU architecture, tensor processors (such as Intel AMX and NVIDIA TensorCore) are increasingly integrating matrix computation units specifically designed to accelerate dense matrix operations while maintaining compatibility with existing instruction sets. Taking the AMX extended instruction set integrated in Intel's latest processors as an example, this technology supports high-throughput mixed-precision matrix multiplication by introducing a block-based register file and a dedicated matrix multiplication unit, significantly accelerating tasks such as deep learning training and inference.

[0003] However, while tensor processors excel in dense matrix computations, applying them to SpMMs presents fundamental challenges: the sparse matrices involved in SpMMs are typically highly irregular, exhibiting non-uniform row lengths and scattered distributions of non-zero elements. This characteristic contradicts the architecture of AMX hardware, which is optimized for regular contiguous memory access patterns, leading to low utilization of computation blocks and registers. Directly mapping sparse computation tasks to AMX hardware not only significantly reduces computational efficiency but may even be inferior to traditional SIMD-based implementations (such as AVX-512) in some scenarios.

[0004] The existing SpMM transport method using tensor processors has the following problems: (1) high idle ratio of computing units; (2) significant additional overhead caused by index decoding and address reconstruction; (3) low cache reuse rate; (4) task contention and write conflict between multiple threads.

[0005] Therefore, how to provide a sparse matrix multiplication acceleration method and system based on tensor processors that can reduce index encoding / decoding overhead and reconstruction overhead is a technical problem that urgently needs to be solved by those skilled in the art. Summary of the Invention

[0006] To address the aforementioned technical problems, the present invention aims to provide a method for accelerating sparse matrix multiplication based on a tensor processor. This invention reduces the overhead of indexing and decoding, reconstruction, and communication of sparse matrices by dividing the sparse matrix into blocks that match the register granularity of the tensor processor, thereby improving the computational performance of sparse matrix multiplication based on a tensor processor.

[0007] The purpose of this invention is to provide a method for accelerating sparse matrix multiplication based on a tensor processor; The technical solution provided by this invention is as follows: A method for accelerating sparse matrix multiplication based on tensor processors includes the following steps: The sparse matrix is ​​divided into blocks to obtain sparse matrix blocks that match the register granularity of the tensor processor. The dense matrix is ​​divided into blocks to obtain dense matrix blocks; Based on the sparse matrix block and the dense matrix block, generate a matrix multiplication task set; The task set is distributed to multiple threads, and the multiple threads are controlled to call the tensor processor to perform matrix multiplication operations respectively. The calculation results of each thread are aggregated to obtain the output matrix.

[0008] Preferably, before assigning the task set to multiple threads, controlling the multiple threads to call the tensor processor to perform matrix multiplication operations respectively, and aggregating the calculation results of each thread to obtain the output matrix, the method further includes: Access frequency analysis is performed on the dense matrix blocks in the task set to identify hot dense matrix blocks, and cache prefetch optimization is performed.

[0009] Preferably, the step of performing access frequency analysis on the dense matrix blocks in the task set, identifying hotspot dense matrix blocks, and performing cache prefetch optimization includes the following steps: Count the number of times each dense matrix block is accessed in the task set; The dense matrix blocks whose access frequency exceeds a preset hotspot threshold are marked as hotspot blocks; Based on the cache capacity, the hot blocks are prefetched into the second or third cache in batches.

[0010] Preferably, the step of allocating the task set to multiple threads, controlling the multiple threads to call the tensor processor to perform matrix multiplication operations respectively, and aggregating the calculation results of each thread to obtain the output matrix specifically involves: A conflict-free task scheduling strategy is used to allocate the task set to multiple threads. These threads then call the tensor processor to perform matrix multiplication operations, and the calculation results from each thread are aggregated to obtain the output matrix. Preferably, when executing the method of allocating the task set to multiple threads using a conflict-free task scheduling strategy, controlling the multiple threads to call the tensor processor to perform matrix multiplication operations respectively, and aggregating the calculation results of each thread to obtain the output matrix, the specific steps include: The task set is allocated to multiple threads using a conflict-free task scheduling strategy. Each thread is controlled to load sparse blocks associated with the current batch of hot blocks, the tensor processor is called to perform matrix multiplication, the calculation results are temporarily stored in the local buffer of the corresponding thread, and after reaching the preset write-back threshold, they are written back to the main memory in batches and aggregated to obtain the output matrix.

[0011] Preferably, the task set further includes: an output matrix block index; The conflict-free task scheduling strategy includes: Pre-analyze the task set and identify the indices of all output matrix blocks; All tasks with the same output matrix block index are assigned to the same thread.

[0012] Preferably, the conflict-free task scheduling strategy further includes a locality-driven optimization strategy, which includes the following steps: Calculate the overlap of accesses to dense matrix blocks among the tasks in the task set; Tasks that meet the preset conditions for access overlap will be preferentially assigned to the same non-uniform memory access node.

[0013] Preferably, the conflict-free task scheduling strategy further includes a load balancing strategy, which includes the following steps: Sort the tasks according to the number of non-zero elements in the sparse block corresponding to the task; A greedy allocation algorithm is used to prioritize the allocation of sorted tasks to the thread with the lowest current load, so as to balance the workload of each thread.

[0014] The second objective of this invention is to provide a sparse matrix multiplication acceleration system based on a tensor processor; The technical solution provided by this invention is as follows: A sparse matrix multiplication acceleration system based on a tensor processor includes: The matrix block module is used to perform block processing on sparse matrices to obtain sparse matrix blocks that match the register granularity of the tensor processor, and to perform block processing on dense matrices to obtain dense matrix blocks. The task generation module is used to generate a matrix multiplication task set based on the sparse matrix block and the dense matrix block; The task scheduling and aggregation module is used to allocate the task set to multiple threads, control the multiple threads to call the tensor processor to perform matrix multiplication operations respectively, and aggregate the calculation results of each thread to obtain the output matrix.

[0015] The third objective of this invention is to provide a computer device; The technical solution provided by this invention is as follows: A computer device includes a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the steps of any of the sparse matrix multiplication acceleration methods described above.

[0016] This invention provides a method for accelerating sparse matrix multiplication based on a tensor processor, comprising the following steps: dividing a sparse matrix into blocks to obtain sparse matrix blocks matching the register granularity of the tensor processor; dividing a dense matrix into blocks to obtain dense matrix blocks; generating a matrix multiplication task set based on the sparse and dense matrix blocks; allocating the task set to multiple threads, controlling the multiple threads to call the tensor processor to perform matrix multiplication operations respectively, and aggregating the calculation results of each thread to obtain an output matrix. Compared with the prior art, this invention reduces the indexing and decoding overhead, reconstruction overhead, and communication overhead of sparse matrices by dividing the sparse matrix into blocks matching the register granularity of the tensor processor; this invention also reduces read / write overhead and improves cache hit rate through cache prefetching optimization, and reduces runtime management and memory access overhead while achieving load balancing through a conflict-free task scheduling strategy, thereby improving the computational performance of sparse matrix multiplication based on a tensor processor.

[0017] The present invention also provides a sparse matrix multiplication acceleration system based on tensor processors. Since this system and the sparse matrix multiplication acceleration method based on tensor processors solve the same technical problem and belong to the same technical concept, they should have the same beneficial effects, and will not be described in detail here. Attached Figure Description

[0018] To more clearly illustrate the technical solutions in the embodiments of this application or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of this application. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0019] Figure 1 This is a flowchart of a sparse matrix multiplication acceleration method based on a tensor processor, as described in an embodiment of the present invention. Figure 2This is a comparison chart of the throughput of the method of the present invention with that of oneDNN and MKL on the CPU; Figure 3 This is a comparison chart of the L2 cache hit rate of the method of this invention and oneDNN; Figure 4 This is a schematic diagram of a sparse matrix multiplication acceleration system based on a tensor processor, as described in an embodiment of the present invention. Figure 5 This is an internal structural diagram of a computer device according to an embodiment of the present invention. Detailed Implementation

[0020] To enable those skilled in the art to better understand the technical solutions in this application, the technical solutions in the embodiments of this application will be clearly and completely described below. Obviously, the described embodiments are only a part of the embodiments of this application, and not all of the embodiments. Based on the embodiments in this application, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this application.

[0021] like Figure 1 As shown, this embodiment of the invention provides a method for accelerating sparse matrix multiplication based on a tensor processor, including the following steps: S1. Divide the sparse matrix into blocks to obtain sparse matrix blocks that match the register granularity of the tensor processor. It should be noted that in the AMX architecture, the basic operation granularity of matrix multiplication is fixed at the tile block size. Based on this, step S1 proposes a sparse coding method closely adapted to the AMX architecture, encoding and reassembling the input matrix at the tile register size granularity. Step S1 specifically includes the following steps: (1) Block partitioning The input sparse matrix A[M,K] is divided into several sparse matrix sub-blocks (e.g., 16×32) according to the tensor processor register size. Each sub-block is called a "tile block" and contains several consecutive elements.

[0022] (2) Block index generation For each non-zero block, generate the following metadata: a. Block line index: stores the block line in which the block belongs; b. In-row offset: Stores the column index of the block in this row; c. Block memory starting address: Points to the starting address of all elements in the block stored contiguously in memory.

[0023] (3) Sparse coding structure The metadata of the sparse matrix blocks described above is stored row-by-row to form a block-level sparse index table. This structure has a constant block size and preserves the complete tile layout to accommodate register operations of tensor processors, even if the blocks contain some zero elements.

[0024] Unlike traditional element-oriented sparse coding such as CSR / COO, this invention uses sparse matrix blocks represented by tiles. Even after partitioning, if a block contains zero elements, it is still stored with a fixed layout of complete tiles. This "redundant storage" may seem to increase zero-padding, but it offers the following advantages: 1. Sparse blocks can be directly aligned with AMX instructions, avoiding complex boundary handling logic. 2. Sparse blocks are stored contiguously in memory as rows, which facilitates sequential multi-threaded partitioning and cache access in subsequent work. 3. Reduce the memory overhead of encoding and decoding. Taking a sparse matrix A[M,K] in bf16 format as an example, the memory overhead required for CSR encoding is (MN+M+1)*idx, which includes the row pointers of (M+1)*idx and the column indices of nnz*idx, where idx is the number of index bytes, N is the number of columns of dense matrix B, and nnz is the number of non-zero elements; while the memory required for AMX sparse encoding to store the index of this matrix is ​​1 / 16*32 of the former.

[0025] Compared to existing technologies, step S1 proposes a sparse matrix block format with granularity equal to the Tile register size of the tensor processor. Each sparse matrix block has a fixed size (the block size is constant, and the complete Tile layout is preserved even if the block contains some zero elements). Through block-level indexing and a fixed-size storage structure, the sparse matrix can be directly mapped to the tensor processor hardware execution unit, avoiding dynamic indexing and decoding, and reducing the indexing and decoding overhead, reconstruction overhead, and communication overhead of the sparse matrix.

[0026] S2. Divide the dense matrix into blocks to obtain dense matrix blocks; In practical applications, the block processing in step S2 involves dividing the dense matrix into blocks with the same structure as the sparse matrix (e.g., 16×32). It is also necessary to generate the metadata of the dense matrix blocks. The generation of the metadata of the dense matrix blocks is the same as that of the sparse matrix blocks, which will not be elaborated here.

[0027] S3. Generate a matrix multiplication task set based on sparse and dense matrix blocks; It should be noted that step S3 divides the entire sparse matrix multiplication task into multiple subtasks, each subtask corresponding to a pair of non-zero block multiplications. Each task in the task set includes: sparse matrix block element information, dense matrix block element information, and output matrix block index. First, the sparse matrix A and dense matrix B are divided into blocks according to steps 1 and 2 respectively, and then their block-level adjacency matrices a and b are extracted. Multiplying a and b yields the corresponding output matrix block c. By tracing the non-zero block positions in the corresponding output matrix block c, and the intermediate indices involved in each non-zero block, we can deduce all task sets Aij, Bjk, Cik.

[0028] In practical applications, generating matrix task sets is not simply a matter of packaging and merging the sparse encodings and corresponding contents of a pair of sparse matrix blocks A and dense matrix blocks B. Instead, A and B are first divided into blocks, and then their block-level adjacency matrices a and b are extracted. Multiplying a and b yields the corresponding output matrix block c. By tracing the non-zero block positions in the corresponding output matrix block c and the intermediate indices involved in each non-zero block, all task sets Aij, Bjk, and Cik are deduced. This ensures that the column indices of A and the row indices of B in the task are naturally consistent, thus completely eliminating redundant index storage. In this way, the entire sparse matrix multiplication operation is "flattened" into a linear task array, which can be scheduled as needed.

[0029] Assuming the dense matrix B has a size of [4*4] after being partitioned into blocks, the corresponding computation task set is:

[0030] Where B[0] represents the 0th column of the dense matrix B. Before calculation, if it is a single core and single thread, all calculation tasks are sent to the AMX calculation unit in ascending order; if it involves multiple threads, these task sets are further divided and then sent to the corresponding threads. This design brings two benefits: (1) Reduced storage overhead: Only the necessary row / column information is retained in the data structure of each task, avoiding unnecessary index storage and reducing the index overhead by half; (2) Improved execution efficiency: No additional index matching and checking is required at runtime, and the task can be directly mapped to the tile-mma instruction call of AMX. The main purpose is to avoid calculating index matching during runtime. The sparse structure of this invention runs on the CPU side. Unlike GPUs or dedicated accelerators, the memory resources on the CPU side are relatively abundant. Therefore, the index structure design is mainly guided by computational efficiency.

[0031] Preferably, before step S5, the method further includes: S4. Analyze the access frequency of dense matrix blocks in the task set, identify hot dense matrix blocks, and perform cache prefetch optimization.

[0032] It's important to note that in sparse matrix multiplication, dense matrix blocks are often accessed repeatedly by multiple sparse matrix blocks. Within a task set, there may be intersections between non-zero block offsets within sparse matrix blocks, meaning that identical dense matrix blocks will participate in subsequent computations. Residenting dense matrix block data in a cache allows for faster exchange with AMX registers. By resident overlapping hotspot blocks in the cache in batches and calculating their related parts, the total read overhead of dense hotspot blocks can be reduced.

[0033] Step S4 includes the following steps: A1. Count the number of times each dense matrix block is accessed in the task set; A2. Mark dense matrix blocks whose access frequency exceeds the preset hotspot threshold as hotspot blocks; In practical applications, the number of times the dense matrix B blocks in each thread task set are accessed is counted and recorded as the hot value. B blocks with high access frequency are marked as "hot blocks"; B blocks with low access frequency are marked as "non-hot blocks".

[0034] A3. Based on the cache capacity, prefetch hot blocks into the second or third cache in batches.

[0035] It should be noted that step A3 specifically involves dividing hot blocks into several batches based on cache capacity. Hot blocks within each batch are preferentially loaded into the L2 / L3 cache and participate in relevant calculations.

[0036] In practical application, assuming that the Tile block is the basic granularity, the size of the sparse matrix A and the dense matrix B are both [4,4]. The specific steps based on step S4 are as follows: (1) Calculate the intersection of the non-zero ordinates of each pair of sparse matrix blocks in the task set of the thread, and take the union of these intersections. They correspond to the dense matrix blocks that need to be read more than once, and are called hot spots; (2) Calculate the intersection of the non-zero ordinates of each sparse matrix block in the thread with the hot spots, so that each sparse matrix block is divided into the part that "intersects with the hot spots" and the complement part; (3) According to the cache capacity, the hot spots data are divided into several batches, each batch containing several rows of complete hot spots data; the data corresponding to the complement part is also divided into several batches, each batch also containing several rows of complete non-hot spots data; (4) When calculating, first calculate the i-th Prefetch the batch of hot spot block data and load it into the cache; then, in turn, send all sparse matrix blocks and corresponding output matrix blocks that have intersection with the batch of hot spot block data into AMX for calculation together with these hot spot block data; (5) Repeat step 4 until all hot spot blocks have been traversed; then start to traverse non-hot spot blocks in batches, prefetch the m-th batch of non-hot spot block data and load it into the cache; then send the sparse matrix blocks and output matrix blocks corresponding to them into AMX in turn for calculation.

[0037] Preferably, the batch prefetching depth in step A3 is dynamically adjusted: The prefetch depth D is determined by the following formula: D = [t_mem / t_task] + margin Where t_mem is the memory access latency, t_task is the average task computation time, and the parameter margin∈[1,3].

[0038] Compared to existing technologies, the cache prefetching optimization of this invention is a caching method that prefetches and reuses hot blocks. By analyzing the access frequency of dense matrix blocks in the task set, hot blocks are identified, and a cache prefetching optimization strategy is applied to these hot blocks, prioritizing their storage in the cache to improve their cache reuse rate. Furthermore, a prefetch depth self-adjustment algorithm is proposed, which automatically adjusts the prefetch batch based on the task execution time to memory latency ratio, thereby hiding memory access latency, enabling prefetching to run parallel to the computation pipeline, and reducing average memory wait time.

[0039] S5. Distribute the task set to multiple threads, control the multiple threads to call the tensor processor to perform matrix multiplication operations, and aggregate the calculation results of each thread to obtain the output matrix.

[0040] Preferably, step S5 specifically includes: A conflict-free task scheduling strategy is adopted to allocate the task set to multiple threads, control multiple threads to call the tensor processor to perform matrix multiplication operations, and aggregate the calculation results of each thread to obtain the output matrix.

[0041] Preferably, when executing a conflict-free task scheduling strategy to allocate the task set to multiple threads, controlling the multiple threads to call the tensor processor to perform matrix multiplication operations, and aggregating the calculation results of each thread to obtain the output matrix, the specific steps include: B1. Use a conflict-free task scheduling strategy to distribute the task set to multiple threads; B2. Control each thread to load the sparse blocks associated with the current batch of hot blocks, call the tensor processor to perform matrix multiplication, temporarily store the calculation results in the local buffer of the corresponding thread, and write them back to the main memory in batches after reaching the preset write-back threshold, and aggregate to obtain the output matrix.

[0042] As one implementation method, when the number of calculation results or the number of task batches reach a preset write-back threshold, the operation of writing the calculation results back to main memory in batches is automatically triggered.

[0043] In practical application, during step B1, the task allocation phase, after determining all rows of tasks assigned to a thread, calculates the intersection of non-zero y-coordinates between each pair of sparse matrix blocks. The union of these intersections corresponds to the dense matrix blocks that need to be read more than once, denoted as hotspot blocks. The intersection of the non-zero y-coordinates of each sparse matrix block in the thread with the hotspot blocks is calculated, thus dividing each sparse matrix block into a hotspot block row intersection part and a complement part. Based on the cache capacity, the hotspot block data is divided into several batches, each batch containing several rows of complete hotspot block data. The data corresponding to the complement part is also divided into several batches, each batch also containing several rows of complete non-hotspot block data. During calculation, the i-th batch of hotspot block data is first prefetched and loaded into the cache. Then, all sparse matrix blocks that intersect with the hotspot blocks in this batch and their corresponding output matrix blocks are sequentially fed into AMX for matrix multiplication. The intermediate results of the output matrix blocks are then written back to memory from the AMX register. Step 4 is repeated. The process continues until all hot blocks have been traversed. Then, non-hot blocks are traversed in batches. The data of the m-th batch of non-hot blocks is prefetched and loaded into the cache. Then, the sparse matrix blocks and output matrix blocks corresponding to them are fed into AMX in turn for calculation. Finally, the calculation result of the C block is written back to memory from the AMX register.

[0044] Preferably, the task set further includes: an output matrix block index; Conflict-free task scheduling strategies include: Pre-analysis task set, identifying the indices of all output matrix blocks; All tasks with the same output matrix block index are assigned to the same thread.

[0045] It should be noted that sparse matrix A and dense matrix B only involve read operations, while only the contents of output matrix C involve both reading and writing. Furthermore, only a certain output matrix block Cij may involve write conflicts during the calculation and accumulation of intermediate results. In other words, only the process of calculating the sum of Ai*Bj involves thread conflicts. Therefore, based on the conflict-free task scheduling strategy, when allocating task sets to different threads, the following principle is followed: task sets whose calculation results are the same output matrix block Cij, i.e., tasks with the same output matrix block index, are assigned to the same thread. In this way, the calculation and accumulation of intermediate results for each output matrix block Cij can be guaranteed to have a strict order, thereby ensuring the correctness of the calculation results. At the same time, this strategy avoids frequent and costly operations such as switching mutex locks or maintaining synchronization queues.

[0046] Preferably, the conflict-free task scheduling strategy further includes a locality-driven optimization strategy, which includes the following steps: Calculate the overlap of accesses to dense matrix blocks among the tasks in the task set; Tasks that meet the preset conditions for access overlap will be preferentially assigned to the same non-uniform memory access node.

[0047] It should be noted that by combining the degree of overlap in access to dense matrix blocks between tasks, tasks with high overlap are assigned to the same non-uniform memory access node (NUMA node) or core cluster for execution, thereby significantly reducing remote memory access.

[0048] Preferably, the conflict-free task scheduling strategy further includes a load balancing strategy, which includes the following steps: Sort the tasks according to the number of non-zero elements in the sparse block corresponding to the task; A greedy allocation algorithm is used to prioritize the allocation of sorted tasks to the thread with the lowest current load, so as to balance the workload of each thread.

[0049] Compared with existing technologies, this invention reduces the overhead of indexing and decoding, reconstruction, and communication of sparse matrices by dividing the sparse matrix into blocks that match the register granularity of the tensor processor. Furthermore, this invention reduces hotspot block read / write overhead and improves cache hit rate through hotspot block identification and cache prefetch optimization strategies. Dynamic prefetching and write-back control strategies dynamically adjust the prefetch depth based on task computation time and memory access latency. When a local result cache or task batch reaches a threshold, a write-back operation to main memory is automatically triggered, achieving memory access latency hiding and a high cache hit rate. Finally, this invention ensures the accuracy of matrix multiplication result transportation through a conflict-free task scheduling strategy, avoiding frequent and costly operations such as switching mutex locks or maintaining synchronization queues. This achieves load balancing while reducing runtime management and memory access overhead, thereby improving the computational performance of sparse matrix multiplication based on tensor processors.

[0050] This paper also demonstrates the effectiveness of the proposed method using experimental data. First, the experimental training platform setup is introduced, followed by detailed experimental comparison results and corresponding visualized quantitative experimental results.

[0051] (1) Baseline: First, the method of this application is defined as SparX, and SparX is compared with the following two works: (1) oneDNN-SpMM, which provides a highly optimized kernel for x86 CPUs. To our knowledge, it is the only publicly available implementation that natively supports SpMM AMX instructions. The operator is configured with the default SpMM algorithm and CSR sparse format. (2) MKL is a widely adopted numerical computation library that provides state-of-the-art dense and sparse linear algebra routines. The SpMM kernel in oneMKL is heavily optimized for AVX-512 and is considered one of the most powerful CPU benchmarks for sparse computation.

[0052] (2) Dataset: We used SuiteSparse, a dataset commonly used in the sparse domain. The SuiteSparse collection is a widely cited compilation of sparse matrix benchmarks from various applications. Furthermore, the dense matrix B, another operand, has randomly generated content, and our experiments covered multiple column counts, ranging from 128 to 1024. The data types we tested were A and B: bf16, and C: fp32.

[0053] (3) Experimental environment: We conducted our experiments on a dual-socket Intel Xeon Gold 6434 server. Each socket contained 8 cores (16 threads) at a clock speed of 3.7 GHz, for a total of 16 cores and 32 hardware threads. The processor was equipped with 768KB L1d, 512KB L1i, 32 MB L2, and 45 MB shared L3 cache. The system contained 2 NUMA nodes and supported 64-bit instruction sets. All experiments were warmed up and run 100 times, with the average value taken. We used two CPUs with different performance levels and calculated their end-to-end speedup ratio.

[0054] We measured the average throughput with N set to 128, 256, and 512 (the number of columns in matrix B). Figure 2The throughput distribution of the three methods on the CPU is shown. For different matrix sizes, the throughput performance of SparX, oneDNN, and MKL differs significantly. For N=128, SparX has a wider throughput range, from 0 to approximately 900 GFLOPS, with about 20% of matrices exceeding 90 GFLOPS, while most medium-sized matrices are concentrated in the 34–90 GFLOPS range. oneDNN's throughput is concentrated around 0–20 GFLOPS, but a few large matrices can reach 19 GFLOPS. MKL's throughput is generally lower than SparX, concentrated in the 0–10 GFLOPS range, with only a few matrices exceeding 8 GFLOPS. It can be seen that for small column-count matrices, SparX can still maintain a high throughput for most matrices, while the throughput of traditional methods is significantly limited by memory access. For N=256, SparX's throughput further increases, with the median matrix ranging from 38 to 100 GFLOPS, reaching a maximum of approximately 794 GFLOPS. The throughput ranges of oneDNN and MKL remain relatively stable between 0–20 GFLOPS and 0–10 GFLOPS, indicating that SparX significantly improves its efficiency in utilizing medium and large-scale matrices as the number of columns increases, while traditional methods are limited by the slower increase in computational density. For N=512, SparX achieves throughput exceeding 100 GFLOPS on most matrices, peaking at nearly 900 GFLOPS, while throughput for small and medium-sized matrices remains in the 17–125 GFLOPS range. The throughput of oneDNN and MKL remains stable around 0–20 GFLOPS, but the throughput for some large matrices slightly increases due to the increase in N. Overall, SparX fully leverages the high computational density of AMXtile instructions on large matrix lengths, while oneDNN and MKL offer relatively limited throughput improvements for large matrices. In summary, the throughput distributions of the three methods differ significantly across different N values. SparX performs exceptionally well on medium to large-scale matrices, fully utilizing AMX's computational resources; the throughput of oneDNN and MKL is limited by traditional instruction sets and memory bandwidth, resulting in a more concentrated and stable distribution, especially on small matrix lengths where their throughput is lower than SparX.

[0055] Table 1 also shows the geometric mean speedup of SparX relative to oneDNN (AMX) and MKL (AVX-512) on the SuiteSparse dataset for N=512. Overall, SparX provides a significant performance improvement for most matrices, with only a small subset showing a decrease in speedup. These cases primarily correspond to smaller matrices or matrices A with relatively few rows, where the overhead of block-based encoding and dynamic scheduling outweighs the computational gain.

[0056] Table 1: Overall performance comparison of SparX with oneDNN and MKL

[0057] Furthermore, according to overall speedup statistics, for matrices with average sizes of [5200, 6091] and [5968, 6927], SparX's speedup is less than twice that of oneDNN and MKL. This is mainly due to differences in operator characteristics. AMX excels at large-scale tensor computations (when M, N, and K are all large), significantly improving computational intensity through tile register blocking and high-throughput matrix-matrix multiplication instructions. However, when M is small (i.e., fewer rows in A), even with a large N / K ratio, the overall computational workload (FLOPs) is insufficient to fully utilize AMX's registers and execution pipeline. Therefore, AMX's instruction-level parallelism is not fully utilized. For small-scale computations, memory bottlenecks dominate, and small-scale SpMM matrices are more prone to memory constraints. The memory access amplification effect of AMX tile register write-back and B row block prefetching is more pronounced. In this scenario, performance is limited by memory bandwidth rather than peak computing power, where AV X512... Lightweight instructions like V NNI are more efficient because they occupy fewer registers and caches, thus reducing memory access pressure. However, there are scheduling and amortization overheads. AMX tile instructions require tile configuration, loading, and write-back, resulting in greater scheduling overhead than single-instruction AVX512. These overheads cannot be amortized for smaller tasks, leading to a decrease in SparX's effective speedup. oneDNN's SpMM operator has almost no scheduling and architecture adaptation overhead, directly using the CSR compression format as the front end and the GEMM kernel as the back end. OneMKL, on the other hand, is very mature in AVX512 tuning. For small matrices, it typically provides performance close to the memory bandwidth limit, thus achieving better performance than SparX within this range.

[0058] Multithreaded performance The experimental data in Table 2 shows the geometric mean speedup of SparX and MKL. In single-threaded mode, SparX is already about 9 times faster than MKL, indicating a significant advantage even in single-threaded scenarios, rather than only showing performance improvement in multi-threaded environments. The speedup increases slightly with the number of threads, peaking at 4 threads (approximately 10.9 times), then decreasing to about 10 times with 16 threads, generally remaining in the 9–11 times range. This slight decrease from 8 to 16 threads is likely due to the increased non-local memory access overhead introduced when running across NUMA nodes, thus affecting the overall speedup. In contrast, the speedup of SparX and oneDNN is highest at 2–8 threads, exceeding 12 times, but drops to about 8 times with 16 threads, suggesting that oneDNN may have a bottleneck in AMX scheduling or additional parallel overhead at high thread counts. The confidence intervals are also relatively wide. For example, in single-threaded scenarios, the CI of SparX relative to oneDNN ranges from approximately 1.9 to 26 times, indicating that matrix sparsity and size have a significant impact on the speedup. Overall, SparX's advantages are more pronounced in large matrix and multi-threaded scenarios.

[0059] Table 2 shows the multithreaded efficiency of different libraries based on the experimental data. SparX's efficiency drops to approximately 0.76 from single-threaded to 16-threaded, demonstrating good scalability. oneDNN's efficiency is 0.70 with the same number of threads, showing a significantly faster decline than SparX. This is presumably because oneDNN's sparse multiplication operator lacks in-depth multithreaded optimization, leading to insufficient load balancing and cache utilization among threads. MKL exhibits the worst scalability, with an efficiency of 0.68 at 16 threads, mainly limited by the AVX512 SIMD bit width and thread scheduling mechanism. Other reasons for the efficiency decline include cache contention, memory bandwidth limitations, and the difficulty in fully utilizing AMX resources for some small-scale matrices. Overall, SparX outperforms MKL and oneDNN in multithreaded scalability and demonstrates a stable performance advantage in large-scale sparse matrix scenarios.

[0060] Table 2. Performance Comparison of SparX, MKL, and oneDNN under Different Numbers of Threads

[0061] Cache optimization performance 1) The effectiveness of caching optimization and prefetching: Figure 3This demonstrates the speedup of SparX under different configurations: no caching optimization, caching only, prefetch only, and full integration of the caching optimizer and prefetcher. On the suitesparse dataset, the fully integrated scheme consistently achieves the highest speedup, exceeding 3.0x on most datasets, reflecting the complementarity between the two optimization modules. The caching only scheme achieves a modest improvement with an average speedup of 1.7x, while the prefetch only scheme has a smaller average speedup of only 1.3x.

[0062] It is worth noting that for smaller matrices or extremely sparse datasets, the prefetch-only approach sometimes outperforms the caching-only approach. This is because small matrices can often be completely cached, reducing the effectiveness of the cache optimizer's hot B-row strategy. In extremely sparse matrices, the read reuse provided by available hot B-rows is almost non-existent. In this case, prefetching can still relatively improve performance by hiding memory latency by asynchronously loading B-blocks before computation. Overall, the results show that the caching optimizer primarily improves the performance of medium to large matrices by reducing repeated accesses to B-rows, while the prefetcher effectively reduces memory latency for small or irregular matrices. The combination of the two strategies yields the most significant speedup, confirming that these optimizations are complementary and together maximize SparX's throughput across a variety of sparse matrix workloads.

[0063] 2) Cache hit rate optimization: We randomly selected 30 representative matrices and compared the L2 cache hit rates of SparX and oneDNN-SPMM, such as... Figure 5As shown in the bar chart, SparX consistently maintains a high L2 cache utilization on most randomly sampled datasets, while oneDNN's L2 cache hit rate is relatively low and fluctuates significantly. The L2 cache is a crucial mediator between the fast but limited L1 cache and the slower but larger L3 cache. High L2 hit rates are essential for maintaining throughput and reducing memory latency during sparse matrix multiplications, as they help buffer frequently accessed data that cannot fully reside in L1. SparX's improved L2 utilization indicates that the operator effectively utilizes temporal and spatial locality, reduces redundant main memory accesses, and minimizes pause cycles caused by cache misses. To more fully demonstrate the effectiveness of the SparX operator in terms of cache utilization, we compared the L1, L2, and L3 cache hit rates with oneDNN NSPMM. As shown in Table 3, SparX consistently maintains a higher cache hit rate across all levels. Specifically, AMX achieved an average L1 hit rate of 96.8%, 7.7 percentage points higher than oneDNN's 89.2%, with geometric mean values ​​of 96.8% and 89.0%, respectively. Similarly, AMX's L2 hit rate was 92.9%, 7.3 percentage points higher than oneDNN's 85.7%, and its L3 hit rate was 75.7%, 5.9 percentage points higher than oneDNN's 69.8%.

[0064] Table 3: Comparison of Cache Hit Rates for SparX and oneDNN

[0065] Furthermore, SparX exhibits extremely high cache utilization: over 97.5% of the datasets show higher hit rates across all cache levels than oneDNN, with only a small portion (approximately 2.5%) of the datasets primarily using oneDNN. The proportion of datasets with similar hit rates (within ±5%) gradually increases from L1 to L3 (33.1% → 48.0%), indicating that although the difference in L3 caching is not significant, AMX consistently maintains high efficiency at both L1 and L2.

[0066] The observed improvements demonstrate that SparX effectively leverages a hierarchical caching architecture, reducing memory access latency and enhancing temporal locality. These results highlight the advantages of SparX's optimization for sparse matrix multiplication, particularly its ability to improve upper-level cache utilization while maintaining performance comparable to L3 cache.

[0067] in conclusion: Experimental results on representative large-scale irregular sparse matrices demonstrate that our SparX implementation significantly outperforms existing CPU implementations, achieving an 8.76x end-to-end geometry speedup compared to oneDNN and a 10.17x end-to-end geometry speedup compared to MKL. Furthermore, our method improves the average L2 cache hit rate by 7.3%, thereby reducing memory access overhead and increasing throughput.

[0068] In one embodiment, such as Figure 4 As shown, a sparse matrix multiplication acceleration system based on a tensor processor includes: The matrix block module is used to perform block processing on sparse matrices to obtain sparse matrix blocks that match the register granularity of the tensor processor, and to perform block processing on dense matrices to obtain dense matrix blocks. The task generation module is used to generate a set of matrix multiplication tasks based on sparse and dense matrix blocks; The task scheduling and aggregation module is used to allocate task sets to multiple threads, control multiple threads to call the tensor processor to perform matrix multiplication operations, and aggregate the calculation results of each thread to obtain the output matrix.

[0069] Preferably, a sparse matrix multiplication acceleration system based on a tensor processor further includes: The cache optimization module is used to analyze the access frequency of dense matrix blocks in the task set, identify hot dense matrix blocks, and perform cache prefetch optimization.

[0070] Specific limitations regarding the tensor processor-based sparse matrix multiplication acceleration system can be found in the above description of the tensor processor-based sparse matrix multiplication acceleration method, and will not be repeated here. Each module in the aforementioned tensor processor-based sparse matrix multiplication acceleration system can be implemented entirely or partially through software, hardware, or a combination thereof. These modules can be embedded in hardware or independent of the processor in a computer device, or stored in software in the computer device's memory, facilitating processor execution of the corresponding operations.

[0071] like Figure 5 As shown, in one embodiment, a computer device is provided, including a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, it performs the following steps: The sparse matrix is ​​divided into blocks to obtain sparse matrix blocks that match the register granularity of the tensor processor. The dense matrix is ​​divided into blocks to obtain dense matrix blocks; Generate a matrix multiplication task set based on sparse and dense matrix blocks; The task set is distributed to multiple threads, and each thread calls the tensor processor to perform matrix multiplication. The calculation results of each thread are then aggregated to obtain the output matrix.

[0072] Those skilled in the art will understand that Figure 3 The structure shown is merely a block diagram of a portion of the structure related to the present application and does not constitute a limitation on the computer device to which the present application is applied. Specific computer devices may include more or fewer components than those shown in the figure, or combine certain components, or have different component arrangements.

[0073] In the embodiments provided in this application, it should be understood that the disclosed methods and systems can be implemented in other ways. The system embodiments described above are merely illustrative. For example, the division of modules is only a logical functional division, and in actual implementation, there may be other division methods, such as: multiple modules or components can be combined, or integrated into another system, or some features can be ignored or not executed. In addition, the coupling, direct coupling, or communication connection between the various components shown or discussed can be through some interfaces, indirect coupling or communication connection of devices or modules, and can be electrical, mechanical, or other forms.

[0074] Furthermore, in the various embodiments of the present invention, each functional module can be fully integrated into a processor, or each module can be a separate device, or two or more modules can be integrated into a device; each functional module in the various embodiments of the present invention can be implemented in hardware or in the form of hardware plus software functional units.

[0075] Those skilled in the art will understand that all or part of the steps of the above method embodiments can be implemented by program instructions and related hardware. The aforementioned program instructions can be stored in a computer-readable storage medium. When the program instructions are executed, they perform the steps of the above method embodiments. The aforementioned storage medium includes various media that can store program code, such as mobile storage devices, read-only memory (ROM), magnetic disks, or optical disks.

[0076] It should be understood that the use of terms such as "system," "apparatus," "unit," and / or "module" in this application is only applicable to distinguishing different components, elements, parts, sections, or assemblies at different levels. However, if other terms can achieve the same purpose, they may be replaced by other expressions.

[0077] As indicated in this application and claims, unless the context clearly indicates otherwise, the words "a," "an," "a," and / or "the" are not specifically singular and may include the plural. Generally, the terms "comprising" and "including" only indicate the inclusion of expressly identified steps and elements, which do not constitute an exclusive list, and the method or apparatus may also include other steps or elements. An element defined by the phrase "comprising an..." does not exclude the presence of other identical elements in the process, method, product, or apparatus that includes the element.

[0078] Furthermore, the terms "first" and "second" are used for descriptive purposes only and should not be construed as indicating or implying relative importance or implicitly specifying the number of technical features indicated. Thus, a feature defined as "first" or "second" may explicitly or implicitly include one or more of that feature. In the description of this application, "a plurality of" or "several" means two or more, unless otherwise explicitly specified.

[0079] If a flowchart is used in this application, it is used to illustrate the operations performed by the system according to embodiments of this application. It should be understood that the preceding or following operations are not necessarily performed in exact order. Instead, the steps can be processed in reverse order or simultaneously. Furthermore, other operations can be added to these processes, or one or more steps can be removed from them.

[0080] The above description of the disclosed embodiments enables those skilled in the art to make or use the invention. Various modifications to these embodiments will be readily apparent to those skilled in the art, and the general principles defined herein may be implemented in other embodiments without departing from the spirit or scope of the invention. Therefore, the invention is not to be limited to the embodiments shown herein, but is to be accorded the widest scope consistent with the principles and novel features disclosed herein.

Claims

1. A sparse matrix multiplication acceleration method based on tensor processors, characterized in that, Includes the following steps: The sparse matrix is ​​divided into blocks to obtain sparse matrix blocks that match the register granularity of the tensor processor. The dense matrix is ​​divided into blocks to obtain dense matrix blocks; Based on the sparse matrix block and the dense matrix block, generate a matrix multiplication task set; The task set is distributed to multiple threads, and the multiple threads are controlled to call the tensor processor to perform matrix multiplication operations respectively. The calculation results of each thread are aggregated to obtain the output matrix.

2. The sparse matrix multiplication acceleration method according to claim 1, characterized in that, Before allocating the task set to multiple threads, controlling the multiple threads to call the tensor processor to perform matrix multiplication operations respectively, and aggregating the calculation results of each thread to obtain the output matrix, the method further includes: Access frequency analysis is performed on the dense matrix blocks in the task set to identify hot dense matrix blocks, and cache prefetch optimization is performed.

3. The sparse matrix multiplication acceleration method according to claim 2, characterized in that, The process of analyzing the access frequency of the dense matrix blocks in the task set, identifying hotspot dense matrix blocks, and performing cache prefetch optimization includes the following steps: Count the number of times each dense matrix block is accessed in the task set; The dense matrix blocks whose access frequency exceeds a preset hotspot threshold are marked as hotspot blocks; Based on the cache capacity, the hot blocks are prefetched into the second or third cache in batches.

4. The sparse matrix multiplication acceleration method according to claim 3, characterized in that, The process involves allocating the task set to multiple threads, controlling each thread to call the tensor processor to perform matrix multiplication, and aggregating the calculation results of each thread to obtain the output matrix. Specifically: A conflict-free task scheduling strategy is used to allocate the task set to multiple threads, and the multiple threads are controlled to call the tensor processor to perform matrix multiplication operations respectively. The calculation results of each thread are aggregated to obtain the output matrix.

5. The sparse matrix multiplication acceleration method according to claim 4, characterized in that, When executing the process of allocating the task set to multiple threads using a conflict-free task scheduling strategy, controlling the multiple threads to call the tensor processor to perform matrix multiplication operations respectively, and aggregating the calculation results of each thread to obtain the output matrix, the specific steps include: The task set is allocated to multiple threads using a conflict-free task scheduling strategy. Each thread is controlled to load sparse blocks associated with the current batch of hot blocks, the tensor processor is called to perform matrix multiplication, the calculation results are temporarily stored in the local buffer of the corresponding thread, and after reaching the preset write-back threshold, they are written back to the main memory in batches and aggregated to obtain the output matrix.

6. The sparse matrix multiplication acceleration method according to claim 5, characterized in that, The task set also includes: output matrix block index; The conflict-free task scheduling strategy includes: Pre-analyze the task set and identify the indices of all output matrix blocks; All tasks with the same output matrix block index are assigned to the same thread.

7. The sparse matrix multiplication acceleration method according to claim 6, characterized in that, The conflict-free task scheduling strategy also includes a locality-driven optimization strategy, which includes the following steps: Calculate the overlap of accesses to dense matrix blocks among the tasks in the task set; Tasks that meet the preset conditions for access overlap will be preferentially assigned to the same non-uniform memory access node.

8. The sparse matrix multiplication acceleration method according to claim 7, characterized in that, The conflict-free task scheduling strategy also includes a load balancing strategy, which includes the following steps: Sort the tasks according to the number of non-zero elements in the sparse block corresponding to the task; A greedy allocation algorithm is used to prioritize the allocation of sorted tasks to the thread with the lowest current load, so as to balance the workload of each thread.

9. A sparse matrix multiplication acceleration system based on a tensor processor, characterized in that, include: The matrix block module is used to perform block processing on sparse matrices to obtain sparse matrix blocks that match the register granularity of the tensor processor, and to perform block processing on dense matrices to obtain dense matrix blocks. The task generation module is used to generate a matrix multiplication task set based on the sparse matrix block and the dense matrix block; The task scheduling and aggregation module is used to allocate the task set to multiple threads, control the multiple threads to call the tensor processor to perform matrix multiplication operations respectively, and aggregate the calculation results of each thread to obtain the output matrix.

10. A computer device, characterized in that, The method includes a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor, when executing the computer program, implements the steps of the sparse matrix multiplication acceleration method as described in any one of claims 1-8.

Images