API Card: supermma#

Note

This chapter is part of the agent-targeted Agent BLAS Performance Guide and is also published as Markdown for agent consumption, in the markdown/ folder of the HTML documentation tree.

supermma is the warp-level tensor-core instruction class based on mma.sync. It is available through both the per-tile BLAS().execute(...) API and the pipeline API. A pipeline can stage operands with TMA or cooperative copies while its computation still uses supermma; choosing a pipeline does not by itself select wgmma or tcgen05. Some GPUs have neither of those instruction classes — Ampere, Ada, and the consumer Blackwell parts SM 120 / SM 121 — so on them every tensor-core GEMM, pipelined or not, uses supermma. Precisions without a supported tensor-core instruction use the separate FMA class.

This card describes the shared-memory operands and register accumulators, then shows both APIs. Two shipped examples show the per-tile API: example/cublasdx/01_gemm_introduction/introduction_example.cu (the three execute forms, plain layouts, double) and example/cublasdx/10_gemm_block_performance/single_gemm_performance.hpp (suggested layouts with the suggested accumulator, half, MaxAlignment(), 128 threads: the configuration this card describes).

Descriptor. Precision, SM, tile shape and thread count determine whether a tensor-core instruction is available. WithPipeline() enables the pipeline API; the additional gates in the Hopper and Blackwell cards determine whether the computation uses wgmma or tcgen05 instead. The descriptor below uses supermma with either API.

  • Precision on tensor cores: half inputs with a half or float accumulator from SM 75; bfloat16 with float, tf32 with float, int8 / uint8 (any pairing) with int32 and double from SM 80; fp8 (e4m3 or e5m2, any pairing) with half or float from SM 89.

    The complex Type of each of those triples is on the table as well. Everything else is FMA: float inputs on every SM (Precision<float> is never promoted to tf32; ask for cublasdx::tfloat32_t and accept the rounding), fp8 below SM 89, bfloat16, tf32 and 8-bit integers below SM 80.

    On SM 1000, 1030 and 1070 real float FMA GEMMs can issue fma.f32x2 (SASS FFMA2, two results per instruction), so accept FFMA2 as well as FFMA in the SASS check below.

  • BlockDim a power of two from 32 to 1024 for the suggested layouts and accumulator to exist; any count in that range compiles, but a non-power-of-two block gets the plain layouts and the default accumulator.

    Without BlockDim the library’s suggested_block_dim applies; for a matrix-vector product (M == 1 or N == 1) it is M * N rounded up to whole warps and clamped to [32, 256].

    With an explicit BlockDim a matrix-vector product uses at most M * N threads (whole warps above 32) and the rest are inactive; this cap does not apply under WithPipeline() or to complex integer types, which keep every BlockDim thread.

  • Size divisible by the warp tiling: the warps form a near-square grid (2 x 2 for 128 threads) and each covers a 16 x 8 instruction, so with 128 threads M must be a multiple of 32 and N of 16 (M of 16 and N of 32 when N > M);

    K a multiple of 16 for 16-bit inputs, of 8 for tf32, of 32 for 8-bit inputs and of 4 for double.

    Scalar FMA is chosen only for M == 1 or N == 1 (matrix-vector) and K == 1 (outer product), for every precision; any other shape, including a small tile, uses the tensor-core instruction when the precision has one.

    The library can adjust the thread tiling to fit M and N, so confirm the generated instruction in the SASS rather than inferring it from divisibility alone.

  • Alignment<16, 16, 16> (= MaxAlignment()) for ldmatrix on A and B, and for stmatrix on C when C has the operands’ element size, on SM 90 and up. The default alignment is the element size, so a half descriptor without it stays supermma but loads shared memory with plain vector loads: a gate on the loads, not on the class.

  • No LeadingDimension and no experimental::Tile: either one silently turns the suggested layouts and the suggested accumulator into the plain ones.

  • Arrangement: free for 16-bit operands (ldmatrix.trans); tf32 and 8-bit operands get ldmatrix only K-major, that is TN. On the FMA path float wants a column-major A and a row-major B, the orientations its shared loads vectorize.

    float and complex<float> FMA with both operands K-major (row-major A, column-major B) and plain shared layouts whose row stride is a multiple of 32 bytes (that is, the leading dimension times the element size) put every row of a tile on the same banks;

    the library re-tiles the thread grid along C’s contiguous axis to avoid the conflicts, and FMA GEMMs with large per-thread fragments are split along K to cap live registers.

    Neither needs an operator; both depend on the leading dimensions you pass, so see the finetuning chapter before padding by hand.

