tile_matmul(): use the cuBLASDx register-accumulator API with suggested shared layouts for Tensor Core precisions
Description
wp.tile_matmul() runs its cuBLASDx GEMM through the shared-memory pointer API (CUBLASDX_API_SMEM in wp_cuda_compile_dot, warp/native/warp.cu): A, B, and C are plain row- or column-major shared-memory tiles, and every call computes C = alpha*A*B + beta*C in place. That has two costs the cuBLASDx documentation calls out for this API: fragment loads from plain layouts suffer shared-memory bank conflicts, and the accumulator is read from and written back to shared memory on every call. cuBLASDx recommends the register (accumulator) API with its suggested shared-memory layouts "in performance-critical environments".
These costs matter most for Tensor Core precisions, where the math itself is cheap. With the precision="tf32" option proposed in #937, TF32 through the shared-memory API is at parity with fp32 or slower on both an RTX 5090 and an H100 PCIe (TF32 reaches under 10% of the H100's TF32 peak; tensor pipe utilization is 7-10%). A standalone spike shows that the register API removes that ceiling.
Spike: same workload, three cuBLASDx execution paths (RTX 5090, cuBLASDx 0.4.0 headers, CUDA 13.3)
Each block loads one A (MxK) and one B (KxN) fp32 tile into shared memory once, then runs 64 back-to-back block GEMMs accumulating into C, over a 2048x2048 grid of output tiles. Best of 3 rounds x 20 launches. TFLOP/s:
| tile, threads | fp32 smem API (Warp today) | fp32 register API | tf32 smem API (Warp with #937) | tf32 register API + suggested layouts |
|---|---|---|---|---|
| 64x64x64, 128 | 58.7 | 62.4 | 50.3 | 106.5 |
| 64x64x64, 256 | 53.3 | 55.1 | 39.1 | 106.9 |
| 64x64x64, 64 | 44.6 | 49.4 | 61.9 | 92.4 |
| 32x32x32, 128 | 48.5 | 51.5 | 32.1 | 118.9 |
| 128x128x32, 128 | 50.8 | 64.0 | 56.8 | 69.3 |
| 128x128x32, 256 | 54.2 | 64.3 | 52.2 | 78.4 |
- TF32 on the register API is 2.1x to 3.7x faster than TF32 on the shared-memory API, and 1.8x to 2.5x faster than fp32 on the shared-memory API. It matches cuBLAS TF32 on the same GPU (91.7 TFLOP/s via PyTorch at N=2048). The RTX 5090's TF32 and FP32 spec peaks are equal (104.8 TFLOP/s); on Hopper and Blackwell data center parts, where the TF32 peak is ~7x the FP32 peak, the gap should be larger.
- fp32 gains only 1.04x to 1.26x from the register API. fp32/fp32/fp32 has no MMA instruction in cuBLASDx, so it stays FFMA-bound whatever the layout. This is a Tensor Core precision lever (TF32, FP16, BF16), not a general one.
- Results on the register path have the expected TF32 error (relative 6.7e-4 against an fp64 reference) and fp32 error (1e-6).
The 128x128 rows are limited by wave quantization at this grid size for all paths.
Proposed change
Add a MathDx GEMM variant built on libmathdx's tensor API (CUBLASDX_API_TENSORS, available in the libmathdx version Warp already links), using:
CUBLASDX_TENSOR_SUGGESTED_SMEM_A/_B: A and B in cuBLASDx's suggested (swizzled) shared layouts.CUBLASDX_TENSOR_SUGGESTED_RMEM_CorCUBLASDX_TENSOR_SUGGESTED_ACCUMULATOR_C: C accumulated in registers across consecutivetile_matmulcalls (the K-loop), materialized once.- libmathdx-generated
copydevice functions to move data between global memory, the suggested shared layouts, and the accumulator.
Design consequences worth agreeing on before implementation:
- The suggested layouts and the accumulator are opaque (
struct { void* ptr; }with unspecified layout), so Warp code cannot index them. A tile in one of these forms is a new storage kind:tile_loadfills it through the generated copy, and it is materialized to a plain shared or register tile when any other operation consumes it. This fits the materialize-on-consume direction of the tile storage work but is distinct from Warp's existing registertile_matmulroute. - The accumulator must persist across calls to pay off, so the pattern
acc = tile_zeros(); for k: tile_matmul(a, b, acc)is the target. Codegen needs to recognize the accumulate-only use ofaccbetween the loop'stile_matmulcalls. - The backward pass adds two GEMMs whose operand roles permute; each needs its own layouts and copies.
- The current shared-memory API remains the fallback for operands that are views or come from
@wp.funcparameters, as with the alignment operator (#1938).
Context
- #937 (
precision="tf32"): correct and verified to use TF32 Tensor Cores, but no speedup through the shared-memory API on either GeForce or H100. - #1938 (Alignment / StaticBlockDim): vectorized the shared-memory API's loads; does not address bank conflicts or the C round trip.
- cuBLASDx docs: "Achieving High Performance" and the execution-methods page recommend the register API and suggested layouts over the shared-memory API.
System Information
RTX 5090 (sm_120), CUDA 13.3, driver 610.43, Warp 1.18.0.dev; cuBLASDx 0.4.0 headers from the nvidia-mathdx PyPI package for the spike; libmathdx 0.4.1 (cuBLASDx 0.7.1) in Warp.
Source: NVIDIA/warp