# 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.

"""Contains GPU-specific Pallas abstractions."""

from __future__ import annotations

import collections
from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence
import contextlib
import dataclasses
import enum
import functools
import itertools as it
import math
from typing import Any, ClassVar, Literal

import jax
from jax._src import api
from jax._src import config
from jax._src import core as jax_core
from jax._src import custom_batching
from jax._src import deprecations
from jax._src import dtypes
from jax._src import effects
from jax._src import frozen_dict
from jax._src import lax
from jax._src import pretty_printer as pp
from jax._src import state
from jax._src import tree_util
from jax._src import util
from jax._src.pallas import core as pallas_core
from jax._src.pallas import helpers as pallas_helpers
from jax._src.pallas import primitives as pallas_primitives
from jax._src.pallas import utils as pallas_utils
from jax._src.state import indexing
from jax._src.state import types as state_types
import jax.experimental.mosaic.gpu as mgpu
from jax.experimental.mosaic.gpu import tcgen05
from jax.experimental.mosaic.gpu import utils as mgpu_utils
from jax.experimental.mosaic.gpu.launch_context import OOBFillMode
import jax.numpy as jnp
from jaxlib.mlir import ir


_Ref = state.AbstractRef | state_types.TransformedRef

DimensionSemantics = Literal["parallel", "sequential"]
TransposeTransform = state_types.TransposeTransform

ScratchShapeTree = pallas_core.ScratchShapeTree

# We align all our SMEM allocations to 1024 bytes. TMA and WGMMA are very
# sensitive to alignment and while this is quite conservative, it gets the job
# done. We should make this more refined in the future.
SMEM_ALIGNMENT = 1024
TMEM_COL_ALIGNMENT = 4


def is_trivial_index(idx, shape) -> bool:
  """Checks if the index selects the entire shape."""

  # Slices that select the entire dimension.
  def _slices(d):
    slices = [slice(b, e, s) for b, e, s in it.product([0, None], [d, None], [1, None])]
    return [indexing.Slice(0, d, 1), *slices]

  if isinstance(idx, tuple):
    return all(i in _slices(d) for d, i in zip(shape, idx))

  return idx is ... or (len(shape) == 1 and idx in _slices(shape[0]))


class TraceScope(enum.Enum):
  WARP = "warp"
  WARPGROUP = "warpgroup"


@dataclasses.dataclass(frozen=True, kw_only=True)
class CompilerParams:
  """Mosaic GPU compiler parameters.

  Attributes:
    approx_math: If True, the compiler is allowed to use approximate
      implementations of some math operations, e.g. ``exp``. Defaults to False.
    dimension_semantics: A list of dimension semantics for each grid dimension
      of the kernel. Either "parallel" for dimensions that can execute in any
      order, or "sequential" for dimensions that must be executed sequentially.
    max_concurrent_steps: The maximum number of sequential stages that are
      active concurrently. Defaults to 1.
    delay_release: The number of steps to wait before reusing the input/output
      references. Defaults to 0, and must be strictly smaller than
      max_concurrent_steps. Generally, you'll want to set it to 1 if you don't
      await the WGMMA in the body.
    unsafe_no_auto_barriers: If True, Pallas will never automatically insert
      barrier instructions that ensure synchronous semantics of loads and
      stores. At the moment, the insertion is done conservatively and might
      regress performance. There are (at least) two conditions that must be
      satisfied for the use of this flag to be safe. First, no memory region is
      ever read *and* written to by the same thread (async copies are performed
      by background threads and do not count towards this rule). Secondly, no
      thread ever calls commit_smem(), reads from the committed SMEM and then
      issues an async copy overwriting that region (this is a very artificial
      and highly unlikely scenario).
    reduction_scratch_bytes: The number of shared memory bytes to reserve as
      scratch space for cross-warp reductions. The higher this value, the more
      registers can be reduced in parallel. 2 * 128 * 6 * 4 = 6144 bytes is
      typically a good value in order to extract most of the potential gains on
      H100 and B200.
    profile_space: The number of profiler events that can be collected in a
      single invocation. It is undefined behavior if a thread collects more
      events than this, unless profile_bounds_check is set.
    profile_dir: The directory to which profiling traces will be written to.
    profile_trace_scope: The scope at which traces are collected (WARP or
      WARPGROUP).
    profile_bounds_check: If True, profiler events past profile_space are
      dropped (the trace is truncated) instead of corrupting SMEM, at the cost
      of a slightly higher per-event profiling overhead.
    skip_device_barrier: If True, skips the cross-device barrier before kernel
      launch. Improper use of this flag can lead to race conditions. !!!Use with
      caution!!! Defaults to False.
    is_multi_process: If True, indicates that the kernel runs in a multi-process
      environment and performs cross-process device communication. This is
      mostly needed when compiling distributed-over-processes code with an
      AbstractMesh, as with a concrete Mesh in standard JIT compilation,
      multi-process execution is detected automatically from the mesh devices.
      Defaults to False.
  """
  approx_math: bool = False
  dimension_semantics: Sequence[DimensionSemantics] | None = None
  max_concurrent_steps: int = 1
  unsafe_no_auto_barriers: bool = False
  reduction_scratch_bytes: int = 128 * 4 * 4
  profile_space: int = 0
  profile_dir: str = ""
  profile_trace_scope: TraceScope = TraceScope.WARPGROUP
  profile_bounds_check: bool = False
  lowering_semantics: mgpu.core.LoweringSemantics = mgpu.core.LoweringSemantics.Warpgroup
  skip_device_barrier: bool = False
  is_multi_process: bool = False

  def __post_init__(self):
    if self.dimension_semantics is not None:
      object.__setattr__(
          self, "dimension_semantics", tuple(self.dimension_semantics)
      )
    if bool(self.profile_space) ^ bool(self.profile_dir):
      raise ValueError(
          "Either both profile_space and profile_dir must be set, or neither."
      )

  replace = dataclasses.replace


class MemorySpace(enum.Enum):
  #: Global memory.
  GMEM = "gmem"
  #: Shared memory.
  SMEM = "smem"
  #: Tensor memory. New addition to Blackwell. Not available on Hopper.
  TMEM = "tmem"
  #: Registers.
  REGS = "regs"

  def __str__(self) -> str:
    return self.value

  def __call__(
      self,
      shape: Sequence[int],
      dtype: jax.typing.DTypeLike,
      *,
      transforms: Sequence[state_types.Transform] = (),
      packed: bool | None = None,
      collective: bool | None = None,
      layout: TMEMLayout | None = None,
  ) -> pallas_core.MemoryRef:
    shape = tuple(shape)
    # TODO(sharadmv): Add HiType constructor support.
    if self == MemorySpace.TMEM:
      if transforms:
        raise ValueError("transforms are not supported for TMEM")
      if collective is None:
        collective = False
      if len(shape) > 2:
        transforms = (CollapseLeadingBatchDimensionsTransform(),)
      if layout is None:
        if packed is None:
          if dtypes.itemsize_bits(dtype) != 32:
            raise ValueError(
                "dtypes narrower than 32-bit require either the packed argument"
                " or an explicit TMEM layout"
            )
          packed = False
        # Ignore batch dimensions for layout inference.
        mgpu_layout = infer_tmem_layout(
            shape[-2:], dtype, packed=packed, collective=collective
        )
      else:
        if packed is not None:
          raise ValueError("packed cannot be specified if layout is specified.")
        mgpu_layout = layout.to_mgpu()
    else:
      if packed is not None or collective is not None or layout is not None:
        raise ValueError("packed, collective and layout arguments are only supported for TMEM.")
      mgpu_layout = None
    return GPUMemoryRef(jax_core.ShapedArray(shape, dtype), memory_space=self,
                        transforms=transforms, layout=mgpu_layout,
                        collective=collective)

  def like(self, shape_dtype_like):
    return self(shape_dtype_like.shape, shape_dtype_like.dtype)


class SemaphoreType(enum.Enum):
  REGULAR = "regular"
  BARRIER = "barrier"

  def __call__(self, shape: tuple[int, ...]):
    dtype: Any
    if self == SemaphoreType.BARRIER:
      dtype = pallas_core.BarrierSemaphore()
    else:
      dtype = pallas_core.Semaphore()
    return pallas_core.MemoryRef(jax_core.ShapedArray(shape, dtype),
                                 MemorySpace.GMEM)

  def get_array_aval(self) -> jax_core.ShapedArray:
    return self(()).get_array_aval()

  def get_ref_aval(self) -> _Ref:
    return self(()).get_ref_aval()

def _check_no_semaphore_outputs(out_type) -> None:
  """Raises if any kernel output is a semaphore or a semaphore Ref.

  Semaphores may only be created by ``alloc_semaphore`` outside of a kernel and
  passed into it as operands.
  """
  for leaf in tree_util.tree_leaves(out_type):
    if isinstance(leaf, SemaphoreType):
      dtype = leaf.get_array_aval().dtype
    else:
      dtype = getattr(leaf, "dtype", None)

    if dtype is not None and (
        isinstance(dtype, pallas_core.AbstractSemaphoreTy)
        or dtypes.issubdtype(dtype, pallas_core.semaphore_dtype)
    ):
      raise ValueError(
          "Kernels cannot return semaphores. Allocate them as scratch if only"
          " used inside this kernel, or via plgpu.alloc_semaphore before the"
          " kernels that use them."
      )


