# Copyright 2024 The JAX Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import contextlib
import dataclasses
import enum
import itertools
import math
import operator
import os
import re
import sys
import tempfile

from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax._src import config
from jax._src import dtypes
from jax._src import test_util as jtu
from jax._src import hypothesis_test_util as htu
from jax._src.interpreters import mlir
from jax._src.lib import version as jaxlib_version
from jax._src.lib.mlir import ir
from jax._src.lib.mlir import passmanager
from jax._src.lib.mlir.dialects import arith
from jax._src.lib.mlir.dialects import cf
from jax._src.lib.mlir.dialects import gpu
from jax._src.lib.mlir.dialects import llvm
from jax._src.lib.mlir.dialects import scf
from jax._src.lib.mlir.dialects import vector
import jax.experimental.mosaic.gpu as mgpu
from jax.experimental.mosaic.gpu import core
from jax.experimental.mosaic.gpu import dialect as mgpu_dialect
from jax.experimental.mosaic.gpu import fragmented_array as fa
from jax.experimental.mosaic.gpu import launch_context
from jax.experimental.mosaic.gpu import layouts
from jax.experimental.mosaic.gpu import profiler
from jax.experimental.mosaic.gpu import tcgen05
from jax.experimental.mosaic.gpu import test_util as mtu
from jax.experimental.mosaic.gpu import utils
from jax.experimental.mosaic.gpu.utils import *  # noqa: F403
import jax.numpy as jnp
import numpy as np

try:
  import jax._src.lib.mosaic_gpu as mosaic_gpu_lib  # noqa: F401
  HAS_MOSAIC_GPU = True
except ImportError:
  mosaic_gpu_lib = None
  HAS_MOSAIC_GPU = False

  class Dimension(enum.IntEnum):  # Just to make parameterized tests expand ok
    x = 0
    y = 1
    z = 2
else:
  Dimension = gpu.Dimension

try:
  import hypothesis as hp
  import hypothesis.strategies as hps
  htu.setup_hypothesis()
except ImportError:
  hp = hps = None


# ruff: noqa: F405
config.parse_flags_with_absl()

def nd_loop(bounds, body, *, _idxs = ()):
  if not bounds:
    body(*_idxs)
    return
  bound, *other_bounds = bounds
  @fori(bound, ())
  def _loop_body(i, _):
    nd_loop(other_bounds, body, _idxs=(*_idxs, i))
    return ()


def copy(src: ir.Value, dst: ir.Value, swizzle: int | None = None):
  index = ir.IndexType.get()
  thread_id = gpu.thread_id(gpu.Dimension.x)
  stride = gpu.block_dim(gpu.Dimension.x)
  for dim in (gpu.Dimension.y, gpu.Dimension.z):
    thread_id = arith.addi(thread_id, arith.muli(gpu.thread_id(dim), stride))
    stride = arith.muli(stride, gpu.block_dim(dim))
  is_first_thread = arith.cmpi(arith.CmpIPredicate.eq, thread_id, c(0, index))

  src_ty = ir.MemRefType(src.type)
  dst_ty = ir.MemRefType(dst.type)
  if src_ty.shape != dst_ty.shape:
    raise ValueError(
        f"src and dst shapes don't match: {src_ty.shape} != {dst_ty.shape}"
    )
  shape = src_ty.shape
  if src_ty.element_type != dst_ty.element_type:
    raise ValueError(
        f"src and dst element types don't match: {src_ty.element_type} !="
        f" {dst_ty.element_type}"
    )
  contig_strides = get_contiguous_strides(shape)
  # If swizzling is on, at least one of the memrefs must be contiguous
  # (simulating a TMA).
  if (swizzle is not None and
      src_ty.get_strides_and_offset()[0] != contig_strides and
      dst_ty.get_strides_and_offset()[0] != contig_strides):
    raise NotImplementedError(src_ty, dst_ty)

  bw = bitwidth(src_ty.element_type)
  if bw < 8:
    assert bw.bit_count() == 1
    packing = 8 // bw
    if shape[-1] % packing:
      raise NotImplementedError
    shape = (*shape[:-1], shape[-1] // packing)
    contig_strides = get_contiguous_strides(shape)
    def bitcast(ref):
      ref_ty = ir.MemRefType(ref.type)
      old_strides = ref_ty.get_strides_and_offset()[0]
      if old_strides[-1] != 1:
        raise NotImplementedError
      new_strides = [s // packing for s in old_strides[:-1]] + [1]
      new_ref_ty = ir.MemRefType.get(
          shape,
          ir.VectorType.get((packing,), src_ty.element_type),  # noqa: F821
          ir.StridedLayoutAttr.get(0, new_strides),
          ref_ty.memory_space,
      )
      # NOTE: memref_ptr applies the offset in case there was any.
      return ptr_as_memref(memref_ptr(ref), new_ref_ty)
    src = bitcast(src)
    dst = bitcast(dst)
    bw = 8
  del src_ty, dst_ty  # If you remove this, update it in the branch above
  dyn_strides = [c(s, index) for s in contig_strides]

  with ir.InsertionPoint(scf.IfOp(is_first_thread).then_block):
    def body(*idx):
      dst_idx = idx
      if swizzle is not None:
        assert swizzle.bit_count() == 1
        assert bw % 8 == 0
        bytes_per_element = bw // 8
        linear_idx = c(0, index)
        for stride, i in zip(dyn_strides, idx):
          linear_idx = arith.addi(linear_idx, arith.muli(i, stride))
        # Swizzle pattern repeats every 128 bytes.
        swizzle_src = arith.remui(
            arith.divui(linear_idx, c(128 // bytes_per_element, index)),
            c(swizzle // 16, index),
        )
        # Swizzle happens in groups of 16 bytes.
        swizzle_shift = 4 - (bytes_per_element.bit_length() - 1)
        dst_linear_idx = arith.xori(
            linear_idx, arith.shli(swizzle_src, c(swizzle_shift, index))
        )
        dst_idx = [
            arith.remui(arith.divui(dst_linear_idx, stride), c(bound, index))
            for stride, bound in zip(dyn_strides, shape)
        ]
      memref.store(memref.load(src, idx), dst, dst_idx)
    nd_loop([c(d, index) for d in shape], body)
    scf.yield_([])
  gpu.barrier()
  nvvm.fence_proxy(nvvm.ProxyKind.async_)


def iota_tensor(m, n, dtype, layout=mgpu.WGMMA_LAYOUT):
  """A tensor with given layout where arr[i, j] = i * N + j."""
  index = ir.IndexType.get()
  mlir_dtype = utils.dtype_to_ir_type(dtype)
  int_ty = ir.IntegerType.get_signless(bitwidth(mlir_dtype))
  ret = mgpu.FragmentedArray.splat(
      llvm.mlir_undef(int_ty), (m, n), is_signed=False
  )
  ret = ret.to_layout(layout)

  def iota_value(_, idx):
    assert len(idx) == 2
    return arith.index_cast(
        int_ty, arith.addi(idx[1], arith.muli(idx[0], c(n, index)))
    )

  ret = ret.foreach(
      iota_value,
      create_array=True,
      is_signed=False,
  )
  return ret.astype(mlir_dtype, is_signed=utils.is_signed(dtype))


class TestCase(parameterized.TestCase):

  def setUp(self, *, artificial_shared_memory_limit=jtu._SMEM_SIZE_BOUND_FOR_TESTS, min_compute_capability="9.0"):
    if jtu.test_device_matches(["rocm"]):
      self.skipTest("Mosaic GPU is not supported on ROCm.")
    if not HAS_MOSAIC_GPU:
      self.skipTest("jaxlib built without Mosaic GPU")
    if (not jtu.test_device_matches(["cuda"]) or
        not jtu.is_cuda_compute_capability_at_least(min_compute_capability)):
      self.skipTest(f"Only works on GPU with capability >= sm{min_compute_capability.replace('.', '')}")
    super().setUp()
    self.prng = np.random.default_rng(1234)
    self.context = mlir.make_ir_context()
    mgpu_dialect.register_dialect(self.context)
    self.enter_context(config.traceback_filtering("off"))
    self.enter_context(self.context)
    self.enter_context(ir.Location.unknown())
    self.enter_context(core.artificial_shared_memory_limit(artificial_shared_memory_limit))

  @contextlib.contextmanager
  def capture_stdout(self):
    if "pytest" in sys.modules:
      self.skipTest("pytest interacts badly with GPU stdout capture")
    if mosaic_gpu_lib is None:
      raise ValueError("Running tests but missing Mosaic GPU extension")
    with jtu.capture_stdout() as stdout:
      yield stdout
      # We need to cudaDeviceSynchronize to make sure printfs are flushed.
      mosaic_gpu_lib._mosaic_gpu_ext._sync_all_devices()


class Sm90ATestCase(TestCase, jtu.CudaArchSpecificTest):

  def setUp(self):
      self.skip_unless_sm90a()
      # No artificially lowered limit for arch-specific tests
      super().setUp(artificial_shared_memory_limit=None)


class TestUtilTest(TestCase):

  def test_copy_basic(self):
    def kernel(ctx, src, dst, _):
      copy(src, dst)
    x = jnp.arange(2 * 3 * 5).reshape(2, 5, 3)
    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, ())(x)
    np.testing.assert_array_equal(y, x)

  def test_copy_swizzle(self):
    def kernel(ctx, src, dst, _):
      copy(src, dst, swizzle=128)
    x = jnp.arange(8 * 32, dtype=jnp.float32).reshape(8, 32)
    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, ())(x)
    expected = np.zeros_like(y)
    for i in range(8):
      for j in range(8):
        js = j ^ i
        expected[i, (j * 4):(j * 4) + 4] = x[i, (js * 4):(js * 4) + 4]
    np.testing.assert_array_equal(y, expected)

  def test_copy_swizzle_noop(self):
    # Two swizzles cancel out
    def kernel(ctx, src, dst, smem):
      copy(src, smem, swizzle=128)
      copy(smem, dst, swizzle=128)
    x = jnp.arange(8 * 32, dtype=jnp.float32).reshape(8, 32)
    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, x)(x)
    np.testing.assert_array_equal(y, x)

  def test_iota_tensor(self):
    m = n = 64
    def kernel(ctx, dst, _):
      index = ir.IndexType.get()
      registers = iota_tensor(m, n, jnp.float32).registers
      assert registers.size == 16, registers.size
      for i, vec_reg in enumerate(registers.flat):
        for j in range(2):
          reg = vector.extract(
              source=vec_reg,
              dynamic_position=[],
              static_position=ir.DenseI64ArrayAttr.get([j]),
          )
          memref.store(
              reg, dst, [gpu.thread_id(gpu.Dimension.x), c(2 * i + j, index)]
          )
    out_shape = jax.ShapeDtypeStruct((128, 32), jnp.float32)
    regs = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    thread_ids = np.arange(128)
    warp_ids = thread_ids // 32
    lane_ids = thread_ids % 32
    thread_rows = warp_ids * 16 + lane_ids // 4
    thread_start_cols = (lane_ids % 4) * 2
    thread_cols = thread_start_cols[:, None] + (np.arange(n // 8)[None] * 8)
    regs = regs.reshape(128, 8, 2, 2)
    for row_half in range(2):
      for col_half in range(2):
        np.testing.assert_array_equal(
            regs[..., row_half, col_half],
            (thread_rows[:, None] + row_half * 8) * n + thread_cols + col_half
        )


class MemRefTest(TestCase):
  @parameterized.product(
      dim=tuple(range(3)),
      strided=(False, True)
  )
  def test_unsqueeze(self, dim, strided):
    def kernel(ctx, inp, out, _):
      if strided:
        for i in range(8):
          s = ds(i, 1)
          out_slice = s if dim != 0 else (slice(None), s)
          copy(
              memref_unsqueeze(memref_slice(inp, s), dim),
              memref_slice(out, out_slice),
          )
      else:
        copy(memref_unsqueeze(inp, dim), out)
    x = np.arange(8 * 16, dtype=jnp.float32).reshape(8, 16)
    out_shape = list(x.shape)
    out_shape.insert(dim, 1)
    out_ty = jax.ShapeDtypeStruct(out_shape, jnp.float32)
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, out_ty, ()
    )(x)
    np.testing.assert_array_equal(y, x.reshape(out_shape))

  @parameterized.product(
      dim=tuple(range(2)),
      strided=(False, True)
  )
  def test_unfold(self, dim, strided):
    in_shape = (8, 16)
    def kernel(ctx, inp, out, _):
      if strided:
        # We slice the dim we don't unfold
        for i in range(in_shape[1 - dim] // 4):
          s = ds(i * 4, 4)
          in_slice = s if dim == 1 else (slice(None), s)
          out_slice = s if dim == 1 else (slice(None),) * 3 + (s,)
          copy(
              memref_unfold(memref_slice(inp, in_slice), dim, (2, 2, None)),
              memref_slice(out, out_slice),
          )
      else:
        copy(memref_unfold(inp, dim, (2, 2, None)), out)
    x = np.arange(np.prod(in_shape), dtype=jnp.float32).reshape(in_shape)
    out_shape = list(in_shape)
    out_shape[dim:dim + 1] = [2, 2, out_shape[dim] // 4]
    out_ty = jax.ShapeDtypeStruct(out_shape, jnp.float32)
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, out_ty, ()
    )(x)
    np.testing.assert_array_equal(y, x.reshape(out_ty.shape))

  @parameterized.product(
      dim=tuple(range(2)),
  )
  def test_fold_not_strided(self, dim):
    def kernel(ctx, inp, out, _):
      copy(memref_fold(inp, dim, 2), out)

    x = np.arange(8 * 2 * 8, dtype=jnp.float32).reshape(8, 2, 8)
    out_ty = jax.ShapeDtypeStruct((16, 8) if dim == 0 else (8, 16), jnp.float32)
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, out_ty, ()
    )(x)
    np.testing.assert_array_equal(y, x.reshape(out_ty.shape))

  @parameterized.named_parameters(
      ("drop_1s", (1, 1, 5, 1, 1, 2, 1, 1), (5, 1, 2)),
      ("add_1s", (5, 1, 2), (1, 1, 5, 1, 1, 2, 1, 1)),
      ("fold", (1, 5, 2, 1,), (1, 10, 1)),
      ("un", (1, 10, 1), (1, 5, 2, 1,)),
      ("to_scalar", (1, 1, 1), ()),
      ("from_scalar", (), (1, 1, 1)),
      ("arbitrary", (2 * 5, 7 * 3), (2, 7, 5, 3)),
  )
  def test_reshape(self, inp_shape, out_shape):
    def kernel(ctx, inp, out, _):
      copy(memref_reshape(inp, out_shape), out)

    x = np.arange(math.prod(inp_shape), dtype=jnp.float32).reshape(inp_shape)
    out_ty = jax.ShapeDtypeStruct(out_shape, jnp.float32)
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, out_ty, ()
    )(x)
    np.testing.assert_array_equal(y, x.reshape(out_shape))

  @parameterized.named_parameters([
      ("packed", (4, 4, 4), (16, 4, 1), 1, 2, False),
      ("strided_end", (4, 4, 4, 4), (256, 64, 16, 4), 1, 2, False),
      ("strided_bot", (4, 4, 4, 4), (256, 16, 4, 1), 1, 2, False),
      ("strided_top", (4, 4, 4, 4), (256, 64, 4, 1), 1, 2, True),
      ("strided_mid", (4, 4, 4, 4), (265, 64, 16, 1), 1, 3, True),
      # TODO(cperivol): Investigate why this is indexing OOB and uncomment.
      # ("overap", (2, 4, 4), (16, 1, 1), 0, 3, True),
  ])
  def test_fold_strided(
      self, shape, strides, dim, fold_rank, throws
  ):
    expanded_shape = get_packed_shape(strides, shape)
    total_size = np.prod(expanded_shape)
    np_inp = np.arange(total_size, dtype=jnp.float32).reshape(expanded_shape)
    index = tuple(slice(0, s) for s in shape)

    # Reference implementation
    def np_fold(inp, dim, fold_rank):
      out_shape = list(inp.shape)
      out_shape[dim : dim + fold_rank] = [
          int(np.prod(inp.shape[dim : dim + fold_rank]))
      ]
      if throws:
        return jax.ShapeDtypeStruct(shape=out_shape, dtype=inp.dtype)
      else:
        return inp.reshape(*out_shape)

    total_size = np.prod(shape) * np.prod(strides)

    def do_test():
      def kernel(ctx, inp, out, _):
        copy(memref_fold(memref_slice(inp, index), dim, fold_rank), out)

      out = np_fold(np_inp[index], dim, fold_rank)
      y = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), np_inp, out, ()
      )(np_inp)
      assert (
          not throws
      ), "If it should have thrown it would during the call."
      np.testing.assert_array_equal(y, out)

    if throws:
      with self.assertRaises(ValueError):
        do_test()
    else:
      do_test()

  @parameterized.parameters(jnp.uint64, jnp.uint32, jnp.uint16, jnp.uint8)
  def test_scalar_argument(self, dtype):
    if dtype == jnp.uint64 and not config.enable_x64.value:
      self.skipTest(
        "64-bit types are disabled: this leads to the input scalar being"
        " traced as a uint32 value, which causes the top 32 bits of the 64-bit"
        " values read from the 32-bit input buffer to sometimes"
        " (nondeterministically) contain garbage.")

    scalar = dtype(42)
    expected = np.full((128, 128), scalar, dtype=dtype)

    def kernel(ctx, inp, out, _):
      del ctx
      inp = memref.load(inp, [])
      mgpu.FragmentedArray.splat(inp, expected.shape, is_signed=True).store_untiled(out)

    res = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        jax.ShapeDtypeStruct(shape=(), dtype=expected.dtype),
        expected,
        (),
    )(scalar)
    np.testing.assert_array_equal(res, expected)

  @jtu.thread_unsafe_test()  # Modifies `os.environ``.
  def test_cluster_id_simplify(self):
    def kernel(ctx, dst, _):
      idx = mgpu.utils.cluster_idx(gpu.Dimension)
      idx = arith.index_cast(ir.IntegerType.get_signless(32), idx)
      memref.store(idx, dst, [gpu.block_id(gpu.Dimension.y)])

    out_ty = jax.ShapeDtypeStruct(shape=(8,), dtype=jnp.int32)
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      f = mgpu.as_gpu_kernel(
          kernel, (1, 8, 1), (128, 1, 1), (), out_ty, (), cluster=(1, 4, 1)
      )
      np.testing.assert_array_equal(f(), np.arange(8) % 4)
    # Make sure MLIR/LLVM have simplified the calculation of cluster_idx and
    # recognized that we only need the Y component.
    self.assertEqual(ptx().count("%cluster_ctaid"), 1)
    self.assertEqual(ptx().count("%ctaid"), 1)
    self.assertEqual(ptx().count("%cluster_ctaid.y"), 1)
    self.assertEqual(ptx().count("%ctaid.y"), 1)

  def test_uses_pdl_mlir(self):
    def body(ctx, src, dst, _):
      del ctx
      copy(src, dst)

    x = jax.ShapeDtypeStruct((128,), jnp.float32)
    module, _, _, _ = core._lower_as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        cluster=(1, 1, 1),
        block=(128, 1, 1),
        in_shapes=(x,),
        out_shape=x,
        inout_shape=(),
        smem_scratch_shape=(),
        lowering_semantics=mgpu.LoweringSemantics.Lane,
        module_name="test_module",
        kernel_name="test_kernel",
        uses_pdl=True,
    )
    self.assertIn("mosaic_gpu.uses_pdl", str(module))

  @parameterized.parameters(gpu.Dimension.x, gpu.Dimension.y)
  def test_cluster_ref_read(self, dim):
    index = ir.IndexType.get()
    dims = (gpu.Dimension.x, gpu.Dimension.y)
    def kernel(ctx, src, dst, scratch):
      smem, barrier = scratch
      cluster_idx = tuple(gpu.cluster_block_id(dim) for dim in dims)
      peer_idx = arith.subi(arith.constant(index, 1), cluster_idx[dim])
      peer_smem = mgpu.get_cluster_ref(smem, dim, peer_idx)
      mgpu.FragmentedArray.load_strided(memref_slice(src, cluster_idx)).store_untiled(smem)
      utils.warpgroup_barrier()
      barrier.arrive()
      barrier.wait()
      mgpu.FragmentedArray.load_strided(peer_smem).store_untiled(memref_slice(dst, cluster_idx))

    barrier = mgpu.ClusterBarrier(collective_dims=(dim,))
    x = np.arange(2 * 2 * 512, dtype=jnp.float32).reshape(2, 2, 512)
    smem = jax.ShapeDtypeStruct(shape=(x.shape[-1],), dtype=jnp.float32)
    f = mgpu.as_gpu_kernel(
        kernel, (2, 2, 1), (128, 1, 1), x, x, (smem, barrier), cluster=(2, 2, 1)
    )
    np.testing.assert_array_equal(f(x), np.flip(x, axis=int(dim)))

  def test_cluster_ref_write(self):
    index = ir.IndexType.get()
    def kernel(ctx, src, dst, scratch):
      smem, barrier = scratch
      cluster_idx = gpu.cluster_block_id(gpu.Dimension.x)
      peer_idx = arith.subi(arith.constant(index, 1), cluster_idx)
      a = mgpu.FragmentedArray.load_tiled(
          memref_slice(src, cluster_idx),
          layout=mgpu.WGMMA_LAYOUT,
          swizzle=128,
          optimized=False,
      )
      peer_barrier = barrier.remap_to_cluster(gpu.Dimension.x, peer_idx)
      peer_barrier.arrive_expect_tx(math.prod(a.shape) * mgpu.bitwidth(a.mlir_dtype) // 8 // 128)
      a.store_tiled_async(
          smem,
          barrier,
          cluster_dim=gpu.Dimension.x,
          cluster_idx=peer_idx,
          swizzle=128,
      )
      barrier.wait()
      mgpu.FragmentedArray.load_strided(smem).store_untiled(memref_slice(dst, cluster_idx))

    barrier = mgpu.Barrier(arrival_count=128)
    x = np.arange(2 * 64 * 32, dtype=jnp.float32).reshape(2, 8, 1, 8, 32)
    smem = jax.ShapeDtypeStruct(shape=x.shape[1:], dtype=jnp.float32)
    f = mgpu.as_gpu_kernel(
        kernel, (2, 1, 1), (128, 1, 1), x, x, (smem, barrier), cluster=(2, 1, 1)
    )
    np.testing.assert_array_equal(f(x), np.flip(x, axis=0))

  @parameterized.parameters(
      ("add", False),
      ("add", True),
      ("min", False),
      ("max", False),
      ("and", False),
      ("or", False),
      ("xor", False),
  )
  def test_cluster_red_async(self, op, is_signed):
    index = ir.IndexType.get()
    jax_dtype = jnp.int32 if is_signed else jnp.uint32
    def kernel(ctx, src, dst, scratch):
      smem, barrier, cluster_barrier = scratch
      cluster_idx = gpu.cluster_block_id(gpu.Dimension.x)
      peer_idx = arith.subi(arith.constant(index, 1), cluster_idx)
      a = mgpu.FragmentedArray.load_tiled(
          memref_slice(src, cluster_idx),
          layout=mgpu.WGMMA_LAYOUT,
          swizzle=128,
          is_signed=is_signed,
          optimized=False,
      )
      identity = {
          "add": 0, "or": 0, "xor": 0,
          "min": np.iinfo(jax_dtype).max,
          "max": np.iinfo(jax_dtype).min,
          "and": -1,
      }[op]
      mlir_dtype = utils.dtype_to_ir_type(jax_dtype)
      mgpu.FragmentedArray.splat(
          c(identity, mlir_dtype), a.shape,
          is_signed=is_signed, layout=a.layout,
      ).store_tiled(smem, swizzle=128)
      cluster_barrier.arrive()
      cluster_barrier.wait()
      tx_bytes = math.prod(a.shape) * mgpu.bitwidth(a.mlir_dtype) // 8 // 128
      barrier.arrive_expect_tx(tx_bytes)
      a.store_tiled_async(
          smem,
          barrier,
          cluster_dim=gpu.Dimension.x,
          cluster_idx=peer_idx,
          swizzle=128,
          atomic=op,
      )
      barrier.wait()
      cluster_barrier.arrive()
      cluster_barrier.wait()
      mgpu.FragmentedArray.load_strided(smem, is_signed=is_signed).store_untiled(
          memref_slice(dst, cluster_idx)
      )

    cluster_barrier = mgpu.ClusterBarrier(collective_dims=(gpu.Dimension.x,))
    barrier = mgpu.Barrier(arrival_count=128)
    x = np.arange(1, 2 * 64 * 32 + 1, dtype=jax_dtype).reshape(2, 8, 1, 8, 32)
    smem = jax.ShapeDtypeStruct(shape=x.shape[1:], dtype=jax_dtype)
    f = mgpu.as_gpu_kernel(
        kernel, (2, 1, 1), (128, 1, 1), x, x, (smem, barrier, cluster_barrier),
        cluster=(2, 1, 1),
    )
    result = f(x)
    np_op = {
        "add": np.add, "min": np.minimum, "max": np.maximum,
        "and": np.bitwise_and, "or": np.bitwise_or, "xor": np.bitwise_xor,
    }[op]
    id_val = {
        "add": 0, "or": 0, "xor": 0,
        "min": np.iinfo(jax_dtype).max,
        "max": np.iinfo(jax_dtype).min,
        "and": np.iinfo(jax_dtype).max,
    }[op]
    id_arr = np.full_like(x[0], id_val, dtype=jax_dtype)
    np.testing.assert_array_equal(result[0], np_op(id_arr, x[1]))
    np.testing.assert_array_equal(result[1], np_op(id_arr, x[0]))


class CtaBarrierTest(TestCase):

  def test_cta_barrier_basic(self):
    i32 = ir.IntegerType.get_signless(32)

    def kernel(ctx, src, dst, scratch):
      del ctx
      smem, barrier = scratch
      wg_id = mgpu.warpgroup_idx(sync=True)
      is_producer = arith.cmpi(arith.CmpIPredicate.eq, wg_id, c(0, i32))
      is_consumer = arith.cmpi(arith.CmpIPredicate.eq, wg_id, c(1, i32))
      with mgpu.when(is_producer):
        mgpu.FragmentedArray.load_strided(src).store_untiled(smem)
        mgpu.commit_shared()
        barrier.arrive()
      with mgpu.when(is_consumer):
        barrier.arrive_and_wait()
        mgpu.FragmentedArray.load_strided(smem).store_untiled(dst)

    barrier = mgpu.CtaBarrier(arrival_count=256)
    x = np.arange(128, dtype=jnp.float32)
    smem = jax.ShapeDtypeStruct(shape=x.shape, dtype=jnp.float32)
    f = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (256, 1, 1),
        x,
        x,
        (smem, barrier),
    )
    np.testing.assert_array_equal(f(x), x)

  def test_cta_barrier_multiple(self):
    i32 = ir.IntegerType.get_signless(32)

    def kernel(ctx, src, dst, scratch):
      del ctx
      smem, barrier = scratch
      wg_id = mgpu.warpgroup_idx(sync=True)
      with mgpu.when(arith.cmpi(arith.CmpIPredicate.eq, wg_id, c(0, i32))):
        mgpu.FragmentedArray.load_strided(src).store_untiled(smem)
        mgpu.commit_shared()
        barrier[1].arrive()
      with mgpu.when(arith.cmpi(arith.CmpIPredicate.eq, wg_id, c(1, i32))):
        barrier[1].arrive_and_wait()
        mgpu.FragmentedArray.load_strided(smem).store_untiled(dst)

    barrier = mgpu.CtaBarrier(arrival_count=256, num_barriers=3)
    x = np.arange(128, dtype=jnp.float32)
    smem = jax.ShapeDtypeStruct(shape=x.shape, dtype=jnp.float32)
    f = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (256, 1, 1), x, x, (smem, barrier)
    )
    np.testing.assert_array_equal(f(x), x)

  def test_cta_barrier_invalid_indexing_raises(self):
    x = jax.ShapeDtypeStruct((128,), jnp.float32)
    barrier = mgpu.CtaBarrier(arrival_count=128, num_barriers=2)

    def unindexed_kernel(ctx, dst, barrier):
      del ctx, dst
      barrier.arrive_and_wait()

    with self.assertRaisesRegex(
        ValueError, "Operation is only valid for a single barrier"
    ):
      mgpu.as_gpu_kernel(
          unindexed_kernel, (1, 1, 1), (128, 1, 1), (), x, barrier
      )

    def oob_kernel(ctx, dst, barrier):
      del ctx, dst
      barrier[2].arrive_and_wait()

    with self.assertRaisesRegex(IndexError, "Barrier offset 2 is out of bounds"):
      mgpu.as_gpu_kernel(oob_kernel, (1, 1, 1), (128, 1, 1), (), x, barrier)

  def test_cta_barrier_oversubscribed_raises(self):
    x = jax.ShapeDtypeStruct((128,), jnp.float32)
    barrier = mgpu.CtaBarrier(arrival_count=256, num_barriers=14)

    def kernel(ctx, dst, barrier):
      del ctx, dst
      barrier[0].arrive_and_wait()

    with self.assertRaisesRegex(
        ValueError,
        "Too many CTA barriers requested: 14 requested, but only 13 available",
    ):
      mgpu.as_gpu_kernel(kernel, (1, 1, 1), (256, 1, 1), (), x, barrier)



def get_packed_shape(strides, shape):
  perm = sorted(range(len(strides)), key=lambda i: strides[i], reverse=True)
  ordered_strides = [strides[i] for i in perm]
  ordered_shape = [shape[i] for i in perm]
  packed_shape = [ordered_shape[-1]]
  packed_shape += [
      stride0 // stride
      for stride0, stride in zip(ordered_strides, ordered_strides[1:])
  ]
  # Invert permutation
  inv_perm = [None] * len(perm)
  for i, p in enumerate(perm):
    inv_perm[p] = i
  return [packed_shape[i] for i in inv_perm]


class WGMMALayoutTest(TestCase):

  @parameterized.product(dtype=[jnp.float16, jnp.float32])
  def test_store_untiled(self, dtype):
    def kernel(ctx, out, _):
      del ctx
      iota_tensor(64, 64, dtype).store_untiled(out, optimized=False)
    expected = np.arange(64 * 64, dtype=dtype).reshape(64, 64)
    iota = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), expected, ()
    )()
    np.testing.assert_array_equal(iota, expected)

  @parameterized.product(
      op=["add", "min", "max", "and", "or", "xor"],
      dtype=[jnp.float32, jnp.float16, jnp.bfloat16, jnp.int32, jnp.uint32],
      gmem=[False, True],
      vec_len=[2, 4, 8],
  )
  def test_atomic_store(self, op, dtype, gmem, vec_len):
    is_float = jnp.issubdtype(dtype, jnp.floating)
    is_16bit_float = jnp.finfo(dtype).bits == 16 if is_float else False
    if is_float and op in ("and", "or", "xor"):
      self.skipTest(f"{op} not supported for float types")
    if dtype == jnp.float32 and op != "add":
      self.skipTest("f32 only supports add with red instruction")
    if is_16bit_float and op in ("min", "max") and not gmem:
      self.skipTest("f16/bf16 min/max only supported on GMEM (vectorized)")
    m, n = 128, 64
    if op == "add" or op == "or" or op == "xor":
      identity = 0
    elif op == "min":
      identity = jnp.finfo(dtype).max if is_float else np.iinfo(dtype).max
    elif op == "max":
      identity = jnp.finfo(dtype).min if is_float else np.iinfo(dtype).min
    elif op == "and":
      identity = np.iinfo(dtype).max
    def kernel(ctx, inp_ref, out, smem):
      del ctx
      index = ir.IndexType.get()
      mlir_dtype = utils.dtype_to_ir_type(dtype)
      layout = fa.tmem_native_layout(vec_len)
      is_signed = utils.is_signed(dtype)
      wg_idx = arith.divui(gpu.thread_id(gpu.Dimension.x), c(128, index))
      offset = arith.muli(wg_idx, c(m, index))
      slice_ref = memref_slice(inp_ref, (ds(offset, m), slice(None)))
      my_fa = mgpu.FragmentedArray.load_untiled(
          slice_ref, layout=layout, is_signed=is_signed, optimized=False,
      )
      atomic_ref = out if gmem else smem
      mgpu.FragmentedArray.splat(
          c(identity, mlir_dtype), (m, n),
          is_signed=is_signed, layout=layout,
      ).store_untiled(atomic_ref, optimized=False)
      gpu.barrier()
      my_fa.store_untiled(atomic_ref, optimized=False, atomic=op)
      if atomic_ref is not out:
        gpu.barrier()
        copy(smem, out)
    x = np.arange(1, m * n + 1, dtype=dtype).reshape(m, n)
    y = np.arange(m * n, 0, -1, dtype=dtype).reshape(m, n)
    inp = np.stack([x, y]).reshape(2 * m, n)
    if op == "add":
      expected = x + y
    elif op == "min":
      expected = np.minimum(x, y)
    elif op == "max":
      expected = np.maximum(x, y)
    else:
      expected = getattr(np, f"bitwise_{op}")(x, y)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (256, 1, 1), inp, expected, expected,
    )(inp)
    np.testing.assert_array_equal(result, expected)

  @parameterized.product(
      op=["add", "min", "max"],
      dtype=[jnp.float32, jnp.float16, jnp.bfloat16],
  )
  def test_atomic_store_strided(self, op, dtype):
    if dtype == jnp.float32 and op != "add":
      self.skipTest("f32 only supports add with red instruction")
    m, n = 128, 64
    if op == "add":
      identity = 0
    elif op == "min":
      identity = jnp.finfo(dtype).max
    elif op == "max":
      identity = jnp.finfo(dtype).min
    def kernel(ctx, inp_ref, out, _):
      del ctx
      index = ir.IndexType.get()
      mlir_dtype = utils.dtype_to_ir_type(dtype)
      is_signed = utils.is_signed(dtype)
      wg_idx = arith.divui(gpu.thread_id(gpu.Dimension.x), c(128, index))
      offset = arith.muli(wg_idx, c(m, index))
      slice_ref = memref_slice(inp_ref, (ds(offset, m), slice(None)))
      my_fa = mgpu.FragmentedArray.load_strided(
          slice_ref, is_signed=is_signed,
      )
      mgpu.FragmentedArray.splat(
          c(identity, mlir_dtype), (m, n), is_signed=is_signed,
      ).store_untiled(out)
      gpu.barrier()
      my_fa.store_untiled(out, atomic=op)
    x = np.arange(1, m * n + 1, dtype=dtype).reshape(m, n)
    y = np.arange(m * n, 0, -1, dtype=dtype).reshape(m, n)
    inp = np.stack([x, y]).reshape(2 * m, n)
    if op == "add":
      expected = x + y
    elif op in ("min", "max"):
      expected = getattr(np, op + "imum")(x, y)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (256, 1, 1), inp, expected, (),
    )(inp)
    np.testing.assert_array_equal(result, expected)

  @parameterized.product(
      dtype=[jnp.float8_e5m2fnuz, jnp.float8_e5m2, jnp.float8_e4m3b11fnuz,
             jnp.float8_e4m3fn, jnp.float8_e4m3fnuz],
      swizzle=(32, 64, 128),
      num_col_tiles=(1, 2, 3),
  )
  def test_load_and_store_tiled_f8(self, dtype, swizzle, num_col_tiles):
    # We use a different test than `test_store_tiled` because converting
    # `iota` to `f8` type requires additional specialized logic that is not
    # yet available.
    col_tiling = swizzle
    m = 128
    n = col_tiling * num_col_tiles
    tiling = (64, col_tiling)
    def kernel(ctx, inp, out, smem):
      del ctx
      smem_inp, smem_out = smem
      copy(inp, smem_inp, swizzle=swizzle)
      arr = mgpu.FragmentedArray.load_tiled(smem_inp, swizzle=swizzle)
      arr.store_tiled(smem_out, swizzle=swizzle)
      copy(smem_out, out, swizzle=swizzle)
    expected = (
        jax.random.randint(
            jax.random.key(42), (m * n,), -16, 15, dtype=jnp.int8
        )
        .reshape(m // tiling[0], tiling[0], n // tiling[1], tiling[1])
        .astype(dtype)
        .transpose(0, 2, 1, 3)
    )
    res = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), expected, expected, (expected,) * 2
    )(expected)
    np.testing.assert_array_equal(res, expected)

  @parameterized.product(
      dtype=[jnp.float32, jnp.float16, jnp.int8],
      swizzle=(32, 64, 128),
      num_col_tiles=(1, 2),
  )
  def test_store_tiled(self, dtype, swizzle, num_col_tiles):
    mlir_dtype = utils.dtype_to_ir_type(dtype)
    if bytewidth(mlir_dtype) > 2 and swizzle == 32:
      self.skipTest("Not implemented")
    col_tiling = swizzle // bytewidth(mlir_dtype)
    m = 128
    n = col_tiling * num_col_tiles
    tiling = (64, col_tiling)
    def kernel(ctx, out, smem):
      del ctx
      iota_tensor(m, n, dtype).store_tiled(smem, swizzle=swizzle)
      copy(smem, out, swizzle=swizzle)
    expected = (
        np.arange(m * n, dtype=dtype)
        .reshape(m // tiling[0], tiling[0], n // tiling[1], tiling[1])
        .transpose(0, 2, 1, 3)
    )
    iota = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), expected, expected
    )()
    np.testing.assert_array_equal(iota, expected)

  @parameterized.product(
      dtype=[jnp.float32, jnp.float16],
      swizzle=(32, 64, 128),
  )
  def test_store_tiled_with_tiling_rank(self, dtype, swizzle):
    mlir_dtype = utils.dtype_to_ir_type(dtype)
    if bytewidth(mlir_dtype) > 2 and swizzle == 32:
      self.skipTest("Not implemented")
    col_tiling = swizzle // bytewidth(mlir_dtype)
    m = 64
    n = col_tiling * 2
    tiling = (64, col_tiling)
    def kernel(ctx, out, smem):
      del ctx
      assert smem.type.rank == 5
      iota_tensor(2 * m, n, dtype).reshape((2, m, n)).store_tiled(smem, swizzle=swizzle, tiling_rank=2)
      copy(smem, out, swizzle=swizzle)
    expected = (
        np.arange(2 * m * n, dtype=dtype)
        .reshape(2, m // tiling[0], tiling[0], n // tiling[1], tiling[1])
        .transpose(0, 1, 3, 2, 4)
    )
    iota = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), expected, expected
    )()
    np.testing.assert_array_equal(iota, expected)

  @parameterized.product(
      jax_dtype_to=(
          jnp.int8, jnp.int16, jnp.int32, jnp.bfloat16, jnp.float8_e4m3fn,
      ),
      # Use different layouts to vary the size of the vector dimension.
      layout=(
          fa.WGMMA_LAYOUT,
          fa.WGMMA_LAYOUT_UPCAST_2X,
          fa.WGMMA_LAYOUT_UPCAST_4X,
      ),
  )
  def test_sub_byte_conversion(self, jax_dtype_to, layout: fa.TiledLayout):
    if jax_dtype_to == jnp.int32 and layout.vector_length == 8:
      self.skipTest(
          "Raises: failed to prove that vector transfers don't cross swizzle"
          " tile boundaries.")
    jax_dtype_from = jnp.int4
    is_signed = utils.is_signed(jax_dtype_to)
    def kernel(ctx, inp, out, smem):
      del ctx  # Unused.
      smem_inp, smem_out = smem
      copy(inp, smem_inp, swizzle=16)
      t = mgpu.FragmentedArray.load_tiled(
          smem_inp, is_signed=True, swizzle=16, layout=layout
      )
      t = t.astype(utils.dtype_to_ir_type(jax_dtype_to), is_signed=is_signed)
      t.store_tiled(smem_out, swizzle=32 * jnp.dtype(jax_dtype_to).itemsize)
      copy(smem_out, out, swizzle=32 * jnp.dtype(jax_dtype_to).itemsize)

    x = self.prng.integers(
        low=-8, high=7, size=(1, 1, 64, 64), dtype=np.int32
    ).astype(jax_dtype_from)
    y = jax.lax.convert_element_type(x, jax_dtype_to)
    f = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, y, (x, y))
    np.testing.assert_array_equal(f(x), y)

  @parameterized.product(
      jax_dtype_from_to=(
          (jnp.int8, jnp.bfloat16),
          (jnp.int4, jnp.bfloat16),
          (jnp.int4, jnp.float8_e4m3fn),
          (jnp.int4, jnp.int8),
          # TODO(apaszke,bchetioui): bf16/f32 -> f8e4m3fn
      ),
      layout_descs=(
          ("WGMMA_LAYOUT", "WGMMA_LAYOUT"),
          ("WGMMA_LAYOUT_8BIT", "WGMMA_LAYOUT_8BIT"),
          ("WGMMA_LAYOUT_UPCAST_2X", "WGMMA_LAYOUT_UPCAST_2X"),
          ("WGMMA_LAYOUT_UPCAST_2X", "WGMMA_LAYOUT"),
          ("WGMMA_LAYOUT", "WGMMA_LAYOUT_UPCAST_2X"),
          ("WGMMA_LAYOUT_UPCAST_4X", "WGMMA_LAYOUT_UPCAST_4X"),
          ("WGMMA_LAYOUT_UPCAST_4X", "WGMMA_LAYOUT_UPCAST_2X"),
          ("WGMMA_LAYOUT_UPCAST_4X", "WGMMA_LAYOUT"),
          ("TMEM_NATIVE_LAYOUT_8", "TMEM_NATIVE_LAYOUT_8"),
      ),
  )
  def test_optimized_conversion(self, jax_dtype_from_to, layout_descs):
    layout_desc_from, layout_desc_to = layout_descs
    if layout_desc_from == "TMEM_NATIVE_LAYOUT_8":
      layout_from = fa.tmem_native_layout(8)
    else:
      layout_from: fa.TiledLayout = getattr(fa, layout_desc_from)
    if layout_desc_to == "TMEM_NATIVE_LAYOUT_8":
      layout_to = fa.tmem_native_layout(8)
    else:
      layout_to: fa.TiledLayout = getattr(fa, layout_desc_to)
    jax_dtype_from, jax_dtype_to = jax_dtype_from_to
    mlir_dtype_from = utils.dtype_to_ir_type(jax_dtype_from)
    mlir_dtype_to = utils.dtype_to_ir_type(jax_dtype_to)
    m = 128
    n = 256
    def kernel(ctx, inp, out, smem):
      del ctx, smem
      t = mgpu.FragmentedArray.load_untiled(
          inp,
          is_signed=utils.is_signed(jax_dtype_from),
          layout=layout_from,
          optimized=False,
      )
      if layout_from != layout_to:
        if (
            layout_from == fa.WGMMA_LAYOUT_UPCAST_4X
            and utils.bitwidth(mlir_dtype_from) != 4
        ):
          self.skipTest("Unimplemented relayout")
        t = t.to_layout(layout_to)
      t = t.astype(mlir_dtype_to, is_signed=utils.is_signed(jax_dtype_to))
      t.store_untiled(out, optimized=False)

    int_sample_dtype = getattr(
        jnp,
        "int" + str(min(bitwidth(mlir_dtype_from), bitwidth(mlir_dtype_to))),
    )
    sample_iinfo = jnp.iinfo(int_sample_dtype)
    values = self.prng.integers(
        low=sample_iinfo.min, high=sample_iinfo.max, size=(m, n), dtype=np.int32
    ).astype(jax_dtype_from)

    expected = values.astype(np.int32).astype(jax_dtype_to)
    @contextlib.contextmanager
    def _maybe_profile():
      yield; return   # Comment to gather statistics.
      with jtu.set_env(MOSAIC_GPU_DUMP_SASS="1"), self.capture_stdout() as sass:
        yield
      log_dir = os.getenv("TEST_UNDECLARED_OUTPUTS_DIR", "/tmp")
      file_path = os.path.join(log_dir, "conversion_stats.csv")
      with open(file_path, "a") as f:
        data = (
            jnp.dtype(jax_dtype_from).name, jnp.dtype(jax_dtype_to).name,
            layout_desc_from, layout_desc_to, sass().count("\n")
        )
        f.write(",".join(map(str, data)) + "\n")
        f.flush()
      self.fail("Disable profiling before submission")
    with _maybe_profile():
      res = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), values, expected, ()
      )(values)
    np.testing.assert_array_equal(res, expected)

  @parameterized.named_parameters(
      ("f32", jnp.float32),
      ("f16", jnp.float16),
      ("i8", jnp.int8),
  )
  def test_load_tiled(self, jax_dtype):
    mlir_dtype = utils.dtype_to_ir_type(jax_dtype)
    m = 128
    n = 256 // bytewidth(mlir_dtype)
    tiling = (64, 128 // bytewidth(mlir_dtype))
    def kernel(ctx, in_, out, smem):
      del ctx
      smem1, smem2 = smem
      copy(in_, smem1, swizzle=128)
      t = mgpu.FragmentedArray.load_tiled(
          smem1, swizzle=128, is_signed=utils.is_signed(jax_dtype)
      )
      t.store_tiled(smem2, swizzle=128)
      copy(smem2, out, swizzle=128)
    expected = (
        np.arange(m * n, dtype=jax_dtype)
        .reshape(m // tiling[0], tiling[0], n // tiling[1], tiling[1])
        .transpose(0, 2, 1, 3)
    )
    iota = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), expected, expected, (expected,) * 2
    )(expected)
    np.testing.assert_array_equal(iota, expected)


class I8Type:
  """A type that represents a 8-bit signed integer.

  This is a workaround to bypass the fact that we don't have a proper 8-bit
  integer type class available in MLIR, and can't instantiate types without a
  MLIR context.
  """

  @staticmethod
  def get():
    return ir.IntegerType.get_signless(8)


# Maps an FP8 MLIR element-type class to the (jax dtype, exponent_bits,
# mantissa_bits) used to quantize reference inputs to that type.
_FP8_DTYPE_INFO = {
    ir.Float8E5M2Type: (jnp.float8_e5m2, 5, 2),
    ir.Float8E4M3FNType: (jnp.float8_e4m3fn, 4, 3),
}


class WGMMATest(TestCase):

  def setUp(self):
    super().setUp()
    if not jtu.is_cuda_compute_capability_equal("9.0"):
      self.skipTest("Only works on GPU with capability sm90a")

  @parameterized.product(
      lhs_transpose=(False, True),
      rhs_transpose=(False, True),
      in_mlir_dtype_cls=(
          ir.F16Type,
          ir.BF16Type,
          ir.F32Type,
          ir.Float8E5M2Type,
          ir.Float8E4M3FNType,
      ),
      m=(64, 128, 192),
      n=(64, 128, 192),
      swizzle=(32, 64, 128),
      jax_out_dtype=(jnp.float16, jnp.float32),
  )
  def test_wgmma_basic_float(
      self,
      lhs_transpose,
      rhs_transpose,
      in_mlir_dtype_cls,
      m,
      n,
      swizzle,
      jax_out_dtype,
  ):
    self._test_wgmma_basic(
        m,
        n,
        k_steps=2,  # Decrease to 1 to simplify debugging.
        in_mlir_dtype_cls=in_mlir_dtype_cls,
        lhs_transpose=lhs_transpose,
        rhs_transpose=rhs_transpose,
        swizzle=swizzle,
        jax_out_dtype=jax_out_dtype,
        lhs_tiling_kind="small+no_transpose" if lhs_transpose else "small",
        rhs_tiling_kind="small+no_transpose" if rhs_transpose else "small",
    )

  @parameterized.product(
      lhs_rhs_dtype_cls=(
          (ir.Float8E4M3FNType, ir.Float8E5M2Type),
          (ir.Float8E5M2Type, ir.Float8E4M3FNType),
      ),
      m=(64, 128),
      n=(128,),
      swizzle=(64, 128),
      jax_out_dtype=(jnp.float16, jnp.float32),
  )
  def test_wgmma_mixed_fp8(self, lhs_rhs_dtype_cls, m, n, swizzle, jax_out_dtype):
    # The e4m3/e5m2 FP8 pair may be mixed across the two operands. FP8 wgmma does
    # not support operand transposes, so the only valid layout is `a` row-major
    # and `b` column-major (rhs_transpose=True).
    lhs_dtype_cls, rhs_dtype_cls = lhs_rhs_dtype_cls
    self._test_wgmma_basic(
        m=m,
        n=n,
        k_steps=2,
        in_mlir_dtype_cls=lhs_dtype_cls,
        rhs_in_mlir_dtype_cls=rhs_dtype_cls,
        lhs_transpose=False,
        rhs_transpose=True,
        swizzle=swizzle,
        jax_out_dtype=jax_out_dtype,
        lhs_tiling_kind="small",
        rhs_tiling_kind="small+no_transpose",
    )

  @parameterized.product(
      in_mlir_dtype_cls=(I8Type,),
      m=(64, 128, 192),
      n=(64, 128, 192),
      swizzle=(32, 64, 128),
      jax_out_dtype=(jnp.int32,),
  )
  def test_wgmma_basic_int(
      self, in_mlir_dtype_cls, m, n, swizzle, jax_out_dtype,
  ):
    self._test_wgmma_basic(
        m,
        n,
        k_steps=2,  # Decrease to 1 to simplify debugging.
        in_mlir_dtype_cls=in_mlir_dtype_cls,
        lhs_transpose=False,
        rhs_transpose=True,
        swizzle=swizzle,
        jax_out_dtype=jax_out_dtype,
        rhs_tiling_kind="small",
        lhs_tiling_kind="small+no_transpose",
    )

  @parameterized.product(
      lhs_transpose=(False, True),
      rhs_transpose=(False, True),
      in_mlir_dtype_cls=(
          ir.F32Type,
          ir.F16Type,
          ir.Float8E5M2Type,
      ),
      swizzle=(32, 64, 128),
      rhs_tiling_kind=("large", "small", "small+no_transpose"),
      lhs_tiling_kind=("large", "small", "small+no_transpose"),
  )
  def test_wgmma_transposes(
      self,
      lhs_transpose,
      rhs_transpose,
      in_mlir_dtype_cls,
      swizzle,
      rhs_tiling_kind,
      lhs_tiling_kind,
  ):
    self._test_wgmma_basic(
        m=128,
        n=192,
        k_steps=2,  # Decrease to 1 to simplify debugging.
        in_mlir_dtype_cls=in_mlir_dtype_cls,
        lhs_transpose=lhs_transpose,
        rhs_transpose=rhs_transpose,
        swizzle=swizzle,
        jax_out_dtype=jnp.float32,
        rhs_tiling_kind=rhs_tiling_kind,
        lhs_tiling_kind=lhs_tiling_kind,
    )

  def _test_wgmma_basic(
      self,
      m,
      n,
      k_steps,
      in_mlir_dtype_cls,
      lhs_transpose,
      rhs_transpose,
      swizzle,
      jax_out_dtype,
      rhs_tiling_kind,
      lhs_tiling_kind,
      rhs_in_mlir_dtype_cls=None,
  ):
    # `b` reuses `a`'s element type unless a distinct one is given (only the
    # e4m3/e5m2 FP8 pair may be mixed).
    if rhs_in_mlir_dtype_cls is None:
      rhs_in_mlir_dtype_cls = in_mlir_dtype_cls
    if jax_out_dtype == jnp.int32 and in_mlir_dtype_cls != I8Type:
      self.skipTest("s32 accumulator only supported with s8 inputs")
    if jax_out_dtype != jnp.int32 and in_mlir_dtype_cls == I8Type:
      self.skipTest("s8 inputs only supported with s32 accumulator")
    if jax_out_dtype == jnp.float16 and in_mlir_dtype_cls in {ir.F32Type, ir.BF16Type}:
      self.skipTest(f"{in_mlir_dtype_cls.get()} does not support f16 output.")
    if swizzle != 128 and lhs_transpose and lhs_tiling_kind == "large":
      self.skipTest("Transpose only supported in 128B swizzled WGMMA")
    if rhs_tiling_kind == "small+no_transpose" and not rhs_transpose:
      self.skipTest("No transpose happening anyway")
    if lhs_tiling_kind == "small+no_transpose" and not lhs_transpose:
      self.skipTest("No transpose happening anyway")

    in_mlir_dtype = in_mlir_dtype_cls.get()
    out_mlir_dtype = utils.dtype_to_ir_type(jax_out_dtype)
    if (lhs_transpose or not rhs_transpose) and bytewidth(in_mlir_dtype) != 2:
      self.skipTest("Transpose only supported in 16-bit WGMMA")
    if isinstance(in_mlir_dtype, ir.F32Type):  # We actually use tf32 instead
      in_jax_dtype = jnp.float32
      exponent_bits, mantissa_bits = 8, 10  # Use tf32
    elif bytewidth(in_mlir_dtype) == 2:
      if n % 64 != 0:
        self.skipTest("16-bit WGMMA only supports n % 64 == 0")
      if isinstance(in_mlir_dtype, ir.F16Type):
        in_jax_dtype = jnp.float16
        exponent_bits, mantissa_bits = 5, 10
      elif isinstance(in_mlir_dtype, ir.BF16Type):
        in_jax_dtype = jnp.bfloat16
        exponent_bits, mantissa_bits = 8, 7
      else:
        raise NotImplementedError(in_mlir_dtype)
    elif in_mlir_dtype_cls in _FP8_DTYPE_INFO:
      in_jax_dtype, exponent_bits, mantissa_bits = (
          _FP8_DTYPE_INFO[in_mlir_dtype_cls]
      )
    elif in_mlir_dtype_cls == I8Type:
      in_jax_dtype = jnp.int8
      exponent_bits = mantissa_bits = None
    else:
      raise NotImplementedError(in_mlir_dtype)

    # `b`'s dtype matches `a`'s except for the mixed e4m3/e5m2 FP8 case. Both FP8
    # types are 1 byte, so the tiling/swizzle math (driven by `in_mlir_dtype`) is
    # identical; only the operand values and `b`'s SMEM dtype differ.
    if rhs_in_mlir_dtype_cls is in_mlir_dtype_cls:
      rhs_in_jax_dtype = in_jax_dtype
      rhs_exponent_bits, rhs_mantissa_bits = exponent_bits, mantissa_bits
    elif rhs_in_mlir_dtype_cls in _FP8_DTYPE_INFO:
      rhs_in_jax_dtype, rhs_exponent_bits, rhs_mantissa_bits = (
          _FP8_DTYPE_INFO[rhs_in_mlir_dtype_cls]
      )
    else:
      raise NotImplementedError(
          "Mixed-dtype WGMMA is only supported for the e4m3/e5m2 FP8 pair, got"
          f" a={in_mlir_dtype_cls} b={rhs_in_mlir_dtype_cls}"
      )
    nk_tile = swizzle // bytewidth(in_mlir_dtype)
    k = nk_tile * k_steps
    if n % nk_tile:
      self.skipTest("tiling does not divide N")
    assert m % 64 == 0 and n % nk_tile == 0

    small_rhs_tile = rhs_tiling_kind != "large"
    transpose_rhs_tiles = rhs_tiling_kind != "small+no_transpose"
    rhs_tiling = (8, nk_tile) if small_rhs_tile else (nk_tile, nk_tile)
    small_lhs_tile = lhs_tiling_kind != "large"
    transpose_lhs_tiles = lhs_tiling_kind != "small+no_transpose"
    lhs_tiling = (8, nk_tile) if small_lhs_tile else (64, nk_tile)

    def kernel(ctx, lhs, rhs, out, scratch):
      lhs_smem, rhs_smem, barriers = scratch
      lhs_transform = (mgpu.TileTransform(lhs_tiling),)
      if lhs_transpose and transpose_lhs_tiles:
        lhs_transform += (mgpu.TransposeTransform((1, 0, 2, 3)),)
      rhs_transform = (mgpu.TileTransform(rhs_tiling),)
      if rhs_transpose and transpose_rhs_tiles:
        rhs_transform += (mgpu.TransposeTransform((1, 0, 2, 3)),)
      ctx.async_copy(
          src_ref=lhs,
          dst_ref=lhs_smem,
          swizzle=swizzle,
          gmem_transform=lhs_transform,
          barrier=barriers[0],
      )
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          swizzle=swizzle,
          gmem_transform=rhs_transform,
          barrier=barriers[1],
      )
      for i in range(2):
        barriers[i].wait()
      is_signed = True if isinstance(in_mlir_dtype, ir.IntegerType) else None
      init_acc = mgpu.WGMMAAccumulator.zero(m=m, n=n, dtype=out_mlir_dtype, is_signed=is_signed)
      if lhs_transpose:
        perm = (0, 1, 3, 2) if transpose_lhs_tiles else (1, 0, 3, 2)
        lhs_smem = memref_transpose(lhs_smem, perm)
      if rhs_transpose:
        perm = (0, 1, 3, 2) if transpose_rhs_tiles else (1, 0, 3, 2)
        rhs_smem = memref_transpose(rhs_smem, perm)
      acc = mgpu.wgmma(init_acc, lhs_smem, rhs_smem, swizzle=swizzle)
      nvvm.wgmma_commit_group_sync_aligned()
      nvvm.wgmma_wait_group_sync_aligned(0)
      acc.value.store_untiled(out, optimized=False)

    x_shape = (k, m) if lhs_transpose else (m, k)
    y_shape = (n, k) if rhs_transpose else (k, n)
    if in_mlir_dtype_cls == I8Type:
      x = self.prng.integers(-128, 127, x_shape).astype(in_jax_dtype)
      y = self.prng.integers(-128, 127, y_shape).astype(rhs_in_jax_dtype)
    else:
      x = jax.lax.reduce_precision(
        self.prng.uniform(-1, 1, x_shape), exponent_bits, mantissa_bits
      ).astype(in_jax_dtype)
      y = jax.lax.reduce_precision(
        self.prng.uniform(-1, 1, y_shape), rhs_exponent_bits, rhs_mantissa_bits
      ).astype(rhs_in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), jax_out_dtype)
    if transpose_rhs_tiles:
      rhs_tiling_t = rhs_tiling[::-1] if rhs_transpose else rhs_tiling
      rhs_smem_shape = (k // rhs_tiling_t[0], n // rhs_tiling_t[1], *rhs_tiling)
    else:
      rhs_smem_shape = tile_shape(y_shape, rhs_tiling)
    if transpose_lhs_tiles:
      lhs_tiling_t = lhs_tiling[::-1] if lhs_transpose else lhs_tiling
      lhs_smem_shape = (m // lhs_tiling_t[0], k // lhs_tiling_t[1], *lhs_tiling)
    else:
      lhs_smem_shape = tile_shape(x_shape, lhs_tiling)
    scratch_shape = [
        jax.ShapeDtypeStruct(lhs_smem_shape, in_jax_dtype),
        jax.ShapeDtypeStruct(rhs_smem_shape, rhs_in_jax_dtype),
        mgpu.Barrier(1, num_barriers=2),
    ]
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (x, y), out_shape, scratch_shape
    )(x, y)
    x32, y32 = x.astype(np.float32), y.astype(np.float32)
    ref = jax.lax.dot(
        x32.T if lhs_transpose else x32,
        y32.T if rhs_transpose else y32,
        precision=jax.lax.Precision.HIGH
    )
    atol = 2e-2 if jax_out_dtype == jnp.float16 else 5e-6
    if isinstance(in_mlir_dtype, ir.IntegerType) and isinstance(out_mlir_dtype, ir.IntegerType):
      atol = 0
    elif utils.bitwidth(in_mlir_dtype) == 8:
      atol = 3e-2
    np.testing.assert_allclose(z, ref, atol=atol)

  # TODO(apaszke): Add support for f32
  @parameterized.product(
      m=(64, 128, 192),
      n=(64, 128, 192),
      k_steps=(1, 2),
      rhs_transpose=(False, True),
      swizzle=(32, 64, 128),
      dtype=[jnp.float16, jnp.bfloat16],
  )
  def test_wgmma_reg_lhs(self, m, n, k_steps, rhs_transpose, swizzle, dtype):
    self._test_wgmma_reg_lhs(m, n, k_steps, rhs_transpose, swizzle, dtype)

  @parameterized.product(
      m=(64, 128, 192),
      n=(64, 128, 192),
      k_steps=(1, 2),
      swizzle=(32, 64, 128),
      dtype=(jnp.int8, jnp.float8_e5m2, jnp.float8_e4m3fn),
  )
  def test_wgmma_reg_lhs_8bit(self, m, n, k_steps, swizzle, dtype):
    # TODO(bchetioui): relax this when ptxas is fixed. As of ptxas 12.8,
    # optimizations eliminate MMA instructions, leading to only the first tile
    # of the result being computed correctly.
    if swizzle == 32 and dtype == jnp.int8:
      self.skipTest("32-bit swizzle not supported for int8")
    self._test_wgmma_reg_lhs(
        m, n, k_steps, rhs_transpose=True, swizzle=swizzle, dtype=dtype
    )

  def _test_wgmma_reg_lhs(self, m, n, k_steps, rhs_transpose, swizzle, dtype):
    index = ir.IndexType.get()
    out_dtype = jnp.int32 if dtype == jnp.int8 else jnp.float32
    bytewidth = jnp.dtype(dtype).itemsize
    nk_tile = swizzle // bytewidth
    k = nk_tile * k_steps
    if n % nk_tile:
      self.skipTest("swizzle must divide N")

    def kernel(ctx, lhs, rhs, out, rhs_smem):
      del ctx
      for ki in range(k_steps):
        for ni in range(n // nk_tile):
          rhs_slice = (
              ds(c(ki * nk_tile, index), nk_tile),
              ds(c(ni * nk_tile, index), nk_tile),
          )
          if rhs_transpose:
            rhs_slice = rhs_slice[::-1]
          copy(
              src=memref_slice(rhs, rhs_slice),
              dst=memref_slice(rhs_smem, (ki, ni)),
              swizzle=swizzle,
          )
      init_acc = mgpu.WGMMAAccumulator.zero(
          m=m, n=n, dtype=utils.dtype_to_ir_type(out_dtype),
          is_signed=True if dtype == jnp.int8 else None,
      )
      layout = fa.WGMMA_LAYOUT_8BIT if dtypes.itemsize_bits(dtype) == 8 else fa.WGMMA_LAYOUT
      lhs_regs = fa.FragmentedArray.load_untiled(
          lhs, layout=layout, optimized=False, is_signed=utils.is_signed(dtype),
      )
      if rhs_transpose:
        rhs_smem = memref_transpose(rhs_smem, (0, 1, 3, 2))
      acc = mgpu.wgmma(init_acc, lhs_regs, rhs_smem, swizzle=swizzle)
      nvvm.wgmma_commit_group_sync_aligned()
      nvvm.wgmma_wait_group_sync_aligned(0)
      acc.value.store_untiled(out, optimized=False)

    y_shape = (n, k) if rhs_transpose else (k, n)
    if dtype == jnp.int8:
      x = np.arange(m * k, dtype=dtype).reshape(m, k)
      y = self.prng.integers(-128, 127, y_shape, dtype=dtype)
    else:
      def quantize_f8(x):
        if dtype not in {jnp.float8_e4m3fn, jnp.float8_e5m2}:
          return x
        if dtype == jnp.float8_e4m3fn:
          exponent_bits, mantissa_bits = 4, 3
        else:
          exponent_bits, mantissa_bits = 5, 2
        return jax.lax.reduce_precision(x, exponent_bits, mantissa_bits)
      x = quantize_f8(self.prng.uniform(-1, 1, (m, k))).astype(dtype)
      y = quantize_f8(self.prng.uniform(-1, 1, y_shape)).astype(dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_dtype)
    scratch_shape = jax.ShapeDtypeStruct(
            (k_steps, n // nk_tile, nk_tile, nk_tile), dtype
    )
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (x, y), out_shape, scratch_shape
    )(x, y)
    ref = jax.lax.dot(
        x, (y.T if rhs_transpose else y), preferred_element_type=out_dtype
    )
    if dtype == jnp.int8:
      atol = rtol = 0
    elif dtype == jnp.float8_e4m3fn:
      atol = rtol = 6e-3
    elif dtype == jnp.float8_e5m2:
      atol = rtol = 3e-3
    else:
      atol, rtol = 1e-6, 5e-4
    np.testing.assert_allclose(z, ref, rtol=rtol, atol=atol)

  @parameterized.product(
      rhs_transpose=(False, True),
      swizzle=(32, 64, 128),
      n=(8, 16),
      small_rhs_tile=(False, True),
  )
  def test_narrow_n(self, rhs_transpose, swizzle, n, small_rhs_tile):
    m, k_steps = 64, 2
    bytewidth = 2
    nk_tile = swizzle // bytewidth
    k = nk_tile * k_steps
    if small_rhs_tile and not rhs_transpose:
      self.skipTest("Small tiles only supported for transposed RHS")

    n_tile = 8 if small_rhs_tile else nk_tile

    def kernel(ctx, rhs, out, smem):
      rhs_smem, barrier = smem
      gmem_slice = (ds(0, k), ds(0, max(n_tile, n)))
      transform = (mgpu.TileTransform((n_tile, nk_tile)),)
      if rhs_transpose:
        gmem_slice = gmem_slice[::-1]
        transform += (mgpu.TransposeTransform((1, 0, 2, 3)),)
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          swizzle=swizzle,
          gmem_slice=gmem_slice,
          gmem_transform=transform,
          barrier=barrier,
      )
      barrier.wait()
      init_acc = mgpu.WGMMAAccumulator.zero(m=m, n=n)
      lhs_regs = iota_tensor(m, k, jnp.float16)
      if rhs_transpose:
        rhs_smem = memref_transpose(rhs_smem, (0, 1, 3, 2))
      if not small_rhs_tile:
        smem_slice = (slice(None), slice(None), slice(None), ds(0, n))
        rhs_smem = memref_slice(rhs_smem, smem_slice)
      acc = mgpu.wgmma(init_acc, lhs_regs, rhs_smem, swizzle=swizzle)
      nvvm.wgmma_commit_group_sync_aligned()
      nvvm.wgmma_wait_group_sync_aligned(0)
      acc.value.store_untiled(out, optimized=False)

    jax_dtype = jnp.float16
    y_shape = (n, k) if rhs_transpose else (k, n)
    y = self.prng.uniform(-1, 1, y_shape).astype(jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), jnp.float32)
    rhs_scratch_shape = jax.ShapeDtypeStruct(
        (k_steps, (n + n_tile - 1) // n_tile, n_tile, nk_tile), jax_dtype
    )
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), y, out_shape, (rhs_scratch_shape, mgpu.Barrier(1)),
    )(y)
    x = np.arange(m * k, dtype=jax_dtype).reshape(m, k)
    ref = jax.lax.dot(
        x, (y.T if rhs_transpose else y), preferred_element_type=jnp.float32
    )
    np.testing.assert_allclose(z, ref, rtol=1e-3, atol=0)


class TCGen05Test(TestCase, jtu.CudaArchSpecificTest):

  def setUp(self):
    self.skip_unless_tcgen05()
    # No artificially lowered limit for arch-specific tests
    super().setUp(artificial_shared_memory_limit=None)

  @parameterized.product(
      dtype_packing=[(jnp.float32, 1), (jnp.float16, 1), (jnp.float16, 2), (jnp.float8_e5m2, 4)],
      reg_tmem_layout=[
          (lambda _c, _p: tcgen05.LAYOUT, lambda _, p: tcgen05.tmem_default_layout(p)),
          (lambda _c, _p: fa.WGMMA_LAYOUT, tcgen05.tmem_half_lane_layout),
          (
              lambda c, _p: tcgen05.fa_m64_collective_layout(c),
              tcgen05.tmem_m64_collective_layout,
          ),
          (
              lambda c, p: tcgen05.tmem_m64_collective_layout(c, p).as_tiled_layout(),
              tcgen05.tmem_m64_collective_layout,
          ),
      ],
  )
  def test_load_store_tmem(self, dtype_packing, reg_tmem_layout):
    dtype, packing = dtype_packing
    reg_layout_f, tmem_layout_f = reg_tmem_layout
    n = 160
    reg_layout = reg_layout_f(n, packing)
    tmem_layout = tmem_layout_f(n, packing)
    m = tmem_layout.base_tile_shape[0]
    is_native_transfer = tmem_layout.as_tiled_layout() == reg_layout
    packed = jnp.finfo(dtype).bits * packing == 32
    if is_native_transfer and not packed:
      self.skipTest("Not implemented yet")
    if not is_native_transfer and dtype == jnp.float8_e5m2:
      self.skipTest("Not implemented yet")

    def kernel(ctx, input, output, tmem):
      del ctx
      tmem.store(fa.FragmentedArray.load_untiled(input, layout=reg_layout, optimized=False))
      tcgen05.commit_tmem()
      tmem.load(reg_layout).store_untiled(output, optimized=False)

    x = self.prng.uniform(-1, 1, (m, n)).astype(dtype)
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, x,
        mgpu.TMEM(x.shape, dtype, layout=tmem_layout),
    )(x)
    np.testing.assert_array_equal(x, y)

  @parameterized.parameters([
      (jnp.float32, 1),
      (jnp.float16, 1),
      (jnp.float16, 2),
      (jnp.float8_e5m2, 4),
      (jnp.float4_e2m1fn, 8),
  ])
  def test_load_store_tmem_native(self, jax_dtype, packing):
    def kernel(ctx, input, output, tmem):
      del ctx
      reg_layout = tcgen05.tmem_default_layout(max(packing, 2)).as_tiled_layout()
      tmem.store(fa.FragmentedArray.load_untiled(input, layout=reg_layout, optimized=False))
      tcgen05.commit_tmem()
      tmem.load(reg_layout).store_untiled(output, optimized=False)

    x = self.prng.uniform(-1, 1, (128, 128)).astype(jax_dtype)
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, x, mgpu.TMEM(x.shape, jax_dtype, packing=packing)
    )(x)
    np.testing.assert_array_equal(x, y)

  @parameterized.parameters([
      (jnp.float32, 1, tcgen05.LAYOUT),
      (jnp.float16, 1, tcgen05.LAYOUT),
      (jnp.float16, 2, tcgen05.LAYOUT),
      (jnp.float8_e5m2, 4, fa.tmem_native_layout(4)),
      (jnp.float4_e2m1fn, 8, fa.tmem_native_layout(8)),
  ])
  def test_load_store_tmem_inference(self, jax_dtype, packing, expected_layout):
    def kernel(ctx, input, output, tmem):
      del ctx
      reg_layout = tcgen05.tmem_default_layout(max(packing, 2)).as_tiled_layout()
      tmem.store(fa.FragmentedArray.load_untiled(input, layout=reg_layout, optimized=False))
      tcgen05.commit_tmem()
      y = tmem.load()
      self.assertEqual(y.layout, expected_layout)
      y.store_untiled(output, optimized=False)

    x = self.prng.uniform(-1, 1, (128, 128)).astype(jax_dtype)
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, x, mgpu.TMEM(x.shape, jax_dtype, packing=packing)
    )(x)
    np.testing.assert_array_equal(x, y)

  @parameterized.product(
      dtype=[jnp.float32, jnp.int32],
      reduction=["min", "max", "absmin", "absmax"],
  )
  def test_load_reduce_tmem_native(self, dtype, reduction):
    if not (
        jtu.is_cuda_compute_capability_at_least("10.1")
        and not jtu.is_cuda_compute_capability_at_least("12.0")
    ):
      self.skipTest("Requires compute capability between 10.1 or 12.0")
    if not jtu.is_cuda_version_at_least(13, 0):
      self.skipTest("Requires CUDA version 13.0 or higher")
    shape = (128, 128)
    if "abs" in reduction and dtype == jnp.int32:
      reduction = reduction[-3:]
      dtype = jnp.uint32

    def kernel(ctx, input_, output, red_output, tmem):
      del ctx
      reg_layout = tcgen05.tmem_default_layout(1).as_tiled_layout()
      input_arr = fa.FragmentedArray.load_untiled(
          input_, layout=reg_layout, optimized=False, is_signed=utils.is_signed(dtype)
      )
      tmem.store(input_arr)
      tcgen05.commit_tmem()
      loaded, reduced = tmem.load(tcgen05.TMEM_NATIVE_LAYOUT, reduce=reduction, is_signed=utils.is_signed(dtype))
      loaded.store_untiled(output, optimized=False)
      reduced.store_untiled(red_output, optimized=False)

    if dtype == jnp.int32:
      x = self.prng.integers(-100000, 100000, shape).astype(dtype)
    elif dtype == jnp.uint32:
      x = self.prng.integers(0, 1000000, shape).astype(dtype)
    else:
      x = self.prng.uniform(-1, 1, shape).astype(dtype)
    red_shape = jax.ShapeDtypeStruct((shape[0],), dtype)
    y, red_y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        x,
        (x, red_shape),
        mgpu.TMEM(shape, dtype),
    )(x)
    np.testing.assert_array_equal(y, x)
    match reduction:
      case "min":
        expected = np.min(x, axis=1)
      case "max":
        expected = np.max(x, axis=1)
      case "absmin":
        expected = np.min(np.abs(x), axis=1)
      case "absmax":
        expected = np.max(np.abs(x), axis=1)
      case _:
        raise ValueError(f"Unsupported reduction: {reduction}")
    np.testing.assert_array_equal(red_y, expected)

  @parameterized.parameters(
      tcgen05.scales_layout(),
      tcgen05.b_scales_m64_collective_layout(),
  )
  def test_load_store_tmem_scales_layout(self, tmem_layout):
    shape = (tmem_layout.base_tile_shape[0], 32)
    reg_layout = tmem_layout.as_tiled_layout()
    def kernel(ctx, input, output, tmem):
      del ctx
      reg = fa.FragmentedArray.load_untiled(
          input, layout=reg_layout, optimized=False
      )
      tmem.store(reg)
      tcgen05.commit_tmem()
      tmem.load(reg_layout).store_untiled(output, optimized=False)

    x = self.prng.uniform(-1, 1, shape).astype(jnp.float8_e5m2)
    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        x,
        x,
        mgpu.TMEM(x.shape, x.dtype, layout=tmem_layout),
    )(x)
    np.testing.assert_array_equal(x, y)

  def test_store_tmem_sparse_meta_layout_raises(self):
    shape = (128, 64)
    tmem_layout = tcgen05.sparse_meta_layout()
    reg_layout = tmem_layout.as_tiled_layout()
    def kernel(ctx, input, output, tmem):
      del ctx
      reg = fa.FragmentedArray.load_untiled(
          input, layout=reg_layout, optimized=False, is_signed=False
      )
      tmem.store(reg)
      # Still calling load on the output, to prevent potential jaxpr DCE.
      tcgen05.commit_tmem()
      tmem.load(reg_layout, is_signed=False).store_untiled(output, optimized=False)

    x = jnp.ones(shape, jnp.uint2)
    with self.assertRaisesRegex(
      NotImplementedError, "Sparse meta layout stores unsupported."
    ):
      mgpu.as_gpu_kernel(
          kernel,
          (1, 1, 1),
          (128, 1, 1),
          x,
          x,
          mgpu.TMEM(x.shape, x.dtype, layout=tmem_layout),
      )(x)

  def test_load_tmem_sparse_meta_layout_raises(self):
    shape = (128, 64)
    tmem_layout = tcgen05.sparse_meta_layout()
    reg_layout = tmem_layout.as_tiled_layout()
    def kernel(ctx, input, output, tmem):
      del ctx
      tmem.load(reg_layout, is_signed=False).store_untiled(output, optimized=False)

    x = jnp.ones(shape, jnp.uint2)
    with self.assertRaisesRegex(
      NotImplementedError, "Sparse meta layout loads unsupported."
    ):
      mgpu.as_gpu_kernel(
          kernel,
          (1, 1, 1),
          (128, 1, 1),
          x,
          x,
          mgpu.TMEM(x.shape, x.dtype, layout=tmem_layout),
      )(x)

  @parameterized.parameters(
      (jnp.float32, tcgen05.LAYOUT, tcgen05.tmem_default_layout(1)),
      (jnp.float16, tcgen05.LAYOUT, tcgen05.tmem_default_layout(1)),
      (jnp.float16, tcgen05.LAYOUT, tcgen05.tmem_default_layout(2)),
      (
        jnp.float16,
        None,
        tcgen05.tmem_half_lane_layout(columns=32, packing=2),
      ),
      (
        jnp.float16,
        None,
        tcgen05.tmem_m64_collective_layout(columns=32, packing=2),
      ),
      (
          jnp.float8_e5m2,
          None,
          tcgen05.scales_layout(),
      ),
      (
          jnp.float8_e5m2,
          None,
          tcgen05.b_scales_m64_collective_layout(),
      ),
  )
  def test_tmem_column_slicing(self, dtype, reg_layout, tmem_layout):
    shape = (tmem_layout.base_tile_shape[0], 64)
    slicing = (slice(None), slice(32, 64))
    if reg_layout is None:
      reg_layout = tmem_layout.as_tiled_layout()

    def kernel(ctx, input, output, tmem):
      del ctx
      reg = fa.FragmentedArray.load_untiled(
          input, layout=reg_layout, optimized=False
      )
      tmem.store(reg)
      tcgen05.commit_tmem()
      tmem_slice = tmem.slice(*slicing)
      tmem_slice.load(reg_layout).store_untiled(output, optimized=False)

    x = self.prng.uniform(-1, 1, shape).astype(dtype)
    expected = x[*slicing]
    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        x,
        expected,
        mgpu.TMEM(x.shape, x.dtype, layout=tmem_layout),
    )(x)
    np.testing.assert_array_equal(expected, y)

  def test_tmem_column_slicing_custom_layout(self):
    # We test the case where tile number of rows is not a divisor of TMEM_ROWS.
    shape, dtype = (26, 512), jnp.float32
    slicing = (slice(None), slice(256, 512))
    packing = 1
    tmem_layout = tcgen05.TMEMLayout(
      fa.Tiling(((26, packing * 128), (fa.WARP_SIZE * packing,), (packing,))),
      warp_dims=(-3,),
      lane_dims=(-2,),
      vector_dim=-1,
    )
    reg_layout = tmem_layout.as_tiled_layout()

    def kernel(ctx, input, output, tmem):
      del ctx
      reg = fa.FragmentedArray.load_untiled(
          input, layout=reg_layout, optimized=False
      )
      tmem.store(reg)
      tcgen05.commit_tmem()
      tmem_slice = tmem.slice(*slicing)
      tmem_slice.load(reg_layout).store_untiled(output, optimized=False)

    x = self.prng.uniform(-1, 1, shape).astype(dtype)
    expected = x[*slicing]
    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        x,
        expected,
        mgpu.TMEM(x.shape, x.dtype, layout=tmem_layout),
    )(x)
    np.testing.assert_array_equal(expected, y)

  def test_tmem_load_single_register(self):
    shape, dtype = (128, 1), jnp.float32
    tmem_layout = tcgen05.tmem_default_layout(packing=1)
    reg_layout = tmem_layout.as_tiled_layout()

    def kernel(ctx, input, output, tmem):
      del ctx
      reg = fa.FragmentedArray.load_untiled(
          input, layout=reg_layout, optimized=False
      )
      tmem.store(reg)
      tcgen05.commit_tmem()
      tmem.load(reg_layout).store_untiled(output, optimized=False)

    x = self.prng.uniform(-1, 1, shape).astype(dtype)
    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        x,
        x,
        mgpu.TMEM(x.shape, x.dtype, layout=tmem_layout),
    )(x)
    np.testing.assert_array_equal(x, y)

  def test_mixed_tmem_allocations_raise(self):
    def body(ctx, out, scratch):
      del ctx, out, scratch

    with self.assertRaisesRegex(
        ValueError,
        "Can't mix collective and non-collective TMEM allocations within the"
        " same kernel.",
    ):
      mgpu.as_gpu_kernel(
          body,
          grid=(1, 1, 1),
          block=(128, 1, 1),
          in_shape=(),
          out_shape=(jax.ShapeDtypeStruct((), jnp.int32),),
          smem_scratch_shape=[
              mgpu.TMEM((128, 128), jnp.float16, collective=True),
              mgpu.TMEM((128, 128), jnp.float16, collective=False),
          ],
      )

  @parameterized.parameters([
      (jnp.float32, 1, "130.0000"),
      (jnp.float16, 1, "130.0000"),
      (jnp.float16, 2, "[132.000000,133.000000]"),
  ])
  @jtu.thread_unsafe_test()
  def test_tmem_debug_print(self, jax_dtype, packing, expected):
    def kernel(ctx, input, output, tmem):
      del ctx, output
      tmem.store(fa.FragmentedArray.load_untiled(input, layout=tcgen05.LAYOUT, optimized=False))
      tcgen05.commit_tmem()
      tmem.slice(slice(None), slice(0, 8))._debug_print()

    x = jnp.arange(128 * 128, dtype=jax_dtype).reshape(128, 128)
    with self.capture_stdout() as stdout:
      mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x, x, mgpu.TMEM(x.shape, jax_dtype, packing=packing),
      )(x).block_until_ready()
    self.assertIn("[1, 2]: " + expected, stdout())

  @parameterized.product(
      lhs_transpose=(False, True),
      rhs_transpose=(False, True),
      in_jax_dtype=(jnp.float16, jnp.bfloat16, jnp.float8_e5m2, jnp.float8_e4m3fn),  # TODO(apaszke): f32
      out_jax_dtype=(jnp.float16, jnp.float32,),
      m=(64, 128,),  # TODO(apaszke): 64, 192, 256
      n=(64, 128, 192, 224, 256, 512),
      swizzle=(32, 64, 128,),
  )
  def test_mma_basic_float(self, **kwargs):
    in_bytewidth = jnp.dtype(kwargs["in_jax_dtype"]).itemsize
    lhs_transpose = kwargs["lhs_transpose"]
    swizzle = kwargs["swizzle"]
    if lhs_transpose and kwargs["m"] * in_bytewidth < swizzle:
      self.skipTest("swizzle too large for input (lhs)")
    n_steps = 2 if kwargs["m"] == 64 else 1
    n_instr_size = kwargs["n"] * in_bytewidth // n_steps
    if n_instr_size < swizzle or n_instr_size % swizzle != 0:
      self.skipTest("swizzle doesn't work with this instruction size")
    if dtypes.itemsize_bits(kwargs["in_jax_dtype"]) <= 8 and kwargs["n"] == swizzle:
      self.skipTest("Only 8-bit and larger inputs are supported for MMA")
    self._basic_mma_test(
        **kwargs,
        k_steps=2,  # Reducing to 1 can be helpful while debugging.
        lhs_transpose_tiles=False,
        rhs_transpose_tiles=False,
    )

  @parameterized.product(
      lhs_transpose=(False, True),
      rhs_transpose=(False, True),
      in_jax_dtype=(jnp.int8,),
      out_jax_dtype=(jnp.int32,),
      m=(64, 128,),  # TODO(apaszke): 192, 256
      n=(64, 128, 160, 192, 256, 512),
      swizzle=(32, 64, 128,),
  )
  def test_mma_basic_int(self, **kwargs):
    self.skip_unless_tcgen05_int8()
    in_bytewidth = jnp.dtype(kwargs["in_jax_dtype"]).itemsize
    lhs_transpose = kwargs["lhs_transpose"]
    swizzle = kwargs["swizzle"]
    if lhs_transpose and kwargs["m"] * in_bytewidth < swizzle:
      self.skipTest("swizzle too large for input (lhs)")
    n_steps = 2 if kwargs["m"] == 64 else 1
    n_instr_size = kwargs["n"] * in_bytewidth // n_steps
    if n_instr_size < swizzle or n_instr_size % swizzle != 0:
      self.skipTest("swizzle doesn't work with this instruction size")
    if dtypes.itemsize_bits(kwargs["in_jax_dtype"]) <= 8 and kwargs["n"] == swizzle:
      self.skipTest("Only 8-bit and larger inputs are supported for MMA")
    self._basic_mma_test(
        **kwargs,
        k_steps=2,  # Reducing to 1 can be helpful while debugging.
        lhs_transpose_tiles=False,
        rhs_transpose_tiles=False,
    )

  @parameterized.product(
      lhs_transpose=(False, True),
      rhs_transpose=(False, True),
      in_jax_dtype=(jnp.float16,),
      out_jax_dtype=(jnp.float32,),
      m=(128,),
      n=(128, 512),
      swizzle=(32, 64, 128,),
      lhs_transpose_tiles=(False, True),
      rhs_transpose_tiles=(False, True),
  )
  def test_mma_transposed_tiles(self, **kwargs):
    if not kwargs["lhs_transpose_tiles"] and not kwargs["rhs_transpose_tiles"]:
      self.skipTest("This is already tested in test_mma_basic")
    self._basic_mma_test(
        **kwargs,
        k_steps=2,  # Reducing to 1 can be helpful while debugging.
    )

  @parameterized.product(
      lhs_transpose=(False, True),
      rhs_transpose=(False, True),
      m=(64, 128,),
      n=(128, 256, 512),
      lhs_swizzle=(32, 64, 128,),
      rhs_swizzle=(32, 64, 128,),
  )
  def test_mma_different_swizzle(self, **kwargs):
    if kwargs["lhs_swizzle"] == kwargs["rhs_swizzle"]:
      self.skipTest("Swizzle is equal")
    self._basic_mma_test(
        in_jax_dtype=jnp.float16,
        out_jax_dtype=jnp.float32,
        swizzle=None,
        k_steps=2,  # Reducing to 1 can be helpful while debugging.
        **kwargs,
    )

  def _basic_mma_test(
      self,
      m,
      n,
      k_steps,
      swizzle,
      lhs_transpose,
      rhs_transpose,
      in_jax_dtype,
      out_jax_dtype,
      rhs_transpose_tiles=False,
      lhs_transpose_tiles=False,
      lhs_swizzle=None,
      rhs_swizzle=None,
  ):
    if lhs_swizzle is None:
      lhs_swizzle = swizzle
    if rhs_swizzle is None:
      rhs_swizzle = swizzle
    swizzle = max(lhs_swizzle, rhs_swizzle)
    if out_jax_dtype != jnp.float32 and (
        in_jax_dtype == jnp.float32 or in_jax_dtype == jnp.bfloat16
    ):
      self.skipTest("Only f32 output is supported for f32 and bf16 input.")

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    swizzle_elems = swizzle // bytewidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    lhs_tiling = (8, lhs_swizzle // bytewidth(in_mlir_dtype))
    rhs_tiling = (8, rhs_swizzle // bytewidth(in_mlir_dtype))

    def kernel(ctx, lhs, rhs, out, scratch):
      lhs_smem, rhs_smem, barriers, mma_barrier, acc = scratch
      lhs_transform = (mgpu.TileTransform(lhs_tiling),)
      if lhs_transpose_tiles:
        lhs_transform += (mgpu.TransposeTransform((1, 0, 2, 3)),)
      rhs_transform = (mgpu.TileTransform(rhs_tiling),)
      if rhs_transpose_tiles:
        rhs_transform += (mgpu.TransposeTransform((1, 0, 2, 3)),)
      ctx.async_copy(
          src_ref=lhs,
          dst_ref=lhs_smem,
          swizzle=lhs_swizzle,
          gmem_transform=lhs_transform,
          barrier=barriers[0],
      )
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          swizzle=rhs_swizzle,
          gmem_transform=rhs_transform,
          barrier=barriers[1],
      )
      barriers[0].wait()
      barriers[1].wait()
      with mgpu.single_thread():
        if lhs_transpose_tiles:
          lhs_smem = memref_transpose(lhs_smem, (1, 0, 2, 3))
        if lhs_transpose:
          lhs_smem = memref_transpose(lhs_smem, (1, 0, 3, 2))
        if rhs_transpose_tiles:
          rhs_smem = memref_transpose(rhs_smem, (1, 0, 2, 3))
        if rhs_transpose:
          rhs_smem = memref_transpose(rhs_smem, (1, 0, 3, 2))
        tcgen05.mma(
            acc, lhs_smem, rhs_smem, a_swizzle=lhs_swizzle, b_swizzle=rhs_swizzle, accumulate=False,
        )
        tcgen05.commit_arrive(mma_barrier)
      mma_barrier.wait(orders_tensor_core=True)
      is_signed = True if jnp.issubdtype(in_jax_dtype, jnp.integer) else None
      acc.load(is_signed=is_signed).store_untiled(out, optimized=False)

    x_shape = (k, m) if lhs_transpose else (m, k)
    x = self.prng.uniform(-1, 1, x_shape).astype(in_jax_dtype)
    y_shape = (n, k) if rhs_transpose else (k, n)
    y = self.prng.uniform(-1, 1, y_shape).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    if y_shape[0] % rhs_tiling[0] != 0 or y_shape[1] % rhs_tiling[1] != 0:
      self.skipTest("rhs tiling must divide y_shape")
    rhs_smem_shape = tile_shape(y_shape, rhs_tiling)
    if rhs_transpose_tiles:
      rhs_smem_shape = (
          rhs_smem_shape[1], rhs_smem_shape[0], *rhs_smem_shape[2:]
      )
    if x_shape[0] % lhs_tiling[0] != 0 or x_shape[1] % lhs_tiling[1] != 0:
      self.skipTest("lhs tiling must divide x_shape")
    lhs_smem_shape = tile_shape(x_shape, lhs_tiling)
    if lhs_transpose_tiles:
      lhs_smem_shape = (
          lhs_smem_shape[1], lhs_smem_shape[0], *lhs_smem_shape[2:]
      )
    scratch_shape = [
        jax.ShapeDtypeStruct(lhs_smem_shape, in_jax_dtype),
        jax.ShapeDtypeStruct(rhs_smem_shape, in_jax_dtype),
        mgpu.Barrier(1, num_barriers=2),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m, n), out_jax_dtype),
    ]
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (x, y), out_shape, scratch_shape
    )(x, y)
    x32, y32 = x.astype(np.float32), y.astype(np.float32)
    ref = jnp.matmul(
        x32.T if lhs_transpose else x32,
        y32.T if rhs_transpose else y32,
        precision="high",
    )
    atol = 2e-2 if out_jax_dtype == jnp.float16 else 2e-5
    rtol = 8e-4 if out_jax_dtype == jnp.float16 else 1e-7
    np.testing.assert_allclose(z, ref, atol=atol, rtol=rtol)

  @parameterized.product(
      in_jax_dtype=(jnp.float16, jnp.bfloat16),  # TODO(apaszke): f32
      out_jax_dtype=(jnp.float16, jnp.float32,),
      m=(128,),  # TODO(apaszke): 64, 192, 256
      n=(64, 160, 128, 256),
  )
  def test_mma_lhs_tmem_float(self, m, n, in_jax_dtype, out_jax_dtype):
    self._basic_mma_lhs_tmem_test(
        m, n, in_jax_dtype, out_jax_dtype, tcgen05.LAYOUT, swizzle=128
    )

  @parameterized.product(
      in_jax_dtype=(jnp.int8, jnp.uint8),
      out_jax_dtype=(jnp.int32,),
      m=(128,),
      n=(64, 128, 256),
  )
  def test_mma_lhs_tmem_integer(self, m, n, in_jax_dtype, out_jax_dtype):
    self.skip_unless_tcgen05_int8()
    self._basic_mma_lhs_tmem_test(
        m, n, in_jax_dtype, out_jax_dtype, fa.tmem_native_layout(vector_length=4),
        swizzle=math.gcd(n, 128)
    )

  def _basic_mma_lhs_tmem_test(
      self, m, n, in_jax_dtype, out_jax_dtype, lhs_layout, swizzle
  ):
    k_steps = 2  # Reducing to 1 can be helpful while debugging.
    if out_jax_dtype == jnp.float16 and in_jax_dtype != jnp.float16:
      self.skipTest("Only f16 input is supported for f16 output.")

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    swizzle_elems = swizzle // bytewidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    rhs_tiling = (8, swizzle_elems)

    def kernel(ctx, lhs, rhs, out, scratch):
      rhs_smem, barrier, mma_barrier, acc, lhs_tmem = scratch
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(rhs_tiling),
          barrier=barrier,
      )
      barrier.wait()
      if jnp.issubdtype(in_jax_dtype, jnp.integer):
        is_signed = jnp.issubdtype(in_jax_dtype, jnp.signedinteger)
      else:
        is_signed = None
      lhs_tmem.store(
          fa.FragmentedArray.load_untiled(
              lhs, layout=lhs_layout, is_signed=is_signed, optimized=False
          )
      )
      tcgen05.commit_tmem()
      with mgpu.single_thread():
        tcgen05.mma(
            acc, lhs_tmem, rhs_smem, a_swizzle=swizzle, b_swizzle=swizzle, accumulate=False,
        )
        tcgen05.commit_arrive(mma_barrier)
      mma_barrier.wait(orders_tensor_core=True)
      acc.load(is_signed=is_signed).store_untiled(out, optimized=False)

    x_shape = (m, k)
    x = self.prng.uniform(-1, 1, x_shape).astype(in_jax_dtype)
    y_shape = (k, n)
    y = self.prng.uniform(-1, 1, y_shape).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    if y_shape[0] % rhs_tiling[0] != 0 or y_shape[1] % rhs_tiling[1] != 0:
      self.skipTest("rhs tiling must divide y_shape")
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape(y_shape, rhs_tiling), in_jax_dtype),
        mgpu.Barrier(1),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((128, n), out_jax_dtype),
        mgpu.TMEM((128, k), in_jax_dtype, packing=4 // bytewidth(in_mlir_dtype)),
    ]
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (x, y), out_shape, scratch_shape
    )(x, y)
    x32, y32 = x.astype(np.float32), y.astype(np.float32)
    ref = jnp.matmul(x32, y32, precision="high")
    atol = 2e-2 if out_jax_dtype == jnp.float16 else 2e-5
    rtol = 8e-4 if out_jax_dtype == jnp.float16 else 1e-7
    np.testing.assert_allclose(z, ref, atol=atol, rtol=rtol)

  def test_tmem_copy_scales(self):
    dtype = jnp.float8_e8m0fnu

    def kernel(ctx, src, out, scratch):
      smem, barrier, tc_barrier, tmem = scratch
      ctx.async_copy(src_ref=src, dst_ref=smem, barrier=barrier)
      barrier.wait()
      with mgpu.single_thread():
        tcgen05.async_copy_scales_smem_to_tmem(smem, tmem)
        tcgen05.commit_arrive(tc_barrier)
      tc_barrier.wait(orders_tensor_core=True)
      # We print as i32, because i8 seems to overflow the CUDA printf buffer and
      # produce a truncated output.
      tcgen05.TMEMRef(
          tmem.address,
          (128, 4),
          ir.IntegerType.get_signless(32),
          tcgen05.tmem_default_layout(),
      )._debug_print()
      copy(src, out)

    shape = (1, 1, 32, 16)
    x = jax.lax.bitcast_convert_type(
        np.arange(math.prod(shape), dtype=np.uint8).reshape(shape), dtype
    )
    scratch_shape = [
        x,
        mgpu.Barrier(1),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((128, 4), dtype, layout=tcgen05.scales_layout()),
    ]
    with self.capture_stdout() as stdout:
      mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x, x, scratch_shape
      )(x)
    matches = 0
    for l in stdout().splitlines():
      if ":" not in l:
        continue
      idxs, value = l.split(":")
      row, col = map(int, idxs[1:-1].split(","))
      base = (row % 32) * 16 + col * 4
      base %= 256  # int8 has very limited range
      expected = base | (base + 1) << 8 | (base + 2) << 16 | (base + 3) << 24
      self.assertEqual(int(value), expected)
      matches += 1
    self.assertEqual(matches, 128 * 4)

  @parameterized.product(
      jax_dtype=(jnp.float4_e2m1fn, jnp.float8_e5m2, jnp.float16, jnp.float32),
      n=(8, 16, 32, 64, 128 + 32, 256),
      swizzle=(16, 32, 64, 128),
  )
  def test_smem_to_tmem_copy(self, jax_dtype, n, swizzle):
    in_mlir_dtype = utils.dtype_to_ir_type(jax_dtype)
    packing = 32 // bitwidth(in_mlir_dtype)
    swizzle_elems = 8 * swizzle // bitwidth(in_mlir_dtype)
    tiling = (8, swizzle_elems)
    m = 128
    if swizzle_elems > n:
      self.skipTest(f"{n=} too short for {swizzle=} and {jax_dtype=}")
    if n % swizzle_elems:
      self.skipTest(f"{n=} not divisible by {swizzle_elems=}")

    def kernel(ctx, src, out, scratch):
      smem, barrier, copy_barrier, tmem = scratch
      ctx.async_copy(
          src_ref=src,
          dst_ref=smem,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
          barrier=barrier,
      )
      barrier.wait()
      with mgpu.single_thread():
        tcgen05.async_copy_smem_to_tmem(smem, tmem, swizzle)
        tcgen05.commit_arrive(copy_barrier)
      copy_barrier.wait(orders_tensor_core=True)
      tmem.load(fa.tmem_native_layout(packing)).store_untiled(out, optimized=False)

    x = self.prng.uniform(-1, 1, (m, n)).astype(jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), jax_dtype)
    smem_shape = tile_shape((m, n), tiling)
    scratch_shape = [
        jax.ShapeDtypeStruct(smem_shape, jax_dtype),
        mgpu.Barrier(1),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m, n), jax_dtype, packing=packing),
    ]
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, out_shape, scratch_shape
    )(x)
    np.testing.assert_array_equal(x, z)

  @parameterized.product(
      jax_dtype=(jnp.float16,),
      n=(64, 256),
      swizzle=(32, 128),
  )
  def test_smem_to_tmem_copy_collective(self, jax_dtype, n, swizzle):
    index = ir.IndexType.get()
    in_mlir_dtype = utils.dtype_to_ir_type(jax_dtype)
    packing = 32 // bitwidth(in_mlir_dtype)
    swizzle_elems = 8 * swizzle // bitwidth(in_mlir_dtype)
    tiling = (8, swizzle_elems)
    m = 256
    m_block = m // 2

    def kernel(ctx, src, out, scratch):
      smem, barrier, copy_barrier, cluster_barrier, tmem = scratch
      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      is_leader_thread = single_thread_predicate()
      is_first_block = arith.cmpi(
          arith.CmpIPredicate.eq, block_id, c(0, index)
      )
      ctx.async_copy(
          src_ref=src,
          dst_ref=smem,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
          barrier=barrier,
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      with when(is_first_block):
        barrier.wait()
      cluster_barrier.arrive(orders_tensor_core=True)
      cluster_barrier.wait(orders_tensor_core=True)
      with when(arith.andi(is_first_block, is_leader_thread)):
        tcgen05.async_copy_smem_to_tmem(
            smem, tmem, swizzle, collective=True
        )
        tcgen05.commit_arrive(copy_barrier, collective=True, ctx=ctx)
      copy_barrier.wait(orders_tensor_core=True)
      m_slice = ds(arith.muli(block_id, c(m_block, index)), m_block)
      tmem.load(fa.tmem_native_layout(packing)).store_untiled(
          memref_slice(out, m_slice), optimized=False
      )

    x = self.prng.uniform(-1, 1, (m, n)).astype(jax_dtype)
    smem_shape = tile_shape((m_block, n), tiling)
    out_shape = jax.ShapeDtypeStruct((m, n), jax_dtype)
    scratch_shape = [
        jax.ShapeDtypeStruct(smem_shape, jax_dtype),
        mgpu.Barrier(1),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.ClusterBarrier(collective_dims=(gpu.Dimension.x,)),
        mgpu.TMEM((m_block, n), jax_dtype, packing=packing, collective=True),
    ]
    z = mgpu.as_gpu_kernel(
        kernel, (2, 1, 1), (128, 1, 1), x, out_shape, scratch_shape,
        cluster=(2, 1, 1),
    )(x)
    np.testing.assert_array_equal(x, z)

  def _sample_scales(self, m, k, n, block_size, scale_jax_dtype):
    ka, kb = jax.random.split(jax.random.key(1234), 2)
    if scale_jax_dtype == jnp.float8_e8m0fnu:
      a_scales = jax.lax.bitcast_convert_type(
          jax.random.randint(ka, (m, k // block_size), 122, 132, dtype=jnp.uint8),
          scale_jax_dtype
      )
      b_scales = jax.lax.bitcast_convert_type(
          jax.random.randint(kb, (n, k // block_size), 122, 132, dtype=jnp.uint8),
          scale_jax_dtype
      )
    elif scale_jax_dtype == jnp.float8_e4m3fn:
      a_scales = jnp.abs(
          jax.random.normal(ka, (m, k // block_size), dtype=jnp.float32).astype(
              scale_jax_dtype
          )
      )
      b_scales = jnp.abs(
          jax.random.normal(kb, (n, k // block_size), dtype=jnp.float32).astype(
              scale_jax_dtype
          )
      )
    else:
      raise ValueError(f"Unsupported scale dtype: {scale_jax_dtype}")
    return a_scales, b_scales

  @parameterized.product(
      in_jax_dtype=(jnp.float8_e5m2, jnp.float8_e4m3fn, jnp.float4_e2m1fn),
      scale_jax_dtype=(jnp.float8_e8m0fnu, jnp.float8_e4m3fn),
      block_size=(16, 32),
      m=(128,),  # TODO(apaszke): 256
      n=(8, 32, 64, 128, 192, 256),
      swizzle=(32, 128),
      lhs_in_tmem=(False, True),
  )
  def test_mma_block_scaled_basic(
      self,
      m,
      n,
      in_jax_dtype,
      scale_jax_dtype,
      block_size,
      swizzle,
      lhs_in_tmem,
  ):
    out_jax_dtype = jnp.float32
    # When swizzle is small, we need to take many steps to make it large enough
    # to make the scale count a multiple of 4.
    k_steps = 4 if swizzle == 32 else 2
    if block_size == 16 and in_jax_dtype != jnp.float4_e2m1fn:
      self.skipTest("Only float4_e2m1fn input is supported for block size 16.")
    if scale_jax_dtype == jnp.float8_e4m3fn and block_size != 16:
      self.skipTest("e4m3fn scale only supports block size 16.")
    if out_jax_dtype == jnp.float16 and in_jax_dtype != jnp.float16:
      self.skipTest("Only f16 input is supported for f16 output.")

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    swizzle_elems = 8 * swizzle // bitwidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    lhs_tiling = rhs_tiling = (8, swizzle_elems)

    def kernel(ctx, lhs, rhs, lhs_scales_gmem, rhs_scales_gmem, out, scratch):
      (
          lhs_smem,
          rhs_smem,
          lhs_scales_smem,
          rhs_scales_smem,
          barriers,
          mma_barrier,
          acc,
          lhs_scales,
          rhs_scales,
          lhs_tmem,
      ) = scratch
      operand_kwargs = dict(
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(lhs_tiling),
      )
      ctx.async_copy(src_ref=lhs, dst_ref=lhs_smem, barrier=barriers[0], **operand_kwargs)
      ctx.async_copy(src_ref=rhs, dst_ref=rhs_smem, barrier=barriers[1], **operand_kwargs)
      ctx.async_copy(src_ref=lhs_scales_gmem, dst_ref=lhs_scales_smem, barrier=barriers[2])
      ctx.async_copy(src_ref=rhs_scales_gmem, dst_ref=rhs_scales_smem, barrier=barriers[3])
      for i in range(4):
        barriers[i].wait()
      with mgpu.single_thread():
        if lhs_in_tmem:
          tcgen05.async_copy_smem_to_tmem(lhs_smem, lhs_tmem, swizzle=swizzle)
        tcgen05.async_copy_scales_smem_to_tmem(lhs_scales_smem, lhs_scales)
        tcgen05.async_copy_scales_smem_to_tmem(rhs_scales_smem, rhs_scales)
        tcgen05.mma(
            acc,
            lhs_tmem if lhs_in_tmem else lhs_smem,
            mgpu.memref_transpose(rhs_smem, (1, 0, 3, 2)),
            a_swizzle=swizzle,
            b_swizzle=swizzle,
            a_scale=lhs_scales,
            b_scale=rhs_scales,
            accumulate=False,
        )
        tcgen05.commit_arrive(mma_barrier)
      mma_barrier.wait(orders_tensor_core=True)
      acc.load().store_untiled(out, optimized=False)

    x_shape = (m, k)
    x = self.prng.uniform(-1, 1, x_shape).astype(in_jax_dtype)
    y_shape = (n, k)
    y = self.prng.uniform(-1, 1, y_shape).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape(x_shape, lhs_tiling), in_jax_dtype),
        jax.ShapeDtypeStruct(tile_shape(y_shape, rhs_tiling), in_jax_dtype),
        jax.ShapeDtypeStruct(
            (m // 128, k // (block_size * 4), 32, 16), scale_jax_dtype
        ),
        jax.ShapeDtypeStruct(
            ((n + 127) // 128, k // (block_size * 4), 32, 16), scale_jax_dtype
        ),
        mgpu.Barrier(1, num_barriers=4),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m, n), out_jax_dtype),
        mgpu.TMEM(
            (m, k // block_size),
            scale_jax_dtype,
            layout=tcgen05.scales_layout(),
        ),
        mgpu.TMEM(
            ((n + 127) // 128 * 128, k // block_size),
            scale_jax_dtype,
            layout=tcgen05.scales_layout(),
        ),
        mgpu.TMEM((m, k), in_jax_dtype, packing=32 // bitwidth(in_mlir_dtype))
        if lhs_in_tmem
        else None,
    ]
    a_scales, b_scales = self._sample_scales(m, k, n, block_size, scale_jax_dtype)
    def format_scales(scales):
      mn, k = scales.shape
      assert mn % 8 == 0 and k % 4 == 0, scales.shape
      pad_mn = (mn + 127) // 128 * 128
      scales = jnp.pad(scales, ((0, pad_mn - mn), (0, 0)))
      return (
          scales.reshape(pad_mn // 128, 4, 32, k // 4, 4)
          .transpose(0, 3, 2, 1, 4)
          .reshape(pad_mn // 128, k // 4, 32, 16)
      )
    a_gpu_scales, b_gpu_scales = map(format_scales, (a_scales, b_scales))
    args = (x, y, a_gpu_scales, b_gpu_scales)
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), args, out_shape, scratch_shape
    )(*args)
    x32, y32 = x.astype(np.float32), y.astype(np.float32)
    a_logical_scales = jnp.repeat(a_scales, block_size, axis=1).astype(jnp.float32)
    b_logical_scales = jnp.repeat(b_scales, block_size, axis=1).astype(jnp.float32)
    ref = jnp.matmul(
        x32 * a_logical_scales,
        (y32 * b_logical_scales).T,
        precision="high",
    )
    np.testing.assert_allclose(z, ref, atol=2e-4, rtol=5e-6)

  @parameterized.product(
    m=(128,),
    n=(64, 128, 192, 256),
    scale_jax_dtype=(jnp.float8_e8m0fnu,),
  )
  def test_mma_block_scaled_collective_m128(self, m, n, scale_jax_dtype):
    in_jax_dtype = jnp.float8_e5m2
    out_jax_dtype = jnp.float32
    scale_block = 32
    swizzle = 128
    k_steps = 2

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    swizzle_elems = 8 * swizzle // bitwidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    lhs_tiling = rhs_tiling = (8, swizzle_elems)

    def kernel(ctx, lhs, rhs, lhs_scales_gmem, rhs_scales_gmem, out, scratch):
      (
          lhs_smem, rhs_smem, lhs_scales_smem, rhs_scales_smem,
          barriers, mma_barrier, acc, lhs_scales, rhs_scales
      ) = scratch
      ctx.async_copy(
          src_ref=lhs,
          dst_ref=lhs_smem,
          barrier=barriers[0],
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(lhs_tiling),
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          barrier=barriers[1],
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(rhs_tiling),
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      ctx.async_copy(
          src_ref=lhs_scales_gmem,
          dst_ref=lhs_scales_smem,
          barrier=barriers[2],
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      ctx.async_copy(
          src_ref=rhs_scales_gmem,
          dst_ref=rhs_scales_smem,
          barrier=barriers[3],
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.REPLICATED,
      )

      is_leader_thread = single_thread_predicate()
      index = ir.IndexType.get()
      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      is_first_block = arith.cmpi(arith.CmpIPredicate.eq, block_id, c(0, index))
      with when(arith.andi(is_first_block, is_leader_thread)):
        for i in range(4):
          barriers[i].wait()
        tcgen05.async_copy_scales_smem_to_tmem(lhs_scales_smem, lhs_scales, collective=True)
        tcgen05.async_copy_scales_smem_to_tmem(rhs_scales_smem, rhs_scales, collective=True)
        tcgen05.mma(
            acc,
            lhs_smem,
            mgpu.memref_transpose(rhs_smem, (1, 0, 3, 2)),
            a_swizzle=swizzle,
            b_swizzle=swizzle,
            accumulate=False,
            collective=True,
            a_scale=lhs_scales,
            b_scale=rhs_scales,
        )
        tcgen05.commit_arrive(mma_barrier, collective=True, ctx=ctx)
      mma_barrier.wait(orders_tensor_core=True)

      m_block_tile = m // 2
      m_slice = ds(arith.muli(block_id, c(m_block_tile, index)), m_block_tile)
      acc.load().store_untiled(memref_slice(out, m_slice), optimized=False)

    x_shape = (m, k)
    x = self.prng.uniform(-1, 1, x_shape).astype(in_jax_dtype)
    y_shape = (n, k)
    y = self.prng.uniform(-1, 1, y_shape).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)

    m_block = m // 2

    def format_a_scales(scales):
      mn, kk = scales.shape
      assert mn % 32 == 0 and kk % 4 == 0, scales.shape
      scales = scales.reshape(2, m_block, kk)
      # Each CTA only has m_block=64, so groups 2 and 3 are padded for the
      # async copy to TMEM.
      scales = jnp.pad(scales, ((0, 0), (0, 128 - m_block), (0, 0)))
      return (
          scales.reshape(2, 4, 32, kk // 4, 4)
          .transpose(0, 3, 2, 1, 4)
          .reshape(2, kk // 4, 32, 16)
      )

    def format_b_scales(scales):
      # Each column holds scale values for 32 N columns. The TMEM layout
      # requires us to split N into two chunks, and pad each chunk to
      # 64x16 so that it occupies exactly 4 TMEM columns for async copy.
      mn, kk = scales.shape
      assert mn % 64 == 0 and kk % 4 == 0, scales.shape
      columns_per_cta = mn // 64
      return jnp.pad(
          scales.reshape(2, columns_per_cta, 32, kk // 4, 4)
          .transpose(3, 0, 2, 1, 4)
          .reshape(1, kk // 4, 64, columns_per_cta * 4),
          ((0, 0), (0, 0), (0, 0), (0, 16 - columns_per_cta * 4)),
      )

    n_block = n // 2

    scratch_shape = [
        jax.ShapeDtypeStruct(
            tile_shape((m_block, k), lhs_tiling), in_jax_dtype
        ),
        jax.ShapeDtypeStruct(
            tile_shape((n_block, k), rhs_tiling), in_jax_dtype
        ),
        jax.ShapeDtypeStruct(
            ((m_block + 127) // 128, k // (scale_block * 4), 32, 16), scale_jax_dtype
        ),
        jax.ShapeDtypeStruct(
            (1, k // (scale_block * 4), 64, 16), scale_jax_dtype
        ),
        mgpu.Barrier(1, num_barriers=4),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m_block, n), out_jax_dtype, collective=True),
        mgpu.TMEM(
            (128, k // scale_block),
            scale_jax_dtype,
            layout=tcgen05.scales_layout(),
            collective=True,
        ),
        mgpu.TMEM(
            (256, k // scale_block),
            scale_jax_dtype,
            layout=tcgen05.b_scales_m64_collective_layout(),
            collective=True,
        ),
    ]

    a_scales, b_scales = self._sample_scales(m, k, n, scale_block, scale_jax_dtype)

    a_gpu_scales = format_a_scales(a_scales)
    b_gpu_scales = format_b_scales(b_scales)
    args = (x, y, a_gpu_scales, b_gpu_scales)
    z = mgpu.as_gpu_kernel(
        kernel, (2, 1, 1), (128, 1, 1), args, out_shape, scratch_shape, cluster=(2, 1, 1),
    )(*args)

    x32, y32 = x.astype(np.float32), y.astype(np.float32)
    a_logical_scales = jnp.repeat(a_scales, scale_block, axis=1).astype(jnp.float32)
    b_logical_scales = jnp.repeat(b_scales, scale_block, axis=1).astype(jnp.float32)
    ref = jnp.matmul(
        x32 * a_logical_scales,
        (y32 * b_logical_scales).T,
        precision="high",
    )
    np.testing.assert_allclose(z, ref, atol=2e-4, rtol=5e-6)

  @parameterized.product(
      m=(256,),
      n=(32, 64, 128, 192, 256),
      scale_jax_dtype=(jnp.float8_e8m0fnu, jnp.float8_e4m3fn),
      block_size=(16, 32),
      lhs_in_tmem=(False, True),
  )
  def test_mma_block_scaled_collective(
      self, m, n, scale_jax_dtype, block_size, lhs_in_tmem
  ):
    if scale_jax_dtype == jnp.float8_e4m3fn and block_size != 16:
      self.skipTest("e4m3fn scale only supports block size 16.")
    in_jax_dtype = jnp.float4_e2m1fn
    out_jax_dtype = jnp.float32
    scale_block = block_size
    swizzle = 128
    k_steps = 2

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    swizzle_elems = 8 * swizzle // bitwidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    lhs_tiling = rhs_tiling = (8, swizzle_elems)

    def kernel(ctx, lhs, rhs, lhs_scales_gmem, rhs_scales_gmem, out, scratch):
      (
          lhs_smem,
          rhs_smem,
          lhs_scales_smem,
          rhs_scales_smem,
          barriers,
          mma_barrier,
          acc,
          lhs_scales,
          rhs_scales,
          lhs_tmem,
      ) = scratch
      ctx.async_copy(
          src_ref=lhs,
          dst_ref=lhs_smem,
          barrier=barriers[0],
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(lhs_tiling),
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          barrier=barriers[1],
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(rhs_tiling),
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      ctx.async_copy(
          src_ref=lhs_scales_gmem,
          dst_ref=lhs_scales_smem,
          barrier=barriers[2],
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      # B scales are replicated! Note that this does not use 2CTA TMA and will
      # need to be awaited in the non-leader CTA or else we will double arrive.
      ctx.async_copy(
          src_ref=rhs_scales_gmem,
          dst_ref=rhs_scales_smem,
          barrier=barriers[3],
          collective=gpu.Dimension.x,
      )

      is_leader_thread = single_thread_predicate()
      index = ir.IndexType.get()
      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      is_first_block = arith.cmpi(arith.CmpIPredicate.eq, block_id, c(0, index))
      with when(arith.andi(is_first_block, is_leader_thread)):
        for i in range(4):
          barriers[i].wait()
        if lhs_in_tmem:
          tcgen05.async_copy_smem_to_tmem(
              lhs_smem, lhs_tmem, swizzle=swizzle, collective=True
          )
        tcgen05.async_copy_scales_smem_to_tmem(lhs_scales_smem, lhs_scales, collective=True)
        tcgen05.async_copy_scales_smem_to_tmem(rhs_scales_smem, rhs_scales, collective=True)
        tcgen05.mma(
            acc,
            lhs_tmem if lhs_in_tmem else lhs_smem,
            mgpu.memref_transpose(rhs_smem, (1, 0, 3, 2)),
            a_swizzle=swizzle,
            b_swizzle=swizzle,
            a_scale=lhs_scales,
            b_scale=rhs_scales,
            accumulate=False,
            collective=True,
        )
        tcgen05.commit_arrive(mma_barrier, collective=True, ctx=ctx)
      mma_barrier.wait(orders_tensor_core=True)
      m_block_tile = m // 2
      m_slice = ds(arith.muli(block_id, c(m_block_tile, index)), m_block_tile)
      acc.load().store_untiled(memref_slice(out, m_slice), optimized=False)

    x_shape = (m, k)
    x = self.prng.uniform(-1, 1, x_shape).astype(in_jax_dtype)
    y_shape = (n, k)
    y = self.prng.uniform(-1, 1, y_shape).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)

    m_block = m // 2
    n_block = n // 2

    scratch_shape = [
        jax.ShapeDtypeStruct(
            tile_shape((m_block, k), lhs_tiling), in_jax_dtype
        ),
        jax.ShapeDtypeStruct(
            tile_shape((n_block, k), rhs_tiling), in_jax_dtype
        ),
        jax.ShapeDtypeStruct(
            (m_block // 128, k // (scale_block * 4), 32, 16), scale_jax_dtype
        ),
        jax.ShapeDtypeStruct(
            ((n + 127) // 128, k // (scale_block * 4), 32, 16), scale_jax_dtype
        ),
        mgpu.Barrier(1, num_barriers=4),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m_block, n), out_jax_dtype, collective=True),
        mgpu.TMEM(
            (m_block, k // scale_block),
            scale_jax_dtype,
            layout=tcgen05.scales_layout(),
            collective=True,
        ),
        mgpu.TMEM(
            ((n + 127) // 128 * 128, k // scale_block),
            scale_jax_dtype,
            layout=tcgen05.scales_layout(),
            collective=True,
        ),
        (
            mgpu.TMEM(
                (m_block, k),
                in_jax_dtype,
                packing=32 // bitwidth(in_mlir_dtype),
                collective=True,
            )
            if lhs_in_tmem
            else None
        ),
    ]

    a_scales, b_scales = self._sample_scales(m, k, n, scale_block, scale_jax_dtype)

    def format_scales(scales):
      mn, k = scales.shape
      assert mn % 32 == 0 and k % 4 == 0, scales.shape
      pad_mn = (mn + 127) // 128 * 128
      scales = jnp.pad(scales, ((0, pad_mn - mn), (0, 0)))
      return (
          scales.reshape(pad_mn // 128, 4, 32, k // 4, 4)
          .transpose(0, 3, 2, 1, 4)
          .reshape(pad_mn // 128, k // 4, 32, 16)
      )

    a_gpu_scales = format_scales(a_scales)
    b_gpu_scales = format_scales(b_scales)
    args = (x, y, a_gpu_scales, b_gpu_scales)
    z = mgpu.as_gpu_kernel(
        kernel, (2, 1, 1), (128, 1, 1), args, out_shape, scratch_shape, cluster=(2, 1, 1),
    )(*args)

    x32, y32 = x.astype(np.float32), y.astype(np.float32)
    a_logical_scales = jnp.repeat(a_scales, scale_block, axis=1).astype(jnp.float32)
    b_logical_scales = jnp.repeat(b_scales, scale_block, axis=1).astype(jnp.float32)
    ref = jnp.matmul(
        x32 * a_logical_scales,
        (y32 * b_logical_scales).T,
        precision="high",
    )
    np.testing.assert_allclose(z, ref, atol=2e-4, rtol=5e-6)

  @parameterized.product(
      lhs_transpose=(False, True),
      rhs_transpose=(False, True),
      ab_dtype=(
          (jnp.float16, jnp.float16),
          (jnp.bfloat16, jnp.bfloat16),
          (jnp.int8, jnp.int8),
          (jnp.float8_e4m3fn, jnp.float8_e4m3fn),
          (jnp.float8_e5m2, jnp.float8_e5m2),
          (jnp.float8_e4m3fn, jnp.float8_e5m2),
          (jnp.float8_e5m2, jnp.float8_e4m3fn),
      ),
      m=(128,),  # TODO(apaszke): 256
      n=(16, 32, 128, 192, 256),
      lhs_swizzle=(32, 64, 128),
      rhs_swizzle=(64, 128),  # 32 is too small and unsupported.
  )
  def test_mma_sparse(
      self,
      m,
      n,
      ab_dtype,
      lhs_swizzle,
      rhs_swizzle,
      lhs_transpose,
      rhs_transpose,
  ):
    a_dtype, b_dtype = ab_dtype
    if a_dtype == jnp.int8:
      self.skip_unless_tcgen05_int8()
    if jnp.issubdtype(a_dtype, jnp.floating):
      out_jax_dtype = jnp.float32
    else:
      out_jax_dtype = jnp.int32
    sparse_meta_dtype = jnp.uint2

    a_mlir_dtype = utils.dtype_to_ir_type(a_dtype)
    b_mlir_dtype = utils.dtype_to_ir_type(b_dtype)
    k = 256
    lhs_tiling = (8, 8 * lhs_swizzle // bitwidth(a_mlir_dtype))
    rhs_tiling = (8, 8 * rhs_swizzle // bitwidth(b_mlir_dtype))
    if not rhs_transpose and n % rhs_tiling[1]:
      self.skipTest(f"{n=} not divisible by {rhs_swizzle=} (for {b_dtype=})")

    def kernel(ctx, lhs, rhs, lhs_sparse_gmem, out, scratch):
      lhs_smem, rhs_smem, lhs_sparse_smem, barriers, mma_barrier, acc, lhs_sparse = scratch
      ctx.async_copy(src_ref=lhs, dst_ref=lhs_smem, barrier=barriers[0], swizzle=lhs_swizzle, gmem_transform=mgpu.TileTransform(lhs_tiling))
      ctx.async_copy(src_ref=rhs, dst_ref=rhs_smem, barrier=barriers[1], swizzle=rhs_swizzle, gmem_transform=mgpu.TileTransform(rhs_tiling))
      ctx.async_copy(src_ref=lhs_sparse_gmem, dst_ref=lhs_sparse_smem, barrier=barriers[2])
      for i in range(3):
        barriers[i].wait()
      with mgpu.single_thread():
        tcgen05.async_copy_sparse_metadata_smem_to_tmem(lhs_sparse_smem, lhs_sparse)
        if lhs_transpose:
          lhs_smem = mgpu.memref_transpose(lhs_smem, (1, 0, 3, 2))
        if rhs_transpose:
          rhs_smem = mgpu.memref_transpose(rhs_smem, (1, 0, 3, 2))
        tcgen05.mma(
            acc,
            lhs_smem,
            rhs_smem,
            a_swizzle=lhs_swizzle,
            b_swizzle=rhs_swizzle,
            a_sparse_metadata=lhs_sparse,
            accumulate=False,
        )
        tcgen05.commit_arrive(mma_barrier)
      mma_barrier.wait(orders_tensor_core=True)
      is_signed = True if jnp.issubdtype(a_dtype, jnp.integer) else None
      acc.load(is_signed=is_signed).store_untiled(out, optimized=False)

    x_shape = (k // 2, m) if lhs_transpose else (m, k // 2)
    y_shape = (n, k) if rhs_transpose else (k, n)
    if jnp.issubdtype(a_dtype, jnp.integer):
      x = jax.random.randint(jax.random.key(1234), x_shape, -64, 64, dtype=a_dtype)
      y = jax.random.randint(jax.random.key(2567), y_shape, -64, 64, dtype=b_dtype)
    else:
      x = self.prng.uniform(-1, 1, x_shape).astype(a_dtype)
      y = self.prng.uniform(-1, 1, y_shape).astype(b_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape(x_shape, lhs_tiling), a_dtype),
        jax.ShapeDtypeStruct(tile_shape(y_shape, rhs_tiling), b_dtype),
        jax.ShapeDtypeStruct((m // 128, k // 128, 128, 64), sparse_meta_dtype),
        mgpu.Barrier(1, num_barriers=3),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m, n), out_jax_dtype),
        mgpu.TMEM((m, k // 2), sparse_meta_dtype, layout=tcgen05.sparse_meta_layout()),
    ]
    index_pairs = np.asarray(np.meshgrid(range(4), range(4))).T.reshape(-1, 2)
    valid_pairs = index_pairs[index_pairs[:, 0] < index_pairs[:, 1]]
    assert len(valid_pairs) == 6
    x_pairs = jax.random.randint(jax.random.key(1234), (m, k // 4), 0, 6, dtype=jnp.uint8)
    x_sparse = valid_pairs[x_pairs]
    assert x_sparse.shape == (m, k // 4, 2)
    def format_sparse_meta(meta):
      mn, k, _2 = meta.shape
      assert _2 == 2
      k *= 2
      if jnp.dtype(a_dtype).itemsize == 1:
        meta_tiled = (
            meta.reshape(mn // 128, 128, k // 64, 64).transpose(0, 2, 1, 3)
        )
      else:
        meta_tiled = (
          meta.reshape(mn // 128, 8, 2, 8, k // 64, 4, 2, 8)
          .transpose(0, 4, 1, 6, 3, 5, 2, 7)
        )
      return (
          meta_tiled.reshape(mn // 128, k // 64, 128, 64)
          .astype(sparse_meta_dtype)
      )
    x_gpu_sparse = format_sparse_meta(x_sparse)
    args = (x, y, x_gpu_sparse)
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), args, out_shape, scratch_shape
    )(*args)
    if lhs_transpose:
      x = x.T
    if rhs_transpose:
      y = y.T
    x_logical = np.zeros_like(x, shape=(m, k // 4, 4))
    np.put_along_axis(x_logical, x_sparse, x.reshape(x_sparse.shape), axis=-1)
    x_logical = x_logical.reshape(m, k)
    ref = jnp.matmul(
        x_logical.astype(jnp.float32), y.astype(jnp.float32), precision="high"
    )
    atol = 2e-2 if out_jax_dtype == jnp.float16 else 7e-5
    rtol = 8e-4 if out_jax_dtype == jnp.float16 else 5e-6
    np.testing.assert_allclose(z, ref, atol=atol, rtol=rtol)

  @parameterized.product(
      in_jax_dtype=(jnp.float16, jnp.bfloat16),
      m=(128,),  # TODO(apaszke): 256
      n=(128, 192, 256),
      lhs_swizzle=(32, 64, 128),
      rhs_swizzle=(64, 128),  # 32 is too small and unsuported.
  )
  def test_mma_sparse_lhs_tmem(
      self, m, n, in_jax_dtype, lhs_swizzle, rhs_swizzle
  ):
    out_jax_dtype = jnp.float32
    sparse_meta_dtype = jnp.uint2

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    k = 256
    rhs_tiling = (8, 8 * rhs_swizzle // bitwidth(in_mlir_dtype))
    if n < rhs_tiling[1] or n % rhs_tiling[1]:
      self.skipTest(f"{n=} not divisible by {rhs_swizzle=} (for {in_jax_dtype=})")

    def kernel(ctx, lhs, rhs, lhs_sparse_gmem, out, scratch):
      (
          rhs_smem,
          lhs_sparse_smem,
          barriers,
          mma_barrier,
          acc,
          lhs_tmem,
          lhs_sparse,
      ) = scratch
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          barrier=barriers[0],
          swizzle=rhs_swizzle,
          gmem_transform=mgpu.TileTransform(rhs_tiling),
      )
      ctx.async_copy(
          src_ref=lhs_sparse_gmem, dst_ref=lhs_sparse_smem, barrier=barriers[1]
      )
      barriers[0].wait()
      barriers[1].wait()
      lhs_tmem.store(
          fa.FragmentedArray.load_untiled(
              lhs, layout=tcgen05.LAYOUT, optimized=False
          )
      )
      tcgen05.commit_tmem()
      with mgpu.single_thread():
        tcgen05.async_copy_sparse_metadata_smem_to_tmem(
            lhs_sparse_smem, lhs_sparse
        )
        tcgen05.mma(
            acc,
            lhs_tmem,
            rhs_smem,
            a_swizzle=lhs_swizzle,
            b_swizzle=rhs_swizzle,
            a_sparse_metadata=lhs_sparse,
            accumulate=False,
        )
        tcgen05.commit_arrive(mma_barrier)
      mma_barrier.wait(orders_tensor_core=True)
      acc.load().store_untiled(out, optimized=False)

    x_shape = (m, k // 2)
    x = self.prng.uniform(-1, 1, x_shape).astype(in_jax_dtype)
    y_shape = (k, n)
    y = self.prng.uniform(-1, 1, y_shape).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape(y_shape, rhs_tiling), in_jax_dtype),
        jax.ShapeDtypeStruct((m // 128, k // 128, 128, 64), sparse_meta_dtype),
        mgpu.Barrier(1, num_barriers=2),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m, n), out_jax_dtype),
        mgpu.TMEM((m, k // 2), in_jax_dtype, packing=2),
        mgpu.TMEM(
            (m, k // 2), sparse_meta_dtype, layout=tcgen05.sparse_meta_layout()
        ),
    ]
    index_pairs = np.asarray(np.meshgrid(range(4), range(4))).T.reshape(-1, 2)
    valid_pairs = index_pairs[index_pairs[:, 0] < index_pairs[:, 1]]
    assert len(valid_pairs) == 6
    x_pairs = jax.random.randint(
        jax.random.key(1234), (m, k // 4), 0, 6, dtype=jnp.uint8
    )
    x_sparse = valid_pairs[x_pairs]
    assert x_sparse.shape == (m, k // 4, 2)

    def format_sparse_meta(meta):
      mn, k, _2 = meta.shape
      assert _2 == 2
      k *= 2
      return (
          meta.reshape(mn // 128, 8, 2, 8, k // 64, 4, 2, 8)
          .transpose(0, 4, 1, 6, 3, 5, 2, 7)
          .reshape(mn // 128, k // 64, 128, 64)
          .astype(sparse_meta_dtype)
      )

    x_gpu_sparse = format_sparse_meta(x_sparse)
    args = (x, y, x_gpu_sparse)
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), args, out_shape, scratch_shape
    )(*args)
    x_logical = np.zeros_like(x, shape=(m, k // 4, 4))
    np.put_along_axis(x_logical, x_sparse, x.reshape(x_sparse.shape), axis=-1)
    x_logical = x_logical.reshape(m, k)
    ref = jnp.matmul(
        x_logical.astype(jnp.float32), y.astype(jnp.float32), precision="high"
    )
    atol = 2e-2 if out_jax_dtype == jnp.float16 else 7e-5
    rtol = 8e-4 if out_jax_dtype == jnp.float16 else 5e-6
    np.testing.assert_allclose(z, ref, atol=atol, rtol=rtol)

  @parameterized.product(
      in_jax_dtype=(jnp.float16, jnp.float8_e4m3fn),
      m=(256,),  # TODO(apaszke): 256
      n=(128, 192, 256),
      lhs_swizzle=(32, 64, 128),
      rhs_swizzle=(64, 128),  # 32 is too small and unsupported.
  )
  def test_mma_sparse_collective(self, m, n, in_jax_dtype, lhs_swizzle, rhs_swizzle):
    out_jax_dtype = jnp.float32
    sparse_meta_dtype = jnp.uint2

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    k = 256
    lhs_tiling = (8, 8 * lhs_swizzle // bitwidth(in_mlir_dtype))
    rhs_tiling = (8, 8 * rhs_swizzle // bitwidth(in_mlir_dtype))
    if (m // 2) % lhs_tiling[1]:
      self.skipTest("LHS not divisible by swizzle")
    if (n // 2) % rhs_tiling[1]:
      self.skipTest("RHS not divisible by swizzle")

    def kernel(ctx, lhs, rhs, lhs_sparse_gmem, out, scratch):
      lhs_smem, rhs_smem, lhs_sparse_smem, barriers, mma_barrier, acc, lhs_sparse = scratch
      ctx.async_copy(
          src_ref=lhs,
          dst_ref=lhs_smem,
          barrier=barriers[0],
          swizzle=lhs_swizzle,
          gmem_transform=mgpu.TileTransform(lhs_tiling),
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          barrier=barriers[1],
          swizzle=rhs_swizzle,
          gmem_transform=mgpu.TileTransform(rhs_tiling),
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(1),
      )
      ctx.async_copy(
          src_ref=lhs_sparse_gmem,
          dst_ref=lhs_sparse_smem,
          barrier=barriers[2],
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      index = ir.IndexType.get()
      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      is_first_block = arith.cmpi(arith.CmpIPredicate.eq, block_id, c(0, index))
      is_leader_thread = single_thread_predicate()
      with when(arith.andi(is_first_block, is_leader_thread)):
        for i in range(3):
          barriers[i].wait()
        tcgen05.async_copy_sparse_metadata_smem_to_tmem(lhs_sparse_smem, lhs_sparse, collective=True)
        tcgen05.mma(
            acc,
            lhs_smem,
            rhs_smem,
            a_swizzle=lhs_swizzle,
            b_swizzle=rhs_swizzle,
            a_sparse_metadata=lhs_sparse,
            accumulate=False,
            collective=True,
        )
        tcgen05.commit_arrive(mma_barrier, collective=True, ctx=ctx)
      mma_barrier.wait(orders_tensor_core=True)
      m_block_tile = m // 2
      m_slice = ds(arith.muli(block_id, c(m_block_tile, index)), m_block_tile)
      acc.load().store_untiled(memref_slice(out, m_slice), optimized=False)

    x_shape = (m, k // 2)
    y_shape = (k, n)
    x = self.prng.uniform(-1, 1, x_shape).astype(in_jax_dtype)
    y = self.prng.uniform(-1, 1, y_shape).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    m_block = m // 2
    n_block = n // 2
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape((m_block, k // 2), lhs_tiling), in_jax_dtype),
        jax.ShapeDtypeStruct(tile_shape((k, n_block), rhs_tiling), in_jax_dtype),
        jax.ShapeDtypeStruct((m_block // 128, k // 128, 128, 64), sparse_meta_dtype),
        mgpu.Barrier(1, num_barriers=3),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m_block, n), out_jax_dtype, collective=True),
        mgpu.TMEM((m_block, k // 2), sparse_meta_dtype, layout=tcgen05.sparse_meta_layout(), collective=True),
    ]
    index_pairs = np.asarray(np.meshgrid(range(4), range(4))).T.reshape(-1, 2)
    valid_pairs = index_pairs[index_pairs[:, 0] < index_pairs[:, 1]]
    assert len(valid_pairs) == 6
    x_pairs = jax.random.randint(jax.random.key(1234), (m, k // 4), 0, 6, dtype=jnp.uint8)
    x_sparse = valid_pairs[x_pairs]
    assert x_sparse.shape == (m, k // 4, 2)
    def format_sparse_meta(meta):
      mn, k, _2 = meta.shape
      assert _2 == 2
      k *= 2
      if jnp.dtype(in_jax_dtype).itemsize == 1:
        meta_tiled = (
            meta.reshape(mn // 128, 128, k // 64, 64).transpose(0, 2, 1, 3)
        )
      else:
        meta_tiled = (
          meta.reshape(mn // 128, 8, 2, 8, k // 64, 4, 2, 8)
          .transpose(0, 4, 1, 6, 3, 5, 2, 7)
        )
      return (
          meta_tiled.reshape(mn // 128, k // 64, 128, 64)
          .astype(sparse_meta_dtype)
      )
    x_gpu_sparse = format_sparse_meta(x_sparse)
    args = (x, y, x_gpu_sparse)
    z = mgpu.as_gpu_kernel(
        kernel, (2, 1, 1), (128, 1, 1), args, out_shape, scratch_shape, cluster=(2, 1, 1)
    )(*args)
    x_logical = np.zeros_like(x, shape=(m, k // 4, 4))
    np.put_along_axis(x_logical, x_sparse, x.reshape(x_sparse.shape), axis=-1)
    x_logical = x_logical.reshape(m, k)
    ref = jnp.matmul(
        x_logical.astype(jnp.float32), y.astype(jnp.float32), precision="high"
    )
    atol = 2e-2 if out_jax_dtype == jnp.float16 else 7e-5
    rtol = 8e-4 if out_jax_dtype == jnp.float16 else 5e-6
    np.testing.assert_allclose(z, ref, atol=atol, rtol=rtol)

  @parameterized.product(
      in_jax_dtype=(jnp.float8_e5m2, jnp.float8_e4m3fn),
      scale_jax_dtype=(jnp.float8_e8m0fnu,),
      m=(128,),
      n=(128, 192, 256),
      swizzle=(128,),
  )
  def test_mma_block_scaled_sparse(self, m, n, in_jax_dtype, scale_jax_dtype, swizzle):
    out_jax_dtype = jnp.float32
    sparse_meta_dtype = jnp.uint2
    block_size = 64
    k_steps = 2

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    swizzle_elems = 8 * swizzle // bitwidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    lhs_tiling = rhs_tiling = (8, swizzle_elems)
    if n % rhs_tiling[1]:
      self.skipTest(f"{n=} too short for {swizzle=} and {in_jax_dtype=}")

    def kernel(ctx, lhs, rhs, lhs_sparse_gmem, lhs_scales_gmem, rhs_scales_gmem, out, scratch):
      (
          lhs_smem, rhs_smem, lhs_sparse_smem,
          lhs_scales_smem, rhs_scales_smem,
          barriers, mma_barrier, acc, lhs_sparse, lhs_scales, rhs_scales,
      ) = scratch
      operand_kwargs = dict(
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(lhs_tiling),
      )
      ctx.async_copy(src_ref=lhs, dst_ref=lhs_smem, barrier=barriers[0], **operand_kwargs)
      ctx.async_copy(src_ref=rhs, dst_ref=rhs_smem, barrier=barriers[1], swizzle=swizzle, gmem_transform=mgpu.TileTransform(rhs_tiling))
      ctx.async_copy(src_ref=lhs_sparse_gmem, dst_ref=lhs_sparse_smem, barrier=barriers[2])
      ctx.async_copy(src_ref=lhs_scales_gmem, dst_ref=lhs_scales_smem, barrier=barriers[3])
      ctx.async_copy(src_ref=rhs_scales_gmem, dst_ref=rhs_scales_smem, barrier=barriers[4])
      for i in range(5):
        barriers[i].wait()
      with mgpu.single_thread():
        tcgen05.async_copy_sparse_metadata_smem_to_tmem(lhs_sparse_smem, lhs_sparse)
        tcgen05.async_copy_scales_smem_to_tmem(lhs_scales_smem, lhs_scales)
        tcgen05.async_copy_scales_smem_to_tmem(rhs_scales_smem, rhs_scales)
        tcgen05.mma(
            acc,
            lhs_smem,
            mgpu.memref_transpose(rhs_smem, (1, 0, 3, 2)),
            a_swizzle=swizzle,
            b_swizzle=swizzle,
            a_scale=lhs_scales,
            b_scale=rhs_scales,
            a_sparse_metadata=lhs_sparse,
            accumulate=False,
        )
        tcgen05.commit_arrive(mma_barrier)
      mma_barrier.wait(orders_tensor_core=True)
      acc.load().store_untiled(out, optimized=False)

    x_shape = (m, k // 2)
    x = self.prng.uniform(-1, 1, x_shape).astype(in_jax_dtype)
    y_shape = (n, k)
    y = self.prng.uniform(-1, 1, y_shape).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape(x_shape, lhs_tiling), in_jax_dtype),
        jax.ShapeDtypeStruct(tile_shape(y_shape, rhs_tiling), in_jax_dtype),
        jax.ShapeDtypeStruct((m // 128, k // 128, 128, 64), sparse_meta_dtype),
        jax.ShapeDtypeStruct((m // 128, k // (block_size * 4), 32, 16), scale_jax_dtype),
        jax.ShapeDtypeStruct((n // 128, k // (block_size * 4), 32, 16), scale_jax_dtype),
        mgpu.Barrier(1, num_barriers=5),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m, n), out_jax_dtype),
        mgpu.TMEM((m, k // 2), sparse_meta_dtype, layout=tcgen05.sparse_meta_layout()),
        mgpu.TMEM((m, k // block_size), scale_jax_dtype, layout=tcgen05.scales_layout()),
        mgpu.TMEM((n, k // block_size), scale_jax_dtype, layout=tcgen05.scales_layout()),
    ]
    index_pairs = np.asarray(np.meshgrid(range(4), range(4))).T.reshape(-1, 2)
    valid_pairs = index_pairs[index_pairs[:, 0] < index_pairs[:, 1]]
    assert len(valid_pairs) == 6
    x_pairs = jax.random.randint(jax.random.key(1234), (m, k // 4), 0, 6, dtype=jnp.uint8)
    x_sparse = valid_pairs[x_pairs]
    assert x_sparse.shape == (m, k // 4, 2)
    def format_sparse_meta(meta):
      mn, k, _2 = meta.shape
      assert _2 == 2
      k *= 2
      meta_tiled = (
          meta.reshape(mn // 128, 128, k // 64, 64).transpose(0, 2, 1, 3)
      )
      return (
          meta_tiled.reshape(mn // 128, k // 64, 128, 64)
          .astype(sparse_meta_dtype)
      )
    x_gpu_sparse = format_sparse_meta(x_sparse)
    a_scales, b_scales = self._sample_scales(m, k, n, block_size, scale_jax_dtype)
    def format_scales(scales):
      mn, k = scales.shape
      assert mn % 128 == 0 and k % 4 == 0, scales.shape
      return (
          scales.reshape(mn // 128, 4, 32, k // 4, 4)
          .transpose(0, 3, 2, 1, 4)
          .reshape(mn // 128, k // 4, 32, 16)
      )
    a_gpu_scales, b_gpu_scales = map(format_scales, (a_scales, b_scales))
    args = (x, y, x_gpu_sparse, a_gpu_scales, b_gpu_scales)
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), args, out_shape, scratch_shape
    )(*args)
    x_logical = np.zeros_like(x, shape=(m, k // 4, 4))
    np.put_along_axis(x_logical, x_sparse, x.reshape(x_sparse.shape), axis=-1)
    x_logical = x_logical.reshape(m, k)
    x32 = x_logical.astype(np.float32)
    y32 = y.astype(np.float32)
    a_logical_scales = jnp.repeat(a_scales, block_size, axis=1).astype(jnp.float32)
    b_logical_scales = jnp.repeat(b_scales, block_size, axis=1).astype(jnp.float32)
    ref = jnp.matmul(
        x32 * a_logical_scales,
        (y32 * b_logical_scales).T,
        precision="high",
    )
    np.testing.assert_allclose(z, ref, atol=2e-4, rtol=5e-6)

  @parameterized.product(
      in_jax_dtype=(jnp.float4_e2m1fn,),
      scale_jax_dtype=(jnp.float8_e8m0fnu, jnp.float8_e4m3fn),
      base_block_size=(16, 32),
      m=(128,),
      n=(128, 256),
      swizzle=(128,),
  )
  def test_mma_block_scaled_sparse_f4(
      self, m, n, in_jax_dtype, scale_jax_dtype, base_block_size, swizzle
  ):
    if scale_jax_dtype == jnp.float8_e4m3fn and base_block_size != 16:
      self.skipTest("e4m3fn scale only supports block size 16.")
    out_jax_dtype = jnp.float32
    sparse_meta_dtype = jnp.uint2
    block_size = base_block_size * 2
    # TODO: increase k to allow sparse K=192 on sm_107a
    k_steps = 2

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    swizzle_elems = 8 * swizzle // bitwidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    lhs_tiling = rhs_tiling = (8, swizzle_elems)

    def kernel(ctx, lhs, rhs, lhs_sparse_gmem, lhs_scales_gmem, rhs_scales_gmem, out, scratch):
      (
          lhs_smem, rhs_smem, lhs_sparse_smem,
          lhs_scales_smem, rhs_scales_smem,
          barriers, mma_barrier, acc, lhs_sparse, lhs_scales, rhs_scales,
      ) = scratch
      operand_kwargs = dict(
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(lhs_tiling),
      )
      ctx.async_copy(src_ref=lhs, dst_ref=lhs_smem, barrier=barriers[0], **operand_kwargs)
      ctx.async_copy(src_ref=rhs, dst_ref=rhs_smem, barrier=barriers[1], swizzle=swizzle, gmem_transform=mgpu.TileTransform(rhs_tiling))
      ctx.async_copy(src_ref=lhs_sparse_gmem, dst_ref=lhs_sparse_smem, barrier=barriers[2])
      ctx.async_copy(src_ref=lhs_scales_gmem, dst_ref=lhs_scales_smem, barrier=barriers[3])
      ctx.async_copy(src_ref=rhs_scales_gmem, dst_ref=rhs_scales_smem, barrier=barriers[4])
      for i in range(5):
        barriers[i].wait()
      with mgpu.single_thread():
        tcgen05.async_copy_sparse_metadata_smem_to_tmem(lhs_sparse_smem, lhs_sparse)
        tcgen05.async_copy_scales_smem_to_tmem(lhs_scales_smem, lhs_scales)
        tcgen05.async_copy_scales_smem_to_tmem(rhs_scales_smem, rhs_scales)
        tcgen05.mma(
            acc,
            lhs_smem,
            mgpu.memref_transpose(rhs_smem, (1, 0, 3, 2)),
            a_swizzle=swizzle,
            b_swizzle=swizzle,
            a_scale=lhs_scales,
            b_scale=rhs_scales,
            a_sparse_metadata=lhs_sparse,
            accumulate=False,
        )
        tcgen05.commit_arrive(mma_barrier)
      mma_barrier.wait(orders_tensor_core=True)
      acc.load().store_untiled(out, optimized=False)

    x_shape = (m, k // 2)
    x = self.prng.uniform(-1, 1, x_shape).astype(in_jax_dtype)
    y_shape = (n, k)
    y = self.prng.uniform(-1, 1, y_shape).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    # 8 (fp4 v0) or 4 (fp4 v1)
    group_elems = tcgen05.sparse_group_elems(in_mlir_dtype)
    n_groups = k // group_elems
    meta_k = n_groups * 2
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape(x_shape, lhs_tiling), in_jax_dtype),
        jax.ShapeDtypeStruct(tile_shape(y_shape, rhs_tiling), in_jax_dtype),
        jax.ShapeDtypeStruct((m // 128, meta_k // 64, 128, 64), sparse_meta_dtype),
        jax.ShapeDtypeStruct((m // 128, k // (block_size * 4), 32, 16), scale_jax_dtype),
        jax.ShapeDtypeStruct((n // 128, k // (block_size * 4), 32, 16), scale_jax_dtype),
        mgpu.Barrier(1, num_barriers=5),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m, n), out_jax_dtype),
        mgpu.TMEM((m, meta_k), sparse_meta_dtype, layout=tcgen05.sparse_meta_layout()),
        mgpu.TMEM((m, k // block_size), scale_jax_dtype, layout=tcgen05.scales_layout()),
        mgpu.TMEM((n, k // block_size), scale_jax_dtype, layout=tcgen05.scales_layout()),
    ]
    index_pairs = np.asarray(np.meshgrid(range(4), range(4))).T.reshape(-1, 2)
    valid_pairs = index_pairs[index_pairs[:, 0] < index_pairs[:, 1]]
    assert len(valid_pairs) == 6
    x_pairs = jax.random.randint(jax.random.key(1234), (m, n_groups), 0, 6, dtype=jnp.uint8)
    x_sparse = valid_pairs[x_pairs]
    assert x_sparse.shape == (m, n_groups, 2)
    def format_sparse_meta(meta):
      mn, groups, _2 = meta.shape
      assert _2 == 2
      k_meta = groups * 2
      meta_tiled = (
          meta.reshape(mn // 128, 128, k_meta // 64, 64).transpose(0, 2, 1, 3)
      )
      return (
          meta_tiled.reshape(mn // 128, k_meta // 64, 128, 64)
          .astype(sparse_meta_dtype)
      )
    x_gpu_sparse = format_sparse_meta(x_sparse)
    a_scales, b_scales = self._sample_scales(m, k, n, block_size, scale_jax_dtype)
    def format_scales(scales):
      mn, k = scales.shape
      assert mn % 128 == 0 and k % 4 == 0, scales.shape
      return (
          scales.reshape(mn // 128, 4, 32, k // 4, 4)
          .transpose(0, 3, 2, 1, 4)
          .reshape(mn // 128, k // 4, 32, 16)
      )
    a_gpu_scales, b_gpu_scales = map(format_scales, (a_scales, b_scales))
    args = (x, y, x_gpu_sparse, a_gpu_scales, b_gpu_scales)
    z = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), args, out_shape, scratch_shape
    )(*args)
    # 4-bit sparse data is filled in sub-chunks of 2 (v0) or 1 (v1) elements.
    sub_chunk = group_elems // 4
    x32 = x.astype(np.float32).reshape(m, n_groups, 2, sub_chunk)
    x_logical32 = np.zeros((m, n_groups, 4, sub_chunk), dtype=np.float32)
    np.put_along_axis(x_logical32, x_sparse[..., np.newaxis], x32, axis=-2)
    x_logical32 = x_logical32.reshape(m, k)
    y32 = y.astype(np.float32)
    a_logical_scales = jnp.repeat(a_scales, block_size, axis=1).astype(jnp.float32)
    b_logical_scales = jnp.repeat(b_scales, block_size, axis=1).astype(jnp.float32)
    ref = jnp.matmul(
        x_logical32 * a_logical_scales,
        (y32 * b_logical_scales).T,
        precision="high",
    )
    np.testing.assert_allclose(z, ref, atol=2e-4, rtol=5e-6)

  @parameterized.product(
      lhs_transpose=(False, True),
      rhs_transpose=(False, True),
      in_jax_dtype=(jnp.float16,),
      out_jax_dtype=(jnp.float32,),
      m=(128, 256),  # TODO(apaszke): 192, 256
      n=(128, 160, 256),
      swizzle=(32, 64, 128,),
  )
  def test_mma_collective(
      self,
      m,
      n,
      swizzle,
      lhs_transpose,
      rhs_transpose,
      in_jax_dtype,
      out_jax_dtype,
  ):
    k_steps = 2  # Reducing to 1 can be helpful while debugging.
    if out_jax_dtype == jnp.float16 and in_jax_dtype != jnp.float16:
      raise self.skipTest("Only f16 input is supported for f16 output.")

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    m_block_tile = m // 2
    n_block_tile = n // 2
    swizzle_elems = swizzle // bytewidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    index = ir.IndexType.get()

    tiling = (8, swizzle_elems)

    def kernel(ctx, lhs, rhs, out, scratch):
      lhs_smem, rhs_smem, barriers, mma_barrier, acc = scratch
      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      ctx.async_copy(
          src_ref=lhs,
          dst_ref=lhs_smem,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
          barrier=barriers[0],
          collective=gpu.Dimension.x,
          # Split non-contracting dim.
          leader_tracked=mgpu.CopyPartition.PARTITIONED(1 if lhs_transpose else 0),
      )
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
          barrier=barriers[1],
          collective=gpu.Dimension.x,
          # Split non-contracting dim.
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0 if rhs_transpose else 1),
      )
      is_leader_thread = single_thread_predicate()
      is_first_block = arith.cmpi(arith.CmpIPredicate.eq, block_id, c(0, index))
      with when(arith.andi(is_first_block, is_leader_thread)):
        barriers[0].wait()
        barriers[1].wait()
        if lhs_transpose:
          lhs_smem = memref_transpose(lhs_smem, (1, 0, 3, 2))
        if rhs_transpose:
          rhs_smem = memref_transpose(rhs_smem, (1, 0, 3, 2))
        tcgen05.mma(
            acc, lhs_smem, rhs_smem, a_swizzle=swizzle, b_swizzle=swizzle, accumulate=False, collective=True
        )
        tcgen05.commit_arrive(mma_barrier, collective=True, ctx=ctx)
      mma_barrier.wait(orders_tensor_core=True)
      m_slice = ds(arith.muli(block_id, c(m_block_tile, index)), m_block_tile)
      acc.load().store_untiled(memref_slice(out, m_slice), optimized=False)

    in_finfo = jnp.finfo(in_jax_dtype)
    exponent_bits, mantissa_bits = in_finfo.nexp, in_finfo.nmant
    def quantize(x):
      # Quantize the input to avoid rounding when feeding the TensorCore
      return jax.lax.reduce_precision(x, exponent_bits, mantissa_bits)

    x_shape = (k, m) if lhs_transpose else (m, k)
    x_block_shape = (k, m_block_tile) if lhs_transpose else (m_block_tile, k)
    x = quantize(self.prng.uniform(-1, 1, x_shape)).astype(in_jax_dtype)
    y_shape = (n, k) if rhs_transpose else (k, n)
    y_block_shape = (n_block_tile, k) if rhs_transpose else (k, n_block_tile)
    y = quantize(self.prng.uniform(-1, 1, y_shape)).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    if any(s % t for s, t in zip(x_block_shape, tiling)):
      self.skipTest("LHS block shape not divisible by tiling.")
    if any(s % t for s, t in zip(y_block_shape, tiling)):
      self.skipTest("RHS block shape not divisible by tiling.")
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape(x_block_shape, tiling), in_jax_dtype),
        jax.ShapeDtypeStruct(tile_shape(y_block_shape, tiling), in_jax_dtype),
        mgpu.Barrier(1, num_barriers=2),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m_block_tile, n), out_jax_dtype, collective=True),
    ]
    z = mgpu.as_gpu_kernel(
        kernel, (2, 1, 1), (128, 1, 1), (x, y), out_shape, scratch_shape, cluster=(2, 1, 1)
    )(x, y)
    x32, y32 = x.astype(np.float32), y.astype(np.float32)
    ref = jnp.matmul(
        (x32.T if lhs_transpose else x32),
        (y32.T if rhs_transpose else y32),
        precision="high",
    )
    atol = 2e-2 if out_jax_dtype == jnp.float16 else 5e-6
    np.testing.assert_allclose(z, ref, atol=atol)

  @parameterized.product(cluster=((2, 2, 1), (2, 1, 2)))
  def test_mma_collective_multidim_cluster(self, cluster):
    in_jax_dtype = jnp.float16
    out_jax_dtype = jnp.float32
    m = n = 128
    swizzle = 128
    k_steps = 2

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    m_block_tile = m // 2
    n_block_tile = n // 2
    swizzle_elems = swizzle // bytewidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    index = ir.IndexType.get()

    tiling = (8, swizzle_elems)

    def kernel(ctx, lhs, rhs, out, scratch):
      lhs_smem, rhs_smem, barriers, cluster_barrier, mma_barrier, acc = scratch
      # CTAs with linear cluster indices 2k and 2k + 1 form the k-th MMA pair.
      block_idx = utils.cluster_idx()
      pair_id = arith.divui(block_idx, c(2, index))
      block_id = arith.remui(block_idx, c(2, index))
      m_start = arith.muli(block_id, c(m_block_tile, index))
      n_start = arith.muli(block_id, c(n_block_tile, index))
      ctx.async_copy(
          src_ref=lhs,
          dst_ref=lhs_smem,
          gmem_slice=(pair_id, ds(m_start, m_block_tile)),
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
          barrier=barriers[0],
      )
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          gmem_slice=(pair_id, slice(None), ds(n_start, n_block_tile)),
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
          barrier=barriers[1],
      )
      barriers[0].wait()
      barriers[1].wait()
      # The MMA also reads the SMEM operands of the other CTA in the pair.
      cluster_barrier.arrive()
      cluster_barrier.wait()
      is_leader_thread = single_thread_predicate()
      is_first_block = arith.cmpi(arith.CmpIPredicate.eq, block_id, c(0, index))
      with when(arith.andi(is_first_block, is_leader_thread)):
        tcgen05.mma(
            acc, lhs_smem, rhs_smem, a_swizzle=swizzle, b_swizzle=swizzle, accumulate=False, collective=True
        )
        tcgen05.commit_arrive(mma_barrier, collective=True, ctx=ctx)
      mma_barrier.wait(orders_tensor_core=True)
      m_slice = ds(m_start, m_block_tile)
      acc.load().store_untiled(memref_slice(out, (pair_id, m_slice)), optimized=False)

    in_finfo = jnp.finfo(in_jax_dtype)
    exponent_bits, mantissa_bits = in_finfo.nexp, in_finfo.nmant
    def quantize(x):
      # Quantize the input to avoid rounding when feeding the TensorCore
      return jax.lax.reduce_precision(x, exponent_bits, mantissa_bits)

    x_block_shape = (m_block_tile, k)
    y_block_shape = (k, n_block_tile)
    # The two pairs compute GEMMs on different data, so an MMA that reads SMEM
    # from a CTA outside its own pair, or a commit that signals another pair's
    # barriers, fails the output comparison instead of passing by accident.
    x = quantize(self.prng.uniform(-1, 1, (2, m, k))).astype(in_jax_dtype)
    y = quantize(self.prng.uniform(-1, 1, (2, k, n))).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((2, m, n), out_jax_dtype)
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape(x_block_shape, tiling), in_jax_dtype),
        jax.ShapeDtypeStruct(tile_shape(y_block_shape, tiling), in_jax_dtype),
        mgpu.Barrier(1, num_barriers=2),
        mgpu.ClusterBarrier(collective_dims=(gpu.Dimension.x,)),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM((m_block_tile, n), out_jax_dtype, collective=True),
    ]
    z = mgpu.as_gpu_kernel(
        kernel, cluster, (128, 1, 1), (x, y), out_shape, scratch_shape, cluster=cluster
    )(x, y)
    ref = jnp.matmul(x.astype(np.float32), y.astype(np.float32), precision="high")
    np.testing.assert_allclose(z, ref, atol=5e-6)

  def test_raises_error_if_collective_arrival_minormost_cluster_dim_not_2(self):
    def kernel(ctx, out, barrier):
      del out
      tcgen05.commit_arrive(barrier, collective=True, ctx=ctx)

    out_shape = jax.ShapeDtypeStruct((128,), jnp.int32)
    with self.assertRaisesRegex(
        ValueError,
        "Collective arrivals require the minormost cluster dimension to have"
        " size 2",
    ):
      mgpu.as_gpu_kernel(
          kernel, (1, 2, 1), (128, 1, 1), (), out_shape, mgpu.Barrier(1), cluster=(1, 2, 1)
      )

  @parameterized.product(
      in_jax_dtype=(jnp.float16,),
      out_jax_dtype=(jnp.float32,),
      m=(256,),  # TODO(apaszke): 64, 192, 256
      n=(128, 192, 224, 256,),
      k_steps=(2,),  # Note: reducing to 1 can be useful for debugging.
      swizzle=(32, 64, 128,),
  )
  def test_mma_collective_lhs_tmem(
      self,
      m,
      n,
      k_steps,
      swizzle,
      in_jax_dtype,
      out_jax_dtype,
  ):
    if out_jax_dtype == jnp.float16 and in_jax_dtype != jnp.float16:
      raise self.skipTest("Only f16 input is supported for f16 output.")

    in_mlir_dtype = utils.dtype_to_ir_type(in_jax_dtype)
    m_block_tile = m // 2
    n_block_tile = n // 2
    swizzle_elems = swizzle // bytewidth(in_mlir_dtype)
    k = swizzle_elems * k_steps
    index = ir.IndexType.get()

    tiling = (8, swizzle_elems)

    def kernel(ctx, lhs, rhs, out, scratch):
      lhs_smem, rhs_smem, barriers, mma_barrier, cluster_barrier, acc, lhs_tmem = scratch
      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      ctx.async_copy(
          src_ref=lhs,
          dst_ref=lhs_smem,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
          barrier=barriers[0],
          collective=gpu.Dimension.x,
          # Split non-contracting dim.
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      ctx.async_copy(
          src_ref=rhs,
          dst_ref=rhs_smem,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
          barrier=barriers[1],
          collective=gpu.Dimension.x,
          # Split non-contracting dim.
          leader_tracked=mgpu.CopyPartition.PARTITIONED(1),
      )

      is_leader_thread = single_thread_predicate()
      is_first_block = arith.cmpi(arith.CmpIPredicate.eq, block_id, c(0, index))

      with when(arith.andi(is_first_block, is_leader_thread)):
        barriers[0].wait()
      gpu.barrier()
      # Because only block 1 waits on the TMA, we need a cluster barrier so
      # that the SMEM updates are visible on block 2.
      cluster_barrier.arrive(orders_tensor_core=True)
      cluster_barrier.wait(orders_tensor_core=True)
      lhs_tmem.store(
          fa.FragmentedArray.load_tiled(
              lhs_smem, swizzle, layout=tcgen05.LAYOUT
          )
      )
      tcgen05.commit_tmem()
      # Make sure TMEM has been loaded on both blocks.
      cluster_barrier.arrive(orders_tensor_core=True)
      cluster_barrier.wait(orders_tensor_core=True)
      with when(arith.andi(is_first_block, is_leader_thread)):
        barriers[1].wait()
        tcgen05.mma(
            acc,
            lhs_tmem,
            rhs_smem,
            a_swizzle=swizzle,
            b_swizzle=swizzle,
            accumulate=False,
            collective=True,
        )
        tcgen05.commit_arrive(mma_barrier, collective=True, ctx=ctx)
      mma_barrier.wait(orders_tensor_core=True)
      m_slice = ds(arith.muli(block_id, c(m_block_tile, index)), m_block_tile)
      acc.load().store_untiled(memref_slice(out, m_slice), optimized=False)

    in_finfo = jnp.finfo(in_jax_dtype)
    exponent_bits, mantissa_bits = in_finfo.nexp, in_finfo.nmant

    def quantize(x):
      # Quantize the input to avoid rounding when feeding the TensorCore
      return jax.lax.reduce_precision(x, exponent_bits, mantissa_bits)

    x_shape = (m, k)
    x_block_shape = (m_block_tile, k)
    x = quantize(self.prng.uniform(-1, 1, x_shape)).astype(in_jax_dtype)
    y_shape = (k, n)
    y_block_shape = (k, n_block_tile)
    y = quantize(self.prng.uniform(-1, 1, y_shape)).astype(in_jax_dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), out_jax_dtype)
    if any(s % t for s, t in zip(x_block_shape, tiling)):
      self.skipTest("LHS block shape not divisible by tiling.")
    if any(s % t for s, t in zip(y_block_shape, tiling)):
      self.skipTest("RHS block shape not divisible by tiling.")
    scratch_shape = [
        jax.ShapeDtypeStruct(tile_shape(x_block_shape, tiling), in_jax_dtype),
        jax.ShapeDtypeStruct(tile_shape(y_block_shape, tiling), in_jax_dtype),
        mgpu.Barrier(1, num_barriers=2),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.ClusterBarrier(collective_dims=(gpu.Dimension.x,)),
        mgpu.TMEM((128, n), out_jax_dtype, collective=True),
        mgpu.TMEM((128, k), in_jax_dtype, collective=True, packing=2),
    ]
    z = mgpu.as_gpu_kernel(
        kernel,
        (2, 1, 1),
        (128, 1, 1),
        (x, y),
        out_shape,
        scratch_shape,
        cluster=(2, 1, 1),
    )(x, y)
    x32, y32 = x.astype(np.float32), y.astype(np.float32)
    ref = jnp.matmul(x32, y32, precision="high")
    atol = 2e-2 if out_jax_dtype == jnp.float16 else 5e-6
    np.testing.assert_allclose(z, ref, atol=atol)

  def test_raises_error_if_tmem_oom(self):
    def kernel(ctx, input, output, scratch):
      del ctx, input, output, scratch

    x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
    scratch_shape = [
        mgpu.TMEM((128, 384), jnp.float32), # Should round up to 512 columns.
        mgpu.TMEM((128, 64), jnp.float32), # Will trigger OOM.
    ]
    with self.assertRaisesRegex(ValueError,
        "Total TMEM allocation exceeds memory limit."):
      mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x, x, scratch_shape
      )(x).block_until_ready()

  def test_raises_error_if_collective_tmem_without_cluster(self):
    def kernel(ctx, input, output, scratch):
      del ctx, input, output, scratch

    x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
    scratch_shape = [mgpu.TMEM((128, 384), jnp.float32, collective=True)]
    with self.assertRaisesRegex(
        ValueError,
        "Collective TMEM allocations are only supported for clusters with an"
        " even number of blocks in them.",
    ):
      mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x, x, scratch_shape
      )(x).block_until_ready()

  @parameterized.parameters((0,), (1,), (2,))
  def test_cluster_launch_control(self, dim):
    # Let's say we have 148 SMs in our gpu. We attempt to schedule 149 blocks on
    # 148 SMs. Only one SM will succeed in stealing the 149th block, and the
    # others will fail. Therefore we test that there is exactly 1 stolen block
    # and the others fail and return -1.
    num_sms = jax.devices()[0].core_count
    num_blocks = num_sms + 1
    grid = [1, 1, 1]
    grid[dim] = num_blocks

    def kernel(ctx, out, scratch):
      del ctx
      cancel_result_ref, barrier, _ = scratch

      is_leader_thread = single_thread_predicate()
      barrier.arrive_expect_tx(16, predicate=is_leader_thread)
      mgpu.try_cluster_cancel(cancel_result_ref, barrier, is_leader_thread)

      barrier.wait()
      *cta_ids, cancelled_launch = mgpu.query_cluster_cancel(cancel_result_ref)
      cta_id = arith.addi(cta_ids[0], arith.addi(cta_ids[1], cta_ids[2]))

      # Store a sentinel value if no work can be scheduled.
      idx = arith.index_cast(ir.IndexType.get(), utils.block_idx())
      sentinel_val = arith.constant(ir.IntegerType.get_signless(32), -1)

      value = arith.select(cancelled_launch, cta_id, sentinel_val)
      memref.store(value, out, [idx])

    cancel_result_ref = jax.ShapeDtypeStruct((16,), jnp.int8)  # 128 bits
    out_ty = jax.ShapeDtypeStruct((num_sms,), jnp.int32)
    scratch = (
        cancel_result_ref,
        mgpu.Barrier(1),
        # Requesting SMEM close to the 228kb limit to ensure that each SM only
        # schedules 1 block.
        jax.ShapeDtypeStruct((220 * 1024,), jnp.int8),
    )
    out = mgpu.as_gpu_kernel(kernel, grid, (128, 1, 1), (), out_ty, scratch)()

    out = np.sort(out)
    out_ref = np.array([-1] * (num_sms - 1) + [num_sms])
    np.testing.assert_array_equal(out, out_ref)


class Sm80Test(TestCase):

  def setUp(self):
    super().setUp(min_compute_capability="8.0")

  def test_barrier_arrive_wait(self):
    i32 = ir.IntegerType.get_signless(32)

    def kernel(ctx, dst, barrier):
      del ctx
      barrier.arrive()
      barrier.wait()
      with mgpu.single_thread(scope=mgpu.ThreadSubset.BLOCK):
        memref.store(arith.constant(i32, 1), dst, [])

    out_shape = jax.ShapeDtypeStruct((), jnp.int32)
    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        (),
        out_shape,
        mgpu.Barrier(arrival_count=128),
    )()
    np.testing.assert_array_equal(y, np.array(1, dtype=np.int32))

  def _test_ldmatrix(self, dtype, m_mult, k_mult, transpose=False, layout=None):
    m, k = 32 * m_mult, 8 * k_mult
    dtype = jnp.dtype(dtype)
    in_shape = (k, m) if transpose else (m, k)
    bitwidth = np.iinfo(dtype).bits
    try:
      swizzle = next(
          s for s in (128, 64, 32, 16)
          if (in_shape[1] * bitwidth // 8) % s == 0
          # For async copy.
          and (((in_shape[0] // 8) * ((in_shape[1] * bitwidth // 8) // s)) % 4 == 0)
      )
    except StopIteration:
      self.skipTest("No suitable swizzle found.")
    packing = 32 // bitwidth
    swizzle_elems = swizzle * 8 // bitwidth
    if layout is not None:
      pass
    elif bitwidth == 8 and transpose:
      # We need each warp to hold a 16x16 submatrix (two 8x16 matrices).
      layout = fa.TiledLayout(
          fa.Tiling(((64, 4 * packing), (16, 4 * packing), (8, 4 * packing), (packing,))),
          warp_dims=(-7,),
          lane_dims=(-3, -2),
          vector_dim=-1,
      )
    else:
      layout = fa.TiledLayout(
          fa.Tiling(((32, 4 * packing), (8, 4 * packing), (packing,))),
          warp_dims=(-5,),
          lane_dims=(-3, -2),
          vector_dim=-1,
      )
    def kernel(ctx, a, out, smem):
      a_smem, barrier = smem
      ctx.async_copy(
          src_ref=a,
          dst_ref=a_smem,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform((8, swizzle_elems)),
          barrier=barrier,
          implementation=mgpu.AsyncCopyImplementation.CP_ASYNC,
      )
      barrier.wait()
      if transpose:
        a_smem = mgpu.memref_transpose(a_smem, (1, 0, 3, 2))
      a_fa = mgpu.FragmentedArray.load_tiled(
          a_smem, layout=layout, swizzle=swizzle, is_signed=True
      )
      a_fa.store_untiled(out, optimized=False)

    x = self.prng.integers(-10000, 10000, in_shape).astype(dtype)
    out_shape = jax.ShapeDtypeStruct((m, k), dtype)
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      y = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x, out_shape,
          (
              jax.ShapeDtypeStruct(tile_shape(in_shape, (8, swizzle_elems)), dtype),
              mgpu.Barrier(1)
          ),
      )(x)
      expected = x.T if transpose else x
      np.testing.assert_array_equal(y, expected)
    if transpose and bitwidth == 8:
      num = next(n for n in (2, 1) if (m_mult * k_mult // 4) % n == 0)
      expected_instr = f"ldmatrix.sync.aligned.m16n16.x{num}.trans.shared"
    else:
      num_per_tile = 8 / (4 * packing)
      num = next(n for n in (4, 2, 1) if int(num_per_tile * m_mult * k_mult) % n == 0)
      expected_instr = (
          f"ldmatrix.sync.aligned.m8n8.x{num}.trans.shared"
          if transpose
          else f"ldmatrix.sync.aligned.m8n8.x{num}.shared"
      )
    self.assertIn(expected_instr, ptx())
    self.assertNotIn("ld.shared", ptx())

  @parameterized.product(
      dtype=(jnp.int8, jnp.int16, jnp.int32),
      m_mult=(1, 2, 3, 4, 6, 7, 9),
      k_mult=(1, 2, 3, 4, 6, 7, 9),
  )
  @jtu.thread_unsafe_test()
  def test_ldmatrix_basic(self, dtype, m_mult, k_mult):
    return self._test_ldmatrix(dtype, m_mult, k_mult, transpose=False)

  @parameterized.product(
      dtype=(jnp.int8, jnp.int16),
      m_mult=(1, 2, 3, 4, 6, 7, 9),
      k_mult=(1, 2, 3, 4, 6, 7, 9),
  )
  @jtu.thread_unsafe_test()
  def test_ldmatrix_transpose(self, dtype, m_mult, k_mult):
    if dtype == jnp.int8 and not jtu.is_cuda_compute_capability_at_least("10.0"):
      raise self.skipTest("int8 transposed ldmatrix requires sm_10.")
    scale = 16 // jax.dtypes.itemsize_bits(dtype)
    return self._test_ldmatrix(dtype, m_mult * scale, k_mult * scale, transpose=True)

  def test_ldmatrix_8b_transpose_wide_row_layout(self):
    if not jtu.is_cuda_compute_capability_at_least("10.0"):
      raise self.skipTest("int8 transposed ldmatrix requires sm_10.")
    packing = 4
    layout = fa.TiledLayout(
        fa.Tiling(((32, 4 * packing), (8, 4 * packing), (packing,))),
        warp_dims=(fa.Replicated(4),),
        lane_dims=(-3, -2),
        vector_dim=-1,
    )
    return self._test_ldmatrix(jnp.int8, 8, 4, transpose=True, layout=layout)


class BarrierTest(TestCase):

  def test_wg_communication(self):
    i32 = ir.IntegerType.get_signless(32)
    def kernel(ctx, dst, scratch):
      tmp, barriers = scratch
      del ctx  # Unused.
      wg_idx = arith.divui(mgpu.warp_idx(), c(4, i32))
      is_first_wg = arith.cmpi(arith.CmpIPredicate.eq, wg_idx, c(0, i32))
      is_second_wg = arith.cmpi(arith.CmpIPredicate.eq, wg_idx, c(1, i32))
      arr = mgpu.FragmentedArray.splat(
          arith.addi(wg_idx, c(1, i32)),
          (128,),
          mgpu.WGStridedFragLayout((128,), 1),
          is_signed=False,
      )
      with ir.InsertionPoint(scf.IfOp(is_first_wg).then_block):
        arr.store_untiled(tmp)
        barriers[0].arrive()  # Signal that tmp is ready.
        barriers[1].wait()  # Wait for the other warp to produce tmp.
        final_arr = arr + mgpu.FragmentedArray.load_strided(
            tmp, is_signed=False
        )
        final_arr.store_untiled(memref_slice(dst, 0), optimized=False)
        scf.yield_([])
      with ir.InsertionPoint(scf.IfOp(is_second_wg).then_block):
        barriers[0].wait()
        final_arr = arr + mgpu.FragmentedArray.load_strided(
            tmp, is_signed=False
        )
        barriers[2].arrive()
        barriers[2].wait()  # Synchronize this warpgroup before we overwrite tmp.
        arr.store_untiled(tmp)
        barriers[1].arrive()  # Signal that tmp is ready.
        final_arr.store_untiled(memref_slice(dst, 1), optimized=False)
        scf.yield_([])
    out_shape = jax.ShapeDtypeStruct((2, 128), jnp.int32)
    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (2 * 128, 1, 1),
        (),
        out_shape,
        (
            jax.ShapeDtypeStruct((128,), jnp.int32),
            mgpu.Barrier(arrival_count=128, num_barriers=3),
        ),
    )()
    np.testing.assert_array_equal(y, np.full_like(y, 3, dtype=np.int32))

  @parameterized.named_parameters(
      (
          f"_{''.join(map(str, collective_dims))}={collective_size}{'_' + ''.join(map(str, noncollective_dims)) if noncollective_dims else ''}{'_group' if group_dims else ''}",
          collective_dims,
          noncollective_dims,
          collective_size,
          group_dims,
      )
      for collective_dims in itertools.chain.from_iterable(
          itertools.combinations(Dimension, n) for n in range(1, 4)
      )
      for noncollective_dims in itertools.chain.from_iterable(
          itertools.combinations(Dimension, n) for n in range(3)
      )
      for collective_size in (1, 2, 4)
      for group_dims in (False,) + ((True,) if len(collective_dims) > 1 else ())
      if all(d not in noncollective_dims for d in collective_dims)
  )
  def test_collective_arrive(self, collective_dims, noncollective_dims, collective_size, group_dims):
    i32 = ir.IntegerType.get_signless(32)
    index = ir.IndexType.get()
    cluster = [1, 1, 1]
    for d in collective_dims:
      cluster[d] = collective_size
    for d in noncollective_dims:
      cluster[d] = 2
    if math.prod(cluster) > jtu.get_cuda_nonportable_max_cluster_size():
      self.skipTest("Cluster too big")
    is_trivial = math.prod(cluster[d] for d in collective_dims) == 1
    def kernel(ctx, dst, mask, collective_barrier):
      cluster_idx = utils.cluster_idx()
      if not is_trivial:
        memref.store(collective_barrier.cluster_mask, mask, [cluster_idx])
      else:
        assert collective_barrier.cluster_mask is None
      collective_barrier.arrive()
      collective_barrier.wait()
      tid = thread_idx()
      linear_idx = arith.index_cast(index, tid)
      stride = c(128, index)
      for d in gpu.Dimension:
        linear_idx = arith.addi(linear_idx, arith.muli(gpu.block_id(d), stride))
        stride = arith.muli(stride, gpu.grid_dim(d))
      memref.store(arith.index_cast(i32, linear_idx), dst, [linear_idx])
    out_shape = jax.ShapeDtypeStruct((math.prod(cluster) * 128,), jnp.int32)
    mask_shape = jax.ShapeDtypeStruct((math.prod(cluster),), jnp.int32)
    barrier_dims = collective_dims
    if group_dims:
      barrier_dims = (collective_dims[:2], *collective_dims[2:])
    scratch = mgpu.ClusterBarrier(barrier_dims)
    y, mask = mgpu.as_gpu_kernel(
        kernel, cluster, (128, 1, 1), (), (out_shape, mask_shape), scratch, cluster=cluster,
    )()
    np.testing.assert_array_equal(
        y, np.arange(math.prod(cluster) * 128, dtype=np.int32)
    )
    if not is_trivial:
      # Verify that the mask is correct. Blocks are column-major, hence the transpose.
      block_bits = 1 << np.arange(math.prod(cluster), dtype=np.int32).reshape(cluster[::-1]).T
      expected_mask = 0
      for bd in barrier_dims:
        if isinstance(bd, gpu.Dimension):
          bd = (bd,)
        least_significant_slice = tuple(
            slice(None) if d in bd else 0 for d in gpu.Dimension
        )
        mask_bits = block_bits[least_significant_slice]
        expected_mask |= np.bitwise_or.reduce(mask_bits, axis=None)
      self.assertEqual(min(mask), expected_mask)

  def test_collective_arrival_count(self):
    i32 = ir.IntegerType.get_signless(32)
    cluster = [2, 1, 1]
    def kernel(ctx, dst, collective_barrier):
      collective_barrier.arrive()
      collective_barrier.arrive()
      collective_barrier.arrive()
      collective_barrier.arrive()
      collective_barrier.wait()
      memref.store(arith.constant(i32, 1), dst, [])
    out_shape = jax.ShapeDtypeStruct((), jnp.int32)
    scratch = mgpu.ClusterBarrier((gpu.Dimension.x,), arrival_count=4)
    y = mgpu.as_gpu_kernel(
        kernel, cluster, (128, 1, 1), (), out_shape, scratch, cluster=cluster,
    )()
    np.testing.assert_array_equal(y, np.ones((), dtype=np.int32))

  @parameterized.product(
      collective_dim=(0, 1, 2),
      cluster=list(itertools.product((1, 2, 4), repeat=3)),
  )
  def test_leader_tracked_cluster_barrier(self, collective_dim, cluster):
    if cluster[collective_dim] == 1:
      self.skipTest("Collective dimension has trivial size")
    if math.prod(cluster) > jtu.get_cuda_nonportable_max_cluster_size():
      self.skipTest("Cluster too big")
    index = ir.IndexType.get()
    def kernel(_, cluster_masks, collective_barrier):
      block_id = utils.cluster_idx()
      collective_dim_id = gpu.cluster_block_id(gpu.Dimension(collective_dim))
      is_leader_block = arith.cmpi(
          arith.CmpIPredicate.eq, collective_dim_id, c(0, index)
      )
      collective_barrier.arrive()
      with when(is_leader_block):
        collective_barrier.wait()
      memref.store(collective_barrier.cluster_mask, cluster_masks, [block_id])
    out_shape = jax.ShapeDtypeStruct((math.prod(cluster),), jnp.int32)
    scratch = mgpu.ClusterBarrier(
        (gpu.Dimension(collective_dim),), arrival_count=1, leader_tracked=True,
    )
    cluster_masks = mgpu.as_gpu_kernel(
        kernel, cluster, (128, 1, 1), (), out_shape, scratch, cluster=cluster,
    )()
    expected_masks = []
    for i in range(math.prod(cluster)):
      # Cluster dimensions are minor-to-major.
      idx = np.unravel_index(i, cluster, order="F")
      mask_shift = 0
      stride = 1
      for dim in range(3):
        if dim == collective_dim:
          stride *= cluster[dim]
          continue
        mask_shift += idx[dim] * stride
        stride *= cluster[dim]
      expected_masks.append(1 << mask_shift)
    np.testing.assert_array_equal(cluster_masks, expected_masks)

  @parameterized.parameters(False, True)
  def test_mbarrier_complete_tx(self, predicated):
    i32 = ir.IntegerType.get_signless(32)

    def kernel(ctx, dst, mbar):
      mbar.arrive_expect_tx(1024 // 128)

      if predicated:
        is_leader = mgpu.single_thread_predicate(mgpu.ThreadSubset.BLOCK)
        mbar.complete_tx(1024, predicate=is_leader)
      else:
        mbar.complete_tx(1024 // 128)

      mbar.wait()

      with mgpu.single_thread(scope=mgpu.ThreadSubset.BLOCK):
        memref.store(arith.constant(i32, 1), dst, [])

    out_shape = jax.ShapeDtypeStruct((), jnp.int32)
    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        (),
        out_shape,
        mgpu.Barrier(arrival_count=128),
    )()

    np.testing.assert_array_equal(y, np.array(1, dtype=np.int32))

  def test_multiple_barriers_requires_explicit_indexing(self):
    def kernel(ctx, out, barrier):
      del ctx, out
      barrier.arrive()

    with self.assertRaisesRegex(
        ValueError, "Operation is only valid for a single barrier"
    ):
      mgpu.as_gpu_kernel(
          kernel,
          (1, 1, 1),
          (128, 1, 1),
          (),
          jax.ShapeDtypeStruct((), jnp.int32),
          mgpu.Barrier(arrival_count=1, num_barriers=3),
      )()


class AsyncCopyTest(TestCase, jtu.CudaArchSpecificTest):

  @parameterized.product(
      swizzle=(None, 32, 64, 128),
      shape=((64, None), (5, None), (2, 3, 5, None)),
      dtype=(jnp.float32, jnp.float16, jnp.int4),
  )
  def test_tma_load_basic(self, swizzle, shape, dtype):
    bw = bitwidth(dtype_to_ir_type(dtype))
    minor_size = 64 if swizzle is None else 8 * swizzle // bw
    shape = (*shape[:-1], minor_size)
    i1 = ir.IntegerType.get_signless(1)
    def kernel(ctx, src, dst, smem):
      tmp, barrier = smem
      ctx.async_copy(src_ref=src, dst_ref=tmp, swizzle=swizzle, barrier=barrier)
      barrier.wait_parity(c(0, i1))
      copy(tmp, dst, swizzle=swizzle)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    smem = (x, mgpu.Barrier(1))
    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)(x)
    np.testing.assert_array_equal(y, x)

  def test_tma_load_oob_over_trivial_gmem_dimension(self):
    src_shape = (64, 1, 64)
    dst_shape = (64, 2, 64)
    i1 = ir.IntegerType.get_signless(1)
    def kernel(ctx: launch_context.LaunchContext, src, dst, smem):
      tmp, barrier = smem
      ctx.async_copy(
          src_ref=src, dst_ref=tmp, swizzle=None, barrier=barrier,
          # Deliberately use a `gmem_slice` that is larger than the source GMEM.
          gmem_slice=(slice(None), slice(0, 2), slice(None)))
      barrier.wait_parity(c(0, i1))
      copy(tmp, dst, swizzle=None)
    x = np.arange(np.prod(src_shape), dtype=jnp.float32).reshape(src_shape)
    out_shape = jax.ShapeDtypeStruct(dst_shape, jnp.float32)
    smem = (out_shape, mgpu.Barrier(1))
    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, out_shape, smem)(x)
    np.testing.assert_array_equal(y[:, 0, :], x[:, 0, :])
    np.testing.assert_array_equal(y[:, 1, :], 0.0)

  def test_tma_load_unrecomputable_peer_without_num_peers_raises(self):
    # A block-derived peer id can't be replayed on the host and, without a
    # multi-device mesh, num_peers == 0, so the per-peer descriptor table
    # can't be sized.
    i32 = ir.IntegerType.get_signless(32)
    def kernel(ctx: launch_context.LaunchContext, src, dst, smem):
      del dst  # Lowering fails before the output is written.
      tmp, barrier = smem
      peer = arith.index_cast(i32, gpu.block_id(gpu.Dimension.x))
      # A swizzle rules out the contiguous bulk copy, forcing a TMA descriptor.
      ctx.async_copy(
          src_ref=src, dst_ref=tmp, swizzle=128, barrier=barrier,
          gmem_peer_id=peer,
      )
    # A 128B swizzle over int32 requires a minor dimension of 32 elements.
    x = np.arange(64 * 32, dtype=np.int32).reshape(64, 32)
    smem = (x, mgpu.Barrier(1))
    with self.assertRaisesRegex(
        NotImplementedError,
        "Peer ids that can't be recomputed on the host are unsupported",
    ):
      mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)

  @parameterized.product(
      rank=(3, 4, 5),
      offset=(0, 1),
      strided=(False, True),
      swizzle=(None, 32, 64, 128),
      dynamic=(False, True),
  )
  # TODO(slebedev): Remove once the minimum jaxlib version is 0.12.0.
  @absltest.skipIf(jaxlib_version < (0, 12, 0), "Requires jaxlib >=0.12.0")
  def test_tma_load_im2col(self, rank, offset, strided, swizzle, dynamic):
    dtype = jnp.int16
    tiled = swizzle is not None
    num_channels = 16 if swizzle is None else 8 * swizzle // 16
    num_rows = 128
    num_spatial = rank - 2
    spatial_sizes = (3, 4, 5)[-num_spatial:]
    grid_shape = (2, *spatial_sizes)
    src_shape = (*grid_shape, 2 * num_channels)
    first_channel = num_channels
    # Bounds, offsets, and strides differ per dimension to pin each one to its
    # dimension. Negative upper bounds are the common case, e.g. (-1, -1) for a
    # 3x3 SAME convolution, so alternate their sign.
    window_start_bounds = tuple(
        (-(d + 1), (-1) ** (d + 1) * (d + 1)) for d in range(num_spatial)
    )
    filter_offsets = tuple(offset * (d + 1) for d in range(num_spatial))
    if strided:
      window_strides = effective_strides = tuple(range(2, num_spatial + 2))
    else:
      window_strides = None
      effective_strides = (1,) * num_spatial
    windows = list(
        itertools.product(
            range(grid_shape[0]),
            *(
                range(lo, size + hi, s)
                for (lo, hi), size, s in zip(
                    window_start_bounds, spatial_sizes, effective_strides
                )
            ),
        )
    )
    # Start halfway through the first image, so that the copy wraps around
    # into the second one and, if fewer than num_rows windows remain, runs past
    # the end of the input.
    first_window = len(windows) // grid_shape[0] // 2
    start_indices = (*windows[first_window], first_channel)

    dst_shape = (num_rows, num_channels)
    if tiled:
      smem_shape = (num_rows // 8, 1, 8, num_channels)
    else:
      smem_shape = dst_shape
    i1 = ir.IntegerType.get_signless(1)

    def kernel(ctx: launch_context.LaunchContext, src, dst, smem):
      tmp, barrier = smem
      indices, offsets = start_indices, filter_offsets
      if dynamic:
        index = ir.IndexType.get()
        indices = [c(v, index) for v in indices]
        offsets = [c(v, index) for v in offsets]
      ctx.async_copy_im2col(
          src_ref=src,
          dst_ref=tmp,
          swizzle=swizzle,
          barrier=barrier,
          window_start_bounds=window_start_bounds,
          start_indices=indices,
          filter_offsets=offsets,
          window_strides=window_strides,
      )
      barrier.wait_parity(c(0, i1))
      copy(tmp, dst, swizzle=swizzle)

    x = np.arange(math.prod(src_shape), dtype=dtype).reshape(src_shape)
    out_shape = jax.ShapeDtypeStruct(smem_shape, dtype)
    smem = (out_shape, mgpu.Barrier(1))
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, out_shape, smem
    )(x).reshape(dst_shape)

    channels = slice(first_channel, first_channel + num_channels)
    expected = np.zeros(dst_shape, dtype=dtype)
    for row, (n, *window) in enumerate(
        windows[first_window : first_window + num_rows]
    ):
      pixel = (n, *(w + o for w, o in zip(window, filter_offsets)))
      if all(0 <= p < s for p, s in zip(pixel, grid_shape)):
        expected[row] = x[(*pixel, channels)]

    np.testing.assert_array_equal(y, expected)

  # TODO(slebedev): Remove once the minimum jaxlib version is 0.12.0.
  @absltest.skipIf(jaxlib_version < (0, 12, 0), "Requires jaxlib >=0.12.0")
  @jtu.thread_unsafe_test()
  def test_tma_load_im2col_reuses_descriptor(self):
    dtype = jnp.int16
    h, w, num_channels = 8, 16, 16
    num_rows = h * w
    valid, same = ((0, 0), (0, 0)), ((-1, -1), (-1, -1))
    copies = (
        (valid, (0, 0, 0, 0), (0, 0), None),
        (valid, (0, 0, 0, 0), (1, 1), (1, 1)),
        (same, (0, -1, -1, 0), (1, 1), None),
    )

    def kernel(ctx: launch_context.LaunchContext, src, dst, smem):
      tmp, barriers = smem
      for i, (bounds, indices, offsets, strides) in enumerate(copies):
        ctx.async_copy_im2col(
            src_ref=src,
            dst_ref=memref_slice(tmp, i),
            barrier=barriers[i],
            window_start_bounds=bounds,
            start_indices=indices,
            filter_offsets=offsets,
            window_strides=strides,
        )
      for i in range(len(copies)):
        barriers[i].wait()
      copy(tmp, dst)

    x = np.arange(1, h * w * num_channels + 1, dtype=dtype).reshape(
        1, h, w, num_channels
    )
    out_shape = jax.ShapeDtypeStruct(
        (len(copies), num_rows, num_channels), dtype
    )
    smem = (out_shape, mgpu.Barrier(1, num_barriers=len(copies)))
    env_vars = {"MOSAIC_GPU_DUMP_HOST_LLVM": "1"}
    with jtu.set_env(**env_vars), self.capture_stdout() as llvm_ir:
      y = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x, out_shape, smem
      )(x)
    self.assertEqual(
        llvm_ir().count("call void @mosaic_gpu_init_tma_im2col_desc("), 2
    )
    image = x[0].reshape(num_rows, num_channels)
    shifted = np.pad(x[0, 1:, 1:], ((0, 1), (0, 1), (0, 0)))
    np.testing.assert_array_equal(y[0], image)
    np.testing.assert_array_equal(
        y[1], shifted.reshape(num_rows, num_channels)
    )
    np.testing.assert_array_equal(y[2], image)

  @parameterized.product(
      swizzle=(None, 32, 64, 128),
      shape=((64, None), (5, None), (2, 3, 5, None)),
      dtype=(jnp.float32, jnp.float16, jnp.int4),
  )
  def test_tma_prefetch_basic(self, swizzle, shape, dtype):
    bw = bitwidth(dtype_to_ir_type(dtype))
    minor_size = 64 if swizzle is None else 8 * swizzle // bw
    shape = (*shape[:-1], minor_size)
    i1 = ir.IntegerType.get_signless(1)
    def kernel(ctx, src, dst, smem):
      tmp, barrier = smem
      ctx.async_prefetch(gmem_ref=src, swizzle=swizzle)
      ctx.async_copy(src_ref=src, dst_ref=tmp, swizzle=swizzle, barrier=barrier)
      barrier.wait_parity(c(0, i1))
      copy(tmp, dst, swizzle=swizzle)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    smem = (x, mgpu.Barrier(1))
    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)(x)
    np.testing.assert_array_equal(y, x)

  @parameterized.product(
      swizzle=(16, 32, 64, 128),
      shape=((64, None),),
      dtype=(jnp.int32, jnp.int16),
      idx_dtype=(jnp.int32, jnp.int8),
      idx_layout=(
          fa.TMA_INDICES_LAYOUT,
          fa.TMA_INDICES_4_LAYOUT,
      )
  )
  def test_tma_gather_basic(self, swizzle, shape, dtype, idx_dtype, idx_layout):
    if not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("TMA gather requires CUDA compute capability 10.0 or higher")
    i1 = ir.IntegerType.get_signless(1)
    swizzle_elems = 8 * swizzle // bitwidth(dtype_to_ir_type(dtype))
    col_slice = swizzle_elems if swizzle != 16 else 128
    shape = (*shape[:-1], 2 * col_slice)
    def kernel(ctx, src, idx, dst, smem):
      tmp, barrier = smem
      idxs = mgpu.FragmentedArray.load_untiled(
          idx, layout=idx_layout, optimized=False, is_signed=False
      )
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          swizzle=swizzle,
          barrier=barrier,
          gmem_slice=(idxs, mgpu.ds(col_slice, col_slice)),
      )
      barrier.wait_parity(c(0, i1))
      copy(tmp, dst, swizzle=swizzle)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    num_items = 12 if idx_layout == fa.TMA_INDICES_4_LAYOUT else 48
    idx = jax.random.permutation(jax.random.key(1234), 48).astype(idx_dtype)[:num_items]
    out_type = jax.ShapeDtypeStruct((num_items, col_slice), dtype)
    smem = (out_type, mgpu.Barrier(1))
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (x, idx), out_type, smem,
    )(x, idx)
    np.testing.assert_array_equal(y, x[idx, slice(col_slice, 2 * col_slice)])

  @parameterized.product(
      swizzle=(16, 32, 64, 128),
      shape=((64, None),),
      dtype=(jnp.int32, jnp.int16),
      transpose_tiles=(False, True),
      idx_layout=(
          fa.TMA_INDICES_LAYOUT,
          fa.TMA_INDICES_4_LAYOUT,
      )
  )
  def test_tma_gather_tiled(self, swizzle, shape, dtype, transpose_tiles, idx_layout):
    if not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("TMA gather requires CUDA compute capability 10.0 or higher")
    i1 = ir.IntegerType.get_signless(1)
    swizzle_elems = 8 * swizzle // bitwidth(dtype_to_ir_type(dtype))
    col_slice = swizzle_elems if swizzle != 16 else 128
    shape = (*shape[:-1], 3 * col_slice)
    # Using (8, swizzle_elems) produces too short transfers (we'd end up with
    # misaligned SMEM addresses).
    tiling = (8, swizzle_elems) if swizzle != 16 else (8, 2 * swizzle_elems)
    if transpose_tiles:
      transforms = (mgpu.TileTransform(tiling), mgpu.TransposeTransform((1, 0, 2, 3)))
    else:
      transforms = mgpu.TileTransform(tiling)
    def kernel(ctx, src, idx, dst, smem):
      tmp, barrier = smem
      idxs = mgpu.FragmentedArray.load_untiled(
          idx, layout=idx_layout, optimized=False, is_signed=False
      )
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          swizzle=swizzle,
          barrier=barrier,
          gmem_slice=(idxs, mgpu.ds(col_slice, 2 * col_slice)),
          gmem_transform=transforms,
      )
      barrier.wait_parity(c(0, i1))
      ctx.async_copy(
          src_ref=tmp,
          dst_ref=dst,
          swizzle=swizzle,
          gmem_transform=transforms,
      )
      ctx.await_async_copy(0)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    num_items = 12 if idx_layout == fa.TMA_INDICES_4_LAYOUT else 48
    padded_num_items = (num_items + tiling[0] - 1) // tiling[0] * tiling[0]
    idx = jax.random.permutation(jax.random.key(1234), padded_num_items).astype(jnp.int32)
    out_type = jax.ShapeDtypeStruct((padded_num_items, 2 * col_slice), dtype)
    smem_shape = tile_shape((padded_num_items, 2 * col_slice), tiling)
    if transpose_tiles:
      smem_shape = (smem_shape[1], smem_shape[0], *smem_shape[2:])
    smem = (
        jax.ShapeDtypeStruct(smem_shape, dtype),
        mgpu.Barrier(1),
    )
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (x, idx), out_type, smem,
    )(x, idx)
    np.testing.assert_array_equal(y[:num_items], x[idx[:num_items], slice(col_slice, 3 * col_slice)])

  @parameterized.product(
      swizzle=(16, 32, 64, 128),
      dtype=(jnp.int32, jnp.int16),
      transpose_tiles=(False, True),
      idx_layout=(
          fa.TMA_INDICES_LAYOUT,
          fa.TMA_INDICES_4_LAYOUT,
      )
  )
  def test_tma_scatter_tiled(self, swizzle, dtype, transpose_tiles, idx_layout):
    self.skip_unless_tcgen05()  # .tile::scatter4 is not supported on sm_120, for example
    swizzle_elems = 8 * swizzle // bitwidth(dtype_to_ir_type(dtype))
    # Using `swizzle_elems` for `swizzle == 16` produces too short transfers
    # that result in misaligned SMEM addresses.
    n_cols = swizzle_elems if swizzle != 16 else 128
    num_items = 12 if idx_layout == fa.TMA_INDICES_4_LAYOUT else 64
    shape = (num_items, n_cols)
    tiling = (8, swizzle_elems) if swizzle != 16 else (8, 2 * swizzle_elems)
    padded_num_items = (num_items + tiling[0] - 1) // tiling[0] * tiling[0]
    shape = (padded_num_items, n_cols)
    if transpose_tiles:
      transforms = (
          mgpu.TileTransform(tiling),
          mgpu.TransposeTransform((1, 0, 2, 3)),
      )
    else:
      transforms = mgpu.TileTransform(tiling)

    def kernel(ctx, src, idx, dst, smem):
      tmp, barrier = smem
      idxs = mgpu.FragmentedArray.load_untiled(
          idx,
          layout=idx_layout,
          optimized=False,
          is_signed=False,
      )
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          swizzle=swizzle,
          barrier=barrier,
          gmem_transform=transforms,
      )
      barrier.wait()
      ctx.async_copy(
          src_ref=tmp,
          dst_ref=dst,
          swizzle=swizzle,
          gmem_slice=(idxs, mgpu.ds(n_cols, n_cols)),
          gmem_transform=transforms,
      )
      ctx.await_async_copy(0)

    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    idx = jax.random.permutation(jax.random.key(1234), 128)
    idx = idx[:padded_num_items].astype(jnp.int32)
    out_shape = (128, 2 * n_cols)
    out_type = jax.ShapeDtypeStruct(out_shape, dtype)
    smem_shape = tile_shape(shape, tiling)
    if transpose_tiles:
      smem_shape = (smem_shape[1], smem_shape[0], *smem_shape[2:])
    smem = (jax.ShapeDtypeStruct(smem_shape, dtype), mgpu.Barrier(1))
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (x, idx), out_type, smem
    )(x, idx)
    np.testing.assert_array_equal(y[idx[:num_items], slice(n_cols, 2 * n_cols)], x[:num_items])

  def test_tma_with_1d_tiling(self):
    swizzle = 128
    dtype = jnp.float16
    shape = (64, 128)
    tiling = (1, swizzle // jnp.dtype(dtype).itemsize)
    def kernel(ctx, dst, smem):
      iota_tensor(*shape, dtype=dtype).store_tiled(smem, swizzle=swizzle)
      ctx.async_copy(
          src_ref=smem,
          dst_ref=dst,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
      )
      ctx.await_async_copy(0)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    smem = jax.ShapeDtypeStruct(utils.tile_shape(shape, tiling), dtype)
    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), (), x, smem)()
    np.testing.assert_array_equal(y, x)

  @parameterized.named_parameters(
      (
          f"_{''.join(map(str, collective_dims))}={collective_size}{'_' + ''.join(map(str, noncollective_dims)) if noncollective_dims else ''}",
          collective_dims,
          noncollective_dims,
          collective_size,
      )
      for collective_dims in itertools.chain.from_iterable(
          itertools.combinations(Dimension, n) for n in range(1, 4)
      )
      for noncollective_dims in itertools.chain.from_iterable(
          itertools.combinations(Dimension, n) for n in range(3)
      )
      for collective_size in (1, 2, 4)
      if all(d not in noncollective_dims for d in collective_dims)
  )
  def test_tma_load_multicast(self, collective_dims, noncollective_dims, collective_dim_size):
    index = ir.IndexType.get()
    swizzle = 128
    dtype = jnp.float16
    cluster = [1, 1, 1]
    for d in collective_dims:
      cluster[d] = collective_dim_size
    for d in noncollective_dims:
      cluster[d] = 2
    if math.prod(cluster) > jtu.get_cuda_nonportable_max_cluster_size():
      self.skipTest("Cluster too big")
    collective_size = math.prod(cluster[d] for d in collective_dims)
    noncollective_size = math.prod(cluster) // collective_size
    # We use the 2 dimension to exercise splitting the collective over
    # multiple dimensions when the cluster is large.
    shape = (noncollective_size, 2, 16 * collective_size, 64)
    minor_size = 64 if swizzle is None else swizzle // jnp.dtype(dtype).itemsize
    shape = (*shape[:-1], minor_size)
    # Note that this kernel does not use the non-collective dimensions in any
    # interesting way and so they don't really have to be part of the cluster.
    # We use them to test that the multicast mask is generated correctly.
    def kernel(ctx, src, dst, scratch):
      tmp, barrier = scratch
      stride = 1
      noncollective_idx = c(0, index)
      for d in noncollective_dims:
        noncollective_idx = arith.addi(
            noncollective_idx,
            arith.muli(gpu.cluster_block_id(d), c(stride, index))
        )
        stride *= cluster[d]
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          gmem_slice=(noncollective_idx,),
          swizzle=swizzle,
          barrier=barrier,
          collective=collective_dims,
      )
      barrier.wait()
      # This is _not_ the real cluster block idx, because it does not consider
      # the column-major ordering of the grid dimensions.
      idx = c(0, index)
      stride = 1
      for d in collective_dims:
        idx = arith.addi(
            idx, arith.muli(gpu.cluster_block_id(d), c(stride, index))
        )
        stride *= cluster[d]
      idx_minor = arith.divui(idx, c(2, index))
      idx_major = arith.remui(idx, c(2, index))
      slc_minor = ds(
          arith.muli(idx_minor, c(16 * 2, index)), 16 * 2
      )
      copy(
          memref_slice(tmp, (idx_major, slc_minor)),
          memref_slice(dst, (noncollective_idx, idx_major, slc_minor)),
          swizzle=swizzle,
      )
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    smem_shape = (jax.ShapeDtypeStruct(shape[1:], dtype), mgpu.Barrier(1))
    y = mgpu.as_gpu_kernel(
        kernel, cluster, (128, 1, 1), x, x, smem_shape, cluster=cluster
    )(x)
    np.testing.assert_array_equal(y, x)

  @parameterized.product(
      swizzle=(None, 128),
      dtype=(jnp.float16, jnp.float32),
      cluster=((2, 1, 1), (2, 2, 1)),
  )
  def test_tma_load_replicated(self, swizzle, dtype, cluster):
    self.skip_unless_tcgen05() # .cta_group::2 is not supported on sm_120, for example
    index = ir.IndexType.get()
    minor_size = 64 if swizzle is None else swizzle // jnp.dtype(dtype).itemsize
    shape = (8, minor_size)
    cluster_y = cluster[1]
    def kernel(ctx, src, dst, scratch):
      tmp, barrier, cluster_barrier = scratch
      block_id_x = gpu.cluster_block_id(gpu.Dimension.x)
      block_id_y = gpu.cluster_block_id(gpu.Dimension.y)
      is_first_block = arith.cmpi(
          arith.CmpIPredicate.eq, block_id_x, c(0, index)
      )
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          gmem_slice=(block_id_y, slice(None), slice(None)),
          swizzle=swizzle,
          barrier=barrier,
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.REPLICATED,
      )
      with when(is_first_block):
        barrier.wait()
      cluster_barrier.arrive()
      cluster_barrier.wait()
      copy(tmp, memref_slice(dst, (block_id_y, block_id_x)), swizzle=swizzle)
    x = np.arange(cluster_y * np.prod(shape), dtype=dtype).reshape(cluster_y, *shape)
    smem_shape = (
        jax.ShapeDtypeStruct(shape, dtype),
        mgpu.Barrier(1),
        mgpu.ClusterBarrier(collective_dims=(gpu.Dimension.x,)),
    )
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      y = mgpu.as_gpu_kernel(
          kernel, cluster, (128, 1, 1), x,
          jax.ShapeDtypeStruct((cluster_y, 2, *shape), dtype), smem_shape,
          cluster=cluster,
      )(x)
    np.testing.assert_array_equal(y[:, 0], x)
    np.testing.assert_array_equal(y[:, 1], x)
    self.assertIn("cta_group::2", ptx())

  @parameterized.product(
      swizzle=(None, 128),
      dtype=(jnp.float16, jnp.float32),
      cluster=((2, 1, 1), (2, 2, 1)),
  )
  def test_tma_load_partitioned(self, swizzle, dtype, cluster):
    self.skip_unless_tcgen05() # .cta_group::2 is not supported on sm_120, for example
    index = ir.IndexType.get()
    minor_size = 64 if swizzle is None else swizzle // jnp.dtype(dtype).itemsize
    shape = (16, minor_size)
    half = shape[0] // 2
    smem_shape_val = (half, minor_size)
    cluster_y = cluster[1]
    def kernel(ctx, src, dst, scratch):
      tmp, barrier, cluster_barrier = scratch
      block_id_x = gpu.cluster_block_id(gpu.Dimension.x)
      block_id_y = gpu.cluster_block_id(gpu.Dimension.y)
      is_first_block = arith.cmpi(
          arith.CmpIPredicate.eq, block_id_x, c(0, index)
      )
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          gmem_slice=(block_id_y, slice(None), slice(None)),
          swizzle=swizzle,
          barrier=barrier,
          collective=gpu.Dimension.x,
          leader_tracked=mgpu.CopyPartition.PARTITIONED(0),
      )
      with when(is_first_block):
        barrier.wait()
      cluster_barrier.arrive()
      cluster_barrier.wait()
      copy(tmp, memref_slice(dst, (block_id_y, block_id_x)), swizzle=swizzle)
    x = np.arange(cluster_y * np.prod(shape), dtype=dtype).reshape(cluster_y, *shape)
    smem_shape = (
        jax.ShapeDtypeStruct(smem_shape_val, dtype),
        mgpu.Barrier(1),
        mgpu.ClusterBarrier(collective_dims=(gpu.Dimension.x,)),
    )
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      y = mgpu.as_gpu_kernel(
          kernel, cluster, (128, 1, 1), x,
          jax.ShapeDtypeStruct((cluster_y, 2, *smem_shape_val), dtype), smem_shape,
          cluster=cluster,
      )(x)
    np.testing.assert_array_equal(y[:, 0], x[:, :half])
    np.testing.assert_array_equal(y[:, 1], x[:, half:])
    self.assertIn("cta_group::2", ptx())


  @parameterized.product(
      swizzle=(None, 128),
      shape=((128, 128), (5, 32, 128)),
      dtype=(jnp.float16, jnp.float32),
  )
  def test_tma_load_tiled(self, swizzle, shape, dtype):
    # TODO(apaszke): ptxas seems to freeze when generating code for copy with
    # swizzle 32 and 64.
    i1 = ir.IntegerType.get_signless(1)
    index = ir.IndexType.get()
    tiling = (32, (swizzle or 128) // jnp.dtype(dtype).itemsize)
    tiled_shape = tile_shape(shape, tiling)[:len(shape)]
    def kernel(ctx, src, dst, scratch):
      tmp, barrier = scratch
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          swizzle=swizzle,
          barrier=barrier,
          gmem_transform=mgpu.TileTransform(tiling),
      )
      barrier.wait_parity(c(0, i1))
      for idxs in np.ndindex(tiled_shape):
        untiled_idxs, tiled_idxs = idxs[:-len(tiling)], idxs[-len(tiling):]
        s = (
            *untiled_idxs,
            *(ds(c(ix * t, index), t) for ix, t in zip(tiled_idxs, tiling)),
        )
        copy(memref_slice(tmp, idxs), memref_slice(dst, s), swizzle=swizzle)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    smem = (
        jax.ShapeDtypeStruct(tile_shape(shape, tiling), dtype),
        mgpu.Barrier(1),
    )
    f = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)
    y = f(x)
    np.testing.assert_array_equal(y, x)

  @parameterized.product(
      swizzle=(None, 128),
      shape=((128, 128), (5, 32, 128)),
      dtype=(jnp.float16, jnp.float32),
  )
  @jtu.thread_unsafe_test()
  def test_tma_prefetch_tiled(self, swizzle, shape, dtype):
    # TODO(apaszke): ptxas seems to freeze when generating code for copy with
    # swizzle 32 and 64.
    i1 = ir.IntegerType.get_signless(1)
    index = ir.IndexType.get()
    tiling = (32, (swizzle or 128) // jnp.dtype(dtype).itemsize)
    tiled_shape = tile_shape(shape, tiling)[:len(shape)]
    def kernel(ctx, src, dst, scratch):
      tmp, barrier = scratch
      ctx.async_prefetch(
          gmem_ref=src, swizzle=swizzle, gmem_transform=mgpu.TileTransform(tiling)
      )
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          swizzle=swizzle,
          barrier=barrier,
          gmem_transform=mgpu.TileTransform(tiling),
      )
      barrier.wait_parity(c(0, i1))
      for idxs in np.ndindex(tiled_shape):
        untiled_idxs, tiled_idxs = idxs[:-len(tiling)], idxs[-len(tiling):]
        s = (
            *untiled_idxs,
            *(ds(c(ix * t, index), t) for ix, t in zip(tiled_idxs, tiling)),
        )
        copy(memref_slice(tmp, idxs), memref_slice(dst, s), swizzle=swizzle)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    smem = (
        jax.ShapeDtypeStruct(tile_shape(shape, tiling), dtype),
        mgpu.Barrier(1),
    )
    env_vars = {
        "MOSAIC_GPU_DUMP_HOST_LLVM": "1",
        "MOSAIC_GPU_DUMP_PTX": "1",
    }
    with jtu.set_env(**env_vars), self.capture_stdout() as ptx_llvm_ir:
      f = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)
      y = f(x)
    # We should only create one descriptor for both prefetch and copy.
    self.assertEqual(ptx_llvm_ir().count("call void @mosaic_gpu_init_tma_desc("), 1)
    self.assertIn("cp.async.bulk.prefetch.tensor", ptx_llvm_ir())
    np.testing.assert_array_equal(y, x)

  @parameterized.product(swizzle=(None, 128))
  def test_tma_load_tiled_rounding(self, swizzle):
    # TODO(apaszke): ptxas seems to freeze when generating code for copy with
    # swizzle 32 and 64.
    shape = (5, 32, 144)
    dtype = jnp.float16
    i1 = ir.IntegerType.get_signless(1)
    index = ir.IndexType.get()
    tiling = (32, (swizzle or 128) // jnp.dtype(dtype).itemsize)
    rounded_shape = (*shape[:-1], shape[-1] // tiling[-1] * tiling[-1])
    tiled_shape = tile_shape(rounded_shape, tiling)[:len(shape)]
    def kernel(ctx, src, dst, scratch):
      tmp, barrier = scratch
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          swizzle=swizzle,
          barrier=barrier,
          gmem_transform=mgpu.TileTransform(
              tiling, rounding=launch_context.Rounding.DOWN
          ),
      )
      barrier.wait_parity(c(0, i1))
      for idxs in np.ndindex(tiled_shape):
        untiled_idxs, tiled_idxs = idxs[:-len(tiling)], idxs[-len(tiling):]
        s = (
            *untiled_idxs,
            *(ds(c(ix * t, index), t) for ix, t in zip(tiled_idxs, tiling)),
        )
        copy(memref_slice(tmp, idxs), memref_slice(dst, s), swizzle=swizzle)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    tmp_shape = jax.ShapeDtypeStruct(tile_shape(rounded_shape, tiling), dtype)
    smem = (tmp_shape, mgpu.Barrier(1))
    out_shape = jax.ShapeDtypeStruct(rounded_shape, dtype)
    f = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, out_shape, smem)
    y = f(x)
    np.testing.assert_array_equal(y, x[..., :rounded_shape[-1]])

  def test_tma_load_indexed_tiled(self):
    shape = (128, 2, 128)
    tiling = mgpu.TileTransform((32, 32))
    def kernel(ctx, src, dst, scratch):
      tmp, barrier = scratch
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          barrier=barrier,
          gmem_transform=tiling,
          gmem_slice=(slice(None), 1, slice(None)),
      )
      barrier.wait()
      ctx.async_copy(src_ref=tmp, dst_ref=dst, gmem_transform=tiling)
      ctx.await_async_copy(0)
    x = np.arange(np.prod(shape), dtype=jnp.float32).reshape(shape)
    smem = (
        jax.ShapeDtypeStruct((4, 4, 32, 32), jnp.float32),
        mgpu.Barrier(1),
    )
    out_shape = jax.ShapeDtypeStruct((128, 128), jnp.float32)
    f = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, out_shape, smem)
    np.testing.assert_array_equal(f(x), x[:, 1, :])

  @parameterized.product(
      swizzle=(None, 128),
      dtype=(jnp.float16, jnp.float32),
  )
  def test_tma_squeeze_indexing(self, swizzle, dtype):
    # TODO(apaszke): ptxas seems to freeze when generating code for copy with
    # swizzle 32 and 64.
    minor_size = 64 if swizzle is None else swizzle // jnp.dtype(dtype).itemsize
    shape = (4, 5, minor_size)
    def kernel(ctx, src, dst, smem):
      tmp, barrier = smem
      for i in range(4):
        ctx.async_copy(
            src_ref=src,
            dst_ref=memref_slice(tmp, i),
            gmem_slice=i,
            swizzle=swizzle,
            barrier=barrier,
        )
        barrier.wait()
      copy(tmp, dst, swizzle=swizzle)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    smem = (x, mgpu.Barrier(1))
    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)(x)
    np.testing.assert_array_equal(y, x)

  def test_parity_tracking(self):
    shape = (16, 64)
    index = ir.IndexType.get()
    def kernel(ctx, src, dst, smem):
      tmp, barrier = smem
      for i in range(shape[0]):
        s = ds(c(i, index), 1)
        ctx.async_copy(
            src_ref=src, dst_ref=tmp, gmem_slice=s, barrier=barrier,
        )
        barrier.wait()
        copy(tmp, memref_slice(dst, s))
    x = np.arange(np.prod(shape), dtype=jnp.float16).reshape(shape)
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, x, (x[0:1], mgpu.Barrier(1))
    )(x)
    np.testing.assert_array_equal(y, x)

  @parameterized.product(
      swizzle=(None, 32, 64, 128),
      shape=((64, None), (5, None), (2, 3, 5, None)),
      dtype=(jnp.float16, jnp.float32),
  )
  def test_tma_store(self, swizzle, shape, dtype):
    minor_size = 64 if swizzle is None else swizzle // jnp.dtype(dtype).itemsize
    shape = (*shape[:-1], minor_size)
    def kernel(ctx, src, dst, tmp):
      copy(src, tmp, swizzle=swizzle)
      ctx.async_copy(src_ref=tmp, dst_ref=dst, swizzle=swizzle)
      ctx.await_async_copy(0)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, x)(x)
    np.testing.assert_array_equal(y, x)

  @parameterized.parameters(0, 1)
  def test_tma_small_tile_load(self, small_dim):
    if small_dim == 0:
      shape = (4, 128)
    elif small_dim == 1:
      shape = (128, 8)
    else:
      raise ValueError("small_dim must be 0 or 1")
    tiled_shape = ((shape[0] + 63) // 64, (shape[1] + 63) // 64, 64, 64)
    padded_shape = (math.prod(tiled_shape[0::2]), math.prod(tiled_shape[1::2]))
    def kernel(ctx, src, dst, smem):
      tmp, barrier = smem
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          swizzle=128,
          gmem_transform=mgpu.TileTransform((64, 64)),
          gmem_slice=(ds(0, padded_shape[0]), ds(0, padded_shape[1])),
          barrier=barrier,
      )
      barrier.wait()
      copy(tmp, dst, swizzle=128)
    x = np.arange(np.prod(shape), dtype=jnp.float16).reshape(shape)
    tiled = jax.ShapeDtypeStruct(tiled_shape, jnp.float16)
    y_tiled = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, tiled, (tiled, mgpu.Barrier(1)),
    )(x)
    y = y_tiled.swapaxes(1, 2).reshape(padded_shape)
    # y should contain x and zero everywhere else.
    np.testing.assert_array_equal(y[:shape[0], :shape[1]], x)
    y_mut = np.asarray(y).copy()
    y_mut[:shape[0], :shape[1]] = 0
    np.testing.assert_array_equal(y_mut, np.zeros_like(y_mut))

  @parameterized.product(small_dim=(0, 1),)
  def test_tma_small_tile_store(self, small_dim):
    if small_dim == 0:
      shape = (4, 128)
    elif small_dim == 1:
      shape = (128, 8)
    else:
      raise ValueError("small_dim must be 0 or 1")
    tiled_shape = ((shape[0] + 63) // 64, (shape[1] + 63) // 64, 64, 64)
    m, n = (math.prod(tiled_shape[0::2]), math.prod(tiled_shape[1::2]))
    def kernel(ctx, dst, tmp):
      vals = iota_tensor(m, n, jnp.float16)
      vals.store_tiled(tmp, swizzle=128)
      ctx.async_copy(
          src_ref=tmp,
          dst_ref=dst,
          swizzle=128,
          gmem_transform=mgpu.TileTransform((64, 64)),
          gmem_slice=(ds(0, m), ds(0, n)),
      )
      ctx.await_async_copy(0)
    tiled = jax.ShapeDtypeStruct(tiled_shape, jnp.float16)
    out = jax.ShapeDtypeStruct(shape, jnp.float16)
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out, tiled,
    )()
    iota = np.arange(m * n, dtype=jnp.float16).reshape([m, n])
    np.testing.assert_array_equal(y, iota[:shape[0], :shape[1]])

  def test_tma_invalid(self):
    def kernel(ctx, src, dst, tmp):
      copy(src, tmp)
      ctx.async_copy(src_ref=tmp, dst_ref=dst)
      ctx.await_async_copy(0)

    def run_kernel(shape):
      x = np.arange(np.prod(shape), dtype=np.int32).reshape(shape)
      _ = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, x)(x)

    with self.assertRaisesRegex(ValueError, "all GMEM strides except the last"):
      run_kernel([1] * 6)

    with self.assertRaisesRegex(
        ValueError, "last dimension to be divisible by 128"
    ):
      run_kernel([23])

  def _test_cp_async(
      self, shape, dtype, swizzle=None, tiling=None, use_barrier=True
  ):

    def kernel(ctx, src, dst, scratch):
      tmp, barrier = scratch
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling) if tiling else (),
          implementation=mgpu.AsyncCopyImplementation.CP_ASYNC,
          barrier=barrier if use_barrier else None,
      )
      if use_barrier:
        barrier.wait()
      else:
        ctx.await_cp_async_copy(0)
      if tiling:
        mgpu.copy_tiled(tmp, dst, swizzle=swizzle)
      else:
        copy(tmp, dst)

    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    smem_shape = mgpu.tile_shape(shape, tiling) if tiling else shape
    smem = jax.ShapeDtypeStruct(smem_shape, dtype)
    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        x,
        x,
        (smem, mgpu.Barrier(arrival_count=1)),
    )(x)
    np.testing.assert_array_equal(y, x)

  @parameterized.product(
      swizzle=(16, 32, 64, 128),
      shape=((64, 128), (128, 32)),
      dtype=(jnp.float32, jnp.float16, jnp.float8_e5m2, jnp.int4),
      use_barrier=(False, True),
  )
  def test_cp_async_tiled(self, swizzle, shape, dtype, use_barrier):
    bw = bitwidth(dtype_to_ir_type(dtype))
    swizzle_elems = 8 * swizzle // bw
    tiling = (8, swizzle_elems)
    if shape[-1] % swizzle_elems:
      self.skipTest("Minor dimension not divisible by swizzle_elems")
    self._test_cp_async(
        shape, dtype, swizzle=swizzle, tiling=tiling, use_barrier=use_barrier
    )

  @parameterized.product(
      swizzle=(32, 128),
      dtype=(jnp.float32, jnp.float16),
      use_barrier=(False, True),
  )
  def test_cp_async_tiled_sliced(self, swizzle, dtype, use_barrier):
    bw = bitwidth(dtype_to_ir_type(dtype))
    swizzle_elems = 8 * swizzle // bw
    tiling = (8, swizzle_elems)

    full_shape = (256, 128)
    slices = (slice(64, 128), slice(32, 96))
    slice_shape = tuple(s.stop - s.start for s in slices)

    def kernel(ctx, src, dst, scratch):
      tmp, barrier = scratch
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          gmem_slice=slices,
          swizzle=swizzle,
          gmem_transform=mgpu.TileTransform(tiling),
          implementation=mgpu.AsyncCopyImplementation.CP_ASYNC,
          barrier=barrier if use_barrier else None,
      )
      if use_barrier:
        barrier.wait()
      else:
        ctx.await_cp_async_copy(0)
      mgpu.copy_tiled(tmp, dst, swizzle=swizzle)

    x = np.arange(np.prod(full_shape), dtype=dtype).reshape(full_shape)
    smem_shape = mgpu.tile_shape(slice_shape, tiling)
    smem = jax.ShapeDtypeStruct(smem_shape, dtype)
    out = jax.ShapeDtypeStruct(slice_shape, dtype)
    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        x,
        out,
        (smem, mgpu.Barrier(arrival_count=1)),
    )(x)
    np.testing.assert_array_equal(y, x[slices])

  @parameterized.product(
      shape=((64, 128), (128, 40), (5, 256)),
      dtype=(jnp.float32, jnp.float16),
      use_barrier=(False, True),
  )
  def test_cp_async_untiled(self, shape, dtype, use_barrier):
    self._test_cp_async(shape, dtype, use_barrier=use_barrier)

  @parameterized.product(
      slice=[(1,), (1, 1)],
      use_barrier=(False, True),
  )
  def test_cp_async_untiled_squeezed_dims(self, slice, use_barrier):
    def kernel(ctx, src, dst, scratch):
      tmp, barrier = scratch
      ctx.async_copy(
          src_ref=src,
          dst_ref=tmp,
          gmem_slice=slice,
          implementation=mgpu.AsyncCopyImplementation.CP_ASYNC,
          barrier=barrier if use_barrier else None,
      )
      if use_barrier:
        barrier.wait()
      else:
        ctx.await_cp_async_copy(0)
      copy(tmp, dst)

    shape = (2, 4, 128)
    dtype = jnp.float32
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    expected = x[slice]

    y = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        x,
        jax.ShapeDtypeStruct(expected.shape, dtype),
        (jax.ShapeDtypeStruct(expected.shape, dtype), mgpu.Barrier(arrival_count=1)),
    )(x)
    np.testing.assert_array_equal(y, expected)


  def test_tma_collective_async_cp_with_no_swizzle(self):
    def body(ctx, src, dst, scratch):
      tmp, barrier = scratch
      ctx.async_copy(
          src_ref=src, dst_ref=tmp, collective=gpu.Dimension.x, barrier=barrier
      )
      barrier.wait()
      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      ctx.async_copy(src_ref=tmp, dst_ref=dst, gmem_slice=block_id)

    dtype = jnp.float32
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(2, 1, 1),
        cluster=(2, 1, 1),
        block=(128, 1, 1),
        in_shape=jax.ShapeDtypeStruct((128,), dtype),
        out_shape=jax.ShapeDtypeStruct((2, 128), dtype),
        smem_scratch_shape=[
            jax.ShapeDtypeStruct((128,), dtype),
            mgpu.Barrier(1),
        ],
    )
    x = jnp.arange(128, dtype=jnp.float32)
    np.testing.assert_array_equal(kernel(x), jnp.stack([x, x], axis=0))


class FragmentedArrayTest(TestCase):

  @parameterized.product(
      op=(
          operator.add,
          operator.mul,
          operator.sub,
          (mgpu.FragmentedArray.min, np.minimum),
          (mgpu.FragmentedArray.max, np.maximum),
      ),
      # TODO(apaszke): Enable float8
      dtype=[jnp.float32, jnp.int32, jnp.uint32, jnp.float16],
      m=(64, 128),
      n=(8, 64, 256),
  )
  @jtu.ignore_warning(
      message="(invalid value|divide by zero)", category=RuntimeWarning
  )
  def test_binary(self, op, dtype, m=64, n=32):
    if isinstance(op, tuple):
      op, np_op = op
    else:
      np_op = op

    for scalar_rhs in [None, 2]:
      def kernel(ctx, lhs, rhs, dst, _):
        mlir_dtype = utils.dtype_to_ir_type(dtype)
        lhs = mgpu.FragmentedArray.load_strided(lhs, is_signed=utils.is_signed(dtype))
        if scalar_rhs is None:
          rhs = mgpu.FragmentedArray.load_strided(rhs, is_signed=utils.is_signed(dtype))
        else:
          rhs = c(scalar_rhs, mlir_dtype)
        op(lhs, rhs).store_untiled(dst, optimized=False)
      if jnp.issubdtype(dtype, jnp.floating):
        x = self.prng.uniform(-1, 1, (m, n)).astype(dtype)
        y = self.prng.uniform(-1, 1, (m, n)).astype(dtype)
      else:
        x = self.prng.integers(-16000, 16000, (m, n)).astype(dtype)
        y = self.prng.integers(-16000, 16000, (m, n)).astype(dtype)
      result = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), (x, y), x, ()
      )(x, y)
      np.testing.assert_array_equal(result, np_op(x, scalar_rhs or y))

  @parameterized.product(
      vec_size=(4, 3, 1),
      dtype=(jnp.float16, jnp.bfloat16, jnp.float8_e4m3fn, jnp.float8_e5m2, jnp.float4_e2m1fn),
  )
  def test_float_abs(self, vec_size, dtype):
    if vec_size * bitwidth(dtype_to_ir_type(dtype)) % 8:
      self.skipTest("Vector length must be multiple of byte size")
    def kernel(ctx, src, dst, _):
      src = fa.FragmentedArray.load_strided(src, vec_size=vec_size)
      src.abs().store_untiled(dst)
    x = jnp.asarray(self.prng.uniform(-4, 4, (12 * 128,))).astype(dtype)
    f = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, x, ()
    )
    expected = jnp.abs(x.astype(np.float32)).astype(x.dtype)
    np.testing.assert_array_equal(f(x), expected)

  def test_minimum_np_compatibility(self):
    one = np.ones((128, 128)).astype(np.float32)
    negz = one * -0.
    posz = one * 0.
    nan = one * np.nan
    expectation = (np.minimum(negz, posz) == negz) & (np.minimum(nan, one) != one)
    assert np.all(expectation), expectation

    def kernel(ctx, dst, _):
      f32 = ir.F32Type.get()
      splat = lambda i: mgpu.FragmentedArray.splat(c(i, f32), (128, 128))
      negz = splat(-0.)
      posz = splat(0.)
      nan = splat(np.nan)
      one = splat(1.)
      res = (negz.min(posz) == negz) & (one.min(nan) != one) & (nan.min(one) != one)
      i8 = ir.IntegerType.get_signless(8)
      res.astype(i8, is_signed=False).store_untiled(dst)

    out_shape = jax.ShapeDtypeStruct((128, 128), np.int8)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    np.testing.assert_array_equal(result == 1, expectation)

  @parameterized.product(
      op=[operator.truediv, operator.floordiv, operator.mod],
      dtype=[jnp.float32, jnp.int32, jnp.uint32],
  )
  def test_division(self, op, dtype, m=64, n=32):
    if jnp.issubdtype(dtype, jnp.integer) and op is operator.truediv:
      self.skipTest("Unsupported for integer types")
    if jnp.issubdtype(dtype, jnp.floating) and op is operator.mod:
      self.skipTest("Unsupported for floating types")

    def kernel(ctx, dst, _):
      iota = iota_tensor(m, n, dtype)
      op(dtype(4.2).item() * iota, iota + 1).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((m, n), dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    iota = np.arange(m * n, dtype=dtype).reshape(m, n)
    np.testing.assert_allclose(
        result, op(dtype(4.2).item() * iota, iota + 1), atol=2e-7
    )

  def test_e8m0_reciprocal(self, m=64, n=128):
    dtype = jnp.float8_e8m0fnu
    layout = fa.WGMMA_LAYOUT

    def kernel(ctx, inp, out, smem):
      del ctx, smem
      x = mgpu.FragmentedArray.load_untiled(inp, layout=layout, optimized=False)
      (1 / x).store_untiled(out, optimized=False)

    bits = self.prng.integers(low=0, high=255, size=(m, n), dtype=np.uint8)
    # Inject a few NaNs
    bits[0, 0] = 255
    bits[1, 3] = 255
    bits[m - 1, n - 1] = 255
    x = jax.lax.bitcast_convert_type(jnp.array(bits, dtype=jnp.uint8), dtype)

    y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, ())(x)
    np.testing.assert_array_equal(y.astype(np.float32), 1 / x.astype(jnp.float32))

  def test_untiled_dim_indexing(self):
    shape = (2, 3, 64, 8)

    def kernel(ctx, src, dst, _):
      x = mgpu.FragmentedArray.load_untiled(
          src, layout=fa.WGMMA_LAYOUT, is_signed=True, optimized=False
      )
      x[1, 1].store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct(shape=shape[2:], dtype=jnp.int32)
    x = np.arange(np.prod(shape), dtype=np.int32).reshape(shape)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, out_shape, ()
    )(x)
    np.testing.assert_array_equal(result, x[1, 1])

  @parameterized.product(
      op=[
          operator.lt,
          operator.le,
          operator.gt,
          operator.ge,
          operator.eq,
          operator.ne,
      ],
      dtype=[jnp.float32, jnp.int32, jnp.uint32],
      rhs_is_literal=[False, True]
  )
  def test_comparison(self, op, dtype, rhs_is_literal, m=64, n=32):
    def kernel(ctx, dst, _):
      i8 = ir.IntegerType.get_signless(8)
      iota = iota_tensor(m, n, dtype)
      rhs = 0 if rhs_is_literal else iota + 1
      res = op(iota, rhs)
      assert not res.is_signed
      res.astype(i8, is_signed=False).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((m, n), jnp.int8)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    iota = np.arange(m * n, dtype=dtype).reshape(m, n)
    rhs = 0 if rhs_is_literal else iota + 1
    np.testing.assert_array_equal(result, op(iota, rhs).astype(jnp.int8))

  @parameterized.product(
      jax_dtype_from=(jnp.float32, jnp.float16, jnp.bfloat16, jnp.float8_e5m2, jnp.float8_e4m3fn, jnp.float8_e8m0fnu),
      jax_dtype_to=(jnp.float32, jnp.float16, jnp.bfloat16, jnp.float8_e5m2, jnp.float8_e4m3fn, jnp.float8_e8m0fnu),
      vec_len=(1, 2, 4, 8),
  )
  def test_conversion_f8_(self, jax_dtype_from, jax_dtype_to, vec_len):
    from_bitwidth = jnp.finfo(jax_dtype_from).bits
    to_bitwidth = jnp.finfo(jax_dtype_to).bits
    if from_bitwidth > 8 and to_bitwidth > 8:
      self.skipTest("At least one of the types should be 8-bit")
    if jax_dtype_from == jax_dtype_to:
      self.skipTest("Identical types, so nothing to test")
    if jnp.float8_e8m0fnu in {
        jax_dtype_from,
        jax_dtype_to,
    } and not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("f8e8m0fnu not supported on pre-Blackwell GPUs")
    if from_bitwidth == to_bitwidth == 8 and {jax_dtype_from, jax_dtype_to} != {
        jnp.float8_e4m3fn, jnp.float8_e5m2,
    }:
      self.skipTest("An unimplemented f8 <-> f8 conversion")
    unimplemented = {
        frozenset((jnp.float8_e4m3fn, jnp.bfloat16)),
        frozenset((jnp.float8_e5m2, jnp.bfloat16)),
        frozenset((jnp.float8_e8m0fnu, jnp.float16)),
    }
    if {jax_dtype_from, jax_dtype_to} in unimplemented:
      self.skipTest("Unimplemented")
    layout = fa.tmem_native_layout(vec_len)
    mlir_dtype_to = utils.dtype_to_ir_type(jax_dtype_to)
    m = 128
    n = 256
    def kernel(ctx, inp, out, smem):
      del ctx, smem
      t = mgpu.FragmentedArray.load_untiled(inp, layout=layout, optimized=False)
      t = t.astype(mlir_dtype_to)
      t.store_untiled(out, optimized=False)

    # For now we only sample the values representable in the narrow type,
    # because XLA and Mosaic disagree about the rounding.
    narrow_type = jax_dtype_from if from_bitwidth < to_bitwidth else jax_dtype_to
    int_sample_dtype = getattr(jnp, "int" + str(jnp.finfo(narrow_type).bits))
    sample_iinfo = jnp.iinfo(int_sample_dtype)
    bits = self.prng.integers(
        low=sample_iinfo.min, high=sample_iinfo.max, size=(m, n), dtype=np.int32
    ).astype(int_sample_dtype)
    values = jax.lax.bitcast_convert_type(bits, narrow_type)
    # A bunch of conversions are only supported for finite values.
    values = values.at[jnp.isinf(values)].set(jnp.finfo(narrow_type).max)
    values = values.astype(jax_dtype_from)

    expected = values.astype(jax_dtype_to)
    res = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), values, expected, ()
    )(values)
    # assert_array_equal fails when both sides are NaN.
    np.testing.assert_allclose(res, expected, atol=0, rtol=0, equal_nan=True)

  @parameterized.product(
      jax_wide_dtype=(jnp.float32, jnp.float16, jnp.bfloat16, jnp.float8_e5m2, jnp.float8_e4m3fn),
      to_narrow=(False, True),
      vec_len=(2, 4, 8),
  )
  def test_conversion_f4_(self, jax_wide_dtype, to_narrow, vec_len):
    if not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("f4 conversions not supported on pre-Blackwell GPUs")

    if to_narrow:
      jax_dtype_from = jax_wide_dtype
      jax_dtype_to = jnp.float4_e2m1fn
    else:
      jax_dtype_from = jnp.float4_e2m1fn
      jax_dtype_to = jax_wide_dtype

    layout = fa.tmem_native_layout(vec_len)
    mlir_dtype_to = utils.dtype_to_ir_type(jax_dtype_to)
    m = 128
    n = 64
    def kernel(ctx, inp, out, smem):
      del ctx, smem
      t = mgpu.FragmentedArray.load_untiled(inp, layout=layout, optimized=False)
      t = t.astype(mlir_dtype_to)
      t.store_untiled(out, optimized=False)

    iota_part = np.tile(np.arange(-8, 8, dtype=np.int8), ((m // 2) * n) // 16).reshape(m // 2, n)
    random_part = self.prng.integers(low=-8, high=8, size=(m // 2, n), dtype=np.int8)
    bits = np.concatenate([iota_part, random_part], axis=0)
    bits = jnp.asarray(bits).astype(jnp.int4)
    f4_values = jax.lax.bitcast_convert_type(bits, jnp.float4_e2m1fn)

    values = f4_values.astype(jax_dtype_from)
    expected = values.astype(jax_dtype_to)
    res = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), values, expected, ()
    )(values)
    np.testing.assert_array_equal(res, expected)

  @parameterized.product(
      jax_wide_dtype=(jnp.float32, jnp.float16, jnp.bfloat16),
      vec_len=(4, 8),
  )
  def test_roundtrip_conversion_f4(self, jax_wide_dtype, vec_len):
    if not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("f4 conversions not supported on pre-Blackwell GPUs")

    def kernel(ctx, inp, out, smem):
      del ctx, smem
      t = mgpu.FragmentedArray.load_untiled(
          inp, layout=fa.tmem_native_layout(vec_len), optimized=False
      )
      t = t.astype(utils.dtype_to_ir_type(jnp.float4_e2m1fn))
      t = t.astype(utils.dtype_to_ir_type(jax_wide_dtype))
      t.store_untiled(out, optimized=False)

    # 1.5 is exactly representable in FP4 e2m1fn.
    values = jnp.full((128, 64), 1.5, dtype=jax_wide_dtype)
    expected = values
    res = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), values, expected, ()
    )(values)
    np.testing.assert_array_equal(res, expected)

  def test_rounding_f8e8m0fnu(self):
    if not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("f8e8m0fnu not supported on pre-Blackwell GPUs")
    layout = fa.tmem_native_layout(2)
    m = 128
    n = 256
    def kernel(rounding, ctx, inp, out, smem):
      del ctx, smem
      t = mgpu.FragmentedArray.load_untiled(inp, layout=layout, optimized=False)
      t = t.astype(ir.Float8E8M0FNUType.get(), rounding=rounding)
      t.store_untiled(out, optimized=False)

    ints = self.prng.integers(
        low=0, high=100000, size=(m, n), dtype=np.int32
    )
    values = ints.astype(jnp.float32)
    next_pow2 = jnp.astype(np.exp2(np.ceil(np.log2(ints))), jnp.float8_e8m0fnu)
    prev_pow2 = jnp.astype(np.exp2(np.floor(np.log2(ints))), jnp.float8_e8m0fnu)

    supported_rounding = (
        mgpu.Rounding.TO_ZERO, mgpu.Rounding.TO_POSITIVE_INFINITY
    )
    for rounding, expected in zip(supported_rounding, (prev_pow2, next_pow2)):
      res = mgpu.as_gpu_kernel(
          functools.partial(kernel, rounding),
          (1, 1, 1), (128, 1, 1), values, expected, ()
      )(values)
      np.testing.assert_array_equal(res, expected)
    with self.assertRaises(ValueError):
      mgpu.as_gpu_kernel(
          functools.partial(kernel, mgpu.Rounding.TO_NEAREST_EVEN),
          (1, 1, 1), (128, 1, 1), values, prev_pow2, ()
      )(values)

  def test_foreach_wgmma_row_array(self):
    def kernel(ctx, out, smem):
      del ctx, smem
      x = iota_tensor(128, 128, jnp.float32)
      row = x.reduce("add", 1)
      # Test returning an array
      row = row.foreach(
          lambda x, _: arith.addf(x, c(1, row.mlir_dtype)), create_array=True
      )
      # Test no array return
      @row.foreach
      def _(v, idx):
        memref.store(v, out, idx)

    result = mgpu.as_gpu_kernel(
        kernel,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct(shape=(128,), dtype=jnp.float32),
        smem_scratch_shape=(),
    )()
    iota = np.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
    np.testing.assert_array_equal(result, iota.sum(axis=1) + 1)

  def test_foreach(self):
    dtype = jnp.int32
    swizzle = 128
    tiling = (8, swizzle // jnp.dtype(dtype).itemsize)
    shape = 128, 192
    mlir_dtype = utils.dtype_to_ir_type(dtype)
    cst = 9999
    def causal(val, idx):
      row, col = idx
      mask = arith.cmpi(arith.CmpIPredicate.uge, row, col)
      return arith.select(mask, val, c(cst, mlir_dtype))

    def kernel(ctx, dst, smem):
      x = iota_tensor(shape[0], shape[1], dtype)
      x.foreach(causal, create_array=True, is_signed=False).store_tiled(smem, swizzle=128)
      mgpu.commit_shared()
      ctx.async_copy(
          src_ref=smem,
          dst_ref=dst,
          gmem_transform=mgpu.TileTransform(tiling),
          swizzle=128,
      )
      ctx.await_async_copy(0)

    iota = np.arange(np.prod(shape), dtype=dtype).reshape(*shape)
    result = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        (),
        jax.ShapeDtypeStruct(shape=shape, dtype=dtype),
        jax.ShapeDtypeStruct(shape=mgpu.tile_shape(shape, tiling), dtype=dtype),
    )()
    expected = jnp.tril(iota) + jnp.triu(jnp.ones(shape), k=1) * cst
    np.testing.assert_array_equal(result, expected)

  @parameterized.product(
      op=[operator.and_, operator.or_, operator.xor, operator.lshift, operator.rshift],
      dtype=[jnp.uint32, jnp.uint8],
  )
  def test_bitwise(self, op, dtype, m=64, n=8):
    is_shift = op in {operator.lshift, operator.rshift}
    if is_shift and (bitwidth := jax.dtypes.itemsize_bits(dtype)) <= 8 and bitwidth != 1:
      self.skipTest("Shift ops are unsupported on 8-bit types.")
    is_8bit = (bitwidth := jax.dtypes.itemsize_bits(dtype)) <= 8 and bitwidth != 1

    # Shift operations are unsupported on 8-bit types. For non-shift bitwise
    # ops on 8-bit types, we use XOR instead of shift to construct the RHS.
    if is_shift:
      rhs_fn = lambda x: x & 0xf
    elif is_8bit:
      rhs_fn = lambda x: x ^ 0x5a
    else:
      rhs_fn = lambda x: x << 2

    def kernel(ctx, dst, _):
      iota = iota_tensor(m, n, dtype)
      rhs = rhs_fn(iota)
      op(iota, rhs).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((m, n), dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    iota = np.arange(m * n, dtype=dtype).reshape(m, n)
    rhs = rhs_fn(iota)
    np.testing.assert_array_equal(result, op(iota, rhs))

  @parameterized.product(
      ops=(
          (lambda x: -x, jax.lax.neg),
          (lambda x: x + 42, lambda x: x + 42),
          (lambda x: x.tanh(), jax.lax.tanh),
          (lambda x: x.abs(), np.abs),
      ),
      dtype=[jnp.float32, jnp.int32, jnp.uint32],
  )
  def test_unary(self, ops, dtype, m=64, n=32):
    op, np_op = ops
    if np_op is jax.lax.tanh and jnp.issubdtype(dtype, jnp.integer):
      raise self.skipTest("Tanh not supported for integer types")

    def kernel(ctx, dst, _):
      iota = iota_tensor(m, n, dtype)
      op(iota).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((m, n), dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    x = np.arange(m * n, dtype=dtype).reshape(m, n)
    np.testing.assert_allclose(result, np_op(x), atol=2e-7, rtol=2e-7)

  def test_select(self, m=64, n=32):

    def kernel(ctx, dst, _):
      iota = iota_tensor(m, n, jnp.int32)
      (iota < 16).select(iota * 2, iota * 3).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((m, n), jnp.int32)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    x = np.arange(m * n, dtype=jnp.int32).reshape(m, n)
    np.testing.assert_array_equal(result, np.where(x < 16, x * 2, x * 3))

  @parameterized.parameters(
      (mgpu.FragmentedArray.exp, np.exp),
      (mgpu.FragmentedArray.exp2, np.exp2),
      (mgpu.FragmentedArray.log, np.log),
      (mgpu.FragmentedArray.log2, np.log2),
      (mgpu.FragmentedArray.sin, np.sin),
      (mgpu.FragmentedArray.cos, np.cos),
      (mgpu.FragmentedArray.tanh, np.tanh),
      (mgpu.FragmentedArray.tanh, np.tanh, jnp.float16),
      (mgpu.FragmentedArray.tanh, np.tanh, jnp.bfloat16),
      (mgpu.FragmentedArray.rsqrt, jax.lax.rsqrt),
      (mgpu.FragmentedArray.sqrt, np.sqrt),
      (mgpu.FragmentedArray.erf, jax.scipy.special.erf),
  )
  @jtu.thread_unsafe_test()  # Modifies ``os.environ``
  @jtu.ignore_warning(message="overflow encountered", category=RuntimeWarning)
  def test_math(self, op, np_op, dtype=jnp.float32, m=64, n=32):
    if jtu.is_running_under_pytest():
      self.skipTest("PTX dump capture fails under pytest")
    def run_test(**kwargs):
      def kernel(ctx, dst, _):
        del ctx
        iota = iota_tensor(m, n, dtype) + 1
        op(iota, **kwargs).store_untiled(dst, optimized=False)
      out_shape = jax.ShapeDtypeStruct((m, n), dtype)
      with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), jtu.capture_stdout() as ptx:
        result = mgpu.as_gpu_kernel(
            kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
        )()
      return result, ptx()
    x = np.arange(m * n, dtype=np.float32).reshape(m, n) + 1
    kwargs = {} if op is mgpu.FragmentedArray.erf else {"approx": False}
    ref = np_op(x).astype(dtype)
    result, ptx = run_test(**kwargs)
    atol = 2e-3 if dtype in (jnp.float16, jnp.bfloat16) else 2e-7
    rtol = 2e-3 if dtype in (jnp.float16, jnp.bfloat16) else 2e-7
    np.testing.assert_allclose(result, ref, atol=atol, rtol=rtol)

    if op is mgpu.FragmentedArray.erf:
      # erf not supported with approximation.
      return

    result_approx, ptx_approx = run_test(approx=True)
    atol = 1e-2 if dtype in (jnp.float16, jnp.bfloat16) else 5e-3
    rtol = 1e-2 if dtype in (jnp.float16, jnp.bfloat16) else 4e-6
    np.testing.assert_allclose(result_approx, ref, atol=atol, rtol=rtol)

    # This is not super precise, but is a generic way of ensuring we take
    # different code generation paths.
    self.assertNotEqual(ptx, ptx_approx)

  def test_atan2(self, m=64, n=32):
    def kernel(ctx, dst, _):
      y = iota_tensor(m, n, jnp.float32) + 1  # Avoid zero
      x = iota_tensor(m, n, jnp.float32) + 2
      y.atan2(x).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((m, n), jnp.float32)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    y = np.arange(m * n, dtype=jnp.float32).reshape(m, n) + 1
    x = np.arange(m * n, dtype=jnp.float32).reshape(m, n) + 2
    np.testing.assert_allclose(result, np.arctan2(y, x), atol=2e-7, rtol=2e-7)

  @parameterized.parameters(jnp.float32, jnp.float16, jnp.bfloat16)
  def test_rpow(self, dtype):
    m, n = 64, 32
    def kernel(ctx, dst, _):
      del ctx
      # Powers of two are exactly representable in every float type, and an
      # exponent below 10 keeps the result small enough for f16 and bf16 too.
      exp = (iota_tensor(m, n, jnp.int32) % 10).astype(
          utils.dtype_to_ir_type(dtype)
      )
      (2 ** exp).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((m, n), dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    ref = (2 ** (np.arange(m * n, dtype=np.int32).reshape(m, n) % 10)).astype(
        dtype
    )
    np.testing.assert_array_equal(result, ref)

  @parameterized.parameters(
      jnp.int16, jnp.int32, jnp.float16, jnp.bfloat16, jnp.float32
  )
  def test_pow(self, dtype):
    m, n = 64, 32
    def kernel(ctx, dst, _):
      del ctx
      # Keep the base below 6 so that its cube is exactly representable in all
      # the tested types (bf16, the coarsest, is exact up to 256).
      base = (iota_tensor(m, n, jnp.int32) % 6).astype(
          utils.dtype_to_ir_type(dtype), is_signed=utils.is_signed(dtype)
      )
      (base ** 3).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((m, n), dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    ref = ((np.arange(m * n, dtype=np.int32).reshape(m, n) % 6) ** 3).astype(
        dtype
    )
    np.testing.assert_array_equal(result, ref)

  def test_strided_copy_noncontig_good(self):
    def kernel(ctx, src, dst, _):
      src_slice = mgpu.memref_slice(src, (slice(None), 1))
      mgpu.FragmentedArray.load_strided(src_slice, is_signed=True, vec_size=4).store_untiled(dst)

    in_shape = jax.ShapeDtypeStruct((32, 2, 32), jnp.int32)
    out_shape = jax.ShapeDtypeStruct((32, 32), jnp.int32)

    kernel_fn = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), in_shape, out_shape, ()
    )
    x = np.arange(math.prod(in_shape.shape), dtype=jnp.int32).reshape(in_shape.shape)
    np.testing.assert_array_equal(kernel_fn(x), x[:, 1])

  def test_strided_copy_noncontig_subbyte(self):
    def kernel(ctx, src, dst, _):
      src_slice = mgpu.memref_slice(src, (slice(None), 1))
      mgpu.FragmentedArray.load_strided(
          src_slice, is_signed=True
      ).store_untiled(dst)

    in_shape = jax.ShapeDtypeStruct((32, 2, 256), jnp.int4)
    out_shape = jax.ShapeDtypeStruct((32, 256), jnp.int4)

    kernel_fn = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), in_shape, out_shape, ()
    )
    x = jax.random.randint(
        jax.random.key(42), in_shape.shape, -128, 127, dtype=in_shape.dtype
    )
    np.testing.assert_array_equal(kernel_fn(x), x[:, 1])

  def test_strided_copy_noncontig_bad(self):
    def kernel(ctx, src, dst, _):
      src_slice = mgpu.memref_slice(src, (slice(None), 1))
      mgpu.FragmentedArray.load_strided(src_slice, is_signed=True, vec_size=2).store_untiled(dst)

    out_shape = jax.ShapeDtypeStruct((256, 7), jnp.int32)

    in_shape = jax.ShapeDtypeStruct((256, 6, 7), jnp.int32)
    msg = (
        "The contiguous dimension of the reference must be a multiple of the"
        " layout's vector size (got 7 and vector size 2)"
    )
    with self.assertRaises(ValueError, msg=msg):
      mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), in_shape, out_shape, ()
      )

    in_shape = jax.ShapeDtypeStruct((256, 5, 7), jnp.int32)
    msg = (
        "Non-contiguous dimension of the reference must have strides that are"
        " multiples of the layout's vector size (got 35 and vector size 2)"
    )
    with self.assertRaises(ValueError, msg=msg):
      mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), in_shape, out_shape, ()
      )

  @parameterized.product(
      dtype=[jnp.float32, jnp.int32],
      m=[128],
      n=[32, 64],
  )
  def test_strided_reduce_sum(self, dtype, m, n):
    def kernel(ctx, src, dst, scratch):
      src = mgpu.FragmentedArray.load_strided(
          src, is_signed=utils.is_signed(dtype)
      )
      acc = src.reduce("add", (0, 1), scratch).broadcast((m,))
      acc.store_untiled(dst, optimized=False)

    in_shape = jax.ShapeDtypeStruct((m, n), dtype)
    out_shape = jax.ShapeDtypeStruct((m,), dtype)
    kernel_fn = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape,
        out_shape,
        smem_scratch_shape=jax.ShapeDtypeStruct((4,), dtype),
    )
    x = np.arange(m * n, dtype=dtype).reshape(m, n)
    np.testing.assert_array_equal(kernel_fn(x), jnp.full((m,), x.sum()))

  def test_dimension_compression_for_vec_size(self):
    def body(ctx, src, dst, _):
      src_arr = mgpu.FragmentedArray.load_strided(
        mgpu.memref_slice(src, (slice(None), slice(4, None))), vec_size=4
      )
      src_arr.store_untiled(dst, optimized=False)
    in_shape = jax.ShapeDtypeStruct((8, 20, 4, 3, 1), jnp.float32)
    out_shape = jax.ShapeDtypeStruct((8, 16, 4, 3, 1), jnp.float32)
    kernel = mgpu.as_gpu_kernel(
        body, (1, 1, 1), (128, 1, 1), in_shape, out_shape, ()
    )
    x = np.arange(math.prod(in_shape.shape), dtype=np.float32).reshape(in_shape.shape)
    np.testing.assert_array_equal(kernel(x), x[:, 4:])

  @parameterized.product(
      dtype=[jnp.float32, jnp.int32],
      m=[128],
      n=[32, 64],
      reduce_both=[False, True],
  )
  def test_splat_reduce_sum(self, dtype, m, n, reduce_both):
    def kernel(ctx, dst, _):
      src = mgpu.FragmentedArray.splat(
          utils.c(1, utils.dtype_to_ir_type(dtype)),
          (m, n),
          is_signed=utils.is_signed(dtype),
      )
      if reduce_both:
        acc = src.reduce("add", (0, 1)).broadcast((m,))
      else:
        acc = src.reduce("add", 1)
      acc.store_untiled(dst, optimized=False)

    kernel_fn = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct((m,), dtype),
        smem_scratch_shape=(),
    )
    result = m * n if reduce_both else n
    np.testing.assert_array_equal(kernel_fn(), jnp.full((m,), result, dtype))

  @parameterized.parameters(jnp.float32, jnp.int32)
  def test_splat_reduce_prod(self, dtype):
    shape = (2, 3)
    def _kernel(ctx, inp, dst, _):
      src = mgpu.FragmentedArray.splat(
          memref.load(inp, []),
          shape,
          is_signed=utils.is_signed(dtype),
      )
      acc = src.reduce("prod", (0, 1))
      acc.store_untiled(dst, optimized=False)

    kernel = mgpu.as_gpu_kernel(
        _kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape=jax.ShapeDtypeStruct((), dtype),
        out_shape=jax.ShapeDtypeStruct((), dtype),
        smem_scratch_shape=(),
    )
    expected = jnp.array(3**math.prod(shape), dtype=dtype)
    np.testing.assert_array_equal(kernel(jnp.array(3, dtype=dtype)), expected)

  @parameterized.named_parameters(
      ("wgmma_row", fa.WGMMA_LAYOUT, fa.WGMMA_ROW_LAYOUT, 1),
      ("wgmma_col", fa.WGMMA_LAYOUT, fa.WGMMA_COL_LAYOUT, 0),
      ("tcgen05_row", tcgen05.LAYOUT, tcgen05.ROW_LAYOUT, 1),
      ("tcgen05_col", tcgen05.LAYOUT, tcgen05.COL_LAYOUT, 0),
  )
  def test_layout_reduction_definition(self, layout, expected_reduced_layout, axis):
    self.assertEqual(layout.reduce((axis,)), expected_reduced_layout)

  def test_layout_reduction_handles_tiles_with_three_different_ranks(self):
    layout = fa.TiledLayout(
        tiling=fa.Tiling(tiles=((1, 2, 64), (2, 16), (8,), (4,), (2,), (1,))),
        warp_dims=(-7,),
        lane_dims=(-6, -5, -4, -3, -2),
        vector_dim=-1,
    )
    self.assertEqual(
        layout.reduce((2,)),
        fa.TiledLayout(
            tiling=fa.Tiling(tiles=((1, 2), (1,))),
            warp_dims=(fa.Replicated(times=4),),
            lane_dims=(
                -2,
                fa.Replicated(times=2),
                fa.Replicated(times=2),
                fa.Replicated(times=2),
                fa.Replicated(times=2),
            ),
            vector_dim=-1,
        ),
    )

  @parameterized.product(
      op=(arith.addf, arith.maximumf),
      m=(64, 128),
      n=(8, 16, 32, 64, 80, 128, 256),
  )
  def test_reduce(self, op, m=64, n=32):
    def kernel(ctx, dst, _):
      iota = iota_tensor(m, n, jnp.float32)
      iota.reduce(op, axis=1).broadcast_in_dim(
          (m, n), (0,), mgpu.WGMMA_LAYOUT
      ).store_untiled(dst, optimized=False)
    out_shape = jax.ShapeDtypeStruct((m, n), jnp.float32)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    x = np.arange(m * n, dtype=jnp.float32).reshape(m, n)
    if op == arith.addf:
      expected = np.broadcast_to(x.sum(axis=1, keepdims=True), x.shape)
    elif op == arith.maximumf:
      expected = np.broadcast_to(x.max(axis=1, keepdims=True), x.shape)
    else:
      raise NotImplementedError(f"Unsupported op: {op}")
    np.testing.assert_array_equal(result, expected)

  @parameterized.named_parameters(
      ("untiled", 0, (2, 64, 32)),
      ("tiled", 2, (2, 64, 32)),
      ("all_tiled", (1, 2), (2, 64, 32)),
      ("all", (0, 1, 2), (2, 64, 32)),
      ("4d_untiled", 0, (2, 4, 64, 16)),
      ("4d_tiled", 3, (2, 4, 64, 16)),
  )
  def test_untiled_reduce(self, axis, shape):
    def kernel(ctx, inp, dst, scratch):
      arr = mgpu.FragmentedArray.load_untiled(
          inp, layout=mgpu.WGMMA_LAYOUT, optimized=False, is_signed=True
      )
      reduced = arr.reduce("add", axis=axis, scratch=scratch)
      reduced.store_untiled(dst, optimized=False)

    in_shape = jax.ShapeDtypeStruct(shape, jnp.int32)
    x = np.arange(math.prod(shape), dtype=jnp.int32).reshape(shape)
    axes = (axis,) if isinstance(axis, int) else axis
    out_shape_tuple = tuple(
        d for i, d in enumerate(shape) if i not in axes
    )
    out_shape = jax.ShapeDtypeStruct(out_shape_tuple, jnp.int32)
    result = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape,
        out_shape,
        smem_scratch_shape=jax.ShapeDtypeStruct((128,), jnp.int32),
    )(x)
    np.testing.assert_array_equal(result, x.sum(axis=axis))

  def test_unreduced_layout_definitions(self):
    def as_unreduced(layout, op):
      """Returns `layout` with its `Replicated` dims marked `Unreduced`."""
      unreduce = lambda ds: tuple(
          fa.Unreduced(d.times) if isinstance(d, fa.Replicated) else d
          for d in ds
      )
      return dataclasses.replace(
          layout,
          warp_dims=unreduce(layout.warp_dims),
          lane_dims=unreduce(layout.lane_dims),
          unreduced_operation=op,
      )

    layout = fa.WGMMA_LAYOUT
    # Full reduction along axis 1 (lane_dims=(-5,), vector_dim)
    r_normal = layout.reduce((1,))
    self.assertFalse(r_normal.has_unreduced_dims)
    self.assertIsNone(r_normal.unreduced_operation)
    self.assertEqual(r_normal, fa.WGMMA_ROW_LAYOUT)

    # Partial reduction along axis 1 leaves the lane dim unreduced.
    r_partial_lane = layout.reduce((1,), local_only=True, op="add")
    self.assertEqual(r_partial_lane, as_unreduced(fa.WGMMA_ROW_LAYOUT, "add"))

    # Partial reduction along axis 0 leaves warp and lane dims unreduced.
    r_partial = layout.reduce((0,), local_only=True, op="add")
    self.assertEqual(r_partial, as_unreduced(layout.reduce((0,)), "add"))
    self.assertIn(fa.Unreduced(times=4), r_partial.warp_dims)
    self.assertIn(fa.Unreduced(times=8), r_partial.lane_dims)

    r_completed = r_partial.reduce(())
    self.assertFalse(r_completed.has_unreduced_dims)
    self.assertEqual(r_completed, layout.reduce((0,)))

    # A full reduction of an unreduced layout also completes the pending
    # partial reduction.
    self.assertEqual(r_partial.reduce((0,)), layout.reduce((0, 1)))

  def test_partial_layout_reduction_of_register_only_dim(self):
    # Dimension 0 is not tiled over warps or lanes, so reducing it partially
    # is already complete: the resulting layout must have no unreduced
    # dimensions and therefore no `unreduced_operation`.
    layout = fa.TiledLayout(
        tiling=fa.Tiling(tiles=((1, 2, 64), (2, 16), (8,), (4,), (2,), (1,))),
        warp_dims=(-7,),
        lane_dims=(-6, -5, -4, -3, -2),
        vector_dim=-1,
    )
    reduced = layout.reduce((0,), local_only=True, op="add")
    self.assertFalse(reduced.has_unreduced_dims)
    self.assertIsNone(reduced.unreduced_operation)
    self.assertEqual(reduced, layout.reduce((0,)))

  @parameterized.named_parameters(
      ("tiled_axis_warp", 0, (64, 32)),
      ("untiled_and_tiled_axes", (0, 1), (2, 64, 32)),
  )
  def test_add_partially_reduced_arrays(self, axis, shape):
    def kernel(ctx, inp1, inp2, dst, scratch):
      arr1 = mgpu.FragmentedArray.load_untiled(
          inp1, layout=mgpu.WGMMA_LAYOUT, optimized=False, is_signed=True
      )
      arr2 = mgpu.FragmentedArray.load_untiled(
          inp2, layout=mgpu.WGMMA_LAYOUT, optimized=False, is_signed=True
      )
      axes = (axis,) if isinstance(axis, int) else axis
      untiled_rank = len(shape) - len(arr1.layout.base_tile_shape)
      tiled_axes = tuple(a - untiled_rank for a in axes if a >= untiled_rank)
      partial_layout = arr1.layout.reduce(tiled_axes, local_only=True, op="add")
      part1 = arr1.reduce("add", axis=axis, target_layout=partial_layout)
      part2 = arr2.reduce("add", axis=axis, target_layout=partial_layout)
      self.assertTrue(part1.is_unreduced)
      self.assertTrue(part2.is_unreduced)
      # Pointwise binary op between identical partially reduced layouts
      part_sum = part1 + part2
      self.assertTrue(part_sum.is_unreduced)
      # Complete the remaining reduction across lanes and warps.
      completed = part_sum.reduce("add", (), scratch)
      self.assertFalse(completed.is_unreduced)
      completed.store_untiled(dst, optimized=False)

    in_shape = jax.ShapeDtypeStruct(shape, jnp.int32)
    x1 = np.arange(math.prod(shape), dtype=jnp.int32).reshape(shape)
    x2 = np.ones(shape, dtype=jnp.int32) * 5
    axes = (axis,) if isinstance(axis, int) else axis
    out_shape_tuple = tuple(d for i, d in enumerate(shape) if i not in axes)
    out_shape = jax.ShapeDtypeStruct(out_shape_tuple, jnp.int32)
    result = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        (in_shape, in_shape),
        out_shape,
        smem_scratch_shape=jax.ShapeDtypeStruct((256,), jnp.int32),
    )(x1, x2)
    expected = (x1 + x2).sum(axis=axis)
    np.testing.assert_array_equal(result, expected)

  def test_partial_reduce_across_lanes_only(self):
    # Reducing axis 1 of WGMMA_LAYOUT only spans lanes, so completing the
    # reduction uses warp shuffles and needs no SMEM scratch.
    shape = (64, 32)
    axis = 1

    def kernel(ctx, inp, dst, _):
      arr = mgpu.FragmentedArray.load_untiled(
          inp, layout=mgpu.WGMMA_LAYOUT, optimized=False, is_signed=True
      )
      partial_layout = arr.layout.reduce((axis,), local_only=True, op="add")
      # Phase 1: partial reduction that leaves the lanes unreduced.
      part = arr.reduce("add", axis=axis, target_layout=partial_layout)
      self.assertTrue(part.is_unreduced)
      # Phase 2: complete the reduction across lanes, without scratch.
      completed = part.reduce("add", ())
      self.assertFalse(completed.is_unreduced)
      completed.store_untiled(dst, optimized=False)

    in_shape = jax.ShapeDtypeStruct(shape, jnp.int32)
    x = np.arange(math.prod(shape), dtype=jnp.int32).reshape(shape)
    out_shape_tuple = tuple(d for i, d in enumerate(shape) if i != axis)
    out_shape = jax.ShapeDtypeStruct(out_shape_tuple, jnp.int32)
    result = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape,
        out_shape,
        smem_scratch_shape=(),
    )(x)
    np.testing.assert_array_equal(result, x.sum(axis=axis))

  def test_partial_reduce_lanes_leaving_warps_unreduced(self):
    # Reducing axis 0 of WGMMA_LAYOUT spans both lanes and warps. Target a
    # layout that reduces the lane dimensions (with warp shuffles) but leaves
    # the warp dimensions unreduced, so that SMEM is only involved once the
    # reduction is completed.
    shape = (64, 32)

    def kernel(ctx, inp, dst, scratch):
      arr = mgpu.FragmentedArray.load_untiled(
          inp, layout=mgpu.WGMMA_LAYOUT, optimized=False, is_signed=True
      )
      local_layout = arr.layout.reduce((0,), local_only=True, op="add")
      full_layout = arr.layout.reduce((0,))
      lanes_reduced_layout = dataclasses.replace(
          local_layout, lane_dims=full_layout.lane_dims
      )
      self.assertTrue(lanes_reduced_layout.has_unreduced_dims)
      # Phase 1: reduce across lanes only. No scratch is required.
      part = arr.reduce("add", axis=0, target_layout=lanes_reduced_layout)
      self.assertTrue(part.is_unreduced)
      self.assertEqual(part.layout, lanes_reduced_layout)
      # Phase 2: complete the reduction across warps.
      completed = part.reduce("add", (), scratch)
      self.assertFalse(completed.is_unreduced)
      completed.store_untiled(dst, optimized=False)

    in_shape = jax.ShapeDtypeStruct(shape, jnp.int32)
    x = np.arange(math.prod(shape), dtype=jnp.int32).reshape(shape)
    out_shape = jax.ShapeDtypeStruct((shape[1],), jnp.int32)
    result = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape,
        out_shape,
        smem_scratch_shape=jax.ShapeDtypeStruct((256,), jnp.int32),
    )(x)
    np.testing.assert_array_equal(result, x.sum(axis=0))

  def test_partial_reduce_chained(self):
    def kernel(ctx, inp, dst, scratch):
      arr = mgpu.FragmentedArray.load_untiled(
          inp, layout=mgpu.WGMMA_LAYOUT, optimized=False, is_signed=True
      )
      # Reduce axis 2 (untiled_rank = 1, tiled axis 1) partially with "max"
      partial_layout_1 = arr.layout.reduce((1,), local_only=True, op="max")
      part_1 = arr.reduce("max", axis=2, target_layout=partial_layout_1)
      self.assertTrue(part_1.is_unreduced)

      # Reduce axis 1 (untiled_rank = 1, tiled axis 0) partially with "max"
      partial_layout_2 = part_1.layout.reduce((0,), local_only=True, op="max")
      part_2 = part_1.reduce("max", axis=1, target_layout=partial_layout_2)
      self.assertTrue(part_2.is_unreduced)

      # Complete reduction across lanes and warps
      completed = part_2.reduce("max", (), scratch)
      self.assertFalse(completed.is_unreduced)
      completed.store_untiled(dst, optimized=False)

    shape = (2, 64, 32)
    in_shape = jax.ShapeDtypeStruct(shape, jnp.int32)
    x = np.arange(math.prod(shape), dtype=jnp.int32).reshape(shape)
    out_shape = jax.ShapeDtypeStruct((2,), jnp.int32)
    result = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape,
        out_shape,
        smem_scratch_shape=jax.ShapeDtypeStruct((256,), jnp.int32),
    )(x)
    expected = x.max(axis=(1, 2))
    np.testing.assert_array_equal(result, expected)

  def test_partial_reduce_scalar_unsupported(self):
    def kernel(ctx, src, scratch):
      del ctx, scratch
      arr = mgpu.FragmentedArray.load_untiled(
          src, layout=mgpu.WGMMA_LAYOUT, optimized=False, is_signed=True
      )
      partial_layout = arr.layout.reduce((0, 1), local_only=True, op="add")
      with self.assertRaises(NotImplementedError):
        arr.reduce("add", axis=(0, 1), target_layout=partial_layout)

    in_shape = jax.ShapeDtypeStruct((64, 32), jnp.int32)
    mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape,
        out_shape=(),
        smem_scratch_shape=(),
    )

  def test_partial_reduce_pointwise_ops(self):
    def kernel(ctx, inp, scratch):
      del ctx, scratch
      arr = mgpu.FragmentedArray.load_untiled(
          inp, layout=mgpu.WGMMA_LAYOUT, optimized=False
      )

      def partial(op):
        layout = arr.layout.reduce((0,), local_only=True, op=op)
        return arr.reduce(op, axis=0, target_layout=layout)

      def splat(value):
        return mgpu.FragmentedArray.splat(
            utils.c(value, arr.mlir_dtype), shape=(arr.shape[1],)
        )

      with self.subTest("add"):
        part = partial("add")
        _ = part + part
        _ = part - part
        _ = -part
        with self.assertRaisesRegex(ValueError, "distribute over"):
          _ = part * part
        with self.assertRaisesRegex(ValueError, "distribute over"):
          _ = part / part
        with self.assertRaisesRegex(ValueError, "distribute over"):
          _ = part.max(part)
        # exp(a + b) = exp(a) * exp(b), which does not preserve the reduction.
        with self.assertRaisesRegex(ValueError, "distribute over"):
          _ = part.exp()
        # Splatting the neutral element is allowed.
        _ = part + 0.0
        _ = splat(0.0) - part
        # Any other splat would be added once per unreduced thread.
        with self.assertRaisesRegex(ValueError, "unreduced dims"):
          _ = part + 1.0
        with self.assertRaisesRegex(ValueError, "unreduced dims"):
          _ = part + splat(1.0)
        with self.assertRaisesRegex(ValueError, "unreduced dims"):
          _ = splat(1.0) - part
        # A splat LHS must not bypass the check on the unreduced RHS.
        with self.assertRaisesRegex(ValueError, "distribute over"):
          _ = splat(1.0) // part

      with self.subTest("prod"):
        part = partial("prod")
        _ = part * part
        _ = part / part
        with self.assertRaisesRegex(ValueError, "distribute over"):
          _ = part + part
        with self.assertRaisesRegex(ValueError, "distribute over"):
          _ = -part
        # Splatting the neutral element is allowed.
        _ = 1.0 / part
        # Any other splat would be multiplied once per unreduced thread.
        with self.assertRaisesRegex(ValueError, "unreduced dims"):
          _ = part * 2.0

      for op in ("max", "min"):
        with self.subTest(op):
          part = partial(op)
          _ = getattr(part, op)(part)
          # Non-decreasing functions distribute over max and min.
          _ = part.exp()
          _ = part.exp(approx=True)
          _ = part.exp2()
          _ = part.tanh()
          _ = part.erf()
          _ = part.round()
          _ = part.round_even()
          with self.assertRaisesRegex(ValueError, "distribute over"):
            _ = getattr(part, "min" if op == "max" else "max")(part)
          with self.assertRaisesRegex(ValueError, "distribute over"):
            _ = part + part
          # -max(a, b) = min(-a, -b) (and vice versa).
          with self.assertRaisesRegex(ValueError, "distribute over"):
            _ = -part
          # `log` is non-decreasing, but it is not defined on negative numbers.
          with self.assertRaisesRegex(ValueError, "distribute over"):
            _ = part.log()
          # `sin` is not monotonic.
          with self.assertRaisesRegex(ValueError, "distribute over"):
            _ = part.sin()
          # max and min are idempotent, so any splat is allowed.
          _ = getattr(part, op)(2.0)
          _ = getattr(splat(2.0), op)(part)

    in_shape = jax.ShapeDtypeStruct((64, 32), jnp.float32)
    mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape,
        out_shape=(),
        smem_scratch_shape=(),
    )

  @parameterized.named_parameters(
      ("neg_f32", "add", jnp.float32, operator.neg, np.negative),
      ("neg_i32", "add", jnp.int32, operator.neg, np.negative),
      ("sub_f32", "add", jnp.float32, operator.sub, np.subtract),
      ("sub_i32", "add", jnp.int32, operator.sub, np.subtract),
      ("truediv", "prod", jnp.float32, operator.truediv, np.divide),
      ("exp_max", "max", jnp.float32, lambda x: x.exp(), np.exp),
      ("exp_min", "min", jnp.float32, lambda x: x.exp(), np.exp),
      ("tanh_max", "max", jnp.float32, lambda x: x.tanh(), np.tanh),
      ("exp_approx_max", "max", jnp.float32, lambda x: x.exp(approx=True),
       np.exp),
      ("round_even_min", "min", jnp.float32, lambda x: x.round_even(), np.rint),
  )
  def test_distributive_pointwise_op_on_partially_reduced_arrays(
      self, reduce_op, dtype, fa_fn, np_fn
  ):
    shape = (64, 32)
    num_operands = np_fn.nin

    def kernel(ctx, *refs):
      del ctx
      *inps, dst, scratch = refs
      parts = []
      for inp in inps:
        arr = mgpu.FragmentedArray.load_untiled(
            inp,
            layout=mgpu.WGMMA_LAYOUT,
            optimized=False,
            is_signed=utils.is_signed(dtype),
        )
        partial_layout = arr.layout.reduce((0,), local_only=True, op=reduce_op)
        part = arr.reduce(reduce_op, axis=0, target_layout=partial_layout)
        self.assertTrue(part.is_unreduced)
        parts.append(part)
      result = fa_fn(*parts)
      self.assertTrue(result.is_unreduced)
      self.assertEqual(result.layout, parts[0].layout)
      completed = result.reduce(reduce_op, (), scratch)
      self.assertFalse(completed.is_unreduced)
      completed.store_untiled(dst, optimized=False)

    if reduce_op == "prod":
      # Keep the products close to 1 to avoid overflow and underflow.
      make = lambda: self.prng.uniform(0.9, 1.1, shape).astype(dtype)
    elif jnp.issubdtype(dtype, jnp.integer):
      make = lambda: self.prng.integers(-1000, 1000, shape).astype(dtype)
    else:
      make = lambda: self.prng.uniform(-4, 4, shape).astype(dtype)
    inputs = [make() for _ in range(num_operands)]
    np_reduce = {
        "add": np.sum, "prod": np.prod, "max": np.max, "min": np.min
    }[reduce_op]
    expected = np_fn(*(np_reduce(x, axis=0) for x in inputs))

    in_shape = jax.ShapeDtypeStruct(shape, dtype)
    result = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        (in_shape,) * num_operands,
        jax.ShapeDtypeStruct((shape[1],), dtype),
        smem_scratch_shape=jax.ShapeDtypeStruct((256,), dtype),
    )(*inputs)
    np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-4)

  def test_splat_rejects_unreduced_layout(self):
    def kernel(ctx, src, scratch):
      del ctx, scratch
      arr = mgpu.FragmentedArray.load_untiled(
          src, layout=mgpu.WGMMA_LAYOUT, optimized=False, is_signed=True
      )
      unreduced_layout = arr.layout.reduce((0,), local_only=True, op="add")
      with self.assertRaisesRegex(ValueError, "unreduced dims"):
        mgpu.FragmentedArray.splat(
            utils.c(1, arr.mlir_dtype),
            shape=(32,),
            layout=unreduced_layout,
            is_signed=True,
        )

    in_shape = jax.ShapeDtypeStruct((64, 32), jnp.int32)
    mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        in_shape,
        out_shape=(),
        smem_scratch_shape=(),
    )

  @parameterized.product(
      op=("add", "prod", "max", "min"),
      dtype=(jnp.float32, jnp.float16, jnp.int32, jnp.int16, jnp.uint32),
  )
  def test_splat_unreduced_layout(self, op, dtype):
    out_shape = (32,)
    is_float = jnp.issubdtype(dtype, jnp.floating)
    match op:
      case "add":
        neutral = 0.0 if is_float else 0
      case "prod":
        neutral = 1.0 if is_float else 1
      case "max":
        neutral = float("-inf") if is_float else int(jnp.iinfo(dtype).min)
      case "min":
        neutral = float("inf") if is_float else int(jnp.iinfo(dtype).max)

    def kernel(ctx, dst, scratch):
      del ctx
      layout = mgpu.WGMMA_LAYOUT
      unreduced_layout = layout.reduce((0,), local_only=True, op=op)
      acc = mgpu.FragmentedArray.splat(
          utils.c(neutral, utils.dtype_to_ir_type(dtype)),
          shape=out_shape,
          layout=unreduced_layout,
          is_signed=utils.is_signed(dtype),
      )
      self.assertTrue(acc.is_unreduced)
      acc = acc.reduce(op, (), scratch)
      acc.store_untiled(dst, optimized=False)

    kernel = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        (),
        jax.ShapeDtypeStruct(out_shape, dtype),
        smem_scratch_shape=jax.ShapeDtypeStruct((256,), dtype),
    )
    expected = np.full(out_shape, neutral, dtype=dtype)
    np.testing.assert_array_equal(kernel(), expected)

  @parameterized.product(
      vec_size=(4, 3, 1),
      dtype=(jnp.float32, jnp.float16, jnp.bfloat16,
             jnp.int32, jnp.int16, jnp.uint32, jnp.uint16),
  )
  @jtu.thread_unsafe_test()
  def test_max(self, vec_size, dtype):
    def kernel(ctx, src, src2, dst, _):
      is_signed = utils.is_signed(dtype)
      src = fa.FragmentedArray.load_strided(src, vec_size=vec_size, is_signed=is_signed)
      src2 = fa.FragmentedArray.load_strided(src2, vec_size=vec_size, is_signed=is_signed)
      src.max(src2).store_untiled(dst)
    x = self.prng.uniform(-1, 1, (12 * 128,)).astype(dtype)
    y = self.prng.uniform(-1, 1, (12 * 128,)).astype(dtype)
    f = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), (x, y), x, ()
      )
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      z = f(x, y).block_until_ready()
    if dtype == jnp.float32:
      dtype_short = "f32"
    elif dtype == jnp.float16:
      dtype_short = "f16"
    elif dtype == jnp.bfloat16:
      dtype_short = "bf16"
    elif jnp.issubdtype(dtype, jnp.signedinteger):
      dtype_short = f"s{dtypes.itemsize_bits(dtype)}"
    elif jnp.issubdtype(dtype, jnp.unsignedinteger):
      dtype_short = f"u{dtypes.itemsize_bits(dtype)}"
    else:
      raise NotImplementedError(f"Unsupported dtype: {dtype}")
    ptx = ptx()
    nan_modifier = ".NaN" if jnp.issubdtype(dtype, jnp.floating) else ""
    instr = f"max{nan_modifier}.{dtype_short} "
    instr_double = f"max{nan_modifier}.{dtype_short}x2 "
    single_converts = ptx.count(instr)
    double_converts = ptx.count(instr_double)
    self.assertEqual(128 * (single_converts + 2 * double_converts), 12 * 128)
    if vec_size % 2:
      self.assertGreater(single_converts, 0)
    elif dtypes.itemsize_bits(dtype) < 32:
      # This, together with the assertion above, implies that all converts
      # happened through doubled operations.
      self.assertEqual(single_converts, 0)
    np.testing.assert_array_equal(z, np.maximum(x, y))

  def test_splat_layout(self):
    m, n = 64, 8
    def kernel(ctx, dst, _):
      iota = iota_tensor(m, n, jnp.float32)
      cte = c(1, iota.mlir_dtype)
      cte_arr = mgpu.FragmentedArray.splat(cte, ())
      cte_arr = cte_arr.reshape((1, 1)).broadcast((m, n))
      (iota + cte_arr).store_untiled(dst, optimized=False)
    out_shape = jax.ShapeDtypeStruct((m, n), jnp.float32)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    expected = np.arange(m * n, dtype=jnp.float32).reshape(m, n) + 1
    np.testing.assert_array_equal(result, expected)

  def test_splat(self):
    def kernel(ctx, dst, _):
      f32 = ir.F32Type.get()
      v = arith.constant(f32, ir.FloatAttr.get(f32, 3.14))
      t = mgpu.FragmentedArray.splat(
          v, (128,), mgpu.WGMMA_ROW_LAYOUT
      )
      t.broadcast_in_dim((128, 32), (0,), mgpu.WGMMA_LAYOUT).store_untiled(dst, optimized=False)
    out_shape = jax.ShapeDtypeStruct((128, 32), jnp.float32)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    np.testing.assert_array_equal(result, np.full((128, 32), 3.14, np.float32))

  def test_splat_binary_ops(self):
    def kernel(ctx, src, dst, _):
      f32 = ir.F32Type.get()
      pi_arr = mgpu.FragmentedArray.load_strided(src)
      assert isinstance(pi_arr.layout, mgpu.WGStridedFragLayout)
      pi_scalar = arith.constant(f32, ir.FloatAttr.get(f32, 3.14))
      pi_splat = mgpu.FragmentedArray.splat(pi_scalar, ())
      assert isinstance(pi_splat.layout, mgpu.WGSplatFragLayout)
      pi_arr_sq = pi_arr * pi_splat.broadcast(pi_arr.shape)
      assert isinstance(pi_arr_sq.layout, mgpu.WGStridedFragLayout)
      pi_arr_cube = pi_splat.broadcast(pi_arr.shape) * pi_arr_sq
      assert isinstance(pi_arr_cube.layout, mgpu.WGStridedFragLayout)
      (pi_arr == pi_arr).select(pi_splat, pi_arr_cube).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((128, 32), jnp.float32)
    inp = jnp.ones_like(out_shape) * 3.14
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), inp, out_shape, ()
    )(inp)
    np.testing.assert_allclose(result, np.full((128, 32), 3.14, np.float32))

  @parameterized.product(in_shape=((128, 128), (128, 64), (64, 128)))
  def test_strided_load_store(self, in_shape):
    def kernel(ctx, *args):
      gmem_input, gmem_output, (smem_input, smem_output) = args
      copy(gmem_input, smem_input)
      t = mgpu.FragmentedArray.load_strided(smem_input)
      t.store_untiled(smem_output)
      copy(smem_output, gmem_output)

    inp = out = self.prng.uniform(-1, 1, in_shape).astype(jnp.float16)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (inp,), out, [inp, out],
    )(inp)
    np.testing.assert_array_equal(inp, result)

  @parameterized.product(
      in_shape=((1024,), (256,), (128,), (64,)),
      dtype=(jnp.float16, jnp.float32),
      swizzle=(16, 32, 64, 128)
  )
  def test_wgmma_row_load_store_with_layout(self, in_shape, dtype, swizzle):
    def kernel(ctx, gmem_input, gmem_output, smem):
      smem_input, smem_output = smem
      copy(gmem_input, smem_input, swizzle=swizzle)
      t = mgpu.FragmentedArray.load_untiled(
          smem_input, layout=mgpu.WGMMA_ROW_LAYOUT, swizzle=swizzle
      )
      t.store_untiled(smem_output)
      copy(smem_output, gmem_output)

    inp = out = self.prng.uniform(-1, 1, in_shape).astype(dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (inp,), out, [inp, out],
    )(inp)
    np.testing.assert_array_equal(inp, result)

  @parameterized.product(
      in_shape=((128,), (64,)),
      dtype=(jnp.float16, jnp.float32),
      swizzle=(16, 32, 64, 128),
  )
  def test_wgmma_col_load_store_with_layout(self, in_shape, dtype, swizzle):
    def kernel(ctx, *args):
      gmem_input, gmem_output, (smem_input, smem_output) = args
      copy(gmem_input, smem_input, swizzle=swizzle)
      t = mgpu.FragmentedArray.load_untiled(
          smem_input, swizzle=swizzle, layout=mgpu.WGMMA_COL_LAYOUT
      )
      t.store_untiled(smem_output)
      copy(smem_output, gmem_output)

    inp = out = self.prng.uniform(-1, 1, in_shape).astype(dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (inp,), out, [inp, out],
    )(inp)
    np.testing.assert_array_equal(result, inp)

  @parameterized.parameters(
      ((1, 128), (4, 128)),
      ((1, 1, 128), (2, 4, 128)),
      ((128,), (4, 128)),
  )
  def test_broadcast_in_dim_major_strided(self, src_shape, dst_shape):
    dtype = jnp.float16
    dims = tuple(range(len(dst_shape) - len(src_shape), len(dst_shape)))

    def kernel(ctx, src, dst, scratch):
      del ctx, scratch
      t = mgpu.FragmentedArray.load_strided(src, vec_size=1)
      layout = mgpu.WGStridedFragLayout(shape=dst_shape, vec_size=1)
      t.broadcast_in_dim(dst_shape, dims, layout).store_untiled(
          dst, optimized=False
      )

    inp = self.prng.uniform(-1, 1, src_shape).astype(dtype)
    result = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        (inp,),
        jax.ShapeDtypeStruct(dst_shape, dtype),
        inp,
    )(inp)
    out_ref = jax.lax.broadcast_in_dim(inp, dst_shape, dims)
    np.testing.assert_array_equal(result, out_ref)

  @parameterized.parameters(
      (128, 128), (128, 64), (64, 128)
  )
  def test_broadcast_in_dim_major_wgmma(self, m, n):
    dtype = jnp.float16

    def kernel(ctx, gmem_input, gmem_output, _):
      t = mgpu.FragmentedArray.load_untiled(
          gmem_input, layout=mgpu.WGMMA_COL_LAYOUT, optimized=False
      )
      t.broadcast_in_dim(
          (m, n), (1,), mgpu.WGMMA_LAYOUT
      ).store_untiled(gmem_output, optimized=False)

    inp = self.prng.uniform(-1, 1, (n,)).astype(dtype)
    out_shape = jax.ShapeDtypeStruct((m, n), dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (inp,), out_shape, inp
    )(inp)
    out_ref = jax.lax.broadcast_in_dim(inp, (m, n), (1,))
    np.testing.assert_array_equal(result, out_ref)

  @parameterized.parameters(
      ((128), (4, 128)),
      ((1, 128), (2, 128)),
      ((1, 128), (4, 128)),
      ((1, 256), (2, 256)),
      ((128, ), (1, 3, 1, 2, 4, 128)),
      ((1, 1, 128,), (1, 3, 1, 2, 4, 128)),
      ((1, 1, 1, 1, 1, 128,), (1, 3, 1, 2, 4, 128)),
      ((2, 4, 128,), (1, 3, 1, 2, 4, 128)),
      ((1, 1, 1, 2, 4, 128,), (1, 3, 1, 2, 4, 128)),
      ((2, 8, 8), (2, 8, 8)),
  )
  def test_broadcast_major_strided(self, in_shape, out_shape):
    dtype = jnp.float16
    def kernel(ctx, gmem_input, gmem_output, _):
      t = mgpu.FragmentedArray.load_strided(gmem_input, vec_size=1)
      t.broadcast(out_shape).store_untiled(gmem_output, optimized=False)
    inp = self.prng.uniform(-1, 1, in_shape).astype(dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (inp,), jax.ShapeDtypeStruct(out_shape, dtype), inp
    )(inp)
    np.testing.assert_array_equal(result, jnp.broadcast_to(inp, out_shape))

  @parameterized.parameters(*mtu.RegisterLayout)
  def test_broadcast_splat(self, layout):
    out_shape = (128, 128)

    def body(ctx, out_ref, scratch):
      del ctx, scratch
      c42 = arith.constant(ir.IntegerType.get_signless(32), 42)
      arr = mgpu.FragmentedArray.splat(c42, (128,), is_signed=True)
      out_layout = layout.to_mgpu(out_shape, jnp.int32)
      result = arr.broadcast_in_dim(out_shape, (0,), out_layout)
      result.store_untiled(out_ref, optimized=False)

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct(out_shape, jnp.int32),
        smem_scratch_shape=[],
    )
    np.testing.assert_array_equal(
        kernel(), np.full(out_shape, 42, dtype=np.int32)
    )

  def test_warp_tree_reduce(self):
    def kernel(ctx, out, *_):
      del ctx
      i32 = ir.IntegerType.get_signless(32)
      tid = gpu.thread_id(gpu.Dimension.x)
      value = arith.index_cast(i32, tid)
      grp = warp_tree_reduce(value, arith.addi, 4)
      memref.store(grp, out, [tid])

    x = np.arange(128, dtype=jnp.int32)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), x, [],
    )()
    for i in range(0, 128, 4):
      x[i:i + 4] = jnp.sum(x[i:i + 4])

    np.testing.assert_array_equal(result, x)

  @parameterized.parameters(2, 4)
  def test_fast_i8_convert(self, reg_length):
    jax_dtype_to = jnp.dtype(jnp.bfloat16)
    jax_dtype_from = jnp.dtype(jnp.int8)
    mlir_dtype_to = utils.dtype_to_ir_type(jax_dtype_to)
    def kernel(ctx, inp, out, smem):
      del ctx, smem
      arr = mgpu.FragmentedArray.load_strided(inp, is_signed=True)
      assert ir.VectorType(arr.registers.flat[0].type).shape == [reg_length]
      arr.astype(mlir_dtype_to).store_untiled(out, optimized=False)

    x = jnp.arange(-128, 128, dtype=jax_dtype_from)
    x = jnp.tile(x, reg_length // 2)
    reference = x.astype(jax_dtype_to)

    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, reference, None,
    )(x)
    np.testing.assert_array_equal(result, reference)

  @parameterized.parameters(
      ([64 * 4], mgpu.WGMMA_ROW_LAYOUT),
      ([64 * 4, 8 * 2], mgpu.WGMMA_LAYOUT),
  )
  def test_splat_relayout(self, shape, new_layout):
    def kernel(ctx, _):
      # No assertions, we are just checking there are no compile-time errors.
      arr = mgpu.FragmentedArray.splat(c(42.0, ir.F32Type.get()), shape)
      arr.to_layout(new_layout)

    _ = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), (), (), None)()

  @parameterized.parameters(
      (mgpu.WGMMA_LAYOUT, mgpu.WGMMA_TRANSPOSED_LAYOUT),
      (mgpu.TCGEN05_LAYOUT, mgpu.TCGEN05_TRANSPOSED_LAYOUT),
      (mgpu.WGMMA_TRANSPOSED_LAYOUT, mgpu.WGMMA_LAYOUT),
      (mgpu.TCGEN05_TRANSPOSED_LAYOUT, mgpu.TCGEN05_LAYOUT),
  )
  def test_transpose_relayout(self, src_layout, dst_layout):
    def is_transposed(layout):
      return (
          layout == mgpu.WGMMA_TRANSPOSED_LAYOUT
          or layout == mgpu.TCGEN05_TRANSPOSED_LAYOUT
      )

    def body(ctx, src, dst, scratch):
      del ctx, scratch
      if is_transposed(src_layout):
        src = utils.memref_transpose(src, (1, 0))
      src_reg = mgpu.FragmentedArray.load_untiled(
          src, layout=src_layout, optimized=False
      )
      dst_reg = src_reg.to_layout(dst_layout)
      if is_transposed(dst_layout):
        dst = utils.memref_transpose(dst, (1, 0))
      dst_reg.store_untiled(dst, optimized=False)

    shape = (128, 128)
    dtype = jnp.float32
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(jax.ShapeDtypeStruct(shape, dtype),),
        out_shape=jax.ShapeDtypeStruct(shape, dtype),
        smem_scratch_shape=[],
    )
    x = self.prng.uniform(-1, 1, shape).astype(dtype)
    np.testing.assert_array_equal(kernel(x), x.T)

  @parameterized.parameters(
      (jnp.float16, jnp.float16),  # Noop
      (jnp.int16, jnp.bfloat16),
      (jnp.int16, jnp.float16),
      (jnp.uint16, jnp.float16),
      (jnp.float32, jnp.int32),
      (jnp.float32, jnp.uint32),
      (jnp.uint32, jnp.int32),
      (jnp.int32, jnp.uint32),
  )
  def test_bitcast(self, in_dtype, out_dtype):
    out_ir_type = utils.dtype_to_ir_type(out_dtype)
    in_is_signed = utils.is_signed(in_dtype)
    out_is_signed = utils.is_signed(out_dtype)

    def kernel(ctx, inp, out, smem):
      del ctx, smem
      arr = mgpu.FragmentedArray.load_strided(inp, is_signed=in_is_signed)
      arr = arr.bitcast(out_ir_type, output_is_signed=out_is_signed)
      arr.store_untiled(out)

    x = jnp.arange(256, dtype=in_dtype)
    reference = jax.lax.bitcast_convert_type(x, out_dtype)

    result = mgpu.as_gpu_kernel(
        kernel,
        (1, 1, 1),
        (128, 1, 1),
        x,
        reference,
        None,
    )(x)
    np.testing.assert_array_equal(result, reference)

  @parameterized.parameters(jnp.float32, jnp.float16, jnp.bfloat16)
  def test_optimization_barrier(self, dtype):
    def kernel(ctx, inp, out, smem):
      del ctx, smem
      arr = mgpu.FragmentedArray.load_strided(inp)
      arr2 = arr * 2
      arr, arr2 = mgpu.optimization_barrier(arr, arr2)
      (arr + arr2).store_untiled(out)

    x = jnp.arange(256, dtype=dtype)

    f = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, None)
    np.testing.assert_array_equal(f(x), x * 3)

  def test_optimization_barrier_with_single_value(self):
    shape = (64, 64)
    value = 5.0
    dtype = jnp.float32
    def kernel(ctx, out, smem):
      del ctx, smem
      mlir_type = utils.dtype_to_ir_type(dtype)
      arr = mgpu.FragmentedArray.splat(c(value, mlir_type), shape)
      arr = mgpu.optimization_barrier(arr)
      arr.store_untiled(out)

    out_shape = jax.ShapeDtypeStruct(shape, dtype)
    f = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ())
    np.testing.assert_array_equal(f(), jnp.full(shape, value, dtype=dtype))

  def test_convert_bool_to_u8(self):
    m, n = 128, 128
    def kernel(ctx, dst, _):
      i8 = ir.IntegerType.get_signless(8)
      iota = iota_tensor(m, n, jnp.uint16)
      (iota > 10).astype(i8, is_signed=False).store_untiled(dst, optimized=False)

    out_shape = jax.ShapeDtypeStruct((m, n), jnp.int8)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    iota = np.arange(m * n, dtype=jnp.uint16).reshape(m, n)
    np.testing.assert_array_equal(result, (iota > 10).astype(jnp.uint8))

  @parameterized.product(
      dtype=(
          jnp.float64,
          jnp.bfloat16,
          jnp.float16,
          jnp.float8_e4m3fn,
          jnp.float8_e5m2,
          jnp.int8,
          jnp.uint8,
          jnp.int4,
          jnp.uint4,
      ),
      m_warps=(1, 2, 4),
  )
  def test_warp_mma(self, dtype, m_warps):
    m, n, k = 128, 128, 128
    dtype = jnp.dtype(dtype)
    if dtype == jnp.float64 and not config.enable_x64.value:
      self.skipTest("float64 requires x64 to be enabled")
    is_integer = jnp.issubdtype(dtype, jnp.integer)
    if is_integer:
      acc_dtype = jnp.int32
    elif dtype == jnp.float64:
      acc_dtype = jnp.float64
    else:
      acc_dtype = jnp.float32
    def kernel(ctx, acc, a, b, out, _):
      mma_layouts = mgpu.MMALayouts(utils.dtype_to_ir_type(dtype), m_warps=m_warps)
      ab_is_signed = utils.is_signed(dtype)
      a_fa = mgpu.FragmentedArray.load_untiled(
          a, layout=mma_layouts.lhs, is_signed=ab_is_signed, optimized=False)
      b_transposed = utils.memref_transpose(b, (1, 0))
      b_fa = mgpu.FragmentedArray.load_untiled(
          b_transposed, layout=mma_layouts.rhs, is_signed=ab_is_signed,
          optimized=False)
      acc_fa = mgpu.FragmentedArray.load_untiled(
          acc, layout=mma_layouts.acc, is_signed=utils.is_signed(acc_dtype),
          optimized=False)
      mgpu.mma(acc_fa, a_fa, b_fa).store_untiled(out, optimized=False)

    if is_integer:
      iinfo = jnp.iinfo(dtype)
      a = self.prng.integers(iinfo.min, iinfo.max + 1, (m, k)).astype(dtype)
      b = self.prng.integers(iinfo.min, iinfo.max + 1, (n, k)).astype(dtype)
      acc = self.prng.integers(-100, 100, (m, n)).astype(acc_dtype)
    else:
      a = self.prng.uniform(-1, 1, (m, k)).astype(dtype)
      b = self.prng.uniform(-1, 1, (n, k)).astype(dtype)
      acc = self.prng.uniform(-1, 1, (m, n)).astype(acc_dtype)

    expected = acc + jax.lax.dot(
        a.astype(acc_dtype), b.astype(acc_dtype).T, precision="high"
    )
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (acc, a, b), expected, ()
    )(acc, a, b)
    if is_integer:
      np.testing.assert_array_equal(result, expected)
    else:
      atol = rtol = 1e-12 if dtype == jnp.float64 else 1e-5
      np.testing.assert_allclose(result, expected, atol=atol, rtol=rtol)

  @parameterized.product(
      dtype=(jnp.float16, jnp.int8,),
      m_warps=(4, 2, 1),
      mults=((1, 1), (2, 1), (1, 2), (3, 3)),
  )
  @jtu.thread_unsafe_test()
  def test_warp_mma_smem(self, dtype, m_warps, mults):
    n_warps = 4 // m_warps
    m_mult, n_mult = mults
    mlir_dtype = utils.dtype_to_ir_type(dtype)
    bw = utils.bytewidth(mlir_dtype)
    m, n, k = 16 * m_warps * m_mult, 8 * n_warps * n_mult, 32 // bw
    dtype = jnp.dtype(dtype)
    is_integer = jnp.issubdtype(dtype, jnp.integer)
    acc_dtype = jnp.int32 if is_integer else jnp.float32

    a_swizzle = b_swizzle = 16
    a_tiling = b_tiling = (8, a_swizzle // bw)

    def kernel(ctx, a, b, out, smem):
      smem_a, smem_b, barriers = smem
      ctx.async_copy(
          src_ref=a,
          dst_ref=smem_a,
          swizzle=a_swizzle,
          gmem_transform=(mgpu.TileTransform(a_tiling),),
          barrier=barriers[0],
      )
      ctx.async_copy(
          src_ref=b,
          dst_ref=smem_b,
          swizzle=b_swizzle,
          gmem_transform=(mgpu.TileTransform(b_tiling),),
          barrier=barriers[1],
      )
      for i in range(2):
        barriers[i].wait()
      mma_layouts = mgpu.MMALayouts(mlir_dtype, m_warps=m_warps)
      is_signed = utils.is_signed(dtype)
      a = mgpu.FragmentedArray.load_tiled(
          smem_a, layout=mma_layouts.lhs, is_signed=is_signed, swizzle=a_swizzle
      )
      smem_b_t = utils.memref_transpose(smem_b, (1, 0, 3, 2))
      b = mgpu.FragmentedArray.load_tiled(
          smem_b_t, layout=mma_layouts.rhs, is_signed=is_signed, swizzle=b_swizzle
      )
      c0 = arith.constant(
          utils.dtype_to_ir_type(acc_dtype), 0 if is_integer else 0.0
      )
      acc = mgpu.FragmentedArray.splat(
          c0, shape=(m, n), layout=mma_layouts.acc, is_signed=is_signed
      )
      acc = mgpu.mma(acc, a, b)
      acc.store_untiled(out, optimized=False)

    if is_integer:
      iinfo = jnp.iinfo(dtype)
      a = self.prng.integers(iinfo.min, iinfo.max + 1, (m, k)).astype(dtype)
      b = self.prng.integers(iinfo.min, iinfo.max + 1, (n, k)).astype(dtype)
    else:
      a = self.prng.uniform(-1, 1, (m, k)).astype(dtype)
      b = self.prng.uniform(-1, 1, (n, k)).astype(dtype)

    expected = a.astype(acc_dtype) @ b.astype(acc_dtype).T
    smem_shapes = (
        jax.ShapeDtypeStruct(mgpu.tile_shape((m, k), a_tiling), dtype),
        jax.ShapeDtypeStruct(mgpu.tile_shape((n, k), b_tiling), dtype),
        mgpu.Barrier(1, num_barriers=2),
    )

    with jtu.set_env(MOSAIC_GPU_DUMP_SASS="1"), self.capture_stdout() as sass:
      result = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), (a, b), expected, smem_shapes
      )(a, b)

    sass_lines = sass().splitlines()
    first_ldsm = next(i for i, line in enumerate(sass_lines) if "LDSM" in line)
    last_mma = max(i for i, line in enumerate(sass_lines) if "MMA" in line)
    interesting_sass = "\n".join(sass_lines[first_ldsm : last_mma + 1])
    self.assertNotIn("MOV", interesting_sass)

    if is_integer:
      np.testing.assert_array_equal(result, expected)
    else:
      np.testing.assert_allclose(result, expected, atol=1e-5)

  @parameterized.parameters(
      (jnp.uint8, jnp.uint16, 255),
      (jnp.uint8, jnp.int16, 255),
      (jnp.int8, jnp.uint16, -127),
      (jnp.int8, jnp.int16, -127),
  )
  def test_convert_int_uint(self, from_dtype, to_dtype, value):
    m, n = 1, 128
    def kernel(ctx, dst, _):
      from_mlir_dtype = utils.dtype_to_ir_type(from_dtype)
      to_mlir_dtype = utils.dtype_to_ir_type(to_dtype)
      from_arr = mgpu.FragmentedArray.splat(
          c(value, from_mlir_dtype),
          (m, n),
          is_signed=utils.is_signed(from_dtype),
      )
      to_arr = from_arr.astype(to_mlir_dtype, is_signed=utils.is_signed(to_dtype))
      to_arr.store_untiled(dst)

    out_shape = jax.ShapeDtypeStruct((m, n), to_dtype)
    result = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
    )()
    expected = jnp.full((m, n), value, dtype=from_dtype).astype(to_dtype)
    np.testing.assert_array_equal(result, expected)

  @parameterized.product(
      swizzle=(16, 32, 64, 128),
      shape=((128, 128), (8, 128), (128, 32), (48, 64)),
      to_smem=(True, False),
  )
  def test_copy_tiled(self, swizzle, shape, to_smem):
    dtype = jnp.int32
    tiling = (8, 8 * swizzle // jnp.iinfo(dtype).bits)
    def kernel(ctx, src, dst, scratch):
      smem, barrier = scratch
      if to_smem:
        mgpu.copy_tiled(src, smem, swizzle=swizzle)
        mgpu.commit_shared()
        ctx.async_copy(
            src_ref=smem,
            dst_ref=dst,
            gmem_transform=mgpu.TileTransform(tiling),
            swizzle=swizzle,
        )
        ctx.await_async_copy(0)
      else:
        ctx.async_copy(
            src_ref=src,
            dst_ref=smem,
            gmem_transform=mgpu.TileTransform(tiling),
            swizzle=swizzle,
            barrier=barrier,
        )
        barrier.wait()
        mgpu.copy_tiled(smem, dst, swizzle=swizzle)

    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    scratch_shape = [
        jax.ShapeDtypeStruct(mgpu.tile_shape(shape, tiling), dtype),
        mgpu.Barrier(1),
    ]
    y = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, x, scratch_shape
    )(x)
    np.testing.assert_array_equal(y, x)

  @parameterized.parameters(
      ((32, 32), (0, 5)),
      ((32, 128), (3,)),
      ((32, 32, 128), (slice(1, 3), 0)),
  )
  def test_splat_indexing(self, shape, indices):
    def _kernel(ctx, out_ref, scratch):
      del ctx, scratch
      splat = mgpu.FragmentedArray.splat(c(1.0, ir.F32Type.get()), shape)
      splat[indices].store_untiled(out_ref)

    expected = np.ones(shape, dtype=jnp.float32)[indices]
    kernel = mgpu.as_gpu_kernel(
        _kernel,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(),
        out_shape=expected,
        smem_scratch_shape=(),
    )
    np.testing.assert_array_equal(kernel(), expected)

  def test_concatenate_strided(self):
    dtype = jnp.float32
    is_signed = utils.is_signed(dtype)
    shapes = (64, 64), (128, 64)

    def _kernel(ctx, inp1, inp2, dst, _):
      arr1 = mgpu.FragmentedArray.load_strided(inp1, is_signed=is_signed)
      arr2 = mgpu.FragmentedArray.load_strided(inp2, is_signed=is_signed)
      res = fa.concatenate([arr1, arr2], axis=0)
      res.store_untiled(dst, optimized=False)

    x0 = np.arange(math.prod(shapes[0]), dtype=np.float32).reshape(shapes[0])
    x1 = np.arange(math.prod(shapes[1]), dtype=np.float32).reshape(shapes[1])
    expected = np.concatenate([x0, x1], axis=0)
    out_shape = jax.ShapeDtypeStruct(expected.shape, dtype)
    kernel = mgpu.as_gpu_kernel(
        _kernel, (1, 1, 1), (128, 1, 1), (x0, x1), out_shape, ()
    )
    np.testing.assert_array_equal(kernel(x0, x1), expected)

  @parameterized.named_parameters(
      ("axis0", 0, ((64, 64), (128, 64))),
      ("axis1", 1, ((64, 64), (64, 128))),
  )
  def test_concatenate_tiled(self, axis, shapes):
    layout = fa.WGMMA_LAYOUT
    dtype = jnp.float32
    is_signed = utils.is_signed(dtype)

    def _kernel(ctx, inp1, inp2, dst, _):
      x0 = mgpu.FragmentedArray.load_untiled(
          inp1, layout=layout, is_signed=is_signed, optimized=False
      )
      x1 = mgpu.FragmentedArray.load_untiled(
          inp2, layout=layout, is_signed=is_signed, optimized=False
      )
      res = fa.concatenate([x0, x1], axis=axis)
      res.store_untiled(dst, optimized=False)

    x0 = np.arange(math.prod(shapes[0]), dtype=np.float32).reshape(shapes[0])
    x1 = np.arange(math.prod(shapes[1]), dtype=np.float32).reshape(shapes[1])
    expected = np.concatenate([x0, x1], axis=axis)
    out_shape = jax.ShapeDtypeStruct(expected.shape, dtype)
    kernel = mgpu.as_gpu_kernel(
        _kernel, (1, 1, 1), (128, 1, 1), (x0, x1), out_shape, ()
    )
    np.testing.assert_array_equal(kernel(x0, x1), expected)

  @parameterized.product(
      dtype=(jnp.int16,),
      transpose=(False, True),
  )
  @jtu.thread_unsafe_test()
  def test_stmatrix(self, dtype, transpose):
    m, k = 128, 128
    dtype = jnp.dtype(dtype)
    swizzle = 128
    def kernel(ctx, x, out, x_smem):
      mma_layouts = mgpu.MMALayouts(utils.dtype_to_ir_type(dtype))
      x_fa = mgpu.FragmentedArray.load_untiled(
          x, layout=mma_layouts.lhs, is_signed=True, optimized=False
      )
      x_smem_store = x_smem
      if transpose:
        x_smem_store = mgpu.memref_transpose(x_smem, (1, 0, 3, 2))
      x_fa.store_tiled(x_smem_store, swizzle=swizzle)
      mgpu.warpgroup_barrier()
      copy(x_smem, mgpu.TileTransform((8, swizzle // 2)).apply(out), swizzle)

    x = self.prng.integers(-10000, 10000, (m, k)).astype(dtype)
    out_shape = (k, m) if transpose else (m, k)
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      y = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x,
          jax.ShapeDtypeStruct(out_shape, dtype),
          jax.ShapeDtypeStruct(tile_shape(out_shape, (8, swizzle // 2)), dtype),
      )(x)
      expected = x.T if transpose else x
      np.testing.assert_array_equal(y, expected)
    expected_instr = (
        "stmatrix.sync.aligned.m8n8.x4.trans.shared"
        if transpose
        else "stmatrix.sync.aligned.m8n8.x4.shared"
    )
    self.assertIn(expected_instr, ptx())
    self.assertNotIn("st.shared", ptx())

  @parameterized.product(
      dtype=(jnp.int16,),
      transpose=(False, True),
  )
  def test_stmatrix_bank_conflict(self, dtype, transpose):
    m, k = 128, 128
    dtype = jnp.dtype(dtype)
    swizzle = 16
    def kernel(ctx, x, out, x_smem):
      mma_layouts = mgpu.MMALayouts(utils.dtype_to_ir_type(dtype))
      x_fa = mgpu.FragmentedArray.load_untiled(
          x, layout=mma_layouts.lhs, is_signed=True, optimized=False
      )
      x_smem_store = x_smem
      if transpose:
        x_smem_store = mgpu.memref_transpose(x_smem, (1, 0))
      x_fa.store_untiled(x_smem_store, swizzle=swizzle)
      mgpu.warpgroup_barrier()
      copy(x_smem, out, swizzle)

    x = self.prng.integers(-10000, 10000, (m, k)).astype(dtype)
    out_shape = (k, m) if transpose else (m, k)
    with self.assertRaises(fa.TransferPlanDerivationError):
      mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x,
          jax.ShapeDtypeStruct(out_shape, dtype),
          jax.ShapeDtypeStruct(out_shape, dtype),
      )(x)


class ProfilerTest(TestCase, jtu.JaxTestCase):

  def test_profiler(self):
    def body(ctx, input, result, scratch):
      del scratch
      with ctx.named_region("load"):
        reg = mgpu.FragmentedArray.load_strided(input)
      with ctx.named_region("store"):
        reg.store_untiled(result)

    dtype = jnp.bfloat16
    shape = (128, 128)
    jax_shape = jax.ShapeDtypeStruct(shape, dtype)
    with tempfile.TemporaryDirectory() as tmpdir:
      kernel = mgpu.as_gpu_kernel(
          body,
          grid=(1, 1, 1),
          block=(128, 1, 1),
          in_shape=(jax_shape),
          out_shape=jax_shape,
          smem_scratch_shape=[],
          prof_spec=profiler.ProfilerSpec(1024, dump_path=tmpdir),
      )
      param = self.prng.uniform(-1, 1, shape).astype(dtype)
      self.assertArraysEqual(kernel(param), param)
      [name] = os.listdir(tmpdir)
      with open(os.path.join(tmpdir, name)) as f:
        data = f.read()
        self.assertEqual(data.count('"name": "load"'), 2)
        self.assertEqual(data.count('"name": "store"'), 2)

  def test_profiler_per_warp(self):
    def body(ctx, input, result, scratch):
      del scratch
      with ctx.named_region("load"):
        reg = mgpu.FragmentedArray.load_strided(input)
      with ctx.named_region("store"):
        reg.store_untiled(result)

    dtype = jnp.bfloat16
    shape = (128, 128)
    jax_shape = jax.ShapeDtypeStruct(shape, dtype)
    with tempfile.TemporaryDirectory() as tmpdir:
      kernel = mgpu.as_gpu_kernel(
          body,
          grid=(1, 1, 1),
          block=(128, 1, 1),
          in_shape=(jax_shape),
          out_shape=jax_shape,
          smem_scratch_shape=[],
          prof_spec=profiler.ProfilerSpec(
              1024, dump_path=tmpdir, trace_scope=mgpu.ThreadSubset.WARP
          ),
      )
      param = self.prng.uniform(-1, 1, shape).astype(dtype)
      self.assertArraysEqual(kernel(param), param)
      [name] = os.listdir(tmpdir)
      with open(os.path.join(tmpdir, name)) as f:
        data = f.read()
        # 4 warps in a block of 128 threads.
        self.assertEqual(data.count('"name": "load"'), 8)
        self.assertEqual(data.count('"name": "store"'), 8)


class LayoutTest(TestCase):

  @parameterized.named_parameters(
      ("f32", jnp.float32, 256),
      ("f16", jnp.float16, 256),
      ("f16_small", jnp.float16, 128),
  )
  def test_store_untiled_splat(self, jax_dtype, size):
    mlir_dtype = utils.dtype_to_ir_type(jax_dtype)
    def kernel(ctx, out, _):
      del ctx
      arr = mgpu.FragmentedArray.splat(
          c(1.0, mlir_dtype), (size,), is_signed=utils.is_signed(jax_dtype)
      )
      arr.store_untiled(out)
    expected = np.ones((size,), jax_dtype)
    mosaic_ones = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), (), expected, ()
    )()
    np.testing.assert_array_equal(mosaic_ones, expected)

  @parameterized.product(
      shape=((128, 128), (64, 8), (64, 256)),
      dtype=(jnp.int32, jnp.int16, jnp.int8),
  )
  def test_wgmma_tiled_layout(self, shape, dtype):
    def kernel(ctx, dst, _):
      iota = iota_tensor(*shape, dtype)
      tiled = iota.to_layout(fa._tiled_wgmma_layout(shape))
      # Note that WGMMA layouts are always (shape[0] // 64, shape[1] // 8, 2, 1)
      self.assertEqual(
          tiled.registers.shape,
          (shape[0] // 64, shape[1] // 8, 1, 1, 2, 1, 1, 1, 1),
      )
      self.assertEqual(tiled.shape, shape)
      self.assertEqual(tiled.mlir_dtype, iota.mlir_dtype)
      tiled.store_untiled(dst, optimized=False)
    ty = jax.ShapeDtypeStruct(shape, dtype)
    f = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), (), ty, ())
    expected = np.arange(math.prod(shape), dtype=dtype).reshape(shape)
    np.testing.assert_array_equal(f(), expected)

  @parameterized.product(
      dtype=[jnp.int8, jnp.int16, jnp.int32],
      swizzle=[16, 32, 64, 128],
      num_col_tiles=[1, 2, 3],
      row_tiling=[8, 64],
  )
  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_copy_tiled(self, dtype, swizzle, num_col_tiles, row_tiling):
    mlir_dtype = utils.dtype_to_ir_type(dtype)
    bw = bytewidth(mlir_dtype)
    col_tiling = swizzle // bw
    if col_tiling % 8:
      self.skipTest("WGMMA layout requires col_tiling % 8 == 0")
    m, n = 128, col_tiling * num_col_tiles
    tiling = (row_tiling, col_tiling)
    def kernel(ctx, in_, out, smems):
      smem_in, smem_out, barrier = smems
      ctx.async_copy(src_ref=in_, dst_ref=smem_in, swizzle=swizzle, barrier=barrier)
      barrier.wait()
      t = mgpu.FragmentedArray.load_tiled(
          smem_in, swizzle=swizzle, is_signed=True, layout=mgpu.WGMMA_LAYOUT
      )
      t.store_tiled(smem_out, swizzle=swizzle)
      mgpu.commit_shared()
      ctx.async_copy(src_ref=smem_out, dst_ref=out, swizzle=swizzle)
      ctx.await_async_copy(0)
    expected = (
        np.arange(m * n, dtype=dtype)
        .reshape(m // tiling[0], tiling[0], n // tiling[1], tiling[1])
        .transpose(0, 2, 1, 3)
    )

    with jtu.set_env(MOSAIC_GPU_DUMP_SASS="1"), self.capture_stdout() as sass:
      iota = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), expected, expected,
          [expected, expected, mgpu.Barrier(1)],
      )(expected)
    np.testing.assert_array_equal(iota, expected)

    # Verify that we don't use too many registers for the transfers.
    # We verify LDS and STS separately, because they might use two different
    # methods of computing offsets and we don't rely on CSE between them.
    expected_regs = swizzle // bytewidth(mlir_dtype) // 8
    # When the bytewidth is smaller than 2 the swizzle pattern changes every 2
    # column tiles, so we only need half the registers.
    if bytewidth(mlir_dtype) < 2:
      expected_regs //= 2
    for instr in ("STS", "LDS"):
      with self.subTest(instr + " count"):
        addrs = re.findall(instr + r".* \[(.*)\]", sass())
        def get_reg(addr):
          if (pos := addr.find("+")) != -1:
            return addr[:pos]
          return addr
        used_regs = {get_reg(addr) for addr in addrs}
        try:
          self.assertLessEqual(len(used_regs), expected_regs)
        except:
          problematic_device_patterns = ("RTX PRO 6000 Blackwell", "GB10$")
          if match := jtu.device_kind_match(problematic_device_patterns):
            self.skipTest(f"{match} uses more registers for an unknown reason")
          raise

  def test_copy_for_upcast(self):
    dtype = jnp.int8
    swizzle = 128
    col_tiling = swizzle // bytewidth(utils.dtype_to_ir_type(dtype))
    m, n = 128, col_tiling * 2
    tiling = (64, col_tiling)
    layout = fa.WGMMA_LAYOUT_UPCAST_2X
    def kernel(ctx, in_, out, smems):
      smem_in, smem_out, barrier = smems
      ctx.async_copy(src_ref=in_, dst_ref=smem_in, swizzle=swizzle, barrier=barrier)
      barrier.wait()
      t = mgpu.FragmentedArray.load_tiled(
          smem_in, swizzle=swizzle, is_signed=True, layout=layout
      )
      t.store_tiled(smem_out, swizzle=swizzle)
      mgpu.commit_shared()
      ctx.async_copy(src_ref=smem_out, dst_ref=out, swizzle=swizzle)
      ctx.await_async_copy(0)
    x = jax.random.randint(
        jax.random.key(42), tile_shape((m, n), tiling), -128, 127, dtype=dtype
    )
    f = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, x, [x, x, mgpu.Barrier(1)],
    )
    np.testing.assert_array_equal(f(x), x)

  @parameterized.product(
      dtype=[jnp.int16, jnp.int32],  # TODO(apaszke): More dtypes
      swizzle=[16, 32, 64, 128],
      layouts=[
          (fa.WGMMA_LAYOUT, fa.WGMMA_TRANSPOSED_LAYOUT),
          (fa.TCGEN05_LAYOUT, fa.TCGEN05_TRANSPOSED_LAYOUT),
      ],
  )
  def test_transpose_tiled(self, dtype, swizzle, layouts):
    mlir_dtype = utils.dtype_to_ir_type(dtype)
    bw = bytewidth(mlir_dtype)
    col_tiling = swizzle // bw
    if bw == 2:
      m, n = 256, 192
    elif bw == 4:
      m, n = 256, 96
    else:
      raise ValueError(f"Unsupported bitwidth: {bw}")
    if jax.local_devices()[0].shared_memory_per_block_optin == 99 * 1024:
      # Only reduce if needed to fit inside SMEM so as to keep >1 row tile
      # in TCGEN05 tilings on relevant hardware.
      m = 128
    tiling = (8, col_tiling)
    if col_tiling < 8:
      self.skipTest("Swizzle too small")
    layout, transpose_layout = layouts
    def kernel(ctx, in_, out, smems):
      smem_in, smem_out, barrier = smems
      ctx.async_copy(src_ref=in_, dst_ref=smem_in, swizzle=swizzle, barrier=barrier)
      barrier.wait()
      t = mgpu.FragmentedArray.load_tiled(
          smem_in, swizzle=swizzle, is_signed=True, layout=layout
      )
      smem_out_t = memref_transpose(smem_out, (1, 0, 3, 2))
      t.to_layout(transpose_layout).store_tiled(smem_out_t, swizzle=swizzle)
      mgpu.commit_shared()
      ctx.async_copy(src_ref=smem_out, dst_ref=out, swizzle=swizzle)
      ctx.await_async_copy(0)
    x = (
        np.arange(m * n, dtype=dtype)
        .reshape(m // tiling[0], tiling[0], n // tiling[1], tiling[1])
        .transpose(0, 2, 1, 3)
    )
    y_ref = (
        np.arange(m * n, dtype=dtype)
        .reshape(m, n)
        .T.reshape(n // tiling[0], tiling[0], m // tiling[1], tiling[1])
        .transpose(0, 2, 1, 3)
    )
    with core.artificial_shared_memory_limit(None):
      y = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x, y_ref, [x, y_ref, mgpu.Barrier(1)],
      )(x)
    np.testing.assert_array_equal(y, y_ref)

  @parameterized.parameters(
      (fa.WGMMA_LAYOUT_UPCAST_2X, fa.WGMMA_LAYOUT, jnp.int8, jnp.int8, 1),
      (fa.WGMMA_LAYOUT_UPCAST_2X, fa.WGMMA_LAYOUT, jnp.int8, jnp.int8, 1, (2,)),
      (fa.WGMMA_LAYOUT_UPCAST_2X, fa.WGMMA_LAYOUT, jnp.int8, jnp.int16, 1),
      (fa.WGMMA_LAYOUT, fa.WGMMA_LAYOUT_UPCAST_2X, jnp.int8, jnp.int8, 0.5),
      (fa.WGMMA_LAYOUT, fa.WGMMA_LAYOUT_UPCAST_2X, jnp.int8, jnp.int8, 0.5, (2,)),
      (fa.WGMMA_LAYOUT, fa.WGMMA_LAYOUT_UPCAST_2X, jnp.int8, jnp.int16, 0.5),
      (
          fa.WGMMA_LAYOUT_UPCAST_4X,
          fa.WGMMA_LAYOUT_UPCAST_2X,
          jnp.int4,
          jnp.int4,
          1,
      ),
      (fa.WGMMA_LAYOUT_UPCAST_2X, fa.WGMMA_LAYOUT, jnp.int4, jnp.int4, 0.5),
      (fa.WGMMA_LAYOUT, fa.WGMMA_LAYOUT_UPCAST_2X, jnp.int4, jnp.int4, 0.25),
      (fa.WGMMA_LAYOUT_UPCAST_4X, fa.WGMMA_LAYOUT, jnp.int4, jnp.int4, 2),
  )
  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_upcast_to_wgmma(
      self,
      start_layout,
      end_layout,
      in_dtype,
      cast_dtype,
      shfl_per_reg,
      leading_shape=(),
  ):
    in_dtype = jnp.dtype(in_dtype)
    out_dtype = jnp.dtype(jnp.int16)
    out_dtype_mlir = utils.dtype_to_ir_type(out_dtype)
    swizzle = 128
    in_col_tiling = 8 * swizzle // jnp.iinfo(in_dtype).bits
    in_tiling = (8, in_col_tiling)
    out_col_tiling = swizzle // out_dtype.itemsize
    out_tiling = (8, out_col_tiling)
    m, n = 64, in_col_tiling * 2
    regs_per_thread = None
    def kernel(ctx, in_, out, smems):
      nonlocal regs_per_thread
      smem_in, smem_out, barrier = smems
      ctx.async_copy(src_ref=in_, dst_ref=smem_in, swizzle=swizzle, barrier=barrier)
      barrier.wait()
      t = mgpu.FragmentedArray.load_tiled(
          smem_in,
          swizzle=swizzle,
          is_signed=True,
          layout=start_layout,
          tiling_rank=2,
      )
      regs_per_thread = t.registers.size
      t = t.astype(utils.dtype_to_ir_type(cast_dtype), is_signed=True)
      t = t.to_layout(end_layout)
      t = t.astype(out_dtype_mlir, is_signed=True)
      t.store_tiled(smem_out, swizzle=swizzle, tiling_rank=2)
      mgpu.commit_shared()
      ctx.async_copy(src_ref=smem_out, dst_ref=out, swizzle=swizzle)
      ctx.await_async_copy(0)
    def tile(x, tiling):
      return x.reshape(
          *x.shape[:-2],
          x.shape[-2] // tiling[0],
          tiling[0],
          x.shape[-1] // tiling[1],
          tiling[1],
      ).swapaxes(-3, -2)
    in_iinfo = jnp.iinfo(in_dtype)
    x = jax.random.randint(
        jax.random.key(42),
        (*leading_shape, m, n),
        in_iinfo.min,
        in_iinfo.max,
        dtype=jnp.int32,
    ).astype(in_dtype)
    xt = tile(x, in_tiling)
    y = x.astype(out_dtype)
    yt = tile(y, out_tiling)
    f = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), xt, yt, [xt, yt, mgpu.Barrier(1)],
    )
    with jtu.set_env(MOSAIC_GPU_DUMP_SASS="1"), self.capture_stdout() as sass:
      yt_kernel = f(xt)
      jax.block_until_ready(yt_kernel)
    np.testing.assert_array_equal(yt_kernel, yt)
    try:
      self.assertEqual(sass().count("SHFL.BFLY"), regs_per_thread * shfl_per_reg)
    except:
      problematic_device_patterns = ("RTX PRO 6000 Blackwell", "GB10$")
      if match := jtu.device_kind_match(problematic_device_patterns):
        self.skipTest(f"{match} requires more SHFL.BFLY for an unknown reason")
      raise

  @parameterized.product(
      in_length=[1, 2, 4, 8],
      out_length=[1, 2, 4, 8],
  )
  def test_convert_tmem_native_vector_length(self, in_length, out_length):
    dtype = jnp.dtype(jnp.int16)
    def kernel(ctx, in_, out, smems):
      smem_in, smem_out, barrier = smems
      ctx.async_copy(src_ref=in_, dst_ref=smem_in, barrier=barrier)
      barrier.wait()
      t = mgpu.FragmentedArray.load_untiled(
          smem_in, layout=mgpu.tmem_native_layout(in_length),
          is_signed=True, optimized=False
      )
      t = t.to_layout(mgpu.tmem_native_layout(out_length))
      t.store_untiled(smem_out, optimized=False)
      mgpu.commit_shared()
      ctx.async_copy(src_ref=smem_out, dst_ref=out)
      ctx.await_async_copy(0)
    iinfo = jnp.iinfo(dtype)
    x = jax.random.randint(
        jax.random.key(42), (128, 128), iinfo.min, iinfo.max, dtype=jnp.int16
    )
    f = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, x, [x, x, mgpu.Barrier(1)],
    )
    y = f(x)
    np.testing.assert_array_equal(y, x)


@dataclasses.dataclass(frozen=True)
class Tile:
  """Defines a Tile transform in a TestCaseInput.

  Note that we cannot simply alias mgpu_dialect.TileTransformAttr.get, because
  we do not have an MLIR context at the point we define the TestCaseInput.
  """

  tiling: tuple[int, ...]

  def attr(self):
    return mgpu_dialect.TileTransformAttr.get(self.tiling)


@dataclasses.dataclass(frozen=True)
class Swizzle:
  """Defines a Swizzle transform in a TestCaseInput.

  Note that we cannot simply alias mgpu_dialect.SwizzleTransformAttr.get,
  because we do not have an MLIR context at the point we define the
  TestCaseInput.
  """

  swizzle: mgpu_dialect.SwizzlingMode

  def attr(self):
    return mgpu_dialect.SwizzleTransformAttr.get(self.swizzle)


class MosaicGpuDialectTest(TestCase, jtu.JaxTestCase):
  """Device tests with lowering from the MLIR dialect and layout inference."""

  def setUp(self):
    if mgpu_dialect is None:
      raise self.skipTest("Test requires Mosaic GPU dialect")
    super().setUp()

  @parameterized.product(
      layout_tag=tuple(mtu.RegisterLayout),
      dtype=(jnp.bfloat16, jnp.int8),
      optimized=(True, False, None),
  )
  def test_smem_gmem_registers_load_store(
      self, layout_tag: mtu.RegisterLayout, dtype, optimized
  ):
    if layout_tag == mtu.RegisterLayout.WG_SPLAT:
      self.skipTest("WG_SPLAT is not supported for `vector.load`.")
    # We don't infer optimized transfer-compatible transforms for load/store to
    # registers with TCGEN05_TMEM_NATIVE layout.
    load_optimized = optimized != False
    if load_optimized and layout_tag == mtu.RegisterLayout.TCGEN05_TMEM_NATIVE:
      self.skipTest(
          "Optimized loads not supported for TCGEN05_TMEM_NATIVE layout"
      )

    # We use a (128, 128) shape to construct the layout, but specialize the
    # actual test shape to the layout in order to avoid long test times for some
    # layouts.
    layout = layout_tag.to_mgpu((128, 128), dtype)
    if isinstance(layout, fa.TiledLayout):
      # Use a tiling_rank + 1-dimensional shape to exercise the path where
      # tiling does not span all dimensions.
      shape = (1, *layout.base_tile_shape[:-1], 2 * layout.base_tile_shape[-1])
    else:
      shape = (128, 128)
    layout_attr = layouts.to_layout_attr(layout)

    def body(ctx, param: ir.Value, result: ir.Value, smem: list[ir.Value]):
      del ctx

      # GMEM -> Registers
      reg = mgpu_dialect.vector_load(param, optimized=False)
      reg = mgpu_dialect.layout_cast(reg, layout_attr)

      # Registers -> SMEM
      mgpu_dialect.vector_store(reg, smem, optimized=optimized)

      # SMEM -> Registers
      reg = mgpu_dialect.vector_load(smem, optimized=optimized)
      reg = mgpu_dialect.layout_cast(reg, layout_attr)

      # Registers -> GMEM
      mgpu_dialect.vector_store(reg, result, optimized=False)

    jax_shape = jax.ShapeDtypeStruct(shape, dtype)
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=jax_shape,
        out_shape=jax_shape,
        smem_scratch_shape=jax_shape,
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    param = self.prng.uniform(-1, 1, shape).astype(dtype)
    self.assertArraysEqual(kernel(param), param)

  @parameterized.parameters(
      (mgpu.WGMMA_LAYOUT, mgpu.WGMMA_TRANSPOSED_LAYOUT),
      (mgpu.WGMMA_TRANSPOSED_LAYOUT, mgpu.WGMMA_LAYOUT),
  )
  def test_transposed_load_store(self, src_layout, dst_layout):
    def is_transposed(layout):
      return layout == mgpu.WGMMA_TRANSPOSED_LAYOUT

    def body(ctx, src_ref, dst_ref, scratch):
      del ctx, scratch
      if is_transposed(src_layout):
        src_ref = utils.memref_transpose(src_ref, (1, 0))
      if is_transposed(dst_layout):
        dst_ref = utils.memref_transpose(dst_ref, (1, 0))
      src_reg = mgpu_dialect.vector_load(src_ref, optimized=False)
      src_layout_attr = layouts.to_layout_attr(src_layout)
      src_reg = mgpu_dialect.layout_cast(src_reg, src_layout_attr)
      dst_layout_attr = layouts.to_layout_attr(dst_layout)
      dst_reg = mgpu_dialect.layout_cast(src_reg, dst_layout_attr)
      mgpu_dialect.vector_store(dst_reg, dst_ref, optimized=False)

    shape = (128, 128)
    dtype = jnp.float32
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(jax.ShapeDtypeStruct(shape, dtype),),
        out_shape=jax.ShapeDtypeStruct(shape, dtype),
        smem_scratch_shape=[],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    x = self.prng.uniform(-1, 1, shape).astype(dtype)
    np.testing.assert_array_equal(kernel(x), x.T)

  def test_layout_cast_convert_fuses_slice_after_layout_cast(self):
    # When we emit fast converts for slice producers we can fuse these into
    # the conversion routine itself. If we don't do this LLVM emits too many
    # shifts.
    # In this regression test we make sure we hit the fusion path.
    shape = (128, 128)
    in_dtype = jnp.int4
    out_dtype = jnp.bfloat16

    def body(ctx, src_ref, dst_ref, scratch):
      del ctx, scratch
      src_reg = mgpu_dialect.vector_load(src_ref, optimized=False)
      src_layout_attr = layouts.to_layout_attr(fa.WGMMA_LAYOUT_UPCAST_2X)
      src_reg = mgpu_dialect.layout_cast(src_reg, src_layout_attr)
      dst_layout_attr = layouts.to_layout_attr(fa.WGMMA_LAYOUT)
      dst_reg = mgpu_dialect.layout_cast(src_reg, dst_layout_attr)

      bf16 = ir.BF16Type.get()
      out_vec_ty = ir.VectorType.get(ir.VectorType(dst_reg.type).shape, bf16)
      conv_reg = arith.sitofp(out_vec_ty, dst_reg)
      mgpu_dialect.vector_store(conv_reg, dst_ref, optimized=False)

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(jax.ShapeDtypeStruct(shape, in_dtype),),
        out_shape=jax.ShapeDtypeStruct(shape, out_dtype),
        smem_scratch_shape=[],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    x = self.prng.integers(-8, 7, shape).astype(in_dtype)
    expected = jax.lax.convert_element_type(x, out_dtype)

    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      np.testing.assert_array_equal(kernel(x), expected)

    self.assertNotRegex(ptx(), r"bfe\.s32")

  def test_pointwise_kernel(self):
    def add(ctx, a, b, result, smem):
      del ctx, smem

      # GMEM -> registers
      a = mgpu_dialect.vector_load(a, optimized=False)
      b = mgpu_dialect.vector_load(b, optimized=False)

      # Computation
      add = arith.addf(a, b)

      # Registers -> GMEM
      mgpu_dialect.vector_store(add, result, optimized=False)

    dtype = jnp.bfloat16
    shape = (128, 128)
    jax_shape = jax.ShapeDtypeStruct(shape, dtype)
    kernel = mgpu.as_gpu_kernel(
        add,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(jax_shape, jax_shape),
        out_shape=jax_shape,
        smem_scratch_shape=[],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    x = self.prng.uniform(-1, 1, shape).astype(dtype)
    y = self.prng.uniform(-1, 1, shape).astype(dtype)

    self.assertArraysEqual(kernel(x, y), x + y)

  @staticmethod
  def kernel_with_tma_cases(dtype: jnp.dtype):
    @dataclasses.dataclass()
    class TestCaseInput:
      shape: tuple[int, ...]
      shape_sliced: tuple[int, ...] = ()
      slice_indices: tuple[int, ...] = ()
      slice_lengths: tuple[int, ...] = ()
      transforms: tuple[Tile | Swizzle, ...] = ()

      def __post_init__(self):
        if not self.shape_sliced:
          self.shape_sliced = self.shape
        if not self.slice_lengths:
          self.slice_lengths = self.shape_sliced
        if not self.slice_indices:
          self.slice_indices = tuple([0] * len(self.slice_lengths))

    result = []
    for swizzle in mgpu_dialect.SwizzlingMode:
      n = swizzle * 8 // jnp.finfo(dtype).bits
      if swizzle == mgpu_dialect.SwizzlingMode.kNoSwizzle:
        #  We need at least one case with no transforms, as this is handled
        #  differently.
        result.append(TestCaseInput(shape=[128, n]))
      result.extend([
          TestCaseInput(
              shape=[128, n],
              transforms=[Swizzle(swizzle)],
          ),
          TestCaseInput(
              shape=[256, n],
              shape_sliced=[128, n],
              transforms=[Swizzle(swizzle)],
          ),
          TestCaseInput(
              shape=[2, 128, n],
              shape_sliced=[128, n],
              slice_lengths=[-1, 128, n],
              slice_indices=[1, 0, 0],
              transforms=[Swizzle(swizzle)],
          ),
          TestCaseInput(
              shape=[256, n],
              shape_sliced=[128, n],
              transforms=[Tile([64, n]), Swizzle(swizzle)],
          ),
      ])
    return result

  @parameterized.parameters(kernel_with_tma_cases(jnp.bfloat16))
  def test_kernel_with_tma(self, test_case):
    def add(
        ctx: launch_context.LaunchContext,
        in_gmem_ref: ir.Value,
        result_gmem_ref: ir.Value,
        smem: list[ir.Value],
    ):
      del ctx
      smem_ref, tma_barrier = smem

      elt_type = ir.MemRefType(in_gmem_ref.type).element_type
      memref_bytes = utils.bytewidth(elt_type) * math.prod(
          test_case.shape_sliced
      )

      i32 = ir.IntegerType.get_signless(32)
      slice_indices = [arith.constant(i32, i) for i in test_case.slice_indices]

      # GMEM -> SMEM
      tma_barrier.arrive_expect_tx(memref_bytes)
      if test_case.transforms:
        transforms_attr = ir.ArrayAttr.get(
            [t.attr() for t in test_case.transforms]
        )
        smem_ref = mgpu_dialect.with_transforms(smem_ref, transforms_attr)
      mgpu_dialect.async_load(
          source=in_gmem_ref,
          destination=smem_ref,
          barrier=tma_barrier.as_barrier_memref(),
          indices=slice_indices,
          slice_lengths=test_case.slice_lengths,
          collective=ir.ArrayAttr.get([]),
      )
      tma_barrier.wait()

      # SMEM -> GMEM
      zero_index = arith.constant(i32, 0)
      mgpu_dialect.async_store(
          source=smem_ref,
          destination=result_gmem_ref,
          indices=[zero_index] * len(test_case.shape_sliced),
          slice_lengths=test_case.shape_sliced,
      )
      nvvm.cp_async_bulk_wait_group(0)
      utils.warpgroup_barrier()

    dtype = jnp.bfloat16

    jax_shape = jax.ShapeDtypeStruct(test_case.shape, dtype)
    jax_shape_sliced = jax.ShapeDtypeStruct(test_case.shape_sliced, dtype)
    kernel = mgpu.as_gpu_kernel(
        add,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(jax_shape),
        out_shape=jax_shape_sliced,
        smem_scratch_shape=[
            jax_shape_sliced,
            mgpu.Barrier(1),
        ],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    x = self.prng.uniform(-1, 1, test_case.shape).astype(dtype)

    input_slice = tuple(
        slice(i * abs(l), (i + 1) * abs(l))
        for i, l in zip(test_case.slice_indices, test_case.slice_lengths)
    )
    self.assertArraysEqual(
        kernel(x),
        (x[input_slice]).reshape(test_case.shape_sliced),
    )

  def test_pointwise_kernel_with_tma(self):
    def add(
        ctx: launch_context.LaunchContext,
        a_gmem_ref: ir.Value,
        b_gmem_ref: ir.Value,
        result_gmem_ref: ir.Value,
        smem: list[ir.Value],
    ):
      del ctx
      a_smem_ref, b_smem_ref, result_smem_ref, tma_barrier = smem

      memref_type = ir.MemRefType(a_gmem_ref.type)
      shape = memref_type.shape
      elt_type = memref_type.element_type

      zero_i32 = arith.constant(ir.IntegerType.get_signless(32), 0)
      zero_slice_indices = [zero_i32] * memref_type.rank

      # GMEM -> SMEM
      memref_bytes = utils.bytewidth(elt_type) * math.prod(shape)
      tma_barrier.arrive_expect_tx(2 * memref_bytes)
      mgpu_dialect.async_load(
          source=a_gmem_ref,
          destination=a_smem_ref,
          barrier=tma_barrier.as_barrier_memref(),
          indices=zero_slice_indices,
          slice_lengths=shape,
          collective=ir.ArrayAttr.get([]),
      )
      mgpu_dialect.async_load(
          source=b_gmem_ref,
          destination=b_smem_ref,
          barrier=tma_barrier.as_barrier_memref(),
          indices=zero_slice_indices,
          slice_lengths=shape,
          collective=ir.ArrayAttr.get([]),
      )
      tma_barrier.wait()

      # SMEM -> registers
      a = mgpu_dialect.vector_load(a_smem_ref)
      b = mgpu_dialect.vector_load(b_smem_ref)

      # Computation
      add = arith.addf(arith.addf(a, b), b)

      # Registers -> SMEM
      mgpu_dialect.vector_store(add, result_smem_ref)

      # SMEM -> GMEM
      mgpu_dialect.async_store(
          source=result_smem_ref,
          destination=result_gmem_ref,
          indices=zero_slice_indices,
          slice_lengths=shape,
      )
      nvvm.cp_async_bulk_wait_group(0)
      utils.warpgroup_barrier()

    dtype = jnp.bfloat16

    spec = jax.ShapeDtypeStruct((2, 3, 4, 64), dtype)
    kernel = mgpu.as_gpu_kernel(
        add,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(spec, spec),
        out_shape=spec,
        smem_scratch_shape=[
            spec,
            spec,
            spec,
            mgpu.Barrier(1),
        ],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    x = self.prng.uniform(-1, 1, spec.shape).astype(dtype)
    y = self.prng.uniform(-1, 1, spec.shape).astype(dtype)

    self.assertArraysEqual(kernel(x, y), x + y + y)

  @parameterized.parameters(
      ((64,), (64, 128), [0]),
      ((64,), (128, 64), [1]),
  )
  def test_broadcast_in_dim(self, input_shape, output_shape, bcast_dims):
    element_value = 42.0
    layout = fa.WGMMA_ROW_LAYOUT if bcast_dims[0] == 0 else fa.WGMMA_COL_LAYOUT
    def body(ctx, result_gmem_ref, scratch):
      del ctx, scratch

      # Create input in registers
      f32 = ir.F32Type.get()
      x_type = ir.VectorType.get(input_shape, f32)
      c = arith.constant(f32, element_value)
      x = vector.broadcast(x_type, c)

      # Computation
      out_type = ir.VectorType.get(output_shape, f32)
      cast = mgpu_dialect.layout_cast(x, layouts.to_layout_attr(layout))
      expanded = mgpu_dialect.broadcast_in_dim(out_type, cast, bcast_dims)

      # Registers -> GMEM
      mgpu_dialect.vector_store(expanded, result_gmem_ref, optimized=False)

    dtype = jnp.float32
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct(output_shape, dtype),
        smem_scratch_shape=[],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    x = np.full(input_shape, element_value, dtype=dtype)
    self.assertArraysEqual(
        kernel(), jax.lax.broadcast_in_dim(x, output_shape, bcast_dims)
    )

  @parameterized.parameters(
      (jnp.float32, 5.0, 2.0, vector.CombiningKind.ADD),
      (jnp.float32, 5.0, 2.0, vector.CombiningKind.MAXIMUMF),
      (jnp.float32, 5.0, 7.0, vector.CombiningKind.MAXIMUMF),
      (jnp.int32, 5, 2, vector.CombiningKind.MAXSI),
      (jnp.int32, -5, -2, vector.CombiningKind.MAXSI),
      (jnp.int32, -2, -5, vector.CombiningKind.MAXSI),
      (jnp.uint32, 5, 2, vector.CombiningKind.MAXUI),
      (jnp.uint32, 2, 5, vector.CombiningKind.MAXUI),
      #
      # TODO(dasenov): Add tests for wgmma_col_layout output once
      # fragmented_array.reduce supports that.
  )
  def test_vector_multi_dim_reduction(
      self,
      dtype,
      input_value,
      init_value,
      kind,
  ):
    input_shape = (128, 64)
    output_shape = (128,)
    red_dims = [1]

    def body(ctx, result_gmem_ref, scratch):
      del ctx, scratch
      el_type = utils.dtype_to_ir_type(dtype)

      # Create source in registers
      source_type = ir.VectorType.get(input_shape, el_type)
      c = arith.constant(el_type, input_value)
      source = vector.broadcast(source_type, c)

      # Create accumulator in registers
      acc_type = ir.VectorType.get(output_shape, el_type)
      c = arith.constant(el_type, init_value)
      acc = vector.broadcast(acc_type, c)

      # Cast inputs
      source = mgpu_dialect.layout_cast(
          source, layouts.to_layout_attr(fa.WGMMA_LAYOUT)
      )
      acc_layout = (
          fa.WGMMA_ROW_LAYOUT if red_dims[0] == 1 else fa.WGMMA_COL_LAYOUT
      )
      acc = mgpu_dialect.layout_cast(acc, layouts.to_layout_attr(acc_layout))

      # Computation
      reduced = vector.multi_reduction(kind, source, acc, red_dims)

      # Registers -> GMEM
      mgpu_dialect.vector_store(reduced, result_gmem_ref, optimized=False)

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct(output_shape, dtype),
        smem_scratch_shape=[],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    source = np.full(input_shape, input_value, dtype=dtype)
    acc = np.full(output_shape, init_value, dtype=dtype)
    if kind == vector.CombiningKind.ADD:
      red = jax.lax.reduce_sum(source, red_dims)
      red = red + acc
    else:
      red = jax.lax.reduce_max(source, red_dims)
      red = jax.lax.max(red, acc)
    self.assertArraysEqual(kernel(), red)

  @parameterized.parameters(fa.WGMMA_ROW_LAYOUT, fa.WGMMA_COL_LAYOUT)
  def test_wgmma_row_col_store(self, in_layout):
    element_value = 42.0
    shape = (64, )
    def body(ctx, result_gmem_ref, smem):
      del ctx

      # Create input in registers
      f32 = ir.F32Type.get()
      x_type = ir.VectorType.get(shape, f32)
      c = arith.constant(f32, element_value)
      x = vector.broadcast(x_type, c)
      cast = mgpu_dialect.layout_cast(x, layouts.to_layout_attr(in_layout))

      # Registers -> SMEM
      mgpu_dialect.vector_store(cast, smem)
      utils.commit_shared()

      # SMEM -> GMEM
      zero_i32 = arith.constant(ir.IntegerType.get_signless(32), 0)
      mgpu_dialect.async_store(
          source=smem,
          destination=result_gmem_ref,
          indices=[zero_i32],
          slice_lengths=shape,
      )
      nvvm.cp_async_bulk_wait_group(0)
      utils.warpgroup_barrier()

    dtype = jnp.float32
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct(shape, dtype),
        smem_scratch_shape=jax.ShapeDtypeStruct(shape, dtype),
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    x = np.full(shape, element_value, dtype=dtype)
    self.assertArraysEqual(kernel(), x)

  @parameterized.parameters(
      # Positive offsets will be passsed as static offsets.
      # Negative offsets will be converted to positive dynamic offsets.
      dict(
          full_shape=(2, 3, 128, 64),
          sub_shape=(32, 64),
          offsets=[-1, 0, -96, 0],
          tiling=None,
          swizzle=None,
      ),
      dict(
          full_shape=(3, 128, 64),
          sub_shape=(32, 64),
          offsets=[-2, -96, 0],
          tiling=[32, 64],
          swizzle=mgpu_dialect.SwizzlingMode.k128ByteSwizzle,
      ),
      dict(
          full_shape=(128, 128),
          sub_shape=(64,),
          offsets=[-1, 64],
          tiling=[64],
          swizzle=mgpu_dialect.SwizzlingMode.k128ByteSwizzle,
      ),
  )
  def test_subview(
      self,
      full_shape,
      sub_shape,
      offsets,
      tiling,
      swizzle,
  ):
    assert len(sub_shape) <= 2
    sizes = [1] * (len(full_shape) - len(sub_shape)) + list(sub_shape)

    def body(
        ctx: launch_context.LaunchContext,
        full_gmem_ref: ir.Value,
        sub_gmem_ref: ir.Value,
        smem: list[ir.Value],
    ):
      del ctx
      full_smem_ref, tma_barrier = smem

      zero_i32 = arith.constant(ir.IntegerType.get_signless(32), 0)
      # GMEM -> SMEM
      operand_elt_type = ir.MemRefType(full_gmem_ref.type).element_type
      bytes = utils.bytewidth(operand_elt_type) * math.prod(full_shape)
      tma_barrier.arrive_expect_tx(bytes)
      mgpu_dialect.async_load(
          source=full_gmem_ref,
          destination=full_smem_ref,
          barrier=tma_barrier.as_barrier_memref(),
          indices=[zero_i32] * len(full_shape),
          slice_lengths=full_shape,
          collective=ir.ArrayAttr.get([]),
      )
      tma_barrier.wait()

      # SubView
      mixed_offsets = [
          o if o >= 0 else arith.constant(ir.IndexType.get(), -o)
          for o in offsets
      ]

      full_ref_type = ir.MemRefType(full_smem_ref.type)
      dynamic = ir.ShapedType.get_dynamic_stride_or_offset()
      rhs_subview_ref_type = ir.MemRefType.get(
          shape=sub_shape,
          element_type=full_ref_type.element_type,
          layout=ir.StridedLayoutAttr.get(
              dynamic, [full_shape[-1], 1] if len(sub_shape) == 2 else [1]
          ),
          memory_space=full_ref_type.memory_space,
      )
      sub_smem_ref = memref.subview(
          full_smem_ref,
          mixed_offsets,
          sizes,
          strides=[1] * len(sizes),
          result_type=rhs_subview_ref_type,
      )

      transforms = []
      if tiling is not None:
        transforms.append(mgpu_dialect.TileTransformAttr.get(tiling))
      if swizzle is not None:
        transforms.append(mgpu_dialect.SwizzleTransformAttr.get(swizzle))

      if transforms:
        sub_smem_ref = mgpu_dialect.with_transforms(
            sub_smem_ref,
            transforms=ir.ArrayAttr.get(transforms),
        )

      # SMEM -> GMEM
      mgpu_dialect.async_store(
          source=sub_smem_ref,
          destination=sub_gmem_ref,
          indices=[zero_i32] * len(sub_shape),
          slice_lengths=sub_shape,
      )
      nvvm.cp_async_bulk_wait_group(0)

    el_type = jnp.bfloat16
    full_jax_shape = jax.ShapeDtypeStruct(full_shape, el_type)
    result_jax_shape = jax.ShapeDtypeStruct(sub_shape, el_type)

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(full_jax_shape),
        out_shape=result_jax_shape,
        smem_scratch_shape=[full_jax_shape, mgpu.Barrier(1)],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    x = self.prng.uniform(0, 10, full_shape).astype(el_type)
    slicing = tuple(slice(abs(o), abs(o) + s) for o, s in zip(offsets, sizes))
    self.assertArraysEqual(kernel(x), x[slicing].reshape(sub_shape))

  def test_custom_primitive_op(self):
    # This test exercises the following cases:
    # - The lowering handles nested blocks and regions (e.g. `scf.IfOp`).
    # - The lowering updates references to inlined operations.
    def body(ctx, result, scratch):
      del ctx, scratch
      result_ty = ir.MemRefType(result.type)
      el_type = result_ty.element_type
      index = ir.IndexType.get()
      op = mgpu_dialect.CustomPrimitiveOp(
          result=[],
          operands_=[result],
          in_layouts=[],
          in_transforms=[],
          out_layouts=[],
      )
      args_ty = [arg.type for arg in op.operands_]
      block = op.body.blocks.append(*args_ty)
      with ir.InsertionPoint(block):
        is_leader_thread = single_thread_predicate()
        with when(is_leader_thread):
          c5 = arith.constant(el_type, 5)
          memref.store(c5, block.arguments[0], [c(0, index)])
        mgpu_dialect.return_([])

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        cluster=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct((1,), jnp.int32),
        smem_scratch_shape=(),
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    self.assertArraysEqual(kernel(), np.array([5], dtype=np.int32))

  def test_profiler(self):
    def body(ctx, input, result, scratch):
      del scratch
      with ctx.named_region("load"):
        reg = mgpu_dialect.vector_load(input, optimized=False)
      with ctx.named_region("store"):
        mgpu_dialect.vector_store(reg, result, optimized=False)

    dtype = jnp.bfloat16
    shape = (128, 128)
    jax_shape = jax.ShapeDtypeStruct(shape, dtype)
    with tempfile.TemporaryDirectory() as tmpdir:
      kernel = mgpu.as_gpu_kernel(
          body,
          grid=(1, 1, 1),
          block=(128, 1, 1),
          in_shape=(jax_shape),
          out_shape=jax_shape,
          smem_scratch_shape=[],
          prof_spec=profiler.ProfilerSpec(1024, dump_path=tmpdir),
          thread_semantics=mgpu.LoweringSemantics.Warpgroup,
      )
      param = self.prng.uniform(-1, 1, shape).astype(dtype)
      self.assertArraysEqual(kernel(param), param)
      [name] = os.listdir(tmpdir)
      with open(os.path.join(tmpdir, name)) as f:
        data = f.read()
        self.assertEqual(data.count('"name": "load"'), 2)
        self.assertEqual(data.count('"name": "store"'), 2)

  def test_arrive_expect_tx_non_uniform(self):
    def body(ctx, result, scratch):
      del ctx
      [mbar] = scratch
      mbar_ref = mbar.as_barrier_memref()
      i32 = ir.IntegerType.get_signless(32)
      index = ir.IndexType.get()

      mgpu_dialect.arrive_expect_tx(barrier=mbar_ref, expect_tx=c(16, i32))

      # Complete tx with CustomPrimitiveOp because there is no `complete_tx` op
      # in the dialect.
      op = mgpu_dialect.CustomPrimitiveOp(
          result=[],
          operands_=[mbar_ref],
          in_layouts=[],
          in_transforms=[ir.ArrayAttr.get([])],
          out_layouts=[],
      )
      args_ty = [arg.type for arg in op.operands_]
      block = op.body.blocks.append(*args_ty)
      with ir.InsertionPoint(block):
        is_leader = single_thread_predicate()
        tx = arith.constant(i32, 16)
        barrier = utils.DialectBarrierRef.from_barrier_memref(block.arguments[0])
        barrier.barrier_ref.complete_tx(tx, predicate=is_leader)
        mgpu_dialect.return_([])

      mbar.wait()

      is_leader_thread = single_thread_predicate()
      with when(is_leader_thread):
        c1 = arith.constant(i32, 1)
        memref.store(c1, result, [c(0, index)])

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        cluster=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct((1,), jnp.int32),
        smem_scratch_shape=[mgpu.Barrier(1)],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    self.assertArraysEqual(kernel(), np.array([1], dtype=np.int32))

  @parameterized.parameters(((128,),), ((128, 128),))
  def test_tma_collective_async_cp(self, in_shape):
    def body(ctx, src, dst, scratch):
      del ctx
      tmp, barrier = scratch
      i32 = ir.IntegerType.get_signless(32)
      zero_i32 = arith.constant(i32, 0)
      src_type = ir.MemRefType(src.type)
      barrier.arrive_expect_tx(
          utils.bytewidth(src_type.element_type) * math.prod(src_type.shape)
      )
      mgpu_dialect.async_load(
          source=src,
          destination=tmp,
          indices=[zero_i32] * src_type.rank,
          slice_lengths=src_type.shape,
          collective=ir.ArrayAttr.get([
              ir.IntegerAttr.get(i32, mgpu_dialect.Dimension.x),
          ]),
          barrier=barrier.as_barrier_memref(),
      )
      barrier.wait()
      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      block_id = arith.index_cast(i32, block_id)
      mgpu_dialect.async_store(
          source=tmp,
          destination=dst,
          indices=[block_id] + [zero_i32] * src_type.rank,
          slice_lengths=[-1, *src_type.shape],
      )

    dtype = jnp.float32
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(2, 1, 1),
        cluster=(2, 1, 1),
        block=(128, 1, 1),
        in_shape=jax.ShapeDtypeStruct(in_shape, dtype),
        out_shape=jax.ShapeDtypeStruct((2, *in_shape), dtype),
        smem_scratch_shape=[
            jax.ShapeDtypeStruct(in_shape, dtype),
            mgpu.Barrier(1),
        ],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    x = self.prng.uniform(-1, 1, in_shape).astype(dtype)
    self.assertArraysEqual(kernel(x), jnp.stack([x, x], axis=0))

  def test_vector_extract_strided_slice(self):
    def body(ctx, src, dst, scratch):
      del ctx, scratch
      src_vec = mgpu_dialect.vector_load(src, optimized=False)
      src_vec = mgpu_dialect.layout_cast(
          src_vec, layouts.to_layout_attr(fa.WGMMA_LAYOUT)
      )
      dst_type = ir.MemRefType(dst.type)
      dest_vec_type = ir.VectorType.get(dst_type.shape, dst_type.element_type)
      sliced_vec = vector.extract_strided_slice(
          dest_vec_type,
          src_vec,
          offsets=[0, 64],
          sizes=[64, 64],
          strides=[1, 1],
      )
      mgpu_dialect.vector_store(sliced_vec, dst, optimized=False)

    dtype = jnp.float32
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=jax.ShapeDtypeStruct((128, 128), dtype),
        out_shape=jax.ShapeDtypeStruct((64, 64), dtype),
        smem_scratch_shape=[],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    x = self.prng.uniform(-1, 1, (128, 128)).astype(dtype)
    self.assertArraysEqual(kernel(x), x[0:64, 64:128])

  @parameterized.product(
      dtype=(jnp.float32, jnp.int32, jnp.uint32),
      dimension=(0, 1),
  )
  def test_broadcasted_iota(self, dtype, dimension):
    def body(ctx, out, scratch):
      del ctx, scratch
      result_type = ir.VectorType.get(out.type.shape, out.type.element_type)
      iota = mgpu_dialect.broadcasted_iota(result_type, dimension)
      mgpu_dialect.vector_store(iota, out, optimized=False)

    shape = (128, 128)
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct(shape, dtype),
        smem_scratch_shape=[],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    expected = jax.lax.broadcasted_iota(dtype, shape, dimension)
    self.assertArraysEqual(kernel(), expected)

  @parameterized.parameters(
      ((4, 64, 64), [[0], [1], [2]], (4, 64, 64), False),
      ((4, 64, 64), [[0], [1, 2], [3]], (4, 4, 16, 64), False),
      ((4, 8, 16, 64), [[0], [1], [2, 3], [4]], (4, 8, 2, 8, 64), False),
      ((4, 64, 64), [[0, 1], [2], [3]], (2, 2, 64, 64), True),
  )
  def test_memref_expand_shape(
      self, input_shape, reassociation, output_shape, has_transforms
  ):
    def body(
        ctx: launch_context.LaunchContext,
        in_gmem_ref: ir.Value,
        out_gmem_ref: ir.Value,
        smem: list[ir.Value],
    ):
      del ctx
      in_smem_ref, tma_barrier = smem

      zero_i32 = arith.constant(ir.IntegerType.get_signless(32), 0)
      # GMEM -> SMEM
      operand_elt_type = ir.MemRefType(in_gmem_ref.type).element_type
      bytes = utils.bytewidth(operand_elt_type) * math.prod(input_shape)
      tma_barrier.arrive_expect_tx(bytes)
      mgpu_dialect.async_load(
          source=in_gmem_ref,
          destination=in_smem_ref,
          barrier=tma_barrier.as_barrier_memref(),
          indices=[zero_i32] * len(input_shape),
          slice_lengths=input_shape,
          collective=ir.ArrayAttr.get([]),
      )
      tma_barrier.wait()

      # ExpandShape
      expanded_smem_ref = memref.expand_shape(
          result=ir.MemRefType.get(
              output_shape,
              in_smem_ref.type.element_type,
              memory_space=in_smem_ref.type.memory_space,
          ),
          src=in_smem_ref,
          reassociation=reassociation,
          output_shape=[],
          static_output_shape=output_shape,
      )

      if has_transforms:
        transforms = [
            mgpu_dialect.TileTransformAttr.get((32,)),
            mgpu_dialect.SwizzleTransformAttr.get(64),
        ]
        expanded_smem_ref = mgpu_dialect.with_transforms(
            expanded_smem_ref, transforms=ir.ArrayAttr.get(transforms),
        )

      # SMEM -> GMEM
      mgpu_dialect.async_store(
          source=expanded_smem_ref,
          destination=out_gmem_ref,
          indices=[zero_i32] * len(output_shape),
          slice_lengths=output_shape,
      )
      nvvm.cp_async_bulk_wait_group(0)

    el_type = jnp.bfloat16
    in_jax_shape = jax.ShapeDtypeStruct(input_shape, el_type)
    result_jax_shape = jax.ShapeDtypeStruct(output_shape, el_type)

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(in_jax_shape),
        out_shape=result_jax_shape,
        smem_scratch_shape=[in_jax_shape, mgpu.Barrier(1)],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    x = self.prng.uniform(0, 10, input_shape).astype(el_type)
    self.assertArraysEqual(kernel(x), x.reshape(output_shape))

  @parameterized.product(
      dtype=(jnp.int32, jnp.int64, jnp.uint32, jnp.uint64, jnp.float32, jnp.float16, jnp.bfloat16),
      reduction_op=("add", "min", "max", "inc", "dec", "and", "or", "xor"),
  )
  def test_async_store_reduction(self, dtype, reduction_op):

    if not config.enable_x64.value and dtype in (jnp.int64, jnp.uint64):
      self.skipTest("x64 support is disabled")

    if reduction_op in ("min", "max"):
      if dtype in (jnp.int32, jnp.int64):
        reduction_op = "s" + reduction_op
      elif dtype in (jnp.uint32, jnp.uint64):
        reduction_op = "u" + reduction_op

    if (
        not launch_context._is_tma_reduction_op_supported(
            reduction_op,
            utils.dtype_to_ir_type(dtype),
        )
        or (
            dtype in (jnp.uint32, jnp.uint64)
            and reduction_op in ("smin", "smax")
        )
        or (
            dtype in (jnp.int32, jnp.int64) and reduction_op in ("umin", "umax")
        )
        or dtype == jnp.int32 and reduction_op in ("inc", "dec")
    ):
      self.skipTest("TMA does not support this reduction op for this dtype")

    shape = (8, 128)

    def body(ctx, src, dst, smem):
      del ctx
      src_smem_ref, tma_barrier = smem
      i32 = ir.IntegerType.get_signless(32)
      zero = arith.constant(i32, 0)
      indices = [zero, zero]
      slice_lengths = src_smem_ref.type.shape

      tma_barrier.arrive_expect_tx(
          utils.bitwidth(src_smem_ref.type.element_type) * math.prod(shape) // 8
      )

      mgpu_dialect.async_load(
          source=src,
          destination=src_smem_ref,
          barrier=tma_barrier.as_barrier_memref(),
          indices=indices,
          slice_lengths=slice_lengths,
          collective=ir.ArrayAttr.get([]),
      )

      tma_barrier.wait()

      reduction_attr = getattr(
          mgpu_dialect.TMAReduction, reduction_op.capitalize()
      )

      mgpu_dialect.async_store(
          source=src_smem_ref,
          destination=dst,
          indices=indices,
          slice_lengths=slice_lengths,
          reduction_op=reduction_attr,
      )
      nvvm.cp_async_bulk_wait_group(0)

    prng_key = jax.random.key(1234)
    k0, k1 = jax.random.split(prng_key, 2)
    if dtype in (jnp.bfloat16, jnp.float16, jnp.float32):
      src = jax.random.uniform(k0, shape, dtype, -10, 10)
      dst = jax.random.uniform(k1, shape, dtype, -10, 10)
    else:
      src = jax.random.randint(k0, shape, -10, 10).astype(dtype)
      dst = jax.random.randint(k1, shape, -10, 10).astype(dtype)

    if reduction_op == "add":
      expected = src + dst
    elif reduction_op in ("min", "smin", "umin"):
      expected = jnp.minimum(src, dst)
    elif reduction_op in ("max", "smax", "umax"):
      expected = jnp.maximum(src, dst)
    elif reduction_op == "and":
      expected = src & dst
    elif reduction_op == "or":
      expected = src | dst
    elif reduction_op == "xor":
      expected = src ^ dst
    elif reduction_op == "inc":
      expected = jnp.where(dst >= src, 0, dst + 1)
    elif reduction_op == "dec":
      expected = jnp.where((dst == 0) | (dst > src), src, dst - 1)
    else:
      raise ValueError(f"Unsupported reduction op: {reduction_op}")

    jax_shape = jax.ShapeDtypeStruct(shape, dtype)
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(jax_shape),
        out_shape=(),
        inout_shape=(jax_shape,),
        smem_scratch_shape=[jax_shape, mgpu.Barrier(1)],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    np.testing.assert_array_equal(kernel(src, dst)[0], expected)


class MosaicGpuDialectSm90ATest(Sm90ATestCase, jtu.JaxTestCase):

  def setUp(self):
    if mgpu_dialect is None:
      raise self.skipTest("Test requires Mosaic GPU dialect")
    super().setUp()

  @parameterized.product(
      swizzle=tuple(mgpu_dialect.SwizzlingMode),
      transpose_lhs=(False, True),
      transpose_rhs=(False, True),
      lhs_in_registers=(False, True),
  )
  def test_wgmma_kernel_with_tma(
      self, swizzle, transpose_lhs, transpose_rhs, lhs_in_registers
  ):
    if swizzle == mgpu_dialect.SwizzlingMode.kNoSwizzle:
      self.skipTest("No swizzle is not supported by wgmma")

    if transpose_lhs and lhs_in_registers:
      self.skipTest("The A operand can only be transposed if it is in SMEM.")

    swizzle_elems = swizzle // np.dtype(jnp.bfloat16).itemsize
    tiling_m, tiling_n, tiling_k = 64, swizzle_elems, swizzle_elems

    groups_m, groups_n, groups_k = 4, 1, 1
    m, n, k = groups_m * tiling_m, groups_n * tiling_n, groups_k * tiling_k

    lhs_shape = (k, m) if transpose_lhs else (m, k)
    rhs_shape = (n, k) if transpose_rhs else (k, n)
    out_shape = (m, n)

    def matmul(
        ctx: launch_context.LaunchContext,
        lhs_gmem_ref: ir.Value,
        rhs_gmem_ref: ir.Value,
        result_gmem_ref: ir.Value,
        smem: list[ir.Value],
    ):
      del ctx
      lhs_smem_ref, rhs_smem_ref, result_smem_ref, tma_barrier = smem

      operand_elt_type = ir.MemRefType(lhs_gmem_ref.type).element_type
      bytes_a = utils.bytewidth(operand_elt_type) * math.prod(lhs_shape)
      bytes_b = utils.bytewidth(operand_elt_type) * math.prod(rhs_shape)

      # GMEM -> SMEM
      zero_i32 = arith.constant(ir.IntegerType.get_signless(32), 0)
      tma_barrier.arrive_expect_tx(bytes_a + bytes_b)
      mgpu_dialect.async_load(
          source=lhs_gmem_ref,
          destination=lhs_smem_ref,
          barrier=tma_barrier.as_barrier_memref(),
          indices=[zero_i32] * len(lhs_shape),
          slice_lengths=lhs_shape,
          collective=ir.ArrayAttr.get([]),
      )
      mgpu_dialect.async_load(
          source=rhs_gmem_ref,
          destination=rhs_smem_ref,
          barrier=tma_barrier.as_barrier_memref(),
          indices=[zero_i32] * len(rhs_shape),
          slice_lengths=rhs_shape,
          collective=ir.ArrayAttr.get([]),
      )
      tma_barrier.wait()

      # Computation
      shape_result = ir.MemRefType(result_gmem_ref.type).shape
      result_elt_type = ir.MemRefType(result_gmem_ref.type).element_type
      acc_elt_type = ir.F32Type.get()
      acc_type = ir.VectorType.get(shape_result, acc_elt_type)
      zero_acc = arith.constant(
          result_elt_type, ir.FloatAttr.get(acc_elt_type, 0.0)
      )
      accumulator = vector.broadcast(acc_type, zero_acc)

      if transpose_lhs:
        lhs_smem_ref = utils.memref_transpose(lhs_smem_ref, (1, 0))
      if transpose_rhs:
        rhs_smem_ref = utils.memref_transpose(rhs_smem_ref, (1, 0))

      if lhs_in_registers:
        # SMEM -> Registers
        lhs_operand = mgpu_dialect.vector_load(lhs_smem_ref)
      else:
        lhs_operand = lhs_smem_ref

      result = mgpu_dialect.wgmma(
          accumulator,
          lhs_operand,
          rhs_smem_ref,
      )

      nvvm.wgmma_commit_group_sync_aligned()
      nvvm.wgmma_wait_group_sync_aligned(0)

      # Registers -> SMEM
      mgpu_dialect.vector_store(result, result_smem_ref)

      # SMEM -> GMEM
      mgpu_dialect.async_store(
          source=result_smem_ref,
          destination=result_gmem_ref,
          indices=[zero_i32, zero_i32],
          slice_lengths=shape_result,
      )
      nvvm.cp_async_bulk_wait_group(0)

    operand_type = jnp.bfloat16
    acctype = jnp.float32
    lhs_jax_shape = jax.ShapeDtypeStruct(lhs_shape, operand_type)
    rhs_jax_shape = jax.ShapeDtypeStruct(rhs_shape, operand_type)
    result_jax_shape = jax.ShapeDtypeStruct(out_shape, acctype)
    kernel = mgpu.as_gpu_kernel(
        matmul,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(lhs_jax_shape, rhs_jax_shape),
        out_shape=result_jax_shape,
        smem_scratch_shape=[
            lhs_jax_shape,
            rhs_jax_shape,
            result_jax_shape,
            mgpu.Barrier(1),
        ],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    prng_key = jax.random.key(1234)
    k0, k1 = jax.random.split(prng_key, 2)

    x = jax.random.randint(k0, lhs_shape, 0, 2).astype(operand_type)
    y = jax.random.randint(k1, rhs_shape, 0, 2).astype(operand_type)

    transpose = lambda x, t: x.T if t else x
    self.assertArraysAllClose(
        kernel(x, y),
        np.matmul(transpose(x, transpose_lhs), transpose(y, transpose_rhs)),
        atol=0,
        rtol=0,
    )

  @parameterized.product(
      dtype=(jnp.int8, jnp.uint8),
      lhs_in_smem=(False, True),
  )
  def test_integer_wgmma(self, dtype, lhs_in_smem):
    m, k, n = 64, 128, 64

    def body(ctx, lhs_gmem, rhs_gmem, result_gmem, scratch):
      del ctx
      lhs_smem, rhs_smem, tma_barrier = scratch

      i32 = ir.IntegerType.get_signless(32)
      zero = arith.constant(i32, 0)

      tma_barrier.arrive_expect_tx(m * k + k * n)
      mgpu_dialect.async_load(
          source=lhs_gmem,
          destination=lhs_smem,
          barrier=tma_barrier.as_barrier_memref(),
          indices=[zero, zero],
          slice_lengths=lhs_smem.type.shape,
          collective=ir.ArrayAttr.get([]),
      )
      mgpu_dialect.async_load(
          source=rhs_gmem,
          destination=rhs_smem,
          barrier=tma_barrier.as_barrier_memref(),
          indices=[zero, zero],
          slice_lengths=rhs_smem.type.shape,
          collective=ir.ArrayAttr.get([]),
      )
      tma_barrier.wait()

      acc_type = ir.VectorType.get((m, n), i32)
      acc = vector.broadcast(acc_type, zero)
      lhs = lhs_smem if lhs_in_smem else mgpu_dialect.vector_load(lhs_smem)
      # Only f16 WGMMA supports transposes
      rhs_smem = utils.memref_transpose(rhs_smem, (1, 0))
      result = mgpu_dialect.wgmma(acc, lhs, rhs_smem)
      nvvm.wgmma_commit_group_sync_aligned()
      nvvm.wgmma_wait_group_sync_aligned(0)
      mgpu_dialect.vector_store(result, result_gmem, optimized=False)

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(
            jax.ShapeDtypeStruct((m, k), dtype),
            jax.ShapeDtypeStruct((n, k), dtype),
        ),
        out_shape=jax.ShapeDtypeStruct((m, n), jnp.int32),
        smem_scratch_shape=[
            jax.ShapeDtypeStruct((m, k), dtype),
            jax.ShapeDtypeStruct((n, k), dtype),
            mgpu.Barrier(1),
        ],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    # Use small values to avoid overflow, [0, 8) for u8 and (-8, 8) for s8.
    is_signed = jnp.issubdtype(dtype, jnp.signedinteger)
    low, high = (-8, 8) if is_signed else (0, 8)
    lhs = self.prng.uniform(low, high, (m, k)).astype(dtype)
    rhs = self.prng.uniform(low, high, (n, k)).astype(dtype)
    self.assertArraysEqual(
        kernel(lhs, rhs),
        np.matmul(lhs.astype(jnp.int32), rhs.astype(jnp.int32).T),
    )


class MosaicGpuDialectTCGen05Test(TestCase, jtu.JaxTestCase, jtu.CudaArchSpecificTest):

  def setUp(self):
    self.skip_unless_tcgen05()
    # No artificially lowered limit for arch-specific tests
    super().setUp(artificial_shared_memory_limit=None)

  @parameterized.named_parameters(
      ("unpacked", (128, 77), jnp.bfloat16, 1, False),
      ("packed", (128, 128), jnp.bfloat16, 2, False),
      ("collective", (128, 64), jnp.bfloat16, 1, True),
  )
  def test_tmem_alloc_dealloc(self, shape, dtype, packing, collective):
    tmem_type = ir.MemRefType.get(
        shape,
        utils.dtype_to_ir_type(dtype),
        memory_space=utils.tmem(),
    )

    def body(
        ctx: launch_context.LaunchContext, x: ir.Value, smem: list[ir.Value]
    ):
      # We need to have a result `x` otherwise the kernel will not be generated.
      del ctx, x
      tmem_ref = mgpu_dialect.tmem_alloc(
          result=tmem_type,
          smem_ptr=smem,
          collective=collective,
          packing=packing,
      )

      mgpu_dialect.tmem_relinquish_alloc_permit(collective=collective)
      mgpu_dialect.tmem_dealloc(tmem_ref)

    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      mgpu.as_gpu_kernel(
          body,
          grid=(2 if collective else 1, 1, 1),
          cluster=(2 if collective else 1, 1, 1),
          block=(128, 1, 1),
          in_shape=(),
          out_shape=(jax.ShapeDtypeStruct((), jnp.int32),),
          smem_scratch_shape=jax.ShapeDtypeStruct((), jnp.int32),
          thread_semantics=mgpu.LoweringSemantics.Warpgroup,
      )()
    [alloc] = re.findall(
        r"tcgen05.alloc.cta_group::([12]).sync.aligned.shared::cta.b32",
        ptx(),
    )
    self.assertEqual(alloc[0], '2' if collective else '1')

    [ld] = re.findall(
        r"ld.shared.b32\s+([%\w]+),\s+\[__dynamic_shmem__0\];",
        ptx(),
    )
    [dealloc] = re.findall(
        r"tcgen05.dealloc.cta_group::([12]).sync.aligned.b32\s+([%\w]+),",
        ptx(),
    )
    self.assertEqual(dealloc[0], '2' if collective else '1')
    self.assertEqual(dealloc[1], ld)
    [relinquish] = re.findall(
        r"tcgen05.relinquish_alloc_permit.cta_group::([12]).sync.aligned;",
        ptx(),
    )
    self.assertEqual(relinquish[0], "2" if collective else "1")

  @parameterized.named_parameters(
      ("unpacked", 1, None),
      ("packed", 2, None),
      ("custom layout", None, tcgen05.tmem_default_layout(packing=1)),
  )
  def test_tmem_load_store(self, packing, layout):
    # TODO(bchetioui): add layout inference logic to handle packed/unpacked int8s.
    dtype = jnp.bfloat16
    shape = (128, 128)

    def body(
        ctx: launch_context.LaunchContext,
        input: ir.Value,
        result: ir.Value,
        tmem: list[ir.Value],
    ):
      del ctx

      # GMEM -> registers
      r_in = mgpu_dialect.vector_load(input, optimized=False)

      # registers -> TMEM
      mgpu_dialect.async_store_tmem(r_in, tmem)
      tcgen05.commit_tmem()

      # TMEM ->registers
      r_out = mgpu_dialect.async_load_tmem(tmem)
      # no need to wait in this case, see:
      # https://docs.jax.dev/en/latest/pallas/gpu/reference.html#allocating-the-accumulator-using-tmem

      # Registers -> GMEM
      mgpu_dialect.vector_store(r_out, result, optimized=False)

    jax_shape = jax.ShapeDtypeStruct(shape, dtype)
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        cluster=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=jax_shape,
        out_shape=jax_shape,
        smem_scratch_shape=mgpu.TMEM(
            shape, dtype, packing=packing, layout=layout
        ),
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    key = jax.random.key(1234)
    x = jax.random.randint(key, shape, -10, 10).astype(dtype)
    self.assertArraysEqual(kernel(x), x)

  def test_async_store_smem_to_tmem(self):
    dtype = jnp.float32
    shape = (128, 128)

    def body(ctx, gmem_in, gmem_out, scratch):
      del ctx
      smem, copy_barrier, tmem = scratch
      reg_in = mgpu_dialect.vector_load(gmem_in, optimized=False)
      layout = layouts.to_layout_attr(fa.WGMMA_LAYOUT)
      reg_in = mgpu_dialect.layout_cast(reg_in, layout)
      mgpu_dialect.vector_store(reg_in, smem, optimized=False)

      mgpu_dialect.async_store_smem_to_tmem(smem, tmem)
      mgpu_dialect.tcgen05_commit_arrive(copy_barrier.as_barrier_memref())
      copy_barrier.wait(orders_tensor_core=True)

      reg_out = mgpu_dialect.async_load_tmem(tmem)
      mgpu_dialect.vector_store(reg_out, gmem_out, optimized=False)

    jax_shape = jax.ShapeDtypeStruct(shape, dtype)
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        cluster=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=jax_shape,
        out_shape=jax_shape,
        smem_scratch_shape=[
            jax_shape,
            mgpu.Barrier(1, orders_tensor_core=True),
            mgpu.TMEM(shape, dtype, packing=1),
        ],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    key = jax.random.key(1234)
    x = jax.random.normal(key, shape).astype(dtype)
    self.assertArraysEqual(kernel(x), x)

  def test_async_store_smem_to_tmem_collective(self):
    dtype = jnp.float32
    shape = (128, 128)

    def body(ctx, gmem_in, gmem_out, scratch):
      smem, copy_barrier, cluster_barrier, tmem = scratch

      block_id = gpu.cluster_block_id(gpu.Dimension.x)

      m_index = arith.muli(block_id, arith.constant(ir.IndexType.get(), 128))

      reg_in = mgpu_dialect.vector_load(memref_slice(gmem_in, ds(m_index, 128)), optimized=False)
      layout = layouts.to_layout_attr(fa.WGMMA_LAYOUT)
      reg_in = mgpu_dialect.layout_cast(reg_in, layout)

      mgpu_dialect.vector_store(reg_in, smem, optimized=False)

      is_first_block = arith.cmpi(
          arith.CmpIPredicate.eq, block_id, c(0, ir.IndexType.get())
      )

      # Ensure SMEM is initialized in both blocks before copying in CTA0.
      cluster_barrier.arrive()
      cluster_barrier.wait()

      with when(is_first_block):
        mgpu_dialect.async_store_smem_to_tmem(smem, tmem, collective=True)
        mgpu_dialect.tcgen05_commit_arrive(
            copy_barrier.as_barrier_memref(), collective=True
        )
      copy_barrier.wait(orders_tensor_core=True)

      reg_out = mgpu_dialect.async_load_tmem(tmem)
      sliced_result_gmem = memref_slice(gmem_out, ds(m_index, 128))
      mgpu_dialect.vector_store(reg_out, sliced_result_gmem, optimized=False)

    gmem_shape = (256, 128)
    jax_shape = jax.ShapeDtypeStruct(gmem_shape, dtype)

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(2, 1, 1),
        cluster=(2, 1, 1),
        block=(128, 1, 1),
        in_shape=jax_shape,
        out_shape=jax_shape,
        smem_scratch_shape=[
            jax.ShapeDtypeStruct(shape, dtype),
            mgpu.Barrier(1, orders_tensor_core=True),
            mgpu.ClusterBarrier(collective_dims=(gpu.Dimension.x,)),
            mgpu.TMEM(shape, dtype, collective=True, packing=1),
        ],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    key = jax.random.key(1234)
    x = jax.random.normal(key, gmem_shape).astype(dtype)
    self.assertArraysEqual(kernel(x), x)


  def _test_tcgen05_mma(
      self, m, n, swizzle, a_type, b_type, acc_type, a_in_tmem
  ):
    if acc_type == jnp.float16 and (a_type != jnp.float16 or b_type != jnp.float16):
      self.skipTest("Only f16 input is supported for f16 output.")
    if a_in_tmem and m != 128:
      self.skipTest("Only M=128 is supported for MMA with A in TMEM.")

    swizzle_elems = swizzle // np.dtype(b_type).itemsize
    groups_k = 2
    k = swizzle_elems * groups_k
    a_packing = 4 // np.dtype(a_type).itemsize
    tmem_cols = tcgen05.tmem_alloc_exact_ncols(n, exact=False)
    if a_in_tmem:
      tmem_cols += tcgen05.tmem_alloc_exact_ncols(k // a_packing, exact=False)
    if tmem_cols > 512:
      self.skipTest(
          f"Number of TMEM colums ({tmem_cols}) exceeds the limit of 512"
          " columns."
      )
    a_shape = (m, k)
    b_shape = (k, n)
    bytes_a = np.dtype(a_type).itemsize * math.prod(a_shape)
    bytes_b = np.dtype(b_type).itemsize * math.prod(b_shape)
    acc_shape = (m, n)

    def matmul(ctx, a_gmem, b_gmem, result_gmem, scratch):
      del ctx
      a_smem, b_smem, tma_barrier, mma_barrier, acc_tmem, a_tmem = scratch

      zero_i32 = arith.constant(ir.IntegerType.get_signless(32), 0)

      # GMEM -> SMEM
      tma_barrier.arrive_expect_tx(bytes_b)
      mgpu_dialect.async_load(
          source=b_gmem,
          destination=b_smem,
          barrier=tma_barrier.as_barrier_memref(),
          indices=[zero_i32] * len(b_shape),
          slice_lengths=b_shape,
          collective=ir.ArrayAttr.get([]),
      )
      tma_barrier.wait()

      if a_in_tmem:
        # GMEM -> Registers -> TMEM
        reg = mgpu_dialect.vector_load(a_gmem, optimized=False)
        mgpu_dialect.async_store_tmem(reg, a_tmem)
        tcgen05.commit_tmem()
      else:
        # GMEM -> SMEM
        tma_barrier.arrive_expect_tx(bytes_a)
        mgpu_dialect.async_load(
            source=a_gmem,
            destination=a_smem,
            barrier=tma_barrier.as_barrier_memref(),
            indices=[zero_i32] * len(a_shape),
            slice_lengths=a_shape,
            collective=ir.ArrayAttr.get([]),
        )
        tma_barrier.wait()

      mgpu_dialect.tcgen05_mma(
          accumulator=acc_tmem,
          a=a_tmem if a_in_tmem else a_smem,
          b=b_smem,
          accumulate=arith.constant(ir.IntegerType.get_signless(1), False),
      )
      mgpu_dialect.tcgen05_commit_arrive(mma_barrier.as_barrier_memref())
      mma_barrier.wait(orders_tensor_core=True)

      # TMEM -> Registers -> GMEM
      r_out = mgpu_dialect.async_load_tmem(acc_tmem)
      mgpu_dialect.vector_store(r_out, result_gmem, optimized=False)

    # Required order: SMEM -> Barrier -> TMEM.
    scratch_shape = [
        jax.ShapeDtypeStruct(a_shape, a_type) if not a_in_tmem else None,
        jax.ShapeDtypeStruct(b_shape, b_type),
        mgpu.Barrier(1),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM(acc_shape, acc_type),
        mgpu.TMEM(a_shape, a_type, packing=a_packing) if a_in_tmem else None,
    ]
    kernel = mgpu.as_gpu_kernel(
        matmul,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(
            jax.ShapeDtypeStruct(a_shape, a_type),
            jax.ShapeDtypeStruct(b_shape, b_type),
        ),
        out_shape=jax.ShapeDtypeStruct(acc_shape, acc_type),
        smem_scratch_shape=scratch_shape,
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    a = self.prng.uniform(-1, 1, a_shape).astype(a_type)
    b = self.prng.uniform(-1, 1, b_shape).astype(b_type)

    if acc_type == jnp.float32 and (
        a_type in (jnp.float8_e5m2, jnp.float8_e4m3fn)
        or b_type in (jnp.float8_e5m2, jnp.float8_e4m3fn)
    ):
      atol, rtol = 1e-2, 1e-2
    else:
      atol = 2e-2 if acc_type == jnp.float16 else 2e-5
      rtol = 8e-4 if acc_type == jnp.float16 else 1e-7
    self.assertArraysAllClose(
        kernel(a, b),
        np.matmul(a.astype(acc_type), b.astype(acc_type)),
        atol=atol,
        rtol=rtol,
    )

  @parameterized.product(
      m=(64, 128),
      n=(128, 256, 512),
      # TODO(allanrenucci): Add 32-byte swizzle once implemented.
      swizzle=(64, 128),
      ab_type=(jnp.float16, jnp.bfloat16),
      acc_type=(jnp.float16, jnp.float32),
      a_in_tmem=(False, True),
  )
  def test_tcgen05_mma(self, m, n, swizzle, ab_type, acc_type, a_in_tmem):
    self._test_tcgen05_mma(
        m, n, swizzle, ab_type, ab_type, acc_type, a_in_tmem
    )

  @parameterized.parameters(
      (jnp.float8_e5m2, jnp.float8_e5m2),
      (jnp.float8_e4m3fn, jnp.float8_e4m3fn),
      (jnp.float8_e5m2, jnp.float8_e4m3fn),
      (jnp.float8_e4m3fn, jnp.float8_e5m2),
  )
  def test_tcgen05_mma_mixed_fp8(self, a_type, b_type):
    self._test_tcgen05_mma(
        m=128,
        n=64,
        swizzle=128,
        a_type=a_type,
        b_type=b_type,
        acc_type=jnp.float32,
        a_in_tmem=False,
    )

  @parameterized.product(
      m=(128,),
      n=(32, 128, 192),
      ab_type=(jnp.float16, jnp.bfloat16, jnp.int8, jnp.float8_e4m3fn),
  )
  def test_tcgen05_mma_sparse(self, m, n, ab_type):
    if ab_type == jnp.int8:
      self.skip_unless_tcgen05_int8()

    acc_type = jnp.float32 if jnp.issubdtype(ab_type, jnp.floating) else jnp.int32
    sparse_meta_dtype = jnp.uint2

    k = 256
    a_shape = (m, k // 2)
    b_shape = (k, n)
    meta_shape = (m // 128, k // 128, 128, 64)
    acc_shape = (m, n)

    def cp(src, dst, barrier, shape, dtype):
      transfer_size_bits = dtypes.itemsize_bits(dtype) * math.prod(shape)
      assert transfer_size_bits % 8 == 0, f"{transfer_size_bits=} not byte aligned"
      transfer_size_bytes = transfer_size_bits // 8
      barrier.arrive_expect_tx(transfer_size_bytes)
      zero_i32 = arith.constant(ir.IntegerType.get_signless(32), 0)
      mgpu_dialect.async_load(
          source=src,
          destination=dst,
          barrier=barrier.as_barrier_memref(),
          indices=[zero_i32] * len(shape),
          slice_lengths=shape,
          collective=ir.ArrayAttr.get([]),
      )
      barrier.wait()

    def matmul(ctx, a_gmem, b_gmem, meta_gmem, result_gmem, scratch):
      del ctx
      (
          a_smem,
          b_smem,
          meta_smem,
          tma_barrier,
          mma_barrier,
          acc_tmem,
          meta_tmem,
      ) = scratch

      # GMEM -> SMEM
      cp(b_gmem, b_smem, tma_barrier, b_shape, ab_type)
      cp(a_gmem, a_smem, tma_barrier, a_shape, ab_type)
      cp(meta_gmem, meta_smem, tma_barrier, meta_shape, sparse_meta_dtype)

      mgpu_dialect.async_store_sparse_metadata_smem_to_tmem(meta_smem, meta_tmem)
      mgpu_dialect.tcgen05_mma(
          accumulator=acc_tmem,
          a=a_smem,
          b=b_smem,
          a_sparse_metadata=meta_tmem,
          accumulate=arith.constant(ir.IntegerType.get_signless(1), False),
      )
      mgpu_dialect.tcgen05_commit_arrive(mma_barrier.as_barrier_memref())
      mma_barrier.wait(orders_tensor_core=True)

      # TMEM -> Registers -> GMEM
      r_out = mgpu_dialect.async_load_tmem(acc_tmem)
      mgpu_dialect.vector_store(r_out, result_gmem, optimized=False)

    scratch_shape = [
        jax.ShapeDtypeStruct(a_shape, ab_type),
        jax.ShapeDtypeStruct(b_shape, ab_type),
        jax.ShapeDtypeStruct(meta_shape, sparse_meta_dtype),
        mgpu.Barrier(1),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM(acc_shape, acc_type),
        mgpu.TMEM((m, k // 2), sparse_meta_dtype),
    ]

    kernel = mgpu.as_gpu_kernel(
        matmul,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(
            jax.ShapeDtypeStruct(a_shape, ab_type),
            jax.ShapeDtypeStruct(b_shape, ab_type),
            jax.ShapeDtypeStruct(meta_shape, sparse_meta_dtype),
        ),
        out_shape=jax.ShapeDtypeStruct(acc_shape, acc_type),
        smem_scratch_shape=scratch_shape,
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    index_pairs = np.asarray(np.meshgrid(range(4), range(4))).T.reshape(-1, 2)
    valid_pairs = index_pairs[index_pairs[:, 0] < index_pairs[:, 1]]
    x_pairs = jax.random.randint(
        jax.random.key(1234), (m, k // 4), 0, 6, dtype=jnp.uint8
    )
    x_sparse = valid_pairs[x_pairs]

    def format_sparse_meta(meta):
      mn, k, _ = meta.shape
      k *= 2
      if np.dtype(ab_type).itemsize == 1:
        meta_tiled = meta.reshape(mn // 128, 128, k // 64, 64).transpose(0, 2, 1, 3)
      else:
        meta_tiled = meta.reshape(mn // 128, 8, 2, 8, k // 64, 4, 2, 8).transpose(0, 4, 1, 6, 3, 5, 2, 7)
      return meta_tiled.reshape(mn // 128, k // 64, 128, 64).astype(sparse_meta_dtype)

    x_gpu_sparse = format_sparse_meta(x_sparse)
    if jnp.issubdtype(ab_type, jnp.integer):
      a = self.prng.integers(-64, 64, size=a_shape).astype(ab_type)
      b = self.prng.integers(-64, 64, size=b_shape).astype(ab_type)
    else:
      a = self.prng.uniform(-1, 1, a_shape).astype(ab_type)
      b = self.prng.uniform(-1, 1, b_shape).astype(ab_type)

    a_logical = np.zeros_like(a, shape=(m, k // 4, 4))
    np.put_along_axis(a_logical, x_sparse, a.reshape(x_sparse.shape), axis=-1)
    a_logical = a_logical.reshape(m, k)

    ref = np.matmul(a_logical.astype(acc_type), b.astype(acc_type))
    atol, rtol = (0, 0) if ab_type == jnp.int8 else (2e-5, 1e-7)

    self.assertArraysAllClose(
        kernel(a, b, x_gpu_sparse),
        ref,
        atol=atol,
        rtol=rtol,
    )

  @parameterized.product(
      m=(256,),
      n=(128, 192),
      ab_type=(jnp.float16, jnp.bfloat16, jnp.int8, jnp.float8_e4m3fn),
  )
  def test_tcgen05_mma_sparse_collective(self, m, n, ab_type):
    if ab_type == jnp.int8:
      self.skip_unless_tcgen05_int8()

    acc_type = jnp.float32 if jnp.issubdtype(ab_type, jnp.floating) else jnp.int32
    sparse_meta_dtype = jnp.uint2

    k = 256
    a_shape = (m, k // 2)
    a_block_shape = (m // 2, k // 2)
    b_shape = (k, n)
    b_block_shape = (k, n // 2)
    meta_shape = (m // 128, k // 128, 128, 64)
    meta_block_shape = (m // 256, k // 128, 128, 64)
    acc_shape = (m, n)
    acc_block_shape = (m // 2, n)

    def cp(src, dst, barrier, shape, dtype, indices):
      transfer_size_bits = dtypes.itemsize_bits(dtype) * math.prod(shape)
      assert transfer_size_bits % 8 == 0, f"{transfer_size_bits=} not byte aligned"
      transfer_size_bytes = transfer_size_bits // 8
      barrier.arrive_expect_tx(transfer_size_bytes)
      mgpu_dialect.async_load(
          source=src,
          destination=dst,
          barrier=barrier.as_barrier_memref(),
          indices=indices,
          slice_lengths=shape,
          collective=ir.ArrayAttr.get([]),
      )
      barrier.wait()

    def matmul(ctx, a_gmem, b_gmem, meta_gmem, result_gmem, scratch):
      (
          a_smem,
          b_smem,
          meta_smem,
          tma_barrier,
          mma_barrier,
          cluster_barrier,
          acc_tmem,
          meta_tmem,
      ) = scratch

      i32_type = ir.IntegerType.get_signless(32)
      zero_i32 = arith.constant(i32_type, 0)

      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      block_id_i32 = arith.index_cast(i32_type, block_id)

      m_index = arith.muli(block_id, arith.constant(ir.IndexType.get(), m // 2))
      m_index_i32 = arith.muli(block_id_i32, arith.constant(i32_type, m // 2))
      n_index_i32 = arith.muli(block_id_i32, arith.constant(i32_type, n // 2))

      # GMEM -> SMEM
      cp(b_gmem, b_smem, tma_barrier, b_block_shape, ab_type, indices=[zero_i32, n_index_i32])
      cp(a_gmem, a_smem, tma_barrier, a_block_shape, ab_type, indices=[m_index_i32, zero_i32])
      cp(meta_gmem, meta_smem, tma_barrier, meta_block_shape, sparse_meta_dtype, indices=[block_id_i32, zero_i32, zero_i32, zero_i32])

      cluster_barrier.arrive(orders_tensor_core=True)
      cluster_barrier.wait(orders_tensor_core=True)

      is_first_block = arith.cmpi(
          arith.CmpIPredicate.eq, block_id, c(0, ir.IndexType.get())
      )

      with when(is_first_block):
        mgpu_dialect.async_store_sparse_metadata_smem_to_tmem(meta_smem, meta_tmem, collective=True)
        mgpu_dialect.tcgen05_mma(
            accumulator=acc_tmem,
            a=a_smem,
            b=b_smem,
            a_sparse_metadata=meta_tmem,
            accumulate=arith.constant(ir.IntegerType.get_signless(1), False),
            collective=True,
        )
        mgpu_dialect.tcgen05_commit_arrive(
            mma_barrier.as_barrier_memref(), collective=True
        )

      mma_barrier.wait(orders_tensor_core=True)

      # TMEM -> Registers -> GMEM
      r_out = mgpu_dialect.async_load_tmem(acc_tmem)
      sliced_result_gmem = memref_slice(result_gmem, ds(m_index, m // 2))
      mgpu_dialect.vector_store(r_out, sliced_result_gmem, optimized=False)

    scratch_shape = [
        jax.ShapeDtypeStruct(a_block_shape, ab_type),
        jax.ShapeDtypeStruct(b_block_shape, ab_type),
        jax.ShapeDtypeStruct(meta_block_shape, sparse_meta_dtype),
        mgpu.Barrier(1),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.ClusterBarrier(collective_dims=(gpu.Dimension.x,)),
        mgpu.TMEM(acc_block_shape, acc_type, collective=True),
        mgpu.TMEM((m // 2, k // 2), sparse_meta_dtype, collective=True),
    ]

    kernel = mgpu.as_gpu_kernel(
        matmul,
        grid=(2, 1, 1),
        cluster=(2, 1, 1),
        block=(128, 1, 1),
        in_shape=(
            jax.ShapeDtypeStruct(a_shape, ab_type),
            jax.ShapeDtypeStruct(b_shape, ab_type),
            jax.ShapeDtypeStruct(meta_shape, sparse_meta_dtype),
        ),
        out_shape=jax.ShapeDtypeStruct(acc_shape, acc_type),
        smem_scratch_shape=scratch_shape,
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    index_pairs = np.asarray(np.meshgrid(range(4), range(4))).T.reshape(-1, 2)
    valid_pairs = index_pairs[index_pairs[:, 0] < index_pairs[:, 1]]
    x_pairs = jax.random.randint(
        jax.random.key(1234), (m, k // 4), 0, 6, dtype=jnp.uint8
    )
    x_sparse = valid_pairs[x_pairs]

    def format_sparse_meta(meta):
      mn, k, _ = meta.shape
      k *= 2
      if np.dtype(ab_type).itemsize == 1:
        meta_tiled = meta.reshape(mn // 128, 128, k // 64, 64).transpose(0, 2, 1, 3)
      else:
        meta_tiled = meta.reshape(mn // 128, 8, 2, 8, k // 64, 4, 2, 8).transpose(0, 4, 1, 6, 3, 5, 2, 7)
      return meta_tiled.reshape(mn // 128, k // 64, 128, 64).astype(sparse_meta_dtype)

    x_gpu_sparse = format_sparse_meta(x_sparse)
    if jnp.issubdtype(ab_type, jnp.integer):
      a = self.prng.integers(-64, 64, size=a_shape).astype(ab_type)
      b = self.prng.integers(-64, 64, size=b_shape).astype(ab_type)
    else:
      a = self.prng.uniform(-1, 1, a_shape).astype(ab_type)
      b = self.prng.uniform(-1, 1, b_shape).astype(ab_type)

    a_logical = np.zeros_like(a, shape=(m, k // 4, 4))
    np.put_along_axis(a_logical, x_sparse, a.reshape(x_sparse.shape), axis=-1)
    a_logical = a_logical.reshape(m, k)

    ref = np.matmul(a_logical.astype(acc_type), b.astype(acc_type))
    atol, rtol = (0, 0) if ab_type == jnp.int8 else (2e-5, 1e-7)

    self.assertArraysAllClose(
        kernel(a, b, x_gpu_sparse),
        ref,
        atol=atol,
        rtol=rtol,
    )

  @parameterized.product(
      m=(128, 256),
      n=(128, 256),
      # TODO(allanrenucci): Add 32-byte swizzle once implemented.
      swizzle=(64, 128),
      ab_type=(jnp.float16, jnp.bfloat16),
      acc_type=(jnp.float16, jnp.float32),
      a_in_tmem=(False, True),
  )
  def test_tcgen05_collective_mma(self, m, n, swizzle, ab_type, acc_type, a_in_tmem):
    if acc_type == jnp.float16 and ab_type != jnp.float16:
      self.skipTest("Only f16 input is supported for f16 output.")
    if a_in_tmem and m != 256:
      self.skipTest("Only M=256 is supported for MMA with A in TMEM.")

    swizzle_elems = swizzle // np.dtype(ab_type).itemsize
    groups_k = 2
    k = swizzle_elems * groups_k
    a_shape = (m, k)
    a_block_shape = (m // 2, k)
    a_packing = 4 // np.dtype(ab_type).itemsize
    b_shape = (k, n)
    b_block_shape = (k, n // 2)
    bytes_a = np.dtype(ab_type).itemsize * math.prod(a_block_shape)
    bytes_b = np.dtype(ab_type).itemsize * math.prod(b_block_shape)
    acc_shape = (m, n)
    acc_block_shape = (m // 2, n)

    def matmul(ctx, a_gmem, b_gmem, result_gmem, scratch):
      (a_smem, b_smem, tma_barrier, mma_barrier, cluster_barrier, acc_tmem, a_tmem) = scratch

      i32_type = ir.IntegerType.get_signless(32)
      zero_i32 = arith.constant(i32_type, 0)

      block_id = gpu.cluster_block_id(gpu.Dimension.x)
      block_id_i32 = arith.index_cast(i32_type, block_id)

      m_index = arith.muli(block_id, arith.constant(ir.IndexType.get(), m // 2))
      m_index_i32 = arith.muli(block_id_i32, arith.constant(i32_type, m // 2))
      n_index_i32 = arith.muli(block_id_i32, arith.constant(i32_type, n // 2))

      # GMEM -> SMEM
      tma_barrier.arrive_expect_tx(bytes_b)
      mgpu_dialect.async_load(
          source=b_gmem,
          destination=b_smem,
          barrier=tma_barrier.as_barrier_memref(),
          indices=[zero_i32, n_index_i32],
          slice_lengths=b_block_shape,
          collective=ir.ArrayAttr.get([]),
      )
      tma_barrier.wait()

      if a_in_tmem:
        # GMEM -> Registers -> TMEM
        sliced_a_gmem = memref_slice(a_gmem, ds(m_index, m // 2))
        reg = mgpu_dialect.vector_load(sliced_a_gmem, optimized=False)
        mgpu_dialect.async_store_tmem(reg, a_tmem)
        tcgen05.commit_tmem()
      else:
        # GMEM -> SMEM
        tma_barrier.arrive_expect_tx(bytes_a)
        mgpu_dialect.async_load(
            source=a_gmem,
            destination=a_smem,
            barrier=tma_barrier.as_barrier_memref(),
            indices=[m_index_i32, zero_i32],
            slice_lengths=a_block_shape,
            collective=ir.ArrayAttr.get([]),
        )
        tma_barrier.wait()

      # Make sure operands have been loaded on both blocks.
      cluster_barrier.arrive(orders_tensor_core=True)
      cluster_barrier.wait(orders_tensor_core=True)

      is_first_block = arith.cmpi(
          arith.CmpIPredicate.eq, block_id, c(0, ir.IndexType.get())
      )
      with when(is_first_block):
        mgpu_dialect.tcgen05_mma(
            accumulator=acc_tmem,
            a=a_tmem if a_in_tmem else a_smem,
            b=b_smem,
            accumulate=arith.constant(ir.IntegerType.get_signless(1), False),
            collective=True,
        )
        mgpu_dialect.tcgen05_commit_arrive(
            mma_barrier.as_barrier_memref(), collective=True
        )

      mma_barrier.wait(orders_tensor_core=True)

      # TMEM -> Registers -> GMEM
      r_out = mgpu_dialect.async_load_tmem(acc_tmem)
      sliced_result_gmem = memref_slice(result_gmem, ds(m_index, m // 2))
      mgpu_dialect.vector_store(r_out, sliced_result_gmem, optimized=False)

    # Required order: SMEM -> Barrier -> TMEM.
    scratch_shape = [
        jax.ShapeDtypeStruct(a_block_shape, ab_type) if not a_in_tmem else None,
        jax.ShapeDtypeStruct(b_block_shape, ab_type),
        mgpu.Barrier(1),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.ClusterBarrier(collective_dims=(gpu.Dimension.x,)),
        mgpu.TMEM(acc_block_shape, acc_type, collective=True),
        mgpu.TMEM(a_block_shape, ab_type, collective=True, packing=a_packing)
        if a_in_tmem
        else None,
    ]
    kernel = mgpu.as_gpu_kernel(
        matmul,
        grid=(2, 1, 1),
        cluster=(2, 1, 1),
        block=(128, 1, 1),
        in_shape=(
            jax.ShapeDtypeStruct(a_shape, ab_type),
            jax.ShapeDtypeStruct(b_shape, ab_type),
        ),
        out_shape=jax.ShapeDtypeStruct(acc_shape, acc_type),
        smem_scratch_shape=scratch_shape,
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    a = self.prng.uniform(-1, 1, a_shape).astype(ab_type)
    b = self.prng.uniform(-1, 1, b_shape).astype(ab_type)

    atol = 2e-2 if acc_type == jnp.float16 else 2e-5
    rtol = 8e-4 if acc_type == jnp.float16 else 1e-7
    self.assertArraysAllClose(
        kernel(a, b),
        np.matmul(a.astype(acc_type), b.astype(acc_type)),
        atol=atol,
        rtol=rtol,
    )

  @parameterized.product(
      m=(128,),
      n=(128, 256),
      scale_type=(jnp.float8_e8m0fnu, jnp.float8_e4m3fn),
      ab_type=(
          (jnp.float8_e5m2, jnp.float8_e5m2),
          (jnp.float8_e4m3fn, jnp.float8_e4m3fn),
          (jnp.float8_e5m2, jnp.float8_e4m3fn),
          (jnp.float8_e4m3fn, jnp.float8_e5m2),
          (jnp.float4_e2m1fn, jnp.float4_e2m1fn),
      ),
  )
  def test_simple_scaled_matmul(self, m, n, scale_type, ab_type):
    a_type, b_type = ab_type
    acc_type = jnp.float32

    block_size = 32
    if scale_type == jnp.float8_e4m3fn:
      if a_type != jnp.float4_e2m1fn or b_type != jnp.float4_e2m1fn:
        self.skipTest("Only float4_e2m1fn input is supported for e4m3fn scale.")
      block_size = 16

    k = 256
    a_shape = (m, k)
    b_shape = (n, k)
    k_scales = k // block_size
    scale_shape_a = (m // 128, k_scales // 4, 32, 16)
    scale_shape_b = (n // 128, k_scales // 4, 32, 16)
    acc_shape = (m, n)

    # TODO(olechwierowicz): Move these sync cp helpers in the top of the file.
    def cp(src, dst, barrier, shape, dtype):
      transfer_size_bits = dtypes.itemsize_bits(dtype) * math.prod(shape)
      assert transfer_size_bits % 8 == 0, f"{transfer_size_bits=} not byte aligned"
      transfer_size_bytes = transfer_size_bits // 8
      barrier.arrive_expect_tx(transfer_size_bytes)
      zero_i32 = arith.constant(ir.IntegerType.get_signless(32), 0)
      mgpu_dialect.async_load(
          source=src,
          destination=dst,
          barrier=barrier.as_barrier_memref(),
          indices=[zero_i32] * len(shape),
          slice_lengths=shape,
          collective=ir.ArrayAttr.get([]),
      )
      barrier.wait()

    def matmul(ctx, a_gmem, b_gmem, a_scale_gmem, b_scale_gmem, result_gmem, scratch):
      del ctx
      (
          a_smem,
          b_smem,
          a_scale_smem,
          b_scale_smem,
          tma_barrier,
          mma_barrier,
          acc_tmem,
          a_scale_tmem,
          b_scale_tmem,
      ) = scratch

      # GMEM -> SMEM
      cp(b_gmem, b_smem, tma_barrier, b_shape, b_type)
      cp(a_gmem, a_smem, tma_barrier, a_shape, a_type)
      cp(a_scale_gmem, a_scale_smem, tma_barrier, scale_shape_a, scale_type)
      cp(b_scale_gmem, b_scale_smem, tma_barrier, scale_shape_b, scale_type)

      mgpu_dialect.async_store_scales_smem_to_tmem(a_scale_smem, a_scale_tmem)
      mgpu_dialect.async_store_scales_smem_to_tmem(b_scale_smem, b_scale_tmem)
      mgpu_dialect.tcgen05_mma(
          accumulator=acc_tmem,
          a=a_smem,
          b=mgpu.memref_transpose(b_smem, (1,0)),
          a_scale=a_scale_tmem,
          b_scale=b_scale_tmem,
          accumulate=arith.constant(ir.IntegerType.get_signless(1), False),
      )
      mgpu_dialect.tcgen05_commit_arrive(mma_barrier.as_barrier_memref())
      mma_barrier.wait(orders_tensor_core=True)

      # TMEM -> Registers -> GMEM
      r_out = mgpu_dialect.async_load_tmem(acc_tmem)
      mgpu_dialect.vector_store(r_out, result_gmem, optimized=False)

    scratch_shape = [
        jax.ShapeDtypeStruct(a_shape, a_type),
        jax.ShapeDtypeStruct(b_shape, b_type),
        jax.ShapeDtypeStruct(scale_shape_a, scale_type),
        jax.ShapeDtypeStruct(scale_shape_b, scale_type),
        mgpu.Barrier(1),
        mgpu.Barrier(1, orders_tensor_core=True),
        mgpu.TMEM(acc_shape, acc_type),
        mgpu.TMEM((m, k_scales), scale_type),
        mgpu.TMEM((n, k_scales), scale_type),
    ]

    kernel = mgpu.as_gpu_kernel(
        matmul,
        grid=(1, 1, 1),
        cluster=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(
            jax.ShapeDtypeStruct(a_shape, a_type),
            jax.ShapeDtypeStruct(b_shape, b_type),
            jax.ShapeDtypeStruct(scale_shape_a, scale_type),
            jax.ShapeDtypeStruct(scale_shape_b, scale_type),
        ),
        out_shape=jax.ShapeDtypeStruct(acc_shape, acc_type),
        smem_scratch_shape=scratch_shape,
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    x = self.prng.uniform(-1, 1, (m, k)).astype(a_type)
    y = self.prng.uniform(-1, 1, (n, k)).astype(b_type)

    ksx, ksy = jax.random.split(jax.random.key(1234), 2)
    def scale_data(key, dim):
      if scale_type == jnp.float8_e4m3fn:
        return jnp.abs(
          jax.random.normal(key, (dim, k_scales), dtype=jnp.float32).astype(
             scale_type
          )
        )
      else:
        assert scale_type == jnp.float8_e8m0fnu
        return jax.lax.bitcast_convert_type(
            jax.random.randint(key, (dim, k_scales), 122, 132, dtype=jnp.uint8),
            scale_type
        )

    x_scale = scale_data(ksx, m)
    y_scale = scale_data(ksy, n)

    def format_scales(scales):
      mn, k = scales.shape
      assert mn % 128 == 0 and k % 4 == 0
      return (
          scales.reshape(mn // 128, 4, 32, k // 4, 4)
          .transpose(0, 3, 2, 1, 4)
          .reshape(mn // 128, k // 4, 32, 16)
      ).astype(scale_type)

    x_scale_fmt = format_scales(x_scale)
    y_scale_fmt = format_scales(y_scale)

    result = kernel(x, y, x_scale_fmt, y_scale_fmt)

    x_logical_scale = np.repeat(x_scale, block_size, axis=1).astype(np.float32)
    y_logical_scale = np.repeat(y_scale, block_size, axis=1).astype(np.float32)

    expected = np.dot(
        np.array(x, dtype=np.float32) * x_logical_scale,
        (np.array(y, dtype=np.float32) * y_logical_scale).T,
    )
    self.assertArraysAllClose(result, expected, atol=2e-4, rtol=5e-6)

  def test_slice_tmem(self):
    def tmem_type(ref: ir.Value):
      return ir.MemRefType.get(
          ref.type.shape, ref.type.element_type, memory_space=utils.tmem()
      )

    def body(ctx, x, y, x_out, y_out, tmem):
      del ctx
      x_tmem = mgpu_dialect.slice_tmem(tmem_type(x), tmem, offset=0)
      y_tmem = mgpu_dialect.slice_tmem(tmem_type(y), tmem, offset=128)
      x_layout = layouts.to_layout_attr(tcgen05.tmem_default_layout(packing=2))
      x_tmem = mgpu_dialect.tmem_layout_cast(x_tmem, x_layout)
      y_layout = layouts.to_layout_attr(tcgen05.tmem_default_layout(packing=1))
      y_tmem = mgpu_dialect.tmem_layout_cast(y_tmem, y_layout)

      # GMEM -> Registers -> TMEM
      x_reg = mgpu_dialect.vector_load(x, optimized=False)
      y_reg = mgpu_dialect.vector_load(y, optimized=False)
      mgpu_dialect.async_store_tmem(x_reg, x_tmem)
      mgpu_dialect.async_store_tmem(y_reg, y_tmem)
      tcgen05.commit_tmem()

      # TMEM -> Registers -> GMEM
      x_reg = mgpu_dialect.async_load_tmem(x_tmem)
      y_reg = mgpu_dialect.async_load_tmem(y_tmem)
      mgpu_dialect.vector_store(x_reg, x_out, optimized=False)
      mgpu_dialect.vector_store(y_reg, y_out, optimized=False)

    in_out_shapes = (
        jax.ShapeDtypeStruct((128, 128), jnp.bfloat16),
        jax.ShapeDtypeStruct((128, 64), jnp.int32),
    )
    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=in_out_shapes,
        out_shape=in_out_shapes,
        smem_scratch_shape=mgpu.TMEM((128, 512), jnp.int32),
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    x = self.prng.uniform(-100, 100, (128, 128)).astype(jnp.bfloat16)
    y = self.prng.uniform(-100, 100, (128, 64)).astype(jnp.int32)
    x_out, y_out = kernel(x, y)
    self.assertArraysEqual(x_out, x)
    self.assertArraysEqual(y_out, y)

  def test_tmem_subview(self):
    def body(ctx, in_ref, out_ref, tmem):
      del ctx
      # GMEM -> Registers -> TMEM
      in_reg = mgpu_dialect.vector_load(in_ref, optimized=False)
      slice_in = memref.subview(
          tmem, offsets=[0, 8], sizes=[128, 200], strides=[1, 1]
      )
      slice_in = memref.subview(
          slice_in, offsets=[0, 0], sizes=[128, 128], strides=[1, 1]
      )
      mgpu_dialect.async_store_tmem(in_reg, slice_in)
      tcgen05.commit_tmem()

      def dynamic_idx(idx: int) -> ir.Value:
        idx_type = ir.IndexType.get()
        return arith.constant(idx_type, idx)

      # TMEM -> Registers -> GMEM
      slice_out = memref.subview(
          tmem,
          offsets=[dynamic_idx(0), dynamic_idx(8)],
          sizes=[128, 128],
          strides=[1, 1],
      )
      out_reg = mgpu_dialect.async_load_tmem(slice_out)
      mgpu_dialect.vector_store(out_reg, out_ref, optimized=False)

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=jax.ShapeDtypeStruct((128, 128), jnp.float32),
        out_shape=jax.ShapeDtypeStruct((128, 128), jnp.float32),
        smem_scratch_shape=mgpu.TMEM((128, 256), jnp.float32),
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )
    x = self.prng.uniform(-100, 100, (128, 128)).astype(jnp.float32)
    self.assertArraysEqual(kernel(x), x)

  @parameterized.parameters(jnp.int32, jnp.int16, jnp.int8)
  def test_tma_gather(self, index_dtype):
    dtype = jnp.bfloat16
    src_shape = (128, 64)
    dst_shape = (32, 64)
    indices_shape = (32,)

    def body(ctx, src, indices, dst, smem):
      del ctx
      smem_ref, tma_barrier = smem

      # Load indices into registers
      indices_vec = mgpu_dialect.vector_load(indices, optimized=False)
      i32 = ir.IntegerType.get_signless(32)
      zero = arith.constant(i32, 0)

      slice_lengths = (32, 64)

      # Load + Gather
      expected_bytes = math.prod(slice_lengths) * np.dtype(dtype).itemsize
      tma_barrier.arrive_expect_tx(expected_bytes)
      mgpu_dialect.async_load(
          source=src,
          destination=smem_ref,
          barrier=tma_barrier.as_barrier_memref(),
          indices=[indices_vec, zero],
          slice_lengths=slice_lengths,
          collective=ir.ArrayAttr.get([]),
      )
      tma_barrier.wait()

      # Store
      mgpu_dialect.async_store(
          source=smem_ref,
          destination=dst,
          indices=[zero, zero],
          slice_lengths=slice_lengths,
      )
      nvvm.cp_async_bulk_wait_group(0)

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(
            jax.ShapeDtypeStruct(src_shape, dtype),
            jax.ShapeDtypeStruct(indices_shape, index_dtype),
        ),
        out_shape=jax.ShapeDtypeStruct(dst_shape, dtype),
        smem_scratch_shape=[
            jax.ShapeDtypeStruct((32, 64), dtype),
            mgpu.Barrier(1),
        ],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    src = self.prng.uniform(-1, 1, src_shape).astype(dtype)
    indices = jax.random.permutation(jax.random.key(0), 128)[:32].astype(index_dtype)
    result = kernel(src, indices)

    # Verification
    np.testing.assert_array_equal(result, src[indices.astype(jnp.int32)])

  @parameterized.parameters(
      (0, (1, 1, 1)),
      (1, (1, 1, 1)),
      (2, (1, 1, 1)),
      (0, (1, 1, 2)),
      (0, (1, 2, 1)),
      (0, (2, 1, 1)),
  )
  def test_cluster_launch_control(self, dim, cluster):
    # Let's say we have 148 SMs in our gpu. The cluster size is either 1 or 2.
    # We attempt to schedule exactly one more cluster than what fits on the GPU.
    # So either 149 clusters if the cluster size is 1, or 75 if the cluster size
    # is 2. Only one cluster will succeed in stealing the extra cluster's work,
    # and the others will fail. Therefore we test that there is exactly 1 stolen
    # cluster and the others fail and return -1.

    num_sms = jax.devices()[0].core_count
    cluster_size = math.prod(cluster)
    assert num_sms % cluster_size == 0

    grid = [1, 1, 1]
    grid[dim] = num_sms // cluster_size + 1
    grid = tuple(g * c for g, c in zip(grid, cluster))

    # Encode the last CTA ID as a single int xxxyyyzzz.
    last_cta_id = (
        (grid[0] - cluster[0]) * 1000000
        + (grid[1] - cluster[1]) * 1000
        + (grid[2] - cluster[2])
    )

    def single_wg_predicate():
      wg_indx_in_cta = utils.warpgroup_idx()
      wg_type = wg_indx_in_cta.type
      return arith.cmpi(arith.CmpIPredicate.eq, wg_indx_in_cta, c(0, wg_type))

    def single_cta_predicate():
      is_first_cta = c(1, ir.IntegerType.get_signless(1))
      index_type = ir.IndexType.get()
      for dim in gpu.Dimension:
        is_first_cta = arith.andi(
            is_first_cta,
            arith.cmpi(arith.CmpIPredicate.eq, utils.cluster_idx(dim), c(0, index_type)),
        )
      return is_first_cta

    def body(
        ctx: launch_context.LaunchContext,
        out_gmem_ref: ir.Value,
        smem: list[ir.Value],
    ):
      del ctx
      cancel_result_ref, barrier, _ = smem

      idx_type = ir.IndexType.get()
      i32 = ir.IntegerType.get_signless(32)

      single_wg = single_wg_predicate()
      if_op = scf.IfOp(single_wg, has_else=True)
      with ir.InsertionPoint(if_op.then_block):
        barrier.arrive_expect_tx(16)
        scf.YieldOp([])
      with ir.InsertionPoint(if_op.else_block):
        barrier.arrive_expect_tx(0)
        scf.YieldOp([])

      single_cta_and_wg = arith.andi(single_cta_predicate(), single_wg)
      mgpu_dialect.try_cluster_cancel(
          cancel_result_ref,
          barrier.as_barrier_memref(),
          predicate=single_cta_and_wg)
      barrier.wait()

      x, y, z, success = mgpu_dialect.query_cluster_cancel(cancel_result_ref)
      # Encode x, y, and z into one int xxxyyyzzz in order to compare with the
      # expected result.
      cta_id = z
      cta_id = arith.addi(cta_id, arith.muli(c(1000, i32), y))
      cta_id = arith.addi(cta_id, arith.muli(c(1000000, i32), x))

      minus_one = arith.constant(i32, -1)
      value = arith.select(success, cta_id, minus_one)

      idx_i32 = arith.index_cast(i32, gpu.block_id(gpu.Dimension.x))
      idy_i32 = arith.index_cast(i32, gpu.block_id(gpu.Dimension.y))
      idz_i32 = arith.index_cast(i32, gpu.block_id(gpu.Dimension.z))

      grid_y = c(grid[1], i32)
      grid_z = c(grid[2], i32)

      idx_yz = arith.addi(arith.muli(idy_i32, grid_z), idz_i32)
      grid_idx = arith.addi(arith.muli(idx_i32, arith.muli(grid_y, grid_z)), idx_yz)
      grid_idx_index = arith.index_cast(idx_type, grid_idx)

      memref.store(value, out_gmem_ref, [grid_idx_index])

    kernel = mgpu.as_gpu_kernel(
        body,
        grid=grid,
        cluster=cluster,
        block=(256, 1, 1),
        in_shape=(),
        out_shape=jax.ShapeDtypeStruct((num_sms,), jnp.int32),
        smem_scratch_shape=[
            jax.ShapeDtypeStruct((16,), jnp.int8),
            mgpu.Barrier(2),
            # Requesting SMEM close to the 228kb limit to ensure that each SM
            # only schedules 1 block.
            jax.ShapeDtypeStruct((220 * 1024,), jnp.int8),
        ],
        thread_semantics=mgpu.LoweringSemantics.Warpgroup,
    )

    result = np.sort(kernel())
    expected = np.array([-1] * (num_sms - cluster_size) + [last_cta_id] * cluster_size)
    np.testing.assert_equal(result, expected)


class UtilsTest(TestCase):
  @parameterized.parameters(
      (1,),
      (slice(2), slice(3),),
      (slice(1), slice(1, 3)),
      *([(utils.DynamicSlice(0, 2),)] if HAS_MOSAIC_GPU else []),
  )
  def test_parse_indices(self, *indices):
    # We are simply making sure this does not raise.
    _, _, _ = utils.parse_indices(indices, (2, 3, 4))

  @parameterized.parameters(
      (42,),
      (slice(42),),
      (slice(0, 42),),
      *([(utils.DynamicSlice(0, 4),)] if HAS_MOSAIC_GPU else []),
  )
  def test_parse_indices_oob(self, indices):
    with self.assertRaisesRegex(IndexError, "out of bounds"):
      utils.parse_indices(indices, (2, 3, 4))

  @parameterized.parameters(
      (-1,),
      (-42,),
      (slice(-1, 0),),
      (slice(-42, -1),),
  )
  def test_parse_indices_negative(self, indices):
    with self.assertRaisesRegex(NotImplementedError, "has negative bounds"):
      utils.parse_indices(indices, (2, 3, 4))

  def test_parse_indices_dynamic_slice(self):
    i = c(0, ir.IndexType.get())
    for s in (slice(i), slice(0, i), slice(0, 2, i)):
      with self.subTest(s=s):
        with self.assertRaisesRegex(NotImplementedError, "has dynamic bounds"):
          utils.parse_indices(s, (2, 3, 4))

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_assert(self):
    if cf is None:
      self.skipTest("``cf`` is not available")

    def kernel(ctx: mgpu.LaunchContext, x_ref, out, scratch) -> None:
      del ctx, out  # Unused.
      # TODO(b/408271232): Use a False condition once the bug is fixed.
      x = mgpu.FragmentedArray.load_strided(x_ref)
      cond = x.reduce("add", 0, *scratch) != 42.0
      cf.assert_(cond.registers.item(), "OOOPS")

    f = mgpu.as_gpu_kernel(
        kernel,
        grid=(1, 1, 1),
        block=(128, 1, 1),
        in_shape=(jax.ShapeDtypeStruct((128,), jnp.float32),),
        out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
        smem_scratch_shape=(jax.ShapeDtypeStruct((4,), jnp.float32),),
    )

    with jtu.set_env(MOSAIC_GPU_DUMP_SASS="1"), self.capture_stdout() as sass:
      jax.block_until_ready(f(jnp.ones((128,), jnp.float32)))

    # SASS doesn't seem to include the assertion message, so we are just
    # checking that __assertfail appears in the symbol table for the kernel.
    self.assertIn("__assertfail", sass())


class EndToEndTest(TestCase):

  def test_kernel_arguments(self):
    dtype = jnp.float32
    swizzle = 128
    bw = bitwidth(dtype_to_ir_type(dtype))
    shape = (128, 8 * swizzle // bw)
    i1 = ir.IntegerType.get_signless(1)
    def kernel(ctx, src, dst, smem):
      tmp, barrier = smem
      ctx.async_copy(src_ref=src, dst_ref=tmp, swizzle=swizzle, barrier=barrier)
      barrier.wait_parity(c(0, i1))
      copy(tmp, dst, swizzle=swizzle)
    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    smem = (x, mgpu.Barrier(1))
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)(x)
      np.testing.assert_array_equal(y, x)
    self.assertEqual(ptx().count("\t.param"), 2)


class SerializationTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    if jtu.test_device_matches(["rocm"]):
      self.skipTest("Mosaic GPU is not supported on ROCm.")

  def test_pass_is_registered(self):
    ctx = mlir.make_ir_context()
    ctx.allow_unregistered_dialects = True
    with ir.Location.unknown(ctx):
      module = ir.Module.create()
      pipeline = passmanager.PassManager.parse(
          "builtin.module(mosaic_gpu-serde{serialize=true})",
          ctx,
      )
      pipeline.run(module.operation)


class ApiTest(TestCase):

  def test_inout(self):
    def kernel(ctx, src, inout, dst, smem):
      val = memref.load(inout, [])
      gpu.barrier()
      new_val = arith.constant(ir.IntegerType.get_signless(32), 42)
      memref.store(new_val, inout, [])
      x = mgpu.FragmentedArray.load_strided(src, is_signed=True)
      (x + val).store_untiled(dst)
    x = jnp.arange(128, dtype=jnp.int32)
    y = jnp.asarray(2.0, dtype=jnp.int32)
    kernel = mgpu.as_gpu_kernel(
        kernel, (1, 1, 1), (128, 1, 1), x, x, (), inout_shape=y,
    )
    xo, yo = kernel(x, y)
    np.testing.assert_array_equal(xo, x + 2.0)
    np.testing.assert_array_equal(yo, jnp.asarray(42, dtype=jnp.int32))

  def test_serialize_uses_bytecode_format(self):
    def kernel(ctx, src, dst, smem):
      del ctx, smem
      x = mgpu.FragmentedArray.load_strided(src, is_signed=True)
      (x + 1).store_untiled(dst)
    x = jnp.arange(128, dtype=jnp.int32)
    with self.subTest("bytecode"):
      f = mgpu.as_gpu_kernel(
          kernel, (1, 1, 1), (128, 1, 1), x, x, (),
      )
      bytecode_stablehlo = jax.jit(f).lower(x).as_text()
      self.assertIn("module = \"ML\\EFR", bytecode_stablehlo)

if hp is not None:
  @hps.composite
  def tiled_layouts(
      draw, initial_tile, vector_transfer: bool = False
  ) -> fa.TiledLayout:
    assert all(t.bit_count() == 1 for t in initial_tile)
    assert math.prod(initial_tile) >= 128
    tiles = [initial_tile]
    dim_offset = len(initial_tile)
    if draw(hps.booleans()):
      warp_dims = [fa.Replicated(2) if draw(hps.booleans()) else None for _ in range(2)]
    else:
      warp_dims = [fa.Replicated(4) if draw(hps.booleans()) else None]
    for i, dim in enumerate(warp_dims):
      if isinstance(dim, fa.Replicated):
        continue
      dim_size = 4 // len(warp_dims)
      warp_dim = draw(
          hps.sampled_from(
              [i for i, t in enumerate(tiles[-1]) if t % dim_size == 0]
          )
      )
      warp_tile = list(tiles[-1])
      warp_tile[warp_dim] //= dim_size
      warp_dims[i] = dim_offset + warp_dim
      tiles.append(warp_tile)
      dim_offset += len(warp_tile)
    lane_dims = [fa.Replicated(2) if draw(hps.booleans()) else None for _ in range(5)]
    for i, dim in enumerate(lane_dims):
      if isinstance(dim, fa.Replicated):
        continue
      lane_dim = draw(hps.sampled_from(
          [i for i, t in enumerate(tiles[-1]) if t % 2 == 0]
      ))
      lane_tile = list(tiles[-1])
      lane_tile[lane_dim] //= 2
      lane_dims[i] = dim_offset + lane_dim
      tiles.append(lane_tile)
      dim_offset += len(lane_tile)
    # Permute lane dims so that they don't always partition the data in order.
    lane_dims = draw(hps.permutations(lane_dims))
    if vector_transfer:
      min_vector_dim = len(tiles[-1]) - 1
    else:
      min_vector_dim = 0
    vector_dim = draw(hps.integers(min_vector_dim, len(tiles[-1]) - 1))
    vector_size = 2 ** draw(
        hps.integers(0, tiles[-1][vector_dim].bit_length() - 1)
    )
    vector_tile = list(tiles[-1])
    assert vector_tile[vector_dim] % vector_size == 0
    vector_tile[vector_dim] //= vector_size
    tiles.append(vector_tile)
    dim_offset += len(vector_tile)
    vector_dim += dim_offset
    dim_offset += len(vector_tile)  # This is the remainder after tiling!

    warp_dims = tuple(
        d if isinstance(d, fa.Replicated) else d - dim_offset
        for d in warp_dims
    )
    lane_dims = tuple(
        d if isinstance(d, fa.Replicated) else d - dim_offset
        for d in lane_dims
    )
    vector_dim = vector_dim - dim_offset
    return fa.TiledLayout(
        tiling=fa.Tiling(tuple(map(tuple, tiles))),
        warp_dims=warp_dims,
        lane_dims=lane_dims,
        vector_dim=vector_dim,
        _check_canonical=False,
    ).canonicalize()

  @hps.composite
  def shape_and_tiled_layout(
      draw, vector_transfer: bool = False
  ) -> tuple[tuple[int, ...], fa.TiledLayout]:
    rank = draw(hps.integers(2, 3))
    initial_tile = tuple(
        draw(hps.sampled_from([1, 2, 4, 8, 16, 32, 64, 128]))
        for _ in range(rank)
    )
    hp.assume(128 <= math.prod(initial_tile) < 128 * 32)
    shape = tuple(t * draw(hps.integers(1, 5)) for t in initial_tile)
    hp.assume(math.prod(shape) <= 128 * 128)
    layout = draw(tiled_layouts(initial_tile, vector_transfer=vector_transfer))
    return shape, layout

  class HypothesisTest(TestCase, jtu.CudaArchSpecificTest):

    def test_reduce(self):
      @hps.composite
      def strategy(draw):
        shape, layout = draw(shape_and_tiled_layout(vector_transfer=True))
        rank = len(shape)
        reduced_dims = draw(hps.sets(hps.integers(0, rank - 1), min_size=1))
        dtype = draw(hps.sampled_from([jnp.int32, jnp.int16]))
        op = draw(hps.sampled_from(["add", "max"]))
        return shape, layout, tuple(reduced_dims), dtype, op

      warp_replicated_major = fa.TiledLayout(
          fa.Tiling(((2,), (1,))), (fa.Replicated(2,), -2), (fa.Replicated(32,),), -1
      )
      warp_replicated_minor = fa.TiledLayout(
          fa.Tiling(((2,), (1,))), (-2, fa.Replicated(2,)), (fa.Replicated(32,),), -1
      )
      warp_row_col_layout = fa.TiledLayout(
          fa.Tiling(((2, 2), (1,))), (-3, -2), (fa.Replicated(32,),), -1
      )
      even_lane_split_layout = fa.TiledLayout(
          fa.Tiling(((8,), (4,), (2,), (1,))),
          (fa.Replicated(4),),
          (-4, fa.Replicated(2), -3, fa.Replicated(2), -2),
          -1,
      )
      odd_lane_split_layout = fa.TiledLayout(
          fa.Tiling(((4,), (2,), (1,))),
          (fa.Replicated(4),),
          (fa.Replicated(2), -3, fa.Replicated(2), -2, fa.Replicated(2)),
          -1,
      )

      @hp.given(strategy())
      @hp.example(((16,), warp_replicated_major, (0,), jnp.int32, "add"))
      @hp.example(((16,), warp_replicated_minor, (0,), jnp.int32, "add"))
      @hp.example(((16, 16), warp_row_col_layout, (0,), jnp.int32, "add"))
      @hp.example(((16, 16), warp_row_col_layout, (1,), jnp.int32, "add"))
      @hp.example(((256,), even_lane_split_layout, (0,), jnp.float32, "max"))
      @hp.example(((256,), odd_lane_split_layout, (0,), jnp.int32, "max"))
      def run(args):
        shape, layout, reduced_dims, dtype, op = args
        out_shape = list(shape)
        for d in sorted(reduced_dims, reverse=True):
          del out_shape[d]
        def kernel(ctx, src, dst, scratch):
          del ctx
          arr = fa.FragmentedArray.load_untiled(src, layout=layout, optimized=False,
                                                is_signed=utils.is_signed(dtype))
          arr.reduce(op, reduced_dims, scratch).store_untiled(dst, optimized=False)
        if jnp.issubdtype(dtype, jnp.integer):
          x = jax.random.randint(jax.random.key(1234), shape, -1000, 1000, dtype)
        else:
          x = jax.random.normal(jax.random.key(1234), shape, dtype)
        out_type = jax.ShapeDtypeStruct(out_shape, dtype)
        scratch_type = jax.ShapeDtypeStruct((2048,), dtype)
        hp.assume(layout.vector_length <= 16)  # Otherwise we run out of scratch
        try:
          result = mgpu.as_gpu_kernel(
              kernel, (1, 1, 1), (128, 1, 1), x, out_type, scratch_type
          )(x)
        except NotImplementedError:
          hp.assume(False)
          return
        if op == "max":
          ref = x.max(reduced_dims)
        else:
          ref = x.sum(reduced_dims, dtype=dtype)
        np.testing.assert_array_equal(result, ref)
      run()

    def test_slice(self):
      i32 = ir.IntegerType.get_signless(32)
      index = ir.IndexType.get()

      @hps.composite
      def strategy(draw):
        shape, layout = draw(shape_and_tiled_layout(vector_transfer=True))
        tiling = layout.base_tile_shape
        tiled_shape = mgpu.tile_shape(shape, tiling)[:len(shape)]
        def draw_slice(size, tile):
          start = draw(hps.integers(0, size - 1))
          length = draw(hps.integers(1, size - start))
          return slice(start * tile, (start + length) * tile)
        slices = tuple(map(draw_slice, tiled_shape, tiling))
        return shape, layout, slices

      basic_slices = (slice(128, 256), slice(16, 16 + 32))
      @hp.given(strategy())
      @hp.example(((256, 256), fa.WGMMA_LAYOUT, basic_slices))
      @hp.example(((256, 256), tcgen05.LAYOUT, basic_slices))
      @hp.example(((256, 256), tcgen05.TMEM_NATIVE_LAYOUT, basic_slices))
      def run(args):
        shape, layout, slices = args
        def kernel(ctx, dst, _):
          def linear_index(*idxs):
            total = arith.constant(index, 0)
            stride = 1
            for i, size in zip(idxs[::-1], shape[::-1]):
              total = arith.addi(total, arith.muli(i, c(stride, index)))
              stride *= size
            return arith.index_cast(i32, total)
          x = mgpu.FragmentedArray.build(
              shape, layout, linear_index, is_signed=True
          )
          x[slices].store_untiled(dst, optimized=False)

        slice_shape = tuple(len(range(size)[s]) for s, size in zip(slices, shape))
        out_shape = jax.ShapeDtypeStruct(shape=slice_shape, dtype=jnp.int32)
        result = mgpu.as_gpu_kernel(
            kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
        )()
        iota = np.arange(np.prod(shape), dtype=jnp.int32).reshape(*shape)
        np.testing.assert_array_equal(result, iota[slices])
      run()

    def test_broadcast(self):
      @hps.composite
      def strategy(draw):
        shape, layout = draw(shape_and_tiled_layout(vector_transfer=True))
        rank = len(shape)
        broadcast_dims = draw(
            hps.sets(hps.integers(0, rank - 1), min_size=1, max_size=rank - 1)
        )
        dtype = draw(hps.sampled_from([jnp.float32, jnp.bfloat16]))
        return shape, layout, tuple(broadcast_dims), dtype

      @hp.given(strategy())
      def run(args):
        out_shape, out_layout, broadcast_dims, dtype = args
        in_shape = list(out_shape)
        for d in sorted(broadcast_dims, reverse=True):
          del in_shape[d]
        in_layout = out_layout.reduce(broadcast_dims)
        dims = tuple(d for d in range(len(out_shape)) if d not in broadcast_dims)
        def kernel(ctx, src, dst, scratch):
          del ctx, scratch  # Unused.
          arr = fa.FragmentedArray.load_untiled(src, layout=in_layout, optimized=False)
          arr.broadcast_in_dim(out_shape, dims, out_layout).store_untiled(dst, optimized=False)
        x = jax.random.normal(jax.random.key(1234), in_shape, dtype)
        out_type = jax.ShapeDtypeStruct(out_shape, dtype)
        try:
          result = mgpu.as_gpu_kernel(
              kernel, (1, 1, 1), (128, 1, 1), x, out_type, ()
          )(x)
        except NotImplementedError:
          hp.assume(False)
          return
        np.testing.assert_array_equal(result, jax.lax.broadcast_in_dim(x, out_shape, dims))
      run()

    @hp.given(hps.data())
    def test_canonicalize_trivial_dims(self, data):
      layout = data.draw(tiled_layouts((128, 1)))
      trivial_dims = [
          i
          for i, d in fa.enumerate_negative(layout.tiled_tiling_shape)
          if d == 1 and i != layout.vector_dim
      ]
      if not trivial_dims:
        hp.assume(False)
      # That should not happen in canonical layouts.
      self.assertNoCommonElements(trivial_dims, layout.partitioned_warp_dims)
      self.assertNoCommonElements(trivial_dims, layout.partitioned_lane_dims)
      # vector_dim can be trivial.
      canonical_layout = layout
      use_trivial_dim = data.draw(
          hps.lists(hps.booleans(), min_size=len(trivial_dims), max_size=len(trivial_dims))
      )
      hp.assume(any(use_trivial_dim))
      for d, use in zip(trivial_dims, use_trivial_dim):
        if not use:
          continue
        if data.draw(hps.booleans()):  # Should we put it in warp or lane dims?
          new_warp_dims = list(layout.warp_dims)
          position = data.draw(hps.integers(0, len(layout.warp_dims)))
          new_warp_dims.insert(position, d)
          layout = fa.TiledLayout(
              layout.tiling,
              warp_dims=tuple(new_warp_dims),
              lane_dims=layout.lane_dims,
              vector_dim=layout.vector_dim,
              _check_canonical=False,
          )
        else:
          new_lane_dims = list(layout.lane_dims)
          position = data.draw(hps.integers(0, len(layout.lane_dims)))
          new_lane_dims.insert(position, d)
          layout = fa.TiledLayout(
              layout.tiling,
              warp_dims=layout.warp_dims,
              lane_dims=tuple(new_lane_dims),
              vector_dim=layout.vector_dim,
              _check_canonical=False,
          )
      self.assertNotEqual(layout, canonical_layout)
      self.assertEqual(layout.canonicalize(), canonical_layout)

    def test_copy_tiled(self):
      @hps.composite
      def strategy(draw):
        swizzle = draw(hps.sampled_from([16, 32, 64, 128]))
        dtype = draw(hps.sampled_from([jnp.int32, jnp.int16, jnp.int8]))
        tiling = (8, swizzle // jnp.dtype(dtype).itemsize)
        shape = [draw(hps.integers(1, 6)) for t in tiling]
        while math.prod(shape) % 4:
          shape[draw(hps.booleans())] *= 2
        shape = [s * t for s, t in zip(shape, tiling)]
        to_smem = draw(hps.booleans())
        return shape, dtype, swizzle, to_smem

      @hp.given(strategy())
      @hp.example(((48, 64), jnp.int32, 16, False))
      @hp.example(((48, 64), jnp.int32, 32, False))
      @hp.example(((48, 64), jnp.int32, 64, False))
      @hp.example(((48, 64), jnp.int32, 128, False))
      @hp.example(((64, 4), jnp.int32, 16, False))
      def run(args):
        shape, dtype, swizzle, to_smem = args
        tiling = (8, 8 * swizzle // jnp.iinfo(dtype).bits)
        def kernel(ctx, src, dst, scratch):
          smem, barrier = scratch
          if to_smem:
            mgpu.copy_tiled(src, smem, swizzle=swizzle)
            mgpu.commit_shared()
            ctx.async_copy(
                src_ref=smem,
                dst_ref=dst,
                gmem_transform=mgpu.TileTransform(tiling),
                swizzle=swizzle,
            )
            ctx.await_async_copy(0)
          else:
            ctx.async_copy(
                src_ref=src,
                dst_ref=smem,
                gmem_transform=mgpu.TileTransform(tiling),
                swizzle=swizzle,
                barrier=barrier,
            )
            barrier.wait()
            mgpu.copy_tiled(smem, dst, swizzle=swizzle)

        x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
        scratch_shape = [
            jax.ShapeDtypeStruct(mgpu.tile_shape(shape, tiling), dtype),
            mgpu.Barrier(1),
        ]
        y = mgpu.as_gpu_kernel(
            kernel, (1, 1, 1), (128, 1, 1), x, x, scratch_shape
        )(x)
        np.testing.assert_array_equal(y, x)
      run()

    def test_dialect_vector_load_store(self):
      @hps.composite
      def strategy(draw):
        shape, layout = draw(shape_and_tiled_layout(vector_transfer=True))
        return shape, layout

      @hp.given(strategy())
      @hp.example(((128, 128), fa.WGMMA_LAYOUT))
      @hp.example(((128, 128), fa.TCGEN05_LAYOUT))
      @hp.example(((128, 128), fa.TMEM_NATIVE_LAYOUT))
      def run(args):
        shape, layout = args
        # Attempt to avoid spills by keeping the shape small enough.
        hp.assume(math.prod(shape) <= 128 * 128)
        dtype = jnp.float32
        layout_attr = layouts.to_layout_attr(layout)

        def body(ctx, input, result, smem):
          del ctx
          # GMEM -> Registers
          reg = mgpu_dialect.vector_load(input, optimized=False)
          reg = mgpu_dialect.layout_cast(reg, layout_attr)
          # Registers -> SMEM
          # Use a non-optimized store to ensure that we don't fail due to
          # unoptimizable stores.
          mgpu_dialect.vector_store(reg, smem, optimized=False)
          # SMEM -> Registers
          # Use a non-optimized load to ensure that we don't fail due to
          # unoptimizable loads.
          reg = mgpu_dialect.vector_load(smem, optimized=False)
          reg = mgpu_dialect.layout_cast(reg, layout_attr)
          # Registers -> GMEM
          mgpu_dialect.vector_store(reg, result, optimized=False)

        jax_shape = jax.ShapeDtypeStruct(shape, dtype)
        kernel = mgpu.as_gpu_kernel(
            body,
            grid=(1, 1, 1),
            block=(128, 1, 1),
            in_shape=jax_shape,
            out_shape=jax_shape,
            smem_scratch_shape=jax_shape,
            thread_semantics=mgpu.LoweringSemantics.Warpgroup,
        )

        input = self.prng.uniform(-1, 1, shape).astype(dtype)
        np.testing.assert_array_equal(kernel(input), input)

      run()

    def test_tmem_column_slicing(self):
      self.skip_unless_tcgen05()

      @hps.composite
      def strategy(draw):
        dtype = draw(
            hps.sampled_from([jnp.float32, jnp.float16, jnp.float8_e5m2])
        )
        initial_tile = (
            draw(hps.sampled_from([32, 64, 128, 256])),
            draw(hps.sampled_from([1, 2, 4, 8, 16])),
        )
        hp.assume(math.prod(initial_tile) >= 128)
        tiled_layout = draw(tiled_layouts(initial_tile, vector_transfer=True))
        bitwidth = jax.dtypes.itemsize_bits(dtype)
        # Ensure we can copy from/to registers.
        hp.assume(tiled_layout.vector_length * bitwidth == 32)
        tmem_layout = tcgen05.TMEMLayout(
            tiled_layout.tiling,
            tiled_layout.warp_dims,
            tiled_layout.lane_dims,
            tiled_layout.vector_dim,
        )
        tile_shape = tmem_layout.base_tile_shape
        shape = (tile_shape[0], 4 * tile_shape[1])
        hp.assume(tmem_layout.cols_in_shape(shape, bitwidth) <= 512)
        offset_steps = draw(hps.integers(0, 3))
        offset = tile_shape[1] * offset_steps
        slice_steps = draw(hps.integers(1, min(3, 4 - offset_steps)))
        slice_size = tile_shape[1] * slice_steps
        slicing = (slice(None), slice(offset, offset + slice_size))
        return dtype, tmem_layout, tiled_layout, shape, slicing

      @hp.given(strategy())
      def run(args):
        dtype, tmem_layout, reg_layout, shape, slicing = args

        def kernel(ctx, input, output, tmem):
          del ctx
          reg = fa.FragmentedArray.load_untiled(
              input, layout=reg_layout, optimized=False
          )
          tmem.store(reg)
          tcgen05.commit_tmem()
          tmem_slice = tmem.slice(*slicing)
          tmem_slice.load(reg_layout).store_untiled(output, optimized=False)

        x = self.prng.uniform(-1, 1, shape).astype(dtype)
        expected = x[*slicing]
        y = mgpu.as_gpu_kernel(
            kernel,
            (1, 1, 1),
            (128, 1, 1),
            x,
            expected,
            mgpu.TMEM(x.shape, x.dtype, layout=tmem_layout),
        )(x)
        np.testing.assert_array_equal(expected, y)

      run()


if __name__ == "__main__":
  absltest.main(argv=["python"], testLoader=jtu.JaxTestLoader())