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

import triton
import triton.language as tl

from triton._internal_testing import (
    is_compile_warmup,
    is_ampere_or_newer,
    is_blackwell,
    is_blackwell_ultra,
    is_rubin,
    is_hip_rdna,
    is_hip_rdna3,
    is_hip_rdna4,
    is_hip_cdna,
    is_hip_cdna3,
    is_hip_cdna4,
    is_hopper_or_newer,
    is_hopper,
)
from triton.compiler import max_shared_mem
from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor
from triton.experimental import gluon
from triton.experimental.gluon import language as ttgl
from triton.experimental.gluon.language.nvidia.ampere import async_copy, mma_v2
from triton.experimental.gluon.language.nvidia.hopper import tma, mbarrier, fence_async_shared
from triton.experimental.gluon.language.nvidia import blackwell
from triton.experimental.gluon.language.nvidia import hopper
from triton.experimental.gluon.language.nvidia import rubin
from triton.experimental.gluon.language.nvidia.blackwell import tma as blackwell_tma
from triton.experimental.gluon.language.amd.cdna4 import async_copy as cdna4_async_copy
from triton.experimental.gluon.language.extra import libdevice
from triton.experimental.gluon.language.nvidia.blackwell import (
    TensorMemoryLayout,
    TensorMemoryScalesLayout,
    allocate_tensor_memory,
    tcgen05_mma_barrier_count,
    tcgen05_mma,
    tcgen05_mma_scaled,
    tcgen05_commit,
    tcgen05_copy,
    clc,
)
from triton.experimental.gluon.nvidia.hopper import TensorDescriptor
from triton._C.libtriton.gluon_ir import make_cga_layout

THREADS_PER_WARP = triton.runtime.driver.active.get_current_target().warp_size


@gluon.jit
def copy_kernel(Out, In, numel, XBLOCK: ttgl.constexpr, layout: ttgl.constexpr):
    xbase = ttgl.program_id(0) * XBLOCK
    xoffset = xbase + ttgl.arange(0, XBLOCK, layout=layout)
    xmask = xoffset < numel
    data = ttgl.load(In + xoffset, xmask)
    ttgl.store(Out + xoffset, data, xmask)


@pytest.mark.parametrize("layout", [
    ttgl.BlockedLayout(size_per_thread=[1], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[4], order=[0]),
    ttgl.BlockedLayout(size_per_thread=[2], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[4], order=[0]),
    ttgl.BlockedLayout(size_per_thread=[4], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[4], order=[0]),
    ttgl.BlockedLayout(size_per_thread=[8], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[4], order=[0]),
    ttgl.BlockedLayout(size_per_thread=[1], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[8], order=[0]),
    ttgl.BlockedLayout(size_per_thread=[2], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[8], order=[0]),
    ttgl.BlockedLayout(size_per_thread=[4], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[8], order=[0]),
    ttgl.BlockedLayout(size_per_thread=[8], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[8], order=[0]),
])
@pytest.mark.parametrize("XBLOCK", [128, 256, 512, 1024, 2048])
def test_copy_kernel(layout, XBLOCK):
    inp = torch.randn(XBLOCK * 4 - 7, device="cuda")
    out = torch.empty_like(inp)

    copy_kernel[(4, )](out, inp, inp.numel(), XBLOCK, layout, num_warps=layout.warps_per_cta[0])
    torch.testing.assert_close(out, inp)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_copy_kernel_multi_cta():
    XBLOCK = 2048
    layout = ttgl.BlockedLayout(size_per_thread=[8], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[8], order=[0],
                                cga_layout=[[1]])

    inp = torch.randn(XBLOCK * 4 - 7, device="cuda")
    out = torch.empty_like(inp)
    copy_kernel[(4, )](out, inp, inp.numel(), XBLOCK, layout, num_warps=layout.warps_per_cta[0], num_ctas=2)
    torch.testing.assert_close(out, inp)


@gluon.jit
def local_store_transposed_cga_kernel(inp, out, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr):
    src_layout: ttgl.constexpr = ttgl.BlockedLayout(
        size_per_thread=[1, BLOCK_N],
        threads_per_warp=[32, 1],
        warps_per_cta=[4, 1],
        order=[0, 1],
        cga_layout=((1, 0), ),
    )
    dst_layout: ttgl.constexpr = ttgl.BlockedLayout(
        size_per_thread=[1, BLOCK_M // 4],
        threads_per_warp=[8, 4],
        warps_per_cta=[4, 1],
        order=[1, 0],
        cga_layout=((0, 1), ),
    )
    smem_layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(
        vec=1,
        per_phase=1,
        max_phase=1,
        order=[1, 0],
        cga_layout=((1, 0), ),
    )

    src_m = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, src_layout))[:, None]
    src_n = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, src_layout))[None, :]
    src_ptrs = inp + src_m * BLOCK_N + src_n
    value = ttgl.load(src_ptrs)
    transposed = value.trans()

    smem = ttgl.allocate_shared_memory(ttgl.float32, [BLOCK_N, BLOCK_M], smem_layout)
    smem.store(transposed)
    ttgl.barrier(cluster=True)

    dst_n = ttgl.arange(0, BLOCK_N)[:, None]
    dst_m = ttgl.arange(0, BLOCK_M)[None, :]
    result = smem.load(dst_layout)
    dst_ptrs = out + dst_n * BLOCK_M + dst_m
    ttgl.store(ttgl.set_auto_layout(dst_ptrs, dst_layout), result)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_local_store_transposed_cga_to_non_transposed_alloc():
    block_m = 256
    block_n = 64
    inp = torch.arange(block_m * block_n, device="cuda", dtype=torch.float32).reshape(block_m, block_n)
    out = torch.empty((block_n, block_m), device="cuda", dtype=torch.float32)

    compiled = local_store_transposed_cga_kernel[(1, )](inp, out, block_m, block_n, num_warps=4, num_ctas=2)
    ptx = compiled.asm["ptx"]
    assert "st.shared::cluster" in ptx
    assert "ld.shared::cluster" in ptx
    torch.testing.assert_close(out, inp.T, atol=0, rtol=0)


@gluon.jit
def async_shared_store_kernel(out, BLOCK: ttgl.constexpr):
    layout: ttgl.constexpr = ttgl.BlockedLayout([1], [32], [4], [0], cga_layout=[[0]])
    shared_layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, order=[0], cga_layout=[[0]])

    offsets = ttgl.arange(0, BLOCK, layout=layout)
    values = offsets.to(ttgl.int32)
    smem = ttgl.allocate_shared_memory(ttgl.int32, [BLOCK], shared_layout)
    bar = mbarrier.allocate_mbarrier()
    mbarrier.init(bar, count=1)
    mbarrier.expect(bar, smem.nbytes_per_cta)
    hopper.async_store(smem, values, bar)
    mbarrier.wait(bar, phase=0, deps=[smem])
    result = smem.load(layout)
    mbarrier.invalidate(bar)
    ttgl.store(out + offsets, result)


@gluon.jit
def async_shared_store_f16_kernel(out, BLOCK: ttgl.constexpr):
    layout: ttgl.constexpr = ttgl.BlockedLayout([2], [32], [4], [0], cga_layout=[[0]])
    shared_layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, order=[0], cga_layout=[[0]])

    offsets = ttgl.arange(0, BLOCK, layout=layout)
    values = offsets.to(ttgl.float16)
    smem = ttgl.allocate_shared_memory(ttgl.float16, [BLOCK], shared_layout)
    bar = mbarrier.allocate_mbarrier()
    mbarrier.init(bar, count=1)
    mbarrier.expect(bar, smem.nbytes_per_cta)
    hopper.async_store(smem, values, bar)
    mbarrier.wait(bar, phase=0, deps=[smem])
    result = smem.load(layout)
    mbarrier.invalidate(bar)
    ttgl.store(out + offsets, result)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_async_shared_store():
    block = 128
    out = torch.empty((block, ), device="cuda", dtype=torch.int32)

    compiled = async_shared_store_kernel[(1, )](out, block, num_warps=4, num_ctas=2)

    assert "st.async.weak.shared::cluster.mbarrier::complete_tx::bytes" in compiled.asm["ptx"]
    torch.testing.assert_close(out, torch.arange(block, device="cuda", dtype=torch.int32))


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_async_shared_store_packed_f16():
    block = 256
    out = torch.empty((block, ), device="cuda", dtype=torch.float16)

    compiled = async_shared_store_f16_kernel[(1, )](out, block, num_warps=4, num_ctas=2)

    assert "st.async.weak.shared::cluster.mbarrier::complete_tx::bytes.b32" in compiled.asm["ptx"]
    torch.testing.assert_close(out, torch.arange(block, device="cuda", dtype=torch.float16))


@gluon.jit
def tma_kernel(desc):
    layout: ttgl.constexpr = ttgl.BlockedLayout([1, 2], [4, 8], [4, 1], [1, 0])
    value = ttgl.full(desc.block_shape, 0, desc.dtype, layout)
    alloc = ttgl.allocate_shared_memory(desc.dtype, desc.block_shape, desc.layout, value)
    tma.async_store(desc, [0, 0], alloc)
    tma.store_wait(0)
    alloc._keep_alive()


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_tma():
    out = torch.ones((16, 16), dtype=torch.float16, device="cuda")
    layout = ttgl.NVMMASharedLayout(
        swizzle_byte_width=32,
        element_bitwidth=16,
        rank=2,
        transposed=False,
        fp4_padded=False,
    )

    desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(out, [16, 16], layout)
    tma_kernel[(1, )](desc)
    torch.testing.assert_close(out, torch.zeros_like(out))


@gluon.jit
def tma_im2col_kernel(in_desc, out_desc, MULTICAST: ttgl.constexpr):
    smem = ttgl.allocate_shared_memory(in_desc.dtype, in_desc.block_shape, in_desc.layout)
    bar = mbarrier.allocate_mbarrier()
    mbarrier.init(bar, count=1)
    mbarrier.expect(bar, in_desc.block_type.nbytes)
    tma.async_load_im2col(in_desc, [0, 0, 0, 0], [0, 0], bar, smem, multicast=MULTICAST)
    mbarrier.wait(bar, phase=0)
    mbarrier.invalidate(bar)
    tma.async_store(out_desc, [0, 0], smem)
    tma.store_wait(pendings=0)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
@pytest.mark.parametrize("pixels_per_column", [32, 256, 512, 1024])
@pytest.mark.parametrize("channels_per_pixel", [32])
@pytest.mark.parametrize("swizzle_byte_width", [32])
@pytest.mark.parametrize("multicast", [False, True], ids=["unicast", "multicast"])
def test_tma_im2col(pixels_per_column, channels_per_pixel, swizzle_byte_width, multicast):
    smem_bytes = pixels_per_column * channels_per_pixel * 4 + 8192  # block + mbarrier overhead
    if smem_bytes > 200000:
        pytest.skip(f"Skipping: shared memory {smem_bytes} exceeds limit")

    inp = torch.arange(pixels_per_column * channels_per_pixel, device="cuda", dtype=torch.float32)
    inp = inp.reshape(1, 1, pixels_per_column, channels_per_pixel)
    out = torch.zeros(pixels_per_column, channels_per_pixel, device="cuda", dtype=torch.float32)

    block_shape = [pixels_per_column, channels_per_pixel]
    layout = ttgl.NVMMASharedLayout(
        swizzle_byte_width=swizzle_byte_width,
        element_bitwidth=32,
        rank=2,
        transposed=False,
        fp4_padded=False,
        cga_layout=((0, 0), ) if multicast else (),
    )

    in_desc = gluon.nvidia.hopper.TensorDescriptorIm2Col(
        base=inp,
        shape=list(inp.shape),
        strides=list(inp.stride()),
        block_shape=block_shape,
        layout=layout,
        padding="zero",
        element_strides=[1, 1, 1, 1],
        pixel_box_lower_corner=[0, 0],
        pixel_box_upper_corner=[0, 0],
    )
    out_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(out, block_shape, layout)
    compiled = tma_im2col_kernel[(1, )](in_desc, out_desc, multicast, num_warps=1, num_ctas=2 if multicast else 1)
    assert (".im2col.mbarrier::complete_tx::bytes.multicast::cluster" in compiled.asm["ptx"]) == multicast
    torch.testing.assert_close(out, inp.reshape(pixels_per_column, channels_per_pixel), atol=0, rtol=0)


@gluon.jit
def tma_round_f32_to_tf32_kernel(in_desc, out_desc):
    smem = ttgl.allocate_shared_memory(in_desc.dtype, in_desc.block_shape, in_desc.layout)
    bar = mbarrier.allocate_mbarrier()
    mbarrier.init(bar, count=1)
    mbarrier.expect(bar, in_desc.nbytes_per_cta)
    tma.async_load(in_desc, [0, 0], bar, smem)
    mbarrier.wait(bar, phase=0, deps=[smem])
    mbarrier.invalidate(bar)
    tma.async_store(out_desc, [0, 0], smem)
    tma.store_wait(0)
    smem._keep_alive()


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_gluon_tma_round_f32_to_tf32():

    def round_to_tf32(x: torch.Tensor) -> torch.Tensor:
        bits = x.view(torch.int32)
        bits_i64 = bits.to(torch.int64) & 0xFFFFFFFF
        exp_mask = 0x7F800000
        is_special = (bits_i64 & exp_mask) == exp_mask
        round_bias = ((bits_i64 >> 13) & 1) + 0x00000FFF
        rounded = (bits_i64 + round_bias) & 0xFFFFE000
        out_bits = torch.where(is_special, bits_i64, rounded)
        return (out_bits & 0xFFFFFFFF).to(torch.int32).view(torch.float32)

    torch.manual_seed(17)
    inp = torch.randn((16, 16), device="cuda", dtype=torch.float32)
    out = torch.empty_like(inp)

    layout = ttgl.NVMMASharedLayout(
        swizzle_byte_width=32,
        element_bitwidth=32,
        rank=2,
        transposed=False,
        fp4_padded=False,
    )
    in_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(inp, [16, 16], layout, round_f32_to_tf32=True)
    out_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(out, [16, 16], layout)

    tma_round_f32_to_tf32_kernel[(1, )](in_desc, out_desc)

    expected = round_to_tf32(inp)
    torch.testing.assert_close(out, expected, rtol=0, atol=0)


@gluon.jit
def tma_multicast_copy_kernel(in_desc, out_desc):
    smem = ttgl.allocate_shared_memory(in_desc.dtype, in_desc.block_shape, in_desc.layout)

    bar = mbarrier.allocate_mbarrier()
    mbarrier.init(bar, count=1)

    mbarrier.expect(bar, in_desc.nbytes_per_cta)
    tma.async_load(in_desc, [0, 0], bar, smem, multicast=True)
    mbarrier.wait(bar, phase=0, deps=[smem])

    tma.async_store(out_desc, [0, 0], smem)
    tma.store_wait(0)

    mbarrier.invalidate(bar)
    smem._keep_alive()


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
@pytest.mark.parametrize("ctas_per_cga", [[2, 1], [1, 4], [4, 4]])
def test_tma_multicast_copy(ctas_per_cga):
    cga_split_num = [min(ctas_per_cga[0], 2), min(ctas_per_cga[1], 2)]
    cga_layout = make_cga_layout(ctas_per_cga, cga_split_num, [1, 0])

    BLOCK_M, BLOCK_N = 16, 16
    BLOCK_M *= cga_split_num[0]
    BLOCK_N *= cga_split_num[1]

    inp = torch.randn((BLOCK_M, BLOCK_N), dtype=torch.float16, device="cuda")
    out = torch.empty_like(inp)

    layout = ttgl.NVMMASharedLayout.get_default_for(
        [BLOCK_M, BLOCK_N],
        ttgl.float16,
        cga_layout=cga_layout,
    )

    in_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(inp, [BLOCK_M, BLOCK_N], layout)
    out_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(out, [BLOCK_M, BLOCK_N], layout)
    num_ctas = ctas_per_cga[0] * ctas_per_cga[1]
    compiled = tma_multicast_copy_kernel[(1, )](
        in_desc,
        out_desc,
        num_warps=4,
        num_ctas=num_ctas,
    )
    expect_multicast = any(ctas_per_cga[i] > cga_split_num[i] for i in range(len(ctas_per_cga)))
    assert (".multicast::cluster" in compiled.asm["ptx"]) == expect_multicast
    torch.testing.assert_close(out, inp, atol=0, rtol=0)


@gluon.jit
def tma_gather_scatter_kernel(in_desc, gather_out_desc, scatter_out_desc, gather_idx_ptr, scatter_idx_ptr,
                              BLOCK_M: ttgl.constexpr, x_offsets_layout: ttgl.constexpr):
    smem = ttgl.allocate_shared_memory(in_desc.dtype, [BLOCK_M, gather_out_desc.block_shape[1]], gather_out_desc.layout)

    bar = mbarrier.allocate_mbarrier()
    mbarrier.init(bar, count=1)

    gather_offsets = ttgl.load(gather_idx_ptr + ttgl.arange(0, BLOCK_M, layout=x_offsets_layout))
    mbarrier.expect(bar, smem.nbytes_per_cta)
    blackwell_tma.async_gather(in_desc, gather_offsets, 0, bar, smem, multicast=True)
    mbarrier.wait(bar, phase=0, deps=[smem])

    mbarrier.invalidate(bar)

    scatter_offsets = ttgl.load(scatter_idx_ptr + ttgl.arange(0, BLOCK_M, layout=x_offsets_layout))
    tma.async_store(gather_out_desc, [0, 0], smem)
    blackwell_tma.async_scatter(scatter_out_desc, scatter_offsets, 0, smem)
    tma.store_wait(0)

    smem._keep_alive()


def get_split_dim(cga_layout, dim):
    return 1 << sum(basis[dim] != 0 for basis in cga_layout)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
