import re
import math
import itertools
import pytest
import torch
from itertools import product

import triton
import triton.language as tl
from triton.backends.compiler import GPUTarget
from triton._internal_testing import is_hip_gfx1250, str_to_triton_dtype, numpy_random, to_triton, unwrap_tensor, float_dtypes, int_dtypes, uint_dtypes
from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor
from triton.experimental import gluon
import triton.experimental.gluon.language as ttgl
from triton.experimental.gluon.language.amd.cdna5 import get_wmma_scale_layout, PartitionedSharedLayout, _valid_dtype_combinations
from triton._C.libtriton.gluon_ir import make_cga_layout


@gluon.jit
def gemm_kernel(a_ptr, b_ptr, c_ptr,  #
                M, N, K,  #
                stride_am, stride_ak,  #
                stride_bk, stride_bn,  #
                stride_cm, stride_cn,  #
                BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,  #
                INSTR_SHAPE_K: ttgl.constexpr, K_WIDTH: ttgl.constexpr):

    BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0])
    WMMA_LAYOUT: ttgl.constexpr = ttgl.amd.AMDWMMALayout(3, True, [[0, 1], [1, 0]], [], [16, 16, INSTR_SHAPE_K])

    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    offs_am = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
    offs_ak = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
    offs_a = offs_am[:, None] * stride_am + offs_ak[None, :] * stride_ak

    offs_bk = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
    offs_bn = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
    offs_b = offs_bk[:, None] * stride_bk + offs_bn[None, :] * stride_bn

    accumulator = ttgl.zeros((BLOCK_M, BLOCK_N), dtype=c_ptr.type.element_ty, layout=WMMA_LAYOUT)
    for k in range(0, ttgl.cdiv(K, BLOCK_K)):
        mask_a = (offs_ak[None, :] < K - k * BLOCK_K) & (offs_am[:, None] < M)
        mask_b = (offs_bk[:, None] < K - k * BLOCK_K) & (offs_bn[None, :] < N)

        a = ttgl.load(a_ptr + offs_a, mask=mask_a, other=0.0)
        b = ttgl.load(b_ptr + offs_b, mask=mask_b, other=0.0)

        a = ttgl.convert_layout(a, ttgl.DotOperandLayout(0, WMMA_LAYOUT, K_WIDTH))
        b = ttgl.convert_layout(b, ttgl.DotOperandLayout(1, WMMA_LAYOUT, K_WIDTH))
        accumulator = ttgl.amd.cdna5.wmma(a, b, accumulator)

        offs_a += BLOCK_K * stride_ak
        offs_b += BLOCK_K * stride_bk

    offs_cm = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, WMMA_LAYOUT))
    offs_cn = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, WMMA_LAYOUT))
    offs_c = stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    mask_c = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    ttgl.store(c_ptr + offs_c, accumulator, mask=mask_c)


def get_test_gemm_block_mnk():
    return [
        (m, n, k) for (m, n) in [(32, 32), (64, 64)] \
                  for k in [32, 64, 128, 256]
    ]


def get_test_gemm_variants():
    return  [
        # float32 * float32 -> float32
        ("float32", "float32", 4),
        # bfloat16/float16 * bfloat16/float16 -> float32
        *[(a, a, 32) for a in ["bfloat16", "float16"]],
        # float8e4m3/float8e5m2 * float8e4m3/float8e5m2 -> float32/float16
        *[(a, b, k) for a in ["float8_e4m3fn", "float8_e5m2"] \
                       for b in ["float8_e4m3fn", "float8_e5m2"] \
                       for k in [64, 128]],
    ]


def get_test_gemm_shapes():
    return [
        (256, 256, 256),
        (250, 250, 250),
    ]


@pytest.mark.parametrize("a_dtype,b_dtype,k_dim", get_test_gemm_variants())
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", get_test_gemm_block_mnk())
def test_compile_gemm(a_dtype, b_dtype, k_dim, BLOCK_M, BLOCK_N, BLOCK_K):
    if BLOCK_K < k_dim:
        pytest.skip("Skip tests where BLOCK_K < k_dim")

    a_dtype = str_to_triton_dtype(a_dtype).name
    b_dtype = str_to_triton_dtype(b_dtype).name

    signature = {
        "a_ptr": f"*{a_dtype}", "b_ptr": f"*{b_dtype}", "c_ptr": "*fp32",  #
        "M": "i32", "N": "i32", "K": "i32",  #
        "stride_am": "i32", "stride_ak": "i32",  #
        "stride_bk": "i32", "stride_bn": "i32",  #
        "stride_cm": "i32", "stride_cn": "i32",  #
        "BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "BLOCK_K": "constexpr",  #
        "INSTR_SHAPE_K": "constexpr", "K_WIDTH": "constexpr"
    }
    constexprs = {
        "BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "BLOCK_K": BLOCK_K,  #
        "INSTR_SHAPE_K": k_dim, "K_WIDTH": 2 if a_dtype == "fp32" else 8
    }
    fn = gemm_kernel

    k = triton.compile(src=gluon._runtime.GluonASTSource(fn, signature, constexprs),
                       target=GPUTarget("hip", 'gfx1250', 32))
    amdgcn = k.asm["amdgcn"]

    wmma_pattern = "v_wmma_"
    wmma_pattern += "f32_"
    wmma_pattern += "16x16x" + str(k_dim) + "_"
    if a_dtype == "fp32":
        wmma_pattern += "f32"
    if a_dtype in ("fp16", "bf16"):
        a_ty = "f16" if a_dtype == "fp16" else "bf16"
        wmma_pattern += a_ty
    if a_dtype in ("fp8e4nv", "fp8e5"):
        a_ty = "fp8" if a_dtype == "fp8e4nv" else "bf8"
        b_ty = "fp8" if b_dtype == "fp8e4nv" else "bf8"
        # NOTE: we always use transposed=True for wmma layout, which will swap A and B
        wmma_pattern += b_ty + "_" + a_ty
    assert re.search(wmma_pattern, amdgcn)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("compact_scale", [
    False,
    True,
], ids=["expanded_scale", "compact_scale"])
@pytest.mark.parametrize("BLOCK_K", [64, 128, 256], ids=lambda v: f"BLOCK_K{v}")
def test_runtime_scaled_upcast_fp4(compact_scale, BLOCK_K):

    @gluon.jit
    def scaled_upcast_fp4_kernel(x_ptr, scale_ptr, y_ptr, BLOCK_M: ttgl.constexpr, BLOCK_K: ttgl.constexpr,
                                 SCALE_FACTOR: ttgl.constexpr, COMPACT_SCALE: ttgl.constexpr):
        packed_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 4], [8, 4], [4, 1], [1, 0])
        compact_layout: ttgl.constexpr = ttgl.BlockedLayout([1, BLOCK_K // SCALE_FACTOR], [8, 4], [4, 1], [1, 0])
        unpacked_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [8, 4], [4, 1], [1, 0])

        offs_m = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, packed_layout))
        offs_k_packed = ttgl.arange(0, BLOCK_K // 2, layout=ttgl.SliceLayout(0, packed_layout))
        x_offsets = offs_m[:, None] * (BLOCK_K // 2) + offs_k_packed[None, :]
        x = ttgl.load(x_ptr + x_offsets)

        if COMPACT_SCALE:
            scale_layout: ttgl.constexpr = compact_layout
            scale_k: ttgl.constexpr = BLOCK_K // SCALE_FACTOR
        else:
            scale_layout: ttgl.constexpr = unpacked_layout
            scale_k: ttgl.constexpr = BLOCK_K
        offs_scale_m = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, scale_layout))
        offs_scale_k = ttgl.arange(0, scale_k, layout=ttgl.SliceLayout(0, scale_layout))
        scale_offsets = offs_scale_m[:, None] * scale_k + offs_scale_k[None, :]
        scale = ttgl.load(scale_ptr + scale_offsets)

        y = ttgl.amd.cdna5.scaled_upcast(x, scale, ttgl.bfloat16, axis=1)
        out_layout: ttgl.constexpr = y.type.layout
        offs_out_m = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, out_layout))
        offs_out_k = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(0, out_layout))
        out_offsets = offs_out_m[:, None] * BLOCK_K + offs_out_k[None, :]
        ttgl.store(y_ptr + out_offsets, y)

    BLOCK_M = 16
    SCALE_FACTOR = 32
    torch.manual_seed(42)

    x, x_ref = create_mxfp_operand(0, BLOCK_M, BLOCK_K, "e2m1")
    scale, scale_ref = create_mxfp_scale(0, BLOCK_M, BLOCK_K, "e8m0", SCALE_FACTOR)
    if not compact_scale:
        scale = scale.repeat_interleave(SCALE_FACTOR, dim=1).contiguous()
    y = torch.empty((BLOCK_M, BLOCK_K), dtype=torch.bfloat16, device="cuda")

    pgm = scaled_upcast_fp4_kernel[(1, )](x.cuda(), scale.cuda(), y, BLOCK_M, BLOCK_K, SCALE_FACTOR, compact_scale,
                                          num_warps=4)

    assert pgm.asm["amdgcn"].count("v_cvt_scale_pk8_bf16_fp4") == BLOCK_K // 32
    y_ref = (x_ref * scale_ref).to(torch.bfloat16)
    torch.testing.assert_close(y.cpu(), y_ref, atol=0, rtol=0)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
def test_runtime_scaled_upcast_fp4_bf16_gemm_layouts():
    # Test packed / compact-scale DistributedLinearLayouts used when upcasting an fp4 weight tile to bf16 for a WMMA GEMM.
    @gluon.jit
    def scaled_upcast_fp4_bf16_gemm_kernel(x_ptr, scale_ptr, y_ptr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,
                                           SCALE_FACTOR: ttgl.constexpr):
        packed_layout: ttgl.constexpr = ttgl.DistributedLinearLayout(
            reg_bases=[[0, 1], [0, 2], [0, 8], [0, 16], [0, 32], [0, 64], [0, 128], [64, 0]],
            lane_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [0, 4]],
            warp_bases=[[16, 0], [32, 0]],
            block_bases=[],
            shape=[BLOCK_N, BLOCK_K // 2],
        )
        compact_scale_layout: ttgl.constexpr = ttgl.DistributedLinearLayout(
            reg_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [64, 0]],
            lane_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [0, 0]],
            warp_bases=[[16, 0], [32, 0]],
            block_bases=[],
            shape=[BLOCK_N, BLOCK_K // SCALE_FACTOR],
        )

        offs_n = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(1, packed_layout))
        offs_k_packed = ttgl.arange(0, BLOCK_K // 2, layout=ttgl.SliceLayout(0, packed_layout))
        x_offsets = offs_n[:, None] * (BLOCK_K // 2) + offs_k_packed[None, :]
        x = ttgl.load(x_ptr + x_offsets)

        scale_k: ttgl.constexpr = BLOCK_K // SCALE_FACTOR
        offs_scale_n = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(1, compact_scale_layout))
        offs_scale_k = ttgl.arange(0, scale_k, layout=ttgl.SliceLayout(0, compact_scale_layout))
        scale_offsets = offs_scale_n[:, None] * scale_k + offs_scale_k[None, :]
        scale = ttgl.load(scale_ptr + scale_offsets)

        y = ttgl.amd.cdna5.scaled_upcast(x, scale, ttgl.bfloat16, axis=1)
        out_layout: ttgl.constexpr = y.type.layout
        offs_out_n = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(1, out_layout))
        offs_out_k = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(0, out_layout))
        out_offsets = offs_out_n[:, None] * BLOCK_K + offs_out_k[None, :]
        ttgl.store(y_ptr + out_offsets, y)

    # Shapes match a typical fp4 weight tile: packed [128, 256] i8 == [128, 512] fp4,
    # compact scale [128, 16] (one e8m0 per 32 elements along K).
    BLOCK_N = 128
    BLOCK_K = 512
    SCALE_FACTOR = 32
    torch.manual_seed(42)

    x, x_ref = create_mxfp_operand(0, BLOCK_N, BLOCK_K, "e2m1")
    scale, scale_ref = create_mxfp_scale(0, BLOCK_N, BLOCK_K, "e8m0", SCALE_FACTOR)
    y = torch.empty((BLOCK_N, BLOCK_K), dtype=torch.bfloat16, device="cuda")

    pgm = scaled_upcast_fp4_bf16_gemm_kernel[(1, )](x.cuda(), scale.cuda(), y, BLOCK_N, BLOCK_K, SCALE_FACTOR,
                                                    num_warps=4)

    assert "v_cvt_scale_pk8_bf16_fp4" in pgm.asm["amdgcn"]
    y_ref = (x_ref * scale_ref).to(torch.bfloat16)
    torch.testing.assert_close(y.cpu(), y_ref, atol=0, rtol=0)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("fp8_dtype", ["e4m3", "e5m2"])
@pytest.mark.parametrize("out_dtype", ["bf16", "f16"], ids=lambda v: f"out_{v}")
def test_runtime_scaled_upcast_fp8(fp8_dtype, out_dtype):
    # Cover all four v_cvt_scale_pk8 fp8 instructions: {bf16,f16} output x
    # {fp8 (e4m3), bf8 (e5m2)} input.
    torch_fp8 = {"e4m3": torch.float8_e4m3fn, "e5m2": torch.float8_e5m2}[fp8_dtype]
    ttgl_out = {"bf16": ttgl.bfloat16, "f16": ttgl.float16}[out_dtype]
    torch_out = {"bf16": torch.bfloat16, "f16": torch.float16}[out_dtype]
    in_mnemonic = {"e4m3": "fp8", "e5m2": "bf8"}[fp8_dtype]

    @gluon.jit
    def scaled_upcast_fp8_kernel(x_ptr, scale_ptr, y_ptr, BLOCK_M: ttgl.constexpr, BLOCK_K: ttgl.constexpr,
                                 OUT_DTYPE: ttgl.constexpr):
        # FP8 v_cvt_scale_pk8 (opSel=0) reads the same Vscale byte for both
        # output-lane halves, so it requires a scale layout that is broadcast
        # across the lane^16 split (lane_bases[4] = [0, 0]).
        layout: ttgl.constexpr = ttgl.DistributedLinearLayout(
            reg_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [0, 32]],
            lane_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [0, 0]],
            warp_bases=[],
            block_bases=[],
            shape=[BLOCK_M, BLOCK_K],
        )

        offs_m = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, layout))
        offs_k = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(0, layout))
        offsets = offs_m[:, None] * BLOCK_K + offs_k[None, :]
        x = ttgl.load(x_ptr + offsets)
        scale = ttgl.load(scale_ptr + offsets)

        y = ttgl.amd.cdna5.scaled_upcast(x, scale, OUT_DTYPE)
        ttgl.store(y_ptr + offsets, y)

    BLOCK_M = 16
    BLOCK_K = 64
    SCALE_FACTOR = 32
    torch.manual_seed(42)

    x, x_ref = create_mxfp_operand(0, BLOCK_M, BLOCK_K, fp8_dtype)
    x = x.view(torch_fp8)
    scale, scale_ref = create_mxfp_scale(0, BLOCK_M, BLOCK_K, "e8m0", SCALE_FACTOR)
    scale = scale.repeat_interleave(SCALE_FACTOR, dim=1).contiguous()
    y = torch.empty((BLOCK_M, BLOCK_K), dtype=torch_out, device="cuda")

    pgm = scaled_upcast_fp8_kernel[(1, )](x.cuda(), scale.cuda(), y, BLOCK_M, BLOCK_K, ttgl_out, num_warps=1)

    assert f"v_cvt_scale_pk8_{out_dtype}_{in_mnemonic}" in pgm.asm["amdgcn"]
    y_ref = (x_ref * scale_ref).to(torch_out)
    torch.testing.assert_close(y.cpu(), y_ref, atol=0, rtol=0)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
def test_runtime_scaled_upcast_fp8_non_broadcast_block16():

    @gluon.jit
    def scaled_upcast_fp8_kernel(x_ptr, scale_ptr, y_ptr, BLOCK_M: ttgl.constexpr, BLOCK_K: ttgl.constexpr):
        # This layout is not broadcast across lane^16 (lane_bases[4] = [16, 0]).
        # FP8 lowering uses v_cvt_scale_pk8 Block16 mode (opSel=8) and packs
        # lane (j^16)'s scale into byte 1 via a cross-lane exchange.
        layout: ttgl.constexpr = ttgl.DistributedLinearLayout(
            reg_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [0, 32]],
            lane_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0]],
            warp_bases=[],
            block_bases=[],
            shape=[BLOCK_M, BLOCK_K],
        )

        offs_m = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, layout))
        offs_k = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(0, layout))
        offsets = offs_m[:, None] * BLOCK_K + offs_k[None, :]
        x = ttgl.load(x_ptr + offsets)
        scale = ttgl.load(scale_ptr + offsets)

        y = ttgl.amd.cdna5.scaled_upcast(x, scale, ttgl.bfloat16)
        ttgl.store(y_ptr + offsets, y)

    BLOCK_M = 32
    BLOCK_K = 64
    SCALE_FACTOR = 32
    torch.manual_seed(42)

    x, x_ref = create_mxfp_operand(0, BLOCK_M, BLOCK_K, "e4m3")
    x = x.view(torch.float8_e4m3fn)
    scale, scale_ref = create_mxfp_scale(0, BLOCK_M, BLOCK_K, "e8m0", SCALE_FACTOR)
    scale = scale.repeat_interleave(SCALE_FACTOR, dim=1).contiguous()
    y = torch.empty((BLOCK_M, BLOCK_K), dtype=torch.bfloat16, device="cuda")

    pgm = scaled_upcast_fp8_kernel[(1, )](x.cuda(), scale.cuda(), y, BLOCK_M, BLOCK_K, num_warps=1)

    assert "v_cvt_scale_pk8_bf16_fp8" in pgm.asm["amdgcn"]
    y_ref = (x_ref * scale_ref).to(torch.bfloat16)
    torch.testing.assert_close(y.cpu(), y_ref, atol=0, rtol=0)


@pytest.mark.parametrize("a_dtype,b_dtype,k_dim", get_test_gemm_variants())
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", get_test_gemm_block_mnk())
@pytest.mark.parametrize("M,N,K", get_test_gemm_shapes())
def test_runtime_gemm(a_dtype, b_dtype, k_dim, BLOCK_M, BLOCK_N, BLOCK_K, M, N, K):
    if BLOCK_K < k_dim:
        pytest.skip("Skip tests where BLOCK_K < k_dim")

    torch.manual_seed(42)

    def create_operand(shape, dtype):
        if dtype in (torch.float16, torch.bfloat16, torch.float32):
            return torch.randn(shape, dtype=dtype)
        else:
            assert dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
            return torch.randint(20, 40, shape, dtype=torch.uint8).view(dtype)

    a_dtype = getattr(torch, a_dtype)
    b_dtype = getattr(torch, b_dtype)

    a = create_operand((M, K), a_dtype)
    b = create_operand((K, N), b_dtype)
    c = torch.zeros((M, N), dtype=torch.float32)
    stride_am, stride_ak = a.stride(0), a.stride(1)
    stride_bk, stride_bn = b.stride(0), b.stride(1)
    stride_cm, stride_cn = c.stride(0), c.stride(1)

    a_device = a.cuda()
    b_device = b.cuda()
    c_device = c.cuda()
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
    gemm_kernel[grid](
        a_device, b_device, c_device,  #
        M, N, K,  #
        stride_am, stride_ak,  #
        stride_bk, stride_bn,  #
        stride_cm, stride_cn,  #
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,  #
        INSTR_SHAPE_K=k_dim, K_WIDTH=2 if a_dtype == torch.float32 else 8)

    c_triton = c_device.cpu()
    c_torch = a.to(torch.float32) @ b.to(torch.float32)
    torch.testing.assert_close(c_triton, c_torch, rtol=1e-4, atol=1e-4)


@gluon.jit
def gemm_3d_kernel(a_ptr, b_ptr, c_ptr,  #
                   B, M, N, K,  #
                   stride_ab, stride_am, stride_ak,  #
                   stride_bb, stride_bk, stride_bn,  #
                   stride_cb, stride_cm, stride_cn,  #
                   BLOCK_B: ttgl.constexpr, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,
                   INSTR_SHAPE_K: ttgl.constexpr, K_WIDTH: ttgl.constexpr):
    load_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1, 8], [1, 4, 8], [1, 4, 1], [2, 1, 0])
    wmma_layout: ttgl.constexpr = ttgl.amd.AMDWMMALayout(version=3, transposed=True, warp_bases=[[0, 0, 1], [0, 1, 0]],
                                                         reg_bases=[], instr_shape=[16, 16, INSTR_SHAPE_K], rank=3)

    load_dim0_layout: ttgl.constexpr = ttgl.SliceLayout(1, ttgl.SliceLayout(2, load_layout))
    load_dim1_layout: ttgl.constexpr = ttgl.SliceLayout(0, ttgl.SliceLayout(2, load_layout))
    load_dim2_layout: ttgl.constexpr = ttgl.SliceLayout(0, ttgl.SliceLayout(1, load_layout))

    wmma_dim0_layout: ttgl.constexpr = ttgl.SliceLayout(1, ttgl.SliceLayout(2, wmma_layout))
    wmma_dim1_layout: ttgl.constexpr = ttgl.SliceLayout(0, ttgl.SliceLayout(2, wmma_layout))
    wmma_dim2_layout: ttgl.constexpr = ttgl.SliceLayout(0, ttgl.SliceLayout(1, wmma_layout))

    pid_b = ttgl.program_id(axis=0)
    pid_m = ttgl.program_id(axis=1)
    pid_n = ttgl.program_id(axis=2)

    offs_ab = ttgl.arange(0, BLOCK_B, layout=load_dim0_layout) + (pid_b * BLOCK_B)
    offs_am = ttgl.arange(0, BLOCK_M, layout=load_dim1_layout) + (pid_m * BLOCK_M)
    offs_ak = ttgl.arange(0, BLOCK_K, layout=load_dim2_layout)
    offs_a = stride_ab * offs_ab[:, None, None] + \
             stride_am * offs_am[None, :, None] + \
             stride_ak * offs_ak[None, None, :]

    offs_bb = ttgl.arange(0, BLOCK_B, layout=load_dim0_layout) + (pid_b * BLOCK_B)
    offs_bk = ttgl.arange(0, BLOCK_K, layout=load_dim1_layout)
    offs_bn = ttgl.arange(0, BLOCK_N, layout=load_dim2_layout) + (pid_n * BLOCK_N)
    offs_b = stride_bb * offs_bb[:, None, None] + \
             stride_bk * offs_bk[None, :, None] + \
             stride_bn * offs_bn[None, None, :]

    accumulator = ttgl.zeros((BLOCK_B, BLOCK_M, BLOCK_N), dtype=c_ptr.type.element_ty, layout=wmma_layout)

    for k in range(0, ttgl.cdiv(K, BLOCK_K)):
        mask_a = (offs_ak[None, None, :] + k * BLOCK_K < K) & (offs_am[None, :, None] < M)
        mask_b = (offs_bk[None, :, None] + k * BLOCK_K < K) & (offs_bn[None, None, :] < N)

        a = ttgl.load(a_ptr + offs_a, mask=mask_a, other=0.0)
        b = ttgl.load(b_ptr + offs_b, mask=mask_b, other=0.0)

        a = ttgl.convert_layout(a, ttgl.DotOperandLayout(0, wmma_layout, K_WIDTH))
        b = ttgl.convert_layout(b, ttgl.DotOperandLayout(1, wmma_layout, K_WIDTH))
        accumulator = ttgl.amd.cdna5.wmma(a, b, accumulator)

        offs_a += BLOCK_K * stride_ak
        offs_b += BLOCK_K * stride_bk

    offs_cb = ttgl.arange(0, BLOCK_B, layout=wmma_dim0_layout) + (pid_b * BLOCK_B)
    offs_cm = ttgl.arange(0, BLOCK_M, layout=wmma_dim1_layout) + (pid_m * BLOCK_M)
    offs_cn = ttgl.arange(0, BLOCK_N, layout=wmma_dim2_layout) + (pid_n * BLOCK_N)
    offs_c = stride_cb * offs_cb[:, None, None] + \
             stride_cm * offs_cm[None, :, None] + \
             stride_cn * offs_cn[None, None, :]

    mask_c = (offs_cm[None, :, None] < M) & (offs_cn[None, None, :] < N)
    ttgl.store(c_ptr + offs_c, accumulator, mask=mask_c)


@pytest.mark.parametrize("a_dtype,b_dtype,k_dim", [
    ("float16", "float16", 32),
])
@pytest.mark.parametrize("BLOCK_B,BLOCK_M,BLOCK_N,BLOCK_K", [(4, 32, 32, 32)])
def test_compile_gemm_3d(a_dtype, b_dtype, k_dim, BLOCK_B, BLOCK_M, BLOCK_N, BLOCK_K):
    if BLOCK_K < k_dim:
        pytest.skip("Skip tests where BLOCK_K < k_dim")

    a_dtype = str_to_triton_dtype(a_dtype).name
    b_dtype = str_to_triton_dtype(b_dtype).name

    signature = {
        "a_ptr": f"*{a_dtype}", "b_ptr": f"*{b_dtype}", "c_ptr": "*fp32",  #
        "B": "i32", "M": "i32", "N": "i32", "K": "i32",  #
        "stride_ab": "i32", "stride_am": "i32", "stride_ak": "i32",  #
        "stride_bb": "i32", "stride_bk": "i32", "stride_bn": "i32",  #
        "stride_cb": "i32", "stride_cm": "i32", "stride_cn": "i32",  #
        "BLOCK_B": "constexpr", "BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "BLOCK_K": "constexpr",  #
        "INSTR_SHAPE_K": "constexpr", "K_WIDTH": "constexpr"
    }
    constexprs = {
        "BLOCK_B": BLOCK_B, "BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "BLOCK_K": BLOCK_K,  #
        "INSTR_SHAPE_K": k_dim, "K_WIDTH": 8
    }
    fn = gemm_3d_kernel

    k = triton.compile(src=gluon._runtime.GluonASTSource(fn, signature, constexprs),
                       target=GPUTarget("hip", 'gfx1250', 32))
    amdgcn = k.asm["amdgcn"]

    wmma_pattern = "v_wmma_f32_16x16x32_f16"
    assert re.search(wmma_pattern, amdgcn)


@pytest.mark.parametrize("k_dim", [32])
@pytest.mark.parametrize("BLOCK_B,BLOCK_M,BLOCK_N,BLOCK_K", [(4, 32, 32, 32)])
@pytest.mark.parametrize("B,M,N,K", [(16, 256, 256, 256), (16, 250, 250, 250)])
def test_runtime_gemm_3d(k_dim, BLOCK_B, BLOCK_M, BLOCK_N, BLOCK_K, B, M, N, K):
    assert BLOCK_K >= k_dim
    assert B % BLOCK_B == 0

    torch.manual_seed(42)

    a = torch.randn((B, M, K), dtype=torch.float16)
    b = torch.randn((B, K, N), dtype=torch.float16)
    c = torch.zeros((B, M, N), dtype=torch.float32)
    stride_ab, stride_am, stride_ak = a.stride(0), a.stride(1), a.stride(2)
    stride_bb, stride_bk, stride_bn = b.stride(0), b.stride(1), b.stride(2)
    stride_cb, stride_cm, stride_cn = c.stride(0), c.stride(1), c.stride(2)

    a_device = a.cuda()
    b_device = b.cuda()
    c_device = c.cuda()
    grid = (triton.cdiv(B, BLOCK_B), triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
    gemm_3d_kernel[grid](
        a_device, b_device, c_device,  #
        B, M, N, K,  #
        stride_ab, stride_am, stride_ak,  #
        stride_bb, stride_bk, stride_bn,  #
        stride_cb, stride_cm, stride_cn,  #
        BLOCK_B=BLOCK_B, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,  #
        INSTR_SHAPE_K=k_dim, K_WIDTH=8)

    c_triton = c_device.cpu()
    c_torch = a.to(torch.float32) @ b.to(torch.float32)
    torch.testing.assert_close(c_triton, c_torch, rtol=1e-4, atol=1e-4)


@gluon.jit
def gemm_async_pipelined_kernel(a_ptr, b_ptr, c_ptr,  #
                                M, N, K,  #
                                stride_am, stride_ak,  #
                                stride_bk, stride_bn,  #
                                stride_cm, stride_cn,  #
                                BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,  #
                                NUM_BUFFERS: ttgl.constexpr, USE_TDM: ttgl.constexpr, IS_B_K_CONTIG: ttgl.constexpr,
                                RESOLVE_PARTITION_CONFLICTS: ttgl.constexpr):
    a_dtype: ttgl.constexpr = a_ptr.type.element_ty
    b_dtype: ttgl.constexpr = b_ptr.type.element_ty
    ttgl.static_assert(a_dtype.is_fp16() or a_dtype.is_bf16(), "Only fp16/bf16 supported for A")
    ttgl.static_assert(b_dtype.is_fp16() or b_dtype.is_bf16(), "Only fp16/bf16 supported for B")
    ttgl.static_assert(NUM_BUFFERS >= 2, "NUM_BUFFERS must be at least 2")

    BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0])

    padded_a: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK_K, 8]], [BLOCK_M, BLOCK_K], [1, 0])
    if IS_B_K_CONTIG:
        padded_b: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK_K, 8]], [BLOCK_N, BLOCK_K], [1, 0])
    else:
        padded_b: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK_N, 8]], [BLOCK_K, BLOCK_N], [1, 0])

    if RESOLVE_PARTITION_CONFLICTS:
        # The padded layouts above are reinterpreted as per-piece sublayouts here.
        _DOT_LAYOUTS: ttgl.constexpr = ttgl.amd.cdna5.make_partitioned_dot_layouts(BLOCK_M, BLOCK_N, padded_a, padded_b,
                                                                                   4, [16, 16, 32], a_transposed=False,
                                                                                   b_transposed=IS_B_K_CONTIG)
    else:
        _DOT_LAYOUTS: ttgl.constexpr = (padded_a, padded_b,
                                        ttgl.amd.AMDWMMALayout(3, True, [[0, 1], [1, 0]], [], [16, 16, 32]))
    SHARED_LAYOUT_A: ttgl.constexpr = _DOT_LAYOUTS[0]
    SHARED_LAYOUT_B: ttgl.constexpr = _DOT_LAYOUTS[1]
    WMMA_LAYOUT: ttgl.constexpr = _DOT_LAYOUTS[2]

    OPERAND_LAYOUT_A: ttgl.constexpr = ttgl.DotOperandLayout(0, WMMA_LAYOUT, 8)
    OPERAND_LAYOUT_B: ttgl.constexpr = ttgl.DotOperandLayout(1, WMMA_LAYOUT, 8)

    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    # Descriptors for TDM
    a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(  #
        base=a_ptr + pid_m * BLOCK_M * stride_am,  #
        shape=(M, K),  #
        strides=(stride_am, stride_ak),  #
        block_shape=(BLOCK_M, BLOCK_K),  #
        layout=SHARED_LAYOUT_A)
    if IS_B_K_CONTIG:
        # TDM always needs the most-minor dimension to be contiguous, so for K-contiguous B we
        # describe the tile as (BLOCK_N, BLOCK_K) and permute it before the dot.
        b_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(  #
            base=b_ptr + pid_n * BLOCK_N * stride_bn,  #
            shape=(N, K),  #
            strides=(stride_bn, stride_bk),  #
            block_shape=(BLOCK_N, BLOCK_K),  #
            layout=SHARED_LAYOUT_B)
    else:
        b_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(  #
            base=b_ptr + pid_n * BLOCK_N * stride_bn,  #
            shape=(K, N),  #
            strides=(stride_bk, stride_bn),  #
            block_shape=(BLOCK_K, BLOCK_N),  #
            layout=SHARED_LAYOUT_B)

    # Pointers for AsyncCopy
    offs_ak = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
    offs_am = (pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))) % M
    a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_ak[None, :] * stride_ak

    if IS_B_K_CONTIG:
        # Tile is [BLOCK_N, BLOCK_K]: N along rows, K along cols (contiguous).
        offs_bk = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
        offs_bn = (pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))) % N
        b_ptrs = b_ptr + offs_bn[:, None] * stride_bn + offs_bk[None, :] * stride_bk
    else:
        # Tile is [BLOCK_K, BLOCK_N]: K along rows, N along cols (contiguous).
        offs_bk = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
        offs_bn = (pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))) % N
        b_ptrs = b_ptr + offs_bk[:, None] * stride_bk + offs_bn[None, :] * stride_bn

    a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, shape=[NUM_BUFFERS] + a_desc.block_shape, layout=a_desc.layout)
    b_buffer = ttgl.allocate_shared_memory(b_desc.dtype, shape=[NUM_BUFFERS] + b_desc.block_shape, layout=b_desc.layout)

    load_idx = 0
    wmma_idx = 0
    accumulator = ttgl.zeros((BLOCK_M, BLOCK_N), dtype=c_ptr.type.element_ty, layout=WMMA_LAYOUT)

    for _ in ttgl.static_range(NUM_BUFFERS - 1):
        if USE_TDM:
            ttgl.amd.cdna5.tdm.async_load(a_desc, [0, load_idx * BLOCK_K],  #
                                          a_buffer.index(load_idx % NUM_BUFFERS))
            if IS_B_K_CONTIG:
                ttgl.amd.cdna5.tdm.async_load(b_desc, [0, load_idx * BLOCK_K],  #
                                              b_buffer.index(load_idx % NUM_BUFFERS))
            else:
                ttgl.amd.cdna5.tdm.async_load(b_desc, [load_idx * BLOCK_K, 0],  #
                                              b_buffer.index(load_idx % NUM_BUFFERS))
        else:
            mask_a = offs_ak[None, :] < K - load_idx * BLOCK_K
            ttgl.amd.cdna5.async_copy.global_to_shared(a_buffer.index(load_idx % NUM_BUFFERS), a_ptrs, mask_a,
                                                       other=0.0)

            if IS_B_K_CONTIG:
                mask_b = offs_bk[None, :] < K - load_idx * BLOCK_K
            else:
                mask_b = offs_bk[:, None] < K - load_idx * BLOCK_K
            ttgl.amd.cdna5.async_copy.global_to_shared(b_buffer.index(load_idx % NUM_BUFFERS), b_ptrs, mask_b,
                                                       other=0.0)
            ttgl.amd.cdna5.async_copy.commit_group()

        load_idx += 1
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    for _ in range(0, ttgl.cdiv(K, BLOCK_K) - (NUM_BUFFERS - 1)):
        if USE_TDM:
            ttgl.amd.cdna5.tdm.async_load(a_desc, [0, load_idx * BLOCK_K],  #
                                          a_buffer.index(load_idx % NUM_BUFFERS))
            if IS_B_K_CONTIG:
                ttgl.amd.cdna5.tdm.async_load(b_desc, [0, load_idx * BLOCK_K],  #
                                              b_buffer.index(load_idx % NUM_BUFFERS))
            else:
                ttgl.amd.cdna5.tdm.async_load(b_desc, [load_idx * BLOCK_K, 0],  #
                                              b_buffer.index(load_idx % NUM_BUFFERS))
        else:
            mask_a = offs_ak[None, :] < K - load_idx * BLOCK_K
            ttgl.amd.cdna5.async_copy.global_to_shared(a_buffer.index(load_idx % NUM_BUFFERS), a_ptrs, mask_a,
                                                       other=0.0)

            if IS_B_K_CONTIG:
                mask_b = offs_bk[None, :] < K - load_idx * BLOCK_K
            else:
                mask_b = offs_bk[:, None] < K - load_idx * BLOCK_K
            ttgl.amd.cdna5.async_copy.global_to_shared(b_buffer.index(load_idx % NUM_BUFFERS), b_ptrs, mask_b,
                                                       other=0.0)
            ttgl.amd.cdna5.async_copy.commit_group()

        load_idx += 1
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

        if USE_TDM:
            ttgl.amd.cdna5.tdm.async_wait((NUM_BUFFERS - 1) * 2)
        else:
            ttgl.amd.cdna5.async_copy.wait_group((NUM_BUFFERS - 1))

        a = a_buffer.index(wmma_idx % NUM_BUFFERS).load(layout=OPERAND_LAYOUT_A)
        if IS_B_K_CONTIG:
            b = b_buffer.index(wmma_idx % NUM_BUFFERS).permute([1, 0]).load(layout=OPERAND_LAYOUT_B)
        else:
            b = b_buffer.index(wmma_idx % NUM_BUFFERS).load(layout=OPERAND_LAYOUT_B)
        accumulator = ttgl.amd.cdna5.wmma(a, b, accumulator)
        wmma_idx += 1

    for i in ttgl.static_range(NUM_BUFFERS - 1):
        if USE_TDM:
            ttgl.amd.cdna5.tdm.async_wait((NUM_BUFFERS - 2 - i) * 2)
        else:
            ttgl.amd.cdna5.async_copy.wait_group((NUM_BUFFERS - 2 - i))

        a = a_buffer.index(wmma_idx % NUM_BUFFERS).load(layout=OPERAND_LAYOUT_A)
        if IS_B_K_CONTIG:
            b = b_buffer.index(wmma_idx % NUM_BUFFERS).permute([1, 0]).load(layout=OPERAND_LAYOUT_B)
        else:
            b = b_buffer.index(wmma_idx % NUM_BUFFERS).load(layout=OPERAND_LAYOUT_B)
        accumulator = ttgl.amd.cdna5.wmma(a, b, accumulator)
        wmma_idx += 1

    offs_cm = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, WMMA_LAYOUT))
    offs_cn = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, WMMA_LAYOUT))
    offs_c = stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    mask_c = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    ttgl.store(c_ptr + offs_c, accumulator, mask=mask_c)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", [(m, n, k) for (m, n) in [(32, 32), (64, 64)] \
                                                               for k in [32, 64]])
