Fitting Into cuBLASDx Idioms

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.

Find every row below that is true of your kernel — expect several; the axes are independent, and a solution is a choice on each one. Each idiom has one directory of measurable variants under example/cublasdx/performance_guide/; each variant’s header comment states what it isolates.

Idiom Catalogue

#	If your kernel…	Example directory
01	holds A / B in a different type than it multiplies in — for example __half data, float accumulation	01_decoupled_precision_inputs
02	writes results in a different type than it accumulates in	02_decoupled_precision_output
03	applies a pure function of each input element’s own value before multiplying	03_elementwise_transform_before_blas
04	preprocesses in a way that needs other elements — row normalization, transpose, FFT — so threads must cooperate	04_cooperative_transform_before_blas
05	combines an input element with a second tensor before the multiply	05_external_data_elementwise_transform_input
06	as 05, but the transform over the external data itself needs cross-thread cooperation	06_external_data_cooperative_transform
07	post-processes each output element using its own global position — mask, bias, positional scaling	07_elementwise_transform_after_blas_indices
08	post-processes needing the whole row, column or tile — max, sum, normalization	08_cooperative_transform_after_blas
09	as 08, but the epilogue also reads a separate tensor (e.g. a statistic carried between tiles)	09_cooperative_transform_after_blas_external_data
10	loads A / B from global memory and feeds them to the GEMM — the most common entry point; the answer is usually the pipeline API, which overlaps loads with the multiply and unlocks TMA and the Hopper/Blackwell MMA instructions	10_overlapping_compute_and_copy
11	multiplies a matrix by its own transpose (A * A^T) — both operand slots, same buffer, same layout	11_overlapping_a_b_same_layout
12	multiplies a matrix by itself (A * A) — same buffer, but the two operand slots want different layouts	12_overlapping_a_b_different_layout
13	feeds one GEMM’s result into another without a round trip through global memory	13_output_as_input
14	has an operand computed rather than stored — generated from its indices, alone or scaled by a small side vector	14_virtual_tensors_a_b
15	accumulates several GEMMs into one output before a single epilogue	15_accumulative_postprocessing_multi_gemm
16	calls the GEMM for every pair of a slice-pair loop nest — products of sliced operands with per-slice scale factors accumulated into one result, as in Ozaki-scheme emulation; the operand of the outer loop can stay staged across the inner loop, and the factorized scales are applied while staging	16_gemm_in_nested_loop
17	runs many small problems, each with its own leading dimension read from an array	17_batched_ld_from_array
18	runs problems of differing shapes or precisions in one kernel	18_batched_different_sizes
19	needs accumulation more accurate than plain fp32 — a separate fp32 register accumulator, periodic fp64 promotion, per-K-tile fp64 promotion, scored against an fp64 ground truth	19_enhanced_precision

Idioms Feed Back Into The Descriptor

Several rows are choices about operators, so expect one pass back through Finding Problem Shape:

Idiom	What it changes in the descriptor
01, 02	Precision<PA, PB, PC> is the idiom — give the accumulation type separately from the input types.
10	WithPipeline(), and a Size whose K spans several tiles (a single-K-tile problem has nothing to overlap). On Hopper the wgmma gates are SM<900, arch_specific>() + Alignment<16, 16, 16>() + WithPipeline() + EnableInputStreaming(), a block of 128, 256, 512 or 1024 threads (BlockDim<128> or <256> in practice), a supported real Precision and a Size divisible by the warpgroup tiling; Arrangement is free for 16-bit inputs and must be TN only for tf32, int8 / uint8 and fp8; at the suggest_pipeline call the global tensors must carry exactly the descriptor’s value types — see the Hopper card (Blackwell: the Blackwell card); any gate missing silently leaves the wgmma path.
11, 12	Arrangement decides whether one buffer can serve both operand slots.
17	A runtime leading dimension passed to get_layout_smem_*() at the call site — not the LeadingDimension operator, which disables suggested layouts.
18	Several descriptors, selected at run time.

Expressing It The Library’s Way

• Tensors, not pointers. The tensor execute(...) overloads give you custom layouts, suggested layouts and the register accumulator. Build inputs with cublasdx::make_tensor and a layout from the descriptor.

• Layouts and accumulators pair up. suggest_layout_smem_*() = swizzled from the tuning database, get_layout_smem_*() = plain; the accumulator must match (mispairing: compile error one way, silent slowdown the other). The execute(a, b) overload that returns an accumulator picks the matching one — use it.

• Copy with the library, then wait. cublasdx::copy / cublasdx::copy_fragment pick vectorization and the async path; cublasdx::copy_wait() afterwards is required (__syncthreads() is not sufficient on SM >= 800). Copy width follows alignment. cublasdx::copy moves bytes only — a transform goes in a load op into the MMA, a shared-to-shared pass, or a hand-written load-transform-store loop; choosing among the three is idiom 03.

• Custom-layout operands keep their smem tag — cute::make_tensor(cute::make_smem_ptr(ptr), layout). With a raw cute::Layout the failure is a template error at execute() that names neither tensor nor layout.

Epilogues

A second kernel re-reads the output from global memory — for a memory-bound epilogue that costs about as much as the GEMM. Act in registers instead: execute(a, b) returns the accumulator, which carries:

Method	Use it when
process_and_store(dst, op)	Each output element is a pure function of itself.
reduce_and_store(dst, op)	The result depends on position or a second tensor; the lambda gets the accumulator element, the output element and the (m, n) coordinate — tile-local: add the tile origin for a global index.
axpby(alpha, beta, dst, load_op, store_op)	C = alpha * A * B + beta * C, with optional transforms in and out.
partition_and_store(dst)	No transform — write the fragment out.

Register output pays whenever the epilogue is elementwise or a reduction that can start per-fragment; it stops paying when the epilogue needs values a thread does not own (transpose, ordered scan, anything non-commutative) — those force the tile through shared memory anyway, so store first and work there. That split is idiom 08’s two main variants.