# Copyright 2024 The JAX Authors.
#
# 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
#
#     https://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.

from collections.abc import Sequence
import contextlib
import dataclasses
import functools
import itertools
import json
import math
import operator
import os
import re
import sys
import tempfile
import traceback
import types
from typing import ClassVar, TYPE_CHECKING
from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax import export
from jax import lax
from jax import sharding
from jax._src import config
from jax._src import core as jax_core
from jax._src import dtypes
from jax._src import test_util as jtu
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import arith as arith_dialect
from jax._src.lib.mlir.dialects import gpu as gpu_dialect
from jax._src.lib.mlir.dialects import memref as memref_dialect
from jax._src.pallas.mosaic.error_handling import VerificationError
from jax._src.pallas.mosaic_gpu import core as gpu_core
from jax._src.pallas.mosaic_gpu import lowering as mgpu_lowering
from jax._src.pallas.mosaic_gpu import pipeline as mgpu_pipeline
from jax._src.state import types as state_types
from jax.experimental import pallas as pl
import jax.experimental.mosaic.gpu as mgpu
from jax.experimental.pallas import mosaic_gpu as _plgpu
import jax.numpy as jnp
import numpy as np

try:
  from jax._src.lib import mosaic_gpu as mosaic_gpu_lib
except ImportError:
  mosaic_gpu_lib = None

jax.config.parse_flags_with_absl()


# We don't want the user to call plgpu.kernel directly, so we monkey patch
# it in `plgpu`.
def do_not_call_me_directly(*args, **kwargs):
  raise RuntimeError("Use self.kernel instead of plgpu.kernel.")

if TYPE_CHECKING:
  plgpu = _plgpu
else:
  # Clone the module locally because the functions may be called from other
  # modules.
  plgpu = types.ModuleType("_plgpu_local")
  plgpu.__dict__.update(_plgpu.__dict__)

  _kernel = _plgpu.kernel
  del _plgpu

  plgpu.kernel = do_not_call_me_directly


def _fori_loop(force_while: bool, lb, ub, body, init):
  if force_while:
    # using jnp.asarray make the matcher for while or scan to think
    # that the bounds are dynamic and forces the use of the while
    # primitive.
    lb, ub = jnp.asarray(lb), jnp.asarray(ub)
  return jax.lax.fori_loop(lb, ub, body, init)


def _sum_same_dtype(x):
  # TODO(slebedev): Remove this once ``FragmentedArray`` supports
  # ``reduce_sum`` for non-32-bit types.
  return jnp.sum(x, dtype=x.dtype)


def _get_linearized_cuda_grid_index():
  shape = ()
  layout = plgpu.Layout.WG_SPLAT

  @plgpu.inline_mgpu(
      arg_types=(),
      return_type=plgpu.ShapeDtypeStruct(shape, jnp.int32, layout=layout),
  )
  def fn(_):
    grid_x = gpu_dialect.grid_dim(gpu_dialect.Dimension.x)
    grid_y = gpu_dialect.grid_dim(gpu_dialect.Dimension.y)
    block_x = gpu_dialect.block_id(gpu_dialect.Dimension.x)
    block_y = gpu_dialect.block_id(gpu_dialect.Dimension.y)
    block_z = gpu_dialect.block_id(gpu_dialect.Dimension.z)

    grid_idx = arith_dialect.addi(
        block_x,
        arith_dialect.addi(
            arith_dialect.muli(block_y, grid_x),
            arith_dialect.muli(block_z, arith_dialect.muli(grid_x, grid_y)),
        ),
    )

    return mgpu.FragmentedArray.splat(
        arith_dialect.index_cast(ir.IntegerType.get_signless(32), grid_idx),
        shape=shape,
        layout=layout.to_mgpu(),
        is_signed=False
    )
  return fn()


def _array_splat(value, shape: tuple[int, ...]):
  """Same as `jnp.full(shape, value, jnp.float32)` but implemented using `inline_mgpu`.

  This is useful to prevent the result from being optimized away.
  """
  @plgpu.inline_mgpu(
      return_type=plgpu.ShapeDtypeStruct(
          shape, jnp.float32, layout=plgpu.Layout.WG_SPLAT(shape)
      ),
  )
  def fn(_):
    ir_value = mgpu.c(value, ir.F32Type.get())
    return mgpu.FragmentedArray.splat(ir_value, shape)
  return fn()


class PallasTestMetaclass(parameterized.TestGeneratorMetaclass):

  def __new__(mcs, *args, lowering_semantics=plgpu.LoweringSemantics.Lane):
    cls = super().__new__(mcs, *args)
    cls.LOWERING_SEMANTICS = lowering_semantics
    return cls


class MonkeyPatchTest(jtu.JaxTestCase):

  def test_calling_kernel_directly_raises(self):
    with self.assertRaises(RuntimeError):
      plgpu.kernel()


def run_on_sm80(method):
  """A marker that allows a test case to run on Ampere GPUs."""
  method._min_capability = "8.0"
  return method


class PallasTest(jtu.JaxTestCase, metaclass=PallasTestMetaclass):
  LOWERING_SEMANTICS: ClassVar[plgpu.LoweringSemantics]

  def setUp(
      self, *, artificial_shared_memory_limit=jtu._SMEM_SIZE_BOUND_FOR_TESTS
  ):
    if jtu.test_device_matches(["rocm"]):
      self.skipTest("Mosaic GPU is not supported on ROCm.")
    min_capability = getattr(
        getattr(type(self), self._testMethodName, None),
        "_min_capability",
        "9.0",
    )
    if not jtu.is_cuda_compute_capability_at_least(min_capability):
      self.skipTest(f"Only works on a GPU with capability >= sm{min_capability}")

    super().setUp()
    self.enter_context(mgpu.core.artificial_shared_memory_limit(artificial_shared_memory_limit))

  def is_wg_semantics(self):
    return self.LOWERING_SEMANTICS == plgpu.LoweringSemantics.Warpgroup

  def skip_if_wg_semantics(self):
    if self.is_wg_semantics():
      self.skipTest("Not supported under WG semantics")

  def kernel(self, *args, **kwargs):
    compiler_params = dataclasses.replace(
        kwargs.pop("compiler_params", plgpu.CompilerParams()),
        lowering_semantics=self.LOWERING_SEMANTICS,
    )
    return _kernel(*args, compiler_params=compiler_params, **kwargs)

  def pallas_call(
      self,
      fn,
      *,
      out_shape,
      grid=(),
      in_specs=(),
      out_specs=(),
      scratch_shapes=(),
      compiler_params=plgpu.CompilerParams(),
  ):
    compiler_params = dataclasses.replace(
        compiler_params,
        lowering_semantics=self.LOWERING_SEMANTICS,
    )
    from jax._src.pallas.mosaic_gpu import pallas_call
    return pallas_call.pallas_call(
        fn,
        out_shape=out_shape,
        grid=grid,
        in_specs=in_specs,
        out_specs=out_specs,
        scratch_shapes=scratch_shapes,
        compiler_params=compiler_params,
        kernel_fn=self.kernel,
    )

  @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()

  def default_transforms(
      self, *, swizzle: int = 128, dtype: jnp.dtype
  ) -> Sequence[plgpu.Transform]:
    swizzle_elems = 8 * swizzle // dtypes.itemsize_bits(dtype)
    return (
        plgpu.TilingTransform((8, swizzle_elems)),
        plgpu.SwizzleTransform(swizzle),
    )


class PallasSm90ATest(PallasTest, 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 PallasTCGen05Test(PallasTest, jtu.CudaArchSpecificTest):

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


class PallasCallTest(PallasTest, jtu.CudaArchSpecificTest):

  def test_memory_space_like(self):
    x = jax.ShapeDtypeStruct((2, 3), jnp.float32)
    ref = gpu_core.MemorySpace.GMEM.like(x)
    self.assertEqual(ref.inner_aval.shape, (2, 3))
    self.assertEqual(ref.inner_aval.dtype, jnp.float32)
    self.assertEqual(ref.memory_space, gpu_core.MemorySpace.GMEM)

  def test_custom_jvp_and_vjp(self):
    @jax.custom_jvp
    def my_jvp_fn(x):
      def body(smem_ref):
        smem_ref[...] = jax.nn.relu(x)
        return smem_ref[...]
      return pl.run_scoped(body, plgpu.SMEM(x.shape, x.dtype))

    @my_jvp_fn.defjvp
    def my_jvp_fn_jvp(primals, tangents):
      return my_jvp_fn(*primals), tangents[0]

    @jax.custom_vjp
    def my_vjp_fn(x):
      return x * 2.0

    my_vjp_fn.defvjp(lambda x: (my_vjp_fn(x), None), lambda _, g: (g * 2.0,))

    @self.kernel(out_type=jax.ShapeDtypeStruct((128,), jnp.float32))
    def kernel(x_ref, o_ref):
      o_ref[...] = my_vjp_fn(my_jvp_fn(x_ref[...]))

    x = jnp.linspace(-5.0, 5.0, 128, dtype=jnp.float32)
    np.testing.assert_allclose(kernel(x), jax.nn.relu(x) * 2.0)

  def test_barrier_ref_shape(self):
    # A single barrier is referred to by a scalar reference, so that uses of a
    # barrier array can be told apart from uses of a single barrier.
    self.assertEqual(plgpu.Barrier().get_ref_aval().shape, ())
    self.assertEqual(plgpu.Barrier(num_barriers=1).get_ref_aval().shape, (1,))
    self.assertEqual(plgpu.Barrier(num_barriers=(2, 3)).get_ref_aval().shape,
                     (2, 3))

  def test_io_aliasing(self):
    @jax.jit
    def f():
      x = jnp.arange(128, dtype=jnp.float32)
      x_ref = jax.new_ref(x)
      @self.kernel()
      def kernel():
        x_ref[...] += 1.0
      kernel()
      return jax.freeze(x_ref)
    self.assertIn("output_to_operand_aliasing={{}: (0, {})}",
                  f.lower().as_text("hlo"))
    np.testing.assert_array_equal(f(), jnp.arange(128, dtype=jnp.float32) + 1)

  def test_multiple_of(self):
    shape = (128, 64)
    out_shape = (64, 64)

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(out_shape, jnp.float32),
        scratch_types=[
            plgpu.SMEM(
                out_shape,
                jnp.float32,
                transforms=self.default_transforms(dtype=jnp.float32),
            ),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_gmem, index_gmem, o_gmem, smem_ref, barrier):
      dyn_index = index_gmem[0]
      dyn_index = pl.multiple_of(dyn_index, (7, 8))
      gmem_slice = pl.ds(dyn_index, 64)
      plgpu.copy_gmem_to_smem(x_gmem.at[gmem_slice, :], smem_ref, barrier)
      plgpu.barrier_wait(barrier)
      plgpu.store(o_gmem, smem_ref[...], optimized=False)

    x = jnp.arange(math.prod(shape), dtype=jnp.float32).reshape(shape)
    index = jnp.array([0], dtype=jnp.int32)

    res = kernel(x, index)
    np.testing.assert_array_equal(res, x[:64])

  def test_jitted_function_containing_multiple_pallas_calls(self):
    # This test aims to ensure that execution works correctly inside CUDA
    # graphs. This is complementary to the test in
    # jaxlib/mosaic/gpu/custom_call_test.cc that checks that such jitted
    # functions do invoke CUDA graphs.
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], jnp.float32),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = x_ref[...] + 1

    @jax.jit
    def f(x):
      # Run the kernel 10 times because CUDA graphs only trigger for >= 5 ops.
      for _ in range(10):
        x = kernel(x)
      return x

    x = jnp.arange(256).astype(jnp.float32)
    np.testing.assert_array_equal(f(x), x + 10)

  @parameterized.product(
      op=[
          lax.neg,
          lax.bitwise_not,
          lax.logistic,
          lax.exp,
          lambda x: x**2,
          lambda x: x**5,
          lax.rsqrt,
          lax.sqrt,
          lax.tanh,
          lax.log,
          jax.nn.gelu,
          lax.abs,
          lax.round,
          lax.clz,
          lambda x: lax.round(x, lax.RoundingMethod.TO_NEAREST_EVEN),
      ],
      approx_math=[True, False],
  )
  def test_unary_op(self, op, approx_math):
    dtype = jnp.int32 if op in (lax.bitwise_not, lax.clz) else jnp.float32

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], dtype),
        compiler_params=plgpu.CompilerParams(approx_math=approx_math),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = op(x_ref[...])

    x = jnp.arange(256).astype(dtype)
    np.testing.assert_allclose(
        kernel(x), op(x), rtol=1e-5 if approx_math else 3e-7
    )

  @parameterized.parameters(False, True)
  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_sqrt(self, approx_math):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], jnp.float32),
        compiler_params=plgpu.CompilerParams(approx_math=approx_math),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = lax.sqrt(x_ref[...])

    x = jnp.arange(256).astype(jnp.float32)
    if approx_math:
      with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as out:
        y = jax.block_until_ready(kernel(x))
      self.assertIn("sqrt.approx.f32", out())
    else:
      y = kernel(x)
    np.testing.assert_allclose(y, lax.sqrt(x), rtol=1e-5)

  @parameterized.product(
      approx_math=(False, True),
      dtype=(jnp.float32, jnp.float16, jnp.bfloat16),
  )
  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_tanh(self, approx_math, dtype):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], dtype),
        compiler_params=plgpu.CompilerParams(approx_math=approx_math),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = lax.tanh(x_ref[...])

    x = (jnp.arange(256).astype(dtype) - 128) / 32
    if approx_math:
      with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as out:
        y = jax.block_until_ready(kernel(x))
      instr = {
          jnp.float32: "tanh.approx.f32",
          jnp.float16: "tanh.approx.f16",
          jnp.bfloat16: "tanh.approx.bf16",
      }[dtype]
      self.assertIn(instr, out())
    else:
      y = kernel(x)
    atol = 1e-2 if approx_math else 2e-3
    rtol = 1e-2 if approx_math else 2e-3
    np.testing.assert_allclose(y, lax.tanh(x), atol=atol, rtol=rtol)

  @parameterized.parameters(jnp.float32, jnp.int32, jnp.uint32)
  def test_sign(self, dtype):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], dtype),
    )
    def kernel(src_ref, dst_ref):
      dst_ref[...] = lax.sign(src_ref[...])

    # Use values that include negative, zero, and positive.
    src = np.arange(256, dtype=dtype) - 128
    np.testing.assert_array_equal(kernel(src), lax.sign(src))

  @parameterized.product(
      op=[
          operator.add,
          lambda x, _: x + 1,  # for int->vector conversion
          operator.sub,
          operator.mul,
          lax.div,
          jnp.minimum,
          jnp.maximum,
          lax.shift_left,
          lax.shift_right_logical,
          lax.shift_right_arithmetic,
      ],
      dtype=[jnp.float32, jnp.int32, jnp.uint32],
  )
  def test_binary_op(self, op, dtype):
    is_int_op = op in (lax.shift_left, lax.shift_right_logical, lax.shift_right_arithmetic)
    if is_int_op and jnp.issubdtype(dtype, jnp.floating):
      self.skipTest("shift ops do not support floating point types")

    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([256], dtype)
    )
    def kernel(x_ref, y_ref, o_ref):
      o_ref[...] = op(x_ref[...], y_ref[...])

    key0, key1 = jax.random.split(jax.random.key(0), 2)
    if is_int_op:
      info = jnp.iinfo(dtype)
      x = jnp.linspace(
          jnp.array(info.min, dtype=dtype),
          jnp.array(info.max, dtype=dtype),
          256,
          dtype=dtype,
      )
      y = jnp.linspace(0, (jax.dtypes.itemsize_bits(dtype)) - 1, 256, dtype=dtype)
    else:
      x = (jax.random.uniform(key0, [256]) * 42 - 21).astype(dtype)
      y = (jax.random.uniform(key1, [256]) * 42 - 21).astype(dtype)

    np.testing.assert_array_equal(kernel(x, y), op(x, y))

  def test_lax_div_int_rounding(self):
    dtype = jnp.int32
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([256], dtype)
    )
    def kernel(x_ref, y_ref, o_ref):
      o_ref[...] = lax.div(x_ref[...], y_ref[...])

    key = jax.random.key(0)
    key_x, key_y, key_sign = jax.random.split(key, 3)
    x = jax.random.randint(key_x, [256], -100, 100, dtype=dtype)
    y = jax.random.randint(key_y, [256], 1, 10, dtype=dtype)
    signs = jax.random.choice(key_sign, jnp.array([-1, 1], dtype=dtype), [256])
    y = y * signs

    np.testing.assert_array_equal(kernel(x, y), lax.div(x, y))

  @parameterized.product(
      op=[
          lax.eq,
          operator.ne,
          operator.lt,
          operator.le,
          operator.gt,
          operator.ge,
      ],
      # TODO(slebedev): Support integral types.
      dtype=[jnp.float32, jnp.int32, jnp.uint32],
  )
  def test_comparison_op(self, op, dtype):
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([256], dtype)
    )
    def kernel(o_ref):
      o_ref[...] = jnp.broadcast_to(
          op(dtype(42), dtype(24)).astype(dtype), o_ref.shape
      )

    np.testing.assert_array_equal(kernel(), jnp.full([256], op(42, 24), dtype))

  @parameterized.product(
      op=(jnp.sum, jnp.max, jnp.min),
      shape=((128,), (128, 64)),
      dtype=(jnp.float32, jnp.int32, jnp.uint32),
  )
  def test_scalar_reduce(self, op, shape, dtype):
    @self.kernel(out_type=jax.ShapeDtypeStruct((), dtype))
    def kernel(x_ref, o_ref):
      o_ref[...] = op(x_ref[...])

    x = jnp.arange(math.prod(shape)).reshape(shape).astype(dtype)
    if dtype != jnp.uint32:
      x = x - x.size // 2  # include negative values.
    np.testing.assert_array_equal(kernel(x), op(x))

  def test_reshape(self):
    shape1, shape2 = (128,), (2, 16, 4)

    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct(shape2, jnp.float32)
    )
    def kernel(x_ref, out_ref):
      x_ref_reshaped = x_ref.reshape(shape2)
      self.assertEqual(x_ref.shape, shape1)
      self.assertEqual(x_ref_reshaped.shape, shape2)
      out_ref[...] = x_ref_reshaped[...]

    x = jnp.arange(math.prod(shape1)).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x.reshape(shape2))

  @parameterized.parameters(jnp.float32, jnp.bfloat16)
  def test_optimized_store_to_gmem(self, dtype):
    shape = (128, 128)
    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)

    def run_kernel(optimized):

      @self.kernel(out_type=jax.ShapeDtypeStruct(shape, dtype))
      def kernel(x_ref, out_ref):
        x = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
        plgpu.store(out_ref, x + 1, optimized=optimized)

      return kernel(x)

    if dtype == jnp.float32:
      np.testing.assert_array_equal(run_kernel(optimized=True), x + 1)
    else:
      # Storing WGMMA bf16 (16B/row) to GMEM touches partial 32B sectors and
      # cannot use an optimized store transfer.
      with self.assertRaisesRegex(Exception, "GMEM transfer does not access"):
        run_kernel(optimized=True)
      np.testing.assert_array_equal(run_kernel(optimized=False), x + 1)

  def test_reshape_tiled(self):
    shape1, shape2 = (6 * 64, 8), (2, 3, 64, 8)

    @self.kernel(out_type=jax.ShapeDtypeStruct(shape2, jnp.float32))
    def kernel(x_ref, out_ref):
      x = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.store(out_ref, x.reshape(shape2), optimized=False)

    x = jnp.arange(math.prod(shape1)).reshape(shape1).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x.reshape(shape2))

  def test_reshape_tiled_into_tiled(self):
    shape1, shape2 = (6 * 64, 8), (2, 3, 64, 8)

    transforms = (plgpu.TilingTransform((8, 8)), plgpu.SwizzleTransform(32))

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape2, jnp.float32),
        scratch_types=[plgpu.SMEM(shape2, jnp.float32, transforms=transforms)],
    )
    def kernel(x_ref, out_ref, scratch_ref):
      x = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      # Here we actually want to test that we can store a tiled value to a ref
      # where the tiling transform rank is not the same as the ref's rank.
      scratch_ref[...] = x.reshape(shape2)
      plgpu.store(out_ref, scratch_ref[...], optimized=False)

    x = jnp.arange(math.prod(shape1)).reshape(shape1).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x.reshape(shape2))

  def test_reshape_tiled_with_indexing(self):
    # This is a regression test from a Tokamax GLU kernel.
    dtype = jnp.float16
    shape1, shape2 = (2, 64, 2, 64), (64, 128)
    swizzle = 128
    swizzle_elems = 8 * swizzle // dtypes.itemsize_bits(dtype)
    transforms = (
        plgpu.TilingTransform((8, 1, swizzle_elems)),
        plgpu.SwizzleTransform(swizzle),
    )

    @functools.partial(
        self.kernel,
        scratch_types=[
            plgpu.SMEM(shape1, dtype, transforms=transforms),
            plgpu.Barrier(),
        ],
        out_type=jax.ShapeDtypeStruct(shape2, dtype)
    )
    def kernel(src_ref, dst_ref, smem_ref, barrier):
      plgpu.copy_gmem_to_smem(src_ref, smem_ref, barrier)
      plgpu.barrier_wait(barrier)
      plgpu.store(dst_ref, smem_ref.at[1].reshape(shape2)[...], optimized=False)

    x = jnp.arange(math.prod(shape1), dtype=dtype).reshape(shape1)
    np.testing.assert_array_equal(kernel(x), x[1].reshape(shape2))

  def test_reshape_splat(self):
    shape = (1, 1, 1)

    @self.kernel(out_type=jax.ShapeDtypeStruct(shape, jnp.float32))
    def kernel(out_ref):
      x = jnp.array(42, dtype=jnp.float32)
      out_ref[...] = x.reshape(shape)

    np.testing.assert_array_equal(
        kernel(), jnp.array(42, dtype=jnp.float32).reshape(shape)
    )

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

    @self.kernel(out_type=jax.ShapeDtypeStruct(shape[2:], jnp.float32))
    def kernel(x_ref, out_ref):
      x = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.store(out_ref, x[1, 1], optimized=False)

    x = jnp.arange(math.prod(shape)).reshape(shape).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x[1, 1])

  def test_slice_array_with_static_ds(self):
    shape = (3, 64, 8)

    @self.kernel(out_type=jax.ShapeDtypeStruct((2, 64, 8), jnp.float32))
    def kernel(x_ref, out_ref):
      x = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.store(out_ref, x[pl.ds(1, 2)], optimized=False)

    x = jnp.arange(math.prod(shape)).reshape(shape).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x[1:3])

  @parameterized.named_parameters(
      ("axis0", 0, ((64, 64), (128, 64))),
      ("axis1", 1, ((64, 64), (64, 128))),
  )
  def test_concatenate(self, axis, shapes):
    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, jnp.float32)

    @self.kernel(out_type=out_shape)
    def kernel(inp1_ref, inp2_ref, out_ref):
      x0 = plgpu.load(inp1_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      x1 = plgpu.load(inp2_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.store(
          out_ref, lax.concatenate([x0, x1], dimension=axis), optimized=False
      )

    np.testing.assert_array_equal(kernel(x0, x1), expected)

  def test_squeeze_to_scalar(self):

    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.float32))
    def kernel(out_ref):
      x  = _array_splat(42, (1, 1, 1))
      out_ref[...] = lax.squeeze(x, dimensions=(0, 1, 2))

    np.testing.assert_array_equal(kernel(), jnp.array(42, dtype=jnp.float32))

  def test_add_one_grid_with_scratch(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([128 * 2], jnp.float32),
        in_specs=[pl.BlockSpec((128,), lambda *i: i)],
        out_specs=pl.BlockSpec((128,), lambda *i: i),
        scratch_shapes=[plgpu.SMEM((128,), jnp.float32)],
        grid=2,
    )
    def kernel(x_ref, o_ref, scratch_ref):
      scratch_ref[...] = x_ref[...] + 1
      o_ref[...] = scratch_ref[...]

    x = jnp.arange(256).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + 1.0)

  @parameterized.product(max_concurrent_steps=[1, 2, 3, 4, 16])
  def test_add_one_grid_pipelined(self, max_concurrent_steps):

    @functools.partial(
        self.pallas_call,
        in_specs=[pl.BlockSpec((128, 16), lambda i, j: (i, j))],
        out_specs=pl.BlockSpec((128, 16), lambda i, j: (i, j)),
        out_shape=jax.ShapeDtypeStruct([128 * 2, 64], jnp.float32),
        compiler_params=plgpu.CompilerParams(
            dimension_semantics=["parallel", "sequential"],
            max_concurrent_steps=max_concurrent_steps,
        ),
        grid=(2, 4),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = x_ref[...] + 1.0

    x = jnp.arange(128 * 2 * 64).reshape((128 * 2, 64)).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + 1.0)

  def test_add_one_grid_pipelined_with_leading_sequential_dimension(self):
    @functools.partial(
        self.pallas_call,
        in_specs=[pl.BlockSpec((128, 16), lambda i, j: (i, j))],
        out_specs=pl.BlockSpec((128, 16), lambda i, j: (i, j)),
        out_shape=jax.ShapeDtypeStruct([128 * 2, 64], jnp.float32),
        compiler_params=plgpu.CompilerParams(
            dimension_semantics=["sequential", "parallel"],
        ),
        grid=(2, 4),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = x_ref[...] + 1.0

    x = jnp.arange(128 * 2 * 64).reshape((128 * 2, 64)).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + 1.0)

  def test_add_one_grid_pipelined_program_id(self):

    @functools.partial(
        self.pallas_call,
        out_specs=pl.BlockSpec((16, 16), lambda i, j: (i, j)),
        out_shape=jax.ShapeDtypeStruct([16, 64], jnp.int32),
        compiler_params=plgpu.CompilerParams(
            dimension_semantics=["parallel", "sequential"],
            max_concurrent_steps=2,
        ),
        grid=(4, 4),
    )
    def kernel(o_ref):
      o_ref[...] = jnp.broadcast_to(pl.program_id(1), o_ref.shape)

    np.testing.assert_array_equal(
        kernel(),
        jnp.repeat(jnp.repeat(jnp.arange(4), 16)[None], 16, axis=0),
    )

  def test_add_one_grid_pipelined_sequential_invariant_output(self):

    @functools.partial(
        self.pallas_call,
        in_specs=[pl.BlockSpec((32, 16), lambda i, j: (i, j))],
        out_specs=pl.BlockSpec((32, 16), lambda i, j: (i, 0)),
        out_shape=jax.ShapeDtypeStruct([32 * 2, 64], jnp.float32),
        compiler_params=plgpu.CompilerParams(
            dimension_semantics=["parallel", "sequential"],
            max_concurrent_steps=2,
        ),
        grid=(2, 4),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = x_ref[...] + 1.0

    x = jnp.arange(32 * 2 * 64).reshape((32 * 2, 64)).astype(jnp.float32)
    y = jnp.empty_like(x)
    for i in range(2):
      i_slice = slice(32 * i, 32 * (i + 1))
      for j in range(4):
        j_slice = slice(16 * j, 16 * (j + 1))
        y = y.at[i_slice, :16].set(x[i_slice, j_slice] + 1)

    # We only compare the elements in the first 16 columns, because the rest
    # are never written to.
    np.testing.assert_array_equal(kernel(x)[:, :16], y[:, :16])

  @parameterized.parameters(jnp.float32, jnp.int32, jnp.uint32)
  def test_iota(self, dtype):
    dimension = 1

    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct((128, 128), dtype)
    )
    def kernel(o_ref):
      value = plgpu.broadcasted_iota(
          dtype, o_ref.shape, dimension, layout=plgpu.Layout.WGMMA
      )
      plgpu.store(o_ref, value, optimized=False)

    np.testing.assert_array_equal(
        kernel(), jax.lax.broadcasted_iota(dtype, (128, 128), dimension)
    )

  @parameterized.parameters(jnp.bfloat16, jnp.int16, jnp.uint16)
  def test_inline_mgpu(self, jnp_type):
    dtype = jnp.dtype(jnp_type)
    is_signed = mgpu.utils.is_signed(dtype)
    shape = (128, 128)
    tile = (64, 128 // dtype.itemsize)
    tiled_shape = list(mgpu.tile_shape(shape, tile))

    key = jax.random.key(0)
    x = jax.random.uniform(key, (2, *shape), minval=-10, maxval=10).astype(
        dtype
    )

    transforms = (
        plgpu.TilingTransform(tile),
        plgpu.SwizzleTransform(128),
    )

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, dtype),
        scratch_types=[
            plgpu.SMEM(
                x.shape,
                dtype,
                transforms=transforms,
            ),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref, o_ref, smem_ref, barrier):
      plgpu.copy_gmem_to_smem(x_ref, smem_ref, barrier)
      plgpu.barrier_wait(barrier)
      # Add an indexer at the end.
      sliced_smem_ref = smem_ref.at[0]
      @plgpu.inline_mgpu(
          arg_types=(plgpu.RefType((
              plgpu.TilingTransform(tile),
              plgpu.SwizzleTransform(128),
          )),),
          return_type=plgpu.ShapeDtypeStruct(
              shape, dtype, layout=plgpu.Layout.WGMMA
          ),
      )
      def foo(ctx, smem_ref):
        del ctx
        assert smem_ref.type.shape == tiled_shape, (smem_ref.type, tiled_shape)
        x = mgpu.FragmentedArray.load_tiled(
            smem_ref, swizzle=128, is_signed=is_signed
        )
        y = mgpu.FragmentedArray.splat(
            mgpu.c(1, x.mlir_dtype),
            shape=x.shape,
            layout=x.layout,
            is_signed=is_signed,
        )
        return (x + x + y)

      arr = foo(sliced_smem_ref)
      @plgpu.inline_mgpu(arg_types=(plgpu.Layout.WGMMA, plgpu.RefType(transforms), plgpu.RefType()))
      def store(ctx, arr, smem_ref, o_ref):
        sliced_smem_ref = mgpu.memref_slice(smem_ref, (0,))
        arr.store_tiled(sliced_smem_ref, swizzle=128)
        mgpu.commit_shared()
        ctx.async_copy(
            src_ref=sliced_smem_ref,
            dst_ref=o_ref,
            swizzle=128,
            gmem_transform=(mgpu.TileTransform(tile)),
        )
        ctx.await_async_copy(0)

        # A dummy if statement to make sure we inline nested blocks correctly.
        is_leader_thread = mgpu.utils.single_thread_predicate()
        with mgpu.utils.when(is_leader_thread):
          pass

      # This time we slice inside the inline_mgpu body.
      store(arr, smem_ref, o_ref)

    np.testing.assert_array_equal(kernel(x), x[0] + x[0] + 1)

  @parameterized.parameters(
      plgpu.Layout.WGMMA,
      plgpu.Layout.WGMMA_UPCAST_2X,
      plgpu.Layout.WGMMA_UPCAST_4X,
      plgpu.Layout.TCGEN05,
  )
  def test_inline_mgpu_layout_args(self, layout: gpu_core.SomeLayout):
    quant_dtype = jnp.int8
    dtype = jnp.bfloat16
    mgpu_layout = layout.to_mgpu()
    shape = (128, 128)

    rngs = list(jax.random.split(jax.random.key(0)))
    x = jax.random.randint(rngs.pop(), shape, minval=-10, maxval=10).astype(
        quant_dtype
    )
    x_s = jax.random.uniform(
        rngs.pop(), shape[0], minval=-100, maxval=100
    ).astype(dtype)

    @self.kernel(out_type=jax.ShapeDtypeStruct(shape, dtype))
    def kernel(x_ref, x_scale_ref, o_ref):
      x = plgpu.load(x_ref, layout=layout, optimized=False).astype(x_scale_ref.dtype)
      x_s = plgpu.load(x_scale_ref, layout=layout.reduce(1), optimized=False)

      @plgpu.inline_mgpu(
          arg_types=(layout, layout.reduce(1)),
          return_type=plgpu.ShapeDtypeStruct(shape, dtype, layout=layout),
      )
      def custom_broadcast(ctx, x_fa, xs_fa):
        del ctx
        return xs_fa.broadcast_in_dim(shape, [0], layout=mgpu_layout) * x_fa

      plgpu.store(o_ref, custom_broadcast(x, x_s), optimized=False)

    np.testing.assert_array_equal(
        kernel(x, x_s),
        x.astype(dtype) * jnp.broadcast_to(x_s[:, None], x.shape),
    )

  def test_inline_mgpu_scalar_args(self):

    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.float32))
    def kernel(o_ref):
      @plgpu.inline_mgpu(
          arg_types=(plgpu.RefType(), plgpu.Layout.WG_SPLAT),
      )
      def write(ctx, ref, value):
        del ctx
        memref_dialect.store(value.registers.item(), ref, [])
      write(o_ref, jnp.array(42.0, dtype=jnp.float32))
    np.testing.assert_array_equal(kernel(), 42.0)

  def test_sync_copy(self):
    shape = (128, 128)
    transforms = self.default_transforms(dtype=jnp.float32)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, jnp.float32),
        scratch_types=[plgpu.SMEM(shape, jnp.float32, transforms=transforms)],
    )
    def kernel(x_ref, y_ref, scratch_ref):
      layout = plgpu.Layout.SMEM_GMEM_COPY(shape, jnp.float32, swizzle=128)
      # GMEM loads require optimized=False, because we can't prove coalescing.
      # But with this layout they should be fast.
      scratch_ref[...] = plgpu.load(x_ref, layout=layout, optimized=False)
      plgpu.store(
          y_ref, plgpu.layout_cast(scratch_ref[...], layout), optimized=False
      )

    x = jnp.arange(math.prod(shape), dtype=jnp.float32).reshape(shape)
    np.testing.assert_array_equal(kernel(x), x)

  @parameterized.product(indexer=[..., slice(128), slice(None, 128)])
  def test_copy_smem_to_gmem(self, indexer):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], jnp.float32),
        scratch_types=[plgpu.SMEM((256,), jnp.float32)],
    )
    def kernel(x_ref, o_ref_gmem, scratch_ref):
      scratch_ref[...] = x_ref[...] + 1
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(scratch_ref.at[indexer], o_ref_gmem.at[indexer])
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(256).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x)[indexer], x[indexer] + 1.0)

  @parameterized.parameters([True, False])
  def test_copy_smem_to_gmem_predicate(self, predicate):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], jnp.float32),
        scratch_types=[plgpu.SMEM((256,), jnp.float32)],
    )
    def kernel(o_ref_gmem, scratch_ref):
      o_ref_gmem[...] = jnp.zeros([256], dtype=jnp.float32)
      scratch_ref[...] = jnp.ones([256], dtype=jnp.float32)
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(scratch_ref, o_ref_gmem, predicate=predicate)
      plgpu.wait_smem_to_gmem(0)

    if predicate:
      np.testing.assert_array_equal(kernel(), jnp.ones([256], dtype=jnp.float32))
    else:
      np.testing.assert_array_equal(kernel(), jnp.zeros([256], dtype=jnp.float32))

  @parameterized.parameters(True, False)
  def test_copy_smem_to_gmem_explicit_commit_group(self, use_single_warp):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], jnp.float32),
        scratch_types=[plgpu.SMEM((256,), jnp.float32)],
    )
    def kernel(x_ref, o_ref_gmem, scratch_ref):
      scratch_ref[...] = x_ref[...] + 1
      plgpu.commit_smem()

      def copy_to_gmem():
        plgpu.copy_smem_to_gmem(scratch_ref, o_ref_gmem, commit_group=False)
        plgpu.commit_smem_to_gmem_group()
        plgpu.wait_smem_to_gmem(0)

      if use_single_warp:
        plgpu.warp_map(lambda warp_id: pl.when(warp_id == 0)(copy_to_gmem))
      else:
        copy_to_gmem()

    x = jnp.arange(256).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + 1.0)

  @parameterized.parameters(jnp.bfloat16, jnp.float16, jnp.float32)
  def test_copy_smem_to_gmem_reduction(self, dtype):

    @jax.jit
    def run(x):
      x_ref = jax.new_ref(x, memory_space=plgpu.GMEM)
      o_ref = jax.new_ref(jnp.zeros(128, dtype=dtype), memory_space=plgpu.GMEM)

      @self.kernel(scratch_types=[plgpu.SMEM((128,), dtype)])
      def kernel(scratch_ref):
        def pipeline_body(_, x_smem):
          scratch_ref[...] = x_smem[...]
          plgpu.commit_smem()
          plgpu.copy_smem_to_gmem(
              scratch_ref.at[...], o_ref.at[...], reduction_op="add"
          )
          plgpu.wait_smem_to_gmem(0)

        plgpu.emit_pipeline(
            pipeline_body,
            in_specs=[pl.BlockSpec((128,), lambda i: (i,))],
            grid=(200,),
        )(x_ref)

      kernel()
      return jax.freeze(o_ref)

    x = jnp.ones(200 * 128).astype(dtype)
    output_val = x.reshape(-1, 128).sum(axis=0)
    np.testing.assert_array_equal(run(x), output_val)

  @parameterized.product(
      case=[
          ((256,), (...,)),
          ((64, 128,), (...,)),
          ((64, 128,), (slice(2, 3), slice(0, 128))),
          ((3, 64, 128), (...,)),
          ((3, 64, 1, 128), (0, slice(0, 32), 0, slice(0, 128))),
          ((3, 64, 1, 128), (...,)),
          ((5, 5, 512,), (4, 4)),
          ((5, 1024,), (4,)),
          ((8192,), (...,)),
          ((8192,), (slice(4096, 8192),)),
      ],
      oob_mode=[
          plgpu.OOBFillMode.UNDEFINED,
          plgpu.OOBFillMode.PROMISE_IN_BOUNDS,
          plgpu.OOBFillMode.ZEROS,
      ],
      dtype=[jnp.bfloat16, jnp.float32],
  )
  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  @jtu.skip_under_pytest("Test fails under pytest in CI")
  def test_copy_gmem_to_smem_contiguous(self, case, oob_mode, dtype):
    shape, indexer = case

    if max(shape) > 256 and oob_mode == plgpu.OOBFillMode.ZEROS:
      self.skipTest(
          "`OOBFillMode.ZEROS` not supported for shapes with max dimension > 256"
      )

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, dtype),
        scratch_types=[plgpu.SMEM(shape, dtype), plgpu.Barrier()],
    )
    def kernel(x_ref_gmem, o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(
          x_ref_gmem.at[indexer],
          scratch_ref.at[indexer],
          barrier_ref,
          oob_mode=oob_mode,
      )
      plgpu.barrier_wait(barrier_ref)
      scratch_ref[indexer] = scratch_ref[indexer] + 1
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(scratch_ref.at[indexer], o_ref.at[indexer])
      plgpu.wait_smem_to_gmem(0)

    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), jtu.capture_stdout() as output:
      x = jax.random.normal(jax.random.key(0), shape, dtype=dtype)
      np.testing.assert_allclose(kernel(x)[indexer], x[indexer] + 1.0)

    ptx = output()
    # The copy out is always a simple contiguous transfer in this test.
    self.assertIn("cp.async.bulk.global.shared::cta.bulk_group", ptx)

    if oob_mode == plgpu.OOBFillMode.ZEROS:
      self.assertNotIn("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes", ptx)
      self.assertIn("cp.async.bulk.tensor", ptx)
    else:
      self.assertIn("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes", ptx)
      self.assertNotIn("cp.async.bulk.tensor", ptx)

      if oob_mode == plgpu.OOBFillMode.PROMISE_IN_BOUNDS:
        self.assertNotIn("min.u", ptx)

  @parameterized.parameters(
      plgpu.OOBFillMode.UNDEFINED, plgpu.OOBFillMode.PROMISE_IN_BOUNDS
  )
  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_copy_gmem_to_smem_clamped_arrivals(self, oob_mode):
    # This test makes sure OOB reads don't deadlock when the transfer size is
    # clamped.
    gmem_shape = (8192,)
    smem_shape = (4096,)
    dtype = jnp.float32

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(gmem_shape, dtype),
        scratch_types=[plgpu.SMEM(smem_shape, dtype), plgpu.Barrier()],
    )
    def kernel(x_ref_gmem, offset, o_ref, scratch_ref, barrier_ref):
      gmem_indexer = (pl.dslice(offset[...], 4096),)
      plgpu.copy_gmem_to_smem(
          x_ref_gmem.at[gmem_indexer],
          scratch_ref,
          barrier_ref,
          oob_mode=oob_mode,
      )
      plgpu.barrier_wait(barrier_ref)

      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(scratch_ref, o_ref.at[gmem_indexer])
      plgpu.wait_smem_to_gmem(0)

    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as output:
      x = jax.numpy.ones(gmem_shape, dtype=dtype)
      jax.jit(kernel).lower(x, 6000).compile()

    ptx = output()

    if oob_mode == plgpu.OOBFillMode.PROMISE_IN_BOUNDS:
      self.assertNotIn("mbarrier.complete_tx", ptx)
    else:
      self.assertIn("mbarrier.complete_tx", ptx)

  @parameterized.named_parameters(
      {"testcase_name": "1d_none",
       "shape": (256,), "indexers": (slice(0, 128), slice(None, 32))},
      {"testcase_name": "1d_offset",
       "shape": (256,), "indexers": (slice(32, 96), slice(0, 32))},
      {"testcase_name": "2d_extract",
       "shape": (64, 64), "indexers": (4, slice(0, 64))},
      )
  def test_copy_smem_to_gmem_with_multiple_gmem_indexers(self, shape, indexers):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, jnp.float32),
        scratch_types=[plgpu.SMEM(shape, jnp.float32)],
    )
    def kernel(x_ref, o_ref_gmem, scratch_ref):
      scratch_ref[...] = x_ref[...] + 1
      plgpu.commit_smem()
      for indexer in indexers:
        scratch_ref = scratch_ref.at[indexer]
        o_ref_gmem = o_ref_gmem.at[indexer]
      plgpu.copy_smem_to_gmem(scratch_ref, o_ref_gmem)
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(np.prod(shape)).astype(jnp.float32).reshape(*shape)
    result = kernel(x)
    ref = x + 1.0
    for indexer in indexers:
      result = result[indexer]
      ref = ref[indexer]
    np.testing.assert_array_equal(result, ref)

  @parameterized.product(indexer=[..., slice(128), slice(None, 128)])
  @run_on_sm80
  def test_copy_gmem_to_smem(self, indexer):
    needs_barrier = jtu.is_cuda_compute_capability_at_least("9.0")

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], jnp.float32),
        scratch_types=[
            plgpu.SMEM((256,), jnp.float32),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref_gmem, o_ref, scratch_ref, barrier_ref):
      barrier = barrier_ref if needs_barrier else None
      plgpu.copy_gmem_to_smem(
          x_ref_gmem.at[indexer],
          scratch_ref.at[indexer],
          barrier,
          oob_mode=plgpu.OOBFillMode.PROMISE_IN_BOUNDS,
      )
      if barrier is not None:
        plgpu.barrier_wait(barrier)
      else:
        plgpu.wait_gmem_to_smem(0)
      o_ref[...] = scratch_ref[...] + 1

    x = jnp.arange(256).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x)[indexer], x[indexer] + 1.0)

  @run_on_sm80
  def test_copy_gmem_to_smem_dynamic_slice(self):
    needs_barrier = jtu.is_cuda_compute_capability_at_least("9.0")

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], jnp.float32),
        scratch_types=[
            plgpu.SMEM((128,), jnp.float32),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref_gmem, o_ref, scratch_ref, barrier_ref):
      assert o_ref.size % scratch_ref.size == 0
      barrier = barrier_ref if needs_barrier else None

      @pl.loop(0, o_ref.size // scratch_ref.size)
      def _(i):
        s = pl.ds(i * 128, 128)
        plgpu.copy_gmem_to_smem(
            x_ref_gmem.at[s],
            scratch_ref,
            barrier,
            oob_mode=plgpu.OOBFillMode.PROMISE_IN_BOUNDS,
        )
        if barrier is not None:
          plgpu.barrier_wait(barrier)
        else:
          plgpu.wait_gmem_to_smem(0)
        o_ref[s] = scratch_ref[...] + 1

    x = jnp.arange(256).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + 1.0)

  @run_on_sm80
  def test_copy_gmem_to_smem_squeeze(self):
    needs_barrier = jtu.is_cuda_compute_capability_at_least("9.0")

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([4, 8, 32], jnp.float32),
        scratch_types=[
            plgpu.SMEM((8, 32), jnp.float32),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref_gmem, o_ref, scratch_ref, barrier_ref):
      barrier = barrier_ref if needs_barrier else None

      @pl.loop(0, o_ref.shape[0])
      def _(i):
        plgpu.copy_gmem_to_smem(
            x_ref_gmem.at[i],
            scratch_ref,
            barrier,
            oob_mode=plgpu.OOBFillMode.PROMISE_IN_BOUNDS,
        )
        if barrier is not None:
          plgpu.barrier_wait(barrier)
        else:
          plgpu.wait_gmem_to_smem(0)
        o_ref[i] = scratch_ref[...] + 1

    x = jnp.arange(4 * 8 * 32).reshape(4, 8, 32).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + 1.0)

  @run_on_sm80
  def test_copy_gmem_to_smem_raises_on_mismatched_shapes(self):
    dtype = jnp.bfloat16

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([128], dtype),
        scratch_types=[
            plgpu.SMEM((256,), dtype),
            plgpu.Barrier(),
        ],
    )
    def kernel(o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(o_ref, scratch_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)

    with self.assertRaisesRegex(
        ValueError,
        "Expected source shape to be \\(256,\\), but got \\(128,\\).",
    ):
      jax.jit(kernel).lower()

  @run_on_sm80
  def test_copy_smem_to_gmem_raises_on_mismatched_shapes(self):
    dtype = jnp.bfloat16

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([128], dtype),
        scratch_types=[
            plgpu.SMEM((256,), dtype),
        ],
    )
    def kernel(o_ref, scratch_ref):
      plgpu.copy_smem_to_gmem(scratch_ref, o_ref)
      plgpu.wait_smem_to_gmem(0)

    with self.assertRaisesRegex(
        ValueError,
        "Expected shapes to match, but src has shape \\(256,\\) and dst has shape \\(128,\\)",
    ):
      jax.jit(kernel).lower()

  def test_copy_gmem_to_smem_raises_on_mismatched_shapes_partitioned(self):
    dtype = jnp.bfloat16

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([127], dtype),
        scratch_types=[
            plgpu.SMEM((8,), dtype),
            plgpu.Barrier(),
        ],
        cluster=(2,),
        cluster_names=("x",),
    )
    def kernel(o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(
          o_ref,
          scratch_ref,
          barrier_ref,
          collective_axes="x",
          leader_tracked=plgpu.CopyPartition.PARTITIONED(0),
      )
      plgpu.barrier_wait(barrier_ref)

    with self.assertRaisesRegex(
        ValueError,
        (
            "Expected source shape to be \\(16,\\), but got \\(127,\\). Dim 0"
            " is partitioned over an axis of size 2."
        ),
    ):
      jax.jit(kernel).lower()

  def test_copy_gmem_to_smem_raises_on_mismatched_shapes_partitioned_composite_index(self):
    dtype = jnp.bfloat16
    @self.kernel(
        out_type=jax.ShapeDtypeStruct([128], dtype),
        scratch_types=[
            plgpu.SMEM((32,), dtype),
            plgpu.Barrier(),
        ],
        cluster=(2,),
        cluster_names=(("x", "y"),),
    )
    def kernel(o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(
          o_ref,
          scratch_ref,
          barrier_ref,
          collective_axes=(("x", "y"),),
          leader_tracked=plgpu.CopyPartition.PARTITIONED(0),
      )
      plgpu.barrier_wait(barrier_ref)

    with self.assertRaisesRegex(
        ValueError,
        (
            "Expected source shape to be \\(64,\\), but got \\(128,\\). Dim 0"
            " is partitioned over an axis of size 2."
        ),
    ):
      jax.jit(kernel).lower()

  def test_copy_gmem_to_smem_raises_on_mismatched_dtypes(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], jnp.bfloat16),
        scratch_types=[
            plgpu.SMEM((256,), jnp.float32),
            plgpu.Barrier(),
        ],
    )
    def kernel(o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(o_ref, scratch_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)

    with self.assertRaisesRegex(
        ValueError,
        "Expected dtypes to match.*src.*bfloat16.*dst.*float32",
    ):
      jax.jit(kernel).lower()

  def test_copy_smem_to_gmem_raises_on_mismatched_dtypes(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], jnp.bfloat16),
        scratch_types=[
            plgpu.SMEM((256,), jnp.float32),
        ],
    )
    def kernel(o_ref, scratch_ref):
      plgpu.copy_smem_to_gmem(scratch_ref, o_ref)
      plgpu.wait_smem_to_gmem(0)

    with self.assertRaisesRegex(
        ValueError,
        "Expected dtypes to match.*src.*float32.*dst.*bfloat16",
    ):
      jax.jit(kernel).lower()

  def test_wait_smem_to_gmem_raises_on_non_integer_operand(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct([256], jnp.float32))
    def kernel(o_ref):
      del o_ref  # Unused.
      plgpu.wait_smem_to_gmem(0.0)

    with self.assertRaisesRegex(
        ValueError, "n must be a non-negative integer, got 0.0"
    ):
      jax.eval_shape(kernel)

  def test_wgmma_wait_raises_on_negative_operand(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct([256], jnp.float32))
    def kernel(o_ref):
      del o_ref  # Unused.
      plgpu.wgmma_wait(-1)

    with self.assertRaisesRegex(
        ValueError, "n must be a non-negative integer, got -1"
    ):
      jax.eval_shape(kernel)

  def test_set_max_registers_raises_on_invalid_action(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct([256], jnp.float32))
    def kernel(o_ref):
      del o_ref  # Unused.
      plgpu.set_max_registers(40, action="Increase")

    with self.assertRaisesRegex(
        ValueError, "action must be 'increase' or 'decrease', got 'Increase'"
    ):
      jax.eval_shape(kernel)

  def test_set_max_registers_raises_on_invalid_register_count(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct([256], jnp.float32))
    def kernel(o_ref):
      del o_ref  # Unused.
      plgpu.set_max_registers(41, action="increase")

    with self.assertRaisesRegex(
        ValueError, "n must be a positive multiple of 8, got 41"
    ):
      jax.eval_shape(kernel)

  @parameterized.parameters(False, True)
  def test_copy_gmem_to_smem_predicate(self, warp_level):
    @self.kernel(
        out_type=(
            jax.ShapeDtypeStruct([256], jnp.float32),
            jax.ShapeDtypeStruct([2], jnp.int32),
        ),
        scratch_types=[
            plgpu.SMEM([256], jnp.float32),
            plgpu.Barrier(),
        ],
        grid=(2,),
        grid_names=("block",),
    )
    def kernel(x_ref, o_ref, completed_ref, scratch_ref, barrier_ref):
      scratch_ref[...] = jnp.zeros([256], dtype=jnp.float32)
      block_id = jax.lax.axis_index("block")
      idx = pl.ds(block_id * 128, 128)
      if warp_level:
        @plgpu.warp_map
        def _(warp_id):
          plgpu.copy_gmem_to_smem(
              x_ref.at[idx],
              scratch_ref.at[idx],
              barrier_ref,
              predicate=jnp.logical_and(block_id == 0, warp_id == 0),
          )
      else:
        plgpu.copy_gmem_to_smem(
            x_ref.at[idx],
            scratch_ref.at[idx],
            barrier_ref,
            predicate=block_id == 0,
        )
      @pl.when(block_id == 0)
      def _():
        plgpu.barrier_wait(barrier_ref)
      completed_ref[block_id] = plgpu.barrier_test(barrier_ref).astype(
          jnp.int32
      )
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(scratch_ref.at[idx], o_ref.at[idx])
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(256).astype(jnp.float32)
    output, completed = kernel(x)
    np.testing.assert_array_equal(output[:128], x[:128])
    np.testing.assert_array_equal(output[128:], jnp.zeros((128,)))
    np.testing.assert_array_equal(completed, jnp.zeros((2,), dtype=jnp.int32))

  def test_collective_copy_gmem_to_smem(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((2, 128), jnp.float32),
        scratch_types=dict(
            smem_ref=plgpu.SMEM((128,), jnp.float32),
            barrier_ref=plgpu.Barrier(),
        ),
        cluster=(2,),
        cluster_names=("cluster",),
    )
    def kernel(x_ref, y_ref, smem_ref, barrier_ref):
      # Specifying collective_axes will enable TMA multicast automatically.
      plgpu.copy_gmem_to_smem(
          x_ref, smem_ref, barrier_ref, collective_axes="cluster"
      )
      plgpu.barrier_wait(barrier_ref)
      plgpu.copy_smem_to_gmem(smem_ref, y_ref.at[jax.lax.axis_index("cluster")])
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(128, dtype=jnp.float32)
    y = kernel(x)
    # Each block gets the same data and writes it out.
    np.testing.assert_array_equal(y, jnp.stack([x, x], axis=0))

  @parameterized.parameters(True, False)
  def test_kernel_scratch_kwargs(self, vmap):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
        scratch_types=dict(
            smem_ref=plgpu.SMEM((128,), jnp.float32),
            barrier_ref=plgpu.Barrier(),
        ),
    )
    def kernel(x_ref, y_ref, *, smem_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(x_ref, smem_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      plgpu.copy_smem_to_gmem(smem_ref, y_ref)
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(128, dtype=jnp.float32)

    if vmap:
      x = jnp.stack([x, x], axis=0)
      kernel = jax.vmap(kernel)

    np.testing.assert_array_equal(kernel(x), x)

  @parameterized.product(
      indexer=[..., slice(128), slice(None, 128)],
      predicate=[None, True, False],
  )
  def test_async_prefetch(self, indexer, predicate):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([256], jnp.float32),
        scratch_types=[
            plgpu.SMEM((256,), jnp.float32),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref_gmem, o_ref, scratch_ref, barrier_ref):
      plgpu.async_prefetch(x_ref_gmem.at[indexer], predicate=predicate)
      plgpu.copy_gmem_to_smem(
          x_ref_gmem.at[indexer], scratch_ref.at[indexer], barrier_ref
      )
      plgpu.barrier_wait(barrier_ref)
      o_ref[...] = scratch_ref[...] + 1

    x = jnp.arange(256).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x)[indexer], x[indexer] + 1.0)

  def test_barrier_repeated_indexing(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
        scratch_shapes=[plgpu.Barrier(num_barriers=4)],
    )
    def kernel(o_ref, barrier_ref):
      b = barrier_ref.at[2:4].at[1:2]
      plgpu.barrier_arrive(b)
      plgpu.barrier_wait(barrier_ref.at[3])
      o_ref[...] = jnp.ones_like(o_ref)

    np.testing.assert_array_equal(kernel(), np.ones((128,), dtype=jnp.float32))

  @parameterized.parameters(True, False)
  def test_barrier_test(self, use_single_warp):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((5,), jnp.int32),
        scratch_types=[plgpu.Barrier(num_barriers=2)],
    )
    def kernel(o_ref, barrier_ref):
      def test_barrier():
        o_ref[0] = plgpu.barrier_test(barrier_ref.at[0]).astype(jnp.int32)
        o_ref[1] = plgpu.barrier_test(barrier_ref.at[1]).astype(jnp.int32)
        plgpu.barrier_arrive(barrier_ref.at[0])
        o_ref[2] = plgpu.barrier_test(barrier_ref.at[0]).astype(jnp.int32)
        o_ref[3] = plgpu.barrier_test(barrier_ref.at[1]).astype(jnp.int32)
        plgpu.barrier_arrive(barrier_ref.at[1])
        o_ref[4] = plgpu.barrier_test(barrier_ref.at[1]).astype(jnp.int32)

      if use_single_warp:
        plgpu.warp_map(lambda warp_id: pl.when(warp_id == 0)(test_barrier))
      else:
        test_barrier()

    np.testing.assert_array_equal(kernel(), np.array([0, 0, 1, 0, 1]))

  def test_barrier_arrive_and_wait(self):
    @functools.partial(
        self.kernel,
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
        scratch_types=[
            plgpu.SMEM((128,), jnp.float32), plgpu.Barrier(num_arrivals=2)
        ],
        num_threads=2,
        thread_name="threads",
    )
    def kernel(o_ref, smem_ref, barrier):
      t_idx = lax.axis_index("threads")

      @pl.when(t_idx == 0)
      def _():
        smem_ref[...] = jnp.full((128,), 42.0, dtype=jnp.float32)
        plgpu.barrier_arrive(barrier)

      @pl.when(t_idx == 1)
      def _():
        plgpu.barrier_arrive_and_wait(barrier)
        o_ref[...] = smem_ref[...]

    np.testing.assert_array_equal(
        kernel(), np.full((128,), 42.0, dtype=jnp.float32)
    )

  @parameterized.named_parameters(
      {
          "testcase_name": "1d_none",
          "shape": (256,),
          "indexers": (slice(0, 128), slice(None, 32)),
      },
      {
          "testcase_name": "1d_offset",
          "shape": (256,),
          "indexers": (slice(32, 96), slice(0, 32)),
      },
      {
          "testcase_name": "2d_extract_static",
          "shape": (64, 64),
          "indexers": (4, slice(0, 64)),
      },
      {
          "testcase_name": "2d_extract_dyn",
          "shape": (64, 64),
          "indexers": lambda in_dev: (
              lax.axis_index("x") + 4 if in_dev else jnp.array(4),
              slice(0, 64),
          ),
      },
  )
  def test_copy_gmem_to_smem_with_multiple_gmem_indexers(self, shape, indexers):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, jnp.float32),
        scratch_types=[
            plgpu.SMEM(shape, jnp.float32),
            plgpu.Barrier(),
        ],
        grid=(1,),
        grid_names=("x",),
    )
    def kernel(x_ref_gmem, o_ref, scratch_ref, barrier_ref):
      scratch_ref_sliced = scratch_ref
      for indexer in indexers(True) if callable(indexers) else indexers:
        scratch_ref_sliced = scratch_ref_sliced.at[indexer]
        x_ref_gmem = x_ref_gmem.at[indexer]
      plgpu.copy_gmem_to_smem(
          x_ref_gmem, scratch_ref_sliced, barrier_ref
      )
      plgpu.barrier_wait(barrier_ref)
      o_ref[...] = scratch_ref[...] + 1

    x = jnp.arange(np.prod(shape)).astype(jnp.float32).reshape(*shape)
    result = kernel(x)
    ref = x + 1.0
    for indexer in indexers(False) if callable(indexers) else indexers:
      result = result[indexer]
      ref = ref[indexer]
    np.testing.assert_array_equal(result, ref)

  def test_gmem_to_smem_with_multiple_smem_indexers(self):
    x = jax.random.uniform(jax.random.key(0), (2, 64, 64), dtype=jnp.float32)
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([64, 64], jnp.float32),
        in_specs=(pl.BlockSpec(memory_space=plgpu.GMEM),),
        scratch_shapes=[
            plgpu.SMEM(x.shape, jnp.float32),
            plgpu.Barrier(),
        ],
    )
    def extract_x0(x_ref_gmem, o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(x_ref_gmem, scratch_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      x_sliced = scratch_ref.at[0, :, :]  # shape=(64, 64)
      o_ref[pl.ds(0, 32), :] = x_sliced[pl.ds(0, 32), :]
      o_ref[pl.ds(32, 32), :] = x_sliced[pl.ds(32, 32), :]
    np.testing.assert_array_equal(extract_x0(x), x[0])

  def test_gmem_to_smem_with_multiple_smem_indexers_and_transforms(self):
    transforms = self.default_transforms(dtype=jnp.int32)
    x = jnp.arange(512 * 512, dtype=jnp.int32).reshape(512, 512)
    @functools.partial(
        self.pallas_call,
        grid=(4, 4),
        out_shape=jax.ShapeDtypeStruct((256, 128), jnp.int32),
        in_specs=(
            plgpu.BlockSpec(
                block_shape=(128, 128),
                index_map=lambda i, j: (i, j),
                memory_space=plgpu.SMEM,
                transforms=transforms,
            ),
        ),
        out_specs=(
            plgpu.BlockSpec(
                block_shape=(64, 32),
                index_map=lambda i, j: (i, j),
                memory_space=plgpu.SMEM,
            )
        ),
    )
    def kernel(x_ref, o_ref):
      x_sliced = x_ref.at[0:64, 32:96].at[:, 0:32]  # get x_ref[0:64, 32:64]
      plgpu.store(o_ref, x_sliced[...], optimized=False)
    ref = jnp.concatenate([x[blk:blk+64, :] for blk in range(0, 512, 128)])
    ref = jnp.concatenate(
        [ref[:, blk+32:blk+64] for blk in range(0, 512, 128)], axis=1)
    np.testing.assert_array_equal(kernel(x), ref)

  @parameterized.parameters(
      *(dict(indexer=i, num_barriers=4) for i in [0, 1, 2, 3]),
      dict(indexer=(0, 0), num_barriers=(2, 2)),
      dict(indexer=(1, 1), num_barriers=(2, 2)),
  )
  def test_copy_gmem_to_smem_with_indexed_barrier(self, indexer, num_barriers):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([128], jnp.float32),
        scratch_types=[
            plgpu.SMEM((128,), jnp.float32),
            plgpu.Barrier(num_barriers=num_barriers),
        ],
    )
    def kernel(x_ref_gmem, o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(
          x_ref_gmem, scratch_ref, barrier_ref.at[indexer]
      )
      plgpu.barrier_wait(barrier_ref.at[indexer])
      o_ref[...] = scratch_ref[...] + 1

    x = jnp.arange(128).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + 1.0)

  @parameterized.named_parameters(("_g2s", False), ("_s2g", True))
  def test_copy_with_transforms(self, to_smem):
    transforms = self.default_transforms(dtype=jnp.float32)

    def kernel(x_ref, o_ref, barrier_ref):
      if to_smem:
        plgpu.copy_gmem_to_smem(x_ref, o_ref, barrier_ref)
        plgpu.barrier_wait(barrier_ref)
      else:
        plgpu.commit_smem()
        plgpu.copy_smem_to_gmem(x_ref, o_ref)
        plgpu.wait_smem_to_gmem(0)

    in_spec = pl.BlockSpec(memory_space=plgpu.GMEM)
    out_spec = plgpu.BlockSpec(
        transforms=transforms,
        memory_space=plgpu.SMEM,
    )
    if not to_smem:
      in_spec, out_spec = out_spec, in_spec
    f = self.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct([128, 128], jnp.float32),
        in_specs=(in_spec,),
        out_specs=out_spec,
        scratch_shapes=[plgpu.Barrier()],
    )
    x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
    np.testing.assert_array_equal(f(x), x)

  @parameterized.parameters(32, 64, 128)
  def test_copy_with_sole_swizzle_lowers_correctly(self, swizzle):
    dtype = jnp.float32
    trailing_dim = 8 * swizzle // 32
    shape = [128, trailing_dim]

    def kernel(x_ref, o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(x_ref, scratch_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      plgpu.copy_smem_to_gmem(scratch_ref, o_ref)
      plgpu.wait_smem_to_gmem(0)

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct(shape, dtype),
        scratch_types=(
            plgpu.SMEM(
                tuple(shape),
                dtype,
                transforms=(plgpu.SwizzleTransform(swizzle),),
            ),
            plgpu.Barrier(),
        ),
    )
    x = jnp.ones(shape, dtype=dtype).reshape(shape)
    jax.block_until_ready(f(x))

  def test_copy_7d_tiling_with_unit_middle_dims(self):
    # 7D GMEM/SMEM refs with a 7D TilingTransform results in a 14D physical
    # shape once tiling is applied.  The unit-sized dimensions are dropped
    # when lowering async copies, in order to bring the shape within TMA's
    # 5-dimension limit.
    shape = (2, 64, 1, 1, 1, 1, 128)
    dtype = jnp.float32
    transforms = (
        plgpu.TilingTransform((2, 1, 1, 1, 1, 1, 32)),
    )

    def kernel(x_ref, o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(x_ref, scratch_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      plgpu.copy_smem_to_gmem(scratch_ref, o_ref)
      plgpu.wait_smem_to_gmem(0)

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct(shape, dtype),
        scratch_types=(
            plgpu.SMEM(shape, dtype, transforms=transforms),
            plgpu.Barrier()
        ),
    )
    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    np.testing.assert_array_equal(f(x), x, strict=True)

  def test_scoped_copy_with_transforms(self):
    ts = self.default_transforms(dtype=jnp.float32)
    def kernel(x_ref, o_ref, barrier_ref):
      def body(tmp_ref):
        plgpu.copy_gmem_to_smem(x_ref, tmp_ref, barrier_ref)
        plgpu.barrier_wait(barrier_ref)
        o_ref[...] = tmp_ref[...] * 2
      pl.run_scoped(body, plgpu.SMEM((128, 64), jnp.float32, transforms=ts))

    in_spec = pl.BlockSpec(memory_space=plgpu.GMEM)
    out_spec = plgpu.BlockSpec(transforms=ts, memory_space=plgpu.SMEM)
    f = self.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct([128, 64], jnp.float32),
        in_specs=(in_spec,),
        out_specs=out_spec,
        scratch_shapes=[plgpu.Barrier()],
    )
    x = jnp.arange(128 * 64, dtype=jnp.float32).reshape(128, 64)
    np.testing.assert_array_equal(f(x), x * 2)

  def test_copy_with_transforms_and_indexing(self):
    def kernel(x_ref, o_ref, scratch_ref, barrier_ref):
      transposed_scratch = scratch_ref.transpose((1, 0, 2, 3))
      for i in range(2):
        plgpu.copy_gmem_to_smem(x_ref, transposed_scratch.at[:, i], barrier_ref)
        plgpu.barrier_wait(barrier_ref)
      plgpu.copy_smem_to_gmem(transposed_scratch, o_ref)
      plgpu.wait_smem_to_gmem(0)

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct([1, 2, 64, 64], jnp.float32),
        scratch_types=(
            plgpu.SMEM((2, 1, 64, 64), jnp.float32,
                       transforms=self.default_transforms(dtype=jnp.float32)),
            plgpu.Barrier()
        ),
    )
    x = jnp.arange(64 * 64, dtype=jnp.float32).reshape(1, 64, 64)
    np.testing.assert_array_equal(f(x), np.stack([x, x], axis=1))

  @parameterized.parameters(
      ((), plgpu.Layout.TMA_INDICES),
      ((), plgpu.Layout.TMA_INDICES_4),
      ((plgpu.TilingTransform((8, 32)), plgpu.SwizzleTransform(128)), plgpu.Layout.TMA_INDICES),
  )
  def test_copy_gmem_to_smem_gather(self, transforms, idxs_layout):
    if not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("Only works on a GPU with capability >= sm100")
    dtype = jnp.int32
    out_shape = (64 if idxs_layout == plgpu.Layout.TMA_INDICES else 12, 128)
    shape = (128, 64 + out_shape[-1])
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(out_shape, dtype),
        out_specs=plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=transforms),
        in_specs=(
            pl.BlockSpec(memory_space=plgpu.GMEM),
            pl.BlockSpec(memory_space=plgpu.GMEM),
        ),
        scratch_shapes=[plgpu.Barrier()],
    )
    def kernel(x_ref_gmem, idx_ref, o_ref, barrier_ref):
      idxs = plgpu.load(idx_ref, layout=idxs_layout, optimized=False)
      plgpu.copy_gmem_to_smem(x_ref_gmem.at[idxs, 64:], o_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)

    x = jnp.arange(math.prod(shape)).reshape(shape).astype(dtype)
    idx = jax.random.permutation(jax.random.key(1234), out_shape[0]).astype(jnp.uint32)
    np.testing.assert_array_equal(kernel(x, idx), x[idx, 64:])

  @parameterized.parameters(
      ((), plgpu.Layout.TMA_INDICES),
      ((), plgpu.Layout.TMA_INDICES_4),
      ((plgpu.TilingTransform((8, 32)), plgpu.SwizzleTransform(128)), plgpu.Layout.TMA_INDICES),
  )
  def test_copy_smem_to_gmem_scatter(self, transforms, idxs_layout):
    self.skip_unless_tcgen05()
    # Make sure we can infer the layout in WG
    indices_layout = tokens_layout = None
    if self.LOWERING_SEMANTICS == plgpu.LoweringSemantics.Lane:
      indices_layout = idxs_layout
      tokens_layout = plgpu.Layout.WGMMA if idxs_layout == plgpu.Layout.TMA_INDICES else None
    dtype = jnp.int32
    num_indices = 64 if idxs_layout == plgpu.Layout.TMA_INDICES else 12
    shape = (num_indices, 128)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, dtype),
        scratch_types=[plgpu.SMEM(shape, dtype, transforms=transforms)],
    )
    def kernel(tokens_ref, perm_ref, o_ref, smem_ref):
      plgpu.store(
          smem_ref,
          plgpu.load(tokens_ref, layout=tokens_layout, optimized=False),
          optimized=False,
      )
      plgpu.commit_smem()
      idxs = plgpu.load(perm_ref, layout=indices_layout, optimized=False)
      plgpu.copy_smem_to_gmem(smem_ref, o_ref.at[idxs, :])
      plgpu.wait_smem_to_gmem(0)

    key = jax.random.key(0)
    tokens = jax.random.randint(key, shape, 0, 100, dtype=dtype)
    perm = jax.random.permutation(key, shape[0]).astype(jnp.uint32)
    expected = jnp.zeros_like(tokens).at[perm].set(tokens)
    np.testing.assert_array_equal(kernel(tokens, perm), expected)

  @parameterized.product(
      src_transposed=(False, True), shape=((128, 128), (1, 128, 128))
  )
  def test_transposed_load_store(self, src_transposed, shape):
    dtype = jnp.float32
    permutation = (0, 2, 1) if len(shape) == 3 else (1, 0)
    @self.kernel(out_type=jax.ShapeDtypeStruct(shape, dtype))
    def kernel(src_ref, dst_ref):
      if src_transposed:
        src_ref = src_ref.transpose(permutation)
        src_layout = plgpu.Layout.WGMMA_TRANSPOSED
        dst_layout = plgpu.Layout.WGMMA
      else:
        dst_ref = dst_ref.transpose(permutation)
        src_layout = plgpu.Layout.WGMMA
        dst_layout = plgpu.Layout.WGMMA_TRANSPOSED
      src = plgpu.load(src_ref, layout=src_layout, optimized=False)
      dst = plgpu.layout_cast(src, dst_layout)
      plgpu.store(dst_ref, dst, optimized=False)

    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    np.testing.assert_array_equal(kernel(x), jnp.transpose(x, permutation))

  def test_wgmma_load_from_transposed_smem(self):
    shape = (128, 128)
    dtype = jnp.float16

    transforms = ()
    if not self.is_wg_semantics():
      transforms = self.default_transforms(swizzle=128, dtype=dtype)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape[::-1], dtype),
        scratch_types=[
            plgpu.SMEM(shape, dtype, transforms=transforms), plgpu.Barrier()
        ],
    )
    def kernel(x_ref, o_ref, smem_ref, barrier):
      plgpu.copy_gmem_to_smem(x_ref, smem_ref, barrier)
      plgpu.barrier_wait(barrier)
      smem_ref_t = smem_ref.transpose((1, 0))
      # TODO(b/540761884): If we skip the layout, we end up inferring WGStrided
      # and emitting vector.load along a non-contiguous dimension and failing
      # the MLIR verifier
      x = plgpu.load(smem_ref_t, layout=plgpu.Layout.WGMMA)
      plgpu.store(o_ref, x, optimized=False)

    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    np.testing.assert_array_equal(kernel(x), x.T)

  def test_wgmma_store_to_transposed_smem(self):
    shape = (128, 128)
    dtype = jnp.float16

    transforms = ()
    if not self.is_wg_semantics():
      transforms = self.default_transforms(swizzle=128, dtype=dtype)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape[::-1], dtype),
        scratch_types=[plgpu.SMEM(shape[::-1], dtype, transforms=transforms)],
    )
    def kernel(x_ref, o_ref, smem_ref):
      x = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      smem_ref.transpose((1, 0))[...] = x
      plgpu.store(o_ref, smem_ref[...], optimized=False)

    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    np.testing.assert_array_equal(kernel(x), x.T)

  @parameterized.product(
      src_memory_space=[plgpu.SMEM, plgpu.GMEM],
      layout=[plgpu.Layout.WG_STRIDED((128,), vec_size=1), None,
      ]
  )
  def test_load_to_strided_layout_with_indexing(self, src_memory_space, layout):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([2, 128], jnp.float32),
        in_specs=[pl.BlockSpec(memory_space=src_memory_space)],
        out_specs=plgpu.BlockSpec(memory_space=plgpu.SMEM),
    )
    def kernel(x_ref, o_ref):
      for i in range(2):
        x = plgpu.load(x_ref.at[i], layout=layout)
        o_ref[i, ...] = x

    x = jnp.arange(2 * 128, dtype=jnp.float32).reshape(2, 128)
    np.testing.assert_array_equal(kernel(x), x)

  @parameterized.parameters(None, plgpu.TilingTransform((32,)))
  def test_smem_gmem_transposed_copies(self, tiling):
    shape = (2, 2, 64)
    transforms = (tiling,) if tiling is not None else ()

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, jnp.float32),
        scratch_types=[
            plgpu.SMEM(shape, jnp.float32, transforms=transforms),
            plgpu.Barrier(),
        ],
    )
    def kernel(src_ref, dst_ref, smem_ref, barrier_ref):
      smem_ref = smem_ref.transpose((1, 0, 2))
      plgpu.copy_gmem_to_smem(src_ref, smem_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      plgpu.copy_smem_to_gmem(smem_ref, dst_ref)

    x = jnp.arange(math.prod(shape), dtype=jnp.float32).reshape(shape)
    np.testing.assert_array_equal(kernel(x), x)

  def test_indexing_before_transpose(self):
    def kernel(x_ref, o_ref, barrier_ref):
      for i in range(2):
        plgpu.copy_gmem_to_smem(
            x_ref, o_ref.at[i].transpose((1, 0, 2)), barrier_ref
        )
        plgpu.barrier_wait(barrier_ref)

    in_spec = pl.BlockSpec(memory_space=plgpu.GMEM)
    out_spec = plgpu.BlockSpec(memory_space=plgpu.SMEM)
    f = self.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct([2, 32, 2, 128], jnp.float32),
        in_specs=(in_spec,),
        out_specs=out_spec,
        scratch_shapes=[plgpu.Barrier()],
    )
    x = jnp.arange(2 * 32 * 128, dtype=jnp.float32).reshape(2, 32, 128)
    xt = x.transpose((1, 0, 2))
    np.testing.assert_array_equal(f(x), np.stack([xt, xt], axis=0))

  def test_indexing_after_transpose(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct([2, 32, 2, 128], jnp.float32),
        scratch_types=[
            plgpu.SMEM([2, 32, 2, 128], jnp.float32),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref, o_ref, smem_ref, barrier_ref):
      smem_transposed = smem_ref.transpose((1, 0, 2, 3))
      for i in range(2):
        plgpu.copy_gmem_to_smem(x_ref, smem_transposed.at[:, i], barrier_ref)
        plgpu.barrier_wait(barrier_ref)
      plgpu.copy_smem_to_gmem(smem_ref, o_ref)

    x = jnp.arange(32 * 2 * 128, dtype=jnp.float32).reshape(32, 2, 128)
    np.testing.assert_array_equal(kernel(x), np.stack([x, x], axis=0))

  def test_copy_gmem_to_smem_in_run_scoped(self):

    @self.kernel(out_type=jax.ShapeDtypeStruct([256], jnp.float32))
    def kernel(x_ref_gmem, o_ref):
      def body(barrier_ref):
        def inner_body(scratch_ref):
          plgpu.copy_gmem_to_smem(x_ref_gmem, scratch_ref, barrier_ref)
          plgpu.barrier_wait(barrier_ref)
          o_ref[...] = scratch_ref[...] + 1
        pl.run_scoped(inner_body, plgpu.SMEM((256,), jnp.float32))
      pl.run_scoped(body, plgpu.Barrier())

    x = jnp.arange(256).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + 1.0)

  def test_add_doubled_sum(self):

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([128], jnp.float32),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = x_ref[...] + jnp.sum(x_ref[...]) + jnp.sum(x_ref[...])

    x = jnp.arange(128).astype(jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + x.sum()*2)

  @parameterized.product(input_factor=[0.001, 1, 10, 100, 100])
  def test_layer_norm(self, input_factor):
    eps = 1e-5
    gamma = 1.0
    beta = 1.0

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], jnp.float32),
    )
    def layer_norm(x_ref, o_ref):
      x_mean = jnp.mean(x_ref[...])
      x_centered = x_ref[...] - x_mean
      o_ref[...] = (
          x_centered * jax.lax.rsqrt(jnp.mean(x_centered**2) + eps) * gamma
          + beta
      )

    def layer_norm_np(x):
      x_mean = np.mean(x)
      x_centered = x - x_mean
      return (x_centered / np.sqrt(np.mean(x_centered**2) + eps) * gamma) + beta

    # Ones are always fully precise
    x = jnp.ones((256,)).astype(jnp.float32) * input_factor
    np.testing.assert_allclose(layer_norm(x), layer_norm_np(x))

    # random (and anything else is not)
    x = (
        jax.random.uniform(jax.random.key(42), shape=(256,), dtype=jnp.float32)
        * input_factor
    )
    np.testing.assert_allclose(layer_norm(x), layer_norm_np(x), rtol=6e-5)

  def test_print(self):

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], jnp.float32),
    )
    def kernel(x_ref, o_ref):
      del x_ref, o_ref
      pl.debug_print("It works!")

    x = jnp.arange(256).astype(jnp.float32)
    with self.capture_stdout() as output:
      jax.block_until_ready(kernel(x))
    self.assertEqual(output(), "It works!\n")

  def test_print_wgmma_tiled_layout(self):
    # TODO(b/462499936): Re-enable when test passes on MIG partition.
    if jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("Test fails when run on MIG partition")

    # The default printf buffer on some smaller GPUs (e.g. Thor) only has space for
    # 4096 threads to printf (short) messages. Keep this shape below that.
    shape = (128, 32)
    size = math.prod(shape)

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(shape, jnp.float32),
        in_specs=[
            plgpu.BlockSpec(
                transforms= self.default_transforms(dtype=jnp.float32),
            )
        ],
    )
    def kernel(x_ref, o_ref):
      del o_ref  # Unused.
      pl.debug_print("prefix {}", x_ref[...])

    x = jnp.arange(size, dtype=jnp.float32).reshape(shape)
    with self.capture_stdout() as get_output:
      jax.block_until_ready(kernel(x))

    output = get_output()
    results = re.findall(r"prefix \[(\d+), (\d+)\]: (\d+).?\d*", output)
    self.assertLen(results, size, output)
    for i, j, v in results:
      i, j, v = map(int, (i, j, v))
      self.assertEqual(v, i * shape[1] + j)

  def test_print_scalar(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], jnp.int32),
    )
    def kernel(x_ref, o_ref):
      del o_ref
      pl.debug_print("x.sum() = {}", _sum_same_dtype(x_ref[...]))

    x = jnp.arange(256, dtype=jnp.int32)
    with self.capture_stdout() as output:
      jax.block_until_ready(kernel(x))

    self.assertIn(f"x.sum() = {x.sum()}", output())

  def test_print_scalar_array(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], jnp.int32),
    )
    def kernel(x_ref, o_ref):
      del o_ref
      pl.debug_print("x.sum() = {}", _sum_same_dtype(x_ref[...]) + 1)

    x = jnp.arange(256, dtype=jnp.int32)
    with self.capture_stdout() as output:
      jax.block_until_ready(kernel(x))

    self.assertIn(f"x.sum() = {x.sum() + 1}", output())

  def test_print_array(self):
    # The default printf buffer on some smaller GPUs (e.g. Thor) only has space for
    # 4096 threads to printf (short) messages. Keep this shape below that.
    in_shape = [2, 1, 64, 32]

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(in_shape, jnp.int32),
    )
    def kernel(x_ref, o_ref):
      del o_ref
      pl.debug_print("x: {}", x_ref[...])

    x = jnp.arange(math.prod(in_shape), dtype=jnp.int32).reshape(in_shape)
    with self.capture_stdout() as output:
      jax.block_until_ready(kernel(x))

    self.assertIn("x: [1, 0, 43, 23]: 3447\n", output())

  def test_print_layout(self):
    shape = (128,)

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(shape, jnp.bfloat16),
    )
    def kernel(x_ref, o_ref):
      del o_ref
      x = plgpu.layout_cast(x_ref[...], plgpu.Layout.WGMMA.reduce(1))
      plgpu.print_layout("x: {}", x)

    x = jnp.arange(math.prod(shape), dtype=jnp.bfloat16).reshape(shape)
    with self.capture_stdout() as output:
      jax.block_until_ready(kernel(x))

    self.assertIn("x: WGMMA_ROW\n", output())

  @parameterized.parameters(
      ((),),
      ((plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128)),),
      ((plgpu.TilingTransform((8, 64)),),),
      ((plgpu.SwizzleTransform(128),),),
  )
  def test_print_layout_smem(self, transforms):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((), jnp.bfloat16),
        scratch_types=[
            plgpu.SMEM((128, 64), jnp.bfloat16, transforms=transforms)
        ],
    )
    def kernel(o_ref, smem_ref):
      del o_ref
      plgpu.print_layout("smem: {}", smem_ref)

    with self.capture_stdout() as output:
      jax.jit(kernel).lower()
    transforms_str = f"({', '.join(str(t) for t in transforms)})"
    self.assertIn(f"smem: {transforms_str}\n", output())

  def test_print_layout_dynamic_slice(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((), jnp.bfloat16),
        scratch_types=[plgpu.SMEM((128, 64), jnp.bfloat16)],
    )
    def kernel(idx_ref, o_ref, smem_ref):
      del o_ref
      idx = idx_ref[...]
      plgpu.print_layout("smem: {}", smem_ref.at[pl.ds(idx, 64), :32])

    with self.capture_stdout() as output:
      jax.jit(kernel).lower(jnp.int32(0))
    self.assertIn("smem: ()\n", output())

  @parameterized.parameters(False, True)
  def test_fp8_relayout(self, from_narrow):
    shape = (128, 32)
    dtype = jnp.float8_e4m3fn
    if from_narrow:
      from_, to = 4, 8
    else:
      from_, to = 8, 4

    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct(shape, dtype)
    )
    def kernel(x_ref, o_ref):
      x = plgpu.load(
          x_ref,
          layout=plgpu.Layout.TCGEN05_TMEM_NATIVE(from_),
          optimized=False
      )
      plgpu.store(
          o_ref,
          plgpu.layout_cast(x, plgpu.Layout.TCGEN05_TMEM_NATIVE(to)),
          optimized=False,
      )

    x = jax.random.normal(jax.random.key(10), shape).astype(dtype)
    np.testing.assert_array_equal(kernel(x), x)

  @parameterized.parameters(
          (plgpu.TilingTransform((1, 32)), plgpu.SwizzleTransform(128)),
          (plgpu.TilingTransform((8, 32)), plgpu.SwizzleTransform(128)),
          (),
  )
  def test_get_swap_with_transforms(self, *transforms):
    shape = (128, 128)

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, jnp.int32),
        scratch_types=[
            plgpu.SMEM(shape, jnp.int32, transforms=tuple(transforms)),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref, o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(x_ref, scratch_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      scratch_ref[...] = scratch_ref[...] * 2
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(scratch_ref, o_ref)
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(math.prod(shape), dtype=jnp.int32).reshape(shape)
    np.testing.assert_array_equal(kernel(x), x * 2)

  @parameterized.parameters(False, True)
  def test_get_swap_with_swizzle_transform(self, transposed):
    dtype = jnp.bfloat16
    swizzle = 128
    swizzle_elems = swizzle // jnp.dtype(dtype).itemsize
    shape = (swizzle_elems, swizzle_elems)
    transforms = (plgpu.SwizzleTransform(swizzle),)

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, dtype),
        scratch_types=[
            plgpu.SMEM(shape, dtype, transforms=transforms),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref, o_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(x_ref, scratch_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      if transposed:
        tmp = plgpu.load(scratch_ref.T, layout=plgpu.Layout.WGMMA_TRANSPOSED)
        scratch_ref.T[...] = tmp * 2
      else:
        tmp = plgpu.load(scratch_ref, layout=plgpu.Layout.WGMMA)
        scratch_ref[...] = tmp * 2
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(scratch_ref, o_ref)
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    np.testing.assert_array_equal(kernel(x), x * 2)

  def test_swap_scalar_constant(self):

    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.int32))
    def kernel(o_ref):
      o_ref[...] = jnp.array(42)

    np.testing.assert_array_equal(kernel(), jnp.array(42, jnp.int32))

  def test_swizzled_dimension_dynamic_slice(self):
    shape = (64, 64)
    dtype = jnp.int32
    swizzle = 128
    swizzle_elems = swizzle // jnp.dtype(dtype).itemsize
    transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)

    @functools.partial(
        self.pallas_call,
        in_specs=[
            plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=transforms),
            plgpu.BlockSpec(memory_space=plgpu.GMEM),
        ],
        out_specs=plgpu.BlockSpec(memory_space=plgpu.GMEM),
        out_shape=jax.ShapeDtypeStruct((64, swizzle_elems), dtype),
    )
    def kernel(src_ref, offset_ref, dst_ref):
      dyn_idx = offset_ref[...] * swizzle_elems
      sliced_ref = src_ref.at[:, pl.ds(dyn_idx, swizzle_elems)]
      plgpu.store(dst_ref, sliced_ref[...], optimized=False)

    src = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    offset = np.array(1, dtype=np.int32)
    expected = src[:, swizzle_elems:swizzle_elems*2]
    np.testing.assert_array_equal(kernel(src, offset), expected)

  def test_swizzled_dimension_dynamic_slice_raises_for_non_divisible_offset(
      self,
  ):
    # TODO(allanrenucci): `UntilingTransform.commute_ndindexer` does not do
    # divisibility checks on dynamic offsets.
    if self.LOWERING_SEMANTICS == mgpu.LoweringSemantics.Lane:
      self.skipTest("Divisibility check missing.")
    shape = (64, 64)
    dtype = jnp.int32
    swizzle = 128
    swizzle_elems = swizzle // jnp.dtype(dtype).itemsize
    transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)

    @functools.partial(
        self.pallas_call,
        in_specs=[
            plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=transforms),
            plgpu.BlockSpec(memory_space=plgpu.GMEM),
        ],
        out_specs=plgpu.BlockSpec(memory_space=plgpu.GMEM),
        out_shape=jax.ShapeDtypeStruct((64, swizzle_elems), dtype),
    )
    def kernel(src_ref, offset_ref, dst_ref):
      dyn_idx = offset_ref[...]
      sliced_ref = src_ref.at[:, pl.ds(dyn_idx, swizzle_elems)]
      dst_ref[...] = sliced_ref[...]

    src = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    offset = np.array(1, dtype=np.int32)
    with self.assertRaisesRegex(
        ValueError, "Failed to infer a possible set of layouts."
    ):
      jax.block_until_ready(kernel(src, offset))

  def test_check(self):
    self.enter_context(config.jax_pallas_enable_debug_checks(True))

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], jnp.int32),
    )
    def kernel(x_ref, o_ref):
      pl.debug_check(_sum_same_dtype(x_ref[...]) > 0, "x.sum() is negative")
      o_ref[...] = x_ref[...]

    x = jnp.arange(256, dtype=jnp.int32)
    np.testing.assert_array_equal(kernel(x), x)

  def test_scalar_load_store(self):

    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.int32))
    def kernel(src_ref, dst_ref):
      dst_ref[...] = plgpu.load(src_ref, layout=plgpu.Layout.WG_SPLAT)

    src = jnp.array(42, dtype=jnp.int32)
    np.testing.assert_array_equal(kernel(src), src)

  def test_scalar_broadcast(self):

    @self.kernel(out_type=jax.ShapeDtypeStruct((128,), jnp.int32))
    def kernel(o_ref):
      o_ref[...] = jnp.broadcast_to(jnp.array(42, dtype=jnp.int32), (128,))

    np.testing.assert_array_equal(kernel(), jnp.full((128,), 42, dtype=jnp.int32))

  def test_run_scoped(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
    )
    def kernel(x_ref, o_ref):
      def body(tmp_ref):
        self.assertEqual(tmp_ref.shape, (8, 128))
        tmp_ref[...] = x_ref[...] + 1.0
        return tmp_ref[...]

      tmp = pl.run_scoped(body, plgpu.SMEM((8, 128), jnp.float32))
      self.assertEqual(tmp.shape, (8, 128))
      o_ref[...] = tmp

    x = np.ones((8, 128), jnp.float32)
    np.testing.assert_array_equal(kernel(x), x + 1.0)

  def test_run_scoped_in_cond(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], jnp.int32),
        in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
        out_specs=pl.BlockSpec(memory_space=plgpu.SMEM),
    )
    def kernel(x_ref_gmem, o_ref):
      def scoped_kernel(barrier_ref):
        plgpu.copy_gmem_to_smem(x_ref_gmem, o_ref, barrier_ref)
        plgpu.barrier_wait(barrier_ref)

      def branch():
        pl.run_scoped(scoped_kernel, plgpu.Barrier())

      jax.lax.cond(x_ref_gmem[0] % 2 == 0, branch, branch)

    x = jnp.full((256,), 1234, dtype=jnp.int32)
    np.testing.assert_array_equal(kernel(x), x)

  def test_program_id(self):
    @functools.partial(
        self.pallas_call,
        in_specs=(),
        out_specs=pl.BlockSpec((128,), lambda *i: i),
        out_shape=jax.ShapeDtypeStruct([128 * 2], jnp.int32),
        grid=2,
    )
    def kernel(o_ref):
      o_ref[...] = jnp.full(o_ref.shape, pl.program_id(0))

    np.testing.assert_array_equal(
        kernel(),
        jnp.array([0] * 128 + [1] * 128, dtype=jnp.int32),
    )

  def test_program_id_in_squashed_grid(self):
    # Tests whether a grid with >3 logical dimensions is correctly squashed to
    # 3 CUDA grid dimensions.
    grid = (2, 3, 4, 5)
    @functools.partial(
        self.pallas_call,
        in_specs=(),
        out_specs=pl.BlockSpec((1,) * len(grid) + (128,), lambda *i: (*i, 0)),
        out_shape=jax.ShapeDtypeStruct([*grid, 128], jnp.int32),
        grid=grid,
    )
    def kernel(o_ref):
      mult = 1
      idx = 0
      for axis in range(len(grid)-1, -1, -1):
        idx += pl.program_id(axis) * mult
        mult *= pl.num_programs(axis)
      o_ref[...] = jnp.full(o_ref.shape, idx)

    np.testing.assert_array_equal(
        kernel()[:, :, :, :, 0],
        jnp.arange(math.prod(grid), dtype=jnp.int32).reshape(*grid)
    )

  @parameterized.parameters(
      ((2, 3), ("a", "b"), (), ()),
      ((2, 3), ("a", "b"), (2,), ("x",)),
      ((2, 3, 4), ("a", "b", "c"), (), ()),
      ((2, 3, 4), ("a", "b", "c"), (2,), ("x",)),
      ((2, 3, 4), ("a", "b", "c"), (2, 3), ("x", "y")),
      ((2, 3, 4, 5), ("a", "b", "c", "d"), (), ()),
      ((2, 3, 4, 5), ("a", "b", "c", "d"), (2,), ("x",)),
  )
  def test_axis_indices_in_grid(self, grid, grid_names, cluster, cluster_names):

    @self.kernel(
        out_type=[
            jax.ShapeDtypeStruct([*cluster, *grid, 128], jnp.int32),
            jax.ShapeDtypeStruct([*cluster, *grid, 128], jnp.int32),
        ],
        grid=grid,
        grid_names=grid_names,
        cluster=cluster,
        cluster_names=cluster_names,
    )
    def kernel(out1_ref, out2_ref):
      pallas_grid_idx = lax.axis_index(grid_names)
      cuda_grid_idx = _get_linearized_cuda_grid_index()

      out_indices = [lax.axis_index(ax) for ax in (*cluster_names, *grid_names)]
      out1_ref[*out_indices] = jnp.full((128,), pallas_grid_idx)
      out2_ref[*out_indices] = jnp.full((128,), cuda_grid_idx)
    out1, out2 = kernel()

    out_per_cta = jnp.arange(math.prod(grid), dtype=jnp.int32).reshape(grid)
    out1_ref = jnp.broadcast_to(out_per_cta[..., None], (*cluster, *grid, 128))
    np.testing.assert_array_equal(out1, out1_ref)

    padded_cluster = (1,) * (len(grid) - len(cluster)) + cluster
    scaled_grid = tuple(g * c for g, c in zip(grid, padded_cluster))
    original = jnp.arange(math.prod(scaled_grid), dtype=jnp.int32).reshape(
        scaled_grid
    )

    # Untile the scaled grid to get the per-cluster grid.
    interleaved_shape = tuple(val for pair in zip(grid, padded_cluster) for val in pair)
    perm = tuple(range(1, 2 * len(grid), 2)) + tuple(range(0, 2 * len(grid), 2))

    out2_ref = original.reshape(interleaved_shape).transpose(perm).squeeze()
    out2_ref = jnp.broadcast_to(out2_ref[..., None], out2_ref.shape + (128,))

    np.testing.assert_array_equal(out2, out2_ref)

  def test_program_id_in_block_spec(self):
    @functools.partial(
        self.pallas_call,
        in_specs=(pl.BlockSpec((2, 128), lambda i: (pl.program_id(0), i)),),
        out_specs=pl.BlockSpec((2, 128), lambda i: (pl.program_id(0), i)),
        out_shape=jax.ShapeDtypeStruct([2, 128], jnp.int32),
        grid=2,
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = x_ref[...]

    x = jnp.arange(2 * 128, dtype=jnp.int32).reshape([2, 128])
    np.testing.assert_array_equal(kernel(x), x)

  @parameterized.parameters(
      (0, plgpu.OOBFillMode.PROMISE_IN_BOUNDS),
      (80, plgpu.OOBFillMode.PROMISE_IN_BOUNDS),
      (0, plgpu.OOBFillMode.ZEROS),
      (112, plgpu.OOBFillMode.ZEROS),
      (192, plgpu.OOBFillMode.ZEROS),
      (208, plgpu.OOBFillMode.ZEROS),
      (0, plgpu.OOBFillMode.UNDEFINED),
      (112, plgpu.OOBFillMode.UNDEFINED),
      (192, plgpu.OOBFillMode.UNDEFINED),
      (208, plgpu.OOBFillMode.UNDEFINED),
  )
  def test_out_of_bounds_tma_reads(self, start_offset, oob_mode):
    dtype = jnp.int8
    gmem_size = 208
    smem_size = 128

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((smem_size,), dtype),
        scratch_types=[plgpu.SMEM((smem_size,), dtype), plgpu.Barrier()],
    )
    def kernel(start_idx, x_ref, o_ref, scratch_ref, barrier):
      # Initialize SMEM with -1s to indicate OOB values.
      scratch_ref[...] = jnp.full(scratch_ref.shape, -1, dtype)
      plgpu.commit_smem()

      idx = pl.ds(start_idx[...], smem_size)
      plgpu.copy_gmem_to_smem(
          x_ref.at[idx], scratch_ref, barrier, oob_mode=oob_mode
      )
      plgpu.barrier_wait(barrier)
      o_ref[...] = scratch_ref[...]

    x = jnp.arange(gmem_size, dtype=dtype)
    start_idx_val = jnp.array(start_offset, dtype=jnp.int32)

    out = kernel(start_idx_val, x)

    valid_slice = x[start_offset : start_offset + smem_size]
    np.testing.assert_array_equal(out[:valid_slice.size], valid_slice)

    if oob_mode == plgpu.OOBFillMode.ZEROS:
      np.testing.assert_array_equal(out[valid_slice.size:], 0)

  def test_num_programs(self):
    @functools.partial(
        self.pallas_call,
        in_specs=(),
        out_specs=pl.BlockSpec((128,), lambda *i: i),
        out_shape=jax.ShapeDtypeStruct([128 * 2], jnp.int32),
        grid=2,
    )
    def kernel(o_ref):
      o_ref[...] = jnp.full(o_ref.shape, pl.num_programs(0), o_ref.dtype)

    np.testing.assert_array_equal(
        kernel(),
        jnp.full([256], 2, dtype=jnp.int32),
    )

  def test_swizzled_blockspec_shapes(self):
    spec = plgpu.BlockSpec(
        (128, 64),
        lambda *i: i,
        transforms=self.default_transforms(dtype=jnp.float16),
    )
    @functools.partial(
        self.pallas_call,
        in_specs=[spec],
        out_specs=spec,
        out_shape=jax.ShapeDtypeStruct((128, 128), jnp.float16),
        grid=(2, 2),
    )
    def kernel(x_ref, o_ref):
      assert x_ref.shape == (128, 64), x_ref.shape
      o_ref[...] = x_ref[...]

    x = jnp.arange(128 * 128).astype(jnp.float16).reshape(128, 128)
    np.testing.assert_array_equal(kernel(x), x)

  @parameterized.product(force_while=[False, True])
  def test_fori_loop_array(self, force_while):
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([256], jnp.int32)
    )
    def kernel(x_ref, o_ref):
      # Equivalent to x_ref[...] + 2 + 3.
      o_ref[...] = _fori_loop(
          force_while, 2, 4, lambda i, x: x + i, x_ref[...]
      )

    x = jnp.arange(256, dtype=jnp.int32)
    np.testing.assert_array_equal(kernel(x), x + 2 + 3)

  @parameterized.product(unroll=[1, 2, 4])
  def test_fori_loop_array_unrolled(self, unroll):
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([256], jnp.int32)
    )
    def kernel(x_ref, o_ref):
      # Equivalent to x_ref[...] + 2 + 3 + 4 + 5.
      o_ref[...] = lax.fori_loop(
          2, 6, lambda i, x: x + i, x_ref[...], unroll=unroll
      )

    x = jnp.arange(256, dtype=jnp.int32)
    np.testing.assert_array_equal(kernel(x), x + 2 + 3 + 4 + 5)

  @parameterized.product(force_while=[False, True])
  def test_fori_loop_scalar(self, force_while):
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([256], jnp.int32)
    )
    def kernel(o_ref):
      # Equivalent to 2 + 3.
      o_ref[...] = jax.lax.broadcast(
          _fori_loop(force_while, 2, 4, lambda i, x: x + i, jnp.int32(0)),
          o_ref.shape,
      )

    np.testing.assert_array_equal(kernel(), jnp.full([256], 5, jnp.int32))

  def test_fori_loop_dynamic_bounds(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], jnp.int32),
        grid=(1,)
    )
    def kernel(o_ref):
      zero = pl.program_id(0)
      # Equivalent to 2 + 3.
      o_ref[...] = jax.lax.broadcast(
          jax.lax.fori_loop(2 + zero, 4 + zero, lambda i, x: x + i, 0), o_ref.shape
      )

    np.testing.assert_array_equal(kernel(), jnp.full([256], 5, dtype=jnp.int32))

  @parameterized.product(force_while=[False, True])
  def test_fori_loop_tuple(self, force_while):
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([256], jnp.int32)
    )
    def kernel(o_ref):
      def body(step, xs):
        return tuple(
            jax.lax.cond(step % 2 == 0, lambda x: x + 1, lambda x: x, x)
            for x in xs
        )

      # Equivalent to 3 * (0 + 1).
      o_ref[...] = jax.lax.broadcast(
          sum(_fori_loop(force_while, 2, 4, body, (jnp.int32(0),) * 3)),
          o_ref.shape,
      )

    np.testing.assert_array_equal(
        kernel(), jnp.full([256], 3 * (0 + 1), jnp.int32)
    )

  @parameterized.product(force_while=[False, True])
  def test_fori_loop_indexed_store(self, force_while):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([4, 128], jnp.float32),
    )
    def kernel(x_ref, y_ref, o_ref):
      def body(idx, _):
        o_ref[idx] = x_ref[idx] + y_ref[idx]
        return ()

      _fori_loop(force_while, 0, 4, body, ())

    x = jnp.arange(4 * 128).reshape(4, 128).astype(jnp.float32)
    y = x + 1
    np.testing.assert_array_equal(kernel(x, y), x + y)

  def test_while_loop(self):
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([128], jnp.int32)
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = jnp.zeros(o_ref.shape, dtype=jnp.int32)

      def cond(acc):
        _, last_o = acc
        return _sum_same_dtype(last_o) < 128*10

      def body(acc):
        i, _ = acc
        o_ref[...] += x_ref[i]
        return i+1, o_ref[...]

      _ = jax.lax.while_loop(cond, body, (0, o_ref[...]))

    np.testing.assert_array_equal(
        kernel(jnp.ones([128, 128], jnp.int32)), jnp.full([128], 10, jnp.int32)
    )

  def test_while_loop_layout_mismatch(self):
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([128], jnp.int32)
    )
    def kernel(o_ref):
      def cond(acc):
        return _sum_same_dtype(acc) < 128

      def body(acc):
        del acc  # Unused.
        o_ref[...] = o_ref[...]  # side-effect to prevent DCE

        # We deliberately do a cast here to trigger a layout mismatch.
        return plgpu.layout_cast(
            jnp.zeros(o_ref.shape, o_ref.dtype), plgpu.Layout.WGMMA.reduce(1)
        )
      # Cast explicitly to cause the mismatch, otherwise layout inference will
      # succeed at constructing a working program.
      strided_input = plgpu.layout_cast(
          o_ref[...], plgpu.Layout.WG_STRIDED(shape=(128,), vec_size=1)
      )
      _ = jax.lax.while_loop(cond, body, strided_input)

    if self.is_wg_semantics():
      with self.assertRaisesRegex(
          ValueError, "Failed to infer a possible set of layouts",
      ):
        kernel()
    else:
      with self.assertRaisesRegex(
          ValueError, "has layout .*, when it should be"
      ):
        kernel()

  def test_cond(self):
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([256], jnp.int32)
    )
    def kernel(x_ref, o_ref):
      jax.lax.cond(
          x_ref[0] % 2 == 0,
          lambda: pl.debug_print("acc % 2"),
          lambda: pl.debug_print("acc"),
      )
      o_ref[...] = jnp.broadcast_to(jnp.asarray(0, dtype=o_ref.dtype), o_ref.shape)

    x = jnp.full((256,), 1234, dtype=jnp.int32)
    with self.capture_stdout() as output:
      jax.block_until_ready(kernel(x))

    self.assertIn("acc % 2", output())

  def test_cond_returning_array(self):
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct([256], jnp.int32)
    )
    def kernel(x_ref, o_ref):
      acc_sum = _sum_same_dtype(x_ref[...])
      acc2, acc = jax.lax.cond(
          acc_sum % 2 == 0,
          lambda: (acc_sum * 2, x_ref[...]),
          lambda: (acc_sum, x_ref[...]),
      )
      o_ref[...] = jnp.broadcast_to(_sum_same_dtype(acc) + acc2, o_ref.shape)

    x = jnp.arange(256, dtype=jnp.int32)
    np.testing.assert_array_equal(kernel(x), jnp.broadcast_to(jnp.sum(x) * 3, [256]))

  def test_cond_constant_output(self):
    # This is a regression test for a type error during lowering, due to
    # returning a FragmentedArray on one side and a scalar on the other.
    def kernel(x_ref, o_ref):
      scale = _array_splat(1.0, ())
      scale_one = jnp.float32(1.0)

      o_ref[...] = jax.lax.cond(
          x_ref[...] < 1.0,
          lambda: scale_one,
          lambda: scale,
      )

    f = self.kernel(
        kernel, out_type=jax.ShapeDtypeStruct((), jnp.float32)
    )
    self.assertArraysEqual(f(jnp.float32(0)), jnp.float32(1.0))

  def test_cond_get_arch(self):
    shape = (64, 256)

    @self.kernel(out_type=jax.ShapeDtypeStruct(shape, jnp.bfloat16))
    def kernel(x_gmem, o_gmem):
      x = plgpu.load(x_gmem, layout=plgpu.Layout.WGMMA, optimized=False)

      def write_out1():
        plgpu.store(o_gmem, x.astype(jnp.bfloat16) + 1, optimized=False)

      def write_out2():
        plgpu.store(o_gmem, x.astype(jnp.bfloat16) + 2, optimized=False)

      jax.lax.cond(x_gmem[0, 0] != 0, write_out1, write_out2)

    x = jnp.full(shape, 7, dtype=jnp.int4)
    with mock.patch.object(mgpu.utils, "get_arch", wraps=mgpu.utils.get_arch) as mock_get_arch:
      np.testing.assert_array_equal(kernel(x), x.astype(jnp.bfloat16) + 1)
      mock_get_arch.assert_called()

  def test_tile_slicing(self):
    shape = (256, 128)
    block_spec = plgpu.BlockSpec(
        transforms=self.default_transforms(dtype=jnp.uint16)
    )
    @functools.partial(
        self.pallas_call,
        in_specs=[block_spec],
        out_specs=block_spec,
        out_shape=jax.ShapeDtypeStruct((64, 64), jnp.uint16),
    )
    def kernel(x_ref, o_ref):
      def sum_tiles(row, acc):
        row_slice = pl.ds(row * 64, 64)
        for col in range(128 // 64):
          acc += x_ref[row_slice, pl.ds(col * 64, 64)]
        return acc
      acc = plgpu.layout_cast(jnp.zeros((64, 64), jnp.uint16), plgpu.Layout.WGMMA)
      o_ref[...] = _fori_loop(False, 0, 256 // 64, sum_tiles, acc)

    x = jnp.arange(math.prod(shape), dtype=jnp.uint16).reshape(shape)
    y = x.reshape(256 // 64, 64, 128 // 64, 64).sum(axis=(0, 2), dtype=jnp.uint16)
    np.testing.assert_array_equal(kernel(x), y)

  def test_slicing(self):
    left = upper = slice(None, 64)
    right = lower = slice(64, None)
    # We rotate the four quadrants of the input clockwise.
    def rotate(src, dst):
      dst[upper, left] = src[lower, left]
      dst[upper, right] = src[upper, left]
      dst[lower, right] = src[upper, right]
      dst[lower, left] = src[lower, right]

    x = jnp.arange(128 * 128).astype(jnp.float16).reshape(128, 128)
    spec = plgpu.BlockSpec(
        transforms=self.default_transforms(dtype=jnp.float16)
    )
    f = self.pallas_call(rotate, out_shape=x, in_specs=[spec], out_specs=spec)
    expected = np.empty_like(x)
    rotate(x, expected)
    np.testing.assert_array_equal(f(x), expected)

  def test_layout_cast(self, shape=(256, 64)):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(shape, jnp.float32),
    )
    def kernel(o_ref):
      value = plgpu.layout_cast(
          jnp.full(shape, 42.0, jnp.float32), plgpu.Layout.WGMMA
      )
      plgpu.store(o_ref, value, optimized=False)

    x = jnp.full(shape, 42.0, jnp.float32)
    np.testing.assert_array_equal(kernel(), x)

  @parameterized.product(
      layouts=[
          (plgpu.Layout.WGMMA, plgpu.Layout.WGMMA_TRANSPOSED),
          (plgpu.Layout.TCGEN05, plgpu.Layout.TCGEN05_TRANSPOSED),
      ],
  )
  def test_transposed_layout(self, layouts):
    layout, transposed_layout = layouts
    dtype = jnp.dtype(jnp.float16)
    shape = (256, 192)
    transforms = self.default_transforms(dtype=dtype)
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(shape[::-1], dtype),
        out_specs=plgpu.BlockSpec(transforms=transforms),
    )
    def kernel(o_ref):
      iota = plgpu.broadcasted_iota(dtype, shape, 0, layout=layout)
      iota *= shape[1]
      iota += plgpu.broadcasted_iota(dtype, shape, 1, layout=layout)
      o_ref_t = o_ref.transpose((1, 0))
      o_ref_t[...] = plgpu.layout_cast(iota, transposed_layout)

    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape).T
    np.testing.assert_array_equal(kernel(), x)

  @parameterized.parameters(plgpu.Layout.WGMMA, plgpu.Layout.TCGEN05)
  @jtu.thread_unsafe_test()
  def test_transposed_stmatrix(self, layout):
    dtype = jnp.float16
    shape = (256, 192)
    transforms = self.default_transforms(dtype=dtype)

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(shape[::-1], dtype),
        out_specs=plgpu.BlockSpec(transforms=transforms),
    )
    def kernel(o_ref):
      iota = plgpu.broadcasted_iota(dtype, shape, 0, layout=layout)
      iota *= shape[1]
      iota += plgpu.broadcasted_iota(dtype, shape, 1, layout=layout)
      o_ref.T[...] = iota

    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape).T
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      np.testing.assert_array_equal(kernel(), x)
    self.assertIn("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16", ptx())
    self.assertNotIn("shfl.sync.bfly", ptx())

  def test_profiler(self):
    def kernel(x_ref, o_ref):
      with jax.named_scope("add"):
        with jax.named_scope("load"):
          x = x_ref[...]
        o = x + x
      with jax.named_scope("store"):
        o_ref[...] = o
    with tempfile.TemporaryDirectory() as tmpdir:
      x = jnp.arange(256).astype(jnp.float32)
      y = self.pallas_call(
          kernel,
          out_shape=jax.ShapeDtypeStruct([256], jnp.float32),
          compiler_params=plgpu.CompilerParams(
              profile_space=16, profile_dir=tmpdir
          ),
      )(x)
      jax.block_until_ready(y)
      jax.effects_barrier()
      [name] = os.listdir(tmpdir)
      with open(os.path.join(tmpdir, name)) as f:
        data = f.read()
        self.assertEqual(data.count('"name": "add"'), 2)
        self.assertEqual(data.count('"name": "load"'), 2)
        self.assertEqual(data.count('"name": "store"'), 2)
      np.testing.assert_array_equal(y, x + x)

  def test_profiler_warp_scope(self):
    def kernel(x_ref, o_ref):
      with jax.named_scope("add"):
        with jax.named_scope("load"):
          x = x_ref[...]
        o = x + x
      with jax.named_scope("store"):
        o_ref[...] = o
    with tempfile.TemporaryDirectory() as tmpdir:
      x = jnp.arange(256).astype(jnp.float32)
      y = self.pallas_call(
          kernel,
          out_shape=jax.ShapeDtypeStruct([256], jnp.float32),
          compiler_params=plgpu.CompilerParams(
              profile_space=16,
              profile_dir=tmpdir,
              profile_trace_scope=plgpu.TraceScope.WARP,
          ),
      )(x)
      jax.block_until_ready(y)
      jax.effects_barrier()
      [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": "add"'), 8)
        self.assertEqual(data.count('"name": "load"'), 8)
        self.assertEqual(data.count('"name": "store"'), 8)
      np.testing.assert_array_equal(y, x + x)

  def test_profiler_bounds_check_truncates(self):
    def kernel(x_ref, o_ref):
      with jax.named_scope("a"):
        with jax.named_scope("b"):
          with jax.named_scope("c"):
            o = x_ref[...] + 1
      o_ref[...] = o

    with tempfile.TemporaryDirectory() as tmpdir:
      x = jnp.arange(256).astype(jnp.float32)
      y = self.pallas_call(
          kernel,
          out_shape=jax.ShapeDtypeStruct([256], jnp.float32),
          compiler_params=plgpu.CompilerParams(
              profile_space=3,
              profile_dir=tmpdir,
              profile_bounds_check=True,
          ),
      )(x)
      jax.block_until_ready(y)
      jax.effects_barrier()
      [name] = os.listdir(tmpdir)
      with open(os.path.join(tmpdir, name)) as f:
        data = f.read()
      # The buffer holds 3 * 2 * 2 - 3 = 9 usable entries, enough for 4 of the
      # 6 events (B:a, B:b, B:c, E:c); the rest are dropped instead of
      # overflowing the buffer. The two ranges left open by truncation (b, a)
      # get synthetic end events, so begins and ends stay balanced.
      self.assertEqual(data.count('"ph": "B"'), 3)
      self.assertEqual(data.count('"ph": "E"'), 3)
      np.testing.assert_array_equal(y, x + 1)

  def test_profiler_computes_correct_allocation_size(self):
    def kernel(x_ref, o_ref):
      with jax.named_scope("copy"):
        o_ref[...] = x_ref[...]
    with tempfile.TemporaryDirectory() as tmpdir:
      # For each range that we record, we need 2xi32 for each of the start and
      # end of the range---i.e. 4xi32 in total per event.
      num_profiler_ranges = 200_000
      expected_allocation_size_bytes = num_profiler_ranges * 4 * 4
      with self.assertRaisesRegex(
          ValueError,
          f"exceeds available shared memory: smem_bytes={expected_allocation_size_bytes}"
      ):
        self.kernel(
            kernel,
            out_type=jax.ShapeDtypeStruct([256], jnp.float32),
            compiler_params=plgpu.CompilerParams(
                profile_space=num_profiler_ranges, profile_dir=tmpdir
            ),
        )(jnp.arange(256).astype(jnp.float32))

  @parameterized.product(
      dtypes=[
          (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_convert_type(self, dtypes):
    in_dtype, out_dtype = dtypes
    m, n = 16, 8
    out_shape = jax.ShapeDtypeStruct((m, n), out_dtype)

    @functools.partial(self.pallas_call, out_shape=out_shape)
    def convert(x_ref, y_ref):
      y_ref[...] = jax.lax.bitcast_convert_type(x_ref[...], out_shape)

    x = jnp.arange(m * n, dtype=in_dtype).reshape((m, n))
    np.testing.assert_array_equal(
        convert(x), jax.lax.bitcast_convert_type(x, out_dtype)
    )

  def test_optimization_barrier(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = lax.optimization_barrier(x_ref[...])

    x = jax.lax.iota(jnp.float32, 128)
    np.testing.assert_array_equal(kernel(x), x)

  def test_optimization_barrier_constant(self):
    @functools.partial(
        self.kernel,
        out_type=jax.ShapeDtypeStruct((), jnp.float32),
    )
    def kernel(o_ref):
      c = jnp.float32(3)
      o_ref[...] = lax.optimization_barrier(c)

    np.testing.assert_array_equal(kernel(), 3)

  def test_optimization_barrier_multiple_inputs(self):
    @functools.partial(
        self.kernel,
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
    )
    def kernel(x_ref, y_ref, o_ref):
      c = jnp.float32(5)
      x, c, y = lax.optimization_barrier([x_ref[...], c, y_ref[...]])
      o_ref[...] = x + c + y

    x = jax.lax.iota(jnp.float32, 128)
    y = jax.lax.iota(jnp.float32, 128) * 3
    np.testing.assert_array_equal(kernel(x, y), x + 5 + y)

  def test_smem_aliasing_works_basic(self):
    in_shape = (2, 256)
    transforms = (plgpu.TilingTransform((64,)),)

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([128], jnp.float32),
        in_specs=[pl.BlockSpec(in_shape)],
        out_specs=pl.BlockSpec((128,), memory_space=plgpu.GMEM),
        scratch_shapes=[
            plgpu.RefUnion(
                # Note: this test exposes internals that we don't particularly
                # want to phold for the sake of testing the functionality of the
                # API. It's expected that this test might end up breaking in the
                # future, e.g. if we decide to change our alignment requirements
                # on SMEM refs---and that's OK. Users should explicitly NOT rely
                # on this exact behaviour.
                #
                # Use a value larger than the number of bytes used for SMEM
                # alignment (1024) in order to make sure that the second ref
                # in the second group aliases the single ref in the first group.
                plgpu.SMEM(in_shape, jnp.float32),
                [
                    plgpu.SMEM((256,), jnp.bfloat16),
                    # Add an arbitrary level of nesting to make sure that we
                    # support PyTrees.
                    [plgpu.SMEM((128,), jnp.float32, transforms=transforms)]
                ],
            )
        ],
    )
    def kernel(x_ref, o_ref128, aliased_ref):
      smem_ref256, [_, [smem_ref128]] = aliased_ref
      # Ensure that extraction via index works the same as unfolding.
      smem_ref128_2 = aliased_ref[1][1][0]
      self.assertIsInstance(smem_ref128, state_types.TransformedRef)
      self.assertIsInstance(smem_ref128_2, state_types.TransformedRef)
      self.assertIs(smem_ref128.ref, smem_ref128_2.ref)
      self.assertEqual(smem_ref128.transforms, smem_ref128_2.transforms)

      # Ensure that the transforms provided in the scratch shapes have been
      # passed correctly.
      extract_alias_transform, tile_transform = smem_ref128.transforms
      self.assertIsInstance(extract_alias_transform, gpu_core.ExtractAliasedRef)
      self.assertIsInstance(tile_transform, gpu_core.UntilingTransform)

      smem_ref256[...] = x_ref[...] + 1
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(smem_ref128, o_ref128)

    x = jnp.arange(512).astype(jnp.float32)
    np.testing.assert_array_equal(
        kernel(x.reshape(in_shape)).reshape((128,)), x[256 : 256 + 128] + 1
    )

  def test_smem_aliasing_refunions_in_wg_semantics_can_have_distinct_transforms(self):
    # A smoke test to ensure that we have different IDs for the two RefUnion
    # allocations in WG semantics.
    if not self.is_wg_semantics():
      self.skipTest("This test is only for WG semantics.")

    in_shape = (2, 256)
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([128], jnp.float32),
        in_specs=[pl.BlockSpec(in_shape)],
        out_specs=pl.BlockSpec((128,), memory_space=plgpu.GMEM),
        scratch_shapes=[
            plgpu.RefUnion(
                plgpu.SMEM(in_shape, jnp.float32),
                [
                    plgpu.SMEM((256,), jnp.bfloat16),
                    [plgpu.SMEM((128,), jnp.float32)]
                ],
            ),
            plgpu.RefUnion(
                plgpu.SMEM(in_shape, jnp.float32),
                [
                    plgpu.SMEM((256,), jnp.bfloat16),
                    [plgpu.SMEM((128,), jnp.float32)]
                ],
            ),
        ],
    )
    def kernel(x_ref, o_ref128, aliased_ref_1, aliased_ref_2):
      def _with_transform(ref, transforms):
        @plgpu.inline_mgpu(arg_types=(plgpu.RefType(transforms),))
        def f(ctx, smem_ref):
          del ctx, smem_ref
        f(ref)

      ref_1, [_, [smem_ref128]] = aliased_ref_1
      _with_transform(ref_1, ())
      ref_2, [_, [_]] = aliased_ref_2
      _with_transform(
          ref_2, (plgpu.TilingTransform((1, 32)), plgpu.SwizzleTransform(128))
      )

      ref_1[...] = x_ref[...] + 1
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(smem_ref128, o_ref128)

    x = jnp.arange(512).astype(jnp.float32)
    np.testing.assert_array_equal(
        kernel(x.reshape(in_shape)).reshape((128,)), x[256 : 256 + 128] + 1
    )

  def test_smem_aliasing_works_with_subbyte_dtypes(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([256], jnp.uint4),
        in_specs=[pl.BlockSpec((128,))],
        out_specs=pl.BlockSpec((256,), memory_space=plgpu.GMEM),
        scratch_shapes=[
            plgpu.RefUnion(
                # Note: this test exposes internals that we don't particularly
                # want to phold for the sake of testing the functionality of the
                # API. It's expected that this test might end up breaking in the
                # future, e.g. if we decide to change our alignment requirements
                # on SMEM refs---and that's OK. Users should explicitly NOT rely
                # on this exact behaviour.
                #
                # This allocation scheme is a bit complicated, but serves to
                # test that
                #   1. Refs are aligned correctly (currently to 1024 bytes);
                #   2. (u)int4 references are not allocated more than 1 byte per
                #      2 elements.
                # The first group of refs serves to create two allocations, each
                # aligned to 1024 bytes. The second group serves to create two
                # allocations where the first one is exactly 1024 bytes,
                # assuming 1 byte per 2 uint4 elements. As a result, if our
                # implementation is correct, the second allocation of the second
                # group should exactly alias the second allocation of the first
                # group.
                [
                    plgpu.SMEM((128,), jnp.int8),
                    plgpu.SMEM((128,), jnp.int8),
                ],
                [plgpu.SMEM((2048,), jnp.uint4), plgpu.SMEM((256,), jnp.uint4)],
            )
        ],
    )
    def kernel(x_ref, o_refi4, aliased_ref):
      [_, smem_refi8], [_, smem_refi4] = aliased_ref
      smem_refi8[...] = x_ref[...]
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(smem_refi4, o_refi4)

    def unpack_i4_as_i8(x):
      x = x.reshape((128, 1))
      x_high = x >> 4
      x_low = x & 0xF
      return jnp.concatenate([x_low, x_high], axis=-1).reshape((256,))

    x = jnp.arange(128).astype(jnp.int8)
    test_as_i8 = jax.lax.convert_element_type(kernel(x), new_dtype=jnp.int8)
    np.testing.assert_array_equal(test_as_i8[:256], unpack_i4_as_i8(x))

  def test_smem_aliasing_works_for_quantization(self):
    shape = (64, 256)
    large_ty, small_ty = jnp.bfloat16, jnp.uint4

    def kernel(x_gmem, o_gmem):
      return pl.run_scoped(
          functools.partial(scoped_kernel, x_gmem, o_gmem),
          plgpu.RefUnion(
              plgpu.SMEM(shape, large_ty, transforms=self.default_transforms(dtype=large_ty)),
              plgpu.SMEM(shape, small_ty, transforms=self.default_transforms(dtype=small_ty, swizzle=32))
          ),
          plgpu.Barrier(num_barriers=1),
      )

    def scoped_kernel(x_gmem, o_gmem, aliased_ref, barrier):
      ref_large_ty, ref_small_ty = aliased_ref
      plgpu.copy_gmem_to_smem(x_gmem, ref_small_ty, barrier=barrier)
      plgpu.barrier_wait(barrier)
      ref_large_ty[...] = ref_small_ty[...].astype(ref_large_ty.dtype) * 3
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(ref_large_ty, o_gmem)
      plgpu.wait_smem_to_gmem(0, wait_read_only=True)

    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct(shape, large_ty),
    )
    key = jax.random.key(42)
    x = jax.random.randint(key, shape, 0, 4).astype(small_ty)
    expected = x * 3
    np.testing.assert_array_equal(kernel_fn(x), expected)

  def test_ref_union_caching(self):
    @self.kernel(
      out_type=(
        jax.ShapeDtypeStruct((64, 32), jnp.float32),
        jax.ShapeDtypeStruct((64, 64), jnp.float16)
      ),
    )
    def kernel(o_gmem_f32, o_gmem_f16):
      def _scope(aliased_ref):
        [smem] = aliased_ref
        smem[...] = jnp.ones_like(smem)
        plgpu.commit_smem()
        if smem.dtype == jnp.float32:
          plgpu.copy_smem_to_gmem(smem, o_gmem_f32)
        else:
          plgpu.copy_smem_to_gmem(smem, o_gmem_f16)
        plgpu.wait_smem_to_gmem(0)

      pl.run_scoped(
          _scope,
          plgpu.RefUnion(plgpu.SMEM((64, 32), jnp.float32)),
      )
      pl.run_scoped(
          _scope,
          plgpu.RefUnion(plgpu.SMEM((64, 64), jnp.float16)),
      )
    of32, of16 = kernel()
    np.testing.assert_array_equal(of32, 1.0)
    np.testing.assert_array_equal(of16, 1.0)

  @parameterized.parameters(True, False)
  def test_smem_aliased_refs_collision(self, use_ref_union):
    # This is a regression test for layout inference of aliased refs.
    # Aliased refs with the same shape and dtype should be treated as distinct
    # references. I.e. they can have different transforms.
    shape, dtype = (64, 32), jnp.float32

    @self.kernel(
        out_type=(
            jax.ShapeDtypeStruct(shape, dtype),
            jax.ShapeDtypeStruct(shape, dtype),
        ),
    )
    def kernel(src_ref, o_ref0, o_ref1):
      def _scope(o_ref, aliased_ref):
        if use_ref_union:
          [smem] = aliased_ref
        else:
          smem = aliased_ref
        smem[...] = plgpu.load(
            src_ref, layout=plgpu.Layout.WGMMA, optimized=False
        )
        plgpu.commit_smem()
        plgpu.copy_smem_to_gmem(smem, o_ref)
        plgpu.wait_smem_to_gmem(0)

      transforms = self.default_transforms(dtype=dtype, swizzle=32)
      alloc = plgpu.SMEM(shape, dtype, transforms=transforms)
      if use_ref_union:
        alloc = plgpu.RefUnion(alloc)
      pl.run_scoped(functools.partial(_scope, o_ref0), alloc)
      transforms = self.default_transforms(dtype=dtype, swizzle=128)
      alloc = plgpu.SMEM(shape, dtype, transforms=transforms)
      if use_ref_union:
        alloc = plgpu.RefUnion(alloc)
      pl.run_scoped(functools.partial(_scope, o_ref1), alloc)

    src = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    out0, out1 = kernel(src)
    np.testing.assert_array_equal(out0, src)
    np.testing.assert_array_equal(out1, src)

  @parameterized.parameters(True, False)
  def test_tmem_aliased_refs_collision(self, use_ref_union):
    # This is a regression test for layout inference of aliased refs.
    # Aliased refs with the same shape and dtype should be treated as distinct
    # references. I.e. they can have different TMEM layouts.
    shape, dtype = (128, 128), jnp.float16

    @self.kernel(out_type=jax.ShapeDtypeStruct((), dtype))
    def kernel(o_ref):
      del o_ref
      def _scope(fmt_str, aliased_ref):
        if use_ref_union:
          [tmem_ref] = aliased_ref
        else:
          tmem_ref = aliased_ref
        plgpu.print_layout(fmt_str, tmem_ref)

      alloc = plgpu.TMEM(shape, dtype, packed=False)
      if use_ref_union:
        alloc = plgpu.RefUnion(alloc)
      pl.run_scoped(functools.partial(_scope, "ref0: {}"), alloc)
      alloc = plgpu.TMEM(shape, dtype, packed=True)
      if use_ref_union:
        alloc = plgpu.RefUnion(alloc)
      pl.run_scoped(functools.partial(_scope, "ref1: {}"), alloc)

    with self.capture_stdout() as output:
      jax.jit(kernel).lower()
    self.assertIn("ref0: TMEM_DEFAULT(packing=1)", output())
    self.assertIn("ref1: TMEM_DEFAULT(packing=2)", output())

  @parameterized.product(
      small_ty=[jnp.int4, jnp.uint4],
      large_ty=[jnp.float8_e4m3fn, jnp.float16, jnp.bfloat16, jnp.float32, jnp.int8, jnp.int16, jnp.int32],
  )
  def test_subbyte_upcast(self, small_ty, large_ty):
    shape = (64, 256)

    @self.kernel(out_type=jax.ShapeDtypeStruct(shape, large_ty))
    def kernel(x_gmem, o_gmem):
      x = plgpu.load(x_gmem, layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.store(o_gmem, x.astype(large_ty), optimized=False)

    key = jax.random.key(42)
    iinfo = dtypes.iinfo(small_ty)
    x = jax.random.randint(key, shape, minval=iinfo.min, maxval=iinfo.max).astype(small_ty)
    expected = x.astype(large_ty)
    np.testing.assert_array_equal(kernel(x), expected)

  def test_assigning_to_ref_union_raises(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([128], jnp.float32),
        in_specs=[pl.BlockSpec((128,))],
        out_specs=pl.BlockSpec((128,), memory_space=plgpu.GMEM),
        scratch_shapes=[plgpu.RefUnion(plgpu.SMEM((128,), jnp.float32))],
    )
    def kernel(x_ref, o_ref128, aliased_ref):
      aliased_ref[...] = x_ref[...] + 1
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(aliased_ref, o_ref128)

    with self.assertRaisesRegex(ValueError, "can't be assigned to"):
      kernel(jnp.arange(128).astype(jnp.float32))

  def test_loading_from_ref_union_works(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([128], jnp.float32),
        in_specs=[pl.BlockSpec((128,))] * 2,
        out_specs=pl.BlockSpec((128,), memory_space=plgpu.GMEM),
        scratch_shapes=[plgpu.RefUnion(plgpu.SMEM((128,), jnp.float32)),
                        plgpu.SMEM((128,), jnp.float32)],
    )
    def kernel(x_ref, y_ref, o_ref128, ref_union, o_smem):
      [aliased_ref] = ref_union
      aliased_ref[...] = x_ref[...]
      plgpu.commit_smem()
      load_ref = lambda r: plgpu.load(r, layout=plgpu.Layout.TCGEN05.reduce(1))
      # This is a regression test for b/423697560, where we used to fail to
      # transform the dtype correctly when processing an aliased ref.
      o_smem[...] = load_ref(aliased_ref) + load_ref(y_ref)
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(o_smem, o_ref128)

    x, y = (jnp.arange(128).astype(jnp.float32) for _ in range(2))
    np.testing.assert_array_equal(kernel(x, y), x + y)

  @parameterized.product(
      sm_steps=(1, 2, 3),
      dynamic_sm_steps=(False, True),
  )
  def test_nd_loop_with_carry(self, sm_steps, dynamic_sm_steps):

    @self.kernel(
        out_type=(
            jax.ShapeDtypeStruct((sm_steps, 132, 128), jnp.int32),
            jax.ShapeDtypeStruct((132,), jnp.int32),
        ),
        grid=(132,),
        grid_names=("sm",),
    )
    def kernel(sm_steps_ref, o_ref, steps_ref):
      sm_steps_val = sm_steps_ref[0] if dynamic_sm_steps else sm_steps
      def body(loop_info, carry):
        idx = loop_info.index
        assert len(idx) == 3
        flat_idx = 0
        for i, d in zip(idx, (sm_steps_val, 4, 33)):
          flat_idx = flat_idx * d + jnp.clip(i, 0, d - 1)
        flat_idx = jnp.int32(flat_idx)
        sm_step = lax.div(
            flat_idx, lax.convert_element_type(lax.axis_size("sm"), jnp.int32)
        )
        o_ref[sm_step, lax.axis_index("sm")] = lax.broadcast(
            flat_idx, o_ref.shape[-1:]
        )
        return carry + 1

      steps_ref[lax.axis_index("sm")] = plgpu.nd_loop(
          (sm_steps_val, 4, 33), collective_axes="sm", init_carry=0
      )(body)

    result, steps = kernel(jnp.array([sm_steps], dtype=jnp.int32))
    for sm_step in range(sm_steps):
      np.testing.assert_array_equal(steps, jnp.full((132,), sm_steps))

      np.testing.assert_array_equal(
          result[sm_step],
          jnp.tile(
              (132 * sm_step + jnp.arange(132))[:, None],
              128,
          ),
      )

  @parameterized.product(
      sm_steps=(1, 2, 3),
      tiling=(None, 1, 2, 4),
      dynamic_sm_steps=(False, True),
  )
  def test_nd_loop(
      self, sm_steps: int, tiling: int | None, dynamic_sm_steps: bool
  ):
    if tiling is not None:
      tiling = (sm_steps, tiling, 33)

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((sm_steps, 132, 128), jnp.int32),
        grid=(132,),
        grid_names=("sm",),
    )
    def kernel(sm_steps_ref, o_ref):
      sm_steps_val = sm_steps_ref[0] if dynamic_sm_steps else sm_steps

      @plgpu.nd_loop((sm_steps_val, 4, 33), tiling=tiling, collective_axes="sm")
      def _(loop_info):
        idx = loop_info.index
        assert len(idx) == 3
        # We need to use `mode="clip"`, because the indices are not static.
        grid = (sm_steps_val, 4, 33)
        if tiling:
          # Reconstruct the tiled grid and index.
          tiled_grid = tuple(g // t for g, t in zip(grid, tiling))
          grid = tiled_grid + tiling
          tile_idx = tuple(
              lax.div(idx, jnp.int32(t)) for idx, t in zip(idx, tiling))
          subtile_idx = tuple(
              lax.rem(idx, jnp.int32(t)) for idx, t in zip(idx, tiling))
          idx = tile_idx + subtile_idx
        flat_idx = 0
        for i, d in zip(idx, grid):
          flat_idx = flat_idx * d + jnp.clip(i, 0, d - 1)
        sm_step = lax.div(
            flat_idx, lax.convert_element_type(lax.axis_size("sm"), jnp.int32)
        )
        o_ref[sm_step, lax.axis_index("sm")] = lax.broadcast(
            flat_idx, o_ref.shape[-1:]
        )

    result = kernel(jnp.array([sm_steps], dtype=jnp.int32))
    for sm_step in range(sm_steps):
      np.testing.assert_array_equal(
          result[sm_step],
          jnp.tile((132 * sm_step + jnp.arange(132))[:, None], 128),
      )

  def test_lowering_error_context(self):
    def body(x_ref, y_ref, barrier):
      plgpu.copy_gmem_to_smem(x_ref, y_ref, barrier)
      plgpu.barrier_wait(barrier)

    x = jnp.arange(127, dtype=jnp.int4)  # Size is not a multiple of bytes
    offending_line = "plgpu.copy_gmem_to_smem(x_ref, y_ref, barrier)"
    try:
      self.pallas_call(
          body,
          in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
          out_specs=pl.BlockSpec(memory_space=plgpu.SMEM),
          out_shape=x,
          scratch_shapes=[plgpu.Barrier()],
      )(x)
    except:
      # assertRaisesRegex raises does not let us match the traceback.
      self.assertIn(offending_line, traceback.format_exc())
    else:
      self.fail("Should have raised an exception")

  def test_lower_with_abstract_mesh(self):
    def kernel(y_ref, sem):
      plgpu.semaphore_signal_multicast(sem, collective_axes='x')
      # Wait for the multicast signal (each device gets signaled by all devices)
      pl.semaphore_wait(sem, 2)  # Wait for signals from both devices
      y_ref[...] = jnp.ones_like(y_ref)

    kernel_jax = self.pallas_call(
        kernel,
        out_specs=pl.BlockSpec(memory_space=plgpu.GMEM),
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        scratch_shapes=[plgpu.SemaphoreType.REGULAR],
    )
    abstract_mesh = jax.sharding.AbstractMesh((2,), ('x',))
    with jax.use_abstract_mesh(abstract_mesh):
      jax.jit(jax.shard_map(kernel_jax, mesh=abstract_mesh, in_specs=(),
                            out_specs=jax.P(), check_vma=False)
              ).trace().lower(lowering_platforms=('gpu',))  # doesn't crash

  @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("xyz", n) for n in range(1, 4)
    )
    for noncollective_dims in itertools.chain.from_iterable(
        itertools.combinations("xyz", 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):
    """
      1. Broadcast a GMEM slice to SMEM across collective CTAs.
      2. Send a SMEM slice from each collective CTA to reconstruct the GMEM slice.
        It's not strictly necessary to use every collective CTA, but we use them
        to test that the cluster axes are used correctly.
    """

    dtype = jnp.float16
    cluster = [1, 1, 1]
    for d in collective_dims:
      cluster["xyz".index(d)] = collective_dim_size
    for d in noncollective_dims:
      cluster["xyz".index(d)] = 2
    if math.prod(cluster) > jtu.get_cuda_nonportable_max_cluster_size():
      self.skipTest("Cluster is too big.")

    collective_size = math.prod(cluster["xyz".index(d)] for d in collective_dims)
    noncollective_size = math.prod(cluster) // collective_size

    swizzle = 128
    swizzle_elems = swizzle // jnp.dtype(dtype).itemsize
    transforms = (
        plgpu.TilingTransform((8, swizzle_elems)),
        plgpu.SwizzleTransform(swizzle),
    )
    shape = (noncollective_size, collective_size * 8, swizzle_elems)

    def body(x_gmem, out_gmem, smem, tma_barrier):
      # Compute the index in a subset of the cluster.
      def cluster_id(axes):
        idx, stride = 0, 1
        for d in sorted(axes):
          idx += lax.axis_index(d) * stride
          stride *= lax.axis_size(d)
        return idx

      noncollective_idx = cluster_id(noncollective_dims)
      collective_idx = cluster_id(collective_dims)

      plgpu.copy_gmem_to_smem(
            x_gmem.at[noncollective_idx],
            smem,
            tma_barrier,
            collective_axes=collective_dims)
      plgpu.barrier_wait(tma_barrier)

      plgpu.commit_smem()
      collective_slice = pl.ds(8 * collective_idx, 8)
      plgpu.copy_smem_to_gmem(
          smem.at[collective_slice],
          out_gmem.at[noncollective_idx, collective_slice, :],
      )
      plgpu.wait_smem_to_gmem(0)

    x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
    kernel = self.kernel(
      body,
      grid=cluster,
      grid_names=("grid_x", "grid_y", "grid_z"),
      cluster=cluster,
      cluster_names=("x", "y", "z"),
      out_type=jax.ShapeDtypeStruct(shape, dtype),
      scratch_types=(
        plgpu.SMEM(shape[1:], dtype, transforms=transforms),
        plgpu.Barrier(),
      )
    )
    np.testing.assert_array_equal(kernel(x), x)

  @parameterized.parameters(
      (plgpu.Layout.WGMMA, 0),
      (plgpu.Layout.WGMMA, 1),
      # To have some layout with vector length of 1
      (plgpu.Layout.TCGEN05_TMEM_NATIVE(1), 0),
      (plgpu.Layout.TCGEN05_TMEM_NATIVE(1), 1),
      # To have some layout with vector length > 2
      (plgpu.Layout.TCGEN05_TMEM_NATIVE(4), 0),
      (plgpu.Layout.TCGEN05_TMEM_NATIVE(4), 1),
  )
  def test_reduction_is_batch_invariant(self, layout, axis):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
    )
    def kernel(x_ref, out_ref):
      x = plgpu.load(x_ref, layout=layout, optimized=False)
      plgpu.store(out_ref, jax.lax.reduce_sum(x, axes=(axis,)), optimized=False)

    row = jax.random.uniform(jax.random.key(0), shape=(128,), dtype=jnp.float32)
    x = jnp.stack([row for _ in range(128)])
    x = x.T if axis == 0 else x
    result = kernel(x)
    expected = result[0]
    self.assertAllClose(expected, jnp.sum(row))
    self.assertArraysEqual(result, jax.lax.broadcast_in_dim(expected, (128,), ()))

  def test_reduction_fails_on_too_little_scratch_bytes_for_cross_warp_reduction(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
        compiler_params=plgpu.CompilerParams(reduction_scratch_bytes=0),
    )
    def kernel(x_ref, y_ref):
      x_val = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      y_ref[...] = jnp.sum(x_val, axis=0)

    with self.assertRaisesRegex(
        ValueError,
        r"Required reduction scratch size \(1024 bytes\) is larger than the "
        r"available scratch size \(0 bytes\)"
    ):
      kernel(jnp.zeros((128, 128), dtype=jnp.float32))

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_reduction_with_more_scratch_uses_less_synchronization(self):
    def run_kernel(x, scratch_bytes):
      def kernel(x_ref, y_ref):
        x_val = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
        y_ref[...] = jnp.sum(x_val, axis=0)
      return self.pallas_call(
          kernel,
          out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
          compiler_params=plgpu.CompilerParams(reduction_scratch_bytes=scratch_bytes)
      )(x)

    x = jax.random.uniform(jax.random.key(0), shape=(128, 128), dtype=jnp.float32)
    with jtu.set_env(MOSAIC_GPU_DUMP_SASS="1"), self.capture_stdout() as sass0:
      out0 = run_kernel(x, 1024).block_until_ready()

    with jtu.set_env(MOSAIC_GPU_DUMP_SASS="1"), self.capture_stdout() as sass1:
      out1 = run_kernel(x, 2 * 1024).block_until_ready()

    self.assertAllClose(out0, jnp.sum(x, axis=0))
    self.assertArraysEqual(out0, out1)

    syncs0 = re.findall(r"BAR.SYNC", sass0())
    syncs1 = re.findall(r"BAR.SYNC", sass1())
    self.assertLess(len(syncs1), len(syncs0))

  @parameterized.product(
      layout=(
          plgpu.Layout.WGMMA,
          plgpu.Layout.TCGEN05,
          plgpu.Layout.TCGEN05_TMEM_NATIVE,
          plgpu.Layout.TCGEN05_M64_COLLECTIVE(128),
          plgpu.Layout.TILED(  # WGMMA, but defined as a custom tiling.
              plgpu.Tiling(((64, 8), (16, 8), (8, 8), (2,))),
              warp_dims=(-7,),
              lane_dims=(-3, -2),
              vector_dim=-1,
          ),
          # To have some layout with vector length of 1
          plgpu.Layout.TCGEN05_TMEM_NATIVE(1),
          # To have some layout with vector length > 2
          plgpu.Layout.TCGEN05_TMEM_NATIVE(4),
      ),
      op=(jnp.sum, jnp.max, jnp.min),
      # TODO(apaszke): Add support for f8 (MLIR/LLVM barfs at the moment).
      dtype=(jnp.float32, jnp.float16, jnp.bfloat16, jnp.int32, jnp.uint32),
  )
  def test_reduce_with_layout(self, layout, op, dtype):
    axis = -1
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128,), dtype),
    )
    def kernel(x_ref, y_ref):
      x_val = plgpu.load(x_ref, layout=layout, optimized=False)
      plgpu.store(
          y_ref, op(x_val, axis=axis).astype(y_ref.dtype), optimized=False
      )

    if jnp.issubdtype(dtype, jnp.integer):
      x = jnp.arange(128 * 128).reshape((128, 128)).astype(dtype)
      if dtype == jnp.int32:
        x = x - x.size // 2  # include negative values.
    else:
      x = jax.random.uniform(jax.random.key(0), shape=(128, 128), dtype=dtype)
    x_result = jax.block_until_ready(kernel(x))
    np.testing.assert_allclose(x_result, op(x, axis=axis), atol=5e-5)

  @parameterized.named_parameters(
      ("negative_ilp", -1, ValueError),
      ("zero_ilp", 0, ValueError),
      ("float_ilp", 2.5, ValueError),
      ("string_ilp", "4", ValueError),
      ("bool_ilp", True, ValueError),
  )
  def test_reduction_accumulator_ilp_validation(self, ilp, expected_error):
    def make_kernel(op):
      @self.kernel(out_type=jax.ShapeDtypeStruct((128,), jnp.float32))
      def kernel(x_ref, y_ref):
        x_val = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
        plgpu.store(
            y_ref, op(x_val, axis=-1, accumulator_ilp=ilp), optimized=False
        )

      return kernel

    x = jnp.ones((128, 128), dtype=jnp.float32)
    for op in (plgpu.sum, plgpu.max, plgpu.min, plgpu.prod):
      with self.assertRaises(expected_error):
        make_kernel(op)(x)

  @parameterized.product(
      op_info=(
          (plgpu.sum, jnp.sum),
          (plgpu.max, jnp.max),
          (plgpu.min, jnp.min),
          (plgpu.prod, jnp.prod),
      ),
      ilp=(1, 2, 4),
  )
  def test_reduction_accumulator_ilp(self, op_info, ilp):
    pl_op, jnp_op = op_info
    axis = -1

    @self.kernel(out_type=jax.ShapeDtypeStruct((128,), jnp.float32))
    def kernel(x_ref, y_ref):
      x_val = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.store(
          y_ref, pl_op(x_val, axis=axis, accumulator_ilp=ilp), optimized=False
      )

    x = jax.random.uniform(jax.random.key(0), shape=(128, 128), dtype=jnp.float32)
    if jnp_op == jnp.prod:
      x = x * 0.1 + 0.95  # Avoid overflow / underflow for product reduction.
    with mock.patch.object(
        mgpu.FragmentedArray,
        "reduce",
        autospec=True,
        side_effect=mgpu.FragmentedArray.reduce,
    ) as mock_reduce:
      x_result = jax.block_until_ready(kernel(x))
      mock_reduce.assert_called()
      for call in mock_reduce.call_args_list:
        self.assertEqual(call.kwargs.get("acc_ilp"), ilp)
    np.testing.assert_allclose(x_result, jnp_op(x, axis=axis), atol=1e-4, rtol=1e-4)

  @parameterized.product(
      op_info=(
          (plgpu.sum, jnp.sum),
          (plgpu.max, jnp.max),
          (plgpu.min, jnp.min),
          (plgpu.prod, jnp.prod),
      ),
      keepdims=(True, False),
  )
  def test_reduction_keepdims(self, op_info, keepdims):
    pl_op, jnp_op = op_info
    axis = 0
    out_shape = (1, 128) if keepdims else (128,)

    @self.kernel(out_type=jax.ShapeDtypeStruct(out_shape, jnp.float32))
    def kernel(x_ref, y_ref):
      x_val = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.store(
          y_ref, pl_op(x_val, axis=axis, keepdims=keepdims), optimized=False
      )

    x = jax.random.uniform(jax.random.key(0), shape=(128, 128), dtype=jnp.float32)
    if jnp_op == jnp.prod:
      x = x * 0.1 + 0.95
    if keepdims:
      with self.assertRaises(NotImplementedError):
        kernel(x)
    else:
      x_result = jax.block_until_ready(kernel(x))
      np.testing.assert_allclose(
          x_result,
          jnp_op(x, axis=axis, keepdims=keepdims),
          atol=1e-4,
          rtol=1e-4,
      )

  @parameterized.product(
      op_info=(
          (plgpu.sum, jnp.sum),
          (plgpu.max, jnp.max),
          (plgpu.min, jnp.min),
          (plgpu.prod, jnp.prod),
      ),
      axis=(-1, 1, (-1,), (1,), -2, 0, (-2,), (0,)),
  )
  def test_reduction_axis_canonicalization(self, op_info, axis):
    pl_op, jnp_op = op_info

    @self.kernel(out_type=jax.ShapeDtypeStruct((128,), jnp.float32))
    def kernel(x_ref, y_ref):
      x_val = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.store(y_ref, pl_op(x_val, axis=axis), optimized=False)

    x = jax.random.uniform(jax.random.key(0), shape=(128, 128), dtype=jnp.float32)
    if jnp_op == jnp.prod:
      x = x * 0.1 + 0.95
    x_result = jax.block_until_ready(kernel(x))
    np.testing.assert_allclose(
        x_result, jnp_op(x, axis=axis), atol=1e-4, rtol=1e-4
    )

  @parameterized.named_parameters(
      ("out_of_bounds_pos", 2, ValueError),
      ("out_of_bounds_neg", -3, ValueError),
      ("duplicate_axis", (0, 0), ValueError),
      ("duplicate_canonical_axis", (1, -1), ValueError),
  )
  def test_reduction_axis_canonicalization_validation(
      self, axis, expected_error
  ):
    def make_kernel(op):
      @self.kernel(out_type=jax.ShapeDtypeStruct((128,), jnp.float32))
      def kernel(x_ref, y_ref):
        x_val = plgpu.load(x_ref, layout=plgpu.Layout.WGMMA, optimized=False)
        plgpu.store(y_ref, op(x_val, axis=axis), optimized=False)

      return kernel

    x = jnp.ones((128, 128), dtype=jnp.float32)
    for op in (plgpu.sum, plgpu.max, plgpu.min, plgpu.prod):
      with self.assertRaises(expected_error):
        make_kernel(op)(x)

  def test_cross_warp_reduction(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
    )
    def kernel(x_ref, y_ref):
      layout = plgpu.Layout.TCGEN05_TMEM_NATIVE(4)
      x_val = plgpu.load(x_ref, layout=layout, optimized=False)
      plgpu.store(y_ref, jnp.sum(x_val, axis=0), optimized=False)

    x = jax.random.uniform(jax.random.key(0), shape=(128, 128), dtype=jnp.float32)
    np.testing.assert_allclose(kernel(x), jnp.sum(x, axis=0), atol=5e-5)

  def _test_broadcast_in_dim_base(self, shape, layout, *, axis, hint):
    assert len(shape) == 2

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, jnp.float32),
    )
    def kernel(x_ref, y_ref):
      reduced_layout = layout.reduce(axis)
      reduced = plgpu.load(x_ref, layout=reduced_layout, optimized=False)
      broadcasted = lax.broadcast_in_dim(reduced, shape, [1 - axis])
      if hint:
        broadcasted = plgpu.layout_cast(broadcasted, layout)
      # Note that without the hint, the layout of broadcasted is not guaranteed
      # to be the same as the layout argument!
      plgpu.store(y_ref, broadcasted, optimized=False)

    x = jax.random.uniform(jax.random.key(0), shape=(128,), dtype=jnp.float32)
    x_result = jax.block_until_ready(kernel(x))
    expected = jnp.expand_dims(x, axis=axis)
    expected = jnp.broadcast_to(expected, shape)
    np.testing.assert_array_equal(x_result, expected)

  @parameterized.product(
      layout=(
          plgpu.Layout.WGMMA,
          plgpu.Layout.TCGEN05,
          plgpu.Layout.TCGEN05_TMEM_NATIVE,
          plgpu.Layout.TCGEN05_M64_COLLECTIVE(128),
      ),
      axis=(0, 1),
      hint=(True, False),
  )
  def test_broadcast_in_dim(self, layout, axis, hint):
    self._test_broadcast_in_dim_base((128, 128), layout, axis=axis, hint=hint)

  # Regression test for a crash when using a small shape.
  def test_broadcast_in_dim_does_not_crash_on_small_shape(self):
    shape = (128, 4)
    self._test_broadcast_in_dim_base(
        shape, plgpu.Layout.TCGEN05_TMEM_NATIVE, axis=1, hint=False
    )

  def test_broadcast_in_dim_wg_strided_majormost_dim(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct((64, 128), jnp.float32),
    )
    def kernel(x_ref, y_ref):
      to_be_broadcasted = plgpu.load(
          x_ref, layout=plgpu.Layout.WG_STRIDED((128,), 1)
      )
      broadcasted = lax.broadcast_in_dim(to_be_broadcasted, (64, 128), (1,))
      y_ref[...] = broadcasted

    result = jax.random.uniform(jax.random.key(0), shape=(128,), dtype=jnp.float32)
    np.testing.assert_array_equal(kernel(result), jnp.broadcast_to(result[None,:], (64, 128)))

  @parameterized.parameters(
      ((4, 128),),
      ((2, 4, 128),),
  )
  def test_broadcast_wg_strided_majormost_dim(self, out_shape):
    dtype = jnp.float32
    @functools.partial(
        self.pallas_call, out_shape=jax.ShapeDtypeStruct(out_shape, dtype)
    )
    def kernel(x_ref, side_load_ref, y_ref):
      x_strided = plgpu.load(
          x_ref, layout=plgpu.Layout.WG_STRIDED((128,), vec_size=1)
      )
      side_load_strided = plgpu.load(
          side_load_ref, layout=plgpu.Layout.WG_STRIDED(out_shape, vec_size=1)
      )
      for _ in range(len(out_shape) - 1):
        x_strided = x_strided[None, ...]
      y_ref[...] = x_strided + side_load_strided[...]

    inp = jax.random.uniform(jax.random.key(0), (128,), dtype)
    side_load = jax.random.uniform(jax.random.key(1), out_shape, dtype)
    np.testing.assert_array_equal(kernel(inp, side_load),
                                  jnp.broadcast_to(inp, out_shape) + side_load)

  def test_broadcast_in_dim_tcgen05_native_layout(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.float32),
    )
    def kernel(x_ref, y_ref):
      reduced = plgpu.load(x_ref, layout=plgpu.Layout.TCGEN05_TMEM_NATIVE.reduce(1), optimized=False)
      broadcasted = lax.broadcast_in_dim(reduced, (128, 128), [0])
      broadcasted = plgpu.layout_cast(broadcasted, plgpu.Layout.TCGEN05_TMEM_NATIVE)
      plgpu.store(y_ref, broadcasted, optimized=False)

    x = jax.random.uniform(jax.random.key(0), shape=(128,), dtype=jnp.float32)
    np.testing.assert_array_equal(kernel(x), jnp.broadcast_to(x[:, None], (128, 128)))

  def test_broadcast_in_dim_size_one_dimension(self):

    @self.kernel(out_type=jax.ShapeDtypeStruct((2, 4, 128), jnp.float32))
    def kernel(src_ref, dst_ref):
      src = plgpu.load(src_ref, layout=plgpu.Layout.WG_STRIDED((1, 128), 1))
      dst = lax.broadcast_in_dim(src, (2, 4, 128), (1, 2))
      dst_ref[...] = dst

    src = jnp.arange(128, dtype=jnp.float32).reshape((1, 128))
    dst = lax.broadcast_in_dim(src, (2, 4, 128), (1, 2))
    np.testing.assert_array_equal(kernel(src), dst)

  @parameterized.parameters("col", "row")
  def test_broadcast_in_dim_with_expanded_dimension(self, mode):
    if not self.is_wg_semantics():
      self.skipTest("Layout inference is only enabled under WG semantics.")

    shape = (64, 64)

    @self.kernel(out_type=jax.ShapeDtypeStruct(shape, jnp.float32))
    def kernel(mask_ref, acc_ref, o_ref):
      mask = plgpu.load(mask_ref, optimized=False)
      acc = plgpu.load(acc_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      if mode == "col":
        plgpu.store(o_ref, acc * mask[:, None], optimized=False)
      elif mode == "row":
        plgpu.store(o_ref, acc * mask[None, :], optimized=False)
      else:
        raise ValueError(f"Unknown mode: {mode}")

    mask = jnp.arange(64, dtype=jnp.float32)
    acc = jnp.arange(64 * 64, dtype=jnp.float32).reshape(shape)
    expected = acc * (mask[:, None] if mode == "col" else mask[None, :])
    np.testing.assert_array_equal(kernel(mask, acc), expected)

  @parameterized.named_parameters((l.name.lower(), l) for l in plgpu.Layout)
  def test_copy_layout(self, layout):
    if layout in {
        plgpu.Layout.WG_SPLAT,
        plgpu.Layout.WGMMA_TRANSPOSED,
        plgpu.Layout.TCGEN05_TRANSPOSED,
        plgpu.Layout.TILED,
        plgpu.Layout.MMA_LHS,
        plgpu.Layout.MMA_RHS,
        plgpu.Layout.MMA_ACC,
    }:
      self.skipTest("Not the right layout for this test")

    # We don't infer optimized transfer-compatible transforms for load to
    # registers with TCGEN05_TMEM_NATIVE layout.
    # TODO(allanrenucci): Manually specify transforms when supported for WG
    # lowering semantic.
    optimized = (
        self.LOWERING_SEMANTICS == plgpu.LoweringSemantics.Lane
        or layout != plgpu.Layout.TCGEN05_TMEM_NATIVE
    ) and layout != plgpu.Layout.TCGEN05_M64_COLLECTIVE_NATIVE

    if "tcgen05" in layout.name.lower():
      self.skip_unless_tcgen05()
      limit = None
      shape = (128, 128)
    else:
      limit = jtu._SMEM_SIZE_BOUND_FOR_TESTS
      shape = (64, 128)
    dtype = jnp.float32
    swizzle = 128
    if layout in (plgpu.Layout.WGMMA_UPCAST_4X, plgpu.Layout.WGMMA_UPCAST_2X):
      dtype = jnp.float8_e5m2
      swizzle = 64
    transforms = self.default_transforms(dtype=dtype, swizzle=swizzle)

    if layout == plgpu.Layout.TCGEN05_M64_COLLECTIVE:
      layout = plgpu.Layout.TCGEN05_M64_COLLECTIVE(128)
    elif layout == plgpu.Layout.TCGEN05_M64_COLLECTIVE_NATIVE:
      self.skipTest("Need to add support for optimized= for stores")
    elif layout == plgpu.Layout.WG_STRIDED:
      layout = plgpu.Layout.WG_STRIDED(shape, 2)
      transforms = ()
    elif layout == plgpu.Layout.SMEM_GMEM_COPY:
      layout = plgpu.Layout.SMEM_GMEM_COPY(shape, jnp.float32, swizzle=128)

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(shape, dtype),
        in_specs=[plgpu.BlockSpec(transforms=transforms)],
        out_specs=plgpu.BlockSpec(transforms=transforms),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = plgpu.load(x_ref, layout=layout, optimized=optimized)

    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    with mgpu.core.artificial_shared_memory_limit(limit):
      np.testing.assert_array_equal(kernel(x), x)

  @parameterized.parameters(
      (((0, 0),), (128, 128), (128, 128)),
      (((0, 1),), (128, 128), (128, 128)),
      (((1, None),), (128, 128), (128,)),
      (((0, 0),), (128, 128), (128, 128)),
      (((0, 0), (0, 0)), (128, 128), (128, 128)),
  )
  def test_vmap_kernel(self, vmap_axes, x_shape, y_shape):
    rng0, rng1 = jax.random.split(jax.random.key(0))
    x = jax.random.uniform(rng0, x_shape, jnp.float32)
    y = jax.random.uniform(rng1, y_shape, jnp.float32)

    out_shape = list(x_shape)
    for x_axis, _ in vmap_axes:
      del out_shape[x_axis]
    out_shape = jax.ShapeDtypeStruct(out_shape, jnp.float32)

    @self.kernel(out_type=out_shape)
    def f(x_ref, y_ref, o_ref):
      o_ref[...] = x_ref[...] + y_ref[...]

    f_ref = lambda x, y: x + y
    for in_axes in vmap_axes:
      f = jax.vmap(f, in_axes)
      f_ref = jax.vmap(f_ref, in_axes)

    np.testing.assert_array_equal(f(x, y), f_ref(x, y))

  def test_discharge_comms_effect(self):
    def body(out, sem):
      pl.semaphore_signal(sem, device_id=jnp.asarray(2, jnp.int32))

    f = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((), jnp.int32),
        scratch_types=[plgpu.SemaphoreType.REGULAR],
    )
    jax_core.check_jaxpr(jax.make_jaxpr(f)().jaxpr)

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  @jtu.skip_under_pytest("Test fails under pytest in CI")
  def test_line_info(self):
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), jtu.capture_stdout() as output:
      @functools.partial(
        self.kernel,
        out_type=jax.ShapeDtypeStruct([256], jnp.float32),
      )
      def kernel(x_ref, o_ref):
        o_ref[...] = x_ref[...] + x_ref[0]

      jax.block_until_ready(kernel(jnp.arange(256, dtype=jnp.float32)))

    ptx = output()
    self.assertIn(".file", ptx)
    self.assertIn(".loc", ptx)
    [path] = re.findall(r'.file\s+\d+\s+"(.+)"', ptx)
    self.assertEndsWith(__file__, path)

  @jtu.thread_unsafe_test()
  @jtu.skip_under_pytest("Test fails under pytest in CI")
  def test_e8m0_reciprocal_ptx(self):
    shape = (64, 128)
    transforms = self.default_transforms(swizzle=128, dtype=jnp.dtype(jnp.float8_e8m0fnu))
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(shape, jnp.float8_e8m0fnu),
        in_specs=[plgpu.BlockSpec(transforms=transforms)],
        out_specs=plgpu.BlockSpec(transforms=transforms),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = 1 / x_ref[...]

    x_uint8 = (np.arange(np.prod(shape), dtype=np.uint8).reshape(shape) % 253) + 1
    x = jax.lax.bitcast_convert_type(
        jnp.array(x_uint8), jnp.float8_e8m0fnu
    )
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      result = jax.block_until_ready(kernel(x))
    ptx_str = ptx()
    self.assertRegex(ptx_str, r"sub\b")
    self.assertNotRegex(ptx_str, r"div\.f")
    result_uint8 = np.asarray(
        jax.lax.bitcast_convert_type(result, jnp.uint8)
    )
    np.testing.assert_array_equal(result_uint8, 254 - x_uint8)

  @jtu.thread_unsafe_test()
  @jtu.skip_under_pytest("Test fails under pytest in CI")
  def test_clz_ptx(self):
    shape = (128,)
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(shape, jnp.int32),
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = lax.clz(x_ref[...])

    x = jnp.arange(128, dtype=jnp.int32)
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      jax.block_until_ready(kernel(x))
    self.assertRegex(ptx(), r"clz\.b32\b")

  def test_collective_arrival_count(self):
    def kernel(dst, collective_barrier):
      plgpu.barrier_arrive(collective_barrier)
      plgpu.barrier_arrive(collective_barrier)
      plgpu.barrier_arrive(collective_barrier)
      plgpu.barrier_arrive(collective_barrier)
      plgpu.barrier_wait(collective_barrier)
      dst[...] = jnp.ones_like(dst)
    y = self.kernel(
      kernel,
      out_type=jax.ShapeDtypeStruct((128,), jnp.int32),
      scratch_types=[plgpu.ClusterBarrier(collective_axes=("x",), num_arrivals=4)],
      cluster=(2,),
      cluster_names=("x",)
    )()
    np.testing.assert_array_equal(y, np.ones((), dtype=np.int32))

  def test_leader_tracked_cluster_barrier(self):
    def kernel(dst_ref, collective_barrier):
      plgpu.barrier_arrive(collective_barrier)
      @pl.when(lax.axis_index("x") == 0)
      def _write():
        plgpu.barrier_wait(collective_barrier)
        dst_ref[...] = 42

    self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((), jnp.int32),
        scratch_types=[
            plgpu.ClusterBarrier(
                collective_axes=("x",), num_arrivals=1, leader_tracked=True
            ),
        ],
        cluster=(2,),
        cluster_names=("x",),
    )()

  @parameterized.parameters("x", "y")
  def test_cluster_ref_read(self, axis_name):

    def kernel(src_ref, dst_ref, scratch_ref, copy_barrier, cluster_barrier):
      my_idx = lax.axis_index(axis_name)
      src_slice = src_ref.at[my_idx]
      plgpu.copy_gmem_to_smem(src_slice, scratch_ref, copy_barrier)
      plgpu.barrier_wait(copy_barrier)
      plgpu.barrier_arrive(cluster_barrier)
      plgpu.barrier_wait(cluster_barrier)
      peer_scratch = plgpu.cluster_ref(scratch_ref, {axis_name: 1 - my_idx})
      for i in range(2):
        s = pl.ds(i * 64, 64)
        plgpu.store(
            dst_ref.at[my_idx, s], peer_scratch.at[s][...], optimized=False
        )

    cluster_names = ("x", "y")
    if axis_name == "x":
      cluster = (2, 1)
    else:
      cluster = (1, 2)

    x = jnp.arange(2 * 128 * 32, dtype=jnp.float32).reshape(2, 128, 32)
    transforms = self.default_transforms(dtype=jnp.float32)

    y = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((2, 128, 32), jnp.float32),
        scratch_types=[
            plgpu.SMEM((128, 32), jnp.float32, transforms=transforms),
            plgpu.Barrier(),
            plgpu.ClusterBarrier(collective_axes=(axis_name,)),
        ],
        cluster=cluster,
        cluster_names=cluster_names,
    )(x)

    expected = jnp.flip(x, axis=0)
    np.testing.assert_array_equal(y, expected)

  @parameterized.parameters(False, True)
  def test_cluster_ref_write(self, orders_tensor_core):
    if orders_tensor_core:
      self.skip_unless_tcgen05()
    dtype = jnp.float32
    logical_shape = (64, 32)
    x_shape = (2, 64, 32)
    transforms = self.default_transforms(dtype=dtype)

    @self.kernel(
        out_type=(
            jax.ShapeDtypeStruct(x_shape, dtype),
            jax.ShapeDtypeStruct((2,), jnp.int32),
        ),
        scratch_types=[
            plgpu.SMEM(logical_shape, dtype, transforms=transforms),
            plgpu.Barrier(num_arrivals=2, orders_tensor_core=orders_tensor_core),
        ],
        cluster=(2,),
        cluster_names=("c",),
    )
    def kernel(x_ref, o_ref, completed_early_ref, smem_ref, barrier):
      my_idx = jax.lax.axis_index("c")
      x = plgpu.load(x_ref.at[my_idx], layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.async_store_smem(
        x, smem_ref, barrier, cluster_idx=1 - my_idx, cluster_dim="c"
      )
      completed_early_ref[my_idx] = plgpu.barrier_test(barrier).astype(jnp.int32)
      plgpu.barrier_arrive(barrier)
      plgpu.barrier_wait(barrier)
      plgpu.store(
          o_ref.at[my_idx],
          plgpu.load(smem_ref, layout=plgpu.Layout.WGMMA),
          optimized=False,
      )

    x = jnp.arange(2 * 64 * 32, dtype=dtype).reshape(x_shape)
    out, completed_early = kernel(x)
    np.testing.assert_array_equal(out, np.flip(x, axis=0))
    # Regression test to ensure we don't arrive more than expected when
    # orders_tensor_core=True.
    np.testing.assert_array_equal(
        completed_early, jnp.zeros((2,), dtype=jnp.int32)
    )

  def test_cluster_ref_write_atomic(self):
    dtype = jnp.int32
    logical_shape = (64, 32)
    x_shape = (2, 64, 32)
    transforms = self.default_transforms(dtype=dtype)

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(x_shape, dtype),
        scratch_types=[
            plgpu.SMEM(logical_shape, dtype, transforms=transforms),
            plgpu.ClusterBarrier(collective_axes=("c",), num_arrivals=1),
            plgpu.Barrier(num_arrivals=1),
            plgpu.Barrier(num_arrivals=1),
        ],
        cluster=(2,),
        cluster_names=("c",),
    )
    def kernel(x_ref, init_ref, o_ref, smem_ref, cluster_barrier, init_barrier, store_barrier):
      my_idx = jax.lax.axis_index("c")
      peer_idx = 1 - my_idx
      plgpu.copy_gmem_to_smem(init_ref.at[my_idx], smem_ref, init_barrier)
      plgpu.barrier_wait(init_barrier)
      plgpu.barrier_arrive(cluster_barrier)
      plgpu.barrier_wait(cluster_barrier)
      x = plgpu.load(x_ref.at[my_idx], layout=plgpu.Layout.WGMMA, optimized=False)
      plgpu.async_store_smem(
          x,
          smem_ref,
          store_barrier,
          cluster_idx=peer_idx,
          cluster_dim="c",
          atomic="add",
      )
      plgpu.barrier_wait(store_barrier)
      plgpu.store(
          o_ref.at[my_idx],
          plgpu.load(smem_ref, layout=plgpu.Layout.WGMMA),
          optimized=False,
      )

    x = jnp.arange(2 * 64 * 32, dtype=dtype).reshape(x_shape)
    init_data = jnp.ones(x_shape, dtype=dtype) * 10
    np.testing.assert_array_equal(kernel(x, init_data), 10 + x[::-1])

  def test_replicated_layout(self):
    shape = (32,)
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct(shape, jnp.float32),
    )
    def kernel(src_ref, dst_ref):
      layout = plgpu.Layout.TILED(
        plgpu.Tiling(((32,), (1,))),
        warp_dims=(plgpu.Replicated(4),),
        lane_dims=(-2,),
        vector_dim=-1,
      )
      dst_ref[...] = plgpu.load(src_ref, layout=layout)
    src = jnp.arange(shape[0], dtype=jnp.float32)
    np.testing.assert_array_equal(kernel(src), src)

  @parameterized.product(
      dtype=[jnp.float32, jnp.int32],
      layout=[plgpu.Layout.WGMMA, None]
  )
  def test_atomic_add(self, dtype, layout):
    m, n = 128, 64
    if layout is not None:
      transforms = self.default_transforms(swizzle=64, dtype=dtype)
    else:
      transforms = ()
    def body(inp_ref, out_ref, smem_ref, barrier):
      wg_idx = lax.axis_index("wg")
      @pl.when(wg_idx == 0)
      def _zero():
        zeros = jnp.zeros_like(smem_ref)
        if layout is not None:
          zeros = plgpu.layout_cast(zeros, layout)
        smem_ref[...] = zeros
      plgpu.barrier_arrive(barrier)
      plgpu.barrier_wait(barrier)
      plgpu.atomic_add(
          smem_ref,
          plgpu.load(inp_ref.at[pl.ds(wg_idx * m, m)], layout=layout, optimized=False),
      )
      plgpu.barrier_arrive(barrier)
      plgpu.barrier_wait(barrier)
      @pl.when(wg_idx == 0)
      def _copy_out():
        plgpu.store(out_ref, smem_ref[...], optimized=False)
    x = jnp.arange(1, m * n + 1, dtype=dtype).reshape(m, n)
    y = jnp.arange(m * n, 0, -1, dtype=dtype).reshape(m, n)
    inp = jnp.concatenate([x, y], axis=0)
    result = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct([m, n], dtype),
        num_threads=2,
        thread_name="wg",
        scratch_types=[
            plgpu.SMEM((m, n), dtype, transforms=transforms),
            plgpu.Barrier(num_arrivals=2),
        ],
    )(inp)
    np.testing.assert_array_equal(result, x + y)

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_atomic_add_gmem(self):
    m, n = 128, 64
    x = jnp.arange(1, m * n + 1, dtype=jnp.float32).reshape(m, n)

    def body(inp_ref, out_ref):
      val = plgpu.load(
          inp_ref, layout=plgpu.Layout.WGMMA, optimized=False
      )
      out_ref[...] = jnp.zeros_like(out_ref)
      plgpu.atomic_add(out_ref, val, optimized=True)

    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      result = self.kernel(
          body,
          out_type=jax.ShapeDtypeStruct([m, n], jnp.float32),
      )(x)
      jax.block_until_ready(result)
    self.assertArraysEqual(result, x)
    self.assertIn("red.global", ptx())

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_griddepcontrol(self):
    @jax.jit
    def f(x):
      def kernel_body(x_ref, o_ref):
        plgpu.griddepcontrol_wait()
        o_ref[...] = x_ref[...] + 1.0
        plgpu.griddepcontrol_launch_dependents()

      return self.kernel(
          kernel_body,
          out_type=jax.ShapeDtypeStruct(x.shape, x.dtype),
      )(x)

    x = jnp.arange(128).astype(jnp.float32)

    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      out = jax.block_until_ready(f(x))

    np.testing.assert_allclose(out, x + 1.0)

    ptx_output = ptx()
    self.assertIn("griddepcontrol.wait;", ptx_output)
    self.assertIn("griddepcontrol.launch_dependents;", ptx_output)

  # TODO(slebedev): Support unsigned int types.
  @parameterized.product(
      dtype=(
          jnp.bfloat16,
          jnp.float16,
          jnp.float64,
          jnp.float8_e4m3fn,
          jnp.float8_e5m2,
          jnp.int8,
          jnp.int4,
      ),
  )
  def test_mma(self, dtype):
    m = k = 128
    n = 8
    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

    @functools.partial(
        self.kernel,
        out_type=jax.ShapeDtypeStruct((m, n), acc_dtype),
    )
    def kernel(a_ref, b_ref, o_ref):
      acc = plgpu.layout_cast(
          jnp.zeros((m, n), acc_dtype), plgpu.Layout.MMA_ACC(dtype)
      )
      a = plgpu.load(a_ref, layout=plgpu.Layout.MMA_LHS(dtype), optimized=False)
      b = plgpu.load(
          b_ref.T, layout=plgpu.Layout.MMA_RHS(dtype), optimized=False
      )
      plgpu.store(o_ref, plgpu.mma(acc, a, b), optimized=False)

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

    res = kernel(a, b)
    ref = jnp.dot(a.astype(acc_dtype), b.T.astype(acc_dtype), precision="high")
    if is_integer:
      np.testing.assert_array_equal(res, ref)
    else:
      atol = rtol = 1e-12 if dtype == jnp.float64 else 1e-2
      np.testing.assert_allclose(res, ref, atol=atol, rtol=rtol)

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_griddepcontrol_warp_mesh(self):
    @jax.jit
    def f(x):
      def kernel_a(x_ref, out_ref):
        @plgpu.warp_map
        def _(warp_id):
          out_ref[warp_id] = x_ref[warp_id] + 1.0
          @pl.when(warp_id == 0)
          def _():
            plgpu.griddepcontrol_launch_dependents()

      def kernel_b(in_ref, out_ref):
        @plgpu.warp_map
        def _(warp_id):
          @pl.when(warp_id == 0)
          def _():
            plgpu.griddepcontrol_wait()
          out_ref[warp_id] = in_ref[warp_id] * 2.0

      intermediate = self.kernel(
          kernel_a,
          out_type=jax.ShapeDtypeStruct(x.shape, x.dtype),
      )(x)
      return self.kernel(
          kernel_b,
          out_type=jax.ShapeDtypeStruct(x.shape, x.dtype),
      )(intermediate)

    x = jnp.arange(4).astype(jnp.float32)

    with (
        jtu.set_env(MOSAIC_GPU_DUMP_PTX="1", MOSAIC_GPU_DUMP_HOST_LLVM="1"),
        self.capture_stdout() as output,
    ):
      out = jax.block_until_ready(f(x))

    np.testing.assert_allclose(out, (x + 1.0) * 2.0)

    stdout = output()
    self.assertIn("griddepcontrol.wait;", stdout)
    self.assertIn("griddepcontrol.launch_dependents;", stdout)

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_griddepcontrol_multi_kernel(self):
    # Tests that we can launch two kernels that communicate via griddepcontrol.
    @jax.jit
    def f(x):
      def kernel_a(x_ref, out_ref):
        out_ref[...] = x_ref[...] + 1.0
        plgpu.griddepcontrol_launch_dependents()

      def kernel_b(in_ref, out_ref):
        plgpu.griddepcontrol_wait()
        out_ref[...] = in_ref[...] * 2.0

      intermediate = self.kernel(
          kernel_a,
          out_type=jax.ShapeDtypeStruct(x.shape, x.dtype),
      )(x)
      return self.kernel(
          kernel_b,
          out_type=jax.ShapeDtypeStruct(x.shape, x.dtype),
      )(intermediate)

    x = jnp.arange(128).astype(jnp.float32)
    with (
        jtu.set_env(MOSAIC_GPU_DUMP_PTX="1", MOSAIC_GPU_DUMP_HOST_LLVM="1"),
        self.capture_stdout() as output,
    ):
      out = jax.block_until_ready(f(x))

    np.testing.assert_allclose(out, (x + 1.0) * 2.0)

    # Static verification of the dynamic launch (PDL) pipeline
    ptx_output = output()

    # 1. Verify device-side hardware synchronization instructions are present
    self.assertIn("griddepcontrol.launch_dependents;", ptx_output)
    self.assertIn("griddepcontrol.wait;", ptx_output)

    # 2. Verify host-side selective PDL launch attributes (uses_pdl flag)
    import re
    # kernel_a (producer) need not enable PDL (uses_pdl should be i32 0)
    kernel_a_pattern = (
        r"define void @kernel_a_mosaic_gpu\(.*?{.*?call void"
        r" @mosaic_gpu_build_kernel_spec\(.*?, i32 0, i32 \d+, ptr %\w+, ptr %\w+\)"
    )
    self.assertRegex(ptx_output, re.compile(kernel_a_pattern, re.DOTALL))

    # kernel_b (consumer) MUST enable PDL (uses_pdl should be i32 1)
    kernel_b_pattern = (
        r"define void @kernel_b_mosaic_gpu\(.*?{.*?call void"
        r" @mosaic_gpu_build_kernel_spec\(.*?, i32 1, i32 \d+, ptr %\w+, ptr %\w+\)"
    )
    self.assertRegex(ptx_output, re.compile(kernel_b_pattern, re.DOTALL))

  def test_mlir_error_includes_mlir_error_callsites(self):
    class MockAttr:
      def __init__(self, loc):
        self.location = loc

      def __str__(self):
        return str(self.location)

    class MockLoc:
      def __init__(self, location):
        self.attr = MockAttr(location)

    class MockDiagnostic:
      def __init__(self, message, loc):
        self.message = message
        self.location = MockLoc(loc)
        self.notes = []

    class MockMLIRError(ir.MLIRError):
      def __init__(self, message, diagnostics):
        super().__init__(message)
        self._diagnostics = diagnostics

      @property
      def error_diagnostics(self):
        return self._diagnostics

    def mock_verify_fail(self) -> bool:
      loc_str = 'loc("kernel"("fn1"("file1.py":10:20) at "fn2"("file2.py":30:40)))'
      raise MockMLIRError("err", [MockDiagnostic("diagnostic", loc_str)])

    with mock.patch.object(ir.Operation, "verify", mock_verify_fail):
      try:
        @self.kernel(out_type=jax.ShapeDtypeStruct((16,), jnp.bfloat16))
        def kernel(out):
          del out
        jax.jit(kernel).lower()
        self.fail("unreachable")
      except VerificationError as e:
        frames = traceback.extract_tb(e.__traceback__)
        frame_caller = frames[-2]
        self.assertEqual(frame_caller.filename, "file2.py")
        self.assertEqual(frame_caller.lineno, 30)
        self.assertEqual(frame_caller.name, "fn2")
        frame_callee = frames[-1]
        self.assertEqual(frame_callee.filename, "file1.py")
        self.assertEqual(frame_callee.lineno, 10)
        self.assertEqual(frame_callee.name, "fn1")


class PallasCallWarpPrimitiveSemanticsTest(PallasTest):
  def setUp(self):
    super().setUp()

  def test_axis_index(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct((2, 128), jnp.int32))
    def kernel(y_ref):
      def scope(ones_smem_ref, threes_smem_ref):
        # Prepare data to copy.
        ones_smem_ref[:] = jnp.ones((1, 128), jnp.int32)
        threes_smem_ref[:] = jnp.ones((1, 128), jnp.int32) * 3
        plgpu.commit_smem()
        @plgpu.warp_map
        def _(warp_id):
          # We cannot load/store inside of core_map, so we issue async
          # copies instead to produce a testable result.
          @pl.when(warp_id == 1)
          def _():
            plgpu.async_prefetch(y_ref.at[1:2])
            plgpu.copy_smem_to_gmem(ones_smem_ref, y_ref.at[0:1])
          @pl.when(warp_id == 3)
          def _():
            plgpu.copy_smem_to_gmem(threes_smem_ref, y_ref.at[1:2])
        plgpu.wait_smem_to_gmem(0)
      pl.run_scoped(scope,
                    plgpu.SMEM((1, 128), jnp.int32),
                    plgpu.SMEM((1, 128), jnp.int32)
                    )
    result = kernel()
    expected = jnp.stack((jnp.ones((128,), jnp.int32),
                          jnp.ones((128,), jnp.int32) * 3), axis=0)
    np.testing.assert_array_equal(result, expected)

  def test_axis_index_mpmd_map(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct((2, 128), jnp.int32))
    def kernel(y_ref):
      def scope(ones_smem_ref, threes_smem_ref):
        # Prepare data to copy.
        ones_smem_ref[:] = jnp.ones((1, 128), jnp.int32)
        threes_smem_ref[:] = jnp.ones((1, 128), jnp.int32) * 3
        plgpu.commit_smem()

        @pl.kernel(mesh=gpu_core.WarpMesh(axis_name='foo'))
        def kernel():
          warp_id = lax.axis_index('foo')
          @pl.when(warp_id == 1)
          def _():
            plgpu.async_prefetch(y_ref.at[1:2])
            plgpu.copy_smem_to_gmem(ones_smem_ref, y_ref.at[0:1])
          @pl.when(warp_id == 3)
          def _():
            plgpu.copy_smem_to_gmem(threes_smem_ref, y_ref.at[1:2])

        kernel()
        plgpu.wait_smem_to_gmem(0)

      pl.run_scoped(scope,
                    plgpu.SMEM((1, 128), jnp.int32),
                    plgpu.SMEM((1, 128), jnp.int32)
                    )
    result = kernel()
    expected = jnp.stack((jnp.ones((128,), jnp.int32),
                          jnp.ones((128,), jnp.int32) * 3), axis=0)
    np.testing.assert_array_equal(result, expected)

  def test_axis_index_multi_warpgroup(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((8,), jnp.int32),
        num_threads=2,
        thread_name="wg",
    )
    def kernel(out_ref):
      wg_idx = lax.axis_index("wg")
      @plgpu.warp_map
      def _(warp_id):
        out_ref[wg_idx * 4 + warp_id] = warp_id
    result = kernel()
    expected = jnp.concatenate([jnp.arange(4, dtype=jnp.int32)] * 2)
    np.testing.assert_array_equal(result, expected)

  def test_scalar_load(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.int32))
    def kernel(x_ref, y_ref):
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 1)
        def _():
          y_ref[...] = x_ref[...]
    np.testing.assert_array_equal(kernel(4), 4)

  def test_scalar_plgpu_load(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.int32))
    def kernel(x_ref, y_ref):
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 1)
        def _():
          y_ref[...] = plgpu.load(x_ref)
    np.testing.assert_array_equal(kernel(jnp.int32(4)), 4)

  def test_non_scalar_load_raises(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct((2,), jnp.int32))
    def kernel(x_ref, y_ref):
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 1)
        def _():
          y_ref[...] = x_ref[...]
    with self.assertRaisesRegex(ValueError, "Can only load scalars",):
      kernel(jnp.ones((2,), jnp.int32))

  def test_scalar_store(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.int32))
    def kernel(x_ref, y_ref):
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 1)
        def _():
          plgpu.store(y_ref, x_ref[...])
    np.testing.assert_array_equal(kernel(jnp.int32(4)), 4)

  @parameterized.parameters(
    lax.add, lax.sub, lax.mul, lax.div, lax.rem, lax.bitwise_and,
    lax.bitwise_or, lax.bitwise_xor, lax.max, lax.min,
    lax.gt, lax.lt, lax.ge, lax.le, lax.eq, lax.ne,
  )
  def test_scalar_binary_op(self, op):
    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.int32))
    def kernel(y_ref):
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 1)
        def _():
          x = jnp.array(1234, dtype=jnp.int32)
          y = jnp.array(6543, dtype=jnp.int32)
          y_ref[...] = op(x, y).astype(jnp.int32)
    result = kernel()
    x = jnp.array(1234, dtype=jnp.int32)
    y = jnp.array(6543, dtype=jnp.int32)
    np.testing.assert_array_equal(result, op(x, y).astype(jnp.int32))

  def test_optimization_barrier(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.int32))
    def kernel(x_ref, y_ref):
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 1)
        def _():
          c = lax.optimization_barrier(x_ref[...])
          y_ref[...] = c
    x = jnp.array(4, dtype=jnp.int32)
    np.testing.assert_array_equal(kernel(x), 4)

  def test_optimization_barrier_multiple_inputs(self):
    @self.kernel(out_type=jax.ShapeDtypeStruct((), jnp.int32))
    def kernel(x_ref, y_ref):
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 1)
        def _():
          c = jnp.int32(5)
          x, c, w = lax.optimization_barrier((x_ref[...], c, warp_id))
          y_ref[...] = x + c + w
    x = jnp.array(4, dtype=jnp.int32)
    np.testing.assert_array_equal(kernel(x), 4 + 5 + 1)

  def test_errors_when_closing_over_array(self):
    # We currently do not allow closing over arrays when mapping over
    # a mesh, since we would need to present a view of the array local
    # to each warp.
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((32, 32), jnp.float32),
        scratch_types=[plgpu.SMEM((32, 32), jnp.float32)],
    )
    def kernel(out_ref, smem_ref):
      arr = jnp.ones((32, 32), dtype=jnp.float32)
      @plgpu.warp_map
      def _(warp_id):
        smem_ref[...] = arr + 1
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(smem_ref, out_ref)
      plgpu.wait_smem_to_gmem(0)
    with self.assertRaisesRegex(
        ValueError, "captures non-scalar array constants"
    ):
      kernel()

  def test_inline_mgpu_in_warp(self):

    @self.kernel(
        out_type=(
            jax.ShapeDtypeStruct((), jnp.int32),
            jax.ShapeDtypeStruct((), jnp.int32),
        ),
    )
    def kernel(y_ref, z_ref):
      @plgpu.warp_map
      def _(warp_id):
        @plgpu.inline_mgpu(
            arg_types=(plgpu.RefType(), plgpu.Layout.WG_SPLAT),
            return_type=plgpu.ShapeDtypeStruct(
                (), jnp.int32, layout=plgpu.Layout.WG_SPLAT
            ),
        )
        def write(ctx, ref, wid):
          del ctx
          memref_dialect.store(wid.registers.item(), ref, [])
          i32 = ir.IntegerType.get_signless(32)
          ten = mgpu.FragmentedArray.splat(
              mgpu.c(10, i32), (), is_signed=True
          )
          return wid + ten
        @pl.when(warp_id == 1)
        def _():
          result = write(y_ref, warp_id)
          z_ref[...] = result
    y, z = kernel()
    np.testing.assert_array_equal(y, 1)
    np.testing.assert_array_equal(z, 11)

  def test_inline_mgpu_in_warp_rejects_non_scalar(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
        out_specs=pl.BlockSpec(memory_space=plgpu.SMEM),
    )
    def kernel(y_ref):
      @plgpu.warp_map
      def _(warp_id):
        @plgpu.inline_mgpu(
            arg_types=(),
            return_type=plgpu.ShapeDtypeStruct(
                (128,), jnp.float32, layout=plgpu.Layout.WGMMA
            ),
        )
        def bad(ctx):
          del ctx
          return mgpu.FragmentedArray.splat(mgpu.c(0.0, ir.F32Type.get()), (128,))
        y_ref[...] = bad()
    with self.assertRaisesRegex(
        Exception,
        r"inline_mgpu in a single-warp context only supports scalar return"
        r" types\. Got shape=\(128,\)\.",
    ):
      kernel()

  @parameterized.parameters(True, False)
  def test_single_warp_loop(self, force_while):
    @self.kernel(out_type=jax.ShapeDtypeStruct((10, 128), jnp.int32))
    def kernel(y_ref):
      def scope(smem_ref):
        # Prepare data to copy.
        for i in range(10):
          smem_ref[i, :] = jnp.ones_like(smem_ref.at[i]) * i
        plgpu.commit_smem()
        @plgpu.warp_map
        def _(warp_id):
          @pl.when(warp_id == 0)
          def _():
            def loop_body(i, _):
              _slice = pl.ds(i, 1)
              plgpu.copy_smem_to_gmem(smem_ref.at[_slice], y_ref.at[_slice])
            _fori_loop(force_while, 0, 10, loop_body, None)
        plgpu.wait_smem_to_gmem(0)
      pl.run_scoped(scope, plgpu.SMEM((10, 128), jnp.int32))
    result = kernel()
    expected = jnp.stack(
        [jnp.ones((128,), jnp.int32) * i for i in range(10)], axis=0)
    np.testing.assert_array_equal(result, expected)

  def test_nd_loop_under_warp(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((2, 32), jnp.int32),
        scratch_types=[
            plgpu.SMEM(
                (32,), jnp.int32, transforms=(plgpu.SwizzleTransform(128),)
            ),
            plgpu.Barrier()
        ],
        grid=(2,),
        grid_names=("sm",)
    )
    def kernel(in_ref, out_ref, scratch_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(in_ref, scratch_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      plgpu.commit_smem()
      def _per_warp(warp_id):
        def store(loop_info: plgpu.NDLoopInfo):
          plgpu.copy_smem_to_gmem(scratch_ref, out_ref.at[loop_info.index])
          plgpu.wait_smem_to_gmem(0)
        pl.when(warp_id == 0)(
            lambda: plgpu.nd_loop(grid=(2,), collective_axes=("sm",))(store)
        )
      plgpu.warp_map(_per_warp)

    x = jnp.arange(32, dtype=jnp.int32)
    np.testing.assert_array_equal(kernel(x), jnp.stack((x, x), axis=0))

  def test_debug_print(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128,), np.int32),
    )
    def kernel(ref):
      ref[...] = ref[...]  # Prevent kernel from being DCE'd
      @plgpu.warp_map
      def _(warp_id):
        pl.debug_print("warp: {}", warp_id)

    with self.capture_stdout() as output:
      jax.block_until_ready(kernel())
    self.assertEqual(
        set(output().splitlines()),
        {
            "warp: 0",
            "warp: 1",
            "warp: 2",
            "warp: 3",
        },
    )

  @parameterized.parameters(False, True)
  def test_copy_gmem_to_smem_from_different_warps(self,
                                                  wait_smem_to_gmem_in_warp):
    # In this test, we issue a copy from from warp 0 and await it in warp 1.
    @self.kernel(out_type=jax.ShapeDtypeStruct((32, 32), jnp.float32))
    def kernel(x_ref, y_ref):
      def scope(smem_ref, tma_barrier):
        @plgpu.warp_map
        def _(warp_id):
          @pl.when(warp_id == 0)
          def _():
            plgpu.copy_gmem_to_smem(x_ref.at[32:64], smem_ref, tma_barrier)

          @pl.when(warp_id == 1)
          def _():
            plgpu.barrier_wait(tma_barrier)
            plgpu.copy_smem_to_gmem(smem_ref, y_ref)
            if wait_smem_to_gmem_in_warp:
              plgpu.wait_smem_to_gmem(0)
        if not wait_smem_to_gmem_in_warp:
          plgpu.wait_smem_to_gmem(0)
      pl.run_scoped(scope,
                    smem_ref=plgpu.SMEM((32, 32), jnp.float32),
                    tma_barrier=plgpu.Barrier())
    x = jax.random.uniform(jax.random.key(42), (64, 32), jnp.float32)
    result = kernel(x)
    np.testing.assert_array_equal(result, x[32:64])

  def test_warp_specialized_transpose(self):
    h, w, s = 2, 4, 64
    dtype = jnp.float32

    def kernel(in_ref, out_ref, smem):
      smem[...] = in_ref[...]

      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 0)
        def _():
          t = smem.transpose((1, 0, 2))
          plgpu.copy_smem_to_gmem(t, out_ref)
          plgpu.wait_smem_to_gmem(0)

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((w, h, s), dtype),
        scratch_types=(plgpu.SMEM((h, w, s), dtype),),
    )
    x = jax.random.uniform(jax.random.key(0), shape=(h, w, s), dtype=dtype)
    np.testing.assert_array_equal(f(x), jnp.transpose(x, (1, 0, 2)))

  def test_warp_specialized_smem_slice(self):
    shape = (2, 64, 128)
    dtype = jnp.float32

    def kernel(in_ref, out_ref, smem):
      smem[...] = in_ref[...]

      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 0)
        def _():
          plgpu.copy_smem_to_gmem(smem.at[0], out_ref.at[0])
          plgpu.copy_smem_to_gmem(smem.at[1], out_ref.at[1])
          plgpu.wait_smem_to_gmem(0)

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct(shape, dtype),
        scratch_types=(plgpu.SMEM(shape, dtype),),
    )
    x = jax.random.uniform(jax.random.key(0), shape=shape, dtype=dtype)
    np.testing.assert_array_equal(f(x), x)


class PallasCallWarpPrimitiveSemanticsWGTest(
    PallasCallWarpPrimitiveSemanticsTest,
    lowering_semantics=plgpu.LoweringSemantics.Warpgroup,
):
  ...


class PallasCallWGTest(
    PallasCallTest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...

  def test_missing_primitive_lowerings_are_tracked(self):
    # This test is a way to keep track of which primitives need to be adapted
    # to using warpgroup semantics. Once the set is empty, we should be able to
    # enable warpgroup semantics by default (assuming we haven't overspecialized
    # lowerings).
    rules = mgpu_lowering.mosaic_lowering_rules
    wg_wg_lowered_primitives = set(
        rules[(plgpu.LoweringSemantics.Warpgroup,
         gpu_core.PrimitiveSemantics.Warpgroup)])
    lane_wg_lowered_primitives = set(rules[
        (plgpu.LoweringSemantics.Lane, gpu_core.PrimitiveSemantics.Warpgroup)])

    actual_missing_primitives = (lane_wg_lowered_primitives -
                                 wg_wg_lowered_primitives)
    expected_missing_primitives = set()

    self.assertSetEqual(actual_missing_primitives, expected_missing_primitives)

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_dump_layout_inference(self):
    x = jnp.ones((64, 64), dtype=jnp.float32)

    @self.kernel(out_type=jax.ShapeDtypeStruct(x.shape, x.dtype))
    def kernel(x_gmem, o_gmem):
      o_gmem[...] = plgpu.load(x_gmem, optimized=False)

    with tempfile.TemporaryDirectory() as dump_dir:
      with jtu.set_env(MOSAIC_GPU_DUMP_TO=dump_dir):
        jax.jit(kernel).lower(x)
      files = os.listdir(dump_dir)
      self.assertTrue(
          any(f.endswith(".before_layout_inference.txt") for f in files)
      )
      self.assertTrue(
          any(f.endswith(".after_layout_inference.txt") for f in files)
      )

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  @parameterized.parameters(None, plgpu.TraceScope.WARP, plgpu.TraceScope.WARPGROUP)
  def test_dump_resources(self, profile_trace_scope):
    x = jax.ShapeDtypeStruct((64, 64), jnp.float32)

    compiler_params = plgpu.CompilerParams()
    if profile_trace_scope is not None:
      compiler_params = plgpu.CompilerParams(
        profile_space=128,
        profile_dir="sponge",
        profile_trace_scope=profile_trace_scope,
      )

    @self.kernel(
        out_type=x,
        scratch_types=[
            plgpu.SMEM((64, 64), jnp.float32),
            plgpu.SMEM((32, 32), jnp.float32),
            plgpu.TMEM((128, 32), jnp.float32),
            plgpu.TMEM((128, 1), jnp.float32),
        ],
        compiler_params=compiler_params,
    )
    def kernel(x_gmem, o_gmem, *scratch_refs):
      del scratch_refs
      o_gmem[...] = plgpu.load(x_gmem, optimized=False)

    expected_smem_bytes = 64 * 64 * 4 + 32 * 32 * 4
    expected_tmem_cols = 40  # 32 + 1 = 33 padded to next multiple of 8

    if profile_trace_scope is not None:
      # For each range that we record, we need 2xi32 for each of the start and
      # end of the range---i.e. 4xi32 in total per event.
      profiler_smem_bytes = compiler_params.profile_space * 16
      if profile_trace_scope == plgpu.TraceScope.WARP:
        profiler_smem_bytes *= 4
      align = 8
      profiler_start = (expected_smem_bytes + align - 1) & ~(align - 1)
      expected_smem_bytes = profiler_start + profiler_smem_bytes

    with (jtu.set_env(MOSAIC_GPU_DUMP_RESOURCES="1"),
          self.capture_stdout() as rs):
      jax.jit(kernel).lower(x)
    self.assertIn(f"smem_scratch_bytes={expected_smem_bytes}", rs())
    self.assertIn(f"tmem_scratch_cols={expected_tmem_cols}", rs())

  def test_identical_kernels_share_a_kernel_hash(self):
    x = jnp.ones((64, 64), dtype=jnp.float32)

    # Use a wrapper to avoid hitting a high-level lowering cache.
    def make_kernel():
      @self.kernel(out_type=jax.ShapeDtypeStruct(x.shape, x.dtype))
      def kernel(x_gmem, o_gmem):
        o_gmem[...] = plgpu.load(x_gmem, optimized=False)
      return kernel

    k1, k2 = make_kernel(), make_kernel()
    ir_text = str(jax.jit(lambda x: (k1(x), k2(x))).lower(x).compiler_ir())

    kernel_hashes = re.findall(
        r'kernel_hash\s*=\s*"((?:[^"\\]|\\.)*)"', ir_text
    )
    self.assertLen(kernel_hashes, 2)
    self.assertLen(set(kernel_hashes), 1)


class PallasCallSm90ATest(PallasSm90ATest):

  @parameterized.parameters(False, True)
  def test_fori_loop_accumulator(self, force_while):
    transforms = self.default_transforms(dtype=jnp.float16)

    @functools.partial(
        self.pallas_call,
        in_specs=[plgpu.BlockSpec((64, 64), transforms=transforms)],
        out_shape=jax.ShapeDtypeStruct((64, 64), jnp.float16),
        out_specs=plgpu.BlockSpec((64, 64)),
    )
    def kernel(i_ref, o_ref):
      def scope(acc_ref):
        return _fori_loop(force_while, 0, 4, lambda _, v: v + acc_ref[...], acc_ref[...])
      plgpu.store(
          o_ref,
          pl.run_state(scope)(plgpu.ACC.init(i_ref[...])),
          optimized=False,
      )

    acc_ini = jnp.ones((64, 64), dtype=jnp.float16)
    np.testing.assert_array_equal(kernel(acc_ini), jnp.full((64, 64), 5, dtype=jnp.float16))

  @parameterized.product(lhs_transpose=[False, True], rhs_transpose=[False, True])
  def test_realistic_matmul(self, lhs_transpose, rhs_transpose):
    dtype = jnp.float16
    swizzle = 128
    elems_128b = swizzle // jnp.dtype(dtype).itemsize
    grid_m, grid_k, grid_n = 132, 10, 4
    tile_m = tile_n = 128
    tile_k = elems_128b
    m, k, n = grid_m * tile_m, grid_k * tile_k, grid_n * tile_n
    def kernel(a_ref, b_ref, o_ref, acc_ref):
      # Make sure tiling does not alter the shape of references
      if lhs_transpose:
        a_ref = a_ref.transpose((1, 0))
      assert a_ref.shape == (tile_m, tile_k)
      if rhs_transpose:
        b_ref = b_ref.transpose((1, 0))
      assert b_ref.shape == (tile_k, tile_n)
      assert o_ref.shape == acc_ref.shape == (tile_m, tile_n)
      plgpu.wgmma(acc_ref, a_ref, b_ref)
      is_last_step = pl.program_id(2) == grid_k - 1
      @pl.when(is_last_step)
      def _epilogue():
        o_ref[...] = acc_ref[...].astype(dtype)
      plgpu.wgmma_wait(1)  # We don't await the last WGMMA, hence delay_release=1

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a_shape = (k, m) if lhs_transpose else (m, k)
    a = jax.random.uniform(key1, shape=a_shape, dtype=dtype)
    b_shape = (n, k) if rhs_transpose else (k, n)
    b = jax.random.uniform(key2, shape=b_shape, dtype=dtype)

    transforms = self.default_transforms(dtype=dtype)

    if lhs_transpose:
      lhs_spec = plgpu.BlockSpec(
          (tile_k, tile_m),
          lambda m, n, k: (k, m),
          delay_release=1,
          transforms=transforms,
      )
    else:
      lhs_spec = plgpu.BlockSpec(
          (tile_m, tile_k),
          lambda m, n, k: (m, k),
          delay_release=1,
          transforms=transforms,
      )
    if rhs_transpose:
      rhs_spec = plgpu.BlockSpec(
          (tile_n, tile_k),
          lambda m, n, k: (n, k),
          delay_release=1,
          transforms=transforms,
      )
    else:
      rhs_spec = plgpu.BlockSpec(
          (tile_k, tile_n),
          lambda m, n, k: (k, n),
          delay_release=1,
          transforms=transforms,
      )
    out_spec = plgpu.BlockSpec(
        (tile_m, tile_n),
        lambda m, n, k: (m, n),
        transforms=transforms,
    )

    res = self.pallas_call(
        kernel,
        in_specs=[lhs_spec, rhs_spec],
        out_specs=out_spec,
        out_shape=jax.ShapeDtypeStruct((m, n), jnp.float16),
        scratch_shapes=[plgpu.ACC((tile_m, tile_n), jnp.float32)],
        grid=(grid_m, grid_n, grid_k),
        compiler_params=plgpu.CompilerParams(
            dimension_semantics=["parallel", "parallel", "sequential"],
            max_concurrent_steps=2,
        ),
    )(a, b)
    np.testing.assert_allclose(
        res,
        jnp.matmul(
            (a.T if lhs_transpose else a),
            (b.T if rhs_transpose else b),
            precision="high",
        ),
        rtol=1e-3,
    )

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_wgmma_fence_ptx(self):
    in_dtype, out_dtype = jnp.float16, jnp.float32
    shape = (64, 64)
    transforms = self.default_transforms(dtype=in_dtype)

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, out_dtype),
        scratch_types=[
            plgpu.SMEM(shape, in_dtype, transforms=transforms),
            plgpu.SMEM(shape, in_dtype, transforms=transforms),
        ],
    )
    def kernel(o_ref, a_smem, b_smem):
      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_smem, b_smem)
        return acc_ref[...]

      plgpu.store(
          o_ref,
          pl.run_scoped(scope, plgpu.ACC(shape, out_dtype)),
          optimized=False,
      )

    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      jax.block_until_ready(kernel())

    self.assertEqual(ptx().count("wgmma.fence.sync.aligned"), 1)

  @parameterized.parameters(jnp.float16, jnp.float32)
  def test_wgmma(self, dtype):
    # TensorCores can only fuse transposes of 16-bit values, and RHS
    # is expected to be column major by default.
    rhs_transpose = jnp.dtype(dtype).itemsize != 2
    def kernel(a_ref, b_ref, o_ref):
      if rhs_transpose:
        b_ref = b_ref.transpose((1, 0))
      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_ref, b_ref)
        return acc_ref[...]

      plgpu.store(
          o_ref,
          pl.run_scoped(scope, plgpu.ACC((64, 192), jnp.float32)),
          optimized=False,
      )

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(64, 128), dtype=dtype)
    b_shape = (128, 192)
    if rhs_transpose:
      b_shape = b_shape[::-1]
    b = jax.random.uniform(key2, shape=b_shape, dtype=dtype)

    transforms = self.default_transforms(dtype=dtype)
    res = self.pallas_call(
        kernel,
        in_specs=[
            plgpu.BlockSpec(
                (64, 128),
                lambda i, j: (i, j),
                transforms=transforms,
            ),
            plgpu.BlockSpec(
                b_shape,
                lambda *i: i,
                transforms=transforms,
            ),
        ],
        out_specs=plgpu.BlockSpec((64, 192), lambda *i: i),
        out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float32),
        grid=(1, 1),
    )(a, b)
    np.testing.assert_allclose(
        res,
        jnp.matmul(a, b.T if rhs_transpose else b, precision="high"),
        rtol=1e-3,
    )

  @parameterized.parameters(jnp.int8, jnp.uint8)
  def test_wgmma_integer(self, dtype):
    m, k, n = 64, 128, 64

    is_signed = jnp.issubdtype(dtype, jnp.signedinteger)
    acc_type = jnp.int32

    def kernel(a_ref, b_ref, o_ref):

      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_ref, b_ref.transpose((1, 0)))
        return acc_ref[...]

      plgpu.store(
          o_ref,
          pl.run_scoped(scope, plgpu.ACC((m, n), acc_type)),
          optimized=False,
      )

    # use small values to avoid overflow, [0, 8) for u8 and (-8, 8) for s8
    random_int_input = lambda key, shape: jax.random.randint(
        key, minval=-8 * is_signed, maxval=8, shape=shape, dtype=dtype
    )

    a = random_int_input(jax.random.key(0), shape=(m, k))
    b = random_int_input(jax.random.key(1), shape=(n, k))

    transforms = self.default_transforms(dtype=dtype)
    res = self.pallas_call(
        kernel,
        in_specs=[
            plgpu.BlockSpec(
                (m, k),
                lambda i, j: (i, j),
                transforms=transforms,
            ),
            plgpu.BlockSpec(
                (n, k),
                lambda *i: i,
                transforms=transforms,
            ),
        ],
        out_specs=plgpu.BlockSpec((m, n), lambda *i: i),
        out_shape=jax.ShapeDtypeStruct((m, n), acc_type),
        grid=(1, 1),
    )(a, b)
    np.testing.assert_array_equal(
        res, a.astype(acc_type) @ b.T.astype(acc_type)
    )

  @parameterized.parameters(
      (jnp.float8_e4m3fn, jnp.float8_e5m2),
      (jnp.float8_e5m2, jnp.float8_e4m3fn),
  )
  def test_wgmma_mixed_fp8(self, lhs_dtype, rhs_dtype):
    m, k, n = 64, 128, 64

    def kernel(a_ref, b_ref, o_ref):

      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_ref, b_ref.transpose((1, 0)))
        return acc_ref[...]

      plgpu.store(
          o_ref,
          pl.run_scoped(scope, plgpu.ACC((m, n), jnp.float32)),
          optimized=False,
      )

    # Small integers are exact in both e4m3 and e5m2, so the reference matmul is
    # exact and the comparison is insensitive to FP8 rounding.
    a = jax.random.randint(jax.random.key(0), (m, k), -8, 8).astype(lhs_dtype)
    b = jax.random.randint(jax.random.key(1), (n, k), -8, 8).astype(rhs_dtype)

    # Both FP8 types are 1 byte, so the LHS and RHS share the same transforms.
    transforms = self.default_transforms(dtype=lhs_dtype)
    res = self.pallas_call(
        kernel,
        in_specs=[
            plgpu.BlockSpec(
                (m, k),
                lambda i, j: (i, j),
                transforms=transforms,
            ),
            plgpu.BlockSpec(
                (n, k),
                lambda *i: i,
                transforms=transforms,
            ),
        ],
        out_specs=plgpu.BlockSpec((m, n), lambda *i: i),
        out_shape=jax.ShapeDtypeStruct((m, n), jnp.float32),
        grid=(1, 1),
    )(a, b)
    np.testing.assert_array_equal(
        res,
        jnp.matmul(
            a.astype(jnp.float32), b.T.astype(jnp.float32), precision="high"
        ),
    )

  def test_wgmma_sliced_acc_flip(self):
    dtype = jnp.float16

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(64, 128), dtype=dtype)
    b = jax.random.uniform(key2, shape=(128, 256), dtype=dtype)

    def kernel(a_ref, b_ref, o_ref):
      def scope(acc_ref):
        plgpu.wgmma(acc_ref.at[:, :128], a_ref, b_ref.at[:, 128:])
        plgpu.wgmma(acc_ref.at[:, 128:], a_ref, b_ref.at[:, :128])
        return acc_ref[...]

      plgpu.store(
          o_ref,
          pl.run_scoped(scope, plgpu.ACC((64, 256), jnp.float32)),
          optimized=False,
      )

    transforms = self.default_transforms(dtype=dtype)
    res = self.pallas_call(
        kernel,
        in_specs=[plgpu.BlockSpec(transforms=transforms)] * 2,
        out_shape=jax.ShapeDtypeStruct((64, 256), jnp.float32),
    )(a, b)

    def flip_halves(x):
      y = x.reshape(*x.shape[:-1], 2, x.shape[-1] // 2)
      y = y[..., ::-1, :]
      return y.reshape(x.shape)

    np.testing.assert_allclose(
        res, jnp.matmul(a, flip_halves(b), precision="high"), rtol=1e-3
    )

  def test_wgmma_registers(self):
    def kernel(a_ref, b_ref, o_ref):
      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_ref[...], b_ref)
        return acc_ref[...]
      plgpu.store(
          o_ref,
          pl.run_scoped(scope, plgpu.ACC((64, 192), jnp.float32)),
          optimized=False,
      )

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(64, 128), dtype=jnp.float16)
    b = jax.random.uniform(key2, shape=(128, 192), dtype=jnp.float16)

    transforms = self.default_transforms(dtype=jnp.float16)
    res = self.pallas_call(
        kernel,
        in_specs=[
            plgpu.BlockSpec(transforms=transforms),
            plgpu.BlockSpec(transforms=transforms),
        ],
        out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float32),
    )(a, b)
    np.testing.assert_allclose(
        res, jnp.matmul(a, b, precision="high"), rtol=1e-3
    )

  @parameterized.parameters(jnp.int8, jnp.float8_e4m3fn, jnp.float8_e5m2)
  def test_wgmma_registers_8bit(self, input_dtype):
    if jnp.issubdtype(input_dtype, jnp.integer):
      out_dtype = jnp.int32
    else:
      out_dtype = jnp.float32
    def kernel(a_ref, b_ref, o_ref):
      def scope(acc_ref):
        a_regs = plgpu.load(a_ref, layout=plgpu.Layout.WGMMA_8BIT)
        plgpu.wgmma(acc_ref, a_regs, b_ref.transpose((1, 0)))
        return acc_ref[...]
      plgpu.store(
          o_ref,
          pl.run_scoped(scope, plgpu.ACC((64, 192), out_dtype)),
          optimized=False,
      )

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    m = 64
    k = 128
    n = 192
    if input_dtype == jnp.int8:
      a = jax.random.randint(key1, shape=(m, k), minval=-128, maxval=127, dtype=jnp.int8)
      b = jax.random.randint(key2, shape=(n, k), minval=-128, maxval=127, dtype=jnp.int8)
    else:
      assert jnp.issubdtype(input_dtype, jnp.floating)
      a = jax.random.uniform(key1, shape=(m, k), dtype=input_dtype)
      b = jax.random.uniform(key2, shape=(n, k), dtype=input_dtype)

    transforms = self.default_transforms(swizzle=64, dtype=input_dtype)
    res = self.pallas_call(
        kernel,
        in_specs=[
            plgpu.BlockSpec(transforms=transforms),
            plgpu.BlockSpec(transforms=transforms),
        ],
        out_shape=jax.ShapeDtypeStruct((64, 192), out_dtype),
    )(a, b)
    ref = jnp.matmul(
        a.astype(out_dtype), b.T.astype(out_dtype), precision="high"
    )
    if input_dtype == jnp.int8:
      np.testing.assert_array_equal(res, ref)
    else:
      np.testing.assert_allclose(res, ref)

  def test_wgmma_registers_init(self):
    def kernel(a_ref, b_ref, i_ref, o_ref):
      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_ref[...], b_ref)
      plgpu.store(
          o_ref,
          pl.run_state(scope)(plgpu.ACC.init(i_ref[...])),
          optimized=False,
      )

    key1, key2, key3 = jax.random.split(jax.random.key(42), 3)
    a = jax.random.uniform(key1, shape=(64, 128), dtype=jnp.float16)
    b = jax.random.uniform(key2, shape=(128, 192), dtype=jnp.float16)
    i = jax.random.uniform(key3, shape=(64, 192), dtype=jnp.float16) * 10

    transforms = self.default_transforms(dtype=jnp.float16)
    res = self.pallas_call(
        kernel,
        in_specs=[
            plgpu.BlockSpec(transforms=transforms),
            plgpu.BlockSpec(transforms=transforms),
            plgpu.BlockSpec(transforms=transforms),
        ],
        out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float16),
    )(a, b, i)
    np.testing.assert_allclose(
        res, i + jnp.matmul(a, b, precision="high"), rtol=2e-3
    )

  def test_run_state_discharge_regression_test(self):
    # This regression test voluntarily allocates an outer accumulator and a
    # SMEM ref using `run_scoped``, and then uses them both inside of
    # `run_state`. This used to trigger two bugs:
    # 1. The SMEM ref would end up being unexpectedly discharged inside of
    #    `run_state`;
    # 2. The dereferencing of the outer accumulator would be improperly handled
    #    in the results of `run_state`.
    transforms = self.default_transforms(dtype=jnp.float16)

    def kernel(a_ref, b_ref, o_ref, barrier_ref):
      def outer_scope(b_smem, acc):
        plgpu.copy_gmem_to_smem(b_ref, b_smem, barrier_ref)
        plgpu.barrier_wait(barrier_ref)

        def inner_scope(inner_acc):
          del inner_acc
          a_regs = plgpu.load(a_ref, layout=plgpu.Layout.WGMMA, optimized=False)
          plgpu.wgmma(acc, a_regs, b_smem)
        # Use an inner accumulator because `run_state` lowering expects it.
        pl.run_state(inner_scope)(plgpu.ACC.init(
            plgpu.layout_cast(jnp.zeros((64, 64), jnp.float32), plgpu.Layout.WGMMA)))
        plgpu.store(o_ref, acc[...], optimized=False)
      pl.run_scoped(outer_scope,
                    plgpu.SMEM(b_ref.shape, jnp.float16, transforms=transforms),
                    plgpu.ACC((64, 192), jnp.float32))

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(64, 128), dtype=jnp.float16)
    b = jax.random.uniform(key2, shape=(128, 192), dtype=jnp.float16)

    res = self.kernel(
        kernel,
        scratch_types=[plgpu.Barrier()],
        out_type=jax.ShapeDtypeStruct((64, 192), jnp.float32),
    )(a, b)
    np.testing.assert_allclose(
        res, jnp.matmul(a, b, precision="high"), rtol=2e-3
    )

  def test_wgmma_registers_init_with_update(self):
    def kernel(a_ref, b_ref, i_ref, o_ref):
      def scope(acc_ref):
        acc_ref[...] *= 2
        plgpu.wgmma(acc_ref, a_ref[...], b_ref)
      plgpu.store(
          o_ref,
          pl.run_state(scope)(plgpu.ACC.init(i_ref[...])),
          optimized=False,
      )

    key1, key2, key3 = jax.random.split(jax.random.key(42), 3)
    a = jax.random.uniform(key1, shape=(64, 128), dtype=jnp.float16)
    b = jax.random.uniform(key2, shape=(128, 192), dtype=jnp.float16)
    i = jax.random.uniform(key3, shape=(64, 192), dtype=jnp.float16) * 10

    transforms = self.default_transforms(dtype=jnp.float16)
    res = self.pallas_call(
        kernel,
        in_specs=[
            plgpu.BlockSpec(transforms=transforms),
            plgpu.BlockSpec(transforms=transforms),
            plgpu.BlockSpec(transforms=transforms),
        ],
        out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float16),
    )(a, b, i)
    np.testing.assert_allclose(
        res, i * 2 + jnp.matmul(a, b, precision="high"), rtol=2e-3
    )

  def test_wgmma_sliced_ref(self):
    def kernel(a_ref, b_ref, o_ref):
      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_ref.at[0], b_ref.at[0])
        return acc_ref[...]

      plgpu.store(
          o_ref,
          pl.run_scoped(scope, plgpu.ACC((64, 192), jnp.float32)),
          optimized=False,
      )

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(2, 64, 128), dtype=jnp.float16)
    b = jax.random.uniform(key2, shape=(2, 128, 192), dtype=jnp.float16)

    transforms = self.default_transforms(dtype=jnp.float16)
    res = self.pallas_call(
        kernel,
        in_specs=[
            plgpu.BlockSpec(transforms=transforms),
            plgpu.BlockSpec(transforms=transforms),
        ],
        out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float32),
    )(a, b)
    np.testing.assert_allclose(
        res, jnp.matmul(a[0], b[0], precision="high"), rtol=1e-3
    )

  def test_collapsing_several_non_unit_tiled_dimensions_raises(self):
    def kernel(out, smem):
      out[...] = smem.reshape(128, 128)[...]

    with self.assertRaisesRegex(
        ValueError,
        "any of the dimensions being collapsed other than the minormost one "
        "has non-unit tiling"
    ):
      self.kernel(
          kernel,
          scratch_types=[
              plgpu.SMEM(
                  (128, 4, 32),
                  jnp.float16,
                  transforms=(
                      plgpu.TilingTransform((2, 16)),
                      plgpu.SwizzleTransform(32),
                  ),
              ),
          ],
          out_type=jax.ShapeDtypeStruct((128, 128), jnp.float16),
      )()

  def test_wgmma_no_valid_tiling_raises(self):
    if not self.is_wg_semantics():
      self.skipTest("This test is only for WG semantics.")

    @self.kernel(
        scratch_types=[
            plgpu.SMEM((64, 12), jnp.float16),
            plgpu.SMEM((12, 64), jnp.float16),
        ],
        out_type=jax.ShapeDtypeStruct((64, 64), jnp.float32),
    )
    def kernel(out, a_smem, b_smem):
      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_smem, b_smem)
        return acc_ref[...]

      out[...] = pl.run_scoped(scope, plgpu.ACC((64, 64), jnp.float32))

    with self.assertRaisesRegex(
        ValueError,
        "No valid tiling found for SMEM reference. \\(12, 64\\) not divisible"
        " by candidate tiling \\(8, 16\\)",
    ):
      jax.jit(kernel).lower()

  def test_wgmma_reshaped_rhs_unit_tiling(self):
    def kernel(lhs_ref, rhs_ref, out_ref, lhs_smem, rhs_smem, barrier_ref):
      plgpu.copy_gmem_to_smem(lhs_ref, lhs_smem, barrier_ref)
      plgpu.copy_gmem_to_smem(rhs_ref, rhs_smem, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      def scope(acc_ref):
        plgpu.wgmma(acc_ref, lhs_smem, rhs_smem.reshape(128, 128))
        return acc_ref[...]
      plgpu.store(
          out_ref,
          pl.run_scoped(scope, plgpu.ACC((64, 128), jnp.float32)),
          optimized=False,
      )

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(64, 128), dtype=jnp.float16)
    b = jax.random.uniform(key2, shape=(128, 2, 64), dtype=jnp.float16)

    lhs_transforms = self.default_transforms(dtype=jnp.float16)
    rhs_transforms = (plgpu.TilingTransform((8, 1, 64)), plgpu.SwizzleTransform(128))

    res = self.kernel(
        kernel,
        scratch_types=[
            plgpu.SMEM((64, 128), jnp.float16, transforms=lhs_transforms),
            plgpu.SMEM((128, 2, 64), jnp.float16, transforms=rhs_transforms),
            plgpu.Barrier(num_arrivals=2),
        ],
        out_type=jax.ShapeDtypeStruct((64, 128), jnp.float32),
    )(a, b)
    np.testing.assert_allclose(
        res, jnp.matmul(a, b.reshape(128, 128), precision="high"), rtol=1e-3
    )

  def test_wgmma_sliced_acc_read(self):
    def kernel(a_ref, b_ref, o_ref):
      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_ref, b_ref)
        return acc_ref[:, :64], acc_ref[:, 64:]

      out0, out1 = pl.run_scoped(scope, plgpu.ACC((64, 128), jnp.float32))
      plgpu.store(o_ref.at[:, :64], out0, optimized=False)
      plgpu.store(o_ref.at[:, 64:], out1, optimized=False)

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(64, 128), dtype=jnp.float16)
    b = jax.random.uniform(key2, shape=(128, 128), dtype=jnp.float16)
    transforms = self.default_transforms(dtype=jnp.float16)
    res = self.pallas_call(
        kernel,
        in_specs=[
            plgpu.BlockSpec((64, 128), lambda *ij: ij, transforms=transforms),
            plgpu.BlockSpec((128, 128), lambda *ij: ij, transforms=transforms),
        ],
        out_specs=plgpu.BlockSpec((64, 128), lambda *ij: ij),
        out_shape=jax.ShapeDtypeStruct((64, 128), jnp.float32),
        grid=(1, 1),
    )(a, b)
    np.testing.assert_allclose(
        res, jnp.matmul(a, b, precision="high"), rtol=1e-3
    )

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_wgmma_accumulator_load_with_custom_wait(self):
    def kernel(a_ref, b_ref, o_ref0, o_ref1):
      def scope(acc0_ref, acc1_ref):
        plgpu.wgmma(acc0_ref, a_ref, b_ref)
        plgpu.wgmma(acc1_ref, a_ref, b_ref)
        acc0 = plgpu.wgmma_accumulator_load(acc0_ref, wait_n=1)
        acc1 = plgpu.wgmma_accumulator_load(acc1_ref, wait_n=0)
        return acc0, acc1
      out0, out1 = pl.run_scoped(
          scope, plgpu.ACC((64, 64), jnp.float32), plgpu.ACC((64, 64), jnp.float32)
      )
      plgpu.store(o_ref0, out0, optimized=False)
      plgpu.store(o_ref1, out1, optimized=False)

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(64, 64), dtype=jnp.float16)
    b = jax.random.uniform(key2, shape=(64, 64), dtype=jnp.float16)
    transforms = self.default_transforms(dtype=jnp.float16)
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      out0, out1 = jax.block_until_ready(self.pallas_call(
          kernel,
          in_specs=[
              plgpu.BlockSpec(transforms=transforms),
              plgpu.BlockSpec(transforms=transforms),
          ],
          out_shape=(jax.ShapeDtypeStruct((64, 64), jnp.float32),
                     jax.ShapeDtypeStruct((64, 64), jnp.float32)),
          grid=(1, 1),
      )(a, b))
    ptx_str = ptx()
    wgmma_waits = re.findall(r"wgmma.wait_group.sync.aligned\s+(\d+);", ptx_str)
    self.assertLen(wgmma_waits, 2)
    self.assertEqual(int(wgmma_waits[0]), 1)
    self.assertEqual(int(wgmma_waits[1]), 0)

    self.assertArraysEqual(out0, out1)
    self.assertAllClose(
        out0,
        jnp.dot(a, b, preferred_element_type=jnp.float32, precision="high"),
        rtol=1e-3,
    )

  @parameterized.product(
      src_memory_space=[plgpu.SMEM, plgpu.GMEM],
      layout=[plgpu.Layout.WGMMA.reduce(1), plgpu.Layout.WGMMA.reduce(0)],
      m=[64, 128, 192],
  )
  def test_load_to_wgmma_row_col_layout_with_indexing(self, src_memory_space, layout, m):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([2, m], jnp.float32),
        in_specs=[pl.BlockSpec(memory_space=src_memory_space)],
        out_specs=plgpu.BlockSpec(memory_space=plgpu.SMEM),
    )
    def kernel(x_ref, o_ref):
      for i in range(2):
        x = plgpu.load(
            x_ref.at[i], layout=layout, optimized=src_memory_space == plgpu.SMEM
        )
        o_ref[i, ...] = x

    x = jnp.arange(2 * m, dtype=jnp.float32).reshape(2, m)
    np.testing.assert_array_equal(kernel(x), x)

  @parameterized.product(
      src_memory_space=[plgpu.SMEM],
      layout=[plgpu.Layout.WGMMA.reduce(1), plgpu.Layout.WGMMA.reduce(0)],
  )
  def test_load_row_input_to_wgmma_with_transforms(self, src_memory_space, layout):
    m, k, n = 64, 128, 192
    key1, key2 = jax.random.split(jax.random.key(42), 2)
    if layout == plgpu.Layout.WGMMA.reduce(1):
      input_shape = (m,)
      broadcast_dim = 0
      expand_dim = 1
    else:
      input_shape = (k,)
      broadcast_dim = 1
      expand_dim = 0
    a = jax.random.uniform(key1, shape=input_shape, dtype=jnp.float16)
    b = jax.random.uniform(key2, shape=(k, n), dtype=jnp.float16)
    def kernel(x_ref, y_ref, o_ref):
      x = plgpu.load(x_ref, layout=layout)
      x = lax.broadcast_in_dim(x, (m, k), [broadcast_dim])

      def compute(acc_ref):
        plgpu.wgmma(acc_ref, x, y_ref)
        return acc_ref[...]

      out = pl.run_scoped(compute, plgpu.ACC((m, n), jnp.float32))
      plgpu.store(o_ref, out, optimized=False)
    f = self.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct([m, n], jnp.float32),
        in_specs=(
            pl.BlockSpec(memory_space=src_memory_space),
            plgpu.BlockSpec(
                transforms=self.default_transforms(dtype=jnp.float16),
            ),
        ),
        out_specs=plgpu.BlockSpec(memory_space=plgpu.SMEM),
    )

    out_ref = jnp.matmul(
        jnp.broadcast_to(jnp.expand_dims(a, axis=expand_dim), (m, k)),
        b,
        precision="high",
    )
    np.testing.assert_allclose(f(a, b), out_ref, rtol=1e-3)

  def test_load_store_wgmma_transposed(self):
    transforms = self.default_transforms(swizzle=64, dtype=jnp.float32)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct([8, 64], jnp.float32),
        scratch_types=[
            plgpu.SMEM((8, 64), jnp.float32, transforms=transforms),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_gmem, o_ref, x_smem, barrier):
      plgpu.copy_gmem_to_smem(x_gmem, x_smem, barrier)
      plgpu.barrier_wait(barrier)
      x = plgpu.load(x_smem.T, layout=plgpu.Layout.WGMMA_TRANSPOSED)
      x_smem.T[...] = x + 1
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(x_smem, o_ref)
      plgpu.wait_smem_to_gmem(0)

    x = jax.random.uniform(jax.random.key(42), shape=(8, 64), dtype=jnp.float32)
    result = kernel(x)
    np.testing.assert_array_equal(result, x + 1)


class PallasCallSm90AWGTest(
    PallasCallSm90ATest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


class PallasCallTCGen05Test(PallasTCGen05Test):

  @parameterized.product(
      orders_tensor_core=(False, True),
      single_warp=(False, True),
  )
  def test_copy_gmem_to_smem_orders_tc(self, orders_tensor_core, single_warp):
    @self.kernel(
        out_type=(
            jax.ShapeDtypeStruct([256], jnp.float32),
            jax.ShapeDtypeStruct([1], jnp.int32),
        ),
        scratch_types=[
            plgpu.SMEM((256,), jnp.float32),
            plgpu.Barrier(num_arrivals=2, orders_tensor_core=orders_tensor_core),
        ],
    )
    def kernel(x_ref, o_ref, completed_early_ref, scratch_ref, barrier_ref):
      def copy_to_smem():
        plgpu.copy_gmem_to_smem(x_ref, scratch_ref, barrier_ref)

      if single_warp:
        plgpu.warp_map(lambda warp_id: pl.when(warp_id == 0)(copy_to_smem))
      else:
        copy_to_smem()
      completed_early_ref[0] = plgpu.barrier_test(barrier_ref).astype(jnp.int32)
      plgpu.barrier_arrive(barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      o_ref[...] = scratch_ref[...] + 1

    x = jnp.arange(256).astype(jnp.float32)
    out, completed_early = kernel(x)
    np.testing.assert_array_equal(out, x + 1.0)
    # Regression test to ensure we don't arrive more than expected when
    # orders_tensor_core=True.
    np.testing.assert_array_equal(
        completed_early, jnp.zeros((1,), dtype=jnp.int32)
    )

  def test_warp_specialized_tmem_slice(self):
    dtype = jnp.float32

    def kernel(in_ref, out_ref, smem, tmem, barrier):
      smem[...] = plgpu.load(in_ref, layout=plgpu.Layout.WGMMA, optimized=False)
      # This is a regression test which requires a dynamic, foldable offset.
      # This allows us to produce a dynamic offset under the hood that later
      # gets simplified away during MLIR canonicalization passes.
      c0 = lax.axis_index("x") * 0

      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 0)
        def _():
          plgpu.async_copy_smem_to_tmem(smem, tmem.at[:, pl.ds(c0, 32)])
          plgpu.tcgen05_commit_arrive(barrier)
      plgpu.barrier_wait(barrier)

      plgpu.store(
          out_ref,
          plgpu.async_load_tmem(tmem.at[:, pl.ds(c0, 32)]),
          optimized=False,
      )

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((128, 32), dtype),
        grid=(1,),
        grid_names=("x",),
        scratch_types=(
            plgpu.SMEM((128, 32), dtype,
                       transforms=self.default_transforms(dtype=dtype)),
            plgpu.TMEM((128, 64), dtype),
            plgpu.Barrier(orders_tensor_core=True)),
    )
    x = jax.random.uniform(jax.random.key(0), shape=(128, 32), dtype=dtype)
    np.testing.assert_array_equal(f(x), x)

  def test_print_layout_tmem(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((), jnp.bfloat16),
        scratch_types=[plgpu.TMEM((128, 64), jnp.bfloat16, packed=True)],
    )
    def kernel(o_ref, tmem_ref):
      del o_ref
      plgpu.print_layout("tmem: {}", tmem_ref)

    with self.capture_stdout() as output:
      jax.jit(kernel).lower()
    self.assertIn("tmem: TMEM_DEFAULT(packing=2)\n", output())

  def test_print_layout_tmem_with_transform(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((), jnp.bfloat16),
        scratch_types=[plgpu.TMEM((128, 256), jnp.bfloat16, packed=True)],
    )
    def kernel(o_ref, tmem_ref):
      del o_ref
      # Slicing TMEM to make sure we handle transforms correctly.
      plgpu.print_layout("tmem: {}", tmem_ref.at[:, :128])

    with self.capture_stdout() as output:
      jax.jit(kernel).lower()
    self.assertIn("tmem: TMEM_DEFAULT(packing=2)\n", output())

  def test_print_layout_tmem_with_transforms(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((), jnp.float32),
        scratch_types=[
            plgpu.RefUnion(
                plgpu.TMEM((128, 64), jnp.float16, packed=False),
                plgpu.TMEM((128, 128), jnp.float16, packed=True),
            ),
        ],
    )
    def kernel(o_ref, aliased_ref):
      del o_ref
      [ref0, ref1] = aliased_ref
      plgpu.print_layout("ref0: {}", ref0)
      plgpu.print_layout("ref1: {}", ref1)

    with self.capture_stdout() as output:
      kernel().block_until_ready()

    self.assertIn("ref0: TMEM_DEFAULT(packing=1)", output())
    self.assertIn("ref1: TMEM_DEFAULT(packing=2)", output())

  def test_mixed_tmem_allocations_raise(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct((), jnp.float32),
        scratch_shapes=[
            plgpu.TMEM((128, 128), jnp.float32, collective=True),
            plgpu.TMEM((128, 128), jnp.float32, collective=False),
        ],
    )
    def kernel(out_ref, tmem_ref0, tmem_ref1):
      del out_ref, tmem_ref0, tmem_ref1

    with self.assertRaisesRegex(
        ValueError,
        "Can't mix collective and non-collective TMEM allocations within the"
        " same kernel.",
    ):
      kernel()

  def test_transposed_tmem_ref_raises(self):
    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct((128, 128), jnp.float32),
        scratch_shapes=[
            plgpu.TMEM((128, 128), jnp.float32),
            plgpu.SMEM((128, 128), jnp.float32),
        ],
    )
    def kernel(out, tmem_ref, smem_ref):
      del out
      smem_ref[...] = plgpu.async_load_tmem(tmem_ref.transpose((1, 0)))

    with self.assertRaisesRegex(
        NotImplementedError, "Unimplemented transforms for TMEM refs"
    ):
      kernel()

  @parameterized.parameters((False,), (True,))
  def test_tmem(self, collective):
    transforms = self.default_transforms(dtype=jnp.float32)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.float32),
        scratch_types=[
            plgpu.TMEM((128, 128), jnp.float32, collective=collective),
            plgpu.TMEM((128, 128), jnp.float32, collective=collective),
            plgpu.SMEM((128, 128), jnp.float32, transforms=transforms),
            plgpu.Barrier(),
        ],
        cluster=(2,) if collective else (),
        cluster_names=("x",) if collective else (),
    )
    def kernel(x_ref, y_ref, tmem_ref, tmem_ref2, smem_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(x_ref, smem_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      # Exercise TMEM by roundtripping SMEM -> TMEM -> TMEM -> SMEM.
      x_val = plgpu.load(smem_ref, layout=plgpu.Layout.TCGEN05)
      plgpu.async_store_tmem(tmem_ref, x_val + 1)
      plgpu.commit_tmem()
      #  We don't await the load, because we never overwrite tmem_ref
      tmem_read = plgpu.async_load_tmem(tmem_ref)
      plgpu.async_store_tmem(tmem_ref2, tmem_read)
      plgpu.commit_tmem()
      #  We don't await the load, because we never overwrite tmem_ref2
      smem_ref[...] = plgpu.async_load_tmem(tmem_ref2)
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(smem_ref, y_ref)
      plgpu.wait_smem_to_gmem(0)

    x = jax.random.uniform(
        jax.random.key(0), shape=(128, 128), dtype=jnp.float32)
    x_result = jax.block_until_ready(kernel(x))
    np.testing.assert_array_equal(x_result, x + 1)

  def test_tmem_allocation_estimation(self):
    """Make sure that we don't overestimate the TMEM allocation.

    All of the refs below are packed and should fit into TMEM at once.
    """
    transforms = self.default_transforms(dtype=jnp.bfloat16)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128, 256), jnp.bfloat16),
        scratch_types=[
            plgpu.TMEM((128, 256), jnp.bfloat16, packed=True),
            plgpu.TMEM((128, 256), jnp.bfloat16, packed=True),
            plgpu.TMEM((128, 256), jnp.bfloat16, packed=True),
            plgpu.SMEM((128, 256), jnp.bfloat16, transforms=transforms),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref, y_ref, tmem_ref1, tmem_ref2, tmem_ref3, smem_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(x_ref, smem_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      x_val = plgpu.load(smem_ref, layout=plgpu.Layout.TCGEN05)
      plgpu.async_store_tmem(tmem_ref1, x_val + 1)
      plgpu.commit_tmem()
      x_val = plgpu.async_load_tmem(tmem_ref1)
      plgpu.async_store_tmem(tmem_ref2, x_val + 1)
      plgpu.commit_tmem()
      x_val = plgpu.async_load_tmem(tmem_ref2)
      plgpu.async_store_tmem(tmem_ref3, x_val + 1)
      plgpu.commit_tmem()
      smem_ref[...] = plgpu.async_load_tmem(tmem_ref3)
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(smem_ref, y_ref)
      plgpu.wait_smem_to_gmem(0)

    x = jax.random.uniform(jax.random.key(0), shape=(128, 256), dtype=jnp.bfloat16)
    x_result = jax.block_until_ready(kernel(x))
    np.testing.assert_array_equal(x_result, x + 3)

  def test_tmem_ref_aliasing(self):
    transforms = self.default_transforms(dtype=jnp.float32)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.float32),
        scratch_types=[
            plgpu.RefUnion(
                [
                    plgpu.TMEM((128, 32), jnp.float32),
                    plgpu.TMEM((128, 32), jnp.float32),
                ],
                plgpu.TMEM((128, 64), jnp.float32),
            ),
            plgpu.SMEM((128, 128), jnp.float32, transforms=transforms),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref, y_ref, aliased_ref, smem_ref, barrier_ref):
      [tmem_128x32a, tmem_128x32b], tmem_128x64 = aliased_ref
      plgpu.copy_gmem_to_smem(x_ref, smem_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      # Test tmem_128x32 a and b
      x_val = plgpu.load(smem_ref.at[:, 0:32], layout=plgpu.Layout.TCGEN05)
      plgpu.async_store_tmem(tmem_128x32a, x_val + 1)
      plgpu.commit_tmem()
      smem_ref[:, 0:32] = plgpu.async_load_tmem(tmem_128x32a)
      plgpu.wait_load_tmem()  # Make sure the load is done before we write to TMEM again.

      x_val = plgpu.load(smem_ref.at[:, 32:64], layout=plgpu.Layout.TCGEN05)
      plgpu.async_store_tmem(tmem_128x32b, x_val + 1)
      plgpu.commit_tmem()
      smem_ref[:, 32:64] = plgpu.async_load_tmem(tmem_128x32b)
      plgpu.wait_load_tmem()  # Make sure the load is done before we write to TMEM again.

      # Test tmem_128x64
      x_val = plgpu.load(smem_ref.at[:, 64:128], layout=plgpu.Layout.TCGEN05)
      plgpu.async_store_tmem(tmem_128x64, x_val + 1)
      plgpu.commit_tmem()
      smem_ref[:, 64:128] = plgpu.async_load_tmem(tmem_128x64)

      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(smem_ref, y_ref)
      plgpu.wait_smem_to_gmem(0)
    x = jax.random.uniform(
        jax.random.key(0), shape=(128, 128), dtype=jnp.float32)
    x_result = jax.block_until_ready(kernel(x))
    np.testing.assert_array_equal(x_result, x + 1)

  def test_tmem_ref_union_3d(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128, 32), jnp.float32),
        scratch_types=[
            plgpu.RefUnion(
                plgpu.TMEM((2, 128, 16), jnp.float32),
                plgpu.TMEM((128, 32), jnp.float32),
            ),
        ],
    )
    def kernel(o_ref, tmem_ref_union):
      tmem0, tmem1 = tmem_ref_union
      # Layout cast is necessary for the test to pass under lane semantics.
      regs = plgpu.layout_cast(jnp.ones_like(tmem0), plgpu.Layout.TCGEN05)
      plgpu.async_store_tmem(tmem0, regs)
      plgpu.store(o_ref, plgpu.async_load_tmem(tmem1), optimized=False)

    np.testing.assert_array_equal(kernel(), jnp.ones((128, 32), dtype=jnp.float32))

  def test_tmem_ref_union_multi_allocation(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128, 512), jnp.float16),
        scratch_types=[
            plgpu.TMEM((128, 256), jnp.float16, packed=False),
            plgpu.RefUnion(
                plgpu.TMEM((128, 256), jnp.float16, packed=False),
            ),
        ],
    )
    def kernel(x_ref, y_ref, tmem_ref1, aliased_tmem_ref2):
      [tmem_ref2] = aliased_tmem_ref2
      x_val1 = plgpu.load(x_ref.at[:, 0:256], layout=plgpu.Layout.TCGEN05, optimized=False)
      x_val2 = plgpu.load(x_ref.at[:, 256:512], layout=plgpu.Layout.TCGEN05, optimized=False)
      plgpu.async_store_tmem(tmem_ref1, x_val1 + 1)
      plgpu.async_store_tmem(tmem_ref2, x_val2 + 2)
      plgpu.commit_tmem()

      plgpu.store(
          y_ref.at[:, 0:256], plgpu.async_load_tmem(tmem_ref1), optimized=False
      )
      plgpu.store(
          y_ref.at[:, 256:512],
          plgpu.async_load_tmem(tmem_ref2),
          optimized=False,
      )

    x = jax.random.uniform(
        jax.random.key(0), shape=(128, 512), dtype=jnp.float16)
    expected = jnp.concatenate([x[:, 0:256] + 1, x[:, 256:512] + 2], axis=1)
    np.testing.assert_array_equal(kernel(x), expected)

  @parameterized.parameters(
      plgpu.Layout.TCGEN05, plgpu.Layout.TCGEN05_TMEM_NATIVE
  )
  def test_tmem_load_layout(self, layout):
    transforms = self.default_transforms(dtype=jnp.float32)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.float32),
        scratch_types=[
            plgpu.TMEM((128, 128), jnp.float32),
            plgpu.SMEM((128, 128), jnp.float32, transforms=transforms),
            plgpu.Barrier(),
        ],
    )
    def kernel(x_ref, y_ref, tmem_ref, smem_ref, barrier_ref):
      plgpu.copy_gmem_to_smem(x_ref, smem_ref, barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      optimized = layout != plgpu.Layout.TCGEN05_TMEM_NATIVE
      x_val = plgpu.load(smem_ref, layout=layout, optimized=optimized)
      plgpu.async_store_tmem(tmem_ref, x_val + 1)
      plgpu.commit_tmem()
      # We don't wait for the load to complete, because we never overwrite
      # tmem_ref.
      smem_ref[...] = plgpu.async_load_tmem(tmem_ref, layout=layout)
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(smem_ref, y_ref)
      plgpu.wait_smem_to_gmem(0)

    x = jax.random.uniform(
        jax.random.key(0), shape=(128, 128), dtype=jnp.float32)
    x_result = jax.block_until_ready(kernel(x))
    np.testing.assert_array_equal(x_result, x + 1)

  @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

    @self.kernel(
        out_type=(
            jax.ShapeDtypeStruct(shape, dtype),
            jax.ShapeDtypeStruct((shape[0],), dtype),
        ),
        scratch_types=[plgpu.TMEM(shape, dtype)],
    )
    def kernel(x_ref, y_ref, red_y_ref, tmem_ref):
      x_val = plgpu.load(x_ref, layout=plgpu.Layout.TCGEN05, optimized=False)
      plgpu.async_store_tmem(tmem_ref, x_val)
      plgpu.commit_tmem()
      loaded, reduced = plgpu.async_load_tmem(
          tmem_ref,
          layout=plgpu.Layout.TCGEN05_TMEM_NATIVE,
          reduce=reduction,
      )
      plgpu.store(y_ref, loaded, optimized=False)
      plgpu.store(red_y_ref, reduced, optimized=False)

    prng = np.random.default_rng(0)
    if dtype == jnp.int32:
      x = prng.integers(-100000, 100000, shape).astype(dtype)
    elif dtype == jnp.uint32:
      x = prng.integers(0, 1000000, shape).astype(dtype)
    else:
      x = prng.uniform(-1, 1, shape).astype(dtype)
    y, red_y = jax.block_until_ready(kernel(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(
      plgpu.Layout.TCGEN05_M64_COLLECTIVE(160),
      plgpu.Layout.TCGEN05_M64_COLLECTIVE_NATIVE(160)
  )
  def test_tmem_store_load_collective(self, layout):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((64, 160), jnp.float32),
        cluster=(2,),
        cluster_names=("cluster",),
        scratch_types=[
            plgpu.TMEM(
                (64, 160),
                jnp.float32,
                collective=True,
                layout=plgpu.TMEMLayout.M64_COLLECTIVE_LAYOUT(160),
            ),
        ],
    )
    def kernel(x_ref, y_ref, tmem_ref):
      x_val = plgpu.load(x_ref, layout=layout, optimized=False)
      plgpu.async_store_tmem(tmem_ref, x_val + 1)
      plgpu.commit_tmem()
      # We don't wait for the load to complete, because we never overwrite
      # tmem_ref.
      plgpu.store(
          y_ref, plgpu.async_load_tmem(tmem_ref, layout=layout), optimized=False
      )

    x = jax.random.uniform(
        jax.random.key(0), shape=(64, 160), dtype=jnp.float32)
    x_result = jax.block_until_ready(kernel(x))
    np.testing.assert_array_equal(x_result, x + 1)

  def test_tmem_column_slicing(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.float32),
        scratch_types=[
            plgpu.TMEM((128, 256), jnp.float32),
        ],
    )
    def kernel(x_ref, y_ref, tmem_ref):
      x_val = plgpu.load(x_ref, layout=plgpu.Layout.TCGEN05, optimized=False)
      tmem_slice = tmem_ref.at[:, 8:208].at[:, 0:128]
      plgpu.async_store_tmem(tmem_slice, x_val + 1)
      plgpu.commit_tmem()
      plgpu.store(
          y_ref, plgpu.async_load_tmem(tmem_ref.at[:, 8:136]), optimized=False
      )

    x = jax.random.uniform(
        jax.random.key(0), shape=(128, 128), dtype=jnp.float32)
    np.testing.assert_array_equal(kernel(x), (x + 1)[:, 0:128])

  def test_tmem_batch_dimension_indexing(self):
    b1, b2, m, n = 2, 3, 128, 4
    shape = (b1, b2, m, n)
    dtype = jnp.float8_e5m2
    tmem_layout = plgpu.TMEMLayout.SCALES_LAYOUT
    reg_layout = plgpu.Layout.TILED(  # tmem_layout.as_tiled_layout()
        plgpu.Tiling(((128, 4), (32, 1))),
        warp_dims=(plgpu.Replicated(4),),
        lane_dims=(-2,),
        vector_dim=-3,
    )

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((m, n), dtype),
        scratch_types=[
            plgpu.TMEM(shape, dtype, layout=tmem_layout),
        ],
    )
    def kernel(src_ref, idx1_ref, idx2_ref, dst_ref, tmem_ref):
      src = plgpu.load(src_ref, layout=reg_layout, optimized=False)
      plgpu.async_store_tmem(tmem_ref, src)
      plgpu.commit_tmem()
      i = idx1_ref[...]
      j = idx2_ref[...]
      tmem_slice = tmem_ref.at[i, j]
      plgpu.store(dst_ref, plgpu.async_load_tmem(tmem_slice), optimized=False)

    idx1, idx2 = 1, 2
    src = jax.random.uniform(jax.random.key(42), shape, dtype, -1, 1)
    np.testing.assert_array_equal(kernel(src, idx1, idx2), src[idx1, idx2])

  @parameterized.product(
      m=[64, 128],
      n=[64, 128, 256],
      swizzle=[64, 32],
      dtype=[jnp.int8, jnp.uint8],
      lhs_tmem=[False, True],
  )
  def test_integer_matmul(self, m, n, swizzle, dtype, lhs_tmem):
    self.skip_unless_tcgen05_int8()
    if n * jnp.dtype(dtype).itemsize <= swizzle:
      self.skipTest("swizzle too big")
    if lhs_tmem and m == 64:
      self.skipTest("m=64 not supported for LHS in TMEM")
    k = 128
    is_signed = jnp.issubdtype(dtype, jnp.signedinteger)
    o_dtype = jnp.int32

    in_transforms = self.default_transforms(dtype=dtype, swizzle=swizzle)

    def kernel(
        a_smem, b_smem, out_ref, acc_tmem, barrier_ref, a_tmem_ref
    ):
      if lhs_tmem:
        lhs_ref = a_tmem_ref
        layout = plgpu.Layout.TCGEN05_TMEM_NATIVE(4)
        plgpu.async_store_tmem(lhs_ref, plgpu.load(a_smem, layout=layout, optimized=False))
        plgpu.commit_tmem()
      else:
        lhs_ref = a_smem

      plgpu.tcgen05_mma(
          acc_tmem, lhs_ref, b_smem, barrier_ref, accumulate=False
      )
      plgpu.barrier_wait(barrier_ref)
      plgpu.store(out_ref, plgpu.async_load_tmem(acc_tmem), optimized=False)

    scratch_shapes = [
        plgpu.TMEM((m, n), o_dtype, packed=False),
        plgpu.Barrier(orders_tensor_core=True),
    ]
    if lhs_tmem:
      scratch_shapes.append(plgpu.TMEM((m, k), dtype, packed=True))
    else:
      scratch_shapes.append(None)

    f = self.pallas_call(
        kernel,
        in_specs=(
            plgpu.BlockSpec(transforms=in_transforms, memory_space=plgpu.SMEM),
            plgpu.BlockSpec(transforms=in_transforms, memory_space=plgpu.SMEM),
        ),
        out_specs=plgpu.BlockSpec(memory_space=plgpu.GMEM),
        out_shape=jax.ShapeDtypeStruct((m, n), o_dtype),
        scratch_shapes=scratch_shapes,
    )
    # use small values to avoid overflow, [0, 8) for u8 and (-8, 8) for s8
    random_int_input = lambda key, shape: jax.random.randint(
        key, minval=-8 * is_signed, maxval=8, shape=shape, dtype=dtype
    )

    x = random_int_input(jax.random.key(0), shape=(m, k))
    y = random_int_input(jax.random.key(1), shape=(k, n))

    result = f(x, y)
    expected = x.astype(o_dtype) @ y.astype(o_dtype)
    np.testing.assert_array_equal(result, expected)

  @parameterized.product(
      m=[64, 128],
      n=[64, 128, 256],
      swizzle=[128, 64, 32],
      ab_type=[
          jnp.float16,
          jnp.bfloat16,
          (jnp.float8_e5m2, jnp.float8_e5m2),
          (jnp.float8_e4m3fn, jnp.float8_e4m3fn),
          (jnp.float8_e5m2, jnp.float8_e4m3fn),
          (jnp.float8_e4m3fn, jnp.float8_e5m2),
      ],
      lhs_tmem=[False, True],
      transpose_rhs=[False, True],
      transpose_lhs=[False, True],
  )
  def test_simple_matmul(
      self, m, n, swizzle, ab_type, lhs_tmem, transpose_lhs, transpose_rhs
  ):
    a_type, b_type = (
        ab_type if isinstance(ab_type, tuple) else (ab_type, ab_type)
    )
    is_fp8 = a_type in (jnp.float8_e5m2, jnp.float8_e4m3fn) or b_type in (
        jnp.float8_e5m2,
        jnp.float8_e4m3fn,
    )
    if is_fp8:
      if transpose_lhs or transpose_rhs:
        self.skipTest("Skipping FP8 transpose tests in simple sweep")
    if transpose_lhs and lhs_tmem:
      self.skipTest("TMEM transpose not supported")
    if n * jnp.dtype(b_type).itemsize <= swizzle:
      self.skipTest("swizzle too big")
    if lhs_tmem and m == 64:
      self.skipTest("m=64 not supported for LHS in TMEM")
    k = 128
    a_transforms = self.default_transforms(dtype=a_type, swizzle=swizzle)
    b_transforms = self.default_transforms(dtype=b_type, swizzle=swizzle)

    def kernel(a_smem, b_smem, out_ref, acc_tmem, barrier_ref, a_tmem_ref):
      if transpose_lhs:
        a_smem = a_smem.transpose((1, 0))
      if transpose_rhs:
        b_smem = b_smem.transpose((1, 0))
      if lhs_tmem:
        lhs_ref = a_tmem_ref
        layout = plgpu.Layout.TCGEN05_TMEM_NATIVE(32 // jnp.finfo(a_type).bits)
        plgpu.async_store_tmem(lhs_ref, plgpu.load(a_smem, layout=layout, optimized=False))
        plgpu.commit_tmem()
      else:
        lhs_ref = a_smem
      plgpu.tcgen05_mma(
          acc_tmem, lhs_ref, b_smem, barrier_ref, accumulate=False
      )
      plgpu.barrier_wait(barrier_ref)
      plgpu.store(out_ref, plgpu.async_load_tmem(acc_tmem).astype(
          jnp.float32 if is_fp8 else a_type
      ), optimized=False)

    scratch_shapes = [
        plgpu.TMEM((m, n), jnp.float32, packed=False),
        plgpu.Barrier(orders_tensor_core=True),
    ]
    if lhs_tmem:
      scratch_shapes.append(plgpu.TMEM((m, k), a_type, packed=True))
    else:
      scratch_shapes.append(None)

    out_dtype = jnp.float32 if is_fp8 else a_type
    f = self.pallas_call(
        kernel,
        in_specs=(
            plgpu.BlockSpec(transforms=a_transforms, memory_space=plgpu.SMEM),
            plgpu.BlockSpec(transforms=b_transforms, memory_space=plgpu.SMEM),
        ),
        out_specs=plgpu.BlockSpec(memory_space=plgpu.GMEM),
        out_shape=jax.ShapeDtypeStruct((m, n), out_dtype),
        scratch_shapes=scratch_shapes,
    )
    lhs_shape = (k, m) if transpose_lhs else (m, k)
    rhs_shape = (n, k) if transpose_rhs else (k, n)
    x = jax.random.uniform(jax.random.key(0), shape=lhs_shape, dtype=a_type)
    y = jax.random.uniform(jax.random.key(1), shape=rhs_shape, dtype=b_type)
    result = f(x, y)
    if transpose_lhs:
      x = jnp.transpose(x, (1, 0))
    if transpose_rhs:
      y = jnp.transpose(y, (1, 0))
    if is_fp8:
      expected = x.astype(jnp.float32) @ y.astype(jnp.float32)
      atol, rtol = 1e-2, 1e-2
    else:
      expected = x @ y
      atol, rtol = 0, 1e-3
    np.testing.assert_allclose(result, expected, atol=atol, rtol=rtol)

  def test_matmul_alignment(self):
    m = k = n = 128
    dtype = jnp.float16
    transforms = self.default_transforms(dtype=dtype)

    def kernel(a_smem, b_smem, out_ref, _, acc_tmem, barrier_ref):
      plgpu.tcgen05_mma(acc_tmem, a_smem, b_smem, barrier_ref, accumulate=False)
      plgpu.barrier_wait(barrier_ref)
      # We don't await the load because acc_tmem is never modified again.
      out_ref[...] = plgpu.async_load_tmem(acc_tmem).astype(dtype)

    spec = plgpu.BlockSpec(transforms=transforms, memory_space=plgpu.SMEM)
    f = self.pallas_call(
        kernel,
        in_specs=(spec, spec),
        out_specs=spec,
        out_shape=jax.ShapeDtypeStruct((m, n), dtype),
        # Add a one column space to test if we align the accumulator.
        scratch_shapes=(
            plgpu.TMEM((128, 1), jnp.float32),
            plgpu.TMEM((m, n), jnp.float32),
            plgpu.Barrier(orders_tensor_core=True),
        ),
    )
    lhs_shape = (m, k)
    rhs_shape = (k, n)
    x = jax.random.uniform(jax.random.key(0), shape=lhs_shape, dtype=dtype)
    y = jax.random.uniform(jax.random.key(1), shape=rhs_shape, dtype=dtype)
    result = f(x, y)
    expected = jnp.matmul(x, y, precision="high")
    np.testing.assert_allclose(result, expected, rtol=1e-3)

  @parameterized.product(
      m=[128],
      n=[128, 256],
      ab_type=[
          jnp.float8_e5m2,
          jnp.float8_e4m3fn,
          (jnp.float8_e5m2, jnp.float8_e4m3fn),
          (jnp.float8_e4m3fn, jnp.float8_e5m2),
          jnp.float4_e2m1fn,
      ],
      block_size=[16, 32],
  )
  def test_simple_scaled_matmul(self, m, n, ab_type, block_size):
    a_type, b_type = (
        ab_type if isinstance(ab_type, tuple) else (ab_type, ab_type)
    )
    if block_size == 16 and (
        a_type != jnp.float4_e2m1fn or b_type != jnp.float4_e2m1fn
    ):
      self.skipTest("Only float4_e2m1fn supports block size 16.")
    # TODO(apaszke): Add support for single-buffering in pallas_call.
    causes_oom = jnp.finfo(a_type).bits == 8 and n == 256
    k = 128 if causes_oom else 256
    swizzle = 128
    a_transforms = self.default_transforms(swizzle=swizzle, dtype=a_type)
    b_transforms = self.default_transforms(swizzle=swizzle, dtype=b_type)
    out_transforms = self.default_transforms(dtype=jnp.float32)

    def kernel(a_smem, b_smem, a_scale_smem, b_scale_smem, out_ref,
               barrier_ref, acc_tmem, a_scale_tmem, b_scale_tmem):
      plgpu.async_copy_scales_to_tmem(a_scale_smem, a_scale_tmem)
      plgpu.async_copy_scales_to_tmem(b_scale_smem, b_scale_tmem)
      # We don't have to await the copy because it's only used by the MMA.
      plgpu.tcgen05_mma(acc_tmem,
                        a_smem,
                        b_smem.transpose((1, 0)),
                        a_scale=a_scale_tmem,
                        b_scale=b_scale_tmem,
                        accumulate=False)
      plgpu.tcgen05_commit_arrive(barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      # We don't await the load because acc_tmem is never modified again.
      out_ref[...] = plgpu.async_load_tmem(acc_tmem)

    scratch_shapes = [
        plgpu.Barrier(orders_tensor_core=True),
        plgpu.TMEM((m, n), jnp.float32),
        plgpu.TMEM(
            (m, k // block_size),
            jnp.float8_e8m0fnu,
            layout=plgpu.TMEMLayout.SCALES_LAYOUT,
        ),
        plgpu.TMEM(
            (n, k // block_size),
            jnp.float8_e8m0fnu,
            layout=plgpu.TMEMLayout.SCALES_LAYOUT,
        ),
    ]

    f = self.pallas_call(
        kernel,
        in_specs=(
            plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=a_transforms),
            plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=b_transforms),
            plgpu.BlockSpec(memory_space=plgpu.SMEM),
            plgpu.BlockSpec(memory_space=plgpu.SMEM),
        ),
        out_shape=jax.ShapeDtypeStruct((m, n), jnp.float32),
        out_specs=plgpu.BlockSpec(transforms=out_transforms),
        scratch_shapes=scratch_shapes,
    )
    x = jax.random.uniform(
        jax.random.key(1), shape=(m, k), dtype=jnp.float32
    ).astype(a_type)
    y = jax.random.uniform(
        jax.random.key(2), shape=(n, k), dtype=jnp.float32
    ).astype(b_type)
    ksx, ksy = jax.random.split(jax.random.key(1234), 2)
    x_scale = jax.lax.bitcast_convert_type(
        jax.random.randint(
            ksx, (m, k // block_size), 122, 132, dtype=jnp.uint8
        ),
        jnp.float8_e8m0fnu,
    )
    y_scale = jax.lax.bitcast_convert_type(
        jax.random.randint(
            ksy, (n, k // block_size), 122, 132, dtype=jnp.uint8
        ),
        jnp.float8_e8m0fnu,
    )
    def format_scales(scales):
      mn, kk = scales.shape
      assert mn % 128 == 0 and kk % 4 == 0
      return (
          scales.reshape(mn // 128, 4, 32, kk // 4, 4)
          .transpose(0, 3, 2, 1, 4)
          .reshape(mn // 128, kk // 4, 32, 16)
      )
    result = f(x, y, format_scales(x_scale), format_scales(y_scale))
    x_logical_scale = jnp.repeat(x_scale, block_size, axis=1).astype(
        jnp.float32
    )
    y_logical_scale = jnp.repeat(y_scale, block_size, axis=1).astype(
        jnp.float32
    )
    expected = jnp.dot(
        x.astype(jnp.float32) * x_logical_scale,
        (y.astype(jnp.float32) * y_logical_scale).T,
    )
    np.testing.assert_allclose(result, expected, rtol=1e-3)

  def test_simple_warp_specialized_matmul(self):
    m = n = 128
    dtype = jnp.float16
    k = swizzle = 128
    transforms = self.default_transforms(dtype=dtype, swizzle=swizzle)

    def kernel(a_ref, b_ref, out_ref, a_smem, b_smem, acc_tmem, tma_barrier, mma_barrier):
      plgpu.copy_gmem_to_smem(a_ref, a_smem, tma_barrier)
      plgpu.copy_gmem_to_smem(b_ref, b_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)

      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 0)
        def _():
          plgpu.tcgen05_mma(acc_tmem,
                            a_smem,
                            b_smem,
                            mma_barrier,
                            accumulate=False)
      plgpu.barrier_wait(mma_barrier)
      # We don't await the load because acc_tmem is never modified again.
      plgpu.store(
          out_ref,
          plgpu.async_load_tmem(acc_tmem).astype(dtype),
          optimized=False,
      )

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((m, n), dtype),
        scratch_types=(
            plgpu.SMEM((m, k), dtype, transforms=transforms),
            plgpu.SMEM((k, n), dtype, transforms=transforms),
            plgpu.TMEM((m, n), jnp.float32, packed=False),
            plgpu.Barrier(num_arrivals=2),
            plgpu.Barrier(orders_tensor_core=True),
        ),
    )
    x = jax.random.uniform(jax.random.key(0), shape=(m, k), dtype=dtype)
    y = jax.random.uniform(jax.random.key(1), shape=(k, n), dtype=dtype)
    np.testing.assert_allclose(
        f(x, y), jnp.matmul(x, y, precision="high"), rtol=1e-3
    )

  @parameterized.product(
      m=[256],
      n=[128, 256],
      scale_jax_dtype=[jnp.float8_e8m0fnu, jnp.float8_e4m3fn],
  )
  def test_collective_scaled_matmul(self, m, n, scale_jax_dtype):
    in_jax_dtype = jnp.float4_e2m1fn
    out_jax_dtype = jnp.float32
    scale_block = 32 if scale_jax_dtype == jnp.float8_e8m0fnu else 16
    swizzle = 128
    k_steps = 2
    swizzle_elems = 8 * swizzle // dtypes.itemsize_bits(in_jax_dtype)
    k = swizzle_elems * k_steps
    transforms = self.default_transforms(swizzle=swizzle, dtype=in_jax_dtype)

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

    def kernel(lhs_gmem, rhs_gmem, lhs_scales_gmem, rhs_scales_gmem, out_gmem,
               lhs_smem, rhs_smem, lhs_scales_smem, rhs_scales_smem,
               tma_barrier, mma_barrier,
               acc_tmem, lhs_scales_tmem, rhs_scales_tmem):
      plgpu.copy_gmem_to_smem(lhs_gmem, lhs_smem, tma_barrier,
                              collective_axes="x", leader_tracked=plgpu.CopyPartition.PARTITIONED(0))
      plgpu.copy_gmem_to_smem(rhs_gmem, rhs_smem, tma_barrier,
                              collective_axes="x", leader_tracked=plgpu.CopyPartition.PARTITIONED(0))
      plgpu.copy_gmem_to_smem(lhs_scales_gmem, lhs_scales_smem, tma_barrier,
                              collective_axes="x", leader_tracked=plgpu.CopyPartition.PARTITIONED(0))
      # RHS scales are replicated (multicast)
      plgpu.copy_gmem_to_smem(rhs_scales_gmem, rhs_scales_smem, tma_barrier,
                              collective_axes="x")
      cluster_idx = lax.axis_index("x")

      @pl.when(cluster_idx == 0)
      def _leader_block():
        plgpu.barrier_wait(tma_barrier)
        plgpu.async_copy_scales_to_tmem(lhs_scales_smem, lhs_scales_tmem, collective_axis="x")
        plgpu.async_copy_scales_to_tmem(rhs_scales_smem, rhs_scales_tmem, collective_axis="x")
        plgpu.tcgen05_mma(
            acc_tmem,
            lhs_smem,
            rhs_smem.transpose((1, 0)),
            mma_barrier,
            a_scale=lhs_scales_tmem,
            b_scale=rhs_scales_tmem,
            accumulate=False,
            collective_axis="x"
        )
      plgpu.barrier_wait(mma_barrier)

      slice_out = pl.ds(cluster_idx * m_block, m_block)
      plgpu.store(
          out_gmem.at[slice_out, :],
          plgpu.async_load_tmem(acc_tmem),
          optimized=False,
      )

    scratch_types = [
        plgpu.SMEM((m_block, k), in_jax_dtype, transforms=transforms),
        plgpu.SMEM((n_block, k), in_jax_dtype, transforms=transforms),
        plgpu.SMEM((m_block // 128, k // (scale_block * 4), 32, 16), scale_jax_dtype),
        plgpu.SMEM((n // 128, k // (scale_block * 4), 32, 16), scale_jax_dtype),
        plgpu.Barrier(num_arrivals=4),
        plgpu.Barrier(orders_tensor_core=True),
        plgpu.TMEM((m_block, n), out_jax_dtype, collective=True),
        plgpu.TMEM((m_block, k // scale_block), scale_jax_dtype,
                   layout=plgpu.TMEMLayout.SCALES_LAYOUT, collective=True),
        plgpu.TMEM((n, k // scale_block), scale_jax_dtype,
                   layout=plgpu.TMEMLayout.SCALES_LAYOUT, collective=True),
    ]

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((m, n), out_jax_dtype),
        cluster=(2,),
        cluster_names=("x",),
        scratch_types=scratch_types,
    )

    x = jax.random.uniform(jax.random.key(1), shape=(m, k), dtype=jnp.float32).astype(in_jax_dtype)
    y = jax.random.uniform(jax.random.key(2), shape=(n, k), dtype=jnp.float32).astype(in_jax_dtype)

    ka, kb = jax.random.split(jax.random.key(1234), 2)
    if scale_jax_dtype == jnp.float8_e8m0fnu:
      x_scale = jax.lax.bitcast_convert_type(
          jax.random.randint(ka, (m, k // scale_block), 122, 132, dtype=jnp.uint8),
          scale_jax_dtype
      )
      y_scale = jax.lax.bitcast_convert_type(
          jax.random.randint(kb, (n, k // scale_block), 122, 132, dtype=jnp.uint8),
          scale_jax_dtype
      )
    else:
      x_scale = jnp.abs(
          jax.random.normal(ka, (m, k // scale_block), dtype=jnp.float32).astype(scale_jax_dtype)
      )
      y_scale = jnp.abs(
          jax.random.normal(kb, (n, k // scale_block), dtype=jnp.float32).astype(scale_jax_dtype)
      )

    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)
      )

    result = f(x, y, format_scales(x_scale), format_scales(y_scale))

    x_logical_scale = jnp.repeat(x_scale, scale_block, axis=1).astype(jnp.float32)
    y_logical_scale = jnp.repeat(y_scale, scale_block, axis=1).astype(jnp.float32)
    expected = jnp.dot(
        x.astype(jnp.float32) * x_logical_scale,
        (y.astype(jnp.float32) * y_logical_scale).T,
    )
    np.testing.assert_allclose(result, expected, rtol=1e-3)

  @parameterized.parameters(64, 128, 192, 256)
  def test_collective_scaled_matmul_m128(self, n):
    m = 128
    scale_jax_dtype = jnp.float8_e8m0fnu
    in_jax_dtype = jnp.float8_e5m2
    out_jax_dtype = jnp.float32
    scale_block = 32
    swizzle = 128
    k_steps = 2
    swizzle_elems = 8 * swizzle // dtypes.itemsize_bits(in_jax_dtype)
    k = swizzle_elems * k_steps
    transforms = self.default_transforms(swizzle=swizzle, dtype=in_jax_dtype)

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

    def kernel(lhs_gmem, rhs_gmem, lhs_scales_gmem, rhs_scales_gmem, out_gmem,
               lhs_smem, rhs_smem, lhs_scales_smem, rhs_scales_smem,
               tma_barrier, mma_barrier,
               acc_tmem, lhs_scales_tmem, rhs_scales_tmem):
      def copy(src, dst, leader_tracked):
        plgpu.copy_gmem_to_smem(
            src,
            dst,
            tma_barrier,
            collective_axes="x",
            leader_tracked=leader_tracked
        )
      copy(lhs_gmem, lhs_smem, plgpu.CopyPartition.PARTITIONED(0))
      copy(rhs_gmem, rhs_smem, plgpu.CopyPartition.PARTITIONED(0))
      copy(lhs_scales_gmem, lhs_scales_smem, plgpu.CopyPartition.PARTITIONED(0))
      # RHS scales are replicated (multicast)
      copy(rhs_scales_gmem, rhs_scales_smem, plgpu.CopyPartition.REPLICATED)
      cluster_idx = lax.axis_index("x")

      @pl.when(cluster_idx == 0)
      def _leader_block():
        plgpu.barrier_wait(tma_barrier)
        plgpu.async_copy_scales_to_tmem(
            lhs_scales_smem, lhs_scales_tmem, collective_axis="x"
        )
        plgpu.async_copy_scales_to_tmem(
            rhs_scales_smem, rhs_scales_tmem, collective_axis="x"
        )
        plgpu.tcgen05_mma(
            acc_tmem,
            lhs_smem,
            rhs_smem.transpose((1, 0)),
            mma_barrier,
            a_scale=lhs_scales_tmem,
            b_scale=rhs_scales_tmem,
            accumulate=False,
            collective_axis="x"
        )
      plgpu.barrier_wait(mma_barrier)

      slice_out = pl.ds(cluster_idx * m_block, m_block)
      plgpu.store(
          out_gmem.at[slice_out, :],
          plgpu.async_load_tmem(acc_tmem),
          optimized=False,
      )

    scratch_types = dict(
        lhs_smem=plgpu.SMEM((m_block, k), in_jax_dtype, transforms=transforms),
        rhs_smem=plgpu.SMEM((n_block, k), in_jax_dtype, transforms=transforms),
        lhs_scales_smem=plgpu.SMEM(
            ((m_block + 127) // 128, k // (scale_block * 4), 32, 16),
            scale_jax_dtype,
        ),
        rhs_scales_smem=plgpu.SMEM(
            (1, k // (scale_block * 4), 64, 16), scale_jax_dtype
        ),
        tma_barrier=plgpu.Barrier(num_arrivals=4),
        mma_barrier=plgpu.Barrier(orders_tensor_core=True),
        acc_tmem=plgpu.TMEM((m_block, n), out_jax_dtype, collective=True),
        lhs_scales_tmem=plgpu.TMEM(
            (128, k // scale_block),
            scale_jax_dtype,
            layout=plgpu.TMEMLayout.SCALES_LAYOUT,
            collective=True,
        ),
        rhs_scales_tmem=plgpu.TMEM(
            (256, k // scale_block),
            scale_jax_dtype,
            layout=plgpu.TMEMLayout.SCALES_M64_COLLECTIVE_LAYOUT,
            collective=True,
        ),
    )

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((m, n), out_jax_dtype),
        cluster=(2,),
        cluster_names=("x",),
        scratch_types=scratch_types,
    )

    x = jax.random.uniform(jax.random.key(1), shape=(m, k), dtype=jnp.float32).astype(in_jax_dtype)
    y = jax.random.uniform(jax.random.key(2), shape=(n, k), dtype=jnp.float32).astype(in_jax_dtype)

    ka, kb = jax.random.split(jax.random.key(1234), 2)
    x_scale = jax.lax.bitcast_convert_type(
        jax.random.randint(ka, (m, k // scale_block), 122, 132, dtype=jnp.uint8),
        scale_jax_dtype
    )
    y_scale = jax.lax.bitcast_convert_type(
        jax.random.randint(kb, (n, k // scale_block), 122, 132, dtype=jnp.uint8),
        scale_jax_dtype
    )

    def format_lhs_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_rhs_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)),
      )

    result = f(x, y, format_lhs_scales(x_scale), format_rhs_scales(y_scale))

    x_logical_scale = jnp.repeat(x_scale, scale_block, axis=1).astype(jnp.float32)
    y_logical_scale = jnp.repeat(y_scale, scale_block, axis=1).astype(jnp.float32)
    expected = jnp.dot(
        x.astype(jnp.float32) * x_logical_scale,
        (y.astype(jnp.float32) * y_logical_scale).T,
    )
    np.testing.assert_allclose(result, expected, atol=2e-4, rtol=5e-6)

  @parameterized.product(
      m=[128],
      n=[128, 256],
      dtype=[jnp.float16],
  )
  def test_simple_sparse_matmul(self, m, n, dtype):
    k = 128
    swizzle = 128 // jnp.dtype(dtype).itemsize
    transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)
    out_transforms = self.default_transforms(dtype=jnp.float32)

    def kernel(a_smem, b_smem, a_sparse_smem, out_ref,
               barrier_ref, acc_tmem, a_sparse_tmem):
      plgpu.async_copy_sparse_metadata_to_tmem(a_sparse_smem, a_sparse_tmem)
      # We don't have to await the copy because it's only used by the MMA.
      plgpu.tcgen05_mma(acc_tmem,
                        a_smem,
                        b_smem.transpose((1, 0)),
                        a_sparse_metadata=a_sparse_tmem,
                        accumulate=False)
      plgpu.tcgen05_commit_arrive(barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      # We don't await the load because acc_tmem is never modified again.
      out_ref[...] = plgpu.async_load_tmem(acc_tmem)

    scratch_shapes = [
        plgpu.Barrier(orders_tensor_core=True),
        plgpu.TMEM((m, n), jnp.float32),
        plgpu.TMEM((m, k // 2), jnp.uint2, layout=plgpu.TMEMLayout.SPARSE_METADATA_LAYOUT),
    ]

    f = self.pallas_call(
        kernel,
        in_specs=(
            plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=transforms),
            plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=transforms),
            plgpu.BlockSpec(memory_space=plgpu.SMEM),
        ),
        out_shape=jax.ShapeDtypeStruct((m, n), jnp.float32),
        out_specs=plgpu.BlockSpec(transforms=out_transforms),
        scratch_shapes=scratch_shapes,
    )
    x = jax.random.uniform(jax.random.key(1), shape=(m, k // 2), dtype=dtype)
    y = jax.random.uniform(jax.random.key(2), shape=(n, k), dtype=dtype)
    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)
    z = f(x, y, plgpu.format_tcgen05_sparse_metadata(x_sparse.astype(jnp.uint2), dtype))
    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.T.astype(jnp.float32), precision="high"
    )
    np.testing.assert_allclose(z, ref, atol=7e-5, rtol=5e-6)

  @parameterized.product(
      m=[128],
      n=[128, 256],
  )
  def test_block_scaled_sparse_matmul(self, m, n):
    in_dtype = jnp.float8_e5m2
    scale_dtype = jnp.float8_e8m0fnu
    block_size = 64
    swizzle = 128
    k = 256
    transforms = self.default_transforms(swizzle=swizzle, dtype=in_dtype)
    out_transforms = self.default_transforms(dtype=jnp.float32)

    def kernel(a_smem, b_smem, a_sparse_smem, a_scale_smem, b_scale_smem,
               out_ref, barrier_ref, acc_tmem, a_sparse_tmem,
               a_scale_tmem, b_scale_tmem):
      plgpu.async_copy_sparse_metadata_to_tmem(a_sparse_smem, a_sparse_tmem)
      plgpu.async_copy_scales_to_tmem(a_scale_smem, a_scale_tmem)
      plgpu.async_copy_scales_to_tmem(b_scale_smem, b_scale_tmem)
      plgpu.tcgen05_mma(acc_tmem,
                        a_smem,
                        b_smem.transpose((1, 0)),
                        a_scale=a_scale_tmem,
                        b_scale=b_scale_tmem,
                        a_sparse_metadata=a_sparse_tmem,
                        accumulate=False)
      plgpu.tcgen05_commit_arrive(barrier_ref)
      plgpu.barrier_wait(barrier_ref)
      out_ref[...] = plgpu.async_load_tmem(acc_tmem)

    scratch_shapes = [
        plgpu.Barrier(orders_tensor_core=True),
        plgpu.TMEM((m, n), jnp.float32),
        plgpu.TMEM((m, k // 2), jnp.uint2, layout=plgpu.TMEMLayout.SPARSE_METADATA_LAYOUT),
        plgpu.TMEM((m, k // block_size), scale_dtype, layout=plgpu.TMEMLayout.SCALES_LAYOUT),
        plgpu.TMEM((n, k // block_size), scale_dtype, layout=plgpu.TMEMLayout.SCALES_LAYOUT),
    ]

    f = self.pallas_call(
        kernel,
        in_specs=(
            plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=transforms),
            plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=transforms),
            plgpu.BlockSpec(memory_space=plgpu.SMEM),
            plgpu.BlockSpec(memory_space=plgpu.SMEM),
            plgpu.BlockSpec(memory_space=plgpu.SMEM),
        ),
        out_shape=jax.ShapeDtypeStruct((m, n), jnp.float32),
        out_specs=plgpu.BlockSpec(transforms=out_transforms),
        scratch_shapes=scratch_shapes,
    )
    x = jax.random.uniform(jax.random.key(1), shape=(m, k // 2), dtype=jnp.float32).astype(in_dtype)
    y = jax.random.uniform(jax.random.key(2), shape=(n, k), dtype=jnp.float32).astype(in_dtype)
    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)
    ksx, ksy = jax.random.split(jax.random.key(5678), 2)
    x_scale = jax.lax.bitcast_convert_type(
        jax.random.randint(ksx, (m, k // block_size), 122, 132, dtype=jnp.uint8),
        scale_dtype
    )
    y_scale = jax.lax.bitcast_convert_type(
        jax.random.randint(ksy, (n, k // block_size), 122, 132, dtype=jnp.uint8),
        scale_dtype
    )
    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)
      )
    def format_sparse_meta(meta):
      mn, k, _2 = meta.shape
      assert _2 == 2
      k *= 2
      return (
          meta.reshape(mn // 128, 128, k // 64, 64)
          .transpose(0, 2, 1, 3)
          .astype(jnp.uint2)
      )
    z = f(
        x, y,
        format_sparse_meta(x_sparse),
        format_scales(x_scale), format_scales(y_scale),
    )
    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(x_scale, block_size, axis=1).astype(jnp.float32)
    b_logical_scales = jnp.repeat(y_scale, 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.parameters(
      (128, jnp.float16)
  )
  def test_manual_tcgen05_commit_arrive(self, swizzle, dtype):
    shape = (128, 128)
    transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)

    def kernel(a_gmem, b_gmem, out_gmem,
        a_smem, b_smem, tma_barrier, mma_barrier, acc_tmem):
      plgpu.copy_gmem_to_smem(a_gmem, a_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)
      plgpu.copy_gmem_to_smem(b_gmem, b_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)

      plgpu.commit_tmem()
      # Check that `predicate` is respected.
      plgpu.tcgen05_commit_arrive(mma_barrier, predicate=False)
      # Don't pass a barrier directly into tcgen05_mma and arrive manually.
      plgpu.tcgen05_mma(acc_tmem, a_smem, b_smem, accumulate=False)
      plgpu.tcgen05_commit_arrive(mma_barrier, predicate=True)
      plgpu.barrier_wait(mma_barrier)
      # We don't await the load because acc_tmem is never modified again.
      acc = plgpu.async_load_tmem(acc_tmem).astype(dtype)
      plgpu.store(out_gmem, acc, optimized=False)

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct(shape, dtype),
        scratch_types=[
          plgpu.SMEM(shape, dtype, transforms=transforms),  # a_smem
          plgpu.SMEM(shape, dtype, transforms=transforms),  # b_smem
          plgpu.Barrier(),  # tma_barrier
          plgpu.Barrier(orders_tensor_core=True),  # mma_barrier
          plgpu.TMEM((128, 128), jnp.float32),  # acc
        ],
    )
    x = jax.random.uniform(jax.random.key(0), shape=shape, dtype=dtype)
    y = jax.random.uniform(jax.random.key(1), shape=shape, dtype=dtype)
    result = f(x, y)
    np.testing.assert_allclose(
        result, jnp.matmul(x, y, precision="high"), rtol=1e-3
    )

  def test_async_copy_smem_to_tmem_warp_semantics(self):
    shape, dtype = (128, 64), jnp.int32
    transforms = self.default_transforms(dtype=dtype)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct(shape, dtype),
        scratch_types=[
            plgpu.SMEM(shape, dtype, transforms=transforms),
            plgpu.TMEM(shape, dtype),
            plgpu.Barrier(),
            plgpu.Barrier(orders_tensor_core=True),
        ],
    )
    def kernel(src_gmem, dst_gmem, smem, tmem, tma_barrier, tc_barrier):
      plgpu.copy_gmem_to_smem(src_gmem, smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 0)
        def _():
          plgpu.async_copy_smem_to_tmem(smem, tmem)
          plgpu.tcgen05_commit_arrive(tc_barrier)
      plgpu.barrier_wait(tc_barrier)
      plgpu.store(dst_gmem, plgpu.async_load_tmem(tmem), optimized=False)

    x = jnp.arange(math.prod(shape), dtype=dtype).reshape(shape)
    self.assertArraysEqual(kernel(x), x)

  def test_async_copy_scales_to_tmem_warp_semantics_lowers_correctly(self):
    tmem_shape, dtype = (128, 4), jnp.float8_e8m0fnu
    smem_shape = (1, 1, 32, 16)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct(smem_shape, dtype),
        scratch_types=[
            plgpu.SMEM(smem_shape, dtype),
            plgpu.TMEM(tmem_shape, dtype, layout=plgpu.TMEMLayout.SCALES_LAYOUT),
            plgpu.Barrier(orders_tensor_core=True)
        ],
    )
    def kernel(out_gmem, smem, tmem, tc_barrier):
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 0)
        def _():
          plgpu.async_copy_scales_to_tmem(smem, tmem)
          plgpu.tcgen05_commit_arrive(tc_barrier)
      plgpu.barrier_wait(tc_barrier)
      plgpu.copy_smem_to_gmem(smem, out_gmem)
      plgpu.wait_smem_to_gmem(0)

    jax.jit(kernel).lower()

  @parameterized.parameters(128, None)
  def test_async_copy_smem_to_tmem(self, swizzle):
    dtype = jnp.float16
    m, n = 128, 128
    if swizzle is None:
      transforms = (plgpu.TilingTransform((8, 8)),)
    else:
      transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)

    def kernel(x_gmem, y_gmem, smem, tma_barrier, mma_barrier, tmem):
      plgpu.copy_gmem_to_smem(x_gmem, smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)
      plgpu.async_copy_smem_to_tmem(smem, tmem)
      plgpu.tcgen05_commit_arrive(mma_barrier)
      plgpu.barrier_wait(mma_barrier)
      plgpu.store(y_gmem, plgpu.async_load_tmem(tmem), optimized=False)

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((m, n), dtype),
        scratch_types=[
            plgpu.SMEM((m, n), dtype, transforms=transforms),
            plgpu.Barrier(),
            plgpu.Barrier(orders_tensor_core=True),
            plgpu.TMEM((m, n), dtype, packed=True),
        ],
    )
    x = jax.random.uniform(jax.random.key(0), shape=(m, n), dtype=dtype)
    result = jax.block_until_ready(f(x))
    np.testing.assert_array_equal(result, x)

  def test_matmul_with_sliced_accumulator(self):
    dtype = jnp.bfloat16
    shape = (128, 128)
    tmem_shape = (128, 2 * 128)
    swizzle = 128

    # Test a matmul with a single block.
    transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)

    def kernel(a_smem, b_smem, out_ref, acc_tmem, barrier_ref):
      acc_tmem_slice = acc_tmem.at[slice(None), pl.dslice(0, 128)]
      plgpu.tcgen05_mma(acc_tmem_slice,
                        a_smem,
                        b_smem,
                        barrier_ref,
                        accumulate=False)
      plgpu.barrier_wait(barrier_ref)
      # We don't await the load because acc_tmem is never modified again.
      plgpu.store(
          out_ref,
          plgpu.async_load_tmem(acc_tmem_slice).astype(dtype),
          optimized=False,
      )

    scratch_shapes = [
        plgpu.TMEM(tmem_shape, jnp.float32, packed=False),
        plgpu.Barrier(orders_tensor_core=True),
    ]

    f = self.pallas_call(
        kernel,
        in_specs=(
            plgpu.BlockSpec(transforms=transforms, memory_space=plgpu.SMEM),
            plgpu.BlockSpec(transforms=transforms, memory_space=plgpu.SMEM),
        ),
        out_specs=plgpu.BlockSpec(memory_space=plgpu.GMEM),
        out_shape=jax.ShapeDtypeStruct(shape, dtype),
        scratch_shapes=scratch_shapes,
    )
    x = jax.random.uniform(jax.random.key(0), shape=shape, dtype=dtype)
    y = jax.random.uniform(jax.random.key(1), shape=shape, dtype=dtype)
    result = f(x, y)
    expected = jnp.matmul(x, y, precision="high")
    np.testing.assert_allclose(result, expected, rtol=1e-3)

  def test_tmem_alloc_within_warp_map(self):
    m, k = 128, 128
    out_shape = (m, k // 2)
    in_shape = (m, k)
    @self.kernel(
        out_type=jax.ShapeDtypeStruct(out_shape, jnp.uint16),
        scratch_types=[
            plgpu.SMEM(
                in_shape,
                jnp.uint8,
                transforms=self.default_transforms(dtype=jnp.uint8, swizzle=64),
            ),
            plgpu.RefUnion(
                plgpu.TMEM(in_shape, jnp.uint8, packed=True),
                plgpu.TMEM(out_shape, jnp.uint16, packed=True),
            ),
        ],
    )
    def kernel(_, out_gmem, smem, tmem_alias):
      tmem, tmem_b = tmem_alias
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 0)
        def _():
          plgpu.async_copy_smem_to_tmem(smem, tmem)
      # using out_gmem to prevent potential DCE on jaxprs
      plgpu.store(
          out_gmem,
          plgpu.async_load_tmem(tmem_b, layout=plgpu.Layout.TCGEN05),
          optimized=False,
      )

    # For simplicity reasons, we only check that we can lower correctly.
    x = jnp.ones(in_shape, jnp.uint8)
    jax.block_until_ready(kernel(x))

  def test_async_copy_sparse_metadata_smem_to_tmem_warp_semantics(self):
    m, k = 128, 64
    dtype = jnp.uint2
    layout = plgpu.TMEMLayout.SPARSE_METADATA_LAYOUT
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((1, 1, m, k), dtype),
        scratch_types=[
            plgpu.SMEM((m // 128, k // 64, 128, 64), dtype),
            plgpu.TMEM((m, k), dtype, layout=layout),
        ],
    )
    def kernel(src_gmem, out_gmem, smem, tmem):
      plgpu.store(
          smem,
          plgpu.load(src_gmem, layout=layout, optimized=False),
          optimized=False,
      )
      @plgpu.warp_map
      def _(warp_id):
        @pl.when(warp_id == 0)
        def _():
          plgpu.async_copy_sparse_metadata_to_tmem(smem, tmem)
      plgpu.copy_smem_to_gmem(smem, out_gmem)

    # For simplicity reasons, we only check that we can lower correctly.
    x = jnp.ones((1, 1, m, k), dtype)
    jax.block_until_ready(kernel(x))

  @parameterized.product(
      m=[256],
      n=[128, 192],
      ab_type=(jnp.float16, jnp.bfloat16, jnp.int8, jnp.float8_e4m3fn)
  )
  def test_sparse_matmul_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
    k = 128
    transforms = self.default_transforms(swizzle=64, dtype=ab_type)
    m_block = m // 2
    n_block = n // 2

    def kernel(
        a_gmem,
        b_gmem,
        a_sparse_gmem,
        out_gmem,
        a_smem,
        b_smem,
        a_sparse_smem,
        tma_barrier,
        mma_barrier,
        cluster_barrier,
        acc_tmem,
        a_sparse_tmem,
    ):
      cluster_idx = lax.axis_index("x")
      slice_a = pl.ds(cluster_idx * m_block, m_block)
      slice_b = pl.ds(cluster_idx * n_block, n_block)
      slice_sparse = pl.ds(cluster_idx, m_block // 128)

      # GMEM -> SMEM
      plgpu.copy_gmem_to_smem(a_gmem.at[slice_a, :], a_smem, tma_barrier)
      plgpu.copy_gmem_to_smem(b_gmem.at[slice_b, :], b_smem, tma_barrier)
      plgpu.copy_gmem_to_smem(a_sparse_gmem.at[slice_sparse, :], a_sparse_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)

      plgpu.barrier_arrive(cluster_barrier)
      plgpu.barrier_wait(cluster_barrier)

      # SMEM -> TMEM
      plgpu.async_copy_sparse_metadata_to_tmem(a_sparse_smem, a_sparse_tmem, collective_axis="x")
      # MMA
      plgpu.tcgen05_mma(
          acc_tmem,
          a_smem,
          b_smem.transpose((1, 0)),
          mma_barrier,
          a_sparse_metadata=a_sparse_tmem,
          accumulate=False,
          collective_axis="x",
      )

      plgpu.barrier_wait(mma_barrier)
      slice_out = pl.ds(cluster_idx * m_block, m_block)
      plgpu.store(
          out_gmem.at[slice_out, :],
          plgpu.async_load_tmem(acc_tmem),
          optimized=False,
      )

    scratch_types = dict(
        a_smem=plgpu.SMEM((m_block, k // 2), ab_type, transforms=transforms),
        b_smem=plgpu.SMEM((n_block, k), ab_type, transforms=transforms),
        a_sparse_smem=plgpu.SMEM(
            (m_block // 128, k // 128, 128, 64), jnp.uint2
        ),
        tma_barrier=plgpu.Barrier(num_arrivals=3),
        mma_barrier=plgpu.Barrier(orders_tensor_core=True),
        cluster_barrier=plgpu.ClusterBarrier(collective_axes=("x",)),
        acc_tmem=plgpu.TMEM((m_block, n), acc_type, collective=True),
        a_sparse_tmem=plgpu.TMEM(
            (m_block, k // 2),
            jnp.uint2,
            layout=plgpu.TMEMLayout.SPARSE_METADATA_LAYOUT,
            collective=True,
        ),
    )

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((m, n), acc_type),
        cluster=(2,),
        cluster_names=("x",),
        scratch_types=scratch_types,
    )

    a_shape = (m, k // 2)
    b_shape = (n, k)
    if jnp.issubdtype(ab_type, jnp.integer):
      x = jax.random.randint(jax.random.key(1), shape=a_shape, minval=-64, maxval=64, dtype=ab_type)
      y = jax.random.randint(jax.random.key(2), shape=b_shape, minval=-64, maxval=64, dtype=ab_type)
    else:
      x = jax.random.uniform(jax.random.key(1), shape=a_shape, dtype=ab_type)
      y = jax.random.uniform(jax.random.key(2), shape=b_shape, dtype=ab_type)

    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]
    in_sparse = plgpu.format_tcgen05_sparse_metadata(x_sparse.astype(jnp.uint2),
                                                     ab_type)
    z = f(x, y, in_sparse)

    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(acc_type), y.T.astype(acc_type), precision="high"
    )
    atol = rtol = 0
    if jnp.issubdtype(ab_type, jnp.floating):
      atol = 7e-5
      rtol = 5e-6
    np.testing.assert_allclose(z, ref, atol=atol, rtol=rtol)

  @parameterized.product(
      m_n_k=[
          (256, 256, 256),
          (256, 128, 128),
          (256, 256, 64),
          (128, 64, 128),
          (128, 64, 128),
      ],
      swizzle=[128, 64, 32],
      ab_type=[
          jnp.float16,
          jnp.bfloat16,
          (jnp.float8_e5m2, jnp.float8_e5m2),
          (jnp.float8_e4m3fn, jnp.float8_e4m3fn),
          (jnp.float8_e5m2, jnp.float8_e4m3fn),
          (jnp.float8_e4m3fn, jnp.float8_e5m2),
      ],
      lhs_tmem=[False, True],
  )
  def test_simple_collective_matmul(self, m_n_k, swizzle, ab_type, lhs_tmem):
    a_type, b_type = (
        ab_type if isinstance(ab_type, tuple) else (ab_type, ab_type)
    )
    is_fp8 = a_type in (jnp.float8_e5m2, jnp.float8_e4m3fn) or b_type in (
        jnp.float8_e5m2,
        jnp.float8_e4m3fn,
    )
    m, n, k = m_n_k
    if min(k, n // 2) * jnp.dtype(a_type).itemsize < swizzle:
      self.skipTest("Swizzle is too big")
    full_lhs_shape = (m, k)
    full_rhs_shape = (k, n)
    full_acc_shape = (m, n)
    block_acc_shape = (m // 2, n)
    block_lhs_shape = (m // 2, k)
    block_rhs_shape = (k, n // 2)
    # Test a collective (paired CTA) matmul on a single block.
    a_transforms = self.default_transforms(swizzle=swizzle, dtype=a_type)
    b_transforms = self.default_transforms(swizzle=swizzle, dtype=b_type)
    if lhs_tmem and m == 128:
      self.skipTest("m=128 not supported for LHS in TMEM")
    out_dtype = jnp.float32 if is_fp8 else a_type

    def kernel(a_gmem, b_gmem, out_gmem, a_smem, b_smem,
               acc_tmem, tma_barrier, mma_barrier,
               cluster_barrier, lhs_tmem_ref):
      cluster_idx = lax.axis_index("x")
      slice_lhs = pl.ds(cluster_idx * block_lhs_shape[0], block_lhs_shape[0])
      slice_rhs = pl.ds(cluster_idx * block_rhs_shape[1], block_rhs_shape[1])

      plgpu.copy_gmem_to_smem(a_gmem.at[slice_lhs, :], a_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)
      plgpu.copy_gmem_to_smem(b_gmem.at[:, slice_rhs], b_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)

      if lhs_tmem:
        lhs_ref = lhs_tmem_ref
        layout = plgpu.Layout.TCGEN05_TMEM_NATIVE(32 // jnp.finfo(a_type).bits)
        plgpu.async_store_tmem(lhs_ref, plgpu.load(a_smem, layout=layout, optimized=False))
        plgpu.commit_tmem()
      else:
        lhs_ref = a_smem

      plgpu.barrier_arrive(cluster_barrier)
      plgpu.barrier_wait(cluster_barrier)

      plgpu.tcgen05_mma(
          acc_tmem,
          lhs_ref,
          b_smem,
          mma_barrier,
          accumulate=False,
          collective_axis="x",
      )
      plgpu.barrier_wait(mma_barrier)
      if m == 128:
        layout = plgpu.Layout.TCGEN05_M64_COLLECTIVE(n)
      else:
        layout = plgpu.Layout.TCGEN05
      # We don't await the load because acc_tmem is never modified again.
      plgpu.store(out_gmem.at[slice_lhs, :], plgpu.async_load_tmem(
          acc_tmem, layout=layout
      ).astype(out_dtype), optimized=False)

    scratch_types = [
        plgpu.SMEM(block_lhs_shape, a_type, transforms=a_transforms),
        plgpu.SMEM(block_rhs_shape, b_type, transforms=b_transforms),
        plgpu.TMEM(block_acc_shape, jnp.float32, collective=True),
        plgpu.Barrier(),
        plgpu.Barrier(orders_tensor_core=True),
        plgpu.ClusterBarrier(collective_axes=("x",)),
    ]
    if lhs_tmem:
      scratch_types.append(
          plgpu.TMEM(block_lhs_shape, a_type, collective=True, packed=True)
      )
    else:
      scratch_types.append(None)

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct(full_acc_shape, out_dtype),
        cluster=(2,),
        cluster_names=("x",),
        scratch_types=scratch_types,
    )
    x = jax.random.uniform(jax.random.key(0), shape=full_lhs_shape, dtype=a_type)
    y = jax.random.uniform(jax.random.key(1), shape=full_rhs_shape, dtype=b_type)
    result = f(x, y)
    expected = x.astype(jnp.float32) @ y.astype(jnp.float32)
    atol, rtol = (1e-2, 1e-2) if is_fp8 else (0, 4e-3)
    np.testing.assert_allclose(result, expected, atol=atol, rtol=rtol)

  @parameterized.product(
      cluster=(((2, 2), ("y", "x")), ((2, 1, 2), ("z", "y", "x")))
  )
  def test_collective_matmul_multidim_cluster(self, cluster):
    cluster_shape, cluster_names = cluster
    m, n, k = 128, 128, 128
    swizzle = 128
    dtype = jnp.float16
    block_lhs_shape = (m // 2, k)
    block_rhs_shape = (k, n // 2)
    block_acc_shape = (m // 2, n)
    a_transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)
    b_transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)
    pair_axis = cluster_names[0]

    def kernel(a_gmem, b_gmem, out_gmem, a_smem, b_smem, acc_tmem,
               tma_barrier, mma_barrier, cluster_barrier):
      pair_idx = lax.axis_index(pair_axis)
      cluster_idx = lax.axis_index("x")
      slice_lhs = pl.ds(cluster_idx * block_lhs_shape[0], block_lhs_shape[0])
      slice_rhs = pl.ds(cluster_idx * block_rhs_shape[1], block_rhs_shape[1])

      plgpu.copy_gmem_to_smem(a_gmem.at[pair_idx, slice_lhs, :], a_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)
      plgpu.copy_gmem_to_smem(b_gmem.at[pair_idx, :, slice_rhs], b_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)

      plgpu.barrier_arrive(cluster_barrier)
      plgpu.barrier_wait(cluster_barrier)

      plgpu.tcgen05_mma(
          acc_tmem,
          a_smem,
          b_smem,
          mma_barrier,
          accumulate=False,
          collective_axis="x",
      )
      plgpu.barrier_wait(mma_barrier)
      plgpu.store(out_gmem.at[pair_idx, slice_lhs, :], plgpu.async_load_tmem(
          acc_tmem, layout=plgpu.Layout.TCGEN05_M64_COLLECTIVE(n)
      ), optimized=False)

    scratch_types = [
        plgpu.SMEM(block_lhs_shape, dtype, transforms=a_transforms),
        plgpu.SMEM(block_rhs_shape, dtype, transforms=b_transforms),
        plgpu.TMEM(block_acc_shape, jnp.float32, collective=True),
        plgpu.Barrier(),
        plgpu.Barrier(orders_tensor_core=True),
        plgpu.ClusterBarrier(collective_axes=("x",)),
    ]
    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((2, m, n), jnp.float32),
        cluster=cluster_shape,
        cluster_names=cluster_names,
        scratch_types=scratch_types,
    )
    # 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 = jax.random.uniform(jax.random.key(0), shape=(2, m, k), dtype=dtype)
    y = jax.random.uniform(jax.random.key(1), shape=(2, k, n), dtype=dtype)
    result = f(x, y)
    expected = x.astype(jnp.float32) @ y.astype(jnp.float32)
    np.testing.assert_allclose(result, expected, atol=0, rtol=4e-3)

  @parameterized.parameters(
      (128, jnp.float16)
  )
  def test_matmul_with_smem_aliasing(self, swizzle, dtype):
    # Perform a 128x128 @ 128x128 matmul and a 128x64 @ 64x128 matmul
    # using aliased Refs pointing to the same SMEM address.
    shape = (128, 128)
    transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)

    def kernel(a_gmem, b_gmem, out_gmem128, out_gmem64,
        a_aliased, b_aliased, tma_barrier, mma_barrier, acc_tmem):
      # Note: We directly copy into 128-sized refs assuming that both aliased
      # refs point to the same address, so we can skip the copy for
      # the 64-sized ref. We transpose the LHS Ref so that the 64-sized Ref
      # receives the correct slice of data from this TMA.
      # As this is implementation dependent, this test may break if we change
      # the underlying aliasing behavior.
      a_smem_128, a_smem_64 = a_aliased
      plgpu.copy_gmem_to_smem(a_gmem, a_smem_128, tma_barrier)
      plgpu.barrier_wait(tma_barrier)
      b_smem_128, b_smem_64 = b_aliased
      plgpu.copy_gmem_to_smem(b_gmem, b_smem_128, tma_barrier)
      plgpu.barrier_wait(tma_barrier)

      # Do 128x128 @ 128x128 matmul
      plgpu.tcgen05_mma(acc_tmem,
                        a_smem_128.transpose((1, 0)),
                        b_smem_128,
                        mma_barrier,
                        accumulate=False)
      plgpu.barrier_wait(mma_barrier)
      plgpu.store(
          out_gmem128,
          plgpu.async_load_tmem(acc_tmem).astype(dtype),
          optimized=False,
      )

      # Do 128x64 @ 64x128 matmul
      plgpu.wait_load_tmem()  # Make sure the loads are complete
      plgpu.tcgen05_mma(acc_tmem,
                        a_smem_64.transpose((1, 0)),
                        b_smem_64,
                        mma_barrier,
                        accumulate=False)
      plgpu.barrier_wait(mma_barrier)
      plgpu.store(
          out_gmem64,
          plgpu.async_load_tmem(acc_tmem).astype(dtype),
          optimized=False,
      )

    f = self.kernel(
        kernel,
        out_type=[jax.ShapeDtypeStruct(shape, dtype),
                   jax.ShapeDtypeStruct(shape, dtype)],
        scratch_types=[
          plgpu.RefUnion(   # aliased a_smem
            plgpu.SMEM(shape, dtype, transforms=transforms),
            plgpu.SMEM((64, 128), dtype, transforms=transforms),
          ),
          plgpu.RefUnion(   # aliased b_smem
            plgpu.SMEM(shape, dtype, transforms=transforms),
            plgpu.SMEM((64, 128), dtype, transforms=transforms),
          ),
          plgpu.Barrier(),  # tma_barrier
          plgpu.Barrier(orders_tensor_core=True),  # mma_barrier
          plgpu.TMEM(shape, jnp.float32),  # acc
        ],
    )
    x = jax.random.uniform(jax.random.key(0), shape=shape, dtype=dtype)
    y = jax.random.uniform(jax.random.key(1), shape=shape, dtype=dtype)
    result_128, result_64 = f(x.T, y)
    np.testing.assert_allclose(
        result_128, jnp.matmul(x, y, precision="high"), rtol=1e-3
    )
    np.testing.assert_allclose(
        result_64, jnp.matmul(x[:, :64], y[:64, :], precision="high"), rtol=1e-3
    )

  @parameterized.parameters(
      (128, jnp.float16)
  )
  def test_matmul_with_tmem_aliasing(self, swizzle, dtype):
    # Perform a 128x128 @ 128x128 matmul and a 128x64 @ 64x128 matmul
    # using aliased Refs pointing to the same TMEM address.
    shape = (128, 128)
    transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)

    def kernel(a_gmem, b_gmem, out_gmem128, out_gmem64,
        a_smem, b_smem, tma_barrier, mma_barrier, aliased_refs):
      plgpu.copy_gmem_to_smem(a_gmem, a_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)
      plgpu.copy_gmem_to_smem(b_gmem, b_smem, tma_barrier)
      plgpu.barrier_wait(tma_barrier)
      [acc_128, lhs_128], [lhs_64, acc_64], _ = aliased_refs

      # Do 128x128 @ 128x128 matmul
      plgpu.async_store_tmem(lhs_128, plgpu.load(a_smem, layout=plgpu.Layout.TCGEN05))
      plgpu.commit_tmem()
      plgpu.tcgen05_mma(acc_128,
                        lhs_128,
                        b_smem,
                        mma_barrier,
                        accumulate=False)
      plgpu.barrier_wait(mma_barrier)
      plgpu.store(
          out_gmem128,
          plgpu.async_load_tmem(acc_128).astype(dtype),
          optimized=False,
      )

      # Do 128x64 @ 64x128 matmul
      plgpu.wait_load_tmem()  # Make sure the loads have completed
      plgpu.async_store_tmem(
          lhs_64,
          plgpu.load(a_smem.at[:, 0:64], layout=plgpu.Layout.TCGEN05),
      )
      plgpu.commit_tmem()
      plgpu.tcgen05_mma(acc_64,
                        lhs_64,
                        b_smem.at[0:64, :],
                        mma_barrier,
                        accumulate=False)
      plgpu.barrier_wait(mma_barrier)
      # We don't await the load because TMEM is never modified again.
      plgpu.store(
          out_gmem64,
          plgpu.async_load_tmem(acc_64).astype(dtype),
          optimized=False,
      )

    f = self.kernel(
        kernel,
        out_type=[
            jax.ShapeDtypeStruct(shape, dtype),
            jax.ShapeDtypeStruct(shape, dtype),
        ],
        scratch_types=[
            plgpu.SMEM(shape, dtype, transforms=transforms),  # a_smem
            plgpu.SMEM(shape, dtype, transforms=transforms),  # b_smem
            plgpu.Barrier(),  # tma_barrier
            plgpu.Barrier(orders_tensor_core=True),  # mma_barrier
            plgpu.RefUnion(  # aliased_refs
                [
                    plgpu.TMEM((128, 128), jnp.float32),  # acc
                    plgpu.TMEM((128, 128), dtype, packed=True),  # lhs
                ],
                [
                    plgpu.TMEM((128, 64), dtype, packed=True),  # lhs
                    plgpu.TMEM((128, 128), jnp.float32),  # acc
                ],
                plgpu.TMEM((128, 128), jnp.float32),  # unused
            ),
        ],
    )
    x = jax.random.uniform(jax.random.key(0), shape=shape, dtype=dtype)
    y = jax.random.uniform(jax.random.key(1), shape=shape, dtype=dtype)
    result_128, result_64 = f(x, y)
    np.testing.assert_allclose(
        result_128, jnp.matmul(x, y, precision="high"), rtol=1e-3
    )
    np.testing.assert_allclose(
        result_64, jnp.matmul(x[:, :64], y[:64, :], precision="high"), rtol=1e-3
    )

  @parameterized.parameters(
      dict(barrier_index=(0,), num_barriers=2),
      dict(barrier_index=(1,), num_barriers=2),
      dict(barrier_index=(0, 0), num_barriers=(1, 2)),
      dict(barrier_index=(0, 1), num_barriers=(1, 2)),
      dict(barrier_index=(1, 0), num_barriers=(2, 1)),
  )
  def test_mma_barrier_indexing(
      self,
      barrier_index,
      num_barriers,
      shape=(128, 128),
      swizzle=128,
      dtype=jnp.float16,
  ):
    transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)

    def kernel(a_smem, b_smem, out_ref, acc_tmem, barrier_ref):
      plgpu.tcgen05_mma(
          acc_tmem,
          a_smem,
          b_smem,
          barrier_ref.at[barrier_index],
          accumulate=False,
      )
      plgpu.barrier_wait(barrier_ref.at[barrier_index])
      plgpu.store(
          out_ref,
          plgpu.async_load_tmem(acc_tmem).astype(dtype),
          optimized=False,
      )

    scratch_shapes = [
        plgpu.TMEM(shape, jnp.float32, packed=False),
        plgpu.Barrier(num_barriers=num_barriers, orders_tensor_core=True),
    ]
    f = self.pallas_call(
        kernel,
        in_specs=(
            plgpu.BlockSpec(transforms=transforms, memory_space=plgpu.SMEM),
            plgpu.BlockSpec(transforms=transforms, memory_space=plgpu.SMEM),
        ),
        out_specs=plgpu.BlockSpec(memory_space=plgpu.GMEM),
        out_shape=jax.ShapeDtypeStruct(shape, dtype),
        scratch_shapes=scratch_shapes,
    )
    x = jax.random.uniform(jax.random.key(0), shape=shape, dtype=dtype)
    y = jax.random.uniform(jax.random.key(1), shape=shape, dtype=dtype)
    result = f(x, y)
    expected = jnp.matmul(x, y, precision="high")
    np.testing.assert_allclose(result, expected, rtol=1e-3)

  @parameterized.product(
      warp_level=(True, False),
      squeezed_index=(True, False),
      cluster=((2,), (2, 2)),
  )
  def test_copy_gmem_to_smem_partitioned(self, warp_level, squeezed_index, cluster):
    block_size = (128, 128)
    partitioned_block_size = (block_size[0] // 2, block_size[1])
    cluster_y = cluster[0] if len(cluster) > 1 else 1
    a = jax.random.uniform(
        jax.random.key(0), shape=(cluster_y, *block_size), dtype=jnp.float32)
    if squeezed_index:
      a = a.reshape(1, cluster_y, *block_size)
    b = jax.random.uniform(
        jax.random.key(1), shape=(cluster_y, *block_size), dtype=jnp.float32)
    def kernel(a_gmem, b_gmem, out_gmem,
              a_smem, b_smem,
              a_tma_barrier, b_tma_barrier, cluster_barrier):
      if squeezed_index:
        a_gmem = a_gmem.at[0]
      cluster_y_idx = lax.axis_index("y") if len(cluster) > 1 else 0
      a_gmem = a_gmem.at[cluster_y_idx]
      b_gmem = b_gmem.at[cluster_y_idx]
      cluster_idx = lax.axis_index("x")
      out_slice = pl.ds(cluster_idx * partitioned_block_size[0],
                        partitioned_block_size[0])

      if warp_level:
        @plgpu.warp_map
        def _per_warp(warp_id):
          @pl.when(warp_id == 0)
          def _():
            plgpu.copy_gmem_to_smem(
                a_gmem,
                a_smem,
                a_tma_barrier,
                collective_axes="x",
                leader_tracked=plgpu.CopyPartition.PARTITIONED(0),
            )
            plgpu.copy_gmem_to_smem(
                b_gmem,
                b_smem,
                b_tma_barrier,
                collective_axes="x",
                leader_tracked=plgpu.CopyPartition.PARTITIONED(0),
            )
      else:
        plgpu.copy_gmem_to_smem(
            a_gmem,
            a_smem,
            a_tma_barrier,
            collective_axes="x",
            leader_tracked=plgpu.CopyPartition.PARTITIONED(0),
        )
        plgpu.copy_gmem_to_smem(
            b_gmem,
            b_smem,
            b_tma_barrier,
            collective_axes="x",
            leader_tracked=plgpu.CopyPartition.PARTITIONED(0),
        )
      @pl.when(cluster_idx == 0)
      def _():
        plgpu.barrier_wait(a_tma_barrier)
        plgpu.barrier_wait(b_tma_barrier)
      plgpu.barrier_arrive(cluster_barrier)
      plgpu.barrier_wait(cluster_barrier)
      out_gmem[cluster_y_idx, out_slice] = a_smem[...] + b_smem[...]
    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((cluster_y, *block_size), jnp.float32),
        cluster_names=("y", "x") if len(cluster) > 1 else ("x",),
        cluster=cluster,
        scratch_types=(
            plgpu.SMEM(partitioned_block_size, jnp.float32),
            plgpu.SMEM(partitioned_block_size, jnp.float32),
            plgpu.Barrier(num_arrivals=1),
            plgpu.Barrier(num_arrivals=1),
            plgpu.ClusterBarrier(collective_axes=("x",)),
        ),
    )
    result = f(a, b)
    if squeezed_index:
      a = a[0]
    np.testing.assert_array_equal(result, a + b)

  @parameterized.parameters(((2,),), ((2, 2),))
  def test_copy_gmem_to_smem_replicated(self, cluster):
    block_size = (64, 64)
    cluster_y = cluster[0] if len(cluster) > 1 else 1
    def kernel(a_gmem, out_gmem, a_smem, tma_barrier, cluster_barrier):
      cluster_y_idx = lax.axis_index("y") if len(cluster) > 1 else 0
      cluster_idx = lax.axis_index("x")
      plgpu.copy_gmem_to_smem(
          a_gmem.at[cluster_y_idx], a_smem, tma_barrier,
          collective_axes="x", leader_tracked=plgpu.CopyPartition.REPLICATED,
      )
      @pl.when(cluster_idx == 0)
      def _():
        plgpu.barrier_wait(tma_barrier)
      plgpu.barrier_arrive(cluster_barrier)
      plgpu.barrier_wait(cluster_barrier)
      out_gmem[cluster_y_idx] = a_smem[...]
    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((cluster_y, *block_size), jnp.float32),
        cluster_names=("y", "x") if len(cluster) > 1 else ("x",),
        cluster=cluster,
        scratch_types=(
            plgpu.SMEM(block_size, jnp.float32),
            plgpu.Barrier(num_arrivals=1),
            plgpu.ClusterBarrier(collective_axes=("x",)),
        ),
    )
    a = jax.random.uniform(jax.random.key(0), shape=(cluster_y, *block_size),
                           dtype=jnp.float32)
    result = f(a)
    np.testing.assert_array_equal(result, a)

  def test_arrive_wait_on_tc_barrier(self):
    def kernel(out_ref, barrier):
      plgpu.barrier_arrive(barrier)
      plgpu.barrier_wait(barrier)
      out_ref[...] = jnp.ones_like(out_ref)

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
        scratch_types=(
            plgpu.Barrier(num_arrivals=1, orders_tensor_core=True),
        ),
    )
    np.testing.assert_array_equal(f(), np.ones((128,), np.float32))

  @parameterized.parameters(
      (0, (1,), False),
      (0, (1,), True),
      (1, (1,), False),
      (2, (1,), False),
      (0, (1, 2,), False),
      (0, (2, 1,), False),
  )
  def test_cluster_launch_control(self, dim, cluster, with_indexing):
    # We attempt to schedule 1 more CTA than can be scheduled at once. Only
    # one CTA will succeed in stealing the last 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
    cluster_size = math.prod(cluster)

    grid = [1, 1, 1]
    grid[dim] = num_sms // cluster_size + 1

    grid_names = tuple("xyz"[: len(grid)])
    cluster_names = tuple("abc"[: len(cluster)])

    def kernel(out_ref, cancel_result_ref, barrier, _):
      if with_indexing:
        cancel_result_ref = cancel_result_ref.at[0]
      plgpu.try_cluster_cancel(
          cancel_result_ref, barrier, collective_axes=(*cluster_names, "wg")
      )
      plgpu.barrier_wait(barrier)

      cta_ids, cancelled_launch = plgpu.query_cluster_cancel(
          cancel_result_ref, grid_names=grid_names)
      cta_id = sum(cta_ids)

      # Store a sentinel value if no work can be scheduled.
      value = lax.select(cancelled_launch, cta_id, jnp.int32(-1))

      grid_idx = lax.axis_index(grid_names) * lax.axis_size(
          cluster_names
      ) + lax.axis_index(cluster_names)
      out_ref[grid_idx] = value

    f = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((num_sms,), jnp.int32),
        grid=grid,
        grid_names=grid_names,
        num_threads=2,
        thread_name="wg",
        cluster=cluster,
        cluster_names=cluster_names,
        scratch_types=[
            plgpu.TryClusterCancelResult(2 if with_indexing else None),
            plgpu.Barrier(num_arrivals=2),
            # Requesting SMEM close to the 228kb limit to ensure that each SM
            # only schedules 1 block.
            plgpu.SMEM((220 * 1024,), jnp.int8),
        ],
    )
    result = np.sort(f())
    last_cta_id = math.ceil(num_sms / cluster_size)
    expected = np.array([-1] * (num_sms - cluster_size) + [last_cta_id] * cluster_size)
    np.testing.assert_equal(result, expected)

  @parameterized.parameters(True, False)
  def test_barrier_arrive_warp_mesh(self, orders_tensor_core):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((2,), jnp.int32),
        scratch_types=[
            plgpu.Barrier(num_arrivals=2, orders_tensor_core=orders_tensor_core)
        ],
    )
    def kernel(out_ref, bar):
      @plgpu.warp_map
      def _per_warp(warp_idx):

        @pl.when(warp_idx < 2)
        def _():
          plgpu.barrier_arrive(bar)
          plgpu.barrier_wait(bar)
          out_ref[warp_idx] = jnp.ones_like(out_ref.at[warp_idx])

    np.testing.assert_array_equal(kernel(), np.array([1, 1], dtype=np.int32))

  def test_barrier_arrive_predicated(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((4,), jnp.int32),
        scratch_types=[plgpu.Barrier(num_arrivals=2)],
    )
    def kernel(out_ref, bar):
      @plgpu.warp_map
      def _per_warp(warp_idx):
        plgpu.barrier_arrive(bar, predicate=(warp_idx < 2))
        plgpu.barrier_wait(bar)
        out_ref[warp_idx] = warp_idx + 1

    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      np.testing.assert_array_equal(
          kernel(), np.array([1, 2, 3, 4], dtype=np.int32)
      )
    self.assertIn("@%p2 mbarrier.arrive", ptx())


class PallasCallTCGen05WGTest(
    PallasCallTCGen05Test, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


class PipelineTest(PallasTest):

  def test_pipeline_mode(self):
    def body(x_ref, y_ref, o_ref):
      x = x_ref[:]
      y = y_ref[:]
      o_ref[:] = x + y

    data_size =  64 * 256
    block_size = 256

    x = jnp.arange(data_size, dtype=jnp.float32)
    y = jnp.arange(data_size, dtype=jnp.float32)
    in_specs = [
        pl.BlockSpec((block_size,), lambda *i: i, pipeline_mode=pl.Buffered(2)),
        pl.BlockSpec((block_size,), lambda *i: i, pipeline_mode=pl.Buffered(1))
    ]
    out_specs = pl.BlockSpec((block_size,), lambda *i: i)

    @jax.jit
    def vadd(x, y):
      return self.pallas_call(
          body,
          out_shape=jax.ShapeDtypeStruct.like(x),
          in_specs=in_specs,
          out_specs=out_specs,
          grid=data_size // block_size,
      )(x, y)

    with self.assertRaisesRegex(
        NotImplementedError, "pipeline_mode= is not supported"
    ):
      vadd(x, y)

  def test_manual(self):
    max_concurrent_steps = 2
    num_steps = 4

    def kernel(x_gmem, o_gmem):
      return pl.run_scoped(
          functools.partial(scoped_kernel, x_gmem, o_gmem),
          plgpu.SMEM((max_concurrent_steps, 32, 16), jnp.float32),
          plgpu.SMEM((max_concurrent_steps, 32, 16), jnp.float32),
          plgpu.Barrier(num_barriers=max_concurrent_steps),
      )

    def scoped_kernel(x_gmem, o_gmem, x_smem, o_smem, barrier):
      gmem_slice = pl.ds(lax.axis_index("x") * 32, 32)

      def body(step, _):
        slot = step % max_concurrent_steps

        # Wait for the current GMEM->SMEM copy to complete.
        plgpu.barrier_wait(barrier.at[slot])
        # Wait for the previous output SMEM->GMEM copy to complete.
        plgpu.wait_smem_to_gmem(max_concurrent_steps - 1)

        o_smem.at[slot][...] = x_smem.at[slot][...] + 1.0

        plgpu.commit_smem()
        plgpu.copy_smem_to_gmem(
            o_smem.at[slot], o_gmem.at[gmem_slice, pl.ds(step * 16, 16)]
        )

        fetch_step = step + max_concurrent_steps
        fetch_slot = slot  # (x + y) % y == x % y
        jax.lax.cond(
            fetch_step < num_steps,
            lambda: plgpu.copy_gmem_to_smem(
                x_gmem.at[gmem_slice, pl.ds(fetch_step * 16, 16)],
                x_smem.at[fetch_slot],
                barrier.at[fetch_slot],
            ),
            lambda: None,
        )
        return ()

      # Initialize the pipeline.
      for slot in range(min(max_concurrent_steps, num_steps)):
        plgpu.copy_gmem_to_smem(
            x_gmem.at[gmem_slice, pl.ds(slot * 16, 16)],
            x_smem.at[slot],
            barrier.at[slot],
        )

      jax.lax.fori_loop(0, num_steps, body, ())

      # Finalize the pipeline.
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(32 * 4 * 64).reshape(32 * 4, 64).astype(jnp.float32)
    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct.like(x),
        grid=(4, 1),
        grid_names=("x", "y"),
    )
    np.testing.assert_array_equal(kernel_fn(x), x + 1.0)

  @parameterized.product(
      transforms=(
          (),
          (plgpu.TilingTransform((8, 32)), plgpu.SwizzleTransform(128)),
      ),
      repeats=(1, 10),
  )
  def test_emit(self, transforms, repeats):
    num_steps = 4

    def kernel(x_gmem, o_gmem):
      for _ in range(repeats):
        plgpu.emit_pipeline(
            kernel_body,
            in_specs=[
                plgpu.BlockSpec(
                    (64, 64), lambda i: (0, i), transforms=transforms
                )
            ],
            out_specs=[
                plgpu.BlockSpec(
                    (64, 64), lambda i: (0, i), transforms=transforms
                )
            ],
            grid=(num_steps,),
            max_concurrent_steps=2,
        )(x_gmem, o_gmem)

    def kernel_body(_, x_smem, o_smem):
      # +1 for the indexing done by ``emit_pipeline`.
      self.assertLen(x_smem.transforms, len(transforms) + 1)
      o_smem[...] = x_smem[...] + 1.0

    x = jnp.arange(64 * num_steps * 64)
    x = x.reshape(-1, num_steps * 64).astype(jnp.float32)
    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct.like(x),
    )
    np.testing.assert_array_equal(kernel_fn(x), x + 1.0)

  def test_nested_emit(self):
    num_steps = 4

    def kernel(x_gmem, o_gmem):
      plgpu.emit_pipeline(
          nested_kernel,
          in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
          out_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
          grid=(),
      )(x_gmem, o_gmem)

    def nested_kernel(_, x_gmem, o_gmem):
      plgpu.emit_pipeline(
          nested_kernel_body,
          in_specs=[pl.BlockSpec((32, 16), lambda i: (0, i))],
          out_specs=[pl.BlockSpec((32, 16), lambda i: (0, i))],
          grid=(num_steps,),
          max_concurrent_steps=2,
      )(x_gmem, o_gmem)

    def nested_kernel_body(_, x_smem, o_smem):
      o_smem[...] = x_smem[...] + 1.0

    x = jnp.arange(32 * num_steps * 16)
    x = x.reshape(-1, num_steps * 16).astype(jnp.float32)
    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct.like(x),
    )
    np.testing.assert_array_equal(kernel_fn(x), x + 1.0)

  def test_emit_with_grid_invariant_output(self):
    num_steps = 4

    def kernel(x_gmem, o_gmem):
      plgpu.emit_pipeline(
          kernel_body,
          in_specs=[pl.BlockSpec((32, 16), lambda i: (0, i))],
          out_specs=[pl.BlockSpec((32, 16), lambda i: (0, 0))],
          grid=(num_steps,),
          max_concurrent_steps=2,
      )(x_gmem, o_gmem)

    def kernel_body(_, x_smem, o_smem):
      o_smem[...] = x_smem[...] + 1.0

    x = jnp.arange(32 * num_steps * 16)
    x = x.reshape(-1, num_steps * 16).astype(jnp.float32)
    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct.like(x),
    )
    y = jnp.empty_like(x)
    for i in range(num_steps):
      i_slice = slice(16 * i, 16 * (i + 1))
      y = y.at[:, :16].set(x[:, i_slice] + 1)
    # We only compare the elements in the first 16 columns, because the rest
    # are never written to.
    np.testing.assert_array_equal(kernel_fn(x)[:, :16], y[:, :16])

  @parameterized.parameters(((0, 2),), ((2, -1),))
  def test_emit_with_empty_grid(self, grid):
    def kernel_body(o_smem):
      o_smem[...] = jnp.zeros((128, 128))

    def kernel(o_gmem):
      plgpu.emit_pipeline(
          kernel_body,
          out_specs=[pl.BlockSpec((128, 128), lambda i, j: (i, j))],
          grid=grid,
      )(o_gmem)

    with self.assertRaisesRegex(
        ValueError, 'All elements in the grid must be strictly positive'
    ):
      self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((256, 256), jnp.float32),
      )()

  def test_emit_with_no_output(self):
    m, n = 16, 128

    def kernel(x_gmem, o_gmem):
      def acc_scope(acc_ref):
        acc_ref[...] = jnp.zeros_like(acc_ref)
        def body(_, x_smem):
          acc_ref[...] += x_smem[...]  # Can't += in a lambda...
        in_specs = [plgpu.BlockSpec((1, n), lambda i: (i, 0), delay_release=1)]
        plgpu.emit_pipeline(
            body,
            in_specs=in_specs,
            grid=(m,),
            max_concurrent_steps=2,
        )(x_gmem)
        return acc_ref[...]

      acc = pl.run_scoped(acc_scope, plgpu.SMEM((1, n), dtype=jnp.float32))

      o_gmem[...] = acc[...]

    dtype = jnp.float32
    x = jax.random.uniform(jax.random.key(0), (m, n)).astype(dtype)

    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((1, n), dtype),
    )

    np.testing.assert_allclose(kernel_fn(x), x.sum(0, keepdims=True), rtol=1e-6)

  @parameterized.product(
      max_concurrent_steps=[2, 3, 4],
      oob_fill_mode=[plgpu.OOBFillMode.PROMISE_IN_BOUNDS, plgpu.OOBFillMode.ZEROS],
  )
  @run_on_sm80
  def test_emit_in_specs_only(self, max_concurrent_steps, oob_fill_mode):
    m, n = 16, 128
    dtype = jnp.float32

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((1, n), dtype),
        scratch_types=[plgpu.SMEM((1, n), dtype=jnp.float32)],
    )
    def kernel(x_gmem, o_gmem, acc_ref):
      acc_ref[...] = jnp.zeros_like(acc_ref)

      def body(_, x_smem):
        acc_ref[...] += x_smem[...]

      plgpu.emit_pipeline(
          body,
          in_specs=[
              plgpu.BlockSpec(
                  (1, n), lambda i: (i, 0), oob_fill_mode=oob_fill_mode,
              )
          ],
          grid=(m,),
          max_concurrent_steps=max_concurrent_steps,
      )(x_gmem)

      o_gmem[...] = acc_ref[...]

    x = jax.random.uniform(jax.random.key(0), (m, n)).astype(dtype)
    np.testing.assert_allclose(kernel(x), x.sum(0, keepdims=True), rtol=1e-6)

  @run_on_sm80
  def test_emit_in_specs_only_requires_in_bounds(self):
    if jtu.is_cuda_compute_capability_at_least("9.0"):
      self.skipTest("cp.async OOB constraint is pre-Hopper only")

    n = 128

    @self.kernel()
    def kernel(x_gmem):
      plgpu.emit_pipeline(
          lambda _, x_smem: None,
          in_specs=[plgpu.BlockSpec((4, n), lambda i: (i, 0))],
          grid=(2,),
      )(x_gmem)

    x = jnp.zeros((6, n), dtype=jnp.float32)

    with self.assertRaisesRegex(NotImplementedError, "provably in bounds"):
      kernel(x)

  @parameterized.product(num_steps=[0, 1, 5, 100], max_concurrent_steps=[2, 3])
  @run_on_sm80
  def test_emit_in_specs_only_dynamic_grid(
      self, num_steps, max_concurrent_steps
  ):
    dtype = jnp.float32
    block_size = 128

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((block_size,), dtype),
        scratch_types=[plgpu.SMEM((block_size,), dtype)],
    )
    def kernel(num_steps_gmem, x_gmem, o_gmem, acc_ref):
      acc_ref[...] = jnp.zeros_like(acc_ref)

      def body(_, x_smem):
        acc_ref[...] += x_smem[...]

      plgpu.emit_pipeline(
          body,
          in_specs=[pl.BlockSpec((block_size,), lambda i: (i,))],
          grid=(num_steps_gmem[...],),
          max_concurrent_steps=max_concurrent_steps,
      )(x_gmem)

      o_gmem[...] = acc_ref[...]

    x = jnp.ones(max(num_steps, 1) * block_size, dtype=dtype)
    np.testing.assert_array_equal(
        kernel(jnp.int32(num_steps), x),
        jnp.full(block_size, num_steps, dtype=dtype)
    )

  def test_pipeline_oob_mode(self):
    # This test crashes with the default OOB fill mode of ZEROS because
    # it can't copy large 1D arrays.

    dtype = jnp.int8
    block_size = 512

    def kernel(x_gmem, o_gmem):
      def pipeline_body(_, x_smem, o_smem):
        o_smem[...] = x_smem[...]

      plgpu.emit_pipeline(
          pipeline_body,
          grid=(1,),
          in_specs=[
              plgpu.BlockSpec(
                  (block_size,),
                  lambda i: (i,),
                  oob_fill_mode=plgpu.OOBFillMode.UNDEFINED,
              )
          ],
          out_specs=[plgpu.BlockSpec((block_size,), lambda i: (i,))],
          max_concurrent_steps=1,
      )(x_gmem, o_gmem)

    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((block_size,), dtype),
    )

    x = jnp.arange(block_size, dtype=dtype)
    np.testing.assert_array_equal(kernel_fn(x), x)

  def test_emit_with_parallel_grid(self):
    num_steps1 = 4
    num_steps2 = 5

    def kernel(x_gmem, o_gmem):
      pid = lax.axis_index("x")
      plgpu.emit_pipeline(
          kernel_body,
          in_specs=[pl.BlockSpec((32, 16), lambda i: (pid, i))],
          out_specs=[pl.BlockSpec((32, 16), lambda i: (pid, i))],
          grid=(num_steps2,),
          max_concurrent_steps=2,
      )(x_gmem, o_gmem)

    def kernel_body(_, x_smem, o_smem):
      o_smem[...] = x_smem[...] + 1.0

    x = jnp.arange(num_steps1 * 32 * num_steps2 * 16)
    x = x.reshape(-1, num_steps2 * 16).astype(jnp.float32)
    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct.like(x),
        grid=(num_steps1,),
        grid_names=("x",),
    )
    y = x + 1.0
    np.testing.assert_array_equal(kernel_fn(x), y)

  def test_emit_with_dynamic_grid_smaller_than_concurrent_steps(self):
    block_x = 128
    x = jax.random.randint(jax.random.key(1234), (block_x,),
                           minval=-128, maxval=128, dtype=jnp.int32)

    def body(num_blocks_gmem, x_gmem, o_gmem):
      num_blocks = num_blocks_gmem[...]
      def pipeline_body(_, x_smem, o_smem):
        o_smem[...] = x_smem[...]
      for _ in range(2):
        plgpu.emit_pipeline(
            pipeline_body,
            grid=(num_blocks,),
            in_specs=[plgpu.BlockSpec((block_x,), lambda i: (i,))],
            out_specs=[plgpu.BlockSpec((block_x,), lambda i: (i,))],
            max_concurrent_steps=2,
        )(x_gmem, o_gmem)

    # The test only intends to check that this does not crash/hang.
    self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((block_x,), jnp.int32),
    )(0, x).block_until_ready()

  @parameterized.product(static=[False, True], short=[False, True])
  def test_emit_with_2d_grid(self, static, short):
    num_steps1 = 4
    num_steps2 = 5
    if short:
      num_steps1 = num_steps2 = 1

    def kernel(x_gmem, o_gmem):
      grid = (num_steps1, num_steps2)
      if static:
        grid = jax.tree.map(jnp.asarray, grid)

      plgpu.emit_pipeline(
          kernel_body,
          in_specs=[pl.BlockSpec((32, 16, 8), lambda i, j: (0, i, j))],
          out_specs=[pl.BlockSpec((32, 16, 8), lambda i, j: (0, i, j))],
          grid=grid,
          max_concurrent_steps=2,
      )(x_gmem, o_gmem)

    def kernel_body(_, x_smem, o_smem):
      o_smem[...] = x_smem[...] + 1.0

    x = jnp.arange(32 * num_steps1 * 16 * num_steps2 * 8)
    x = x.reshape(-1, num_steps1 * 16, num_steps2 * 8).astype(jnp.float32)
    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct.like(x),
    )
    np.testing.assert_array_equal(kernel_fn(x), x + 1.0)

  def test_emit_with_carry(self):
    num_steps = 4

    def kernel(o_gmem):
      plgpu.emit_pipeline(
          kernel_body,
          out_specs=[pl.BlockSpec((64, 64), lambda i: (0, i))],
          grid=(num_steps,),
          max_concurrent_steps=2,
          init_carry=0,
      )(o_gmem)

    def kernel_body(_, o_smem, carry):
      o_smem[...] = lax.broadcast(carry, o_smem.shape)
      return carry + 1

    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((64, num_steps * 64), jnp.int32),
    )
    np.testing.assert_array_equal(
        kernel_fn(), jnp.tile(jnp.repeat(jnp.arange(num_steps), 64), (64, 1))
    )

  @parameterized.parameters((pl.Squeezed(),), (None,))
  def test_emit_with_squeezed_dim(self, squeezed_dim):

    shape = (16, 256)
    num_steps = shape[0]

    def kernel(x_gmem, o_gmem):
      plgpu.emit_pipeline(
          kernel_body,
          in_specs=[pl.BlockSpec((squeezed_dim, shape[1]), lambda i: (i, 0))],
          out_specs=[pl.BlockSpec((squeezed_dim, shape[1]), lambda i: (i, 0))],
          grid=(num_steps,),
          max_concurrent_steps=2,
      )(x_gmem, o_gmem)

    def kernel_body(_, in_smem, o_smem):
      assert in_smem.shape == (shape[1],)
      assert o_smem.shape == (shape[1],)
      o_smem[...] = in_smem[...] + 1

    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((16, 256), jnp.int32),
    )
    x = jnp.arange(16 * 256, dtype=jnp.int32).reshape(16, 256)
    np.testing.assert_array_equal(kernel_fn(x), x + 1)

  def test_emit_with_element_dim(self):
    shape = (20, 256)

    @self.kernel()
    def kernel(x_gmem, o_gmem):
      index_map = lambda i: (10 * i + 1, 0)

      @functools.partial(
          plgpu.emit_pipeline,
          in_specs=[pl.BlockSpec((pl.Element(8), shape[1]), index_map)],
          out_specs=[pl.BlockSpec((pl.Element(8), shape[1]), index_map)],
          grid=(2,),
          max_concurrent_steps=2,
      )
      def pipeline(_, x_smem, o_smem):
        assert x_smem.shape == (8, shape[1])
        assert o_smem.shape == (8, shape[1])
        o_smem[...] = x_smem[...] + 1

      pipeline(x_gmem, o_gmem)

    x = jnp.arange(np.prod(shape), dtype=jnp.int32).reshape(*shape)

    @jax.jit
    def run(x):
      x_ref = jax.new_ref(x, memory_space=plgpu.GMEM)
      kernel(x_ref, x_ref)
      return jax.freeze(x_ref)

    active_rows = ((jnp.arange(shape[0]) - 1) % 10 < 8).astype(x.dtype)
    np.testing.assert_array_equal(run(x), x + active_rows[..., None])


class PipelineWGTest(
    PipelineTest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


class PipelineSm90ATest(PallasSm90ATest):

  def test_realistic_matmul(self):
    dtype = jnp.float16
    swizzle = 128
    elems_128b = swizzle // jnp.dtype(dtype).itemsize
    grid_m, grid_k, grid_n = 132, 10, 4
    tile_m = tile_n = 128
    assert tile_m % elems_128b == 0
    tile_k = elems_128b
    m, k, n = grid_m * tile_m, grid_k * tile_k, grid_n * tile_n

    transforms = self.default_transforms(swizzle=swizzle, dtype=dtype)

    def kernel(a_gmem, b_gmem, o_smem, acc):
      def kernel_body(_, a_smem, b_smem):
        assert a_smem.shape == (tile_m, tile_k)
        assert b_smem.shape == (tile_k, tile_n)
        plgpu.wgmma(acc, a_smem, b_smem)
        plgpu.wgmma_wait(1)

      pid_m = pl.program_id(0)
      pid_n = pl.program_id(1)
      in_specs = [
          plgpu.BlockSpec(
              (tile_m, tile_k), lambda k: (pid_m, k), transforms=transforms,
              delay_release=1,
          ),
          plgpu.BlockSpec(
              (tile_k, tile_n), lambda k: (k, pid_n), transforms=transforms,
              delay_release=1,
          ),
      ]
      plgpu.emit_pipeline(
          kernel_body,
          in_specs=in_specs,
          grid=(grid_k,),
          max_concurrent_steps=2,
      )(a_gmem, b_gmem)

      o_smem[...] = acc[...].astype(dtype)

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(m, k), dtype=dtype)
    b = jax.random.uniform(key2, shape=(k, n), dtype=dtype)

    res = self.pallas_call(
        kernel,
        in_specs=[
            pl.BlockSpec(memory_space=plgpu.GMEM),
            pl.BlockSpec(memory_space=plgpu.GMEM),
        ],
        out_specs=plgpu.BlockSpec(
            (tile_m, tile_n), lambda m, n: (m, n), transforms=transforms
        ),
        out_shape=jax.ShapeDtypeStruct((m, n), jnp.float16),
        scratch_shapes=[plgpu.ACC((tile_m, tile_n), jnp.float32)],
        grid=(grid_m, grid_n),
    )(a, b)
    np.testing.assert_array_equal(res, jnp.matmul(a, b, precision="high"))


class PipelineSm90AWGTest(
    PipelineSm90ATest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


class WarpSpecializedPipelineTest(PallasTest):

  @parameterized.product(m=[512], n=[512], repeats=[1, 10],
                         manual_consumed_barriers=[False, True],
                         max_concurrent_steps=[2, 3])
  def test_pipelined_copy(
      self, m, n, repeats, manual_consumed_barriers, max_concurrent_steps
  ):
    x = jax.random.uniform(jax.random.key(0), (m, n), dtype=jnp.float16)
    blk_m = blk_n = 32

    def copy_kernel(_, x_smem, o_smem, o_last_block_smem, *consumed_barriers):
      o_smem[...] = x_smem[...]
      o_last_block_smem[...] = x_smem[...]
      if manual_consumed_barriers:
        [x_barrier] = consumed_barriers
        plgpu.barrier_arrive(x_barrier)

    spec = pl.BlockSpec(
        block_shape=(2 * blk_m, blk_n), index_map=lambda i, j: (i, j)
    )
    def body(*gmem_refs):
      pipeline = mgpu_pipeline.emit_pipeline_warp_specialized(
          copy_kernel,
          grid=(m // (2 * blk_m), n // blk_n),
          memory_registers=40,
          max_concurrent_steps=max_concurrent_steps,
          num_compute_wgs=1,
          wg_axis="wg",
          manual_consumed_barriers=manual_consumed_barriers,
          in_specs=[spec],
          out_specs=[
              spec,
              # Create an index-invariant output.
              pl.BlockSpec(
                  block_shape=(2 * blk_m, blk_n), index_map=lambda i, j: (0, 0)
              ),
          ],
      )
      for _ in range(repeats):
        pipeline(*gmem_refs)  # Make sure we can run the pipeline multiple times
    kernel = self.kernel(
        body,
        out_type=(
            jax.ShapeDtypeStruct((m, n), jnp.float16),
            jax.ShapeDtypeStruct((2 * blk_m, blk_n), jnp.float16),
        ),
        compiler_params=plgpu.CompilerParams(approx_math=True),
        num_threads=2,
        thread_name="wg",
    )
    out, out_last_block = kernel(x)
    np.testing.assert_array_equal(out, x)
    np.testing.assert_array_equal(out_last_block, x[-(2 * blk_m):, -blk_n:])

  @parameterized.product(
      m=[256, 64],
      n=[256, 64],
      num_compute_wgs=[1],  # TODO(apaszke): Use 2WGs once we add support for outputs.
      static=[False, True],
      manual_produced_barriers=[False, True],
      manual_consumed_barriers=[False, True],
      in_tree_template=[(0, 1), ((0, (1,), None))],
  )
  def test_elementwise_add(self, m, n, num_compute_wgs, static,
                           manual_produced_barriers, manual_consumed_barriers, in_tree_template):
    blk_m = blk_n = 64
    if m % (num_compute_wgs * blk_m):
      self.skipTest(f"{m=} must be divisible by {num_compute_wgs=} * {blk_m=}")
    spec = pl.BlockSpec(
        block_shape=(num_compute_wgs * blk_m, blk_n),
        index_map=lambda i, j: (i, j),
    )
    in_treedef = jax.tree.structure(in_tree_template)
    in_specs = jax.tree.unflatten(in_treedef, (spec, spec))

    def tiled_add_kernel(_, *smems):
      flat_smems, _ = jax.tree.flatten(smems)
      x_smem, y_smem, o_smem, *barriers = flat_smems

      if manual_produced_barriers:
        x_produced_barrier, y_produced_barrier, *barriers = barriers
        plgpu.barrier_wait(x_produced_barrier)
        plgpu.barrier_wait(y_produced_barrier)
      wg_idx = lax.axis_index("wg")
      m_slice = pl.ds(wg_idx * blk_m, blk_m)
      o_smem[m_slice] = x_smem[m_slice] + y_smem[m_slice]
      if manual_consumed_barriers:
        x_consumed_barrier, y_consumed_barrier = barriers
        plgpu.barrier_arrive(x_consumed_barrier)
        plgpu.barrier_arrive(y_consumed_barrier)

    def pipeline(*gmem_refs):
      grid = (m // (num_compute_wgs * blk_m), n // blk_n)
      if not static:
        grid = jax.tree.map(jnp.asarray, grid)
      return mgpu_pipeline.emit_pipeline_warp_specialized(
          tiled_add_kernel,
          grid=grid,
          max_concurrent_steps=2,
          num_compute_wgs=num_compute_wgs,
          memory_registers=40,
          wg_axis="wg",
          in_specs=in_specs,
          out_specs=[spec],
          manual_produced_barriers=manual_produced_barriers,
          manual_consumed_barriers=manual_consumed_barriers,
      )(*gmem_refs)

    kernel = self.kernel(
        pipeline,
        out_type=jax.ShapeDtypeStruct((m, n), jnp.float32),
        compiler_params=plgpu.CompilerParams(approx_math=True),
        num_threads=num_compute_wgs + 1,
        thread_name="wg",
    )
    x = jax.random.uniform(jax.random.key(0), (m, n), dtype=jnp.float32)
    y = jax.random.uniform(jax.random.key(1), (m, n), dtype=jnp.float32)
    inputs = jax.tree.unflatten(in_treedef, (x, y))
    np.testing.assert_allclose(kernel(*inputs), x + y, atol=1e-4)

  def test_carry_accumulate(self, m=256, n=256, num_compute_wgs=2):
    blk_m = blk_n = 64

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((blk_m, blk_n), jnp.float32),
        scratch_types=[
            plgpu.SMEM((blk_m, blk_n), jnp.float32),
        ],
        compiler_params=plgpu.CompilerParams(approx_math=True),
        num_threads=num_compute_wgs + 1,
        thread_name="wg",
    )
    def kernel(x_gmem, acc_gmem, acc_smem):
      def _compute_thread(pipeline_fn):
        # Cast the init value to the same layout as x_smem, so the pipeline loop
        # carry has a constant signature.
        o_acc = plgpu.layout_cast(
          jnp.full((blk_m, blk_n,), 0, dtype=jnp.float32),
          plgpu.Layout.WG_STRIDED((blk_m, blk_n), vec_size=2))
        # Pass control to the pipeline emitter and return the final carry.
        o_final = pipeline_fn(o_acc)
        # Note that both compute WGs are doing identical work so the potential
        # race condition on the store here won't affect the result.
        acc_smem[...] = o_final
        plgpu.commit_smem()
        plgpu.copy_smem_to_gmem(acc_smem, acc_gmem)
        plgpu.wait_smem_to_gmem(0)

      def tiled_acc_kernel(_, x_smem, carry):
        new_carry = x_smem[...] + carry
        return new_carry

      pipeline = mgpu_pipeline.emit_pipeline_warp_specialized(
          tiled_acc_kernel,
          grid=(m // blk_m, n // blk_n),
          max_concurrent_steps=2,
          num_compute_wgs=num_compute_wgs,
          memory_registers=40,
          wg_axis="wg",
          compute_context=_compute_thread,
          in_specs=[
              pl.BlockSpec(
                  block_shape=(blk_m, blk_n), index_map=lambda i, j: (i, j)
              )
          ],
          out_specs=[],
      )
      pipeline(x_gmem)

    x = jax.random.uniform(jax.random.key(0), (m, n), dtype=jnp.float32)
    ref = jnp.sum(jnp.stack(np.split(x, m // blk_m, axis=0)), axis=0)
    ref = jnp.sum(jnp.stack(np.split(ref, n // blk_n, axis=1)), axis=0)
    np.testing.assert_allclose(kernel(x), ref, atol=1e-4)

  @parameterized.product(
      num_compute_wgs=[1],  # TODO(apaszke): Use 2WGs once we add support for outputs.
      static=[False, True],
      manual_consumed_barriers=[False, True],
      small_shape=[True, False],
      max_concurrent_steps=[2, 3, 4],
  )
  def test_delay_release(
      self, num_compute_wgs, static, manual_consumed_barriers, small_shape,
      max_concurrent_steps
  ):
    if small_shape:
      m = n = 64
    else:
      m = n = 256
    blk_m, blk_n = 32, 64
    spec = plgpu.BlockSpec(
        block_shape=(num_compute_wgs * blk_m, blk_n),
        index_map=lambda i, j: (i, j),
        delay_release=1,
    )
    out_spec = pl.BlockSpec(
        block_shape=(num_compute_wgs * blk_m, blk_n),
        index_map=lambda i, j: (i, j),
    )

    def tiled_add_kernel(idx, x_smem, y_smem, o_smem, *consumed_barriers):
      wg_idx = lax.axis_index("wg")
      m_slice = pl.ds(wg_idx * blk_m, blk_m)
      o_smem[m_slice] = x_smem[m_slice] + y_smem[m_slice]
      if manual_consumed_barriers:
        @pl.when(jnp.logical_or(idx[0] != 0, idx[1] != 0))
        def _signal_consumed():
          for b in consumed_barriers:
            plgpu.barrier_arrive(b)

    def pipeline(*gmem_refs):
      grid = (m // (num_compute_wgs * blk_m), n // blk_n)
      if not static:
        grid = jax.tree.map(jnp.asarray, grid)
      return mgpu_pipeline.emit_pipeline_warp_specialized(
          tiled_add_kernel,
          grid=grid,
          max_concurrent_steps=max_concurrent_steps,
          manual_consumed_barriers=manual_consumed_barriers,
          num_compute_wgs=num_compute_wgs,
          memory_registers=40,
          wg_axis="wg",
          in_specs=[spec, spec],
          out_specs=[out_spec],
      )(*gmem_refs)

    kernel = self.kernel(
        pipeline,
        out_type=jax.ShapeDtypeStruct((m, n), jnp.float32),
        compiler_params=plgpu.CompilerParams(approx_math=True),
        num_threads=num_compute_wgs + 1,
        thread_name="wg",
    )
    x = jax.random.uniform(jax.random.key(0), (m, n), dtype=jnp.float32)
    y = jax.random.uniform(jax.random.key(1), (m, n), dtype=jnp.float32)
    np.testing.assert_allclose(kernel(x, y), x + y, atol=1e-4)

  def test_different_delay_release(self):
    m, n = 128, 64
    blk_m, blk_n = 32, 32
    in_specs = [
        plgpu.BlockSpec(
            block_shape=(blk_m, blk_n),
            index_map=lambda i, j: (i, j),
            delay_release=delay,
        )
        for delay in range(3)
    ]
    out_spec = pl.BlockSpec(
        block_shape=(blk_m, blk_n),
        index_map=lambda i, j: (i, j),
    )

    def tiled_add_kernel(_, x_smem, y_smem, z_smem, o_smem):
      o_smem[...] = x_smem[...] + y_smem[...] + z_smem[...]

    def pipeline(*gmem_refs):
      grid = (m // blk_m, n // blk_n)
      return mgpu_pipeline.emit_pipeline(
          tiled_add_kernel,
          grid=grid,
          max_concurrent_steps=4,
          in_specs=in_specs,
          out_specs=[out_spec],
      )(*gmem_refs)

    kernel = self.kernel(
        pipeline,
        out_type=jax.ShapeDtypeStruct((m, n), jnp.float32),
    )
    x = jax.random.uniform(jax.random.key(0), (m, n), dtype=jnp.float32)
    y = jax.random.uniform(jax.random.key(1), (m, n), dtype=jnp.float32)
    z = jax.random.uniform(jax.random.key(3), (m, n), dtype=jnp.float32)
    np.testing.assert_allclose(kernel(x, y, z), x + y + z)

  @parameterized.product(
      delay_release=[0, 1],
  )
  def test_repeat(self, delay_release):
    num_steps = 4

    def kernel_body(_, x_smem, o_smem):
      o_smem[...] = x_smem[...] + 1.0

    def kernel(x_gmem, o_gmem):
      in_specs = [
          plgpu.BlockSpec((64, 64), lambda i: (0, i), delay_release=delay_release)
      ]
      out_specs = [plgpu.BlockSpec((64, 64), lambda i: (0, i))]
      for _ in range(3):
        plgpu.emit_pipeline_warp_specialized(
            kernel_body,
            in_specs=in_specs,
            out_specs=out_specs,
            grid=(num_steps,),
            max_concurrent_steps=2,
            num_compute_wgs=1,
            memory_registers=40,
            wg_axis="wg",
        )(x_gmem, o_gmem)

    x = jnp.arange(64 * num_steps * 64)
    x = x.reshape(-1, num_steps * 64).astype(jnp.float32)
    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct.like(x),
        num_threads=2,
        thread_name="wg",
    )
    np.testing.assert_array_equal(kernel_fn(x), x + 1.0)

  @parameterized.parameters((False,), (True,))
  def test_stationary_input(self, flip):
    m = n = 256
    blk_m = blk_n = 64

    def add_kernel(_, x_smem, y_smem, o_smem):
      if flip:
        x_smem, y_smem = y_smem, x_smem
      o_smem[...] = x_smem[...] + y_smem[...]

    def body(*gmem_refs):
      mgpu_pipeline.emit_pipeline_warp_specialized(
          add_kernel,
          grid=(m // blk_m, n // blk_n),
          memory_registers=40,
          max_concurrent_steps=2,
          num_compute_wgs=1,
          wg_axis="wg",
          in_specs=[
              pl.BlockSpec(
                  block_shape=(blk_m, blk_n), index_map=lambda i, j: (i, j)
              ),
              pl.BlockSpec(
                  block_shape=(blk_m, blk_n), index_map=lambda i, j: (0, 0)
              )
          ][::(-1 if flip else 1)],
          out_specs=[
              pl.BlockSpec(
                  block_shape=(blk_m, blk_n), index_map=lambda i, j: (i, j)
              ),
          ],
      )(*gmem_refs)
    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((m, n), jnp.float16),
        num_threads=2,
        thread_name="wg",
    )
    x = jax.random.uniform(jax.random.key(1), (m, n), dtype=jnp.float16)
    y = jax.random.uniform(jax.random.key(2), (blk_m, blk_n), dtype=jnp.float16)
    ref = x + np.tile(y, (m // blk_m, n // blk_n))
    if flip:
      x, y = y, x
    # TODO(apaszke,justinfu): Fix the bug (this test freezes) and remove this restriction.
    with self.assertRaisesRegex(
        NotImplementedError,
        "Only inputs with a dependency on the grid are supported.",
    ):
      out = kernel(x, y)
      np.testing.assert_array_equal(out, ref)

  def test_no_output(self):
    m = n = 256
    blk_m = blk_n = 64

    def body(x_ref, o_ref, o_scratch, barrier):
      @pl.when(lax.axis_index("wg") == 0)
      def _():
        o_scratch[...] = jnp.zeros_like(o_scratch)

      # Wait for scratch to be initialized
      plgpu.barrier_arrive(barrier)
      plgpu.barrier_wait(barrier)

      # Make sure we can run the pipeline many times. This also introduces
      # extra jitter into warp scheduling and has uncovered bugs in the past.
      @pl.loop(0, 10)
      def _pipeline_loop(_):
        def add(_, x_smem):
          slc = pl.ds(lax.axis_index("wg") * (blk_m // 2), blk_m // 2)
          o_scratch[slc] += x_smem[slc]
        mgpu_pipeline.emit_pipeline_warp_specialized(
            add,
            grid=(m // blk_m, n // blk_n),
            memory_registers=40,
            max_concurrent_steps=2,
            num_compute_wgs=2,
            wg_axis="wg",
            in_specs=[
                pl.BlockSpec(
                    block_shape=(blk_m, blk_n), index_map=lambda i, j: (i, j)
                ),
            ]
        )(x_ref)

      # Wait for both compute WGs to finish initializing the output
      plgpu.barrier_arrive(barrier)
      plgpu.barrier_wait(barrier)

      @pl.when(lax.axis_index("wg") == 0)
      def _():
        plgpu.copy_smem_to_gmem(o_scratch, o_ref)
        plgpu.wait_smem_to_gmem(0, wait_read_only=True)

    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((blk_m, blk_n), jnp.float32),
        num_threads=3,
        thread_name="wg",
        scratch_types=[
            plgpu.SMEM((blk_m, blk_n), jnp.float32),
            plgpu.Barrier(num_arrivals=3),
        ],
    )
    x = jax.random.uniform(jax.random.key(1234), (m, n), dtype=jnp.float32)
    ref = 10 * x.reshape(m // blk_m, blk_m, n // blk_n, blk_n).sum((0, 2))
    np.testing.assert_allclose(kernel(x), ref, rtol=5e-6)

  @parameterized.product(manual_consumed_barriers=[False, True])
  def test_pipelined_pipeline(self, manual_consumed_barriers):
    m = n = 512

    x = jax.random.randint(jax.random.key(0), (m, n), -10, 15, dtype=jnp.int32)
    blk_m = blk_n = 64

    def body(x_ref, out_gmem_ref, out_ref):
      wg_idx = jax.lax.axis_index("wg")
      @pl.when(wg_idx == 0)
      def _zero_output():
        out_ref[...] = jnp.zeros_like(out_ref)

      def pipeline_body(_, x_smem, *consumed_barriers):
        out_ref[...] += x_smem[...]
        if manual_consumed_barriers:
          [x_barrier] = consumed_barriers
          plgpu.barrier_arrive(x_barrier)

      spec = pl.BlockSpec(
          block_shape=(blk_m, blk_n), index_map=lambda i, j: (i, j)
      )
      pipeline = functools.partial(
          mgpu_pipeline.emit_pipeline_warp_specialized,
          body=pipeline_body,
          grid=(m // blk_m, n // blk_n),
          memory_registers=40,
          max_concurrent_steps=2,
          num_compute_wgs=1,
          wg_axis="wg",
          manual_consumed_barriers=manual_consumed_barriers,
          in_specs=[spec],
      )

      @functools.partial(
          pl.run_scoped,
          allocs=pipeline(pipeline_state=None).get_allocations(x_ref),
          collective_axes="wg",
      )
      def _pipeline_scope(allocs):
        @pl.loop(0, 2)
        def _outer_loop(_):
          @pl.loop(0, 4)
          def _pipeline_loop(i):
            state = plgpu.PipelinePipeline.START
            state = jnp.where(i > 0, plgpu.PipelinePipeline.STEADY, state)
            state = jnp.where(i == 3, plgpu.PipelinePipeline.STOP, state)
            pipeline(pipeline_state=state)(x_ref, allocations=allocs)
          # Make sure we have properly quiesced the pipeline.
          pipeline(pipeline_state=None)(x_ref, allocations=allocs)

      @pl.when(wg_idx == 0)
      def _store_out():
        out_gmem_ref[...] = out_ref[...]

    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((blk_m, blk_n), jnp.int32),
        compiler_params=plgpu.CompilerParams(approx_math=True),
        scratch_types=[plgpu.SMEM((blk_m, blk_n), jnp.int32)],
        num_threads=2,
        thread_name="wg",
    )
    out = kernel(x)
    np.testing.assert_array_equal(
        out, x.reshape(m // blk_m, blk_m, n // blk_n, blk_n).sum((0, 2)) * 10
    )

  @parameterized.product(manual_consumed_barriers=[False, True])
  def test_pipeline_with_manual_allocation(self, manual_consumed_barriers):
    m = n = 512

    x = jax.random.randint(jax.random.key(4), (m, n), -10, 15, dtype=jnp.int32)
    y = jax.random.randint(jax.random.key(5), (m, n), -10, 15, dtype=jnp.int32)
    blk_m = blk_n = 64

    def body(x_ref, y_ref, out_gmem_ref, out_ref):
      wg_idx = jax.lax.axis_index("wg")
      @pl.when(wg_idx == 0)
      def _zero_output():
        out_ref[...] = jnp.zeros_like(out_ref)

      def pipeline_body(_, x_smem, y_smem, *consumed_barriers):
        out_ref[...] += x_smem[...] + y_smem[...]
        for b in consumed_barriers:
          plgpu.barrier_arrive(b)

      spec = pl.BlockSpec(
          block_shape=(blk_m, blk_n), index_map=lambda i, j: (i, j)
      )
      pipeline = mgpu_pipeline.emit_pipeline_warp_specialized(
          body=pipeline_body,
          grid=(m // blk_m, n // blk_n),
          memory_registers=40,
          max_concurrent_steps=2,
          num_compute_wgs=1,
          wg_axis="wg",
          manual_consumed_barriers=manual_consumed_barriers,
          in_specs=[spec, spec],
      )

      @functools.partial(
          pl.run_scoped,
          allocs=pipeline.get_allocations(x_ref, y_ref),
          collective_axes="wg",
      )
      def _alloc_scope(allocs):
        @pl.loop(0, 4)
        def _outer_loop(_):
          pipeline(x_ref, y_ref, allocations=allocs)

      @pl.when(wg_idx == 0)
      def _store_out():
        out_gmem_ref[...] = out_ref[...]

    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((blk_m, blk_n), jnp.int32),
        compiler_params=plgpu.CompilerParams(approx_math=True),
        scratch_types=[plgpu.SMEM((blk_m, blk_n), jnp.int32)],
        num_threads=2,
        thread_name="wg",
    )
    np.testing.assert_array_equal(
        kernel(x, y),
        (x + y).reshape(m // blk_m, blk_m, n // blk_n, blk_n).sum((0, 2)) * 4,
    )

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_collective(self):
    num_steps = 4

    def kernel(x_gmem, o_gmem):
      cluster_idx = lax.axis_index("cluster")
      in_specs = [
          plgpu.BlockSpec(
              (64, 64), lambda i: (0, i), collective_axes=("cluster",)
          )
      ]
      out_specs = [plgpu.BlockSpec((1, 64, 64), lambda i: (cluster_idx, 0, i))]
      # Run a few times to make sure we leave barriers in a good state.
      for _ in range(3):
        def pipeline_body(_, x_smem, o_smem):
          o_smem[0, ...] = x_smem[...] + 1.0
        plgpu.emit_pipeline_warp_specialized(
            pipeline_body,
            in_specs=in_specs,
            out_specs=out_specs,
            grid=(num_steps,),
            max_concurrent_steps=2,
            num_compute_wgs=1,
            memory_registers=40,
            wg_axis="wg",
        )(x_gmem, o_gmem)

    x = jnp.arange(64 * num_steps * 64)
    x = x.reshape(-1, num_steps * 64).astype(jnp.float32)
    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((2, *x.shape), x.dtype),
        num_threads=2,
        thread_name="wg",
        cluster=(2,),
        cluster_names=("cluster",)
    )
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), self.capture_stdout() as ptx:
      y = jax.block_until_ready(kernel_fn(x))
    self.assertIn(
        "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster",
        ptx(),
    )
    np.testing.assert_array_equal(y, np.stack([x + 1.0, x + 1.0]))

  @parameterized.parameters((pl.Squeezed(),), (None,))
  def test_emit_with_squeezed_dim(self, squeezed_dim):
    shape = (16, 256)
    num_steps = shape[0]

    def kernel(x_gmem, o_gmem):
      plgpu.emit_pipeline_warp_specialized(
          kernel_body,
          in_specs=[pl.BlockSpec((squeezed_dim, shape[1]), lambda i: (i, 0))],
          out_specs=[pl.BlockSpec((squeezed_dim, shape[1]), lambda i: (i, 0))],
          grid=(num_steps,),
          max_concurrent_steps=2,
          num_compute_wgs=1,
          memory_registers=40,
          wg_axis="wg",
      )(x_gmem, o_gmem)

    def kernel_body(_, in_smem, o_smem):
      o_smem[...] = in_smem[...] + 1

    kernel_fn = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((16, 256), jnp.int32),
        num_threads=2,
        thread_name="wg",
    )
    x = jnp.arange(16 * 256, dtype=jnp.int32).reshape(16, 256)
    np.testing.assert_array_equal(kernel_fn(x), x + 1)


class WarpSpecializedPipelineWGTest(
    WarpSpecializedPipelineTest,
    lowering_semantics=plgpu.LoweringSemantics.Warpgroup,
):
  ...


class CoreMapTest(PallasTest, jtu.CudaArchSpecificTest):

  def test_multiple_wg(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((2, 128), np.int32),
        num_threads=2,
        thread_name="wg",
    )
    def kernel(o_ref):
      wg_idx = jax.lax.axis_index("wg")
      o_ref[wg_idx] = jnp.broadcast_to(wg_idx, (128,))

    np.testing.assert_array_equal(
        kernel(), np.repeat(np.arange(2), 128).reshape(2, 128)
    )

  def test_multiple_wg_with_grid(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((4, 2, 128), np.int32),
        grid=(2, 2),
        grid_names=("x", "y"),
        num_threads=2,
        thread_name="wg",
    )
    def kernel(o_ref):
      xy_idx = jax.lax.axis_index(("x", "y"))
      yx_idx = jax.lax.axis_index(("y", "x"))
      wg_idx = jax.lax.axis_index("wg")
      num_wgs = jax.lax.axis_size("wg")
      o_ref[xy_idx, wg_idx] = jnp.broadcast_to(
          yx_idx * num_wgs + wg_idx, (128,)
      )

    np.testing.assert_array_equal(
        kernel(), np.repeat([0, 1, 4, 5, 2, 3, 6, 7], 128).reshape(4, 2, 128)
    )

  def test_multiple_wg_with_squashed_grid(self):
    # Tests whether a grid with >3 logical dimensions is correctly squashed to
    # 3 CUDA grid dimensions.
    b = 4
    x_dim = 3
    y_dim = 5
    z_dim = 7
    num_threads = 2

    @self.kernel(
        out_type=jax.ShapeDtypeStruct(
            (b, x_dim, y_dim, z_dim, num_threads, 128), np.int32
        ),
        grid=(b, x_dim, y_dim, z_dim),
        grid_names=("b", "x", "y", "z"),
        num_threads=num_threads,
        thread_name="wg",
    )
    def kernel(o_ref):
      b_idx = jax.lax.axis_index("b")
      x_idx = jax.lax.axis_index("x")
      y_idx = jax.lax.axis_index("y")
      z_idx = jax.lax.axis_index("z")
      wg_idx = jax.lax.axis_index("wg")
      bxyzw_idx = jax.lax.axis_index(("b", "x", "y", "z", "wg"))
      o_ref[b_idx, x_idx, y_idx, z_idx, wg_idx] = jnp.broadcast_to(
          bxyzw_idx, (128,)
      )

    result = kernel()[:, :, :, :, :, 0]
    ref = np.arange(b * x_dim * y_dim * z_dim * num_threads).reshape(
        result.shape
    )
    np.testing.assert_array_equal(result, ref)

  def test_cross_wg_barrier(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((2, 128), np.int32),
        # Each warpgroup is a single logical thread!
        scratch_types=[plgpu.Barrier(num_arrivals=2)],
        num_threads=2,
        thread_name="wg",
    )
    def kernel(o_ref, barrier):
      plgpu.barrier_arrive(barrier)
      plgpu.barrier_wait(barrier)
      wg_idx = jax.lax.axis_index("wg")
      o_ref[wg_idx] = jnp.broadcast_to(wg_idx, (128,))

    np.testing.assert_array_equal(
        kernel(), np.repeat([0, 1], 128).reshape(2, 128)
    )

  def test_cluster(self):

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128,), np.int32),
        grid=(2,),
        grid_names=("x",),
        cluster=(2,),
        cluster_names=("cluster",),
    )
    def kernel(ref):
      block_idx = jax.lax.axis_index("x")
      cluster_idx = jax.lax.axis_index("cluster")
      pl.debug_print("block: {} cluster: {}", block_idx, cluster_idx)

      ref[...] = ref[...]

    with self.capture_stdout() as output:
      jax.block_until_ready(kernel())
    self.assertEqual(
        set(output().splitlines()),
        {
            "block: 0 cluster: 0",
            "block: 1 cluster: 0",
            "block: 0 cluster: 1",
            "block: 1 cluster: 1",
        },
    )

  def test_realistic_matmul_with_cluster(self):
    self.skip_unless_sm90a()  # Requires WGMMA.

    dtype = jnp.float16
    swizzle = 128
    elems_128b = swizzle // jnp.dtype(dtype).itemsize
    grid_m, grid_k, grid_n = 132, 10, 32
    # TODO(slebedev): Remove ``grid_tile_n`` to simplify the test.
    grid_tile_n = 4
    assert grid_n % grid_tile_n == 0
    cluster_m = 2
    cluster_n = 2
    cluster_tile_n = min(cluster_n, grid_tile_n)
    tile_m = tile_n = 128
    assert tile_m % elems_128b == 0
    tile_k = elems_128b
    m, k, n = grid_m * tile_m, grid_k * tile_k, grid_n * tile_n

    transforms = self.default_transforms(dtype=dtype)

    max_concurrent_steps = 2
    delay_release = 1

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((m, n), dtype),
        scratch_types=[
            plgpu.SMEM(
                (max_concurrent_steps, tile_m, tile_k),
                dtype,
                transforms=transforms,
            ),
            plgpu.SMEM(
                (max_concurrent_steps, tile_k, tile_n),
                dtype,
                transforms=transforms,
            ),
            plgpu.SMEM((tile_m, tile_n), dtype, transforms=transforms),
            plgpu.ACC((tile_m, tile_n), jnp.float32),
            plgpu.Barrier(num_arrivals=2, num_barriers=max_concurrent_steps),
            plgpu.ClusterBarrier(
                collective_axes=(("x", "z"), "y"),
                num_barriers=max_concurrent_steps,
            ),
        ],
        grid=(grid_tile_n, grid_m, grid_n // grid_tile_n),
        grid_names=("tile_n", "m", "n"),
        cluster=(cluster_tile_n, cluster_m, cluster_n // cluster_tile_n),
        cluster_names=("x", "y", "z"),
    )
    def kernel(
        a_gmem,
        b_gmem,
        o_gmem,
        a_smem,
        b_smem,
        o_smem,
        acc,
        barrier,
        cluster_barrier,
    ):
      m_slice = pl.ds(lax.axis_index("m") * tile_m, tile_m)
      n_slice = pl.ds(
          (lax.axis_index("tile_n") + lax.axis_index("n") * grid_tile_n)
          * tile_n,
          tile_n,
      )

      def fetch(step, slot):
        if not isinstance(slot, int):  # Skip in initialization.
          plgpu.barrier_arrive(cluster_barrier.at[slot])
          plgpu.barrier_wait(cluster_barrier.at[slot])

        k_slice = pl.ds(step * tile_k, tile_k)
        plgpu.copy_gmem_to_smem(
            a_gmem.at[m_slice, k_slice],
            a_smem.at[slot],
            barrier.at[slot],
            collective_axes=("x", "z"),
        )
        plgpu.copy_gmem_to_smem(
            b_gmem.at[k_slice, n_slice],
            b_smem.at[slot],
            barrier.at[slot],
            collective_axes="y",
        )

      # Initialize the pipeline.
      for slot in range(min(max_concurrent_steps, grid_k)):
        fetch(slot, slot)

      def body(step, _):
        slot = step % max_concurrent_steps
        plgpu.barrier_wait(barrier.at[slot])

        plgpu.wgmma(acc, a_smem.at[slot], b_smem.at[slot])
        plgpu.wgmma_wait(delay_release)

        fetch_step = step + (max_concurrent_steps - delay_release)
        fetch_slot = lax.rem(fetch_step, max_concurrent_steps)
        jax.lax.cond(
            lax.bitwise_and(step >= delay_release, fetch_step < grid_k),
            lambda: fetch(fetch_step, fetch_slot),
            lambda: None,
        )
        return ()

      jax.lax.fori_loop(0, grid_k, body, ())

      # Finalize the pipeline.
      o_smem[...] = acc[...].astype(dtype)
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(o_smem, o_gmem.at[m_slice, n_slice])
      plgpu.wait_smem_to_gmem(0)

    key1, key2 = jax.random.split(jax.random.key(42), 2)
    a = jax.random.uniform(key1, shape=(m, k), dtype=dtype)
    b = jax.random.uniform(key2, shape=(k, n), dtype=dtype)
    np.testing.assert_array_equal(
        kernel(a, b), jnp.matmul(a, b, precision="high")
    )


class CoreMapWGTest(
    CoreMapTest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


class PrettyPrintingTest(PallasTest):

  def test_load(self):

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct([2, 128], jnp.float32),
        in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
        out_specs=plgpu.BlockSpec(memory_space=plgpu.SMEM),
    )
    def kernel(x_ref, o_ref):
      for i in range(2):
        x = plgpu.load(x_ref.at[i])
        o_ref[i, ...] = x

    _ = str(jax.make_jaxpr(kernel)(jax.ShapeDtypeStruct((2, 128), jnp.float32)))

  def test_copy_primitives(self):
    num_steps = 4

    @self.kernel(out_type=jax.ShapeDtypeStruct((64, 64), jnp.float32))
    def kernel(x_gmem, o_gmem):
      # ``plgpu.emit_pipeline`` is implemented in terms of async copy and
      # synchronization primitives.
      plgpu.emit_pipeline(
          kernel_body,
          in_specs=[pl.BlockSpec((64, 64), lambda i: (0, i))],
          out_specs=[
              pl.BlockSpec(
                  (64, 64),
                  lambda i: (0, i),
              )
          ],
          grid=(num_steps,),
          max_concurrent_steps=2,
      )(x_gmem, o_gmem)

    def kernel_body(_, x_smem, o_smem):
      o_smem[...] = x_smem[...] + 1.0

    _ = str(jax.make_jaxpr(kernel)(jax.ShapeDtypeStruct((64, 64), jnp.float32)))

  def test_wgmma(self):
    transforms = self.default_transforms(dtype=jnp.float16)

    @functools.partial(
        self.pallas_call,
        out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float32),
        in_specs=[
            plgpu.BlockSpec(transforms=transforms),
            plgpu.BlockSpec(transforms=transforms),
        ],
    )
    def kernel(a_ref, b_ref, o_ref):
      def scope(acc_ref):
        plgpu.wgmma(acc_ref, a_ref[...], b_ref)
        return acc_ref[...]

      o_ref[...] = pl.run_scoped(scope, plgpu.ACC((64, 192), jnp.float32))

    _ = str(
        jax.make_jaxpr(kernel)(
            jax.ShapeDtypeStruct((64, 128), jnp.float16),
            jax.ShapeDtypeStruct((128, 192), jnp.float16),
        )
    )


class CostEstimateTest(PallasTest):

  FLOPS = 1234
  TRANSCENDENTALS = 21
  BYTES_ACCESSED = 5678
  REMOTE_BYTES_TRANSFERRED = 91

  def cost_estimate(self):
    return pl.CostEstimate(
        flops=self.FLOPS,
        transcendentals=self.TRANSCENDENTALS,
        bytes_accessed=self.BYTES_ACCESSED,
        remote_bytes_transferred=self.REMOTE_BYTES_TRANSFERRED,
    )

  def add_one_kernel(self, cost_estimate):
    @functools.partial(
        self.kernel,
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.float32),
        cost_estimate=cost_estimate,
    )
    def kernel(x_ref, o_ref):
      o_ref[...] = x_ref[...] + 1.0

    return kernel

  def extract_cost_estimate(self, f, *args) -> dict[str, int] | None:
    hlo = jax.jit(f).lower(*args).compiler_ir("hlo").as_hlo_text()
    match = re.search(r'cost_estimate_json\s*=\s*"((?:[^"\\]|\\.)*)"', hlo)
    if match is None:
      return None
    return json.loads(match.group(1).replace(r"\22", '"').replace(r"\"", '"'))

  def test_cost_estimate_is_attached_to_the_custom_call(self):
    x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
    kernel = self.add_one_kernel(self.cost_estimate())

    self.assertEqual(
        self.extract_cost_estimate(kernel, x),
        dataclasses.asdict(self.cost_estimate()),
    )

  def test_no_cost_estimate_by_default(self):
    x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
    kernel = self.add_one_kernel(None)

    self.assertIsNone(self.extract_cost_estimate(kernel, x))

  def test_cost_estimate_is_scaled_by_the_vmap_axis_size(self):
    axis_size = 3
    x = jnp.arange(axis_size * 128 * 128, dtype=jnp.float32).reshape(
        axis_size, 128, 128
    )
    kernel = self.add_one_kernel(self.cost_estimate())
    expected = pl.CostEstimate(
        flops=axis_size * self.FLOPS,
        transcendentals=axis_size * self.TRANSCENDENTALS,
        bytes_accessed=axis_size * self.BYTES_ACCESSED,
        remote_bytes_transferred=axis_size * self.REMOTE_BYTES_TRANSFERRED,
    )

    self.assertEqual(
        self.extract_cost_estimate(jax.vmap(kernel), x),
        dataclasses.asdict(expected),
    )


class ExportTest(PallasTest):

  def test_export_succeeds(self):
    out_shape = jax.ShapeDtypeStruct([128], jnp.float32)

    @functools.partial(self.pallas_call, out_shape=out_shape)
    def kernel(x_ref, o_ref):
      o_ref[...] = x_ref[...] + 1.0

    _ = export.export(kernel)(out_shape)


class ExamplesTest(PallasTest):

  # Basic
  def test_stage0(self):
    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)

    @self.kernel(out_type=x)
    def kernel(l_ref, r_ref, o_ref):
      o_ref[...] = l_ref[...] + r_ref[...]

    np.testing.assert_allclose(kernel(x, x), x + x)

  # Multi-block kernels
  def test_stage1(self):
    row_block = 64
    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)

    @self.kernel(
        out_type=x,
        grid=(2,),
        grid_names=("rows",),
    )
    def kernel(l_ref, r_ref, o_ref):
      my_slice = pl.ds(lax.axis_index("rows") * row_block, row_block)
      o_ref[my_slice] = l_ref[my_slice] + r_ref[my_slice]

    np.testing.assert_allclose(kernel(x, x), x + x)

  # Async copies
  def test_stage3(self):
    row_block, col_block = 64, 128

    @self.kernel(
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.float16),
        scratch_types=[
            *([plgpu.SMEM((row_block, col_block), jnp.float16)] * 3),
            plgpu.Barrier(num_arrivals=2),
        ],
        grid=(2,),
        grid_names=("rows",),
    )
    def kernel(l_ref, r_ref, o_ref, l_smem, r_smem, o_smem, barrier):
      my_slice = pl.ds(lax.axis_index("rows") * row_block, row_block)
      plgpu.copy_gmem_to_smem(l_ref.at[my_slice], l_smem, barrier)
      plgpu.copy_gmem_to_smem(r_ref.at[my_slice], r_smem, barrier)
      plgpu.barrier_wait(barrier)
      o_smem[...] = l_smem[...] + r_smem[...]
      plgpu.commit_smem()
      plgpu.copy_smem_to_gmem(o_smem, o_ref.at[my_slice])
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
    np.testing.assert_allclose(kernel(x, x), x + x)

  # Pipelining
  def test_stage4(self):
    row_block, col_block = 64, 32
    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)

    @self.kernel(
        out_type=x,
        grid=(2,),
        grid_names=("rows",),
    )
    def kernel(l_ref, r_ref, o_ref):
      def compute(_, l_smem, r_smem, o_smem):
        o_smem[...] = l_smem[...] + r_smem[...]
      r = lax.axis_index("rows")
      block = pl.BlockSpec((row_block, col_block), lambda c: (r, c))
      plgpu.emit_pipeline(
          compute,
          grid=(l_ref.shape[1] // col_block,),
          in_specs=[block] * 2,
          out_specs=[block],
      )(l_ref, r_ref, o_ref)

    np.testing.assert_allclose(kernel(x, x), x + x)

  # Transforms
  def test_stage5(self):
    row_block, col_block = 64, 32
    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)

    @self.kernel(
        out_type=x,
        grid=(2,),
        grid_names=("rows",),
    )
    def kernel(l_ref, r_ref, o_ref):
      def compute(_, l_smem, r_smem, o_smem):
        o_smem[...] = l_smem[...] + r_smem[...]
      r = lax.axis_index("rows")
      block = plgpu.BlockSpec(
          (row_block, col_block),
          lambda c: (r, c),
          transforms=self.default_transforms(swizzle=64, dtype=jnp.float16),
      )
      plgpu.emit_pipeline(
          compute,
          grid=(l_ref.shape[1] // col_block,),
          in_specs=[block] * 2,
          out_specs=[block],
      )(l_ref, r_ref, o_ref)

    np.testing.assert_allclose(kernel(x, x), x + x)


class ExamplesWGTest(
    ExamplesTest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


class SemaphoreTest(PallasTest):

  def test_lowering(self):
    # This is a smoke test until we add support for lowering of semaphore ops.
    def body(i_ref1, i_ref2, o_ref, sem_ref):
      del i_ref2  # Only here to have a different number of inputs and outputs.
      assert sem_ref.shape == (4,)
      assert jnp.issubdtype(sem_ref.dtype, pl.semaphore)
      o_ref[...] = i_ref1[...]
    x = jnp.arange(128, dtype=jnp.float32).reshape((128,))
    kernel = self.pallas_call(
        body,
        out_shape=x,
        scratch_shapes=[plgpu.SemaphoreType.REGULAR((4,))],
    )
    text = jax.jit(kernel).lower(x, x).as_text()
    self.assertIn(
        "output_operand_aliases ="
        " [#stablehlo.output_operand_alias<output_tuple_indices = [1],"
        " operand_index = 2, operand_tuple_indices = []>]",
        text,
    )
    self.assertIn(
        "(tensor<128xf32>, tensor<128xf32>, tensor<4xi32>) ->"
        " (tensor<128xf32>, tensor<4xi32>)",
        text,
    )

  def test_basic(self):
    def body(o_ref, sem_ref):
      assert jnp.issubdtype(sem_ref.dtype, pl.semaphore)
      pl.semaphore_signal(sem_ref)
      o_ref[...] = jnp.ones_like(o_ref)
      pl.semaphore_wait(sem_ref)
    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
        scratch_types=[plgpu.SemaphoreType.REGULAR],
        grid=(2,),
        grid_names=("x",),
    )
    text = jax.jit(kernel).lower().as_text()
    np.testing.assert_array_equal(kernel(), jnp.ones((128,), jnp.float32))
    # The semaphore array is scaled up by the grid size.
    self.assertIn("(tensor<2xi32>) -> (tensor<128xf32>, tensor<2xi32>)", text)

  def test_with_profiler(self):
    # Dealing with profiler and semaphores together is tricky because they both
    # add extra outputs to the HLO op.
    def body(o_ref, sem_ref):
      assert jnp.issubdtype(sem_ref.dtype, pl.semaphore)
      with jax.named_scope("output"):
        o_ref[...] = jnp.ones_like(o_ref)
    with tempfile.TemporaryDirectory() as tmp_dir:
      kernel = self.kernel(
          body,
          out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
          scratch_types=[plgpu.SemaphoreType.REGULAR],
          grid=(2,),
          grid_names=("x",),
          compiler_params=plgpu.CompilerParams(profile_space=32, profile_dir=tmp_dir),
      )
      text = jax.jit(kernel).lower().as_text()
      np.testing.assert_array_equal(kernel(), jnp.ones((128,), jnp.float32))
    self.assertIn(
        "(tensor<2xi32>) -> (tensor<128xf32>, tensor<2xi32>, tensor<256xui32>)",
        text,
    )

  def test_global_semaphore(self):
    # We signal from block 0 and wait on block 1 to test whether the semaphore
    # is globally shared.
    def body(out_ref):
      sem_ref = pl.get_global(plgpu.SemaphoreType.REGULAR)
      block_id = lax.axis_index("x")
      @pl.when(block_id == 0)
      def _():
        pl.semaphore_signal(sem_ref)
      @pl.when(block_id == 1)
      def _():
        pl.semaphore_wait(sem_ref)
        out_ref[...] = jnp.ones_like(out_ref)

    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
        grid=(10,),
        grid_names=("x",),
    )
    result = kernel()
    np.testing.assert_array_equal(result, jnp.ones((128,), jnp.float32))

  def test_global_semaphore_with_multiple_threads(self):
    def body(out_ref):
      sem_ref = pl.get_global(plgpu.SemaphoreType.REGULAR)
      block_id = lax.axis_index("x")
      @pl.when(block_id == 0)
      def _():
        pl.semaphore_signal(sem_ref)
      @pl.when(block_id == 1)
      def _():
        pl.semaphore_wait(sem_ref)
        out_ref[...] = jnp.ones_like(out_ref)

    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
        grid=(10,),
        grid_names=("x",),
        thread_name="wg",
        num_threads=2,
    )
    result = kernel()
    np.testing.assert_array_equal(result, jnp.ones((128,), jnp.float32))

  @parameterized.parameters("sys", "gpu")
  def test_global_semaphore_with_multiple_warps(self, memory_scope):
    def body(out_ref):
      sem_ref = pl.get_global(plgpu.SemaphoreType.REGULAR)
      @plgpu.warp_map
      def _per_warp(warp_idx):
        @pl.when(warp_idx == 0)
        def _():
          if memory_scope == "sys":
            pl.semaphore_signal(sem_ref)
            plgpu.semaphore_signal(sem_ref, memory_scope="sys")
          else:
            plgpu.semaphore_signal(sem_ref, 2, memory_scope="gpu")

        @pl.when(warp_idx == 1)
        def _():
          plgpu.semaphore_wait(
              sem_ref, 2, memory_scope=memory_scope, decrement=True
          )

      out_ref[0] = pl.semaphore_read(sem_ref)

    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((1,), jnp.int32),
        compiler_params=plgpu.CompilerParams(),
    )
    result = kernel()
    np.testing.assert_array_equal(result, jnp.full((1,), 0, jnp.int32))

  def test_multiple_get_global_semaphores(self):
    def body(out_ref):
      sem1 = pl.get_global(plgpu.SemaphoreType.REGULAR)
      sem2 = pl.get_global(plgpu.SemaphoreType.REGULAR)
      block_id = lax.axis_index("x")
      @pl.when(block_id == 0)
      def _():
        pl.semaphore_signal(sem1, inc=5)
        pl.semaphore_signal(sem2, inc=10)
      @pl.when(block_id == 1)
      def _():
        pl.semaphore_wait(sem1, value=5, decrement=False)
        pl.semaphore_wait(sem2, value=10, decrement=False)
        val1 = pl.semaphore_read(sem1)
        val2 = pl.semaphore_read(sem2)
        out_ref[0] = val1
        out_ref[1] = val2

    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((2,), jnp.int32),
        grid=(10,),
        grid_names=("x",),
    )
    result = kernel()
    np.testing.assert_array_equal(result, jnp.array([5, 10], jnp.int32))

  def test_get_global_in_and_outside_control_flow(self):
    def body(out_ref):
      sem_before = pl.get_global(plgpu.SemaphoreType.REGULAR)
      block_id = lax.axis_index("x")

      @pl.when(block_id == 0)
      def _():
        sem_inside = pl.get_global(plgpu.SemaphoreType.REGULAR)
        pl.semaphore_signal(sem_inside, 7)
        pl.semaphore_signal(sem_before, 3)
        val_inside = pl.semaphore_read(sem_inside)
        out_ref[1] = val_inside

      sem_after = pl.get_global(plgpu.SemaphoreType.REGULAR)
      pl.semaphore_signal(sem_after, 11)
      val_before = pl.semaphore_read(sem_before)
      val_after = pl.semaphore_read(sem_after)
      out_ref[0] = val_before
      out_ref[2] = val_after

    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((3,), jnp.int32),
        grid=(1,),
        grid_names=("x",),
    )
    result = kernel()
    np.testing.assert_array_equal(result, jnp.array([3, 7, 11], jnp.int32))

  def test_multiple_semaphore_scopes(self):
    def body(out_ref):
      global_sem = pl.get_global(plgpu.SemaphoreType.REGULAR)

      @functools.partial(pl.run_scoped, block_sem=plgpu.SemaphoreType.REGULAR)
      def _scope2(block_sem):
        block_id = lax.axis_index("x")
        pl.semaphore_signal(block_sem)

        @pl.when(block_id == 0)
        def _():
          pl.semaphore_signal(global_sem)

        @pl.when(block_id == 1)
        def _():
          pl.semaphore_wait(global_sem)
          out_ref[...] = jnp.ones_like(out_ref)

        pl.semaphore_wait(block_sem)

    kernel = self.kernel(
        body,
        out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
        grid=(10,),
        grid_names=("x",),
    )
    result = kernel()
    np.testing.assert_array_equal(result, jnp.ones((128,), jnp.float32))

  def test_first_class_semaphore(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((2,), jnp.int32),
        grid=(2,),
        grid_names=("x",),
        num_threads=1,
        thread_name="wg",
    )
    def kernel(sem_ref, out_ref):
      self.assertTrue(jnp.issubdtype(sem_ref.dtype, pl.semaphore))
      block_id = lax.axis_index("x")
      init_val = pl.semaphore_read(sem_ref.at[0])
      @pl.when(block_id == 0)
      def _():
        pl.semaphore_signal(sem_ref.at[0], inc=3)
        pl.semaphore_signal(sem_ref.at[1], inc=7)
      @pl.when(block_id == 1)
      def _():
        pl.semaphore_wait(sem_ref.at[0], value=3, decrement=False)
        pl.semaphore_wait(sem_ref.at[1], value=7, decrement=False)
        out_ref[0] = pl.semaphore_read(sem_ref.at[0]) + init_val
        out_ref[1] = pl.semaphore_read(sem_ref.at[1])

    @jax.jit
    def run():
      sem_ref = plgpu.alloc_semaphore((2,))
      sem_aval = jax_core.typeof(sem_ref)
      self.assertTrue(jnp.issubdtype(sem_aval.dtype, pl.semaphore))
      return kernel(sem_ref)

    np.testing.assert_array_equal(run(), jnp.array([3, 7], jnp.int32))

  def test_kernel_rejects_semaphore_array_operand(self):
    @self.kernel(
        out_type=jax.ShapeDtypeStruct((2,), jnp.int32),
        grid=(1,),
        grid_names=("x",),
    )
    def kernel(sem_ref, out_ref):
      del sem_ref, out_ref

    @jax.jit
    def run():
      # Reading the ref yields a by-value semaphore array, which is rejected.
      return kernel(plgpu.alloc_semaphore((2,))[...])

    with self.assertRaisesRegex(ValueError, "Cannot pass semaphores into kernels"):
      run()

  def test_kernel_cannot_return_semaphore(self):
    with self.assertRaisesRegex(
        ValueError, "Kernels cannot return semaphores"
    ):

      @self.kernel(
          out_type=plgpu.SemaphoreType.REGULAR((2,)),
          grid=(1,),
          grid_names=("x",),
          num_threads=1,
          thread_name="wg",
      )
      def init_kernel(out_sem):
        del out_sem


class SemaphoreWGTest(
    SemaphoreTest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


class ExamplesSm90ATest(PallasSm90ATest):

  # WGMMA
  def test_stage6(self):
    m_block = n_block = 64
    k_block = 32
    x = jax.random.uniform(jax.random.key(42), shape=(128, 128), dtype=jnp.float16)

    @self.kernel(
        out_type=x,
        grid=(2, 2),
        grid_names=("m", "n"),
    )
    def kernel(l_ref, r_ref, o_ref):
      def compute(idxs, l_smem, r_smem, o_smem):
        def do_wgmma(acc_ref):
          plgpu.wgmma(acc_ref, l_smem, r_smem)
          return acc_ref[...]
        acc = pl.run_scoped(do_wgmma, plgpu.ACC((m_block, n_block), jnp.float16))
        def _store():
          o_smem[...] = acc
        def _acc():
          o_smem[...] += acc
        step, = idxs
        lax.cond(step == 0, _store, _acc)
      m = lax.axis_index("m")
      n = lax.axis_index("n")
      transforms = self.default_transforms(swizzle=64, dtype=jnp.float16)
      plgpu.emit_pipeline(
          compute,
          grid=(l_ref.shape[1] // k_block,),
          in_specs=[
              plgpu.BlockSpec(
                  (m_block, k_block), lambda k: (m, k), transforms=transforms
              ),
              plgpu.BlockSpec(
                  (k_block, n_block), lambda k: (k, n), transforms=transforms
              ),
          ],
          out_specs=[
              plgpu.BlockSpec(
                  (m_block, n_block), lambda k: (m, n), transforms=transforms
              )
          ],
      )(l_ref, r_ref, o_ref)

    np.testing.assert_allclose(
        kernel(x, x), jnp.matmul(x, x, precision="high"), rtol=2e-3
    )

  # TODO(apaszke): Clusters and multicast


class ExamplesSm90AWGTest(
    ExamplesSm90ATest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


class HelpersTest(PallasTest):

  @parameterized.product(
      m=[16],
      n=[7, 16],
      minor_dim=[0, 1],
      tile_width=[1, 3, 4],
      dynamic_shape=(False, True),
  )
  def test_planar_snake(self, m, n, minor_dim, tile_width, dynamic_shape):
    reference = np.full((m, n), -1)
    counter = itertools.count()
    minor_size, major_size = (m, n) if minor_dim == 0 else (n, m)
    for minor_tile in range(pl.cdiv(minor_size, tile_width)):
      for major in range(major_size):
        major = major if minor_tile % 2 == 0 else major_size - 1 - major
        minor_offset = minor_tile * tile_width
        for minor_in_tile in range(min(tile_width, minor_size - minor_offset)):
          minor = minor_offset + minor_in_tile
          idx = (minor, major) if minor_dim == 0 else (major, minor)
          reference[idx] = next(counter)
    results = np.full((m, n), -1)
    for lin in range(m * n):
      shape = (jnp.int32(m), jnp.int32(n)) if dynamic_shape else (m, n)
      results[
          plgpu.planar_snake(np.int32(lin), shape, minor_dim, tile_width)
      ] = lin
    np.testing.assert_array_equal(results, reference)

  def test_planar_snake_golden_with_partial_tile(self):
    m, n = 5, 5
    with self.subTest("minor_dim=0 tile_width=3"):
      results = np.full((m, n), -1)
      for lin in range(m * n):
        results[plgpu.planar_snake(np.int32(lin), (m, n), 0, 3)] = lin
      expected = np.array([
          [0, 3, 6, 9, 12],
          [1, 4, 7, 10, 13],
          [2, 5, 8, 11, 14],
          [23, 21, 19, 17, 15],
          [24, 22, 20, 18, 16]])
      np.testing.assert_array_equal(results, expected)
    with self.subTest("minor_dim=1 tile_width=3"):
      results = np.full((m, n), -1)
      for lin in range(m * n):
        results[plgpu.planar_snake(np.int32(lin), (m, n), 1, 3)] = lin
      expected = np.array([
          [0, 1, 2, 23, 24],
          [3, 4, 5, 21, 22],
          [6, 7, 8, 19, 20],
          [9, 10, 11, 17, 18],
          [12, 13, 14, 15, 16]])
      np.testing.assert_array_equal(results, expected)

  def test_planar_snake_golden(self):
    m, n = 8, 8
    with self.subTest("minor_dim=0 tile_width=2"):
      results = np.full((m, n), -1)
      for lin in range(m * n):
        results[plgpu.planar_snake(np.int32(lin), (m, n), 0, 2)] = lin
      expected = np.array([
          [0, 2, 4, 6, 8, 10, 12, 14],
          [1, 3, 5, 7, 9, 11, 13, 15],
          [30, 28, 26, 24, 22, 20, 18, 16],
          [31, 29, 27, 25, 23, 21, 19, 17],
          [32, 34, 36, 38, 40, 42, 44, 46],
          [33, 35, 37, 39, 41, 43, 45, 47],
          [62, 60, 58, 56, 54, 52, 50, 48],
          [63, 61, 59, 57, 55, 53, 51, 49],
      ])
      np.testing.assert_array_equal(results, expected)
    with self.subTest("minor_dim=1 tile_width=2"):
      results = np.full((m, n), -1)
      for lin in range(m * n):
        results[plgpu.planar_snake(np.int32(lin), (m, n), 1, 2)] = lin
      expected = np.array([
          [0, 1, 30, 31, 32, 33, 62, 63],
          [2, 3, 28, 29, 34, 35, 60, 61],
          [4, 5, 26, 27, 36, 37, 58, 59],
          [6, 7, 24, 25, 38, 39, 56, 57],
          [8, 9, 22, 23, 40, 41, 54, 55],
          [10, 11, 20, 21, 42, 43, 52, 53],
          [12, 13, 18, 19, 44, 45, 50, 51],
          [14, 15, 16, 17, 46, 47, 48, 49],
      ])
      np.testing.assert_array_equal(results, expected)
    with self.subTest("minor_dim=0 tile_width=1"):
      results = np.full((m, n), -1)
      for lin in range(m * n):
        results[plgpu.planar_snake(np.int32(lin), (m, n), 0, 1)] = lin
      expected = np.array([
          [0, 1, 2, 3, 4, 5, 6, 7],
          [15, 14, 13, 12, 11, 10, 9, 8],
          [16, 17, 18, 19, 20, 21, 22, 23],
          [31, 30, 29, 28, 27, 26, 25, 24],
          [32, 33, 34, 35, 36, 37, 38, 39],
          [47, 46, 45, 44, 43, 42, 41, 40],
          [48, 49, 50, 51, 52, 53, 54, 55],
          [63, 62, 61, 60, 59, 58, 57, 56],
      ])
      np.testing.assert_array_equal(results, expected)

  @parameterized.parameters(
      ((100,), ()),  # grid < SM count
      ((300,), ()),  # grid > SM count
      ((3, 3, 3, 3, 3), ()),  #  squashed grid dimensions
      ((50,), (2, 1)),  # small grid w/ cluster
      ((50, 4), (1, 2)),  # large grid w/ cluster
  )
  def test_dynamic_work_scheduling(self, grid, cluster):
    if not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("Only works on a GPU with capability >= sm100a")

    grid_names = tuple(str(i) for i in range(len(grid)))
    cluster_names = tuple("c"+str(i) for i in range(len(cluster)))
    def body(out_gmem, _):
      sm_idx = lax.axis_index(grid_names)
      cluster_idx = ()
      if cluster:
        cluster_idx = tuple(lax.axis_index(axis) for axis in cluster_names)
      @plgpu.dynamic_scheduling_loop(grid_names, cluster_axes=cluster_names)
      def loop_body(loop_info: plgpu.NDLoopInfo):
        out_gmem[*loop_info.index, *cluster_idx] = sm_idx
    out_shape = (*grid, *cluster)
    max_shared_memory = jax.local_devices()[0].shared_memory_per_block_optin
    # Mosaic GPU uses some shared memory implicitly, so we can't
    # explicitly request the full amount.
    large_amount_of_shared_memory = int(0.9 * max_shared_memory)
    with mgpu.core.artificial_shared_memory_limit(None):
      result = self.kernel(body,
                  out_type=jax.ShapeDtypeStruct(out_shape, jnp.int32),
                  grid=grid,
                  grid_names=grid_names,
                  cluster=cluster,
                  cluster_names=cluster_names,
                  # Allocate a large amount of SMEM to prevent multiple blocks
                  # being scheduled on the same SM.
                  scratch_types=[
                    plgpu.SMEM((large_amount_of_shared_memory,), jnp.int8)],
                  )()

    # Result maps grid_idx -> SM that performed the work.
    # Check that each SM had at least 1 block of work.
    cluster_size = int(np.prod(cluster))
    num_sms = min(jax.devices()[0].core_count // cluster_size, np.prod(grid))
    histogram = np.histogram(result, bins=range(num_sms+1))[0]
    self.assertEqual(np.sum(histogram), np.prod(out_shape))
    self.assertGreaterEqual(np.min(histogram), 1)
    # Make sure all blocks > num_sms were stolen.
    self.assertEqual(np.max(result), jnp.int32(num_sms) - 1)

  def test_dynamic_work_scheduling_with_carry(self):
    if not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("Only works on a GPU with capability >= sm100a")

    # In this test we make SM 0 run a the dynamic scheduling loop while all
    # other SMs spin. This means SM 0 should steal all of the work and we
    # keep track of the number of stolen blocks in the carry.
    blocks_to_steal = 100
    sm_count = jax.devices()[0].core_count
    def body(out_gmem, _):
      sm_idx = lax.axis_index("x")
      global_semaphore = pl.get_global(plgpu.SemaphoreType.REGULAR)

      @pl.when(sm_idx == 0)
      def _steal_loop():
        def loop_body(loop_info: plgpu.NDLoopInfo, carry: jax.Array):
          del loop_info
          return carry + jnp.int32(1)

        final_carry = plgpu.dynamic_scheduling_loop(
            ("x",), init_carry=jnp.int32(0)
        )(loop_body)
        out_gmem[0] = final_carry
        pl.semaphore_signal(global_semaphore, inc=sm_count)

      # All SMs wait until SM 0 has finished all blocks.
      pl.semaphore_wait(global_semaphore)

    max_shared_memory = jax.local_devices()[0].shared_memory_per_block_optin
    # Mosaic GPU uses some shared memory implicitly, so we can't
    # explicitly request the full amount.
    large_amount_of_shared_memory = int(0.9 * max_shared_memory)
    with mgpu.core.artificial_shared_memory_limit(None):
      result = self.kernel(body,
                  out_type=jax.ShapeDtypeStruct((1,), jnp.int32),
                  grid=(sm_count + blocks_to_steal,),
                  grid_names=("x",),
                  # Allocate a large amount of SMEM to prevent multiple blocks
                  # being scheduled on the same SM.
                  scratch_types=[
                    plgpu.SMEM((large_amount_of_shared_memory,), jnp.int8)],
                  )()
    self.assertEqual(result[0], blocks_to_steal + 1)


class HelpersWGTest(
    HelpersTest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


class DistributedTest(PallasTest):
  def test_lowering_with_explicit_sharding(self):
    mesh = jax.make_mesh(
      (1, 1),
      ("x", "y"),
      axis_types=(sharding.AxisType.Explicit, sharding.AxisType.Explicit))

    with jax.set_mesh(mesh):
      x = jax.random.normal(jax.random.key(0), shape=(64, 64), dtype=jnp.float32)
      x = sharding.reshard(x, sharding.PartitionSpec("x", "y"))
      @jax.jit
      def sharded_relu(x):
          x_spec = jax.typeof(x).sharding.spec
          @jax.shard_map(in_specs=x_spec, out_specs=x_spec)
          def shmap(x_shard):
              def relu_kernel(x_ref, o_ref, x_smem, barrier):
                  plgpu.copy_gmem_to_smem(x_ref, x_smem, barrier)
                  plgpu.barrier_wait(barrier)
                  x_value = x_smem[...]
                  x_smem[...] = jnp.where(x_value > 0, x_value, 0)
                  plgpu.commit_smem()
                  plgpu.copy_smem_to_gmem(x_smem, o_ref)
                  plgpu.wait_smem_to_gmem(0)
              out_shape = jax.ShapeDtypeStruct(shape=(64, 64), dtype=jnp.float32)
              result = self.kernel(relu_kernel,
                  out_type=out_shape,
                  scratch_types=dict(
                      x_smem=plgpu.SMEM((64, 64), jnp.float32),
                      barrier=plgpu.Barrier()
                  )
              )(x_shard)
              return result
          return shmap(x)
      result = sharded_relu(x)
    np.testing.assert_array_equal(result, jnp.maximum(x, 0))


class DistributedWGTest(
    DistributedTest, lowering_semantics=plgpu.LoweringSemantics.Warpgroup
):
  ...


@jtu.with_config(jax_pallas_poison_buffers=True)
class MosaicGPUPoisonTest(PallasTest, jtu.CudaArchSpecificTest):

  def test_poison_buffers_nyi(self):
    def kernel(o_ref):
      scratch = jax.empty_ref(
          jax.ShapeDtypeStruct((128,), jnp.float32), memory_space=plgpu.SMEM
      )
      o_ref[...] = scratch[...]

    with self.assertRaisesRegex(
        NotImplementedError, "Buffer poisoning is not supported on GPU yet"
    ):
      self.pallas_call(
          kernel,
          out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
      )()


if __name__ == "__main__":
  absltest.main(testLoader=jtu.JaxTestLoader())