#1965·warp

tile_matmul(): use the cuBLASDx register-accumulator API with suggested shared layouts for Tensor Core precisions

Author: daedalus5Created Sep 16, 2026Updated Sep 18, 2026
Labelstile

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_C or CUBLASDX_TENSOR_SUGGESTED_ACCUMULATOR_C: C accumulated in registers across consecutive tile_matmul calls (the K-loop), materialized once.
  • libmathdx-generated copy device 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_load fills 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 register tile_matmul route.
  • 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 of acc between the loop's tile_matmul calls.
  • 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.func parameters, 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.