@pytest.mark.parametrize("cga_layout", [
    [[1, 0]],
    [[0, 0], [1, 0]],
    [[1, 0], [0, 0]],
    [[1, 0], [2, 0]],
])
def test_tma_gather_scatter_multi_cta(cga_layout):
    cga_split_num = [get_split_dim(cga_layout, dim) for dim in range(2)]

    BLOCK_M = 32 * cga_split_num[0]
    BLOCK_N = 128 * cga_split_num[1]

    inp = torch.arange(BLOCK_M * BLOCK_N, dtype=torch.float16, device="cuda").reshape(BLOCK_M, BLOCK_N)
    gather_idx = torch.arange(BLOCK_M - 1, -1, -1, dtype=torch.int32, device="cuda")
    scatter_idx = (torch.arange(0, BLOCK_M, dtype=torch.int32, device="cuda") + 1) % BLOCK_M
    gather_out = torch.empty_like(inp)
    scatter_out = torch.zeros_like(inp)

    layout = ttgl.NVMMASharedLayout.get_default_for(
        [BLOCK_M, BLOCK_N],
        ttgl.float16,
        cga_layout=cga_layout,
    )
    in_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(inp, [1, BLOCK_N // cga_split_num[1]], layout)
    gather_out_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(gather_out, [BLOCK_M, BLOCK_N], layout)
    scatter_out_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(scatter_out, [1, BLOCK_N // cga_split_num[1]],
                                                                        layout)

    offset_layout = ttgl.BlockedLayout([4, 1], [1, 32], [4, 1], [0, 1], cga_layout=cga_layout)
    x_offsets_layout = ttgl.SliceLayout(1, offset_layout)

    num_ctas = 1 << len(cga_layout)
    compiled = tma_gather_scatter_kernel[(1, )](
        in_desc,
        gather_out_desc,
        scatter_out_desc,
        gather_idx,
        scatter_idx,
        BLOCK_M,
        x_offsets_layout,
        num_warps=4,
        num_ctas=num_ctas,
    )

    expected_gather = inp[gather_idx.to(torch.int64)]
    expected_scatter = torch.zeros_like(inp)
    expected_scatter[scatter_idx.to(torch.int64)] = expected_gather
    expect_multicast = any(all(coord == 0 for coord in basis) for basis in cga_layout)
    assert (".multicast::cluster" in compiled.asm["ptx"]) == expect_multicast
    torch.testing.assert_close(gather_out, expected_gather, atol=0, rtol=0)
    torch.testing.assert_close(scatter_out, expected_scatter, atol=0, rtol=0)


@gluon.jit
def tma_gather_scatter_subslice_kernel(in_desc, out_desc, parent_out, KIND: ttgl.constexpr):
    if KIND == "row":
        cga_layout: ttgl.constexpr = ()
    elif KIND == "column_sharded":
        cga_layout: ttgl.constexpr = ((0, 1), )
    elif KIND == "mixed_remote":
        cga_layout: ttgl.constexpr = ((1, 0), (0, 1))
    elif KIND == "remote_multicast":
        cga_layout: ttgl.constexpr = ((1, 0), (0, 0))
    else:
        cga_layout: ttgl.constexpr = ((1, 0), )
    shared_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(128, 16, rank=2, cga_layout=cga_layout)
    parent_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [4, 8], [4, 1], [1, 0], cga_layout=cga_layout)
    indices_layout: ttgl.constexpr = ttgl.SliceLayout(
        1, ttgl.BlockedLayout([4, 1], [1, 32], [4, 1], [0, 1], cga_layout=cga_layout))
    shape: ttgl.constexpr = (64, 256) if KIND == "column" or KIND == "column_sharded" else (128, 128)
    parent = ttgl.allocate_shared_memory(ttgl.float16, shape, shared_layout,
                                         value=ttgl.full(shape, -7, ttgl.float16, layout=parent_layout))
    dim: ttgl.constexpr = 1 if KIND == "column" or KIND == "column_sharded" else 0
    tile = parent.slice(shape[dim] // 2, shape[dim] // 2, dim=dim)

    indices = ttgl.arange(0, 64, layout=indices_layout)
    barrier = mbarrier.allocate_mbarrier()
    mbarrier.init(barrier, count=1)
    mbarrier.expect(barrier, tile.nbytes_per_cta)
    blackwell_tma.async_gather(in_desc, 63 - indices, 0, barrier, tile, multicast=KIND == "remote_multicast")
    mbarrier.wait(barrier, phase=0, deps=[tile])
    if KIND != "row":
        ttgl.barrier(cluster=True)

    if KIND == "row" or KIND == "column":
        blackwell_tma.async_scatter(out_desc, (indices + 1) % 64, 0, tile)
        tma.store_wait(0)

    values = parent.load(parent_layout)
    offsets = ttgl.arange(0, shape[0])[:, None] * shape[1] + ttgl.arange(0, shape[1])[None, :]
    ttgl.store(parent_out + ttgl.set_auto_layout(offsets, values.type.layout), values)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
@pytest.mark.parametrize("kind", ["row", "column", "remote", "column_sharded", "mixed_remote", "remote_multicast"])
def test_tma_gather_scatter_subslice(kind):
    values = torch.arange(64 * 128, dtype=torch.float16, device="cuda").reshape(64, 128)
    scatter_out = torch.zeros_like(values)
    parent_out = torch.empty((64, 256) if kind in ("column", "column_sharded") else (128, 128), dtype=torch.float16,
                             device="cuda")
    cga_layout = {
        "row": (),
        "column": ((1, 0), ),
        "remote": ((1, 0), ),
        "column_sharded": ((0, 1), ),
        "mixed_remote": ((1, 0), (0, 1)),
        "remote_multicast": ((1, 0), (0, 0)),
    }[kind]
    shared_layout = ttgl.NVMMASharedLayout(128, 16, rank=2, cga_layout=cga_layout)
    in_desc = TensorDescriptor.from_tensor(values, [1, 128], shared_layout)
    out_desc = TensorDescriptor.from_tensor(scatter_out, [1, 128], shared_layout)

    compiled = tma_gather_scatter_subslice_kernel[(1, )](in_desc, out_desc, parent_out, kind, num_warps=4,
                                                         num_ctas=1 << len(cga_layout))

    gathered = values.flip(0)
    expected_parent = torch.full_like(parent_out, -7)
    if kind in ("column", "column_sharded"):
        expected_parent[:, 128:] = gathered
    else:
        expected_parent[64:] = gathered
    if kind in ("remote", "column_sharded", "mixed_remote", "remote_multicast"):
        assert "mapa.shared::cluster" in compiled.asm["ptx"]
        assert "tile::gather4.cta_group::2.shared::cluster" in compiled.asm["ptx"]
    else:
        torch.testing.assert_close(scatter_out, gathered.roll(1, 0), atol=0, rtol=0)
    torch.testing.assert_close(parent_out, expected_parent, atol=0, rtol=0)


@gluon.jit
def tcgen05_mma_multicast_commit_kernel(a_desc, b_desc, out_ptrs, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                        acc_tmem_layout: ttgl.constexpr, blocked_c: ttgl.constexpr):
    smem_a = ttgl.allocate_shared_memory(a_desc.dtype, a_desc.block_shape, a_desc.layout)
    smem_b = ttgl.allocate_shared_memory(b_desc.dtype, b_desc.block_shape, b_desc.layout)
    acc_tmem = allocate_tensor_memory(ttgl.float32, [BLOCK_M, BLOCK_N], acc_tmem_layout)

    tma_bar = mbarrier.allocate_mbarrier(two_ctas=acc_tmem_layout.two_ctas)
    mbarrier.init(tma_bar, count=1)
    mma_bar = mbarrier.allocate_mbarrier()
    mbarrier.init(mma_bar, count=tcgen05_mma_barrier_count([smem_a, smem_b], True, acc_tmem.type.layout.two_ctas))

    mbarrier.expect(tma_bar, a_desc.nbytes_per_cta + b_desc.nbytes_per_cta)
    tma.async_load(a_desc, [0, 0], tma_bar, smem_a, multicast=True)
    tma.async_load(b_desc, [0, 0], tma_bar, smem_b, multicast=True)
    mbarrier.wait(tma_bar, phase=0, deps=[smem_a, smem_b])
    mbarrier.invalidate(tma_bar)

    # If it's not in a loop we don't striclty need multicast=True, but we add it to exercise the path in the test
    tcgen05_mma(smem_a, smem_b, acc_tmem, use_acc=False, multicast=True, mbarriers=[mma_bar])
    mbarrier.wait(mma_bar, phase=0, deps=[smem_a, smem_b])
    mbarrier.invalidate(mma_bar)

    out = acc_tmem.load()
    out = ttgl.convert_layout(out, blocked_c)

    out_offs_m = ttgl.arange(0, BLOCK_M)[:, None]
    out_offs_n = ttgl.arange(0, BLOCK_N)[None, :]
    out_ptrs = out_ptrs + out_offs_m * BLOCK_N + out_offs_n
    ttgl.store(out_ptrs, out)


@gluon.jit
def tma_wrong_barrier_cga_layout_kernel(a_desc, b_desc, BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                                        acc_tmem_layout: ttgl.constexpr):
    smem_a = ttgl.allocate_shared_memory(a_desc.dtype, a_desc.block_shape, a_desc.layout)
    smem_b = ttgl.allocate_shared_memory(b_desc.dtype, b_desc.block_shape, b_desc.layout)
    acc_tmem = allocate_tensor_memory(ttgl.float32, [BLOCK_M, BLOCK_N], acc_tmem_layout)

    bad_barrier_layout: ttgl.constexpr = mbarrier.MBarrierLayout(cga_layout=((0, ), (0, )))
    bad_tma_bar = ttgl.allocate_shared_memory(ttgl.int64, [1], bad_barrier_layout)
    mbarrier.init(bad_tma_bar, count=1)
    mbarrier.expect(bad_tma_bar, a_desc.nbytes_per_cta)
    tma.async_load(a_desc, [0, 0], bad_tma_bar, smem_a)

    tcgen05_mma(smem_a, smem_b, acc_tmem, use_acc=False)


def make_2cta_cga_layout(ctas_per_cga, cta_split, cta_order, two_cta_dim):
    ctas_per_cga = list(ctas_per_cga)
    cta_split = list(cta_split)
    assert cta_split[two_cta_dim] > 1
    cta_split[two_cta_dim] //= 2
    ctas_per_cga[two_cta_dim] //= 2
    aux_cga_layout = make_cga_layout(ctas_per_cga, cta_split, cta_order)
    assert two_cta_dim in (0, 1)
    basis = [0, 0]
    basis[two_cta_dim] = 1
    for b in aux_cga_layout:
        b[two_cta_dim] *= 2
    cga_layout = [basis] + aux_cga_layout
    return cga_layout


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
@pytest.mark.parametrize("ctas_per_cga", [[2, 1], [2, 4], [4, 4]])
@pytest.mark.parametrize("two_ctas", [True, False] if is_blackwell() else [False])
def test_tcgen05_mma_multicast_commit(ctas_per_cga, two_ctas):

    if two_ctas:
        ctas_per_cga_b = [ctas_per_cga[0] // 2, 2 * ctas_per_cga[1]]
    else:
        ctas_per_cga_b = ctas_per_cga
    BLOCK_M = 128 * ctas_per_cga[0]
    BLOCK_N = 64 * ctas_per_cga_b[1]
    BLOCK_K = 32

    # multicast into tcgen05_mma
    cta_split_a = [ctas_per_cga[0], 1]
    cta_split_b = [1, ctas_per_cga_b[1]]
    cta_order = [1, 0]

    if two_ctas:
        cga_layout_a = make_2cta_cga_layout(ctas_per_cga, cta_split_a, cta_order, 0)
        cga_layout_b = make_2cta_cga_layout(ctas_per_cga_b, cta_split_b, cta_order, 1)
        cga_layout_c = make_2cta_cga_layout(ctas_per_cga, ctas_per_cga, cta_order, 0)
    else:
        cga_layout_a = make_cga_layout(ctas_per_cga, cta_split_a, cta_order)
        cga_layout_b = make_cga_layout(ctas_per_cga_b, cta_split_b, cta_order)
        cga_layout_c = make_cga_layout(ctas_per_cga, ctas_per_cga, cta_order)

    shared_layout_a = ttgl.NVMMASharedLayout.get_default_for([BLOCK_M, BLOCK_K], ttgl.float16, cga_layout=cga_layout_a)
    shared_layout_b = ttgl.NVMMASharedLayout.get_default_for([BLOCK_K, BLOCK_N], ttgl.float16, cga_layout=cga_layout_b)

    a = torch.randn((BLOCK_M, BLOCK_K), dtype=torch.float16, device="cuda")
    b = torch.randn((BLOCK_K, BLOCK_N), dtype=torch.float16, device="cuda")
    out = torch.empty((BLOCK_M, BLOCK_N), dtype=torch.float32, device="cuda")

    a_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(a, [BLOCK_M, BLOCK_K], shared_layout_a)
    b_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(b, [BLOCK_K, BLOCK_N], shared_layout_b)

    tmem_shape = (128, BLOCK_N // ctas_per_cga[1])
    acc_tmem_layout = TensorMemoryLayout(
        block=tmem_shape,
        col_stride=1,
        two_ctas=two_ctas,
        cga_layout=cga_layout_c,
    )
    blocked_c = ttgl.BlockedLayout([1, 2], [ctas_per_cga[1], 32 // ctas_per_cga[1]], [4, 1], [1, 0],
                                   cga_layout=cga_layout_c)

    compiled = tcgen05_mma_multicast_commit_kernel[(1, )](
        a_desc,
        b_desc,
        out,
        BLOCK_M,
        BLOCK_N,
        acc_tmem_layout,
        blocked_c,
        num_warps=4,
        num_ctas=ctas_per_cga[0] * ctas_per_cga[1],
    )

    assert "tcgen05.commit.cta_group::" + ("2" if two_ctas else "1") in compiled.asm["ptx"]
    if two_ctas:
        assert "fence.mbarrier_init.release.cluster" in compiled.asm["ptx"]
    # For [2, 1] and two_ctas we don't multicast as there are not enough tiles
    # but we do a commit.multicast::cluster so let's grep that one instead
    assert ("multicast::cluster" in compiled.asm["ptx"])
    torch.testing.assert_close(out, torch.matmul(a.to(out.dtype), b.to(out.dtype)))


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tma_rejects_wrong_barrier_cga_layout_in_two_cta_kernel(capfd):
    ctas_per_cga = [2, 2]
    ctas_per_cga_b = [1, 4]
    block_m = 128 * ctas_per_cga[0]
    block_n = 64 * ctas_per_cga_b[1]
    block_k = 32
    cga_layout_a = make_2cta_cga_layout(ctas_per_cga, [ctas_per_cga[0], 1], [1, 0], 0)
    cga_layout_b = make_2cta_cga_layout(ctas_per_cga_b, [1, ctas_per_cga_b[1]], [1, 0], 1)
    cga_layout_c = make_2cta_cga_layout(ctas_per_cga, ctas_per_cga, [1, 0], 0)

    a = torch.empty((block_m, block_k), dtype=torch.float16, device="cuda")
    b = torch.empty((block_k, block_n), dtype=torch.float16, device="cuda")
    a_layout = ttgl.NVMMASharedLayout.get_default_for([block_m, block_k], ttgl.float16, cga_layout=cga_layout_a)
    b_layout = ttgl.NVMMASharedLayout.get_default_for([block_k, block_n], ttgl.float16, cga_layout=cga_layout_b)
    a_desc = TensorDescriptor.from_tensor(a, [block_m, block_k], a_layout)
    b_desc = TensorDescriptor.from_tensor(b, [block_k, block_n], b_layout)
    acc_tmem_layout = TensorMemoryLayout(
        block=(128, block_n // ctas_per_cga[1]),
        col_stride=1,
        two_ctas=True,
        cga_layout=cga_layout_c,
    )

    with pytest.raises(RuntimeError, match="PassManager::run failed"):
        tma_wrong_barrier_cga_layout_kernel.warmup(
            a_desc,
            b_desc,
            block_m,
            block_n,
            acc_tmem_layout,
            grid=(1, ),
            num_warps=4,
            num_ctas=4,
        )
    captured = capfd.readouterr()
    assert "TMA barrier cga_layout must be [[1], [2]] or [[0], [1]], got [[0], [0]]" in (captured.out + captured.err)


@gluon.jit
def tcgen05_mma_scaled_direct_multicast_kernel(a_desc, b_desc, out_ptr, BLOCK_M: ttgl.constexpr,
                                               BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,
                                               NUM_K_TILES: ttgl.constexpr, cga_layout_a: ttgl.constexpr,
                                               cga_layout_b: ttgl.constexpr, cga_layout_c: ttgl.constexpr,
                                               blocked_c: ttgl.constexpr):
    smem_a = ttgl.allocate_shared_memory(a_desc.dtype, a_desc.block_shape, a_desc.layout)
    smem_b = ttgl.allocate_shared_memory(b_desc.dtype, b_desc.block_shape, b_desc.layout)

    a_scale_layout: ttgl.constexpr = TensorMemoryScalesLayout(cga_layout=cga_layout_a)
    b_scale_layout: ttgl.constexpr = TensorMemoryScalesLayout(cga_layout=cga_layout_b)
    a_scale = allocate_tensor_memory(ttgl.uint8, [BLOCK_M, BLOCK_K // 32], a_scale_layout)
    b_scale = allocate_tensor_memory(ttgl.uint8, [BLOCK_N, BLOCK_K // 32], b_scale_layout)
    a_scale.store(ttgl.full([BLOCK_M, BLOCK_K // 32], 127, ttgl.uint8, a_scale.get_reg_layout()))
    b_scale.store(ttgl.full([BLOCK_N, BLOCK_K // 32], 127, ttgl.uint8, b_scale.get_reg_layout()))

    acc_layout: ttgl.constexpr = TensorMemoryLayout([128, 128], col_stride=1, cga_layout=cga_layout_c, two_ctas=True)
    acc = allocate_tensor_memory(ttgl.float32, [BLOCK_M, BLOCK_N], acc_layout)

    tma_bar = mbarrier.allocate_mbarrier(two_ctas=True)
    mbarrier.init(tma_bar, count=1)
    mma_bar = mbarrier.allocate_mbarrier()
    mbarrier.init(mma_bar, count=tcgen05_mma_barrier_count([smem_a, smem_b], True, acc.type.layout.two_ctas))

    phase_tma = 0
    phase_mma = 0
    for k in range(NUM_K_TILES):
        offs_k = k * BLOCK_K
        mbarrier.expect(tma_bar, a_desc.nbytes_per_cta + b_desc.nbytes_per_cta)
        tma.async_load(a_desc, [0, offs_k], tma_bar, smem_a, multicast=True)
        tma.async_load(b_desc, [offs_k, 0], tma_bar, smem_b, multicast=True)
        mbarrier.wait(tma_bar, phase_tma, deps=[smem_a, smem_b])
        tcgen05_mma_scaled(smem_a, smem_b, acc, a_scale, b_scale, "e5m2", "e5m2", use_acc=k != 0, multicast=True,
                           mbarriers=[mma_bar])
        mbarrier.wait(mma_bar, phase_mma, deps=[smem_a, smem_b])
        phase_tma ^= 1
        phase_mma ^= 1

    mbarrier.invalidate(tma_bar)
    mbarrier.invalidate(mma_bar)

    out = ttgl.convert_layout(acc.load(), blocked_c)
    out_offs_m = ttgl.arange(0, BLOCK_M)[:, None]
    out_offs_n = ttgl.arange(0, BLOCK_N)[None, :]
    ttgl.store(out_ptr + out_offs_m * BLOCK_N + out_offs_n, out)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tcgen05_mma_scaled_direct_multicast_barrier():
    num_ctas = 4
    BLOCK_M = 256
    BLOCK_N = 256
    BLOCK_K = 128
    NUM_K_TILES = 4
    cga_layout_a = ((1, 0), (0, 0))
    cga_layout_b = ((0, 1), (0, 2))
    cga_layout_c = ((1, 0), (0, 1))
    shared_layout_a = ttgl.NVMMASharedLayout.get_default_for([BLOCK_M, BLOCK_K], ttgl.float8e5, cga_layout=cga_layout_a)
    shared_layout_b = ttgl.NVMMASharedLayout.get_default_for([BLOCK_K, BLOCK_N], ttgl.float8e5, cga_layout=cga_layout_b)
    blocked_c = ttgl.BlockedLayout([1, 2], [2, 16], [4, 1], [1, 0], cga_layout=cga_layout_c)

    a = torch.randint(20, 40, (BLOCK_M, NUM_K_TILES * BLOCK_K), device="cuda",
                      dtype=torch.uint8).view(torch.float8_e5m2)
    b = torch.randint(20, 40, (NUM_K_TILES * BLOCK_K, BLOCK_N), device="cuda",
                      dtype=torch.uint8).view(torch.float8_e5m2)
    out = torch.empty((BLOCK_M, BLOCK_N), device="cuda", dtype=torch.float32)
    a_desc = TensorDescriptor.from_tensor(a, [BLOCK_M, BLOCK_K], shared_layout_a)
    b_desc = TensorDescriptor.from_tensor(b, [BLOCK_K, BLOCK_N], shared_layout_b)

    tcgen05_mma_scaled_direct_multicast_kernel[(1, )](a_desc, b_desc, out, BLOCK_M, BLOCK_N, BLOCK_K, NUM_K_TILES,
                                                      cga_layout_a, cga_layout_b, cga_layout_c, blocked_c, num_warps=4,
                                                      num_ctas=num_ctas)
    torch.testing.assert_close(out, torch.matmul(a.to(torch.float32), b.to(torch.float32)), atol=5e-2, rtol=1e-3)


@gluon.jit
def async_copy_mbarrier_kernel(out, inp, xnumel, XBLOCK: ttgl.constexpr, YBLOCK: ttgl.constexpr):
    smem = ttgl.allocate_shared_memory(inp.dtype.element_ty, [XBLOCK, YBLOCK],
                                       ttgl.SwizzledSharedLayout(1, 1, 1, order=[1, 0]))
    block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 4], [1, 32], [4, 1], [1, 0])
    xindex = ttgl.arange(0, XBLOCK, ttgl.SliceLayout(1, block_layout))[:, None]
    yindex = ttgl.arange(0, YBLOCK, ttgl.SliceLayout(0, block_layout))[None, :]
    mask = xindex < xnumel
    async_copy.async_load(
        smem,
        inp + xindex * YBLOCK + yindex,
        mask,
    )
    mbar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
    mbarrier.init(mbar, count=1)
    async_copy.mbarrier_arrive(mbar)
    mbarrier.arrive(mbar)
    mbarrier.wait(mbar, 0)

    val = smem.load(block_layout)
    ttgl.store(out + xindex * YBLOCK + yindex, val)


@pytest.mark.skipif(not is_ampere_or_newer(), reason="Requires Ampere")
def test_async_copy_mbarrier():
    tensor_opts = dict(dtype=torch.float, device="cuda")
    out = torch.empty((32, 32), **tensor_opts)
    inp = torch.randn((20, 32), **tensor_opts)
    async_copy_mbarrier_kernel[(1, )](out, inp, inp.shape[0], XBLOCK=32, YBLOCK=32)
    torch.testing.assert_close(out[:20], inp)
    torch.testing.assert_close(out[20:], torch.zeros((12, 32), **tensor_opts))


# Equivalence-class multicast: multicast_cta selects which CTA-ID bits to multicast
# along.  Two CTAs share a class when (id_a & ~mask) == (id_b & ~mask).
#
#   CTA ID   0x0  0x1  0x2  0x3  0x4  0x5  0x6  0x7
#   & ~0x5   0x0  0x0  0x2  0x2  0x0  0x0  0x2  0x2
#   classes: {0,1,4,5}, {2,3,6,7}  (groups of 4)
#
# Multicast requires the identity CGA layout, so two_ctas barriers are not
# supported. All legal multicast_cta values are tested; multicast_cta=0 exercises the
# non-multicast path.
@pytest.mark.parametrize("num_ctas", [2, 4, 8])
@pytest.mark.parametrize("multicast_cta", range(8))
@pytest.mark.skipif(not is_rubin(), reason="Requires Rubin")
def test_mbarrier_arrive_multicast(num_ctas, multicast_cta):
    if multicast_cta >= num_ctas:
        pytest.skip("multicast_cta must be < num_ctas")
    init_count = 1 << multicast_cta.bit_count()

    @gluon.jit
    def kernel(out_ptr, multicast_cta: ttgl.constexpr, init_count: ttgl.constexpr):
        bar = rubin.mbarrier.allocate_mbarrier()
        rubin.mbarrier.init(bar, count=init_count)
        rubin.mbarrier.arrive(bar, multicast_cta=multicast_cta)
        rubin.mbarrier.wait(bar, 0)
        rubin.mbarrier.invalidate(bar)
        ttgl.store(out_ptr + ttgl.program_id(0), ttgl.program_id(0))

    out = torch.zeros(num_ctas, device="cuda", dtype=torch.int32)
    compiled = kernel[(num_ctas, )](out, multicast_cta, init_count, num_ctas=num_ctas, num_warps=4)

    ttgir = compiled.asm["ttgir"]
    if multicast_cta:
        assert f"multicastCTA = {multicast_cta} : i32" in ttgir, "Expected multicastCTA in TTGIR"
    else:
        assert "multicastCTA" not in ttgir, "Expected no multicastCTA in TTGIR"

    torch.testing.assert_close(out, torch.arange(num_ctas, device="cuda", dtype=torch.int32))


@pytest.mark.parametrize("num_ctas", [2, 4, 8])
@pytest.mark.skipif(not is_rubin(), reason="Requires Rubin")
def test_mbarrier_arrive_broadcast_one_cta(num_ctas):

    @gluon.jit
    def kernel(out_ptr, multicast_cta: ttgl.constexpr):
        pid = ttgl.program_id(0)
        cta_rank = ttgl.inline_asm_elementwise(
            "mov.u32 $0, %cluster_ctarank;",
            "=r",
            [],
            dtype=ttgl.int32,
            is_pure=True,
            pack=1,
        )
        bar = rubin.mbarrier.allocate_mbarrier()
        rubin.mbarrier.init(bar, count=1)
        # CTA 0 does a multicast arrive, all CTAs wait.
        rubin.mbarrier.arrive(bar, multicast_cta=multicast_cta, pred=cta_rank == 0)
        rubin.mbarrier.wait(bar, 0)
        rubin.mbarrier.invalidate(bar)
        ttgl.store(out_ptr + pid, pid)

    out = torch.zeros(num_ctas, device="cuda", dtype=torch.int32)
    kernel[(num_ctas, )](out, num_ctas - 1, num_ctas=num_ctas, num_warps=4)

    torch.testing.assert_close(out, torch.arange(num_ctas, device="cuda", dtype=torch.int32))


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_device_tma_load():

    @gluon.jit
    def tma_device_load_kernel(input_ptr, output_ptr, XBLOCK: ttgl.constexpr, smem_layout: ttgl.constexpr):
        input_desc = tma.make_tensor_descriptor(
            input_ptr,
            shape=[XBLOCK, XBLOCK],
            strides=[XBLOCK, 1],
            block_shape=[XBLOCK, XBLOCK],
            layout=smem_layout,
        )

        smem = ttgl.allocate_shared_memory(ttgl.float16, [XBLOCK, XBLOCK], smem_layout)
        bar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
        mbarrier.init(bar, count=1)

        mbarrier.expect(bar, input_desc.nbytes_per_cta)
        tma.async_load(input_desc, [0, 0], bar, smem)
        mbarrier.wait(bar, 0)
        mbarrier.invalidate(bar)

        block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 2], [4, 8], [4, 1], [1, 0])
        xindex = ttgl.arange(0, XBLOCK, ttgl.SliceLayout(1, block_layout))[:, None]
        yindex = ttgl.arange(0, XBLOCK, ttgl.SliceLayout(0, block_layout))[None, :]
        val = smem.load(block_layout)
        ttgl.store(output_ptr + yindex + xindex * XBLOCK, val)

    XBLOCK = 16
    input = torch.zeros((XBLOCK, XBLOCK), device="cuda", dtype=torch.float16)
    output = torch.ones_like(input)
    smem_layout = ttgl.NVMMASharedLayout(
        swizzle_byte_width=32,
        element_bitwidth=16,
        rank=2,
        transposed=False,
        fp4_padded=False,
    )

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

    triton.set_allocator(alloc_fn)

    tma_device_load_kernel[(1, )](input, output, XBLOCK, smem_layout)
    torch.testing.assert_close(input, output)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_device_tma_store():

    @gluon.jit
    def tma_device_store_kernel(out_ptr, XBLOCK: ttgl.constexpr, smem_layout: ttgl.constexpr):
        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 2], [4, 8], [4, 1], [1, 0])
        value = ttgl.full([XBLOCK, XBLOCK], 0, ttgl.float16, layout)
        alloc = ttgl.allocate_shared_memory(ttgl.float16, [XBLOCK, XBLOCK], smem_layout, value)
        out_desc = tma.make_tensor_descriptor(
            out_ptr,
            shape=[XBLOCK, XBLOCK],
            strides=[XBLOCK, 1],
            block_shape=[XBLOCK, XBLOCK],
            layout=smem_layout,
        )
        tma.async_store(out_desc, [0, 0], alloc)
        tma.store_wait(0)
        alloc._keep_alive()

    XBLOCK = 16
    out = torch.ones((XBLOCK, XBLOCK), dtype=torch.float16, device="cuda")
    smem_layout = ttgl.NVMMASharedLayout(
        swizzle_byte_width=32,
        element_bitwidth=16,
        rank=2,
        transposed=False,
        fp4_padded=False,
    )

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

    triton.set_allocator(alloc_fn)

    tma_device_store_kernel[(1, )](out, XBLOCK, smem_layout)
    torch.testing.assert_close(out, torch.zeros_like(out))


@gluon.jit
def mma_kernel(a, b, out, M: ttgl.constexpr, N: ttgl.constexpr, K: ttgl.constexpr, block_layout_a: ttgl.constexpr,
               block_layout_b: ttgl.constexpr, cga_layout_c: ttgl.constexpr, acc_layout: ttgl.constexpr,
               shared_layout_a: ttgl.constexpr, shared_layout_b: ttgl.constexpr, acc_dtype: ttgl.constexpr,
               ASYNC: ttgl.constexpr, USE_TCGEN05: ttgl.constexpr):
    a_offs_m = ttgl.arange(0, M)[:, None]
    a_offs_k = ttgl.arange(0, K)[None, :]
    b_offs_k = ttgl.arange(0, K)[:, None]
    b_offs_n = ttgl.arange(0, N)[None, :]

    operand_dtype = a.dtype.element_ty
    a_ptrs = a + a_offs_m * K + a_offs_k
    b_ptrs = b + b_offs_k * N + b_offs_n
    a_tile = ttgl.load(ttgl.set_auto_layout(a_ptrs, block_layout_a))
    b_tile = ttgl.load(ttgl.set_auto_layout(b_ptrs, block_layout_b))

    smem_a = ttgl.allocate_shared_memory(operand_dtype, [M, K], shared_layout_a, a_tile)
    smem_b = ttgl.allocate_shared_memory(operand_dtype, [K, N], shared_layout_b, b_tile)

    if USE_TCGEN05:
        two_ctas: ttgl.constexpr = acc_layout.two_ctas
        fence_async_shared(cluster=two_ctas)
        mma_barrier = mbarrier.allocate_mbarrier()
        mbarrier.init(mma_barrier, count=1)

        acc_tmem = allocate_tensor_memory(acc_dtype, [M, N], acc_layout)

        tcgen05_mma(smem_a, smem_b, acc_tmem, use_acc=False, mbarriers=[mma_barrier])
        mbarrier.wait(mma_barrier, phase=0, deps=[smem_a, smem_b])
        mbarrier.invalidate(mma_barrier)

        acc = acc_tmem.load()
    else:
        acc = ttgl.zeros([M, N], dtype=acc_dtype, layout=acc_layout)
        acc = hopper.warpgroup_mma(smem_a, smem_b, acc, is_async=ASYNC)

        if ASYNC:
            acc = hopper.warpgroup_mma_wait(num_outstanding=0, deps=[acc])

    out_offs_m = ttgl.arange(0, M)[:, None]
    out_offs_n = ttgl.arange(0, N)[None, :]
    out_ptrs = out + out_offs_m * N + out_offs_n
    ttgl.store(out_ptrs, acc)


@pytest.mark.skipif(not is_hopper(), reason="Requires Hopper")
@pytest.mark.parametrize("ASYNC", [True, False])
def test_warpgroup_mma(ASYNC):
    torch.manual_seed(0)
    M, N, K = 64, 32, 32
    warps = [4, 1]
    block_layout = ttgl.BlockedLayout([1, 1], [1, THREADS_PER_WARP], warps_per_cta=warps, order=[1, 0])
    acc_layout = ttgl.NVMMADistributedLayout(version=[3, 0], warps_per_cta=warps, instr_shape=[16, 32, 16])
    shared_layout_a = ttgl.NVMMASharedLayout.get_default_for([M, K], ttgl.float16)
    shared_layout_b = ttgl.NVMMASharedLayout.get_default_for([K, N], ttgl.float16)
    a = torch.randn((M, K), device="cuda", dtype=torch.float16)
    b = torch.randn((K, N), device="cuda", dtype=torch.float16)
    out = torch.zeros((M, N), device="cuda", dtype=torch.float16)
    mma_kernel[(1, )](
        a,
        b,
        out,
        M,
        N,
        K,
        block_layout,
        block_layout,
        block_layout,
        acc_layout,
        shared_layout_a,
        shared_layout_b,
        ttgl.float16,
        ASYNC,
        False,
        num_warps=warps[0] * warps[1],
    )

    ref = torch.matmul(a, b)

    torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-1)


@gluon.jit
def warpgroup_mma_wait_multi_warpgroup_kernel(a_ptr, b0_ptr, b1_ptr, out_ptr, D: ttgl.constexpr, NBLK: ttgl.constexpr,
                                              WG: ttgl.constexpr):
    M: ttgl.constexpr = 64 * WG
    c_layout: ttgl.constexpr = ttgl.NVMMADistributedLayout([3, 0], warps_per_cta=[4 * WG, 1], instr_shape=[16, 64, 16])
    smem_b: ttgl.constexpr = ttgl.NVMMASharedLayout.get_default_for([D, D], ttgl.bfloat16)
    smem_a: ttgl.constexpr = ttgl.NVMMASharedLayout.get_default_for([M, D], ttgl.bfloat16)
    layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [4, 8], [4 * WG, 1], [1, 0])

    a_smem = ttgl.allocate_shared_memory(ttgl.bfloat16, [M, D], smem_a)
    b0_smem = ttgl.allocate_shared_memory(ttgl.bfloat16, [D, D], smem_b)
    b1_smem = ttgl.allocate_shared_memory(ttgl.bfloat16, [D, D], smem_b)

    rows = ttgl.arange(0, D, layout=ttgl.SliceLayout(1, layout))
    cols = ttgl.arange(0, D, layout=ttgl.SliceLayout(0, layout))
    offs = rows[:, None] * D + cols[None, :]
    aoffs = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, layout))[:, None] * D + cols[None, :]

    acc = ttgl.zeros([M, D], ttgl.float32, c_layout)
    inflight = hopper.warpgroup_mma_init(acc)

    async_copy.async_load(a_smem, a_ptr + aoffs)
    async_copy.commit_group()

    for i in tl.range(0, NBLK):
        async_copy.async_load(b0_smem, b0_ptr + i * D * D + offs)
        acc = hopper.warpgroup_mma_wait(num_outstanding=0, deps=(inflight, ))
        async_copy.async_load(b1_smem, b1_ptr + i * D * D + offs)
        async_copy.commit_group()
        async_copy.wait_group(0)
        hopper.fence_async_shared()
        acc = hopper.warpgroup_mma(a_smem, b0_smem, acc, use_acc=True, is_async=True)
        acc = hopper.warpgroup_mma_wait(num_outstanding=0, deps=(acc, ))
        inflight = hopper.warpgroup_mma(a_smem, b1_smem, acc, use_acc=True, is_async=True)
    acc = hopper.warpgroup_mma_wait(num_outstanding=0, deps=(inflight, ))

    out_r = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, c_layout))
    out_c = ttgl.arange(0, D, layout=ttgl.SliceLayout(0, c_layout))
    pid = ttgl.program_id(0)
    ttgl.store(out_ptr + pid.to(ttgl.int64) * (M * D) + out_r[:, None] * D + out_c[None, :], acc)


@pytest.mark.skipif(not is_hopper(), reason="Requires Hopper")
def test_warpgroup_mma_wait_synchronizes_warpgroups():
    torch.manual_seed(0)
    d, num_warp_groups = 64, 2
    m = 64 * num_warp_groups
    num_blocks, grid = 512, 528

    b0 = torch.randint(-3, 4, (num_blocks, d, d), device="cuda", dtype=torch.float32)
    b1 = torch.randint(-3, 4, (num_blocks, d, d), device="cuda", dtype=torch.float32)
    expected = (b0.sum(dim=(0, 1)) + b1.sum(dim=(0, 1)))[None, :].expand(m, d)
    assert expected.abs().max() < 2**24

    a = torch.ones((m, d), device="cuda", dtype=torch.bfloat16)
    b0 = b0.to(torch.bfloat16)
    b1 = b1.to(torch.bfloat16)

    for _ in range(5):
        out = torch.empty((grid, m, d), dtype=torch.float32, device="cuda")
        warpgroup_mma_wait_multi_warpgroup_kernel[(grid, )](a, b0, b1, out, D=d, NBLK=num_blocks, WG=num_warp_groups,
                                                            num_warps=4 * num_warp_groups)
        torch.testing.assert_close(out, expected[None, :, :].expand_as(out), atol=0, rtol=0)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tcgen05_mma_two_ctas_transposed_lhs_shared_layout():
    torch.manual_seed(0)
    BLOCK_M, N, K = 64, 128, 256
    M = 2 * BLOCK_M
    warps = [4, 1]
    cga_layout_a = ((1, 0), )
    cga_layout_b = ((0, 1), )
    cga_layout_c = ((1, 0), )

    block_layout_a = ttgl.BlockedLayout([1, 8], [1, THREADS_PER_WARP], warps_per_cta=warps, order=[0, 1],
                                        cga_layout=cga_layout_a)
    block_layout_b = ttgl.BlockedLayout([1, 8], [1, THREADS_PER_WARP], warps_per_cta=warps, order=[1, 0],
                                        cga_layout=cga_layout_b)
    shared_layout_a = ttgl.NVMMASharedLayout(
        swizzle_byte_width=128,
        transposed=True,
        element_bitwidth=16,
        rank=2,
        cga_layout=cga_layout_a,
    )
    shared_layout_b = ttgl.NVMMASharedLayout.get_default_for([K, N], ttgl.float16, cga_layout=cga_layout_b)
    acc_layout = TensorMemoryLayout(
        block=(BLOCK_M, N),
        col_stride=1,
        cga_layout=cga_layout_c,
        two_ctas=True,
    )

    device = triton.runtime.driver.active.get_current_device()
    a = torch.randn((M, K), device=device, dtype=torch.float16)
    b = torch.randn((K, N), device=device, dtype=torch.float16)
    out = torch.empty((M, N), device=device, dtype=torch.float32)

    compiled = mma_kernel[(1, )](
        a,
        b,
        out,
        M,
        N,
        K,
        block_layout_a,
        block_layout_b,
        cga_layout_c,
        acc_layout,
        shared_layout_a,
        shared_layout_b,
        ttgl.float32,
        False,
        True,
        num_warps=warps[0] * warps[1],
        num_ctas=2,
    )

    assert "tcgen05.mma.cta_group::2" in compiled.asm["ptx"]
    ref = torch.matmul(a.to(out.dtype), b.to(out.dtype))
    torch.testing.assert_close(out, ref, atol=5e-2, rtol=5e-1)


@gluon.jit
def tma_mma_shared_inputs_kernel(a_desc, b_desc, out_ptr, out_desc, gather_idx_ptr, scatter_idx_ptr,
                                 BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, BLOCK_K: ttgl.constexpr,
                                 NUM_K_TILES: ttgl.constexpr, block_layout_c: ttgl.constexpr,
                                 acc_layout: ttgl.constexpr, acc_tmem_layout: ttgl.constexpr,
                                 use_tcgen05: ttgl.constexpr, multicast: ttgl.constexpr,
                                 use_gather_scatter: ttgl.constexpr):
    smem_a = ttgl.allocate_shared_memory(a_desc.dtype, [BLOCK_M, BLOCK_K], a_desc.layout)
    smem_b = ttgl.allocate_shared_memory(b_desc.dtype, b_desc.block_shape, b_desc.layout)

    two_ctas: ttgl.constexpr = isinstance(acc_tmem_layout, TensorMemoryLayout) and acc_tmem_layout.two_ctas

    tma_bar = mbarrier.allocate_mbarrier(two_ctas=two_ctas)
    mbarrier.init(tma_bar, count=1)
    phase_tma = 0

    if use_tcgen05:
        acc_tmem = allocate_tensor_memory(
            element_ty=ttgl.float32,
            shape=[BLOCK_M, BLOCK_N],
            layout=acc_tmem_layout,
        )
        mma_bar = mbarrier.allocate_mbarrier()
        phase_mma = 0
        mbarrier.init(mma_bar, count=tcgen05_mma_barrier_count([smem_a, smem_b], multicast,
                                                               acc_tmem.type.layout.two_ctas))
    else:
        acc = ttgl.zeros([BLOCK_M, BLOCK_N], dtype=ttgl.float32, layout=acc_layout)

    if use_gather_scatter:
        gather_offsets_layout: ttgl.constexpr = ttgl.SliceLayout(
            1, ttgl.BlockedLayout([4, 1], [1, 32], [ttgl.num_warps(), 1], [0, 1], cga_layout=a_desc.layout.cga_layout))
        gather_offsets = ttgl.load(gather_idx_ptr + ttgl.arange(0, BLOCK_M, layout=gather_offsets_layout))

    for k in range(NUM_K_TILES):
        mbarrier.expect(tma_bar, smem_a.nbytes_per_cta + smem_b.nbytes_per_cta)
        if use_gather_scatter:
            blackwell_tma.async_gather(a_desc, gather_offsets, k * BLOCK_K, tma_bar, smem_a, multicast=multicast)
        else:
            tma.async_load(a_desc, [0, k * BLOCK_K], tma_bar, smem_a, multicast=multicast)
        tma.async_load(b_desc, [k * BLOCK_K, 0], tma_bar, smem_b, multicast=multicast)
        mbarrier.wait(tma_bar, phase=phase_tma, deps=[smem_a, smem_b])
        phase_tma ^= 1

        if use_tcgen05:
            tcgen05_mma(smem_a, smem_b, acc_tmem, use_acc=(k != 0), multicast=multicast, mbarriers=[mma_bar])
            mbarrier.wait(mma_bar, phase=phase_mma, deps=[smem_a, smem_b])
            phase_mma ^= 1
        else:
            acc = hopper.warpgroup_mma(smem_a, smem_b, acc, is_async=False)
            if multicast:
                # multicast into wgmma doesn't make much sense as you need to synchronise all
                # CTAs after the wgmma, as it doesn't provide a finer synchronization mechanism.
                ttgl.barrier(cluster=True)

    mbarrier.invalidate(tma_bar)

    if use_tcgen05:
        mbarrier.invalidate(mma_bar)
        acc = acc_tmem.load()

    acc = ttgl.convert_layout(acc, block_layout_c)
    if use_gather_scatter:
        scatter_offsets_layout: ttgl.constexpr = ttgl.SliceLayout(
            1, ttgl.BlockedLayout([4, 1], [1, 32], [ttgl.num_warps(), 1], [0, 1],
                                  cga_layout=out_desc.layout.cga_layout))
        scatter_offsets = ttgl.load(scatter_idx_ptr + ttgl.arange(0, BLOCK_M, layout=scatter_offsets_layout))
        acc_smem = ttgl.allocate_shared_memory(out_desc.dtype, [BLOCK_M, BLOCK_N], out_desc.layout, acc)
        blackwell_tma.async_scatter(out_desc, scatter_offsets, 0, acc_smem)
        tma.store_wait(0)
        acc_smem._keep_alive()
    else:
        offs_m = ttgl.arange(0, BLOCK_M)[:, None]
        offs_n = ttgl.arange(0, BLOCK_N)[None, :]
        ttgl.store(out_ptr + offs_m * BLOCK_N + offs_n, acc)


@pytest.mark.skipif(not (is_hopper() or is_blackwell()), reason="Requires Hopper or Blackwell")
@pytest.mark.parametrize("warps", ([8, 1], [4, 2], [4, 1]))
@pytest.mark.parametrize("reps", ([1, 1, 1], [2, 2, 2], [1, 4, 2]))
@pytest.mark.parametrize("ctas_per_cga", [[1, 1], [2, 1], [4, 4]])
@pytest.mark.parametrize("two_ctas", [False, True] if is_blackwell() else [False])
@pytest.mark.parametrize("multicast", [False, True])
@pytest.mark.parametrize("use_gather_scatter", [False, True] if is_blackwell() else [False])
def test_tma_mma_shared_inputs(warps, reps, ctas_per_cga, two_ctas, multicast, use_gather_scatter):
    bitwidth = 16
    acc_dtype = torch.float32

    if ctas_per_cga[0] == 1 and two_ctas:
        pytest.skip("Need at least 2 CTAs along M for 2CTA mode")

    cta_order = [1, 0]

    if two_ctas:
        assert ctas_per_cga[0] >= 2, "Need at least 2 CTAs along M for 2CTA mode"
        ctas_per_cga_b = [ctas_per_cga[0] // 2, 2 * ctas_per_cga[1]]
    else:
        ctas_per_cga_b = ctas_per_cga
    cta_split_a = [ctas_per_cga[0], 1]
    cta_split_b = [1, ctas_per_cga_b[1]]

    # M = 128 for blackkwell
    instr_shape = [32 if is_blackwell() else 16, 32, 256 // bitwidth]
    NUM_K_TILES = 4
    BLOCK_M = instr_shape[0] * warps[0] * ctas_per_cga[0] * reps[0]
    BLOCK_N = instr_shape[1] * warps[1] * ctas_per_cga_b[1] * reps[1]
    if is_blackwell() and BLOCK_N >= 256 * ctas_per_cga[1]:
        # tcgen05 doesn't support reps along N
        BLOCK_N = 256 * ctas_per_cga[1]
    BLOCK_K = instr_shape[2] * reps[2]
    K = (256 // bitwidth) * NUM_K_TILES

    if two_ctas:
        cga_layout_a = make_2cta_cga_layout(ctas_per_cga, cta_split_a, cta_order, 0)
        cga_layout_b = make_2cta_cga_layout(ctas_per_cga_b, cta_split_b, cta_order, 1)
        cga_layout_c = make_2cta_cga_layout(ctas_per_cga, ctas_per_cga, cta_order, 0)
    else:
        cga_layout_a = make_cga_layout(ctas_per_cga, cta_split_a, cta_order)
        cga_layout_b = make_cga_layout(ctas_per_cga_b, cta_split_b, cta_order)
        cga_layout_c = make_cga_layout(ctas_per_cga, ctas_per_cga, cta_order)

    block_layout_c = ttgl.BlockedLayout([1, 8], [1, THREADS_PER_WARP], warps_per_cta=warps, order=[1, 0],
                                        cga_layout=cga_layout_c)

    acc_layout = ttgl.NVMMADistributedLayout(version=[3, 0], warps_per_cta=warps, instr_shape=instr_shape,
                                             cga_layout=cga_layout_c)

    tmem_shape = (min(BLOCK_M // ctas_per_cga[0], 128), BLOCK_N // ctas_per_cga[1])
    acc_tmem_layout = TensorMemoryLayout(
        block=tmem_shape,
        col_stride=1,
        cga_layout=cga_layout_c,
        two_ctas=two_ctas,
    )

    def cast(x, dtype):
        if dtype != torch.float32:
            return x.to(dtype)
        # For b16 and fp32 (in both hopper and blackwell it seems)
        # Element-wise multiplication of matrix A and B is performed with specified precision.
        # wgmma.mma_async operation involving type .tf32 will truncate lower 13 bits of the 32-bit
        # input data before multiplication is issued
        x = x.view(torch.int32)
        x = x & ~((1 << 13) - 1)
        return x.view(dtype)

    torch_dtype = torch.float16
    device = triton.runtime.driver.active.get_current_device()
    a = cast(torch.randn((BLOCK_M, K), device=device, dtype=torch.float32), torch_dtype)
    # We transpose b in the kernel
    b = cast(torch.randn((K, BLOCK_N), device=device, dtype=torch.float32), torch_dtype)
    out = torch.empty((BLOCK_M, BLOCK_N), device=device, dtype=acc_dtype)

    gluon_dtype = ttgl.float16
    shared_layout_a = ttgl.NVMMASharedLayout.get_default_for([BLOCK_M, BLOCK_K], gluon_dtype, cga_layout=cga_layout_a)
    shared_layout_b = ttgl.NVMMASharedLayout.get_default_for([BLOCK_K, BLOCK_N], gluon_dtype, cga_layout=cga_layout_b)
    shared_layout_c = ttgl.NVMMASharedLayout.get_default_for([BLOCK_M, BLOCK_N], ttgl.float32, cga_layout=cga_layout_c)
    assert shared_layout_a.swizzle_byte_width != 0
    assert shared_layout_b.swizzle_byte_width != 0
    a_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(a, [1 if use_gather_scatter else BLOCK_M, BLOCK_K],
                                                              shared_layout_a)
    b_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(b, [BLOCK_K, BLOCK_N], shared_layout_b)
    num_warps = warps[0] * warps[1]
    num_ctas = ctas_per_cga[0] * ctas_per_cga[1]
    out_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(out, [1, BLOCK_N], shared_layout_c)
    gather_idx = torch.arange(BLOCK_M - 1, -1, -1, dtype=torch.int32, device=device)
    scatter_idx = (torch.arange(0, BLOCK_M, dtype=torch.int32, device=device) + 1) % BLOCK_M

    try:
        tma_mma_shared_inputs_kernel[(1, )](
            a_desc,
            b_desc,
            out,
            out_desc,
            gather_idx,
            scatter_idx,
            BLOCK_M,
            BLOCK_N,
            BLOCK_K,
            NUM_K_TILES,
            block_layout_c,
            acc_layout,
            acc_tmem_layout,
            is_blackwell(),
            multicast=multicast,
            use_gather_scatter=use_gather_scatter,
            num_warps=num_warps,
            num_ctas=num_ctas,
        )
    except triton.runtime.errors.OutOfResources:
        pytest.skip("Too much shared memory required")

    try:
        allow_tf32 = torch.backends.cuda.matmul.allow_tf32
        torch.backends.cuda.matmul.allow_tf32 = True
        if use_gather_scatter:
            matmul = torch.matmul(a[gather_idx.to(torch.int64)].to(torch.float32), b.to(torch.float32))
            ref = torch.empty_like(matmul)
            # Correct as scatter_idx is a permutation!
            ref[scatter_idx.to(torch.int64)] = matmul
        else:
            ref = torch.matmul(a.to(torch.float32), b.to(torch.float32))
    finally:
        torch.backends.cuda.matmul.allow_tf32 = allow_tf32

    if bitwidth == 8:
        atol, rtol = 8e-2, 8e-1
    elif bitwidth == 16:
        atol, rtol = 5e-2, 5e-1
    else:
        atol, rtol = 8e-4, 8e-3
    torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)


@pytest.mark.skipif(not (is_hopper() or is_blackwell()), reason="Requires Hopper or Blackwell")
@pytest.mark.parametrize("bitwidth, transpose_a, transpose_b, acc_dtype",
                         [(bitwidth, transpose_a, transpose_b, acc_dtype)
                          for bitwidth in [8, 16, 32]
                          for (transpose_a, transpose_b) in product([False, True], repeat=2)
                          for acc_dtype in [torch.float16, torch.float32]
                          if bitwidth == 16 or (acc_dtype == torch.float32 and not transpose_a and transpose_b)])
@pytest.mark.parametrize("warps", ([8, 1], [4, 2], [4, 1]))
@pytest.mark.parametrize("swizzling_a, swizzling_b", product([0, 32, 64, 128], repeat=2))
@pytest.mark.parametrize("instr_m", [64, 128] if is_blackwell() else [64])
@pytest.mark.parametrize("shape_m, shape_n, shape_k", [(1, 1, 1), (2, 4, 1), (2, 2, 4)])
@pytest.mark.parametrize("ctas_per_cga", [[1, 1], [2, 1], [4, 4]])
@pytest.mark.parametrize("two_ctas", [False, True] if is_blackwell() else [False])
@pytest.mark.enable_warmup(min_capability=9, priority=1)
def test_mma_shared_inputs(bitwidth, transpose_a, transpose_b, acc_dtype, warps, swizzling_a, swizzling_b, instr_m,
                           shape_m, shape_n, shape_k, ctas_per_cga, two_ctas):
    # FIXME: Workaround for a bug in PTXAS when the shared layout is transposed and the swizzling is 0
    # This is fixed in PTXAS 13.0.88. Remove once we upgrade
    if bitwidth == 16 and ((transpose_a and swizzling_a == 0 and shape_m > 1) or
                           (not transpose_b and swizzling_b == 0 and shape_n > 1)):
        pytest.skip("Skipped due to a bug in PTXAS when the shared layout is transposed and the swizzling is 0")
    if ctas_per_cga[0] == 1 and two_ctas:
        pytest.skip("Need at least 2 CTAs for 2CTA mode")

    use_tcgen05 = is_blackwell()

    torch_dtype_map = {
        8: torch.float8_e4m3fn,
        16: torch.float16,
        32: torch.float32,
    }
    acc_dtype_map = {
        torch.float16: ttgl.float16,
        torch.float32: ttgl.float32,
    }

    # We'll choose a larger instr shape along N, but sure
    # instr_m is the instruction per warp group so we divide by 4
    instr_shape = [instr_m // 4, 32, 256 // bitwidth]
    M = instr_shape[0] * warps[0]
    N = instr_shape[1] * warps[1]
    K = instr_shape[2]

    if two_ctas:
        assert ctas_per_cga[0] >= 2, "Need at least 2 CTAs along M for 2CTA mode"
        ctas_per_cga_b = [ctas_per_cga[0] // 2, 2 * ctas_per_cga[1]]
    else:
        ctas_per_cga_b = ctas_per_cga
    cta_split_a = [ctas_per_cga[0], 1]
    cta_split_b = [1, ctas_per_cga_b[1]]

    def min_shape(swizzling, dim0, dim1, trans):
        tile_cols = (8 * max(16, swizzling)) // bitwidth
        outer_dim, contig_dim = (dim0, dim1)
        if trans:
            outer_dim, contig_dim = contig_dim, outer_dim
        contig_dim = max(contig_dim, tile_cols)
        outer_dim = max(outer_dim, 8)
        if trans:
            outer_dim, contig_dim = contig_dim, outer_dim
        return outer_dim, contig_dim

    # Get the minimum shape for the given swizzling / transpose
    M, K = min_shape(swizzling_a, M, K, transpose_a)
    K, N = min_shape(swizzling_b, K, N, transpose_b)
    M *= shape_m * ctas_per_cga[0]
    N *= shape_n * ctas_per_cga_b[1]
    K *= shape_k
    instr_shape[1] *= shape_n

    num_warps = warps[0] * warps[1]
    num_ctas = ctas_per_cga[0] * ctas_per_cga[1]

    if is_blackwell():
        # Avoid too many rows in TMEM
        MAX_ROWS = 512
        if M * N // 128 // num_ctas > MAX_ROWS:
            N //= (M * N // 128 // num_ctas // MAX_ROWS)

    total_shmem = (M + N) * K * bitwidth // 8
    device = triton.runtime.driver.active.get_current_device()
    MAX_SHMEM = max_shared_mem(device)
    if total_shmem > MAX_SHMEM:
        pytest.skip(f"Total shared memory {total_shmem} bytes exceeds the maximum allowed {MAX_SHMEM} bytes")

    if two_ctas and N // ctas_per_cga[1] == 512:
        # grep for [Note: numRepN > 1 and two_ctas]
        pytest.skip("grep for [Note: numRepN > 1 and two_ctas]")

    assert M >= 64, "M must be at least 64 for mmav3 and mmav5"

    def log2_int(x):
        return x.bit_length() - 1

    def get_shared_swizzling_zero(M, K, transpose, cga_layout):
        if cga_layout:
            dim_cga = [1, 1]
            for b in cga_layout:
                for i, bi in enumerate(b):
                    if bi != 0:
                        dim_cga[i] *= 2
            cta_shape = (M // dim_cga[0], K // dim_cga[1])
            cta_layout = get_shared_swizzling_zero(cta_shape[0], cta_shape[1], transpose, None)
            cga_bases = list(cga_layout)
            for b in cga_bases:
                for i in range(len(b)):
                    b[i] *= cta_shape[i]
            return ttgl.SharedLinearLayout(cta_layout.offset_bases, cga_bases)
        if transpose:
            assert not cga_layout
            shared = get_shared_swizzling_zero(K, M, False, cga_layout)
            # Transpose the bases
            bases = list(shared.offset_bases)
            for i in range(len(bases)):
                bases[i] = [bases[i][1], bases[i][0]]
            return ttgl.SharedLinearLayout(bases)
        bases = []
        for i in range(log2_int(128 // bitwidth)):
            bases.append([0, 1 << i])
        for i in range(log2_int(M)):
            bases.append([1 << i, 0])
        for i in range(log2_int(K // (128 // bitwidth))):
            offset = int(math.log2(128 // bitwidth)) + i
            bases.append([0, 1 << offset])
        return ttgl.SharedLinearLayout(bases)

    torch_dtype = torch_dtype_map[bitwidth]
    gl_acc_dtype = acc_dtype_map[acc_dtype]
    out_dtype = torch.float32
    cta_order = [1, 0]

    if two_ctas:
        cga_layout_a = make_2cta_cga_layout(ctas_per_cga, cta_split_a, cta_order, 0)
        cga_layout_b = make_2cta_cga_layout(ctas_per_cga_b, cta_split_b, cta_order, 1)
        # The TMEM layout for instr_m == 128 splits along M, the one for instr_m == 64 splits along N
        cga_layout_c = make_2cta_cga_layout(ctas_per_cga, ctas_per_cga, cta_order, 0)
    else:
        cga_layout_a = make_cga_layout(ctas_per_cga, cta_split_a, cta_order)
        cga_layout_b = make_cga_layout(ctas_per_cga_b, cta_split_b, cta_order)
        cga_layout_c = make_cga_layout(ctas_per_cga, ctas_per_cga, cta_order)
    cga_layout_c = tuple(tuple(basis) for basis in cga_layout_c)

    block_layout_a = ttgl.BlockedLayout([1, 8], [1, THREADS_PER_WARP], warps_per_cta=warps, order=[0, 1],
                                        cga_layout=cga_layout_a)
    block_layout_b = ttgl.BlockedLayout([1, 8], [1, THREADS_PER_WARP], warps_per_cta=warps, order=[1, 0],
                                        cga_layout=cga_layout_b)
    if swizzling_a == 0:
        shared_layout_a = get_shared_swizzling_zero(M, K, transpose_a, cga_layout_a)
    else:
        shared_layout_a = ttgl.NVMMASharedLayout(swizzle_byte_width=swizzling_a, element_bitwidth=bitwidth, rank=2,
                                                 transposed=transpose_a, cga_layout=cga_layout_a)
    if swizzling_b == 0:
        shared_layout_b = get_shared_swizzling_zero(K, N, transpose_b, cga_layout_b)
    else:
        shared_layout_b = ttgl.NVMMASharedLayout(swizzle_byte_width=swizzling_b, element_bitwidth=bitwidth, rank=2,
                                                 transposed=transpose_b, cga_layout=cga_layout_b)
    if use_tcgen05:
        tmem_shape = (instr_m, min(N // ctas_per_cga[1], 256))
        acc_layout = TensorMemoryLayout(tmem_shape, col_stride=32 // torch.finfo(acc_dtype).bits,
                                        cga_layout=cga_layout_c, two_ctas=two_ctas)
    else:
        acc_layout = ttgl.NVMMADistributedLayout(version=[3, 0], warps_per_cta=warps, instr_shape=instr_shape,
                                                 cga_layout=cga_layout_c)

    torch.manual_seed(0)

    def cast(x, dtype):
        if dtype != torch.float32:
            return x.to(dtype)
        # For b16 and fp32 (in both hopper and blackwell it seems)
        # Element-wise multiplication of matrix A and B is performed with specified precision.
        # wgmma.mma_async operation involving type .tf32 will truncate lower 13 bits of the 32-bit
        # input data before multiplication is issued
        x = x.view(torch.int32)
        x = x & ~((1 << 13) - 1)
        return x.view(dtype)

    # Sample bf16 as tf32 does not use the full range
    a = cast(torch.randn((M, K), device=device, dtype=torch.float32), torch_dtype)
    b = cast(torch.randn((K, N), device=device, dtype=torch.float32), torch_dtype)
    out = torch.zeros((M, N), device=device, dtype=out_dtype)

    compiled = mma_kernel[(1, )](
        a,
        b,
        out,
        M,
        N,
        K,
        block_layout_a,
        block_layout_b,
        cga_layout_c,
        acc_layout,
        shared_layout_a,
        shared_layout_b,
        gl_acc_dtype,
        False,
        use_tcgen05,
        num_warps=num_warps,
        num_ctas=num_ctas,
    )
    if is_compile_warmup():
        return

    assert two_ctas == ("two_ctas" in compiled.asm["ttgir"])
    if two_ctas:
        assert "fence.mbarrier_init.release.cluster" in compiled.asm["ptx"]

    try:
        allow_tf32 = torch.backends.cuda.matmul.allow_tf32
        torch.backends.cuda.matmul.allow_tf32 = True
        allow_fp16_red = torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction
        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = acc_dtype == torch.float16
        ref = torch.matmul(a.to(acc_dtype), b.to(acc_dtype)).to(out_dtype)
    finally:
        torch.backends.cuda.matmul.allow_tf32 = allow_tf32
        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = allow_fp16_red

    if bitwidth == 8:
        atol, rtol = 5e-2, 5e-1
    elif bitwidth == 16:
        atol, rtol = 5e-2, 5e-1
    else:
        atol, rtol = 5e-4, 5e-3
    torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)


@pytest.mark.skipif(not is_hip_cdna4(), reason="Requires CDNA4")
@pytest.mark.parametrize("use_buffer_load", [True, False])
def test_amd_direct_load_to_shared(use_buffer_load):

    @gluon.jit
    def kernel(a_ptr, b_ptr, use_buffer_load: ttgl.constexpr):
        blocked: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [32, 2], [4, 1], [1, 0])
        shared: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, order=[1, 0])

        smem = ttgl.allocate_shared_memory(a_ptr.dtype.element_ty, [128, 16], shared)
        offsets = ttgl.arange(0, 128, layout=ttgl.SliceLayout(1, blocked))[:, None] * 16 + \
                  ttgl.arange(0, 16, layout=ttgl.SliceLayout(0, blocked))[None, :]
        if use_buffer_load:
            cdna4_async_copy.buffer_load_to_shared(smem, a_ptr, offsets)
        else:
            cdna4_async_copy.global_load_to_shared(smem, a_ptr + offsets)
        cdna4_async_copy.commit_group()

        cdna4_async_copy.wait_group(0)
        a = cdna4_async_copy.load_shared_relaxed(smem, blocked)

        ttgl.store(b_ptr + offsets, a)

    torch.manual_seed(0)
    a = torch.randn((128, 16), dtype=torch.float16, device='cuda')
    b = torch.empty_like(a)
    pgm = kernel[(1, )](a, b, use_buffer_load)

    torch.testing.assert_close(a, b)
    assert re.search(r'ttg\.local_load .* \{ttg\.amdg\.syncedViaAsyncWait = true\}', pgm.asm['ttgir'], re.MULTILINE)
    if use_buffer_load:
        assert re.search(r"buffer_load.*lds$", pgm.asm['amdgcn'], re.MULTILINE)
    else:
        assert re.search(r"global_load_lds", pgm.asm['amdgcn'], re.MULTILINE)
    assert 'vmcnt(0)' in pgm.asm['amdgcn']


@pytest.mark.skipif(not (is_hip_rdna3() or is_hip_rdna4()), reason="Requires RDNA3 or RDNA4")
@pytest.mark.parametrize("M, N, K", [(64, 64, 64)])
@pytest.mark.parametrize("in_dtype", ['float16', 'bfloat16'])
def test_amd_wmma(M, N, K, in_dtype):

    @gluon.jit
    def kernel(a_ptr, b_ptr, c_ptr,  #
               stride_am, stride_ak,  #
               stride_bk, stride_bn,  #
               stride_cm, stride_cn,  #
               BLOCK_SIZE_M: ttgl.constexpr,  #
               BLOCK_SIZE_N: ttgl.constexpr,  #
               BLOCK_SIZE_K: ttgl.constexpr,  #
               BLOCKED_LAYOUT: ttgl.constexpr,  #
               WMMA_LAYOUT: ttgl.constexpr,  #
               K_WIDTH: ttgl.constexpr):
        offs_am = ttgl.arange(0, BLOCK_SIZE_M, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))
        offs_bn = ttgl.arange(0, BLOCK_SIZE_N, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))

        offs_ak = ttgl.arange(0, BLOCK_SIZE_K, layout=ttgl.SliceLayout(0, BLOCKED_LAYOUT))
        offs_bk = ttgl.arange(0, BLOCK_SIZE_K, layout=ttgl.SliceLayout(1, BLOCKED_LAYOUT))

        offs_a = offs_am[:, None] * stride_am + offs_ak[None, :] * stride_ak
        offs_b = offs_bk[:, None] * stride_bk + offs_bn[None, :] * stride_bn

        a = ttgl.load(a_ptr + offs_a)
        b = ttgl.load(b_ptr + offs_b)

        a = ttgl.convert_layout(a, layout=ttgl.DotOperandLayout(0, WMMA_LAYOUT, K_WIDTH))
        b = ttgl.convert_layout(b, layout=ttgl.DotOperandLayout(1, WMMA_LAYOUT, K_WIDTH))

        acc = ttgl.zeros([BLOCK_SIZE_M, BLOCK_SIZE_N], ttgl.float32, WMMA_LAYOUT)
        if WMMA_LAYOUT.version == 1:
            c = ttgl.amd.rdna3.wmma(a, b, acc)
        else:
            ttgl.static_assert(WMMA_LAYOUT.version == 2, "WMMA_LAYOUT.version must be 1 or 2")
            c = ttgl.amd.rdna4.wmma(a, b, acc)
        c = c.to(a_ptr.dtype.element_ty)

        offs_cm = ttgl.arange(0, BLOCK_SIZE_M, layout=ttgl.SliceLayout(1, WMMA_LAYOUT))
        offs_cn = ttgl.arange(0, BLOCK_SIZE_N, layout=ttgl.SliceLayout(0, WMMA_LAYOUT))
        offs_c = offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn
        ttgl.store(c_ptr + offs_c, c)

    elem_type = torch.float16 if in_dtype == 'float16' else torch.bfloat16
    a = torch.randn((M, K), device='cuda', dtype=elem_type)
    b = torch.randn((K, N), device='cuda', dtype=elem_type)
    c = torch.empty((M, N), device=a.device, dtype=elem_type)

    blocked = ttgl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0])
    wmma_version = 1 if is_hip_rdna3() else 2
    k_width = 16 if is_hip_rdna3() else 8
    wmma = ttgl.amd.AMDWMMALayout(wmma_version, True, [[0, 1], [1, 0]])
    kernel[1, 1](a, b, c, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), BLOCK_SIZE_M=M,
                 BLOCK_SIZE_N=N, BLOCK_SIZE_K=K, BLOCKED_LAYOUT=blocked, WMMA_LAYOUT=wmma, K_WIDTH=k_width, num_warps=4)

    ref = torch.matmul(a, b)
    triton_output = c
    torch.testing.assert_close(ref, triton_output)


@pytest.mark.skipif(not (is_hip_cdna3() or is_hip_cdna4()), reason="Requires CDNA3 or CDNA4")
@pytest.mark.parametrize("M, N, K", [(32, 32, 16), (16, 16, 32)])
@pytest.mark.parametrize("in_dtype", ['float16', 'bfloat16'])
@pytest.mark.parametrize("num_warps", [4, 8])
@pytest.mark.parametrize("cdna_version", [3, 4])
def test_amd_mfma(M, N, K, in_dtype, num_warps, cdna_version):
    if is_hip_cdna3() and cdna_version != 3:
        pytest.skip("On CDNA3 target, skip if mfma version is not 3")

    if is_hip_cdna4() and cdna_version != 4:
        pytest.skip("On CDNA4 target, skip if mfma version is not 4")

    @gluon.jit
    def kernel(a_ptr, b_ptr, c_ptr,  #
               stride_am, stride_ak,  #
               stride_bk, stride_bn,  #
               stride_cm, stride_cn,  #
               BLOCK_SIZE_M: ttgl.constexpr, BLOCK_SIZE_N: ttgl.constexpr, BLOCK_SIZE_K: ttgl.constexpr,
               blocked: ttgl.constexpr, k_width: ttgl.constexpr, mfma_layout: ttgl.constexpr):
        dot_a_layout: ttgl.constexpr = ttgl.DotOperandLayout(operand_index=0, parent=mfma_layout, k_width=k_width)
        dot_b_layout: ttgl.constexpr = ttgl.DotOperandLayout(operand_index=1, parent=mfma_layout, k_width=k_width)

        offs_am = ttgl.arange(0, BLOCK_SIZE_M, layout=ttgl.SliceLayout(1, blocked))
        offs_bn = ttgl.arange(0, BLOCK_SIZE_N, layout=ttgl.SliceLayout(0, blocked))

        offs_ak = ttgl.arange(0, BLOCK_SIZE_K, layout=ttgl.SliceLayout(0, blocked))
        offs_bk = ttgl.arange(0, BLOCK_SIZE_K, layout=ttgl.SliceLayout(1, blocked))
        offs_a = offs_am[:, None] * stride_am + offs_ak[None, :] * stride_ak
        offs_b = offs_bk[:, None] * stride_bk + offs_bn[None, :] * stride_bn

        a = ttgl.amd.cdna3.buffer_load(ptr=a_ptr, offsets=offs_a)
        b = ttgl.amd.cdna3.buffer_load(ptr=b_ptr, offsets=offs_b)
        a1 = ttgl.convert_layout(a, layout=dot_a_layout)
        b1 = ttgl.convert_layout(b, layout=dot_b_layout)
        acc = ttgl.zeros([BLOCK_SIZE_M, BLOCK_SIZE_N], ttgl.float32, mfma_layout)
        c = ttgl.amd.cdna3.mfma(a1, b1, acc)
        c = ttgl.convert_layout(c, layout=blocked)
        c = c.to(a_ptr.dtype.element_ty)

        offs_cm = ttgl.arange(0, BLOCK_SIZE_M)
        offs_cn = ttgl.arange(0, BLOCK_SIZE_N)
        offs_c = offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn
        ttgl.amd.cdna3.buffer_store(stored_value=c, ptr=c_ptr, offsets=offs_c)

    elem_type = torch.float16 if in_dtype == 'float16' else torch.bfloat16
    a = torch.randn((M, K), device='cuda', dtype=elem_type) - 0.5
    b = torch.randn((K, N), device='cuda', dtype=elem_type) - 0.5
    c = torch.empty((M, N), device=a.device, dtype=elem_type)
    nonkdim: ttgl.constexpr = 32
    kdim: ttgl.constexpr = 8 if cdna_version == 3 else 16
    k_width: ttgl.constexpr = 4 if cdna_version == 3 else 8
    blocked: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[4, 4], threads_per_warp=[4, 16],
                                                 warps_per_cta=[num_warps, 1], order=[1, 0])
    mfma_layout: ttgl.constexpr = ttgl.amd.AMDMFMALayout(version=cdna_version, instr_shape=[nonkdim, nonkdim, kdim],
                                                         transposed=True, warps_per_cta=[num_warps, 1])

    kernel[1, 1](
        a, b, c,  #
        a.stride(0), a.stride(1),  #
        b.stride(0), b.stride(1),  #
        c.stride(0), c.stride(1),  #
        BLOCK_SIZE_M=M, BLOCK_SIZE_N=N, BLOCK_SIZE_K=K,  #
        blocked=blocked, k_width=k_width, mfma_layout=mfma_layout,  #
        num_warps=num_warps)

    ref = torch.matmul(a, b)
    triton_output = c
    torch.testing.assert_close(ref, triton_output)


@pytest.mark.skipif(not is_hip_cdna4(), reason="Requires CDNA4")
@pytest.mark.parametrize("M, N, K", [(32, 32, 128)])
@pytest.mark.parametrize("a_type, b_type", [(a_type, b_type)
                                            for a_type in ["e2m1", "e4m3", "e5m2"]
                                            for b_type in ["e2m1", "e4m3", "e5m2"]])
@pytest.mark.parametrize("has_scale", [True, False])
def test_amd_mfma_scaled(M, N, K, a_type, b_type, has_scale, device='cuda'):

    @gluon.jit
    def kernel(out_ptr, a_ptr, b_ptr, a_scale_ptr, b_scale_ptr,  #
               M: ttgl.constexpr, N: ttgl.constexpr, K: ttgl.constexpr,  #
               a_type: tl.constexpr, b_type: tl.constexpr):
        DIV_FACTOR_A: tl.constexpr = 2 if a_type == "e2m1" else 1
        DIV_FACTOR_B: tl.constexpr = 2 if b_type == "e2m1" else 1
        K_A: tl.constexpr = K // DIV_FACTOR_A
        K_B: tl.constexpr = K // DIV_FACTOR_B

        mfma_layout: ttgl.constexpr = ttgl.amd.AMDMFMALayout(version=4, instr_shape=[16, 16, 128], transposed=True,
                                                             warps_per_cta=[2, 2])

        a_unpacked_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 16], [8, 8], [4, 1], [1, 0])
        a_packed_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [8, 8], [4, 1], [1, 0])
        a_load_layout: ttgl.constexpr = a_packed_layout if a_type == "e2m1" else a_unpacked_layout
        a_layout: ttgl.constexpr = ttgl.DotOperandLayout(operand_index=0, parent=mfma_layout, k_width=16)
        a_scale_layout: ttgl.constexpr = ttgl.amd.cdna4.get_mfma_scale_layout(a_layout, [M, K // 32])

        b_unpacked_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 16], [32, 2], [4, 1], [1, 0])
        b_packed_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 8], [16, 4], [4, 1], [1, 0])
        b_load_layout: ttgl.constexpr = b_packed_layout if b_type == "e2m1" else b_unpacked_layout
        b_layout: ttgl.constexpr = ttgl.DotOperandLayout(operand_index=1, parent=mfma_layout, k_width=16)
        b_scale_layout: ttgl.constexpr = ttgl.amd.cdna4.get_mfma_scale_layout(b_layout, [N, K // 32])

        a_offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, a_load_layout))[:, None]
        a_offs_k = ttgl.arange(0, K_A, layout=ttgl.SliceLayout(0, a_load_layout))[None, :]
        a = ttgl.amd.cdna4.buffer_load(a_ptr, a_offs_m * K_A + a_offs_k)
        a = ttgl.convert_layout(a, a_layout)

        b_offs_k = ttgl.arange(0, K_B, layout=ttgl.SliceLayout(1, b_load_layout))[:, None]
        b_offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, b_load_layout))[None, :]
        b = ttgl.amd.cdna4.buffer_load(b_ptr, b_offs_k * N + b_offs_n)
        b = ttgl.convert_layout(b, b_layout)

        a_scale = None
        if a_scale_ptr is not None:
            a_scale_offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, a_scale_layout))[:, None]
            a_scale_offs_k = ttgl.arange(0, K // 32, layout=ttgl.SliceLayout(0, a_scale_layout))[None, :]
            a_scale = ttgl.amd.cdna4.buffer_load(a_scale_ptr, a_scale_offs_m * (K // 32) + a_scale_offs_k)

        b_scale = None
        if b_scale_ptr is not None:
            b_scale_offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(1, b_scale_layout))[:, None]
            b_scale_offs_k = ttgl.arange(0, K // 32, layout=ttgl.SliceLayout(0, b_scale_layout))[None, :]
            b_scale = ttgl.amd.cdna4.buffer_load(b_scale_ptr, b_scale_offs_n * (K // 32) + b_scale_offs_k)

        zero = ttgl.zeros([M, N], dtype=ttgl.float32, layout=mfma_layout)
        c = ttgl.amd.cdna4.mfma_scaled(a, a_scale, a_type, b, b_scale, b_type, zero)
        c = c.to(out_ptr.dtype.element_ty)

        out_offs_m = ttgl.arange(0, M)[:, None]
        out_offs_n = ttgl.arange(0, N)[None, :]
        ttgl.amd.cdna4.buffer_store(c, out_ptr, out_offs_m * N + out_offs_n)

    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.to(device), v_ref.to(device)

    def _create_mxfp_scale(operand: int, m: int, n: int):
        size = (m, n // 32)
        scale = MXScaleTensor(size=tuple(size)).random(1 / 32, 32)
        scale_ref = scale.to(torch.float32).repeat_interleave(32, dim=1)
        scale_ref = scale_ref.T.contiguous() if operand == 1 else scale_ref
        return scale.data.to(device), scale_ref.to(device)

    torch.manual_seed(0)
    a, a_ref = _create_mxfp_operand(0, M, K, a_type)
    b, b_ref = _create_mxfp_operand(1, K, N, b_type)

    if has_scale:
        a_scale, a_scale_ref = _create_mxfp_scale(0, M, K)
        b_scale, b_scale_ref = _create_mxfp_scale(1, N, K)
        out = torch.empty((M, N), dtype=torch.float32, device=device)
        compiled = kernel[(1, )](out, a, b, a_scale, b_scale, M, N, K, a_type, b_type, num_warps=4)
        out_ref = torch.matmul(a_ref * a_scale_ref, b_ref * b_scale_ref)
        torch.testing.assert_close(out, out_ref)
    else:
        out = torch.empty((M, N), dtype=torch.float32, device=device)
        compiled = kernel[(1, )](out, a, b, None, None, M, N, K, a_type, b_type, num_warps=4)
        out_ref = torch.matmul(a_ref, b_ref)
        torch.testing.assert_close(out, out_ref)

    assert 'v_mfma_scale_f32_16x16x128_f8f6f4' in compiled.asm['amdgcn']


def test_math_fast_expf():

    @gluon.jit
    def fast_expf_kernel(x_ptr, y_ptr, warp_size: ttgl.constexpr, num_warps: ttgl.constexpr):
        blocked: ttgl.constexpr = ttgl.BlockedLayout([1], [warp_size], [num_warps], [0])

        offs = ttgl.arange(0, warp_size * num_warps, layout=blocked)
        x = ttgl.load(x_ptr + offs)
        y = libdevice.fast_expf(x)
        ttgl.store(y_ptr + offs, y)

    num_warps = 4

    torch.manual_seed(0)
    x = torch.randn(THREADS_PER_WARP * num_warps, device="cuda", dtype=torch.float32)
    y = torch.empty_like(x)
    fast_expf_kernel[(1, )](x, y, THREADS_PER_WARP, num_warps)
    torch.testing.assert_close(y, torch.exp(x), atol=1e-5, rtol=1e-4)


def test_math_fast_dividef():

    @gluon.jit
    def fast_dividef_kernel(x_ptr, y_ptr, z_ptr, warp_size: ttgl.constexpr, num_warps: ttgl.constexpr):
        blocked: ttgl.constexpr = ttgl.BlockedLayout([1], [warp_size], [num_warps], [0])

        offs = ttgl.arange(0, warp_size * num_warps, layout=blocked)
        x = ttgl.load(x_ptr + offs)
        y = ttgl.load(y_ptr + offs)
        z = libdevice.fast_dividef(x, y)
        ttgl.store(z_ptr + offs, z)

    num_warps = 4

    torch.manual_seed(0)
    x = torch.randn(THREADS_PER_WARP * num_warps, device="cuda", dtype=torch.float32)
    y = torch.randn_like(x)
    z = torch.empty_like(x)
    y[y == 0] = 1.0
    fast_dividef_kernel[(1, )](x, y, z, THREADS_PER_WARP, num_warps)
    torch.testing.assert_close(z, torch.div(x, y), atol=1e-5, rtol=1e-4)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tmem_copy_2d():
    device = "cuda"

    smem_h = 64
    smem_w = 16
    num_rows = 128
    num_cols = smem_h * smem_w // 32

    @gluon.jit
    def kernel(in_ptr, out_ptr, smem_h: ttgl.constexpr, smem_w: ttgl.constexpr, num_rows: ttgl.constexpr,
               num_cols: ttgl.constexpr):
        in_ptrs = in_ptr + ttgl.arange(0, smem_h)[:, None] * smem_w + ttgl.arange(0, smem_w)[None, :]
        out_ptrs = out_ptr + ttgl.arange(0, num_rows)[:, None] * num_cols + ttgl.arange(0, num_cols)[None, :]

        blocked: ttgl.constexpr = ttgl.BlockedLayout([1, 4], [32, 1], [4, 1], [1, 0])
        value = ttgl.load(ttgl.set_auto_layout(in_ptrs, blocked))

        smem_layout: ttgl.constexpr = ttgl.SharedLinearLayout(
            offset_bases=[[0, 1], [0, 2], [32, 0], [0, 4], [1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [0, 8]])
        tmem_layout: ttgl.constexpr = TensorMemoryScalesLayout()
        smem = ttgl.allocate_shared_memory(ttgl.int8, (smem_h, smem_w), layout=smem_layout)
        tmem_pool_layout: ttgl.constexpr = TensorMemoryLayout((num_rows, 256), col_stride=1)
        tmem_pool = allocate_tensor_memory(ttgl.float32, (num_rows, 512), layout=tmem_pool_layout)
        tmem = tmem_pool.slice(480, num_cols).reinterpret(ttgl.int8, (smem_h, smem_w), tmem_layout)

        barrier = ttgl.allocate_shared_memory(ttgl.int64, [1], ttgl.constexpr(mbarrier.MBarrierLayout()))
        mbarrier.init(barrier, count=1)

        smem.store(value)
        tcgen05_copy(smem, tmem)
        tcgen05_commit(barrier)
        mbarrier.wait(barrier, phase=0)
        tmem_alias: ttgl.constexpr = TensorMemoryLayout((num_rows, num_cols), col_stride=1)
        tmem = tmem.reinterpret(shape=(num_rows, num_cols), layout=tmem_alias)
        value = tmem.load(blocked)
        ttgl.store(ttgl.set_auto_layout(out_ptrs, blocked), value)

    torch.manual_seed(0)
    x = torch.randint(size=(smem_h, smem_w), low=-100, high=100, dtype=torch.int8).to(device)
    #x = torch.arange(smem_h * smem_w, dtype=torch.int8, device=device).reshape(smem_h, smem_w)
    z_tri = torch.zeros(size=(num_rows, num_cols), dtype=torch.int8).to(device)
    kernel[(1, )](x, z_tri, smem_h, smem_w, num_rows, num_cols)

    # offset_bases=[[0, 1], [0, 2], [32, 0], [0, 4], [1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [0, 8]],
    # Split into contiguous shmem chunks
    x_res = x.reshape(2, 32, 2, 2, 4)
    # Put tmem cols first then rows
    x_res = x_res.permute(1, 2, 3, 0, 4)
    # Reshape as 32xnum_cols
    x_res = x_res.reshape(num_rows // 4, num_cols)

    warps = torch.chunk(z_tri, chunks=4, dim=0)
    for warp in warps:
        torch.testing.assert_close(x_res, warp)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tmem_pipeline_stage_subslice_index_reinterpret():

    @gluon.jit
    def kernel(inp, out):
        parent = allocate_tensor_memory(ttgl.float32, [5, 128, 64], TensorMemoryLayout([128, 64], col_stride=1))
        layout: ttgl.constexpr = parent.index(0).get_reg_layout()
        rows = ttgl.arange(0, 128, layout=ttgl.SliceLayout(1, layout))[:, None]
        cols = ttgl.arange(0, 64, layout=ttgl.SliceLayout(0, layout))[None, :]
        for stage in ttgl.static_range(5):
            parent.index(stage).store(ttgl.load(inp + stage * 8192 + rows * 64 + cols))
        stages = parent.slice(2, 2, dim=0)
        first = stages.reinterpret().index(0)
        first.store(first.load() + 11.0)
        second = stages.slice(1, 1, dim=0).index(0)
        second.store(second.load() + 7.0)
        for stage in ttgl.static_range(5):
            ttgl.store(out + stage * 8192 + rows * 64 + cols, parent.index(stage).load(layout))

    inp = torch.arange(5 * 128 * 64, device="cuda", dtype=torch.int32).remainder(64).float().reshape(5, 128, 64)
    out = torch.empty_like(inp)
    kernel[(1, )](inp, out, num_warps=4)
    expected = inp.clone()
    expected[2] += 11
    expected[3] += 7
    torch.testing.assert_close(out, expected, atol=0, rtol=0)


@pytest.mark.parametrize("M, N, BLOCK_M, BLOCK_N", [(256, 128, 128, 128), (128, 256, 64, 128)])
@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tmem_row_slice(M, N, BLOCK_M, BLOCK_N):

    @gluon.jit
    def _store_view(view, out):
        values = view.load()
        layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[1, 32],
                                                    warps_per_cta=[1, 4], order=[1, 0])
        values = ttgl.convert_layout(values, layout)
        rows = ttgl.arange(0, view.shape[0], layout=ttgl.SliceLayout(1, layout))
        cols = ttgl.arange(0, view.shape[1], layout=ttgl.SliceLayout(0, layout))
        ttgl.store(out + rows[:, None] * view.shape[1] + cols[None, :], values)

    @gluon.jit
    def tmem_row_slice_kernel(inp, out0, out1, out2, M: ttgl.constexpr, N: ttgl.constexpr, BLOCK_M: ttgl.constexpr,
                              BLOCK_N: ttgl.constexpr):
        parent = allocate_tensor_memory(ttgl.float32, [M, N], TensorMemoryLayout([BLOCK_M, BLOCK_N], col_stride=1))
        layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[1, 32],
                                                    warps_per_cta=[1, 4], order=[1, 0])
        rows = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, layout))
        cols = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, layout))
        values = ttgl.load(inp + rows[:, None] * N + cols[None, :])
        parent.store(ttgl.convert_layout(values, parent.get_reg_layout()))

        row_half0 = parent.slice(0, M // 2, dim=0)
        row_half1 = parent.slice(M // 2, M // 2, dim=0)
        _store_view(row_half0, out0)
        _store_view(row_half1, out1)
        _store_view(row_half1.slice(N // 2, N // 2), out2)

    inp = torch.arange(M * N, device="cuda", dtype=torch.float32).reshape(M, N)
    out0 = torch.empty((M // 2, N), device="cuda", dtype=torch.float32)
    out1 = torch.empty_like(out0)
    out2 = torch.empty((M // 2, N // 2), device="cuda", dtype=torch.float32)
    tmem_row_slice_kernel[(1, )](inp, out0, out1, out2, M, N, BLOCK_M, BLOCK_N, num_warps=4)
    torch.testing.assert_close(out0, inp[:M // 2], atol=0, rtol=0)
    torch.testing.assert_close(out1, inp[M // 2:], atol=0, rtol=0)
    torch.testing.assert_close(out2, inp[M // 2:, N // 2:], atol=0, rtol=0)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tmem_source_layout_contiguous_subslice():

    @gluon.jit
    def kernel(inp, out):
        parent = allocate_tensor_memory(ttgl.float32, [128, 512], TensorMemoryLayout([128, 256], col_stride=1))
        parent_layout: ttgl.constexpr = parent.get_reg_layout()
        rows = ttgl.arange(0, 128, layout=ttgl.SliceLayout(1, parent_layout))[:, None]
        cols = ttgl.arange(0, 512, layout=ttgl.SliceLayout(0, parent_layout))[None, :]
        parent.store(ttgl.load(inp + rows * 512 + cols))

        view = parent.slice(208, 256)
        view_layout: ttgl.constexpr = TensorMemoryLayout([128, 256], col_stride=1)
        view = view.reinterpret(ttgl.float32, [128, 256], view_layout)
        view.store(view.load() + 7.0)
        ttgl.store(out + rows * 512 + cols, parent.load(parent_layout))

    inp = torch.arange(128 * 512, device="cuda", dtype=torch.int32).remainder(64)
    inp = inp.to(torch.float32).reshape(128, 512)
    out = torch.empty_like(inp)
    kernel[(1, )](inp, out, num_warps=4)
    ref = inp.clone()
    ref[:, 208:464] += 7
    torch.testing.assert_close(out, ref, atol=0, rtol=0)


@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
@pytest.mark.parametrize("M, N, BLOCK_N, dim", [(256, 256, 64, 0), (256, 128, 128, 1)])
@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tmem_noncontiguous_subslice_load_store(dtype, M, N, BLOCK_N, dim):

    @gluon.jit
    def kernel(inp, out, M: ttgl.constexpr, N: ttgl.constexpr, BLOCK_N: ttgl.constexpr, DIM: ttgl.constexpr):
        parent = allocate_tensor_memory(
            inp.dtype.element_ty, [M, N],
            TensorMemoryLayout([128, BLOCK_N], col_stride=32 // inp.dtype.element_ty.primitive_bitwidth))
        parent_layout: ttgl.constexpr = parent.get_reg_layout()
        rows = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, parent_layout))[:, None]
        cols = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, parent_layout))[None, :]
        parent.store(ttgl.load(inp + rows * N + cols))

        if DIM == 0:
            view = parent.slice(M // 2, M // 2, dim=0)
        else:
            view = parent.slice(N // 2, N // 2, dim=1)
        view.store(view.load() + 7.0)
        ttgl.store(out + rows * N + cols, parent.load(parent_layout))

    inp = torch.arange(M * N, device="cuda", dtype=torch.int32).remainder(64).to(dtype).reshape(M, N)
    out = torch.empty_like(inp)
    kernel[(1, )](inp, out, M, N, BLOCK_N, dim, num_warps=4)
    ref = inp.clone()
    if dim == 0:
        ref[M // 2:] += 7
    else:
        ref[:, N // 2:] += 7
    torch.testing.assert_close(out, ref, atol=0, rtol=0)


@pytest.mark.skipif(not is_blackwell_ultra(), reason="Requires Blackwell Ultra")
def test_tmem_noncontiguous_subslice_load_min():

    @gluon.jit
    def kernel(inp, out, reduced_out):
        parent = allocate_tensor_memory(ttgl.float32, [256, 256], TensorMemoryLayout([128, 64], col_stride=1))
        parent_layout: ttgl.constexpr = parent.get_reg_layout()
        rows = ttgl.arange(0, 256, layout=ttgl.SliceLayout(1, parent_layout))[:, None]
        cols = ttgl.arange(0, 256, layout=ttgl.SliceLayout(0, parent_layout))[None, :]
        parent.store(ttgl.load(inp + rows * 256 + cols))

        values, reduced = parent.slice(128, 128, dim=0).load_min()
        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], [4, 1], [1, 0])
        view_rows = ttgl.arange(0, 128, layout=ttgl.SliceLayout(1, layout))[:, None]
        view_cols = ttgl.arange(0, 256, layout=ttgl.SliceLayout(0, layout))[None, :]
        ttgl.store(out + view_rows * 256 + view_cols, ttgl.convert_layout(values, layout))
        red_layout: ttgl.constexpr = ttgl.SliceLayout(1, layout)
        red_rows = ttgl.arange(0, 128, layout=red_layout)
        ttgl.store(reduced_out + red_rows, ttgl.convert_layout(reduced, red_layout))

    torch.manual_seed(0)
    inp = torch.randn((256, 256), dtype=torch.float32, device="cuda")
    out = torch.empty((128, 256), dtype=torch.float32, device="cuda")
    reduced = torch.empty((128, ), dtype=torch.float32, device="cuda")
    kernel[(1, )](inp, out, reduced, num_warps=4)
    torch.testing.assert_close(out, inp[128:], atol=0, rtol=0)
    torch.testing.assert_close(reduced, inp[128:].amin(dim=1), atol=0, rtol=0)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tmem_subslice_unpacked_one_column():

    @gluon.jit
    def kernel(inp, out):
        parent = allocate_tensor_memory(ttgl.float16, [128, 2], TensorMemoryLayout([128, 1], col_stride=2))
        layout: ttgl.constexpr = parent.get_reg_layout()
        rows = ttgl.arange(0, 128, layout=ttgl.SliceLayout(1, layout))
        cols = ttgl.arange(0, 2, layout=ttgl.SliceLayout(0, layout))
        parent.store(ttgl.load(inp + rows[:, None] * 2 + cols[None, :]))

        view = parent.slice(0, 1)
        view.store(ttgl.full([128, 1], 777, dtype=ttgl.float16, layout=view.get_reg_layout()))
        ttgl.store(out + rows[:, None] * 2 + cols[None, :], parent.load())

    inp = torch.arange(256, device="cuda", dtype=torch.float16).reshape(128, 2)
    out = torch.empty_like(inp)
    kernel[(1, )](inp, out, num_warps=4)
    torch.testing.assert_close(out[:, 0], torch.full_like(out[:, 0], 777), atol=0, rtol=0)
    torch.testing.assert_close(out[:, 1], inp[:, 1], atol=0, rtol=0)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tmem_subslice_block_m_64():

    @gluon.jit
    def kernel(s_ptr, out_ptr):
        BLOCK_M: ttgl.constexpr = 64
        N: ttgl.constexpr = 128
        BLOCK_N: ttgl.constexpr = 64

        tmem_layout: ttgl.constexpr = TensorMemoryLayout((BLOCK_M, BLOCK_N), col_stride=1)
        s_tmem = allocate_tensor_memory(ttgl.float32, (BLOCK_M, N), layout=tmem_layout)
        o_tmem = allocate_tensor_memory(ttgl.float32, (BLOCK_M, N), layout=tmem_layout)

        layout: ttgl.constexpr = s_tmem.get_reg_layout()

        offsets = ttgl.arange(0, BLOCK_M)[:, None] * N + ttgl.arange(0, N)[None, :]
        offsets = ttgl.set_auto_layout(offsets, layout)
        s = ttgl.load(s_ptr + offsets)

        s_tmem.store(s)
        o_tmem.store(s)

        p_tmem_parent = s_tmem.reinterpret(ttgl.float16, [BLOCK_M, 2 * N], tmem_layout)
        p_tmem = p_tmem_parent.slice(0, N)
        p_tmem.store(ttgl.full((BLOCK_M, N), 0.0, dtype=ttgl.float16, layout=layout))

        d1_tmem_layout: ttgl.constexpr = TensorMemoryLayout((BLOCK_M, 2), col_stride=1)

        d1_tmem_parent = s_tmem.reinterpret(layout=d1_tmem_layout)
        m_tmem = d1_tmem_parent.slice(N // 2, 2)
        d1_layout: ttgl.constexpr = m_tmem.get_reg_layout()
        m_tmem.store(ttgl.full((BLOCK_M, 2), 2.0, dtype=ttgl.float32, layout=d1_layout))
        l_tmem = d1_tmem_parent.slice(N // 2 + 4, 2)
        l_tmem.store(ttgl.full((BLOCK_M, 2), 3.0, dtype=ttgl.float32, layout=d1_layout))
        a_tmem = d1_tmem_parent.slice(N // 2 + 8, 2)
        a_tmem.store(ttgl.full((BLOCK_M, 2), 4.0, dtype=ttgl.float32, layout=d1_layout))

        s = s_tmem.load()

        ttgl.store(out_ptr + offsets, s)

    torch.manual_seed(0)
    s = torch.randn((64, 128), dtype=torch.float32, device="cuda")

    out_tri = torch.empty_like(s)
    compiled = kernel[(1, )](s, out_tri)

    ttgir = compiled.asm["ttgir"]
    # Check that we have two 64x128xf32 allocations.
    assert ttgir.count("ttng.tmem_alloc") == 2
    assert ttgir.count("ttng.tmem_alloc : () -> !ttg.memdesc<64x128xf32") == 2

    # Check that we allocated only 128 columns of TMEM.
    llir = compiled.asm["llir"]
    assert llir.count("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [$1], 128")

    # Given TMEM[0:32] is the slice of TMEM for warpgroup 0, the expected layout
    # of S is
    #
    #   TMEM[0:16]  = S[0:16, 0:64]
    #   TMEM[16:32] = S[0:16, 64:128]
    #
    # When slicing S to obtain P, we expect it to overlap with the left half,
    # i.e. S[0:16, 0:32] and S[0:16, 64:96].
    out_ref = s
    out_ref[:, 0:32] = 0.0
    out_ref[:, 64:96] = 0.0

    # Given S = [s0, s1, s2, s3], they are arranged like
    #
    #   TMEM[0:16]  = [s0, s1]
    #   TMEM[16:32] = [s2, s3]
    #
    # Thus the narrow parent view is sliced at offsets that map back to s1.
    out_ref[:, 32:34] = 2.0
    out_ref[:, 34:36] = 3.0
    out_ref[:, 36:38] = 4.0

    torch.testing.assert_close(out_ref, out_tri, atol=0, rtol=0)


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_block_m_64_mma():

    @gluon.jit
    def kernel(a_ptr, b_ptr, c_ptr, d_ptr):
        BLOCK_M: ttgl.constexpr = 64
        N: ttgl.constexpr = 128
        BLOCK_N: ttgl.constexpr = 64

        a_offsets = ttgl.arange(0, BLOCK_M)[:, None] * N + ttgl.arange(0, N)[None, :]
        b_offsets = ttgl.arange(0, N)[:, None] * N + ttgl.arange(0, N)[None, :]

        a_tmem_layout: ttgl.constexpr = TensorMemoryLayout((BLOCK_M, BLOCK_N), col_stride=1)
        acc_tmem_layout: ttgl.constexpr = TensorMemoryLayout((BLOCK_M, BLOCK_N), col_stride=1)
        al_tmem = allocate_tensor_memory(ttgl.float16, (BLOCK_M, N), layout=a_tmem_layout)
        ar_tmem = allocate_tensor_memory(ttgl.float16, (BLOCK_M, N), layout=a_tmem_layout)
        acc_tmem = allocate_tensor_memory(ttgl.float32, (BLOCK_M, N), layout=acc_tmem_layout)
        a_layout: ttgl.constexpr = al_tmem.get_reg_layout(instr_variant="32x32b_splitn")
        b_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], [4, 1], [1, 0])
        a_offsets = ttgl.set_auto_layout(a_offsets, a_layout)
        b_offsets = ttgl.set_auto_layout(b_offsets, b_layout)

        a = ttgl.load(a_ptr + a_offsets)
        b = ttgl.load(b_ptr + b_offsets)
        c = ttgl.load(c_ptr + a_offsets)

        a0, a1 = a.reshape((BLOCK_M, 2, N // 2)).permute(0, 2, 1).split()

        al = ttgl.join(a0, a1).permute(0, 2, 1).reshape((BLOCK_M, N))
        ar = ttgl.join(a1, a0).permute(0, 2, 1).reshape((BLOCK_M, N))

        al_tmem.store(ttgl.convert_layout(al, a_layout, assert_trivial=True))
        ar_tmem.store(ttgl.convert_layout(ar, a_layout, assert_trivial=True))

        b_shared_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=32, element_bitwidth=16, rank=2)
        b_shared = ttgl.allocate_shared_memory(ttgl.float16, [N, N], layout=b_shared_layout)
        b_shared.store(b)

        acc_tmem.store(c)

        bar = ttgl.allocate_shared_memory(ttgl.int64, [1], ttgl.constexpr(mbarrier.MBarrierLayout()))
        mbarrier.init(bar, count=1)

        # This is a manually tiled MMA where LHS is in TMEM with blockM=64,
        # where we circumvent the limitation that LHS and accumulator need to
        # share the same TMEM rows by storing the LHS twice.
        #
        # TMEM      al   ar   c
        # [0, 16)   a0   a1   c0
        # [16, 32)  a1   a0   c1
        #
        # d0 = a0 @ b00 + a1 @ b10 + c0
        # d1 = a0 @ b10 + a1 @ b11 + c1

        N2: ttgl.constexpr = N // 2
        c0 = acc_tmem.slice(0, N2)
        c1 = acc_tmem.slice(N2, N2)

        tcgen05_mma(al_tmem.slice(0, N2), b_shared.slice(0, N2, dim=0).slice(0, N2, dim=1), c0)
        tcgen05_mma(ar_tmem.slice(0, N2), b_shared.slice(N2, N2, dim=0).slice(0, N2, dim=1), c0)
        tcgen05_mma(ar_tmem.slice(N2, N2), b_shared.slice(0, N2, dim=0).slice(N2, N2, dim=1), c1)
        tcgen05_mma(al_tmem.slice(N2, N2), b_shared.slice(N2, N2, dim=0).slice(N2, N2, dim=1), c1)

        tcgen05_commit(bar)
        mbarrier.wait(bar, 0)
        mbarrier.invalidate(bar)

        d = acc_tmem.load(a_layout)
        ttgl.store(d_ptr + a_offsets, d)

    torch.manual_seed(0)
    a = torch.randn((64, 128), dtype=torch.float16, device="cuda")
    b = torch.randn((128, 128), dtype=torch.float16, device="cuda")
    c = torch.randn((64, 128), dtype=torch.float32, device="cuda")

    d_tri = torch.empty_like(c)
    compiled = kernel[(1, )](a, b, c, d_tri)

    ttgir = compiled.asm["ttgir"]
    assert ttgir.count("ttng.tmem_alloc") == 3
    assert ttgir.count("ttng.tmem_alloc : () -> !ttg.memdesc<64x128xf32") == 1
    assert ttgir.count("ttng.tmem_alloc : () -> !ttg.memdesc<64x128xf16") == 2

    llir = compiled.asm["llir"]
    assert llir.count("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [$1], 128")

    d_ref = a @ b + c
    torch.testing.assert_close(d_ref, d_tri, rtol=0.08, atol=0)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_multibuffer_prefix_subslice():

    @gluon.jit
    def kernel(out0, out3):
        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], [4, 1], [1, 0])
        shared_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(128, 32, rank=2)
        parent = ttgl.allocate_shared_memory(ttgl.float32, [7, 32, 32], shared_layout)
        stage0 = parent.index(0)
        stage3 = parent.slice(3, 3, dim=0).index(0)
        stage0.store(ttgl.full([32, 32], 1.0, ttgl.float32, layout))
        stage3.store(ttgl.full([32, 32], 3.0, ttgl.float32, layout))
        ttgl.barrier()

        rows = ttgl.arange(0, 32, layout=ttgl.SliceLayout(1, layout))[:, None]
        cols = ttgl.arange(0, 32, layout=ttgl.SliceLayout(0, layout))[None, :]
        offsets = rows * 32 + cols
        ttgl.store(out0 + offsets, stage0.load(layout))
        ttgl.store(out3 + offsets, stage3.load(layout))

    out0 = torch.empty((32, 32), device="cuda", dtype=torch.float32)
    out3 = torch.empty_like(out0)
    kernel[(1, )](out0, out3, num_warps=4)
    torch.testing.assert_close(out0, torch.ones_like(out0), atol=0, rtol=0)
    torch.testing.assert_close(out3, torch.full_like(out3, 3), atol=0, rtol=0)


def test_slice_reinterpret():
    BLOCK = ttgl.constexpr(2048)
    SPLIT_BLOCK = ttgl.constexpr(BLOCK // 2)
    XBLOCK = ttgl.constexpr(32)
    YBLOCK = ttgl.constexpr(SPLIT_BLOCK // 4 // XBLOCK)
    NUM_THREADS = ttgl.constexpr(THREADS_PER_WARP)

    @gluon.jit
    def kernel(in_ptr, out_ptr):
        smem_layout_1d: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[0])
        smem_layout_2d: ttgl.constexpr = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])
        smem = ttgl.allocate_shared_memory(ttgl.int8, [BLOCK], smem_layout_1d)
        smem_slice0 = smem.slice(0, SPLIT_BLOCK)
        smem_i32 = smem.reinterpret(ttgl.int32, [2 * XBLOCK, YBLOCK], smem_layout_2d)
        smem_slice1 = smem_i32.slice(XBLOCK, XBLOCK, dim=0)

        offs = ttgl.arange(0, XBLOCK)[:, None] * YBLOCK + ttgl.arange(0, YBLOCK)[None, :]
        blocked: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, NUM_THREADS], [1, 4], [1, 0])
        value = ttgl.load(ttgl.set_auto_layout(in_ptr + offs, blocked))

        blocked_1d: ttgl.constexpr = ttgl.BlockedLayout([1], [NUM_THREADS], [4], [0])
        smem_slice1.store(value)
        smem_slice0.store(ttgl.zeros((SPLIT_BLOCK, ), dtype=ttgl.int8, layout=blocked_1d))
        value = smem_slice1.load(blocked)
        ttgl.store(ttgl.set_auto_layout(out_ptr + offs, blocked), value)

    input = torch.randint(0, 100, (XBLOCK, YBLOCK), dtype=torch.int32, device="cuda")
    output = torch.empty_like(input)
    kernel[(1, )](input, output)
    torch.testing.assert_close(input, output, atol=0, rtol=0)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_tma_slice():
    XBLOCK = YBLOCK = ttgl.constexpr(128)

    @gluon.jit
    def kernel(in_desc, out_desc):
        smem = ttgl.allocate_shared_memory(in_desc.dtype, [2 * XBLOCK, YBLOCK], in_desc.layout)
        smem_slice0 = smem.slice(0, XBLOCK)
        smem_slice1 = smem.slice(XBLOCK, XBLOCK)

        bar = ttgl.allocate_shared_memory(ttgl.int64, [1], ttgl.constexpr(mbarrier.MBarrierLayout()))
        mbarrier.init(bar, count=1)

        mbarrier.expect(bar, in_desc.nbytes_per_cta)
        tma.async_load(in_desc, [0, 0], bar, smem_slice1)
        mbarrier.wait(bar, phase=0)

        blocked: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], [1, 4], [1, 0])
        smem_slice0.store(ttgl.zeros((XBLOCK, YBLOCK), dtype=ttgl.float32, layout=blocked))

        tma.async_store(out_desc, [0, 0], smem_slice1)
        tma.store_wait(0)

    input = torch.rand((XBLOCK, YBLOCK), dtype=torch.float32, device="cuda")
    output = torch.empty_like(input)

    block_shape = [XBLOCK.value, YBLOCK.value]
    layout = ttgl.NVMMASharedLayout.get_default_for(block_shape, ttgl.float32)
    in_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(input, block_shape, layout)
    out_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(output, block_shape, layout)
    kernel[(1, )](in_desc, out_desc)
    torch.testing.assert_close(input, output, atol=0, rtol=0)


@pytest.mark.parametrize("swizzle", [32, 64, 128])
@pytest.mark.parametrize("num_warps", [4, 8])
@pytest.mark.parametrize("M, N, BLOCK_N", [(128, 128, 128), (256, 128, 64), (128, 128, 16)])
@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tmem_copy_no_scales(M, N, BLOCK_N, num_warps, swizzle):

    @gluon.jit
    def tmem_copy_no_scales(in_ptr, out_ptr, M: ttgl.constexpr, N: ttgl.constexpr, BLOCK_N: ttgl.constexpr,
                            swizzle: ttgl.constexpr, num_warps: ttgl.constexpr):
        tmem_layout: ttgl.constexpr = TensorMemoryLayout(
            block=(128, BLOCK_N),
            col_stride=32 // in_ptr.dtype.element_ty.primitive_bitwidth,
        )
        tmem = allocate_tensor_memory(
            element_ty=in_ptr.dtype.element_ty,
            shape=[M, N],
            layout=tmem_layout,
        )
        tmem_reg_layout: ttgl.constexpr = tmem.get_reg_layout()
        offs_m = ttgl.arange(0, M, ttgl.SliceLayout(1, tmem_reg_layout))
        offs_n = ttgl.arange(0, N, ttgl.SliceLayout(0, tmem_reg_layout))
        offs = offs_m[:, None] * N + offs_n[None, :]

        input = ttgl.load(in_ptr + offs)

        smem_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=swizzle, element_bitwidth=32, rank=2)
        smem = ttgl.allocate_shared_memory(in_ptr.dtype.element_ty, [M, N], layout=smem_layout)

        smem.store(input)
        bar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
        mbarrier.init(bar, count=1)
        tcgen05_copy(smem, tmem)
        tcgen05_commit(bar)
        mbarrier.wait(bar, phase=0)
        output = tmem.load()
        ttgl.store(out_ptr + offs, output)

    input = torch.arange(M * N, device="cuda").reshape(M, N).to(torch.int32)
    output = torch.empty_like(input)

    tmem_copy_no_scales[(1, )](input, output, M, N, BLOCK_N, swizzle, num_warps=num_warps)
    assert (output == input).all()


@pytest.mark.parametrize("N, BLOCK_N", [(256, 64), (64, 8), (64, 4)])
@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tmem_copy_noncontiguous_subslice(N, BLOCK_N):

    @gluon.jit
    def kernel(inp, out, N: ttgl.constexpr, BLOCK_N: ttgl.constexpr):
        parent = allocate_tensor_memory(ttgl.int32, [256, N], TensorMemoryLayout([128, BLOCK_N], col_stride=1))
        parent_layout: ttgl.constexpr = parent.get_reg_layout()
        parent.store(ttgl.full([256, N], -7, dtype=ttgl.int32, layout=parent_layout))
        view = parent.slice(128, 128, dim=0)

        view_layout: ttgl.constexpr = view.get_reg_layout()
        rows = ttgl.arange(0, 128, layout=ttgl.SliceLayout(1, view_layout))[:, None]
        cols = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, view_layout))[None, :]
        values = ttgl.load(inp + rows * N + cols)
        smem_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=32, rank=2)
        smem = ttgl.allocate_shared_memory(ttgl.int32, [128, N], layout=smem_layout)
        smem.store(values)

        bar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
        mbarrier.init(bar, count=1)
        tcgen05_copy(smem, view)
        tcgen05_commit(bar)
        mbarrier.wait(bar, phase=0)

        parent_rows = ttgl.arange(0, 256, layout=ttgl.SliceLayout(1, parent_layout))[:, None]
        parent_cols = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, parent_layout))[None, :]
        ttgl.store(out + parent_rows * N + parent_cols, parent.load(parent_layout))

    inp = torch.arange(128 * N, dtype=torch.int32, device="cuda").reshape(128, N)
    out = torch.empty((256, N), dtype=torch.int32, device="cuda")
    kernel[(1, )](inp, out, N, BLOCK_N, num_warps=4)
    torch.testing.assert_close(out[:128], torch.full_like(out[:128], -7), atol=0, rtol=0)
    torch.testing.assert_close(out[128:], inp, atol=0, rtol=0)


@gluon.jit
def early_return_kernel(x):
    if x.sum(0).sum(0):
        return x
    x = x + x
    return x


def test_2d_tensor_early_return():
    warp_size = ttgl.constexpr(THREADS_PER_WARP)

    @gluon.jit
    def kernel(N, out):
        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, warp_size], [1, 4], [1, 0])
        BLOCK: ttgl.constexpr = 32

        x0 = ttgl.arange(0, BLOCK, layout=ttgl.SliceLayout(1, layout))
        x1 = ttgl.arange(0, BLOCK, layout=ttgl.SliceLayout(0, layout))
        x = x0[:, None] * x1[None, :]
        for i in range(N):
            x += early_return_kernel(x)
        ttgl.store(out, x.sum(0).sum(0))

    out = torch.empty(1, dtype=torch.int32, device="cuda")
    compiled_kernel = kernel.warmup(N=100, out=out, grid=(1, ))
    assert compiled_kernel.asm["llir"].count("define") == 1


@pytest.mark.skipif(not is_hip_cdna3() and not is_hip_cdna4(), reason="Requires CDNA3 or CDNA4")
def test_inline_with_amdgpu_dialect():

    @gluon.jit
    def buffer_load(x, offsets):
        return ttgl.amd.cdna3.buffer_load(ptr=x, offsets=offsets)

    @gluon.jit
    def kernel(x, y):
        layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1], threads_per_warp=[64], warps_per_cta=[4],
                                                    order=[0])
        offsets = ttgl.arange(0, 64, layout=layout)

        a = buffer_load(x, offsets)
        ttgl.amd.cdna3.buffer_store(stored_value=a, ptr=y, offsets=offsets)

    input = torch.arange(64, device="cuda").to(torch.int32)
    output = torch.empty_like(input)

    compiled_kernel = kernel.warmup(input, output, grid=(1, ))
    assert compiled_kernel.asm["ttgir"].count("tt.func private") == 0


@pytest.mark.parametrize("interval_pairs", [[[32, 4]], [[16, 4]], [[16, 4], [64, 8]]])
@pytest.mark.parametrize(
    "shared_layout",
    [{"order": [0, 1]}, {"order": [1, 0]},
     {"offsets": [[0, 1], [0, 2], [0, 8], [0, 4], [0, 16], [0, 32], [2, 0], [1, 0], [4, 0], [8, 0], [16, 0], [32, 0]]}])
@pytest.mark.parametrize("slice_m_offset, slice_n_offset, slice_m, slice_n", [(48, 16, 16, 16), (32, 48, 32, 16),
                                                                              (48, 32, 16, 32)])
def test_padded_shared_layout_subslice(interval_pairs, shared_layout, slice_m_offset, slice_n_offset, slice_m, slice_n):
    m = 64
    n = 64
    num_warps = 1
    num_warps_cst = ttgl.constexpr(num_warps)
    warp_size_cst = ttgl.constexpr(THREADS_PER_WARP)

    shape = [m, n]
    if "order" in shared_layout:
        order = shared_layout["order"]
        smem_layout = ttgl.constexpr(ttgl.PaddedSharedLayout.with_identity_for(interval_pairs, shape, order))
    elif "offsets" in shared_layout:
        offsets = shared_layout["offsets"]
        blocks = []
        smem_layout = ttgl.constexpr(ttgl.PaddedSharedLayout(interval_pairs, offsets, blocks, shape))

    @gluon.jit
    def kernel(in_ptr, out_ptr, M: ttgl.constexpr, N: ttgl.constexpr, SLICE_M_OFFSET: ttgl.constexpr,
               SLICE_N_OFFSET: ttgl.constexpr, SLICE_M: ttgl.constexpr, SLICE_N: ttgl.constexpr):
        blocked: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [warp_size_cst, 1], [1, num_warps_cst], [1, 0])
        offs_m_load = ttgl.arange(0, M, ttgl.SliceLayout(1, blocked))
        offs_n_load = ttgl.arange(0, N, ttgl.SliceLayout(0, blocked))
        in_offs = offs_m_load[:, None] * N + offs_n_load[None, :]

        in_data = ttgl.load(in_ptr + in_offs)

        smem = ttgl.allocate_shared_memory(ttgl.int32, [M, N], smem_layout)
        smem_slice0 = smem.slice(SLICE_M_OFFSET, SLICE_M, dim=0)
        smem_slice1 = smem_slice0.slice(SLICE_N_OFFSET, SLICE_N, dim=1)

        smem.store(in_data)

        out_data = smem_slice1.load(blocked)

        offs_m_store = ttgl.arange(0, SLICE_M, ttgl.SliceLayout(1, blocked))
        offs_n_store = ttgl.arange(0, SLICE_N, ttgl.SliceLayout(0, blocked))
        out_offs = offs_m_store[:, None] * SLICE_N + offs_n_store[None, :]
        ttgl.store(out_ptr + out_offs, out_data)

    input = torch.arange(m * n, device="cuda").reshape(m, n).to(torch.int32)
    output = torch.zeros((slice_m, slice_n), dtype=torch.int32, device="cuda")
    ref_output = input[slice_m_offset:slice_m_offset + slice_m, slice_n_offset:slice_n_offset + slice_n]

    kernel[(1, )](input, output, m, n, slice_m_offset, slice_n_offset, slice_m, slice_n, num_warps=num_warps)

    assert (output == ref_output).all()


@gluon.jit
def _combine_add2(a, b, c, d):
    parent: ttgl.constexpr = ttgl.BlockedLayout([1, 2], [1, 32], [1, ttgl.num_warps()], [1, 0],
                                                cga_layout=[[0, 0]] * (ttgl.num_ctas().bit_length() - 1))
    layout: ttgl.constexpr = ttgl.SliceLayout(0, parent)
    lanes = ttgl.arange(0, 2, layout=layout)
    lhs = ttgl.where(lanes == 0, a, b)
    rhs = ttgl.where(lanes == 0, c, d)
    return ttgl.split(blackwell.add2(lhs, rhs))


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_packed_arith_reduction(dtype):

    @gluon.jit
    def kernel(a_ptr, b_ptr, out_a_ptr, out_b_ptr):
        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 4], [4, 8], [4, 1], [1, 0])
        offs_m = ttgl.arange(0, 128, ttgl.SliceLayout(1, layout))
        offs_n = ttgl.arange(0, 128, ttgl.SliceLayout(0, layout))
        offsets = offs_m[:, None] * 128 + offs_n[None, :]
        a = ttgl.load(a_ptr + offsets)
        b = ttgl.load(b_ptr + offsets)
        a, b = ttgl.reduce((a, b), axis=1, combine_fn=_combine_add2)
        ttgl.store(out_a_ptr + offs_m, a)
        ttgl.store(out_b_ptr + offs_m, b)

    torch.manual_seed(0)
    a = torch.randint(-1, 2, (128, 128), device="cuda").to(dtype)
    b = torch.randint(-1, 2, (128, 128), device="cuda").to(dtype)
    out_a = torch.empty(128, device="cuda", dtype=dtype)
    out_b = torch.empty_like(out_a)
    compiled = kernel[(1, )](a, b, out_a, out_b)

    torch.testing.assert_close(out_a, a.sum(dim=1), atol=0, rtol=0)
    torch.testing.assert_close(out_b, b.sum(dim=1), atol=0, rtol=0)
    suffix = {torch.float32: "f32x2", torch.float16: "f16x2", torch.bfloat16: "bf16x2"}[dtype]
    assert "ttng.packed_arith add" in compiled.asm["ttgir"]
    assert f"add.{suffix}" in compiled.asm["ptx"]
    if dtype == torch.float32:
        assert "prmt.b32" not in compiled.asm["ptx"]


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("op", ["add", "sub", "mul", "fma", "min", "max"])
def test_packed_arith(op, dtype):
    if dtype == torch.float32 and op in ("min", "max"):
        pytest.skip("packed float32 does not support min or max")

    op_name = op
    BLOCK_M = ttgl.constexpr(128)
    BLOCK_N = ttgl.constexpr(128)
    threads_per_warp = ttgl.constexpr(THREADS_PER_WARP)
    op = ttgl.constexpr(op)

    @gluon.jit
    def kernel(a_ptr, b_ptr, c_ptr, out_ptr):
        layout: ttgl.constexpr = ttgl.BlockedLayout(
            size_per_thread=[1, BLOCK_N],
            threads_per_warp=[threads_per_warp, 1],
            warps_per_cta=[ttgl.num_warps(), 1],
            order=[0, 1],
        )
        offs_m = ttgl.arange(0, BLOCK_M, layout=ttgl.SliceLayout(1, layout))[:, None]
        offs_n = ttgl.arange(0, BLOCK_N, layout=ttgl.SliceLayout(0, layout))[None, :]
        a = ttgl.load(a_ptr + offs_m * BLOCK_N + offs_n)
        b = ttgl.load(b_ptr + offs_m * BLOCK_N + offs_n)
        c = ttgl.load(c_ptr + offs_m * BLOCK_N + offs_n)
        result_dtype: ttgl.constexpr = a.dtype if op == "add" else None

        if op == "add":
            out = blackwell.add2(a, b, dtype=result_dtype)
        elif op == "sub":
            out = blackwell.sub2(a, b, dtype=result_dtype)
        elif op == "mul":
            out = blackwell.mul2(a, b, dtype=result_dtype)
        elif op == "fma":
            out = blackwell.fma2(a, b, c, dtype=result_dtype)
        elif op == "min":
            out = blackwell.min2(a, b, dtype=result_dtype)
        else:
            out = blackwell.max2(a, b, dtype=result_dtype)

        ttgl.store(out_ptr + offs_m * BLOCK_N + offs_n, out)

    torch.manual_seed(0)
    shape = [BLOCK_M.value, BLOCK_N.value]
    a = torch.rand(shape, dtype=dtype, device="cuda")
    b = torch.rand(shape, dtype=dtype, device="cuda")
    c = torch.rand(shape, dtype=dtype, device="cuda")
    out = torch.empty_like(a)

    compiled = kernel[(1, )](a, b, c, out)
    if op == "add":
        ref = a + b
    elif op == "sub":
        ref = a - b
    elif op == "mul":
        ref = a * b
    elif op == "fma":
        ref = torch.addcmul(c, a, b)
    elif op == "min":
        ref = torch.minimum(a, b)
    else:
        ref = torch.maximum(a, b)

    tolerance = {torch.float32: 1e-6, torch.float16: 1e-3, torch.bfloat16: 1e-2}[dtype]
    torch.testing.assert_close(ref, out, atol=tolerance, rtol=tolerance)

    suffix = {torch.float32: "f32x2", torch.float16: "f16x2", torch.bfloat16: "bf16x2"}[dtype]
    modifier = ".rn" if op_name == "fma" else ""
    assert f"{op_name}{modifier}.{suffix}" in compiled.asm["ptx"]
    assert f"ttng.packed_arith {op_name}" in compiled.asm["ttgir"]


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_packed_arith_chains(dtype):

    @gluon.jit
    def kernel(a_ptr, b_ptr, c_ptr, out_ptr):
        layout: ttgl.constexpr = ttgl.BlockedLayout([4], [32], [4], [0])
        offsets = ttgl.arange(0, 256, layout=layout)
        a = ttgl.load(a_ptr + offsets)
        b = ttgl.load(b_ptr + offsets)
        c = ttgl.load(c_ptr + offsets)

        value = blackwell.add2(a, b)
        value = blackwell.sub2(value, b)
        value = blackwell.mul2(value, b)
        value = blackwell.fma2(value, b, c)

        ttgl.store(out_ptr + offsets, value)

    torch.manual_seed(0)
    a = torch.rand(256, dtype=dtype, device="cuda")
    b = torch.rand_like(a)
    c = torch.rand_like(a)
    out = torch.empty_like(a)
    compiled = kernel[(1, )](a, b, c, out)

    reference = torch.addcmul(c, ((a + b) - b) * b, b)
    tolerance = {torch.float32: 1e-6, torch.float16: 1e-3, torch.bfloat16: 1e-2}[dtype]
    torch.testing.assert_close(reference, out, atol=tolerance, rtol=tolerance)

    suffix = {torch.float32: "f32x2", torch.float16: "f16x2", torch.bfloat16: "bf16x2"}[dtype]
    pattern = rf"\b(add|sub|mul|fma)(?:\.rn)?\.{suffix}\s+(%[\w.$]+),\s*([^;]+);"
    instructions = re.findall(pattern, compiled.asm["ptx"])
    groups = [[instruction
               for instruction in instructions
               if instruction[0] == operation]
              for operation in ("add", "sub", "mul", "fma")]
    assert all(len(group) == 2 for group in groups)
    assert [instruction[0] for instruction in instructions
            ] == [operation for operation in ("add", "sub", "mul", "fma") for _ in range(2)]
    for previous, current in zip(groups, groups[1:]):
        registers = {instruction[1] for instruction in previous}
        assert all(registers.intersection(re.findall(r"%[\w.$]+", instruction[2])) for instruction in current)
    assert "prmt.b32" not in compiled.asm["ptx"]