@pytest.mark.parametrize("NUM_BUFFERS", [2, 4])
@pytest.mark.parametrize("ASYNC_LOAD_TYPE", ["ASYNC_COPY", "TDM"])
@pytest.mark.parametrize("B_K_CONTIG", [True, False])
@pytest.mark.parametrize("RESOLVE_PARTITION_CONFLICTS", [True, False])
def test_compile_gemm_async_pipelined(BLOCK_M, BLOCK_N, BLOCK_K, NUM_BUFFERS, ASYNC_LOAD_TYPE, B_K_CONTIG,
                                      RESOLVE_PARTITION_CONFLICTS):
    if RESOLVE_PARTITION_CONFLICTS:
        if ASYNC_LOAD_TYPE == "ASYNC_COPY":
            pytest.skip("AsyncCopy is not supported with partition shared layouts")
        # The kernel uses num_warps=4 and instr_shape=[16, 16, 32], which gives
        # WARP_TILES_M=4 and WARP_TILES_N=2 (see make_partitioned_dot_layouts).
        # The minimum partitionable tile size is WARP_TILES_{M,N} * instr_shape[{0,1}].
        min_block_m = 4 * 16
        min_block_n = 2 * 16
        if BLOCK_M < min_block_m or BLOCK_N < min_block_n:
            pytest.skip(
                f"Minimal block shape for partition conflicts resolution is ({min_block_m}, {min_block_n}) along (M, N)"
            )

    # Inner strides need to be constexpr (1) to get contiguity. Note the compiler frontend does the same for normal dispatches
    signature = {
        "a_ptr": "*fp16", "b_ptr": "*fp16", "c_ptr": "*fp32",  #
        "M": "i32", "N": "i32", "K": "i32",  #
        "stride_am": "i32", "stride_ak": "constexpr",  #
        "stride_bk": "constexpr" if B_K_CONTIG else "i32", "stride_bn": "i32" if B_K_CONTIG else "constexpr",  #
        "stride_cm": "i32", "stride_cn": "constexpr",  #
        "BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "BLOCK_K": "constexpr",  #
        "NUM_BUFFERS": "constexpr", "USE_TDM": "constexpr", "IS_B_K_CONTIG": "constexpr", "RESOLVE_PARTITION_CONFLICTS":
        "constexpr"
    }

    constexprs = {
        "stride_ak": 1, ("stride_bk" if B_K_CONTIG else "stride_bn"): 1, "stride_cn": 1,  #
        "BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "BLOCK_K": BLOCK_K, "NUM_BUFFERS": NUM_BUFFERS,  #
        "USE_TDM": ASYNC_LOAD_TYPE == "TDM", "IS_B_K_CONTIG": B_K_CONTIG,  #
        "RESOLVE_PARTITION_CONFLICTS": RESOLVE_PARTITION_CONFLICTS
    }

    fn = gemm_async_pipelined_kernel

    # AsyncCopy requires >= 32 bits per lane so we have to pass divisibility for arguments used in pointer arithmetic
    attrs = []
    if ASYNC_LOAD_TYPE == "ASYNC_COPY":
        attrs = {k: [["tt.divisibility", 16]] for k in [(x, ) for x in range(11)]}

    k = triton.compile(src=gluon._runtime.GluonASTSource(fn, signature, constexprs, attrs=attrs),
                       target=GPUTarget("hip", 'gfx1250', 32))
    ttgir = k.asm["ttgir"]
    amdgcn = k.asm["amdgcn"]

    assert re.search("v_wmma_f32_16x16x32_f16", amdgcn)

    if ASYNC_LOAD_TYPE == "TDM":
        # LLVM may duplicate the dynamic loop body while optimizing the final
        # assembly. Check the exact TDM copy count before LLVM backend loop
        # transforms, and only require the final ISA to contain the lowering.
        assert len(re.findall("amdg.async_tdm_copy_global_to_local", ttgir)) == NUM_BUFFERS * 2
        for cnt in range(NUM_BUFFERS - 1, -1, -1):
            assert re.search(f"s_wait_tensorcnt 0x{(cnt * 2):x}", amdgcn)
        assert len(re.findall("tensor_load_to_lds", amdgcn)) >= NUM_BUFFERS * 2
    else:
        copy_instr_for_A = BLOCK_M // 4 // 4
        b_rows = BLOCK_N if B_K_CONTIG else BLOCK_K
        copy_isntr_for_B = b_rows // 4 // 4
        copy_instr_per_iter = copy_instr_for_A + copy_isntr_for_B
        assert len(re.findall("ttg.async_copy_global_to_local", ttgir)) == NUM_BUFFERS * 2
        for cnt in range(NUM_BUFFERS - 1, -1, -1):
            assert re.search(f"s_wait_asynccnt 0x{(cnt * copy_instr_per_iter):x}", amdgcn)
        # Each instruction loads 4 rows per warp and we have 4 warps (see BlockedLayout in test)
        # Only treat this as a lower bound because LLVM might unroll the loop
        assert len(re.findall("global_load_async_to_lds", amdgcn)) >= NUM_BUFFERS * copy_instr_per_iter


@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", [(m, n, k) for (m, n) in [(32, 32), (64, 64)] \
                                                               for k in [32, 64]])
@pytest.mark.parametrize("NUM_BUFFERS", [2, 4])
@pytest.mark.parametrize("M,N,K", [(256, 256, 512), (240, 240, 496), (250, 250, 510)])
@pytest.mark.parametrize("ASYNC_LOAD_TYPE", ["ASYNC_COPY", "TDM"])
@pytest.mark.parametrize("B_K_CONTIG", [True, False])
@pytest.mark.parametrize("RESOLVE_PARTITION_CONFLICTS", [True, False])
def test_runtime_gemm_async_pipelined(BLOCK_M, BLOCK_N, BLOCK_K, NUM_BUFFERS, M, N, K, ASYNC_LOAD_TYPE, B_K_CONTIG,
                                      RESOLVE_PARTITION_CONFLICTS):
    if triton.cdiv(K, BLOCK_K) < NUM_BUFFERS:
        pytest.skip("Skip tests where K/BLOCK_K < NUM_BUFFERS")

    if ASYNC_LOAD_TYPE == "ASYNC_COPY" and any([x % 16 != 0 for x in [M, N, K]]):
        pytest.skip("AsyncCopy tests need divisibility==16 to get vectorization information")

    if RESOLVE_PARTITION_CONFLICTS:
        if ASYNC_LOAD_TYPE == "ASYNC_COPY":
            pytest.skip("AsyncCopy is not supported with partition shared layouts")
        # The kernel uses num_warps=4 and instr_shape=[16, 16, 32], which gives
        # WARP_TILES_M=4 and WARP_TILES_N=2 (see make_partitioned_dot_layouts).
        # The minimum partitionable tile size is WARP_TILES_{M,N} * instr_shape[{0,1}].
        min_block_m = 4 * 16
        min_block_n = 2 * 16
        if BLOCK_M < min_block_m or BLOCK_N < min_block_n:
            pytest.skip(
                f"Minimal block shape for partition conflicts resolution is ({min_block_m}, {min_block_n}) along (M, N)"
            )

    torch.manual_seed(42)

    a = torch.randn((M, K), dtype=torch.float16)
    if B_K_CONTIG:
        # Allocate as (N, K) and transpose so that the K dimension is the contiguous one.
        b = torch.randn((N, K), dtype=torch.float16).t()
    else:
        b = torch.randn((K, N), dtype=torch.float16)
    c = torch.zeros((M, N), dtype=torch.float32)

    a_device = a.cuda()
    b_device = b.cuda()
    c_device = c.cuda()

    stride_am, stride_ak = a_device.stride(0), a_device.stride(1)
    stride_bk, stride_bn = b_device.stride(0), b_device.stride(1)
    stride_cm, stride_cn = c_device.stride(0), c_device.stride(1)

    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
    gemm_async_pipelined_kernel[grid](
        a_device, b_device, c_device,  #
        M, N, K,  #
        stride_am, stride_ak,  #
        stride_bk, stride_bn,  #
        stride_cm, stride_cn,  #
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,  #
        NUM_BUFFERS=NUM_BUFFERS, USE_TDM=ASYNC_LOAD_TYPE == "TDM", IS_B_K_CONTIG=B_K_CONTIG,
        RESOLVE_PARTITION_CONFLICTS=RESOLVE_PARTITION_CONFLICTS)

    c_triton = c_device.cpu()
    c_torch = a.to(torch.float32) @ b.to(torch.float32)
    torch.testing.assert_close(c_triton, c_torch, rtol=1e-4, atol=1e-4)


@gluon.jit
def gemm_async_kernel(a_ptr, b_ptr, c_ptr,  #
                      M, N, K,  #
                      stride_am, stride_ak,  #
                      stride_bk, stride_bn,  #
                      stride_cm, stride_cn,  #
                      BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,  #
                      INSTR_SHAPE_K: ttgl.constexpr, K_WIDTH: ttgl.constexpr, USE_TDM: ttgl.constexpr):

    BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0])
    WMMA_LAYOUT: ttgl.constexpr = ttgl.amd.AMDWMMALayout(3, True, [[0, 1], [1, 0]], [], [16, 16, INSTR_SHAPE_K])
    SHARED_LAYOUT_A: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [BLOCK_M, BLOCK_K], [1, 0])
    SHARED_LAYOUT_B: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [BLOCK_K, BLOCK_N], [1, 0])

    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    # Descriptors for TDM
    a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr + pid_m * BLOCK_M * stride_am, shape=(M, K),
                                                       strides=(stride_am, stride_ak), block_shape=(BLOCK_M, BLOCK_K),
                                                       layout=SHARED_LAYOUT_A)
    b_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr + pid_n * BLOCK_N * stride_bn, shape=(K, N),
                                                       strides=(stride_bk, stride_bn), block_shape=(BLOCK_K, BLOCK_N),
                                                       layout=SHARED_LAYOUT_B)

    # Pointers for AsyncCopy
    offs_ak = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
    offs_am = (pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))) % M
    a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_ak[None, :] * stride_ak

    offs_bk = ttgl.arange(0, BLOCK_K, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
    offs_bn = (pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))) % N
    b_ptrs = b_ptr + offs_bk[:, None] * stride_bk + offs_bn[None, :] * stride_bn

    a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, shape=a_desc.block_shape, layout=a_desc.layout)
    b_buffer = ttgl.allocate_shared_memory(b_desc.dtype, shape=b_desc.block_shape, layout=b_desc.layout)

    accumulator = ttgl.zeros((BLOCK_M, BLOCK_N), dtype=c_ptr.type.element_ty, layout=WMMA_LAYOUT)
    for k in range(0, ttgl.cdiv(K, BLOCK_K)):
        if USE_TDM:
            ttgl.amd.cdna5.tdm.async_load(a_desc, [0, k * BLOCK_K], a_buffer)
            ttgl.amd.cdna5.tdm.async_load(b_desc, [k * BLOCK_K, 0], b_buffer)
            ttgl.amd.cdna5.tdm.async_wait(0)
        else:
            mask_a = offs_ak[None, :] < K - k * BLOCK_K
            ttgl.amd.cdna5.async_copy.global_to_shared(a_buffer, a_ptrs, mask_a, other=0.0)

            mask_b = offs_bk[:, None] < K - k * BLOCK_K
            ttgl.amd.cdna5.async_copy.global_to_shared(b_buffer, b_ptrs, mask_b, other=0.0)
            ttgl.amd.cdna5.async_copy.commit_group()
            a_ptrs += BLOCK_K * stride_ak
            b_ptrs += BLOCK_K * stride_bk
            ttgl.amd.cdna5.async_copy.wait_group(0)

        a = a_buffer.load(layout=BLOCKED_LAYOUT)
        b = b_buffer.load(layout=BLOCKED_LAYOUT)

        a = ttgl.convert_layout(a, ttgl.DotOperandLayout(0, WMMA_LAYOUT, K_WIDTH))
        b = ttgl.convert_layout(b, ttgl.DotOperandLayout(1, WMMA_LAYOUT, K_WIDTH))
        accumulator = ttgl.amd.cdna5.wmma(a, b, accumulator)

    offs_cm = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, WMMA_LAYOUT))
    offs_cn = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, WMMA_LAYOUT))
    offs_c = stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    mask_c = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    ttgl.store(c_ptr + offs_c, accumulator, mask=mask_c)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", [(32, 32, 32), (64, 64, 64), (128, 128, 64)])
@pytest.mark.parametrize("a_dtype,b_dtype,k_dim", [
    ("bfloat16", "bfloat16", 32),
    ("float8_e5m2", "float8_e5m2", 64),
])
@pytest.mark.parametrize("ASYNC_LOAD_TYPE", ["ASYNC_COPY", "TDM"])
def test_compile_gemm_async(BLOCK_M, BLOCK_N, BLOCK_K, a_dtype, b_dtype, k_dim, ASYNC_LOAD_TYPE):
    if BLOCK_K < k_dim:
        pytest.skip("Skip tests where BLOCK_K < k_dim")

    a_dtype = str_to_triton_dtype(a_dtype).name
    b_dtype = str_to_triton_dtype(b_dtype).name

    # AsyncCopy requires >= 32 bits per lane so we have to pass divisibility for arguments used in pointer arithmetic
    attrs = []
    if ASYNC_LOAD_TYPE == "ASYNC_COPY":
        attrs = {(k, ): [["tt.divisibility", 16]] for k in [0, 1, 2, 3, 4, 5, 6, 8, 10]}

    k = triton.compile(
        # Inner strides need to be constexpr (1) to get contiguity. Note the compiler frontend does the same for normal dispatches
        gluon._runtime.GluonASTSource(
            fn=gemm_async_kernel, signature={
                "a_ptr": f"*{a_dtype}", "b_ptr": f"*{b_dtype}", "c_ptr": "*fp32",  #
                "M": "i32", "N": "i32", "K": "i32",  #
                "stride_am": "i32", "stride_ak": "constexpr",  #
                "stride_bk": "i32", "stride_bn": "constexpr",  #
                "stride_cm": "i32", "stride_cn": "constexpr",  #
                "BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "BLOCK_K": "constexpr",  #
                "INSTR_SHAPE_K": "constexpr", "K_WIDTH": "constexpr", "USE_TDM": "constexpr"
            }, attrs=attrs, constexprs={
                "stride_ak": 1, "stride_bn": 1, "stride_cn": 1, "BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "BLOCK_K":
                BLOCK_K, "INSTR_SHAPE_K": k_dim, "K_WIDTH": 8, "USE_TDM": ASYNC_LOAD_TYPE == "TDM"
            }), target=GPUTarget("hip", 'gfx1250', 32))
    amdgcn = k.asm["amdgcn"]

    if ASYNC_LOAD_TYPE == "TDM":
        patterns = ("tensor_load_to_lds", "s_wait_tensorcnt 0x0")
    elif ASYNC_LOAD_TYPE == "ASYNC_COPY":
        patterns = ("global_load_async_to_lds", "s_wait_asynccnt 0x0")

    for pattern in patterns:
        assert re.search(pattern, amdgcn), f"Can't find {pattern} in amdgcn"


@pytest.mark.parametrize("M,N,K", [(256, 256, 128), (250, 250, 120)])
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", [(32, 32, 32), (64, 64, 64), (128, 128, 64)])
@pytest.mark.parametrize("a_dtype,b_dtype,k_dim", [
    ("bfloat16", "bfloat16", 32),
    ("float8_e5m2", "float8_e5m2", 64),
])
@pytest.mark.parametrize("ASYNC_LOAD_TYPE", ["ASYNC_COPY", "TDM"])
def test_runtime_gemm_async(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, a_dtype, b_dtype, k_dim, ASYNC_LOAD_TYPE):
    if BLOCK_K < k_dim:
        pytest.skip("Skip tests where BLOCK_K < k_dim")
    if ASYNC_LOAD_TYPE == "ASYNC_COPY" and any([x % 16 != 0 for x in [M, N, K]]):
        pytest.skip("AsyncCopy tests need divisibility==16 to get vectorization information")

    torch.manual_seed(42)

    def create_operand(shape, dtype):
        if dtype == torch.bfloat16:
            return torch.randn(shape, dtype=dtype)
        else:
            assert dtype == torch.float8_e5m2
            return torch.randint(20, 40, shape, dtype=torch.uint8).view(dtype)

    a_dtype = getattr(torch, a_dtype)
    b_dtype = getattr(torch, b_dtype)

    a = create_operand((M, K), a_dtype)
    b = create_operand((K, N), b_dtype)
    c = torch.zeros((M, N), dtype=torch.float32)
    stride_am, stride_ak = a.stride(0), a.stride(1)
    stride_bk, stride_bn = b.stride(0), b.stride(1)
    stride_cm, stride_cn = c.stride(0), c.stride(1)

    a_device = a.cuda()
    b_device = b.cuda()
    c_device = c.cuda()
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
    gemm_async_kernel[grid](
        a_device, b_device, c_device,  #
        M, N, K,  #
        stride_am, stride_ak,  #
        stride_bk, stride_bn,  #
        stride_cm, stride_cn,  #
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,  #
        INSTR_SHAPE_K=k_dim, K_WIDTH=8, USE_TDM=ASYNC_LOAD_TYPE == "TDM")

    c_triton = c_device.cpu()
    c_torch = a.to(torch.float32) @ b.to(torch.float32)
    torch.testing.assert_close(c_triton, c_torch, rtol=1e-4, atol=1e-4)


def torch_gemm_mxfp(a, b, a_scale, b_scale, scale_block, M, N, K):
    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 create_mxfp_operand(operand: int, m: int, n: int, dtype: str):
    size = (m, n)
    if dtype == 'e4m3':
        v = torch.randint(20, 40, size, dtype=torch.uint8)
        v_ref = v.view(torch.float8_e4m3fn).to(torch.float32)
    elif dtype == 'e5m2':
        v = torch.randint(20, 40, size, dtype=torch.uint8)
        v_ref = v.view(torch.float8_e5m2).to(torch.float32)
    else:
        assert dtype == 'e2m1'
        pack_dim = 1 if operand == 0 else 0
        v_mxfp4 = MXFP4Tensor(size=size).random()
        v = v_mxfp4.to_packed_tensor(pack_dim)
        v_ref = v_mxfp4.to(torch.float32)
    return v, v_ref


def create_mxfp_scale(operand: int, m: int, n: int, scale_dtype: str, scale_factor: int):
    pack_dim = 1 if operand == 0 else 0
    size = (m, n // scale_factor) if pack_dim == 1 else (m // scale_factor, n)
    if scale_dtype == "e8m0":
        scale = MXScaleTensor(size=tuple(size)).random(low=1.0, high=scale_factor)
    else:
        assert scale_dtype == "e4m3"
        scale = torch.randint(20, 40, size, dtype=torch.uint8).view(torch.float8_e4m3fn)
    scale_ref = scale.to(torch.float32).repeat_interleave(scale_factor, dim=pack_dim)
    return scale.data, scale_ref


def get_test_mxfp_block_mnk():
    return [(m, n, k) for m, n in [(16, 16), (32, 32), (64, 64)] for k in [64, 128, 256]]


def get_test_mxfp_variants():
    types = ["e2m1", "e4m3", "e5m2"]
    return [(a_type, b_type) for a_type in types for b_type in types]


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("wmma_shape", [(16, 16), (32, 16)])
@pytest.mark.parametrize("transposed", [False, True])
@pytest.mark.parametrize(
    "M, N, K",
    get_test_mxfp_block_mnk() +
    # Add small K testcases
    [(16, 64, 32), (64, 128, 64)] +
    # Add non-square test cases
    [(32, 64, 128), (64, 32, 128)])
@pytest.mark.parametrize("a_type, b_type", get_test_mxfp_variants())
@pytest.mark.parametrize("a_scale_type, b_scale_type", itertools.product(["e8m0", "e4m3"], repeat=2))
@pytest.mark.parametrize("scale_factor", [16, 32])
@pytest.mark.parametrize("with_a_scale, with_b_scale", itertools.product([True, False], repeat=2))
def test_amd_wmma_scaled(wmma_shape, transposed, M, N, K, a_type, b_type, a_scale_type, b_scale_type, scale_factor,
                         with_a_scale, with_b_scale):
    instr_m, instr_n = wmma_shape

    if instr_m == 32:
        if a_type != "e2m1" or b_type != "e2m1":
            pytest.skip("32x16 WMMA only supports FP4 x FP4")

    @gluon.jit
    def kernel(c_ptr, a_ptr, a_scale_ptr, b_ptr, b_scale_ptr,  #
               a_type: ttgl.constexpr, b_type: ttgl.constexpr,  #
               BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,  #
               SCALE_FACTOR: ttgl.constexpr, INSTR_SHAPE_M: ttgl.constexpr, INSTR_SHAPE_N: ttgl.constexpr,  #
               TRANSPOSED: ttgl.constexpr):
        DIV_FACTOR_A: ttgl.constexpr = 2 if a_type == "e2m1" else 1
        DIV_FACTOR_B: ttgl.constexpr = 2 if b_type == "e2m1" else 1

        wmma_layout: ttgl.constexpr = ttgl.amd.AMDWMMALayout(  #
            version=3, transposed=TRANSPOSED, warp_bases=[[0, 1], [1, 0]],
            instr_shape=[INSTR_SHAPE_M, INSTR_SHAPE_N, 128])
        wmma_layout_packed: ttgl.constexpr = ttgl.amd.AMDWMMALayout(  #
            version=3, transposed=TRANSPOSED, warp_bases=[[0, 1], [1, 0]],
            instr_shape=[INSTR_SHAPE_M, INSTR_SHAPE_N, 64])
        a_layout: ttgl.constexpr = ttgl.DotOperandLayout(  #
            0, wmma_layout_packed if a_type == "e2m1" else wmma_layout, k_width=16)
        b_layout: ttgl.constexpr = ttgl.DotOperandLayout(  #
            1, wmma_layout_packed if b_type == "e2m1" else wmma_layout, k_width=16)
        a_scale_layout: ttgl.constexpr = get_wmma_scale_layout(a_layout, [BLOCK_M, BLOCK_K // SCALE_FACTOR],
                                                               SCALE_FACTOR)
        b_scale_layout: ttgl.constexpr = get_wmma_scale_layout(b_layout, [BLOCK_N, BLOCK_K // SCALE_FACTOR],
                                                               SCALE_FACTOR)

        a_offs = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, a_layout))[:, None] * (BLOCK_K // DIV_FACTOR_A) + \
                 ttgl.arange(0, BLOCK_K // DIV_FACTOR_A, layout=ttgl.SliceLayout(0, a_layout))[None, :]
        a = ttgl.load(a_ptr + a_offs)

        b_offs = ttgl.arange(0, BLOCK_K // DIV_FACTOR_B, layout=ttgl.SliceLayout(1, b_layout))[:, None] * BLOCK_N + \
                 ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, b_layout))[None, :]
        b = ttgl.load(b_ptr + b_offs)

        if a_scale_ptr is not None:
            a_scale_offs = \
                ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, a_scale_layout))[:, None] * (BLOCK_K // SCALE_FACTOR) + \
                ttgl.arange(0, BLOCK_K // SCALE_FACTOR, layout=ttgl.SliceLayout(0, a_scale_layout))[None, :]
            a_scale = ttgl.load(a_scale_ptr + a_scale_offs)
        else:
            a_scale = None

        if b_scale_ptr is not None:
            b_scale_offs = \
                ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(1, b_scale_layout))[:, None] * (BLOCK_K // SCALE_FACTOR) + \
                ttgl.arange(0, BLOCK_K // SCALE_FACTOR, layout=ttgl.SliceLayout(0, b_scale_layout))[None, :]
            b_scale = ttgl.load(b_scale_ptr + b_scale_offs)
        else:
            b_scale = None

        zero = ttgl.zeros([BLOCK_M, BLOCK_N], dtype=ttgl.float32, layout=wmma_layout)
        c = ttgl.amd.cdna5.wmma_scaled(a, a_scale, a_type, b, b_scale, b_type, zero)
        c = c.to(c_ptr.dtype.element_ty)

        c_offs = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, wmma_layout))[:, None] * BLOCK_N + \
                 ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, wmma_layout))[None, :]
        ttgl.store(c_ptr + c_offs, c)

    if (a_type, b_type, a_scale_type, b_scale_type) not in _valid_dtype_combinations:
        pytest.skip(f"Invalid type combination: {a_type}, {a_scale_type}, {b_type}, {b_scale_type}")

    torch.manual_seed(42)
    a, a_ref = create_mxfp_operand(0, M, K, a_type)
    b, b_ref = create_mxfp_operand(1, K, N, b_type)
    a_scale, a_scale_ref = create_mxfp_scale(0, M, K, a_scale_type, scale_factor)
    b_scale, b_scale_ref = create_mxfp_scale(1, K, N, b_scale_type, scale_factor)
    b_scale = b_scale.permute(1, 0).contiguous()

    a, a_scale = a.cuda(), a_scale.cuda()
    b, b_scale = b.cuda(), b_scale.cuda()
    if not with_a_scale:
        a_scale = None
        a_scale_ref = 1.0
    if not with_b_scale:
        b_scale = None
        b_scale_ref = 1.0

    c = torch.zeros((M, N), dtype=torch.float32).cuda()
    pgm = kernel[(1, )](c, a, a_scale, b, b_scale, a_type, b_type, M, N, K, scale_factor, instr_m, instr_n, transposed,
                        num_warps=4)

    no_scales = not with_a_scale and not with_b_scale

    if instr_m == 32:
        if scale_factor == 32 or no_scales:
            assert "v_wmma_scale_f32_32x16x128_f4" in pgm.asm["amdgcn"]
        else:
            assert "v_wmma_scale16_f32_32x16x128_f4" in pgm.asm["amdgcn"]
    else:
        if scale_factor == 32 or no_scales:
            assert "v_wmma_scale_f32_16x16x128_f8f6f4" in pgm.asm["amdgcn"]
        else:
            assert "v_wmma_scale16_f32_16x16x128_f8f6f4" in pgm.asm["amdgcn"]

    c_torch = (a_ref * a_scale_ref) @ (b_ref * b_scale_ref)
    torch.testing.assert_close(c.cpu(), c_torch, atol=1e-5, rtol=2e-5)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("M, N, K", get_test_mxfp_block_mnk())
