PaperlayerView exact revision on arXiv

Source-linked AI summary

Hierarchical Shared Memory-Aware Optimization for TRSM on GPU Platforms

Xinzhe Chen, Haowei Li, Lijuan Hu, Wenjing Ma, Fangfang Liu

arXiv:2608.25469v1cs.DC

TL;DR

GPU TRSM remains challenging for left-side lower-triangular systems because forward substitution imposes strict dependencies and shared memory cannot always stage both operands. HSMA-TRSM combines pipelined small-scale solving, diagonal-block decoupling with adaptive blocking, and compile-time platform-specific selection, achieving peak speedups of 2.05× over cuBLAS and 2.06× over rocBLAS.

  • Problem

    Left-side lower-triangular GPU TRSM is constrained by row-wise forward-substitution dependencies and insufficient shared memory for wide data types, especially double complex.

  • Method

    HSMA-TRSM combines compute-memory pipelining and a seven-stage double-complex pipeline for small cases with diagonal-block decoupling, adaptive blocking, and offline-profiled compile-time selection for large cases.

  • Results

    2.05× over cuBLAS and 2.06× over rocBLAS are HSMA-TRSM's peak reported speedups across the evaluated accelerators.

  • Takeaways & Limitations

    Gains are strongest in shared-memory-constrained double-complex small cases and large real-type cases, while mature vendor kernels and some complex large cases leave less headroom.

  • Takeaways & Limitations

    At m=64, the on-chip direct-solve strategy is a boundary-case compromise, and larger float or double problems favor the blocked formulation's higher-throughput GEMM updates.

Abstract

from arXiv · show

Triangular Solve with Multiple Right-hand Sides (TRSM) is a fundamental BLAS Level-3 operation that underpins LU/Cholesky decomposition, sparse direct solvers, and matrix inversion. In the left-side lower-triangular case studied in this paper, efficient GPU implementation remains challenging because forward substitution introduces strict row-wise dependencies, and shared memory is too scarce to hold both operand matrices for wide data types such as double complex. This paper presents HSMA-TRSM, a hierarchical shared memory-aware optimization framework for left-side lower-triangular TRSM on NVIDIA A100, NVIDIA H800, and Hygon DCU Z100 accelerators. For the small-scale regime (m,n<=64), we design a pipelined compute-memory overlap mechanism through loop unrolling and instruction reordering, and propose a dual thread-group seven-stage pipeline strategy to address shared memory constraints for double complex types. For large-scale problems, we introduce a diagonal block decoupling optimization with an O(IB)shared-memory footprint for diagonal block inversion, enabling adaptive block size selection based on matrix scale and hardware characteristics. A compile-time configuration selection framework based on offline profiling and online lookup selects the optimal block size per platform with zero runtime overhead. Evaluated on NVIDIA A100, H800, and Hygon DCU Z100, HSMA-TRSM achieves peak speedups of 2.05xover cuBLAS and 2.06xover rocBLAS. The gains are strongest in shared-memory-constrained double-complex small cases and in large real-type cases where adaptive blocking improves GEMM-dominated updates, while mature vendor kernels leave less optimization headroom in some regimes.

1 Introduction

GPU TRSM is important but difficult to optimize for left-side lower-triangular solves because forward substitution creates row-wise dependencies and shared memory limits operand staging. HSMA-TRSM addresses these constraints with hierarchical pipelines, diagonal-block decoupling, adaptive blocking, and platform-specific configuration selection.

  • TRSM is a BLAS Level-3 operation used in LU and Cholesky decomposition, sparse direct solvers, matrix inversion, and least squares problems.
  • Small matrices with m,n≤64 leave limited shared memory for compute-memory pipelining, especially with double-complex operands.A 64 × 64 double-complex block occupies 64KB, exhausting the DCU budget and exceeding NVIDIA's 48KB per-block budget.
  • Large-scale vendor implementations fix block sizes and conventional diagonal-block processing requires O(I^2_B) shared memory, restricting larger blocks and adaptation across hardware.rocBLAS fixes the block size at 128, while existing optimizations often target a single architecture.
  • HSMA-TRSM overlaps prefetching with computation, uses a dual thread-group seven-stage pipeline for double complex, and decouples diagonal-block inversion to reduce shared-memory use to O(I_B).The framework also selects near-optimal block sizes adaptively and uses compile-time offline profiling with online lookup.
  • 2.05× and 2.06× are the reported peak speedups over cuBLAS and rocBLAS, respectively.

2 Background