The descriptor below satisfies all of them from SM 80 up:

using BLAS = decltype(cublasdx::Size<128, 128, 32>()
    + cublasdx::Precision<__nv_bfloat16, __nv_bfloat16, float>()
    + cublasdx::Type<cublasdx::type::real>() + cublasdx::Function<cublasdx::function::MM>()
    + cublasdx::Arrangement<cublasdx::row_major, cublasdx::col_major, cublasdx::row_major>()  // TN; free for bf16
    + cublasdx::Block() + cublasdx::BlockDim<128>() + cublasdx::StaticBlockDim()              // power of two
    + cublasdx::Alignment<16, 16, 16>()                                                        // ldmatrix needs 16
    + cublasdx::SM<900>());                                                                    // no modifier needed

StaticBlockDim() promises a launch of exactly BLAS::block_dim threads. For this descriptor, every thread owns part of the result and is_thread_active() is true. Verify the selected instruction in the SASS as described below.

Per-Tile API#

Host side. For a per-tile call, size the shared memory, raise the limit and launch with the descriptor’s block.

const unsigned smem = cublasdx::get_shared_storage_size_ab<BLAS>();        // A and B tiles; add scratch after
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(smem + scratch_bytes));
kernel<<<grid, BLAS::block_dim, smem + scratch_bytes, stream>>>(...);

  • The suggested layouts occupy the same bytes as the plain ones, so get_shared_storage_size_ab sizes both (16 KiB for the card’s tile in bf16); get_shared_storage_size<BLAS>() adds the C tile for the shared-memory-output form below, and cublasdx::make_shared_storage_calculator() takes explicit layouts when the tiles are not the descriptor’s.

  • Declare the kernel with __launch_bounds__(BLAS::max_threads_per_block).

Device side. The register-accumulator form, one K tile per execute:

extern __shared__ __align__(16) cublasdx::byte smem[];
auto [smem_a, smem_b] = cublasdx::slice_shared_memory_ab<BLAS>(smem);
auto a_s = cublasdx::make_tensor(smem_a, BLAS::suggest_layout_smem_a());
auto b_s = cublasdx::make_tensor(smem_b, BLAS::suggest_layout_smem_b());
auto acc = BLAS::suggest_accumulator();                                    // zero; pairs with suggested layouts
for (int kt = 0; kt < k_tiles; ++kt) {
    cublasdx::copy<BLAS, cublasdx::alignment_of<BLAS>::a>(a_tile(kt), a_s);  // every block_dim thread
    cublasdx::copy<BLAS, cublasdx::alignment_of<BLAS>::b>(b_tile(kt), b_s);
    cublasdx::copy_wait();                                                  // cp.async wait + __syncthreads
    BLAS().execute(a_s, b_s, acc);                                          // acc += A_kt * B_kt
    __syncthreads();                                                        // before smem_a / smem_b are rewritten
}
if (acc.is_thread_active()) {
    acc.partition_and_store(c_tile);
}

The shared-memory-output form, C = alpha * A * B + beta * C in place with plain layouts:

auto [smem_a, smem_b, smem_c] = cublasdx::slice_shared_memory<BLAS>(smem);
auto c_s = cublasdx::make_tensor(smem_c, BLAS::get_layout_smem_c());       // a_s, b_s, three copies, copy_wait
BLAS().execute(alpha, a_s, b_s, beta, c_s);
__syncthreads();
cublasdx::copy<BLAS, cublasdx::alignment_of<BLAS>::c>(c_s, c_tile);

  • The accumulator is zero when constructed and every execute(a, b, acc) adds to it, so one accumulator carries a whole K loop; acc.clear() restarts it. BLAS().execute(a, b) with no accumulator returns a fresh one, suggested when the layouts are suggested.

  • suggest_accumulator() pairs with the suggested layouts of both operands; get_accumulator() pairs with any layouts. The suggested accumulator with a plain layout is a compile error (Suggested accumulator can be used only with suggested layouts); suggested layouts with the default accumulator compile and run the default configuration, so the swizzle buys nothing.

  • Load functors (execute(a, b, acc, a_op, b_op)) can transform operands on the supermma path. The wgmma and tcgen05 instruction paths require identity loads.

  • cublasdx::copy runs on the descriptor’s block_dim threads and copy_wait() ends with a block barrier; execute ends with none, so put a __syncthreads() between it and the next write to the operands it read.

  • Epilogue accessors have the meaning given in the Hopper card: reduce_and_store, process_and_store, partition_and_store, make_partition_and_copy / partition_and_copy, acc.axpby(alpha, beta, c), get_results() (a copy) and map_fragment_index(i).

    Guard the ones that touch memory with acc.is_thread_active() where your own code touches memory: with StaticBlockDim() and a launch of exactly block_dim threads it is always true.

    Otherwise it is false for every thread outside the GEMM (for example the surplus threads of a matrix-vector product), is_predicated() is also true when some block threads are inactive, and the library accessors (partition_and_*, process_and_*, reduce_and_store, axpby, the copies) are no-ops there, so they need no guard; get_results() is unspecified on inactive threads.

    A block launched with more threads than BlockDim must keep the extra threads out of execute and the copies, and the linear thread index follows BlockDim’s rank, not the launch’s.

Pipeline API#

Add WithPipeline() to the same descriptor and pass global-memory tensors to suggest_pipeline. The computation remains supermma for this descriptor. The shipped example/cublasdx/01_gemm_introduction/introduction_pipeline.cu also demonstrates a supermma pipeline.

using PipelinedBLAS = decltype(BLAS() + cublasdx::WithPipeline());
auto A = cublasdx::make_gmem_tensor<cublasdx::row_major>(a_ptr, M, K, lda);
auto B = cublasdx::make_gmem_tensor<cublasdx::col_major>(b_ptr, K, N, ldb);
auto p = cublasdx::suggest_pipeline<Depth, PipelinedBLAS>(A, B);
if (!p) { /* handle p.error() before launching */ }
const auto smem_bytes = p->buffer_size();                     // add any application scratch separately
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(smem_bytes));
kernel<<<grid, p->get_block_dim(), smem_bytes, stream>>>(..., p->get_device_handle());

Use the pipeline’s block dimensions, buffer size and buffer alignment. The supermma pipeline uses the same block threads for loading, computing and the epilogue; it adds no producer warpgroup. On Hopper and later, TMA is selected when the operand layouts and alignments support it; otherwise the pipeline uses cooperative copies. Neither copy method changes the instruction class.

On the device, obtain the tile and let the pipeline drive the K loop. Take the accumulator from the tile pipeline rather than constructing it with BLAS::suggest_accumulator():

extern __shared__ __align__(device_pipeline.buffer_alignment()) cublasdx::byte smem[];
auto tile = device_pipeline.get_tile(smem, row, col);
tile.execute([&](auto& acc) {
    auto out = cublasdx::get_tile(C, PipelinedBLAS::c_shape, row, col);
    acc.partition_and_store(out);
});

Pass the handle as __grid_constant__ DevicePipeline const device_pipeline and declare the kernel with __launch_bounds__(DevicePipeline::max_threads_per_block, 1). The second argument does not have to be 1: a higher minimum-blocks value lets several blocks share an SM so one block’s epilogue overlaps another’s compute — a supermma-device substitute for warp specialization. It trades away registers, so measure it. For a persistent loop, construct the tile once and call device_pipeline.reset_tile(tile, row, col) for each subsequent tile. Use tile.epilogue_sync() when the epilogue needs to synchronize participating threads. With a reusable accumulator, call finish_accumulation() from every thread that participates in tile.execute(acc) before reading its result.

The pipeline requires whole Size tiles in the global dimensions and matching K extents for A and B. Handle pipeline-creation errors, including alignment errors, before launch. Read or discard each completed accumulator before executing the next tile. Load functors can be passed as tile.execute(epilogue, a_op, b_op); EnableInputStreaming() forbids non-identity load functors.

Checking The Generated Instructions#

Confirm the path in the SASS: HMMA for 16-bit and tf32 inputs (HMMA.16816.F32.BF16 for the card’s descriptor), IMMA for 8-bit integers, with LDSM present when the suggested layouts took effect (LDSM.16.M88.4 here). FMA uses FFMA or FFMA2 for float inputs. Check each kernel separately; these instruction names identify the computation for either API.