def alloc_semaphore(shape: tuple[int, ...] = ()) -> jax.Ref:
  """Returns a zero-initialized GMEM semaphore Ref.

  This function can only be used outside of a kernel; for semaphores scoped to
  a single kernel, use ``scratch_types`` instead. It is the only way to create
  semaphores that outlive a single kernel invocation. It performs *no*
  cross-device synchronization: the caller is responsible for the initial
  sync, e.g. by running a kernel that is data dependent on the returned refs
  with ``skip_device_barrier=False``.

  Example::

    sem = plgpu.alloc_semaphore((num_devices,))
    # The first kernel consuming ``sem`` must not skip the device barrier, so
    # that all devices observe the zero-initialized semaphores before use.
    out = plgpu.kernel(body, out_type=out_type, grid=grid, grid_names=names)(
        sem, x
    )

  Args:
    shape: The shape of the semaphore array.
  """

  zeros = lax.convert_element_type(
      jnp.zeros(tuple(shape), jnp.int32), pallas_core.Semaphore()
  )
  return jax_core.new_ref(zeros, memory_space=MemorySpace.GMEM)

class PrimitiveSemantics(enum.Enum):
  """Thread semantics for a primitives at the Pallas user-level."""

  Warp = enum.auto()
  Warpgroup = enum.auto()


# Convenience constants for (lowering, primitive) thread semantics pairs.
LANExWG_SEMANTICS = (
    mgpu.LoweringSemantics.Lane, PrimitiveSemantics.Warpgroup)
LANExWARP_SEMANTICS = (
    mgpu.LoweringSemantics.Lane, PrimitiveSemantics.Warp)
WGxWG_SEMANTICS = (
    mgpu.LoweringSemantics.Warpgroup, PrimitiveSemantics.Warpgroup)
WGxWARP_SEMANTICS = (
    mgpu.LoweringSemantics.Warpgroup, PrimitiveSemantics.Warp)


def kernel(
    body: Callable[..., None] | api.NotSpecified = api.NotSpecified(),
    *,
    out_type: object = (),
    scratch_types: ScratchShapeTree = (),
    compiler_params: pallas_core.CompilerParams | None = None,
    cost_estimate: pallas_core.CostEstimate | None = None,
    # Mesh kwargs
    grid: tuple[int, ...] = (),
    grid_names: tuple[str, ...] = (),
    cluster: tuple[int, ...] = (),
    cluster_names: tuple[str, ...] = (),
    num_threads: int | None = None,
    thread_name: str | None = None,
    interpret: Any = None,
    debug: bool = False,
    **mesh_kwargs: Any,
) -> Any:
  r"""Entry point for defining a Mosaic GPU kernel.

  Args:
    body: The kernel body, which should take as arguments the input, output, and
      scratch Refs. The number of input Refs is determined by the number of
      arguments passed into kernel returned by this function. The number of
      output and scratch Refs are determined by `out_shape` and `scratch_shapes`
      respectively.
    out_type: The type of the output. Should be a PyTree of
      ``jax.ShapeDtypeStruct`` or JAX types.
    scratch_types: The types of the scratch ``Ref``\s to allocate. Should be a
      PyTree of ``jax.ShapeDtypeStruct`` or JAX types.
    compiler_params: Additional compiler options. See the `CompilerParams`
      dataclass for more details.
    cost_estimate: The cost estimate of the function.
    grid: A tuple of integers specifying the size of the kernel grid.
    grid_names: The axis names of the grid. Must be the same length as `grid`.
    cluster: A tuple of integers specifying the size of the kernel cluster.
    cluster_names: The axis names of the grid. Must be the same length as
      `cluster`.
    num_threads: The number of threads to launch per block. Note that these do
      not correspond to CUDA threads, but rather to warpgroups on Hopper and
      Blackwell GPUs.
    thread_name: The axis name used to query the thread index.
    debug: Whether or not to output helpful debugging information.
    **mesh_kwargs: Additional mesh kwargs. See `Mesh` for more details.

  Returns:
    If ``body`` is provided, returns a function that runs the kernel. It should
    take any number of input operands and returns an output with the same PyTree
    structure as ``out_shape``.

    If ``body`` is omitted, returns a decorator that can be used to annotate
    a kernel body.
  """
  if isinstance(body, api.NotSpecified):
    return lambda fun: kernel(
        fun,
        out_type=out_type,
        scratch_types=scratch_types,
        compiler_params=compiler_params,
        cost_estimate=cost_estimate,
        grid=grid,
        grid_names=grid_names,
        cluster=cluster,
        cluster_names=cluster_names,
        num_threads=num_threads,
        thread_name=thread_name,
        interpret=interpret,
        debug=debug,
        **mesh_kwargs,
    )

  if unwrap_out := not isinstance(out_type, (tuple, list)):
    out_type = (out_type,)

  mesh = Mesh(
      grid=grid,
      grid_names=grid_names,
      cluster=cluster,
      cluster_names=cluster_names,
      num_threads=num_threads,
      thread_name=thread_name,
      **mesh_kwargs,
  )

  _check_no_semaphore_outputs(out_type)

  # TODO(slebedev): Use mesh-specific batching rules in ``mpmd_map`` instead.
  @custom_batching.custom_vmap
  def wrapper(*operands):
    for op in tree_util.tree_leaves(operands):
      if not isinstance(op, state.AbstractRef):
        aval = jax_core.typeof(op)
        if isinstance(aval, jax_core.ShapedArray) and (
            dtypes.issubdtype(aval.dtype, pallas_core.semaphore_dtype)
            or isinstance(aval.dtype, pallas_core.AbstractSemaphoreTy)
        ):
          raise ValueError(
              "Cannot pass semaphores into kernels as arrays. Wrap them in"
              f" jax.new_ref: {op}"
          )
    thread_name = mesh.thread_name if mesh.thread_name is not None else ()

    def kernel_body(*refs):
      # NOTE: We cannot use the ``scratch_types=`` argument of ``pl.kernel``
      # for these, because some scratch types return ``TransformedRef``s in
      # ``get_ref_aval``, which is not yet supported by ``mpmd_map``.
      pallas_primitives.run_scoped(
          functools.partial(body, *refs),
          *scratch_types if isinstance(scratch_types, Sequence) else (),
          collective_axes=thread_name,
          **scratch_types if isinstance(scratch_types, Mapping) else {},
      )

    name = (
        getattr(body, "__name__", "anonymous")
        if mesh.kernel_name is None
        else mesh.kernel_name
    )
    # TODO(slebedev): This is only here for backward compatibility. Remove.
    with config._check_vma(False):
      outs = pallas_helpers.kernel(
          kernel_body,
          out_type=out_type,
          mesh=mesh,
          compiler_params=compiler_params,
          cost_estimate=cost_estimate,
          interpret=interpret,
          name=name,
          debug=debug,
      )(*operands)
    return outs[0] if unwrap_out else outs

  @wrapper.def_vmap
  def _vmap_rule(axis_size, in_batched, *args):
    axis_name = object()

    def batched_body(*refs, **scratch_ref_kwargs):
      idx = lax.axis_index(axis_name)
      lens = (len(args), len(out_type))
      operand_refs, out_refs, scratch_refs = util.split_list(refs, lens)
      slice_ref = lambda r, b=True: (r.at[idx] if b else r)
      operand_refs = tree_util.tree_map(slice_ref, operand_refs, in_batched)
      out_refs = tree_util.tree_map(slice_ref, out_refs)
      return body(*operand_refs, *out_refs, *scratch_refs, **scratch_ref_kwargs)

    out_type_ = out_type[0] if unwrap_out else out_type
    add_batch_dim = lambda x: x.update(shape=(axis_size, *x.shape))
    mesh_kwargs_ = dict(mesh_kwargs)
    batched_cost_estimate = None
    if cost_estimate is not None and jax_core.is_constant_dim(axis_size):
      batched_cost_estimate = dataclasses.replace(
          cost_estimate,
          flops=cost_estimate.flops * axis_size,
          bytes_accessed=cost_estimate.bytes_accessed * axis_size,
          transcendentals=cost_estimate.transcendentals * axis_size,
          remote_bytes_transferred=(
              cost_estimate.remote_bytes_transferred * axis_size
          ),
      )
    out = kernel(
        batched_body,
        out_type=tree_util.tree_map(add_batch_dim, out_type_),
        scratch_types=scratch_types,
        compiler_params=compiler_params,
        cost_estimate=batched_cost_estimate,
        grid=(axis_size,) + grid,
        grid_names=(axis_name,) + grid_names,  # pyrefly: ignore[bad-argument-type]
        cluster=cluster,
        cluster_names=cluster_names,
        num_threads=num_threads,
        thread_name=thread_name,
        interpret=interpret,
        debug=debug,
        **mesh_kwargs_,
    )(*args)
    out_batched = tree_util.tree_map(lambda _: True, out_type_)
    return out, out_batched

  return wrapper


