Source-linked AI summary
Performance Optimization of Deep Learning Sparse Matrix Kernels on Intel Max Series GPU
Mohammad Zubair, Christoph Bauinger
TL;DR
Efficient GPU kernels for machine-learning sparse operations are challenging to optimize near peak with implicit-SIMD programming. This paper implements SPMM, SDDMM, and FusedMM with Intel ESIMD. The implementations outperform oneMKL and a V100 reference, reaching up to 3.4x and 10x, respectively.
Problem
Efficient GPU implementations of machine-learning sparse operations are needed, but implicit-SIMD CUDA or SYCL makes optimization near peak difficult.
Source: §1 · pp. 1–2Method
The paper develops Intel GPU implementations of SPMM, SDDMM, and FusedMM using explicitly vectorized ESIMD kernels.
Source: §1 · p. 2Results
Takeaways & Limitations
FusedMM avoids explicitly storing the SDDMM result and can perform better than applying SDDMM and SPMM sequentially.
Source: FusedMM · p. 3Evaluation focuses on one Intel GPU stack; cross-stack links are slower than HBM and uncached, so the authors recommend explicit two-stack scaling.
Source: §2.2 · p. 4
Abstract
from arXiv · showhide
In this paper, we focus on three sparse matrix operations that are relevant for machine learning applications, namely, the sparse-dense matrix multiplication (SPMM), the sampled dense-dense matrix multiplication (SDDMM), and the composition of the SDDMM with SPMM, also termed as FusedMM. We develop optimized implementations for SPMM, SDDMM, and FusedMM operations utilizing Intel oneAPI's Explicit SIMD (ESIMD) SYCL extension API. In contrast to CUDA or SYCL, the ESIMD API enables the writing of explicitly vectorized kernel code. Sparse matrix algorithms implemented with the ESIMD API achieved performance close to the peak of the targeted Intel Data Center GPU. We compare our performance results to Intel's oneMKL library on Intel GPUs and to a recent CUDA implementation for the sparse matrix operations on NVIDIA's V100 GPU and demonstrate that our implementations for sparse matrix operations outperform either.
1 Introduction
Machine-learning workloads need efficient GPU implementations of sparse operations, and this paper develops Intel GPU kernels using ESIMD. It reports performance gains over a V100 reference and Intel oneMKL.
1 Introduction
- SPMM multiplies sparse by dense matrices, SDDMM computes selected entries of a dense product, and FusedMM composes the two for applications such as sparse transformers.FusedMM uses the sparse SDDMM output as input to the following SPMM; fusing can increase performance.Source: §1 · p. 1
- Unlike scientific matrices often below 1% density, machine-learning matrices can be 10%-30% dense yet remain sparse enough to motivate compressed storage.Source: §1 · pp. 1–2
- Memory bandwidth is the main GPU bottleneck; ESIMD enables explicit vectorization and greater control over register usage and thread divergence, but lacks non-Intel GPU support.Source: §1 · p. 2
- Up to 10x versus the V100 reference [4] and up to 3.4x versus Intel oneMKL are reported for the optimized implementations.Source: §1 · p. 2
2 Background
This background defines SDDMM, SPMM, and FusedMM, then establishes the CSR storage and Intel Max 1550/ESIMD setting for the algorithms.
2.1 Sparse Matrix Operations
- The notation and matrix sizes investigated align with the reference implementation [4] and sizes commonly found in deep-learning networks.Source: §2.1 · p. 2
SDDMM
- SDDMM multiplies dense C by dense B^T and masks the result with I_A, computing dot products only where I_A is nonzero.Source: SDDMM · p. 3
SPMM
- SPMM multiplies sparse A by dense B to produce dense C.Source: SPMM · p. 3
FusedMM
- FusedMM composes SDDMM with SPMM, avoids explicitly storing the SDDMM result, and can outperform applying the operations sequentially.Source: FusedMM · p. 3
Matrix Storage Layout
- Dense matrices use row-major order; CSR stores nonzeros plus row and column index arrays, reducing memory and operation count at the cost of indirect addressing.ia marks row starts, while ja and avalue hold the column indices and values; Fig. 1 illustrates the matrix and its CSR arrays.Source: Matrix Storage Layout · p. 3, Figure 1 · p. 3
2.2 Intel Data Center GPU Max 1550
- The Max 1550 has two Xe stacks, but this study focuses on one because cross-stack access is slower than HBM access and is not cached.The GPU has 128 GB HBM and theoretical peaks of approximately 54 Tflops/s for FP32/FP64 and 832 Tflops/s for bfloat16 matrix engines.Source: §2.2 · pp. 3–4
2.3 Intel oneAPI/ESIMD
- Unlike SYCL subgroups, ESIMD work-items use explicit simd vectorization, enabling control of vector width, registers, memory caching, and divergent branches; kernels launch via SYCL_ESIMD_KERNEL.Source: §2.3 · p. 4
3 Overview of the Algorithms
The algorithms parallelize sparse matrix operations over output rows, with SDDMM optionally distributing a row across threads. FusedMM computes SDDMM values and immediately uses them in SPMM.
3.1 Parallelizing SDDMM
- SDDMM assigns threads to output rows, each computing that row’s nonzeros; multiple threads can divide a row’s nonzeros to increase parallelism.Each nonzero is a dot product of a row of C with a column of Bᵀ.Source: §3.1 · pp. 4–5
3.2 Parallelizing SPMM
- SPMM assigns one thread to each output row, accumulating scaled rows of B selected by the nonzero column indices of the corresponding sparse row.The inner loop is a SAXPY computation over the dense row.Source: §3.2 · pp. 5–6
3.3 Parallelizing FusedMM
- FusedMM computes each sparse row’s SDDMM values and uses them directly for SPMM, without storing those intermediate values in global matrix A.The row-wise algorithm combines the SDDMM and SPMM procedures.Source: §3.3 · pp. 6–7
4 ESIMD Implementation
The implementations target a matrix collection with 72 randomized tests and support N = 32 or N = 128. ESIMD kernels divide work into SIMD-sized chunks, with tuning for occupancy and latency; FusedMM combines SDDMM and SPMM in one kernel.
4.1 Sparse Matrix Collection
- The collection contains 72 randomized tests, and the implementation is guaranteed to work only for N = 32 or N = 128.These sizes are described as typical batch sizes in machine-learning training datasets.Source: §4.1 · p. 7
4.2 ESIMD Implementation of SDDMM
- SDDMM launches M × NT work-items, with each work-group computing one sparse row and dividing its nonzeros across NT threads in chunks of V_LC.Each thread computes ⌈nnzr/NT⌉ nonzeros, further partitioned into ⌈nnzt/V_LC⌉ chunks; the kernel loads dense rows, computes dot products, and stores the resulting values.Source: §4.2 · p. 8
- SDDMM prefetches the next chunk’s column indices and unrolls the inner loop to hide memory latency and use the vector engine.NT is selected as a power of two no larger than 16 to maximize occupancy; the occupancy discussion targets 4096 work-items per GPU stack.Source: §4.2 · pp. 8–9
4.3 Implementation of SPMM
- SPMM assigns one work-item to each output row and processes the sparse row in V_LC-sized chunks, accumulating scaled dense rows into the output.The implementation also uses prefetching and inner-loop unrolling; multiple work-items per work-group can compute multiple output rows.Source: §4.3 · p. 10
4.4 Implementation of FusedMM
- FusedMM uses NT = 1 so its SDDMM and SPMM stages have the same total thread count, then computes the intermediate values without storing them in avalues.The fused kernel reuses the SDDMM and SPMM kernel logic and writes the final E row.Source: §4.4 · pp. 10–11
5 Experiments
The experiments compare ESIMD sparse-matrix kernels with oneMKL and prior CUDA work, then use roofline analysis to examine performance for selected SDDMM and SPMM cases. FusedMM shows the largest reported maximum speedup over oneMKL.
5 Experiments
- Tests used a single stack of an Intel Data Center GPU Max 1550, Intel's icpx compiler 2023.2.0.20230721, and an unreleased engineering driver.The oneMKL tests used the -qmkl compile option.Source: §5 · p. 11
- For SDDMM, performance gains over [4] average approximately 4.5x and reach approximately 10x at most.The maximum occurs for the 32k/8k/32/90% case.Source: §5 · p. 12
- Our SDDMM achieves 1.84x average and 3.4x maximum speedups over oneMKL dense blas::gemm.The comparison is disadvantageous to oneMKL because 70%-90% of dense GEMM's computed flops are irrelevant to SDDMM; oneMKL has no dedicated SDDMM function.Source: §5 · p. 11, Figure 12 · p. 13, Figure 13 · p. 13
- Our SPMM achieves 1.37x average and 2.16x maximum gains over oneMKL sparse::gemm, and 2.4x average and 4x maximum gains over [4].The authors attribute the oneMKL gains partly to specializing for matrix sizes and sparsities relevant to machine learning.Source: §5 · p. 14, Figure 14 · p. 15, Figure 15 · p. 15
- FusedMM achieves 1.3x the throughput of our SDDMM and 1.074x that of our SPMM on average.The authors say its average gain over oneMKL exceeds the gains in either separate operation, underscoring the importance of fusion.Source: §5 · p. 16
- FusedMM is approximately 2x faster on average and 3.25x at most than oneMKL, whose SDDMM comparison excludes the time for sampling.The authors report a peak of 12,647 Gflops/s, or 48% of the device's theoretical peak.Source: §5 · p. 16, Figure 16 · p. 17, Figure 17 · p. 17
5.1 Performance Analysis
- At M=8192, K=8192, N=128, and α=0.7, SDDMM reaches 9,732 Gflops/s, 77% of the 12,600 Gflops/s SP Vector Add Peak.Its L3 arithmetic intensity is 5.6, compared with a stated worst case of 0.49; the plotted measurement reflects cached iterations after the first.Source: Figure 18 · p. 18, §5.1 · p. 18
- At M=8192, K=8192, N=128, and α=0.7, SPMM reaches 37% of its 26,000 Gflops/s SP Vector MAD Peak.The authors report L3 arithmetic intensity of 8.06 and say there is still room to improve SPMM performance.Source: §5.1 · pp. 18–19, Figure 19 · p. 19
6 Conclusion and Future Work
The paper concludes that its optimized sparse-matrix implementations outperform existing approaches, including oneMKL and CUDA work on an NVIDIA V100. It highlights operator fusion and identifies bfloat16 and larger sparse datasets as future directions.
6 Conclusion and Future Work
- The implementations outperform oneMKL by up to 2.16x for specialized SPMM and up to 3.25x for fused SDDMM-SPMM, and run up to 10x faster than optimized CUDA code on an NVIDIA V100.The conclusion describes the SPMM specialization as targeting relatively dense sparse matrices.Source: §6 · p. 19
- The conclusion identifies the lack of dedicated SDDMM functionality in oneMKL and recommends fused operations wherever possible for machine-learning applications.Source: §6 · p. 19
- Future work will explore bfloat16 sparse operations and evaluate performance on larger sparse-matrix datasets from different machine-learning applications.Source: §6 · p. 19