TRSM solves triangular systems with multiple dense right-hand sides, and the studied left-side lower-triangular case proceeds by dependency-bound forward substitution. GPU performance therefore depends on managing serial dependencies, data reuse, and limited shared memory while exploiting asynchronous memory operations.

  • TRSM solves op(A)X=αB, where A is triangular and B is a dense matrix containing multiple right-hand sides.
  • Left-side lower-triangular TRSM computes solution rows by forward substitution, with each row depending on the preceding rows.
  • The computation has complexity O(m^2n), data volume O(m^2 + mn), and O(n) reuse of each triangular-matrix element.
  • GPU platforms differ in execution units and thread organization, using 32-thread warps on NVIDIA and 64-thread wavefronts on Hygon DCU Z100.
  • Shared memory is faster than global memory but limited to 48KB per block on A100/H800 and 64KB on DCU in the target setting.
  • Asynchronous memory instructions can prefetch future data while current data is computed, enabling compute-memory pipelines that hide memory latency.

3 Related Work

Prior TRSM research covers recursive blocking, batching, portability, communication avoidance, and auto-tuning, but target-setting gaps remain in shared-memory-constrained pipelining, diagonal-block resource costs, and low-overhead cross-platform tuning.

  • cuBLAS and rocBLAS use hierarchical TRSM implementations that generally direct-solve small problems and block larger problems into solves and GEMM updates.
  • Small-scale implementations separate loading, computation, and write-back, limiting overlap with forward substitution; double complex further constrains shared-memory staging.
  • rocBLAS fixes the large-scale block size at 128, while fused diagonal-block processing requires 4×I^2_B×sizeof(datatype) shared memory.
  • Smaller GEMM blocks reduce computational density and fail to amortize memory costs, making larger compute-intensive updates important for throughput.
  • Existing vendor optimizations rely heavily on platform-specific hardware features and low-level scheduling strategies that are difficult to migrate across architectures.
  • Prior work improves recursive blocking, batched kernels, communication costs, and auto-tuning, but does not jointly resolve the three target-setting issues identified here.

4 Design

HSMA-TRSM selects TRSM strategies hierarchically by matrix scale and data type, combining direct shared-memory solving for small cases with blocked GEMM-oriented processing for large cases. It addresses shared-memory limits through pipelining and diagonal-block decoupling, then uses adaptive blocking and compile-time lookup to improve platform portability and performance.

  • Overall Framework Design: Small-scale TRSM is memory-bound with O(1) arithmetic intensity, so optimization targets data reuse and memory-latency hiding rather than raw throughput.Large-scale blocked TRSM is increasingly compute-bound because its dominant work is organized as GEMM updates.
  • Overall Framework Design: HSMA-TRSM automatically selects direct solving, a double-complex partitioned pipeline, or block-GEMM processing based on matrix scale and data type.The framework uses 64 as the empirically supported boundary between small-scale and large-scale regimes.
  • Partitioned Pipeline Strategy: The 64 × 64 double-complex case uses a partitioned pipeline because direct solving cannot fit sufficient right-hand-side data on chip, while blocked processing cannot amortize its setup overhead.This is explicitly characterized as a boundary-case compromise rather than a generally preferred strategy.
  • Small-scale TRSM Optimization: The small-scale pipeline overlaps future-row loads with current-row solving and delays write-backs, while column-wise mapping exploits independent right-hand-side columns.For double complex, a dual thread-group seven-stage pipeline assigns memory operations and computation to different groups, achieving full compute-memory overlap despite shared-memory limits.
  • Large-scale TRSM Optimization: Diagonal-block decoupling reduces inversion storage from O(I_B^2) to O(I_B), enabling larger feasible block sizes and more effective GEMM updates.The method retains only the current column, previously solved columns, triangular sub-blocks, identity tiles, and intermediate results in shared memory.
  • Large-scale TRSM Optimization: Adaptive blocking balances launch overhead, GEMM efficiency, load balance, and hardware limits, while offline profiling and runtime lookup select configurations without runtime tuning overhead.Compared with rocBLAS’s fixed block size of 128, the framework selects configurations according to matrix scale and data type.

5 Performance Evaluation

Performance is evaluated on three GPU accelerators against vendor BLAS libraries, using the platforms’ differing architectures and capabilities as the experimental context.

  • The evaluation compares HSMA-TRSM with cuBLAS v13.2.1 on NVIDIA platforms and rocBLAS v5.1 on the Hygon DCU Z100.

5.1 Experimental Setup