@pytest.mark.skipif(not is_hip_cdna4(), reason="Requires CDNA4")
def test_buffer_atomic_rmw_add_bf16():
    BLOCK = 128
    elem_type = torch.bfloat16
    SIZE_PER_THREAD = 8

    @gluon.jit
    def kernel(a, BLOCK: ttgl.constexpr, SIZE_PER_THREAD: ttgl.constexpr):
        blocked: ttgl.constexpr = ttgl.BlockedLayout([SIZE_PER_THREAD], [64], [4], [0])
        offsets = ttgl.arange(0, BLOCK, layout=blocked)
        val = ttgl.full([BLOCK], 1.0, ttgl.bfloat16, layout=blocked)
        ttgl.amd.cdna4.buffer_atomic_add(a, offsets, val, mask=1, scope="cta", sem="relaxed")

    a = torch.randn((BLOCK), dtype=elem_type, device="cuda")
    origin_a = a.clone()
    compiled = kernel[(1, )](a, BLOCK, SIZE_PER_THREAD)

    torch_ref = origin_a + torch.ones((BLOCK, ), device='cuda', dtype=torch.bfloat16)
    torch.testing.assert_close(a, torch_ref)

    ttgir = compiled.asm["ttgir"]
    assert ttgir.count("amdg.buffer_atomic_rmw fadd, relaxed, cta") == 1

    llir = compiled.asm["llir"]
    assert llir.count("tail call <2 x bfloat> @llvm.amdgcn.raw.ptr.buffer.atomic.fadd.v2bf16") == SIZE_PER_THREAD // 2