@dataclasses.dataclass(frozen=True)
class GPUMemoryRef(pallas_core.MemoryRef):
  transforms: Sequence[state_types.Transform] = ()

  layout: tcgen05.TMEMLayout | None = dataclasses.field(default=None, kw_only=True)
  collective: bool | None = dataclasses.field(default=None, kw_only=True)

  def __post_init__(self):
    is_tmem = self.memory_space == MemorySpace.TMEM
    assert (self.layout is not None) == is_tmem
    assert (self.collective is not None) == is_tmem

  def get_ref_aval(self) -> _Ref:
    aval: Any = jax_core.ShapedArray(self.shape, self.dtype)
    physical_aval = state_types.transform_type(self.transforms, aval)
    if self.memory_space == MemorySpace.TMEM:
      aval = AbstractTMEMRef(
          aval, self.memory_space, self.layout, self.collective
      )
      physical_ref_aval = AbstractTMEMRef(
          physical_aval, self.memory_space, self.layout, self.collective
      )
    else:
      aval = state.AbstractRef(aval, memory_space=self.memory_space)
      physical_ref_aval = state.AbstractRef(physical_aval, memory_space=self.memory_space)
    transforms: list[state_types.Transform] = pallas_core.undo_transforms(
        aval, self.transforms
    )
    ref = state_types.TransformedRef(physical_ref_aval, tuple(transforms))
    if not ref.transforms:
      return ref.ref
    return ref


def align_to(x: int, alignment: int):
  if rem := x % alignment:
    return x + alignment - rem
  return x


# A tree of `GPUMemoryRef`s.
_GPUMemoryRefTree = Any


def _ref_group_size(refs: _GPUMemoryRefTree) -> int:
  size = 0
  for ref in jax.tree.leaves(refs):
    # Make sure that the start of each ref is aligned with `SMEM_ALIGNMENT`.
    size = align_to(size, SMEM_ALIGNMENT)
    if jnp.issubdtype(ref.dtype, jnp.integer):
      nbits = jnp.iinfo(ref.dtype).bits
    elif jnp.issubdtype(ref.dtype, jnp.floating):
      nbits = jnp.finfo(ref.dtype).bits
    else:
      raise NotImplementedError(f"Unsupported dtype: {ref.dtype}")
    ref_bits = math.prod(ref.shape) * nbits
    if ref_bits % 8:
      raise ValueError(
          "Only byte-aligned shapes are supported. Got shape:"
          f" {ref.dtype}{ref.shape}"
      )
    size += ref_bits // 8
  return size


def _ref_group_tmem_col_size(refs: _GPUMemoryRefTree) -> int:
  """Returns the total number of TMEM columns used by a group of aliased Refs.
  """
  ncols = 0
  for ref in jax.tree.leaves(refs):
    # Refs with leading batch dimensions are collapsed to 2D in TMEM (see
    # `CollapseLeadingBatchDimensionsTransform`), so we compute the column count
    # on the collapsed shape.
    shape = ref.shape
    if len(shape) > 2:
      shape = (shape[-2], math.prod(shape[:-2]) * shape[-1])
    ref_ncols = ref.layout.cols_in_shape(shape,
                                         dtypes.itemsize_bits(ref.dtype))
    ncols += align_to(ref_ncols, TMEM_COL_ALIGNMENT)
  return ncols


def infer_tmem_layout(
    shape: tuple[int, ...],
    dtype: jax.typing.DTypeLike,
    *,
    packed: bool,
    collective: bool) -> tcgen05.TMEMLayout:
  """Infers the number of columns used and layout for allocating TMEM Refs."""
  if packed:
    packing = 32 // dtypes.itemsize_bits(dtype)
  else:
    packing = 1
  return tcgen05._infer_tmem_layout(shape, collective=collective, packing=packing)


def flatten_ref_union(ref_union: AbstractRefUnion) -> tuple[_Ref, ...]:
  """Flattens a union of trees of references into a tuple of references.

  This is the moral equivalent of `jax.tree.leaves` for aliased references.
  """
  flat_refs = []
  ref_union_mem_space = (ref_union.memory_space
                         if isinstance(ref_union, AbstractRefUnion) else
                         jax_core.typeof(ref_union).memory_space)
  if ref_union_mem_space == SMEM:
    union_bytes = 0
    for group_idx, ref_group in enumerate(ref_union.refs):
      byte_offset = 0
      def unflatten(ref):
        nonlocal byte_offset
        byte_offset = align_to(byte_offset, SMEM_ALIGNMENT)
        assert isinstance(ref, state.AbstractRef) or isinstance(
            ref, pallas_core.TransformedRef
        )
        if not isinstance(ref, pallas_core.TransformedRef):
          ref = pallas_core.TransformedRef(ref, transforms=())
        transform = ExtractAliasedRef.from_transformed_ref(ref, byte_offset, group_idx)
        result = pallas_core.TransformedRef(
            ref_union, transforms=(transform, *ref.transforms)
        )
        if jnp.issubdtype(ref.dtype, jnp.integer):
          nbits = jnp.iinfo(ref.dtype).bits
        elif jnp.issubdtype(ref.dtype, jnp.floating):
          nbits = jnp.finfo(ref.dtype).bits
        else:
          raise NotImplementedError(f"Unsupported dtype: {ref.dtype}")
        ref_bits = math.prod(ref.shape) * nbits
        if ref_bits % 8:
          raise ValueError(
              "Only byte-aligned shapes are supported. Got shape:"
              f" {ref.dtype}{ref.shape}"
          )
        byte_offset += ref_bits // 8
        return result
      flat_refs.append(jax.tree.map(unflatten, ref_group))
      union_bytes = max(union_bytes, byte_offset)
    assert union_bytes == ref_union.shape[0]
  elif ref_union_mem_space == TMEM:
    union_cols = 0
    for group_idx, ref_group in enumerate(ref_union.refs):
      col_offset = 0
      def unflatten(ref):
        nonlocal col_offset
        col_offset = align_to(col_offset, TMEM_COL_ALIGNMENT)
        if not isinstance(ref, pallas_core.TransformedRef):
          ref = pallas_core.TransformedRef(ref, transforms=())
        # `ref.ref.shape` is the physical (collapsed to 2D) TMEM shape, whereas
        # `ref.shape` may carry leading batch dimensions.
        ncols = ref.layout.cols_in_shape(ref.ref.shape,
                                         dtypes.itemsize_bits(ref.dtype))
        transform = ExtractAliasedRef.from_transformed_ref(
            ref, col_offset, group_idx, layout=ref.layout)
        result = pallas_core.TransformedRef(
            ref_union, transforms=(transform, *ref.transforms)
        )
        col_offset += ncols
        return result
      flat_refs.append(jax.tree.map(unflatten, ref_group))
      union_cols = max(union_cols, col_offset)
    assert union_cols == ref_union.shape[1], (union_cols, ref_union.shape[1])
  else:
    raise NotImplementedError("Only SMEM and TMEM refs are supported.")
  return tuple(flat_refs)


class AbstractRefUnion(state.AbstractRef):
  refs: Sequence[_GPUMemoryRefTree]

  def __init__(
      self,
      aval,
      refs: Sequence[_GPUMemoryRefTree],
      memory_space,
  ):
    self.refs = refs
    super().__init__(aval, memory_space=memory_space)

  def _iter(self, tracer):
    return iter(flatten_ref_union(tracer))

  def _getitem(self, tracer, idx):
    return list(iter(tracer))[idx]

  def _setitem(self, tracer, idx, value):
    del tracer, idx, value  # Unused.
    raise ValueError("Ref unions can't be assigned to.")

  def update(self, inner_aval=None, memory_space=None, kind=None):
    ref = super().update(inner_aval, memory_space, kind)
    return AbstractRefUnion(ref.inner_aval, self.refs, self.memory_space)

  @functools.cached_property
  def layout(self) -> tcgen05.TMEMLayout:
    if self.memory_space != TMEM:
      raise ValueError("layout attribute is only defined for TMEM refs")
    return tcgen05.tmem_default_layout(packing=1)

  @functools.cached_property
  def collective(self) -> bool:
    if self.memory_space != TMEM:
      raise ValueError("collective attribute is only defined for TMEM refs")
    ref_leaves = jax.tree.leaves(self.refs)
    first_ref = ref_leaves[0]
    assert all(ref.collective == first_ref.collective for ref in ref_leaves)
    return first_ref.collective

  def __eq__(self, other):
    return (
        type(self) is type(other)
        and self.inner_aval == other.inner_aval
        and self.memory_space == other.memory_space
        and self.refs == other.refs
    )

  def __hash__(self):
    # `flatten_ref_union(self)` creates `TransformedRef`s that refer to `self`,
    # so we extract the transforms from the `TransformedRef`s in order to avoid
    # infinite recursion.
    all_transforms = tuple(
        ref.transforms if isinstance(ref, pallas_core.TransformedRef) else ()
        for ref in flatten_ref_union(self)
    )
    return hash((self.inner_aval, self.memory_space, all_transforms))