@pytest.mark.parametrize("a_type, b_type", get_test_mxfp_variants())
@pytest.mark.parametrize("a_scale_type, b_scale_type", itertools.product(["e8m0", "e4m3"], repeat=2))
@pytest.mark.parametrize("scale_factor", [16, 32])
def test_amd_wmma_scaled_multi_cta(M, N, K, a_type, b_type, a_scale_type, b_scale_type, scale_factor):

    @gluon.constexpr_function
    def _get_wmma_layout(cga_layout=[[0, 1], [1, 0]], packed=False):
        instr_shape = [16, 16, 64] if packed else [16, 16, 128]
        return ttgl.amd.AMDWMMALayout(version=3, transposed=True, warp_bases=[[0, 1], [1, 0]], instr_shape=instr_shape,
                                      cga_layout=cga_layout)

    @gluon.constexpr_function
    def _get_wmma_operand_layout(operand_index, dtype, transposed=False):
        cga_layout = [[0, 0], [1, 0]] if operand_index == 0 else [[0, 1], [0, 0]]
        if transposed:
            cga_layout = [list(reversed(row)) for row in cga_layout]
        wmma_layout = _get_wmma_layout(cga_layout, packed=(dtype == "e2m1"))
        return ttgl.DotOperandLayout(operand_index, wmma_layout, k_width=16)

    @gluon.constexpr_function
    def _get_wmma_scale_layout(operand_index, dtype, shape, scale_factor):
        transposed = True if operand_index == 1 else False
        operand_layout = _get_wmma_operand_layout(operand_index, dtype, transposed)
        return get_wmma_scale_layout(operand_layout, shape, scale_factor)

    @gluon.jit
    def kernel(c_ptr, a_ptr, a_scale_ptr, b_ptr, b_scale_ptr,  #
               a_type: ttgl.constexpr, b_type: ttgl.constexpr,  #
               BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,  #
               SCALE_FACTOR: ttgl.constexpr):
        DIV_FACTOR_A: ttgl.constexpr = 2 if a_type == "e2m1" else 1
        DIV_FACTOR_B: ttgl.constexpr = 2 if b_type == "e2m1" else 1

        acc_layout: ttgl.constexpr = _get_wmma_layout()
        a_layout: ttgl.constexpr = _get_wmma_operand_layout(0, a_type)
        b_layout: ttgl.constexpr = _get_wmma_operand_layout(1, b_type)
        a_scale_layout: ttgl.constexpr = _get_wmma_scale_layout(0, a_type, [BLOCK_M, BLOCK_K // SCALE_FACTOR],
                                                                SCALE_FACTOR)
        b_scale_layout: ttgl.constexpr = _get_wmma_scale_layout(1, b_type, [BLOCK_N, BLOCK_K // SCALE_FACTOR],
                                                                SCALE_FACTOR)

        a_offs = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, a_layout))[:, None] * (BLOCK_K // DIV_FACTOR_A) + \
                 ttgl.arange(0, BLOCK_K // DIV_FACTOR_A, layout=ttgl.SliceLayout(0, a_layout))[None, :]
        a = ttgl.load(a_ptr + a_offs)

        b_offs = ttgl.arange(0, BLOCK_K // DIV_FACTOR_B, layout=ttgl.SliceLayout(1, b_layout))[:, None] * BLOCK_N + \
                 ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, b_layout))[None, :]
        b = ttgl.load(b_ptr + b_offs)

        a_scale_offs = \
            ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, a_scale_layout))[:, None] * (BLOCK_K // SCALE_FACTOR) + \
            ttgl.arange(0, BLOCK_K // SCALE_FACTOR, layout=ttgl.SliceLayout(0, a_scale_layout))[None, :]
        a_scale = ttgl.load(a_scale_ptr + a_scale_offs)

        b_scale_offs = \
            ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(1, b_scale_layout))[:, None] * (BLOCK_K // SCALE_FACTOR) + \
            ttgl.arange(0, BLOCK_K // SCALE_FACTOR, layout=ttgl.SliceLayout(0, b_scale_layout))[None, :]
        b_scale = ttgl.load(b_scale_ptr + b_scale_offs)

        zero = ttgl.zeros([BLOCK_M, BLOCK_N], dtype=ttgl.float32, layout=acc_layout)
        c = ttgl.amd.cdna5.wmma_scaled(a, a_scale, a_type, b, b_scale, b_type, zero)
        c = c.to(c_ptr.dtype.element_ty)

        c_offs = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, acc_layout))[:, None] * BLOCK_N + \
                 ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, acc_layout))[None, :]
        ttgl.store(c_ptr + c_offs, c)

    if (a_type, b_type, a_scale_type, b_scale_type) not in _valid_dtype_combinations:
        pytest.skip(f"Invalid type combination: {a_type}, {a_scale_type}, {b_type}, {b_scale_type}")

    torch.manual_seed(42)
    M, N = M * 2, N * 2
    a, a_ref = create_mxfp_operand(0, M, K, a_type)
    b, b_ref = create_mxfp_operand(1, K, N, b_type)
    a_scale, a_scale_ref = create_mxfp_scale(0, M, K, a_scale_type, scale_factor)
    b_scale, b_scale_ref = create_mxfp_scale(1, K, N, b_scale_type, scale_factor)
    b_scale = b_scale.permute(1, 0).contiguous()

    a, a_scale = a.cuda(), a_scale.cuda()
    b, b_scale = b.cuda(), b_scale.cuda()
    c = torch.zeros((M, N), dtype=torch.float32).cuda()
    pgm = kernel[(1, )](c, a, a_scale, b, b_scale, a_type, b_type, M, N, K, scale_factor, num_warps=4, num_ctas=4)
    if scale_factor == 32:
        assert "v_wmma_scale_f32_16x16x128_f8f6f4" in pgm.asm["amdgcn"]
    else:
        assert "v_wmma_scale16_f32_16x16x128_f8f6f4" in pgm.asm["amdgcn"]

    c_torch = (a_ref * a_scale_ref) @ (b_ref * b_scale_ref)
    torch.testing.assert_close(c.cpu(), c_torch, atol=1e-5, rtol=2e-5)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("B", [4])
@pytest.mark.parametrize("M, N, K", get_test_mxfp_block_mnk())
@pytest.mark.parametrize("a_type, b_type", get_test_mxfp_variants())
@pytest.mark.parametrize("a_scale_type, b_scale_type", itertools.product(["e8m0", "e4m3"], repeat=2))
@pytest.mark.parametrize("scale_factor", [16, 32])
def test_amd_wmma_scaled_batched(B, M, N, K, a_type, b_type, a_scale_type, b_scale_type, scale_factor):

    @gluon.constexpr_function
    def _slice_layout(layout, indices):
        for i in reversed(indices):
            layout = ttgl.SliceLayout(i, layout)
        return layout

    @gluon.jit
    def _offsets(dim0, dim1, dim2, layout):
        return ttgl.arange(0, dim0, _slice_layout(layout, [1, 2]))[:, None, None] * (dim1 * dim2) + \
               ttgl.arange(0, dim1, _slice_layout(layout, [0, 2]))[None, :, None] * dim2 + \
               ttgl.arange(0, dim2, _slice_layout(layout, [0, 1]))[None, None, :]

    @gluon.jit
    def kernel(c_ptr, a_ptr, a_scale_ptr, b_ptr, b_scale_ptr,  #
               a_type: ttgl.constexpr, b_type: ttgl.constexpr,  #
               BLOCK_B: ttgl.constexpr, BLOCK_M: ttgl.constexpr,  #
               BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,  #
               SCALE_FACTOR: ttgl.constexpr):
        DIV_FACTOR_A: ttgl.constexpr = 2 if a_type == "e2m1" else 1
        DIV_FACTOR_B: ttgl.constexpr = 2 if b_type == "e2m1" else 1

        warp_bases: ttgl.constexpr = [[1, 0, 0], [2, 0, 0]]
        wmma_layout: ttgl.constexpr = \
            ttgl.amd.AMDWMMALayout(3, True, warp_bases, instr_shape=[16, 16, 128], rank=3)
        wmma_layout_packed: ttgl.constexpr = \
            ttgl.amd.AMDWMMALayout(3, True, warp_bases, instr_shape=[16, 16, 64], rank=3)
        a_layout: ttgl.constexpr = \
            ttgl.DotOperandLayout(0, wmma_layout_packed if a_type == "e2m1" else wmma_layout, 16)
        b_layout: ttgl.constexpr = \
            ttgl.DotOperandLayout(1, wmma_layout_packed if b_type == "e2m1" else wmma_layout, 16)
        a_scale_layout: ttgl.constexpr = \
            get_wmma_scale_layout(a_layout, [BLOCK_B, BLOCK_M, BLOCK_K // SCALE_FACTOR], SCALE_FACTOR)
        b_scale_layout: ttgl.constexpr = \
            get_wmma_scale_layout(b_layout, [BLOCK_B, BLOCK_N, BLOCK_K // SCALE_FACTOR], SCALE_FACTOR)

        a_offs = _offsets(BLOCK_B, BLOCK_M, BLOCK_K // DIV_FACTOR_A, a_layout)
        a = ttgl.load(a_ptr + a_offs)
        b_offs = _offsets(BLOCK_B, BLOCK_K // DIV_FACTOR_B, BLOCK_N, b_layout)
        b = ttgl.load(b_ptr + b_offs)

        a_scale_offs = _offsets(BLOCK_B, BLOCK_M, BLOCK_K // SCALE_FACTOR, a_scale_layout)
        a_scale = ttgl.load(a_scale_ptr + a_scale_offs)
        b_scale_offs = _offsets(BLOCK_B, BLOCK_N, BLOCK_K // SCALE_FACTOR, b_scale_layout)
        b_scale = ttgl.load(b_scale_ptr + b_scale_offs)

        zero = ttgl.zeros([BLOCK_B, BLOCK_M, BLOCK_N], ttgl.float32, wmma_layout)
        c = ttgl.amd.cdna5.wmma_scaled(a, a_scale, a_type, b, b_scale, b_type, zero)
        c = c.to(c_ptr.dtype.element_ty)

        c_offs = _offsets(BLOCK_B, BLOCK_M, BLOCK_N, wmma_layout)
        ttgl.store(c_ptr + c_offs, c)

    if (a_type, b_type, a_scale_type, b_scale_type) not in _valid_dtype_combinations:
        pytest.skip(f"Invalid type combination: {a_type}, {a_scale_type}, {b_type}, {b_scale_type}")

    torch.manual_seed(42)
    a, a_ref = zip(*[create_mxfp_operand(0, M, K, a_type) for _ in range(B)])
    b, b_ref = zip(*[create_mxfp_operand(1, K, N, b_type) for _ in range(B)])
    a_scale, a_scale_ref = zip(*[create_mxfp_scale(0, M, K, a_scale_type, scale_factor) for _ in range(B)])
    b_scale, b_scale_ref = zip(*[create_mxfp_scale(1, K, N, b_scale_type, scale_factor) for _ in range(B)])

    a = torch.stack(a, dim=0)
    b = torch.stack(b, dim=0)
    a_scale = torch.stack(a_scale, dim=0)
    b_scale = torch.stack(b_scale, dim=0).permute(0, 2, 1).contiguous()

    a_ref = torch.stack(a_ref, dim=0)
    b_ref = torch.stack(b_ref, dim=0)
    a_scale_ref = torch.stack(a_scale_ref, dim=0)
    b_scale_ref = torch.stack(b_scale_ref, dim=0)

    a, a_scale = a.cuda(), a_scale.cuda()
    b, b_scale = b.cuda(), b_scale.cuda()

    c = torch.zeros((B, M, N), dtype=torch.float32).cuda()
    kernel[(1, )](c, a, a_scale, b, b_scale, a_type, b_type, B, M, N, K, scale_factor, num_warps=4)

    c_torch = (a_ref * a_scale_ref) @ (b_ref * b_scale_ref)
    torch.testing.assert_close(c.cpu(), c_torch, atol=1e-5, rtol=2e-5)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("M, N, K", [(16, 16, 128), (32, 32, 128), (32, 32, 256), (32, 32, 512), (64, 64, 128),
                                     (128, 128, 256)])
@pytest.mark.parametrize("mxfp_type", ["e2m1"])
@pytest.mark.parametrize("hasScale", [True, False])
@pytest.mark.parametrize("scale_dtype, scale_factor", [("e8m0", 32), ("e4m3", 16)])
def test_amd_wmma_scaled_tdm(M, N, K, mxfp_type, hasScale, scale_dtype, scale_factor):

    @triton.jit
    def scaled_wmma_tdm_triton_kernel(a_base, stride_am, stride_ak, a_scale, b_base, stride_bk, stride_bn, b_scale, out,
                                      BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
                                      type_a: tl.constexpr, type_b: tl.constexpr, SCALE_FACTOR: tl.constexpr):
        DIV_FACTOR_A: tl.constexpr = 2 if type_a == "e2m1" else 1
        DIV_FACTOR_B: tl.constexpr = 2 if type_b == "e2m1" else 1
        PACKED_BLOCK_K_A: tl.constexpr = BLOCK_K // DIV_FACTOR_A
        PACKED_BLOCK_K_B: tl.constexpr = BLOCK_K // DIV_FACTOR_B
        a_desc = tl.make_tensor_descriptor(base=a_base, shape=(BLOCK_M, PACKED_BLOCK_K_A),
                                           strides=(stride_am, stride_ak), block_shape=(BLOCK_M, PACKED_BLOCK_K_A))
        b_desc = tl.make_tensor_descriptor(base=b_base, shape=(PACKED_BLOCK_K_B, BLOCK_N),
                                           strides=(stride_bk, stride_bn), block_shape=(PACKED_BLOCK_K_B, BLOCK_N))
        a = a_desc.load([0, 0])
        b = b_desc.load([0, 0])
        SCALE_BLOCK_K: tl.constexpr = BLOCK_K // SCALE_FACTOR

        if a_scale is not None:
            scale_a_ptr = a_scale + tl.arange(0, BLOCK_M)[:, None] * SCALE_BLOCK_K + tl.arange(0,
                                                                                               SCALE_BLOCK_K)[None, :]
            a_scale = tl.load(scale_a_ptr)
        if b_scale is not None:
            scale_b_ptr = b_scale + tl.arange(0, BLOCK_N)[:, None] * SCALE_BLOCK_K + tl.arange(0,
                                                                                               SCALE_BLOCK_K)[None, :]
            b_scale = tl.load(scale_b_ptr)
        c = tl.dot_scaled(a, a_scale, type_a, b, b_scale, type_b)
        out_ptr = out + tl.arange(0, BLOCK_M)[:, None] * BLOCK_N + tl.arange(0, BLOCK_N)[None, :]
        tl.store(out_ptr, c)

    @gluon.jit
    def scaled_wmma_tdm_gluon_kernel(a_base, stride_am, stride_ak, a_scale, b_base, stride_bk, stride_bn, b_scale, out,
                                     BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,
                                     type_a: ttgl.constexpr, type_b: ttgl.constexpr, SCALE_FACTOR: ttgl.constexpr):
        DIV_FACTOR_A: ttgl.constexpr = 2 if type_a == "e2m1" else 1
        DIV_FACTOR_B: ttgl.constexpr = 2 if type_b == "e2m1" else 1
        PACKED_BLOCK_K_A: ttgl.constexpr = BLOCK_K // DIV_FACTOR_A
        PACKED_BLOCK_K_B: ttgl.constexpr = BLOCK_K // DIV_FACTOR_B
        SCALE_BLOCK_K: ttgl.constexpr = BLOCK_K // SCALE_FACTOR

        scale_blocked_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [8, 4], [4, 1], [1, 0])
        a_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 16], [8, 4], [4, 1], [1, 0])

        b_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 16], [16, 2], [4, 1], [1, 0])
        SHARED_LAYOUT_A: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]],
                                                                                    [BLOCK_M, PACKED_BLOCK_K_A], [1, 0])
        SHARED_LAYOUT_B: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]],
                                                                                    [PACKED_BLOCK_K_B, BLOCK_N], [1, 0])

        wmma_layout: ttgl.constexpr = ttgl.amd.AMDWMMALayout(version=3, transposed=True, warp_bases=[[0, 1], [1, 0]],
                                                             instr_shape=[16, 16, 128])
        wmma_layout_packed: ttgl.constexpr = ttgl.amd.AMDWMMALayout(version=3, transposed=True, warp_bases=[[0, 1],
                                                                                                            [1, 0]],
                                                                    instr_shape=[16, 16, 64])

        a_operand_layout: ttgl.constexpr = ttgl.DotOperandLayout(
            operand_index=0, parent=wmma_layout_packed if type_a == "e2m1" else wmma_layout, k_width=16)
        b_operand_layout: ttgl.constexpr = ttgl.DotOperandLayout(
            operand_index=1, parent=wmma_layout_packed if type_b == "e2m1" else wmma_layout, k_width=16)
        a_scale_layout: ttgl.constexpr = get_wmma_scale_layout(a_operand_layout, [BLOCK_M, SCALE_BLOCK_K], SCALE_FACTOR)
        b_scale_layout: ttgl.constexpr = get_wmma_scale_layout(b_operand_layout, [BLOCK_N, SCALE_BLOCK_K], SCALE_FACTOR)

        zero = ttgl.zeros([BLOCK_M, BLOCK_N], dtype=ttgl.float32, layout=wmma_layout)

        a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_base, shape=(BLOCK_M, PACKED_BLOCK_K_A),
                                                           strides=(stride_am, stride_ak),
                                                           block_shape=(BLOCK_M, PACKED_BLOCK_K_A),
                                                           layout=SHARED_LAYOUT_A)
        a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, shape=a_desc.block_shape, layout=a_desc.layout)
        ttgl.amd.cdna5.tdm.async_load(a_desc, [0, 0], a_buffer)
        ttgl.amd.cdna5.tdm.async_wait(0)
        a = a_buffer.load(layout=a_layout)
        a = ttgl.convert_layout(a, a_operand_layout)

        b_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=b_base, shape=(PACKED_BLOCK_K_B, BLOCK_N),
                                                           strides=(stride_bk, stride_bn),
                                                           block_shape=(PACKED_BLOCK_K_B, BLOCK_N),
                                                           layout=SHARED_LAYOUT_B)
        b_buffer = ttgl.allocate_shared_memory(b_desc.dtype, shape=b_desc.block_shape, layout=b_desc.layout)
        ttgl.amd.cdna5.tdm.async_load(b_desc, [0, 0], b_buffer)
        ttgl.amd.cdna5.tdm.async_wait(0)
        b = b_buffer.load(layout=b_layout)
        b = ttgl.convert_layout(b, b_operand_layout)

        if a_scale is not None:
            offs_scale_am = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, scale_blocked_layout))
            off_scale_ak = ttgl.arange(0, SCALE_BLOCK_K, layout=ttgl.SliceLayout(0, scale_blocked_layout))
            a_scale_offsets = offs_scale_am[:, None] * SCALE_BLOCK_K + off_scale_ak[None, :]
            scale_a = ttgl.load(a_scale + a_scale_offsets)
            scale_a = ttgl.convert_layout(scale_a, a_scale_layout)
        else:
            scale_a = None

        if b_scale is not None:
            offs_scale_bn = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(1, scale_blocked_layout))
            offs_scale_bk = ttgl.arange(0, SCALE_BLOCK_K, layout=ttgl.SliceLayout(0, scale_blocked_layout))
            b_scale_offsets = offs_scale_bn[:, None] * SCALE_BLOCK_K + offs_scale_bk[None, :]
            scale_b = ttgl.load(b_scale + b_scale_offsets)
            scale_b = ttgl.convert_layout(scale_b, b_scale_layout)
        else:
            scale_b = None

        c = ttgl.amd.cdna5.wmma_scaled(a, scale_a, type_a, b, scale_b, type_b, zero)
        c = c.to(out.dtype.element_ty)

        offs_cm = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, wmma_layout))
        offs_cn = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, wmma_layout))
        out_offsets = offs_cm[:, None] * BLOCK_N + offs_cn[None, :]
        out = out + out_offsets
        ttgl.store(out, c)

    torch.manual_seed(0)

    type_a = mxfp_type
    type_b = mxfp_type

    DIV_FACTOR_A = 2 if type_a == "e2m1" else 1
    DIV_FACTOR_B = 2 if type_b == "e2m1" else 1

    x = torch.randint(20, 40, (M, K // DIV_FACTOR_A), dtype=torch.uint8).cuda()
    y = torch.randint(20, 40, (K // DIV_FACTOR_B, N), dtype=torch.uint8).cuda()

    if hasScale:
        min_scale, max_scale = (0, 142)
        if scale_dtype == "e8m0":
            scale_x = torch.randint(min_scale, max_scale + 1, (M, K // scale_factor), dtype=torch.uint8).cuda()
            scale_y = torch.randint(min_scale, max_scale + 1, (N, K // scale_factor), dtype=torch.uint8).cuda()
        elif scale_dtype == "e4m3":
            scale_x = torch.randint(20, 40, [M, K // scale_factor], dtype=torch.uint8).view(torch.float8_e4m3fn).cuda()
            scale_y = torch.randint(20, 40, [N, K // scale_factor], dtype=torch.uint8).view(torch.float8_e4m3fn).cuda()
    else:
        scale_x = None
        scale_y = None

    def make_finite(x, dtype):
        if dtype not in ("e5m2", "e4m3"):
            return x
        mask = 0x7C if dtype == "e5m2" else 0x7F
        finite = torch.arange(x.numel(), dtype=torch.uint8).cuda().reshape_as(x) % mask
        x_finite = torch.where(x & mask == mask, finite | (0x80 & x), x)
        x.copy_(x_finite)
        return x

    x = make_finite(x, type_a)
    y = make_finite(y, type_b)

    z = torch.zeros((M, N), dtype=torch.float32).cuda()
    pgm = scaled_wmma_tdm_gluon_kernel[(1, )](x, *x.stride(), scale_x, y, *y.stride(), scale_y, z, M, N, K, type_a,
                                              type_b, scale_factor)
    amdgcn = pgm.asm["amdgcn"]

    patterns = (
        "tensor_load_to_lds",
        "s_wait_tensorcnt 0x0",
    )
    for pattern in patterns:
        assert re.search(pattern, amdgcn), f"Can't find {pattern} in amdgcn"

    z_ref = torch.zeros((M, N), dtype=torch.float32).cuda()
    scaled_wmma_tdm_triton_kernel[(1, )](x, *x.stride(), scale_x, y, *y.stride(), scale_y, z_ref, M, N, K, type_a,
                                         type_b, scale_factor)

    torch.testing.assert_close(z.cpu(), z_ref.cpu(), rtol=1e-5, atol=1e-5)


@gluon.jit
def tensor_async_copy_kernel(a_ptr, b_ptr, M, N,  #
                             BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, NUM_BUFFERS: ttgl.constexpr):
    num_warps: ttgl.constexpr = ttgl.num_warps()
    smem_layout: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [BLOCK_M, BLOCK_N], [1, 0])
    block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [num_warps, 1], [1, 0])

    pid_m = ttgl.program_id(axis=0)
    pid_n = ttgl.program_id(axis=1)

    a_buffer = ttgl.allocate_shared_memory(a_ptr.type.element_ty, [NUM_BUFFERS, BLOCK_M, BLOCK_N], smem_layout)

    idx_m = pid_m * BLOCK_M
    for i in ttgl.static_range(0, NUM_BUFFERS):
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N

        offs_am = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, block_layout))
        offs_an = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, block_layout))
        a_ptrs = a_ptr + offs_am[:, None] * N + offs_an[None, :]
        a_mask = (offs_am[:, None] < M) & (offs_an[None, :] < N)
        ttgl.amd.cdna5.async_copy.global_to_shared(a_buffer.index(i), a_ptrs, a_mask, other=0.0)
        ttgl.amd.cdna5.async_copy.commit_group()

    ttgl.amd.cdna5.async_copy.wait_group(0)

    for i in ttgl.static_range(0, NUM_BUFFERS):
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
        a = a_buffer.index(i).load(layout=block_layout)

        offs_bm = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, block_layout))
        offs_bn = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, block_layout))
        offs_b = (offs_bm[:, None] * N) + offs_bn[None, :]
        b_mask = (offs_bm[:, None] < M) & (offs_bn[None, :] < N)
        ttgl.store(b_ptr + offs_b, a, mask=b_mask)


@gluon.jit
def tensor_device_tdm_copy_kernel(a_ptr, b_ptr, M, N,  #
                                  BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, NUM_BUFFERS: ttgl.constexpr,
                                  USE_PARTITIONED: ttgl.constexpr):
    num_warps: ttgl.constexpr = ttgl.num_warps()
    if USE_PARTITIONED:
        NUM_PARTITIONS: ttgl.constexpr = 2
        NUM_GROUPS: ttgl.constexpr = 1
        inner_layout: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]],
                                                                                 [BLOCK_M // NUM_PARTITIONS, BLOCK_N],
                                                                                 [1, 0])
        smem_layout: ttgl.constexpr = PartitionedSharedLayout(NUM_PARTITIONS, NUM_GROUPS, 0, inner_layout)
    else:
        smem_layout: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [BLOCK_M, BLOCK_N], [1, 0])
    block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [num_warps, 1], [1, 0])

    pid_m = ttgl.program_id(axis=0)
    pid_n = ttgl.program_id(axis=1)

    a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(M, N), strides=(N, 1),
                                                       block_shape=(BLOCK_M, BLOCK_N), layout=smem_layout)
    a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, [NUM_BUFFERS] + a_desc.block_shape, a_desc.layout)

    idx_m = pid_m * BLOCK_M
    for i in ttgl.static_range(0, NUM_BUFFERS):
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
        ttgl.amd.cdna5.tdm.async_load(a_desc, [idx_m, idx_n], a_buffer.index(i))

    ttgl.amd.cdna5.tdm.async_wait(0)

    for i in ttgl.static_range(0, NUM_BUFFERS):
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
        a = a_buffer.index(i).load(layout=block_layout)

        offs_bm = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, block_layout))
        offs_bn = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, block_layout))
        offs_b = (offs_bm[:, None] * N) + offs_bn[None, :]
        b_mask = (offs_bm[:, None] < M) & (offs_bn[None, :] < N)
        ttgl.store(b_ptr + offs_b, a, mask=b_mask)


@gluon.jit
def tensor_host_tdm_copy_kernel(a_desc, b_ptr, M, N,  #
                                NUM_BUFFERS: ttgl.constexpr):
    num_warps: ttgl.constexpr = ttgl.num_warps()
    block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [num_warps, 1], [1, 0])

    pid_m = ttgl.program_id(axis=0)
    pid_n = ttgl.program_id(axis=1)

    a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, [NUM_BUFFERS] + a_desc.block_shape, a_desc.layout)
    BLOCK_M: ttgl.constexpr = a_desc.block_shape[0]
    BLOCK_N: ttgl.constexpr = a_desc.block_shape[1]

    idx_m = pid_m * BLOCK_M
    for i in ttgl.static_range(0, NUM_BUFFERS):
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
        ttgl.amd.cdna5.tdm.async_load(a_desc, [idx_m, idx_n], a_buffer.index(i))

    ttgl.amd.cdna5.tdm.async_wait(0)

    for i in ttgl.static_range(0, NUM_BUFFERS):
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
        a = a_buffer.index(i).load(layout=block_layout)

        offs_bm = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, block_layout))
        offs_bn = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, block_layout))
        offs_b = (offs_bm[:, None] * N) + offs_bn[None, :]
        b_mask = (offs_bm[:, None] < M) & (offs_bn[None, :] < N)
        ttgl.store(b_ptr + offs_b, a, mask=b_mask)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(32, 32), (32, 64), (64, 64), (1, 512), (256, 2)])
@pytest.mark.parametrize("NUM_BUFFERS", [2])
@pytest.mark.parametrize("NUM_WARPS", [4, 8])
@pytest.mark.parametrize("ASYNC_LOAD_TYPE", ["ASYNC_COPY", "DEVICE_TDM", "HOST_TDM"])
def test_compile_tensor_copy(BLOCK_M, BLOCK_N, NUM_BUFFERS, ASYNC_LOAD_TYPE, NUM_WARPS):
    attrs = None
    if ASYNC_LOAD_TYPE == "ASYNC_COPY":
        # AsyncCopy requires >= 32 bits per lane so we have to pass divisibility for arguments
        attrs = {k: [["tt.divisibility", 16]] for k in [(x, ) for x in range(4)]}
        fn = tensor_async_copy_kernel
        signature = {
            "a_ptr": "*fp16",
            "b_ptr": "*fp16",
            "M": "i32",
            "N": "i32",
            "BLOCK_M": "constexpr",
            "BLOCK_N": "constexpr",
            "NUM_BUFFERS": "constexpr",
        }
        constexprs = {"BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "NUM_BUFFERS": NUM_BUFFERS}
    elif "DEVICE_TDM":
        fn = tensor_device_tdm_copy_kernel
        signature = {
            "a_ptr": "*fp16",
            "b_ptr": "*fp16",
            "M": "i32",
            "N": "i32",
            "BLOCK_M": "constexpr",
            "BLOCK_N": "constexpr",
            "NUM_BUFFERS": "constexpr",
            "USE_PARTITIONED": "constexpr",
        }
        constexprs = {"BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "NUM_BUFFERS": NUM_BUFFERS}
    else:
        assert ASYNC_LOAD_TYPE == "HOST_TDM"
        fn = tensor_host_tdm_copy_kernel
        smem_layout = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [BLOCK_M, BLOCK_N], [1, 0])
        signature = {
            "a_desc": f"tensordesc<fp16[{BLOCK_M}, {BLOCK_N}],{smem_layout}>",
            "b_ptr": "*fp16",
            "M": "i32",
            "N": "i32",
            "NUM_BUFFERS": "constexpr",
        }
        constexprs = {"NUM_BUFFERS": NUM_BUFFERS}

    k = triton.compile(
        gluon._runtime.GluonASTSource(fn, signature, constexprs, attrs),
        target=GPUTarget("hip", 'gfx1250', 32),
        options={"num_warps": NUM_WARPS},
    )
    amdgcn = k.asm["amdgcn"]

    if ASYNC_LOAD_TYPE in {"DEVICE_TDM", "HOST_TDM"}:
        pattern = {"tensor_load_to_lds", "s_wait_tensorcnt 0x0"}
    else:
        ASYNC_LOAD_TYPE == "ASYNC_COPY"
        pattern = {"global_load_async_to_lds", "s_wait_asynccnt 0x0"}
    for p in pattern:
        assert re.search(p, amdgcn), f"Can't find {p} in amdgcn"


@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(32, 32), (32, 64), (64, 64), (1, 512), (256, 2)])
@pytest.mark.parametrize("NUM_BUFFERS", [2])
@pytest.mark.parametrize("NUM_WARPS", [4, 8])
@pytest.mark.parametrize("ASYNC_LOAD_TYPE", ["ASYNC_COPY", "DEVICE_TDM", "HOST_TDM", "DEVICE_TDM_PARTITIONED"])
@pytest.mark.parametrize("M,N", [(1024, 1024), (1008, 1008)])
def test_runtime_tensor_copy(M, N, BLOCK_M, BLOCK_N, NUM_BUFFERS, ASYNC_LOAD_TYPE, NUM_WARPS):
    if ASYNC_LOAD_TYPE == "ASYNC_COPY" and any([x % 16 != 0 for x in [M, N]]):
        pytest.skip("AsyncCopy tests need divisibility==16 to get vectorization information")

    # PartitionedSharedLayout requires BLOCK_M divisible by num_partitions * num_groups (2*1=2)
    if ASYNC_LOAD_TYPE == "DEVICE_TDM_PARTITIONED" and BLOCK_M % 2 != 0:
        pytest.skip("PartitionedSharedLayout requires BLOCK_M divisible by num_partitions * num_groups")

    torch.manual_seed(42)
    a = torch.randint(0x0, 0xFFFF, (M, N), dtype=torch.uint16)
    b = torch.zeros_like(a)

    a_device = a.cuda()
    b_device = b.cuda()
    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N * NUM_BUFFERS))
    if ASYNC_LOAD_TYPE == "ASYNC_COPY":
        tensor_async_copy_kernel[grid](a_device, b_device, M, N, BLOCK_M, BLOCK_N, NUM_BUFFERS, num_warps=NUM_WARPS)
    elif ASYNC_LOAD_TYPE == "DEVICE_TDM":
        tensor_device_tdm_copy_kernel[grid](a_device, b_device, M, N, BLOCK_M, BLOCK_N, NUM_BUFFERS,
                                            USE_PARTITIONED=False, num_warps=NUM_WARPS)
    elif ASYNC_LOAD_TYPE == "DEVICE_TDM_PARTITIONED":
        tensor_device_tdm_copy_kernel[grid](a_device, b_device, M, N, BLOCK_M, BLOCK_N, NUM_BUFFERS,
                                            USE_PARTITIONED=True, num_warps=NUM_WARPS)
    else:
        assert ASYNC_LOAD_TYPE == "HOST_TDM"
        smem_layout = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [BLOCK_M, BLOCK_N], [1, 0])
        a_desc = gluon.amd.cdna5.TensorDescriptor.from_tensor(a_device, [BLOCK_M, BLOCK_N], layout=smem_layout)
        tensor_host_tdm_copy_kernel[grid](a_desc, b_device, M, N, NUM_BUFFERS, num_warps=NUM_WARPS)

    b_triton = b_device.cpu()
    assert torch.equal(b_triton, a)


# ---------------------------------------------------------------------------
# Partitioned TDM copy tests
#
# These test the partition-aware TDM lowering path, which adjusts the warp
# distribution so each wave's tile stays within a single LDS partition.
# Key scenarios:
#   - partitionDim=0 vs partitionDim=1
#   - Single TDM instruction (warps >= numLogicalPieces)
#   - Multi-instruction split  (warps <  numLogicalPieces)
#   - Different numPartitions / numGroups combinations
# ---------------------------------------------------------------------------


@gluon.jit
def partitioned_tdm_copy_kernel(a_ptr, b_ptr, M, N,  #
                                BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,  #
                                NUM_PARTITIONS: ttgl.constexpr, NUM_GROUPS: ttgl.constexpr,  #
                                PARTITION_DIM: ttgl.constexpr,  #
                                INSTR_SHAPE_K: ttgl.constexpr = 32):
    """TDM load with PartitionedSharedLayout; LDS read uses WMMA dot layout (transpose path)."""
    num_warps: ttgl.constexpr = ttgl.num_warps()

    if PARTITION_DIM == 0:
        inner_shape_m: ttgl.constexpr = BLOCK_M // (NUM_PARTITIONS * NUM_GROUPS)
        inner_shape_n: ttgl.constexpr = BLOCK_N
    else:
        inner_shape_m: ttgl.constexpr = BLOCK_M
        inner_shape_n: ttgl.constexpr = BLOCK_N // (NUM_PARTITIONS * NUM_GROUPS)

    inner_layout: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [inner_shape_m, inner_shape_n],
                                                                             [1, 0])
    smem_layout: ttgl.constexpr = PartitionedSharedLayout(NUM_PARTITIONS, NUM_GROUPS, PARTITION_DIM, inner_layout)

    block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [num_warps, 1], [1, 0])
    WMMA_LAYOUT: ttgl.constexpr = ttgl.amd.AMDWMMALayout(3, True, [[0, 1], [1, 0]], [], [16, 16, INSTR_SHAPE_K])
    DOT_RHS_LAYOUT: ttgl.constexpr = ttgl.DotOperandLayout(1, WMMA_LAYOUT, 8)

    pid_m = ttgl.program_id(axis=0)
    pid_n = ttgl.program_id(axis=1)
    idx_m = pid_m * BLOCK_M
    idx_n = pid_n * BLOCK_N

    a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(M, N), strides=(N, 1),
                                                       block_shape=(BLOCK_M, BLOCK_N), layout=smem_layout)
    a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, a_desc.block_shape, a_desc.layout)

    ttgl.amd.cdna5.tdm.async_load(a_desc, [idx_m, idx_n], a_buffer)
    ttgl.amd.cdna5.tdm.async_wait(0)

    a_dot = a_buffer.load(layout=DOT_RHS_LAYOUT)
    a = ttgl.convert_layout(a_dot, block_layout)

    offs_bm = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, block_layout))
    offs_bn = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, block_layout))
    offs_b = (offs_bm[:, None] * N) + offs_bn[None, :]
    b_mask = (offs_bm[:, None] < M) & (offs_bn[None, :] < N)
    ttgl.store(b_ptr + offs_b, a, mask=b_mask)


