import torch
import pytest
import triton
from triton.experimental import gluon
import triton.experimental.gluon.language as gl
from triton.experimental.gluon.language.amd.cdna5 import tdm
from triton.experimental.gluon.language.amd.cdna5 import async_copy as cp
from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor

# Handle imports for both pytest (module context) and direct execution
try:
    from .cdna5_utils import static_profile, composition
except ImportError:
    from cdna5_utils import static_profile, composition


@gluon.jit
def swiglu_epilogue(acc):
    """SwiGLU: split (M, 2N) into interleaved gate/up pairs and return swish(gate) * up."""
    BLOCK_M: gl.constexpr = acc.shape[0]
    BLOCK_N: gl.constexpr = acc.shape[1] // 2
    gate, up = gl.split(gl.reshape(acc, (BLOCK_M, BLOCK_N, 2)))
    # swish(x) = x * sigmoid(x); sigmoid(x) = 1 / (1 + exp(-x))
    return gate * (1.0 / (1.0 + gl.exp(-gate))) * up


@gluon.constexpr_function
def get_scale_blocked_layout(num_warps: gl.constexpr):
    return gl.BlockedLayout([1, 8], [1, 32], [num_warps // 2, 2], [1, 0])


@gluon.constexpr_function
def get_wmma_layout(num_warps, packed, scale_preshuffle, instr_m=16):
    assert (num_warps in (4, 8))
    assert (instr_m in (16, 32))
    if scale_preshuffle:
        reg_bases = [[0, 1], [1, 0]]
        tiles_per_warp = 2
    else:
        reg_bases = []
        tiles_per_warp = 1

    if num_warps == 4:
        warp_bases = [[0, tiles_per_warp], [tiles_per_warp, 0]]
    else:
        warp_bases = [[0, tiles_per_warp], [tiles_per_warp, 0], [tiles_per_warp * 2, 0]]

    instr_shape = [instr_m, 16, 64] if packed else [instr_m, 16, 128]

    return gl.amd.AMDWMMALayout(3, True, warp_bases, reg_bases, instr_shape)


@gluon.constexpr_function
def reverse_tdm_warp_used_hint(x: gl.constexpr, num_warps: gl.constexpr):
    if x is None:
        return None

    result = 0
    for i in range(num_warps):
        if x & (1 << i):
            result |= 1 << (num_warps - 1 - i)
    return result


@gluon.aggregate
class MXFPGEMMConfig:
    BLOCK_M: gl.constexpr
    BLOCK_N: gl.constexpr
    BLOCK_K: gl.constexpr
    DTYPE_A: gl.constexpr
    DTYPE_B: gl.constexpr
    DIV_FACTOR_A: gl.constexpr
    DIV_FACTOR_B: gl.constexpr
    NUM_BUFFERS: gl.constexpr
    TRANSPOSE_B: gl.constexpr
    WITH_A_SCALE: gl.constexpr
    NUM_LOADS_IN_BATCH: gl.constexpr
    NUM_SUBTILES: gl.constexpr  # (M, N, K)
    NUM_WARPS: gl.constexpr
    TDM_WARP_USED_HINT_A: gl.constexpr
    TDM_WARP_USED_HINT_B: gl.constexpr

    # Layouts
    shared_layout_a: gl.constexpr
    dot_layout_a: gl.constexpr

    shared_layout_b: gl.constexpr
    dot_layout_b: gl.constexpr

    shared_layout_a_scale: gl.constexpr
    layout_a_scale: gl.constexpr

    shared_layout_b_scale: gl.constexpr
    layout_b_scale: gl.constexpr

    acc_layout: gl.constexpr

    # Scales
    SCALE_PRESHUFFLE: gl.constexpr
    PRESHUFFLE_FACTOR: gl.constexpr
    SCALE_KWIDTH: gl.constexpr
    BLOCK_M_PRESHUFFLED: gl.constexpr
    BLOCK_N_PRESHUFFLED: gl.constexpr
    BLOCK_K_SCALE_PRESHUFFLED: gl.constexpr
    SCALE_BLOCK: gl.constexpr
    ASYNC_COPY_SCALE: gl.constexpr
    L2_PREFETCH_DISTANCE: gl.constexpr

    # Fused activation epilogue. For SwiGLU, BLOCK_N/N refer to
    # the packed (doubled) compute-path width. The store reduces it by 2.
    ACTIVATION: gl.constexpr

    TDM_SPLIT: gl.constexpr

    @gluon.constexpr_function
    def __init__(self, BLOCK_M, BLOCK_N, BLOCK_K, DTYPE_A, DTYPE_B, SCALE_BLOCK, NUM_BUFFERS, TRANSPOSE_B, WITH_A_SCALE,
                 SCALE_PRESHUFFLE, NUM_WARPS, TDM_WARP_USED_HINT=None, ASYNC_COPY_SCALE=False, NUM_SUBTILES=(1, 1, 1),
                 L2_PREFETCH_DISTANCE=0, ACTIVATION="", RESOLVE_PARTITION_CONFLICTS=False, TDM_SPLIT=False):
        self.BLOCK_M = gl.constexpr(BLOCK_M)
        self.BLOCK_N = gl.constexpr(BLOCK_N)
        self.BLOCK_K = gl.constexpr(BLOCK_K)
        self.DTYPE_A = gl.constexpr(DTYPE_A)
        self.DTYPE_B = gl.constexpr(DTYPE_B)
        self.NUM_BUFFERS = gl.constexpr(NUM_BUFFERS)
        self.TRANSPOSE_B = gl.constexpr(TRANSPOSE_B)
        self.WITH_A_SCALE = gl.constexpr(WITH_A_SCALE)
        self.SCALE_PRESHUFFLE = gl.constexpr(SCALE_PRESHUFFLE)
        self.SCALE_BLOCK = gl.constexpr(SCALE_BLOCK)
        self.DIV_FACTOR_A = gl.constexpr(2 if DTYPE_A == "e2m1" else 1)
        self.DIV_FACTOR_B = gl.constexpr(2 if DTYPE_B == "e2m1" else 1)
        self.NUM_LOADS_IN_BATCH = gl.constexpr(4 if WITH_A_SCALE else 3)
        self.NUM_WARPS = gl.constexpr(NUM_WARPS)
        self.TDM_WARP_USED_HINT_A = gl.constexpr(TDM_WARP_USED_HINT)
        self.TDM_WARP_USED_HINT_B = gl.constexpr(reverse_tdm_warp_used_hint(TDM_WARP_USED_HINT, NUM_WARPS))
        self.ASYNC_COPY_SCALE = gl.constexpr(ASYNC_COPY_SCALE)
        self.NUM_SUBTILES = gl.constexpr(NUM_SUBTILES)
        self.L2_PREFETCH_DISTANCE = gl.constexpr(L2_PREFETCH_DISTANCE)
        self.ACTIVATION = gl.constexpr(ACTIVATION)
        self.TDM_SPLIT = gl.constexpr(TDM_SPLIT)
        if ACTIVATION == "swiglu":
            assert (BLOCK_N // NUM_SUBTILES[1]) % 2 == 0, \
            "SwiGLU requires (BLOCK_N // NUM_SUBTILES[1]) % 2 == 0"

        NUM_SUBTILES_M = self.NUM_SUBTILES[0]
        NUM_SUBTILES_N = self.NUM_SUBTILES[1]
        NUM_SUBTILES_K = self.NUM_SUBTILES[2]

        BLOCK_K_SCALE = BLOCK_K // SCALE_BLOCK
        self.SCALE_KWIDTH = gl.constexpr(4 if BLOCK_K_SCALE >= 4 else BLOCK_K_SCALE)
        self.PRESHUFFLE_FACTOR = gl.constexpr(128 if SCALE_PRESHUFFLE else 1)

        self.BLOCK_M_PRESHUFFLED = gl.constexpr(BLOCK_M // self.PRESHUFFLE_FACTOR)
        self.BLOCK_N_PRESHUFFLED = gl.constexpr(BLOCK_N // self.PRESHUFFLE_FACTOR)
        self.BLOCK_K_SCALE_PRESHUFFLED = gl.constexpr(BLOCK_K_SCALE * self.PRESHUFFLE_FACTOR)

        INSTR_M: gl.constexpr = 32 if (DTYPE_A == "e2m1" and DTYPE_B == "e2m1") else 16

        BLOCK_K_PACKED_A = BLOCK_K // self.DIV_FACTOR_A
        BLOCK_K_PACKED_B = BLOCK_K // self.DIV_FACTOR_B
        PAD_INTERVAL_A = 256 if BLOCK_K_PACKED_A <= 256 else BLOCK_K_PACKED_A
        PAD_INTERVAL_B = 256 if BLOCK_K_PACKED_B <= 256 else BLOCK_K_PACKED_B

        # The operand shared layouts (and, for the resolve path, the partition
        # split) are built over the per-buffer tile.  Split TDM gives each half
        # its own buffer (A halved along M, B halved along N), so the layout tile
        # is halved on that axis; otherwise it spans the full BLOCK_M / BLOCK_N.
        LAYOUT_BLOCK_M = (BLOCK_M // NUM_SUBTILES_M) if TDM_SPLIT else BLOCK_M
        LAYOUT_BLOCK_N = (BLOCK_N // NUM_SUBTILES_N) if TDM_SPLIT else BLOCK_N

        padded_a = gl.PaddedSharedLayout.with_identity_for([[PAD_INTERVAL_A, 16]], [LAYOUT_BLOCK_M, BLOCK_K_PACKED_A],
                                                           [1, 0])
        if TRANSPOSE_B:
            padded_b = gl.PaddedSharedLayout.with_identity_for([[PAD_INTERVAL_B, 16]],
                                                               [LAYOUT_BLOCK_N, BLOCK_K_PACKED_B], [1, 0])
        else:
            padded_b = gl.PaddedSharedLayout.with_identity_for([[LAYOUT_BLOCK_N, 16]],
                                                               [BLOCK_K_PACKED_B, LAYOUT_BLOCK_N], [1, 0])

        if RESOLVE_PARTITION_CONFLICTS:
            # Split each operand tile along its M (A) / N (B) axis into partitioned
            # pieces and build a partition-aware WMMA layout to avoid LDS partition
            # conflicts.
            SLICE_M = BLOCK_M // NUM_SUBTILES_M
            SLICE_N = BLOCK_N // NUM_SUBTILES_N
            shared_a, shared_b, WMMA_LAYOUT = gl.amd.cdna5.make_partitioned_dot_layouts(
                LAYOUT_BLOCK_M, LAYOUT_BLOCK_N, padded_a, padded_b, NUM_WARPS, [INSTR_M, 16, 128], a_transposed=False,
                b_transposed=TRANSPOSE_B, slice_m=SLICE_M, slice_n=SLICE_N)
            # The packed (fp4) WMMA layout shares the partition-aware warp/register
            # bases but halves the K instruction extent.
            WMMA_LAYOUT_PACKED = gl.amd.AMDWMMALayout(3, True, WMMA_LAYOUT.warp_bases, WMMA_LAYOUT.reg_bases,
                                                      [INSTR_M, 16, 64])
            self.shared_layout_a = gl.constexpr(shared_a)
            self.shared_layout_b = gl.constexpr(shared_b)
        else:
            WMMA_LAYOUT = get_wmma_layout(NUM_WARPS, False, SCALE_PRESHUFFLE, INSTR_M)
            WMMA_LAYOUT_PACKED = get_wmma_layout(NUM_WARPS, True, SCALE_PRESHUFFLE, INSTR_M)
            self.shared_layout_a = gl.constexpr(padded_a)
            self.shared_layout_b = gl.constexpr(padded_b)

        self.dot_layout_a = gl.constexpr(
            gl.DotOperandLayout(operand_index=0, parent=WMMA_LAYOUT_PACKED if DTYPE_A == "e2m1" else WMMA_LAYOUT,
                                k_width=16))
        self.dot_layout_b = gl.constexpr(
            gl.DotOperandLayout(operand_index=1, parent=WMMA_LAYOUT_PACKED if DTYPE_B == "e2m1" else WMMA_LAYOUT,
                                k_width=16))
        self.layout_a_scale = gl.constexpr(
            gl.amd.cdna5.get_wmma_scale_layout(self.dot_layout_a,
                                               [BLOCK_M // NUM_SUBTILES_M, BLOCK_K_SCALE // NUM_SUBTILES_K]))
        self.layout_b_scale = gl.constexpr(
            gl.amd.cdna5.get_wmma_scale_layout(self.dot_layout_b,
                                               [BLOCK_N // NUM_SUBTILES_N, BLOCK_K_SCALE // NUM_SUBTILES_K]))
        self.acc_layout = gl.constexpr(WMMA_LAYOUT)

        self.shared_layout_a_scale = gl.constexpr(
            gl.PaddedSharedLayout.with_identity_for([[256, 8]],
                                                    [self.BLOCK_M_PRESHUFFLED, self.BLOCK_K_SCALE_PRESHUFFLED], [1, 0]))
        self.shared_layout_b_scale = gl.constexpr(
            gl.PaddedSharedLayout.with_identity_for([[256, 8]],
                                                    [self.BLOCK_N_PRESHUFFLED, self.BLOCK_K_SCALE_PRESHUFFLED], [1, 0]))


@gluon.aggregate
class ScaleAsyncCopyDescriptor:
    cfg: MXFPGEMMConfig
    op_idx: gl.constexpr
    ptr: gl.tensor
    offs: gl.tensor
    step_nonk: gl.tensor
    step_k: gl.tensor
    dtype: gl.constexpr
    block_shape: gl.constexpr
    layout: gl.constexpr

    @gluon.constexpr_function
    def __init__(self, cfg: MXFPGEMMConfig, op_idx, ptr, offs, step_nonk, step_k, layout):
        self.cfg = cfg
        self.op_idx = gl.constexpr(op_idx)
        self.ptr = ptr
        self.offs = offs
        self.step_nonk = step_nonk
        self.step_k = step_k
        BLOCK_NONK = cfg.BLOCK_M_PRESHUFFLED if op_idx == 0 else cfg.BLOCK_N_PRESHUFFLED
        self.dtype = gl.constexpr(ptr.dtype.element_ty)
        self.block_shape = gl.constexpr([BLOCK_NONK, cfg.BLOCK_K_SCALE_PRESHUFFLED])
        self.layout = gl.constexpr(layout)

    @gluon.jit
    def initialize(cfg: MXFPGEMMConfig, op_idx: gl.constexpr, ptr, off, stride, layout):
        gl.static_assert(op_idx == 0 or op_idx == 1)
        if op_idx == 0:
            BLOCK_NONK: gl.constexpr = cfg.BLOCK_M_PRESHUFFLED
        else:
            BLOCK_NONK: gl.constexpr = cfg.BLOCK_N_PRESHUFFLED
        BLOCK_K: gl.constexpr = cfg.BLOCK_K_SCALE_PRESHUFFLED

        blocked_layout: gl.constexpr = get_scale_blocked_layout(cfg.NUM_WARPS)
        offs_non_k = gl.arange(0, BLOCK_NONK, gl.SliceLayout(1, blocked_layout))
        offs_k = gl.arange(0, BLOCK_K, gl.SliceLayout(0, blocked_layout))
        offs = off + offs_non_k[:, None] * stride + offs_k[None, :]
        step_nonk = BLOCK_NONK * stride
        step_k = BLOCK_K

        return ScaleAsyncCopyDescriptor(cfg, op_idx, ptr, offs, step_nonk, step_k, layout)

    @gluon.jit
    def issue_async_load(self, idx: int, buffer, pred=1):
        if pred:
            cp.global_to_shared(buffer, self.ptr + idx * self.step_k + self.offs)
            cp.commit_group()


@gluon.aggregate
class MXFPGEMMProgramBase:

    @gluon.constexpr_function
    def __init__(self):
        pass

    @gluon.jit
    def issue_load_a_scale(self, a_scale_desc, load_idx, pred=1, phase=None):
        cfg = self.cfg
        if cfg.WITH_A_SCALE:
            # When phase is provided, derive the buffer slot from it so that loads
            # and local_loads share the same SSA base value for index expressions,
            # enabling membar to prove buffer-slot disjointness.
            if phase is not None:
                slot = phase % cfg.NUM_BUFFERS
            else:
                slot = load_idx % cfg.NUM_BUFFERS
            gl.amd.cdna5.tdm.async_load(a_scale_desc, [0, 0], self.a_scale_buffer.index(slot), pred=pred)
            a_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_scale_desc,
                                                                     add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
        return a_scale_desc

    @gluon.jit
    def issue_load_b_scale(self, b_scale_desc, load_idx, pred=1, phase=None):
        cfg = self.cfg
        if phase is not None:
            slot = phase % cfg.NUM_BUFFERS
        else:
            slot = load_idx % cfg.NUM_BUFFERS
        gl.amd.cdna5.tdm.async_load(b_scale_desc, [0, 0], self.b_scale_buffer.index(slot), pred=pred)
        b_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_scale_desc,
                                                                 add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
        return b_scale_desc

    @gluon.jit
    def issue_load_a_data(self, a_desc, load_idx, pred=1, phase=None):
        cfg = self.cfg
        BLOCK_K: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A
        if phase is not None:
            slot = phase % cfg.NUM_BUFFERS
        else:
            slot = load_idx % cfg.NUM_BUFFERS
        gl.amd.cdna5.tdm.async_load(a_desc, [0, 0], self.a_buffer.index(slot), pred=pred)
        return gl.amd.cdna5.tdm.update_tensor_descriptor(a_desc, add_offsets=[0, BLOCK_K])

    @gluon.jit
    def issue_load_b_data(self, b_desc, load_idx, pred=1, phase=None):
        cfg = self.cfg
        BLOCK_K: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_B
        if phase is not None:
            slot = phase % cfg.NUM_BUFFERS
        else:
            slot = load_idx % cfg.NUM_BUFFERS
        gl.amd.cdna5.tdm.async_load(b_desc, [0, 0], self.b_buffer.index(slot), pred=pred)
        if cfg.TRANSPOSE_B:
            b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[0, BLOCK_K])
        else:
            b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[BLOCK_K, 0])
        return b_desc

    @gluon.jit
    def issue_l2_prefetch_a(self, distance: gl.constexpr, load_idx, pred=True):
        cfg = self.cfg
        if distance < 0:
            return

        prefetch_iteration = load_idx + distance
        BLOCK_K_PACKED_A: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A
        gl.amd.cdna5.tdm.prefetch(self.a_desc, [0, prefetch_iteration * BLOCK_K_PACKED_A], pred=pred)

    @gluon.jit
    def issue_l2_prefetch_b(self, distance: gl.constexpr, load_idx, pred=True):
        cfg = self.cfg
        if distance < 0:
            return

        prefetch_iteration = load_idx + distance
        BLOCK_K_PACKED_B: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_B
        if cfg.TRANSPOSE_B:
            gl.amd.cdna5.tdm.prefetch(self.b_desc, [0, prefetch_iteration * BLOCK_K_PACKED_B], pred=pred)
        else:
            gl.amd.cdna5.tdm.prefetch(self.b_desc, [prefetch_iteration * BLOCK_K_PACKED_B, 0], pred=pred)

    @gluon.jit
    def issue_l2_prefetch_a_scale(self, distance: gl.constexpr, load_idx, pred=True):
        cfg = self.cfg
        if distance < 0:
            return

        if not cfg.ASYNC_COPY_SCALE:
            if cfg.WITH_A_SCALE:
                prefetch_iteration = load_idx + distance
                gl.amd.cdna5.tdm.prefetch(self.a_scale_desc, [0, prefetch_iteration * cfg.BLOCK_K_SCALE_PRESHUFFLED],
                                          pred=pred)

    @gluon.jit
    def issue_l2_prefetch_b_scale(self, distance: gl.constexpr, load_idx, pred=True):
        cfg = self.cfg
        if distance < 0:
            return

        if not cfg.ASYNC_COPY_SCALE:
            prefetch_iteration = load_idx + distance
            gl.amd.cdna5.tdm.prefetch(self.b_scale_desc, [0, prefetch_iteration * cfg.BLOCK_K_SCALE_PRESHUFFLED],
                                      pred=pred)

    @gluon.jit
    def issue_l2_prefetches_scale(self, distance: gl.constexpr, load_idx, pred=True):
        self.issue_l2_prefetch_a_scale(distance, load_idx, pred=pred)
        self.issue_l2_prefetch_b_scale(distance, load_idx, pred=pred)

    @gluon.jit
    def issue_l2_prefetches_data(self, distance: gl.constexpr, load_idx, pred=True):
        self.issue_l2_prefetch_a(distance, load_idx, pred=pred)
        self.issue_l2_prefetch_b(distance, load_idx, pred=pred)

    @gluon.jit
    def issue_l2_prefetches(self, distance: gl.constexpr, load_idx, pred=True):
        self.issue_l2_prefetches_scale(distance, load_idx, pred=pred)
        self.issue_l2_prefetches_data(distance, load_idx, pred=pred)

    @gluon.jit
    def issue_l2_prefetches_prologue(self, load_idx):
        cfg = self.cfg
        L2_PREFETCH_DISTANCE: gl.constexpr = cfg.L2_PREFETCH_DISTANCE
        NUM_BUFFERS: gl.constexpr = cfg.NUM_BUFFERS

        for i in gl.static_range(NUM_BUFFERS, NUM_BUFFERS + L2_PREFETCH_DISTANCE):
            self.issue_l2_prefetches(i, load_idx)

    @gluon.jit
    def issue_local_loads(self, wmma_idx, a_buffer, b_buffer, a_scale_buffer, b_scale_buffer):
        cfg = self.cfg
        BLOCK_K_SCALE: gl.constexpr = cfg.BLOCK_K // cfg.SCALE_BLOCK
        a = a_buffer.index(wmma_idx % cfg.NUM_BUFFERS).load(layout=cfg.dot_layout_a)
        if cfg.TRANSPOSE_B:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).permute([1, 0]).load(layout=cfg.dot_layout_b)
        else:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).load(layout=cfg.dot_layout_b)
        if cfg.WITH_A_SCALE:
            a_scale_buffer_slice = a_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
        b_scale_buffer_slice = b_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
        if cfg.SCALE_PRESHUFFLE:
            if cfg.WITH_A_SCALE:
                a_scale_buffer_slice = a_scale_buffer_slice.reshape((
                    cfg.BLOCK_M_PRESHUFFLED,  #
                    BLOCK_K_SCALE // cfg.SCALE_KWIDTH,  #
                    cfg.PRESHUFFLE_FACTOR // 4,  #
                    4,  #
                    cfg.SCALE_KWIDTH)).permute((0, 3, 2, 1, 4)).reshape((cfg.BLOCK_M, BLOCK_K_SCALE))
            b_scale_buffer_slice = b_scale_buffer_slice.reshape((
                cfg.BLOCK_N_PRESHUFFLED,  #
                BLOCK_K_SCALE // cfg.SCALE_KWIDTH,  #
                cfg.PRESHUFFLE_FACTOR // 4,  #
                4,  #
                cfg.SCALE_KWIDTH)).permute((0, 3, 2, 1, 4)).reshape((cfg.BLOCK_N, BLOCK_K_SCALE))
        if cfg.WITH_A_SCALE:
            scale_a = a_scale_buffer_slice.load(layout=cfg.layout_a_scale)
        else:
            # Use a placeholder to make compiler happy
            scale_a = 0
            scale_a = scale_a.to(gl.uint8)
        scale_b = b_scale_buffer_slice.load(layout=cfg.layout_b_scale)
        return a, b, scale_a, scale_b


@gluon.jit
def apply_activation_and_store(cfg, accumulator, c_desc, off_m, off_n):
    if cfg.ACTIVATION == "swiglu":
        output = gl.convert_layout(swiglu_epilogue(accumulator), cfg.acc_layout)
    else:
        output = accumulator
    BLOCK_N_STORE: gl.constexpr = cfg.BLOCK_N // (2 if cfg.ACTIVATION == "swiglu" else 1)
    acc_buffer = gl.allocate_shared_memory(output.dtype, [cfg.BLOCK_M, BLOCK_N_STORE], c_desc.layout)
    acc_buffer.store(output)
    gl.amd.cdna5.tdm.async_store(c_desc, [off_m, off_n], acc_buffer)
    gl.amd.cdna5.tdm.async_wait(0)


@composition
@gluon.aggregate
class MXFPGEMMPipelinedProgram:
    base: MXFPGEMMProgramBase

    cfg: MXFPGEMMConfig
    a_buffer: gl.shared_memory_descriptor
    b_buffer: gl.shared_memory_descriptor
    a_scale_buffer: gl.shared_memory_descriptor | gl.constexpr
    b_scale_buffer: gl.shared_memory_descriptor

    a_desc: tdm.tensor_descriptor
    b_desc: tdm.tensor_descriptor
    a_scale_desc: tdm.tensor_descriptor | gl.constexpr
    b_scale_desc: tdm.tensor_descriptor

    c_desc: tdm.tensor_descriptor
    c_off_m: gl.tensor
    c_off_n: gl.tensor

    @gluon.constexpr_function
    def __init__(self, cfg: MXFPGEMMConfig, a_buffer, b_buffer, a_scale_buffer, b_scale_buffer, a_desc, b_desc,
                 a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n):
        self.cfg = cfg
        self.a_buffer = a_buffer
        self.b_buffer = b_buffer
        # Have to use constexpr to workaround a compiler issue with optional scale
        if cfg.WITH_A_SCALE:
            self.a_scale_buffer = a_scale_buffer
        else:
            self.a_scale_buffer = gl.constexpr(a_scale_buffer)
        self.b_scale_buffer = b_scale_buffer
        self.a_desc = a_desc
        self.b_desc = b_desc
        if cfg.WITH_A_SCALE:
            self.a_scale_desc = a_scale_desc
        else:
            self.a_scale_desc = gl.constexpr(a_scale_desc)
        self.b_scale_desc = b_scale_desc
        self.c_desc = c_desc
        self.c_off_m = c_off_m
        self.c_off_n = c_off_n

        self.base = MXFPGEMMProgramBase()

    @gluon.jit
    def initialize(cfg: MXFPGEMMConfig, a_desc, b_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n):
        NUM_BUFFERS: gl.constexpr = cfg.NUM_BUFFERS
        a_buffer = gl.allocate_shared_memory(a_desc.dtype, shape=[NUM_BUFFERS] + a_desc.block_shape,
                                             layout=a_desc.layout)
        b_buffer = gl.allocate_shared_memory(b_desc.dtype, shape=[NUM_BUFFERS] + b_desc.block_shape,
                                             layout=b_desc.layout)
        if cfg.WITH_A_SCALE:
            a_scale_buffer = gl.allocate_shared_memory(a_scale_desc.dtype,
                                                       shape=[NUM_BUFFERS] + a_scale_desc.block_shape,
                                                       layout=a_scale_desc.layout)
        else:
            a_scale_buffer = gl.constexpr(0)

        b_scale_buffer = gl.allocate_shared_memory(b_scale_desc.dtype, shape=[NUM_BUFFERS] + b_scale_desc.block_shape,
                                                   layout=b_scale_desc.layout)

        return MXFPGEMMPipelinedProgram(cfg, a_buffer, b_buffer, a_scale_buffer, b_scale_buffer, a_desc, b_desc,
                                        a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n)

    @gluon.jit
    def pipeline(self, K):
        cfg = self.cfg
        load_idx = 0
        wmma_idx = 0

        self.issue_l2_prefetches_prologue(load_idx)

        a_desc = self.a_desc
        b_desc = self.b_desc
        a_scale_desc = self.a_scale_desc
        b_scale_desc = self.b_scale_desc

        # prologue
        for _ in gl.static_range(cfg.NUM_BUFFERS - 1):
            if cfg.WITH_A_SCALE:
                a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx)
            b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx)
            a_desc = self.issue_load_a_data(a_desc, load_idx)
            b_desc = self.issue_load_b_data(b_desc, load_idx)
            load_idx = load_idx + 1

        accumulator = gl.zeros((cfg.BLOCK_M, cfg.BLOCK_N), dtype=gl.float32, layout=self.cfg.acc_layout)
        loop_ub = gl.cdiv(K, cfg.BLOCK_K)
        gl.assume(loop_ub > 0)
        epilogue_lb = loop_ub - (cfg.NUM_BUFFERS - 1)
        for i in range(0, loop_ub):
            pred = i - epilogue_lb
            pred = (pred >> 31) & 1
            phase = wmma_idx + cfg.NUM_BUFFERS - 1
            if cfg.WITH_A_SCALE:
                a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx, pred=pred, phase=phase)
            b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx, pred=pred, phase=phase)
            a_desc = self.issue_load_a_data(a_desc, load_idx, pred=pred, phase=phase)
            b_desc = self.issue_load_b_data(b_desc, load_idx, pred=pred, phase=phase)
            self.issue_l2_prefetches(cfg.L2_PREFETCH_DISTANCE, load_idx)
            load_idx = load_idx + 1

            gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 1) * self.cfg.NUM_LOADS_IN_BATCH)

            a, b, scale_a, scale_b = self.issue_local_loads(wmma_idx, self.a_buffer, self.b_buffer, self.a_scale_buffer,
                                                            self.b_scale_buffer)
            wmma_idx += 1
            accumulator = gl.amd.cdna5.wmma_scaled(a, scale_a, cfg.DTYPE_A, b, scale_b, cfg.DTYPE_B, accumulator)

        apply_activation_and_store(self.cfg, accumulator, self.c_desc, self.c_off_m, self.c_off_n)

    @gluon.jit
    def warp_pipeline(self, K):
        cfg = self.cfg
        load_idx = 0
        wmma_idx = 0

        self.issue_l2_prefetches_prologue(load_idx)

        a_desc = self.a_desc
        b_desc = self.b_desc
        a_scale_desc = self.a_scale_desc
        b_scale_desc = self.b_scale_desc

        # prologue
        for _ in gl.static_range(cfg.NUM_BUFFERS - 1):
            if cfg.WITH_A_SCALE:
                a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx)
            b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx)
            a_desc = self.issue_load_a_data(a_desc, load_idx)
            b_desc = self.issue_load_b_data(b_desc, load_idx)
            load_idx = load_idx + 1

        accumulator = gl.zeros((cfg.BLOCK_M, cfg.BLOCK_N), dtype=gl.float32, layout=self.cfg.acc_layout)
        loop_ub = gl.cdiv(K, cfg.BLOCK_K) - (cfg.NUM_BUFFERS - 1)
        gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 2) * self.cfg.NUM_LOADS_IN_BATCH)
        gl.assume(loop_ub >= 0)
        for _ in range(0, loop_ub):
            with gl.amd.warp_pipeline_stage("tdm+lds", priority=1):
                a, b, scale_a, scale_b = self.issue_local_loads(wmma_idx, self.a_buffer, self.b_buffer,
                                                                self.a_scale_buffer, self.b_scale_buffer)
                wmma_idx += 1
                phase = wmma_idx + cfg.NUM_BUFFERS - 2
                if cfg.WITH_A_SCALE:
                    a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx, phase=phase)
                b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx, phase=phase)
                a_desc = self.issue_load_a_data(a_desc, load_idx, phase=phase)
                b_desc = self.issue_load_b_data(b_desc, load_idx, phase=phase)

            gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 2) * self.cfg.NUM_LOADS_IN_BATCH)
            with gl.amd.warp_pipeline_stage("wmma", priority=0):
                self.issue_l2_prefetches(cfg.L2_PREFETCH_DISTANCE, load_idx)
                load_idx = load_idx + 1
                accumulator = gl.amd.cdna5.wmma_scaled(a, scale_a, cfg.DTYPE_A, b, scale_b, cfg.DTYPE_B, accumulator)

        # epilogue
        for i in gl.static_range(cfg.NUM_BUFFERS - 1):
            gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 2 - i) * self.cfg.NUM_LOADS_IN_BATCH)
            a, b, scale_a, scale_b = self.issue_local_loads(wmma_idx, self.a_buffer, self.b_buffer, self.a_scale_buffer,
                                                            self.b_scale_buffer)
            wmma_idx += 1
            accumulator = gl.amd.cdna5.wmma_scaled(a, scale_a, cfg.DTYPE_A, b, scale_b, cfg.DTYPE_B, accumulator)

        apply_activation_and_store(self.cfg, accumulator, self.c_desc, self.c_off_m, self.c_off_n)