@dataclasses.dataclass(init=False, frozen=True)
class RefUnion(GPUMemoryRef):
  """A sequence of trees of refs that are allowed to reuse the same memory.

  One should not make assumptions as to how each ref will map to the underlying
  memory region, since arbitrary padding may be applied in between different
  refs.

  As such, ref unions are only safe to use when the groups of refs that we
  intend to alias have disjoint lifetimes (i.e. one should never attempt to read
  data using a different ref than the one that was used to write the data).
  """
  refs: Sequence[_GPUMemoryRefTree] = ()

  def __init__(self, *refs: _GPUMemoryRefTree):
    ref_leaves = jax.tree.leaves(refs)
    if all(ref.memory_space == SMEM for ref in ref_leaves):
      object.__setattr__(self, "refs", refs)
      num_bytes = max(map(_ref_group_size, self.refs))
      super().__init__(
          inner_aval=jax_core.ShapedArray(
              (num_bytes,), jnp.int8
          ),
          memory_space=SMEM,
          transforms=(),
      )
    elif all(ref.memory_space == TMEM for ref in ref_leaves):
      object.__setattr__(self, "refs", refs)
      max_cols = max(map(_ref_group_tmem_col_size, self.refs))
      is_collective = ref_leaves[0].collective
      if any(r.collective != is_collective for r in ref_leaves):
        raise ValueError(
            "Some aliased TMEM references are collective and some are not."
        )
      super().__init__(
          inner_aval=jax_core.ShapedArray(
              shape=(128, max_cols,),
              dtype=jnp.int32,
          ),
          memory_space=TMEM,
          transforms=(),
          layout=tcgen05.tmem_default_layout(packing=1),
          collective=all(ref.collective for ref in ref_leaves),
      )
    else:
      raise NotImplementedError(
          "All aliased Refs must have the same memory space (SMEM or TMEM). "
          f"Got {(ref.memory_space for ref in ref_leaves)}.")

  def get_ref_aval(self) -> AbstractRefUnion:
    inner_aval = jax.core.ShapedArray(self.shape, self.dtype)
    refs_aval = jax.tree.map(lambda ref: ref.get_ref_aval(), self.refs)
    return AbstractRefUnion(inner_aval, refs_aval,
                            memory_space=self.memory_space)


Index = mgpu.DynamicSlice | slice | int | ir.Value