@pytest.mark.skipif(not is_ampere_or_newer(), reason="Requires Ampere or newer")
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
def test_mma_v2(dtype):
    torch.manual_seed(42)
    B = ttgl.constexpr(128)
    threads_per_warp = ttgl.constexpr(THREADS_PER_WARP)

    @gluon.jit
    def kernel(a_ptr, b_ptr, c_ptr, out_ptr):
        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [threads_per_warp, 1], [ttgl.num_warps(), 1], [1, 0])
        acc_layout: ttgl.constexpr = ttgl.NVMMADistributedLayout(version=[2, 0], warps_per_cta=[ttgl.num_warps(), 1],
                                                                 instr_shape=[16, 8])
        lhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=acc_layout, operand_index=0, k_width=8)
        rhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=acc_layout, operand_index=1, k_width=8)

        offs_m = ttgl.arange(0, B, layout=ttgl.SliceLayout(1, layout))[:, None]
        offs_n = ttgl.arange(0, B, layout=ttgl.SliceLayout(0, layout))[None, :]
        offs = offs_m * B + offs_n
        a = ttgl.convert_layout(ttgl.load(a_ptr + offs), lhs_layout)
        b = ttgl.convert_layout(ttgl.load(b_ptr + offs), rhs_layout)
        c = ttgl.convert_layout(ttgl.load(c_ptr + offs), acc_layout)
        if c.dtype == ttgl.bfloat16:
            out = mma_v2(a, b, c.to(ttgl.float32), input_precision="tf32").to(ttgl.bfloat16)
        else:
            out = mma_v2(a, b, c, input_precision="tf32")
        ttgl.store(out_ptr + offs, ttgl.convert_layout(out, layout))

    a = torch.randn((B, B), dtype=dtype, device="cuda")
    b = torch.randn((B, B), dtype=dtype, device="cuda")
    c = torch.randn((B, B), dtype=dtype, device="cuda")
    out = torch.empty((B, B), dtype=dtype, device="cuda")
    kernel[(1, )](a, b, c, out)
    torch.testing.assert_close(out, torch.addmm(c, a, b), atol=0.05, rtol=1e-2)