@composition
@gluon.aggregate
class MXFPGEMMSliceKProgram:
    base: MXFPGEMMProgramBase

    cfg: MXFPGEMMConfig
    a_buffer: gl.shared_memory_descriptor
    b_buffer: gl.shared_memory_descriptor
    a_scale_buffer: gl.shared_memory_descriptor | gl.constexpr
    b_scale_buffer: gl.shared_memory_descriptor

    a_desc: tdm.tensor_descriptor
    b_desc: tdm.tensor_descriptor
    a_scale_desc: tdm.tensor_descriptor | gl.constexpr
    b_scale_desc: tdm.tensor_descriptor

    c_desc: tdm.tensor_descriptor
    c_off_m: gl.tensor
    c_off_n: gl.tensor

    @gluon.constexpr_function
    def __init__(self, cfg: MXFPGEMMConfig, a_buffer, b_buffer, a_scale_buffer, b_scale_buffer, a_desc, b_desc,
                 a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n):
        self.cfg = cfg
        self.a_buffer = a_buffer
        self.b_buffer = b_buffer
        # Have to use constexpr to workaround a compiler issue with optional scale
        if cfg.WITH_A_SCALE:
            self.a_scale_buffer = a_scale_buffer
        else:
            self.a_scale_buffer = gl.constexpr(a_scale_buffer)
        self.b_scale_buffer = b_scale_buffer
        self.a_desc = a_desc
        self.b_desc = b_desc
        if cfg.WITH_A_SCALE:
            self.a_scale_desc = a_scale_desc
        else:
            self.a_scale_desc = gl.constexpr(a_scale_desc)
        self.b_scale_desc = b_scale_desc
        self.c_desc = c_desc
        self.c_off_m = c_off_m
        self.c_off_n = c_off_n

        self.base = MXFPGEMMProgramBase()

    @gluon.jit
    def initialize(cfg: MXFPGEMMConfig, a_desc, b_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n):
        NUM_BUFFERS: gl.constexpr = cfg.NUM_BUFFERS
        a_buffer = gl.allocate_shared_memory(a_desc.dtype, shape=[NUM_BUFFERS] + a_desc.block_shape,
                                             layout=a_desc.layout)
        b_buffer = gl.allocate_shared_memory(b_desc.dtype, shape=[NUM_BUFFERS] + b_desc.block_shape,
                                             layout=b_desc.layout)
        if cfg.WITH_A_SCALE:
            a_scale_buffer = gl.allocate_shared_memory(a_scale_desc.dtype,
                                                       shape=[NUM_BUFFERS] + a_scale_desc.block_shape,
                                                       layout=a_scale_desc.layout)
        else:
            a_scale_buffer = gl.constexpr(0)

        b_scale_buffer = gl.allocate_shared_memory(b_scale_desc.dtype, shape=[NUM_BUFFERS] + b_scale_desc.block_shape,
                                                   layout=b_scale_desc.layout)

        return MXFPGEMMSliceKProgram(cfg, a_buffer, b_buffer, a_scale_buffer, b_scale_buffer, a_desc, b_desc,
                                     a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n)

    @gluon.jit
    def issue_subtile_local_loads(self, wmma_idx, subtile_start_idx: gl.constexpr, a_buffer, b_buffer, a_scale_buffer,
                                  b_scale_buffer):
        cfg = self.cfg
        NUM_SUBTILES_K: gl.constexpr = cfg.NUM_SUBTILES[2]
        SUBTILE_LEN: gl.constexpr = cfg.BLOCK_K // NUM_SUBTILES_K
        BLOCK_K_SCALE: gl.constexpr = cfg.BLOCK_K // cfg.SCALE_BLOCK
        SUBTILE_LEN_SCALE: gl.constexpr = SUBTILE_LEN // cfg.SCALE_BLOCK
        subtile_start: gl.constexpr = subtile_start_idx * SUBTILE_LEN
        a = a_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start // cfg.DIV_FACTOR_A,
                                                             SUBTILE_LEN // cfg.DIV_FACTOR_A,
                                                             1).load(layout=cfg.dot_layout_a)
        if cfg.TRANSPOSE_B:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start // cfg.DIV_FACTOR_B,
                                                                 SUBTILE_LEN // cfg.DIV_FACTOR_B,
                                                                 1).permute([1, 0]).load(layout=cfg.dot_layout_b)
        else:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start // cfg.DIV_FACTOR_B,
                                                                 SUBTILE_LEN // cfg.DIV_FACTOR_B,
                                                                 0).load(layout=cfg.dot_layout_b)
        if cfg.WITH_A_SCALE:
            a_scale_buffer_slice = a_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
        b_scale_buffer_slice = b_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
        if cfg.SCALE_PRESHUFFLE:
            if cfg.WITH_A_SCALE:
                a_scale_buffer_slice = a_scale_buffer_slice \
                    .reshape((cfg.BLOCK_M_PRESHUFFLED, BLOCK_K_SCALE // cfg.SCALE_KWIDTH, cfg.PRESHUFFLE_FACTOR // 4, 4, cfg.SCALE_KWIDTH)) \
                    .permute((0, 3, 2, 1, 4)) \
                    .reshape((cfg.BLOCK_M, BLOCK_K_SCALE))
            b_scale_buffer_slice = b_scale_buffer_slice \
                .reshape((cfg.BLOCK_N_PRESHUFFLED, BLOCK_K_SCALE // cfg.SCALE_KWIDTH, cfg.PRESHUFFLE_FACTOR // 4, 4, cfg.SCALE_KWIDTH)) \
                .permute((0, 3, 2, 1, 4)) \
                .reshape((cfg.BLOCK_N, BLOCK_K_SCALE))
        if cfg.WITH_A_SCALE:
            a_scale_buffer_slice = a_scale_buffer_slice.slice(subtile_start // cfg.SCALE_BLOCK, SUBTILE_LEN_SCALE, 1)
            scale_a = a_scale_buffer_slice.load(layout=cfg.layout_a_scale)
        else:
            scale_a = 0
            scale_a = scale_a.to(gl.uint8)

        b_scale_buffer_slice = b_scale_buffer_slice.slice(subtile_start // cfg.SCALE_BLOCK, SUBTILE_LEN_SCALE, 1)
        scale_b = b_scale_buffer_slice.load(layout=cfg.layout_b_scale)

        return a, b, scale_a, scale_b

    @gluon.jit
    def pipeline(self, K):
        cfg = self.cfg
        load_idx = 0
        wmma_idx = 0

        self.issue_l2_prefetches_prologue(load_idx)

        a_desc = self.a_desc
        b_desc = self.b_desc
        a_scale_desc = self.a_scale_desc
        b_scale_desc = self.b_scale_desc

        for _ in gl.static_range(cfg.NUM_BUFFERS - 1):
            if cfg.WITH_A_SCALE:
                a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx)
            b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx)
            a_desc = self.issue_load_a_data(a_desc, load_idx)
            b_desc = self.issue_load_b_data(b_desc, load_idx)
            load_idx = load_idx + 1

        gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 2) * cfg.NUM_LOADS_IN_BATCH)
        a0, b0, scale_a0, scale_b0 = self.issue_subtile_local_loads(wmma_idx, 0, self.a_buffer, self.b_buffer,
                                                                    self.a_scale_buffer, self.b_scale_buffer)

        accumulator = gl.zeros((cfg.BLOCK_M, cfg.BLOCK_N), dtype=gl.float32, layout=self.cfg.acc_layout)
        main_iters = gl.cdiv(K, cfg.BLOCK_K) - (cfg.NUM_BUFFERS - 1)
        gl.assume(main_iters >= 0)
        for _ in range(0, main_iters):
            # iter i
            accumulator = gl.amd.cdna5.wmma_scaled(a0, scale_a0, cfg.DTYPE_A, b0, scale_b0, cfg.DTYPE_B, accumulator)

            a1, b1, scale_a1, scale_b1 = self.issue_subtile_local_loads(wmma_idx, 1, self.a_buffer, self.b_buffer,
                                                                        self.a_scale_buffer, self.b_scale_buffer)
            wmma_idx += 1

            if cfg.WITH_A_SCALE:
                a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx)
            b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx)
            a_desc = self.issue_load_a_data(a_desc, load_idx)
            b_desc = self.issue_load_b_data(b_desc, load_idx)
            self.issue_l2_prefetches(cfg.L2_PREFETCH_DISTANCE, load_idx)
            load_idx = load_idx + 1

            accumulator = gl.amd.cdna5.wmma_scaled(a1, scale_a1, cfg.DTYPE_A, b1, scale_b1, cfg.DTYPE_B, accumulator)

            gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 2) * cfg.NUM_LOADS_IN_BATCH)
            a0, b0, scale_a0, scale_b0 = self.issue_subtile_local_loads(wmma_idx, 0, self.a_buffer, self.b_buffer,
                                                                        self.a_scale_buffer, self.b_scale_buffer)

        for i in gl.static_range(cfg.NUM_BUFFERS - 1):
            if i > 0:
                gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 2 - i) * cfg.NUM_LOADS_IN_BATCH)
                a0, b0, scale_a0, scale_b0 = self.issue_subtile_local_loads(wmma_idx, 0, self.a_buffer, self.b_buffer,
                                                                            self.a_scale_buffer, self.b_scale_buffer)
            accumulator = gl.amd.cdna5.wmma_scaled(a0, scale_a0, cfg.DTYPE_A, b0, scale_b0, cfg.DTYPE_B, accumulator)
            a1, b1, scale_a1, scale_b1 = self.issue_subtile_local_loads(wmma_idx, 1, self.a_buffer, self.b_buffer,
                                                                        self.a_scale_buffer, self.b_scale_buffer)
            wmma_idx += 1
            accumulator = gl.amd.cdna5.wmma_scaled(a1, scale_a1, cfg.DTYPE_A, b1, scale_b1, cfg.DTYPE_B, accumulator)

        apply_activation_and_store(self.cfg, accumulator, self.c_desc, self.c_off_m, self.c_off_n)

    @gluon.jit
    def warp_pipeline(self, K):
        cfg = self.cfg
        load_idx = 0
        wmma_idx = 0
        gl.static_assert(cfg.NUM_BUFFERS == 3)

        self.issue_l2_prefetches_prologue(load_idx)

        a_desc = self.a_desc
        b_desc = self.b_desc
        a_scale_desc = self.a_scale_desc
        b_scale_desc = self.b_scale_desc

        # prologue
        for _ in gl.static_range(cfg.NUM_BUFFERS - 1):
            if cfg.WITH_A_SCALE:
                a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx)
            b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx)
            a_desc = self.issue_load_a_data(a_desc, load_idx)
            b_desc = self.issue_load_b_data(b_desc, load_idx)
            load_idx = load_idx + 1

        accumulator = gl.zeros((cfg.BLOCK_M, cfg.BLOCK_N), dtype=gl.float32, layout=self.cfg.acc_layout)
        loop_ub = gl.cdiv(K, cfg.BLOCK_K) - (cfg.NUM_BUFFERS - 1)
        gl.assume(loop_ub >= 0)
        # wait for the first prefetch
        gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 2) * self.cfg.NUM_LOADS_IN_BATCH)
        for _ in range(0, loop_ub):
            with gl.amd.warp_pipeline_stage("lds0", priority=1):
                a0, b0, scale_a0, scale_b0 = self.issue_subtile_local_loads(wmma_idx, 0, self.a_buffer, self.b_buffer,
                                                                            self.a_scale_buffer, self.b_scale_buffer)

            gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 3) * self.cfg.NUM_LOADS_IN_BATCH)
            with gl.amd.warp_pipeline_stage("tdm+wmma+lds1", priority=0):
                phase = wmma_idx + cfg.NUM_BUFFERS - 1
                if cfg.WITH_A_SCALE:
                    a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx, phase=phase)
                b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx, phase=phase)
                a_desc = self.issue_load_a_data(a_desc, load_idx, phase=phase)
                b_desc = self.issue_load_b_data(b_desc, load_idx, phase=phase)
                self.issue_l2_prefetches(cfg.L2_PREFETCH_DISTANCE, load_idx)
                load_idx = load_idx + 1
                accumulator = gl.amd.cdna5.wmma_scaled(a0, scale_a0, cfg.DTYPE_A, b0, scale_b0, cfg.DTYPE_B,
                                                       accumulator)
                a1, b1, scale_a1, scale_b1 = self.issue_subtile_local_loads(wmma_idx, 1, self.a_buffer, self.b_buffer,
                                                                            self.a_scale_buffer, self.b_scale_buffer)
                wmma_idx += 1
                accumulator = gl.amd.cdna5.wmma_scaled(a1, scale_a1, cfg.DTYPE_A, b1, scale_b1, cfg.DTYPE_B,
                                                       accumulator)

        # epilogue
        for i in gl.static_range(cfg.NUM_BUFFERS - 1):
            gl.amd.cdna5.tdm.async_wait((cfg.NUM_BUFFERS - 1 - i) * self.cfg.NUM_LOADS_IN_BATCH)
            a0, b0, scale_a0, scale_b0 = self.issue_subtile_local_loads(wmma_idx, 0, self.a_buffer, self.b_buffer,
                                                                        self.a_scale_buffer, self.b_scale_buffer)
            accumulator = gl.amd.cdna5.wmma_scaled(a0, scale_a0, cfg.DTYPE_A, b0, scale_b0, cfg.DTYPE_B, accumulator)

            a1, b1, scale_a1, scale_b1 = self.issue_subtile_local_loads(wmma_idx, 1, self.a_buffer, self.b_buffer,
                                                                        self.a_scale_buffer, self.b_scale_buffer)
            accumulator = gl.amd.cdna5.wmma_scaled(a1, scale_a1, cfg.DTYPE_A, b1, scale_b1, cfg.DTYPE_B, accumulator)
            wmma_idx += 1

        apply_activation_and_store(self.cfg, accumulator, self.c_desc, self.c_off_m, self.c_off_n)