_PARTITIONED_TDM_PARAMS = [
    # # --- partitionDim = 0 (rows) ---
    # 2 partitions x 1 group = 2 pieces along dim0 -> 4 warps covers it
    (64, 32, 2, 1, 0),
    # 2 partitions x 2 groups = 4 pieces along dim0 -> 4 warps covers it
    (64, 32, 2, 2, 0),
    # 2 partitions x 4 groups = 8 pieces along dim0 -> 4 warps < 8 -> 2 instructions
    (64, 32, 2, 4, 0),
    # 2 partitions x 8 groups = 16 pieces along dim0 -> 4 warps < 16 -> 4 instructions
    (128, 64, 2, 2, 0),
    # --- partitionDim = 1 (cols) ---
    # 2 partitions x 1 group = 2 pieces along dim1
    (32, 64, 2, 1, 1),
    # 2 partitions x 2 groups = 4 pieces along dim1 -> 4 warps covers it
    (64, 64, 2, 2, 1),
    # 2 partitions x 4 groups = 8 pieces along dim1 -> 4 warps < 8 -> 2 instructions
    (64, 128, 2, 4, 1),
    # 2 partitions x 8 groups = 16 pieces along dim1 -> 4 warps < 16 -> 4 instructions
    (64, 256, 2, 8, 1),
    # 2 partitions x 4 groups = 8 pieces along dim1; partition select lands on a kLane basis (lane remap path)
    (64, 64, 2, 4, 1),
]

_PARTITIONED_TDM_DTYPE_CONFIG = {
    # dtype -> (torch_dtype, triton pointer-type string, WMMA instr K, ds_load_tr* regex)
    "f16": (torch.float16, "fp16", 32, r"ds_load_tr16"),
    "i8": (torch.int8, "i8", 64, r"ds_load_tr8"),
}


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("DTYPE", list(_PARTITIONED_TDM_DTYPE_CONFIG.keys()))
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,NUM_PARTITIONS,NUM_GROUPS,PARTITION_DIM", _PARTITIONED_TDM_PARAMS)
def test_runtime_partitioned_tdm_load(BLOCK_M, BLOCK_N, NUM_PARTITIONS, NUM_GROUPS, PARTITION_DIM, DTYPE):
    """Test correctness of TDM async_load + transpose LDS load with PartitionedSharedEncodingAttr."""
    torch_dtype, _, instr_k, _ = _PARTITIONED_TDM_DTYPE_CONFIG[DTYPE]
    M, N = 256, 256
    num_warps = 4

    torch.manual_seed(42)
    a = torch.randint(-128, 128, (M, N)).to(torch_dtype)
    b = torch.zeros_like(a)

    a_device = a.cuda()
    b_device = b.cuda()

    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
    partitioned_tdm_copy_kernel[grid](a_device, b_device, M, N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
                                      NUM_PARTITIONS=NUM_PARTITIONS, NUM_GROUPS=NUM_GROUPS, PARTITION_DIM=PARTITION_DIM,
                                      INSTR_SHAPE_K=instr_k, num_warps=num_warps)

    b_triton = b_device.cpu()
    mismatched = (b_triton != a).sum().item()
    total = a.numel()
    assert mismatched == 0, (f"Mismatch: {mismatched}/{total} ({100*mismatched/total:.1f}%) elements differ. "
                             f"dtype={DTYPE}, partitionDim={PARTITION_DIM}, "
                             f"numPartitions={NUM_PARTITIONS}, numGroups={NUM_GROUPS}")


_LOAD_SHARED_FP4_REPACKED_CASES = [
    pytest.param(
        (128, 32),
        ttgl.BlockedLayout([2, 16], [16, 2], [4, 1], [1, 0]),
        ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0]),
        ttgl.amd.AMDWMMALayout(3, False, [[0, 1], [1, 0]], [], [16, 16, 64]),
        1,
        id="swizzled",
    ),
    pytest.param(
        (32, 128),
        ttgl.BlockedLayout([16, 2], [2, 16], [1, 4], [0, 1]),
        PartitionedSharedLayout(
            2,
            1,
            0,
            ttgl.PaddedSharedLayout.with_identity_for([[512, 16]], [16, 128], [0, 1]),
        ),
        ttgl.amd.AMDWMMALayout(3, True, [[0, 1], [1, 0]], [], [16, 16, 64]),
        0,
        id="partitioned-padded-dim0",
    ),
    pytest.param(
        (128, 32),
        ttgl.BlockedLayout([2, 16], [16, 2], [4, 1], [1, 0]),
        PartitionedSharedLayout(2, 1, 1, ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0])),
        ttgl.amd.AMDWMMALayout(3, True, [[0, 1], [1, 0]], [], [16, 16, 64]),
        1,
        id="partitioned-swizzled-dim1",
    ),
]


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("src_shape,src_layout,shared_layout,wmma_layout,op_idx", _LOAD_SHARED_FP4_REPACKED_CASES)
def test_load_shared_fp4_repacked(src_shape, src_layout, shared_layout, wmma_layout, op_idx):

    @gluon.jit
    def kernel(
        src_ptr,
        dst_ptr,
        src_rows: ttgl.constexpr,
        src_cols: ttgl.constexpr,
        src_layout: ttgl.constexpr,
        shared_layout: ttgl.constexpr,
        wmma_layout: ttgl.constexpr,
        op_idx: ttgl.constexpr,
    ):
        out_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 32], [16, 2], [4, 1], [1, 0])
        dot_layout: ttgl.constexpr = ttgl.DotOperandLayout(op_idx, wmma_layout, 16)

        offs_m = ttgl.arange(0, src_rows, layout=ttgl.SliceLayout(1, src_layout))
        offs_n = ttgl.arange(0, src_cols, layout=ttgl.SliceLayout(0, src_layout))
        value = ttgl.load(src_ptr + offs_m[:, None] * src_cols + offs_n[None, :])
        smem = ttgl.allocate_shared_memory(src_ptr.type.element_ty, [src_rows, src_cols], shared_layout, value)

        value = ttgl.amd.cdna5.load_shared_fp4_repacked(smem, dot_layout)
        value = ttgl.convert_layout(value, out_layout)
        out_m = ttgl.arange(0, 64, layout=ttgl.SliceLayout(1, out_layout))
        out_n = ttgl.arange(0, 64, layout=ttgl.SliceLayout(0, out_layout))
        ttgl.store(dst_ptr + out_m[:, None] * 64 + out_n[None, :], value)

    src = torch.randint(0, 256, src_shape, dtype=torch.uint8)
    if op_idx == 0:
        logical = torch.empty((64, 128), dtype=torch.uint8)
        logical[0::2, :] = src & 0x0F
        logical[1::2, :] = src >> 4
        expected = logical[:, 0::2] | (logical[:, 1::2] << 4)
    else:
        logical = torch.empty((128, 64), dtype=torch.uint8)
        logical[:, 0::2] = src & 0x0F
        logical[:, 1::2] = src >> 4
        expected = logical[0::2, :] | (logical[1::2, :] << 4)

    result = torch.empty((64, 64), dtype=torch.uint8, device="cuda")
    pgm = kernel[(1, )](src.cuda(), result, *src_shape, src_layout, shared_layout, wmma_layout, op_idx, num_warps=4)

    assert "ds_load_tr4_b64" in pgm.asm["amdgcn"]
    torch.testing.assert_close(result.cpu(), expected, rtol=0, atol=0)


@pytest.mark.parametrize("DTYPE", list(_PARTITIONED_TDM_DTYPE_CONFIG.keys()))
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,NUM_PARTITIONS,NUM_GROUPS,PARTITION_DIM", _PARTITIONED_TDM_PARAMS)
def test_compile_partitioned_tdm_transpose_load(BLOCK_M, BLOCK_N, NUM_PARTITIONS, NUM_GROUPS, PARTITION_DIM, DTYPE):
    """Check that LDS load from PartitionedSharedEncodingAttr lowers to ds_load_tr16 when appropriate."""
    _, triton_dtype, instr_k, ds_load_tr_re = _PARTITIONED_TDM_DTYPE_CONFIG[DTYPE]
    num_warps = 4
    ptr_ty = f"*{triton_dtype}"
    signature = {
        "a_ptr": ptr_ty,
        "b_ptr": ptr_ty,
        "M": "i32",
        "N": "i32",
        "BLOCK_M": "constexpr",
        "BLOCK_N": "constexpr",
        "NUM_PARTITIONS": "constexpr",
        "NUM_GROUPS": "constexpr",
        "PARTITION_DIM": "constexpr",
        "INSTR_SHAPE_K": "constexpr",
    }
    constexprs = {
        "BLOCK_M": BLOCK_M,
        "BLOCK_N": BLOCK_N,
        "NUM_PARTITIONS": NUM_PARTITIONS,
        "NUM_GROUPS": NUM_GROUPS,
        "PARTITION_DIM": PARTITION_DIM,
        "INSTR_SHAPE_K": instr_k,
    }
    k = triton.compile(
        src=gluon._runtime.GluonASTSource(partitioned_tdm_copy_kernel, signature, constexprs),
        target=GPUTarget("hip", 'gfx1250', 32),
        options={"num_warps": num_warps},
    )
    assert re.search(ds_load_tr_re, k.asm["amdgcn"]), \
        f"expected transpose LDS loads ({ds_load_tr_re}) in amdgcn for dtype={DTYPE}"


@gluon.jit
def tensor_device_tdm_multi_cta_load_and_store_kernel(a_ptr, b_ptr, M, N,  #
                                                      BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                                      block_layout: ttgl.constexpr, smem_layout: ttgl.constexpr,
                                                      USE_TDM_LOAD: ttgl.constexpr, USE_TDM_STORE: ttgl.constexpr):
    pid_m = ttgl.program_id(axis=0)
    pid_n = ttgl.program_id(axis=1)
    idx_m = pid_m * BLOCK_M
    idx_n = pid_n * BLOCK_N

    a_buffer = ttgl.allocate_shared_memory(a_ptr.type.element_ty, (BLOCK_M, BLOCK_N), smem_layout)

    # Load data - either using TDM load or async_copy
    if USE_TDM_LOAD:
        a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(M, N), strides=(N, 1),
                                                           block_shape=(BLOCK_M, BLOCK_N), layout=smem_layout)
        ttgl.amd.cdna5.tdm.async_load(a_desc, [idx_m, idx_n], a_buffer)
        ttgl.amd.cdna5.tdm.async_wait(0)
    else:
        offs_am = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, block_layout))
        offs_an = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, block_layout))
        offs_a = (offs_am[:, None] * N) + offs_an[None, :]
        a_mask = (offs_am[:, None] < M) & (offs_an[None, :] < N)
        a_ptrs = a_ptr + offs_a
        ttgl.amd.cdna5.async_copy.global_to_shared(a_buffer, a_ptrs, a_mask)
        ttgl.amd.cdna5.async_copy.commit_group()
        ttgl.amd.cdna5.async_copy.wait_group(0)

    # Store data - either using TDM store or local_load + store
    if USE_TDM_STORE:
        b_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr, shape=(M, N), strides=(N, 1),
                                                           block_shape=(BLOCK_M, BLOCK_N), layout=smem_layout)
        ttgl.amd.cdna5.tdm.async_store(b_desc, [idx_m, idx_n], a_buffer)
        ttgl.amd.cdna5.tdm.async_wait(0)
    else:
        a = a_buffer.load(layout=block_layout)
        offs_bm = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, block_layout))
        offs_bn = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, block_layout))
        offs_b = (offs_bm[:, None] * N) + offs_bn[None, :]
        b_mask = (offs_bm[:, None] < M) & (offs_bn[None, :] < N)
        ttgl.store(b_ptr + offs_b, a, mask=b_mask)


@pytest.mark.parametrize("USE_TDM_LOAD, USE_TDM_STORE", [(True, False), (False, True)])
@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(64, 64), (128, 64)])
@pytest.mark.parametrize("NUM_WARPS", [1, 4])
@pytest.mark.parametrize("M,N", [(64, 64), (576, 576)])
@pytest.mark.parametrize("CGALayout", [
    [[0, 0]],
    [[0, 1]],
    [[1, 0]],
    [[0, 1], [0, 0], [1, 0], [0, 0]],
    [[0, 1], [0, 2], [0, 0], [0, 0]],
    [[1, 0], [2, 0], [4, 0], [0, 0]],
])
def test_runtime_tensor_load_and_store_multi_cta(M, N, BLOCK_M, BLOCK_N, NUM_WARPS, CGALayout, USE_TDM_LOAD,
                                                 USE_TDM_STORE):
    torch.manual_seed(42)
    # Overallocate to catch out of bounds writes, the padding of our input is random and the pading of the output is 0
    # to catch out of bounds writes.
    overallocation_factor = 4
    a = torch.randint(0x0, 0xFFFF, (M * overallocation_factor, N), dtype=torch.uint16)
    out = torch.zeros((M * overallocation_factor, N), dtype=a.dtype)

    a_device = a.cuda()
    out_device = out.cuda()
    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
    num_ctas = 2**len(CGALayout)
    smem_layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0], CGALayout)
    block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [NUM_WARPS, 1], [1, 0], CGALayout)
    tensor_device_tdm_multi_cta_load_and_store_kernel[grid](a_device, out_device, M, N, BLOCK_M, BLOCK_N, block_layout,
                                                            smem_layout, USE_TDM_LOAD, USE_TDM_STORE,
                                                            num_warps=NUM_WARPS, num_ctas=num_ctas)

    out_ref = out
    out_ref[:M, :N] = a[:M, :N]
    assert torch.equal(out_ref, out_device.cpu())


@gluon.jit
def tensor_fill_kernel(a_ptr, M, N, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, NUM_BUFFERS: ttgl.constexpr):
    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0])
    BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0])

    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(M, N), strides=(N, 1),
                                                       block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
    a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, [NUM_BUFFERS] + a_desc.block_shape, a_desc.layout)

    idx_m = pid_m * BLOCK_M
    for i in ttgl.static_range(0, NUM_BUFFERS):
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
        vm = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
        vn = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
        v = (vm[:, None] * N) + vn[None, :]
        v = v.to(a_desc.dtype)
        a_buffer.index(i).store(v)

    for i in ttgl.static_range(0, NUM_BUFFERS):
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
        ttgl.amd.cdna5.tdm.async_store(a_desc, [idx_m, idx_n], a_buffer.index(i))

    ttgl.amd.cdna5.tdm.async_wait(0)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(32, 32), (32, 64), (64, 64)])
@pytest.mark.parametrize("NUM_BUFFERS", [1, 2])
def test_compile_tensor_fill(BLOCK_M, BLOCK_N, NUM_BUFFERS):
    k = triton.compile(
        gluon._runtime.GluonASTSource(
            fn=tensor_fill_kernel, signature={
                "a_ptr": "*fp16", "M": "i32", "N": "i32",  #
                "BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "NUM_BUFFERS": "constexpr"
            }, constexprs={"BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "NUM_BUFFERS": NUM_BUFFERS}),
        target=GPUTarget("hip", 'gfx1250', 32))

    amdgcn = k.asm["amdgcn"]

    for pattern in ("tensor_store_from_lds", "s_wait_tensorcnt 0x0"):
        assert re.search(pattern, amdgcn)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(32, 32), (32, 64), (64, 64)])
@pytest.mark.parametrize("NUM_BUFFERS", [1, 2])
@pytest.mark.parametrize("M,N", [(1024, 1024), (1000, 1000)])
def test_runtime_tensor_fill(M, N, BLOCK_M, BLOCK_N, NUM_BUFFERS):
    a = torch.zeros((M, N), dtype=torch.uint16)

    a_device = a.cuda()
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N * NUM_BUFFERS), 1)
    tensor_fill_kernel[grid](a_device, M, N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, NUM_BUFFERS=NUM_BUFFERS)

    a_triton = a_device.cpu()
    a_ref = torch.arange(M, dtype=torch.int16).unsqueeze(1) * N + \
            torch.arange(N, dtype=torch.int16).unsqueeze(0)
    a_ref = a_ref.to(torch.uint16)
    assert torch.equal(a_triton, a_ref)


@gluon.jit
def tensor_descriptor_load_store_nd_kernel_device_tdm(out_ptr, a_ptr, shape, strides, BLOCK_SHAPE, out_shape,
                                                      out_strides, SHARED_LAYOUT: ttgl.constexpr):
    ndim: ttgl.constexpr = len(BLOCK_SHAPE)
    desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=shape, strides=strides, block_shape=BLOCK_SHAPE,
                                                     layout=SHARED_LAYOUT)

    offs = (0, ) * ndim
    block_shared = ttgl.allocate_shared_memory(desc.dtype, shape=desc.block_shape, layout=desc.layout)
    ttgl.amd.cdna5.tdm.async_load(desc, offs, block_shared)
    ttgl.amd.cdna5.tdm.async_wait(0)

    out_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=out_ptr, shape=out_shape, strides=out_strides,
                                                         block_shape=BLOCK_SHAPE, layout=SHARED_LAYOUT)

    ttgl.amd.cdna5.tdm.async_store(out_desc, offs, block_shared)
    ttgl.amd.cdna5.tdm.async_wait(0)


@gluon.jit
def tensor_descriptor_load_store_nd_kernel_host_tdm(out_desc, inp_desc):
    ndim: ttgl.constexpr = len(inp_desc.block_shape)
    offs = (0, ) * ndim
    block_shared = ttgl.allocate_shared_memory(inp_desc.dtype, shape=inp_desc.block_shape, layout=inp_desc.layout)
    ttgl.amd.cdna5.tdm.async_load(inp_desc, offs, block_shared)
    ttgl.amd.cdna5.tdm.async_wait(0)

    ttgl.amd.cdna5.tdm.async_store(out_desc, offs, block_shared)
    ttgl.amd.cdna5.tdm.async_wait(0)


def _run_tensor_descriptor_load_store_test(dtype_str, ndim, INNER_BLOCK, TDM_TYPE, SHARED_LAYOUT):
    """Utility function to run TDM load/store tests with a given shared layout."""
    alloc_shape = [1, 1, 3, 7, INNER_BLOCK][-ndim:]

    BLOCK_SHAPE = (2, 2, 4, 8, INNER_BLOCK)[-ndim:]
    inp = to_triton(numpy_random(alloc_shape, dtype_str), device="cpu", dst_type=dtype_str)
    inp.data = inp.data[..., :INNER_BLOCK - 3]
    out = inp.new_empty(BLOCK_SHAPE)

    inp = inp.cuda()
    out = out.cuda()

    if TDM_TYPE == "DEVICE_TDM":
        constexpr_block_shape = tuple(ttgl.constexpr(v) for v in BLOCK_SHAPE)
        k = tensor_descriptor_load_store_nd_kernel_device_tdm[(1, )](out, inp, inp.shape,
                                                                     inp.stride(), constexpr_block_shape, out.shape,
                                                                     out.stride(), SHARED_LAYOUT)
    else:
        assert TDM_TYPE == "HOST_TDM"
        inp_desc = gluon.amd.cdna5.TensorDescriptor.from_tensor(inp, list(BLOCK_SHAPE), layout=SHARED_LAYOUT)
        out_desc = gluon.amd.cdna5.TensorDescriptor.from_tensor(out, list(BLOCK_SHAPE), layout=SHARED_LAYOUT)
        k = tensor_descriptor_load_store_nd_kernel_host_tdm[(1, )](out_desc, inp_desc)

    amdgcn = k.asm["amdgcn"]
    for pattern in ("tensor_load_to_lds", "tensor_store_from_lds", "s_wait_tensorcnt 0x0"):
        assert re.search(pattern, amdgcn)

    # Check in-bounds
    actual = unwrap_tensor(out.cpu())
    expect = unwrap_tensor(inp.cpu())
    idx = tuple(slice(None, s) for s in inp.shape)
    assert torch.equal(expect, actual[idx])

    # Check out-of-bounds
    actual[idx].zero_()
    expect = expect.new_zeros(BLOCK_SHAPE)
    assert torch.equal(expect, actual)


@pytest.mark.parametrize("ndim", [1, 2, 3, 4, 5])
@pytest.mark.parametrize("INNER_BLOCK", [4, 8, 16, 32, 64, 128])
@pytest.mark.parametrize("dtype_str", sorted(set(float_dtypes + int_dtypes)))
@pytest.mark.parametrize("TDM_TYPE", ["DEVICE_TDM", "HOST_TDM"])
def test_tensor_descriptor_load_store_nd(dtype_str, ndim, INNER_BLOCK, TDM_TYPE):
    """Test TDM load/store with swizzled shared layout."""

    order = [ndim - 1 - i for i in range(ndim)]
    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=order)

    _run_tensor_descriptor_load_store_test(dtype_str, ndim, INNER_BLOCK, TDM_TYPE, SHARED_LAYOUT)


@pytest.mark.parametrize("ndim", [1, 2, 3, 4, 5])
@pytest.mark.parametrize("INNER_BLOCK", [16, 64])
@pytest.mark.parametrize("dtype_str", ["float16", "int32"])
def test_tensor_descriptor_load_store_nd_with_padding(dtype_str, ndim, INNER_BLOCK):
    """Test TDM load/store with padded shared memory layout.
    TDM store only supports padding when:
    1. There is a single padding interval
    2. The padding interval equals the innermost block dimension
    """
    # Create padded shared layout where padding interval = innermost block dimension
    BLOCK_SHAPE = (2, 2, 4, 8, INNER_BLOCK)[-ndim:]
    order = [ndim - 1 - i for i in range(ndim)]
    padding = [INNER_BLOCK, 8]
    PADDED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([padding], BLOCK_SHAPE, order)

    _run_tensor_descriptor_load_store_test(dtype_str, ndim, INNER_BLOCK, "DEVICE_TDM", PADDED_LAYOUT)


def test_tensor_descriptor_load_store_invalid_blocksize():
    """Test that TDM operations fail when block size exceeds 2^16 (65536)"""
    ndim = 2
    INNER_BLOCK = 2**17  # 131072, exceeds 2^16 limit
    dtype_str = 'float32'

    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1,
                                                              order=[ndim - 1 - i for i in range(ndim)])

    alloc_shape = [7, INNER_BLOCK]
    BLOCK_SHAPE = (8, INNER_BLOCK)

    inp = to_triton(numpy_random(alloc_shape, dtype_str), device="cpu", dst_type=dtype_str)
    inp.data = inp.data[..., :INNER_BLOCK - 3]
    out = inp.new_empty(BLOCK_SHAPE)
    inp = inp.cuda()
    out = out.cuda()

    constexpr_block_shape = tuple(ttgl.constexpr(v) for v in BLOCK_SHAPE)

    # Expect compilation to fail due to block size exceeding maximum
    try:
        tensor_descriptor_load_store_nd_kernel_device_tdm[(1, )](out, inp, inp.shape,
                                                                 inp.stride(), constexpr_block_shape, out.shape,
                                                                 out.stride(), SHARED_LAYOUT)
        pytest.fail(
            f"Expected compilation to fail for block size {INNER_BLOCK} (2^17) > 65536 (2^16), but it succeeded")
    except Exception as e:
        error_msg = str(e)
        assert "error encountered during parsing" in error_msg.lower(), \
            f"Expected parsing error for block size > 65536, but got: {error_msg}"


@gluon.jit
def tensor_descriptor_prefetch_nd_kernel_device_tdm(a_ptr, shape, strides, BLOCK_SHAPE, SHARED_LAYOUT: ttgl.constexpr,
                                                    PREFETCH_SPECULATIVE: ttgl.constexpr):
    ndim: ttgl.constexpr = len(BLOCK_SHAPE)
    desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=shape, strides=strides, block_shape=BLOCK_SHAPE,
                                                     layout=SHARED_LAYOUT)
    offs = (0, ) * ndim
    ttgl.amd.cdna5.tdm.prefetch(desc, offs, speculative=PREFETCH_SPECULATIVE)


@gluon.jit
def tensor_descriptor_prefetch_nd_kernel_host_tdm(inp_desc, SPECULATIVE: ttgl.constexpr):
    ndim: ttgl.constexpr = len(inp_desc.block_shape)
    offs = (0, ) * ndim
    ttgl.amd.cdna5.tdm.prefetch(inp_desc, offs, speculative=SPECULATIVE)


@pytest.mark.parametrize("ndim", [1, 2, 3, 4, 5])
@pytest.mark.parametrize("INNER_BLOCK", [8, 256])
@pytest.mark.parametrize("dtype", ["i8", "fp16", "fp32", "fp64"])
@pytest.mark.parametrize("SPECULATIVE", [True, False])
@pytest.mark.parametrize("TDM_TYPE", ["DEVICE_TDM", "HOST_TDM"])
def test_compile_tensor_descriptor_prefetch_nd(dtype, ndim, INNER_BLOCK, SPECULATIVE, TDM_TYPE):
    SHARED_LAYOUT = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1,
                                              order=[ndim - 1 - i for i in range(ndim)])

    BLOCK_SHAPE = (2, 2, 4, 8, INNER_BLOCK)[-ndim:]
    STRIDES = (1, 2, 4, 8, 16)[:ndim][::-1]
    shape_str = ", ".join(str(s) for s in BLOCK_SHAPE)

    if TDM_TYPE == "DEVICE_TDM":
        fn = tensor_descriptor_prefetch_nd_kernel_device_tdm
        signature = {
            "a_ptr": f"*{dtype}",
            "shape": tuple("i32" for _ in range(ndim)),
            "strides": tuple("constexpr" for _ in range(ndim)),
            "BLOCK_SHAPE": tuple("constexpr" for _ in range(ndim)),
            "SHARED_LAYOUT": "constexpr",
            "PREFETCH_SPECULATIVE": "constexpr",
        }
        constexprs = {
            # For tuples we need to specifiy the parameter index
            **{(2, i): STRIDES[i]
               for i in range(ndim)},
            **{(3, i): BLOCK_SHAPE[i]
               for i in range(ndim)},
            "SHARED_LAYOUT": SHARED_LAYOUT,
            "PREFETCH_SPECULATIVE": SPECULATIVE,
        }
    else:
        assert TDM_TYPE == "HOST_TDM"
        fn = tensor_descriptor_prefetch_nd_kernel_host_tdm
        signature = {
            "inp_desc": f"tensordesc<{dtype}[{shape_str}],{SHARED_LAYOUT}>",
            "SPECULATIVE": "constexpr",
        }
        constexprs = {"SPECULATIVE": SPECULATIVE}

    k = triton.compile(
        gluon._runtime.GluonASTSource(fn, signature, constexprs),
        target=GPUTarget("hip", 'gfx1250', 32),
    )
    amdgcn = k.asm["amdgcn"]

    for pattern in ("global_prefetch_b8", "scope:SCOPE_SE"):
        assert re.search(pattern, amdgcn)
    if not SPECULATIVE:
        assert re.search("th:TH_LOAD_NT", amdgcn)


@pytest.mark.parametrize("ndim", [1, 2, 3, 4, 5])
@pytest.mark.parametrize("INNER_BLOCK", [8, 128, 256])
@pytest.mark.parametrize("dtype_str", ["int8", "float16", "float32", "float64"])
@pytest.mark.parametrize("SPECULATIVE", [True, False])
@pytest.mark.parametrize("TDM_TYPE", ["DEVICE_TDM", "HOST_TDM"])
def test_runtime_tensor_descriptor_prefetch_nd(dtype_str, ndim, INNER_BLOCK, SPECULATIVE, TDM_TYPE):
    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1,
                                                              order=[ndim - 1 - i for i in range(ndim)])

    alloc_shape = [1, 1, 3, 7, INNER_BLOCK][-ndim:]

    BLOCK_SHAPE = (2, 2, 4, 8, INNER_BLOCK)[-ndim:]
    inp = to_triton(numpy_random(alloc_shape, dtype_str), device="cpu", dst_type=dtype_str)
    inp.data = inp.data[..., :INNER_BLOCK - 3]
    # uint_dtypes require special handling because PyTorch only has full native support
    # for uint8. While PyTorch 2.1+ added limited support for uint16, uint32, and uint64,
    # they still lack complete functionality across all PyTorch ops. They are stored as
    # signed tensors with the same bit width and wrapped in TensorWrapper for reinterpretation
    # to unsigned. The .base attribute accesses the underlying signed tensor for CUDA transfer.
    if dtype_str in uint_dtypes:
        inp.base = inp.base.cuda()
    else:
        inp = inp.cuda()

    if TDM_TYPE == "DEVICE_TDM":
        constexpr_block_shape = tuple(ttgl.constexpr(v) for v in BLOCK_SHAPE)
        tensor_descriptor_prefetch_nd_kernel_device_tdm[(1, )](inp, inp.shape, inp.stride(), constexpr_block_shape,
                                                               SHARED_LAYOUT, SPECULATIVE)
    else:
        assert TDM_TYPE == "HOST_TDM"
        inp_desc = gluon.amd.cdna5.TensorDescriptor.from_tensor(inp, list(BLOCK_SHAPE), layout=SHARED_LAYOUT)
        tensor_descriptor_prefetch_nd_kernel_host_tdm[(1, )](inp_desc, SPECULATIVE)


@gluon.jit
def tdm_prefetch_store_back_offsets_kernel(inp_ptr, out_ptr, shape, inp_strides, out_strides, block_shape,
                                           prefetch_block_shape: ttgl.constexpr, rank: ttgl.constexpr):
    pid = (ttgl.program_id(0), ttgl.program_id(1), ttgl.program_id(2))

    if rank == 1:
        layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, [0])
    elif rank == 2:
        layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0])
    else:  # rank == 3
        layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, [2, 1, 0])

    # Compute linear index and starting indices for the tensor descriptor.
    linear_idx = pid[0]
    indices = [pid[0] * block_shape[0]]
    if rank > 1:
        linear_idx = linear_idx * ttgl.num_programs(1) + pid[1]
        indices = [pid[0] * block_shape[0], pid[1] * block_shape[1]]
    if rank > 2:
        linear_idx = linear_idx * ttgl.num_programs(2) + pid[2]
        indices = [pid[0] * block_shape[0], pid[1] * block_shape[1], pid[2] * block_shape[2]]

    desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(inp_ptr, shape=shape, strides=inp_strides, block_shape=block_shape,
                                                     layout=layout)
    prefetch_offsets = ttgl.amd.cdna5.tdm._test_prefetch_with_offsets(desc, indices, pred=True, speculative=False)

    out_layout: ttgl.constexpr = prefetch_offsets.type.layout

    # Create pointer offsets based on rank
    if rank == 1:
        offs_0 = ttgl.arange(0, prefetch_block_shape[0], layout=out_layout)
        out_ptrs = out_ptr + linear_idx * out_strides[0] + offs_0 * out_strides[1]
    elif rank == 2:
        offs_0 = ttgl.arange(0, prefetch_block_shape[0], layout=ttgl.SliceLayout(1, out_layout))
        offs_1 = ttgl.arange(0, prefetch_block_shape[1], layout=ttgl.SliceLayout(0, out_layout))
        out_ptrs = ((out_ptr + (linear_idx * out_strides[0])) + (offs_0[:, None]) * out_strides[1] +
                    (offs_1[None, :]) * out_strides[2])
    else:  # rank == 3
        offs_0 = ttgl.arange(0, prefetch_block_shape[0], layout=ttgl.SliceLayout(1, ttgl.SliceLayout(2, out_layout)))
        offs_1 = ttgl.arange(0, prefetch_block_shape[1], layout=ttgl.SliceLayout(0, ttgl.SliceLayout(2, out_layout)))
        offs_2 = ttgl.arange(0, prefetch_block_shape[2], layout=ttgl.SliceLayout(0, ttgl.SliceLayout(1, out_layout)))
        out_ptrs = ((out_ptr + (linear_idx * out_strides[0])) + (offs_0[:, None, None]) * out_strides[1] +
                    (offs_1[None, :, None]) * out_strides[2] + (offs_2[None, None, :]) * out_strides[3])

    ttgl.store(out_ptrs, prefetch_offsets)