def test_dot_fma():
    torch.manual_seed(42)
    B = ttgl.constexpr(32)
    threads_per_warp = ttgl.constexpr(THREADS_PER_WARP)

    @gluon.jit
    def kernel(a_ptr, b_ptr, c_ptr, out_ptr):
        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [threads_per_warp, 1], [ttgl.num_warps(), 1], [1, 0])
        lhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=layout, operand_index=0, k_width=0)
        rhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=layout, operand_index=1, k_width=0)

        offs_m = ttgl.arange(0, B, layout=ttgl.SliceLayout(1, layout))[:, None]
        offs_n = ttgl.arange(0, B, layout=ttgl.SliceLayout(0, layout))[None, :]
        offs = offs_m * B + offs_n
        a = ttgl.convert_layout(ttgl.load(a_ptr + offs), lhs_layout)
        b = ttgl.convert_layout(ttgl.load(b_ptr + offs), rhs_layout)
        c = ttgl.load(c_ptr + offs)
        out = ttgl.dot_fma(a, b, c)
        ttgl.store(out_ptr + offs, out)

    a = torch.rand((B, B), dtype=torch.float32, device="cuda")
    b = torch.rand((B, B), dtype=torch.float32, device="cuda")
    c = torch.rand((B, B), dtype=torch.float32, device="cuda")
    out = torch.empty((B, B), dtype=torch.float32, device="cuda")
    kernel[(1, )](a, b, c, out)
    torch.testing.assert_close(out, torch.addmm(c, a, b), atol=1e-2, rtol=1e-2)


def test_dot3d_fma():
    torch.manual_seed(42)
    B = ttgl.constexpr(32)
    BATCH = ttgl.constexpr(8)
    threads_per_warp = ttgl.constexpr(THREADS_PER_WARP)

    @gluon.jit
    def kernel(a_ptr, b_ptr, c_ptr, out_ptr):
        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1, 1], [1, threads_per_warp, 1], [ttgl.num_warps(), 1, 1],
                                                    [2, 1, 0])
        lhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=layout, operand_index=0, k_width=0)
        rhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=layout, operand_index=1, k_width=0)

        offs_b = ttgl.arange(0, BATCH, layout=ttgl.SliceLayout(1, ttgl.SliceLayout(2, layout)))[:, None, None]
        offs_m = ttgl.arange(0, B, layout=ttgl.SliceLayout(0, ttgl.SliceLayout(2, layout)))[None, :, None]
        offs_n = ttgl.arange(0, B, layout=ttgl.SliceLayout(0, ttgl.SliceLayout(1, layout)))[None, None, :]
        offs = offs_b * B * B + offs_m * B + offs_n
        a = ttgl.convert_layout(ttgl.load(a_ptr + offs), lhs_layout)
        b = ttgl.convert_layout(ttgl.load(b_ptr + offs), rhs_layout)
        c = ttgl.load(c_ptr + offs)
        out = ttgl.dot_fma(a, b, c)
        ttgl.store(out_ptr + offs, out)

    a = torch.rand((BATCH, B, B), dtype=torch.float32, device="cuda")
    b = torch.rand((BATCH, B, B), dtype=torch.float32, device="cuda")
    c = torch.rand((BATCH, B, B), dtype=torch.float32, device="cuda")
    out = torch.empty((BATCH, B, B), dtype=torch.float32, device="cuda")
    kernel[(1, )](a, b, c, out)
    torch.testing.assert_close(out, torch.matmul(a, b) + c, atol=1e-2, rtol=1e-2)


@gluon.jit
def kernel_auto_layout_constant(threads_per_warp: ttgl.constexpr):
    BLOCK: ttgl.constexpr = 16
    SIZE: ttgl.constexpr = 10

    mask = ttgl.full(
        (BLOCK, BLOCK),
        True,
        ttgl.int1,
        ttgl.BlockedLayout(
            size_per_thread=[1, 1],
            threads_per_warp=[1, threads_per_warp],
            warps_per_cta=[1, 4],
            order=[1, 0],
        ),
    )

    mask &= (ttgl.arange(0, BLOCK, ttgl.AutoLayout()) < SIZE).expand_dims(0)
    mask &= (ttgl.arange(0, BLOCK, ttgl.AutoLayout()) < SIZE).expand_dims(1)


def test_auto_layout_constant():
    kernel_auto_layout_constant.warmup(THREADS_PER_WARP, grid=(1, ))


def test_split_auto_layout_execution():
    threads_per_warp = ttgl.constexpr(THREADS_PER_WARP)

    @gluon.jit
    def kernel(in_ptr, out_ptr, XBLOCK: ttgl.constexpr):
        in_offsets = ttgl.arange(0, 2 * XBLOCK)
        x = ttgl.load(in_ptr + in_offsets).reshape((XBLOCK, 2))
        lhs, rhs = x.split()

        out_layout: ttgl.constexpr = ttgl.BlockedLayout([1], [threads_per_warp], [4], [0])
        diff = ttgl.set_auto_layout(rhs - lhs, out_layout)
        out_offsets = ttgl.arange(0, XBLOCK)
        ttgl.store(out_ptr + out_offsets, diff)

    XBLOCK = 128
    input = torch.randint(-100, 100, (2 * XBLOCK, ), device="cuda", dtype=torch.int32)
    output = torch.empty(XBLOCK, device="cuda", dtype=torch.int32)
    ref = input.view(XBLOCK, 2)[:, 1] - input.view(XBLOCK, 2)[:, 0]

    kernel[(1, )](input, output, XBLOCK, num_warps=4)
    torch.testing.assert_close(output, ref)