@composition
@gluon.aggregate
class MXFPGEMMSliceNKProgram:
    base: MXFPGEMMProgramBase

    cfg: MXFPGEMMConfig
    a_buffer: gl.shared_memory_descriptor
    b_buffer: gl.shared_memory_descriptor
    a_scale_buffer: gl.shared_memory_descriptor | gl.constexpr
    b_scale_buffer: gl.shared_memory_descriptor

    a_desc: tdm.tensor_descriptor
    b_desc: tdm.tensor_descriptor
    a_scale_desc: tdm.tensor_descriptor | ScaleAsyncCopyDescriptor | gl.constexpr
    b_scale_desc: tdm.tensor_descriptor | ScaleAsyncCopyDescriptor

    c_desc: tdm.tensor_descriptor
    c_off_m: gl.tensor
    c_off_n: gl.tensor

    @gluon.constexpr_function
    def __init__(self, cfg: MXFPGEMMConfig, a_buffer, b_buffer, a_scale_buffer, b_scale_buffer, a_desc, b_desc,
                 a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n):
        self.cfg = cfg
        self.a_buffer = a_buffer
        self.b_buffer = b_buffer
        if cfg.WITH_A_SCALE:
            self.a_scale_buffer = a_scale_buffer
        else:
            self.a_scale_buffer = gl.constexpr(a_scale_buffer)

        self.b_scale_buffer = b_scale_buffer
        self.a_desc = a_desc
        self.b_desc = b_desc
        if cfg.WITH_A_SCALE:
            self.a_scale_desc = a_scale_desc
        else:
            self.a_scale_desc = gl.constexpr(a_scale_desc)
        self.b_scale_desc = b_scale_desc
        self.c_desc = c_desc
        self.c_off_m = c_off_m
        self.c_off_n = c_off_n

        self.base = MXFPGEMMProgramBase()

    @gluon.jit
    def initialize(cfg: MXFPGEMMConfig, a_desc, b_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n):
        NUM_BUFFERS: gl.constexpr = cfg.NUM_BUFFERS
        a_buffer = gl.allocate_shared_memory(a_desc.dtype, shape=[NUM_BUFFERS] + a_desc.block_shape,
                                             layout=a_desc.layout)
        b_buffer = gl.allocate_shared_memory(b_desc.dtype, shape=[NUM_BUFFERS] + b_desc.block_shape,
                                             layout=b_desc.layout)
        if cfg.WITH_A_SCALE:
            a_scale_buffer = gl.allocate_shared_memory(a_scale_desc.dtype,
                                                       shape=[NUM_BUFFERS] + a_scale_desc.block_shape,
                                                       layout=a_scale_desc.layout)
        else:
            a_scale_buffer = gl.constexpr(0)

        b_scale_buffer = gl.allocate_shared_memory(b_scale_desc.dtype, shape=[NUM_BUFFERS] + b_scale_desc.block_shape,
                                                   layout=b_scale_desc.layout)

        return MXFPGEMMSliceNKProgram(cfg, a_buffer, b_buffer, a_scale_buffer, b_scale_buffer, a_desc, b_desc,
                                      a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n)

    @gluon.jit
    def issue_local_load_a(self, wmma_idx, subtile_start_idx: gl.constexpr, a_buffer, a_scale_buffer):
        cfg = self.cfg
        NUM_SUBTILES_K: gl.constexpr = cfg.NUM_SUBTILES[2]
        SUBTILE_LEN: gl.constexpr = cfg.BLOCK_K // NUM_SUBTILES_K
        BLOCK_K_SCALE: gl.constexpr = cfg.BLOCK_K // cfg.SCALE_BLOCK
        subtile_start: gl.constexpr = subtile_start_idx * SUBTILE_LEN
        a = a_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start // cfg.DIV_FACTOR_A,
                                                             SUBTILE_LEN // cfg.DIV_FACTOR_A,
                                                             1).load(layout=cfg.dot_layout_a)
        if cfg.WITH_A_SCALE:
            a_scale_buffer_slice = a_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
            if cfg.SCALE_PRESHUFFLE:
                a_scale_buffer_slice = a_scale_buffer_slice.reshape((
                    cfg.BLOCK_M_PRESHUFFLED,  #
                    BLOCK_K_SCALE // cfg.SCALE_KWIDTH,  #
                    cfg.PRESHUFFLE_FACTOR // 4,  #
                    4,  #
                    cfg.SCALE_KWIDTH)).permute((0, 3, 2, 1, 4)).reshape((cfg.BLOCK_M, BLOCK_K_SCALE))
            a_scale_buffer_slice = a_scale_buffer_slice.slice(subtile_start // cfg.SCALE_BLOCK,
                                                              SUBTILE_LEN // cfg.SCALE_BLOCK, 1)
            scale_a = a_scale_buffer_slice.load(layout=cfg.layout_a_scale)
        else:
            scale_a = 0
            scale_a = scale_a.to(gl.uint8)
        return a, scale_a

    @gluon.jit
    def issue_local_load_b(self, wmma_idx, subtile_start_idx_k: gl.constexpr, subtile_start_idx_n: gl.constexpr,
                           b_buffer, b_scale_buffer):
        cfg = self.cfg
        NUM_SUBTILES_N: gl.constexpr = cfg.NUM_SUBTILES[1]
        NUM_SUBTILES_K: gl.constexpr = cfg.NUM_SUBTILES[2]
        SUBTILE_LEN_K: gl.constexpr = cfg.BLOCK_K // NUM_SUBTILES_K
        SUBTILE_LEN_N: gl.constexpr = cfg.BLOCK_N // NUM_SUBTILES_N
        BLOCK_K_SCALE: gl.constexpr = cfg.BLOCK_K // cfg.SCALE_BLOCK
        subtile_start_k: gl.constexpr = subtile_start_idx_k * SUBTILE_LEN_K
        subtile_start_n: gl.constexpr = subtile_start_idx_n * SUBTILE_LEN_N
        if cfg.TRANSPOSE_B:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start_n, SUBTILE_LEN_N, 0) \
            .slice(subtile_start_k // cfg.DIV_FACTOR_B, SUBTILE_LEN_K // cfg.DIV_FACTOR_B, 1) \
            .permute([1, 0]).load(layout=cfg.dot_layout_b)
        else:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start_k // cfg.DIV_FACTOR_B,
                                                                 SUBTILE_LEN_K // cfg.DIV_FACTOR_B, 0) \
            .slice(subtile_start_n, SUBTILE_LEN_N, 1) \
            .load(layout=cfg.dot_layout_b)
        b_scale_buffer_slice = b_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
        if cfg.SCALE_PRESHUFFLE:
            b_scale_buffer_slice = b_scale_buffer_slice.reshape((
                cfg.BLOCK_N_PRESHUFFLED,  #
                BLOCK_K_SCALE // cfg.SCALE_KWIDTH,  #
                cfg.PRESHUFFLE_FACTOR // 4,  #
                4,  #
                cfg.SCALE_KWIDTH)).permute((0, 3, 2, 1, 4)).reshape((cfg.BLOCK_N, BLOCK_K_SCALE))
        b_scale_buffer_slice = b_scale_buffer_slice.slice(subtile_start_n, SUBTILE_LEN_N, 0) \
            .slice(subtile_start_k // cfg.SCALE_BLOCK, SUBTILE_LEN_K // cfg.SCALE_BLOCK, 1)
        scale_b = b_scale_buffer_slice.load(layout=cfg.layout_b_scale)
        return b, scale_b

    @gluon.jit
    def issue_load_a_scale(self, a_scale_desc, load_idx, pred=1):
        cfg = self.cfg
        if cfg.WITH_A_SCALE:
            a_scale_buffer_slice = self.a_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
            if cfg.ASYNC_COPY_SCALE:
                a_scale_desc.issue_async_load(load_idx, a_scale_buffer_slice, pred=pred)
            else:
                gl.amd.cdna5.tdm.async_load(a_scale_desc, [0, 0],  #
                                            a_scale_buffer_slice, pred=pred)
                a_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_scale_desc,
                                                                         add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
        return a_scale_desc

    @gluon.jit
    def issue_load_a_data(self, a_desc, load_idx, pred=1):
        cfg = self.cfg
        BLOCK_K: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A
        gl.amd.cdna5.tdm.async_load(a_desc, [0, 0],  #
                                    self.a_buffer.index(load_idx % cfg.NUM_BUFFERS),  #
                                    pred=pred)
        return gl.amd.cdna5.tdm.update_tensor_descriptor(a_desc, add_offsets=[0, BLOCK_K])

    @gluon.jit
    def issue_load_b_scale(self, b_scale_desc, load_idx, pred=1):
        cfg = self.cfg
        b_scale_buffer_slice = self.b_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
        if cfg.ASYNC_COPY_SCALE:
            b_scale_desc.issue_async_load(load_idx, b_scale_buffer_slice, pred=pred)
        else:
            gl.amd.cdna5.tdm.async_load(b_scale_desc, [0, 0],  #
                                        b_scale_buffer_slice, pred=pred)
            b_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_scale_desc,
                                                                     add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
        return b_scale_desc

    @gluon.jit
    def issue_load_b_data(self, b_desc, load_idx, pred=1):
        cfg = self.cfg
        BLOCK_K: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_B
        gl.amd.cdna5.tdm.async_load(b_desc, [0, 0],  #
                                    self.b_buffer.index(load_idx % cfg.NUM_BUFFERS),  #
                                    pred=pred)
        if cfg.TRANSPOSE_B:
            b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[0, BLOCK_K])
        else:
            b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[BLOCK_K, 0])
        return b_desc

    @gluon.jit
    def async_wait(self, waitcnt: int):
        cfg = self.cfg
        if cfg.ASYNC_COPY_SCALE:
            gl.amd.cdna5.tdm.async_wait(waitcnt)
            cp.wait_group(waitcnt if cfg.WITH_A_SCALE else (waitcnt // 2))
        else:
            gl.amd.cdna5.tdm.async_wait((waitcnt * 2) if cfg.WITH_A_SCALE else (waitcnt + waitcnt // 2))

    @gluon.jit
    def pipeline(self, K):
        cfg = self.cfg
        load_idx = 0
        wmma_idx = 0

        a_desc = self.a_desc
        b_desc = self.b_desc
        a_scale_desc = self.a_scale_desc
        b_scale_desc = self.b_scale_desc

        for _ in gl.static_range(cfg.NUM_BUFFERS - 1):
            if cfg.WITH_A_SCALE:
                a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx)
            b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx)
            a_desc = self.issue_load_a_data(a_desc, load_idx)
            b_desc = self.issue_load_b_data(b_desc, load_idx)
            load_idx = load_idx + 1

        self.async_wait((cfg.NUM_BUFFERS - 2) * 2)
        a0, scale_a0 = self.issue_local_load_a(wmma_idx, 0, self.a_buffer, self.a_scale_buffer)
        b00, scale_b00 = self.issue_local_load_b(wmma_idx, 0, 0, self.b_buffer, self.b_scale_buffer)

        if cfg.WITH_A_SCALE:
            a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx)
        b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx)
        a_desc = self.issue_load_a_data(a_desc, load_idx)
        b_desc = self.issue_load_b_data(b_desc, load_idx)

        c0 = gl.zeros((cfg.BLOCK_M // cfg.NUM_SUBTILES[0], cfg.BLOCK_N // cfg.NUM_SUBTILES[1]), dtype=gl.float32,
                      layout=cfg.acc_layout)
        c1 = gl.zeros((cfg.BLOCK_M // cfg.NUM_SUBTILES[0], cfg.BLOCK_N // cfg.NUM_SUBTILES[1]), dtype=gl.float32,
                      layout=cfg.acc_layout)
        loop_ub = gl.cdiv(K, cfg.BLOCK_K)
        epilogue_lb = loop_ub - (cfg.NUM_BUFFERS - 1)
        gl.assume(loop_ub >= cfg.NUM_BUFFERS)
        for i in range(0, loop_ub):
            # iter i
            c0 = gl.amd.cdna5.wmma_scaled(a0, scale_a0, cfg.DTYPE_A, b00, scale_b00, cfg.DTYPE_B, c0)
            b01, scale_b01 = self.issue_local_load_b(wmma_idx, 0, 1, self.b_buffer, self.b_scale_buffer)

            pred_prefetch = i - epilogue_lb
            pred_prefetch = ((pred_prefetch >> 31) & 1).to(gl.int1)
            self.issue_l2_prefetches(cfg.L2_PREFETCH_DISTANCE, load_idx, pred=pred_prefetch)
            load_idx = load_idx + 1

            # iter i
            c1 = gl.amd.cdna5.wmma_scaled(a0, scale_a0, cfg.DTYPE_A, b01, scale_b01, cfg.DTYPE_B, c1)
            a1, scale_a1 = self.issue_local_load_a(wmma_idx, 1, self.a_buffer, self.a_scale_buffer)
            b10, scale_b10 = self.issue_local_load_b(wmma_idx, 1, 0, self.b_buffer, self.b_scale_buffer)

            # iter i
            c0 = gl.amd.cdna5.wmma_scaled(a1, scale_a1, cfg.DTYPE_A, b10, scale_b10, cfg.DTYPE_B, c0)
            b11, scale_b11 = self.issue_local_load_b(wmma_idx, 1, 1, self.b_buffer, self.b_scale_buffer)

            # iter i
            wmma_idx += 1
            c1 = gl.amd.cdna5.wmma_scaled(a1, scale_a1, cfg.DTYPE_A, b11, scale_b11, cfg.DTYPE_B, c1)

            # iter i + NUM_BUFFERS - 1
            pred_load = i + 1 - epilogue_lb
            pred_load = (pred_load >> 31) & 1
            self.async_wait(0)
            if cfg.WITH_A_SCALE:
                a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx, pred=pred_load)
            b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx, pred=pred_load)
            a_desc = self.issue_load_a_data(a_desc, load_idx, pred=pred_load)
            b_desc = self.issue_load_b_data(b_desc, load_idx, pred=pred_load)

            a0, scale_a0 = self.issue_local_load_a(wmma_idx, 0, self.a_buffer, self.a_scale_buffer)
            b00, scale_b00 = self.issue_local_load_b(wmma_idx, 0, 0, self.b_buffer, self.b_scale_buffer)

        accumulator = gl.join(c0, c1)
        accumulator = accumulator.permute(0, 2, 1).reshape((cfg.BLOCK_M, cfg.BLOCK_N))
        accumulator = gl.convert_layout(accumulator, cfg.acc_layout, assert_trivial=True)

        apply_activation_and_store(self.cfg, accumulator, self.c_desc, self.c_off_m, self.c_off_n)