@pytest.mark.parametrize("shape,block_shape", [
    # 1D
    ((512, ), (512, )),
    ((512, ), (64, )),
    ((512, ), (8, )),
    ((512, ), (1024, )),
    ((522, ), (1024, )),
    ((16, ), (8, )),
    # 2D
    ((16, 16), (16, 16)),
    ((32, 512), (32, 512)),
    ((32, 512), (32, 32)),
    ((32, 512), (32, 1024)),
    ((32, 516), (64, 1024)),
    # 3D
    ((2, 4, 64), (2, 4, 64)),
    ((2, 4, 64), (1, 1, 32)),
    ((2, 4, 64), (4, 8, 128)),
])
def test_tdm_prefetch_offsets(shape, block_shape):
    rank = len(shape)
    grid = tuple(triton.cdiv(shape[i], block_shape[i]) for i in range(rank))

    inp = torch.empty(shape, dtype=torch.int32)
    inp_handle = inp.cuda()

    # Each prefetch loads 256B along the fastest dim; scale that axis accordingly.
    prefetch_byte_width = 256
    elems_per_prefetch = prefetch_byte_width // inp.element_size()
    prefetches_in_fast_dim = max(1, block_shape[-1] // elems_per_prefetch)
    prefetch_block_shape = block_shape[:-1] + (prefetches_in_fast_dim, )

    num_programs = math.prod(grid)
    out_shape = (num_programs, ) + tuple(prefetch_block_shape)
    out = torch.zeros(out_shape, dtype=torch.int64)
    out_handle = out.cuda()

    constexpr_block_shape = tuple(ttgl.constexpr(v) for v in block_shape)
    constexpr_prefetch_block_shape = tuple(ttgl.constexpr(v) for v in prefetch_block_shape)

    tdm_prefetch_store_back_offsets_kernel[grid](inp_handle, out_handle, inp.shape, inp.stride(), out_handle.stride(),
                                                 constexpr_block_shape, constexpr_prefetch_block_shape,
                                                 ttgl.constexpr(rank), num_warps=1, num_ctas=1)

    # Compute reference values for prefetch offsets
    out_ref = torch.zeros(out_shape, dtype=torch.int64)

    # Last dimension steps by prefetch chunk size
    prefetch_strides = inp.stride()[:-1] + (elems_per_prefetch, )

    cta_idx = 0
    # Pad grid and block size to 3D to generalize the loop for 1D - 3D
    grid_3d = (grid + (1, 1))[:3]
    prefetch_block_shape_3d = (tuple(prefetch_block_shape) + (1, 1))[:3]

    # Compute for each CTA it's expected prefetch offsets, see TDMPrefetchOp for more details.
    for pid_x, pid_y, pid_z in product(range(grid_3d[0]), range(grid_3d[1]), range(grid_3d[2])):
        pid = [pid_x, pid_y, pid_z]
        # Compute base offset for the CTA
        base = sum(pid[d] * block_shape[d] * inp.stride()[d] for d in range(rank))

        # Create a flattened view into the nD reference to unify the indexing logic over all dimensions
        cta_ref = out_ref[cta_idx].reshape(-1)
        flat_offset_idx = 0

        for x, y, z in product(range(prefetch_block_shape_3d[0]), range(prefetch_block_shape_3d[1]),
                               range(prefetch_block_shape_3d[2])):
            indices = [x, y, z]
            offset = base + sum(indices[d] * prefetch_strides[d] for d in range(rank))
            # We only mask at the end of the tensor. Rows are allowed to wrap into the next one
            cta_ref[flat_offset_idx] = 0 if offset >= inp.numel() else offset
            flat_offset_idx += 1
        cta_idx += 1

    torch.testing.assert_close(out_handle.cpu(), out_ref)


@gluon.jit
def mxgemm_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: ttgl.constexpr, DTYPE_B: ttgl.constexpr,
                  SCALE_BLOCK: ttgl.constexpr, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                  BLOCK_K: ttgl.constexpr, GROUP_SIZE_M: ttgl.constexpr):
    DIV_FACTOR_A: ttgl.constexpr = 2 if DTYPE_A == "e2m1" else 1
    DIV_FACTOR_B: ttgl.constexpr = 2 if DTYPE_B == "e2m1" else 1
    BLOCK_K_SCALE: ttgl.constexpr = BLOCK_K // SCALE_BLOCK
    BLOCK_K_PACKED_A: ttgl.constexpr = BLOCK_K // DIV_FACTOR_A
    BLOCK_K_PACKED_B: ttgl.constexpr = BLOCK_K // DIV_FACTOR_B

    BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [8, 4], [4, 1], [1, 0])
    A_BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 16], [8, 4], [4, 1], [1, 0])
    B_BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 16], [16, 2], [4, 1], [1, 0])

    WMMA_LAYOUT: ttgl.constexpr = ttgl.amd.AMDWMMALayout(3, transposed=True, warp_bases=[[0, 1], [1, 0]],
                                                         instr_shape=[16, 16, 128])
    WMMA_LAYOUT_PACKED: ttgl.constexpr = ttgl.amd.AMDWMMALayout(3, transposed=True, warp_bases=[[0, 1], [1, 0]],
                                                                instr_shape=[16, 16, 64])

    DOT_LAYOUT_A: ttgl.constexpr = ttgl.DotOperandLayout(
        operand_index=0, parent=WMMA_LAYOUT_PACKED if DTYPE_A == "e2m1" else WMMA_LAYOUT, k_width=16)
    DOT_LAYOUT_B: ttgl.constexpr = ttgl.DotOperandLayout(
        operand_index=1, parent=WMMA_LAYOUT_PACKED if DTYPE_B == "e2m1" else WMMA_LAYOUT, k_width=16)
    A_SCALE_LINEAR_LAYOUT: ttgl.constexpr = get_wmma_scale_layout(DOT_LAYOUT_A, [BLOCK_M, BLOCK_K_SCALE])
    B_SCALE_LINEAR_LAYOUT: ttgl.constexpr = get_wmma_scale_layout(DOT_LAYOUT_B, [BLOCK_N, BLOCK_K_SCALE])

    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    num_pid_n = ttgl.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

    offs_am = (pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, A_BLOCKED_LAYOUT))) % M
    offs_ak = ttgl.arange(0, BLOCK_K_PACKED_A, layout=ttgl.SliceLayout(0, A_BLOCKED_LAYOUT))
    offs_bk = ttgl.arange(0, BLOCK_K_PACKED_B, layout=ttgl.SliceLayout(1, B_BLOCKED_LAYOUT))
    offs_bn = (pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, B_BLOCKED_LAYOUT))) % N

    offs_scale_am = (pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))) % M
    offs_scale_ak = ttgl.arange(0, BLOCK_K_SCALE, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
    offs_scale_bn = (pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))) % N
    offs_scale_bk = ttgl.arange(0, BLOCK_K_SCALE, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))

    a_scale_ptr = a_scale + offs_scale_am[:, None] * stride_scale + offs_scale_ak[None, :]
    b_scale_ptr = b_scale + offs_scale_bn[:, None] * stride_scale + offs_scale_bk[None, :]
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_ak[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_bk[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    accumulator = ttgl.zeros((BLOCK_M, BLOCK_N), dtype=ttgl.float32, layout=WMMA_LAYOUT)
    for k in range(0, ttgl.cdiv(K, BLOCK_K)):
        k_remaining_a = K - k * BLOCK_K_PACKED_A
        k_remaining_b = K - k * BLOCK_K_PACKED_B
        valid_k_a = offs_ak < k_remaining_a
        valid_k_b = offs_bk < k_remaining_b

        scale_a = ttgl.load(a_scale_ptr)
        scale_b = ttgl.load(b_scale_ptr)
        scale_a = ttgl.convert_layout(scale_a, A_SCALE_LINEAR_LAYOUT)
        scale_b = ttgl.convert_layout(scale_b, B_SCALE_LINEAR_LAYOUT)

        a = ttgl.load(a_ptrs, mask=valid_k_a[None, :], other=0.0)
        b = ttgl.load(b_ptrs, mask=valid_k_b[:, None], other=0.0)
        a = ttgl.convert_layout(a, DOT_LAYOUT_A)
        b = ttgl.convert_layout(b, DOT_LAYOUT_B)

        accumulator = ttgl.amd.cdna5.wmma_scaled(a, scale_a, DTYPE_A, b, scale_b, DTYPE_B, accumulator)

        a_ptrs += BLOCK_K_PACKED_A * stride_ak
        b_ptrs += BLOCK_K_PACKED_B * stride_bk

        a_scale_ptr += BLOCK_K_SCALE
        b_scale_ptr += BLOCK_K_SCALE

    offs_cm = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, WMMA_LAYOUT))
    offs_cn = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, WMMA_LAYOUT))
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    ttgl.store(c_ptrs, accumulator, mask=c_mask)


@pytest.mark.parametrize("BLOCK_M, BLOCK_N, BLOCK_K", [(32, 32, 128), (64, 64, 128), (64, 64, 64)])
@pytest.mark.parametrize("DTYPE_A", ["float8_e5m2", "float8_e4m3", "float4"])
@pytest.mark.parametrize("DTYPE_B", ["float8_e5m2", "float8_e4m3", "float4"])
def test_compile_mxgemm(BLOCK_M, BLOCK_N, BLOCK_K, DTYPE_A, DTYPE_B):
    scale_block = 32

    triton_dtype_converter = {'float8_e5m2': "fp8e5", "float8_e4m3": "fp8e4nv", "float4": "u8"}
    dot_scaled_dtype_converter = {'float8_e5m2': "e5m2", "float8_e4m3": "e4m3", "float4": "e2m1"}

    k = triton.compile(
        gluon._runtime.GluonASTSource(
            fn=mxgemm_kernel, signature={
                "a_ptr": f"*{triton_dtype_converter[DTYPE_A]}", "b_ptr": f"*{triton_dtype_converter[DTYPE_B]}", "c_ptr":
                "*fp32", "a_scale": "*u8", "b_scale": "*u8", "M": "i32", "N": "i32", "K": "i32", "stride_am": "i32",
                "stride_ak": "i32", "stride_bk": "i32", "stride_bn": "i32", "stride_cm": "i32", "stride_cn": "i32",
                "stride_scale": "i32", "DTYPE_A": "constexpr", "DTYPE_B": "constexpr", "SCALE_BLOCK": "constexpr",
                "BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "BLOCK_K": "constexpr", "GROUP_SIZE_M": "constexpr"
            }, constexprs={
                "DTYPE_A": dot_scaled_dtype_converter[DTYPE_A], "DTYPE_B": dot_scaled_dtype_converter[DTYPE_B],
                "SCALE_BLOCK": scale_block, "BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "BLOCK_K": BLOCK_K, "GROUP_SIZE_M":
                1
            }), target=GPUTarget("hip", 'gfx1250', 32))

    amdgcn = k.asm["amdgcn"]
    pattern = "v_wmma_scale_f32_16x16x128_f8f6f4"
    assert re.search(pattern, amdgcn), f"Can't find instruction {pattern} in AMDGCN assembly"


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


@pytest.mark.parametrize("M, N, K", [(32, 32, 128), (128, 128, 512), (1, 8192, 512)])
@pytest.mark.parametrize("BLOCK_M, BLOCK_N, BLOCK_K", [(32, 32, 128), (64, 64, 128), (64, 64, 64)])
@pytest.mark.parametrize("DTYPE_A", ["e5m2", "e4m3", "e2m1"])
@pytest.mark.parametrize("DTYPE_B", ["e5m2", "e4m3", "e2m1"])
def test_runtime_mxgemm(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, DTYPE_A, DTYPE_B):
    scale_block = 32

    torch.manual_seed(0)

    a = init_mxfp_data(DTYPE_A, M, K)
    b = init_mxfp_data(DTYPE_B, K, N)
    a_size = (M, (K + scale_block - 1) // scale_block)
    b_size = (N, (K + scale_block - 1) // scale_block)
    a_scale = MXScaleTensor(size=a_size).random(low=1.0, high=32.0)
    b_scale = MXScaleTensor(size=b_size).random(low=1.0, high=32.0)

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

    a_scale = a_scale.data
    b_scale = b_scale.data

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

    c_d = torch.zeros(M, N, dtype=torch.float32).cuda()
    a_d = a.data.contiguous().cuda()
    b_d = b.data.contiguous().cuda()
    a_scale_d = a_scale.cuda()
    b_scale_d = b_scale.cuda()

    stride_am, stride_ak = a_d.stride(0), a_d.stride(1)
    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 = a_scale_d.stride(0)

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

    mxgemm_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_A, DTYPE_B, scale_block, BLOCK_M, BLOCK_N, BLOCK_K,
                        group_size_m, num_warps=4, num_ctas=1)

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


@gluon.jit
def cluster_load_and_write_back_kernel(a_ptr, out_ptr, M, N, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                       blocked_layout: ttgl.constexpr):
    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    offs_m = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, blocked_layout))
    offs_n = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, blocked_layout))

    a_ptrs = a_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

    a = ttgl.load(a_ptrs, mask)

    out_ptrs = out_ptr + offs_m[:, None] * N + offs_n[None, :]
    ttgl.store(out_ptrs, a, mask)


@pytest.mark.parametrize("blocked_layout", [
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 2], order=[1, 0],
                       cga_layout=[[0, 1]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[2, 2], order=[1, 0],
                       cga_layout=[[1, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[4, 8], warps_per_cta=[4, 1], order=[1, 0],
                       cga_layout=[[1, 0], [2, 0], [0, 0], [0, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 2], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[0, 1], [0, 0], [1, 0], [0, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 4], threads_per_warp=[4, 8], warps_per_cta=[2, 2], order=[1, 0],
                       cga_layout=[[1, 0], [2, 0], [0, 0], [0, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 4], order=[1, 0],
                       cga_layout=[[0, 1], [0, 0], [1, 0], [0, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 16], threads_per_warp=[4, 8], warps_per_cta=[2, 2], order=[1, 0],
                       cga_layout=[[0, 1], [0, 2], [0, 4], [0, 0]]),
])
@pytest.mark.parametrize("dtype", [
    # Test from 1 byte -> 8 bytes dtypes
    torch.float64, torch.float32, torch.float16, torch.float8_e4m3fn
])
def test_runtime_cluster_load(blocked_layout, dtype):
    M = 128
    N = 128
    BLOCK_M = 64
    BLOCK_N = 64
    num_ctas = 2**len(blocked_layout.cga_layout)

    if dtype == torch.float8_e4m3fn:
        # range from min normal (0 00001 00) to max normal (0 11110 11)
        a = torch.randint(0x04, 0x7B, (M, N), dtype=torch.uint8).view(dtype)
    else:
        a = torch.rand((M, N), dtype=dtype)
    out = torch.empty_like(a)
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
    num_warps = blocked_layout.warps_per_cta[0] * blocked_layout.warps_per_cta[1]
    out_handle = out.cuda()
    cluster_load_and_write_back_kernel[grid](a.cuda(), out_handle, M, N, BLOCK_M, BLOCK_N, blocked_layout,
                                             num_warps=num_warps, num_ctas=num_ctas)
    out_tri = out_handle.cpu()
    out_ref = a.cpu()
    assert torch.equal(out_tri, out_ref)


@gluon.jit
def async_load_and_write_back_kernel(a_ptr, out_ptr, M, N, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                     blocked_layout: ttgl.constexpr, shared_layout: ttgl.constexpr):
    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    offs_m = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, blocked_layout))
    offs_n = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, blocked_layout))

    a_ptrs = a_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

    buffer = ttgl.allocate_shared_memory(a_ptr.type.element_ty, [BLOCK_M, BLOCK_N], shared_layout)
    ttgl.amd.cdna5.async_copy.global_to_shared(buffer, a_ptrs)
    ttgl.amd.cdna5.async_copy.commit_group()
    ttgl.amd.cdna5.async_copy.wait_group(0)

    res = buffer.load(blocked_layout)

    out_ptrs = out_ptr + offs_m[:, None] * N + offs_n[None, :]
    ttgl.store(out_ptrs, res, mask)


ASYNC_COPY_TEST_PARAM_SIZE = pytest.mark.parametrize("M,N", [(128, 128), (1024, 1024), (1008, 1008)])
# We require the vec size to determine if we can use async_copy (>=4bytes), if it's a coalesced layout just assume 16
ASYNC_COPY_TEST_PARAM_SHARED_LAYOUT = pytest.mark.parametrize("vec_size, shared_layout", [
    (16, ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0])),
    (4, ttgl.SwizzledSharedLayout(4, 2, 4, [1, 0])),
    (8, ttgl.SwizzledSharedLayout(8, 2, 4, [1, 0])),
    (16, ttgl.SwizzledSharedLayout(16, 2, 4, [1, 0])),
    (4, ttgl.PaddedSharedLayout.with_identity_for([[4, 4], [8, 4]], [128, 128], [1, 0])),
    (4,
     ttgl.PaddedSharedLayout([[4, 1]], [[0, 1], [0, 2], [0, 8], [0, 4], [16, 0], [32, 0], [0, 16], [0, 32], [0, 64],
                                        [1, 0], [2, 0], [4, 0], [8, 0], [64, 0]], [], [128, 128])),
    (1, ttgl.SwizzledSharedLayout(1, 1, 1, [0, 1])),
    (1, ttgl.SwizzledSharedLayout(4, 2, 4, [0, 1])),
    (1, ttgl.SwizzledSharedLayout(8, 2, 4, [0, 1])),
    (1, ttgl.SwizzledSharedLayout(16, 2, 4, [0, 1])),
    (1, ttgl.PaddedSharedLayout.with_identity_for([[4, 4]], [128, 128], [0, 1])),
    (1, ttgl.PaddedSharedLayout.with_identity_for([[4, 1]], [128, 128], [0, 1])),
])
ASYNC_COPY_TEST_PARAM_DTYPE = pytest.mark.parametrize("dtype", [
    # Test from 1 byte -> 8 bytes dtypes
    torch.float64, torch.float32, torch.float16, torch.float8_e4m3fn
])


def _test_runtime_async_copy_layouts(M, N, vec_size, shared_layout, dtype, use_mbarrier):
    BLOCK_M = 128
    BLOCK_N = 128

    if dtype == torch.float8_e4m3fn:
        # range from min normal (0 00001 00) to max normal (0 11110 11)
        a = torch.randint(0x04, 0x7B, (M, N), dtype=torch.uint8).view(dtype)
    else:
        a = torch.rand((M, N), dtype=dtype)
    out = torch.empty_like(a)
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
    out_handle = out.cuda()

    if not use_mbarrier:
        blocked_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0])
        run_kernel = lambda: async_load_and_write_back_kernel[grid](a.cuda(), out_handle, M, N, BLOCK_M, BLOCK_N,
                                                                    blocked_layout, shared_layout)
    else:
        run_kernel = lambda: async_copy_mbarrier_kernel[grid](a.cuda(), out_handle, M, N, BLOCK_M, BLOCK_N,
                                                              shared_layout)

    if (vec_size * dtype.itemsize) < 4:
        # If we have less than 4 contiguous bytes we expect to abort compilation
        with pytest.raises(RuntimeError):
            run_kernel()
    else:
        run_kernel()
        out_tri = out_handle.cpu()
        out_ref = a.cpu()
        assert torch.equal(out_tri, out_ref)


@ASYNC_COPY_TEST_PARAM_SIZE
@ASYNC_COPY_TEST_PARAM_SHARED_LAYOUT
@ASYNC_COPY_TEST_PARAM_DTYPE
def test_runtime_async_copy(M, N, vec_size, shared_layout, dtype):
    _test_runtime_async_copy_layouts(M, N, vec_size, shared_layout, dtype, False)


@pytest.mark.parametrize("blocked_layout", [
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[0, 1]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[1, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[0, 1], [0, 2], [0, 0], [0, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[0, 1], [0, 0], [1, 0], [0, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[0, 1], [0, 2], [0, 4], [0, 0]]),
])
def test_runtime_async_copy_layouts_multi_cta(blocked_layout):
    M = 1024
    N = 1024
    BLOCK_M = 128
    BLOCK_N = 128
    num_ctas = 2**len(blocked_layout.cga_layout)

    shared_layout = ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0], blocked_layout.cga_layout)

    a = torch.rand((M, N), dtype=torch.float32)
    out = torch.empty_like(a)
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
    out_handle = out.cuda()
    async_load_and_write_back_kernel[grid](a.cuda(), out_handle, M, N, BLOCK_M, BLOCK_N, blocked_layout, shared_layout,
                                           num_warps=1, num_ctas=num_ctas)
    out_tri = out_handle.cpu()
    out_ref = a.cpu()
    assert torch.equal(out_tri, out_ref)


@gluon.jit
def scaled_wmma_scale_preshuffle(a_base, stride_am, stride_ak, a_scale, b_base, stride_bk, stride_bn, b_scale, out,
                                 stride_scale, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                 BLOCK_K: ttgl.constexpr, type_a: ttgl.constexpr, type_b: ttgl.constexpr,
                                 TRANSPOSED_WMMA: ttgl.constexpr):
    DIV_FACTOR_A: ttgl.constexpr = 2 if type_a == "e2m1" else 1
    DIV_FACTOR_B: ttgl.constexpr = 2 if type_b == "e2m1" else 1
    PACKED_BLOCK_K_A: ttgl.constexpr = BLOCK_K // DIV_FACTOR_A
    PACKED_BLOCK_K_B: ttgl.constexpr = BLOCK_K // DIV_FACTOR_B
    SCALE_BLOCK_K: ttgl.constexpr = BLOCK_K // 32
    SCALE_KWIDTH: ttgl.constexpr = 4 if SCALE_BLOCK_K >= 4 else SCALE_BLOCK_K

    NON_K_PRESHUFFLE_BLOCK_SIZE: ttgl.constexpr = 64

    scale_blocked_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [8, 4], [4, 1], [1, 0])
    a_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 16], [8, 4], [4, 1], [1, 0])
    b_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 16], [16, 2], [4, 1], [1, 0])

    wmma_layout: ttgl.constexpr = ttgl.amd.AMDWMMALayout(version=3, transposed=TRANSPOSED_WMMA, reg_bases=[[0, 1],
                                                                                                           [1, 0]],
                                                         warp_bases=[[0, 2], [2, 0]], instr_shape=[16, 16, 128])
    wmma_layout_packed: ttgl.constexpr = ttgl.amd.AMDWMMALayout(version=3, transposed=TRANSPOSED_WMMA,
                                                                reg_bases=[[0, 1], [1, 0]], warp_bases=[[0, 2], [2, 0]],
                                                                instr_shape=[16, 16, 64])

    operand_a_layout: ttgl.constexpr = ttgl.DotOperandLayout(
        operand_index=0, parent=wmma_layout_packed if type_a == "e2m1" else wmma_layout, k_width=16)
    operand_b_layout: ttgl.constexpr = ttgl.DotOperandLayout(
        operand_index=1, parent=wmma_layout_packed if type_b == "e2m1" else wmma_layout, k_width=16)

    a_scale_linear_layout: ttgl.constexpr = ttgl.amd.cdna5.get_wmma_scale_layout(operand_a_layout,
                                                                                 [BLOCK_M, SCALE_BLOCK_K])
    b_scale_linear_layout: ttgl.constexpr = ttgl.amd.cdna5.get_wmma_scale_layout(operand_b_layout,
                                                                                 [BLOCK_N, SCALE_BLOCK_K])

    zero = ttgl.zeros([BLOCK_M, BLOCK_N], dtype=ttgl.float32, layout=wmma_layout)

    offs_am = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, a_layout))
    offs_ak = ttgl.arange(0, PACKED_BLOCK_K_A, layout=ttgl.SliceLayout(0, a_layout))
    a_offsets = offs_am[:, None] * stride_am + offs_ak[None, :] * stride_ak
    a = ttgl.load(a_base + a_offsets)
    a = ttgl.convert_layout(a, operand_a_layout)

    offs_bk = ttgl.arange(0, PACKED_BLOCK_K_B, layout=ttgl.SliceLayout(1, b_layout))
    offs_bn = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, b_layout))
    b_offsets = offs_bk[:, None] * stride_bk + offs_bn[None, :] * stride_bn
    b = ttgl.load(b_base + b_offsets)
    b = ttgl.convert_layout(b, operand_b_layout)

    offs_scale_am = ttgl.arange(0, BLOCK_M // NON_K_PRESHUFFLE_BLOCK_SIZE,
                                layout=ttgl.SliceLayout(1, scale_blocked_layout))
    off_scale_ak = ttgl.arange(0, SCALE_BLOCK_K * NON_K_PRESHUFFLE_BLOCK_SIZE,
                               layout=ttgl.SliceLayout(0, scale_blocked_layout))
    a_scale_offsets = offs_scale_am[:, None] * stride_scale + off_scale_ak[None, :]
    scale_a = ttgl.load(a_scale + a_scale_offsets)

    offs_scale_bn = ttgl.arange(0, BLOCK_N // NON_K_PRESHUFFLE_BLOCK_SIZE,
                                layout=ttgl.SliceLayout(1, scale_blocked_layout))
    offs_scale_bk = ttgl.arange(0, SCALE_BLOCK_K * NON_K_PRESHUFFLE_BLOCK_SIZE,
                                layout=ttgl.SliceLayout(0, scale_blocked_layout))
    b_scale_offsets = offs_scale_bn[:, None] * stride_scale + offs_scale_bk[None, :]
    scale_b = ttgl.load(b_scale + b_scale_offsets)

    scale_a = scale_a.reshape(BLOCK_M // NON_K_PRESHUFFLE_BLOCK_SIZE, SCALE_BLOCK_K // SCALE_KWIDTH, 16, 4,
                              SCALE_KWIDTH).trans(0, 3, 2, 1, 4).reshape(BLOCK_M, SCALE_BLOCK_K)
    scale_b = scale_b.reshape(BLOCK_N // NON_K_PRESHUFFLE_BLOCK_SIZE, SCALE_BLOCK_K // SCALE_KWIDTH, 16, 4,
                              SCALE_KWIDTH).trans(0, 3, 2, 1, 4).reshape(BLOCK_N, SCALE_BLOCK_K)
    scale_a = ttgl.convert_layout(scale_a, a_scale_linear_layout)
    scale_b = ttgl.convert_layout(scale_b, b_scale_linear_layout)

    c = ttgl.amd.cdna5.wmma_scaled(a, scale_a, type_a, b, scale_b, type_b, zero)
    c = c.to(out.dtype.element_ty)

    offs_cm = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, wmma_layout))
    offs_cn = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, wmma_layout))
    out_offsets = offs_cm[:, None] * BLOCK_N + offs_cn[None, :]
    out = out + out_offsets
    ttgl.store(out, c)


@pytest.mark.parametrize("M, N, K", [(128, 128, 64), (128, 128, 128), (256, 256, 256)])
@pytest.mark.parametrize("type_a", ["e5m2", "e2m1", "e4m3"])
@pytest.mark.parametrize("type_b", ["e5m2", "e2m1", "e4m3"])
@pytest.mark.parametrize("TRANSPOSED_WMMA", [True, False])
def test_compile_wmma_scale_preshuffle(M, N, K, type_a, type_b, TRANSPOSED_WMMA):
    dtype_converter = {'e5m2': "fp8e5", "e4m3": "fp8e4nv", "e2m1": "u8"}

    signature = {
        "a_base": f"*{dtype_converter[type_a]}", "stride_am": "i32", "stride_ak": "i32", "a_scale": "*i8", "b_base":
        f"*{dtype_converter[type_b]}", "stride_bk": "i32", "stride_bn": "i32", "b_scale": "*i8", "out": "*i32",
        "stride_scale": "i32", "BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "BLOCK_K": "constexpr", "type_a":
        "constexpr", "type_b": "constexpr", "TRANSPOSED_WMMA": "constexpr"
    }

    constexprs = {
        "BLOCK_M": M, "BLOCK_N": N, "BLOCK_K": K, "type_a": type_a, "type_b": type_b, "TRANSPOSED_WMMA": TRANSPOSED_WMMA
    }
    k = triton.compile(
        src=gluon._runtime.GluonASTSource(fn=scaled_wmma_scale_preshuffle, signature=signature, constexprs=constexprs),
        target=GPUTarget("hip", 'gfx1250', 32))
    amdgcn = k.asm["amdgcn"]

    instr = "v_wmma_scale_f32_16x16x128_f8f6f4"
    scale_opsel_a = "matrix_a_scale:MATRIX_SCALE_ROW1"
    scale_opsel_b = "matrix_b_scale:MATRIX_SCALE_ROW1"
    for suffix in (scale_opsel_a, scale_opsel_b, f"{scale_opsel_a} {scale_opsel_b}"):
        pattern = f"{instr}.*{suffix}\n"
        assert re.search(pattern, amdgcn), f"Can't find pattern {pattern} in AMDGCN assembly"


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("M, N, K", [(64, 64, 64), (128, 128, 128), (256, 256, 256)])
@pytest.mark.parametrize("type_a", ["e5m2", "e2m1", "e4m3"])
@pytest.mark.parametrize("type_b", ["e5m2", "e2m1", "e4m3"])
@pytest.mark.parametrize("TRANSPOSED_WMMA", [True, False])
def test_runtime_wmma_scale_preshuffle(M, N, K, type_a, type_b, TRANSPOSED_WMMA):

    def pack_scale(x):
        PRESHUFFLE_FACTOR = 64
        NON_K, K_SCALE = x.shape
        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, 16, 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)

    torch.manual_seed(0)

    a = init_mxfp_data(type_a, M, K)
    b = init_mxfp_data(type_b, K, N)
    scale_a_size = (M, (K + 32 - 1) // 32)
    scale_b_size = (N, (K + 32 - 1) // 32)

    scale_a_mxfp4 = MXScaleTensor(size=scale_a_size).random(low=1.0, high=32.0)
    scale_b_mxfp4 = MXScaleTensor(size=scale_b_size).random(low=1.0, high=32.0)

    c_torch = torch_gemm_mxfp(a, b, scale_a_mxfp4, scale_b_mxfp4, 32, M, N, K)

    if type_a == "e2m1":
        a = a.to_packed_tensor(dim=1)

    if type_b == "e2m1":
        b = b.to_packed_tensor(dim=0)
    a = a.data.contiguous().cuda()
    b = b.data.contiguous().cuda()

    scale_a = scale_a_mxfp4.data
    scale_b = scale_b_mxfp4.data

    scale_a = pack_scale(scale_a)
    scale_b = pack_scale(scale_b)

    scale_a = scale_a.cuda()
    scale_b = scale_b.cuda()

    stride_scale = scale_a.stride(0)

    c = torch.zeros((M, N), dtype=torch.float32).cuda()
    scaled_wmma_scale_preshuffle[(1, )](a, *a.stride(), scale_a, b, *b.stride(), scale_b, c, stride_scale, M, N, K,
                                        type_a, type_b, TRANSPOSED_WMMA)

    torch.testing.assert_close(c.cpu(), c_torch, rtol=1e-5, atol=1e-5)


@gluon.jit
def async_copy_mbarrier_kernel(a_ptr, out_ptr, M, N, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                               shared_layout: ttgl.constexpr):

    ASYNC_LOAD_BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0])
    BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [2, 2], [1, 0])
    NUM_WARPS: ttgl.constexpr = 4
    WARP_SIZE: ttgl.constexpr = 32

    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    offs_m = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, ASYNC_LOAD_BLOCKED_LAYOUT))
    offs_n = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, ASYNC_LOAD_BLOCKED_LAYOUT))

    out_offs_m = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
    out_offs_n = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))

    a_ptrs = a_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (out_offs_m[:, None] < M) & (out_offs_n[None, :] < N)

    mbar = ttgl.allocate_shared_memory(ttgl.int64, [1], ttgl.amd.cdna5.mbarrier.MBarrierLayout())
    buffer = ttgl.allocate_shared_memory(a_ptr.type.element_ty, [BLOCK_M, BLOCK_N], shared_layout)
    # NOTE: Setting count = NUM_WARPS * WARP_SIZE * 2 is only for testing purposes, in order to also exercise the ttgl.amd.cdna5.mbarrier.arrive API.
    # In practice, since we know that phase is initialized to 0, we can just set count = NUM_WARPS * WARP_SIZE and call directly ttgl.amd.cdna5.mbarrier.wait(mbar, 0).
    ttgl.amd.cdna5.mbarrier.init(mbar, count=NUM_WARPS * WARP_SIZE * 2)
    ttgl.amd.cdna5.async_copy.global_to_shared(buffer, a_ptrs)
    ttgl.amd.cdna5.async_copy.mbarrier_arrive(mbar)
    prior_phase = ttgl.amd.cdna5.mbarrier.arrive(mbar)
    ttgl.amd.cdna5.mbarrier.wait(mbar, prior_phase)

    res = buffer.load(BLOCKED_LAYOUT)

    out_ptrs = out_ptr + out_offs_m[:, None] * N + out_offs_n[None, :]
    ttgl.store(out_ptrs, res, mask)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(32, 32), (32, 64), (64, 64)])