Experiments use square random matrices across four data types and both scale regimes, while numerical correctness is assessed against cuBLAS on A100.

  • Tests use square matrices with random initialization across float, double, float complex, and double complex types.Both small-scale and large-scale regimes are evaluated.
  • Each test point runs 10 times, with the first run used for warmup and the average of the last 9 runs reported.
  • Numerical correctness is measured against cuBLAS on A100 using the normwise relative Frobenius error.

5.2 Small-Scale TRSM Performance

For m,n≤64, HSMA-TRSM uses shared-memory direct solving with compute-memory overlap, achieving its clearest gains near the upper range and for complex types.

  • The small-scale kernel improves direct solving through loop unrolling, instruction reordering, and compute-memory overlap.
  • At m=64, float reaches 4.64 GFLOPS on DCU and 5.75 GFLOPS on A100, improving over rocBLAS by 18.4% and cuBLAS by 29.2%, respectively.
  • At m=64, float complex reaches 11.61 GFLOPS on DCU, improving over rocBLAS by 86.3%.
  • At m=64, double complex reaches 3.53 GFLOPS on DCU, 15.88 GFLOPS on A100, and 12.47 GFLOPS on H800.These correspond to improvements of 96.1%, 14.5%, and 105.4% over vendor libraries.
  • Double complex shows the largest speedups, reaching 2.05× on H800 and 1.96× on DCU, while float and double mostly remain within 1.1×–1.3×.
  • At m=64 on A100, HSMA-TRSM improves over MAGMA by 28.9% for float, 24.6% for double, 5.1% for float complex, and 41.4% for double complex.

5.3 Large-Scale TRSM Performance

For large problems, diagonal-block decoupling reduces shared-memory requirements and adaptive blocking improves performance as GEMM updates become dominant, though gains vary by platform and data type.

  • Large-scale TRSM decomposes computation into diagonal block inversion and GEMM update phases.
  • Diagonal-block decoupling reduces shared-memory requirements from O(I^2B) to O(IB), extending the feasible block size.
  • The adaptive algorithm increases block size from 256 at m=1024 to 512 at m=8192 in representative float and double cases.
  • On DCU, adaptive blocking improves performance over rocBLAS’s fixed-128 baseline by 39.3% at m=1024 and 23.8% at m=8192.
  • At m=16384 on DCU, speedups remain 1.63× for float and 2.06× for double.Peak performance reaches 6010 GFLOPS for float and 3654 GFLOPS for double.
  • On A100, HSMA-TRSM exceeds cuBLAS and MAGMA at the upper end of the scale range, whereas H800 double-complex gains remain close to parity with cuBLAS.
  • On A100, representative comparisons show 11798 versus 8837 GFLOPS for float at m=4096 and 10577 versus 7344 GFLOPS for double at m=4096.

5.4 Analysis and Discussion

Cross-platform TRSM behavior reflects both hardware throughput and per-block on-chip storage, with small problems benefiting from staging and overlap while large problems favor GEMM-dominated blocked updates.

  • DCU’s 64KB per-block LDS budget favors fine-grained on-chip staging, whereas A100 and H800 use 48KB per-block shared memory.
  • Figures 11 and 12 present measured large-scale performance and speedup across four data types and three platforms.
  • A100 and H800 provide higher absolute compute throughput and memory bandwidth, making larger blocked updates more effective in GEMM-dominated workloads.
  • Small-scale TRSM has low arithmetic intensity and is sensitive to shared-memory reuse and memory-computation overlap.
  • Large-scale blocked TRSM increasingly spends time in GEMM-like updates, while strong vendor-library cases indicate limited remaining optimization headroom.

6 Conclusion and Future Work

HSMA-TRSM targets small- and large-scale left-side lower-triangular TRSM with specialized pipelines and adaptive blocked updates. It reports peak speedups up to 2.05× over vendor libraries, while future work extends the framework beyond its current scope.

  • HSMA-TRSM uses pipelined compute-memory overlap and a seven-stage dual thread-group pipeline for small-scale double-complex TRSM, reaching a peak 2.05× speedup over vendor libraries.
  • For larger problems, diagonal block decoupling reduces inversion storage from O(I^2B) to O(IB), enabling larger blocked updates.
  • HSMA-TRSM achieves up to 51.9% improvement for double precision on DCU and maintains 1.63× to 2.06× speedups at m = 16384 for real types.
  • Gains are more modest for mature vendor kernels and some complex large-scale cases.
  • Future work includes kernel fusion for very small scales, batched TRSM, and extensions to sparse and mixed-precision settings.
Loading 2608.25469v1…