@composition
@gluon.aggregate
class MXFPGEMMSliceMNKProgram:
    base: MXFPGEMMProgramBase

    cfg: MXFPGEMMConfig
    a_buffer: gl.shared_memory_descriptor
    b_buffer: gl.shared_memory_descriptor
    a_scale_buffer: gl.shared_memory_descriptor | gl.constexpr
    b_scale_buffer: gl.shared_memory_descriptor

    a_desc: tdm.tensor_descriptor
    b_desc: tdm.tensor_descriptor
    a_scale_desc: tdm.tensor_descriptor | ScaleAsyncCopyDescriptor
    b_scale_desc: tdm.tensor_descriptor | ScaleAsyncCopyDescriptor

    c_desc: tdm.tensor_descriptor
    c_off_m: gl.tensor
    c_off_n: gl.tensor

    @gluon.constexpr_function
    def __init__(self, cfg: MXFPGEMMConfig, a_buffer, b_buffer, a_scale_buffer, b_scale_buffer, a_desc, b_desc,
                 a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n):
        self.cfg = cfg
        self.a_buffer = a_buffer
        self.b_buffer = b_buffer
        if cfg.WITH_A_SCALE:
            self.a_scale_buffer = a_scale_buffer
        else:
            self.a_scale_buffer = gl.constexpr(a_scale_buffer)

        self.b_scale_buffer = b_scale_buffer
        self.a_desc = a_desc
        self.b_desc = b_desc
        self.a_scale_desc = a_scale_desc
        self.b_scale_desc = b_scale_desc
        self.c_desc = c_desc
        self.c_off_m = c_off_m
        self.c_off_n = c_off_n

        self.base = MXFPGEMMProgramBase()

    @gluon.jit
    def initialize(cfg: MXFPGEMMConfig, a_desc, b_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n):
        NUM_BUFFERS: gl.constexpr = cfg.NUM_BUFFERS
        a_buffer = gl.allocate_shared_memory(a_desc.dtype, shape=[NUM_BUFFERS] + a_desc.block_shape,
                                             layout=a_desc.layout)
        b_buffer = gl.allocate_shared_memory(b_desc.dtype, shape=[NUM_BUFFERS] + b_desc.block_shape,
                                             layout=b_desc.layout)
        if cfg.WITH_A_SCALE:
            a_scale_buffer = gl.allocate_shared_memory(a_scale_desc.dtype,
                                                       shape=[NUM_BUFFERS] + a_scale_desc.block_shape,
                                                       layout=a_scale_desc.layout)
        else:
            a_scale_buffer = gl.constexpr(0)

        b_scale_buffer = gl.allocate_shared_memory(b_scale_desc.dtype, shape=[NUM_BUFFERS] + b_scale_desc.block_shape,
                                                   layout=b_scale_desc.layout)

        return MXFPGEMMSliceMNKProgram(cfg, a_buffer, b_buffer, a_scale_buffer, b_scale_buffer, a_desc, b_desc,
                                       a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n)

    @gluon.jit
    def issue_local_load_a(self, wmma_idx, subtile_start_idx_m: gl.constexpr, subtile_start_idx_k: gl.constexpr,
                           a_buffer, a_scale_buffer):
        cfg = self.cfg
        NUM_SUBTILES_M: gl.constexpr = cfg.NUM_SUBTILES[0]
        NUM_SUBTILES_K: gl.constexpr = cfg.NUM_SUBTILES[2]
        SUBTILE_LEN_M: gl.constexpr = cfg.BLOCK_M // NUM_SUBTILES_M
        SUBTILE_LEN_K: gl.constexpr = cfg.BLOCK_K // NUM_SUBTILES_K
        BLOCK_K_SCALE: gl.constexpr = cfg.BLOCK_K // cfg.SCALE_BLOCK
        subtile_start_m: gl.constexpr = subtile_start_idx_m * SUBTILE_LEN_M
        subtile_start_k: gl.constexpr = subtile_start_idx_k * SUBTILE_LEN_K
        a = a_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start_m, SUBTILE_LEN_M, 0) \
            .slice(subtile_start_k // cfg.DIV_FACTOR_A, SUBTILE_LEN_K // cfg.DIV_FACTOR_A, 1) \
            .load(layout=cfg.dot_layout_a)
        if cfg.WITH_A_SCALE:
            a_scale_buffer_slice = a_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
            if cfg.SCALE_PRESHUFFLE:
                a_scale_buffer_slice = a_scale_buffer_slice.reshape((
                    cfg.BLOCK_M_PRESHUFFLED,  #
                    BLOCK_K_SCALE // cfg.SCALE_KWIDTH,  #
                    cfg.PRESHUFFLE_FACTOR // 4,  #
                    4,  #
                    cfg.SCALE_KWIDTH)).permute((0, 3, 2, 1, 4)).reshape((cfg.BLOCK_M, BLOCK_K_SCALE))
            a_scale_buffer_slice = a_scale_buffer_slice.slice(subtile_start_m, SUBTILE_LEN_M, 0) \
                .slice(subtile_start_k // cfg.SCALE_BLOCK, SUBTILE_LEN_K // cfg.SCALE_BLOCK, 1)
            scale_a = a_scale_buffer_slice.load(layout=cfg.layout_a_scale)
        else:
            scale_a = 0
            scale_a = scale_a.to(gl.uint8)
        return a, scale_a

    @gluon.jit
    def issue_local_load_b(self, wmma_idx, subtile_start_idx_k: gl.constexpr, subtile_start_idx_n: gl.constexpr,
                           b_buffer, b_scale_buffer):
        cfg = self.cfg
        NUM_SUBTILES_N: gl.constexpr = cfg.NUM_SUBTILES[1]
        NUM_SUBTILES_K: gl.constexpr = cfg.NUM_SUBTILES[2]
        SUBTILE_LEN_K: gl.constexpr = cfg.BLOCK_K // NUM_SUBTILES_K
        SUBTILE_LEN_N: gl.constexpr = cfg.BLOCK_N // NUM_SUBTILES_N
        BLOCK_K_SCALE: gl.constexpr = cfg.BLOCK_K // cfg.SCALE_BLOCK
        subtile_start_k: gl.constexpr = subtile_start_idx_k * SUBTILE_LEN_K
        subtile_start_n: gl.constexpr = subtile_start_idx_n * SUBTILE_LEN_N
        if cfg.TRANSPOSE_B:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start_n, SUBTILE_LEN_N, 0) \
            .slice(subtile_start_k // cfg.DIV_FACTOR_B, SUBTILE_LEN_K // cfg.DIV_FACTOR_B, 1) \
            .permute([1, 0]).load(layout=cfg.dot_layout_b)
        else:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start_k // cfg.DIV_FACTOR_B,
                                                                 SUBTILE_LEN_K // cfg.DIV_FACTOR_B, 0) \
            .slice(subtile_start_n, SUBTILE_LEN_N, 1) \
            .load(layout=cfg.dot_layout_b)
        b_scale_buffer_slice = b_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
        if cfg.SCALE_PRESHUFFLE:
            b_scale_buffer_slice = b_scale_buffer_slice.reshape((
                cfg.BLOCK_N_PRESHUFFLED,  #
                BLOCK_K_SCALE // cfg.SCALE_KWIDTH,  #
                cfg.PRESHUFFLE_FACTOR // 4,  #
                4,  #
                cfg.SCALE_KWIDTH)).permute((0, 3, 2, 1, 4)).reshape((cfg.BLOCK_N, BLOCK_K_SCALE))
        b_scale_buffer_slice = b_scale_buffer_slice.slice(subtile_start_n, SUBTILE_LEN_N, 0) \
            .slice(subtile_start_k // cfg.SCALE_BLOCK, SUBTILE_LEN_K // cfg.SCALE_BLOCK, 1)
        scale_b = b_scale_buffer_slice.load(layout=cfg.layout_b_scale)
        return b, scale_b

    @gluon.jit
    def issue_load_a_scale(self, a_scale_desc, load_idx, pred=1):
        cfg = self.cfg
        if cfg.WITH_A_SCALE:
            a_scale_buffer_slice = self.a_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
            if cfg.ASYNC_COPY_SCALE:
                a_scale_desc.issue_async_load(load_idx, a_scale_buffer_slice, pred=pred)
            else:
                gl.amd.cdna5.tdm.async_load(a_scale_desc, [0, 0],  #
                                            a_scale_buffer_slice, pred=pred, warp_used_hint=cfg.TDM_WARP_USED_HINT_A)
                a_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_scale_desc,
                                                                         add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
        return a_scale_desc

    @gluon.jit
    def issue_load_a_data(self, a_desc, load_idx, pred=1):
        cfg = self.cfg
        BLOCK_K: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A
        gl.amd.cdna5.tdm.async_load(a_desc, [0, 0],  #
                                    self.a_buffer.index(load_idx % cfg.NUM_BUFFERS),  #
                                    pred=pred, warp_used_hint=cfg.TDM_WARP_USED_HINT_A)
        return gl.amd.cdna5.tdm.update_tensor_descriptor(a_desc, add_offsets=[0, BLOCK_K])

    @gluon.jit
    def issue_load_b_scale(self, b_scale_desc, load_idx, pred=1):
        cfg = self.cfg
        b_scale_buffer_slice = self.b_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
        if cfg.ASYNC_COPY_SCALE:
            b_scale_desc.issue_async_load(load_idx, b_scale_buffer_slice, pred=pred)
        else:
            gl.amd.cdna5.tdm.async_load(b_scale_desc, [0, 0],  #
                                        b_scale_buffer_slice, pred=pred, warp_used_hint=cfg.TDM_WARP_USED_HINT_B)
            b_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_scale_desc,
                                                                     add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
        return b_scale_desc

    @gluon.jit
    def issue_load_b_data(self, b_desc, load_idx, pred=1):
        cfg = self.cfg
        BLOCK_K: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_B
        gl.amd.cdna5.tdm.async_load(b_desc, [0, 0],  #
                                    self.b_buffer.index(load_idx % cfg.NUM_BUFFERS),  #
                                    pred=pred, warp_used_hint=cfg.TDM_WARP_USED_HINT_B)
        if cfg.TRANSPOSE_B:
            b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[0, BLOCK_K])
        else:
            b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[BLOCK_K, 0])
        return b_desc

    @gluon.jit
    def issue_load_scales(self, a_scale_desc, b_scale_desc, load_idx, pred=1):
        cfg = self.cfg
        if cfg.WITH_A_SCALE:
            if not cfg.ASYNC_COPY_SCALE and cfg.TDM_WARP_USED_HINT_A:
                a_dest = self.a_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
                b_dest = self.b_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
                a_scale_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_scale_desc, add_offsets=[0, 0],
                                                                              pred=pred, clamp_bounds=True)
                b_scale_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_scale_desc, add_offsets=[0, 0],
                                                                              pred=pred, clamp_bounds=True)
                gl.amd.cdna5.tdm.async_load_fused([(a_scale_load_desc, a_dest, cfg.TDM_WARP_USED_HINT_A),
                                                   (b_scale_load_desc, b_dest, cfg.TDM_WARP_USED_HINT_B)])
                a_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_scale_desc,
                                                                         add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
                b_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_scale_desc,
                                                                         add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
                return a_scale_desc, b_scale_desc
            a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx, pred=pred)
        b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx, pred=pred)
        return a_scale_desc, b_scale_desc

    @gluon.jit
    def issue_load_data(self, a_desc, b_desc, load_idx, pred=1):
        cfg = self.cfg
        if cfg.TDM_WARP_USED_HINT_A:
            # Partial TDM: fuse the A/B data loads into one explicit TDM
            # instruction instead of two separately-hinted loads.
            BLOCK_K_A: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A
            BLOCK_K_B: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_B
            a_dest = self.a_buffer.index(load_idx % cfg.NUM_BUFFERS)
            b_dest = self.b_buffer.index(load_idx % cfg.NUM_BUFFERS)
            a_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_desc, add_offsets=[0, 0], pred=pred,
                                                                    clamp_bounds=True)
            b_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[0, 0], pred=pred,
                                                                    clamp_bounds=True)
            gl.amd.cdna5.tdm.async_load_fused([(a_load_desc, a_dest, cfg.TDM_WARP_USED_HINT_A),
                                               (b_load_desc, b_dest, cfg.TDM_WARP_USED_HINT_B)])
            a_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_desc, add_offsets=[0, BLOCK_K_A])
            if cfg.TRANSPOSE_B:
                b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[0, BLOCK_K_B])
            else:
                b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[BLOCK_K_B, 0])
        else:
            a_desc = self.issue_load_a_data(a_desc, load_idx, pred=pred)
            b_desc = self.issue_load_b_data(b_desc, load_idx, pred=pred)
        return a_desc, b_desc

    @gluon.jit
    def async_wait(self, waitcnt: int):
        cfg = self.cfg
        # A fused TDM load counts as a single outstanding op regardless of its
        # member count, unlike separately-hinted loads which count individually.
        DATA_OPS: gl.constexpr = 1 if cfg.TDM_WARP_USED_HINT_A else 2
        SCALE_OPS: gl.constexpr = 1 if cfg.WITH_A_SCALE and cfg.TDM_WARP_USED_HINT_A else (2 if cfg.WITH_A_SCALE else 1)
        if cfg.ASYNC_COPY_SCALE:
            gl.amd.cdna5.tdm.async_wait(waitcnt * DATA_OPS // 2)
            cp.wait_group(waitcnt if cfg.WITH_A_SCALE else (waitcnt // 2))
        else:
            gl.amd.cdna5.tdm.async_wait(waitcnt * (DATA_OPS + SCALE_OPS) // 2)

    @gluon.jit
    def pipeline(self, K):
        cfg = self.cfg
        load_idx = 0
        wmma_idx = 0

        a_desc = self.a_desc
        b_desc = self.b_desc
        a_scale_desc = self.a_scale_desc
        b_scale_desc = self.b_scale_desc

        for _ in gl.static_range(cfg.NUM_BUFFERS - 1):
            a_scale_desc, b_scale_desc = self.issue_load_scales(a_scale_desc, b_scale_desc, load_idx)
            a_desc, b_desc = self.issue_load_data(a_desc, b_desc, load_idx)
            load_idx = load_idx + 1

        for d in gl.static_range(cfg.NUM_BUFFERS, cfg.NUM_BUFFERS + cfg.L2_PREFETCH_DISTANCE):
            self.issue_l2_prefetches_data(d, 0)
        for d in gl.static_range(cfg.NUM_BUFFERS, cfg.NUM_BUFFERS + cfg.L2_PREFETCH_DISTANCE):
            self.issue_l2_prefetches_scale(d, 0)

        self.async_wait((cfg.NUM_BUFFERS - 2) * 2)
        a00, scale_a00 = self.issue_local_load_a(wmma_idx, 0, 0, self.a_buffer, self.a_scale_buffer)
        b00, scale_b00 = self.issue_local_load_b(wmma_idx, 0, 0, self.b_buffer, self.b_scale_buffer)

        a_scale_desc, b_scale_desc = self.issue_load_scales(a_scale_desc, b_scale_desc, load_idx)
        a_desc, b_desc = self.issue_load_data(a_desc, b_desc, load_idx)

        SUBTILE_M: gl.constexpr = cfg.BLOCK_M // cfg.NUM_SUBTILES[0]
        SUBTILE_N: gl.constexpr = cfg.BLOCK_N // cfg.NUM_SUBTILES[1]
        c00 = gl.zeros((SUBTILE_M, SUBTILE_N), dtype=gl.float32, layout=cfg.acc_layout)
        c01 = gl.zeros((SUBTILE_M, SUBTILE_N), dtype=gl.float32, layout=cfg.acc_layout)
        c10 = gl.zeros((SUBTILE_M, SUBTILE_N), dtype=gl.float32, layout=cfg.acc_layout)
        c11 = gl.zeros((SUBTILE_M, SUBTILE_N), dtype=gl.float32, layout=cfg.acc_layout)

        loop_ub = gl.cdiv(K, cfg.BLOCK_K)
        epilogue_lb = loop_ub - (cfg.NUM_BUFFERS - 1)
        gl.assume(loop_ub >= cfg.NUM_BUFFERS)
        for i in range(0, loop_ub):
            c00 = gl.amd.cdna5.wmma_scaled(a00, scale_a00, cfg.DTYPE_A, b00, scale_b00, cfg.DTYPE_B, c00)

            pred_prefetch = i - epilogue_lb
            pred_prefetch = ((pred_prefetch >> 31) & 1).to(gl.int1)

            b01, scale_b01 = self.issue_local_load_b(wmma_idx, 0, 1, self.b_buffer, self.b_scale_buffer)

            self.issue_l2_prefetches_scale(cfg.L2_PREFETCH_DISTANCE, load_idx, pred=pred_prefetch)

            c01 = gl.amd.cdna5.wmma_scaled(a00, scale_a00, cfg.DTYPE_A, b01, scale_b01, cfg.DTYPE_B, c01)
            a10, scale_a10 = self.issue_local_load_a(wmma_idx, 1, 0, self.a_buffer, self.a_scale_buffer)

            self.issue_l2_prefetches_data(cfg.L2_PREFETCH_DISTANCE, load_idx, pred=pred_prefetch)
            load_idx = load_idx + 1

            c10 = gl.amd.cdna5.wmma_scaled(a10, scale_a10, cfg.DTYPE_A, b00, scale_b00, cfg.DTYPE_B, c10)
            b10, scale_b10 = self.issue_local_load_b(wmma_idx, 1, 0, self.b_buffer, self.b_scale_buffer)

            c11 = gl.amd.cdna5.wmma_scaled(a10, scale_a10, cfg.DTYPE_A, b01, scale_b01, cfg.DTYPE_B, c11)
            a01, scale_a01 = self.issue_local_load_a(wmma_idx, 0, 1, self.a_buffer, self.a_scale_buffer)

            c00 = gl.amd.cdna5.wmma_scaled(a01, scale_a01, cfg.DTYPE_A, b10, scale_b10, cfg.DTYPE_B, c00)
            b11, scale_b11 = self.issue_local_load_b(wmma_idx, 1, 1, self.b_buffer, self.b_scale_buffer)

            c01 = gl.amd.cdna5.wmma_scaled(a01, scale_a01, cfg.DTYPE_A, b11, scale_b11, cfg.DTYPE_B, c01)
            a11, scale_a11 = self.issue_local_load_a(wmma_idx, 1, 1, self.a_buffer, self.a_scale_buffer)
            wmma_idx += 1

            c10 = gl.amd.cdna5.wmma_scaled(a11, scale_a11, cfg.DTYPE_A, b10, scale_b10, cfg.DTYPE_B, c10)

            c11 = gl.amd.cdna5.wmma_scaled(a11, scale_a11, cfg.DTYPE_A, b11, scale_b11, cfg.DTYPE_B, c11)

            pred_load = i + 1 - epilogue_lb
            pred_load = (pred_load >> 31) & 1

            a_scale_desc, b_scale_desc = self.issue_load_scales(a_scale_desc, b_scale_desc, load_idx, pred=pred_load)

            a_desc, b_desc = self.issue_load_data(a_desc, b_desc, load_idx, pred=pred_load)

            self.async_wait((cfg.NUM_BUFFERS - 1) * 2)

            a00, scale_a00 = self.issue_local_load_a(wmma_idx, 0, 0, self.a_buffer, self.a_scale_buffer)
            b00, scale_b00 = self.issue_local_load_b(wmma_idx, 0, 0, self.b_buffer, self.b_scale_buffer)

        acc_top = gl.join(c00, c01).permute(0, 2, 1).reshape((SUBTILE_M, cfg.BLOCK_N))
        acc_bot = gl.join(c10, c11).permute(0, 2, 1).reshape((SUBTILE_M, cfg.BLOCK_N))
        accumulator = gl.join(acc_top, acc_bot).permute(2, 0, 1).reshape((cfg.BLOCK_M, cfg.BLOCK_N))
        accumulator = gl.convert_layout(accumulator, cfg.acc_layout)

        apply_activation_and_store(self.cfg, accumulator, self.c_desc, self.c_off_m, self.c_off_n)


@composition
@gluon.aggregate
class MXFPGEMMSliceMNKTDMSplitProgram:
    """Split-TDM variant of the sliceMNK schedule.

    Identical compute pipeline to MXFPGEMMSliceMNKProgram but the A and B
    operand tiles are split into two halves each so that every K-block issues
    four (instead of two) data TDM loads: A is split along M (a0/a1) and B along
    N (b0/b1).  Each half has its own LDS buffer and TDM descriptor, which lets
    the half-loads overlap independently.
    """
    base: MXFPGEMMProgramBase

    cfg: MXFPGEMMConfig
    a0_buffer: gl.shared_memory_descriptor
    a1_buffer: gl.shared_memory_descriptor
    b0_buffer: gl.shared_memory_descriptor
    b1_buffer: gl.shared_memory_descriptor
    a_scale_buffer: gl.shared_memory_descriptor | gl.constexpr
    b_scale_buffer: gl.shared_memory_descriptor

    a0_desc: tdm.tensor_descriptor
    a1_desc: tdm.tensor_descriptor
    b0_desc: tdm.tensor_descriptor
    b1_desc: tdm.tensor_descriptor
    a_scale_desc: tdm.tensor_descriptor | ScaleAsyncCopyDescriptor
    b_scale_desc: tdm.tensor_descriptor | ScaleAsyncCopyDescriptor

    c_desc: tdm.tensor_descriptor
    c_off_m: gl.tensor
    c_off_n: gl.tensor

    @gluon.constexpr_function
    def __init__(self, cfg: MXFPGEMMConfig, a0_buffer, a1_buffer, b0_buffer, b1_buffer, a_scale_buffer, b_scale_buffer,
                 a0_desc, a1_desc, b0_desc, b1_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m, c_off_n):
        self.cfg = cfg
        self.a0_buffer = a0_buffer
        self.a1_buffer = a1_buffer
        self.b0_buffer = b0_buffer
        self.b1_buffer = b1_buffer
        if cfg.WITH_A_SCALE:
            self.a_scale_buffer = a_scale_buffer
        else:
            self.a_scale_buffer = gl.constexpr(a_scale_buffer)

        self.b_scale_buffer = b_scale_buffer
        self.a0_desc = a0_desc
        self.a1_desc = a1_desc
        self.b0_desc = b0_desc
        self.b1_desc = b1_desc
        self.a_scale_desc = a_scale_desc
        self.b_scale_desc = b_scale_desc
        self.c_desc = c_desc
        self.c_off_m = c_off_m
        self.c_off_n = c_off_n

        self.base = MXFPGEMMProgramBase()

    @gluon.jit
    def initialize(cfg: MXFPGEMMConfig, a0_desc, a1_desc, b0_desc, b1_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m,
                   c_off_n):
        NUM_BUFFERS: gl.constexpr = cfg.NUM_BUFFERS
        a0_buffer = gl.allocate_shared_memory(a0_desc.dtype, shape=[NUM_BUFFERS] + a0_desc.block_shape,
                                              layout=a0_desc.layout)
        a1_buffer = gl.allocate_shared_memory(a1_desc.dtype, shape=[NUM_BUFFERS] + a1_desc.block_shape,
                                              layout=a1_desc.layout)
        b0_buffer = gl.allocate_shared_memory(b0_desc.dtype, shape=[NUM_BUFFERS] + b0_desc.block_shape,
                                              layout=b0_desc.layout)
        b1_buffer = gl.allocate_shared_memory(b1_desc.dtype, shape=[NUM_BUFFERS] + b1_desc.block_shape,
                                              layout=b1_desc.layout)
        if cfg.WITH_A_SCALE:
            a_scale_buffer = gl.allocate_shared_memory(a_scale_desc.dtype,
                                                       shape=[NUM_BUFFERS] + a_scale_desc.block_shape,
                                                       layout=a_scale_desc.layout)
        else:
            a_scale_buffer = gl.constexpr(0)

        b_scale_buffer = gl.allocate_shared_memory(b_scale_desc.dtype, shape=[NUM_BUFFERS] + b_scale_desc.block_shape,
                                                   layout=b_scale_desc.layout)

        return MXFPGEMMSliceMNKTDMSplitProgram(cfg, a0_buffer, a1_buffer, b0_buffer, b1_buffer, a_scale_buffer,
                                               b_scale_buffer, a0_desc, a1_desc, b0_desc, b1_desc, a_scale_desc,
                                               b_scale_desc, c_desc, c_off_m, c_off_n)

    @gluon.jit
    def issue_local_load_a(self, wmma_idx, subtile_start_idx_m: gl.constexpr, subtile_start_idx_k: gl.constexpr,
                           a_scale_buffer):
        cfg = self.cfg
        NUM_SUBTILES_M: gl.constexpr = cfg.NUM_SUBTILES[0]
        NUM_SUBTILES_K: gl.constexpr = cfg.NUM_SUBTILES[2]
        SUBTILE_LEN_M: gl.constexpr = cfg.BLOCK_M // NUM_SUBTILES_M
        SUBTILE_LEN_K: gl.constexpr = cfg.BLOCK_K // NUM_SUBTILES_K
        BLOCK_K_SCALE: gl.constexpr = cfg.BLOCK_K // cfg.SCALE_BLOCK
        subtile_start_m: gl.constexpr = subtile_start_idx_m * SUBTILE_LEN_M
        subtile_start_k: gl.constexpr = subtile_start_idx_k * SUBTILE_LEN_K
        if subtile_start_idx_m == 0:
            a_buffer = self.a0_buffer
        else:
            gl.static_assert(subtile_start_idx_m == 1)
            a_buffer = self.a1_buffer
        a = a_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(0, SUBTILE_LEN_M, 0) \
            .slice(subtile_start_k // cfg.DIV_FACTOR_A, SUBTILE_LEN_K // cfg.DIV_FACTOR_A, 1) \
            .load(layout=cfg.dot_layout_a)
        if cfg.WITH_A_SCALE:
            a_scale_buffer_slice = a_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
            if cfg.SCALE_PRESHUFFLE:
                a_scale_buffer_slice = a_scale_buffer_slice.reshape((
                    cfg.BLOCK_M_PRESHUFFLED,  #
                    BLOCK_K_SCALE // cfg.SCALE_KWIDTH,  #
                    cfg.PRESHUFFLE_FACTOR // 4,  #
                    4,  #
                    cfg.SCALE_KWIDTH)).permute((0, 3, 2, 1, 4)).reshape((cfg.BLOCK_M, BLOCK_K_SCALE))
            a_scale_buffer_slice = a_scale_buffer_slice.slice(subtile_start_m, SUBTILE_LEN_M, 0) \
                .slice(subtile_start_k // cfg.SCALE_BLOCK, SUBTILE_LEN_K // cfg.SCALE_BLOCK, 1)
            scale_a = a_scale_buffer_slice.load(layout=cfg.layout_a_scale)
        else:
            scale_a = 0
            scale_a = scale_a.to(gl.uint8)
        return a, scale_a

    @gluon.jit
    def issue_local_load_b(self, wmma_idx, subtile_start_idx_k: gl.constexpr, subtile_start_idx_n: gl.constexpr,
                           b_scale_buffer):
        cfg = self.cfg
        NUM_SUBTILES_N: gl.constexpr = cfg.NUM_SUBTILES[1]
        NUM_SUBTILES_K: gl.constexpr = cfg.NUM_SUBTILES[2]
        SUBTILE_LEN_K: gl.constexpr = cfg.BLOCK_K // NUM_SUBTILES_K
        SUBTILE_LEN_N: gl.constexpr = cfg.BLOCK_N // NUM_SUBTILES_N
        BLOCK_K_SCALE: gl.constexpr = cfg.BLOCK_K // cfg.SCALE_BLOCK
        subtile_start_k: gl.constexpr = subtile_start_idx_k * SUBTILE_LEN_K
        subtile_start_n: gl.constexpr = subtile_start_idx_n * SUBTILE_LEN_N
        if subtile_start_idx_n == 0:
            b_buffer = self.b0_buffer
        else:
            gl.static_assert(subtile_start_idx_n == 1)
            b_buffer = self.b1_buffer
        if cfg.TRANSPOSE_B:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(0, SUBTILE_LEN_N, 0) \
            .slice(subtile_start_k // cfg.DIV_FACTOR_B, SUBTILE_LEN_K // cfg.DIV_FACTOR_B, 1) \
            .permute([1, 0]).load(layout=cfg.dot_layout_b)
        else:
            b = b_buffer.index(wmma_idx % cfg.NUM_BUFFERS).slice(subtile_start_k // cfg.DIV_FACTOR_B,
                                                                 SUBTILE_LEN_K // cfg.DIV_FACTOR_B, 0) \
            .slice(0, SUBTILE_LEN_N, 1) \
            .load(layout=cfg.dot_layout_b)
        b_scale_buffer_slice = b_scale_buffer.index(wmma_idx % cfg.NUM_BUFFERS)
        if cfg.SCALE_PRESHUFFLE:
            b_scale_buffer_slice = b_scale_buffer_slice.reshape((
                cfg.BLOCK_N_PRESHUFFLED,  #
                BLOCK_K_SCALE // cfg.SCALE_KWIDTH,  #
                cfg.PRESHUFFLE_FACTOR // 4,  #
                4,  #
                cfg.SCALE_KWIDTH)).permute((0, 3, 2, 1, 4)).reshape((cfg.BLOCK_N, BLOCK_K_SCALE))
        b_scale_buffer_slice = b_scale_buffer_slice.slice(subtile_start_n, SUBTILE_LEN_N, 0) \
            .slice(subtile_start_k // cfg.SCALE_BLOCK, SUBTILE_LEN_K // cfg.SCALE_BLOCK, 1)
        scale_b = b_scale_buffer_slice.load(layout=cfg.layout_b_scale)
        return b, scale_b

    @gluon.jit
    def issue_load_a_scale(self, a_scale_desc, load_idx, pred=1):
        cfg = self.cfg
        if cfg.WITH_A_SCALE:
            a_scale_buffer_slice = self.a_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
            if cfg.ASYNC_COPY_SCALE:
                a_scale_desc.issue_async_load(load_idx, a_scale_buffer_slice, pred=pred)
            else:
                gl.amd.cdna5.tdm.async_load(a_scale_desc, [0, 0],  #
                                            a_scale_buffer_slice, pred=pred, warp_used_hint=cfg.TDM_WARP_USED_HINT_A)
                a_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_scale_desc,
                                                                         add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
        return a_scale_desc

    @gluon.jit
    def issue_load_b_scale(self, b_scale_desc, load_idx, pred=1):
        cfg = self.cfg
        b_scale_buffer_slice = self.b_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
        if cfg.ASYNC_COPY_SCALE:
            b_scale_desc.issue_async_load(load_idx, b_scale_buffer_slice, pred=pred)
        else:
            gl.amd.cdna5.tdm.async_load(b_scale_desc, [0, 0],  #
                                        b_scale_buffer_slice, pred=pred, warp_used_hint=cfg.TDM_WARP_USED_HINT_B)
            b_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_scale_desc,
                                                                     add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
        return b_scale_desc

    @gluon.jit
    def issue_load_a_data(self, a_desc, load_idx, half: gl.constexpr, pred=1):
        cfg = self.cfg
        BLOCK_K: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A
        if half == 0:
            a_buffer = self.a0_buffer
        else:
            a_buffer = self.a1_buffer
        gl.amd.cdna5.tdm.async_load(a_desc, [0, 0],  #
                                    a_buffer.index(load_idx % cfg.NUM_BUFFERS),  #
                                    pred=pred, warp_used_hint=cfg.TDM_WARP_USED_HINT_A)
        return gl.amd.cdna5.tdm.update_tensor_descriptor(a_desc, add_offsets=[0, BLOCK_K])

    @gluon.jit
    def issue_load_b_data(self, b_desc, load_idx, half: gl.constexpr, pred=1):
        cfg = self.cfg
        BLOCK_K: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_B
        if half == 0:
            b_buffer = self.b0_buffer
        else:
            b_buffer = self.b1_buffer
        gl.amd.cdna5.tdm.async_load(b_desc, [0, 0],  #
                                    b_buffer.index(load_idx % cfg.NUM_BUFFERS),  #
                                    pred=pred, warp_used_hint=cfg.TDM_WARP_USED_HINT_B)
        if cfg.TRANSPOSE_B:
            b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[0, BLOCK_K])
        else:
            b_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_desc, add_offsets=[BLOCK_K, 0])
        return b_desc

    @gluon.jit
    def issue_load_scales(self, a_scale_desc, b_scale_desc, load_idx, pred=1):
        cfg = self.cfg
        if cfg.WITH_A_SCALE:
            if not cfg.ASYNC_COPY_SCALE and cfg.TDM_WARP_USED_HINT_A:
                a_dest = self.a_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
                b_dest = self.b_scale_buffer.index(load_idx % cfg.NUM_BUFFERS)
                a_scale_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_scale_desc, add_offsets=[0, 0],
                                                                              pred=pred, clamp_bounds=True)
                b_scale_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_scale_desc, add_offsets=[0, 0],
                                                                              pred=pred, clamp_bounds=True)
                gl.amd.cdna5.tdm.async_load_fused([(a_scale_load_desc, a_dest, cfg.TDM_WARP_USED_HINT_A),
                                                   (b_scale_load_desc, b_dest, cfg.TDM_WARP_USED_HINT_B)])
                a_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a_scale_desc,
                                                                         add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
                b_scale_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b_scale_desc,
                                                                         add_offsets=[0, cfg.BLOCK_K_SCALE_PRESHUFFLED])
                return a_scale_desc, b_scale_desc
            a_scale_desc = self.issue_load_a_scale(a_scale_desc, load_idx, pred=pred)
        b_scale_desc = self.issue_load_b_scale(b_scale_desc, load_idx, pred=pred)
        return a_scale_desc, b_scale_desc

    @gluon.jit
    def issue_load_data(self, a0_desc, a1_desc, b0_desc, b1_desc, load_idx, pred=1):
        cfg = self.cfg
        if cfg.TDM_WARP_USED_HINT_A:
            # Partial TDM: fuse each (a-half, b-half) pair into one explicit
            # TDM instruction. a0/a1 (resp. b0/b1) share the same hint, so they
            # can't be fused together; keep the a0+b0, then a1+b1 interleave
            # order to match the original TDM queue ordering.
            BLOCK_K_A: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A
            BLOCK_K_B: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_B
            a0_dest = self.a0_buffer.index(load_idx % cfg.NUM_BUFFERS)
            a1_dest = self.a1_buffer.index(load_idx % cfg.NUM_BUFFERS)
            b0_dest = self.b0_buffer.index(load_idx % cfg.NUM_BUFFERS)
            b1_dest = self.b1_buffer.index(load_idx % cfg.NUM_BUFFERS)
            a0_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a0_desc, add_offsets=[0, 0], pred=pred,
                                                                     clamp_bounds=True)
            b0_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b0_desc, add_offsets=[0, 0], pred=pred,
                                                                     clamp_bounds=True)
            a1_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a1_desc, add_offsets=[0, 0], pred=pred,
                                                                     clamp_bounds=True)
            b1_load_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b1_desc, add_offsets=[0, 0], pred=pred,
                                                                     clamp_bounds=True)
            gl.amd.cdna5.tdm.async_load_fused([(a0_load_desc, a0_dest, cfg.TDM_WARP_USED_HINT_A),
                                               (b0_load_desc, b0_dest, cfg.TDM_WARP_USED_HINT_B)])
            gl.amd.cdna5.tdm.async_load_fused([(a1_load_desc, a1_dest, cfg.TDM_WARP_USED_HINT_A),
                                               (b1_load_desc, b1_dest, cfg.TDM_WARP_USED_HINT_B)])
            a0_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a0_desc, add_offsets=[0, BLOCK_K_A])
            a1_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(a1_desc, add_offsets=[0, BLOCK_K_A])
            if cfg.TRANSPOSE_B:
                b0_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b0_desc, add_offsets=[0, BLOCK_K_B])
                b1_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b1_desc, add_offsets=[0, BLOCK_K_B])
            else:
                b0_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b0_desc, add_offsets=[BLOCK_K_B, 0])
                b1_desc = gl.amd.cdna5.tdm.update_tensor_descriptor(b1_desc, add_offsets=[BLOCK_K_B, 0])
        else:
            # Interleave the four half-loads a0, b0, a1, b1 so that A and B halves
            # alternate on the TDM in-order queue.
            a0_desc = self.issue_load_a_data(a0_desc, load_idx, 0, pred=pred)
            b0_desc = self.issue_load_b_data(b0_desc, load_idx, 0, pred=pred)
            a1_desc = self.issue_load_a_data(a1_desc, load_idx, 1, pred=pred)
            b1_desc = self.issue_load_b_data(b1_desc, load_idx, 1, pred=pred)
        return a0_desc, a1_desc, b0_desc, b1_desc

    @gluon.jit
    def issue_l2_prefetch_a(self, distance: gl.constexpr, load_idx, pred=True):
        cfg = self.cfg
        if distance < 0:
            return
        prefetch_iteration = load_idx + distance
        BLOCK_K_PACKED_A: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A
        gl.amd.cdna5.tdm.prefetch(self.a0_desc, [0, prefetch_iteration * BLOCK_K_PACKED_A], pred=pred)
        gl.amd.cdna5.tdm.prefetch(self.a1_desc, [0, prefetch_iteration * BLOCK_K_PACKED_A], pred=pred)

    @gluon.jit
    def issue_l2_prefetch_b(self, distance: gl.constexpr, load_idx, pred=True):
        cfg = self.cfg
        if distance < 0:
            return
        prefetch_iteration = load_idx + distance
        BLOCK_K_PACKED_B: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_B
        if cfg.TRANSPOSE_B:
            gl.amd.cdna5.tdm.prefetch(self.b0_desc, [0, prefetch_iteration * BLOCK_K_PACKED_B], pred=pred)
            gl.amd.cdna5.tdm.prefetch(self.b1_desc, [0, prefetch_iteration * BLOCK_K_PACKED_B], pred=pred)
        else:
            gl.amd.cdna5.tdm.prefetch(self.b0_desc, [prefetch_iteration * BLOCK_K_PACKED_B, 0], pred=pred)
            gl.amd.cdna5.tdm.prefetch(self.b1_desc, [prefetch_iteration * BLOCK_K_PACKED_B, 0], pred=pred)

    @gluon.jit
    def issue_l2_prefetches_data(self, distance: gl.constexpr, load_idx, pred=True):
        self.issue_l2_prefetch_a(distance, load_idx, pred=pred)
        self.issue_l2_prefetch_b(distance, load_idx, pred=pred)

    @gluon.jit
    def async_wait(self, waitcnt: int):
        cfg = self.cfg
        # A fused TDM load counts as a single outstanding op regardless of its
        # member count, unlike separately-hinted loads which count individually.
        DATA_OPS: gl.constexpr = 2 if cfg.TDM_WARP_USED_HINT_A else 4
        SCALE_OPS: gl.constexpr = 1 if cfg.WITH_A_SCALE and cfg.TDM_WARP_USED_HINT_A else (2 if cfg.WITH_A_SCALE else 1)
        if cfg.ASYNC_COPY_SCALE:
            gl.amd.cdna5.tdm.async_wait(waitcnt * DATA_OPS // 2)
            cp.wait_group(waitcnt if cfg.WITH_A_SCALE else (waitcnt // 2))
        else:
            gl.amd.cdna5.tdm.async_wait(waitcnt * (DATA_OPS + SCALE_OPS) // 2)

    @gluon.jit
    def pipeline(self, K):
        cfg = self.cfg
        load_idx = 0
        wmma_idx = 0

        a0_desc = self.a0_desc
        a1_desc = self.a1_desc
        b0_desc = self.b0_desc
        b1_desc = self.b1_desc
        a_scale_desc = self.a_scale_desc
        b_scale_desc = self.b_scale_desc

        for _ in gl.static_range(cfg.NUM_BUFFERS - 1):
            a_scale_desc, b_scale_desc = self.issue_load_scales(a_scale_desc, b_scale_desc, load_idx)
            a0_desc, a1_desc, b0_desc, b1_desc = self.issue_load_data(a0_desc, a1_desc, b0_desc, b1_desc, load_idx)
            load_idx = load_idx + 1

        for d in gl.static_range(cfg.NUM_BUFFERS, cfg.NUM_BUFFERS + cfg.L2_PREFETCH_DISTANCE):
            self.issue_l2_prefetches_data(d, 0)
        for d in gl.static_range(cfg.NUM_BUFFERS, cfg.NUM_BUFFERS + cfg.L2_PREFETCH_DISTANCE):
            self.issue_l2_prefetches_scale(d, 0)

        self.async_wait((cfg.NUM_BUFFERS - 2) * 2)
        a00, scale_a00 = self.issue_local_load_a(wmma_idx, 0, 0, self.a_scale_buffer)
        b00, scale_b00 = self.issue_local_load_b(wmma_idx, 0, 0, self.b_scale_buffer)

        a_scale_desc, b_scale_desc = self.issue_load_scales(a_scale_desc, b_scale_desc, load_idx)
        a0_desc, a1_desc, b0_desc, b1_desc = self.issue_load_data(a0_desc, a1_desc, b0_desc, b1_desc, load_idx)

        SUBTILE_M: gl.constexpr = cfg.BLOCK_M // cfg.NUM_SUBTILES[0]
        SUBTILE_N: gl.constexpr = cfg.BLOCK_N // cfg.NUM_SUBTILES[1]
        c00 = gl.zeros((SUBTILE_M, SUBTILE_N), dtype=gl.float32, layout=cfg.acc_layout)
        c01 = gl.zeros((SUBTILE_M, SUBTILE_N), dtype=gl.float32, layout=cfg.acc_layout)
        c10 = gl.zeros((SUBTILE_M, SUBTILE_N), dtype=gl.float32, layout=cfg.acc_layout)
        c11 = gl.zeros((SUBTILE_M, SUBTILE_N), dtype=gl.float32, layout=cfg.acc_layout)

        loop_ub = gl.cdiv(K, cfg.BLOCK_K)
        epilogue_lb = loop_ub - (cfg.NUM_BUFFERS - 1)
        gl.assume(loop_ub >= cfg.NUM_BUFFERS)
        for i in range(0, loop_ub):
            c00 = gl.amd.cdna5.wmma_scaled(a00, scale_a00, cfg.DTYPE_A, b00, scale_b00, cfg.DTYPE_B, c00)

            pred_prefetch = i - epilogue_lb
            pred_prefetch = ((pred_prefetch >> 31) & 1).to(gl.int1)

            b01, scale_b01 = self.issue_local_load_b(wmma_idx, 0, 1, self.b_scale_buffer)

            self.issue_l2_prefetches_scale(cfg.L2_PREFETCH_DISTANCE, load_idx, pred=pred_prefetch)

            c01 = gl.amd.cdna5.wmma_scaled(a00, scale_a00, cfg.DTYPE_A, b01, scale_b01, cfg.DTYPE_B, c01)
            a10, scale_a10 = self.issue_local_load_a(wmma_idx, 1, 0, self.a_scale_buffer)

            self.issue_l2_prefetches_data(cfg.L2_PREFETCH_DISTANCE, load_idx, pred=pred_prefetch)
            load_idx = load_idx + 1

            c10 = gl.amd.cdna5.wmma_scaled(a10, scale_a10, cfg.DTYPE_A, b00, scale_b00, cfg.DTYPE_B, c10)
            b10, scale_b10 = self.issue_local_load_b(wmma_idx, 1, 0, self.b_scale_buffer)

            c11 = gl.amd.cdna5.wmma_scaled(a10, scale_a10, cfg.DTYPE_A, b01, scale_b01, cfg.DTYPE_B, c11)
            a01, scale_a01 = self.issue_local_load_a(wmma_idx, 0, 1, self.a_scale_buffer)

            c00 = gl.amd.cdna5.wmma_scaled(a01, scale_a01, cfg.DTYPE_A, b10, scale_b10, cfg.DTYPE_B, c00)
            b11, scale_b11 = self.issue_local_load_b(wmma_idx, 1, 1, self.b_scale_buffer)

            c01 = gl.amd.cdna5.wmma_scaled(a01, scale_a01, cfg.DTYPE_A, b11, scale_b11, cfg.DTYPE_B, c01)
            a11, scale_a11 = self.issue_local_load_a(wmma_idx, 1, 1, self.a_scale_buffer)
            wmma_idx += 1

            c10 = gl.amd.cdna5.wmma_scaled(a11, scale_a11, cfg.DTYPE_A, b10, scale_b10, cfg.DTYPE_B, c10)

            c11 = gl.amd.cdna5.wmma_scaled(a11, scale_a11, cfg.DTYPE_A, b11, scale_b11, cfg.DTYPE_B, c11)

            pred_load = i + 1 - epilogue_lb
            pred_load = (pred_load >> 31) & 1

            a_scale_desc, b_scale_desc = self.issue_load_scales(a_scale_desc, b_scale_desc, load_idx, pred=pred_load)

            a0_desc, a1_desc, b0_desc, b1_desc = self.issue_load_data(a0_desc, a1_desc, b0_desc, b1_desc, load_idx,
                                                                      pred=pred_load)

            self.async_wait((cfg.NUM_BUFFERS - 1) * 2)

            a00, scale_a00 = self.issue_local_load_a(wmma_idx, 0, 0, self.a_scale_buffer)
            b00, scale_b00 = self.issue_local_load_b(wmma_idx, 0, 0, self.b_scale_buffer)

        acc_top = gl.join(c00, c01).permute(0, 2, 1).reshape((SUBTILE_M, cfg.BLOCK_N))
        acc_bot = gl.join(c10, c11).permute(0, 2, 1).reshape((SUBTILE_M, cfg.BLOCK_N))
        accumulator = gl.join(acc_top, acc_bot).permute(2, 0, 1).reshape((cfg.BLOCK_M, cfg.BLOCK_N))
        accumulator = gl.convert_layout(accumulator, cfg.acc_layout)

        apply_activation_and_store(self.cfg, accumulator, self.c_desc, self.c_off_m, self.c_off_n)


@gluon.jit
def create_tensor_descriptor(cfg: MXFPGEMMConfig, a_ptr, a_offs, b_ptr, b_offs, a_scale_ptr, a_scale_offs, b_scale_ptr,
                             b_scale_offs, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_scale):
    SCALE_BLOCK: gl.constexpr = cfg.SCALE_BLOCK
    PRESHUFFLE_FACTOR: gl.constexpr = cfg.PRESHUFFLE_FACTOR
    if cfg.TDM_SPLIT:
        a_desc = gl.constexpr(0)
        b_desc = gl.constexpr(0)
    else:
        a_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr + a_offs,  #
                                                         shape=(M, K // cfg.DIV_FACTOR_A),  #
                                                         strides=(stride_am, stride_ak),  #
                                                         block_shape=(cfg.BLOCK_M, cfg.BLOCK_K // cfg.DIV_FACTOR_A),  #
                                                         layout=cfg.shared_layout_a)

        if cfg.TRANSPOSE_B:
            b_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr + b_offs,  #
                                                             shape=(N, K // cfg.DIV_FACTOR_B),  #
                                                             strides=(stride_bn, stride_bk),  #
                                                             block_shape=(cfg.BLOCK_N,
                                                                          cfg.BLOCK_K // cfg.DIV_FACTOR_B),  #
                                                             layout=cfg.shared_layout_b)
        else:
            b_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr + b_offs,  #
                                                             shape=(K // cfg.DIV_FACTOR_B, N),  #
                                                             strides=(stride_bk, stride_bn),  #
                                                             block_shape=(cfg.BLOCK_K // cfg.DIV_FACTOR_B,
                                                                          cfg.BLOCK_N),  #
                                                             layout=cfg.shared_layout_b)

    if cfg.ASYNC_COPY_SCALE:
        if cfg.WITH_A_SCALE:
            a_scale_desc = ScaleAsyncCopyDescriptor.initialize(cfg, 0, a_scale_ptr, a_scale_offs, stride_scale,
                                                               cfg.shared_layout_a_scale)
        else:
            a_scale_desc = ScaleAsyncCopyDescriptor.initialize(cfg, 0, b_scale_ptr, b_scale_offs, stride_scale,
                                                               cfg.shared_layout_a_scale)
        b_scale_desc = ScaleAsyncCopyDescriptor.initialize(cfg, 1, b_scale_ptr, b_scale_offs, stride_scale,
                                                           cfg.shared_layout_b_scale)
    else:
        if cfg.WITH_A_SCALE:
            a_scale_base = a_scale_ptr + a_scale_offs
        else:
            a_scale_base = b_scale_ptr + b_scale_offs
        a_scale_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(
            base=a_scale_base,  #
            shape=(M // PRESHUFFLE_FACTOR, K // SCALE_BLOCK * PRESHUFFLE_FACTOR),  #
            strides=(stride_scale, 1),  #
            block_shape=(cfg.BLOCK_M_PRESHUFFLED, cfg.BLOCK_K_SCALE_PRESHUFFLED),  #
            layout=cfg.shared_layout_a_scale)

        b_scale_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(
            base=b_scale_ptr + b_scale_offs,  #
            shape=(N // PRESHUFFLE_FACTOR, K // SCALE_BLOCK * PRESHUFFLE_FACTOR),  #
            strides=(stride_scale, 1),  #
            block_shape=(cfg.BLOCK_N_PRESHUFFLED, cfg.BLOCK_K_SCALE_PRESHUFFLED),  #
            layout=cfg.shared_layout_b_scale)

    return a_desc, b_desc, a_scale_desc, b_scale_desc


@gluon.jit
def create_split_data_tensor_descriptor(cfg: MXFPGEMMConfig, a_ptr, a_offs, b_ptr, b_offs, M, N, K, stride_am,
                                        stride_ak, stride_bk, stride_bn):
    HALF_M: gl.constexpr = cfg.BLOCK_M // 2
    HALF_N: gl.constexpr = cfg.BLOCK_N // 2
    BLOCK_K_PACKED_A: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A
    BLOCK_K_PACKED_B: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_B

    a0_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr + a_offs,  #
                                                      shape=(M, K // cfg.DIV_FACTOR_A),  #
                                                      strides=(stride_am, stride_ak),  #
                                                      block_shape=(HALF_M, BLOCK_K_PACKED_A),  #
                                                      layout=cfg.shared_layout_a)
    a1_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr + a_offs + HALF_M * stride_am,  #
                                                      shape=(M, K // cfg.DIV_FACTOR_A),  #
                                                      strides=(stride_am, stride_ak),  #
                                                      block_shape=(HALF_M, BLOCK_K_PACKED_A),  #
                                                      layout=cfg.shared_layout_a)

    if cfg.TRANSPOSE_B:
        b0_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr + b_offs,  #
                                                          shape=(N, K // cfg.DIV_FACTOR_B),  #
                                                          strides=(stride_bn, stride_bk),  #
                                                          block_shape=(HALF_N, BLOCK_K_PACKED_B),  #
                                                          layout=cfg.shared_layout_b)
        b1_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr + b_offs + HALF_N * stride_bn,  #
                                                          shape=(N, K // cfg.DIV_FACTOR_B),  #
                                                          strides=(stride_bn, stride_bk),  #
                                                          block_shape=(HALF_N, BLOCK_K_PACKED_B),  #
                                                          layout=cfg.shared_layout_b)
    else:
        b0_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr + b_offs,  #
                                                          shape=(K // cfg.DIV_FACTOR_B, N),  #
                                                          strides=(stride_bk, stride_bn),  #
                                                          block_shape=(BLOCK_K_PACKED_B, HALF_N),  #
                                                          layout=cfg.shared_layout_b)
        b1_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr + b_offs + HALF_N * stride_bn,  #
                                                          shape=(K // cfg.DIV_FACTOR_B, N),  #
                                                          strides=(stride_bk, stride_bn),  #
                                                          block_shape=(BLOCK_K_PACKED_B, HALF_N),  #
                                                          layout=cfg.shared_layout_b)

    return a0_desc, a1_desc, b0_desc, b1_desc


@gluon.jit
def mxgemm_tdm_pipelined_kernel(a_ptr, b_ptr, c_ptr, a_scale, b_scale, M, N, K, stride_am, stride_ak, stride_bk,
                                stride_bn, stride_cm, stride_cn, stride_scale, DTYPE_A: gl.constexpr,
                                DTYPE_B: gl.constexpr, SCALE_BLOCK: gl.constexpr, BLOCK_M: gl.constexpr,
                                BLOCK_N: gl.constexpr, BLOCK_K: gl.constexpr, GROUP_SIZE_M: gl.constexpr,
                                TRANSPOSE_B: gl.constexpr, NUM_BUFFERS: gl.constexpr, SCALE_PRESHUFFLE: gl.constexpr,
                                ASYNC_COPY_SCALE: gl.constexpr, WITH_A_SCALE: gl.constexpr, SCHEDULE: gl.constexpr,
                                NUM_WARPS: gl.constexpr, PINGPONG: gl.constexpr, L2_PREFETCH_DISTANCE: gl.constexpr = 0,
                                ACTIVATION: gl.constexpr = "", PARTIAL_TDM: gl.constexpr = False,
                                RESOLVE_PARTITION_CONFLICTS: gl.constexpr = False, TDM_SPLIT: gl.constexpr = False):

    if PINGPONG:
        gl.static_assert(NUM_WARPS == 8 and (SCHEDULE == 'baseline' or SCHEDULE == 'sliceK'))

    if SCHEDULE == 'sliceMNK':
        NUM_SUBTILES: gl.constexpr = (2, 2, 2)
    elif SCHEDULE == 'sliceNK':
        NUM_SUBTILES: gl.constexpr = (1, 2, 2)
    elif SCHEDULE == 'sliceK':
        NUM_SUBTILES: gl.constexpr = (1, 1, 2)
    else:
        gl.static_assert(SCHEDULE == 'baseline')
        NUM_SUBTILES: gl.constexpr = (1, 1, 1)

    # For SwiGLU the compute pipeline runs on the packed (2x) N, only the store
    # path uses the reduced output width.
    ACTIVATION_REDUCTION_N: gl.constexpr = 2 if ACTIVATION == "swiglu" else 1
    BLOCK_N_PACKED: gl.constexpr = BLOCK_N * ACTIVATION_REDUCTION_N
    N_PACKED = N * ACTIVATION_REDUCTION_N

    # The TDM warp-used hint determines which warps participate in TDM loads.
    # For partial TDM, the hint indicates that only half of the warps are used
    # for the TDM loads of A/B and AScale/BScale. TDM_WARP_USED_HINT applies to
    # A and AScale, while reverse_tdm_warp_used_hint provides the hint for B and BScale.
    if PARTIAL_TDM:
        if NUM_WARPS == 8:
            TDM_WARP_USED_HINT: gl.constexpr = 0b01010101
        else:
            gl.static_assert(NUM_WARPS == 4)
            TDM_WARP_USED_HINT: gl.constexpr = 0b00000101
    else:
        TDM_WARP_USED_HINT: gl.constexpr = None

    cfg = MXFPGEMMConfig(BLOCK_M, BLOCK_N_PACKED, BLOCK_K, DTYPE_A, DTYPE_B, SCALE_BLOCK, NUM_BUFFERS, TRANSPOSE_B,
                         WITH_A_SCALE, SCALE_PRESHUFFLE, NUM_WARPS, TDM_WARP_USED_HINT, ASYNC_COPY_SCALE, NUM_SUBTILES,
                         L2_PREFETCH_DISTANCE, ACTIVATION, RESOLVE_PARTITION_CONFLICTS, TDM_SPLIT)

    pid = gl.program_id(axis=0)
    num_pid_m = gl.cdiv(M, BLOCK_M)
    num_pid_n = gl.cdiv(N, BLOCK_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    a_offs = pid_m * BLOCK_M * stride_am
    b_offs = pid_n * BLOCK_N_PACKED * stride_bn
    a_scale_offs = pid_m * cfg.BLOCK_M_PRESHUFFLED * stride_scale
    b_scale_offs = pid_n * cfg.BLOCK_N_PRESHUFFLED * stride_scale
    a_desc, b_desc, a_scale_desc, b_scale_desc = create_tensor_descriptor(cfg, a_ptr, a_offs, b_ptr, b_offs, a_scale,
                                                                          a_scale_offs, b_scale, b_scale_offs, M,
                                                                          N_PACKED, K, stride_am, stride_ak, stride_bk,
                                                                          stride_bn, stride_scale)

    shared_layout_acc: gl.constexpr = gl.SwizzledSharedLayout(1, 1, 1, [1, 0])
    c_desc = gl.amd.cdna5.tdm.make_tensor_descriptor(base=c_ptr, shape=(M, N), strides=(stride_cm, stride_cn),
                                                     block_shape=(BLOCK_M, BLOCK_N), layout=shared_layout_acc)
    c_off_m = pid_m * BLOCK_M
    c_off_n = pid_n * BLOCK_N

    if SCHEDULE == 'sliceMNK':
        if TDM_SPLIT:
            a0_desc, a1_desc, b0_desc, b1_desc = create_split_data_tensor_descriptor(
                cfg, a_ptr, a_offs, b_ptr, b_offs, M, N_PACKED, K, stride_am, stride_ak, stride_bk, stride_bn)
            pgm = MXFPGEMMSliceMNKTDMSplitProgram.initialize(cfg, a0_desc, a1_desc, b0_desc, b1_desc, a_scale_desc,
                                                             b_scale_desc, c_desc, c_off_m, c_off_n)
        else:
            pgm = MXFPGEMMSliceMNKProgram.initialize(cfg, a_desc, b_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m,
                                                     c_off_n)
    elif SCHEDULE == 'sliceNK':
        pgm = MXFPGEMMSliceNKProgram.initialize(cfg, a_desc, b_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m,
                                                c_off_n)
    elif SCHEDULE == 'sliceK':
        pgm = MXFPGEMMSliceKProgram.initialize(cfg, a_desc, b_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m,
                                               c_off_n)
    else:
        gl.static_assert(SCHEDULE == 'baseline')
        pgm = MXFPGEMMPipelinedProgram.initialize(cfg, a_desc, b_desc, a_scale_desc, b_scale_desc, c_desc, c_off_m,
                                                  c_off_n)

    if PINGPONG:
        pgm.warp_pipeline(K)
    else:
        pgm.pipeline(K)


def torch_gemm_mxfp(a, b, a_scale, b_scale, scale_block, M, N, K):
    if a_scale is None:
        a_scale_f32 = torch.full((M, K), 1.0, dtype=torch.float32)
    else:
        a_scale_f32 = a_scale.to(torch.float32).repeat_interleave(scale_block, dim=1)[:M, :K]
    b_scale_f32 = b_scale.to(torch.float32).repeat_interleave(scale_block, dim=1).T.contiguous()[:K, :N]

    a_f32 = a.to(torch.float32)
    b_f32 = b.to(torch.float32)

    return torch.matmul(a_f32 * a_scale_f32, b_f32 * b_scale_f32).to(torch.float32)


def init_data(dtype, d0: int, d1: int):
    if dtype == 'float4':
        return MXFP4Tensor(size=(d0, d1)).random()
    elif dtype == "float8_e5m2":
        return torch.randint(20, 40, (d0, d1), dtype=torch.uint8).view(torch.float8_e5m2)
    elif dtype == "float8_e4m3":
        return torch.randint(20, 40, (d0, d1), dtype=torch.uint8).view(torch.float8_e4m3fn)
    else:
        raise NotImplementedError(f"NYI: unsupported dtype: {dtype}")


def pack_scale(x):
    if x is None:
        return x
    NON_K, K_SCALE = x.shape
    preshuffle_factor = 128
    num_chunk_m = NON_K // preshuffle_factor
    SCALE_KWIDTH = 4 if K_SCALE >= 4 else K_SCALE
    num_chunk_k = K_SCALE // SCALE_KWIDTH

    x = x.view(num_chunk_m, 4, preshuffle_factor // 4, num_chunk_k, SCALE_KWIDTH)
    x = x.permute(0, 3, 2, 1, 4).contiguous()
    return x.view(NON_K // preshuffle_factor, K_SCALE * preshuffle_factor)


def interleave_b_columns(w_gate, w_up, DTYPE_B):
    """Column-interleave two (K, N) B weights into a (K, 2*N) [gate, up, ...] tensor."""
    if DTYPE_B == 'float4':
        K, N = w_gate.data.shape
        interleaved = torch.stack([w_gate.data, w_up.data], dim=-1).reshape(K, 2 * N).contiguous()
        combined = MXFP4Tensor(size=(K, 2 * N))
        combined.data = interleaved
        return combined
    K, N = w_gate.shape
    u_gate = w_gate.view(torch.uint8)
    u_up = w_up.view(torch.uint8)
    interleaved = torch.stack([u_gate, u_up], dim=-1).reshape(K, 2 * N).contiguous()
    return interleaved.view(w_gate.dtype)


def interleave_b_scale_rows(s_gate, s_up):
    """Interleave two (N, K_scale) scale tensors along N into (2*N, K_scale)."""
    data_gate = s_gate.data
    data_up = s_up.data
    N, K_scale = data_gate.shape
    interleaved = torch.stack([data_gate, data_up], dim=1).reshape(2 * N, K_scale).contiguous()
    combined = MXScaleTensor(size=(2 * N, K_scale))
    combined.data = interleaved
    return combined


@pytest.mark.parametrize(
    "DTYPE_A, DTYPE_B",
    [['float8_e5m2', 'float4'], ['float4', 'float8_e4m3'], ['float8_e4m3', 'float8_e5m2'], ['float4', 'float4']])
@pytest.mark.parametrize("M,N,K", [(512, 512, 512)])
@pytest.mark.parametrize("SCHEDULE,BLOCK_M,BLOCK_N,BLOCK_K", [('baseline', 256, 256, 128), ('sliceK', 128, 256, 256)])
@pytest.mark.parametrize("TRANSPOSE_B", [True])
@pytest.mark.parametrize("SCALE_PRESHUFFLE", [True])
@pytest.mark.parametrize("WITH_A_SCALE", [True, False])
@pytest.mark.parametrize("ASYNC_COPY_SCALE", [False])
@pytest.mark.parametrize("NUM_BUFFERS", [2, 3])
@pytest.mark.parametrize("GROUP_SIZE_M", [8])
@pytest.mark.parametrize("PINGPONG", [True, False])
@pytest.mark.parametrize("L2_PREFETCH_DISTANCE", [-1, 0, 2])
def test_runtime_mxgemm_tdm_8warps_pipeline(DTYPE_A, DTYPE_B, M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, TRANSPOSE_B,
                                            NUM_BUFFERS, SCALE_PRESHUFFLE, WITH_A_SCALE, SCHEDULE, ASYNC_COPY_SCALE,
                                            GROUP_SIZE_M, PINGPONG, L2_PREFETCH_DISTANCE,
                                            RESOLVE_PARTITION_CONFLICTS=False, PARTIAL_TDM=False, BENCHMARK_MODE=None,
                                            BENCHMARK_NUM_ITERS=32):
    SCALE_BLOCK = 32
    numWarps = 8
    numCtas = 1

    torch.manual_seed(0)

    def is_fp8(dtype):
        return dtype in ['float8_e5m2', 'float8_e4m3']

    if SCHEDULE == 'sliceK' and NUM_BUFFERS == 3:
        problem_size = BLOCK_M * BLOCK_N * BLOCK_K
        if is_fp8(DTYPE_A) and is_fp8(DTYPE_B):
            if problem_size >= 128 * 256 * 256:
                pytest.skip(
                    'Large block size will exceed lds limit with fp8 inputs and 3 buffers. Please use 128x128x256')

    if SCHEDULE == 'sliceNK' and (PINGPONG or NUM_BUFFERS == 3):
        pytest.skip('NYI: Skipping pingpong or 3 buffers in sliceNK schedule')

    if PINGPONG and NUM_BUFFERS != 3:
        pytest.skip('Pingpong only supports 3 buffers')

    a = init_data(DTYPE_A, M, K)
    b = init_data(DTYPE_B, K, N)
    a_scale_size = (M, (K + SCALE_BLOCK - 1) // SCALE_BLOCK)
    b_scale_size = (N, (K + SCALE_BLOCK - 1) // SCALE_BLOCK)
    if WITH_A_SCALE:
        a_scale = MXScaleTensor(size=a_scale_size).random(low=1.0, high=32.0)
    else:
        a_scale = None
    b_scale = MXScaleTensor(size=b_scale_size).random(low=1.0, high=32.0)

    c_ref = torch_gemm_mxfp(a, b, a_scale, b_scale, SCALE_BLOCK, M, N, K)

    if WITH_A_SCALE:
        a_scale = a_scale.data
    b_scale = b_scale.data

    if SCALE_PRESHUFFLE:
        a_scale = pack_scale(a_scale)
        b_scale = pack_scale(b_scale)

    # mxfp4 input needs packed along the k dim, i.e., two mxfp4 are packed in one uint8
    if DTYPE_A in ['float4', 'float6_e2m3', 'float6_e3m2']:
        a = a.to_packed_tensor(dim=1)
    if DTYPE_B in ['float4', 'float6_e2m3', 'float6_e3m2']:
        b = b.to_packed_tensor(dim=0)

    c_d = torch.zeros(M, N, dtype=torch.float32).cuda()
    a_d = a.data.contiguous().cuda()
    if TRANSPOSE_B:
        b_d = b.data.T.contiguous().cuda()
    else:
        b_d = b.data.contiguous().cuda()
    if WITH_A_SCALE:
        a_scale_d = a_scale.cuda()
    else:
        a_scale_d = None
    b_scale_d = b_scale.cuda()

    stride_am, stride_ak = a_d.stride(0), a_d.stride(1)
    if TRANSPOSE_B:
        stride_bk, stride_bn = b_d.stride(1), b_d.stride(0)
    else:
        stride_bk, stride_bn = b_d.stride(0), b_d.stride(1)
    stride_cm, stride_cn = c_d.stride(0), c_d.stride(1)
    stride_scale = b_scale_d.stride(0)

    numBlocks = triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N)
    grid = [numBlocks, 1, 1]

    dtype_converter = {'float8_e5m2': "e5m2", "float8_e4m3": "e4m3", "float4": "e2m1"}

    fn = lambda: mxgemm_tdm_pipelined_kernel[grid](
        a_d, b_d, c_d, a_scale_d, b_scale_d, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
        stride_scale, dtype_converter[DTYPE_A], dtype_converter[DTYPE_B], SCALE_BLOCK, BLOCK_M, BLOCK_N, BLOCK_K,
        GROUP_SIZE_M, TRANSPOSE_B, NUM_BUFFERS, SCALE_PRESHUFFLE, ASYNC_COPY_SCALE, WITH_A_SCALE, SCHEDULE, numWarps,
        PINGPONG, L2_PREFETCH_DISTANCE=L2_PREFETCH_DISTANCE, PARTIAL_TDM=PARTIAL_TDM, RESOLVE_PARTITION_CONFLICTS=
        RESOLVE_PARTITION_CONFLICTS, num_warps=numWarps, num_ctas=numCtas, waves_per_eu=(numWarps // 4))

    if BENCHMARK_MODE == 'graph':
        time = triton.testing.do_bench_cudagraph(fn, rep=BENCHMARK_NUM_ITERS)
        tflops = 2 * M * N * K / (time * 1e-3) / 1e12
        print(f'execution time: {time} ms, {tflops:.2f} TFLOPS')
    elif BENCHMARK_MODE == 'eager':
        time = triton.testing.do_bench(fn, warmup=10, rep=BENCHMARK_NUM_ITERS)
        tflops = 2 * M * N * K / (time * 1e-3) / 1e12
        print(f'execution time: {time} ms, {tflops:.2f} TFLOPS')
    else:
        k = fn()
        static_profile(k)

        if TRANSPOSE_B:
            assert 'ds_load_u8' not in k.asm['amdgcn']

        if L2_PREFETCH_DISTANCE >= 0:
            assert 'global_prefetch_b8' in k.asm['amdgcn']
        else:
            assert 'global_prefetch_b8' not in k.asm['amdgcn']

        torch.testing.assert_close(c_d.cpu(), c_ref.cpu(), rtol=1e-5, atol=1e-8)
        print('✅Pass')


@pytest.mark.parametrize(
    "DTYPE_A, DTYPE_B",
    [['float8_e5m2', 'float4'], ['float4', 'float8_e4m3'], ['float8_e4m3', 'float8_e5m2'], ['float4', 'float4']])
@pytest.mark.parametrize("M,N,K", [(128, 128, 128), (256, 256, 512), (256, 256, 1024)])
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K,NUM_BUFFERS,ACTIVATION", [
    (64, 64, 64, 2, 'swiglu'),
    (64, 64, 64, 4, ''),
    (128, 128, 128, 2, 'swiglu'),
    (128, 128, 128, 4, ''),
    (256, 256, 256, 2, 'swiglu'),
    (128, 256, 256, 3, ''),
    (256, 256, 128, 4, ''),
])
@pytest.mark.parametrize("TRANSPOSE_B", [True, False])
@pytest.mark.parametrize("SCALE_PRESHUFFLE", [True, False])
@pytest.mark.parametrize("WITH_A_SCALE", [True, False])
@pytest.mark.parametrize("SCHEDULE", ['sliceNK', 'sliceK', 'baseline', 'sliceMNK'])
@pytest.mark.parametrize("ASYNC_COPY_SCALE", [True, False])
@pytest.mark.parametrize("GROUP_SIZE_M", [8])
@pytest.mark.parametrize("L2_PREFETCH_DISTANCE", [-1, 0, 2])
@pytest.mark.parametrize("PARTIAL_TDM", [True, False])
@pytest.mark.parametrize("RESOLVE_PARTITION_CONFLICTS", [True, False])
@pytest.mark.parametrize("TDM_SPLIT", [True, False])
def test_runtime_mxgemm_tdm_pipelined(DTYPE_A, DTYPE_B, M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, TRANSPOSE_B, NUM_BUFFERS,
                                      SCALE_PRESHUFFLE, WITH_A_SCALE, SCHEDULE, ASYNC_COPY_SCALE, GROUP_SIZE_M,
                                      L2_PREFETCH_DISTANCE, ACTIVATION, PARTIAL_TDM, RESOLVE_PARTITION_CONFLICTS,
                                      TDM_SPLIT, BENCHMARK_MODE=None, BENCHMARK_NUM_ITERS=32):
    """
    Pipelined mxfp GEMM with optional fused SwiGLU epilogue.

    When ACTIVATION == "swiglu", B is (K, 2N) with interleaved [gate, up] columns and
    the output is (M, N) = swish(gate) * up.
    """
    SCALE_BLOCK = 32
    numWarps = 4
    numCtas = 1
    IS_SWIGLU = ACTIVATION == "swiglu"
    EFFECTIVE_BLOCK_N = BLOCK_N * 2 if IS_SWIGLU else BLOCK_N
    is_fp4fp4 = DTYPE_A == "float4" and DTYPE_B == "float4"

    if RESOLVE_PARTITION_CONFLICTS and SCHEDULE != "sliceMNK":
        pytest.skip("Only test RESOLVE_PARTITION_CONFLICTS in sliceMNK schedule")

    if RESOLVE_PARTITION_CONFLICTS and is_fp4fp4:
        pytest.skip("NYI: RESOLVE_PARTITION_CONFLICTS not supported with fp4xfp4")

    if SCHEDULE != "sliceMNK" and TDM_SPLIT:
        pytest.skip("Only test TDM_SPLIT in sliceMNK schedule")

    if PARTIAL_TDM and RESOLVE_PARTITION_CONFLICTS and not TDM_SPLIT:
        pytest.skip("PARTIAL_TDM with RESOLVE_PARTITION_CONFLICTS requires TDM_SPLIT")

    if SCALE_PRESHUFFLE:
        if BLOCK_M < 128 or EFFECTIVE_BLOCK_N < 128 or (SCHEDULE != "baseline" and BLOCK_K < 256):
            pytest.skip("Skipping block sizes too small for preshuffling")

    if not WITH_A_SCALE and DTYPE_A == "float4":
        pytest.skip("Skip fp4 x mxfp gemm to reduce test cases.")

    if SCHEDULE != 'baseline' and not (SCALE_PRESHUFFLE and TRANSPOSE_B):
        pytest.skip('Only test with SCALE_PRESHUFFLE and TRANSPOSE_B in sliceK and sliceNK schedules')

    is_fp8_a = DTYPE_A in ('float8_e5m2', 'float8_e4m3')
    is_fp8_b = DTYPE_B in ('float8_e5m2', 'float8_e4m3')

    if NUM_BUFFERS >= 3 and (BLOCK_M >= 256 or EFFECTIVE_BLOCK_N >= 256 or BLOCK_K >= 256) and is_fp8_a and is_fp8_b:
        pytest.skip("Skip large block size with >=3 buffers to not exceed lds limit")

    if SCHEDULE == 'sliceNK':
        if BLOCK_K < 256 or EFFECTIVE_BLOCK_N < 256:
            pytest.skip('BLOCK_K and BLOCK_N are too small for sliceNK schedule')
        if M < 256:
            pytest.skip('Skip small problem size to reduce test cases')
    else:
        if ASYNC_COPY_SCALE:
            pytest.skip('Only use ASYNC_COPY_SCALE in sliceNK schedule')

    if SCHEDULE == 'sliceK' and (K + BLOCK_K - 1) // BLOCK_K < NUM_BUFFERS - 1:
        pytest.skip('sliceK pipeline requires K_iters >= NUM_BUFFERS - 1')

    if SCHEDULE != 'baseline' and NUM_BUFFERS >= 3:
        problem_size = BLOCK_M * EFFECTIVE_BLOCK_N * BLOCK_K
        if is_fp8_a and is_fp8_b and problem_size >= 128 * 256 * 256:
            pytest.skip('Large block size with fp8 inputs and >=3 buffers will exceed lds limit')

    if ASYNC_COPY_SCALE and (M < BLOCK_M or N < BLOCK_N or K < BLOCK_K):
        pytest.skip('NYI: Skipping small problem sizes for async copy scale')

    if IS_SWIGLU and BLOCK_M >= 256 and BLOCK_N >= 256:
        pytest.skip("SwiGLU uses 2N-packed B, large block size will exceed lds limit")

    torch.manual_seed(0)

    a = init_data(DTYPE_A, M, K)
    a_scale_size = (M, (K + SCALE_BLOCK - 1) // SCALE_BLOCK)
    b_scale_size = (N, (K + SCALE_BLOCK - 1) // SCALE_BLOCK)
    if WITH_A_SCALE:
        a_scale = MXScaleTensor(size=a_scale_size).random(low=1.0, high=32.0)
    else:
        a_scale = None

    if IS_SWIGLU:
        w_gate = init_data(DTYPE_B, K, N)
        w_up = init_data(DTYPE_B, K, N)
        b_scale_gate = MXScaleTensor(size=b_scale_size).random(low=1.0, high=32.0)
        b_scale_up = MXScaleTensor(size=b_scale_size).random(low=1.0, high=32.0)

        gate_ref = torch_gemm_mxfp(a, w_gate, a_scale, b_scale_gate, SCALE_BLOCK, M, N, K)
        up_ref = torch_gemm_mxfp(a, w_up, a_scale, b_scale_up, SCALE_BLOCK, M, N, K)
        c_ref = gate_ref * torch.sigmoid(gate_ref) * up_ref

        b = interleave_b_columns(w_gate, w_up, DTYPE_B)
        b_scale = interleave_b_scale_rows(b_scale_gate, b_scale_up)
    else:
        b = init_data(DTYPE_B, K, N)
        b_scale = MXScaleTensor(size=b_scale_size).random(low=1.0, high=32.0)
        c_ref = torch_gemm_mxfp(a, b, a_scale, b_scale, SCALE_BLOCK, M, N, K)

    if WITH_A_SCALE:
        a_scale = a_scale.data
    b_scale = b_scale.data

    if SCALE_PRESHUFFLE:
        a_scale = pack_scale(a_scale)
        b_scale = pack_scale(b_scale)

    # mxfp4 input needs packed along the k dim, i.e., two mxfp4 are packed in one uint8
    if DTYPE_A in ['float4', 'float6_e2m3', 'float6_e3m2']:
        a = a.to_packed_tensor(dim=1)
    if DTYPE_B in ['float4', 'float6_e2m3', 'float6_e3m2']:
        b = b.to_packed_tensor(dim=0)

    c_d = torch.zeros(M, N, dtype=torch.float32).cuda()
    a_d = a.data.contiguous().cuda()
    if TRANSPOSE_B:
        b_d = b.data.T.contiguous().cuda()
    else:
        b_d = b.data.contiguous().cuda()
    if WITH_A_SCALE:
        a_scale_d = a_scale.cuda()
    else:
        a_scale_d = None
    b_scale_d = b_scale.cuda()

    stride_am, stride_ak = a_d.stride(0), a_d.stride(1)
    if TRANSPOSE_B:
        stride_bk, stride_bn = b_d.stride(1), b_d.stride(0)
    else:
        stride_bk, stride_bn = b_d.stride(0), b_d.stride(1)
    stride_cm, stride_cn = c_d.stride(0), c_d.stride(1)
    stride_scale = b_scale_d.stride(0)

    numBlocks = triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N)
    grid = [numBlocks, 1, 1]

    dtype_converter = {'float8_e5m2': "e5m2", "float8_e4m3": "e4m3", "float4": "e2m1"}

    fn = lambda: mxgemm_tdm_pipelined_kernel[
        grid](a_d, b_d, c_d, a_scale_d, b_scale_d, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm,
              stride_cn, stride_scale, dtype_converter[DTYPE_A], dtype_converter[DTYPE_B], SCALE_BLOCK, BLOCK_M,
              BLOCK_N, BLOCK_K, GROUP_SIZE_M, TRANSPOSE_B, NUM_BUFFERS, SCALE_PRESHUFFLE, ASYNC_COPY_SCALE,
              WITH_A_SCALE, SCHEDULE, NUM_WARPS=numWarps, PINGPONG=False, L2_PREFETCH_DISTANCE=L2_PREFETCH_DISTANCE,
              ACTIVATION=ACTIVATION, PARTIAL_TDM=PARTIAL_TDM, RESOLVE_PARTITION_CONFLICTS=RESOLVE_PARTITION_CONFLICTS,
              TDM_SPLIT=TDM_SPLIT, num_warps=numWarps, num_ctas=numCtas, waves_per_eu=numWarps // 4)

    if BENCHMARK_MODE == 'graph':
        time = triton.testing.do_bench_cudagraph(fn, rep=BENCHMARK_NUM_ITERS)
        tflops = 2 * M * N * K / (time * 1e-3) / 1e12
        print(f'execution time: {time} ms, {tflops:.2f} TFLOPS')
    elif BENCHMARK_MODE == 'eager':
        time = triton.testing.do_bench(fn, warmup=30, rep=BENCHMARK_NUM_ITERS)
        tflops = 2 * M * N * K / (time * 1e-3) / 1e12
        print(f'execution time: {time} ms, {tflops:.2f} TFLOPS')
    else:
        k = fn()
        static_profile(k)

        if TRANSPOSE_B:
            assert 'ds_load_u8' not in k.asm['amdgcn']

        if L2_PREFETCH_DISTANCE >= 0:
            assert 'global_prefetch_b8' in k.asm['amdgcn']
        else:
            assert 'global_prefetch_b8' not in k.asm['amdgcn']

        is_fp4fp4 = DTYPE_A == "float4" and DTYPE_B == "float4"
        if is_fp4fp4:
            assert 'v_wmma_scale_f32_32x16x128_f4' in k.asm['amdgcn']
        else:
            assert 'v_wmma_scale_f32_16x16x128_f8f6f4' in k.asm['amdgcn']

        # Relaxed tolerance for SwiGLU because of sigmoid/exp in the epilogue.
        if IS_SWIGLU:
            torch.testing.assert_close(c_d.cpu(), c_ref.cpu(), rtol=1e-2, atol=1e-2)
        else:
            torch.testing.assert_close(c_d.cpu(), c_ref.cpu(), rtol=1e-5, atol=1e-8)
        print('✅Pass')


if __name__ == '__main__':
    import argparse

    supported_dtypes = ['float8_e4m3', 'float8_e5m2', 'float4']

    parser = argparse.ArgumentParser()
    parser.add_argument('-M', type=int, default=8192, help='problem M size')
    parser.add_argument('-N', type=int, default=8192, help='problem N size')
    parser.add_argument('-K', type=int, default=1024, help='problem K size')
    parser.add_argument('-BM', type=int, default=256, help='BLOCK_M')
    parser.add_argument('-BN', type=int, default=256, help='BLOCK_N')
    parser.add_argument('-BK', type=int, default=128, help='BLOCK_K')
    parser.add_argument('--num_warps', type=int, default=4, choices=[4, 8])
    parser.add_argument('--num_buffers', type=int, default=2, choices=[2, 3, 4])
    parser.add_argument('--group_size_m', type=int, default=8, choices=[1, 2, 4, 8])
    parser.add_argument('--scale_preshuffled', action='store_true')
    parser.add_argument('--with_a_scale', action='store_true')
    parser.add_argument('--async_copy_scale', action='store_true')
    parser.add_argument('--schedule', type=str, choices=['sliceMNK', 'sliceNK', 'sliceK', 'baseline'],
                        default='sliceNK')
    parser.add_argument('--dtype_a', type=str, default='float8_e4m3', choices=supported_dtypes)
    parser.add_argument('--dtype_b', type=str, default='float8_e4m3', choices=supported_dtypes)
    parser.add_argument('--pingpong', action='store_true')
    parser.add_argument('--l2_prefetch_distance', type=int, default=-1, choices=[-1, 0, 1, 2, 3, 4, 5],
                        help='Prefetch distance (in iterations) for operands into L2. -1 disables L2 prefetch.')
    parser.add_argument('--activation', type=str, default='', choices=['', 'swiglu'],
                        help='Optional fused activation epilogue')
    parser.add_argument('--partial_tdm', action='store_true')
    parser.add_argument(
        '--resolve_partition_conflicts', action='store_true',
        help='Use partitioned shared layouts and a partition-aware WMMA layout to avoid LDS '
        'partition conflicts')
    parser.add_argument(
        '--tdm_split', action='store_true',
        help='For the sliceMNK schedule, split A along M and B along N into two TDM halves each '
        '(four data loads per K-block) using MXFPGEMMSliceMNKTDMSplitProgram')
    parser.add_argument(
        "--benchmark-mode",
        choices=("graph", "eager", "none"),
        default="none",
        help="Timing method. `graph` uses triton.testing.do_bench_cudagraph.",
    )
    parser.add_argument(
        "--benchmark-num-iters",
        type=int,
        default=32,
        help="Number of iterations (rep) to run when benchmarking.",
    )
    args = parser.parse_args()
    BENCHMARK_MODE = None if args.benchmark_mode == "none" else args.benchmark_mode

    if args.pingpong:
        assert (args.num_warps == 8 and (args.schedule == 'baseline' or args.schedule == 'sliceK'))

    if args.num_warps == 8:
        assert (args.num_buffers in (3, 4) and not args.async_copy_scale)
        test_runtime_mxgemm_tdm_8warps_pipeline(args.dtype_a, args.dtype_b,  #
                                                args.M, args.N, args.K,  #
                                                args.BM, args.BN, args.BK,  #
                                                TRANSPOSE_B=True,  #
                                                NUM_BUFFERS=args.num_buffers,  #
                                                SCALE_PRESHUFFLE=args.scale_preshuffled,  #
                                                WITH_A_SCALE=args.with_a_scale,  #
                                                SCHEDULE=args.schedule,  #
                                                ASYNC_COPY_SCALE=False,  #
                                                GROUP_SIZE_M=args.group_size_m,  #
                                                PINGPONG=args.pingpong,  #
                                                L2_PREFETCH_DISTANCE=args.l2_prefetch_distance,  #
                                                PARTIAL_TDM=args.partial_tdm,  #
                                                RESOLVE_PARTITION_CONFLICTS=args.resolve_partition_conflicts,  #
                                                BENCHMARK_MODE=BENCHMARK_MODE,  #
                                                BENCHMARK_NUM_ITERS=args.benchmark_num_iters)
    else:
        assert (args.num_buffers in (2, 3, 4))
        test_runtime_mxgemm_tdm_pipelined(args.dtype_a, args.dtype_b,  #
                                          args.M, args.N, args.K,  #
                                          args.BM, args.BN, args.BK,  #
                                          TRANSPOSE_B=True,  #
                                          NUM_BUFFERS=args.num_buffers,  #
                                          SCALE_PRESHUFFLE=args.scale_preshuffled,  #
                                          WITH_A_SCALE=args.with_a_scale,  #
                                          SCHEDULE=args.schedule,  #
                                          ASYNC_COPY_SCALE=args.async_copy_scale,  #
                                          GROUP_SIZE_M=args.group_size_m,  #
                                          L2_PREFETCH_DISTANCE=args.l2_prefetch_distance,  #
                                          ACTIVATION=args.activation,  #
                                          PARTIAL_TDM=args.partial_tdm,  #
                                          RESOLVE_PARTITION_CONFLICTS=args.resolve_partition_conflicts,  #
                                          BENCHMARK_MODE=BENCHMARK_MODE,  #
                                          BENCHMARK_NUM_ITERS=args.benchmark_num_iters,  #
                                          TDM_SPLIT=args.tdm_split)