def test_compile_async_copy_mbarrier(BLOCK_M, BLOCK_N):
    SHARED_LAYOUT = ttgl.SwizzledSharedLayout(8, 2, 4, [1, 0])
    signature = {
        "a_ptr": "*fp16", "out_ptr": "*fp16", "M": "i32", "N": "i32",  #
        "BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "shared_layout": "constexpr"
    }
    constexprs = {"BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "shared_layout": SHARED_LAYOUT}
    # AsyncCopy requires >= 32 bits per lane so we have to pass divisibility for arguments used in pointer arithmetic
    attrs = {k: [["tt.divisibility", 16]] for k in [(x, ) for x in range(4)]}
    k = triton.compile(
        gluon._runtime.GluonASTSource(fn=async_copy_mbarrier_kernel, signature=signature, attrs=attrs,
                                      constexprs=constexprs), target=GPUTarget("hip", 'gfx1250', 32))
    pattern = ("global_load_async_to_lds", "ds_atomic_async_barrier_arrive_b64", "ds_atomic_barrier_arrive_rtn_b64",
               "s_sleep")

    amdgcn = k.asm["amdgcn"]
    for pattern in pattern:
        assert re.search(pattern, amdgcn)

    assert not re.search("s_wait_asynccnt 0x0", amdgcn)


@ASYNC_COPY_TEST_PARAM_SIZE
@ASYNC_COPY_TEST_PARAM_SHARED_LAYOUT
@ASYNC_COPY_TEST_PARAM_DTYPE
def test_runtime_async_copy_mbarrier(M, N, vec_size, shared_layout, dtype):
    _test_runtime_async_copy_layouts(M, N, vec_size, shared_layout, dtype, True)


@gluon.jit
def tensor_async_copy_mbarrier_kernel(a_ptr, b_ptr, M, N,  #
                                      BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, NUM_BUFFERS: ttgl.constexpr,
                                      BLOCKED_LAYOUT: ttgl.constexpr, NUM_WARPS: ttgl.constexpr):
    SHARED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [BLOCK_M, BLOCK_N], [1, 0])
    WARP_SIZE: ttgl.constexpr = 32
    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(M, N), strides=(N, 1),
                                                       block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
    bars = ttgl.allocate_shared_memory(ttgl.int64, [NUM_BUFFERS, 1], ttgl.amd.cdna5.mbarrier.MBarrierLayout())
    a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, [NUM_BUFFERS] + a_desc.block_shape, a_desc.layout)

    # NOTE: barrier count takes into account both warp count (NUM_WARPS which is used for TDM) + thread count (NUM_WARPS * WARP_SIZE which is used for mbarrier.arrive)
    # NOTE: Setting count = NUM_WARPS + NUM_WARPS * WARP_SIZE is only for testing purposes, in order to also exercise the ttgl.amd.cdna5.mbarrier.arrive API.
    # In practice, since we know that phase is initialized to 0, we can just set count = NUM_WARPS and call directly ttgl.amd.cdna5.mbarrier.wait(bars.index(i), 0).

    for i in ttgl.static_range(0, NUM_BUFFERS):
        ttgl.amd.cdna5.mbarrier.init(bars.index(i), count=NUM_WARPS + NUM_WARPS * WARP_SIZE)

    idx_m = pid_m * BLOCK_M
    for i in ttgl.static_range(0, NUM_BUFFERS):
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
        ttgl.amd.cdna5.tdm.async_load(a_desc, [idx_m, idx_n], a_buffer.index(i), mbarrier=bars.index(i))

    for i in ttgl.static_range(0, NUM_BUFFERS):
        prior_phase = ttgl.amd.cdna5.mbarrier.arrive(bars.index(i))
        ttgl.amd.cdna5.mbarrier.wait(bars.index(i), prior_phase)
        idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
        a = a_buffer.index(i).load(layout=BLOCKED_LAYOUT)

        offs_bm = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
        offs_bn = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
        offs_b = (offs_bm[:, None] * N) + offs_bn[None, :]
        mask_b = (offs_bm[:, None] < M) & (offs_bn[None, :] < N)

        ttgl.store(b_ptr + offs_b, a, mask=mask_b)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(32, 32), (32, 64), (64, 64)])
@pytest.mark.parametrize("NUM_BUFFERS", [1, 2])
@pytest.mark.parametrize("NUM_WARPS", [4])
def test_compile_tensor_copy_mbarrier(BLOCK_M, BLOCK_N, NUM_BUFFERS, NUM_WARPS):
    BLOCKED_LAYOUT = ttgl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0])
    signature = {
        "a_ptr": "*fp16", "b_ptr": "*fp16", "M": "i32", "N": "i32",  #
        "BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "NUM_BUFFERS": "constexpr", "BLOCKED_LAYOUT": "constexpr",
        "NUM_WARPS": "constexpr"
    }
    constexprs = {
        "BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "NUM_BUFFERS": NUM_BUFFERS, "BLOCKED_LAYOUT": BLOCKED_LAYOUT,
        "NUM_WARPS": NUM_WARPS
    }
    attrs = []
    k = triton.compile(
        gluon._runtime.GluonASTSource(fn=tensor_async_copy_mbarrier_kernel, signature=signature, attrs=attrs,
                                      constexprs=constexprs), target=GPUTarget("hip", 'gfx1250', 32))
    pattern = ("tensor_load_to_lds", "ds_atomic_barrier_arrive_rtn_b64", "s_sleep")

    amdgcn = k.asm["amdgcn"]
    for pattern in pattern:
        assert re.search(pattern, amdgcn)

    assert not re.search("s_wait_tensorcnt 0x0", amdgcn)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(32, 32), (32, 64), (64, 64), (1, 512), (256, 2)])
@pytest.mark.parametrize("NUM_BUFFERS", [1, 2])
@pytest.mark.parametrize("NUM_WARPS", [4, 8])
@pytest.mark.parametrize("M,N", [(1024, 1024), (1008, 1008), (1000, 1000)])
def test_runtime_tensor_copy_mbarrier(M, N, BLOCK_M, BLOCK_N, NUM_BUFFERS, NUM_WARPS):
    torch.manual_seed(42)
    a = torch.randint(0x0, 0xFFFF, (M, N), dtype=torch.uint16)
    b = torch.zeros_like(a)

    a_device = a.cuda()
    b_device = b.cuda()
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N * NUM_BUFFERS), 1)

    blocked_layout = ttgl.BlockedLayout([1, 8], [4, 8], [NUM_WARPS, 1], [1, 0])

    tensor_async_copy_mbarrier_kernel[grid](a_device, b_device, M, N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
                                            NUM_BUFFERS=NUM_BUFFERS, BLOCKED_LAYOUT=blocked_layout, NUM_WARPS=NUM_WARPS,
                                            num_warps=NUM_WARPS)

    b_triton = b_device.cpu()
    assert torch.equal(b_triton, a)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
def test_tdm_load_pred():

    @gluon.jit
    def kernel(a_ptr, b_ptr):
        shared_layout: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [16, 32], [1, 0])
        reg_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 4], [4, 8], [4, 1], [1, 0])

        desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(16, 64), strides=(64, 1),
                                                         block_shape=(16, 32), layout=shared_layout)
        smem = ttgl.allocate_shared_memory(desc.dtype, shape=desc.block_shape, layout=desc.layout)
        b_offs_m = ttgl.arange(0, 16, layout=ttgl.SliceLayout(1, reg_layout))
        b_offs_n = ttgl.arange(0, 32, layout=ttgl.SliceLayout(0, reg_layout))
        b_ptrs = b_ptr + b_offs_m[:, None] * 64 + b_offs_n[None, :]

        ttgl.amd.cdna5.tdm.async_load(desc, [0, 0], smem, pred=0)
        ttgl.amd.cdna5.tdm.async_wait(0)
        tile1 = smem.load(reg_layout)
        ttgl.store(b_ptrs, tile1)

        ttgl.amd.cdna5.tdm.async_load(desc, [0, 32], smem, pred=1)
        ttgl.amd.cdna5.tdm.async_wait(0)
        tile2 = smem.load(reg_layout)
        ttgl.store(b_ptrs + 32, tile2)

    a = torch.randint(0x0, 0xFFFF, (16, 64), dtype=torch.uint16)
    b = torch.zeros_like(a)

    a_device = a.cuda()
    b_device = b.cuda()
    kernel[(1, )](a_device, b_device)

    b = b_device.cpu()
    assert torch.equal(a[:, 32:], b[:, 32:]) and not torch.equal(a[:, :32], b[:, :32])


# Check that negative TDM offsets are treated as unsigned so they will mask (zero-fill) out the whole tile.
@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
def test_tdm_load_negative_offset():

    @gluon.jit
    def tdm_load_negative_offset_kernel(a_ptr, b_ptr):
        shared_layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])
        reg_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 4], [4, 8], [4, 1], [1, 0])

        desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(16, 64), strides=(64, 1),
                                                         block_shape=(16, 64), layout=shared_layout)
        smem = ttgl.allocate_shared_memory(desc.dtype, shape=desc.block_shape, layout=desc.layout)

        ttgl.amd.cdna5.tdm.async_load(desc, [-1, -1], smem)
        ttgl.amd.cdna5.tdm.async_wait(0)

        b_offs_m = ttgl.arange(0, 16, layout=ttgl.SliceLayout(1, reg_layout))
        b_offs_n = ttgl.arange(0, 64, layout=ttgl.SliceLayout(0, reg_layout))
        b_ptrs = b_ptr + b_offs_m[:, None] * 64 + b_offs_n[None, :]
        tile = smem.load(reg_layout)
        ttgl.store(b_ptrs, tile)

    a = torch.randint(0x0, 0xFFFF, (16, 64), dtype=torch.uint16)
    b_device = torch.randint(0x0, 0xFFFF, (16, 64), dtype=torch.uint16).cuda()
    tdm_load_negative_offset_kernel[(1, )](a.cuda(), b_device)

    assert torch.all(b_device.cpu() == 0)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("XBLOCK", [128])
def test_ws_store_wait_load(XBLOCK):
    """
    Tests warp specialization with mbarrier synchronization on GFX1250.

    This test validates the mbarrier wait/arrive mechanism for synchronizing data flow
    between two specialized warp groups using helper variables ready_bar and done_bar:
    - ws_producer (worker) partition: Stores data to shared memory and signals completion via ready_bar
    - ws_consumer (default) partition: Waits on ready_bar, loads the data, processes it, stores to
      a different shared memory location, and signals completion via done_bar

    The main kernel (executed by default warps) then waits for done_bar, loads the final result, and stores
    it to global memory. The test verifies data integrity by comparing the output with an expected
    arange pattern.
    """

    import os
    if os.environ.get("TRITON_INSTRUMENTATION_MODE") == "consan":
        from triton.runtime._allocation import set_profile_allocator

        def alloc_fn(size: int, alignment: int, stream):
            return torch.empty(size, device="cuda", dtype=torch.int8)

        set_profile_allocator(alloc_fn)

    @gluon.jit
    def ws_consumer(smem, ready_bar, done_bar, layout: ttgl.constexpr):
        ttgl.amd.cdna5.mbarrier.wait(ready_bar, phase=0)
        val = smem.index(0).load(layout)
        smem.index(1).store(val)
        ttgl.amd.cdna5.mbarrier.arrive(done_bar, count=1)

    @gluon.jit
    def ws_producer(smem, ready_bar, XBLOCK: ttgl.constexpr, layout: ttgl.constexpr):
        smem.index(0).store(ttgl.arange(0, XBLOCK, layout).to(ttgl.float16))
        ttgl.amd.cdna5.mbarrier.arrive(ready_bar, count=1)

    @gluon.jit
    def ws_kernel(output, XBLOCK: ttgl.constexpr):
        WARP_SIZE: ttgl.constexpr = 32
        smem_layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[0])
        blocked_layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1], threads_per_warp=[32],
                                                            warps_per_cta=[4], order=[0])
        smem = ttgl.allocate_shared_memory(ttgl.float16, [2, XBLOCK], smem_layout)
        bar = ttgl.allocate_shared_memory(ttgl.int64, [2, 1], ttgl.amd.cdna5.mbarrier.MBarrierLayout())
        for i in range(2):
            # we have 4 default warps and 4 worker warps and arrive on barrier once per thread
            ttgl.amd.cdna5.mbarrier.init(bar.index(i), count=4 * WARP_SIZE)
        ready_bar = bar.index(0)
        done_bar = bar.index(1)
        # NOTE: We have 8 warps in total. worker_num_warps = [4] (num warps for ws_producer partition) and num_warps = 4 (num warps for consumer partition)
        ttgl.warp_specialize([
            (ws_consumer, (smem, ready_bar, done_bar, blocked_layout)),
            (ws_producer, (smem, ready_bar, XBLOCK, blocked_layout)),
        ], [4])
        ttgl.amd.cdna5.mbarrier.wait(done_bar, phase=0)
        val = smem.index(1).load(blocked_layout)
        output_ptrs = output + ttgl.arange(0, XBLOCK, blocked_layout)
        ttgl.store(output_ptrs, val)

    output = torch.empty((XBLOCK, ), dtype=torch.float16).cuda()
    ws_kernel[(1, )](output, XBLOCK=XBLOCK, num_warps=4)
    torch_output = torch.arange(0, XBLOCK, dtype=torch.float16)
    output_ref = output.cpu()
    assert torch.equal(output_ref, torch_output)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("XBLOCK", [128])
@pytest.mark.parametrize("NUM_ITERS", [10])
def test_ws_store_wait_load_loop(XBLOCK, NUM_ITERS):
    """
    Tests warp specialization with mbarrier synchronization in a loop and phase tracking on GFX1250.

    This test validates iterative producer-consumer synchronization using three mbarriers:
    - ready_bar: Signals that the producer has written data to shared memory
    - done_bar: Signals that the consumer has finished all iterations
    - empty_bar: Signals that the consumer has consumed data and buffer is empty

    - ws_producer (worker) partition: Waits for empty_bar, writes data, signals via ready_bar (loops NUM_ITERS times)
    - ws_consumer (default) partition: Waits for ready_bar, reads and accumulates data, signals via empty_bar (loops NUM_ITERS times)

    Both partitions track phases (1-bit parity phase which toggles between 0 for even and 1 for odd). After all iterations, the main kernel
    (executed by default warps) waits for done_bar, loads the accumulated result, and stores it to global memory.
    The test verifies that the output equals the expected arange pattern.
    """

    import os
    if os.environ.get("TRITON_INSTRUMENTATION_MODE") == "consan":
        from triton.runtime._allocation import set_profile_allocator

        def alloc_fn(size: int, alignment: int, stream):
            return torch.empty(size, device="cuda", dtype=torch.int8)

        set_profile_allocator(alloc_fn)

    @gluon.jit
    def ws_consumer(smem, ready_bar, done_bar, empty_bar, XBLOCK: ttgl.constexpr, NUM_ITERS: ttgl.constexpr,
                    layout: ttgl.constexpr):
        acc = ttgl.zeros([XBLOCK], ttgl.float16, layout)
        phase = 0
        for _ in ttgl.static_range(NUM_ITERS):
            ttgl.amd.cdna5.mbarrier.wait(ready_bar, phase=phase)
            phase = phase ^ 1
            val = smem.index(0).load(layout)
            acc += val
            ttgl.amd.cdna5.mbarrier.arrive(empty_bar, count=1)

        smem.index(1).store(acc)
        ttgl.amd.cdna5.mbarrier.arrive(done_bar, count=1)

    @gluon.jit
    def ws_producer(smem, ready_bar, empty_bar, XBLOCK: ttgl.constexpr, NUM_ITERS: ttgl.constexpr,
                    layout: ttgl.constexpr):
        val = ttgl.arange(0, XBLOCK, layout).to(ttgl.float16)
        phase = 0
        for _ in ttgl.static_range(NUM_ITERS):
            ttgl.amd.cdna5.mbarrier.wait(empty_bar, phase=phase)
            phase = phase ^ 1
            smem.index(0).store(val)
            ttgl.amd.cdna5.mbarrier.arrive(ready_bar, count=1)

    @gluon.jit
    def ws_kernel(output, XBLOCK: ttgl.constexpr, NUM_ITERS: ttgl.constexpr):
        WARP_SIZE: ttgl.constexpr = 32
        smem_layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[0])
        blocked_layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1], threads_per_warp=[32],
                                                            warps_per_cta=[4], order=[0])
        smem = ttgl.allocate_shared_memory(ttgl.float16, [2, XBLOCK], smem_layout)
        bar = ttgl.allocate_shared_memory(ttgl.int64, [3, 1], ttgl.amd.cdna5.mbarrier.MBarrierLayout())
        for i in ttgl.static_range(3):
            # we have 4 default warps and 4 worker warps and arrive on barrier once per thread
            ttgl.amd.cdna5.mbarrier.init(bar.index(i), count=4 * WARP_SIZE)
        ready_bar = bar.index(0)
        done_bar = bar.index(1)
        empty_bar = bar.index(2)

        ttgl.amd.cdna5.mbarrier.arrive(empty_bar, count=1)
        # NOTE: We have 8 warps in total. worker_num_warps = [4] (num warps for ws_producer partition) and num_warps = 4 (num warps for consumer partition)
        ttgl.warp_specialize([
            (ws_consumer, (smem, ready_bar, done_bar, empty_bar, XBLOCK, NUM_ITERS, blocked_layout)),
            (ws_producer, (smem, ready_bar, empty_bar, XBLOCK, NUM_ITERS, blocked_layout)),
        ], [4])
        ttgl.amd.cdna5.mbarrier.wait(done_bar, phase=0)
        val = smem.index(1).load(blocked_layout)
        output_ptrs = output + ttgl.arange(0, XBLOCK, blocked_layout)
        ttgl.store(output_ptrs, val)

    output = torch.empty((XBLOCK, ), dtype=torch.float16).cuda()
    ws_kernel[(1, )](output, XBLOCK=XBLOCK, NUM_ITERS=NUM_ITERS, num_warps=4)
    torch_output = NUM_ITERS * torch.arange(0, XBLOCK, dtype=torch.float16)
    output_ref = output.cpu()
    assert torch.equal(output_ref, torch_output)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(32, 32), (32, 64), (64, 64), (1, 512), (256, 2)])
@pytest.mark.parametrize("NUM_BUFFERS", [1, 2])
@pytest.mark.parametrize("NUM_TOTAL_WARPS", [8, 16])
@pytest.mark.parametrize("M,N", [(32, 32), (1024, 1024), (1008, 1008), (1000, 1000)])
def test_runtime_ws_tensor_async_load_store_mbarrier(M, N, BLOCK_M, BLOCK_N, NUM_BUFFERS, NUM_TOTAL_WARPS):
    """
    Tests warp specialization with tensor descriptor async load/store operations coordinated by mbarriers on GFX1250.

    This test validates the producer-consumer pattern using TDM async operations
    with multiple buffers, where each buffer has its own dedicated mbarrier for synchronization:
    - ws_producer (worker) partition: Asynchronously loads data from global memory to shared memory buffers
      using TDM async_load, with each load operation automatically signaling its corresponding mbarrier
    - ws_consumer (default) partition: Waits on each buffer's mbarrier, then asynchronously stores data
      from shared memory to global memory using TDM async_store

    The synchronization pattern uses one mbarrier per buffer (bars.index(i)), ensuring that the consumer
    only accesses a buffer after the producer has completed loading into it.

    The test verifies that the output matches the input, confirming that async load/store operations are correctly coordinated by mbarriers.
    """

    import os
    if os.environ.get("TRITON_INSTRUMENTATION_MODE") == "consan":
        from triton.runtime._allocation import set_profile_allocator

        def alloc_fn(size: int, alignment: int, stream):
            return torch.empty(size, device="cuda", dtype=torch.int8)

        set_profile_allocator(alloc_fn)

    @gluon.jit
    def ws_producer(a_desc, a_buffer, bars, pid_n, idx_m, BLOCK_N: ttgl.constexpr, NUM_BUFFERS: ttgl.constexpr):
        for i in ttgl.static_range(0, NUM_BUFFERS):
            idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
            ttgl.amd.cdna5.tdm.async_load(a_desc, [idx_m, idx_n], a_buffer.index(i), mbarrier=bars.index(i))

    @gluon.jit
    def ws_consumer(b_desc, a_buffer, bars, pid_n, idx_m, BLOCK_N: ttgl.constexpr, NUM_BUFFERS: ttgl.constexpr):
        for i in ttgl.static_range(0, NUM_BUFFERS):
            ttgl.amd.cdna5.mbarrier.wait(bars.index(i), 0)
            idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
            ttgl.amd.cdna5.tdm.async_store(b_desc, [idx_m, idx_n], a_buffer.index(i))

        ttgl.amd.cdna5.tdm.async_wait(0)

    @gluon.jit
    def ws_tensor_async_load_store_mbarrier_kernel(a_ptr, b_ptr, M, N,  #
                                                   BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                                   NUM_BUFFERS: ttgl.constexpr, NUM_WARPS: ttgl.constexpr):
        SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0])
        PRODUCER_WARPS: ttgl.constexpr = NUM_WARPS // 2
        pid = ttgl.program_id(axis=0)
        num_pid_m = ttgl.cdiv(M, BLOCK_M)
        pid_m = pid % num_pid_m
        pid_n = pid // num_pid_m

        a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(M, N), strides=(N, 1),
                                                           block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
        b_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr, shape=(M, N), strides=(N, 1),
                                                           block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
        bars = ttgl.allocate_shared_memory(ttgl.int64, [NUM_BUFFERS, 1], ttgl.amd.cdna5.mbarrier.MBarrierLayout())
        a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, [NUM_BUFFERS] + a_desc.block_shape, a_desc.layout)

        for i in ttgl.static_range(0, NUM_BUFFERS):
            ttgl.amd.cdna5.mbarrier.init(bars.index(i), count=PRODUCER_WARPS)

        idx_m = pid_m * BLOCK_M

        ttgl.warp_specialize([
            (ws_consumer, (b_desc, a_buffer, bars, pid_n, idx_m, BLOCK_N, NUM_BUFFERS)),
            (ws_producer, (a_desc, a_buffer, bars, pid_n, idx_m, BLOCK_N, NUM_BUFFERS)),
        ], [PRODUCER_WARPS])

    torch.manual_seed(42)
    a = torch.randint(0x0, 0xFFFF, (M, N), dtype=torch.uint16)
    b = torch.zeros_like(a)

    a_device = a.cuda()
    b_device = b.cuda()
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N * NUM_BUFFERS), 1)

    ws_tensor_async_load_store_mbarrier_kernel[grid](a_device, b_device, M, N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
                                                     NUM_BUFFERS=NUM_BUFFERS, NUM_WARPS=NUM_TOTAL_WARPS,
                                                     num_warps=NUM_TOTAL_WARPS // 2)

    b_triton = b_device.cpu()
    assert torch.equal(b_triton, a)


@pytest.mark.parametrize("BLOCK_M,BLOCK_N", [(32, 32), (32, 64), (64, 64), (1, 512), (256, 2)])
@pytest.mark.parametrize("NUM_BUFFERS", [1, 2])
@pytest.mark.parametrize("NUM_TOTAL_WARPS", [8, 16])
@pytest.mark.parametrize("M,N", [(32, 32), (1024, 1024), (1008, 1008), (1000, 1000)])
def test_runtime_ws_tensor_copy_mbarrier(M, N, BLOCK_M, BLOCK_N, NUM_BUFFERS, NUM_TOTAL_WARPS):
    """
    Tests warp specialization with mixed async/sync operations coordinated by mbarriers on GFX1250.

    This test validates the producer-consumer pattern using a combination of TDM async loads and
    synchronous stores with multiple buffers, where each buffer has its own dedicated mbarrier:
    - ws_producer (worker) partition: Asynchronously loads data from global memory to shared memory buffers
      using TDM async_load, with each load operation automatically signaling its corresponding mbarrier
    - ws_consumer (default) partition: Waits on each buffer's mbarrier, loads data from shared memory
      into registers using regular loads, then stores to global memory using regular synchronous stores

    The synchronization pattern uses one mbarrier per buffer (bars.index(i)), ensuring that the consumer
    only accesses a buffer after the producer has completed loading into it.

    NOTE: This test showcases that tensors (here: b_ptr) can be passed as arguments to the default partition
    (here: ws_consumer), which is not supported for worker partitions.

    The test verifies that the output matches the input, confirming correct synchronization.
    """

    import os
    if os.environ.get("TRITON_INSTRUMENTATION_MODE") == "consan":
        from triton.runtime._allocation import set_profile_allocator

        def alloc_fn(size: int, alignment: int, stream):
            return torch.empty(size, device="cuda", dtype=torch.int8)

        set_profile_allocator(alloc_fn)

    @gluon.jit
    def ws_producer(a_desc, a_buffer, bars, pid_n, idx_m, BLOCK_N: ttgl.constexpr, NUM_BUFFERS: ttgl.constexpr):
        for i in ttgl.static_range(0, NUM_BUFFERS):
            idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
            ttgl.amd.cdna5.tdm.async_load(a_desc, [idx_m, idx_n], a_buffer.index(i), mbarrier=bars.index(i))

    @gluon.jit
    def ws_consumer(a_buffer, b_ptr, bars, pid_n, idx_m, M, N, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                    NUM_BUFFERS: ttgl.constexpr, BLOCKED_LAYOUT: ttgl.constexpr):
        for i in ttgl.static_range(0, NUM_BUFFERS):
            ttgl.amd.cdna5.mbarrier.wait(bars.index(i), 0)
            idx_n = pid_n * (BLOCK_N * NUM_BUFFERS) + i * BLOCK_N
            a = a_buffer.index(i).load(layout=BLOCKED_LAYOUT)
            offs_bm = idx_m + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
            offs_bn = idx_n + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
            offs_b = (offs_bm[:, None] * N) + offs_bn[None, :]
            mask_b = (offs_bm[:, None] < M) & (offs_bn[None, :] < N)
            ttgl.store(b_ptr + offs_b, a, mask=mask_b)

    @gluon.jit
    def ws_tensor_async_copy_mbarrier_kernel(a_ptr, b_ptr, M, N,  #
                                             BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                             NUM_BUFFERS: ttgl.constexpr, BLOCKED_LAYOUT: ttgl.constexpr,
                                             NUM_WARPS: ttgl.constexpr):
        SHARED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[32, 4]], [BLOCK_M, BLOCK_N], [1, 0])
        PRODUCER_WARPS: ttgl.constexpr = NUM_WARPS // 2
        pid = ttgl.program_id(axis=0)
        num_pid_m = ttgl.cdiv(M, BLOCK_M)
        pid_m = pid % num_pid_m
        pid_n = pid // num_pid_m

        a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(M, N), strides=(N, 1),
                                                           block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
        bars = ttgl.allocate_shared_memory(ttgl.int64, [NUM_BUFFERS, 1], ttgl.amd.cdna5.mbarrier.MBarrierLayout())
        a_buffer = ttgl.allocate_shared_memory(a_desc.dtype, [NUM_BUFFERS] + a_desc.block_shape, a_desc.layout)

        for i in ttgl.static_range(0, NUM_BUFFERS):
            # TDM arrives on barrier once per warp, so use producer warp count
            ttgl.amd.cdna5.mbarrier.init(bars.index(i), count=PRODUCER_WARPS)

        idx_m = pid_m * BLOCK_M

        ttgl.warp_specialize([
            (ws_consumer, (a_buffer, b_ptr, bars, pid_n, idx_m, M, N, BLOCK_M, BLOCK_N, NUM_BUFFERS, BLOCKED_LAYOUT)),
            (ws_producer, (a_desc, a_buffer, bars, pid_n, idx_m, BLOCK_N, NUM_BUFFERS)),
        ], [PRODUCER_WARPS])

    torch.manual_seed(42)
    a = torch.randint(0x0, 0xFFFF, (M, N), dtype=torch.uint16)
    b = torch.zeros_like(a)

    a_device = a.cuda()
    b_device = b.cuda()
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N * NUM_BUFFERS), 1)

    blocked_layout = ttgl.BlockedLayout([1, 8], [4, 8], [NUM_TOTAL_WARPS // 2, 1], [1, 0])

    ws_tensor_async_copy_mbarrier_kernel[grid](a_device, b_device, M, N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
                                               NUM_BUFFERS=NUM_BUFFERS, BLOCKED_LAYOUT=blocked_layout,
                                               NUM_WARPS=NUM_TOTAL_WARPS, num_warps=NUM_TOTAL_WARPS // 2)

    b_triton = b_device.cpu()
    assert torch.equal(b_triton, a)


@pytest.mark.parametrize("M,N", [(128, 128), (1024, 1024)])
@pytest.mark.parametrize("shared_layout", [
    ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0]),
    ttgl.SwizzledSharedLayout(16, 2, 4, [1, 0]),
    ttgl.PaddedSharedLayout.with_identity_for([[4, 1]], [128, 128], [1, 0]),
])
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.float8_e4m3fn])
@pytest.mark.parametrize("NUM_TOTAL_WARPS", [8])
def test_runtime_ws_async_copy_mbarrier(M, N, shared_layout, dtype, NUM_TOTAL_WARPS):
    """
    Tests warp specialization with async_copy operations and mbarrier synchronization on GFX1250.

    This test validates the producer-consumer pattern using async_copy with two mbarriers:
    - ready_bar: Signals that ws_producer has completed copying data to the input buffer
    - done_bar: Signals that ws_consumer has completed processing and writing to the output buffer

    - ws_producer (default) partition: Copies data from global memory to shared memory
      then signals completion via mbarrier_arrive on ready_bar.
    - ws_consumer (worker) partition: Waits on ready_bar, loads data from the input shared memory buffer,
      stores it to an output shared memory buffer, then signals done_bar.

    The main kernel (executed by default warps) waits on done_bar, then loads data
    from the output buffer and stores it to global memory.

    NOTE: This test showcases that tensors (here: a_ptrs) can be passed as arguments to
    the default partition (here: ws_producer), which is not supported for worker partitions.

    The test verifies that the output matches the input, confirming correct synchronization.
    """

    @gluon.jit
    def ws_producer(a_ptrs, buffer, ready_bar):
        ttgl.amd.cdna5.async_copy.global_to_shared(buffer, a_ptrs)
        ttgl.amd.cdna5.async_copy.mbarrier_arrive(ready_bar)

    @gluon.jit
    def ws_consumer(in_buffer, out_buffer, ready_bar, done_bar, BLOCKED_LAYOUT: ttgl.constexpr):
        ttgl.amd.cdna5.mbarrier.wait(ready_bar, 0)
        val = in_buffer.load(BLOCKED_LAYOUT)
        out_buffer.store(val)
        ttgl.amd.cdna5.mbarrier.arrive(done_bar, count=1)

    @gluon.jit
    def ws_async_copy_mbarrier_kernel(a_ptr, out_ptr, M, N, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                      shared_layout: ttgl.constexpr, NUM_WARPS: ttgl.constexpr):
        PARTITION_WARPS: ttgl.constexpr = NUM_WARPS // 2
        ASYNC_LOAD_BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [PARTITION_WARPS, 1], [1, 0])
        BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout(
            [1, 8], [4, 8], [PARTITION_WARPS // 2, PARTITION_WARPS // (PARTITION_WARPS // 2)], [1, 0])
        WARP_SIZE: ttgl.constexpr = 32

        pid = ttgl.program_id(axis=0)
        num_pid_m = ttgl.cdiv(M, BLOCK_M)
        pid_m = pid % num_pid_m
        pid_n = pid // num_pid_m

        offs_m = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, ASYNC_LOAD_BLOCKED_LAYOUT))
        offs_n = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, ASYNC_LOAD_BLOCKED_LAYOUT))

        a_ptrs = a_ptr + offs_m[:, None] * N + offs_n[None, :]

        mbar = ttgl.allocate_shared_memory(ttgl.int64, [2, 1], ttgl.amd.cdna5.mbarrier.MBarrierLayout())
        buffer = ttgl.allocate_shared_memory(a_ptr.type.element_ty, [BLOCK_M, BLOCK_N], shared_layout)
        out_buffer = ttgl.allocate_shared_memory(out_ptr.type.element_ty, [BLOCK_M, BLOCK_N], shared_layout)

        ready_bar = mbar.index(0)
        done_bar = mbar.index(1)

        # TDM arrives on barrier once per warp, so use partition warp count
        ttgl.amd.cdna5.mbarrier.init(ready_bar, count=PARTITION_WARPS * WARP_SIZE)
        ttgl.amd.cdna5.mbarrier.init(done_bar, count=PARTITION_WARPS * WARP_SIZE)

        ttgl.warp_specialize([
            (ws_producer, (a_ptrs, buffer, ready_bar)),
            (ws_consumer, (buffer, out_buffer, ready_bar, done_bar, BLOCKED_LAYOUT)),
        ], [PARTITION_WARPS])

        out_offs_m = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
        out_offs_n = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
        mask = (out_offs_m[:, None] < M) & (out_offs_n[None, :] < N)
        ttgl.amd.cdna5.mbarrier.wait(done_bar, 0)
        res = out_buffer.load(BLOCKED_LAYOUT)
        out_ptrs = out_ptr + out_offs_m[:, None] * N + out_offs_n[None, :]
        ttgl.store(out_ptrs, res, mask)

    BLOCK_M = 128
    BLOCK_N = 128

    if dtype == torch.float8_e4m3fn:
        # range from min normal (0 00001 00) to max normal (0 11110 11)
        a = torch.randint(0x04, 0x7B, (M, N), dtype=torch.uint8).view(dtype)
    else:
        a = torch.rand((M, N), dtype=dtype)
    out = torch.empty_like(a)
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
    out_handle = out.cuda()

    ws_async_copy_mbarrier_kernel[grid](a.cuda(), out_handle, M, N, BLOCK_M, BLOCK_N, shared_layout,
                                        NUM_WARPS=NUM_TOTAL_WARPS, num_warps=NUM_TOTAL_WARPS // 2)
    out_tri = out_handle.cpu()
    out_ref = a.cpu()
    assert torch.equal(out_tri, out_ref)


# ==============================================================================
# Test async_copy shared_to_global with various layouts and vectorization
# ==============================================================================


@gluon.jit
def async_store_and_write_back_kernel(a_ptr, out_ptr, M, N, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                      blocked_layout: ttgl.constexpr, shared_layout: ttgl.constexpr):
    """
    Test kernel for async_copy.shared_to_global with 2D tensors.
    Loads from global -> shared (regular), then stores from shared -> global (async).
    """
    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    offs_m = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, blocked_layout))
    offs_n = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, blocked_layout))

    a_ptrs = a_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

    buffer = ttgl.allocate_shared_memory(a_ptr.type.element_ty, [BLOCK_M, BLOCK_N], shared_layout)

    # Regular load from global and store to shared
    value = ttgl.load(a_ptrs, mask=mask)
    buffer.store(value)

    # Async store from shared to global
    out_ptrs = out_ptr + offs_m[:, None] * N + offs_n[None, :]
    ttgl.amd.cdna5.async_copy.shared_to_global(out_ptrs, buffer, mask=mask)
    ttgl.amd.cdna5.async_copy.commit_group()
    ttgl.amd.cdna5.async_copy.wait_group(0)