@dataclasses.dataclass(frozen=True)
class TilingTransform(state_types.Transform):
  """Represents a tiling transformation for memory refs.

  A tiling of (X, Y) on an array of shape (M, N) will result in a transformed
  shape of (M // X, N // Y, X, Y). Ex. A (256, 256) block that is tiled with a
  tiling of (64, 32) will be tiled as (4, 8, 64, 32).
  """
  tiling: tuple[int, ...]

  def transform_type(self, x):
    match x:
      case jax_core.ShapedArray():
        shape = x.shape
        if shape is None:
          return x
        leading_dims = shape[: -len(self.tiling) :]
        tiled_dims = shape[-len(self.tiling) :]
        assert all(d % t == 0 for d, t in zip(tiled_dims, self.tiling))
        num_tiles = [d // t for d, t in zip(tiled_dims, self.tiling)]
        new_shape = (*leading_dims, *num_tiles, *self.tiling)
        return x.update(shape=new_shape)
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case _:
        raise TypeError(f"Cannot transform type: {x}")

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    return UntilingTransform(self.tiling)

@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class UntilingTransform(state_types.Transform):
  tiling: tuple[int, ...] = jax.tree.static()

  def transform_type(self, x):
    match x:
      case jax_core.ShapedArray():
        shape = x.shape
        if shape is None:
          return x
        assert shape[-len(self.tiling) :] == self.tiling, (shape, self.tiling)
        shape = shape[: -len(self.tiling)]  # Drop tiling
        new_shape = shape[: -len(self.tiling)] + tuple(
            block_dim * tiling_dim
            for block_dim, tiling_dim in zip(
                shape[-len(self.tiling) :], self.tiling
            )
        )
        return x.update(shape=new_shape)
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case _:
        raise TypeError(f"Cannot transform type: {x}")

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    return TilingTransform(self.tiling)

  def commute_transpose(
      self, _: jax_core.AbstractValue,
      transpose: state_types.TransposeTransform,
  ) -> tuple[state_types.TransposeTransform, UntilingTransform]:
    # The transpose in question is applied to the untiled ref so we
    # need to translate it by duplicating and offsetting the last part.
    perm = transpose.permutation
    off = len(perm)
    new_suffix = [i + off for i in perm[-len(self.tiling) :]]
    if set(new_suffix) != set(range(off, off + len(self.tiling))):
      raise ValueError(
          "Transpose cannot be moved before a tiling transform when it changes"
          f" the set of tiled dimensions. (permutation: {perm}, tiling:"
          f" {self.tiling})"
      )

    new_tiling = tuple(self.tiling[i - off] for i in new_suffix)
    new_transpose = state_types.TransposeTransform((*perm, *new_suffix))
    return new_transpose, dataclasses.replace(self, tiling=new_tiling)

  def commute_ndindexer(
      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, UntilingTransform]:
    del aval
    idxs = indexer.indices
    indexer_shape = indexer.shape
    untiled_idxs = idxs[: -len(self.tiling)]
    tiled_idxs = idxs[-len(self.tiling) :]
    idxs_after_tiling: list[indexing.Slice] = []
    leading_shape, untiled_shape = (
        indexer_shape[: -len(self.tiling)],
        indexer_shape[-len(self.tiling) :],
    )
    for idx, tile, dim in zip(tiled_idxs, self.tiling, untiled_shape):
      match idx:
        case slice() | indexing.Slice():
          if isinstance(idx, slice):
            ds = indexing.Slice.from_slice(idx, dim)
          else:
            ds = idx
          if ds.stride is not None and ds.stride != 1:
            raise NotImplementedError(
                f"Strided slices unsupported. Got stride: {ds.stride}"
            )
          start, size = ds.start, ds.size
          if (
              start is not None and isinstance(start, int) and start % tile
          ) or (size is not None and isinstance(size, int) and size % tile):
            raise ValueError(
                f"Expected slice start ({start}) and slice size ({size})"
                f" to be divisible by the tile size ({tile})"
            )
          def _maybe_cdiv_with_cast(x, y):
            if x is None:
              return None
            if isinstance(x, jax.Array):
              # If x is an int32, we need to make sure y is an int32 to avoid
              # a dtype mismatch.
              y = jnp.array(y, x.dtype)
            return pallas_utils.cdiv(x, y)
          new_start = _maybe_cdiv_with_cast(start, tile)
          new_size = _maybe_cdiv_with_cast(size, tile)
          idxs_after_tiling.append(indexing.Slice(new_start, new_size))
        case _:
          raise TypeError(f"Unsupported index type: {type(idx)}")
    assert all(a % b == 0 for a, b in zip(untiled_shape, self.tiling))
    tiled_shape = [
        *(a // b for a, b in zip(untiled_shape, self.tiling)),
        *self.tiling,
    ]
    new_indexer = indexing.NDIndexer.from_indices_shape(
        indices=(*untiled_idxs, *idxs_after_tiling),
        shape=(*leading_shape, *tiled_shape)
    )
    return new_indexer, self

  def commute_reshape(
      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform
  ) -> tuple[state_types.ReshapeTransform, UntilingTransform]:
    if not transform.shape:
      raise NotImplementedError(
          "Commuting a `UntilingTransform` with a `ReshapeTransform` is not "
          "supported when the target shape has 0 dimensions"
      )
    if not self.tiling:
      raise NotImplementedError(
          "Commuting a `UntilingTransform` with a `ReshapeTransform` is not "
          "supported when the tiling is empty"
      )
    untiled_aval = self.transform_type(aval)
    assert isinstance(untiled_aval, jax_core.ShapedArray)
    components = [[]]
    # We assume that we support only folds here for the moment. Therefore, we
    # can gather a number of consecutive dimensions such that their product
    # equals the dimension currently being processed in the reshaped shape.
    for d in untiled_aval.shape:
      reshaped_dim_size = transform.shape[len(components) - 1]
      components[-1].append(d)
      component_size = math.prod(components[-1])
      if component_size == reshaped_dim_size:
        components.append([])
      elif component_size > reshaped_dim_size:
        raise NotImplementedError(
            "Unfolding dimensions is not supported when commuting an "
            " `UntilingTransform` with a `ReshapeTransform`"
        )
    assert not components[-1]
    components.pop()
    assert len(components) == len(transform.shape)

    rev_tiling_to_process = list(self.tiling)[::-1]
    rev_shape_to_process = untiled_aval.shape[-len(self.tiling):][::-1]
    rev_new_tiling: list[int] = []
    rev_new_tiled_dims: list[int] = []
    for component in components[::-1]:
      # The construction above should guarantee that there is never an empty
      # component, which simplifies indexing below.
      assert component
      ndim = len(component)
      if len(rev_tiling_to_process) < ndim:
        raise NotImplementedError(
            "Folding tiled dimensions into untiled dimensions is not supported"
        )
      rev_tiling_slice = rev_tiling_to_process[:ndim]
      rev_shape_slice = rev_shape_to_process[:ndim]
      new_tiling_dim = math.prod(rev_tiling_slice)
      num_elems = math.prod(rev_shape_slice)
      assert num_elems % new_tiling_dim == 0
      new_tiled_dim = num_elems // new_tiling_dim
      # If any other dimension than the minormost one has non-unit tiling, then
      # we cannot commute the reshape and untile transforms.
      #
      # Note that we could also support the case where we are collapsing
      # trailing tiled dimensions where the tile size is the dimension size
      # (i.e. there is a single tile).
      if any(t != 1 for t in rev_tiling_slice[1:]):
        before = (
            *[s // t for s, t in zip(rev_shape_slice, rev_tiling_slice)],
            *rev_tiling_slice,
        )
        after = (new_tiled_dim, new_tiling_dim)
        raise ValueError(
            "Cannot commute `UntilingTransform` with `ReshapeTransform` when "
            "any of the dimensions being collapsed other than the minormost "
            "one has non-unit tiling. Attempted to reshape tiled slice of "
            f"shape {before} into shape {after}"
        )
      rev_new_tiling.append(new_tiling_dim)
      rev_new_tiled_dims.append(new_tiled_dim)
      rev_tiling_to_process = rev_tiling_to_process[ndim:]
      rev_shape_to_process = rev_shape_to_process[ndim:]
      if not rev_tiling_to_process:
        break
    assert not rev_tiling_to_process
    assert not rev_shape_to_process
    new_tiling = tuple(rev_new_tiling[::-1])
    new_tiled_dims = tuple(rev_new_tiled_dims[::-1])
    new_shape = (
        *transform.shape[:len(components) - len(rev_new_tiling)],
        *new_tiled_dims,
        *new_tiling,
    )

    return state_types.ReshapeTransform(new_shape), UntilingTransform(new_tiling)

  def pretty_print(self, context: jax_core.JaxprPpContext) -> pp.Doc:
    return pp.text(f"{{untile({list(self.tiling)})}}")


def batch_transform(
    transform: state_types.Transform, leading_rank: int
) -> state_types.Transform:
  match transform:
    case TransposeTransform() as t:
      return TransposeTransform(
          (*range(leading_rank), *(d + leading_rank for d in t.permutation))
      )
    case TilingTransform() | SwizzleTransform() as t:
      return t
    case _:
      raise NotImplementedError(f"Unsupported transform: {type(transform)}")


def to_gpu_transform(
    transform: state_types.Transform,
) -> mgpu.MemRefTransform:
  match transform:
    case TransposeTransform(permutation):
      return mgpu.TransposeTransform(permutation)
    case TilingTransform(tiling):
      return mgpu.TileTransform(tiling)
    case _:
      raise TypeError(f"Unsupported transform: {type(transform)}")


def to_transform_attr(
    transform: state_types.Transform,
) -> ir.Attribute:
  match transform:
    case SwizzleTransform(swizzle):
      return mgpu.dialect.SwizzleTransformAttr.get(swizzle)
    case _:
      return to_gpu_transform(transform).to_attr()


# TODO(sharadmv): upstream into pallas core
def commute_transpose_indexer(
    _: jax_core.AbstractValue,
    transpose: state_types.TransposeTransform,
    indexer: indexing.NDIndexer,
) -> tuple[indexing.NDIndexer, state_types.TransposeTransform]:
  idxs = indexer.indices
  removed_dims = [
      i
      for i, idx in enumerate(idxs)
      if not isinstance(idx, (slice, indexing.Slice))
  ]
  new_perm = tuple(
      p - sum(d < p for d in removed_dims)
      for p in transpose.permutation
      if p not in removed_dims
  )
  new_shape = tuple(
      indexer.shape[i] for i in state_types._perm_inverse(transpose.permutation)
  )
  new_idxs = tuple(
      idxs[i] for i in state_types._perm_inverse(transpose.permutation)
  )
  new_indexer = indexing.NDIndexer.from_indices_shape(
      indices=new_idxs, shape=new_shape
  )
  return new_indexer, state_types.TransposeTransform(new_perm)


@tree_util.register_dataclass
@dataclasses.dataclass
class PeerMemRef(state_types.Transform):
  device_id: Any
  device_id_type: pallas_primitives.DeviceIdType = jax.tree.static()

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    raise NotImplementedError()

  def transform_type(self, x):
    return x

  def commute_ndindexer(
      self, _: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, PeerMemRef]:
    return indexer, self


@tree_util.register_dataclass
@dataclasses.dataclass
class MulticastRef(state_types.Transform):
  collective_axes: tuple[Hashable, ...] = jax.tree.static()

  def transform_type(self, x):
    return x

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    raise NotImplementedError()

  def commute_ndindexer(
      self, _: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, MulticastRef]:
    return indexer, self


def remote_ref(
    ref: _Ref,
    device_id: pallas_primitives.DeviceId,
    device_id_type: pallas_primitives.DeviceIdType | None = None,
) -> pallas_core.TransformedRef:
  """Translate memref to a symmetric memref on a peer device."""
  if device_id_type is not None:
    deprecations.warn(
        "jax-pallas-device-id-type",
        "device_id_type is deprecated and will be removed in a future release.",
        stacklevel=2,
    )
  else:
    device_id_type = pallas_primitives.DeviceIdType.MESH
  device_id = pallas_primitives.canonicalize_device_id(device_id)
  if not isinstance(ref, pallas_core.TransformedRef):
    if not isinstance(jax_core.typeof(ref), state_types.AbstractRef):
      raise TypeError("ref must be a reference")
    ref = pallas_core.TransformedRef(ref, transforms=())
  if any(isinstance(t, MulticastRef) for t in ref.transforms):
    raise ValueError("Can't make a multicast reference into a peer reference.")
  return pallas_core.TransformedRef(
      ref.ref, (*ref.transforms, PeerMemRef(device_id, device_id_type)),
  )


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class ClusterRefTransform(state_types.Transform):
  dims: tuple[jax_core.AxisName, ...] = jax.tree.static()
  idxs: tuple[Any, ...]

  def __post_init__(self):
    if len(self.dims) != len(self.idxs):
      raise ValueError("dims and idxs must have the same length")

  def transform_type(self, x):
    return x

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    raise NotImplementedError()

  def commute_ndindexer(
      self, _: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, ClusterRefTransform]:
    return indexer, self


def cluster_ref(
    ref: _Ref,
    block_id: dict[jax_core.AxisName, Any],
) -> pallas_core.TransformedRef:
  """Translate memref to a peer memref in the cluster."""
  if not isinstance(ref, pallas_core.TransformedRef):
    if not isinstance(jax_core.typeof(ref), state_types.AbstractRef):
      raise TypeError("ref must be a reference")
    ref = pallas_core.TransformedRef(ref, transforms=())
  dims = tuple(block_id.keys())
  idxs = tuple(block_id.values())
  return pallas_core.TransformedRef(
      ref.ref, (*ref.transforms, ClusterRefTransform(dims, idxs)),
  )


def multicast_ref(
    ref: _Ref,
    collective_axes: Hashable | tuple[Hashable, ...],
) -> pallas_core.TransformedRef:
  """Return a multicast reference for cross-device operations.

  Args:
    ref: The reference to transform.
    collective_axes: The JAX mesh axes indicating the devices to operate on.
  """
  if not isinstance(collective_axes, tuple):
    collective_axes = (collective_axes,)
  if not isinstance(ref, pallas_core.TransformedRef):
    if not isinstance(jax_core.typeof(ref), state_types.AbstractRef):
      raise TypeError("ref must be a reference")
    ref = pallas_core.TransformedRef(ref, transforms=())
  if any(isinstance(t, PeerMemRef) for t in ref.transforms):
    raise ValueError("Can't make a peer reference into a multicast reference.")
  return pallas_core.TransformedRef(
      ref.ref, (*ref.transforms, MulticastRef(collective_axes)),
  )


def transpose_ref(
    ref: pallas_core.TransformedRef | Any,
    permutation: tuple[int, ...],
) -> pallas_core.TransformedRef:
  assert hasattr(ref, "memory_space")
  if ref.memory_space == MemorySpace.TMEM:
    raise ValueError("Can't transpose a TMEM reference.")
  return ref.transpose(permutation)


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class ExtractAliasedRef(state_types.Transform):
  """Bitcasts the underlying ref at the given offset to the given shape and dtype."""
  dtype: dtypes.DType = jax.tree.static()
  shape: tuple[int, ...] = jax.tree.static()
  offset: int = jax.tree.static()

  # The index of the group of this aliased ref within the input RefUnion.
  alias_group_idx: int = jax.tree.static()

  # TMEM-specific params
  layout: tcgen05.TMEMLayout | None = jax.tree.static()

  @classmethod
  def from_transformed_ref(
      cls,
      ref: pallas_core.TransformedRef,
      byte_offset: int,
      alias_group_idx: int,
      layout: tcgen05.TMEMLayout | None = None,
  ):
    return cls(dtypes.dtype(ref.dtype), ref.ref.shape, byte_offset, alias_group_idx, layout)

  def transform_type(self, x):
    match x:
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case jax_core.ShapedArray():
        return x.update(shape=self.shape, dtype=self.dtype)
      case _:
        raise TypeError(f"Unsupported type: {x}")


@dataclasses.dataclass(frozen=True)
class SwizzleTransform(state_types.Transform):
  swizzle: int

  def __post_init__(self):
    if self.swizzle not in {32, 64, 128}:
      raise ValueError(
          f"Swizzle {self.swizzle} is not supported. Only 32, 64 and 128 are"
          " accepted."
      )

  def transform_type(
      self, x: jax_core.AbstractValue
  ) -> jax_core.AbstractValue:
    match x:
      case jax_core.ShapedArray():
        swizzle_elems = (self.swizzle * 8) // dtypes.itemsize_bits(x.dtype)
        if swizzle_elems != x.shape[-1]:
          raise ValueError(
              f"Swizzle {self.swizzle} requires the trailing dimension to be of"
              f" size {swizzle_elems}, but got shape: {x.shape}"
          )
        return x
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case _:
        raise NotImplementedError(f"Unsupported type: {x}")

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    return UnswizzleRef(self.swizzle)


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class UnswizzleRef(state_types.Transform):
  swizzle: int = jax.tree.static()

  def transform_type(self, x: jax_core.AbstractValue) -> jax_core.AbstractValue:
    # Swizzling preserves the type
    return x

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    return SwizzleTransform(self.swizzle)

  def swizzle_elems(self, dtype: jax.typing.DTypeLike | ir.Type) -> int:
    if not isinstance(dtype, ir.Type):
      dtype = mgpu_utils.dtype_to_ir_type(dtype)
    return (self.swizzle * 8) // mgpu.bitwidth(dtype)

  def commute_transpose(
      self, _: jax_core.AbstractValue, transpose: state_types.TransposeTransform
  ) -> tuple[state_types.TransposeTransform, UnswizzleRef]:
    perm = transpose.permutation
    if perm[-1] != len(perm) - 1:
      raise ValueError("Can't transpose the swizzled dimension.")
    return transpose, self

  def commute_reshape(
      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform
  ) -> tuple[state_types.ReshapeTransform, UnswizzleRef]:
    shape = aval.shape
    if shape[-1] != self.swizzle_elems(aval.dtype):
      raise ValueError(
          f"Reshape shape {shape} is not divisible by swizzle elements"
          f" {self.swizzle_elems(aval.dtype)}"
      )
    return transform, self

  def commute_ndindexer(
      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, UnswizzleRef]:
    if not hasattr(aval, "dtype"):
      raise ValueError(
          f"Cannot commute unswizzle and indexer with {aval}, which does not"
          " have a dtype"
      )
    dtype = aval.dtype
    swizzle_elems = self.swizzle_elems(dtype)
    idxs = indexer.indices
    if not idxs:
      return indexer, self
    if not all(isinstance(idx, (slice, indexing.Slice)) for idx in idxs[-2:]):
      raise NotImplementedError(
          f"Non-slice indices are not supported in 2 minormost dims: {idxs}"
      )
    last_idx = idxs[-1]
    if isinstance(last_idx, indexing.Slice):
      if last_idx.start != 0 or last_idx.size != swizzle_elems:
        raise ValueError("Swizzled dims cannot be sliced")
    else:
      assert isinstance(last_idx, slice)
      if (
          (last_idx.step is not None and last_idx.step != 1)
          or (last_idx.start is not None and last_idx.start != 0)
          or (last_idx.stop is not None and last_idx.stop != swizzle_elems)
      ):
        raise ValueError("Swizzled dims cannot be sliced")
    return indexer, self

  def pretty_print(self, context: jax_core.JaxprPpContext) -> pp.Doc:
    return pp.text(f"{{unswizzle({self.swizzle})}}")


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class CollapseLeadingBatchDimensionsTransform(state_types.Transform):
  """A transform that collapses leading batch dimensions into the minor dimension.

  Specifically, it maps `(*batch_shape, m, n)` to `(m, math.prod(batch_shape) *
  n)`.
  """

  def transform_type(
      self, x: jax_core.AbstractValue
  ) -> state_types.AbstractRef:
    match x:
      case jax_core.ShapedArray():
        if x.ndim < 2:
          raise ValueError(f"Unsupported ndim: {x.ndim}")
        batch_size = math.prod(x.shape[:-2])
        transformed_shape = (x.shape[-2], batch_size * x.shape[-1])
        return x.update(shape=transformed_shape)
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case _:
        raise TypeError(f"Unsupported type: {x}")

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    assert hasattr(x, "shape")
    return ExpandLeadingBatchDimensionsTransform(x.shape[:-2])


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class ExpandLeadingBatchDimensionsTransform(state_types.Transform):
  """The inverse of CollapseLeadingBatchDimensionsTransform.

  Specifically, it maps `(m, math.prod(batch_shape) * n)` to `(*batch_shape, m,
  n)`.
  """

  batch_shape: tuple[int, ...] = jax.tree.static()

  def transform_type(
      self, x: jax_core.AbstractValue
  ) -> state_types.AbstractRef:
    match x:
      case jax_core.ShapedArray():
        if x.ndim != 2:
          raise ValueError(f"Unsupported shape: {x.shape}")
        batch_size = math.prod(self.batch_shape)
        if x.shape[1] % batch_size != 0:
          raise ValueError(
              f"Second dimension {x.shape[1]} must be divisible by batch_size"
              f" {batch_size}"
          )
        transformed_shape = self.batch_shape + (
            x.shape[0],
            x.shape[1] // batch_size,
        )
        return x.update(shape=transformed_shape)
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case _:
        raise TypeError(f"Unsupported type: {x}")

  def commute_ndindexer(
      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, state_types.Transform]:
    del aval
    batch_shape = self.batch_shape
    k = len(batch_shape)
    if len(indexer.indices) != k + 2:
      raise ValueError(
          f"Expected indexer to have exactly {k + 2} dimensions, "
          f"but got {len(indexer.indices)}."
      )
    batch_indices = indexer.indices[:-2]
    row_idx = indexer.indices[-2]
    col_idx = indexer.indices[-1]

    for idx in batch_indices:
      if isinstance(idx, indexing.Slice):
        raise NotImplementedError("Slicing batch dimensions is not supported.")

    batch_size = math.prod(batch_shape)
    m, n = indexer.shape[-2], indexer.shape[-1]
    physical_shape = (m, batch_size * n)

    batch_idx = 0
    for idx, size in zip(batch_indices, batch_shape):
      assert not isinstance(idx, indexing.Slice)
      batch_idx = batch_idx * size + idx

    if isinstance(col_idx, indexing.Slice):
      # We shift the column slice by batch_idx * n.
      new_col_idx = indexing.Slice(
          batch_idx * n + col_idx.start, col_idx.size, col_idx.stride
      )
    else:
      new_col_idx = batch_idx * n + col_idx

    new_indexer = indexing.NDIndexer.from_indices_shape(
        (row_idx, new_col_idx), physical_shape
    )
    return new_indexer, IdentityTransform()


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class IdentityTransform(state_types.Transform):
  """An identity transform."""

  def transform_type(self, x: jax_core.AbstractValue) -> jax_core.AbstractValue:
    return x

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    return self


@dataclasses.dataclass
class BlockSpec(pallas_core.BlockSpec):
  r"""A GPU-specific ``BlockSpec``.

  Attributes:
    transforms: A sequence of transforms that will be applied to the
      reference.
    delay_release: used during pipelining to delay the release of
      resources of a slot after it is used in the computation.
    collective_axes: When set, all blocks along the specified axes must execute
      the same sequence of pipeline operations (with the only exception being
      the index_map in non-collective ``BlockSpec``\ s), and all of them must
      return the same block from the index_map for this operand. This enables
      the pipelining helpers to use collective async copies, which can improve
      performance.
    oob_fill_mode: Determines the behavior for out-of-bounds accesses during
      pipelined GMEM-to-SMEM copies. Defaults to ``OOBFillMode.ZEROS``.
  """
  transforms: Sequence[state_types.Transform] = ()
  delay_release: int = 0
  collective_axes: tuple[Hashable, ...] | None = None
  oob_fill_mode: OOBFillMode = OOBFillMode.ZEROS

  def to_block_mapping(
      self,
      origin: pallas_core.OriginStr,
      array_aval: jax_core.ShapedArray,
      *,
      index_map_avals: Sequence[jax_core.AbstractValue],
      index_map_tree: tree_util.PyTreeDef,
      grid: pallas_core.GridMappingGrid,
      vmapped_dims: tuple[int, ...],
      allow_captured_consts: bool = False,
      debug: bool = False,
  ) -> pallas_core.BlockMapping:
    if self.collective_axes:
      raise ValueError(
          "collective_axes is not supported in pallas_call. Use plgpu.kernel"
          " with plgpu.emit_pipeline_warp_specialized instead."
      )
    bm = super().to_block_mapping(
        origin,
        array_aval,
        index_map_avals=index_map_avals,
        index_map_tree=index_map_tree,
        grid=grid,
        vmapped_dims=vmapped_dims,
        allow_captured_consts=allow_captured_consts,
        debug=debug,
    )
    block_inner_aval = bm.block_aval.inner_aval
    for t in self.transforms:
      block_inner_aval = t.transform_type(block_inner_aval)
    return bm.replace(
        transformed_block_aval=bm.block_aval.update(
            inner_aval=block_inner_aval
        ),
        transforms=self.transforms,
    )


GMEM = MemorySpace.GMEM
SMEM = MemorySpace.SMEM
TMEM = MemorySpace.TMEM
REGS = MemorySpace.REGS


class barrier_dtype(dtypes.extended):
  pass


@dataclasses.dataclass(frozen=True)
class BarrierType(dtypes.ExtendedDType):
  type: ClassVar[Any] = barrier_dtype  # pyrefly: ignore[bad-override]
  name: ClassVar[str] = "barrier"

  num_arrivals: int
  orders_tensor_core: bool

  def __str__(self):
    return self.name


@dataclasses.dataclass(frozen=True)
class ClusterBarrierType(dtypes.ExtendedDType):
  type: ClassVar[Any] = barrier_dtype  # pyrefly: ignore[bad-override]
  name: ClassVar[str] = "cluster_barrier"

  collective_axes: tuple[str | tuple[str, ...], ...]
  num_arrivals: int
  orders_tensor_core: bool
  leader_tracked: bool = False

  def __str__(self):
    return self.name


@dataclasses.dataclass(frozen=True, kw_only=True)
class Barrier:
  """Describes a barrier reference.

  Attributes:
    num_arrivals: The number of arrivals that will be recorded by this barrier.
    num_barriers: The shape of the barrier array to create. By default a single
      barrier is created, and the reference to it has an empty shape. Individual
      barriers of an array can be accessed by indexing into the barrier Ref.
    orders_tensor_core: If False, a successful wait from one thread does not
      guarantee that the TensorCore-related operations in other threads have
      completed. Similarly, when False any TensorCore operation in the waiting
      thread is allowed to begin before the wait succeeds.
  """
  num_arrivals: int = 1
  num_barriers: int | Sequence[int] = ()
  orders_tensor_core: bool = False

  def __post_init__(self):
    if (n := self.num_arrivals) < 1:
      raise ValueError(f"Num arrivals must be at least 1, but got {n}")

    if isinstance(self.num_barriers, int):
      object.__setattr__(self, "num_barriers", (self.num_barriers,))
    else:
      object.__setattr__(self, "num_barriers", tuple(self.num_barriers))

  def get_array_aval(self) -> jax_core.ShapedArray:
    raise ValueError("Barriers are not arrays")

  def get_ref_aval(self) -> state.AbstractRef:
    ty = BarrierType(self.num_arrivals, self.orders_tensor_core)
    return state.AbstractRef(jax_core.ShapedArray(self.num_barriers, ty), SMEM)


@dataclasses.dataclass(frozen=True, kw_only=True)
class ClusterBarrier:
  collective_axes: tuple[str | tuple[str, ...], ...]
  num_barriers: int | Sequence[int] = ()
  num_arrivals: int = 1
  orders_tensor_core: bool = False
  leader_tracked: bool = False

  def __post_init__(self):
    if (n := self.num_arrivals) < 1:
      raise ValueError(f"Num arrivals must be at least 1, but got {n}")

    if isinstance(self.num_barriers, int):
      object.__setattr__(self, "num_barriers", (self.num_barriers,))
    else:
      object.__setattr__(self, "num_barriers", tuple(self.num_barriers))

  def get_array_aval(self) -> jax_core.ShapedArray:
    raise ValueError("Cluster barriers are not arrays")

  def get_ref_aval(self) -> state.AbstractRef:
    ty = ClusterBarrierType(
        collective_axes=self.collective_axes,
        num_arrivals=self.num_arrivals,
        orders_tensor_core=self.orders_tensor_core,
        leader_tracked=self.leader_tracked,
    )
    return state.AbstractRef(jax_core.ShapedArray(self.num_barriers, ty), SMEM)


@dataclasses.dataclass(frozen=True)
class WGMMAAccumulatorRef:
  shape: tuple[int, int]
  dtype: jax.typing.DTypeLike = jnp.float32
  _init: Any = state_types.uninitialized

  def get_ref_aval(self) -> state.AbstractRef:
    if self._init is not state_types.uninitialized:
      raise ValueError(
          "Preinitialized WGMMAAccumulatorRef only supported in pl.run_state."
      )
    return WGMMAAbstractAccumulatorRef(
        jax_core.ShapedArray(shape=self.shape, dtype=self.dtype), MemorySpace.REGS
    )

  @staticmethod
  def init(array):
    return WGMMAAccumulatorRef(array.shape, array.dtype, _init=array)


def _wgmma_ref_type_mapping(ref: WGMMAAccumulatorRef):
  aval = WGMMAAbstractAccumulatorRef(
      jax_core.ShapedArray(shape=ref.shape, dtype=ref.dtype), MemorySpace.REGS
  )
  return aval, ref._init
state_types._ref_type_aval_mappings[WGMMAAccumulatorRef] = _wgmma_ref_type_mapping


class WGMMAAbstractAccumulatorRef(state.AbstractRef):
  __slots__ = ["inner_aval", "memory_space"]

  def __repr__(self) -> str:
    return f'Accumulator{{{self.inner_aval.str_short()}}}'

  def update(self, inner_aval=None, memory_space=None, kind=None):
    ref = super().update(inner_aval, memory_space, kind)
    return WGMMAAbstractAccumulatorRef(
        inner_aval=ref.inner_aval,
        memory_space=ref.memory_space,
    )

  def _getitem(self, tracer, idx):
    from jax._src.pallas.mosaic_gpu.primitives import wgmma_accumulator_load  # pyrefly: ignore[missing-import]
    arr = wgmma_accumulator_load(tracer, wait_n=0)
    if not is_trivial_index(idx, tracer.shape):
      arr = arr[idx]

    return arr

  def _setitem(self, tracer, idx, value):
    from jax._src.pallas.mosaic_gpu.primitives import wgmma_accumulator_store  # pyrefly: ignore[missing-import]
    if not is_trivial_index(idx, tracer.shape):
      raise NotImplementedError(
          "Non-trivial indexing on WGMMAAbstractAccumulatorRef is not supported"
          " for stores."
      )
    wgmma_accumulator_store(tracer, value)


class AbstractTMEMRef(state.AbstractRef):
  __slots__ = ["inner_aval", "memory_space", "layout", "collective"]

  def __init__(self, inner_aval, memory_space, layout, collective):
    super().__init__(inner_aval, memory_space)
    self.layout = layout
    self.collective = collective

  def __repr__(self) -> str:
    return f'TMEM({self.inner_aval.str_short()}, layout={self.layout}, collective={self.collective})'

  def update(self, inner_aval=None, memory_space=None, kind=None):
    ref = super().update(inner_aval, memory_space, kind)
    return AbstractTMEMRef(
        ref.inner_aval, ref.memory_space, self.layout, self.collective
    )


_WARPGROUP_AXIS_NAME = object()

@dataclasses.dataclass(frozen=True, kw_only=True)
class Mesh(pallas_core.Mesh):
  grid: Sequence[int] = ()
  grid_names: Sequence[str] = ()
  cluster: Sequence[int] = ()
  cluster_names: Sequence[str] = ()
  # Those are NOT CUDA threads. On Hopper they correspond to warpgroups.
  num_threads: int | None = None
  thread_name: str | None = None
  kernel_name: str | None = None

  def __post_init__(self):
    if len(self.cluster) > 3:
      raise ValueError(f"cluster= must be at most 3D, got {self}.")
    if len(self.grid_names) != len(self.grid):
      raise ValueError(
          f"grid_names must have the same length as grid, got {self}."
      )
    if len(self.cluster_names) != len(self.cluster):
      raise ValueError(
          f"cluster_names must have the same length as cluster, got {self}."
      )
    if (self.thread_name is None) != (self.num_threads is None):
      raise ValueError(
          "num_threads and thread_name must be either both set or both None,"
          f" got {self}"
      )
    max_mosaic_threads = 2048 // 128
    if self.num_threads is not None and self.num_threads > max_mosaic_threads:
      raise ValueError(
          "Requested too many CUDA threads per block. Each Mosaic thread"
          f" corresponds to 128 CUDA threads. At most {max_mosaic_threads}"
          f" are supported, got {self}"
      )
    object.__setattr__(self, "grid", tuple(self.grid))
    object.__setattr__(self, "grid_names", tuple(self.grid_names))
    object.__setattr__(self, "cluster", tuple(self.cluster))
    object.__setattr__(self, "cluster_names", tuple(self.cluster_names))

  @property
  def default_memory_space(self) -> MemorySpace:
    return MemorySpace.GMEM

  @property
  def shape(self) -> collections.OrderedDict[object, int]:
    pairs: Iterable[tuple[object, int]]
    if self.num_threads is not None:
      pairs = zip(
          (*self.grid_names, *self.cluster_names, self.thread_name),
          (*self.grid, *self.cluster, self.num_threads),
      )
    else:
      pairs = zip(
          (*self.grid_names, *self.cluster_names),
          (*self.grid, *self.cluster),
      )
    return collections.OrderedDict(pairs)

  def discharges_effect(self, effect: jax_core.Effect) -> bool:
    return effect is _wgmma_pipeline_effect or effect is _memory_effect

  def check_is_compatible_with(self, other_mesh):
    raise NotImplementedError()

  @property
  def core_type(self) -> str:
    return "gpu"

  @property
  def supported_memory_spaces(self) -> Sequence[MemorySpace]:
    return [*MemorySpace]

  @contextlib.contextmanager
  def tracing_context(self):
    # This is needed to support program_id inside of plgpu kernels.
    with pallas_core.tracing_grid_env(
        tuple(self.shape.values()), mapped_dims=()
    ):
      yield


@dataclasses.dataclass(frozen=True, kw_only=True)
class WarpMesh(pallas_core.Mesh):
  """Represents a mesh over individual warps within a warpgroup.

  When used in conjunction with `core_map`, the warp ID will be visible
  within the body of the wrapped scope by querying `lax.axis_index` with
  the specified axis name.
  """

  _NUM_WARPS_PER_WARPGROUP: ClassVar[int] = 4
  axis_name: jax_core.AxisName

  @property
  def shape(self):
    return collections.OrderedDict([
        (self.axis_name, self._NUM_WARPS_PER_WARPGROUP),
    ])

  @property
  def default_memory_space(self) -> MemorySpace:
    raise NotImplementedError

  def discharges_effect(self, effect: jax_core.Effect) -> Literal[False]:
    del effect
    return False

  def check_is_compatible_with(self, other_mesh):
    raise NotImplementedError()

  @property
  def core_type(self) -> str:
    return "gpu"

  @property
  def supported_memory_spaces(self) -> Sequence[MemorySpace]:
    return [GMEM, SMEM, TMEM]

  @contextlib.contextmanager
  def tracing_context(self):
    yield


class MemoryEffect(jax_core.Effect):
  pass


pallas_core.kernel_local_effects.add_type(MemoryEffect)
effects.control_flow_allowed_effects.add_type(MemoryEffect)
_memory_effect = MemoryEffect()


class _WGMMAPipelineEffect(effects.Effect):
  pass


pallas_core.kernel_local_effects.add_type(_WGMMAPipelineEffect)
effects.control_flow_allowed_effects.add_type(_WGMMAPipelineEffect)
_wgmma_pipeline_effect = _WGMMAPipelineEffect()


class PdlEffect(jax_core.Effect):
  """Indicates that a kernel uses Programmatic Dependency Launch (PDL).

  This effect is associated with primitives that perform grid-level
  synchronization (e.g. waiting for dependent grids). Its presence in a
  JAXPR instructs the compiler/runtime to enable programmatic stream
  serialization, allowing the CUDA driver to launch dependent kernels
  concurrently in the same stream.
  """


pallas_core.kernel_local_effects.add_type(PdlEffect)
effects.control_flow_allowed_effects.add_type(PdlEffect)
_pdl_effect = PdlEffect()


# We define the layout_cast primitive here, because it needs to be available in
# the lowering code (to provide layout hints to the rules).
#
# TODO(bchetioui): unify layout_cast_p with relayout_p
layout_cast_p = jax_core.Primitive("layout_cast")


@layout_cast_p.def_abstract_eval
def _layout_cast_abstract_eval(x, new_layout):
  del new_layout  # Unused.
  return x


def layout_cast(x: Any, new_layout: SomeLayout):
  """Casts the layout of the given array."""
  return layout_cast_p.bind(x, new_layout=new_layout)


class SomeLayout:

  def reduce(self, axes: int | Sequence[int]) -> SomeLayout:
    if isinstance(axes, int):
      axes = (axes,)
    return ReducedLayout(self, axes)

  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
    raise NotImplementedError


@dataclasses.dataclass(frozen=True)
class ParameterizedLayout(SomeLayout):
  layout_cls: Layout | TMEMLayout
  args: Sequence[Any]
  kwargs: Any

  def __post_init__(self):
    object.__setattr__(self, "args", tuple(self.args))
    object.__setattr__(self, "kwargs", frozen_dict.FrozenDict(self.kwargs))

  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
    if args or kwargs:
      raise ValueError(f"Can't instantiate {self} with arguments.")
    return self.layout_cls.to_mgpu(*self.args, **self.kwargs)


@dataclasses.dataclass(frozen=True)
class ReducedLayout(SomeLayout):
  layout: SomeLayout
  axes: Sequence[int]

  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
    if args or kwargs:
      raise ValueError(f"Can't instantiate {self} with arguments.")
    layout = self.layout.to_mgpu()
    if not isinstance(layout, mgpu.TiledLayout):
      raise ValueError("Only TiledLayout supports reductions.")
    return layout.reduce(self.axes)


class Layout(SomeLayout, enum.Enum):
  #: [m, n] matrix, where m % 64 == 0 == n % 8.
  WGMMA = enum.auto()
  WGMMA_8BIT = enum.auto()
  WGMMA_UPCAST_2X = enum.auto()
  WGMMA_UPCAST_4X = enum.auto()
  WGMMA_TRANSPOSED = enum.auto()

  WG_SPLAT = enum.auto()
  WG_STRIDED = enum.auto()

  TILED = enum.auto()

  TCGEN05 = enum.auto()
  TCGEN05_TRANSPOSED = enum.auto()
  TCGEN05_M64_COLLECTIVE = enum.auto()
  TCGEN05_TMEM_NATIVE = enum.auto()
  TCGEN05_M64_COLLECTIVE_NATIVE = enum.auto()

  SMEM_GMEM_COPY = enum.auto()
  TMA_INDICES = enum.auto()
  TMA_INDICES_4 = enum.auto()

  # TODO(b/435159109): Remove this once LLVM regression is addressed.
  _WGMMA_ACC_32BIT = enum.auto()  # Temporarily exposed to work around LLVM bugs

  #: LHS operand layout for SM8x ``mma.sync``
  MMA_LHS = enum.auto()
  #: RHS operand layout for SM8x ``mma.sync``
  MMA_RHS = enum.auto()
  #: Accumulator's layout for SM8x ``mma.sync``
  MMA_ACC = enum.auto()

  def __call__(self, *args, **kwargs) -> ParameterizedLayout:
    return ParameterizedLayout(self, args, kwargs)

  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
    def check_no_args():
      if args or kwargs:
        raise ValueError(f"Can't instantiate {self} with arguments.")

    match self:
      case Layout.WGMMA_TRANSPOSED:
        check_no_args()
        return mgpu.WGMMA_TRANSPOSED_LAYOUT
      case Layout.WGMMA:
        check_no_args()
        return mgpu.WGMMA_LAYOUT
      case Layout.WGMMA_8BIT:
        check_no_args()
        return mgpu.WGMMA_LAYOUT_8BIT
      case Layout.WGMMA_UPCAST_2X:
        check_no_args()
        return mgpu.WGMMA_LAYOUT_UPCAST_2X
      case Layout.WGMMA_UPCAST_4X:
        check_no_args()
        return mgpu.WGMMA_LAYOUT_UPCAST_4X
      case Layout._WGMMA_ACC_32BIT:
        check_no_args()
        return mgpu.fragmented_array.WGMMA_LAYOUT_ACC_32BIT
      case Layout.WG_SPLAT:
        return mgpu.WGSplatFragLayout(*args, **kwargs)
      case Layout.WG_STRIDED:
        return mgpu.WGStridedFragLayout(*args, **kwargs)
      case Layout.TILED:
        return mgpu.TiledLayout(*args, **kwargs)
      case Layout.TCGEN05:
        check_no_args()
        return mgpu.TCGEN05_LAYOUT
      case Layout.TCGEN05_TRANSPOSED:
        check_no_args()
        return mgpu.TCGEN05_TRANSPOSED_LAYOUT
      case Layout.TCGEN05_TMEM_NATIVE:
        if args or kwargs:
          return mgpu.tmem_native_layout(*args, **kwargs)
        return mgpu.TMEM_NATIVE_LAYOUT
      case Layout.TCGEN05_M64_COLLECTIVE:
        return tcgen05.fa_m64_collective_layout(*args, **kwargs)
      case Layout.TCGEN05_M64_COLLECTIVE_NATIVE:
        return tcgen05.tmem_m64_collective_layout(*args, **kwargs).as_tiled_layout()
      case Layout.SMEM_GMEM_COPY:
        normalize_args = lambda shape, dtype, swizzle: (shape, dtype, swizzle)
        shape, dtype, swizzle = normalize_args(*args, **kwargs)
        bitwidth = dtypes.itemsize_bits(dtype)
        tiling = (8, 8 * swizzle // bitwidth)
        row_tiles, col_tiles = mgpu.tile_shape(shape, tiling)[-4:-2]
        return mgpu.fragmented_array.tiled_copy_smem_gmem_layout(
            row_tiles, col_tiles, swizzle, bitwidth
        )
      case Layout.TMA_INDICES:
        return mgpu.TMA_INDICES_LAYOUT
      case Layout.MMA_LHS:
        normalize_args = lambda dtype, *, m_warps=4: (dtype, m_warps)
        dtype, m_warps = normalize_args(*args, **kwargs)
        return mgpu.MMALayouts(dtype, m_warps=m_warps).lhs
      case Layout.MMA_RHS:
        normalize_args = lambda dtype, *, m_warps=4: (dtype, m_warps)
        dtype, m_warps = normalize_args(*args, **kwargs)
        return mgpu.MMALayouts(dtype, m_warps=m_warps).rhs
      case Layout.MMA_ACC:
        normalize_args = lambda dtype, *, m_warps=4: (dtype, m_warps)
        dtype, m_warps = normalize_args(*args, **kwargs)
        return mgpu.MMALayouts(dtype, m_warps=m_warps).acc
      case Layout.TMA_INDICES_4:
        return mgpu.TMA_INDICES_4_LAYOUT


class TMEMLayout(enum.Enum):
  """Layout for TMEM references."""
  # TODO(apaszke): Remove the layout suffix.
  SCALES_LAYOUT = enum.auto()
  SPARSE_METADATA_LAYOUT = enum.auto()
  M64_COLLECTIVE_LAYOUT = enum.auto()
  SCALES_M64_COLLECTIVE_LAYOUT = enum.auto()

  def __call__(self, *args, **kwargs) -> ParameterizedLayout:
    return ParameterizedLayout(self, args, kwargs)

  def to_mgpu(self, *args, **kwargs) -> tcgen05.TMEMLayout:
    match self:
      case TMEMLayout.SCALES_LAYOUT:
        return tcgen05.scales_layout(*args, **kwargs)
      case TMEMLayout.SPARSE_METADATA_LAYOUT:
        return tcgen05.sparse_meta_layout(*args, **kwargs)
      case TMEMLayout.M64_COLLECTIVE_LAYOUT:
        return tcgen05.tmem_m64_collective_layout(*args, **kwargs)
      case TMEMLayout.SCALES_M64_COLLECTIVE_LAYOUT:
        return tcgen05.b_scales_m64_collective_layout(*args, **kwargs)
    raise ValueError(f"Invalid TMEMLayout: {self}")


def TryClusterCancelResult(
    num_buffers: int | None = None) -> pallas_core.MemoryRef:
  """Helper function to create Refs for cluster launch control results.

  Args:
    num_buffers: Optional argument for specifying the number of buffers
      to allocate. If None, will return a single 16-byte buffer. If specified,
      will return a (num_buffers, 16)-shaped buffer.

  Returns:
    A MemoryRef with the correct shape for holding the opaque cluster launch
    control result.
  """
  if num_buffers is None:
    return SMEM((16,), jnp.int8)
  else:
    return SMEM((num_buffers, 16), jnp.int8)