def fp8e8m0_to_float32(scale):
    scale = scale.view(torch.uint8)
    scale = scale.to(torch.int32)
    scale = scale << 23
    scale = scale.view(torch.float32)
    return scale


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tcgen05_mma_scaled_minimal():
    M = 128
    N = 128
    K = 128
    threads_per_warp = ttgl.constexpr(THREADS_PER_WARP)

    @gluon.jit
    def kernel(out_ptr, M: ttgl.constexpr, N: ttgl.constexpr, K: ttgl.constexpr, a, b, a_scale, b_scale):
        # Simple register layout for creating constants and storing results
        reg_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [threads_per_warp, 1], [ttgl.num_warps(), 1], [1, 0])

        # Shared-memory layouts for MMA operands
        nvmma_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=128, transposed=False,
                                                              element_bitwidth=8, rank=2)
        # Allocate zero operands in shared memory (values don't matter since scales are zero)
        block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], warps_per_cta=[ttgl.num_warps(), 1],
                                                          order=[1, 0])
        a_offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, block_layout))[:, None]
        a_offs_k = ttgl.arange(0, K, layout=ttgl.SliceLayout(0, block_layout))[None, :]
        b_offs_k = ttgl.arange(0, K, layout=ttgl.SliceLayout(1, block_layout))[:, None]
        b_offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, block_layout))[None, :]

        a_tile = ttgl.load(a + a_offs_m * K + a_offs_k)
        b_tile = ttgl.load(b + b_offs_k * N + b_offs_n)
        a_smem = ttgl.allocate_shared_memory(ttgl.float8e5, [M, K], nvmma_layout, a_tile)
        b_smem = ttgl.allocate_shared_memory(ttgl.float8e5, [K, N], nvmma_layout, b_tile)

        # Accumulator in TMEM initialized to ones
        acc_tmem_layout: ttgl.constexpr = TensorMemoryLayout([M, N], col_stride=1)
        acc_tmem = allocate_tensor_memory(ttgl.float32, [M, N], acc_tmem_layout)
        tmem_reg_layout: ttgl.constexpr = acc_tmem.get_reg_layout()
        acc_init = ttgl.zeros([M, N], ttgl.float32, layout=tmem_reg_layout)
        acc_tmem.store(acc_init)

        # Zero scales in TMEM
        scale_layout: ttgl.constexpr = TensorMemoryScalesLayout()
        a_scale_tmem = allocate_tensor_memory(a_scale.dtype.element_ty, [M, K // 32], scale_layout)
        b_scale_tmem = allocate_tensor_memory(b_scale.dtype.element_ty, [N, K // 32], scale_layout)
        scale_reg_layout_m: ttgl.constexpr = a_scale_tmem.get_reg_layout()
        scale_reg_layout_n: ttgl.constexpr = b_scale_tmem.get_reg_layout()
        scale_offs_k = ttgl.arange(0, (K // 32), layout=ttgl.SliceLayout(0, scale_reg_layout_m))[None, :]
        scale_offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, scale_reg_layout_m))[:, None]
        scale_offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(1, scale_reg_layout_n))[:, None]
        a_scale_init = ttgl.load(a_scale + scale_offs_m * (K // 32) + scale_offs_k)
        b_scale_init = ttgl.load(b_scale + scale_offs_n * (K // 32) + scale_offs_k)
        a_scale_tmem.store(a_scale_init)
        b_scale_tmem.store(b_scale_init)

        # Issue a single scaled MMA with a completion barrier.
        bar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
        mbarrier.init(bar, count=1)
        tcgen05_mma_scaled(a_smem, b_smem, acc_tmem, a_scale_tmem, b_scale_tmem, "e5m2", "e5m2", use_acc=True,
                           mbarriers=[bar])
        mbarrier.wait(bar, phase=0)

        # Load result from TMEM and store to global
        out_reg = acc_tmem.load()
        store_layout: ttgl.constexpr = reg_layout
        offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, store_layout))[:, None]
        offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, store_layout))[None, :]
        offs = offs_m * N + offs_n
        ttgl.store(out_ptr + offs, ttgl.convert_layout(out_reg, store_layout))

    out = torch.empty((M, N), dtype=torch.float32, device="cuda")
    a = torch.randint(20, 40, (M, K), dtype=torch.uint8, device="cuda").view(torch.float8_e5m2)
    b = torch.randint(20, 40, (K, N), dtype=torch.uint8, device="cuda").view(torch.float8_e5m2)
    a_scale = torch.randint(64, 130, (M, K // 32), dtype=torch.uint8, device="cuda")
    b_scale = torch.randint(64, 130, (N, K // 32), dtype=torch.uint8, device="cuda")
    compiled = kernel[(1, )](out, M, N, K, a, b, a_scale, b_scale)
    A = a.to(torch.float32)
    B = b.to(torch.float32)
    a_scale_f32 = fp8e8m0_to_float32(a_scale)
    b_scale_f32 = fp8e8m0_to_float32(b_scale)
    a_scale_f32 = a_scale_f32.repeat_interleave(32, dim=1)
    b_scale_f32 = b_scale_f32.repeat_interleave(32, dim=1)
    b_scale_f32 = b_scale_f32.T.contiguous()
    A = A * a_scale_f32
    B = B * b_scale_f32
    ref = torch.matmul(A, B)
    torch.testing.assert_close(out, ref, atol=1e-6, rtol=1e-6)
    ttgir = compiled.asm["ttgir"]
    assert "ttng.tc_gen5_mma_scaled" in ttgir


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tcgen05_mma_scaled_noncontiguous_accumulator():

    @gluon.jit
    def kernel(a, b, out):
        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], [4, 1], [1, 0])
        smem_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=128, transposed=False,
                                                             element_bitwidth=8, rank=2)
        rows = ttgl.arange(0, 128, layout=ttgl.SliceLayout(1, layout))[:, None]
        cols = ttgl.arange(0, 128, layout=ttgl.SliceLayout(0, layout))[None, :]
        a_smem = ttgl.allocate_shared_memory(ttgl.float8e5, [128, 128], smem_layout, ttgl.load(a + rows * 128 + cols))
        b_smem = ttgl.allocate_shared_memory(ttgl.float8e5, [128, 128], smem_layout, ttgl.load(b + rows * 128 + cols))

        parent = allocate_tensor_memory(ttgl.float32, [256, 128], TensorMemoryLayout([128, 64], col_stride=1))
        parent_layout: ttgl.constexpr = parent.get_reg_layout()
        parent.store(ttgl.full([256, 128], -7.0, dtype=ttgl.float32, layout=parent_layout))
        acc = parent.slice(128, 128, dim=0)

        scale_layout: ttgl.constexpr = TensorMemoryScalesLayout()
        a_scale = allocate_tensor_memory(ttgl.uint8, [128, 4], scale_layout)
        b_scale = allocate_tensor_memory(ttgl.uint8, [128, 4], scale_layout)
        a_scale.store(ttgl.full([128, 4], 127, dtype=ttgl.uint8, layout=a_scale.get_reg_layout()))
        b_scale.store(ttgl.full([128, 4], 127, dtype=ttgl.uint8, layout=b_scale.get_reg_layout()))
        bar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
        mbarrier.init(bar, count=1)
        tcgen05_mma_scaled(a_smem, b_smem, acc, a_scale, b_scale, "e5m2", "e5m2", use_acc=False, mbarriers=[bar])
        mbarrier.wait(bar, phase=0)

        parent_rows = ttgl.arange(0, 256, layout=ttgl.SliceLayout(1, parent_layout))[:, None]
        parent_cols = ttgl.arange(0, 128, layout=ttgl.SliceLayout(0, parent_layout))[None, :]
        ttgl.store(out + parent_rows * 128 + parent_cols, parent.load(parent_layout))

    torch.manual_seed(0)
    a = torch.randint(20, 40, (128, 128), dtype=torch.uint8, device="cuda").view(torch.float8_e5m2)
    b = torch.randint(20, 40, (128, 128), dtype=torch.uint8, device="cuda").view(torch.float8_e5m2)
    out = torch.empty((256, 128), dtype=torch.float32, device="cuda")
    kernel[(1, )](a, b, out, num_warps=4)
    torch.testing.assert_close(out[:128], torch.full_like(out[:128], -7), atol=0, rtol=0)
    torch.testing.assert_close(out[128:], a.float() @ b.float(), atol=1e-6, rtol=1e-6)


@pytest.mark.parametrize(
    "a_format, a_torch_dtype, b_format, b_torch_dtype, a_is_fp4, b_is_fp4, a_fp4_padded, scale_vec_size, nvfp4",
    [
        pytest.param("e4m3", torch.float8_e4m3fn, "e4m3", torch.float8_e4m3fn, False, False, False, 32, False,
                     id="fp8_e4m3_e4m3"),
        pytest.param("e4m3", torch.float8_e4m3fn, "e5m2", torch.float8_e5m2, False, False, False, 32, False,
                     id="fp8_e4m3_e5m2"),
        pytest.param("e5m2", torch.float8_e5m2, "e4m3", torch.float8_e4m3fn, False, False, False, 32, False,
                     id="fp8_e5m2_e4m3"),
        pytest.param("e5m2", torch.float8_e5m2, "e5m2", torch.float8_e5m2, False, False, False, 32, False,
                     id="fp8_e5m2_e5m2"),
        pytest.param("e2m1", None, "e4m3", torch.float8_e4m3fn, True, False, True, 32, False, id="padded_fp4_e4m3"),
        pytest.param("e2m1", None, "e5m2", torch.float8_e5m2, True, False, True, 32, False, id="padded_fp4_e5m2"),
        pytest.param("e2m1", None, "e2m1", None, True, True, False, 32, False, id="mxfp4_mxfp4"),
        pytest.param("e2m1", None, "e2m1", None, True, True, False, 16, True, id="nvfp4_nvfp4"),
    ],
)
@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_tcgen05_mma_scaled_lhs_tmem(a_format, a_torch_dtype, b_format, b_torch_dtype, a_is_fp4, b_is_fp4, a_fp4_padded,
                                     scale_vec_size, nvfp4):
    torch.manual_seed(0)
    M = 128
    N = 128
    K = 256 if b_is_fp4 else 128
    threads_per_warp = ttgl.constexpr(THREADS_PER_WARP)

    @gluon.jit
    def kernel(out_ptr, M: ttgl.constexpr, N: ttgl.constexpr, K: ttgl.constexpr, a, b, a_scale, b_scale,
               A_FORMAT: ttgl.constexpr, B_FORMAT: ttgl.constexpr, A_IS_FP4: ttgl.constexpr, B_IS_FP4: ttgl.constexpr,
               A_FP4_PADDED: ttgl.constexpr, SCALE_VEC_SIZE: ttgl.constexpr):
        store_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [threads_per_warp, 1], [ttgl.num_warps(), 1], [1, 0])

        A_K_INPUT: ttgl.constexpr = K // 2 if A_IS_FP4 else K
        A_K_TMEM: ttgl.constexpr = K if A_FP4_PADDED else A_K_INPUT
        B_K_STORAGE: ttgl.constexpr = K // 2 if B_IS_FP4 else K
        a_tmem_layout: ttgl.constexpr = TensorMemoryLayout([M, A_K_TMEM], col_stride=1)
        a_tmem = allocate_tensor_memory(a.dtype.element_ty, [M, A_K_TMEM], a_tmem_layout)
        a_reg_layout: ttgl.constexpr = a_tmem.get_reg_layout()
        a_offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, a_reg_layout))[:, None]
        a_offs_k = ttgl.arange(0, A_K_TMEM, layout=ttgl.SliceLayout(0, a_reg_layout))[None, :]
        if A_FP4_PADDED:
            a_packed = ttgl.load(a + a_offs_m * A_K_INPUT + (a_offs_k // 2))
            # mxf8f6f4 TMEM E2M1 stores S:E2:M1 in bits [5:2].
            a_lo = (a_packed & 0x0F) << 2
            a_hi = (a_packed & 0xF0) >> 2
            a_values = ttgl.where((a_offs_k & 1) == 0, a_lo, a_hi)
        else:
            a_values = ttgl.load(a + a_offs_m * A_K_INPUT + a_offs_k)
        a_tmem.store(a_values)

        if B_IS_FP4:
            b_layout: ttgl.constexpr = ttgl.NVMMASharedLayout.get_default_for([N, B_K_STORAGE], ttgl.uint8)
            b_reg_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], [ttgl.num_warps(), 1], [1, 0])
            b_offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(1, b_reg_layout))[:, None]
            b_offs_k = ttgl.arange(0, B_K_STORAGE, layout=ttgl.SliceLayout(0, b_reg_layout))[None, :]
            b_smem = ttgl.allocate_shared_memory(ttgl.uint8, [N, B_K_STORAGE], b_layout,
                                                 ttgl.load(b + b_offs_n * B_K_STORAGE + b_offs_k))
        else:
            b_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=128, transposed=False,
                                                              element_bitwidth=8, rank=2)
            b_reg_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], [ttgl.num_warps(), 1], [1, 0])
            b_offs_k = ttgl.arange(0, K, layout=ttgl.SliceLayout(1, b_reg_layout))[:, None]
            b_offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, b_reg_layout))[None, :]
            b_smem = ttgl.allocate_shared_memory(b.dtype.element_ty, [K, N], b_layout,
                                                 ttgl.load(b + b_offs_k * N + b_offs_n))

        acc_layout: ttgl.constexpr = TensorMemoryLayout([M, N], col_stride=1)
        acc_tmem = allocate_tensor_memory(ttgl.float32, [M, N], acc_layout)
        acc_reg_layout: ttgl.constexpr = acc_tmem.get_reg_layout()
        acc_tmem.store(ttgl.zeros([M, N], ttgl.float32, layout=acc_reg_layout))

        scale_layout: ttgl.constexpr = TensorMemoryScalesLayout()
        a_scale_tmem = allocate_tensor_memory(a_scale.dtype.element_ty, [M, K // SCALE_VEC_SIZE], scale_layout)
        b_scale_tmem = allocate_tensor_memory(b_scale.dtype.element_ty, [N, K // SCALE_VEC_SIZE], scale_layout)
        scale_layout_m: ttgl.constexpr = a_scale_tmem.get_reg_layout()
        scale_layout_n: ttgl.constexpr = b_scale_tmem.get_reg_layout()
        scale_offs_k_m = ttgl.arange(0, K // SCALE_VEC_SIZE, layout=ttgl.SliceLayout(0, scale_layout_m))[None, :]
        scale_offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, scale_layout_m))[:, None]
        scale_offs_k_n = ttgl.arange(0, K // SCALE_VEC_SIZE, layout=ttgl.SliceLayout(0, scale_layout_n))[None, :]
        scale_offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(1, scale_layout_n))[:, None]
        a_scale_tmem.store(ttgl.load(a_scale + scale_offs_m * (K // SCALE_VEC_SIZE) + scale_offs_k_m))
        b_scale_tmem.store(ttgl.load(b_scale + scale_offs_n * (K // SCALE_VEC_SIZE) + scale_offs_k_n))

        bar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
        mbarrier.init(bar, count=1)
        if A_FP4_PADDED:
            padded_layout: ttgl.constexpr = TensorMemoryLayout([M, A_K_INPUT], col_stride=1, fp4_padded=True)
            a_tmem = a_tmem.reinterpret(a.dtype.element_ty, [M, A_K_INPUT], padded_layout)
        if B_IS_FP4:
            b_smem = b_smem.permute((1, 0))
        tcgen05_mma_scaled(a_tmem, b_smem, acc_tmem, a_scale_tmem, b_scale_tmem, A_FORMAT, B_FORMAT, use_acc=False,
                           mbarriers=[bar])
        mbarrier.wait(bar, phase=0, deps=[b_smem])

        out = acc_tmem.load()
        offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, store_layout))[:, None]
        offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, store_layout))[None, :]
        ttgl.store(out_ptr + offs_m * N + offs_n, ttgl.convert_layout(out, store_layout))

    out = torch.empty((M, N), dtype=torch.float32, device="cuda")
    if nvfp4:
        a_scale = torch.ones((M, K // scale_vec_size), dtype=torch.float32, device="cuda").to(torch.float8_e4m3fn)
        b_scale = torch.ones((N, K // scale_vec_size), dtype=torch.float32, device="cuda").to(torch.float8_e4m3fn)
        a_scale_ref = a_scale.to(torch.float32)
        b_scale_ref = b_scale.to(torch.float32)
    else:
        a_scale = None
        b_scale = None
        a_scale_ref = None
        b_scale_ref = None
    if a_is_fp4:
        if a_scale is None:
            a_scale = torch.full((M, K // scale_vec_size), 127, dtype=torch.uint8, device="cuda")
            a_scale_ref = fp8e8m0_to_float32(a_scale)
        a_mx = MXFP4Tensor(size=(M, K), device="cuda").random()
        a = a_mx.to_packed_tensor(dim=1)
        a_ref = a_mx.to(torch.float32)
    else:
        a_scale = torch.randint(64, 130, (M, K // scale_vec_size), dtype=torch.uint8, device="cuda")
        a_scale_ref = fp8e8m0_to_float32(a_scale)
        a = torch.randint(20, 40, (M, K), dtype=torch.uint8, device="cuda").view(a_torch_dtype)
        a_ref = a.to(torch.float32)

    if b_is_fp4:
        if b_scale is None:
            b_scale = torch.full((N, K // scale_vec_size), 127, dtype=torch.uint8, device="cuda")
            b_scale_ref = fp8e8m0_to_float32(b_scale)
        b_mx = MXFP4Tensor(size=(N, K), device="cuda").random()
        b = b_mx.to_packed_tensor(dim=1)
        b_ref = b_mx.to(torch.float32)
        b_ref = b_ref * b_scale_ref.repeat_interleave(scale_vec_size, dim=1)
        b_ref = b_ref.T.contiguous()
    else:
        b_scale = torch.randint(64, 130, (N, K // scale_vec_size), dtype=torch.uint8, device="cuda")
        b_scale_ref = fp8e8m0_to_float32(b_scale)
        b = torch.randint(20, 40, (K, N), dtype=torch.uint8, device="cuda").view(b_torch_dtype)
        b_ref = b.to(torch.float32) * b_scale_ref.repeat_interleave(scale_vec_size, dim=1).T.contiguous()

    a_ref = a_ref * a_scale_ref.repeat_interleave(scale_vec_size, dim=1)
    ref = torch.matmul(a_ref, b_ref)
    atol, rtol = (1e-3, 1e-3) if a_is_fp4 or b_is_fp4 else (1e-6, 1e-6)

    kernel[(1, )](out, M, N, K, a, b, a_scale, b_scale, a_format, b_format, a_is_fp4, b_is_fp4, a_fp4_padded,
                  scale_vec_size)
    torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)


@pytest.mark.skipif(not is_ampere_or_newer(), reason="Requires Ampere or newer")
def test_coalesced_layout():

    @gluon.jit
    def kernel(in_ptr, out_ptr,  #
               xnumel, ynumel, xstride_in, ystride_in, xstride_out, ystride_out,  #
               XBLOCK: ttgl.constexpr, YBLOCK: ttgl.constexpr):
        pid_x = ttgl.program_id(0)
        pid_y = ttgl.program_id(1)
        indices_x = pid_x * XBLOCK + ttgl.arange(0, XBLOCK, ttgl.CoalescedLayout())
        indices_y = pid_y * YBLOCK + ttgl.arange(0, YBLOCK, ttgl.CoalescedLayout())

        in_offsets = xstride_in * indices_x[:, None] + ystride_in * indices_y[None, :]
        out_offsets = xstride_out * indices_x[:, None] + ystride_out * indices_y[None, :]

        # MASK
        mask = (indices_x[:, None] < xnumel) & (indices_y[None, :] < ynumel)

        # IN PTR
        in_ptrs = in_ptr + in_offsets
        value = ttgl.load(in_ptrs, mask=mask)
        value = ttgl.sin(value)
        value = ttgl.maximum(value, 0.0)

        # OUT PTR
        out_ptrs = out_ptr + out_offsets
        ttgl.store(out_ptrs, value, mask=mask)

    XBLOCK = 128
    YBLOCK = 256
    xnumel = 1000
    ynumel = 2000
    input = torch.randn((xnumel, ynumel), device="cuda")
    output = torch.zeros_like(input)
    ref = torch.maximum(torch.sin(input), torch.tensor(0.0, device="cuda"))

    grid = (triton.cdiv(xnumel, XBLOCK), triton.cdiv(ynumel, YBLOCK))
    kernel[grid](  #
        input, output, xnumel, ynumel,  #
        *input.stride(), *output.stride(),  #
        XBLOCK, YBLOCK, num_warps=4)

    torch.testing.assert_close(output, ref)


@pytest.mark.skipif(not is_ampere_or_newer(), reason="Requires Ampere or newer")
def test_convert_auto_layout_to_coalesced_layout():

    @gluon.jit
    def kernel(in_ptr, out_ptr,  #
               xnumel, ynumel, xstride_in, ystride_in, xstride_out, ystride_out,  #
               XBLOCK: ttgl.constexpr, YBLOCK: ttgl.constexpr):
        pid_x = ttgl.program_id(0)
        pid_y = ttgl.program_id(1)
        indices_x = pid_x * XBLOCK + ttgl.arange(0, XBLOCK, ttgl.AutoLayout())
        indices_y = pid_y * YBLOCK + ttgl.arange(0, YBLOCK, ttgl.AutoLayout())

        in_offsets = xstride_in * indices_x[:, None] + ystride_in * indices_y[None, :]
        out_offsets = xstride_out * indices_x[:, None] + ystride_out * indices_y[None, :]

        # MASK
        mask = (indices_x[:, None] < xnumel) & (indices_y[None, :] < ynumel)  # auto layout

        # IN PTR
        in_ptrs = ttgl.set_auto_layout(in_ptr + in_offsets, ttgl.CoalescedLayout())
        value = ttgl.load(in_ptrs, mask=mask)

        # OUT PTR
        out_ptrs = ttgl.set_auto_layout(out_ptr + out_offsets, ttgl.CoalescedLayout())
        out_mask_layouted = ttgl.set_auto_layout(mask, ttgl.CoalescedLayout())
        ttgl.store(out_ptrs, value, mask=out_mask_layouted)

    XBLOCK = 128
    YBLOCK = 256
    xnumel = 1000
    ynumel = 2000
    input = torch.ones((xnumel, ynumel), device="cuda")
    output = torch.zeros_like(input)
    ref = torch.ones_like(input)

    grid = (triton.cdiv(xnumel, XBLOCK), triton.cdiv(ynumel, YBLOCK))
    kernel[grid](  #
        input, output, xnumel, ynumel,  #
        *input.stride(), *output.stride(),  #
        XBLOCK, YBLOCK, num_warps=4)

    torch.testing.assert_close(output, ref)


@gluon.jit
def in_thread_transpose_roundtrip_kernel(input, output, M: ttgl.constexpr, N: ttgl.constexpr,
                                         first_layout: ttgl.constexpr, second_layout: ttgl.constexpr,
                                         shared_layout: ttgl.constexpr):
    offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, first_layout))[:, None]
    offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, first_layout))[None, :]

    load_data = ttgl.load(input + offs_m * N + offs_n)
    converted_data = ttgl.convert_layout(load_data, second_layout)
    smem = ttgl.allocate_shared_memory(input.dtype.element_ty, [M, N], shared_layout, converted_data)
    out_data = smem.load(first_layout)
    ttgl.store(output + offs_m * N + offs_n, out_data)


@pytest.mark.skipif(not (is_hip_cdna() or is_hip_rdna()),
                    reason="Correctness tests for special cases on AMD architectures")
@pytest.mark.parametrize("src_reg_bases, dst_reg_bases", [
    ([[0, 1], [0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], [[1, 0], [2, 0], [4, 0], [0, 1], [0, 2], [0, 4]]),
    ([[0, 1], [0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], [[0, 1], [0, 4], [0, 2], [1, 0], [2, 0], [4, 0]]),
    ([[0, 1], [0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], [[0, 2], [0, 1], [0, 4], [1, 0], [4, 0], [2, 0]]),
    ([[0, 0], [0, 1], [0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], [[0, 2], [0, 1], [0, 4], [1, 0], [4, 0], [2, 0]]),
    ([[0, 1], [0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], [[0, 2], [0, 0], [0, 1], [0, 4], [1, 0], [4, 0], [2, 0]]),
    ([[0, 1], [0, 0], [0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], [[0, 0], [0, 2], [0, 1], [0, 4], [1, 0], [4, 0], [2, 0]
                                                                ]),
    ([[0, 1], [0, 2], [0, 4], [0, 0], [1, 0], [2, 0], [4, 0]], [[0, 2], [0, 1], [0, 0], [0, 4], [1, 0], [4, 0], [2, 0]
                                                                ]),
])
def test_in_thread_convert_layout_8bit(src_reg_bases, dst_reg_bases):
    torch.manual_seed(0)
    dtype = torch.int8
    M = 8
    N = 8 * THREADS_PER_WARP

    numLaneBases = int(math.log2(THREADS_PER_WARP))
    lane_bases = [[0, 8 * (2**baseNo)] for baseNo in range(numLaneBases)]
    warp_bases = []
    first_layout = ttgl.DistributedLinearLayout(reg_bases=src_reg_bases, lane_bases=lane_bases, warp_bases=warp_bases,
                                                block_bases=[], shape=[M, N])

    second_layout = ttgl.DistributedLinearLayout(reg_bases=dst_reg_bases, lane_bases=lane_bases, warp_bases=warp_bases,
                                                 block_bases=[], shape=[M, N])

    shared_layout = ttgl.SwizzledSharedLayout(1, 1, 1, order=[0, 1])
    input_buffer = (torch.randn((M, N), device="cuda") * 100).to(dtype)
    output_buffer = torch.zeros((M, N), device="cuda", dtype=dtype)
    pgm = in_thread_transpose_roundtrip_kernel[(1, )](input_buffer, output_buffer, M, N, first_layout, second_layout,
                                                      shared_layout, num_warps=1)

    assert re.search(r"v_perm", pgm.asm['amdgcn'], re.MULTILINE)
    torch.testing.assert_close(input_buffer, output_buffer, atol=1e-3, rtol=1e-3)


@gluon.jit
def descriptor_shape_kernel(desc, expect_shape):
    for i in ttgl.static_range(len(expect_shape)):
        ttgl.device_assert(desc.shape[i] == expect_shape[i])


@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_descriptor_shape():
    t = torch.randint(0, 256, (512, 512), dtype=torch.uint8)

    for fp4_padded in [True]:
        layout = ttgl.NVMMASharedLayout.get_default_for([128, 64], ttgl.uint8, fp4_padded=fp4_padded)
        desc = TensorDescriptor.from_tensor(t, [128, 64], layout)
        descriptor_shape_kernel[(1, )](desc, t.shape, num_warps=1, debug=True)
        torch.cuda.synchronize()


@gluon.jit
def shared_gather_kernel(
    matrix_ptr,
    indices_ptr,
    output_ptr,
    N: ttgl.constexpr,
    M: ttgl.constexpr,
    layout_2d: ttgl.constexpr,
    layout_1d: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
):
    """Test shared memory gather using smem.gather() with axis-based API."""
    # Load the matrix from global memory into registers
    indices_x = ttgl.arange(0, N, layout=ttgl.SliceLayout(dim=1, parent=layout_2d))
    indices_y = ttgl.arange(0, M, layout=ttgl.SliceLayout(dim=0, parent=layout_2d))
    offsets_2d = indices_x[:, None] * M + indices_y[None, :]
    matrix_data = ttgl.load(matrix_ptr + offsets_2d)

    # Allocate 2D shared memory and store the matrix
    smem_2d = ttgl.allocate_shared_memory(ttgl.float32, [N, M], layout=shared_layout)
    smem_2d.store(matrix_data)

    # Reshape to 1D to test gather along axis 0
    smem_1d = smem_2d.reshape([N * M])

    # Load the gather indices (diagonal elements: 0, M+1, 2*(M+1), ...)
    offsets_1d = ttgl.arange(0, N, layout=layout_1d)
    indices = ttgl.load(indices_ptr + offsets_1d)

    # Gather using axis-based API: result[i] = smem_1d[indices[i]]
    gathered = smem_1d.gather(indices, axis=0)

    # Store result to global memory
    ttgl.store(output_ptr + offsets_1d, gathered)


@pytest.mark.parametrize("N,M", [(32, 32), (64, 64), (128, 128)])
def test_shared_gather(N, M):
    """Test gathering from 1D reshaped shared memory (diagonal of 2D matrix)."""
    device = torch.device("cuda")

    # Create a test matrix with known values
    matrix = torch.arange(N * M, dtype=torch.float32, device=device).reshape(N, M)

    # Create gather indices for diagonal elements: 0, M+1, 2*(M+1), ...
    indices = torch.arange(N, dtype=torch.int32, device=device) * (M + 1)

    output = torch.zeros(N, dtype=torch.float32, device=device)

    # Compute expected result: diagonal elements
    expected = matrix.flatten()[indices]

    # Create layouts dynamically based on THREADS_PER_WARP
    layout_2d = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[THREADS_PER_WARP // 4, 4],
                                   warps_per_cta=[1, 1], order=[1, 0])
    layout_1d = ttgl.BlockedLayout(size_per_thread=[1], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[1],
                                   order=[0])
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    # Launch kernel
    shared_gather_kernel[(1, )](
        matrix,
        indices,
        output,
        N=N,
        M=M,
        layout_2d=layout_2d,
        layout_1d=layout_1d,
        shared_layout=shared_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


@gluon.jit
def shared_gather_scatter_two_ctas_kernel(
    inp,
    out,
    layout: ttgl.constexpr,
    GATHER: ttgl.constexpr,
):
    shared_layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, order=[1, 0], cga_layout=[[0, 1]])
    rows = ttgl.arange(0, 2, layout=ttgl.SliceLayout(1, layout))
    cols = ttgl.arange(0, 32, layout=ttgl.SliceLayout(0, layout))
    offsets = rows[:, None] * 32 + cols[None, :]
    values = ttgl.load(inp + offsets)
    smem = ttgl.allocate_shared_memory(ttgl.int32, [2, 32], shared_layout, value=values)
    ttgl.barrier(cluster=True)

    peer_cols = (cols ^ 1)[None, :] + rows[:, None] * 0
    if GATHER:
        result = smem.gather(peer_cols, axis=1)
    else:
        smem.scatter(values, peer_cols, axis=1)
        ttgl.barrier(cluster=True)
        result = smem.load(layout)
    ttgl.store(out + offsets, result)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
@pytest.mark.parametrize("gather", [True, False], ids=["gather", "scatter"])
def test_shared_gather_scatter_two_ctas(gather):
    inp = torch.arange(64, dtype=torch.int32, device="cuda").reshape(2, 32)
    out = torch.empty_like(inp)
    layout = ttgl.BlockedLayout([1, 1], [1, 32], [1, 4], [1, 0], cga_layout=[[0, 1]])

    shared_gather_scatter_two_ctas_kernel[(1, )](
        inp,
        out,
        layout,
        GATHER=gather,
        num_warps=4,
        num_ctas=2,
    )

    torch.testing.assert_close(out, inp.reshape(2, 16, 2).flip(-1).reshape(2, 32))


def _shared_gather_cga_cases():
    cases = [pytest.param(1, 0, id="1-cta-local")]
    for num_ctas in (2, 4, 8, 16):
        cases.append(pytest.param(num_ctas, 0, id=f"{num_ctas}-cta-sharded"))
        if num_ctas >= 4:
            # Keep bit 0 sharded for a remote read and replicate alternating
            # higher bits, including the noncontiguous mask 0b1010 at 16 CTAs.
            broadcast_mask = sum(1 << bit for bit in range(1, num_ctas.bit_length() - 1, 2))
            cases.append(pytest.param(num_ctas, broadcast_mask, id=f"{num_ctas}-cta-partial"))
        cases.append(pytest.param(num_ctas, num_ctas - 1, id=f"{num_ctas}-cta-broadcast"))
    return cases


def _shared_gather_layouts(num_ctas, broadcast_mask):
    num_cta_bits = num_ctas.bit_length() - 1
    full_cga = tuple((1 << bit, ) for bit in range(num_cta_bits))
    shared_cga = []
    shard_bit = 0
    for bit in range(num_cta_bits):
        if broadcast_mask & (1 << bit):
            shared_cga.append((0, ))
        else:
            shared_cga.append((1 << shard_bit, ))
            shard_bit += 1
    value_layout = ttgl.BlockedLayout([1], [32], [1], [0], cga_layout=full_cga)
    seed_layout = ttgl.SwizzledSharedLayout(1, 1, 1, order=[0], cga_layout=full_cga)
    shared_layout = ttgl.SwizzledSharedLayout(1, 1, 1, order=[0], cga_layout=shared_cga)
    return value_layout, seed_layout, shared_layout, shard_bit


@gluon.jit
def shared_gather_cga_kernel(
    out,
    value_layout: ttgl.constexpr,
    seed_layout: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
    BROADCAST_MASK: ttgl.constexpr,
    NUM_SHARD_BITS: ttgl.constexpr,
    CTA_TILE: ttgl.constexpr,
):
    num_ctas: ttgl.constexpr = ttgl.num_ctas()
    num_cta_bits: ttgl.constexpr = num_ctas.bit_length() - 1
    block: ttgl.constexpr = CTA_TILE * num_ctas
    offsets = ttgl.arange(0, block, layout=value_layout)
    # Seed each physical CTA with distinct layout-derived values, then view the
    # same allocation through the CGA layout under test.
    seed = ttgl.allocate_shared_memory(ttgl.int32, [block], seed_layout, value=offsets)
    smem = seed.reinterpret(shape=[CTA_TILE << NUM_SHARD_BITS], layout=shared_layout)
    ttgl.barrier(cluster=True)

    cta = offsets // CTA_TILE
    local = offsets & (CTA_TILE - 1)
    # Compress the non-broadcast CTA bits into a shard ID. Flipping its low bit
    # forces a remote read while leaving the replica bits unchanged.
    shard = local * 0
    shard_bit = 0
    for bit in range(num_cta_bits):
        if (BROADCAST_MASK & (1 << bit)) == 0:
            shard |= ((cta >> bit) & 1) << shard_bit
            shard_bit += 1
    if NUM_SHARD_BITS:
        indices = (shard ^ 1) * CTA_TILE + local
    else:
        indices = local ^ 1
    result = smem.gather(indices, axis=0)
    ttgl.store(out + offsets, result)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
@pytest.mark.parametrize("num_ctas,broadcast_mask", _shared_gather_cga_cases())
def test_shared_gather_cga(num_ctas, broadcast_mask):
    cta_tile = 32
    out = torch.empty(cta_tile * num_ctas, dtype=torch.int32, device="cuda")
    value_layout, seed_layout, shared_layout, num_shard_bits = _shared_gather_layouts(num_ctas, broadcast_mask)

    shared_gather_cga_kernel[(1, )](
        out,
        value_layout,
        seed_layout,
        shared_layout,
        BROADCAST_MASK=broadcast_mask,
        NUM_SHARD_BITS=num_shard_bits,
        CTA_TILE=cta_tile,
        num_warps=1,
        num_ctas=num_ctas,
    )

    offsets = torch.arange(out.numel(), dtype=torch.int32, device="cuda")
    if broadcast_mask == num_ctas - 1:
        expected = (offsets // cta_tile) * cta_tile + ((offsets % cta_tile) ^ 1)
    else:
        expected = ((offsets // cta_tile) ^ 1) * cta_tile + offsets % cta_tile
    torch.testing.assert_close(out, expected)


@gluon.jit
def shared_subslice_two_ctas_kernel(
    inp,
    values,
    result_out,
    final_out,
    OP: ttgl.constexpr,
    alloc_layout: ttgl.constexpr,
    tile_layout: ttgl.constexpr,
):
    shared_layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, order=[1, 0], cga_layout=[[1, 0]])
    alloc_rows = ttgl.arange(0, 4, layout=ttgl.SliceLayout(1, alloc_layout))
    alloc_cols = ttgl.arange(0, 32, layout=ttgl.SliceLayout(0, alloc_layout))
    alloc_offsets = alloc_rows[:, None] * 32 + alloc_cols[None, :]
    initial = ttgl.load(inp + alloc_offsets)
    smem = ttgl.allocate_shared_memory(ttgl.int32, [4, 32], shared_layout, value=initial)
    ttgl.barrier(cluster=True)

    tile = smem.slice(2, 2, dim=0)
    tile_rows = ttgl.arange(0, 2, layout=ttgl.SliceLayout(1, tile_layout))
    tile_cols = ttgl.arange(0, 32, layout=ttgl.SliceLayout(0, tile_layout))
    tile_offsets = tile_rows[:, None] * 32 + tile_cols[None, :]

    if OP == "load":
        result = tile.load(tile_layout)
        ttgl.store(result_out + tile_offsets, result)
    elif OP == "gather":
        peer_cols = (tile_cols ^ 1)[None, :] + tile_rows[:, None] * 0
        result = tile.gather(peer_cols, axis=1)
        ttgl.store(result_out + tile_offsets, result)
    else:
        update = ttgl.load(values + tile_offsets)
        if OP == "store":
            tile.store(update)
        elif OP == "scatter":
            peer_cols = (tile_cols ^ 1)[None, :] + tile_rows[:, None] * 0
            tile.scatter(update, peer_cols, axis=1)
        else:
            indices = tile_cols[None, :] + tile_rows[:, None] * 0
            old = tile.atomic_scatter_add(update, indices, axis=1)
            ones = ttgl.full([2, 32], 1, ttgl.int32, tile_layout)
            tile.atomic_scatter_add(ones, indices, axis=1)
            ttgl.store(result_out + tile_offsets, old)
        ttgl.barrier(cluster=True)
        final = smem.load(alloc_layout)
        ttgl.store(final_out + alloc_offsets, final)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
@pytest.mark.parametrize("op", ["load", "store", "gather", "scatter", "atomic"])
def test_shared_subslice_two_ctas(op):
    inp = torch.arange(4 * 32, dtype=torch.int32, device="cuda").reshape(4, 32)
    values = torch.arange(2 * 32, dtype=torch.int32, device="cuda").reshape(2, 32) + 1000
    if op == "atomic":
        values.fill_(1)
    result = torch.empty((2, 32), dtype=torch.int32, device="cuda")
    final = torch.empty_like(inp)
    alloc_layout = ttgl.BlockedLayout([1, 1], [1, 32], [1, 4], [1, 0], cga_layout=[[1, 0]])
    tile_layout = ttgl.BlockedLayout([1, 1], [1, 32], [1, 4], [1, 0], cga_layout=[[0, 0]])

    compiled = shared_subslice_two_ctas_kernel[(1, )](
        inp,
        values,
        result,
        final,
        OP=op,
        alloc_layout=alloc_layout,
        tile_layout=tile_layout,
        num_warps=4,
        num_ctas=2,
    )

    ptx = compiled.asm["ptx"]
    if op in ("load", "gather"):
        assert "ld.shared::cluster" in ptx
        expected = inp[2:4]
        if op == "gather":
            expected = expected.reshape(2, 16, 2).flip(-1).reshape(2, 32)
        torch.testing.assert_close(result, expected, atol=0, rtol=0)
    else:
        expected = inp.clone()
        if op == "store":
            expected[2:4] = values
            assert "st.shared::cluster" in ptx
        elif op == "scatter":
            expected[2:4] = values.reshape(2, 16, 2).flip(-1).reshape(2, 32)
            assert "st.shared::cluster" in ptx
        else:
            expected[2:4] += values + 1
            assert "atom.shared::cluster.cluster.relaxed.add.u32" in ptx
            assert "red.shared::cluster.cluster.relaxed.inc.u32" in ptx
            torch.testing.assert_close(result, inp[2:4], atol=0, rtol=0)
        torch.testing.assert_close(final, expected, atol=0, rtol=0)


@gluon.jit
def cross_cta_tma_kernel(desc, out, cga_layout: ttgl.constexpr, MODE: ttgl.constexpr):
    alloc_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [4, 8], [4, 1], [1, 0], cga_layout=cga_layout)
    shared_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(128, 32, rank=2, cga_layout=cga_layout)
    parent = ttgl.allocate_shared_memory(ttgl.int32, [256, 64], shared_layout,
                                         value=ttgl.full([256, 64], -5, ttgl.int32, layout=alloc_layout))
    tile = parent.slice(128, 128, dim=0)
    if MODE == "load" or MODE == "multicast":
        bar = mbarrier.allocate_mbarrier()
        mbarrier.init(bar, count=1)
        mbarrier.expect(bar, desc.nbytes_per_cta)
        tma.async_load(desc, [0, 0], bar, tile, multicast=MODE == "multicast")
        mbarrier.wait(bar, phase=0)
        ttgl.barrier(cluster=True)
        result = parent.load(alloc_layout)
        offsets = ttgl.arange(0, 256)[:, None] * 64 + ttgl.arange(0, 64)[None, :]
        offsets = ttgl.set_auto_layout(offsets, result.type.layout)
        ttgl.store(out + offsets, result)
    elif MODE == "store":
        tma.async_store(desc, [0, 0], tile)
    else:
        tma.async_atomic_add(desc, [0, 0], tile)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
@pytest.mark.parametrize("mode,cga_layout", [
    ("load", ((1, 0), )),
    ("load", ((1, 0), (0, 0))),
    ("multicast", ((1, 0), (0, 0))),
    ("load", ((0, 0), (1, 0))),
    ("store", ((1, 0), )),
    ("reduce", ((1, 0), )),
], ids=["two-ctas", "four-ctas", "multicast", "nonpeer", "store", "reduce"])
def test_cross_cta_tma(mode, cga_layout, capfd):
    if mode in ("load", "multicast") and not is_blackwell():
        pytest.skip("Requires Blackwell CTA-pair completion")
    inp = torch.arange(128 * 64, device="cuda", dtype=torch.int32).reshape(128, 64)
    out = torch.full((256, 64), -1, device="cuda", dtype=torch.int32)
    layout = ttgl.NVMMASharedLayout(128, 32, rank=2, cga_layout=cga_layout)
    desc = TensorDescriptor.from_tensor(inp, [128, 64], layout)

    if mode not in ("load", "multicast"):
        with pytest.raises(RuntimeError, match="error encountered during parsing"):
            cross_cta_tma_kernel.warmup(desc, out, cga_layout, MODE=mode, grid=(1, ), num_warps=4, num_ctas=2)
        assert "source subview may have an origin in another CTA" in "".join(capfd.readouterr())
    elif cga_layout[0] == (0, 0):
        with pytest.raises(RuntimeError, match="error encountered during parsing"):
            cross_cta_tma_kernel.warmup(desc, out, cga_layout, MODE=mode, grid=(1, ), num_warps=4, num_ctas=4)
        assert "TMA destination and completion barrier must belong to the same CTA or a CTA pair" in "".join(
            capfd.readouterr())
    else:
        compiled = cross_cta_tma_kernel[(1, )](desc, out, cga_layout, mode, num_warps=4, num_ctas=1 << len(cga_layout))
        assert "mapa.shared::cluster" in compiled.asm["ptx"]
        assert "cp.async.bulk.tensor.2d.cta_group::2.shared::cluster.global" in compiled.asm["ptx"]
        torch.testing.assert_close(out, torch.cat((torch.full_like(inp, -5), inp)), atol=0, rtol=0)


@gluon.jit
def shared_subslice_swizzled_register_block_kernel(inp, out, alloc_layout: ttgl.constexpr, tile_layout: ttgl.constexpr,
                                                   shared_layout: ttgl.constexpr):
    rows = ttgl.arange(0, 4, layout=ttgl.SliceLayout(1, alloc_layout))
    cols = ttgl.arange(0, 32, layout=ttgl.SliceLayout(0, alloc_layout))
    offsets = rows[:, None] * 32 + cols[None, :]
    initial = ttgl.load(inp + offsets)
    smem = ttgl.allocate_shared_memory(ttgl.int32, [4, 32], shared_layout, value=initial)
    ttgl.barrier(cluster=True)

    tile = smem.slice(2, 2, dim=0)
    loaded = tile.load(tile_layout)
    tile_rows = ttgl.arange(0, 2, layout=ttgl.SliceLayout(1, tile_layout))
    tile_cols = ttgl.arange(0, 32, layout=ttgl.SliceLayout(0, tile_layout))
    tile_offsets = tile_rows[:, None] * 32 + tile_cols[None, :]
    ttgl.store(out + tile_offsets, loaded)


@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
def test_shared_subslice_swizzled_register_block():
    # The shared layout maps logical row bit 0 to an offset bit and block bit 0,
    # while logical row bit 1 maps to block bit 0. Slicing at row 2 therefore
    # makes the affine block offset overlap the tile's register block basis.
    shared_layout = ttgl.SharedLinearLayout(
        offset_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [3, 0]],
        block_bases=[[2, 0]],
    )
    alloc_layout = ttgl.DistributedLinearLayout(
        reg_bases=[],
        lane_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16]],
        warp_bases=[[1, 0], [2, 0]],
        block_bases=[[0, 0]],
        shape=[4, 32],
    )
    tile_layout = ttgl.DistributedLinearLayout(
        reg_bases=[[1, 0]],
        lane_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16]],
        warp_bases=[[0, 0], [0, 0]],
        block_bases=[[0, 0]],
        shape=[2, 32],
    )

    inp = torch.arange(4 * 32, dtype=torch.int32, device="cuda").reshape(4, 32)
    out = torch.empty((2, 32), dtype=torch.int32, device="cuda")
    shared_subslice_swizzled_register_block_kernel[(1, )](
        inp,
        out,
        alloc_layout=alloc_layout,
        tile_layout=tile_layout,
        shared_layout=shared_layout,
        num_warps=4,
        num_ctas=2,
    )
    torch.testing.assert_close(out, inp[2:4], atol=0, rtol=0)


@gluon.jit
def shared_scatter_kernel(
    indices_ptr,
    values_ptr,
    output_ptr,
    N: ttgl.constexpr,
    M: ttgl.constexpr,
    layout_2d: ttgl.constexpr,
    layout_1d: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
    use_broadcast: ttgl.constexpr,
):
    """Test shared memory scatter using smem.scatter() with axis-based API."""
    # Allocate 2D shared memory initialized to zero
    smem = ttgl.allocate_shared_memory(ttgl.float32, [N, M], layout=shared_layout)

    # Initialize shared memory to zero
    indices_x = ttgl.arange(0, N, layout=ttgl.SliceLayout(dim=1, parent=layout_2d))
    indices_y = ttgl.arange(0, M, layout=ttgl.SliceLayout(dim=0, parent=layout_2d))
    offsets_2d = indices_x[:, None] * M + indices_y[None, :]
    zeros = ttgl.zeros([N, M], ttgl.float32, layout=layout_2d)
    smem.store(zeros)

    if use_broadcast:
        # values has shape [N, 1] and indices has shape [1, M]. scatter broadcasts
        # both to [N, M], so each row's scalar value is written across columns:
        # smem[i, indices[0, j]] = values[i, 0].
        values = ttgl.load(values_ptr + indices_x)[:, None]
        indices = indices_y[None, :]
        smem.scatter(values, indices, axis=1)
    else:
        # Reshape to 1D to test scatter along axis 0
        smem_1d = smem.reshape([N * M])

        # Load the scatter indices and values (diagonal elements: 0, M+1, 2*(M+1), ...)
        offsets_1d = ttgl.arange(0, N, layout=layout_1d)
        indices = ttgl.load(indices_ptr + offsets_1d)
        values = ttgl.load(values_ptr + offsets_1d)

        # Scatter using axis-based API: smem_1d[indices[i]] = values[i]
        smem_1d.scatter(values, indices, axis=0)

    # Read back the full matrix from shared memory
    matrix_data = smem.load(layout=layout_2d)

    # Store result to global memory
    ttgl.store(output_ptr + offsets_2d, matrix_data)


@pytest.mark.parametrize("N,M", [(32, 32), (64, 64), (128, 128)])
def test_shared_scatter(N, M):
    """Test scattering to 1D reshaped shared memory (diagonal of 2D matrix)."""
    device = torch.device("cuda")

    # Create scatter indices for diagonal elements: 0, M+1, 2*(M+1), ...
    indices = torch.arange(N, dtype=torch.int32, device=device) * (M + 1)

    # Create values to scatter
    values = torch.arange(N, dtype=torch.float32, device=device) + 100.0

    output = torch.zeros((N, M), dtype=torch.float32, device=device)

    # Compute expected result: matrix starts at zero, then diagonal gets values
    expected = torch.zeros((N, M), dtype=torch.float32, device=device)
    for i in range(N):
        expected[i, i] = values[i]

    # Create layouts dynamically based on THREADS_PER_WARP
    layout_2d = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[THREADS_PER_WARP // 4, 4],
                                   warps_per_cta=[1, 1], order=[1, 0])
    layout_1d = ttgl.BlockedLayout(size_per_thread=[1], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[1],
                                   order=[0])
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    # Launch kernel
    shared_scatter_kernel[(1, )](
        indices,
        values,
        output,
        N=N,
        M=M,
        layout_2d=layout_2d,
        layout_1d=layout_1d,
        shared_layout=shared_layout,
        use_broadcast=False,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)

    broadcast_output = torch.empty((N, M), dtype=torch.float32, device=device)
    shared_scatter_kernel[(1, )](
        indices,
        values,
        broadcast_output,
        N=N,
        M=M,
        layout_2d=layout_2d,
        layout_1d=layout_1d,
        shared_layout=shared_layout,
        use_broadcast=True,
        num_warps=1,
    )

    torch.testing.assert_close(broadcast_output, values[:, None].expand(N, M).contiguous())


@gluon.jit
def shared_atomic_scatter_rmw_kernel(
    values_ptr,
    indices_ptr,
    mask_ptr,
    old_ptr,
    final_ptr,
    N: ttgl.constexpr,
    M: ttgl.constexpr,
    axis: ttgl.constexpr,
    op: ttgl.constexpr,
    dtype: ttgl.constexpr,
    use_mask: ttgl.constexpr,
    use_constant_values: ttgl.constexpr,
    init_value: ttgl.constexpr,
    RHS_N: ttgl.constexpr,
    RHS_M: ttgl.constexpr,
    layout_2d: ttgl.constexpr,
    layout_rhs: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
):
    indices_x = ttgl.arange(0, N, layout=ttgl.SliceLayout(dim=1, parent=layout_2d))
    indices_y = ttgl.arange(0, M, layout=ttgl.SliceLayout(dim=0, parent=layout_2d))
    offsets_2d = indices_x[:, None] * M + indices_y[None, :]
    rhs_x = ttgl.arange(0, RHS_N, layout=ttgl.SliceLayout(dim=1, parent=layout_rhs))
    rhs_y = ttgl.arange(0, RHS_M, layout=ttgl.SliceLayout(dim=0, parent=layout_rhs))
    offsets_rhs = rhs_x[:, None] * RHS_M + rhs_y[None, :]

    smem = ttgl.allocate_shared_memory(dtype, [N, M], layout=shared_layout)
    smem.store(ttgl.full([N, M], init_value, dtype, layout=layout_2d))

    indices = ttgl.load(indices_ptr + offsets_rhs)
    if use_constant_values:
        # Cover the all-ones RHS path; int32 add can lower to atom.inc on NVIDIA.
        values = ttgl.full([RHS_N, RHS_M], 1, dtype, layout_rhs)
    else:
        values = ttgl.load(values_ptr + offsets_rhs)

    if use_mask:
        mask = ttgl.load(mask_ptr + offsets_rhs)
        if op == "add":
            old = smem.atomic_scatter_add(values, indices, axis=axis, mask=mask)
        elif op == "max":
            old = smem.atomic_scatter_max(values, indices, axis=axis, mask=mask)
        elif op == "min":
            old = smem.atomic_scatter_min(values, indices, axis=axis, mask=mask)
        elif op == "and":
            old = smem.atomic_scatter_and(values, indices, axis=axis, mask=mask)
        elif op == "or":
            old = smem.atomic_scatter_or(values, indices, axis=axis, mask=mask)
        elif op == "xor":
            old = smem.atomic_scatter_xor(values, indices, axis=axis, mask=mask)
        else:
            old = smem.atomic_scatter_xchg(values, indices, axis=axis, mask=mask)
    else:
        if op == "add":
            old = smem.atomic_scatter_add(values, indices, axis=axis)
        elif op == "max":
            old = smem.atomic_scatter_max(values, indices, axis=axis)
        elif op == "min":
            old = smem.atomic_scatter_min(values, indices, axis=axis)
        elif op == "and":
            old = smem.atomic_scatter_and(values, indices, axis=axis)
        elif op == "or":
            old = smem.atomic_scatter_or(values, indices, axis=axis)
        elif op == "xor":
            old = smem.atomic_scatter_xor(values, indices, axis=axis)
        else:
            old = smem.atomic_scatter_xchg(values, indices, axis=axis)
    ttgl.store(old_ptr + offsets_rhs, old)

    final = smem.load(layout=layout_2d)
    ttgl.store(final_ptr + offsets_2d, final)


def _expected_shared_atomic_scatter_rmw(op, init_value, indices, values, mask, axis, shape):
    device = values.device
    expected = torch.full(shape, init_value, dtype=values.dtype, device=device)
    rhs_n, rhs_m = values.shape

    def combine(old, val):
        if op == "add":
            return old + val
        if op == "max":
            return torch.maximum(old, val)
        if op == "min":
            return torch.minimum(old, val)
        if op == "and":
            return old & val
        if op == "or":
            return old | val
        if op == "xor":
            return old ^ val
        # xchg
        return val

    if axis == 1:
        row_ids = torch.arange(rhs_n, device=device)
        for j in range(rhs_m):
            active = mask[:, j]
            rows = row_ids[active]
            cols = indices[active, j].long()
            expected[rows, cols] = combine(expected[rows, cols], values[active, j])
    else:
        col_ids = torch.arange(rhs_m, device=device)
        for i in range(rhs_n):
            active = mask[i, :]
            rows = indices[i, active].long()
            cols = col_ids[active]
            expected[rows, cols] = combine(expected[rows, cols], values[i, active])
    return expected


def _shared_atomic_scatter_rmw_cases():
    int32_dtype_cases = [
        ("unmasked_int32", False, torch.int32, ttgl.int32),
        ("masked_int32", True, torch.int32, ttgl.int32),
    ]
    add_dtype_cases = [
        *int32_dtype_cases,
        ("unmasked_float16", False, torch.float16, ttgl.float16),
        ("unmasked_float32", False, torch.float32, ttgl.float32),
        ("masked_float32", True, torch.float32, ttgl.float32),
    ]
    xchg_dtype_cases = [
        *int32_dtype_cases,
        ("unmasked_float32", False, torch.float32, ttgl.float32),
        ("masked_float32", True, torch.float32, ttgl.float32),
    ]
    value_cases = [
        ("loaded_values", False),
        ("constant_values", True),
    ]
    shape_cases = [
        ("axis1_rhs_cols", 16, 32, 1, (16, 2)),
        ("axis0_rhs_rows", 32, 16, 0, (2, 16)),
    ]
    op_cases = [
        ("add", 0, add_dtype_cases),
        ("max", 0, int32_dtype_cases),
        ("min", 99, int32_dtype_cases),
        ("and", 15, int32_dtype_cases),
        ("or", 0, int32_dtype_cases),
        ("xor", 0, int32_dtype_cases),
        ("xchg", 0, xchg_dtype_cases),
    ]

    cases = []
    for op, init_value, dtype_cases in op_cases:
        for dtype_id, use_mask, torch_dtype, gluon_dtype in dtype_cases:
            for values_id, use_constant_values in value_cases:
                for shape_id, N, M, axis, rhs_shape in shape_cases:
                    id_parts = [op, dtype_id]
                    if values_id:
                        id_parts.append(values_id)
                    if shape_id:
                        id_parts.append(shape_id)
                    case_id = "_".join(id_parts)
                    cases.append(
                        pytest.param(
                            op,
                            init_value,
                            use_mask,
                            torch_dtype,
                            gluon_dtype,
                            use_constant_values,
                            N,
                            M,
                            axis,
                            rhs_shape,
                            id=case_id,
                        ))

    return cases


@pytest.mark.parametrize(
    "op,init_value,use_mask,torch_dtype,gluon_dtype,use_constant_values,N,M,axis,rhs_shape",
    _shared_atomic_scatter_rmw_cases(),
)
def test_shared_atomic_scatter_rmw(op, init_value, use_mask, torch_dtype, gluon_dtype, use_constant_values, N, M, axis,
                                   rhs_shape):
    device = torch.device("cuda")
    rhs_n, rhs_m = rhs_shape

    if use_constant_values:
        values = torch.ones(rhs_shape, dtype=torch_dtype, device=device)
    else:
        values = torch.arange(rhs_n * rhs_m, dtype=torch.int32, device=device).reshape(rhs_shape)
        values = (values % 7 + 1).to(torch_dtype)
    if op == "xchg":
        # make xchg deterministic with broadcasting
        if axis == 1:
            base_indices = torch.arange(0, rhs_m, dtype=torch.int32, device=device)[None, :]
        else:
            base_indices = torch.arange(0, rhs_n, dtype=torch.int32, device=device)[:, None]
        indices = base_indices.expand(rhs_n, rhs_m).contiguous()
    else:
        torch.manual_seed(23 if use_mask else 17)
        upper = M if axis == 1 else N
        indices = torch.randint(0, upper, rhs_shape, dtype=torch.int32, device=device)

    if use_mask:
        mask = (torch.arange(rhs_n * rhs_m, device=device).reshape(rhs_n, rhs_m) % 3) != 0
    else:
        mask = torch.ones((rhs_n, rhs_m), dtype=torch.bool, device=device)

    old = torch.empty((rhs_n, rhs_m), dtype=torch_dtype, device=device)
    final = torch.empty((N, M), dtype=torch_dtype, device=device)
    expected = _expected_shared_atomic_scatter_rmw(op, init_value, indices, values, mask, axis, (N, M))

    layout_2d = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[THREADS_PER_WARP // 4, 4],
                                   warps_per_cta=[4, 1], order=[1, 0])
    layout_rhs = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[THREADS_PER_WARP // 4, 4],
                                    warps_per_cta=[4, 1], order=[1, 0])
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    shared_atomic_scatter_rmw_kernel[(1, )](
        values,
        indices,
        mask,
        old,
        final,
        N=N,
        M=M,
        axis=axis,
        op=op,
        dtype=gluon_dtype,
        use_mask=use_mask,
        use_constant_values=use_constant_values,
        init_value=init_value,
        RHS_N=rhs_n,
        RHS_M=rhs_m,
        layout_2d=layout_2d,
        layout_rhs=layout_rhs,
        shared_layout=shared_layout,
        num_warps=4,
    )

    torch.testing.assert_close(final, expected, atol=0, rtol=0)


@gluon.jit
def shared_atomic_scatter_rmw_broadcast_kernel(
    values_ptr,
    old_ptr,
    final_ptr,
    N: ttgl.constexpr,
    M: ttgl.constexpr,
    layout_2d: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
):
    indices_x = ttgl.arange(0, N, layout=ttgl.SliceLayout(dim=1, parent=layout_2d))
    indices_y = ttgl.arange(0, M, layout=ttgl.SliceLayout(dim=0, parent=layout_2d))
    offsets_2d = indices_x[:, None] * M + indices_y[None, :]

    smem = ttgl.allocate_shared_memory(ttgl.float32, [N, M], layout=shared_layout)
    smem.store(ttgl.zeros([N, M], ttgl.float32, layout=layout_2d))

    # values has shape [N, 1], indices has shape [1, M], and mask has shape
    # [N, 1]. atomic_scatter_add broadcasts all three operands to [N, M].
    values = ttgl.load(values_ptr + indices_x)[:, None]
    indices = indices_y[None, :]
    mask = (indices_x >= 0)[:, None]
    old = smem.atomic_scatter_add(values, indices, axis=1, mask=mask)
    ttgl.store(old_ptr + offsets_2d, old)

    final = smem.load(layout=layout_2d)
    ttgl.store(final_ptr + offsets_2d, final)


def test_shared_atomic_scatter_rmw_broadcast():
    device = torch.device("cuda")
    N, M = 16, 32
    values = torch.arange(N, dtype=torch.float32, device=device) + 1.0
    old = torch.empty((N, M), dtype=torch.float32, device=device)
    final = torch.empty((N, M), dtype=torch.float32, device=device)

    layout_2d = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[THREADS_PER_WARP // 4, 4],
                                   warps_per_cta=[4, 1], order=[1, 0])
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    shared_atomic_scatter_rmw_broadcast_kernel[(1, )](
        values,
        old,
        final,
        N=N,
        M=M,
        layout_2d=layout_2d,
        shared_layout=shared_layout,
        num_warps=4,
    )

    expected = values[:, None].expand(N, M).contiguous()
    torch.testing.assert_close(final, expected, atol=0, rtol=0)
    torch.testing.assert_close(old, torch.zeros_like(old), atol=0, rtol=0)


# ============================================================================
# Multi-warp Tests
# ============================================================================


@pytest.mark.parametrize("N,M,num_warps", [(64, 64, 2), (128, 128, 4)])
def test_scatter_gather_multiwarp(N, M, num_warps):
    """Test scatter and gather with multiple warps."""
    device = torch.device("cuda")

    # Create layouts with multiple warps (shared across both tests)
    layout_2d = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[THREADS_PER_WARP // 4, 4],
                                   warps_per_cta=[num_warps, 1], order=[1, 0])
    layout_1d = ttgl.BlockedLayout(size_per_thread=[1], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[num_warps],
                                   order=[0])
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    # Test gather
    matrix = torch.arange(N * M, dtype=torch.float32, device=device).reshape(N, M)
    gather_indices = torch.arange(N, dtype=torch.int32, device=device) * (M + 1)
    gather_output = torch.zeros(N, dtype=torch.float32, device=device)
    gather_expected = matrix.flatten()[gather_indices]

    shared_gather_kernel[(1, )](
        matrix,
        gather_indices,
        gather_output,
        N=N,
        M=M,
        layout_2d=layout_2d,
        layout_1d=layout_1d,
        shared_layout=shared_layout,
        num_warps=num_warps,
    )

    torch.testing.assert_close(gather_output, gather_expected)

    # Test scatter
    scatter_indices = torch.arange(N, dtype=torch.int32, device=device) * (M + 1)
    scatter_values = torch.arange(N, dtype=torch.float32, device=device) + 100.0
    scatter_output = torch.zeros((N, M), dtype=torch.float32, device=device)
    scatter_expected = torch.zeros((N, M), dtype=torch.float32, device=device)
    for i in range(N):
        scatter_expected[i, i] = scatter_values[i]

    shared_scatter_kernel[(1, )](
        scatter_indices,
        scatter_values,
        scatter_output,
        N=N,
        M=M,
        layout_2d=layout_2d,
        layout_1d=layout_1d,
        shared_layout=shared_layout,
        use_broadcast=False,
        num_warps=num_warps,
    )

    torch.testing.assert_close(scatter_output, scatter_expected)


# ============================================================================
# 2D Native Gather/Scatter Tests
# ============================================================================


@gluon.jit
def gather_2d_kernel(
    matrix_ptr,
    indices_ptr,
    output_ptr,
    N: ttgl.constexpr,
    M: ttgl.constexpr,
    axis: ttgl.constexpr,
    layout_2d: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
):
    """Test 2D gather along specified axis."""
    # Load the matrix from global memory [N, M]
    indices_x = ttgl.arange(0, N, layout=ttgl.SliceLayout(dim=1, parent=layout_2d))
    indices_y = ttgl.arange(0, M, layout=ttgl.SliceLayout(dim=0, parent=layout_2d))
    offsets_2d = indices_x[:, None] * M + indices_y[None, :]
    matrix_data = ttgl.load(matrix_ptr + offsets_2d)

    # Store in shared memory
    smem = ttgl.allocate_shared_memory(ttgl.float32, [N, M], layout=shared_layout)
    smem.store(matrix_data)

    # Load indices [N, M] - same rank as source
    indices = ttgl.load(indices_ptr + offsets_2d)

    # Gather along specified axis
    gathered = smem.gather(indices, axis=axis)

    # Store result
    ttgl.store(output_ptr + offsets_2d, gathered)


@pytest.mark.parametrize("N,M,axis", [(32, 32, 0), (32, 32, 1), (64, 64, 0), (64, 64, 1)])
def test_gather_2d_native(N, M, axis):
    """Test 2D gather along different axes."""
    device = torch.device("cuda")

    # Create a test matrix [N, M]
    matrix = torch.arange(N * M, dtype=torch.float32, device=device).reshape(N, M)

    # Create indices [N, M] - each position specifies where to gather from along the axis
    if axis == 0:
        # Each column gathers from a shifted row pattern
        indices = torch.arange(M, dtype=torch.int32, device=device)[None, :].expand(N, M)
        indices = (indices + torch.arange(N, dtype=torch.int32, device=device)[:, None]) % N
        # Expected: result[i, j] = matrix[indices[i, j], j]
        expected = torch.gather(matrix, 0, indices.long())
    else:  # axis == 1
        # Each row gathers from a shifted column pattern
        indices = torch.arange(N, dtype=torch.int32, device=device)[:, None].expand(N, M)
        indices = (indices + torch.arange(M, dtype=torch.int32, device=device)[None, :]) % M
        # Expected: result[i, j] = matrix[i, indices[i, j]]
        expected = torch.gather(matrix, 1, indices.long())

    output = torch.zeros((N, M), dtype=torch.float32, device=device)

    # Create layouts dynamically based on THREADS_PER_WARP
    layout_2d = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[THREADS_PER_WARP // 4, 4],
                                   warps_per_cta=[1, 1], order=[1, 0])
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    gather_2d_kernel[(1, )](
        matrix,
        indices,
        output,
        N=N,
        M=M,
        axis=axis,
        layout_2d=layout_2d,
        shared_layout=shared_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


@gluon.jit
def scatter_2d_kernel(
    indices_ptr,
    values_ptr,
    output_ptr,
    N: ttgl.constexpr,
    M: ttgl.constexpr,
    axis: ttgl.constexpr,
    layout_2d: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
):
    """Test 2D scatter along specified axis."""
    # Initialize shared memory to zero
    smem = ttgl.allocate_shared_memory(ttgl.float32, [N, M], layout=shared_layout)

    indices_x = ttgl.arange(0, N, layout=ttgl.SliceLayout(dim=1, parent=layout_2d))
    indices_y = ttgl.arange(0, M, layout=ttgl.SliceLayout(dim=0, parent=layout_2d))
    offsets_2d = indices_x[:, None] * M + indices_y[None, :]
    zeros = ttgl.zeros([N, M], ttgl.float32, layout=layout_2d)
    smem.store(zeros)

    # Load indices [N, M] and values [N, M]
    indices = ttgl.load(indices_ptr + offsets_2d)
    values = ttgl.load(values_ptr + offsets_2d)

    # Scatter along specified axis
    smem.scatter(values, indices, axis=axis)

    # Read back the result
    result = smem.load(layout=layout_2d)
    ttgl.store(output_ptr + offsets_2d, result)


@pytest.mark.parametrize("N,M,axis", [(32, 32, 0), (32, 32, 1)])
def test_scatter_2d_native(N, M, axis):
    """Test 2D scatter along different axes."""
    device = torch.device("cuda")

    # Create indices [N, M] - reverse pattern for scatter
    if axis == 0:
        indices = torch.arange(M, dtype=torch.int32, device=device)[None, :].expand(N, M)
        indices = (N - 1 - indices - torch.arange(N, dtype=torch.int32, device=device)[:, None]) % N
    else:  # axis == 1
        indices = torch.arange(N, dtype=torch.int32, device=device)[:, None].expand(N, M)
        indices = (M - 1 - indices - torch.arange(M, dtype=torch.int32, device=device)[None, :]) % M

    # Create values to scatter
    values = torch.arange(N * M, dtype=torch.float32, device=device).reshape(N, M) + 100.0

    output = torch.zeros((N, M), dtype=torch.float32, device=device)

    # Expected: scatter values according to indices
    expected = torch.zeros((N, M), dtype=torch.float32, device=device)
    expected.scatter_(axis, indices.long(), values)

    # Create layouts dynamically based on THREADS_PER_WARP
    layout_2d = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[THREADS_PER_WARP // 4, 4],
                                   warps_per_cta=[1, 1], order=[1, 0])
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    scatter_2d_kernel[(1, )](
        indices,
        values,
        output,
        N=N,
        M=M,
        axis=axis,
        layout_2d=layout_2d,
        shared_layout=shared_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


# ============================================================================
# 3D Gather/Scatter Tests
# ============================================================================


@gluon.jit
def gather_3d_kernel(
    tensor_ptr,
    indices_ptr,
    output_ptr,
    N: ttgl.constexpr,
    M: ttgl.constexpr,
    P: ttgl.constexpr,
    axis: ttgl.constexpr,
    layout_3d: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
):
    """Test 3D gather along specified axis."""
    # Load the tensor from global memory [N, M, P]
    idx_n = ttgl.arange(0, N)[:, None, None]
    idx_m = ttgl.arange(0, M)[None, :, None]
    idx_p = ttgl.arange(0, P)[None, None, :]

    offsets_3d = idx_n * (M * P) + idx_m * P + idx_p
    offsets_3d = ttgl.set_auto_layout(offsets_3d, layout_3d)

    tensor_data = ttgl.load(tensor_ptr + offsets_3d)

    # Store in shared memory
    smem = ttgl.allocate_shared_memory(ttgl.float32, [N, M, P], layout=shared_layout)
    smem.store(tensor_data)

    # Load indices [N, M, P] - same rank as source
    indices_data = ttgl.load(indices_ptr + offsets_3d)

    # Gather along specified axis
    gathered = smem.gather(indices_data, axis=axis)

    # Store result
    ttgl.store(output_ptr + offsets_3d, gathered)


@pytest.mark.parametrize("N,M,P,axis", [(16, 8, 4, 0), (16, 8, 4, 1), (16, 8, 4, 2)])
def test_gather_3d_native(N, M, P, axis):
    """Test 3D gather along different axes."""
    device = torch.device("cuda")

    # Create a test tensor [N, M, P]
    tensor = torch.arange(N * M * P, dtype=torch.float32, device=device).reshape(N, M, P)

    # Create indices [N, M, P] - each position specifies where to gather from along the axis
    if axis == 0:
        # Pattern for gathering along first dimension
        base = torch.arange(M * P, dtype=torch.int32, device=device).reshape(1, M, P)
        offset = torch.arange(N, dtype=torch.int32, device=device).reshape(N, 1, 1)
        indices = (base + offset) % N
    elif axis == 1:
        # Pattern for gathering along second dimension
        base = torch.arange(N, dtype=torch.int32, device=device).reshape(N, 1, 1)
        offset = torch.arange(P, dtype=torch.int32, device=device).reshape(1, 1, P)
        indices = ((base + offset) % M).expand(N, M, P).contiguous()
    else:  # axis == 2
        # Pattern for gathering along third dimension
        base = torch.arange(N * M, dtype=torch.int32, device=device).reshape(N, M, 1)
        indices = (base % P).expand(N, M, P).contiguous()

    # Ensure indices is contiguous in C-style layout
    indices = indices.contiguous()

    # Compute expected result using torch.gather
    expected = torch.gather(tensor, axis, indices.long())

    output = torch.zeros((N, M, P), dtype=torch.float32, device=device)

    # Create layouts dynamically based on THREADS_PER_WARP
    layout_3d = ttgl.BlockedLayout(size_per_thread=[1, 1, 1], threads_per_warp=[4, 4, THREADS_PER_WARP // 16],
                                   warps_per_cta=[1, 1, 1], order=[2, 1, 0])
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[2, 1, 0])

    gather_3d_kernel[(1, )](
        tensor,
        indices,
        output,
        N=N,
        M=M,
        P=P,
        axis=axis,
        layout_3d=layout_3d,
        shared_layout=shared_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


@gluon.jit
def scatter_3d_kernel(
    indices_ptr,
    values_ptr,
    output_ptr,
    N: ttgl.constexpr,
    M: ttgl.constexpr,
    P: ttgl.constexpr,
    axis: ttgl.constexpr,
    layout_3d: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
):
    """Test 3D scatter along specified axis."""
    idx_n = ttgl.arange(0, N)[:, None, None]
    idx_m = ttgl.arange(0, M)[None, :, None]
    idx_p = ttgl.arange(0, P)[None, None, :]

    offsets_3d = idx_n * (M * P) + idx_m * P + idx_p
    offsets_3d = ttgl.set_auto_layout(offsets_3d, layout_3d)

    # Initialize shared memory to zero
    smem = ttgl.allocate_shared_memory(ttgl.float32, [N, M, P], layout=shared_layout)
    zeros = ttgl.full([N, M, P], 0.0, ttgl.float32, layout=layout_3d)
    smem.store(zeros)

    # Load indices [N, M, P] and values [N, M, P]
    indices_data = ttgl.load(indices_ptr + offsets_3d)
    values_data = ttgl.load(values_ptr + offsets_3d)

    # Scatter along specified axis
    smem.scatter(values_data, indices_data, axis=axis)

    # Read back the result
    result = smem.load(layout=layout_3d)
    ttgl.store(output_ptr + offsets_3d, result)


@pytest.mark.parametrize("N,M,P,axis", [(16, 8, 4, 0), (16, 8, 4, 1), (16, 8, 4, 2)])
def test_scatter_3d_native(N, M, P, axis):
    """Test 3D scatter along different axes."""
    device = torch.device("cuda")

    # Create indices [N, M, P] that form a permutation along the scatter axis
    if axis == 0:
        # For axis 0: permute N dimension, keeping (M, P) coordinates fixed
        # Each (j, k) position has a unique permutation of N indices
        base = torch.arange(M * P, dtype=torch.int32, device=device).reshape(1, M, P)
        offset = torch.arange(N, dtype=torch.int32, device=device).reshape(N, 1, 1)
        indices = ((N - 1 - base - offset) % N).contiguous()
    elif axis == 1:
        # For axis 1: permute M dimension, keeping (N, P) coordinates fixed
        # Each (i, k) position has a unique permutation of M indices
        base = torch.arange(N * P, dtype=torch.int32, device=device).reshape(N, 1, P)
        offset = torch.arange(M, dtype=torch.int32, device=device).reshape(1, M, 1)
        indices = ((M - 1 - base - offset) % M).contiguous()
    else:  # axis == 2
        # For axis 2: permute P dimension, keeping (N, M) coordinates fixed
        # Each (i, j) position has a unique permutation of P indices
        base = torch.arange(N * M, dtype=torch.int32, device=device).reshape(N, M, 1)
        offset = torch.arange(P, dtype=torch.int32, device=device).reshape(1, 1, P)
        indices = ((P - 1 - base - offset) % P).contiguous()

    # Ensure indices is contiguous
    indices = indices.contiguous()

    # Create values to scatter
    values = (torch.arange(N * M * P, dtype=torch.float32, device=device).reshape(N, M, P) + 200.0).contiguous()

    output = torch.zeros((N, M, P), dtype=torch.float32, device=device)

    # Expected: scatter values according to indices
    expected = torch.zeros((N, M, P), dtype=torch.float32, device=device)
    expected.scatter_(axis, indices.long(), values)

    # Create layouts dynamically based on THREADS_PER_WARP
    layout_3d = ttgl.BlockedLayout(size_per_thread=[1, 1, 1], threads_per_warp=[4, 4, THREADS_PER_WARP // 16],
                                   warps_per_cta=[1, 1, 1], order=[2, 1, 0])
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[2, 1, 0])

    scatter_3d_kernel[(1, )](
        indices,
        values,
        output,
        N=N,
        M=M,
        P=P,
        axis=axis,
        layout_3d=layout_3d,
        shared_layout=shared_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


# =============================================================================
# Subslice Tests (2D slicing along individual dimensions)
# =============================================================================


@gluon.jit
def gather_subslice_2d_kernel(
    matrix_ptr,
    indices_ptr,
    output_ptr,
    M: ttgl.constexpr,
    N: ttgl.constexpr,
    SLICE_M_OFFSET: ttgl.constexpr,
    SLICE_N_OFFSET: ttgl.constexpr,
    SLICE_M: ttgl.constexpr,
    SLICE_N: ttgl.constexpr,
    layout_full: ttgl.constexpr,
    layout_slice: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
):
    """Gather from a 2D subsliced shared memory descriptor."""
    # Load full matrix into shared memory
    offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, layout_full))[:, None]
    offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, layout_full))[None, :]
    in_offs = offs_m * N + offs_n
    in_data = ttgl.load(matrix_ptr + in_offs)

    smem = ttgl.allocate_shared_memory(ttgl.float32, [M, N], layout=shared_layout)
    smem.store(in_data)

    # Create 2D subslice
    smem_slice = smem.slice(SLICE_M_OFFSET, SLICE_M, dim=0).slice(SLICE_N_OFFSET, SLICE_N, dim=1)

    # Load indices for gathering within the slice
    slice_offs_m = ttgl.arange(0, SLICE_M, layout=ttgl.SliceLayout(1, layout_slice))[:, None]
    slice_offs_n = ttgl.arange(0, SLICE_N, layout=ttgl.SliceLayout(0, layout_slice))[None, :]
    idx_offs = slice_offs_m * SLICE_N + slice_offs_n
    indices = ttgl.load(indices_ptr + idx_offs)

    # Gather along axis 0: result[i, j] = smem_slice[indices[i, j], j]
    gathered = smem_slice.gather(indices, axis=0)

    # Store result
    ttgl.store(output_ptr + idx_offs, gathered)


@pytest.mark.parametrize("M,N,slice_m_offset,slice_n_offset,slice_m,slice_n", [
    # Offset must be a multiple of tile (slice) size for each dimension
    (64, 64, 48, 16, 16, 16),  # offset 48 % 16 == 0, offset 16 % 16 == 0
    (64, 64, 32, 48, 32, 16),  # offset 32 % 32 == 0, offset 48 % 16 == 0
    (64, 64, 48, 32, 16, 32),  # offset 48 % 16 == 0, offset 32 % 32 == 0
])
def test_gather_subslice_2d(M, N, slice_m_offset, slice_n_offset, slice_m, slice_n):
    """Test gathering from a 2D subsliced shared memory descriptor."""
    device = torch.device("cuda")

    # Create input matrix
    matrix = torch.arange(M * N, dtype=torch.float32, device=device).reshape(M, N)

    # Create indices for gather (within the slice dimensions)
    # Each position gathers from a shifted row
    indices = torch.arange(slice_n, dtype=torch.int32, device=device)[None, :].expand(slice_m, slice_n)
    indices = (indices + torch.arange(slice_m, dtype=torch.int32, device=device)[:, None]) % slice_m

    output = torch.zeros((slice_m, slice_n), dtype=torch.float32, device=device)

    # Expected: gather from the subslice
    subslice = matrix[slice_m_offset:slice_m_offset + slice_m, slice_n_offset:slice_n_offset + slice_n]
    expected = torch.gather(subslice, 0, indices.long())

    # Layouts
    layout_full = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    layout_slice = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    # Use non-swizzled layout for subslicing
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    gather_subslice_2d_kernel[(1, )](
        matrix,
        indices,
        output,
        M=M,
        N=N,
        SLICE_M_OFFSET=slice_m_offset,
        SLICE_N_OFFSET=slice_n_offset,
        SLICE_M=slice_m,
        SLICE_N=slice_n,
        layout_full=layout_full,
        layout_slice=layout_slice,
        shared_layout=shared_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


@gluon.jit
def scatter_subslice_2d_kernel(
    indices_ptr,
    values_ptr,
    output_ptr,
    M: ttgl.constexpr,
    N: ttgl.constexpr,
    SLICE_M_OFFSET: ttgl.constexpr,
    SLICE_N_OFFSET: ttgl.constexpr,
    SLICE_M: ttgl.constexpr,
    SLICE_N: ttgl.constexpr,
    layout_full: ttgl.constexpr,
    layout_slice: ttgl.constexpr,
    shared_layout: ttgl.constexpr,
):
    """Scatter to a 2D subsliced shared memory descriptor."""
    # Initialize shared memory with -1
    offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, layout_full))[:, None]
    offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, layout_full))[None, :]
    full_offs = offs_m * N + offs_n
    init_data = ttgl.full([M, N], -1.0, dtype=ttgl.float32, layout=layout_full)

    smem = ttgl.allocate_shared_memory(ttgl.float32, [M, N], layout=shared_layout)
    smem.store(init_data)

    # Create 2D subslice
    smem_slice = smem.slice(SLICE_M_OFFSET, SLICE_M, dim=0).slice(SLICE_N_OFFSET, SLICE_N, dim=1)

    # Load indices and values for scattering within the slice
    slice_offs_m = ttgl.arange(0, SLICE_M, layout=ttgl.SliceLayout(1, layout_slice))[:, None]
    slice_offs_n = ttgl.arange(0, SLICE_N, layout=ttgl.SliceLayout(0, layout_slice))[None, :]
    idx_offs = slice_offs_m * SLICE_N + slice_offs_n
    indices = ttgl.load(indices_ptr + idx_offs)
    values = ttgl.load(values_ptr + idx_offs)

    # Scatter along axis 0: smem_slice[indices[i, j], j] = values[i, j]
    smem_slice.scatter(values, indices, axis=0)

    # Load back full matrix
    result = smem.load(layout=layout_full)
    ttgl.store(output_ptr + full_offs, result)


@pytest.mark.parametrize("M,N,slice_m_offset,slice_n_offset,slice_m,slice_n", [
    # Offset must be a multiple of tile (slice) size for each dimension
    (64, 64, 48, 16, 16, 16),  # offset 48 % 16 == 0, offset 16 % 16 == 0
    (64, 64, 32, 48, 32, 16),  # offset 32 % 32 == 0, offset 48 % 16 == 0
])
def test_scatter_subslice_2d(M, N, slice_m_offset, slice_n_offset, slice_m, slice_n):
    """Test scattering to a 2D subsliced shared memory descriptor."""
    device = torch.device("cuda")

    # Create indices (reverse pattern for scatter)
    indices = torch.arange(slice_n, dtype=torch.int32, device=device)[None, :].expand(slice_m, slice_n)
    indices = (slice_m - 1 - indices - torch.arange(slice_m, dtype=torch.int32, device=device)[:, None]) % slice_m

    # Create values to scatter
    values = torch.arange(slice_m * slice_n, dtype=torch.float32, device=device).reshape(slice_m, slice_n) + 100.0

    output = torch.zeros((M, N), dtype=torch.float32, device=device)

    # Expected: -1 everywhere, then scatter into the subslice region
    expected = torch.full((M, N), -1.0, dtype=torch.float32, device=device)
    subslice_expected = torch.zeros((slice_m, slice_n), dtype=torch.float32, device=device)
    subslice_expected.scatter_(0, indices.long(), values)
    expected[slice_m_offset:slice_m_offset + slice_m, slice_n_offset:slice_n_offset + slice_n] = subslice_expected

    # Layouts
    layout_full = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    layout_slice = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    shared_layout = ttgl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0])

    scatter_subslice_2d_kernel[(1, )](
        indices,
        values,
        output,
        M=M,
        N=N,
        SLICE_M_OFFSET=slice_m_offset,
        SLICE_N_OFFSET=slice_n_offset,
        SLICE_M=slice_m,
        SLICE_N=slice_n,
        layout_full=layout_full,
        layout_slice=layout_slice,
        shared_layout=shared_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


# =============================================================================
# Padded Layout Tests
# =============================================================================


@gluon.jit
def gather_padded_kernel(
    matrix_ptr,
    indices_ptr,
    output_ptr,
    M: ttgl.constexpr,
    N: ttgl.constexpr,
    layout_2d: ttgl.constexpr,
    padded_layout: ttgl.constexpr,
):
    """Gather from shared memory with a padded layout."""
    # Load matrix into padded shared memory
    offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, layout_2d))[:, None]
    offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, layout_2d))[None, :]
    in_offs = offs_m * N + offs_n
    in_data = ttgl.load(matrix_ptr + in_offs)

    smem = ttgl.allocate_shared_memory(ttgl.float32, [M, N], layout=padded_layout)
    smem.store(in_data)

    # Load indices
    indices = ttgl.load(indices_ptr + in_offs)

    # Gather along axis 0
    gathered = smem.gather(indices, axis=0)

    ttgl.store(output_ptr + in_offs, gathered)


@pytest.mark.parametrize("M,N", [(64, 64)])
@pytest.mark.parametrize("interval_pairs", [[[32, 4]], [[16, 4]], [[16, 4], [64, 8]]])
@pytest.mark.parametrize("order", [[0, 1], [1, 0]])
def test_gather_padded(M, N, interval_pairs, order):
    """Test gathering from shared memory with a padded layout."""
    device = torch.device("cuda")

    # Create input matrix
    matrix = torch.arange(M * N, dtype=torch.float32, device=device).reshape(M, N)

    # Create indices for gather along axis 0
    indices = torch.arange(N, dtype=torch.int32, device=device)[None, :].expand(M, N)
    indices = (indices + torch.arange(M, dtype=torch.int32, device=device)[:, None]) % M

    output = torch.zeros((M, N), dtype=torch.float32, device=device)

    # Expected: gather along axis 0
    expected = torch.gather(matrix, 0, indices.long())

    # Layouts
    layout_2d = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    padded_layout = ttgl.PaddedSharedLayout.with_identity_for(interval_pairs, [M, N], order)

    gather_padded_kernel[(1, )](
        matrix,
        indices,
        output,
        M=M,
        N=N,
        layout_2d=layout_2d,
        padded_layout=padded_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


@gluon.jit
def scatter_padded_kernel(
    indices_ptr,
    values_ptr,
    output_ptr,
    M: ttgl.constexpr,
    N: ttgl.constexpr,
    layout_2d: ttgl.constexpr,
    padded_layout: ttgl.constexpr,
):
    """Scatter to shared memory with a padded layout."""
    # Initialize padded shared memory with zeros
    offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, layout_2d))[:, None]
    offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, layout_2d))[None, :]
    full_offs = offs_m * N + offs_n
    zeros = ttgl.zeros([M, N], ttgl.float32, layout=layout_2d)

    smem = ttgl.allocate_shared_memory(ttgl.float32, [M, N], layout=padded_layout)
    smem.store(zeros)

    # Load indices and values
    indices = ttgl.load(indices_ptr + full_offs)
    values = ttgl.load(values_ptr + full_offs)

    # Scatter along axis 0
    smem.scatter(values, indices, axis=0)

    # Load back
    result = smem.load(layout=layout_2d)
    ttgl.store(output_ptr + full_offs, result)


@pytest.mark.parametrize("M,N", [(64, 64)])
@pytest.mark.parametrize("interval_pairs", [[[32, 4]], [[16, 4]]])
@pytest.mark.parametrize("order", [[0, 1], [1, 0]])
def test_scatter_padded(M, N, interval_pairs, order):
    """Test scattering to shared memory with a padded layout."""
    device = torch.device("cuda")

    # Create indices (reverse pattern)
    indices = torch.arange(N, dtype=torch.int32, device=device)[None, :].expand(M, N)
    indices = (M - 1 - indices - torch.arange(M, dtype=torch.int32, device=device)[:, None]) % M

    # Create values
    values = torch.arange(M * N, dtype=torch.float32, device=device).reshape(M, N) + 100.0

    output = torch.zeros((M, N), dtype=torch.float32, device=device)

    # Expected: scatter along axis 0
    expected = torch.zeros((M, N), dtype=torch.float32, device=device)
    expected.scatter_(0, indices.long(), values)

    # Layouts
    layout_2d = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    padded_layout = ttgl.PaddedSharedLayout.with_identity_for(interval_pairs, [M, N], order)

    scatter_padded_kernel[(1, )](
        indices,
        values,
        output,
        M=M,
        N=N,
        layout_2d=layout_2d,
        padded_layout=padded_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


# =============================================================================
# Padded Layout with Subslice Tests
# =============================================================================


@gluon.jit
def gather_padded_subslice_kernel(
    matrix_ptr,
    indices_ptr,
    output_ptr,
    M: ttgl.constexpr,
    N: ttgl.constexpr,
    SLICE_M_OFFSET: ttgl.constexpr,
    SLICE_N_OFFSET: ttgl.constexpr,
    SLICE_M: ttgl.constexpr,
    SLICE_N: ttgl.constexpr,
    layout_full: ttgl.constexpr,
    layout_slice: ttgl.constexpr,
    padded_layout: ttgl.constexpr,
):
    """Gather from a subsliced padded shared memory descriptor."""
    # Load full matrix into padded shared memory
    offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, layout_full))[:, None]
    offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, layout_full))[None, :]
    in_offs = offs_m * N + offs_n
    in_data = ttgl.load(matrix_ptr + in_offs)

    smem = ttgl.allocate_shared_memory(ttgl.float32, [M, N], layout=padded_layout)
    smem.store(in_data)

    # Create 2D subslice
    smem_slice = smem.slice(SLICE_M_OFFSET, SLICE_M, dim=0).slice(SLICE_N_OFFSET, SLICE_N, dim=1)

    # Load indices for gathering within the slice
    slice_offs_m = ttgl.arange(0, SLICE_M, layout=ttgl.SliceLayout(1, layout_slice))[:, None]
    slice_offs_n = ttgl.arange(0, SLICE_N, layout=ttgl.SliceLayout(0, layout_slice))[None, :]
    idx_offs = slice_offs_m * SLICE_N + slice_offs_n
    indices = ttgl.load(indices_ptr + idx_offs)

    # Gather along axis 0
    gathered = smem_slice.gather(indices, axis=0)

    ttgl.store(output_ptr + idx_offs, gathered)