@gluon.jit
def async_copy_shared_to_global_multi_cta_kernel(a_ptr, out_ptr, M, N, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                                 blocked_layout: ttgl.constexpr, shared_layout: ttgl.constexpr):
    """
    Test kernel for async_copy.shared_to_global with multi-CTA and 2D tensors.
    """
    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    offs_m = pid_m * BLOCK_M + ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, blocked_layout))
    offs_n = pid_n * BLOCK_N + ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, blocked_layout))

    a_ptrs = a_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

    buffer = ttgl.allocate_shared_memory(a_ptr.type.element_ty, [BLOCK_M, BLOCK_N], shared_layout)

    # Regular load from global and store to shared
    value = ttgl.load(a_ptrs, mask=mask)
    buffer.store(value)

    # Async store from shared to global
    out_ptrs = out_ptr + offs_m[:, None] * N + offs_n[None, :]
    ttgl.amd.cdna5.async_copy.shared_to_global(out_ptrs, buffer, mask=mask)
    ttgl.amd.cdna5.async_copy.commit_group()
    ttgl.amd.cdna5.async_copy.wait_group(0)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@ASYNC_COPY_TEST_PARAM_SIZE
@ASYNC_COPY_TEST_PARAM_SHARED_LAYOUT
@ASYNC_COPY_TEST_PARAM_DTYPE
def test_runtime_async_store(M, N, vec_size, shared_layout, dtype):
    """Test async_copy.shared_to_global with various layouts, sizes, and dtypes."""
    BLOCK_M = 128
    BLOCK_N = 128
    blocked_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0])

    if dtype == torch.float8_e4m3fn:
        # range from min normal (0 00001 00) to max normal (0 11110 11)
        a = torch.randint(0x04, 0x7B, (M, N), dtype=torch.uint8).view(dtype)
    else:
        a = torch.rand((M, N), dtype=dtype)
    out = torch.empty_like(a)
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
    out_handle = out.cuda()

    run_kernel = lambda: async_store_and_write_back_kernel[grid](a.cuda(), out_handle, M, N, BLOCK_M, BLOCK_N,
                                                                 blocked_layout, shared_layout)

    run_kernel()
    out_tri = out_handle.cpu()
    out_ref = a.cpu()
    assert torch.equal(out_tri, out_ref)


@pytest.mark.parametrize("blocked_layout", [
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[0, 1]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[1, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[0, 1], [0, 2], [0, 0], [0, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[0, 1], [0, 0], [1, 0], [0, 0]]),
    ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[4, 8], warps_per_cta=[1, 1], order=[1, 0],
                       cga_layout=[[0, 1], [0, 2], [0, 4], [0, 0]]),
])
def test_async_copy_shared_to_global_multi_cta(blocked_layout):
    """Test async_copy.shared_to_global with multi-CTA configurations."""
    M, N = 1024, 1024
    BLOCK_M, BLOCK_N = 128, 128
    num_ctas = 2**len(blocked_layout.cga_layout)

    shared_layout = ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0], blocked_layout.cga_layout)

    a = torch.rand((M, N), dtype=torch.float32)
    out = torch.empty_like(a)

    a_d = a.cuda()
    out_d = out.cuda()

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

    async_copy_shared_to_global_multi_cta_kernel[grid](a_d, out_d, M, N, BLOCK_M, BLOCK_N, blocked_layout,
                                                       shared_layout, num_warps=1, num_ctas=num_ctas)
    out_tri = out_d.cpu()
    assert torch.equal(out_tri, a)


@gluon.jit
def cluster_barrier_arrive_kernel():
    ttgl.amd.cdna5.cluster.arrive()


@gluon.jit
def cluster_barrier_wait_kernel():
    ttgl.amd.cdna5.cluster.wait()


def test_compile_cluster_barrier_arrive():
    """Test that cluster barrier arrive operation compiles correctly."""
    k = triton.compile(src=gluon._runtime.GluonASTSource(cluster_barrier_arrive_kernel, {}, {}),
                       target=GPUTarget("hip", 'gfx1250', 32), options={"num_ctas": 2})
    amdgcn = k.asm["amdgcn"]
    # Check that the ROCDL barrier signal instruction is present in the assembly
    assert "s_barrier_signal -3" in amdgcn


def test_compile_cluster_barrier_wait():
    """Test that cluster barrier wait operation compiles correctly."""
    k = triton.compile(src=gluon._runtime.GluonASTSource(cluster_barrier_wait_kernel, {}, {}),
                       target=GPUTarget("hip", 'gfx1250', 32), options={"num_ctas": 2})
    amdgcn = k.asm["amdgcn"]
    # Check that the ROCDL barrier wait instruction is present in the assembly
    assert "s_barrier_wait -3" in amdgcn


@gluon.jit
def cluster_barrier_arrive_and_wait_kernel():
    ttgl.amd.cdna5.cluster.arrive()
    ttgl.amd.cdna5.cluster.wait()


def test_runtime_cluster_barrier_arrive_and_wait():
    # Ensure that arrive and wait don't hang
    cluster_barrier_arrive_and_wait_kernel[(4, )](num_warps=8, num_ctas=2)
    torch.cuda.synchronize()


# =============================================================================
# TDM Scatter Mode Tests
# =============================================================================


@gluon.jit
def tdm_scatter_kernel(inp_ptr, out_ptr, dst_row_indices_ptr, M_out, N_out, stride_m, stride_n, BLOCK_M: ttgl.constexpr,
                       BLOCK_N: ttgl.constexpr, NUM_INDICES: ttgl.constexpr, DST_COL_OFFSET: ttgl.constexpr):
    """Kernel that uses TDM scatter to write non-contiguous rows."""
    num_warps: ttgl.constexpr = ttgl.num_warps()
    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    # For scatter indices, use a layout where all indices are available per thread.
    # NUM_INDICES must be a power of 2.
    IDX_BASE_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([NUM_INDICES, 1], [1, 32], [1, num_warps], [1, 0])
    IDX_LAYOUT: ttgl.constexpr = ttgl.SliceLayout(1, IDX_BASE_LAYOUT)

    # Allocate shared memory for the data to scatter
    smem = ttgl.allocate_shared_memory(inp_ptr.type.element_ty, (BLOCK_M, BLOCK_N), SHARED_LAYOUT)

    # Load data from global to shared memory using TDM
    inp_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=inp_ptr, shape=(BLOCK_M, BLOCK_N), strides=(BLOCK_N, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.tdm.async_load(inp_desc, [0, 0], smem)
    ttgl.amd.cdna5.tdm.async_wait(0)

    # Create tensor descriptor for output
    out_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=out_ptr, shape=(M_out, N_out), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)

    idx_offs = ttgl.arange(0, NUM_INDICES, layout=IDX_LAYOUT)
    dst_row_indices = ttgl.load(dst_row_indices_ptr + idx_offs)

    # Scatter the data to non-contiguous rows starting at DST_COL_OFFSET
    out_desc = ttgl.amd.cdna5.tdm.update_tensor_descriptor(out_desc, add_offsets=[0, DST_COL_OFFSET], clamp_bounds=True)
    ttgl.amd.cdna5.tdm.async_scatter(out_desc, dst_row_indices, smem)
    ttgl.amd.cdna5.tdm.async_wait(0)


@pytest.mark.parametrize("NUM_INDICES", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("BLOCK_M", [16, 32])
@pytest.mark.parametrize("BLOCK_N", [64, 128])
@pytest.mark.parametrize("dtype", ["fp16", "bf16", "fp8e5", "fp8e4nv"])
@pytest.mark.parametrize("index_dtype", ["i16", "i32"])
def test_compile_tdm_scatter(NUM_INDICES, BLOCK_M, BLOCK_N, dtype, index_dtype):
    """Test that TDM scatter compiles correctly."""
    # i16 indices: up to 16 per instruction, i32 indices: up to 8 per instruction
    max_indices_per_instr = 16 if index_dtype == "i16" else 8
    if NUM_INDICES > max_indices_per_instr:
        pytest.skip(f"NUM_INDICES={NUM_INDICES} exceeds max {max_indices_per_instr} for {index_dtype} indices")

    dtype_str = dtype
    signature = {
        "inp_ptr": f"*{dtype_str}",
        "out_ptr": f"*{dtype_str}",
        "dst_row_indices_ptr": f"*{index_dtype}",
        "M_out": "i32",
        "N_out": "i32",
        "stride_m": "i32",
        "stride_n": "i32",
        "BLOCK_M": "constexpr",
        "BLOCK_N": "constexpr",
        "NUM_INDICES": "constexpr",
        "DST_COL_OFFSET": "constexpr",
    }
    constexprs = {
        "BLOCK_M": BLOCK_M,
        "BLOCK_N": BLOCK_N,
        "NUM_INDICES": NUM_INDICES,
        "DST_COL_OFFSET": 0,
    }

    k = triton.compile(gluon._runtime.GluonASTSource(tdm_scatter_kernel, signature, constexprs),
                       target=GPUTarget("hip", 'gfx1250', 32))
    amdgcn = k.asm["amdgcn"]

    # Verify the scatter uses tensor store from lds
    assert re.search("tensor_store_from_lds", amdgcn), "Expected tensor_store_from_lds instruction for TDM scatter"
    assert re.search("s_wait_tensorcnt 0x0", amdgcn), "Expected s_wait_tensorcnt instruction"


def _create_scatter_test_data(shape, dtype):
    """Create random test data for TDM scatter tests, handling special dtypes."""
    if dtype == torch.float8_e5m2:
        # float8_e5m2: range from min normal to max normal
        return torch.randint(0x04, 0x7B, shape, dtype=torch.uint8).view(dtype)
    elif dtype == torch.float8_e4m3fn:
        # float8_e4m3fn: range from min normal to max normal
        return torch.randint(0x08, 0x77, shape, dtype=torch.uint8).view(dtype)
    else:
        # float16, bfloat16, float32, etc.
        return torch.randn(shape, dtype=dtype)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("NUM_INDICES", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("BLOCK_M", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("BLOCK_N", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("dst_col_offset", [0, 32])
@pytest.mark.parametrize("dtype", [
    torch.float16,
    torch.bfloat16,
    torch.float8_e5m2,
    torch.float8_e4m3fn,
])
@pytest.mark.parametrize("index_dtype", [torch.int16, torch.int32])
def test_runtime_tdm_scatter(NUM_INDICES, BLOCK_M, BLOCK_N, dst_col_offset, dtype, index_dtype):
    """Test TDM scatter correctness at runtime."""
    torch.manual_seed(42)

    # i16 indices: up to 16 per instruction, i32 indices: up to 8 per instruction
    max_indices_per_instr = 16 if index_dtype == torch.int16 else 8
    if NUM_INDICES > max_indices_per_instr:
        pytest.skip(f"NUM_INDICES={NUM_INDICES} exceeds max {max_indices_per_instr} for {index_dtype} indices")

    M_out = 2048
    # Output width needs to accommodate dst_col_offset + BLOCK_N
    N_out = dst_col_offset + BLOCK_N

    # Create input data (contiguous block)
    inp = _create_scatter_test_data((BLOCK_M, BLOCK_N), dtype)

    # Create output tensor (larger, with zeros)
    out = torch.zeros((M_out, N_out), dtype=dtype)

    # Create dst_row_indices (row indices for non-contiguous rows in output)
    # NUM_INDICES must be a power of 2
    dst_row_indices = torch.arange(NUM_INDICES, dtype=index_dtype)

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = dst_row_indices.cuda()

    h = tdm_scatter_kernel[(1, )](inp_d, out_d, indices_d, M_out=M_out, N_out=N_out, stride_m=out_d.stride(0),
                                  stride_n=out_d.stride(1), BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, NUM_INDICES=NUM_INDICES,
                                  DST_COL_OFFSET=dst_col_offset)

    # Assert that only a single tensor_store_from_lds is emitted
    amdgcn = h.asm["amdgcn"]
    actual_count = len(re.findall(r"tensor_store_from_lds", amdgcn))
    assert actual_count == 1, \
        f"Expected 1 tensor_store_from_lds instruction, but found {actual_count}"

    out_result = out_d.cpu()

    # Build expected output: zeros with scattered rows at the right positions
    # Use uint8 view for indexing since float8 types don't support index operations in the latest stable pytorch yet
    ref_out = torch.zeros_like(out)
    elem_size = ref_out.element_size()
    ref_out_bytes = ref_out.view(torch.uint8).reshape(M_out, -1)
    inp_bytes = inp.view(torch.uint8).reshape(BLOCK_M, -1)
    ref_out_bytes[dst_row_indices.long(),
                  dst_col_offset * elem_size:(dst_col_offset + BLOCK_N) * elem_size] = inp_bytes[:NUM_INDICES]
    torch.testing.assert_close(out_result, ref_out)


@pytest.mark.parametrize("BLOCK_M", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("BLOCK_N", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("index_dtype", ["i16", "i32"])
def test_compile_tdm_scatter_multiple_instructions(BLOCK_M, BLOCK_N, index_dtype):
    """Test that TDM scatter emits the correct number of instructions for large BLOCK_M."""
    # i16 indices: up to 16 per instruction, i32 indices: up to 8 per instruction
    max_per_instr = 16 if index_dtype == "i16" else 8
    if BLOCK_M <= max_per_instr:
        pytest.skip("This test is for cases requiring multiple TDM instructions")
    dtype = "fp32"
    dtype_str = str(dtype).split('.')[-1]

    # For multiple instructions test, use BLOCK_M as NUM_INDICES
    signature = {
        "inp_ptr": f"*{dtype_str}",
        "out_ptr": f"*{dtype_str}",
        "dst_row_indices_ptr": f"*{index_dtype}",
        "M_out": "i32",
        "N_out": "i32",
        "stride_m": "i32",
        "stride_n": "i32",
        "BLOCK_M": "constexpr",
        "BLOCK_N": "constexpr",
        "NUM_INDICES": "constexpr",
        "DST_COL_OFFSET": "constexpr",
    }
    constexprs = {
        "BLOCK_M": BLOCK_M,
        "BLOCK_N": BLOCK_N,
        "NUM_INDICES": BLOCK_M,
        "DST_COL_OFFSET": 0,
    }

    k = triton.compile(gluon._runtime.GluonASTSource(tdm_scatter_kernel, signature, constexprs),
                       target=GPUTarget("hip", 'gfx1250', 32))
    amdgcn = k.asm["amdgcn"]

    expected_num_instructions = math.ceil(BLOCK_M / max_per_instr)
    actual_count = len(re.findall(r"tensor_store_from_lds", amdgcn))

    assert actual_count == expected_num_instructions, \
        f"Expected {expected_num_instructions} tensor_store_from_lds instructions for BLOCK_M={BLOCK_M} " \
        f"with {index_dtype} indices, but found {actual_count}"


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("BLOCK_M", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("BLOCK_N", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("dst_col_offset", [0, 16])
@pytest.mark.parametrize("dtype", [
    torch.float16,
    torch.bfloat16,
    torch.float8_e5m2,
    torch.float8_e4m3fn,
])
@pytest.mark.parametrize("index_dtype", [torch.int16, torch.int32])
def test_runtime_tdm_scatter_multiple_instructions(BLOCK_M, BLOCK_N, dst_col_offset, dtype, index_dtype):
    """Test TDM scatter with more rows than can fit in a single instruction."""
    torch.manual_seed(42)

    # i16 indices: up to 16 per instruction, i32 indices: up to 8 per instruction
    max_per_instr = 16 if index_dtype == torch.int16 else 8
    if BLOCK_M <= max_per_instr:
        pytest.skip("This test is for cases requiring multiple TDM instructions")

    M_out = 2048  # Large enough to hold all scattered rows
    N_out = dst_col_offset + BLOCK_N

    # Create input data
    inp = _create_scatter_test_data((BLOCK_M, BLOCK_N), dtype)

    # Create output tensor
    out = torch.zeros((M_out, N_out), dtype=dtype)

    # Create randomized non-contiguous dst_row_indices
    dst_row_indices = torch.randperm(M_out, dtype=torch.int32)[:BLOCK_M]

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = dst_row_indices.cuda()

    tdm_scatter_kernel[(1, )](inp_d, out_d, indices_d, M_out=M_out, N_out=N_out, stride_m=N_out, stride_n=1,
                              BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, NUM_INDICES=BLOCK_M, DST_COL_OFFSET=dst_col_offset)

    out_result = out_d.cpu()

    # Build expected output: zeros with scattered rows at the right positions
    # Use uint8 view for indexing since float8 types don't support index operations in the latest stable pytorch yet
    ref_out = torch.zeros_like(out)
    elem_size = ref_out.element_size()
    ref_out_bytes = ref_out.view(torch.uint8).reshape(M_out, -1)
    inp_bytes = inp.view(torch.uint8).reshape(BLOCK_M, -1)
    ref_out_bytes[dst_row_indices.long(), dst_col_offset * elem_size:(dst_col_offset + BLOCK_N) * elem_size] = inp_bytes
    torch.testing.assert_close(out_result, ref_out)


@gluon.jit
def tdm_scatter_multi_col_kernel(inp_ptr, out_ptr, dst_row_indices_ptr, M, N, stride_m, BLOCK_M: ttgl.constexpr,
                                 BLOCK_N: ttgl.constexpr):
    """TDM scatter kernel that processes multiple column blocks."""
    num_warps: ttgl.constexpr = ttgl.num_warps()
    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])
    IDX_BASE_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([BLOCK_M, 1], [1, 32], [1, num_warps], [1, 0])
    IDX_LAYOUT: ttgl.constexpr = ttgl.SliceLayout(1, IDX_BASE_LAYOUT)

    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    smem = ttgl.allocate_shared_memory(inp_ptr.type.element_ty, (BLOCK_M, BLOCK_N), SHARED_LAYOUT)
    inp_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=inp_ptr, shape=(M, N), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.tdm.async_load(inp_desc, [pid_m * BLOCK_M, pid_n * BLOCK_N], smem)
    ttgl.amd.cdna5.tdm.async_wait(0)

    out_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=out_ptr, shape=(M, N), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)

    idx_offs = ttgl.arange(0, BLOCK_M, layout=IDX_LAYOUT)
    idx_mask = (pid_m * BLOCK_M + idx_offs) < M
    dst_row_indices = ttgl.load(dst_row_indices_ptr + pid_m * BLOCK_M + idx_offs, mask=idx_mask, other=M)

    col_offset = pid_n * BLOCK_N
    out_desc = ttgl.amd.cdna5.tdm.update_tensor_descriptor(out_desc, add_offsets=[0, col_offset], clamp_bounds=True)
    ttgl.amd.cdna5.tdm.async_scatter(out_desc, dst_row_indices, smem)
    ttgl.amd.cdna5.tdm.async_wait(0)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("N", [16, 32, 64, 100, 128, 140, 200, 250, 256, 300, 384, 400, 500])
@pytest.mark.parametrize("num_warps", [4, 8])
@pytest.mark.parametrize("index_dtype", [torch.int16, torch.int32])
def test_runtime_tdm_scatter_partial_column_block(N, num_warps, index_dtype):
    """Test TDM scatter with partial last column block (N not multiple of BLOCK_N)."""
    torch.manual_seed(42)

    BLOCK_M = 128
    BLOCK_N = 128
    M = 256

    inp = _create_scatter_test_data((M, N), torch.float16)
    out = torch.zeros((M, N), dtype=torch.float16)
    dst_row_indices = torch.randperm(M, dtype=index_dtype)

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = dst_row_indices.cuda()

    grid_n = triton.cdiv(N, BLOCK_N)
    grid = (triton.cdiv(M, BLOCK_M) * grid_n, )

    tdm_scatter_multi_col_kernel[grid](inp_d, out_d, indices_d, M=M, N=N, stride_m=N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
                                       num_warps=num_warps)

    out_result = out_d.cpu()

    # Build reference - scatter each source row to its destination
    ref_out = torch.zeros_like(out)
    for src_row in range(M):
        dst_row = dst_row_indices[src_row].item()
        if dst_row < M:
            ref_out[dst_row] = inp[src_row]

    torch.testing.assert_close(out_result, ref_out)


@gluon.jit
def _tdm_scatter_padded_kernel(inp_ptr, out_ptr, dst_row_indices_ptr, M_out, N_out, stride_m, BLOCK_M: ttgl.constexpr,
                               BLOCK_N: ttgl.constexpr, NUM_INDICES: ttgl.constexpr, DST_COL_OFFSET: ttgl.constexpr,
                               SHARED_LAYOUT: ttgl.constexpr, IDX_LAYOUT: ttgl.constexpr):
    """TDM scatter kernel using a padded shared layout."""
    smem = ttgl.allocate_shared_memory(inp_ptr.type.element_ty, (BLOCK_M, BLOCK_N), SHARED_LAYOUT)

    inp_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=inp_ptr, shape=(BLOCK_M, BLOCK_N), strides=(BLOCK_N, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.tdm.async_load(inp_desc, [0, 0], smem)
    ttgl.amd.cdna5.tdm.async_wait(0)

    out_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=out_ptr, shape=(M_out, N_out), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)

    idx_offs = ttgl.arange(0, NUM_INDICES, layout=IDX_LAYOUT)
    dst_row_indices = ttgl.load(dst_row_indices_ptr + idx_offs)

    out_desc = ttgl.amd.cdna5.tdm.update_tensor_descriptor(out_desc, add_offsets=[0, DST_COL_OFFSET], clamp_bounds=True)
    ttgl.amd.cdna5.tdm.async_scatter(out_desc, dst_row_indices, smem)
    ttgl.amd.cdna5.tdm.async_wait(0)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("NUM_INDICES", [4, 8, 16])
@pytest.mark.parametrize("BLOCK_N", [16, 64, 128])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("index_dtype", [torch.int16, torch.int32])
def test_runtime_tdm_scatter_padded(NUM_INDICES, BLOCK_N, dtype, index_dtype):
    """Test TDM scatter with padded shared layout.

    Verifies the tile-widening + OOB workaround: tile_dim0 is widened to
    include padding so LDS reads use the correct stride, while tensor_dim0
    is clamped to the original column width so OOB drops padding elements.
    """
    torch.manual_seed(42)

    BLOCK_M = NUM_INDICES
    SHARED_LAYOUT = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK_N, 8]], [BLOCK_M, BLOCK_N], [1, 0])

    num_warps = 4
    IDX_BASE = ttgl.BlockedLayout([NUM_INDICES, 1], [1, 32], [1, num_warps], [1, 0])
    IDX_LAYOUT = ttgl.SliceLayout(1, IDX_BASE)

    M_out = 2048
    N_out = BLOCK_N

    inp = _create_scatter_test_data((BLOCK_M, BLOCK_N), dtype)
    out = torch.zeros((M_out, N_out), dtype=dtype)
    dst_row_indices = torch.arange(NUM_INDICES, dtype=index_dtype)

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = dst_row_indices.cuda()

    _tdm_scatter_padded_kernel[(1, )](inp_d, out_d, indices_d, M_out=M_out, N_out=N_out, stride_m=out_d.stride(0),
                                      BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, NUM_INDICES=NUM_INDICES, DST_COL_OFFSET=0,
                                      SHARED_LAYOUT=SHARED_LAYOUT, IDX_LAYOUT=IDX_LAYOUT, num_warps=num_warps)

    out_result = out_d.cpu()
    ref_out = torch.zeros_like(out)
    for i in range(NUM_INDICES):
        ref_out[dst_row_indices[i].item()] = inp[i]
    torch.testing.assert_close(out_result, ref_out)


# =============================================================================
# TDM Gather Mode Tests
# =============================================================================


@gluon.jit
def tdm_gather_kernel(inp_ptr, out_ptr, src_row_indices_ptr, M_inp, N_inp, stride_m, stride_n, BLOCK_M: ttgl.constexpr,
                      BLOCK_N: ttgl.constexpr, NUM_INDICES: ttgl.constexpr, SRC_COL_OFFSET: ttgl.constexpr,
                      SHARED_LAYOUT: ttgl.constexpr):
    """Kernel that uses TDM gather to read non-contiguous rows."""
    num_warps: ttgl.constexpr = ttgl.num_warps()

    # For gather indices, use a layout where all indices are available per thread.
    # NUM_INDICES must be a power of 2.
    IDX_BASE_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([NUM_INDICES, 1], [1, 32], [1, num_warps], [1, 0])
    IDX_LAYOUT: ttgl.constexpr = ttgl.SliceLayout(1, IDX_BASE_LAYOUT)

    # Allocate shared memory for the gathered data
    smem = ttgl.allocate_shared_memory(inp_ptr.type.element_ty, (BLOCK_M, BLOCK_N), SHARED_LAYOUT)

    # Create tensor descriptor for input (source)
    inp_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=inp_ptr, shape=(M_inp, N_inp), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)

    idx_offs = ttgl.arange(0, NUM_INDICES, layout=IDX_LAYOUT)
    src_row_indices = ttgl.load(src_row_indices_ptr + idx_offs)

    inp_desc = ttgl.amd.cdna5.tdm.update_tensor_descriptor(inp_desc, add_offsets=[0, SRC_COL_OFFSET], clamp_bounds=True)
    ttgl.amd.cdna5.tdm.async_gather(inp_desc, src_row_indices, smem)
    ttgl.amd.cdna5.tdm.async_wait(0)

    # Store gathered data to output using TDM
    out_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=out_ptr, shape=(BLOCK_M, BLOCK_N), strides=(BLOCK_N, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.tdm.async_store(out_desc, [0, 0], smem)
    ttgl.amd.cdna5.tdm.async_wait(0)


@gluon.jit
def tdm_gather_multi_cta_kernel(inp_ptr, out_ptr, src_row_indices_ptr, M_inp, N_inp, stride_m, BLOCK_M: ttgl.constexpr,
                                BLOCK_N: ttgl.constexpr, SRC_COL_OFFSET: ttgl.constexpr, SHARED_LAYOUT: ttgl.constexpr,
                                BLOCK_LAYOUT: ttgl.constexpr, IDX_LAYOUT: ttgl.constexpr):
    """Kernel that uses TDM gather with a multi-CTA shared-memory layout."""
    smem = ttgl.allocate_shared_memory(inp_ptr.type.element_ty, (BLOCK_M, BLOCK_N), SHARED_LAYOUT)

    inp_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=inp_ptr, shape=(M_inp, N_inp), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)

    idx_offs = ttgl.arange(0, BLOCK_M, layout=IDX_LAYOUT)
    src_row_indices = ttgl.load(src_row_indices_ptr + idx_offs)

    ttgl.amd.cdna5.tdm.async_gather(inp_desc, src_row_indices, smem)
    ttgl.amd.cdna5.tdm.async_wait(0)

    gathered = smem.load(layout=BLOCK_LAYOUT)
    offs_m = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, BLOCK_LAYOUT))
    offs_n = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, BLOCK_LAYOUT))
    ttgl.store(out_ptr + offs_m[:, None] * BLOCK_N + offs_n[None, :], gathered)


@pytest.mark.parametrize("NUM_INDICES", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("BLOCK_M", [16, 32])
@pytest.mark.parametrize("BLOCK_N", [64, 128])
@pytest.mark.parametrize("dtype", ["fp16", "bf16", "fp8e5", "fp8e4nv"])
@pytest.mark.parametrize("index_dtype", ["i16", "i32"])
@pytest.mark.parametrize("shared_layout", ["padded", "swizzled"])
def test_compile_tdm_gather(NUM_INDICES, BLOCK_M, BLOCK_N, dtype, index_dtype, shared_layout):
    """Test that TDM gather compiles correctly."""
    # i16 indices: up to 16 per instruction, i32 indices: up to 8 per instruction
    max_indices_per_instr = 16 if index_dtype == "i16" else 8
    if NUM_INDICES > max_indices_per_instr:
        pytest.skip(f"NUM_INDICES={NUM_INDICES} exceeds max {max_indices_per_instr} for {index_dtype} indices")

    if shared_layout == "padded":
        SHARED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK_N, 8]], [BLOCK_M, BLOCK_N],
                                                                                  [1, 0])
    else:
        SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    dtype_str = dtype
    signature = {
        "inp_ptr": f"*{dtype_str}",
        "out_ptr": f"*{dtype_str}",
        "src_row_indices_ptr": f"*{index_dtype}",
        "M_inp": "i32",
        "N_inp": "i32",
        "stride_m": "i32",
        "stride_n": "i32",
        "BLOCK_M": "constexpr",
        "BLOCK_N": "constexpr",
        "NUM_INDICES": "constexpr",
        "SRC_COL_OFFSET": "constexpr",
        "SHARED_LAYOUT": "constexpr",
    }
    constexprs = {
        "BLOCK_M": BLOCK_M,
        "BLOCK_N": BLOCK_N,
        "NUM_INDICES": NUM_INDICES,
        "SRC_COL_OFFSET": 0,
        "SHARED_LAYOUT": SHARED_LAYOUT,
    }

    k = triton.compile(gluon._runtime.GluonASTSource(tdm_gather_kernel, signature, constexprs),
                       target=GPUTarget("hip", 'gfx1250', 32))
    amdgcn = k.asm["amdgcn"]
    # Verify the gather uses tensor load to lds
    assert re.search("tensor_load_to_lds", amdgcn), "Expected tensor_load_to_lds instruction for TDM gather"
    assert re.search("s_wait_tensorcnt 0x0", amdgcn), "Expected s_wait_tensorcnt instruction"


@pytest.mark.parametrize("CGALayout", [[[0, 0]],  # 1x1 cluster (single-cta)
                                       [[0, 1]],  # 1x2 cluster
                                       [[1, 0], [2, 0]],  # 4x1 cluster
                                       [[0, 1], [0, 2]],  # 1x4 cluster
                                       [[1, 0], [0, 1]],  # 2x2 cluster
                                       ])
