Add wide_stores support for MXFP4 (4wave) GEMM (llvm backend) (#1282)

## Wide epilogue stores for bf16 MXFP4 GEMM via `v_permlane16_swap`

In MXFP4 GEMM kernels with bf16 output, the MMA accumulator layout
(`F32_16x16x128_F8F6F4`) distributes results across threads in a
non-contiguous pattern each thread holds 4 values spanning 4 rows but
only 1 column. Without optimization, the epilogue emits 192 scalar
`buffer_store_short` (2 bytes each), underutilizing memory bandwidth.

This PR adds a `wide_stores` that swaps MFMA operands (B as LHS, A as
RHS) and uses `v_permlane16_swap_b32` to coalesce epilogue stores into
`buffer_store_dwordx4` (16 bytes each), reducing store count from 192 to
48.

### How it works
**Swapped MFMA operands (`wide_stores=True`):** The kernel computes C^T
= B × A^T (and stores C) instead of C = A × B^T, making the
accumulator's 4-contiguous values align with the output memory's
stride-1 dimension. The source/target write syntax maps `Register[N, M]`
to `Memory[M, N]` so the output is C [M, N] directly.

**`wide_store_coalescing` pass:** An always-on graph pass that
identifies eligible bf16 global Write nodes using the source/target
dimension remapping pattern (indicating swapped operands). Only tags
writes that satisfy all conditions: global address space, bf16 dtype,
and source/target syntax. Safe to run unconditionally, non-`wide_stores`
kernels are unaffected.

**`_write_permlane_pack_to_global` codegen:** At code generation time,
tagged writes emit `v_permlane16_swap_b32` to exchange each thread's 4
bf16 values (packed as 2 i32 dwords) with a partner lane 16 positions
apart. Both lane halves write identical 8-bf16 vectors to the same
address (benign duplicate store), avoiding divergent control flow. The
buffer descriptor's `valid_bytes` handles out-of-bounds suppression for
dynamic shapes.

### Assembly comparison (shape 1024×3072×8192, block 256×192×256)

| Instruction | Baseline | Wide Stores | Change |
|---|---|---|---|
| `buffer_store_short` | 192 | 0 | eliminated |
| `buffer_store_dwordx4` | 0 | 48 | 16B/store |
| `v_permlane16_swap` | 0 | 96 | lane exchange |
| `v_cvt_pk_bf16_f32` | 192 | 96 | 2x reduction |
| Total asm lines | 12,115 | 9,322 | 23% fewer |

LLVM now emits `v_cvt_pk_bf16_f32` natively (bf16 softening has been
fixed upstream), so the paired f32->bf16 conversion is a single
instruction.

Compute-bound shapes (large K) show 6-9% gains; small-K shapes show
13-27% gains

| Shape (M, N, K) | Baseline TFLOPS | Wide Stores TFLOPS | Speedup |
|---|---|---|---|
| (315904, 384, 1792) | 1169.51 | 1557.86 | 1.332x |
| (20480, 1152, 28928) | 2714.36 | 3015.93 | 1.111x |
| (5888, 1920, 3328) | 1582.35 | 2104.81 | 1.330x |
| (46336, 2688, 8448) | 2495.16 | 2940.50 | 1.179x |
| (512, 3072, 428288) | 590.71 | 648.15 | 1.097x |
| (256, 4224, 102656) | 409.79 | 458.95 | 1.120x |
| (512, 4224, 55808) | 794.47 | 893.25 | 1.124x |
| (4864, 4608, 2560) | 1435.43 | 2123.30 | 1.479x |
| (4864, 4608, 5888) | 2083.50 | 2669.39 | 1.281x |
| (4608, 4992, 7424) | 2227.52 | 2813.12 | 1.263x |
| (2048, 5376, 63232) | 3112.73 | 3405.23 | 1.094x |
| (2048, 5760, 6400) | 2230.97 | 2690.92 | 1.206x |
| (3584, 6912, 548608) | 3272.81 | 3531.95 | 1.079x |
| (4864, 7680, 7936) | 2406.89 | 3036.57 | 1.262x |
| (13056, 10368, 15360) | 2936.09 | 3352.67 | 1.142x |
| (14848, 12672, 10496) | 2781.61 | 3254.58 | 1.170x |
| (1792, 19968, 60928) | 3115.72 | 3392.02 | 1.089x |

---------

Signed-off-by: xintin <gaurav.verma@amd.com>

Author: xintin

examples/python/7.1_schedule.py

@@ -12,36 +12,37 @@
 """
 
 import torch
-import wave_lang.kernel.lang as tkl
+from utils import list_tests, parse_args, run_test
 
-from wave_lang.kernel.wave.compile import wave_compile
-from wave_lang.kernel.wave.utils.run_utils import set_default_run_config
-from wave_lang.kernel.wave.templates import (
-    get_tagged_mxfp4_gemm,
-    get_tagged_mxfp4_gemm_preshuffle_b,
-    get_tagged_mxfp4_gemm_preshuffle_scales,
-    get_tagged_mxfp4_gemm_preshuffle_scales_and_B,
+import wave_lang.kernel.lang as tkl
+from wave_lang.kernel.lang.global_symbols import (
+    GLOBAL_ADDRESS_SPACE,
+    SHARED_ADDRESS_SPACE,
 )
+from wave_lang.kernel.wave.compile import wave_compile
 from wave_lang.kernel.wave.schedules import (
-    get_mxfp4_dbuf_schedule,
-    get_mxfp4_dbuf_pingpong_schedule,
-    get_mxfp4_dbuf_mixed_pingpong_schedule,
     get_mxfp4_asymmetric_schedule,
+    get_mxfp4_dbuf_mixed_pingpong_schedule,
     get_mxfp4_dbuf_mixed_pingpong_shuffle_schedule,
+    get_mxfp4_dbuf_pingpong_schedule,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds,
+    get_mxfp4_dbuf_schedule,
+)
+from wave_lang.kernel.wave.templates import (
+    get_tagged_mxfp4_gemm,
+    get_tagged_mxfp4_gemm_preshuffle_b,
+    get_tagged_mxfp4_gemm_preshuffle_b_wide_store,
+    get_tagged_mxfp4_gemm_preshuffle_scales,
+    get_tagged_mxfp4_gemm_preshuffle_scales_and_B,
 )
 from wave_lang.kernel.wave.utils.mxfp_utils import (
-    generate_gemm_afp4wfp4_inputs,
-    torchScaledGemmMXFP4,
     b_preshuffle,
     e8m0_shuffle,
+    generate_gemm_afp4wfp4_inputs,
+    torchScaledGemmMXFP4,
 )
-from wave_lang.kernel.lang.global_symbols import (
-    GLOBAL_ADDRESS_SPACE,
-    SHARED_ADDRESS_SPACE,
-)
-from utils import parse_args, list_tests, run_test
+from wave_lang.kernel.wave.utils.run_utils import set_default_run_config
 
 
 def _run_mxfp_gemm(gemm, shape):
@@ -424,6 +425,44 @@ def test_dbuf_4wave_mxfp_dynamic_preshuffle_b_gemm(
     print("MXFP GEMM preshuffle-B 4-wave dynamic M, N, K (LLVM backend) test passed!")
 
 
+def test_dbuf_4wave_mxfp_dynamic_preshuffle_b_gemm_wide_stores(
+    is_debug=False,
+    shape=(1024, 3072, 8192),
+    block=(256, 192, 256),
+    eliminate_epilogue=False,
+):
+    """Preshuffle-B MXFP4 GEMM with dynamic M, N, K and wide epilogue stores.
+
+    Uses the wide_store variant to swap MFMA operands (B as LHS, A as RHS),
+    aligning the accumulator's contiguous values with the output's stride-1
+    dimension. The coalesce_wide_stores pass emits v_permlane16_swap_b32
+    + buffer_store_dwordx4 (8 bf16 per store) instead of buffer_store_short.
+    """
+    gemm, options = get_tagged_mxfp4_gemm_preshuffle_b_wide_store(
+        shape,
+        block,
+        wave_shape=(2, 2),
+        reorder_workgroups=True,
+    )
+    dynamic_symbols = [tkl.sym.M, tkl.sym.N, tkl.sym.K]
+    for sym in dynamic_symbols:
+        del options.subs[sym]
+    options.dynamic_symbols = dynamic_symbols
+    options.use_buffer_ops = True
+    options.backend = "llvm"
+    options.wave_runtime = True
+    options.eliminate_epilogue = eliminate_epilogue
+    schedule = get_mxfp4_asymmetric_schedule(
+        eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
+    )
+    options.print_ir_after = "all" if is_debug else []
+    options = set_default_run_config(options)
+    gemm = wave_compile(options, gemm, schedule)
+
+    _run_mxfp_gemm_preshuffle(gemm, shape, all=True, output_dtype=torch.bfloat16)
+    print("MXFP GEMM preshuffle-B 4-wave dynamic M, N, K (wide stores) test passed!")
+
+
 def test_dbuf_4wave_mxfp_dynamic_preshuffle_b_gemm_asm(
     is_debug=False,
     shape=(1024, 1024, 8192),

lit_tests/kernel/wave/mlir_roundtrip_pipeline.py

@@ -304,6 +304,7 @@ def mxfp4_gemm_progressive_roundtrip():
             "guard_g2s_with_bounds_check",
             "schedule_reordering",
             "minimize_shared_allocs",
+            "coalesce_wide_stores",
             "add_shared_memory_barriers",
             "add_cluster_barriers",
             "compute_shared_memory_usage",

lit_tests/kernel/wave/wide_stores_mxfp4.py

@@ -0,0 +1,82 @@
+# RUN: python %s | FileCheck %s
+
+"""
+Test wide store coalescing for preshuffle-B MXFP4 GEMM with bf16 output.
+
+The wide_store variant kernel swaps MFMA operands (B as LHS, A as RHS)
+so the accumulator's 4-contiguous values align with the output's stride-1
+dimension. The coalesce_wide_stores pass tags eligible bf16 global
+writes, and the codegen emits v_permlane16_swap_b32 to exchange data
+between lane pairs 16 apart, producing 8 consecutive bf16 values written
+as a single buffer_store_dwordx4.
+
+Key structural invariants verified:
+  1. Function signature accepts dynamic index arguments for M, N, K.
+  2. Register type is [N, M] (swapped from standard [M, N]).
+  3. scaled_mfma has B as LHS and A as RHS (swapped operands).
+  4. rocdl.permlane16.swap used for lane exchange.
+  5. vector.store of vector<8xbf16> for wide stores.
+  6. arith.truncf for f32 -> bf16 conversion.
+"""
+
+import wave_lang.kernel.lang as tkl
+from wave_lang.kernel.wave.compile import wave_compile
+from wave_lang.kernel.wave.constraints import ScaledMMAType
+from wave_lang.kernel.wave.schedules import get_mxfp4_asymmetric_schedule
+from wave_lang.kernel.wave.templates import (
+    get_tagged_mxfp4_gemm_preshuffle_b_wide_store,
+)
+from wave_lang.kernel.wave.utils.general_utils import run_test
+
+
+@run_test
+def test_wide_stores_preshuffle_b_mxfp4():
+    shape = (1024, 3072, 8192)
+    block = (256, 192, 256)
+    kernel, options = get_tagged_mxfp4_gemm_preshuffle_b_wide_store(
+        shape,
+        block,
+        wave_shape=(2, 2),
+        reorder_workgroups=True,
+        mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4,
+    )
+    dynamic_symbols = [tkl.sym.M, tkl.sym.N, tkl.sym.K]
+    for sym in dynamic_symbols:
+        del options.subs[sym]
+    options.dynamic_symbols = dynamic_symbols
+    schedule = get_mxfp4_asymmetric_schedule(is_bscale_shuffled=True)
+    options.use_buffer_ops = True
+    options.compile_to_mlir = True
+    options.device = "hip"
+    options.target = "gfx950"
+    result = wave_compile(options, kernel, schedule)
+    print(result.asm)
+
+    # CHECK-LABEL: test_wide_stores_preshuffle_b_mxfp4
+
+    # 1. Dynamic index arguments for M, N, K in function signature.
+    # CHECK: func.func @gemm(%arg0: {{.*}}, %arg1: {{.*}}, %arg2: {{.*}}, %arg3: {{.*}}, %arg4: {{.*}}, %arg5: index, %arg6: index, %arg7: index)
+
+    # 2. f32 -> bf16 conversion in the epilogue.
+    # CHECK: arith.truncf %{{.*}} : vector<4xf32> to vector<4xbf16>
+
+    # 3. vector.bitcast from bf16 to i32 for permlane swap.
+    # CHECK: vector.bitcast %{{.*}} : vector<4xbf16> to vector<2xi32>
+
+    # 4. rocdl.permlane16.swap for lane exchange.
+    # CHECK: rocdl.permlane16.swap
+
+    # 5. llvm.extractvalue to get the swapped value.
+    # CHECK: llvm.extractvalue %{{.*}}[0]
+
+    # 6. arith.select to choose between original and swapped values.
+    # CHECK: arith.select %{{.*}}, %{{.*}}, %{{.*}} : i32
+
+    # 7. vector.from_elements to pack 4 i32 into vector<4xi32>.
+    # CHECK: vector.from_elements %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : vector<4xi32>
+
+    # 8. vector.bitcast from i32 to bf16 for the wide store.
+    # CHECK: vector.bitcast %{{.*}} : vector<4xi32> to vector<8xbf16>
+
+    # 9. Wide store of 8 bf16 values.
+    # CHECK: vector.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<{{.*}}xbf16, {{.*}}>, vector<8xbf16>

tests/kernel/wave_gemm_mxfp_test.py

@@ -29,6 +29,7 @@
 from wave_lang.kernel.wave.templates import (
     get_tagged_mxfp4_gemm,
     get_tagged_mxfp4_gemm_preshuffle_b,
+    get_tagged_mxfp4_gemm_preshuffle_b_wide_store,
     get_tagged_mxfp4_gemm_preshuffle_scales,
     get_tagged_mxfp4_gemm_preshuffle_scales_and_B,
 )
@@ -1028,6 +1029,62 @@ def testScaledGemmMXFP4PreshuffleBDynamicWaveRuntime(
     torch.testing.assert_close(torch_out, out, check_dtype=False)
 
 
+@require_e2e
+@require_cdna4
+@pytest.mark.timeout(900)
+@pytest.mark.parametrize(
+    "shape",
+    [(1024, 3072, 8192)],
+)
+@pytest.mark.parametrize(
+    "block_shape,wave_shape",
+    [((256, 192, 256), (2, 2))],
+)
+@pytest.mark.parametrize(
+    "mfma_variant",
+    [ScaledMMAType.F32_16x16x128_F8F6F4],
+)
+def testScaledGemmMXFP4PreshuffleBWideStores(
+    shape: tuple[int, int, int],
+    block_shape: tuple[int, int, int],
+    wave_shape: tuple[int, int],
+    mfma_variant: ScaledMMAType,
+):
+    """End-to-end test for MXFP4 GEMM with wide epilogue stores (dwordx4).
+
+    Uses the wide_store variant to swap MFMA operands and emit
+    buffer_store_dwordx4 via v_permlane16_swap_b32 for bf16 output.
+    """
+    gemm, options = get_tagged_mxfp4_gemm_preshuffle_b_wide_store(
+        shape,
+        block_shape,
+        wave_shape=wave_shape,
+        mfma_variant=mfma_variant,
+        reorder_workgroups=True,
+    )
+    dynamic_symbols = [tkl.sym.M, tkl.sym.N, tkl.sym.K]
+    for sym in dynamic_symbols:
+        del options.subs[sym]
+    options.dynamic_symbols = dynamic_symbols
+    schedule = get_mxfp4_asymmetric_schedule(is_bscale_shuffled=True)
+    options.use_buffer_ops = True
+    options = set_default_run_config(options)
+    gemm = wave_compile(options, gemm, schedule)
+
+    x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs(shape)
+    torch_out = torchScaledGemmMXFP4(x, w, x_scales, w_scales)
+
+    w_t = w.T.contiguous()
+    w_t_ps = b_preshuffle(w_t)
+    x_scales_ps = e8m0_shuffle(x_scales)
+    w_scales_ps = e8m0_shuffle(w_scales)
+
+    out = device_zeros(x.shape[0], w_t_ps.shape[0], dtype=torch.bfloat16)
+    gemm(x, x_scales_ps, w_t_ps, w_scales_ps, out)
+
+    torch.testing.assert_close(torch_out, out, check_dtype=False)
+
+
 MACROTILES_PRESHUFFLE_8WAVE_PINGPONG = [
     (256, 160, 256),
     (256, 224, 256),