@pytest.mark.parametrize("interval_pairs", [[[32, 4]], [[16, 4]]])
@pytest.mark.parametrize("order", [[0, 1], [1, 0]])
@pytest.mark.parametrize("slice_m_offset,slice_n_offset,slice_m,slice_n", [
    (48, 16, 16, 16),
    (32, 48, 32, 16),
    (48, 32, 16, 32),
])
def test_gather_padded_subslice(interval_pairs, order, slice_m_offset, slice_n_offset, slice_m, slice_n):
    """Test gathering from a subsliced padded shared memory descriptor."""
    M, N = 64, 64
    device = torch.device("cuda")

    # Create input matrix
    matrix = torch.arange(M * N, dtype=torch.float32, device=device).reshape(M, N)

    # Create indices for gather within the slice
    indices = torch.arange(slice_n, dtype=torch.int32, device=device)[None, :].expand(slice_m, slice_n)
    indices = (indices + torch.arange(slice_m, dtype=torch.int32, device=device)[:, None]) % slice_m

    output = torch.zeros((slice_m, slice_n), dtype=torch.float32, device=device)

    # Expected: gather from the subslice
    subslice = matrix[slice_m_offset:slice_m_offset + slice_m, slice_n_offset:slice_n_offset + slice_n]
    expected = torch.gather(subslice, 0, indices.long())

    # Layouts
    layout_full = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    layout_slice = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    padded_layout = ttgl.PaddedSharedLayout.with_identity_for(interval_pairs, [M, N], order)

    gather_padded_subslice_kernel[(1, )](
        matrix,
        indices,
        output,
        M=M,
        N=N,
        SLICE_M_OFFSET=slice_m_offset,
        SLICE_N_OFFSET=slice_n_offset,
        SLICE_M=slice_m,
        SLICE_N=slice_n,
        layout_full=layout_full,
        layout_slice=layout_slice,
        padded_layout=padded_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


@gluon.jit
def scatter_padded_subslice_kernel(
    indices_ptr,
    values_ptr,
    output_ptr,
    M: ttgl.constexpr,
    N: ttgl.constexpr,
    SLICE_M_OFFSET: ttgl.constexpr,
    SLICE_N_OFFSET: ttgl.constexpr,
    SLICE_M: ttgl.constexpr,
    SLICE_N: ttgl.constexpr,
    layout_full: ttgl.constexpr,
    layout_slice: ttgl.constexpr,
    padded_layout: ttgl.constexpr,
):
    """Scatter to a subsliced padded shared memory descriptor."""
    # Initialize padded shared memory with -1
    offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, layout_full))[:, None]
    offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, layout_full))[None, :]
    full_offs = offs_m * N + offs_n
    init_data = ttgl.full([M, N], -1.0, dtype=ttgl.float32, layout=layout_full)

    smem = ttgl.allocate_shared_memory(ttgl.float32, [M, N], layout=padded_layout)
    smem.store(init_data)

    # Create 2D subslice
    smem_slice = smem.slice(SLICE_M_OFFSET, SLICE_M, dim=0).slice(SLICE_N_OFFSET, SLICE_N, dim=1)

    # Load indices and values for scattering within the slice
    slice_offs_m = ttgl.arange(0, SLICE_M, layout=ttgl.SliceLayout(1, layout_slice))[:, None]
    slice_offs_n = ttgl.arange(0, SLICE_N, layout=ttgl.SliceLayout(0, layout_slice))[None, :]
    idx_offs = slice_offs_m * SLICE_N + slice_offs_n
    indices = ttgl.load(indices_ptr + idx_offs)
    values = ttgl.load(values_ptr + idx_offs)

    # Scatter along axis 0
    smem_slice.scatter(values, indices, axis=0)

    # Load back full matrix
    result = smem.load(layout=layout_full)
    ttgl.store(output_ptr + full_offs, result)