def test_compile_tdm_gather_multi_cta(CGALayout):
    """Test that TDM gather compiles for multi-CTA layouts."""
    BLOCK_M = 16
    BLOCK_N = 64
    NUM_WARPS = 1
    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0], CGALayout)
    BLOCK_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [NUM_WARPS, 1], [1, 0], CGALayout)
    IDX_CGA_LAYOUT: ttgl.constexpr = [[basis[0], 0] for basis in CGALayout]
    IDX_PARENT_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([BLOCK_M, 1], [1, 32], [1, NUM_WARPS], [1, 0],
                                                           IDX_CGA_LAYOUT)
    IDX_LAYOUT: ttgl.constexpr = ttgl.SliceLayout(1, IDX_PARENT_LAYOUT)
    num_ctas = 2**len(CGALayout)

    signature = {
        "inp_ptr": "*fp16",
        "out_ptr": "*fp16",
        "src_row_indices_ptr": "*i32",
        "M_inp": "i32",
        "N_inp": "i32",
        "stride_m": "i32",
        "BLOCK_M": "constexpr",
        "BLOCK_N": "constexpr",
        "SRC_COL_OFFSET": "constexpr",
        "SHARED_LAYOUT": "constexpr",
        "BLOCK_LAYOUT": "constexpr",
        "IDX_LAYOUT": "constexpr",
    }
    constexprs = {
        "BLOCK_M": BLOCK_M,
        "BLOCK_N": BLOCK_N,
        "SRC_COL_OFFSET": 0,
        "SHARED_LAYOUT": SHARED_LAYOUT,
        "BLOCK_LAYOUT": BLOCK_LAYOUT,
        "IDX_LAYOUT": IDX_LAYOUT,
    }

    k = triton.compile(gluon._runtime.GluonASTSource(tdm_gather_multi_cta_kernel, signature, constexprs),
                       target=GPUTarget("hip", 'gfx1250', 32), options={"num_warps": NUM_WARPS, "num_ctas": num_ctas})
    assert re.search("tensor_load_to_lds", k.asm["amdgcn"])


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("NUM_INDICES", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("BLOCK_M", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("BLOCK_N", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("src_col_offset", [0, 32])
@pytest.mark.parametrize("dtype", [
    torch.float16,
    torch.bfloat16,
    torch.float8_e5m2,
    torch.float8_e4m3fn,
])
@pytest.mark.parametrize("index_dtype", [torch.int16, torch.int32])
@pytest.mark.parametrize("shared_layout", ["padded", "swizzled"])
def test_runtime_tdm_gather(NUM_INDICES, BLOCK_M, BLOCK_N, src_col_offset, dtype, index_dtype, shared_layout):
    """Test TDM gather correctness at runtime."""
    torch.manual_seed(42)

    if shared_layout == "padded":
        SHARED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK_N, 8]], [BLOCK_M, BLOCK_N],
                                                                                  [1, 0])
    else:
        SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    # i16 indices: up to 16 per instruction, i32 indices: up to 8 per instruction
    max_indices_per_instr = 16 if index_dtype == torch.int16 else 8
    if NUM_INDICES > max_indices_per_instr:
        pytest.skip(f"NUM_INDICES={NUM_INDICES} exceeds max {max_indices_per_instr} for {index_dtype} indices")

    M_inp = 2048
    # Input width needs to accommodate src_col_offset + BLOCK_N
    N_inp = src_col_offset + BLOCK_N
    inp = _create_scatter_test_data((M_inp, N_inp), dtype)

    out = torch.zeros((BLOCK_M, BLOCK_N), dtype=dtype)
    src_row_indices = torch.arange(NUM_INDICES, dtype=index_dtype)

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = src_row_indices.cuda()

    h = tdm_gather_kernel[(1, )](inp_d, out_d, indices_d, M_inp=M_inp, N_inp=N_inp, stride_m=inp_d.stride(0),
                                 stride_n=inp_d.stride(1), BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, NUM_INDICES=NUM_INDICES,
                                 SRC_COL_OFFSET=src_col_offset, SHARED_LAYOUT=SHARED_LAYOUT)

    amdgcn = h.asm["amdgcn"]
    actual_count = len(re.findall(r"tensor_load_to_lds", amdgcn))
    assert actual_count == 1, \
        f"Expected 1 tensor_load_to_lds instruction, but found {actual_count}"

    out_result = out_d.cpu()

    # Only compare the gathered rows (first NUM_INDICES rows)
    # Rows beyond NUM_INDICES contain uninitialized LDS data
    gathered_out = out_result[:NUM_INDICES]
    # Use uint8 view for indexing since float8 types don't support index operations in the latest stable pytorch yet
    elem_size = inp.element_size()
    inp_bytes = inp.view(torch.uint8).reshape(M_inp, -1)
    ref_bytes = inp_bytes[src_row_indices.long(), src_col_offset * elem_size:(src_col_offset + BLOCK_N) * elem_size]
    torch.testing.assert_close(gathered_out.view(torch.uint8), ref_bytes)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("CGALayout", [[[0, 0]],  # 1x1 cluster (single-cta)
                                       [[0, 1]],  # 1x2 cluster
                                       [[1, 0], [2, 0]],  # 4x1 cluster
                                       [[0, 1], [0, 2]],  # 1x4 cluster
                                       [[1, 0], [0, 1]],  # 2x2 cluster
                                       ])
@pytest.mark.parametrize("src_col_offset,physical_tail_cols", [(0, 0), (16, 16)])
def test_runtime_tdm_gather_multi_cta(CGALayout, src_col_offset, physical_tail_cols):
    """Test TDM gather correctness with multi-CTA and multicast layouts."""
    torch.manual_seed(42)

    BLOCK_M = 16
    BLOCK_N = 64
    NUM_WARPS = 1
    num_ctas = 2**len(CGALayout)
    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0], CGALayout)
    BLOCK_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [NUM_WARPS, 1], [1, 0], CGALayout)
    IDX_CGA_LAYOUT: ttgl.constexpr = [[basis[0], 0] for basis in CGALayout]
    IDX_PARENT_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([BLOCK_M, 1], [1, 32], [1, NUM_WARPS], [1, 0],
                                                           IDX_CGA_LAYOUT)
    IDX_LAYOUT: ttgl.constexpr = ttgl.SliceLayout(1, IDX_PARENT_LAYOUT)

    M_inp = 128
    N_inp = BLOCK_N
    physical_N_inp = N_inp + physical_tail_cols
    inp = _create_scatter_test_data((M_inp, physical_N_inp), torch.float16)
    out = torch.zeros((BLOCK_M, BLOCK_N), dtype=torch.float16)
    src_row_indices = torch.randperm(M_inp, dtype=torch.int32)[:BLOCK_M]

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = src_row_indices.cuda()

    tdm_gather_multi_cta_kernel[(1, )](inp_d, out_d, indices_d, M_inp=M_inp, N_inp=N_inp, stride_m=inp_d.stride(0),
                                       BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, SHARED_LAYOUT=SHARED_LAYOUT,
                                       SRC_COL_OFFSET=src_col_offset, BLOCK_LAYOUT=BLOCK_LAYOUT, IDX_LAYOUT=IDX_LAYOUT,
                                       num_warps=NUM_WARPS, num_ctas=num_ctas)

    valid_cols = N_inp - src_col_offset
    ref = torch.zeros_like(out)
    ref[:, :valid_cols] = inp[src_row_indices.long(), src_col_offset:N_inp]
    torch.testing.assert_close(out_d.cpu().view(torch.uint8), ref.view(torch.uint8))


@pytest.mark.parametrize("BLOCK_M", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("BLOCK_N", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("index_dtype", ["i16", "i32"])
@pytest.mark.parametrize("shared_layout", ["padded", "swizzled"])
def test_compile_tdm_gather_multiple_instructions(BLOCK_M, BLOCK_N, index_dtype, shared_layout):
    """Test that TDM gather emits the correct number of instructions for large BLOCK_M."""
    # i16 indices: up to 16 per instruction, i32 indices: up to 8 per instruction
    max_per_instr = 16 if index_dtype == "i16" else 8
    if BLOCK_M <= max_per_instr:
        pytest.skip("This test is for cases requiring multiple TDM instructions")
    dtype = "fp32"
    dtype_str = str(dtype).split('.')[-1]

    if shared_layout == "padded":
        SHARED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK_N, 8]], [BLOCK_M, BLOCK_N],
                                                                                  [1, 0])
    else:
        SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    # For multiple instructions test, use BLOCK_M as NUM_INDICES
    signature = {
        "inp_ptr": f"*{dtype_str}",
        "out_ptr": f"*{dtype_str}",
        "src_row_indices_ptr": f"*{index_dtype}",
        "M_inp": "i32",
        "N_inp": "i32",
        "stride_m": "i32",
        "stride_n": "i32",
        "BLOCK_M": "constexpr",
        "BLOCK_N": "constexpr",
        "NUM_INDICES": "constexpr",
        "SRC_COL_OFFSET": "constexpr",
        "SHARED_LAYOUT": "constexpr",
    }
    constexprs = {
        "BLOCK_M": BLOCK_M,
        "BLOCK_N": BLOCK_N,
        "NUM_INDICES": BLOCK_M,
        "SRC_COL_OFFSET": 0,
        "SHARED_LAYOUT": SHARED_LAYOUT,
    }

    k = triton.compile(gluon._runtime.GluonASTSource(tdm_gather_kernel, signature, constexprs),
                       target=GPUTarget("hip", 'gfx1250', 32))
    amdgcn = k.asm["amdgcn"]

    expected_num_instructions = math.ceil(BLOCK_M / max_per_instr)
    actual_count = len(re.findall(r"tensor_load_to_lds", amdgcn))

    assert actual_count == expected_num_instructions, \
        f"Expected {expected_num_instructions} tensor_load_to_lds instructions for BLOCK_M={BLOCK_M} " \
        f"with {index_dtype} indices, but found {actual_count}"


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("BLOCK_M", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("BLOCK_N", [16, 32, 64, 128, 256])
@pytest.mark.parametrize("src_col_offset", [0, 16])
@pytest.mark.parametrize("dtype", [
    torch.float16,
    torch.bfloat16,
    torch.float8_e5m2,
    torch.float8_e4m3fn,
])
@pytest.mark.parametrize("index_dtype", [torch.int16, torch.int32])
@pytest.mark.parametrize("shared_layout", ["padded", "swizzled"])
def test_runtime_tdm_gather_multiple_instructions(BLOCK_M, BLOCK_N, src_col_offset, dtype, index_dtype, shared_layout):
    """Test TDM gather with more rows than can fit in a single instruction."""
    torch.manual_seed(42)

    if shared_layout == "padded":
        SHARED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK_N, 8]], [BLOCK_M, BLOCK_N],
                                                                                  [1, 0])
    else:
        SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    # i16 indices: up to 16 per instruction, i32 indices: up to 8 per instruction
    max_per_instr = 16 if index_dtype == torch.int16 else 8
    if BLOCK_M <= max_per_instr:
        pytest.skip("This test is for cases requiring multiple TDM instructions")

    M_inp = 2048  # Large enough to hold all source rows
    N_inp = src_col_offset + BLOCK_N

    inp = _create_scatter_test_data((M_inp, N_inp), dtype)

    out = torch.zeros((BLOCK_M, BLOCK_N), dtype=dtype)

    # Create randomized non-contiguous src_row_indices
    src_row_indices = torch.randperm(M_inp, dtype=torch.int32)[:BLOCK_M]

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = src_row_indices.cuda()

    tdm_gather_kernel[(1, )](inp_d, out_d, indices_d, M_inp=M_inp, N_inp=N_inp, stride_m=N_inp, stride_n=1,
                             BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, NUM_INDICES=BLOCK_M, SRC_COL_OFFSET=src_col_offset,
                             SHARED_LAYOUT=SHARED_LAYOUT)

    out_result = out_d.cpu()

    # Build expected output: gathered rows from input
    # Use uint8 view for indexing since float8 types don't support index operations in the latest stable pytorch yet
    elem_size = inp.element_size()
    inp_bytes = inp.view(torch.uint8).reshape(M_inp, -1)
    ref_bytes = inp_bytes[src_row_indices.long(), src_col_offset * elem_size:(src_col_offset + BLOCK_N) * elem_size]
    torch.testing.assert_close(out_result.view(torch.uint8), ref_bytes)


@gluon.jit
def tdm_gather_multi_col_kernel(inp_ptr, out_ptr, src_row_indices_ptr, M, N, stride_m, BLOCK_M: ttgl.constexpr,
                                BLOCK_N: ttgl.constexpr):
    """TDM gather kernel that processes multiple column blocks."""
    num_warps: ttgl.constexpr = ttgl.num_warps()
    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])
    IDX_BASE_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([BLOCK_M, 1], [1, 32], [1, num_warps], [1, 0])
    IDX_LAYOUT: ttgl.constexpr = ttgl.SliceLayout(1, IDX_BASE_LAYOUT)

    pid = ttgl.program_id(axis=0)
    num_pid_m = ttgl.cdiv(M, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    smem = ttgl.allocate_shared_memory(inp_ptr.type.element_ty, (BLOCK_M, BLOCK_N), SHARED_LAYOUT)

    inp_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=inp_ptr, shape=(M, N), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)

    idx_offs = ttgl.arange(0, BLOCK_M, layout=IDX_LAYOUT)
    idx_mask = (pid_m * BLOCK_M + idx_offs) < M
    src_row_indices = ttgl.load(src_row_indices_ptr + pid_m * BLOCK_M + idx_offs, mask=idx_mask, other=M)

    col_offset = pid_n * BLOCK_N
    inp_desc = ttgl.amd.cdna5.tdm.update_tensor_descriptor(inp_desc, add_offsets=[0, col_offset], clamp_bounds=True)
    ttgl.amd.cdna5.tdm.async_gather(inp_desc, src_row_indices, smem)
    ttgl.amd.cdna5.tdm.async_wait(0)

    out_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=out_ptr, shape=(M, N), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.tdm.async_store(out_desc, [pid_m * BLOCK_M, pid_n * BLOCK_N], smem)
    ttgl.amd.cdna5.tdm.async_wait(0)


@gluon.jit
def _tdm_gather_layout_kernel(inp_ptr, out_ptr, src_row_indices_ptr, M_inp, N_inp, stride_m, stride_n,
                              BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, NUM_INDICES: ttgl.constexpr,
                              SRC_COL_OFFSET: ttgl.constexpr, SHARED_LAYOUT: ttgl.constexpr,
                              IDX_LAYOUT: ttgl.constexpr):
    """Generic TDM gather kernel parameterized by index layout.

    The caller supplies IDX_LAYOUT (a SliceLayout over a 2D BlockedLayout)
    which determines how indices are distributed across warps.  The LLVM
    lowering uses LinearLayout + freeVarMasks to generically handle any
    valid distribution (replicated, partitioned, or mixed).
    """
    smem = ttgl.allocate_shared_memory(inp_ptr.type.element_ty, (BLOCK_M, BLOCK_N), SHARED_LAYOUT)

    inp_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=inp_ptr, shape=(M_inp, N_inp), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)

    idx_offs = ttgl.arange(0, NUM_INDICES, layout=IDX_LAYOUT)
    src_row_indices = ttgl.load(src_row_indices_ptr + idx_offs)

    inp_desc = ttgl.amd.cdna5.tdm.update_tensor_descriptor(inp_desc, add_offsets=[0, SRC_COL_OFFSET], clamp_bounds=True)
    ttgl.amd.cdna5.tdm.async_gather(inp_desc, src_row_indices, smem)
    ttgl.amd.cdna5.tdm.async_wait(0)

    out_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=out_ptr, shape=(BLOCK_M, BLOCK_N), strides=(BLOCK_N, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.tdm.async_store(out_desc, [0, 0], smem)
    ttgl.amd.cdna5.tdm.async_wait(0)


def _tdm_gather_scatter_index_layouts():
    """Yield (num_indices, num_warps, idx_layout) for TDM gather and scatter layout testing.

    Covers five layout categories:
      - Replicated:       warpsPerCTA=[1, W]   — all warps hold all indices, only warp 0 fires.
      - Partitioned:      warpsPerCTA=[W, 1]   — each warp owns a disjoint subset, all fire.
      - Mixed:            warpsPerCTA=[A, R]   — A warps partition, R redundant copies predicated off.
      - Redundant regs:   sizePerThread=[ni, 2] — each thread holds duplicate values (regMask != 0).
      - Wrapping:         sizePerThread < numIndices/numWarps — register dim tiles, non-contiguous positions.
    """
    cases = []
    for ni in [4, 8, 16]:
        # Replicated
        for nw in [2, 4]:
            parent = ttgl.BlockedLayout([ni, 1], [1, 32], [1, nw], [1, 0])
            cases.append(pytest.param(ni, nw, ttgl.SliceLayout(1, parent), id=f"replicated-ni{ni}-w{nw}"))

        # Partitioned
        for nw in [2, 4]:
            if ni >= nw and ni % nw == 0:
                parent = ttgl.BlockedLayout([ni // nw, 1], [1, 32], [nw, 1], [1, 0])
                cases.append(pytest.param(ni, nw, ttgl.SliceLayout(1, parent), id=f"partitioned-ni{ni}-w{nw}"))

        # Mixed (partially redundant warps)
        for active, redundant in [(2, 2), (1, 4)]:
            nw = active * redundant
            if ni >= active and ni % active == 0:
                parent = ttgl.BlockedLayout([ni // active, 1], [1, 32], [active, redundant], [1, 0])
                cases.append(
                    pytest.param(ni, nw, ttgl.SliceLayout(1, parent), id=f"mixed_{active}x{redundant}-ni{ni}-w{nw}"))

        # Redundant registers: sizePerThread has >1 in the sliced dim,
        # so each thread holds duplicate index values (regMask != 0).
        # E.g., sizePerThread=[ni, 2] → 2*ni register slots, only ni unique.
        for nw in [1, 2]:
            parent = ttgl.BlockedLayout([ni, 2], [1, 32], [1, nw], [1, 0])
            cases.append(pytest.param(ni, nw, ttgl.SliceLayout(1, parent), id=f"redundant_reg-ni{ni}-w{nw}"))

    # Wrapping: sizePerThread[0] < numIndices/numWarps, so the register
    # dimension tiles (wraps) and produces non-contiguous dim0 positions.
    # E.g., 32 indices with spt=4 and 4 warps: register bases {1,2,16}.
    for ni, nw, spt in [(16, 4, 2), (16, 2, 4), (32, 4, 4)]:
        if ni >= nw and ni % nw == 0 and (ni // nw) > spt:
            parent = ttgl.BlockedLayout([spt, 1], [1, 32], [nw, 1], [1, 0])
            cases.append(pytest.param(ni, nw, ttgl.SliceLayout(1, parent), id=f"wrapping-ni{ni}-w{nw}-spt{spt}"))

    return cases


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("num_indices, num_warps, idx_layout", _tdm_gather_scatter_index_layouts())
@pytest.mark.parametrize("BLOCK_N", [64, 128])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("index_dtype", [torch.int16, torch.int32])
def test_runtime_tdm_gather_layouts(num_indices, num_warps, idx_layout, BLOCK_N, dtype, index_dtype):
    """Test TDM gather correctness across different index layouts.

    A single test function exercises replicated, partitioned, and mixed warp
    layouts by injecting the index layout from _tdm_gather_scatter_index_layouts().
    Mirrors the pattern used by NVIDIA's test_gather_layouts in test_lowerings.py.
    """
    torch.manual_seed(42)

    BLOCK_M = num_indices
    SHARED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK_N, 8]], [BLOCK_M, BLOCK_N],
                                                                              [1, 0])

    M_inp = 2048
    N_inp = BLOCK_N
    inp = _create_scatter_test_data((M_inp, N_inp), dtype)
    out = torch.zeros((BLOCK_M, BLOCK_N), dtype=dtype)
    src_row_indices = torch.arange(num_indices, dtype=index_dtype)

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = src_row_indices.cuda()

    _tdm_gather_layout_kernel[(1, )](inp_d, out_d, indices_d, M_inp=M_inp, N_inp=N_inp, stride_m=inp_d.stride(0),
                                     stride_n=inp_d.stride(1), BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
                                     NUM_INDICES=num_indices, SRC_COL_OFFSET=0, SHARED_LAYOUT=SHARED_LAYOUT,
                                     IDX_LAYOUT=idx_layout, num_warps=num_warps)

    out_result = out_d.cpu()
    gathered_out = out_result[:num_indices]
    elem_size = inp.element_size()
    inp_bytes = inp.view(torch.uint8).reshape(M_inp, -1)
    ref_bytes = inp_bytes[src_row_indices.long(), :BLOCK_N * elem_size]
    torch.testing.assert_close(gathered_out.view(torch.uint8), ref_bytes)


@gluon.jit
def _tdm_scatter_layout_kernel(inp_ptr, out_ptr, dst_row_indices_ptr, M_out, N_out, stride_m, stride_n,
                               BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, NUM_INDICES: ttgl.constexpr,
                               DST_COL_OFFSET: ttgl.constexpr, SHARED_LAYOUT: ttgl.constexpr,
                               IDX_LAYOUT: ttgl.constexpr):
    """Generic TDM scatter kernel parameterized by index layout.

    Mirrors _tdm_gather_layout_kernel but for the scatter direction:
    data flows from LDS to non-contiguous global memory rows.
    """
    smem = ttgl.allocate_shared_memory(inp_ptr.type.element_ty, (BLOCK_M, BLOCK_N), SHARED_LAYOUT)

    inp_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=inp_ptr, shape=(BLOCK_M, BLOCK_N), strides=(BLOCK_N, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.tdm.async_load(inp_desc, [0, 0], smem)
    ttgl.amd.cdna5.tdm.async_wait(0)

    out_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=out_ptr, shape=(M_out, N_out), strides=(stride_m, 1),
                                                         block_shape=(BLOCK_M, BLOCK_N), layout=SHARED_LAYOUT)

    idx_offs = ttgl.arange(0, NUM_INDICES, layout=IDX_LAYOUT)
    dst_row_indices = ttgl.load(dst_row_indices_ptr + idx_offs)

    out_desc = ttgl.amd.cdna5.tdm.update_tensor_descriptor(out_desc, add_offsets=[0, DST_COL_OFFSET], clamp_bounds=True)
    ttgl.amd.cdna5.tdm.async_scatter(out_desc, dst_row_indices, smem)
    ttgl.amd.cdna5.tdm.async_wait(0)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("num_indices, num_warps, idx_layout", _tdm_gather_scatter_index_layouts())
@pytest.mark.parametrize("BLOCK_N", [64, 128])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("index_dtype", [torch.int16, torch.int32])
def test_runtime_tdm_scatter_layouts(num_indices, num_warps, idx_layout, BLOCK_N, dtype, index_dtype):
    """Test TDM scatter correctness across different index layouts.

    Use _tdm_gather_scatter_index_layouts() to exercise replicated, partitioned,
    and mixed warp layouts for the scatter direction. Verifies that the
    LinearLayout-based warp predication and LDS offset computation produce
    correct results for all layout configurations.
    """
    torch.manual_seed(42)

    BLOCK_M = num_indices
    SHARED_LAYOUT: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    M_out = 2048
    N_out = BLOCK_N

    inp = _create_scatter_test_data((BLOCK_M, BLOCK_N), dtype)
    out = torch.zeros((M_out, N_out), dtype=dtype)
    dst_row_indices = torch.arange(num_indices, dtype=index_dtype)

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = dst_row_indices.cuda()

    _tdm_scatter_layout_kernel[(1, )](inp_d, out_d, indices_d, M_out=M_out, N_out=N_out, stride_m=out_d.stride(0),
                                      stride_n=out_d.stride(1), BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
                                      NUM_INDICES=num_indices, DST_COL_OFFSET=0, SHARED_LAYOUT=SHARED_LAYOUT,
                                      IDX_LAYOUT=idx_layout, num_warps=num_warps)

    out_result = out_d.cpu()
    ref_out = torch.zeros_like(out)
    elem_size = ref_out.element_size()
    ref_out_bytes = ref_out.view(torch.uint8).reshape(M_out, -1)
    inp_bytes = inp.view(torch.uint8).reshape(BLOCK_M, -1)
    ref_out_bytes[dst_row_indices.long(), :BLOCK_N * elem_size] = inp_bytes[:num_indices]
    torch.testing.assert_close(out_result, ref_out)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("N", [16, 32, 64, 100, 128, 140, 200, 250, 256, 300, 384, 400, 500])
@pytest.mark.parametrize("num_warps", [4, 8])
@pytest.mark.parametrize("index_dtype", [torch.int16, torch.int32])
def test_runtime_tdm_gather_partial_column_block(N, num_warps, index_dtype):
    """Test TDM gather with partial last column block (N not multiple of BLOCK_N)."""
    torch.manual_seed(42)

    BLOCK_M = 128
    BLOCK_N = 128
    M = 256

    inp = _create_scatter_test_data((M, N), torch.float16)
    out = torch.zeros((M, N), dtype=torch.float16)
    src_row_indices = torch.randperm(M, dtype=index_dtype)

    inp_d = inp.cuda()
    out_d = out.cuda()
    indices_d = src_row_indices.cuda()

    grid_n = triton.cdiv(N, BLOCK_N)
    grid = (triton.cdiv(M, BLOCK_M) * grid_n, )

    tdm_gather_multi_col_kernel[grid](inp_d, out_d, indices_d, M=M, N=N, stride_m=N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
                                      num_warps=num_warps)

    out_result = out_d.cpu()

    # Build reference - gather each source row to its destination
    ref_out = torch.zeros_like(out)
    for dst_row in range(M):
        src_row = src_row_indices[dst_row].item()
        if src_row < M:
            ref_out[dst_row] = inp[src_row]

    torch.testing.assert_close(out_result, ref_out)


@gluon.jit
def buffer_load_store_roundtrip_kernel(a_ptr, b_ptr, BLOCK: ttgl.constexpr, loadCM: ttgl.constexpr,
                                       storeCM: ttgl.constexpr):
    BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([8], [32], [1], [0])
    pid = ttgl.program_id(axis=0)
    offs = pid * BLOCK + ttgl.arange(0, BLOCK, layout=BLOCKED_LAYOUT)
    data = ttgl.amd.cdna5.buffer_load(ptr=a_ptr, offsets=offs, cache=loadCM)
    ttgl.amd.cdna5.buffer_store(stored_value=data, ptr=b_ptr, offsets=offs, cache=storeCM)


@gluon.jit
def async_load_store_roundtrip_kernel(a_ptr, b_ptr, BLOCK: ttgl.constexpr, loadCM: ttgl.constexpr,
                                      storeCM: ttgl.constexpr):
    BLOCKED_LAYOUT: ttgl.constexpr = ttgl.BlockedLayout([8], [32], [1], [0])
    SHARED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK // 2, 8]], [BLOCK], [0])
    pid = ttgl.program_id(axis=0)
    offs = pid * BLOCK + ttgl.arange(0, BLOCK, layout=BLOCKED_LAYOUT)
    buffer = ttgl.allocate_shared_memory(ttgl.float16, shape=[BLOCK], layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.async_copy.global_to_shared(buffer, a_ptr + offs, cache_modifier=loadCM)
    ttgl.amd.cdna5.async_copy.commit_group()
    ttgl.amd.cdna5.async_copy.wait_group(0)
    ttgl.amd.cdna5.async_copy.shared_to_global(b_ptr + offs, buffer, cache_modifier=storeCM)


@gluon.jit
def tdm_load_store_roundtrip_kernel(a_ptr, b_ptr, BLOCK: ttgl.constexpr, loadCM: ttgl.constexpr,
                                    storeCM: ttgl.constexpr):
    SHARED_LAYOUT: ttgl.constexpr = ttgl.PaddedSharedLayout.with_identity_for([[BLOCK, 8]], [BLOCK], [0])
    pid = ttgl.program_id(axis=0)

    a_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=a_ptr, shape=(BLOCK, ), strides=(1, ),
                                                       block_shape=(BLOCK, ), layout=SHARED_LAYOUT)
    buffer = ttgl.allocate_shared_memory(ttgl.float16, shape=[BLOCK], layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.tdm.async_load(a_desc, [pid * BLOCK], buffer, cache_modifier=loadCM)
    ttgl.amd.cdna5.tdm.async_wait(0)

    b_desc = ttgl.amd.cdna5.tdm.make_tensor_descriptor(base=b_ptr, shape=(BLOCK, ), strides=(1, ),
                                                       block_shape=(BLOCK, ), layout=SHARED_LAYOUT)
    ttgl.amd.cdna5.tdm.async_store(b_desc, [pid * BLOCK], buffer, cache_modifier=storeCM)
    ttgl.amd.cdna5.tdm.async_wait(0)


@pytest.mark.parametrize("loadCM, storeCM", [(".ca", ".wb"), (".cg", ".cg"), (".cs", ".cs"), (".cv", ".wt")])
@pytest.mark.parametrize(
    "test_kernel",
    [buffer_load_store_roundtrip_kernel, async_load_store_roundtrip_kernel, tdm_load_store_roundtrip_kernel])
def test_cache_modifier(loadCM, storeCM, test_kernel):
    BLOCK = 256
    N = 256

    src = torch.rand((N, ), dtype=torch.float16).cuda()
    dst = torch.empty((N, ), dtype=torch.float16).cuda()
    pgm = test_kernel[(triton.cdiv(N, BLOCK), )](src, dst, BLOCK, loadCM, storeCM, num_warps=1)

    torch.testing.assert_close(dst, src)

    amdgcn = pgm.asm["amdgcn"]

    load_found = False
    store_found = False
    for line in amdgcn.split("\n"):
        if "buffer_load_b128" in line or "global_load_async_to_lds_b128" in line or "tensor_load_to_lds" in line:
            load_found = True
            if loadCM == ".ca":
                assert "scope" not in line and "th" not in line
            if loadCM == ".cg":
                assert "scope:SCOPE_DEV" in line and "th" not in line
            if loadCM == ".cs":
                assert "scope" not in line and "th:TH_LOAD_NT" in line
            if loadCM == ".cv":
                assert "scope:SCOPE_SYS" in line and "th:TH_LOAD_BYPASS" in line
        if "buffer_store_b128" in line or "global_store_async_from_lds_b128" in line or "tensor_store_from_lds" in line:
            store_found = True
            if storeCM == ".wb":
                assert "scope" not in line and "th" not in line
            if storeCM == ".cg":
                assert "scope:SCOPE_DEV" in line and "th" not in line
            if storeCM == ".cs":
                assert "scope" not in line and "th:TH_STORE_NT" in line
            if storeCM == ".wt":
                assert "scope:SCOPE_SYS" in line and "th:TH_STORE_BYPASS" in line

    assert load_found
    assert store_found


# Multi-CTA batched matmul: each CTA in the cluster computes one batch of A @ B.
@gluon.jit
def gemm_3d_cga_split_kernel(a_ptr, b_ptr, c_ptr, M, N, K,  #
                             BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,
                             NUM_CTAS: ttgl.constexpr, K_WIDTH: ttgl.constexpr, load_layout: ttgl.constexpr,
                             wmma_layout: ttgl.constexpr):
    load_dim0_layout: ttgl.constexpr = ttgl.SliceLayout(1, ttgl.SliceLayout(2, load_layout))
    load_dim1_layout: ttgl.constexpr = ttgl.SliceLayout(0, ttgl.SliceLayout(2, load_layout))
    load_dim2_layout: ttgl.constexpr = ttgl.SliceLayout(0, ttgl.SliceLayout(1, load_layout))
    wmma_dim0_layout: ttgl.constexpr = ttgl.SliceLayout(1, ttgl.SliceLayout(2, wmma_layout))
    wmma_dim1_layout: ttgl.constexpr = ttgl.SliceLayout(0, ttgl.SliceLayout(2, wmma_layout))
    wmma_dim2_layout: ttgl.constexpr = ttgl.SliceLayout(0, ttgl.SliceLayout(1, wmma_layout))

    ob = ttgl.arange(0, NUM_CTAS, layout=load_dim0_layout)
    om = ttgl.arange(0, BLOCK_M, layout=load_dim1_layout)
    ok = ttgl.arange(0, BLOCK_K, layout=load_dim2_layout)
    okb = ttgl.arange(0, BLOCK_K, layout=load_dim1_layout)
    on = ttgl.arange(0, BLOCK_N, layout=load_dim2_layout)

    offs_a = ob[:, None, None] * (M * K) + om[None, :, None] * K + ok[None, None, :]
    a = ttgl.load(a_ptr + offs_a)
    offs_b = ob[:, None, None] * (K * N) + okb[None, :, None] * N + on[None, None, :]
    b = ttgl.load(b_ptr + offs_b)

    a = ttgl.convert_layout(a, ttgl.DotOperandLayout(0, wmma_layout, K_WIDTH))
    b = ttgl.convert_layout(b, ttgl.DotOperandLayout(1, wmma_layout, K_WIDTH))
    acc = ttgl.zeros((NUM_CTAS, BLOCK_M, BLOCK_N), dtype=ttgl.float32, layout=wmma_layout)
    acc = ttgl.amd.cdna5.wmma(a, b, acc)

    sb = ttgl.arange(0, NUM_CTAS, layout=wmma_dim0_layout)
    sm = ttgl.arange(0, BLOCK_M, layout=wmma_dim1_layout)
    sn = ttgl.arange(0, BLOCK_N, layout=wmma_dim2_layout)
    offs_c = sb[:, None, None] * (BLOCK_M * BLOCK_N) + sm[None, :, None] * BLOCK_N + sn[None, None, :]
    ttgl.store(c_ptr + offs_c, acc)


@pytest.mark.skipif(not is_hip_gfx1250(), reason="Requires CDNA5")
@pytest.mark.parametrize("num_ctas", [2, 4])
def test_runtime_gemm_3d_multi_cta(num_ctas):
    BLOCK_M = BLOCK_N = 32
    BLOCK_K = 64
    num_warps = 4
    M, N, K = BLOCK_M, BLOCK_N, BLOCK_K

    cga = make_cga_layout([num_ctas, 1, 1], [num_ctas, 1, 1], [2, 1, 0])
    n_warps = 2
    m_warps = num_warps // n_warps
    sM = BLOCK_M // (m_warps * 8)
    sN = BLOCK_N // (n_warps * 4)
    load_layout = ttgl.BlockedLayout([1, sM, sN], [1, 8, 4], [1, m_warps, n_warps], [2, 1, 0], cga)
    warp_bases = [[0, 0, 1]]
    for i in range(int(math.log2(num_warps // 2))):
        warp_bases.append([0, 1 << i, 0])
    wmma_layout = ttgl.amd.AMDWMMALayout(version=3, transposed=True, warp_bases=warp_bases, reg_bases=[],
                                         instr_shape=[16, 16, 32], cga_layout=cga, rank=3)

    torch.manual_seed(0)
    a = torch.randn((num_ctas, M, K), dtype=torch.float16)
    b = torch.randn((num_ctas, K, N), dtype=torch.float16)
    c = torch.zeros((num_ctas, M, N), dtype=torch.float32)
    a_device, b_device, c_device = a.cuda(), b.cuda(), c.cuda()

    # Launching with num_ctas > 1 runs the multi-CTA CGA-layout verifier on the rank-3 dot.
    gemm_3d_cga_split_kernel[(1, )](a_device, b_device, c_device, M, N, K, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
                                    BLOCK_K=BLOCK_K, NUM_CTAS=num_ctas, K_WIDTH=8, load_layout=load_layout,
                                    wmma_layout=wmma_layout, num_warps=num_warps, num_ctas=num_ctas)

    torch.testing.assert_close(c_device.cpu(), torch.bmm(a.float(), b.float()), rtol=2e-2, atol=2e-2)
