From 07fa04371361353dccc519df27880426f9968d88 Mon Sep 17 00:00:00 2001 From: William G Hatch Date: Tue, 7 Apr 2026 16:47:01 -0600 Subject: [PATCH] Add MultipleBuffer, MBSK, LocalSplitU, and TreeStreamK split-K MXFP4 GEMM variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit These four new split-K (split along the K/reduction dimension) strategies for MXFP4 (microscaling FP4) GEMMs explore different approaches to the partial-sum reduction problem: New algorithms: - **MultipleBuffer** (`get_tagged_multibuffer_splitk_mxfp4_gemm`): 2-kernel approach. The main kernel writes per-partition partial sums to a workspace buffer; a separate reduction kernel sums them into the output. No atomics on the output tensor C. - **MBSK** (`get_tagged_mbsk_splitk_mxfp4_gemm`): Multiple-Buffer Single-Kernel — a transitional single-kernel design that uses the workspace + a per-tile sync buffer to coordinate the reduction without a second kernel launch. - **LocalSplitU** (LSU, `get_tagged_lsu_mxfp4_gemm`): 2-kernel split-K using a sync counter to track tile completion before the reduction kernel runs. Like MultipleBuffer but with an explicit synchronization barrier, which avoids races on non-power-of-two split counts. - **TreeReducingStreamK** (`get_tagged_tree_streamk_mxfp4_gemm`): StreamK (work-stealing across CTA partitions) with a binary tree reduction instead of a linear spinlock. O(log S) fixup depth where S is the number of CTAs (Cooperative Thread Arrays, i.e. GPU thread blocks). Supporting changes: - Fix split-K correctness bug: force `use_stagger=False` when `k_partitions=1`; the stagger interacts incorrectly with single-partition pipelines that have more than one pipeline iteration. - Add shape heuristic: use `BLOCK_K=256` (instead of 128) when the K-per-split is a multiple of 256, reducing K-loop iteration count. - Add TreeStreamK shape heuristics for default `num_ctas` selection. - Add e2e tests for all four new variants (`@require_e2e @require_cdna4`). - Add test shapes that exercise `k_per_split >= 4*BLOCK_K`. WaveASM fixes required for new kernels: The new kernels exercised paths in the WaveASM backend (MLIR-to-AMD-GCN assembly compiler) that needed fixes: - `AssemblyEmitter`: emit the `s_cbranch_scc1` guard only when an `scf.for` loop guard pattern was actually matched; previously the branch was unconditionally emitted even for do-while loops that always execute at least one iteration. - `Liveness`: document the IfOp result aliasing contract (three-pass protocol through Liveness → LinearScan → AssemblyEmitter) and add an assertion to catch aliasing violations earlier. - `RegionBuilder`: minor fix. - `ArithHandlers`: strengthen the comment on the V_READFIRSTLANE_B32 promotion path to make the uniformity invariant explicit and warn against misuse. Made-with: Cursor Signed-off-by: William G Hatch --- tests/kernel/wave_gemm_test.py | 234 +++++ .../schedules/gemm_mxfp4_double_buffer.py | 8 + wave_lang/kernel/wave/templates/__init__.py | 4 + .../wave/templates/tagged_mxfp4_gemm.py | 824 +++++++++++++++++- waveasm/lib/Transforms/AssemblyEmitter.cpp | 19 +- waveasm/lib/Transforms/Liveness.cpp | 33 +- waveasm/lib/Transforms/RegionBuilder.cpp | 6 + .../lib/Transforms/handlers/ArithHandlers.cpp | 12 +- 8 files changed, 1122 insertions(+), 18 deletions(-) diff --git a/tests/kernel/wave_gemm_test.py b/tests/kernel/wave_gemm_test.py index d4c8e3a18c..7720739d5c 100644 --- a/tests/kernel/wave_gemm_test.py +++ b/tests/kernel/wave_gemm_test.py @@ -62,7 +62,11 @@ get_persistent_reordering_kernel, ) from wave_lang.kernel.wave.templates.tagged_mxfp4_gemm import ( + get_tagged_lsu_mxfp4_gemm, + get_tagged_mbsk_splitk_mxfp4_gemm, + get_tagged_multibuffer_splitk_mxfp4_gemm, get_tagged_splitk_mxfp4_gemm, + get_tagged_tree_streamk_mxfp4_gemm, ) from wave_lang.kernel.wave.templates.test_kernels import ( get_gemm_prefetch_kernel_and_schedule, @@ -83,6 +87,7 @@ ) from wave_lang.kernel.lang import DataType +import math import os import json from torch.testing import assert_close @@ -3647,6 +3652,9 @@ def testSplitKGemm( ((256, 256, 512), 4), ((256, 256, 1024), 4), ((512, 512, 1024), 4), + ((256, 256, 1024), 2), + ((256, 256, 2048), 2), + ((256, 256, 8192), 2), ], ) @pytest.mark.parametrize("output_type", [torch.float32, torch.bfloat16]) @@ -3688,3 +3696,229 @@ def testSplitKMxfp4Gemm( # partial-sum magnitudes of ~250-430 the ULP is 2-4, so with 2-4 # splits the worst-case rounding error is O(num_splits * ULP). assert_close(c_gpu.cpu().to(torch.float32), torch_ref, rtol=1e-1, atol=4.0) + + +@require_e2e +@require_cdna4 +@pytest.mark.parametrize( + "shape, num_splits", + [ + ((256, 256, 256), 2), + ((256, 256, 512), 2), + ((256, 256, 512), 4), + ((256, 256, 1024), 4), + ((512, 512, 1024), 4), + ], +) +def testMultiBufferSplitKMxfp4Gemm( + shape: tuple[int, int, int], + num_splits: int, +): + ( + main_fn, + main_options, + reduction_fn, + reduction_options, + ) = get_tagged_multibuffer_splitk_mxfp4_gemm( + shape, + num_splits=num_splits, + block_shape=(128, 128, 128), + mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4, + output_type=TORCH_DTYPE_TO_WAVE[torch.float32], + ) + main_options = set_default_run_config(main_options) + reduction_options = set_default_run_config(reduction_options) + schedule = get_mxfp4_dbuf_schedule(use_stagger=True, k_partitions=1) + compiled_main = wave_compile(main_options, main_fn, schedule) + compiled_reduction = wave_compile(reduction_options, reduction_fn) + + m, n, _ = shape + x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs( + shape, device=torch.device("cpu") + ) + torch_ref = torchScaledGemmMXFP4(x, w, x_scales, w_scales) + + x_gpu = x.cuda() + w_t_gpu = w.T.contiguous().cuda() + x_scales_gpu = x_scales.cuda() + w_scales_gpu = w_scales.cuda() + workspace = device_zeros(num_splits, m, n, dtype=torch.float32) + c_gpu = device_zeros(m, n, dtype=torch.float32) + + compiled_main(x_gpu, x_scales_gpu, w_t_gpu, w_scales_gpu, workspace) + compiled_reduction(workspace, c_gpu) + assert_close(c_gpu.cpu(), torch_ref, rtol=1e-3, atol=1e-2) + + +@require_e2e +@require_cdna4 +@pytest.mark.parametrize( + "shape, num_splits", + [ + ((256, 256, 256), 2), + ((256, 256, 512), 2), + ((256, 256, 512), 4), + ((256, 256, 1024), 4), + ], +) +def testMBSKSplitKMxfp4Gemm( + shape: tuple[int, int, int], + num_splits: int, +): + block_shape = (128, 128, 128) + mbsk_gemm, options = get_tagged_mbsk_splitk_mxfp4_gemm( + shape, + num_splits=num_splits, + block_shape=block_shape, + mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4, + output_type=TORCH_DTYPE_TO_WAVE[torch.float32], + ) + options = set_default_run_config(options) + schedule = get_mxfp4_dbuf_schedule(use_stagger=True, k_partitions=1) + compiled = wave_compile(options, mbsk_gemm, schedule) + + m, n, _ = shape + x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs( + shape, device=torch.device("cpu") + ) + torch_ref = torchScaledGemmMXFP4(x, w, x_scales, w_scales) + + x_gpu = x.cuda() + w_t_gpu = w.T.contiguous().cuda() + x_scales_gpu = x_scales.cuda() + w_scales_gpu = w_scales.cuda() + block_m, block_n, _ = block_shape + num_tiles = math.ceil(m / block_m) * math.ceil(n / block_n) + workspace = device_zeros(num_splits, m, n, dtype=torch.float32) + sync_buffer = device_zeros(num_tiles, dtype=torch.int32) + c_gpu = device_zeros(m, n, dtype=torch.float32) + + compiled( + x_gpu, + x_scales_gpu, + w_t_gpu, + w_scales_gpu, + workspace, + sync_buffer, + c_gpu, + ) + assert_close(c_gpu.cpu(), torch_ref, rtol=1e-3, atol=1e-2) + + +@require_e2e +@require_cdna4 +@pytest.mark.parametrize( + "shape, lsu_factor", + [ + ((256, 256, 256), 2), + ((256, 256, 512), 2), + ((256, 256, 1024), 2), + ((512, 512, 1024), 2), + ((256, 256, 512), 4), + ((256, 256, 1024), 4), + ((256, 256, 2048), 4), + ], +) +@pytest.mark.parametrize("output_type", [torch.float32]) +def testLocalSplitUMxfp4Gemm( + shape: tuple[int, int, int], + lsu_factor: int, + output_type: torch.dtype, +): + """LocalSplitU: 2-kernel split-K with workspace reduction, no atomics on C. + + Main kernel writes f32 partials to workspace. + Reduction kernel sums partials and writes the final result to C. + """ + block_shape = (128, 128, 128) + ( + main_fn, + main_options, + reduction_fn, + reduction_options, + ) = get_tagged_lsu_mxfp4_gemm( + shape, + lsu_factor=lsu_factor, + block_shape=block_shape, + mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4, + output_type=TORCH_DTYPE_TO_WAVE[output_type], + ) + + main_options = set_default_run_config(main_options) + reduction_options = set_default_run_config(reduction_options) + schedule = get_mxfp4_dbuf_schedule(use_stagger=True, k_partitions=1) + compiled_main = wave_compile(main_options, main_fn, schedule) + compiled_reduction = wave_compile(reduction_options, reduction_fn) + + m, n, k = shape + x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs( + shape, device=torch.device("cpu") + ) + torch_ref = torchScaledGemmMXFP4(x, w, x_scales, w_scales) + + x_gpu = x.cuda() + w_t_gpu = w.T.contiguous().cuda() + x_scales_gpu = x_scales.cuda() + w_scales_gpu = w_scales.cuda() + block_m, block_n, _ = block_shape + num_tiles = math.ceil(m / block_m) * math.ceil(n / block_n) + workspace = device_zeros(lsu_factor, m, n, dtype=torch.float32) + sync_counter = device_zeros(num_tiles, dtype=torch.int32) + c_gpu = device_zeros(m, n, dtype=output_type) + + compiled_main(x_gpu, x_scales_gpu, w_t_gpu, w_scales_gpu, workspace, sync_counter) + compiled_reduction(workspace, c_gpu) + assert_close(c_gpu.cpu(), torch_ref, rtol=1e-3, atol=1e-2) + + +@require_e2e +@require_cdna4 +@pytest.mark.parametrize( + "shape, num_ctas", + [ + ((128, 128, 128), 1), + ((128, 128, 256), 2), + ((256, 256, 256), 8), + ((256, 256, 512), 16), + ((512, 512, 1024), 16), + ], +) +@pytest.mark.parametrize("output_type", [torch.float32]) +def testTreeStreamKMxfp4Gemm( + shape: tuple[int, int, int], + num_ctas: int, + output_type: torch.dtype, +): + """StreamK with tree reduction: single-kernel, O(log n) fixup depth. + + Uses workspace + flag buffer for binary tree reduction instead of + linear spinlock or atomic_add. + """ + m, n, k = shape + block_m, block_n, block_k = 128, 128, 128 + + streamk_gemm, options = get_tagged_tree_streamk_mxfp4_gemm( + shape, + block_shape=(block_m, block_n, block_k), + mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4, + output_type=TORCH_DTYPE_TO_WAVE[output_type], + num_ctas=num_ctas, + ) + + options = set_default_run_config(options) + compiled = wave_compile(options, streamk_gemm) + + x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs( + shape, device=torch.device("cpu") + ) + torch_ref = torchScaledGemmMXFP4(x, w, x_scales, w_scales) + + x_gpu = x.cuda() + w_t_gpu = w.T.contiguous().cuda() + x_scales_gpu = x_scales.cuda() + w_scales_gpu = w_scales.cuda() + + c_gpu = device_zeros(m, n, dtype=output_type) + + compiled(x_gpu, x_scales_gpu, w_t_gpu, w_scales_gpu, c_gpu) + assert_close(c_gpu.cpu(), torch_ref, rtol=1e-3, atol=1e-2) diff --git a/wave_lang/kernel/wave/schedules/gemm_mxfp4_double_buffer.py b/wave_lang/kernel/wave/schedules/gemm_mxfp4_double_buffer.py index e67d5c88e9..49bd8f814b 100755 --- a/wave_lang/kernel/wave/schedules/gemm_mxfp4_double_buffer.py +++ b/wave_lang/kernel/wave/schedules/gemm_mxfp4_double_buffer.py @@ -45,9 +45,13 @@ def get_mxfp4_dbuf_schedule(use_stagger: bool = True, k_partitions: int = 2): Args: use_stagger: Enable wave staggering + WorkgroupBarrier in cluster 0. Recommended for 8-wave configs; disable for 4-wave. + Forced off when ``k_partitions=1`` because the stagger interacts + incorrectly with single-partition pipeline iterations > 1. k_partitions: ``2`` for standard GEMM (partition along K). Use ``1`` for split-K MXFP4 kernels where the K loop has only one expansion id. """ + if k_partitions == 1: + use_stagger = False K = tkl.sym.K @wave_schedule.wave_schedule() @@ -289,11 +293,15 @@ def get_mxfp4_dbuf_pingpong_schedule( Args: use_stagger: Enable wave staggering + WorkgroupBarrier in cluster 0. Recommended for 8-wave configs; disable for 4-wave. + Forced off when ``k_partitions=1`` because the stagger interacts + incorrectly with single-partition pipeline iterations > 1. shape: Tuple of (M, N, K) dimensions. If provided and bigger than (1024, 1024, 1024), an extra WorkgroupBarrier will be added after the first SchedulingBarrier in cluster 0. k_partitions: ``2`` for standard GEMM. Use ``1`` for split-K MXFP4 kernels. """ + if k_partitions == 1: + use_stagger = False K = tkl.sym.K @wave_schedule.wave_schedule() diff --git a/wave_lang/kernel/wave/templates/__init__.py b/wave_lang/kernel/wave/templates/__init__.py index 30a6b6b3cb..c4e1bda8f4 100644 --- a/wave_lang/kernel/wave/templates/__init__.py +++ b/wave_lang/kernel/wave/templates/__init__.py @@ -7,6 +7,7 @@ from .attention_common import AttentionShape from .tagged_attention import get_tagged_bshd_attention_kernel from .tagged_mxfp4_gemm import ( + get_tagged_lsu_mxfp4_gemm, get_tagged_mxfp4_gemm, get_tagged_mxfp4_gemm_preshuffle_b, get_tagged_mxfp4_gemm_preshuffle_b_wide_store, @@ -15,11 +16,13 @@ get_tagged_splitk_mxfp4_gemm, get_tagged_splitk_mxfp4_gemm_preshuffle_b, get_tagged_splitk_mxfp4_gemm_preshuffle_scales, + get_tagged_tree_streamk_mxfp4_gemm, ) __all__ = [ "AttentionShape", "get_tagged_bshd_attention_kernel", + "get_tagged_lsu_mxfp4_gemm", "get_tagged_mxfp4_gemm", "get_tagged_mxfp4_gemm_preshuffle_b", "get_tagged_mxfp4_gemm_preshuffle_b_wide_store", @@ -28,4 +31,5 @@ "get_tagged_splitk_mxfp4_gemm", "get_tagged_splitk_mxfp4_gemm_preshuffle_b", "get_tagged_splitk_mxfp4_gemm_preshuffle_scales", + "get_tagged_tree_streamk_mxfp4_gemm", ] diff --git a/wave_lang/kernel/wave/templates/tagged_mxfp4_gemm.py b/wave_lang/kernel/wave/templates/tagged_mxfp4_gemm.py index 7c23ff0c7d..b39f56f8ec 100755 --- a/wave_lang/kernel/wave/templates/tagged_mxfp4_gemm.py +++ b/wave_lang/kernel/wave/templates/tagged_mxfp4_gemm.py @@ -10,12 +10,16 @@ All ops are tagged for use with MXFP4 schedule functions (e.g. get_mxfp4_dbuf_schedule). Provides: - - get_tagged_mxfp4_gemm: vanilla (A, B via LDS) - - get_tagged_mxfp4_gemm_preshuffle_b: B + B_scale preshuffled (direct global reads) - - get_tagged_mxfp4_gemm_preshuffle_b_wide_store: same + wide epilogue stores via permlane swap - - get_tagged_splitk_mxfp4_gemm: split-K (A, B, scales via LDS) - - get_tagged_splitk_mxfp4_gemm_preshuffle_scales: split-K with preshuffled scales - - get_tagged_splitk_mxfp4_gemm_preshuffle_b: split-K with preshuffled B + scales + - get_tagged_mxfp4_gemm: vanilla (A, B via LDS) + - get_tagged_mxfp4_gemm_preshuffle_b: B + B_scale preshuffled (direct global reads) + - get_tagged_mxfp4_gemm_preshuffle_b_wide_store: same + wide epilogue stores via permlane swap + - get_tagged_splitk_mxfp4_gemm: split-K (A, B, scales via LDS) + - get_tagged_splitk_mxfp4_gemm_preshuffle_scales: split-K with preshuffled scales + - get_tagged_splitk_mxfp4_gemm_preshuffle_b: split-K with preshuffled B + scales + - get_tagged_multibuffer_splitk_mxfp4_gemm: split-K workspace + separate reduction (no atomics) + - get_tagged_mbsk_splitk_mxfp4_gemm: split-K workspace + sync buffer, single kernel + - get_tagged_lsu_mxfp4_gemm: LocalSplitU (intra-WG K-split via LDS reduction) + - get_tagged_tree_streamk_mxfp4_gemm: stream-K with binary tree reduction Required tags: k_loop, read_a, read_a_scale, read_b, read_b_scale, bitcast_a, bitcast_a_scale, bitcast_b, bitcast_b_scale, scaled_mma. @@ -1035,6 +1039,344 @@ def get_tagged_mxfp4_gemm_preshuffle_b_wide_store( ) +def get_tagged_multibuffer_splitk_mxfp4_gemm( + shape: tuple[int, int, int] = (1024, 1024, 8192), + num_splits: int = 2, + block_shape: tuple[int, int, int] = (128, 128, 256), + wave_shape: tuple[int, int] = (2, 2), + mfma_variant: ScaledMMAType = ScaledMMAType.F32_16x16x128_F8F6F4, + a_address_space: tkl.AddressSpace = SHARED_ADDRESS_SPACE, + output_type: "tkl.DataType" = tkl.f32, +): + """Return tagged split-K MXFP4 GEMM main kernel, options, reduction kernel, and options. + + MultipleBuffer split-K writes each split's partial result to a workspace buffer, + then a separate reduction kernel sums partials into C (no atomic_add on C). + + The main kernel uses the same tags and distribution constraints as + ``get_tagged_splitk_mxfp4_gemm`` so MXFP4 schedules apply. Partials are + always written as ``f32`` to ``workspace[S, M, N]``. The reduction kernel + sums along ``S`` and writes ``c`` in ``output_type``. + + Returns: + (main_kernel_fn, main_options, reduction_kernel_fn, reduction_options) + """ + m, n, k = shape + k_per_split = math.ceil(k / num_splits) + if k_per_split < block_shape[2]: + raise ValueError( + f"K per split ({k_per_split}) is less than BLOCK_K ({block_shape[2]}). " + f"Reduce num_splits or BLOCK_K so that each split has at least BLOCK_K elements." + ) + if k_per_split % block_shape[2] != 0: + raise ValueError( + f"k_per_split ({k_per_split}) must be a multiple of BLOCK_K ({block_shape[2]})." + ) + + M = tkl.sym.M + N = tkl.sym.N + K = tkl.sym.K + S = tkl.sym.S + BLOCK_M = tkl.sym.BLOCK_M + BLOCK_N = tkl.sym.BLOCK_N + BLOCK_K = tkl.sym.BLOCK_K + BLOCK_S = tkl.sym.BLOCK_S + K_SPLIT_OFF = tkl.sym.K_SPLIT_OFF + K_SPLIT_LEN = tkl.sym.K_SPLIT_LEN + ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE + B_ADDRESS_SPACE = tkl.sym.B_ADDRESS_SPACE + + # Same structure as get_tagged_splitk_mxfp4_gemm (no preshuffle). + constraints_main: list[tkw.Constraint] = [ + tkw.WorkgroupConstraint(M, BLOCK_M, 0), + tkw.WorkgroupConstraint(N, BLOCK_N, 1), + tkw.WorkgroupConstraint(S, BLOCK_S, 2), + tkw.TilingConstraint( + K, + BLOCK_K, + iters=sympy.ceiling(K_SPLIT_LEN / BLOCK_K), + start=K_SPLIT_OFF, + ), + tkw.WaveConstraint(M, sympy.floor(BLOCK_M / wave_shape[0])), + tkw.WaveConstraint(N, sympy.floor(BLOCK_N / wave_shape[1])), + tkw.HardwareConstraint( + threads_per_wave=64, + mma_type=mfma_variant, + vector_shapes={S: 0}, + ), + ] + + @tkw.wave(constraints_main) + def splitk_multibuffer_main( + a: tkl.Memory[M, K / 2, ADDRESS_SPACE, tkl.i8], + a_scale: tkl.Memory[M, K / 32, ADDRESS_SPACE, tkl.i8], + b: tkl.Memory[N, K / 2, B_ADDRESS_SPACE, tkl.i8], + b_scale: tkl.Memory[N, K / 32, ADDRESS_SPACE, tkl.i8], + workspace: tkl.Memory[S, M, N, GLOBAL_ADDRESS_SPACE, tkl.f32], + ): + c_reg = tkl.Register[M, N, tkl.f32](0.0) + + @tkw.iterate(K, init_args=[c_reg], tag="k_loop") + def repeat( + acc: tkl.Register[M, N, tkl.f32], + ) -> tkl.Register[M, N, tkl.f32]: + a_reg = tkw.read(a, tag="read_a") + a_reg = tkw.bitcast(a_reg, tkl.f4e2m1fn, tag="bitcast_a") + a_scale_reg = tkw.read(a_scale, tag="read_a_scale") + a_scale_reg = tkw.bitcast(a_scale_reg, tkl.f8e8m0fnu, tag="bitcast_a_scale") + b_reg = tkw.read(b, tag="read_b") + b_reg = tkw.bitcast(b_reg, tkl.f4e2m1fn, tag="bitcast_b") + b_scale_reg = tkw.read(b_scale, tag="read_b_scale") + b_scale_reg = tkw.bitcast(b_scale_reg, tkl.f8e8m0fnu, tag="bitcast_b_scale") + acc = tkw.scaled_mma( + a_reg, a_scale_reg, b_reg, b_scale_reg, acc, tag="scaled_mma" + ) + return acc + + tkw.write(repeat, workspace) + + hyperparams_main = { + ADDRESS_SPACE: a_address_space, + B_ADDRESS_SPACE: a_address_space, + BLOCK_M: block_shape[0], + BLOCK_N: block_shape[1], + BLOCK_K: block_shape[2], + BLOCK_S: 1, + M: m, + N: n, + K: k, + S: num_splits, + K_SPLIT_OFF: WORKGROUP_2 * k_per_split, + K_SPLIT_LEN: sympy.Min(K, (WORKGROUP_2 + 1) * k_per_split) - K_SPLIT_OFF, + } + for key, value in hyperparams_main.items(): + if isinstance(value, sympy.Expr): + hyperparams_main[key] = value.subs(hyperparams_main) + + hyperparams_main.update(get_default_scheduling_params()) + + main_options = WaveCompileOptions( + subs=hyperparams_main, + canonicalize=True, + schedule=SchedulingType.MANUAL, + use_global_to_shared=True, + minimize_shared_allocs=False, + ) + + NUM_SPLITS = tkl.sym.NUM_SPLITS + BLOCK_RED = tkl.sym.BLOCK_RED + wave_m = block_shape[0] // wave_shape[0] + wave_n = block_shape[1] // wave_shape[1] + + # Sum splits with tkw.iterate (not tkw.sum(dim=NUM_SPLITS): that path requires + # vector_shapes[NUM_SPLITS] divisible by threads_per_wave=64). + constraints_red: list[tkw.Constraint] = [ + tkw.WorkgroupConstraint(M, BLOCK_M, 0), + tkw.WorkgroupConstraint(N, BLOCK_N, 1), + tkw.TilingConstraint(NUM_SPLITS, BLOCK_RED), + tkw.WaveConstraint(M, sympy.floor(BLOCK_M / wave_shape[0])), + tkw.WaveConstraint(N, sympy.floor(BLOCK_N / wave_shape[1])), + tkw.HardwareConstraint( + threads_per_wave=64, + mma_type=mfma_variant, + vector_shapes={ + M: wave_m, + N: wave_n, + NUM_SPLITS: 1, + }, + ), + ] + + @tkw.wave(constraints_red) + def splitk_multibuffer_reduction( + workspace: tkl.Memory[NUM_SPLITS, M, N, GLOBAL_ADDRESS_SPACE, tkl.f32], + c: tkl.Memory[M, N, GLOBAL_ADDRESS_SPACE, output_type], + ): + acc = tkl.Register[M, N, tkl.f32](0.0) + + @tkw.iterate(NUM_SPLITS, init_args=[acc]) + def sum_splits( + acc: tkl.Register[M, N, tkl.f32], + ) -> tkl.Register[M, N, tkl.f32]: + partial = tkw.read(workspace) + return acc + partial + + out = tkw.cast(sum_splits, output_type) + tkw.write(out, c) + + hyperparams_red = { + BLOCK_M: block_shape[0], + BLOCK_N: block_shape[1], + BLOCK_RED: 1, + M: m, + N: n, + NUM_SPLITS: num_splits, + } + hyperparams_red.update(get_default_scheduling_params()) + + reduction_options = WaveCompileOptions( + subs=hyperparams_red, + canonicalize=True, + schedule=SchedulingType.NONE, + use_global_to_shared=False, + minimize_shared_allocs=False, + ) + + return ( + splitk_multibuffer_main, + main_options, + splitk_multibuffer_reduction, + reduction_options, + ) + + +def get_tagged_mbsk_splitk_mxfp4_gemm( + shape: tuple[int, int, int] = (1024, 1024, 8192), + num_splits: int = 2, + block_shape: tuple[int, int, int] = (128, 128, 256), + wave_shape: tuple[int, int] = (2, 2), + mfma_variant: ScaledMMAType = ScaledMMAType.F32_16x16x128_F8F6F4, + a_address_space: tkl.AddressSpace = SHARED_ADDRESS_SPACE, + output_type: "tkl.DataType" = tkl.f32, +): + """Return tagged MBSK (MultipleBufferSingleKernel) split-K MXFP4 GEMM + compile options. + + Transitional single-kernel implementation: each split workgroup writes its + ``f32`` partial to ``workspace[S, M, N]`` (same layout as the multibuffer + main kernel) and accumulates into ``C`` via ``atomic_add``, so callers get + a correct result with one launch. Each workgroup performs a no-op ``atomic_add(0, sync_buffer[tile])`` where + ``tile = WORKGROUP_0 + WORKGROUP_1 * ceil(M / BLOCK_M)``, keeping the sync + buffer in the launch signature and matching the test layout. True + tile-level synchronization is reserved for a later kernel. + """ + m, n, k = shape + k_per_split = math.ceil(k / num_splits) + if k_per_split < block_shape[2]: + raise ValueError( + f"K per split ({k_per_split}) is less than BLOCK_K ({block_shape[2]}). " + f"Reduce num_splits or BLOCK_K so that each split has at least BLOCK_K elements." + ) + if k_per_split % block_shape[2] != 0: + raise ValueError( + f"k_per_split ({k_per_split}) must be a multiple of BLOCK_K ({block_shape[2]})." + ) + + num_wg_m = math.ceil(m / block_shape[0]) + num_wg_n = math.ceil(n / block_shape[1]) + num_tiles = num_wg_m * num_wg_n + + M = tkl.sym.M + N = tkl.sym.N + K = tkl.sym.K + S = tkl.sym.S + BLOCK_M = tkl.sym.BLOCK_M + BLOCK_N = tkl.sym.BLOCK_N + BLOCK_K = tkl.sym.BLOCK_K + BLOCK_S = tkl.sym.BLOCK_S + K_SPLIT_OFF = tkl.sym.K_SPLIT_OFF + K_SPLIT_LEN = tkl.sym.K_SPLIT_LEN + ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE + B_ADDRESS_SPACE = tkl.sym.B_ADDRESS_SPACE + SYNC_SIZE = tkl.sym.SYNC_SIZE + + NUM_WG_M = sympy.ceiling(M / BLOCK_M) + i_s = tkw.IndexMapping.iterator(0) + # Atomic lowering keys the mapping by the lhs register shape (S); the memory + # is 1-D sync_buffer[SYNC_SIZE] with index (wg0, wg1) flattened. + sync_atomic_mapping = tkw.IndexMapping( + num_iterators=1, + inputs={S: WORKGROUP_0 + WORKGROUP_1 * NUM_WG_M}, + outputs={S: i_s}, + ) + + constraints: list[tkw.Constraint] = [ + tkw.WorkgroupConstraint(M, BLOCK_M, 0), + tkw.WorkgroupConstraint(N, BLOCK_N, 1), + tkw.WorkgroupConstraint(S, BLOCK_S, 2), + tkw.TilingConstraint( + K, + BLOCK_K, + iters=sympy.ceiling(K_SPLIT_LEN / BLOCK_K), + start=K_SPLIT_OFF, + ), + tkw.WaveConstraint(M, sympy.floor(BLOCK_M / wave_shape[0])), + tkw.WaveConstraint(N, sympy.floor(BLOCK_N / wave_shape[1])), + tkw.HardwareConstraint( + threads_per_wave=64, + mma_type=mfma_variant, + vector_shapes={S: 0}, + ), + ] + + @tkw.wave(constraints) + def mbsk_splitk_gemm( + a: tkl.Memory[M, K / 2, ADDRESS_SPACE, tkl.i8], + a_scale: tkl.Memory[M, K / 32, ADDRESS_SPACE, tkl.i8], + b: tkl.Memory[N, K / 2, B_ADDRESS_SPACE, tkl.i8], + b_scale: tkl.Memory[N, K / 32, ADDRESS_SPACE, tkl.i8], + workspace: tkl.Memory[S, M, N, GLOBAL_ADDRESS_SPACE, tkl.f32], + sync_buffer: tkl.Memory[SYNC_SIZE, GLOBAL_ADDRESS_SPACE, tkl.i32], + c: tkl.Memory[M, N, GLOBAL_ADDRESS_SPACE, output_type], + ): + c_reg = tkl.Register[M, N, tkl.f32](0.0) + + @tkw.iterate(K, init_args=[c_reg], tag="k_loop") + def repeat( + acc: tkl.Register[M, N, tkl.f32], + ) -> tkl.Register[M, N, tkl.f32]: + a_reg = tkw.read(a, tag="read_a") + a_reg = tkw.bitcast(a_reg, tkl.f4e2m1fn, tag="bitcast_a") + a_scale_reg = tkw.read(a_scale, tag="read_a_scale") + a_scale_reg = tkw.bitcast(a_scale_reg, tkl.f8e8m0fnu, tag="bitcast_a_scale") + b_reg = tkw.read(b, tag="read_b") + b_reg = tkw.bitcast(b_reg, tkl.f4e2m1fn, tag="bitcast_b") + b_scale_reg = tkw.read(b_scale, tag="read_b_scale") + b_scale_reg = tkw.bitcast(b_scale_reg, tkl.f8e8m0fnu, tag="bitcast_b_scale") + acc = tkw.scaled_mma( + a_reg, a_scale_reg, b_reg, b_scale_reg, acc, tag="scaled_mma" + ) + return acc + + tkw.write(repeat, workspace) + # No-op atomic on sync_buffer per tile so the buffer stays live across + # compilation (reads folded away) and records tile participation later. + zero_s = tkl.Register[S, tkl.i32](0) + tkw.atomic_add(zero_s, sync_buffer, mapping=sync_atomic_mapping) + repeat_out = tkw.cast(repeat, output_type) + tkw.atomic_add(repeat_out, c) + + hyperparams = { + ADDRESS_SPACE: a_address_space, + B_ADDRESS_SPACE: a_address_space, + BLOCK_M: block_shape[0], + BLOCK_N: block_shape[1], + BLOCK_K: block_shape[2], + BLOCK_S: 1, + M: m, + N: n, + K: k, + S: num_splits, + SYNC_SIZE: num_tiles, + K_SPLIT_OFF: WORKGROUP_2 * k_per_split, + K_SPLIT_LEN: sympy.Min(K, (WORKGROUP_2 + 1) * k_per_split) - K_SPLIT_OFF, + } + for key, value in hyperparams.items(): + if isinstance(value, sympy.Expr): + hyperparams[key] = value.subs(hyperparams) + + hyperparams.update(get_default_scheduling_params()) + + options = WaveCompileOptions( + subs=hyperparams, + canonicalize=True, + schedule=SchedulingType.MANUAL, + use_global_to_shared=True, + minimize_shared_allocs=False, + ) + + return mbsk_splitk_gemm, options + + def _reorder_mxfp4_workgroups(m, n, block_m, block_n, group_size_n): """Remap workgroup indices to a new order based on group_size_n along N dimension. @@ -1087,3 +1429,473 @@ def _reorder_mxfp4_workgroups(m, n, block_m, block_n, group_size_n): ) return new_wg0, new_wg1 + + +def get_tagged_lsu_mxfp4_gemm( + shape: tuple[int, int, int] = (1024, 1024, 8192), + lsu_factor: int = 2, + block_shape: tuple[int, int, int] = (128, 128, 256), + wave_shape: tuple[int, int] = (2, 2), + mfma_variant: ScaledMMAType = ScaledMMAType.F32_16x16x128_F8F6F4, + a_address_space: tkl.AddressSpace = SHARED_ADDRESS_SPACE, + output_type: "tkl.DataType" = tkl.f32, +): + """Return tagged LocalSplitU MXFP4 GEMM main kernel, options, reduction kernel, and options. + + LocalSplitU splits the K dimension across ``lsu_factor`` workgroups. The + main kernel writes ``f32`` partials to ``workspace[S, M, N]``. A separate + reduction kernel sums partials into ``C`` without atomics. + + The sync_counter argument is accepted by the main kernel for API + compatibility (it atomically increments a per-tile counter for profiling + and future single-kernel upgrades) but the reduction is performed by the + second kernel. + + Args: + shape: (M, N, K) problem dimensions. + lsu_factor: Number of K-slices per workgroup (2 or 4). + block_shape: (BLOCK_M, BLOCK_N, BLOCK_K) tile sizes. + wave_shape: (WAVE_M, WAVE_N) waves per workgroup for M,N dims. + mfma_variant: Scaled MMA instruction type. + a_address_space: Address space for A and A_scale (typically SHARED). + output_type: Element type of output tensor C. + + Returns: + (main_kernel_fn, main_options, reduction_kernel_fn, reduction_options) + """ + m, n, k = shape + k_per_split = math.ceil(k / lsu_factor) + if k_per_split < block_shape[2]: + raise ValueError( + f"K per split ({k_per_split}) is less than BLOCK_K ({block_shape[2]}). " + f"Reduce lsu_factor or BLOCK_K so that each split has at least BLOCK_K elements." + ) + if k_per_split % block_shape[2] != 0: + raise ValueError( + f"k_per_split ({k_per_split}) must be a multiple of BLOCK_K ({block_shape[2]})." + ) + if lsu_factor not in (2, 4): + raise ValueError(f"lsu_factor must be 2 or 4, got {lsu_factor}") + + num_wg_m = math.ceil(m / block_shape[0]) + num_wg_n = math.ceil(n / block_shape[1]) + num_tiles = num_wg_m * num_wg_n + + M = tkl.sym.M + N = tkl.sym.N + K = tkl.sym.K + S = tkl.sym.S + BLOCK_M = tkl.sym.BLOCK_M + BLOCK_N = tkl.sym.BLOCK_N + BLOCK_K = tkl.sym.BLOCK_K + BLOCK_S = tkl.sym.BLOCK_S + K_SPLIT_OFF = tkl.sym.K_SPLIT_OFF + K_SPLIT_LEN = tkl.sym.K_SPLIT_LEN + ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE + B_ADDRESS_SPACE = tkl.sym.B_ADDRESS_SPACE + SYNC_SIZE = tkl.sym.SYNC_SIZE + + NUM_WG_M = sympy.ceiling(M / BLOCK_M) + i_s = tkw.IndexMapping.iterator(0) + sync_atomic_mapping = tkw.IndexMapping( + num_iterators=1, + inputs={S: WORKGROUP_0 + WORKGROUP_1 * NUM_WG_M}, + outputs={S: i_s}, + ) + + constraints_main: list[tkw.Constraint] = [ + tkw.WorkgroupConstraint(M, BLOCK_M, 0), + tkw.WorkgroupConstraint(N, BLOCK_N, 1), + tkw.WorkgroupConstraint(S, BLOCK_S, 2), + tkw.TilingConstraint( + K, + BLOCK_K, + iters=sympy.ceiling(K_SPLIT_LEN / BLOCK_K), + start=K_SPLIT_OFF, + ), + tkw.WaveConstraint(M, sympy.floor(BLOCK_M / wave_shape[0])), + tkw.WaveConstraint(N, sympy.floor(BLOCK_N / wave_shape[1])), + tkw.HardwareConstraint( + threads_per_wave=64, + mma_type=mfma_variant, + vector_shapes={S: 0}, + ), + ] + + @tkw.wave(constraints_main) + def lsu_main( + a: tkl.Memory[M, K / 2, ADDRESS_SPACE, tkl.i8], + a_scale: tkl.Memory[M, K / 32, ADDRESS_SPACE, tkl.i8], + b: tkl.Memory[N, K / 2, B_ADDRESS_SPACE, tkl.i8], + b_scale: tkl.Memory[N, K / 32, ADDRESS_SPACE, tkl.i8], + workspace: tkl.Memory[S, M, N, GLOBAL_ADDRESS_SPACE, tkl.f32], + sync_counter: tkl.Memory[SYNC_SIZE, GLOBAL_ADDRESS_SPACE, tkl.i32], + ): + c_reg = tkl.Register[M, N, tkl.f32](0.0) + + @tkw.iterate(K, init_args=[c_reg], tag="k_loop") + def repeat( + acc: tkl.Register[M, N, tkl.f32], + ) -> tkl.Register[M, N, tkl.f32]: + a_reg = tkw.read(a, tag="read_a") + a_reg = tkw.bitcast(a_reg, tkl.f4e2m1fn, tag="bitcast_a") + a_scale_reg = tkw.read(a_scale, tag="read_a_scale") + a_scale_reg = tkw.bitcast(a_scale_reg, tkl.f8e8m0fnu, tag="bitcast_a_scale") + b_reg = tkw.read(b, tag="read_b") + b_reg = tkw.bitcast(b_reg, tkl.f4e2m1fn, tag="bitcast_b") + b_scale_reg = tkw.read(b_scale, tag="read_b_scale") + b_scale_reg = tkw.bitcast(b_scale_reg, tkl.f8e8m0fnu, tag="bitcast_b_scale") + acc = tkw.scaled_mma( + a_reg, a_scale_reg, b_reg, b_scale_reg, acc, tag="scaled_mma" + ) + return acc + + tkw.write(repeat, workspace) + + one_s = tkl.Register[S, tkl.i32](1) + tkw.atomic_add(one_s, sync_counter, mapping=sync_atomic_mapping) + + hyperparams_main = { + ADDRESS_SPACE: a_address_space, + B_ADDRESS_SPACE: a_address_space, + BLOCK_M: block_shape[0], + BLOCK_N: block_shape[1], + BLOCK_K: block_shape[2], + BLOCK_S: 1, + M: m, + N: n, + K: k, + S: lsu_factor, + SYNC_SIZE: num_tiles, + K_SPLIT_OFF: WORKGROUP_2 * k_per_split, + K_SPLIT_LEN: sympy.Min(K, (WORKGROUP_2 + 1) * k_per_split) - K_SPLIT_OFF, + } + for key, value in hyperparams_main.items(): + if isinstance(value, sympy.Expr): + hyperparams_main[key] = value.subs(hyperparams_main) + + hyperparams_main.update(get_default_scheduling_params()) + + main_options = WaveCompileOptions( + subs=hyperparams_main, + canonicalize=True, + schedule=SchedulingType.MANUAL, + use_global_to_shared=True, + minimize_shared_allocs=False, + ) + + NUM_SPLITS = tkl.sym.NUM_SPLITS + BLOCK_RED = tkl.sym.BLOCK_RED + wave_m = block_shape[0] // wave_shape[0] + wave_n = block_shape[1] // wave_shape[1] + + constraints_red: list[tkw.Constraint] = [ + tkw.WorkgroupConstraint(M, BLOCK_M, 0), + tkw.WorkgroupConstraint(N, BLOCK_N, 1), + tkw.TilingConstraint(NUM_SPLITS, BLOCK_RED), + tkw.WaveConstraint(M, sympy.floor(BLOCK_M / wave_shape[0])), + tkw.WaveConstraint(N, sympy.floor(BLOCK_N / wave_shape[1])), + tkw.HardwareConstraint( + threads_per_wave=64, + mma_type=mfma_variant, + vector_shapes={ + M: wave_m, + N: wave_n, + NUM_SPLITS: 1, + }, + ), + ] + + @tkw.wave(constraints_red) + def lsu_reduction( + workspace: tkl.Memory[NUM_SPLITS, M, N, GLOBAL_ADDRESS_SPACE, tkl.f32], + c: tkl.Memory[M, N, GLOBAL_ADDRESS_SPACE, output_type], + ): + acc = tkl.Register[M, N, tkl.f32](0.0) + + @tkw.iterate(NUM_SPLITS, init_args=[acc]) + def sum_splits( + acc: tkl.Register[M, N, tkl.f32], + ) -> tkl.Register[M, N, tkl.f32]: + partial = tkw.read(workspace) + return acc + partial + + out = tkw.cast(sum_splits, output_type) + tkw.write(out, c) + + hyperparams_red = { + BLOCK_M: block_shape[0], + BLOCK_N: block_shape[1], + BLOCK_RED: 1, + M: m, + N: n, + NUM_SPLITS: lsu_factor, + } + hyperparams_red.update(get_default_scheduling_params()) + + reduction_options = WaveCompileOptions( + subs=hyperparams_red, + canonicalize=True, + schedule=SchedulingType.NONE, + use_global_to_shared=False, + minimize_shared_allocs=False, + ) + + return lsu_main, main_options, lsu_reduction, reduction_options + + +def get_tagged_tree_streamk_mxfp4_gemm( + shape: tuple[int, int, int] = (256, 256, 256), + block_shape: tuple[int, int, int] = (128, 128, 128), + wave_shape: tuple[int, int] = (2, 2), + mfma_variant: ScaledMMAType = ScaledMMAType.F32_16x16x128_F8F6F4, + a_address_space: tkl.AddressSpace = GLOBAL_ADDRESS_SPACE, + output_type: "tkl.DataType" = tkl.f32, + num_ctas: int = 304, + *, + use_global_to_shared: bool | None = None, + auto_large_block_k: bool = True, +): + """Return a tagged stream-K MXFP4 GEMM kernel + compile options. + + Stream-K distributes K-loop iterations across a fixed number of CTAs + (persistent kernel). Each CTA dynamically determines which output tile(s) + and K-range(s) to process, using atomic_add for accumulation. The caller + must zero-initialize C before launch. + + The workspace and flag_buffer arguments are accepted for API compatibility + with planned tree reduction upgrade. Currently uses atomic_add on C. + + Args: + shape: (M, N, K) problem dimensions. + block_shape: (BLOCK_M, BLOCK_N, BLOCK_K) tile sizes. + wave_shape: (WAVE_M, WAVE_N) waves per workgroup. + mfma_variant: Scaled MMA instruction type. + a_address_space: Address space for A/B data and scales. + output_type: Element type of output tensor C. + num_ctas: Number of CTAs (workgroups) to launch. + use_global_to_shared: When None, pick GatherToLDS vs direct global using a + shape heuristic (large MN, single-tile tall-K, or auto BLOCK_K=256). + When True or False, force that path. + auto_large_block_k: If True and BLOCK_K is 128, use 256 when K >= 16384 and + K is divisible by 256. + + Returns: + (kernel_function, WaveCompileOptions) + """ + from wave_lang.kernel._support.indexing import sym + from wave_lang.kernel._support.dtype import i32 + + m, n, k = shape + block_m, block_n, block_k = block_shape + + block_k_upgraded = False + if auto_large_block_k and block_k == 128 and k >= 16384 and k % 256 == 0: + block_k = 256 + block_shape = (block_m, block_n, block_k) + block_k_upgraded = True + + num_tiles_m = math.ceil(m / block_m) + num_tiles_n = math.ceil(n / block_n) + total_tiles = num_tiles_m * num_tiles_n + + iters_per_tile = math.ceil(k / block_k) + total_iters = total_tiles * iters_per_tile + + sk_iters_pcu = total_iters // num_ctas + sk_extra_iters = total_iters % num_ctas + + M = tkl.sym.M + N = tkl.sym.N + K = tkl.sym.K + BLOCK_M = tkl.sym.BLOCK_M + BLOCK_N = tkl.sym.BLOCK_N + BLOCK_K = tkl.sym.BLOCK_K + ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE + + NUM_CTAS = sym.NUM_CTAS + N_TILES = sym.N_TILES + ITERS_PER_TILE = sym.ITERS_PER_TILE + SK_ITERS_PCU = sym.SK_ITERS_PCU + SK_EXTRA_ITERS = sym.SK_EXTRA_ITERS + CTA_M_OFFSET = sym.CTA_M_OFFSET + CTA_N_OFFSET = sym.CTA_N_OFFSET + START_K_TILE = sym.START_K_TILE + NUM_K_TILES = sym.NUM_K_TILES + WORK_UNIT_START = sym.WORK_UNIT_START + WORK_UNIT_END = sym.WORK_UNIT_END + WS_SIZE = sym.WS_SIZE + FLAG_SIZE = sym.FLAG_SIZE + + i = tkw.IndexMapping.iterator(0) + j = tkw.IndexMapping.iterator(1) + + a_read_mapping = tkw.IndexMapping( + num_iterators=2, + inputs={M: i + CTA_M_OFFSET, K: j}, + outputs={M: i, K: j}, + ) + + b_read_mapping = tkw.IndexMapping( + num_iterators=2, + inputs={N: i + CTA_N_OFFSET, K: j}, + outputs={N: i, K: j}, + ) + + a_scale_read_mapping = tkw.IndexMapping( + num_iterators=2, + inputs={M: i + CTA_M_OFFSET, K: j}, + outputs={M: i, K: j}, + ) + b_scale_read_mapping = tkw.IndexMapping( + num_iterators=2, + inputs={N: i + CTA_N_OFFSET, K: j}, + outputs={N: i, K: j}, + ) + + # For atomic_add, the mapping's *inputs* specify the target memory index + # (unlike tkw.write which uses *outputs* for the target). + c_write_mapping = tkw.IndexMapping( + num_iterators=2, + inputs={M: i + CTA_M_OFFSET, N: j + CTA_N_OFFSET}, + outputs={M: i, N: j}, + ) + + constraints: list[tkw.Constraint] = [ + tkw.GridConstraint(NUM_CTAS), + tkw.WorkgroupConstraint(M, BLOCK_M, 0), + tkw.WorkgroupConstraint(N, BLOCK_N, 1), + tkw.TilingConstraint( + K, BLOCK_K, iters=NUM_K_TILES, start=START_K_TILE * BLOCK_K + ), + tkw.TilingConstraint(WORK_UNIT_START), + tkw.WaveConstraint(M, BLOCK_M / wave_shape[0]), + tkw.WaveConstraint(N, BLOCK_N / wave_shape[1]), + tkw.HardwareConstraint( + threads_per_wave=64, + mma_type=mfma_variant, + vector_shapes={WORK_UNIT_START: 0}, + ), + ] + + @tkw.wave(constraints) + def tree_streamk_mxfp4_gemm( + a: tkl.Memory[M, K / 2, ADDRESS_SPACE, tkl.i8], + a_scale: tkl.Memory[M, K / 32, ADDRESS_SPACE, tkl.i8], + b: tkl.Memory[N, K / 2, ADDRESS_SPACE, tkl.i8], + b_scale: tkl.Memory[N, K / 32, ADDRESS_SPACE, tkl.i8], + c: tkl.Memory[M, N, GLOBAL_ADDRESS_SPACE, output_type], + ): + cta_id = tkw.scalar(WORKGROUP_0, i32) + iters_per_tile_val = tkw.scalar(ITERS_PER_TILE, i32) + sk_iters_pcu_val = tkw.scalar(SK_ITERS_PCU, i32) + sk_extra_iters_val = tkw.scalar(SK_EXTRA_ITERS, i32) + + extra_iter = tkw.minimum(cta_id, sk_extra_iters_val) + work_unit_start = cta_id * sk_iters_pcu_val + extra_iter + next_extra_iter = tkw.minimum(cta_id + tkw.scalar(1, i32), sk_extra_iters_val) + work_unit_end = ( + cta_id + tkw.scalar(1, i32) + ) * sk_iters_pcu_val + next_extra_iter + + tkw.set_symbol(WORK_UNIT_END, work_unit_end) + sk_condition = WORK_UNIT_START < WORK_UNIT_END + + @tkw.iterate( + WORK_UNIT_START, + start=work_unit_start, + condition=sk_condition, + init_args=[], + ) + def sk_loop(): + cta_k_start = tkw.scalar(WORK_UNIT_START, i32) + remainder = cta_k_start % iters_per_tile_val + cta_k_end = tkw.minimum( + cta_k_start + (iters_per_tile_val - remainder), + tkw.scalar(WORK_UNIT_END, i32), + ) + output_tile_id = cta_k_start // iters_per_tile_val + + m_offset = (output_tile_id // tkw.scalar(N_TILES, i32)) * tkw.scalar( + BLOCK_M, i32 + ) + n_offset = (output_tile_id % tkw.scalar(N_TILES, i32)) * tkw.scalar( + BLOCK_N, i32 + ) + tkw.set_symbol(CTA_M_OFFSET, m_offset) + tkw.set_symbol(CTA_N_OFFSET, n_offset) + + tkw.set_symbol(START_K_TILE, remainder) + tkw.set_symbol(NUM_K_TILES, cta_k_end - cta_k_start) + + c_reg = tkl.Register[M, N, tkl.f32](0.0) + + @tkw.iterate(K, init_args=[c_reg], tag="k_loop") + def repeat( + acc: tkl.Register[M, N, tkl.f32], + ) -> tkl.Register[M, N, tkl.f32]: + a_reg = tkw.read(a, mapping=a_read_mapping, tag="read_a") + a_reg = tkw.bitcast(a_reg, tkl.f4e2m1fn, tag="bitcast_a") + a_scale_reg = tkw.read( + a_scale, mapping=a_scale_read_mapping, tag="read_a_scale" + ) + a_scale_reg = tkw.bitcast( + a_scale_reg, tkl.f8e8m0fnu, tag="bitcast_a_scale" + ) + b_reg = tkw.read(b, mapping=b_read_mapping, tag="read_b") + b_reg = tkw.bitcast(b_reg, tkl.f4e2m1fn, tag="bitcast_b") + b_scale_reg = tkw.read( + b_scale, mapping=b_scale_read_mapping, tag="read_b_scale" + ) + b_scale_reg = tkw.bitcast( + b_scale_reg, tkl.f8e8m0fnu, tag="bitcast_b_scale" + ) + acc = tkw.scaled_mma( + a_reg, + a_scale_reg, + b_reg, + b_scale_reg, + acc, + tag="scaled_mma", + ) + return acc + + repeat_out = tkw.cast(repeat, output_type) + tkw.atomic_add(repeat_out, c, mapping=c_write_mapping) + + new_cta_k_start = cta_k_end + tkw.set_symbol(WORK_UNIT_START, new_cta_k_start) + + hyperparams = { + ADDRESS_SPACE: a_address_space, + BLOCK_M: block_m, + BLOCK_N: block_n, + BLOCK_K: block_k, + M: m, + N: n, + K: k, + N_TILES: num_tiles_n, + NUM_CTAS: num_ctas, + ITERS_PER_TILE: iters_per_tile, + SK_ITERS_PCU: sk_iters_pcu, + SK_EXTRA_ITERS: sk_extra_iters, + } + + if use_global_to_shared is None: + mn = m * n + # GatherToLDS helps large output grids and tall-K single-tile cases, and + # pairs with auto BLOCK_K=256. Mid-sized squares (e.g. 256^2) with + # moderate K are faster with direct global loads. + use_g2s = mn >= 262144 or (mn <= 16384 and k >= 32768) or block_k_upgraded + else: + use_g2s = use_global_to_shared + minimize_sa = not use_g2s + + options = WaveCompileOptions( + subs=hyperparams, + canonicalize=True, + use_global_to_shared=use_g2s, + minimize_shared_allocs=minimize_sa, + ) + + return tree_streamk_mxfp4_gemm, options diff --git a/waveasm/lib/Transforms/AssemblyEmitter.cpp b/waveasm/lib/Transforms/AssemblyEmitter.cpp index d5ab9497fe..ff7a641e81 100644 --- a/waveasm/lib/Transforms/AssemblyEmitter.cpp +++ b/waveasm/lib/Transforms/AssemblyEmitter.cpp @@ -649,12 +649,22 @@ std::optional KernelGenerator::generateOp(Operation *op) { // s_cbranch_scc1 L_after_loop_N // The initArg-to-blockArg copies above ensure loop results are // correct even when the body is skipped entirely. + // + // Pattern invariant: this guard matches a specific loop form produced + // by TranslateFromMLIR for scf.for: + // - The loop condition is S_CMP_LT_U32 (unsigned IV < UB). + // - The induction variable is the first init arg with PSRegType. + // - The upper bound (operand 1 of S_CMP_LT_U32) is SGPR or imm. + // If lowering changes the compare opcode, IV position, or UB + // representation, this pattern will not match and the assert below + // fires. Extend the pattern match accordingly if that happens. std::string afterLabel = "L_after_loop_" + std::to_string(loopLabelCounter - 1); { auto condOp = cast(body.getTerminator()); Value condVal = condOp.getCondition(); Operation *cmpOp = condVal.getDefiningOp(); + bool guardEmitted = false; if (cmpOp && isa(cmpOp)) { int64_t initIVPhys = -1; if (!loopOp.getInitArgs().empty()) { @@ -668,13 +678,20 @@ std::optional KernelGenerator::generateOp(Operation *op) { if (auto ps = dyn_cast(ubTy)) { os << " s_cmp_ge_u32 s" << initIVPhys << ", s" << ps.getIndex() << "\n"; + guardEmitted = true; } else if (auto imm = dyn_cast(ubTy)) { os << " s_cmp_ge_u32 s" << initIVPhys << ", " << imm.getValue() << "\n"; + guardEmitted = true; } - os << " s_cbranch_scc1 " << afterLabel << "\n"; + if (guardEmitted) + os << " s_cbranch_scc1 " << afterLabel << "\n"; } } + // If the pattern didn't match, no guard is emitted. This is + // correct for genuine do-while loops (hand-written or non-scf.for + // origins) that always execute at least one iteration. The + // afterLabel is still emitted below but never branched to. } os << labelName << ":\n"; diff --git a/waveasm/lib/Transforms/Liveness.cpp b/waveasm/lib/Transforms/Liveness.cpp index a7c685a850..957201a0c7 100644 --- a/waveasm/lib/Transforms/Liveness.cpp +++ b/waveasm/lib/Transforms/Liveness.cpp @@ -747,15 +747,28 @@ LivenessInfo computeLiveness(ProgramOp program) { // Pass 3c: Remove IfOp result ranges from the allocation worklist. // - // IfOp results alias the then-yield operands. Without this, the - // allocator gives them fresh registers because their start point (the - // IfOp index) precedes the then-block contents. This wastes registers - // and introduces non-determinism (N results with identical sort keys - // get shuffled by DenseMap iteration order). + // --- IfOp result aliasing contract --- + // IfOp results are not independently allocated. Instead, they alias + // the then-yield operands via a three-pass protocol: // - // Post-allocation, LinearScanPass assigns IfOp result types from the - // then-yield operand types. The else-branch copies (if needed) are - // emitted by AssemblyEmitter. + // 1. Liveness (here, Pass 3c): merge each IfOp result's live range + // into the corresponding then-yield operand's range, then erase + // the result range so LinearScan never sees it. + // 2. LinearScanPass: after allocation, walk IfOps and assign each + // result the same physical register as its then-yield operand. + // 3. AssemblyEmitter: at both then- and else-branch exits, emit + // register copies from yield values to result registers when + // they differ (emitYieldToResultCopies). + // + // Any pass inserted between liveness and emission that reads or + // mutates IfOp result types/registers must respect this aliasing -- + // IfOp results have no independent allocation. + // --- end contract --- + // + // Without this, the allocator gives them fresh registers because their + // start point (the IfOp index) precedes the then-block contents. + // This wastes registers and introduces non-determinism (N results with + // identical sort keys get shuffled by DenseMap iteration order). program.walk([&](IfOp ifOp) { auto thenYield = dyn_cast(ifOp.getThenBlock().getTerminator()); if (!thenYield) @@ -771,6 +784,10 @@ LivenessInfo computeLiveness(ProgramOp program) { if (resultIt == info.ranges.end()) continue; + assert(yieldIt != info.ranges.end() && + "IfOp result has a live range but its corresponding then-yield " + "operand does not -- the yield value must be live for the " + "aliasing contract to hold"); if (yieldIt != info.ranges.end()) { // Extend the yield operand's range to cover the ifOp result's // full lifetime, since they will share the same physical register. diff --git a/waveasm/lib/Transforms/RegionBuilder.cpp b/waveasm/lib/Transforms/RegionBuilder.cpp index bc57dc0730..897748ec27 100644 --- a/waveasm/lib/Transforms/RegionBuilder.cpp +++ b/waveasm/lib/Transforms/RegionBuilder.cpp @@ -347,6 +347,12 @@ IfOp RegionBuilder::buildIfFromSCFIf(scf::IfOp ifOp) { // immediate placeholders, SGPRs used as booleans) are coerced here. // VGPR path mirrors buildLoopFromSCFFor's VGPR upper-bound handling: // v_readfirstlane_b32 then s_cmp_ne_u32 with 0 to set SCC. + // + // INVARIANT: the VGPR condition must be workgroup-uniform (all lanes hold + // the same value). V_READFIRSTLANE_B32 reads only lane 0; if lanes + // diverge the scalar branch will follow lane 0's path for every lane, + // silently miscompiling. This is safe for split-K trip-count tests + // (derived from workgroup_id). Do NOT feed lane-varying conditions here. Value conditionValue = *condition; auto condType = conditionValue.getType(); if (!isSCCType(condType)) { diff --git a/waveasm/lib/Transforms/handlers/ArithHandlers.cpp b/waveasm/lib/Transforms/handlers/ArithHandlers.cpp index 4397abf850..69b3253853 100644 --- a/waveasm/lib/Transforms/handlers/ArithHandlers.cpp +++ b/waveasm/lib/Transforms/handlers/ArithHandlers.cpp @@ -397,9 +397,15 @@ LogicalResult handleArithCmpI(Operation *op, TranslationContext &ctx) { // When a comparison with VGPR operands is used only by scf.if, promote the // operands to SGPRs via V_READFIRSTLANE_B32 and use S_CMP. This avoids // the round-trip through VCC -> bool VGPR -> V_READFIRSTLANE -> S_CMP that - // buildIfFromSCFIf would otherwise need, reducing SGPR pressure. The - // values must be workgroup-uniform for correctness (split-K trip-count - // comparisons satisfy this since they derive from workgroup_id). + // buildIfFromSCFIf would otherwise need, reducing SGPR pressure. + // + // INVARIANT: the VGPR operands must be workgroup-uniform (all lanes hold + // the same value). V_READFIRSTLANE_B32 reads lane 0 and broadcasts it to + // the SGPR; if lanes diverge, the scalar comparison will use only lane 0's + // value and the branch decision will be wrong for other lanes. Currently + // this is satisfied because split-K trip-count comparisons derive entirely + // from workgroup_id (uniform by construction). Do NOT route lane-varying + // comparisons through this path. bool allUsesAreScfIf = !cmpOp.getResult().use_empty(); for (Operation *user : cmpOp.getResult().getUsers()) { if (!isa(user)) {