@pytest.mark.parametrize("interval_pairs", [[[32, 4]], [[16, 4]]])
@pytest.mark.parametrize("order", [[0, 1], [1, 0]])
@pytest.mark.parametrize("slice_m_offset,slice_n_offset,slice_m,slice_n", [
    (48, 16, 16, 16),
    (32, 48, 32, 16),
])
def test_scatter_padded_subslice(interval_pairs, order, slice_m_offset, slice_n_offset, slice_m, slice_n):
    """Test scattering to a subsliced padded shared memory descriptor."""
    M, N = 64, 64
    device = torch.device("cuda")

    # Create indices (reverse pattern)
    indices = torch.arange(slice_n, dtype=torch.int32, device=device)[None, :].expand(slice_m, slice_n)
    indices = (slice_m - 1 - indices - torch.arange(slice_m, dtype=torch.int32, device=device)[:, None]) % slice_m

    # Create values
    values = torch.arange(slice_m * slice_n, dtype=torch.float32, device=device).reshape(slice_m, slice_n) + 100.0

    output = torch.zeros((M, N), dtype=torch.float32, device=device)

    # Expected: -1 everywhere, then scatter into the subslice region
    expected = torch.full((M, N), -1.0, dtype=torch.float32, device=device)
    subslice_expected = torch.zeros((slice_m, slice_n), dtype=torch.float32, device=device)
    subslice_expected.scatter_(0, indices.long(), values)
    expected[slice_m_offset:slice_m_offset + slice_m, slice_n_offset:slice_n_offset + slice_n] = subslice_expected

    # Layouts
    layout_full = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    layout_slice = ttgl.BlockedLayout(
        size_per_thread=[1, 1],
        threads_per_warp=[THREADS_PER_WARP // 4, 4],
        warps_per_cta=[1, 1],
        order=[1, 0],
    )
    padded_layout = ttgl.PaddedSharedLayout.with_identity_for(interval_pairs, [M, N], order)

    scatter_padded_subslice_kernel[(1, )](
        indices,
        values,
        output,
        M=M,
        N=N,
        SLICE_M_OFFSET=slice_m_offset,
        SLICE_N_OFFSET=slice_n_offset,
        SLICE_M=slice_m,
        SLICE_N=slice_n,
        layout_full=layout_full,
        layout_slice=layout_slice,
        padded_layout=padded_layout,
        num_warps=1,
    )

    torch.testing.assert_close(output, expected)


# --- TMEM Load with Reduction Tests ---


@gluon.jit
def tmem_reduction_kernel(
    in_ptr,
    out_ptr,
    red_ptr,
    M: ttgl.constexpr,
    N: ttgl.constexpr,
    RED_OP: ttgl.constexpr,
    USE_ABS: ttgl.constexpr,
    PROPAGATE_NAN: ttgl.constexpr,
    num_warps: ttgl.constexpr,
):
    """Kernel to test TMEM load with hardware reduction."""
    global_memory_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], [1, num_warps], [1, 0])
    global_memory_layout_1d: ttgl.constexpr = ttgl.BlockedLayout([1], [32], [num_warps], [0])

    # Offsets for 2D tensor
    offs_m = ttgl.arange(0, M, ttgl.SliceLayout(1, global_memory_layout))
    offs_n = ttgl.arange(0, N, ttgl.SliceLayout(0, global_memory_layout))
    offs_2d = offs_m[:, None] * N + offs_n[None, :]

    # Load input from global memory
    input_data = ttgl.load(in_ptr + offs_2d)

    # Setup TMEM layout - blockN must match N for single reduction value per row
    tmem_layout: ttgl.constexpr = TensorMemoryLayout(block=(128, N), col_stride=1,  # packed for f32
                                                     )

    # Allocate TMEM
    tmem = allocate_tensor_memory(
        element_ty=in_ptr.dtype.element_ty,
        shape=[M, N],
        layout=tmem_layout,
    )

    # Get register layout for TMEM access
    tmem_reg_layout: ttgl.constexpr = tmem.get_reg_layout()

    # Store input to TMEM
    input_data = ttgl.convert_layout(input_data, tmem_reg_layout)
    tmem.store(input_data)

    # Load from TMEM with reduction
    if RED_OP == "min":
        output, reduced = tmem.load_min(abs=USE_ABS, propagate_nan=PROPAGATE_NAN)
    elif RED_OP == "max":
        output, reduced = tmem.load_max(abs=USE_ABS, propagate_nan=PROPAGATE_NAN)

    # Store full output
    output = ttgl.convert_layout(output, global_memory_layout)
    ttgl.store(out_ptr + offs_2d, output)

    # Store reduced output (1D tensor of shape [M])
    offs_1d = ttgl.arange(0, M, global_memory_layout_1d)
    reduced = ttgl.convert_layout(reduced, global_memory_layout_1d)
    ttgl.store(red_ptr + offs_1d, reduced)


@pytest.mark.skipif(not (is_blackwell_ultra() or is_rubin()), reason="Requires Blackwell Ultra or Rubin")
@pytest.mark.parametrize("red_op", ["min", "max"])
@pytest.mark.parametrize("use_abs", [False, True])
@pytest.mark.parametrize("propagate_nan", [tl.PropagateNan.NONE, tl.PropagateNan.ALL])
@pytest.mark.parametrize(
    "M, N, num_warps",
    [(128, 32, 4), (128, 64, 4), (128, 128, 4), (128, 256, 4), (256, 128, 8)],
)
def test_tmem_reduction(red_op, use_abs, propagate_nan, M, N, num_warps):
    """Test TMEM load with hardware reduction on MxN tile

    Note: With M=128, only 4 warps can be used (warpsPerCTA=[4,1]) since all
    warps must fit in the M dimension for reduction. 8 warps would require
    M=256 (8*32=256). The N=256 case tests partial reduction combining where
    4 hardware reductions are combined via llvm.minnum/maxnum.
    """

    # Create test input with some negative values
    input_tensor = torch.randn(M, N, dtype=torch.float32, device="cuda")

    # Inject NaN for testing if needed
    use_nan = False if propagate_nan == tl.PropagateNan.NONE else True
    if use_nan:
        input_tensor[10, 5] = float("nan")
        input_tensor[50, 15] = float("nan")

    # Output tensors
    output = torch.empty_like(input_tensor)
    red_output = torch.empty(M, dtype=torch.float32, device="cuda")

    # Run kernel
    tmem_reduction_kernel[(1, )](
        input_tensor,
        output,
        red_output,
        M,
        N,
        red_op,
        use_abs,
        propagate_nan,
        num_warps=num_warps,
    )

    # Verify full output matches input (tmem store/load roundtrip)
    # Use equal_nan=True when we have NaN values in the input
    torch.testing.assert_close(input_tensor, output, atol=0, rtol=0, equal_nan=use_nan)

    # Compute expected reduction
    ref_input = torch.abs(input_tensor) if use_abs else input_tensor
    torch_red = torch.min if red_op == "min" else torch.max
    expected_red = torch_red(ref_input, dim=1).values

    # Verify reduction output
    # Use equal_nan=True when testing NaN propagation
    torch.testing.assert_close(expected_red, red_output, atol=1e-5, rtol=1e-5, equal_nan=use_nan)


@pytest.mark.parametrize("num_ctas", [1, 2])
@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_clc_basic(num_ctas):
    # Launch a grid with 2x the number of CTAs as the number of SMs
    # And ask to allocate a big chunk of smem per block (almost all
    # the smem minus 32 i64 elements to make room for the barriers),
    # so that we force 1 block per SM

    @gluon.jit
    def clc_kernel(WasLaunched, IsCancelled, ProgramId, smem_size: ttgl.constexpr):
        # Allocate clc_mbar before clc_result to make sure that we are indeed aligning
        # clc_result correctly after a i64 element.
        clc_mbar = mbarrier.allocate_mbarrier()
        cga_layout: ttgl.constexpr = [[0]] * (ttgl.num_ctas().bit_length() - 1)
        layout: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, order=[0], cga_layout=cga_layout)
        clc_result = ttgl.allocate_shared_memory(ttgl.int64, [2], layout)
        mbarrier.init(clc_mbar, count=1)

        # Large shared memory allocation to force 1 block per SM
        dummy = ttgl.allocate_shared_memory(ttgl.int64, [smem_size // 8 - 32, 1], clc_mbar.layout)

        clc.try_cancel(clc_result, clc_mbar)
        mbarrier.expect(clc_mbar, 16, from_cta=0x0)
        mbarrier.wait(clc_mbar, 0)

        response = clc.load_result(clc_result)
        pid = ttgl.program_id(0)
        ttgl.store(WasLaunched + pid, True)
        ttgl.store(IsCancelled + pid, response.is_canceled())
        ttgl.store(ProgramId + pid, response.program_id(0))
        dummy._keep_alive()

    dev_props = torch.cuda.get_device_properties("cuda")
    num_sms = dev_props.multi_processor_count
    device = triton.runtime.driver.active.get_current_device()
    smem_size = max_shared_mem(device)
    grid = 2 * (num_sms // num_ctas)

    was_launched = torch.zeros([grid], dtype=torch.bool, device="cuda")
    is_cancelled = torch.zeros([grid], dtype=torch.bool, device="cuda")
    program_ids = torch.zeros([grid], dtype=torch.int32, device="cuda")
    clc_kernel[(grid, )](was_launched, is_cancelled, program_ids, smem_size, num_ctas=num_ctas)

    num_launched = torch.sum(was_launched).item()
    assert num_launched < grid

    num_cancelled = torch.sum(is_cancelled).item()
    assert num_launched + num_cancelled == grid

    for pid in range(grid):
        if is_cancelled[pid]:
            assert not was_launched[program_ids[pid]]


def align_to(a, b):
    return triton.cdiv(a, b) * b


def make_operand_descriptor(value, BLOCK_MN, BLOCK_K, MIXED_PREC, cga_layout=None):
    IS_FP4 = value.dtype == torch.uint8
    ELEM_PER_BYTE = 2 if IS_FP4 else 1
    IS_MIXED_PREC_FP4 = MIXED_PREC and IS_FP4
    layout = ttgl.NVMMASharedLayout.get_default_for(
        [BLOCK_MN, BLOCK_K // ELEM_PER_BYTE],
        ttgl.uint8 if IS_FP4 else ttgl.float8e4nv,
        fp4_padded=IS_MIXED_PREC_FP4,
        cga_layout=cga_layout,
    )
    return TensorDescriptor.from_tensor(value, [BLOCK_MN, BLOCK_K // ELEM_PER_BYTE], layout)


def make_output_descriptor(M, N, dtype, BLOCK_M, BLOCK_N, cga_layout=None):
    C = torch.empty(M, N, device="cuda", dtype=dtype)
    C_dtype = getattr(ttgl, str(dtype).split('.')[1])
    C_desc_layout = ttgl.NVMMASharedLayout.get_default_for([BLOCK_M, BLOCK_N], C_dtype, cga_layout=cga_layout)
    return TensorDescriptor.from_tensor(C, [BLOCK_M, BLOCK_N], C_desc_layout)


def random_quantized_tensor(MN, K, format):
    assert format in ["mxfp4", "mxfp8", "nvfp4"]
    VEC_SIZE = 16 if format == "nvfp4" else 32
    base = MXFP4Tensor(size=(MN, K), device="cuda").random()
    scale = MXScaleTensor(size=(MN, K // VEC_SIZE), device="cuda").random(low=1 / 128, high=2.0)
    ref = base.to(torch.float32)
    scale_ref = scale.to(torch.float32)
    value = ref * scale_ref.repeat_interleave(VEC_SIZE, dim=1)
    if format == "mxfp8":
        return ref.to(torch.float8_e4m3fn), scale.data, value
    elif format == "mxfp4":
        return base.to_packed_tensor(dim=1), scale.data, value
    else:
        return base.to_packed_tensor(dim=1), scale_ref.to(torch.float8_e4m3fn), value


def swizzle_scales_packed_block(scales, VEC_SIZE):
    PAD_MN = align_to(scales.shape[0], 128) - scales.shape[0]
    PAD_K = align_to(scales.shape[1], 4) - scales.shape[1]
    scales = torch.nn.functional.pad(scales, (0, PAD_K, 0, PAD_MN))
    MN, SCALE_K = scales.shape[0], scales.shape[1]
    REP_MN = MN // 128
    REP_K = SCALE_K // 4
    scales = scales.reshape(REP_MN, 4, 32, REP_K, 4)
    scales = scales.permute(0, 3, 2, 1, 4)
    return scales.contiguous()


def make_scales_descriptor(scales, BLOCK_MN, BLOCK_K, VEC_SIZE, cga_layout=None):
    REP_MN = BLOCK_MN // 128
    REP_K = BLOCK_K // (VEC_SIZE * 4)
    block_shape = [1, REP_MN, REP_K, 2, 256]
    scales = scales.reshape(1, scales.shape[0], scales.shape[1], 2, 256)
    IS_NVFP4 = scales.dtype == torch.float8_e4m3fn
    layout = ttgl.NVMMASharedLayout.get_default_for(block_shape, ttgl.float8e4nv if IS_NVFP4 else ttgl.uint8,
                                                    cga_layout=cga_layout)
    return TensorDescriptor.from_tensor(scales, block_shape, layout)


def get_cga_layout(layout, op_idx):
    """Based on get_cga_layout from 03-matmul-multicta.py.
    Difference: B is N×K (transposed), so offsets are in dim 0 not dim 1."""
    assert op_idx in (0, 1)
    if not layout:
        return layout

    # 2CTA performs an outer product so bases are [1, 0] and [0, 1]
    assert list(layout[0]) == [1, 0]
    first = (1, 0)

    # Broadcast along K (the reduction dimension)
    # We multiply by 2 for op_idx == 1, as we have added the (1, 0) basis.
    def broadcast(b):
        return (b[0], 0) if op_idx == 0 else (2 * b[1], 0)

    return (first, *map(broadcast, layout[1:]))


@gluon.jit
def unswizzle_scales_shared_memory(smem, BLOCK_MN: ttgl.constexpr, BLOCK_K: ttgl.constexpr, VEC_SIZE: ttgl.constexpr):
    smem = smem.reshape((smem.shape[1], smem.shape[2], 32, 4, 4))
    smem = smem.permute((0, 3, 2, 1, 4))
    return smem.reshape((BLOCK_MN, BLOCK_K // VEC_SIZE))


@gluon.jit
def mma_scaled_tcgen05_copy_kernel(a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc, VEC_SIZE: ttgl.constexpr,
                                   block_layout_c: ttgl.constexpr, a_scale_layout_tmem: ttgl.constexpr,
                                   b_scale_layout_tmem: ttgl.constexpr, ctas_per_cga: ttgl.constexpr,
                                   multicast: ttgl.constexpr):
    A_IS_FP4: ttgl.constexpr = a_desc.dtype == ttgl.uint8
    B_IS_FP4: ttgl.constexpr = b_desc.dtype == ttgl.uint8
    A_ELEM_PER_BYTE: ttgl.constexpr = 2 if A_IS_FP4 else 1
    B_ELEM_PER_BYTE: ttgl.constexpr = 2 if B_IS_FP4 else 1
    BLOCK_M: ttgl.constexpr = c_desc.block_type.shape[0]
    BLOCK_N: ttgl.constexpr = c_desc.block_type.shape[1]
    BLOCK_K: ttgl.constexpr = a_desc.block_type.shape[1] * A_ELEM_PER_BYTE
    K = a_desc.shape[1] * A_ELEM_PER_BYTE

    a_smem = ttgl.allocate_shared_memory(a_desc.dtype, a_desc.block_type.shape, a_desc.layout)
    b_smem = ttgl.allocate_shared_memory(b_desc.dtype, b_desc.block_type.shape, b_desc.layout)

    num_ctas: ttgl.constexpr = ttgl.num_ctas()
    multi_cta: ttgl.constexpr = num_ctas > 1
    scale_layout_a: ttgl.constexpr = TensorMemoryScalesLayout(cga_layout=a_scale_layout_tmem)
    scale_layout_b: ttgl.constexpr = TensorMemoryScalesLayout(cga_layout=b_scale_layout_tmem)
    a_scale_tmem = allocate_tensor_memory(a_scale_desc.dtype, [BLOCK_M, BLOCK_K // VEC_SIZE], scale_layout_a)
    b_scale_tmem = allocate_tensor_memory(b_scale_desc.dtype, [BLOCK_N, BLOCK_K // VEC_SIZE], scale_layout_b)
    tmem_layout: ttgl.constexpr = TensorMemoryLayout(
        [BLOCK_M // ctas_per_cga[0], BLOCK_N // ctas_per_cga[1]],
        col_stride=1,
        cga_layout=c_desc.layout.cga_layout,
        two_ctas=multi_cta,
    )
    acc_tmem = allocate_tensor_memory(ttgl.float32, [BLOCK_M, BLOCK_N], tmem_layout)

    tma_bar = mbarrier.allocate_mbarrier(two_ctas=multi_cta)
    mbarrier.init(tma_bar, count=1)

    mma_bar = mbarrier.allocate_mbarrier()
    mma_bar_count: ttgl.constexpr = tcgen05_mma_barrier_count([a_smem, b_smem], multicast=multicast,
                                                              two_ctas=acc_tmem.type.layout.two_ctas)
    mbarrier.init(mma_bar, count=mma_bar_count)

    phase_tma = 0
    phase_mma = 0
    pid_m = ttgl.program_id(0)
    pid_n = ttgl.program_id(1)
    off_m = pid_m * BLOCK_M
    off_n = pid_n * BLOCK_N

    a_scale_smem = ttgl.allocate_shared_memory(a_scale_desc.dtype, a_scale_desc.block_type.shape, a_scale_desc.layout)
    b_scale_smem = ttgl.allocate_shared_memory(b_scale_desc.dtype, b_scale_desc.block_type.shape, b_scale_desc.layout)
    REP_M: ttgl.constexpr = a_scale_desc.block_type.shape[1]
    REP_N: ttgl.constexpr = b_scale_desc.block_type.shape[1]
    A_REP_K: ttgl.constexpr = a_scale_desc.block_type.shape[2]
    B_REP_K: ttgl.constexpr = b_scale_desc.block_type.shape[2]
    off_m_a_scale = pid_m * REP_M
    off_n_b_scale = pid_n * REP_N

    for k in range(0, K, BLOCK_K):
        off_k_a = k // A_ELEM_PER_BYTE
        off_k_b = k // B_ELEM_PER_BYTE
        off_k_a_scale = (k // BLOCK_K) * A_REP_K
        off_k_b_scale = (k // BLOCK_K) * B_REP_K

        EXPECTED_BYTES: ttgl.constexpr = (a_desc.nbytes_per_cta + b_desc.nbytes_per_cta + a_scale_desc.nbytes_per_cta +
                                          b_scale_desc.nbytes_per_cta)
        mbarrier.expect(tma_bar, EXPECTED_BYTES)
        tma.async_load(a_desc, [off_m, off_k_a], tma_bar, a_smem, multicast=multicast)
        tma.async_load(b_desc, [off_n, off_k_b], tma_bar, b_smem, multicast=multicast)
        tma.async_load(a_scale_desc, [0, off_m_a_scale, off_k_a_scale, 0, 0], tma_bar, a_scale_smem,
                       multicast=multicast)
        tma.async_load(b_scale_desc, [0, off_n_b_scale, off_k_b_scale, 0, 0], tma_bar, b_scale_smem,
                       multicast=multicast)
        mbarrier.wait(tma_bar, phase_tma, deps=[a_smem, b_smem, a_scale_smem, b_scale_smem])
        phase_tma ^= 1

        a_scale = unswizzle_scales_shared_memory(a_scale_smem, BLOCK_M, BLOCK_K, VEC_SIZE)
        b_scale = unswizzle_scales_shared_memory(b_scale_smem, BLOCK_N, BLOCK_K, VEC_SIZE)
        tcgen05_copy(a_scale, a_scale_tmem)
        tcgen05_copy(b_scale, b_scale_tmem)

        a_format: ttgl.constexpr = "e2m1" if A_IS_FP4 else "e4m3"
        b_format: ttgl.constexpr = "e2m1" if B_IS_FP4 else "e4m3"
        tcgen05_mma_scaled(a_smem, b_smem.permute((1, 0)), acc_tmem, a_scale_tmem, b_scale_tmem, a_format, b_format,
                           use_acc=(k != 0), multicast=multicast, mbarriers=[mma_bar])
        mbarrier.wait(mma_bar, phase_mma)
        phase_mma ^= 1

    mbarrier.invalidate(tma_bar)
    mbarrier.invalidate(mma_bar)
    acc = acc_tmem.load()
    if multi_cta:
        acc = ttgl.convert_layout(acc, block_layout_c)
    acc = acc.to(c_desc.dtype)
    acc_smem = ttgl.allocate_shared_memory(c_desc.dtype, c_desc.block_type.shape, c_desc.layout)
    acc_smem.store(acc)
    tma.async_store(c_desc, [off_m, off_n], acc_smem)
    tma.store_wait(0)


def mma_scaled_tcgen05_copy(A, B, A_scale, B_scale, VEC_SIZE, BLOCK_M, BLOCK_N, BLOCK_K, ctas_per_cga, multicast,
                            out_dtype=torch.float16):
    from dataclasses import replace
    M, N = A.shape[0], B.shape[0]
    MIXED_PREC = A.dtype != B.dtype
    num_ctas = ctas_per_cga[0] * ctas_per_cga[1]
    multi_cta = num_ctas > 1
    warps = [4, 1]
    num_warps = warps[0] * warps[1]

    cta_order = (1, 0) if multi_cta else None
    cga_layout_c = make_2cta_cga_layout(ctas_per_cga, ctas_per_cga, cta_order, 0) if multi_cta else None
    cga_layout_a = get_cga_layout(cga_layout_c, 0)
    cga_layout_b = get_cga_layout(cga_layout_c, 1)

    # Scale CGA layouts derived from C layout:
    # A scales project M (c[0]), B scales project N (c[1])
    # CGA layouts are 5D (matching scale SMEM shape [1, REP_MN, REP_K, 2, 256])
    # TMEM layouts are 2D (matching scale TMEM shape [BLOCK_MN, BLOCK_K // VEC_SIZE])
    cga_layout_a_scale = [[0, c[0], 0, 0, 0] for c in cga_layout_c] if multi_cta else None
    cga_layout_b_scale = [[0, c[1], 0, 0, 0] for c in cga_layout_c] if multi_cta else None
    a_scale_layout_tmem = tuple((c[0], 0) for c in cga_layout_c) if multi_cta else ()
    b_scale_layout_tmem = tuple((c[1], 0) for c in cga_layout_c) if multi_cta else ()

    A_desc = make_operand_descriptor(A, BLOCK_M, BLOCK_K, MIXED_PREC, cga_layout=cga_layout_a)
    B_desc = make_operand_descriptor(B, BLOCK_N, BLOCK_K, MIXED_PREC, cga_layout=cga_layout_b)
    C_desc = make_output_descriptor(M, N, out_dtype, BLOCK_M, BLOCK_N, cga_layout=cga_layout_c)
    A_scale_desc = make_scales_descriptor(A_scale, BLOCK_M, BLOCK_K, VEC_SIZE, cga_layout=cga_layout_a_scale)
    B_scale_desc = make_scales_descriptor(B_scale, BLOCK_N, BLOCK_K, VEC_SIZE, cga_layout=cga_layout_b_scale)

    a_scale_layout = ttgl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5,
                                            cga_layout=cga_layout_a_scale)
    b_scale_layout = ttgl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5,
                                            cga_layout=cga_layout_b_scale)
    A_scale_desc = replace(A_scale_desc, layout=a_scale_layout)
    B_scale_desc = replace(B_scale_desc, layout=b_scale_layout)

    block_layout_c = ttgl.BlockedLayout([1, 8], [1, 32], warps_per_cta=warps, order=[1, 0],
                                        cga_layout=cga_layout_c) if multi_cta else None

    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
    mma_scaled_tcgen05_copy_kernel[grid](A_desc, B_desc, C_desc, A_scale_desc, B_scale_desc, VEC_SIZE, block_layout_c,
                                         a_scale_layout_tmem, b_scale_layout_tmem, ctas_per_cga, num_warps=num_warps,
                                         num_ctas=num_ctas, multicast=multicast)
    return C_desc.base


@pytest.mark.parametrize("M, N, K", [(2048, 2048, 4096)])
@pytest.mark.parametrize("BLOCK_K", [128, 256])
@pytest.mark.parametrize("a_format, b_format, VEC_SIZE", [
    pytest.param("mxfp8", "mxfp8", 32, id="mxfp8_mxfp8_block32"),
    pytest.param("nvfp4", "nvfp4", 16, id="nvfp4_nvfp4_block16"),
    pytest.param("mxfp4", "mxfp8", 32, id="mxfp4_mxfp8_block32"),
    pytest.param("mxfp8", "mxfp4", 32, id="mxfp8_mxfp4_block32"),
])
@pytest.mark.parametrize("ctas_per_cga", [(1, 1), (2, 1), (4, 1), (2, 2), (4, 2), (4, 4)])
@pytest.mark.parametrize("multicast", [True, False])
@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_mma_scaled_tcgen05_copy(M, N, K, BLOCK_K, a_format, b_format, VEC_SIZE, ctas_per_cga, multicast):
    BLOCK_M = 128 * ctas_per_cga[0]
    BLOCK_N = 128 * ctas_per_cga[1]
    torch.manual_seed(0)
    A, A_scale, A_ref = random_quantized_tensor(M, K, a_format)
    B, B_scale, B_ref = random_quantized_tensor(N, K, b_format)
    A_scale = swizzle_scales_packed_block(A_scale, VEC_SIZE)
    B_scale = swizzle_scales_packed_block(B_scale, VEC_SIZE)
    C_ref = A_ref @ B_ref.T
    C = mma_scaled_tcgen05_copy(A, B, A_scale, B_scale, VEC_SIZE, BLOCK_M, BLOCK_N, BLOCK_K, ctas_per_cga, multicast)
    torch.testing.assert_close(C_ref, C.to(torch.float32), atol=1e-3, rtol=1e-3)


@gluon.jit
def tmem288k_simple_kernel(in_ptr, out_ptr, M: ttgl.constexpr, Ncol1: ttgl.constexpr, Ncol2: ttgl.constexpr,
                           num_warps: ttgl.constexpr):
    global_memory_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], [1, num_warps], [1, 0])

    offs_m = ttgl.arange(0, M, ttgl.SliceLayout(1, global_memory_layout))
    offs_n128 = ttgl.arange(0, 128, ttgl.SliceLayout(0, global_memory_layout))
    offs_n32 = ttgl.arange(0, 32, ttgl.SliceLayout(0, global_memory_layout))

    N: ttgl.constexpr = Ncol1 + Ncol2

    # Allocate some tensor memory.
    tmem_layout: ttgl.constexpr = TensorMemoryLayout(
        block=(128, 32),
        col_stride=32 // in_ptr.dtype.element_ty.primitive_bitwidth,
    )

    tmem1 = allocate_tensor_memory(
        element_ty=in_ptr.dtype.element_ty,
        shape=[M, Ncol1],
        layout=tmem_layout,
    )
    if Ncol2 > 0:
        tmem2 = allocate_tensor_memory(
            element_ty=in_ptr.dtype.element_ty,
            shape=[M, Ncol2],
            layout=tmem_layout,
        )

    # Get the register layout needed to access tensor memory from descriptors.
    tmem_reg_layout1: ttgl.constexpr = tmem1.slice(0, 128).get_reg_layout(num_warps=num_warps)
    if Ncol2 > 0:
        tmem_reg_layout2: ttgl.constexpr = tmem2.slice(0, 32).get_reg_layout(num_warps=num_warps)

    assert Ncol1 == 128 or Ncol1 == 256 or Ncol1 == 512

    n1a: ttgl.constexpr = 0
    offs = offs_m[:, None] * N + n1a + offs_n128[None, :]
    input = ttgl.load(in_ptr + offs)
    input = ttgl.convert_layout(input, tmem_reg_layout1)
    tmem_slice = tmem1.slice(n1a, 128)
    tmem_slice.store(input)

    if Ncol1 > 128:
        n1b: ttgl.constexpr = 128
        offs = offs_m[:, None] * N + n1b + offs_n128[None, :]
        input = ttgl.load(in_ptr + offs)
        input = ttgl.convert_layout(input, tmem_reg_layout1)
        tmem_slice = tmem1.slice(n1b, 128)
        tmem_slice.store(input)

    if Ncol1 > 256:
        n1c: ttgl.constexpr = 256
        offs = offs_m[:, None] * N + n1c + offs_n128[None, :]
        input = ttgl.load(in_ptr + offs)
        input = ttgl.convert_layout(input, tmem_reg_layout1)
        tmem_slice = tmem1.slice(n1c, 128)
        tmem_slice.store(input)

        n1d: ttgl.constexpr = 384
        offs = offs_m[:, None] * N + n1d + offs_n128[None, :]
        input = ttgl.load(in_ptr + offs)
        input = ttgl.convert_layout(input, tmem_reg_layout1)
        tmem_slice = tmem1.slice(n1d, 128)
        tmem_slice.store(input)

    if Ncol2 > 0:
        n2a: ttgl.constexpr = 0
        offs = offs_m[:, None] * N + Ncol1 + n2a + offs_n32[None, :]
        input = ttgl.load(in_ptr + offs)
        input = ttgl.convert_layout(input, tmem_reg_layout2)
        tmem_slice = tmem2.slice(n2a, 32)
        tmem_slice.store(input)

    if Ncol2 > 32:
        n2b: ttgl.constexpr = 32
        offs = offs_m[:, None] * N + Ncol1 + n2b + offs_n32[None, :]
        input = ttgl.load(in_ptr + offs)
        input = ttgl.convert_layout(input, tmem_reg_layout2)
        tmem_slice = tmem2.slice(n2b, 32)
        tmem_slice.store(input)

    offs = offs_m[:, None] * N + n1a + offs_n128[None, :]
    tmem_slice = tmem1.slice(n1a, 128)
    output = tmem_slice.load(tmem_reg_layout1)
    output = ttgl.convert_layout(output, global_memory_layout)
    ttgl.store(out_ptr + offs, output)

    if Ncol1 > 128:
        offs = offs_m[:, None] * N + n1b + offs_n128[None, :]
        tmem_slice = tmem1.slice(n1b, 128)
        output = tmem_slice.load(tmem_reg_layout1)
        output = ttgl.convert_layout(output, global_memory_layout)
        ttgl.store(out_ptr + offs, output)

    if Ncol1 > 256:
        offs = offs_m[:, None] * N + n1c + offs_n128[None, :]
        tmem_slice = tmem1.slice(n1c, 128)
        output = tmem_slice.load(tmem_reg_layout1)
        output = ttgl.convert_layout(output, global_memory_layout)
        ttgl.store(out_ptr + offs, output)

        offs = offs_m[:, None] * N + n1d + offs_n128[None, :]
        tmem_slice = tmem1.slice(n1d, 128)
        output = tmem_slice.load(tmem_reg_layout1)
        output = ttgl.convert_layout(output, global_memory_layout)
        ttgl.store(out_ptr + offs, output)

    if Ncol2 > 0:
        offs = offs_m[:, None] * N + Ncol1 + n2a + offs_n32[None, :]
        tmem_slice = tmem2.slice(n2a, 32)
        output = tmem_slice.load(tmem_reg_layout2)
        output = ttgl.convert_layout(output, global_memory_layout)
        ttgl.store(out_ptr + offs, output)

    if Ncol2 > 32:
        offs = offs_m[:, None] * N + Ncol1 + n2b + offs_n32[None, :]
        tmem_slice = tmem2.slice(n2b, 32)
        output = tmem_slice.load(tmem_reg_layout2)
        output = ttgl.convert_layout(output, global_memory_layout)
        ttgl.store(out_ptr + offs, output)


@pytest.mark.skipif(not is_rubin(), reason="Requires Rubin")
@pytest.mark.parametrize("Ncol1", [256, 512])
@pytest.mark.parametrize("Ncol2", [32, 64])
def test_tmem288k_simple(Ncol1: int, Ncol2: int):
    M = 128
    N = Ncol1 + Ncol2
    input = torch.randn(M, N, dtype=torch.float32).to(device="cuda")
    output = torch.empty_like(input)

    num_warps = 4
    tmem288k_simple_kernel[(1, )](input, output, M, Ncol1, Ncol2, num_warps=num_warps)
    torch.testing.assert_close(input.cpu(), output.cpu(), atol=0, rtol=0)
