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

"""GPU-specific Pallas primitives."""

from __future__ import annotations

from collections.abc import Callable, Hashable, Sequence
import contextlib
import dataclasses
import enum
import functools
import itertools
import math
from typing import Any, Literal, assert_never, overload

import jax
from jax._src import core as jax_core
from jax._src import debugging
from jax._src import dtypes
from jax._src import lax
from jax._src import literals
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.interpreters import batching
from jax._src.lax import utils as lax_utils
from jax._src.layout import get_layout_mode, LayoutMode
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 builtin as builtin_dialect
from jax._src.lib.mlir.dialects import gpu as gpu_dialect
from jax._src.lib.mlir.dialects import llvm as llvm_dialect
from jax._src.lib.mlir.dialects import vector as vector_dialect
from jax._src.pallas import core as pallas_core
from jax._src.pallas import primitives as pallas_primitives
from jax._src.pallas import utils as pallas_utils
from jax._src.pallas.mosaic_gpu import core as gpu_core
from jax._src.pallas.mosaic_gpu import lowering
from jax._src.pallas.mosaic_gpu.core import state_types
from jax._src.state import discharge
from jax._src.state import indexing
from jax._src.state import primitives as state_primitives
from jax.experimental.mosaic import gpu as mgpu
from jax.experimental.mosaic.gpu import layouts as mgpu_layouts
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 AsyncCopyImplementation
from jax.experimental.mosaic.gpu.launch_context import CopyPartition
from jax.experimental.mosaic.gpu.launch_context import OOBFillMode
from jax.experimental.mosaic.gpu.launch_context import TileTransform
import jax.numpy as jnp
import numpy as np


AxisName = jax_core.AxisName
WARP_SIZE = 32
WARPGROUP_SIZE = 128


_Ref = state.AbstractRef | state_types.TransformedRef
SomeLayout = gpu_core.SomeLayout

def _check_ref(
    aval: object, name: str, memory_space: gpu_core.MemorySpace
) -> None:
  if not isinstance(aval, state_types.AbstractRef):
    raise TypeError(f"{name} must be a reference, got {aval}")
  aval_memory_space = getattr(aval, "memory_space", None)
  if (
      aval_memory_space is None
      or aval_memory_space is pallas_core.MemorySpace.DEFAULT
  ):
    aval_memory_space = gpu_core.GMEM
  if aval_memory_space is not memory_space:
    raise ValueError(
        f"{name} must be a {memory_space.name.upper()} reference, got {aval}"
    )


def _check_layout_mode():
  layout_mode = get_layout_mode()
  if layout_mode not in (LayoutMode.PALLAS_GPU, LayoutMode.AUTO):
    raise ValueError(
        "Layout mode must be PALLAS_GPU or AUTO, when tracing a Pallas/Mosaic "
        f"GPU kernel, but got {layout_mode}"
    )


def _andi_maybe_none(a: ir.Value | None, b: ir.Value | None) -> ir.Value | None:
  if a is None:
    return b
  if b is None:
    return a
  return arith_dialect.andi(a, b)


print_layout_p = jax_core.Primitive("print_layout")
print_layout_p.multiple_results = True


@print_layout_p.def_effectful_abstract_eval
def _print_layout_abstract_eval(aval_in, *_, **params):
  del aval_in, params  # Unused.
  return (), {debugging.debug_effect}


@lowering.register_lowering_rule(print_layout_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(
    print_layout_p, mgpu.LoweringSemantics.Warpgroup
)
def _print_layout_lowering(
    ctx: lowering.LoweringRuleContext,
    x: mgpu.FragmentedArray | tcgen05.TMEMRef | ir.Value,
    *transforms_leaves,
    fmt: str,
    transforms_tree
):
  if transforms_tree is not None:
    assert isinstance(ctx.avals_in[0], state_types.AbstractRef)
    transform_avals = transforms_tree.unflatten(ctx.avals_in[1:])
    x, _, remaining_transforms = lowering._handle_transforms(  # pyrefly: ignore[bad-specialization]
        ctx, ctx.avals_in[0], x, transform_avals,
        transforms_tree.unflatten(transforms_leaves),
    )
  elif isinstance(ctx.avals_in[0], state_types.AbstractRef):
    x, _, remaining_transforms = lowering._handle_transforms(  # pyrefly: ignore[bad-specialization]
        ctx, ctx.avals_in[0], x, (), ()
    )
  else:
    remaining_transforms = []
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    if (
        isinstance(x, ir.Value)
        and isinstance(x.type, ir.MemRefType)
        and mgpu_utils.is_smem_ref(x)
    ):
      swizzle = None
      inverse_transforms: list[TileTransform] = []
      for t in remaining_transforms:
        match t:
          case gpu_core.UnswizzleRef(s):
            swizzle = s
          case gpu_core.UntilingTransform(tiling):
            inverse_transforms.append(TileTransform(tiling))
          case _:
            raise NotImplementedError(f"Unsupported transform: {t}")
      layout_str = mgpu.dialect_lowering.pprint_transforms(
          tuple(inverse_transforms[::-1]), swizzle
      )
    else:
      if remaining_transforms:
        raise NotImplementedError(
            f"Unsupported transforms {remaining_transforms}."
        )
      layout_str = mgpu.dialect_lowering.pprint_layout(x)  # pyrefly: ignore[bad-argument-type]
    print(fmt.format(layout_str))
  else:
    assert isinstance(x, ir.Value)
    if remaining_transforms:
      raise NotImplementedError(
          f"Unsupported transforms {remaining_transforms}."
      )
    mgpu.dialect.print_layout(fmt, x)
  return ()


def print_layout(fmt: str, x: jax.typing.ArrayLike | _Ref) -> None:
  """Prints the layout chosen by Mosaic GPU for a given array, SMEM, or TMEM reference.

  This is evaluated at compile-time and has no incidence on the runtime behavior
  of the program.

  Args:
    fmt: The format string to use for printing the layout.
    x: The array, SMEM, or TMEM reference to print the layout of.
  """
  if isinstance(x, pallas_core.TransformedRef):
    transforms_leaves, transforms_tree = jax.tree.flatten(x.transforms)
    x = x.ref
  else:
    transforms_leaves, transforms_tree = [], None
  print_layout_p.bind(
      x,
      *transforms_leaves,
      fmt=fmt,
      transforms_tree=transforms_tree,
  )


copy_smem_to_gmem_p = jax_core.Primitive("copy_smem_to_gmem")
copy_smem_to_gmem_p.multiple_results = True


@copy_smem_to_gmem_p.def_effectful_abstract_eval
def _copy_smem_to_gmem_abstract_eval(src, dst, *args, **params):
  _check_ref(src, "src", gpu_core.SMEM)
  _check_ref(dst, "dst", gpu_core.GMEM)
  src_transforms_treedef = params["src_transforms_treedef"]
  dst_transforms_treedef = params["dst_transforms_treedef"]
  flat_src_transforms, flat_dst_transforms = util.split_list(
      args,
      [
          src_transforms_treedef.num_leaves,
          dst_transforms_treedef.num_leaves,
      ],
  )[:2]
  src_transforms = src_transforms_treedef.unflatten(flat_src_transforms)
  dst_transforms = dst_transforms_treedef.unflatten(flat_dst_transforms)

  src_ref = pallas_core.TransformedRef(src, src_transforms)
  dst_ref = pallas_core.TransformedRef(dst, dst_transforms)

  if src_ref.dtype != dst_ref.dtype:
    raise ValueError(
        "Expected dtypes to match but src had type "
        f"{src_ref.dtype} and dst had type {dst_ref.dtype}"
    )
  src_shape = src_ref.shape
  dst_shape = dst_ref.shape

  if any(s != d for s, d in zip(src_shape, dst_shape, strict=True)):
    raise ValueError(
        f"Expected shapes to match, but src has shape {src_shape} and dst has "
        f"shape {dst_shape}"
    )

  return (), {state.ReadEffect(0), state.WriteEffect(1)}


def _copy_smem_to_gmem_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  src, dst, *flat_args = eqn.invars
  src_transforms_treedef = eqn.params["src_transforms_treedef"]
  dst_transforms_treedef = eqn.params["dst_transforms_treedef"]
  pp_params = {}
  if not (commit_group := eqn.params["commit_group"]):
    pp_params["commit_group"] = commit_group
  if eqn.params["has_user_predicate"]:
    flat_args, user_predicate = flat_args[:-1], flat_args[-1]
    pp_params["user_predicate"] = user_predicate.pretty_print(context)
  if reduction_op := eqn.params["reduction_op"]:
    pp_params["reduction_op"] = reduction_op
  flat_src_transforms, flat_dst_transforms = util.split_list(
      flat_args,
      [src_transforms_treedef.num_leaves],
  )
  src_transforms = src_transforms_treedef.unflatten(flat_src_transforms)
  dst_transforms = dst_transforms_treedef.unflatten(flat_dst_transforms)
  return pp.concat([
      pp.text("copy_smem_to_gmem"),
      jax_core.pp_kv_pairs(pp_params.items(), context, settings),
      pp.text(" "),
      state_primitives.pp_ref_transforms(context, src, src_transforms),
      pp.text(" -> "),
      state_primitives.pp_ref_transforms(context, dst, dst_transforms),
  ])


jax_core.pp_eqn_rules[copy_smem_to_gmem_p] = _copy_smem_to_gmem_pp_eqn


@lowering.register_lowering_rule(
    copy_smem_to_gmem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(
    copy_smem_to_gmem_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(
    copy_smem_to_gmem_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(
    copy_smem_to_gmem_p, *gpu_core.WGxWARP_SEMANTICS
)
def _copy_smem_to_gmem_lowering(
    ctx: lowering.LoweringRuleContext,
    src,
    dst,
    *flat_args,
    src_transforms_treedef,
    dst_transforms_treedef,
    has_user_predicate,
    commit_group,
    reduction_op,
):
  if has_user_predicate:
    flat_args, user_predicate = flat_args[:-1], flat_args[-1]
    predicate = lowering._ensure_ir_value(user_predicate, jnp.bool)
  else:
    predicate = None

  flat_src_transforms, flat_dst_transforms = util.split_list(
      flat_args,
      [src_transforms_treedef.num_leaves],
  )
  src_transforms = src_transforms_treedef.unflatten(flat_src_transforms)
  dst_transforms = dst_transforms_treedef.unflatten(flat_dst_transforms)
  handle_transposes = (
      ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup
  )
  src_aval = ctx.avals_in[0]
  assert isinstance(src_aval, state_types.AbstractRef)
  src_transform_avals = src_transforms_treedef.unflatten(
      ctx.avals_in[2 : 2 + src_transforms_treedef.num_leaves]
  )
  dst_transform_avals = dst_transforms_treedef.unflatten(
      ctx.avals_in[
          2
          + src_transforms_treedef.num_leaves : 2
          + src_transforms_treedef.num_leaves
          + dst_transforms_treedef.num_leaves
      ]
  )
  src, src_aval, src_transforms = lowering._handle_transforms(
      ctx, src_aval, src, src_transform_avals, src_transforms,
      handle_transposes=handle_transposes
  )
  copy_params = {
      **_extract_gmem_copy_params(ctx, dst_transforms, dst_transform_avals, supports_multicast=True),
      **_extract_smem_copy_params(src_aval, src_transforms),
  }
  is_scatter = False
  if gmem_slice := copy_params.get("gmem_slice", ()):
    first_idx = gmem_slice[0]
    if isinstance(first_idx, mgpu.FragmentedArray) and first_idx.shape:
      is_scatter = True

  if is_scatter:
    if predicate is not None:
      raise NotImplementedError("Gather/scatter TMA does not support predicates yet.")
    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
      raise ValueError("Gather/scatter operations are not supported in a warp context.")

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    if not is_scatter:
      lane_pred = ctx.module_ctx.single_lane_predicate
      predicate = _andi_maybe_none(predicate, lane_pred)

    ctx.launch_ctx.async_copy(
        src_ref=src,
        dst_ref=dst,
        arrive=commit_group,
        reduction_op=reduction_op,
        oob_mode=OOBFillMode.UNDEFINED,
        **copy_params,
        **(dict(predicate=predicate) if predicate is not None else {}),  # pyrefly: ignore[bad-argument-type]
    )
    return ()

  if gmem_slice := copy_params.get("gmem_slice", ()):
    indices, slice_lengths = _split_gmem_slice(gmem_slice)
  else:
    i32 = ir.IntegerType.get_signless(32)
    slice_lengths = ir.MemRefType(src.type).shape
    indices = [mgpu.utils.c(0, i32)] * len(slice_lengths)

  assert copy_params.get("swizzle") is None
  peer_id = copy_params.get("gmem_peer_id")
  if peer_id is mgpu.GLOBAL_BROADCAST:
    is_global_broadcast = True
    peer_id = None
  else:
    is_global_broadcast = False

  assert not copy_params.get("gmem_transform")
  if reduction_op is not None:
    reduction_op_attr = getattr(mgpu.dialect.TMAReduction, reduction_op.capitalize())
  else:
    reduction_op_attr = None

  mgpu.dialect.async_store(
      src,
      dst,
      indices,
      slice_lengths,
      predicate=predicate,
      commit_group=commit_group,
      reduction_op=reduction_op_attr,
      gmem_peer_id=peer_id,
      is_global_broadcast=is_global_broadcast,
  )
  return ()


def _split_gmem_slice(gmem_slice):
  i32 = ir.IntegerType.get_signless(32)
  indices = []
  slice_lengths = []
  for idx in gmem_slice:
    match idx:
      case slice():
        indices.append(mgpu.utils.c(idx.start, i32))
        slice_lengths.append(idx.stop - idx.start)
      case mgpu.DynamicSlice():
        indices.append(arith_dialect.index_cast(i32, idx.base))  # pyrefly: ignore[bad-argument-type]
        slice_lengths.append(idx.length)
      case ir.Value() if isinstance(idx.type, ir.IndexType):
        indices.append(arith_dialect.index_cast(i32, idx))
        slice_lengths.append(-1)
      case ir.Value() if isinstance(idx.type, ir.IntegerType):
        indices.append(idx)
        slice_lengths.append(-1)
      case ir.Value() if isinstance(idx.type, ir.VectorType):
        indices.append(idx)
        [length] = ir.VectorType(idx.type).shape
        slice_lengths.append(length)
      case _:
        raise NotImplementedError(f"Unsupported GMEM slice: {idx}")
  return indices, slice_lengths


def _extract_gmem_copy_params(
    ctx, transforms, transform_avals, supports_multicast=False
):
  if not transforms:
    return {}
  peer_id = None
  indexers = []
  for transform, transform_aval in zip(
      transforms, transform_avals, strict=True
  ):
    if isinstance(transform, gpu_core.PeerMemRef):
      peer_id = lowering._device_id_to_logical(
          ctx,
          transform.device_id,
          transform.device_id_type,
          transform_aval.device_id,
      )
      peer_id = lowering._ensure_ir_value(peer_id, jnp.int32)
      continue
    elif isinstance(transform, gpu_core.MulticastRef):
      if not supports_multicast:
        raise ValueError(
            "Multicast refs are not supported by this primitive."
        )
      if (mesh_info := ctx.module_ctx.mesh_info) is None:
        raise ValueError(
            "JAX device mesh is required by multicast copies, but not defined."
            " Use jax.set_mesh."
        )
      if set(transform.collective_axes) != set(mesh_info.axis_names):
        raise NotImplementedError(
            "Only collective_axes that include all JAX device mesh  axes are"
            f" supported, but got {transform.collective_axes}. Make sure to"
            f" pass collective_axes={mesh_info.axis_names}"
        )
      peer_id = mgpu.GLOBAL_BROADCAST
      continue
    elif isinstance(transform, indexing.NDIndexer):
      indexers.append(transform)
    else:
      raise NotImplementedError(
          "Non-indexing transforms on GMEM refs are not implemented.")
  if indexers:
    indexer = lowering.merge_indexers(indexers)
    gmem_slice = lowering._ndindexer_indices(indexer, allow_arrays=True)
  else:
    gmem_slice = ()
  return dict(
      gmem_slice=gmem_slice,
      gmem_peer_id=peer_id,
  )


def _extract_smem_copy_params(aval, transforms):
  if not transforms:
    return {}
  # Split off swizzling, if present
  match transforms:
    case [gpu_core.UnswizzleRef(swizzle), *transforms]:
      pass
    case _:
      swizzle = None
  reversed_transforms = pallas_core.undo_transforms(aval, transforms)
  gpu_transforms = tuple(
      gpu_core.to_gpu_transform(t) for t in reversed_transforms
  )
  return dict(
      gmem_transform=gpu_transforms,
      swizzle=swizzle,
  )


def copy_smem_to_gmem(
    src: _Ref,
    dst: _Ref,
    predicate: bool | jax.Array | None = None,
    *,
    commit_group: bool = True,
    reduction_op: mgpu.TMAReductionOp | None = None,
) -> None:
  """Asynchronously copies a SMEM reference to a GMEM reference.

  Args:
    src: The SMEM reference to copy from.
    dst: The GMEM reference to copy to.
    predicate: A boolean indicating whether the copy should be performed. If
      ``None``, the copy is always performed.
    commit_group: If ``True``, this and any previously uncommitted copies are
      committed to a group and can be awaited jointly via
      :func:`jax.experimental.pallas.mosaic_gpu.wait_smem_to_gmem`.
    reduction_op: If set, perform the specified reduction operation when storing
      to GMEM. For example, using ``"add"`` is conceptually equivalent to
      doing ``src += dst``.

  See also:
    :func:`jax.experimental.pallas.mosaic_gpu.wait_smem_to_gmem`
    :func:`jax.experimental.pallas.mosaic_gpu.commit_smem`
  """
  src, src_transforms = state_primitives.get_ref_and_transforms(
      src, None, "copy_smem_to_gmem"
  )
  dst, dst_transforms = state_primitives.get_ref_and_transforms(
      dst, None, "copy_smem_to_gmem"
  )
  flat_src_transforms, src_transforms_treedef = tree_util.tree_flatten(
      src_transforms
  )
  flat_dst_transforms, dst_transforms_treedef = tree_util.tree_flatten(
      dst_transforms
  )
  copy_smem_to_gmem_p.bind(
      src,
      dst,
      *flat_src_transforms,
      *flat_dst_transforms,
      *[] if predicate is None else [predicate],
      src_transforms_treedef=src_transforms_treedef,
      dst_transforms_treedef=dst_transforms_treedef,
      has_user_predicate=predicate is not None,
      commit_group=commit_group,
      reduction_op=reduction_op,
  )
  return None


async_store_smem_p = jax_core.Primitive("async_store_smem")
async_store_smem_p.multiple_results = True


@async_store_smem_p.def_effectful_abstract_eval
def _async_store_smem_abstract_eval(
    src,
    ref,
    barrier,
    cluster_idx,
    *flat_transforms_avals,
    ref_transforms_treedef,
    barrier_transforms_treedef,
    **_,
):
  del cluster_idx  # Unused.
  _check_ref(ref, "ref", gpu_core.SMEM)
  _check_ref(barrier, "barrier", gpu_core.SMEM)
  flat_ref_transforms_avals, flat_barrier_transforms_avals = util.split_list(
      flat_transforms_avals,
      [ref_transforms_treedef.num_leaves],
  )
  ref_transform_avals = ref_transforms_treedef.unflatten(
      flat_ref_transforms_avals
  )
  barrier_transform_avals = barrier_transforms_treedef.unflatten(
      flat_barrier_transforms_avals
  )
  transformed_ref = pallas_core.TransformedRef(ref, ref_transform_avals)
  if src.shape != transformed_ref.shape:
    raise TypeError(
        f"The stored value has shape {src.shape}, but the target reference has"
        f" shape {transformed_ref.shape}"
    )
  if src.dtype != transformed_ref.dtype:
    raise TypeError(
        f"The stored value has dtype {src.dtype}, but the target reference has"
        f" dtype {transformed_ref.dtype}"
    )
  transformed_barrier = pallas_core.TransformedRef(barrier, barrier_transform_avals)
  if transformed_barrier.size != 1:
    raise TypeError(
        "Expected a single barrier, got a barrier reference with shape"
        f" {transformed_barrier.shape}"
    )

  effs = {gpu_core._memory_effect, state.WriteEffect(1)}
  return (), effs


@lowering.register_lowering_rule(async_store_smem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(async_store_smem_p, mgpu.LoweringSemantics.Warpgroup)
def _async_store_smem_lowering(
    ctx: lowering.LoweringRuleContext,
    src,
    ref,
    barrier,
    cluster_idx,
    *flat_transforms,
    ref_transforms_treedef,
    barrier_transforms_treedef,
    cluster_dim,
    optimized,
    atomic,
):
  flat_ref_transforms, flat_barrier_transforms = util.split_list(
      flat_transforms,
      [ref_transforms_treedef.num_leaves],
  )
  flat_ref_transforms_avals, _ = util.split_list(
      ctx.avals_in[4:],
      [ref_transforms_treedef.num_leaves],
  )

  ref_transform_avals = ref_transforms_treedef.unflatten(
      flat_ref_transforms_avals
  )
  ref_aval = ctx.avals_in[1]
  barrier_ref_aval = ctx.avals_in[2]
  assert isinstance(ref_aval, state_types.AbstractRef)
  assert isinstance(barrier_ref_aval, state_types.AbstractRef)

  ref_transforms = ref_transforms_treedef.unflatten(flat_ref_transforms)
  barrier_transforms = barrier_transforms_treedef.unflatten(flat_barrier_transforms)

  ref_smem, _, remaining_ref_transforms = lowering._handle_transforms(
      ctx,
      ref_aval,
      ref,
      ref_transform_avals,
      ref_transforms,
      handle_transposes=True,
  )

  base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)
  if base_index is not None:
    barrier = barrier[base_index]

  cluster_idx_val = lowering._as_index(cluster_idx)
  gpu_cluster_dim = lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, cluster_dim)

  shape = ctx.avals_in[0].shape
  dtype = ctx.avals_in[0].dtype
  if not shape:
    raise NotImplementedError("Scalars are not supported in async_store_smem")

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    if remaining_ref_transforms:
      raise ValueError(f"Unexpected unhandled transforms: {remaining_ref_transforms}")
    assert isinstance(barrier, mgpu.DialectBarrierRef)
    cluster_idx_i32 = arith_dialect.index_cast(
        ir.IntegerType.get_signless(32), cluster_idx_val
    )
    atomic_type = None
    if atomic is not None:
      atomic_type = _atomic_op_type_to_int(AtomicOpType(atomic))
    mgpu.dialect.async_store_smem(
        src,
        ref_smem,
        barrier.as_barrier_memref(),
        gpu_cluster_dim.value,
        cluster_idx_i32,
        atomic_type=atomic_type,
        optimized=optimized,
    )
    return ()

  match remaining_ref_transforms:
    case (gpu_core.UnswizzleRef(swizzle), gpu_core.UntilingTransform(tiling)):
      pass
    case _:
      raise NotImplementedError("async_store_smem requires a tiled and swizzled ref")

  total_bits = math.prod(shape) * dtypes.itemsize_bits(dtype)
  if total_bits % 8:
    raise ValueError(
        f"Can only transfer integer bytes (shape={shape}, dtype={dtype})"
    )
  total_bytes = total_bits // 8
  if total_bytes % WARPGROUP_SIZE:
    raise NotImplementedError(f"Transfer is not a multiple of {WARPGROUP_SIZE} bytes")

  peer_barrier = barrier.remap_to_cluster(gpu_cluster_dim, cluster_idx_val)
  orders_tensor_core = getattr(
      barrier_ref_aval.inner_aval.dtype, "orders_tensor_core", False  # pyrefly: ignore[missing-attribute]
  )
  if orders_tensor_core:
    tx_bytes = total_bytes
    predicate = ctx.module_ctx.single_lane_predicate
    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
      scope = mgpu_utils.ThreadSubset.WARP
    else:
      scope = mgpu_utils.ThreadSubset.WARPGROUP
  else:
    tx_bytes = total_bytes // WARPGROUP_SIZE
    predicate = None
    scope = None
  peer_barrier.arrive_expect_tx(
      tx_bytes, predicate=predicate, tensor_core_order_scope=scope
  )

  lowering._ensure_fa(src, dtype).store_tiled_async(
      ref_smem,
      barrier,
      cluster_dim=gpu_cluster_dim,
      cluster_idx=cluster_idx_val,
      swizzle=swizzle,
      optimized=optimized,
      tiling_rank=len(tiling),
      atomic=atomic,
  )
  return ()


def async_store_smem(
    src: jax.Array,
    ref: _Ref,
    barrier: _Ref,
    *,
    cluster_idx: jax.Array,
    cluster_dim: Hashable | int,
    optimized: bool = True,
    atomic: Literal["add", "max", "min", "and", "or", "xor"] | None = None,
) -> None:
  """Asynchronously stores an array to a SMEM reference within the cluster.

  Args:
    src: The array containing the data to be stored.
    ref: The SMEM reference to store to.
    barrier: The barrier to update when the copy has completed (in destination block).
    cluster_idx: The index of the target cluster block within cluster_dim.
    cluster_dim: The cluster axis of cluster_idx.
    optimized: If True, the store is guaranteed not to cause any bank conflicts.
    atomic: The reduction operation to apply instead of overwriting the data.
  """
  ref, ref_transforms = state_primitives.get_ref_and_transforms(
      ref, None, "async_store_smem"
  )
  barrier, barrier_transforms = state_primitives.get_ref_and_transforms(
      barrier, None, "async_store_smem"
  )
  flat_ref_transforms, ref_transforms_treedef = tree_util.tree_flatten(
      ref_transforms
  )
  flat_barrier_transforms, barrier_transforms_treedef = tree_util.tree_flatten(
      barrier_transforms
  )
  async_store_smem_p.bind(
      src,
      ref,
      barrier,
      cluster_idx,
      *flat_ref_transforms,
      *flat_barrier_transforms,
      ref_transforms_treedef=ref_transforms_treedef,
      barrier_transforms_treedef=barrier_transforms_treedef,
      cluster_dim=cluster_dim,
      optimized=optimized,
      atomic=atomic,
  )
  return None


copy_gmem_to_smem_p = jax_core.Primitive("copy_gmem_to_smem")
copy_gmem_to_smem_p.multiple_results = True


@copy_gmem_to_smem_p.def_effectful_abstract_eval
def _copy_gmem_to_smem_abstract_eval(src, dst, *args, has_barrier, **params):
  _check_ref(src, "src", gpu_core.GMEM)
  _check_ref(dst, "dst", gpu_core.SMEM)
  if has_barrier:
    barrier, *args = args
    _check_ref(barrier, "barrier", gpu_core.SMEM)
  src_transforms_treedef = params["src_transforms_treedef"]
  dst_transforms_treedef = params["dst_transforms_treedef"]
  flat_src_transforms, flat_dst_transforms = util.split_list(
      args,
      [
          src_transforms_treedef.num_leaves,
          dst_transforms_treedef.num_leaves,
      ],
  )[:2]
  src_transforms = src_transforms_treedef.unflatten(flat_src_transforms)
  dst_transforms = dst_transforms_treedef.unflatten(flat_dst_transforms)

  src_ref = pallas_core.TransformedRef(src, src_transforms)
  dst_ref = pallas_core.TransformedRef(dst, dst_transforms)

  if src_ref.dtype != dst_ref.dtype:
    raise ValueError(
        "Expected dtypes to match but src had type "
        f"{src_ref.dtype} and dst had type {dst_ref.dtype}"
    )

  src_shape = src_ref.shape
  dst_shape = dst_ref.shape

  leader_tracked = params.get("leader_tracked", None)
  partition_axis = None
  if isinstance(leader_tracked, CopyPartition.PARTITIONED):
    partition_axis = leader_tracked.axis

  collective_axes = params.get("collective_axes", None)
  axis_size = 1
  if partition_axis is not None and collective_axes is not None:
    if isinstance(collective_axes, str):
      collective_axes = (collective_axes,)
    axis_size = math.prod(
        jax_core.get_axis_env().axis_size(axis) for axis in collective_axes
    )

  expected_src_shape = tuple(
      d * axis_size if i == partition_axis else d
      for i, d in enumerate(dst_shape)
  )

  if src_shape != expected_src_shape:
    if partition_axis is not None:
      partitioned_desc = (
          f"Dim {partition_axis} is partitioned over an axis of size {axis_size}."
      )
    else:
      partitioned_desc = ""
    raise ValueError(
        f"Expected source shape to be {expected_src_shape}, but got {src_shape}."
        f" {partitioned_desc}"
    )

  return (), {state.ReadEffect(0), state.WriteEffect(1)}


def _copy_gmem_to_smem_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  if eqn.params["has_barrier"]:
    src, dst, barrier, *flat_args = eqn.invars
  else:
    src, dst, *flat_args = eqn.invars
    barrier = None
  src_transforms_treedef = eqn.params["src_transforms_treedef"]
  dst_transforms_treedef = eqn.params["dst_transforms_treedef"]
  barrier_transforms_treedef = eqn.params["barrier_transforms_treedef"]
  pp_params = {}
  if collective_axes := eqn.params["collective_axes"]:
    pp_params["collective_axes"] = collective_axes
  if eqn.params["has_user_predicate"]:
    flat_args, user_predicate = flat_args[:-1], flat_args[-1]
    pp_params["user_predicate"] = user_predicate.pretty_print(context)
  flat_src_transforms, flat_dst_transforms, flat_barrier_transforms = (
      util.split_list(
          flat_args,
          [
              src_transforms_treedef.num_leaves,
              dst_transforms_treedef.num_leaves,
          ],
      )
  )
  src_transforms = src_transforms_treedef.unflatten(flat_src_transforms)
  dst_transforms = dst_transforms_treedef.unflatten(flat_dst_transforms)
  parts = [
      pp.text("copy_gmem_to_smem"),
      jax_core.pp_kv_pairs(pp_params.items(), context, settings),
      pp.text(" "),
      state_primitives.pp_ref_transforms(context, src, src_transforms),
      pp.text(" -> "),
      state_primitives.pp_ref_transforms(context, dst, dst_transforms),
  ]
  if barrier is not None:
    barrier_transforms = barrier_transforms_treedef.unflatten(
        flat_barrier_transforms
    )
    parts.extend([
        pp.text(" using "),
        state_primitives.pp_ref_transforms(
            context, barrier, barrier_transforms
        ),
    ])
  return pp.concat(parts)


jax_core.pp_eqn_rules[copy_gmem_to_smem_p] = _copy_gmem_to_smem_pp_eqn


@lowering.register_lowering_rule(
    copy_gmem_to_smem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(
    copy_gmem_to_smem_p, *gpu_core.LANExWARP_SEMANTICS
)
@lowering.register_lowering_rule(
    copy_gmem_to_smem_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(
    copy_gmem_to_smem_p, *gpu_core.WGxWARP_SEMANTICS
)
def _copy_gmem_to_smem_lowering(
    ctx: lowering.LoweringRuleContext,
    src,
    dst,
    *flat_args,
    src_transforms_treedef,
    dst_transforms_treedef,
    barrier_transforms_treedef,
    collective_axes,
    leader_tracked,
    oob_mode,
    has_barrier,
    has_user_predicate,
):
  if has_barrier:
    barrier, *flat_args = flat_args
    num_non_transform_args = 3  # src, dst, barrier
  else:
    barrier = None
    num_non_transform_args = 2  # src, dst
  del has_barrier
  if has_user_predicate:
    flat_args, user_predicate = flat_args[:-1], flat_args[-1]
    predicate = lowering._ensure_ir_value(user_predicate, jnp.bool)
  else:
    predicate = None

  flat_src_transforms, flat_dst_transforms, flat_barrier_transforms = (
      util.split_list(
          flat_args,
          [
              src_transforms_treedef.num_leaves,
              dst_transforms_treedef.num_leaves,
          ],
      )
  )
  flat_src_transforms_avals, flat_dst_transforms_avals, _ = (
      util.split_list(
          ctx.avals_in[num_non_transform_args:],
          [
              src_transforms_treedef.num_leaves,
              dst_transforms_treedef.num_leaves,
          ],
      )
  )
  src_transform_avals = src_transforms_treedef.unflatten(
      flat_src_transforms_avals
  )
  src_ref_aval = ctx.avals_in[0]
  dst_ref_aval = ctx.avals_in[1]
  src_transforms = src_transforms_treedef.unflatten(flat_src_transforms)
  dst_transforms = dst_transforms_treedef.unflatten(flat_dst_transforms)
  handle_transposes = (
      ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup
  )
  assert isinstance(src_ref_aval, state_types.AbstractRef)
  assert isinstance(dst_ref_aval, state_types.AbstractRef)

  dst_transform_avals = dst_transforms_treedef.unflatten(
      flat_dst_transforms_avals)
  assert len(dst_transform_avals) == len(dst_transforms)
  dst, dst_ref_aval, dst_transforms = lowering._handle_transforms(
      ctx, dst_ref_aval, dst, dst_transform_avals, dst_transforms,
      handle_transposes=handle_transposes)

  copy_params = {
      **_extract_smem_copy_params(dst_ref_aval, dst_transforms),
      **_extract_gmem_copy_params(ctx, src_transforms, src_transform_avals),
  }
  orders_tensor_core = False
  if barrier is not None:
    barrier_ref_aval = ctx.avals_in[2]
    assert isinstance(barrier_ref_aval, state_types.AbstractRef)
    orders_tensor_core = getattr(
        barrier_ref_aval.inner_aval.dtype, "orders_tensor_core", False  # pyrefly: ignore[missing-attribute]
    )
    base_index = _get_barrier_base_index(
        barrier_ref_aval,
        barrier_transforms_treedef.unflatten(flat_barrier_transforms),
    )
    if base_index is not None:
      barrier = barrier[base_index]
  collective = None
  if collective_axes is not None:
    collective = tuple(
        lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, axis)
        for axis in collective_axes
    )

  is_leader_tracked_copy = collective and leader_tracked is not None
  dst_ty = ir.MemRefType(dst.type)
  bits = math.prod(dst_ty.shape) * mgpu.bitwidth(dst_ty.element_type)
  if bits % 8:
    raise ValueError(
        f"Can only transfer integer bytes (shape={dst_ty.shape},"
        f" dtype={dst_ty.element_type})"
    )
  bytes = bits // 8

  if is_leader_tracked_copy:
    # Leader receives the completion messages from both CTAs.
    bytes *= 2
    if len(collective) != 1:
      raise ValueError(
          f"Expected exactly one collective axis, got {collective_axes=}"
      )
    if (
        collective != (gpu_dialect.Dimension.x,)
        or ctx.launch_ctx.cluster_size[0] != 2
    ):
      expected_axis = ctx.module_ctx.axis_names.cluster[-1]
      row_major_cluster_size = tuple(ctx.launch_ctx.cluster_size[::-1])
      raise NotImplementedError(
          "Partitioned loads only supported along the last cluster dimension "
          f"{expected_axis} if it has size 2. Got {collective_axes=} and "
          f"cluster size {row_major_cluster_size}."
      )

  # TMA is only available on Hopper and newer. On older architectures we fall
  # back to the cp.async implementation.
  if is_cp_async := mgpu.utils.get_arch().major < 9:
    if barrier is not None:
      raise ValueError(
          "copy_gmem_to_smem with a barrier is only supported Hopper and newer"
          " GPUs, which use the TMA implementation"
      )
    if collective_axes is not None:
      raise ValueError("Only the TMA implementation supports collective copies")
    if leader_tracked is not None:
      raise ValueError(
          "Only the TMA implementation supports leader_tracked copies"
      )
    # cp.async does not predicate out-of-bounds accesses, so the caller has to
    # guarantee that the copy stays in bounds.
    if oob_mode != OOBFillMode.PROMISE_IN_BOUNDS:
      raise ValueError(
          "The cp.async implementation only supports "
          "oob_mode=OOBFillMode.PROMISE_IN_BOUNDS"
      )
    if has_user_predicate:
      raise NotImplementedError(
          "The cp.async implementation does not support user-defined predicates"
      )
  else:
    if oob_mode is None:
      oob_mode = OOBFillMode.ZEROS

    if barrier is None:
      raise ValueError(
          "copy_gmem_to_smem without a barrier is only supported on pre-Hopper"
          " GPUs, which use the cp.async implementation"
      )

  i32 = ir.IntegerType.get_signless(32)
  arrive_predicate = predicate
  if is_leader_tracked_copy:
    first_block = arith_dialect.cmpi(
        arith_dialect.CmpIPredicate.eq,
        mgpu.utils.cluster_idx(collective[0]),
        mgpu.c(0, ir.IndexType.get()),
    )
    arrive_predicate = _andi_maybe_none(arrive_predicate, first_block)

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    if (
        ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warpgroup
        and ctx.module_ctx.auto_barriers
    ):
      mgpu.warpgroup_barrier()  # Make sure all reads have completed.

    if not is_cp_async:
      assert barrier is not None
      if not orders_tensor_core and bytes % WARPGROUP_SIZE:
        raise NotImplementedError(
            "Only copies transferring a number of bytes divisible by the"
            f" warpgroup size are supported. Got {bytes=} but warpgroup size is"
            f" {WARPGROUP_SIZE}"
        )
      if orders_tensor_core:
        if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
          scope = mgpu_utils.ThreadSubset.WARP
        else:
          scope = mgpu_utils.ThreadSubset.WARPGROUP
        barrier.arrive_expect_tx(
            bytes,
            predicate=_andi_maybe_none(
                arrive_predicate, ctx.module_ctx.single_lane_predicate
            ),
            tensor_core_order_scope=scope,
        )
      elif ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warpgroup:
        # We arrive uniformly from each thread in the WG, so we need to divide
        # the number of bytes by the number of threads in the WG.
        # TODO: apaszke - Relax this. We can just select the WG leader and
        # have it arrive with the whole transfer size, while everyone else
        # arrives with 0.
        # But we should continue using this scheme as it's likely to be
        # faster.
        bytes //= WARPGROUP_SIZE
        barrier.arrive_expect_tx(bytes, predicate=arrive_predicate)
      else:
        # In Warp-level lowering, we arrive on each CUDA thread in a warp, but
        # the barrier still expects a full 128 arrivals so we arrive 4 times
        # on each CUDA thread instead.
        # TODO(justinfu): The arrival counts are wrong if called outside of a
        # single warp. Figure out how to guard against this in user code.
        bytes = bytes // WARP_SIZE
        if arrive_predicate is not None:
          arrive_ctx = mgpu.when(arrive_predicate)
        else:
          arrive_ctx = contextlib.nullcontext()
        with arrive_ctx:
          barrier.arrive(arrival_count=3, can_complete=False)
          barrier.arrive_expect_tx(bytes)

    lane_pred = ctx.module_ctx.single_lane_predicate
    predicate = _andi_maybe_none(predicate, lane_pred)
    predicate_kwarg = (
        {}
        if is_cp_async
        else dict(predicate=predicate)
    )
    # Gathers are a warpgroup-level collective and can't take a predicate.
    if gmem_slice := copy_params.get("gmem_slice", ()):
      first_idx = gmem_slice[0]
      if isinstance(first_idx, mgpu.FragmentedArray) and first_idx.shape:
        if has_user_predicate:
          raise NotImplementedError("Gather/scatter TMA does not support predicates yet.")
        predicate_kwarg = {}
    ctx.launch_ctx.async_copy(
        src_ref=src,
        dst_ref=dst,
        barrier=barrier,
        arrive=False,
        collective=collective,
        leader_tracked=leader_tracked,
        oob_mode=oob_mode,
        implementation=(
            AsyncCopyImplementation.CP_ASYNC
            if is_cp_async
            else AsyncCopyImplementation.TMA
        ),
        **copy_params,
        **predicate_kwarg,  # pyrefly: ignore[bad-argument-type]
    )
    return ()
  if "gmem_slice" not in copy_params or not copy_params["gmem_slice"]:
    slice_lengths = ir.MemRefType(src.type).shape
    indices = [mgpu.utils.c(0, i32)] * len(slice_lengths)
  else:
    indices, slice_lengths = _split_gmem_slice(copy_params["gmem_slice"])
  assert copy_params.get("swizzle") is None
  assert not copy_params.get("gmem_transform")

  if is_cp_async:
    mgpu.dialect.async_load(
        src,
        dst,
        barrier=None,  # pyrefly: ignore[bad-argument-type]
        indices=indices,
        slice_lengths=slice_lengths,
        collective=ir.ArrayAttr.get([]),
        predicate=predicate,
        oob_fill_mode=ir.IntegerAttr.get(i32, oob_mode.value),
    )
    return ()

  match leader_tracked:
    case CopyPartition.REPLICATED:
      leader_tracked_attr = mgpu.dialect.CopyReplicatedAttr.get()
    case CopyPartition.PARTITIONED(axis):
      leader_tracked_attr = mgpu.dialect.CopyPartitionedAttr.get(axis)
    case _:
      leader_tracked_attr = None

  assert barrier is not None
  barrier_ref = barrier.as_barrier_memref()

  if arrive_predicate is not None:
    arrive_ctx = mgpu.when(arrive_predicate)
  else:
    arrive_ctx = contextlib.nullcontext()

  with arrive_ctx:
    mgpu.dialect.arrive_expect_tx(barrier_ref, mgpu.c(bytes, i32))

  mgpu.dialect.async_load(
      src,
      dst,
      indices,  # pyrefly: ignore[bad-argument-type]
      slice_lengths,  # pyrefly: ignore[bad-argument-type]
      barrier=barrier_ref,  # pyrefly: ignore[bad-keyword-argument]
      predicate=predicate,
      collective=ir.ArrayAttr.get(  # pyrefly: ignore[bad-keyword-argument]
          [ir.IntegerAttr.get(i32, axis) for axis in collective or []]
      ),
      leader_tracked=leader_tracked_attr,
      oob_fill_mode=ir.IntegerAttr.get(i32, oob_mode.value),
      gmem_peer_id=copy_params.get("gmem_peer_id"),
  )
  return ()

def copy_gmem_to_smem(
    src: _Ref,
    dst: _Ref,
    barrier: _Ref | None = None,
    *,
    collective_axes: str | tuple[str, ...] | None = None,
    leader_tracked: CopyPartition | None = None,
    oob_mode: OOBFillMode | None = None,
    predicate: bool | jax.Array | None = None,
) -> None:
  """Asynchronously copies a GMEM reference to a SMEM reference.

  The underlying copy implementation is selected automatically based on the
  target architecture: TMA on Hopper and newer, and ``cp.async`` on
  older architectures.

  When ``collective_axes`` is specified, the copy involves multiple CUDA blocks
  in the cluster. The value of ``leader_tracked`` determines the behavior:

  * ``None`` (**multicast**): all CUDA blocks sharing the same index along the
    collective axes receive the same data from ``src``.
  * ``CopyPartition.PARTITIONED(axis)``: each CUDA block in the cluster
    receives a ``transfer_size // cluster_size`` tile of ``src``. E.g. for
    ``src`` of shape ``(256, 256)`` with cluster size 2 along axis 0: block 0
    gets ``src[0:128, :]``, block 1 gets ``src[128:256, :]``.
  * ``CopyPartition.REPLICATED``: all CUDA blocks in the cluster load the same
    data, but only the first block tracks progress via barrier arrivals.

  .. note::

    For leader-tracked copies, only the first CUDA block in the cluster arrives
    on the barrier. If other blocks need to consume the copied data, an
    additional cluster barrier is necessary to ensure all blocks have
    finished the copy.

  Args:
    src: The source Ref. Must be in GMEM.
    dst: The destination Ref. Must be in SMEM.
    barrier: The barrier to use for tracking completion of the copy. Required on
      Hopper and newer GPUs, which use the TMA implementation, and must be
      omitted on pre-Hopper GPUs, which use the ``cp.async`` implementation.
      When omitted, completion of the copy has to be awaited with
      :func:`jax.experimental.pallas.mosaic_gpu.wait_gmem_to_smem` instead of
      :func:`jax.experimental.pallas.mosaic_gpu.barrier_wait`.
    collective_axes: The collective axes to use for the copy. Only a single
      collective axis is supported when ``leader_tracked`` is specified (but
      the axis can be composite). Only supported with the TMA implementation
      (Hopper and newer).
    leader_tracked: If specified, only the leader block in the cluster will
      observe the completion of the copy. If
      ``CopyPartition.PARTITIONED(axis)``, performs a partitioned collective
      copy along the given axis. If ``CopyPartition.REPLICATED``, all blocks
      load the same data. Only supported with the TMA implementation (Hopper
      and newer).
    oob_mode: The optional out-of-bounds fill mode. Can be
      ``OOBFillMode.UNDEFINED``, ``OOBFillMode.PROMISE_IN_BOUNDS`` or
      ``OOBFillMode.ZEROS`` for the TMA implementation, and only
      ``OOBFillMode.PROMISE_IN_BOUNDS`` for the ``cp.async`` one.
    predicate: A boolean indicating whether the copy should be performed. If
      ``None``, the copy is always performed. If ``False``, neither the copy nor
      the barrier arrival is performed.

  See also:
    :func:`jax.experimental.pallas.mosaic_gpu.barrier_arrive`
    :func:`jax.experimental.pallas.mosaic_gpu.barrier_wait`
    :func:`jax.experimental.pallas.mosaic_gpu.wait_gmem_to_smem`
  """
  src, src_transforms = state_primitives.get_ref_and_transforms(
      src, None, "copy_gmem_to_smem"
  )
  dst, dst_transforms = state_primitives.get_ref_and_transforms(
      dst, None, "copy_gmem_to_smem"
  )
  flat_src_transforms, src_transforms_treedef = tree_util.tree_flatten(
      src_transforms
  )
  flat_dst_transforms, dst_transforms_treedef = tree_util.tree_flatten(
      dst_transforms
  )
  has_barrier = barrier is not None
  if has_barrier:
    barrier, barrier_transforms = state_primitives.get_ref_and_transforms(
        barrier, None, "copy_gmem_to_smem"
    )
    barrier_operands = [barrier]
  else:
    barrier_transforms = []
    barrier_operands = []
  flat_barrier_transforms, barrier_transforms_treedef = tree_util.tree_flatten(
      barrier_transforms
  )
  if isinstance(collective_axes, str):
    collective_axes = (collective_axes,)
  if leader_tracked is not None and collective_axes is None:
    raise ValueError(
        "`collective_axes` must be specified when `leader_tracked` is set"
    )
  copy_gmem_to_smem_p.bind(
      src,
      dst,
      *barrier_operands,
      *flat_src_transforms,
      *flat_dst_transforms,
      *flat_barrier_transforms,
      *[] if predicate is None else [predicate],
      src_transforms_treedef=src_transforms_treedef,
      dst_transforms_treedef=dst_transforms_treedef,
      barrier_transforms_treedef=barrier_transforms_treedef,
      collective_axes=collective_axes,
      leader_tracked=leader_tracked,
      oob_mode=oob_mode,
      has_barrier=has_barrier,
      has_user_predicate=predicate is not None,
  )
  return None


wait_gmem_to_smem_p = jax_core.Primitive("wait_gmem_to_smem")
wait_gmem_to_smem_p.multiple_results = True


@wait_gmem_to_smem_p.def_effectful_abstract_eval
def _wait_gmem_to_smem_abstract_eval(*, n: int):
  if not isinstance(n, int) or n < 0:
    raise ValueError(f"n must be a non-negative integer, got {n}")
  return (), {gpu_core._memory_effect}


@lowering.register_lowering_rule(
    wait_gmem_to_smem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(
    wait_gmem_to_smem_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(
    wait_gmem_to_smem_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(
    wait_gmem_to_smem_p, *gpu_core.WGxWARP_SEMANTICS)
def _wait_gmem_to_smem_lowering(ctx: lowering.LoweringRuleContext, *, n: int):
  if mgpu.utils.get_arch().major >= 9:
    raise ValueError(
        "wait_gmem_to_smem is only supported on pre-Hopper GPUs, which use"
        " cp.async for GMEM->SMEM copies."
    )
  ctx.launch_ctx.await_cp_async_copy(allow_groups=n)
  return ()


def wait_gmem_to_smem(n: int) -> None:
  """Waits until at most ``n`` ``cp.async`` GMEM->SMEM copies are in flight.

  .. note:: This waiting mechanism is only supported on pre-Hopper GPUs, which
            use the ``cp.async`` implementation of
            :func:`jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem`.

  Args:
    n: The maximum number of copies allowed to remain in flight.
  """
  wait_gmem_to_smem_p.bind(n=n)


async_prefetch_p = jax_core.Primitive("async_prefetch")
async_prefetch_p.multiple_results = True

@async_prefetch_p.def_effectful_abstract_eval
def _async_prefetch_abstract_eval(ref, *args, **params):
  del args, params  # Unused.
  _check_ref(ref, "ref", gpu_core.GMEM)
  return (), {state.ReadEffect(0)}


def _async_prefetch_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  ref, *flat_args = eqn.invars
  pp_params = {}
  if eqn.params["has_user_predicate"]:
    *flat_args, user_predicate = flat_args
    pp_params["user_predicate"] = user_predicate.pretty_print(context)
  ref_transforms_treedef = eqn.params["ref_transforms_treedef"]
  ref_transforms = ref_transforms_treedef.unflatten(flat_args)
  return pp.concat([
      pp.text("async_prefetch"),
      jax_core.pp_kv_pairs(pp_params.items(), context, settings),
      pp.text(" "),
      state_primitives.pp_ref_transforms(context, ref, ref_transforms),
  ])


jax_core.pp_eqn_rules[async_prefetch_p] = _async_prefetch_pp_eqn


@lowering.register_lowering_rule(async_prefetch_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(async_prefetch_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(async_prefetch_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(async_prefetch_p, *gpu_core.WGxWARP_SEMANTICS)
def _async_prefetch_lowering(
    ctx: lowering.LoweringRuleContext,
    ref,
    *flat_args,
    ref_transforms_treedef,
    collective_axes,
    leader_tracked,
    has_user_predicate: bool = False,
):
  if has_user_predicate:
    *flat_args, user_predicate = flat_args
    predicate = lowering._ensure_ir_value(user_predicate, jnp.bool)  # pylint: disable=protected-access
  else:
    predicate = None

  ref_transforms = ref_transforms_treedef.unflatten(flat_args)
  ref_transform_avals = ref_transforms_treedef.unflatten(
      ctx.avals_in[1 : 1 + ref_transforms_treedef.num_leaves]
  )
  copy_params = _extract_gmem_copy_params(
      ctx, ref_transforms, ref_transform_avals
  )
  collective = None
  if collective_axes is not None:
    collective = tuple(
        lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, axis)
        for axis in collective_axes
    )

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    lane_pred = ctx.module_ctx.single_lane_predicate
    predicate = _andi_maybe_none(predicate, lane_pred)
    predicate_kwarg: dict[str, Any] = dict(predicate=predicate)
    if gmem_slice := copy_params.get("gmem_slice", ()):
      first_idx = gmem_slice[0]
      # Gathers are a warpgroup-level collective and can't take a predicate.
      if isinstance(first_idx, mgpu.FragmentedArray) and first_idx.shape:
        if has_user_predicate:
          raise NotImplementedError(
              "Gather/scatter TMA does not support predicates yet."
          )
        predicate_kwarg = {}

    ctx.launch_ctx.async_prefetch(
        gmem_ref=ref,
        collective=collective,
        leader_tracked=leader_tracked,
        **copy_params,
        **predicate_kwarg,
    )
    return ()

  if "gmem_slice" not in copy_params:
    i32 = ir.IntegerType.get_signless(32)
    slice_lengths = ir.MemRefType(ref.type).shape
    indices = [mgpu.utils.c(0, i32)] * len(slice_lengths)
  else:
    indices, slice_lengths = _split_gmem_slice(copy_params["gmem_slice"])
  assert copy_params.get("swizzle") is None
  assert not copy_params.get("gmem_transform")
  if copy_params.get("gmem_peer_id", None) is not None:
    raise NotImplementedError(
        "GMEM refs with peer ids are not supported in warpgroup lowering."
    )
  mgpu.dialect.async_prefetch(
      ref,
      indices,
      slice_lengths,
      collective=ir.ArrayAttr.get([]),
      predicate=predicate,
  )
  return ()


def async_prefetch(
    ref: _Ref,
    *,
    collective_axes: str | tuple[str, ...] | None = None,
    leader_tracked: CopyPartition | None = None,
    predicate: bool | jax.Array | None = None,
) -> None:
  """Asynchronously prefetches a GMEM reference to the L2 cache.

  If collective_axes is specified, each CUDA block only prefetches a part of
  the ``ref``, with other parts covered by blocks that share the same index
  along the collective axis.

  Specifying leader_tracked and collective_axes doesn't change the semantics
  of the prefetch, but if it's followed by a GMEM to SMEM copy, then it allows
  us to reuse the same TMA descriptor.

  Args:
    ref: The source Ref. Must be in GMEM.
    collective_axes: The collective axes to use for the prefetch.
    leader_tracked: The partitioning to use for the prefetch.
    predicate: A boolean indicating whether the prefetch should be performed. If
      ``None``, the prefetch is always performed.
  """
  ref, ref_transforms = state_primitives.get_ref_and_transforms(
      ref, None, "async_prefetch"
  )
  flat_ref_transforms, ref_transforms_treedef = tree_util.tree_flatten(
      ref_transforms
  )
  if isinstance(collective_axes, str):
    collective_axes = (collective_axes,)
  async_prefetch_p.bind(
      ref,
      *flat_ref_transforms,
      *() if predicate is None else (predicate,),
      ref_transforms_treedef=ref_transforms_treedef,
      collective_axes=collective_axes,
      leader_tracked=leader_tracked,
      has_user_predicate=predicate is not None,
  )
  return None


def _get_barrier_base_index(aval, transforms) -> ir.Value | None:
  if not transforms:
    return None
  strides = list(pallas_utils.strides_from_shape(aval.shape))
  base_index: ir.Value | None = None
  while transforms:
    match transforms:
      case [indexing.NDIndexer() as indexer, *transforms]:
        num_int_idxs = 0
        for i, (idx, stride) in enumerate(zip(indexer.indices, strides[:])):
          if isinstance(idx, indexing.Slice):
            if idx.stride != 1:
              raise NotImplementedError(
                  "Barrier does not support slice with `stride != 1`"
              )
            idx = idx.start
          else:
            # This dimension is absent for any corresponding `NDIndexer`s, so
            # we remove the corresponding stride.
            strides.pop(i - num_int_idxs)
            num_int_idxs += 1

          if isinstance(
              idx, (int, ir.Value, mgpu.FragmentedArray, literals.TypedNdArray)
          ):
            idx = lowering._as_index(idx)
          else:
            raise ValueError(
                "Barrier can only be indexed with integers or slices, got"
                f" {idx}"
            )

          idx = arith_dialect.muli(idx, lowering._as_index(stride))
          if base_index is None:
            base_index = idx
          else:
            base_index = arith_dialect.addi(base_index, idx)
      case _:
        raise ValueError("Barrier does not support arbitrary transforms")
  return base_index


barrier_arrive_p = jax_core.Primitive("barrier_arrive")
barrier_arrive_p.multiple_results = True


@barrier_arrive_p.def_effectful_abstract_eval
def _barrier_arrive_abstract_eval(barrier, *args, **params):
  del args, params  # Unused.
  _check_ref(barrier, "barrier", gpu_core.SMEM)
  return (), {gpu_core._memory_effect}


def _barrier_arrive_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  barrier, *flat_args = eqn.invars
  pp_params = {}
  if eqn.params["has_user_predicate"]:
    *flat_args, user_predicate = flat_args
    pp_params["user_predicate"] = user_predicate.pretty_print(context)
  transforms_treedef = eqn.params["transforms_treedef"]
  transforms = transforms_treedef.unflatten(flat_args)
  return pp.concat([
      pp.text("barrier_arrive"),
      jax_core.pp_kv_pairs(pp_params.items(), context, settings),
      pp.text(" "),
      state_primitives.pp_ref_transforms(context, barrier, transforms),
  ])


jax_core.pp_eqn_rules[barrier_arrive_p] = _barrier_arrive_pp_eqn


@lowering.register_lowering_rule(barrier_arrive_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(barrier_arrive_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(barrier_arrive_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(barrier_arrive_p, *gpu_core.WGxWARP_SEMANTICS)
def _barrier_arrive_lowering(
    ctx: lowering.LoweringRuleContext,
    barrier,
    *flat_args,
    transforms_treedef,
    has_user_predicate: bool = False,
):
  if has_user_predicate:
    *flat_args, user_predicate = flat_args
    predicate = lowering._ensure_ir_value(user_predicate, jnp.bool)  # pylint: disable=protected-access
  else:
    predicate = None

  transforms = transforms_treedef.unflatten(flat_args)
  barrier_aval = ctx.avals_in[0]
  assert isinstance(barrier_aval, state_types.AbstractRef)
  base_index = _get_barrier_base_index(barrier_aval, transforms)
  if base_index is not None:
    barrier = barrier[base_index]
  sem_dtype = barrier_aval.inner_aval.dtype  # pyrefly: ignore[missing-attribute]
  orders_tensor_core = getattr(sem_dtype, "orders_tensor_core", False)

  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    scope = mgpu_utils.ThreadSubset.WARP
  else:
    scope = mgpu_utils.ThreadSubset.WARPGROUP

  if isinstance(barrier, mgpu.CollectiveBarrierRef):
    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
      raise NotImplementedError(
          "Arriving on a collective barrier is not supported in a warp context"
      )
    barrier.arrive(orders_tensor_core=orders_tensor_core, predicate=predicate)
  elif ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    barrier.arrive(orders_tensor_core=orders_tensor_core, predicate=predicate)
  else:
    if scope == mgpu_utils.ThreadSubset.WARP and not orders_tensor_core:
      arrival_count = 4
    else:
      arrival_count = 1

    pred = ctx.module_ctx.single_lane_predicate if orders_tensor_core else None
    barrier.arrive(
        arrival_count=arrival_count,
        orders_tensor_core=orders_tensor_core,
        predicate=_andi_maybe_none(predicate, pred),
        scope=scope,
    )
  return ()


def barrier_arrive(
    barrier: state.AbstractRef,
    *,
    predicate: bool | jax.Array | None = None,
) -> None:
  """Arrives at the given barrier."""
  barrier, transforms = state_primitives.get_ref_and_transforms(
      barrier, None, "barrier_arrive"
  )
  flat_transforms, transforms_treedef = tree_util.tree_flatten(transforms)
  barrier_arrive_p.bind(
      barrier,
      *flat_transforms,
      *() if predicate is None else (predicate,),
      transforms_treedef=transforms_treedef,
      has_user_predicate=predicate is not None,
  )


barrier_arrive_and_wait_p = jax_core.Primitive("barrier_arrive_and_wait")
barrier_arrive_and_wait_p.multiple_results = True


@barrier_arrive_and_wait_p.def_effectful_abstract_eval
def _barrier_arrive_and_wait_abstract_eval(barrier, *args, **params):
  del args, params  # Unused.
  _check_ref(barrier, "barrier", gpu_core.SMEM)
  return (), {gpu_core._memory_effect}


def _barrier_arrive_and_wait_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  del settings
  barrier, *flat_transforms = eqn.invars
  transforms_treedef = eqn.params["transforms_treedef"]
  transforms = transforms_treedef.unflatten(flat_transforms)
  return pp.concat([
      pp.text("barrier_arrive_and_wait"),
      pp.text(" "),
      state_primitives.pp_ref_transforms(context, barrier, transforms),
  ])


jax_core.pp_eqn_rules[barrier_arrive_and_wait_p] = (
    _barrier_arrive_and_wait_pp_eqn
)


@lowering.register_lowering_rule(
    barrier_arrive_and_wait_p, mgpu.LoweringSemantics.Lane
)
@lowering.register_lowering_rule(
    barrier_arrive_and_wait_p, *gpu_core.LANExWARP_SEMANTICS
)
@lowering.register_lowering_rule(
    barrier_arrive_and_wait_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(
    barrier_arrive_and_wait_p, *gpu_core.WGxWARP_SEMANTICS
)
def _barrier_arrive_and_wait_lowering(
    ctx: lowering.LoweringRuleContext,
    barrier,
    *flat_transforms,
    transforms_treedef,
):
  _barrier_arrive_lowering(
      ctx, barrier, *flat_transforms, transforms_treedef=transforms_treedef
  )
  _barrier_wait_lowering(
      ctx, barrier, *flat_transforms, transforms_treedef=transforms_treedef
  )
  return ()


def barrier_arrive_and_wait(barrier: state.AbstractRef) -> None:
  """Arrives at and waits on the given barrier."""
  barrier, transforms = state_primitives.get_ref_and_transforms(
      barrier, None, "barrier_arrive_and_wait"
  )
  flat_transforms, transforms_treedef = tree_util.tree_flatten(transforms)
  barrier_arrive_and_wait_p.bind(
      barrier, *flat_transforms, transforms_treedef=transforms_treedef
  )


barrier_test_p = jax_core.Primitive("barrier_test")
barrier_test_p.multiple_results = False


@barrier_test_p.def_effectful_abstract_eval
def _barrier_test_abstract_eval(barrier, *args, **params):
  _check_ref(barrier, "barrier", gpu_core.SMEM)
  del args, params  # Unused.
  return jax_core.ShapedArray((), bool), {gpu_core._memory_effect}


def _barrier_test_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  del settings
  barrier, *flat_transforms = eqn.invars
  transforms_treedef = eqn.params["transforms_treedef"]
  transforms = transforms_treedef.unflatten(flat_transforms)
  return pp.concat([
      pp.text("barrier_test"),
      pp.text(" "),
      state_primitives.pp_ref_transforms(context, barrier, transforms),
  ])


jax_core.pp_eqn_rules[barrier_test_p] = _barrier_test_pp_eqn


@lowering.register_lowering_rule(barrier_test_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(barrier_test_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(
    barrier_test_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(barrier_test_p, *gpu_core.WGxWARP_SEMANTICS)
def _barrier_test_lowering(
    ctx: lowering.LoweringRuleContext,
    barrier,
    *flat_transforms,
    transforms_treedef,
):
  barrier_aval = ctx.avals_in[0]
  assert isinstance(barrier_aval, state_types.AbstractRef)
  transforms = transforms_treedef.unflatten(flat_transforms)
  orders_tensor_core = getattr(
      barrier_aval.inner_aval.dtype, "orders_tensor_core", False  # pyrefly: ignore[missing-attribute]
  )

  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    scope = mgpu_utils.ThreadSubset.WARP
  else:
    scope = mgpu_utils.ThreadSubset.WARPGROUP

  base_index = _get_barrier_base_index(barrier_aval, transforms)
  if base_index is not None:
    barrier = barrier[base_index]
  wait_complete = barrier.test(
      orders_tensor_core=orders_tensor_core, scope=scope
  )
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    return wait_complete
  return mgpu.FragmentedArray.splat(wait_complete, shape=(), is_signed=False)


def barrier_test(barrier: state.AbstractRef) -> jax.Array:
  """Tests the given barrier.

  This is a non-blocking equivalent of `barrier_wait`, which returns a boolean
  indicating whether or not the current barrier phase is complete.

  `barrier_test` is only supported within a warp context.
  """
  barrier, transforms = state_primitives.get_ref_and_transforms(
      barrier, None, "barrier_test"
  )
  flat_transforms, transforms_treedef = tree_util.tree_flatten(transforms)
  return barrier_test_p.bind(
      barrier, *flat_transforms, transforms_treedef=transforms_treedef,
  )


barrier_wait_p = jax_core.Primitive("barrier_wait")
barrier_wait_p.multiple_results = True


@barrier_wait_p.def_effectful_abstract_eval
def _barrier_wait_abstract_eval(barrier, *args, **params):
  _check_ref(barrier, "barrier", gpu_core.SMEM)
  del args, params  # Unused.
  return (), {gpu_core._memory_effect}


def _barrier_wait_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  del settings
  barrier, *flat_transforms = eqn.invars
  transforms_treedef = eqn.params["transforms_treedef"]
  transforms = transforms_treedef.unflatten(flat_transforms)
  return pp.concat([
      pp.text("barrier_wait"),
      pp.text(" "),
      state_primitives.pp_ref_transforms(context, barrier, transforms),
  ])


jax_core.pp_eqn_rules[barrier_wait_p] = _barrier_wait_pp_eqn


@lowering.register_lowering_rule(barrier_wait_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(barrier_wait_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(
    barrier_wait_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(barrier_wait_p, *gpu_core.WGxWARP_SEMANTICS)
def _barrier_wait_lowering(
    ctx: lowering.LoweringRuleContext,
    barrier,
    *flat_transforms,
    transforms_treedef,
):
  barrier_aval = ctx.avals_in[0]
  assert isinstance(barrier_aval, state_types.AbstractRef)
  transforms = transforms_treedef.unflatten(flat_transforms)
  orders_tensor_core = getattr(
      barrier_aval.inner_aval.dtype, "orders_tensor_core", False  # pyrefly: ignore[missing-attribute]
  )
  base_index = _get_barrier_base_index(barrier_aval, transforms)
  if base_index is not None:
    barrier = barrier[base_index]
  barrier.wait(orders_tensor_core=orders_tensor_core)
  return ()


def barrier_wait(barrier: state.AbstractRef) -> None:
  """Waits on the given barrier."""
  barrier, transforms = state_primitives.get_ref_and_transforms(
      barrier, None, "barrier_wait"
  )
  flat_transforms, transforms_treedef = tree_util.tree_flatten(transforms)
  barrier_wait_p.bind(
      barrier, *flat_transforms, transforms_treedef=transforms_treedef,
  )


wait_smem_to_gmem_p = jax_core.Primitive("wait_smem_to_gmem")
wait_smem_to_gmem_p.multiple_results = True


@wait_smem_to_gmem_p.def_effectful_abstract_eval
def _wait_smem_to_gmem_abstract_eval(*, n: int, wait_read_only):
  del wait_read_only  # Unused.
  if not isinstance(n, int) or n < 0:
    raise ValueError(f"n must be a non-negative integer, got {n}")
  return (), {gpu_core._memory_effect}


@lowering.register_lowering_rule(
    wait_smem_to_gmem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(
    wait_smem_to_gmem_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(
    wait_smem_to_gmem_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(
    wait_smem_to_gmem_p, *gpu_core.WGxWARP_SEMANTICS)
def _wait_smem_to_gmem_lowering(
    ctx: lowering.LoweringRuleContext, *, n: int, wait_read_only
):
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    scope = mgpu_utils.ThreadSubset.WARP
  else:
    scope = mgpu_utils.ThreadSubset.WARPGROUP
  ctx.launch_ctx.await_async_copy(
      allow_groups=n, await_read_only=wait_read_only,
      scope=scope
  )
  return ()


def wait_smem_to_gmem(n: int, wait_read_only: bool = False) -> None:
  """Waits until no more than the most recent ``n`` SMEM->GMEM copies issued by the calling thread are in flight.

  Args:
    n: The maximum number of copies in flight to wait for.
    wait_read_only: If ``True``, wait for the in flight copies to finish
      reading from SMEM. The writes to GMEM are not waited for.
  """
  wait_smem_to_gmem_p.bind(n=n, wait_read_only=wait_read_only)


commit_group_p = jax_core.Primitive("commit_group")
commit_group_p.multiple_results = True


@commit_group_p.def_effectful_abstract_eval
def _commit_group_abstract_eval():
  return (), {gpu_core._memory_effect}


@lowering.register_lowering_rule(commit_group_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(commit_group_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(
    commit_group_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(commit_group_p, *gpu_core.WGxWARP_SEMANTICS)
def _commit_group_lowering(ctx: lowering.LoweringRuleContext):
  del ctx  # Unused.
  llvm_dialect.inline_asm(
      ir.Type.parse("!llvm.void"),
      [],
      "cp.async.bulk.commit_group;",
      "",
      has_side_effects=True,
  )
  return ()


def commit_smem_to_gmem_group() -> None:
  """Commits all issued but uncommitted SMEM->GMEM copies to a group."""
  commit_group_p.bind()


# WGMMA on an accumulator reference
wgmma_ref_p = jax_core.Primitive("wgmma_ref")
wgmma_ref_p.multiple_results = True


def wgmma(acc: gpu_core.WGMMAAbstractAccumulatorRef, a, b) -> None:
  """Performs an asynchronous warp group matmul-accumulate on the given references.

  Conceptually, this is equivalent to doing ``acc[...] += a[...] @ b[...]``,
  except that the computation is performed asynchronously.

  Args:
    acc: The accumulator reference. Needs to be allocated via
      :func:`jax.experimental.pallas.run_scoped` called with a
      :func:`jax.experimental.pallas.mosaic_gpu.WGMMAAccumulatorRef`.
    a: The left hand side operand reference.
    b: The right hand side operand reference.

  See also:
    :func:`jax.experimental.pallas.mosaic_gpu.wgmma_wait`
  """
  m, n = acc.shape
  m2, k = a.shape
  k2, n2 = b.shape

  if m != m2 or n != n2 or k != k2:
    raise ValueError(
        f"Incompatible shapes for matrix multiplication: lhs={a.shape},"
        f" rhs={b.shape=}, acc={acc.shape}"
    )

  # A and B must share a dtype, except that the e4m3/e5m2 FP8 pair may be mixed:
  # `wgmma` takes independent `.atype`/`.btype` operands for FP8.
  fp8_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2)
  both_fp8 = a.dtype in fp8_dtypes and b.dtype in fp8_dtypes
  if a.dtype != b.dtype and not both_fp8:
    raise ValueError(
        "Mixed input dtypes for matrix multiplication unsupported: "
        f"lhs={a.dtype}, rhs={b.dtype}"
    )

  acc_transforms_leaves: list
  if isinstance(acc, pallas_core.TransformedRef):
    acc_transforms_leaves, acc_transforms_tree = jax.tree.flatten(acc.transforms)
    acc = acc.ref
  else:
    acc_transforms_leaves, acc_transforms_tree = [], None

  if isinstance(a, pallas_core.TransformedRef):
    a_transforms_leaves, a_transforms_tree = jax.tree.flatten(a.transforms)
    a = a.ref
  else:
    a_transforms_leaves, a_transforms_tree = [], None

  if isinstance(b, pallas_core.TransformedRef):
    b_transforms_leaves, b_transforms_tree = jax.tree.flatten(b.transforms)
    b = b.ref
  else:
    b_transforms_leaves, b_transforms_tree = [], None

  wgmma_ref_p.bind(
      acc,
      a,
      b,
      *acc_transforms_leaves,
      *a_transforms_leaves,
      *b_transforms_leaves,
      acc_transforms_tree=acc_transforms_tree,
      a_transforms_tree=a_transforms_tree,
      b_transforms_tree=b_transforms_tree,
  )


@wgmma_ref_p.def_effectful_abstract_eval
def _wgmma_ref_effectful_abstract_eval(acc_aval, a_aval, b_aval, *_, **params):
  del b_aval, params
  if not isinstance(acc_aval, gpu_core.WGMMAAbstractAccumulatorRef):
    raise TypeError(f"Expected WGMMAAbstractAccumulatorRef got {acc_aval}")
  return (), {
      gpu_core._wgmma_pipeline_effect,
      state.WriteEffect(0),
      state.ReadEffect(0),
      state.ReadEffect(2),
      *([state.ReadEffect(1)] if isinstance(a_aval, state.AbstractRef) else [])
  }


def _wgmma_ref_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  del settings
  acc, a, b, *leaves = eqn.invars
  transform_treedefs = [
      eqn.params["acc_transforms_tree"],
      eqn.params["a_transforms_tree"],
      eqn.params["b_transforms_tree"],
  ]
  transform_leaves = util.split_list(
      leaves, [getattr(tree, "num_leaves", 0) for tree in transform_treedefs]
  )
  acc_transforms, a_transforms, b_transforms = (
      () if treedef is None else treedef.unflatten(leaves)
      for treedef, leaves in zip(transform_treedefs, transform_leaves)
  )
  return pp.concat([
      pp.text("wgmma_ref"),
      pp.text(" "),
      state_primitives.pp_ref_transforms(context, acc, acc_transforms),
      pp.text(" <- "),
      state_primitives.pp_ref_transforms(context, a, a_transforms),
      pp.text(" @ "),
      state_primitives.pp_ref_transforms(context, b, b_transforms),
  ])


jax_core.pp_eqn_rules[wgmma_ref_p] = _wgmma_ref_pp_eqn


@discharge.register_discharge_rule(wgmma_ref_p)
def _wgmma_ref_discharge(ctx, *args, **kwargs):
  del ctx
  return (wgmma_p.bind(*args, **kwargs), *([None] * (len(args) - 1))), []


# Functional WGMMA, returns a shaped array. Internal.
wgmma_p = jax_core.Primitive("wgmma")


@lowering.register_lowering_rule(wgmma_p, mgpu.LoweringSemantics.Lane)
def _wgmma_lowering(
    ctx: lowering.LoweringRuleContext,
    acc,
    a,
    b,
    *transforms_leaves,
    acc_transforms_tree,
    a_transforms_tree,
    b_transforms_tree,
):
  lhs_swizzle: int | None = None
  transform_treedefs = [
      acc_transforms_tree, a_transforms_tree, b_transforms_tree
  ]
  transform_leaves = util.split_list(
      transforms_leaves, [getattr(tree, "num_leaves", 0) for tree in transform_treedefs]
  )
  acc_transforms, a_transforms, b_transforms = (
      None if treedef is None else treedef.unflatten(leaves)
      for treedef, leaves in zip(transform_treedefs, transform_leaves)
  )

  acc_indices = None
  if acc_transforms is not None:
    if not all(isinstance(t, indexing.NDIndexer) for t in acc_transforms):
      raise ValueError("WGMMA accumulator only supports indexing transforms")
    acc_indexer = lowering.merge_indexers(acc_transforms)
    if acc_indexer.int_indexer_shape:
      raise NotImplementedError("int_indexer_shape non-empty")
    acc_indices = lowering._ndindexer_indices(acc_indexer)

  transform_avals_list = util.split_list(
      ctx.avals_in[3:], [getattr(tree, "num_leaves", 0) for tree in transform_treedefs]
  )
  a_aval = ctx.avals_in[1]
  if not isinstance(a, mgpu.FragmentedArray):
    if not isinstance(a_aval, state_types.AbstractRef):
      assert isinstance(a_aval, jax_core.ShapedArray), (type(a_aval),)
      a_ref_aval = state.AbstractRef(a_aval)
    else:
      a_ref_aval = a_aval
    if a_transforms is not None:
      assert transform_treedefs[1] is not None
      a_transform_avals = transform_treedefs[1].unflatten(transform_avals_list[1])
    else:
      a_transforms = ()
      a_transform_avals = ()
    a, _, a_transforms = lowering._handle_transforms(
        ctx, a_ref_aval, a, a_transform_avals, a_transforms,
        handle_transposes=False, handle_reshapes=False)
    match a_transforms:
      case (gpu_core.UnswizzleRef(lhs_swizzle), gpu_core.UntilingTransform(tiling)):
        lhs_transpose = False
      case (
          gpu_core.UnswizzleRef(lhs_swizzle),
          gpu_core.UntilingTransform(tiling),
          state_types.TransposeTransform((1, 0)),
      ):
        lhs_transpose = True
      case ():
        raise ValueError(
            "When WGMMA lhs is passed in as a ref, it must be transformed by"
            " swizzling and tiling appropriately."
        )
      case _:
        raise ValueError(f"WGMMA lhs has unsupported transforms: {a_transforms}.")
    a_mlir_dtype = ir.MemRefType(a.type).element_type
    swizzle_elems = lhs_swizzle // mgpu_utils.bytewidth(a_mlir_dtype)
    if tiling != (8, swizzle_elems):
      raise NotImplementedError(
          f"WGMMA lhs tiling does not fit swizzle. Got {tiling=}, expected (8, {swizzle_elems})"
      )
  else:
    lhs_transpose = False

  b_aval = ctx.avals_in[2]
  if not isinstance(b_aval, state_types.AbstractRef):
    assert isinstance(b_aval, jax_core.ShapedArray)
    b_ref_aval = state_types.AbstractRef(b_aval)
  else:
    b_ref_aval = b_aval
  if b_transforms is not None:
    assert transform_treedefs[2] is not None
    b_transform_avals = transform_treedefs[2].unflatten(transform_avals_list[2])
  else:
    b_transforms = ()
    b_transform_avals = ()
  b, _, b_transforms = lowering._handle_transforms(
      ctx, b_ref_aval, b, b_transform_avals, b_transforms,
      handle_transposes=False)

  match b_transforms:
    case (gpu_core.UnswizzleRef(rhs_swizzle), gpu_core.UntilingTransform(rhs_tiling)):
      rhs_transpose = False
    case (
        gpu_core.UnswizzleRef(rhs_swizzle),
        gpu_core.UntilingTransform(rhs_tiling),
        state_types.TransposeTransform((1, 0)),
    ):
      rhs_transpose = True
    case _:
      raise ValueError(f"WGMMA rhs has unsupported transforms: {b_transforms}.")

  if lhs_swizzle is not None:
    b_mlir_dtype = ir.MemRefType(b.type).element_type
    swizzle_elems = rhs_swizzle // mgpu_utils.bytewidth(b_mlir_dtype)
    if rhs_swizzle != lhs_swizzle:
      raise NotImplementedError("WGMMA rhs swizzle must match lhs swizzle")
    if rhs_tiling != (8, swizzle_elems):
      raise NotImplementedError("WGMMA rhs tiling does not fit swizzle")

  if lhs_transpose:
    assert isinstance(a, ir.Value)
    a = mgpu.memref_transpose(a, (1, 0, 3, 2))
  if rhs_transpose:
    b = mgpu.memref_transpose(b, (1, 0, 3, 2))
  acc_in = acc
  if acc_indices is not None:
    acc_in = mgpu.WGMMAAccumulator(
        _value=acc._value[acc_indices],
        _original_layout=acc._original_layout,
        _sync=False,
    )
  acc_out = mgpu.wgmma(acc_in, a, b, swizzle=rhs_swizzle)
  if acc_indices is not None:
    acc_value = acc._value.copy()
    acc_value[acc_indices] = acc_out._value
    acc_out = mgpu.WGMMAAccumulator(
        _value=acc_value, _original_layout=acc._original_layout, _sync=False
    )
  void = ir.Type.parse("!llvm.void")
  llvm_dialect.inline_asm(
      void, [], "wgmma.commit_group.sync.aligned;", "", has_side_effects=True
  )
  return acc_out


@lowering.register_lowering_rule(wgmma_p, mgpu.LoweringSemantics.Warpgroup)
def _wgmma_warpgroup_lowering(
    ctx: lowering.LoweringRuleContext,
    acc,
    a,
    b,
    *transforms_leaves,
    acc_transforms_tree,
    a_transforms_tree,
    b_transforms_tree,
):
  transform_treedefs = [
      acc_transforms_tree, a_transforms_tree, b_transforms_tree
  ]
  transform_leaves = util.split_list(
      transforms_leaves, [getattr(tree, "num_leaves", 0) for tree in transform_treedefs]
  )
  acc_transforms, a_transforms, b_transforms = (
      None if treedef is None else treedef.unflatten(leaves)
      for treedef, leaves in zip(transform_treedefs, transform_leaves)
  )

  if acc_transforms is not None:
    if not all(isinstance(t, indexing.NDIndexer) for t in acc_transforms):
      raise ValueError("WGMMA accumulator only supports indexing transforms")
    acc_indexer = lowering.merge_indexers(acc_transforms)
    if acc_indexer.int_indexer_shape:
      raise NotImplementedError("int_indexer_shape non-empty")
    acc_indices = lowering._ndindexer_indices(acc_indexer)
  else:
    acc_indices = ()

  transform_avals_list = util.split_list(
      ctx.avals_in[3:], [getattr(tree, "num_leaves", 0) for tree in transform_treedefs]
  )

  a_aval = ctx.avals_in[1]
  if isinstance(a_aval, state_types.AbstractRef):
    if a_transforms is not None:
      a_transform_avals = a_transforms_tree.unflatten(transform_avals_list[1])
    else:
      a_transforms = ()
      a_transform_avals = ()
    a, _, a_transforms = lowering._handle_transforms(
        ctx, a_aval, a, a_transform_avals, a_transforms
    )
    if a_transforms:
      raise ValueError(
          f"WGMMA lhs has unsupported transforms: {a_transforms}."
      )

  b_aval = ctx.avals_in[2]
  if isinstance(b_aval, state_types.AbstractRef):
    if b_transforms is not None:
      b_transform_avals = b_transforms_tree.unflatten(transform_avals_list[2])
    else:
      b_transforms = ()
      b_transform_avals = ()
    b, _, b_transforms = lowering._handle_transforms(
        ctx, b_aval, b, b_transform_avals, b_transforms
    )
    if b_transforms:
      raise ValueError(
          f"WGMMA rhs has unsupported transforms: {b_transforms}."
      )

  offsets = []
  strides = []
  sizes = []
  for s in acc_indices:
    if not isinstance(s, slice):
      raise NotImplementedError("Only static slices are supported")
    if s.start is None or s.stop is None:
      raise NotImplementedError("Static slices must have start and stop")
    if s.step is not None and s.step != 1:
      raise NotImplementedError("Static slices must have step 1")
    offsets.append(s.start)
    sizes.append(s.stop - s.start)
    strides.append(s.step or 1)

  if acc_indices:
    acc_type = ir.VectorType.get(sizes, ir.VectorType(acc.type).element_type)
    acc_in = vector_dialect.extract_strided_slice(
        acc_type,
        source=acc,
        offsets=offsets,
        sizes=sizes,
        strides=strides,
    )
  else:
    acc_in = acc

  new_acc = mgpu.dialect.wgmma(acc_in, a, b)

  if acc_indices:
    assert offsets and strides
    new_acc = vector_dialect.insert_strided_slice(
        new_acc, acc, offsets, strides
    )

  void = ir.Type.parse("!llvm.void")
  llvm_dialect.inline_asm(
      void, [], "wgmma.commit_group.sync.aligned;", "", has_side_effects=True
  )
  return new_acc


@wgmma_p.def_effectful_abstract_eval
def _wgmma_effectful_abstract_eval(acc, lhs_ref, *args, **kwargs):
  del args, kwargs
  return acc, {
      gpu_core._wgmma_pipeline_effect,
      state.ReadEffect(2),
      *([state.ReadEffect(1)] if isinstance(lhs_ref, state.AbstractRef) else [])
  }


def mma(acc: jax.Array, a: jax.Array, b: jax.Array, /) -> jax.Array:
  """Computes ``acc + a @ b`` synchronously using Ampere MMA instructions."""
  return mma_p.bind(acc, a, b)


mma_p = jax_core.Primitive("mma")


@mma_p.def_abstract_eval
def _mma_abstract_eval(acc, a, b):
  m, n = acc.shape
  m2, k = a.shape
  k2, n2 = b.shape
  if m != m2 or n != n2 or k != k2:
    raise ValueError(
        f"Incompatible shapes for matrix multiplication: lhs={a.shape},"
        f" rhs={b.shape}, acc={acc.shape}"
    )
  if a.dtype != b.dtype:
    raise TypeError(f"Operand dtypes must match: {a.dtype} != {b.dtype}")
  if jnp.issubdtype(a.dtype, jnp.integer) or a.dtype == jnp.bool_:
    if acc.dtype != jnp.int32:
      raise NotImplementedError(
          "Only int32 accumulator supported for integer and boolean operands."
          f" Got {acc.dtype}"
      )
  elif jnp.issubdtype(a.dtype, jnp.floating):
    if acc.dtype not in (jnp.float64, jnp.float32, jnp.float16):
      raise NotImplementedError(
          "Only float64, float32 and float16 accumulators supported for"
          f" floating operands. Got {acc.dtype}"
      )
  else:
    raise NotImplementedError(f"Unsupported operand type: {a.dtype}")

  _check_layout_mode()
  if get_layout_mode() == LayoutMode.PALLAS_GPU:
    out_layout = lax_utils.call_layout_rule(
        mma_p, in_avals=(acc, a, b), out_avals=acc, layout_rule=None
    )
    if out_layout != acc.layout:
      raise ValueError(
          f"Accumulator was provided with layout {acc.layout} but expected "
          f"layout {out_layout}"
      )
  return acc


@lowering.register_lowering_rule(mma_p, mgpu.LoweringSemantics.Lane)
def _mma_lowering(ctx: lowering.LoweringRuleContext, acc, a, b):
  del ctx  # Unused.
  return mgpu.mma(acc, a, b)


@lowering.register_lowering_rule(mma_p, mgpu.LoweringSemantics.Warpgroup)
def _mma_warpgroup_lowering(ctx: lowering.LoweringRuleContext, acc, a, b):
  del ctx  # Unused.
  return mgpu.dialect.mma(acc, a, b)


wgmma_wait_p = jax_core.Primitive("wgmma_wait")
wgmma_wait_p.multiple_results = True


def wgmma_wait(n: int):
  """Waits until there is no more than ``n`` WGMMA operations in flight."""
  return wgmma_wait_p.bind(n=n)


@wgmma_wait_p.def_effectful_abstract_eval
def wgmma_wait_effectful_abstract_eval(*, n: int):
  if not isinstance(n, int) or n < 0:
    raise ValueError(f"n must be a non-negative integer, got {n}")
  return [], {gpu_core._wgmma_pipeline_effect}


@lowering.register_lowering_rule(wgmma_wait_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(wgmma_wait_p, mgpu.LoweringSemantics.Warpgroup)
def _wgmma_wait_lowering(ctx: lowering.LoweringRuleContext, *, n: int):
  del ctx
  void = ir.Type.parse("!llvm.void")
  llvm_dialect.inline_asm(
      void, [], f"wgmma.wait_group.sync.aligned {n};", "", has_side_effects=True,
  )
  return ()


wgmma_accumulator_deref_p = jax_core.Primitive("wgmma_accumulator_deref")


def wgmma_accumulator_load(acc, *, wait_n: int | None = 0):
  """Dereferences an accumulator register.

  Before dereferencing, this operation waits until there is no more than
  ``wait_n`` WGMMA operations in flight. If ``wait_n`` is None, no
  synchronization is performed.
  """
  if wait_n is not None and wait_n < 0:
    raise ValueError(f"wait_n must be non-negative, got {wait_n=}")

  if not isinstance(acc.aval, gpu_core.WGMMAAbstractAccumulatorRef):
    raise TypeError(f"acc must be a WGMMAAccumulatorAbstractRef, got {acc.aval=}")

  return wgmma_accumulator_deref_p.bind(acc, wait_n=wait_n)


@wgmma_accumulator_deref_p.def_effectful_abstract_eval
def _wgmma_accumulator_deref_abstract_eval(acc, **_):
  # Dereferencing implies flushing so we have a wgmma pipeline effect.
  ret = acc.inner_aval if isinstance(acc, state.AbstractRef) else acc
  assert isinstance(ret, jax_core.ShapedArray), acc
  return ret, {gpu_core._wgmma_pipeline_effect}


@discharge.register_discharge_rule(wgmma_accumulator_deref_p)
def _wgmma_accumulator_deref_discharge(ctx, acc, *, wait_n):
  del ctx
  return (None,), wgmma_accumulator_deref_p.bind(acc, wait_n=wait_n)


@lowering.register_lowering_rule(
    wgmma_accumulator_deref_p, mgpu.LoweringSemantics.Lane
)
@lowering.register_lowering_rule(
    wgmma_accumulator_deref_p, mgpu.LoweringSemantics.Warpgroup
)
def _wgmma_accumulator_deref_lowering(
    ctx: lowering.LoweringRuleContext, acc, *, wait_n: int | None
):
  if wait_n is not None:
    void = ir.Type.parse("!llvm.void")
    llvm_dialect.inline_asm(
        void, [], f"wgmma.wait_group.sync.aligned {int(wait_n)};", "", has_side_effects=True,
    )
  return (
      acc.value
      if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane
      else acc
  )


# This primitive stores a value into a WGMMA accumulator ref. It is the
# counterpart to wgmma_accumulator_deref_p. Before discharge, it writes to an
# accumulator ref and the return value is unused. After discharge, the ref is
# replaced by a regular value: the discharge rule re-binds the primitive, and
# the lowering wraps the value into a WGMMAAccumulator.
wgmma_accumulator_store_p = jax_core.Primitive("wgmma_accumulator_store")


def wgmma_accumulator_store(acc_ref, val):
  if not isinstance(acc_ref.aval, gpu_core.WGMMAAbstractAccumulatorRef):
    raise TypeError(f"acc must be a WGMMAAccumulatorAbstractRef, got {acc_ref.aval=}")
  wgmma_accumulator_store_p.bind(acc_ref, val)


@wgmma_accumulator_store_p.def_effectful_abstract_eval
def _wgmma_accumulator_store_abstract_eval(acc, val):
  # Before discharge acc is a WGMMAAbstractAccumulatorRef. After discharge,
  # the discharge rule re-binds the primitive and acc becomes a ShapedArray.
  if isinstance(acc, gpu_core.WGMMAAbstractAccumulatorRef):
    inner = acc.inner_aval
    assert isinstance(inner, jax_core.ShapedArray)
  elif isinstance(acc, jax_core.ShapedArray):
    inner = acc
  else:
    raise TypeError(f"Expected WGMMAAbstractAccumulatorRef or ShapedArray, got {type(acc)}")
  if inner.shape != val.shape:
    raise ValueError(
        f"Accumulator shape {inner.shape} does not match value shape {val.shape}"
    )
  if inner.dtype != val.dtype:
    raise ValueError(
        f"Accumulator dtype {inner.dtype} does not match value dtype {val.dtype}"
    )
  effects: set[jax_core.Effect] = {gpu_core._wgmma_pipeline_effect}
  if isinstance(acc, gpu_core.WGMMAAbstractAccumulatorRef):
    effects.add(state.WriteEffect(0))
  return inner, effects


@discharge.register_discharge_rule(wgmma_accumulator_store_p)
def _wgmma_accumulator_store_discharge(ctx, acc, val):
  del ctx
  return (wgmma_accumulator_store_p.bind(acc, val), None), []


@lowering.register_lowering_rule(
    wgmma_accumulator_store_p, mgpu.LoweringSemantics.Lane
)
def _wgmma_accumulator_store_lowering(
    ctx: lowering.LoweringRuleContext, acc, val
):
  del ctx, acc
  return mgpu.WGMMAAccumulator.from_registers(val)


@lowering.register_lowering_rule(
    wgmma_accumulator_store_p, mgpu.LoweringSemantics.Warpgroup
)
def _wgmma_accumulator_store_warpgroup_lowering(
    ctx: lowering.LoweringRuleContext, acc, val
):
  del ctx, acc
  val = mgpu.dialect.optimization_barrier([val])
  lowering._wgmma_fence_aligned()
  return val

# MMA for TensorCore gen 5.
tcgen05_mma_p = jax_core.Primitive("tcgen05_mma")
tcgen05_mma_p.multiple_results = True

def tcgen05_mma(acc: _Ref,
                a: _Ref,
                b: _Ref,
                barrier: _Ref | None = None,
                *,
                a_scale: _Ref | None = None,
                b_scale: _Ref | None = None,
                a_sparse_metadata: _Ref | None = None,
                accumulate: bool | jax.Array = True,
                collective_axis: str | None = None):
  """Asynchronous matrix-multiply accumulate for TensorCore gen 5 (Blackwell).

  If run in collective mode, ``acc``, ``a`` (LHS), and ``b`` (RHS) should
  correspond to half of the total inputs to the MMA, where ``acc`` and ``a``
  (LHS) are split in half along the rows and ``b`` (RHS) is split along the
  columns like so::

    -----------    -----------   -----------
    |  ACC1   |    |  LHS1   |   |    |    |
    ----------- += ----------- @ |RHS1|RHS2|
    |  ACC2   |    |  LHS2   |   |    |    |
    -----------    -----------   -----------

  To use the block-scaled matrix-multiply, provide ``a_scale`` and ``b_scale``
  operands (they must be both present or both unspecified).

  Args:
    acc: The accumulator. Must be a TMEM Ref.
    a: The left-hand side. Must be a TMEM/SMEM Ref.
    b: The right-hand side. Must be an SMEM Ref.
    barrier: Optional barrier Ref for synchronizing with the tensor core.
      Must have orders_tensor_core set to True. If not specified, the MMA
      completion should be explicitly observed by calling
      :func:`jax.experimental.pallas.mosaic_gpu.tcgen05_commit_arrive`
    a_scale: An optional scale for the ``a`` operand. Must be a TMEM Ref if present.
    b_scale: An optional scale for the ``b`` operand. Must be a TMEM Ref if present.
    a_sparse_metadata: An optional sparse metadata for the ``a`` operand.
      Must be a TMEM Ref if present.
    accumulate: Whether to accumulate into acc or overwrite it.
    collective_axis: The name of the cluster axis along which to perform
      a collective MMA. The cluster axis should have a size of exactly 2,
      and must be on the minormost cluster axis.
  """
  acc_m, acc_n = acc.shape
  lhs_m, lhs_k = a.shape
  rhs_k, rhs_n = b.shape
  is_sparse = a_sparse_metadata is not None

  if acc_m != lhs_m:
    raise ValueError(
        "Accumulator and LHS have incompatible shapes. Expected LHS to have"
        " shape (m, k) and accumulator to have shape (m, n). Accumulator:"
        f" {acc.shape}. LHS: {a.shape}."
    )

  if collective_axis is not None:
    if acc_n != rhs_n * 2:
      raise ValueError(
          "Accumulator and RHS have incompatible shapes. Expected RHS to have "
          "shape (k, n) and accumulator to have shape (m, n * 2) in "
          f"collective mode. Accumulator: {acc.shape}. RHS: {b.shape}."
      )
  elif acc_n != rhs_n:
    raise ValueError(
        "Accumulator and RHS have incompatible shapes. Expected RHS to have"
        " shape (k, n) and accumulator to have shape (m, n). Accumulator:"
        f" {acc.shape}. RHS: {b.shape}."
    )

  if (lhs_k * (1 + is_sparse)) != rhs_k:
    raise ValueError(
        f"LHS and RHS have incompatible shapes. LHS: {a.shape}. RHS: {b.shape}.")

  if isinstance(acc, pallas_core.TransformedRef):
    acc_transforms_leaves, acc_transforms_tree = jax.tree.flatten(
        acc.transforms)
    acc = acc.ref
  else:
    acc_transforms_leaves, acc_transforms_tree = [], None

  if isinstance(a, pallas_core.TransformedRef):
    a_transforms_leaves, a_transforms_tree = jax.tree.flatten(a.transforms)
    a = a.ref
  else:
    a_transforms_leaves, a_transforms_tree = [], None

  if isinstance(b, pallas_core.TransformedRef):
    b_transforms_leaves, b_transforms_tree = jax.tree.flatten(b.transforms)
    b = b.ref
  else:
    b_transforms_leaves, b_transforms_tree = [], None

  if (is_scaled := a_scale is not None) != (b_scale is not None):
    raise ValueError("a_scale and b_scale must both be present or absent.")
  scales = []
  if isinstance(a_scale, pallas_core.TransformedRef):
    a_scale_transforms_leaves, a_scale_transforms_tree = jax.tree.flatten(
        a_scale.transforms
    )
    scales.append(a_scale.ref)
  else:
    a_scale_transforms_leaves, a_scale_transforms_tree = [], None
    scales.append(a_scale)
  if isinstance(b_scale, pallas_core.TransformedRef):
    b_scale_transforms_leaves, b_scale_transforms_tree = jax.tree.flatten(
        b_scale.transforms
    )
    scales.append(b_scale.ref)
  else:
    b_scale_transforms_leaves, b_scale_transforms_tree = [], None
    scales.append(b_scale)
  if not is_scaled:
    scales = []

  if isinstance(a_sparse_metadata, pallas_core.TransformedRef):
    a_sparse_metadata_transforms_leaves, a_sparse_metadata_transforms_tree = jax.tree.flatten(
        a_sparse_metadata.transforms
    )
    sparse_metadata = [a_sparse_metadata.ref]
  else:
    a_sparse_metadata_transforms_leaves, a_sparse_metadata_transforms_tree = [], None
    sparse_metadata = [a_sparse_metadata] if is_sparse else []

  if isinstance(barrier, pallas_core.TransformedRef):
    barrier_transforms_leaves, barrier_transforms_tree = jax.tree.flatten(
        barrier.transforms
    )
    barrier = barrier.ref
  else:
    barrier_transforms_leaves, barrier_transforms_tree = [], None

  if barrier is not None:
    barrier_ref = [barrier]
    arrive = True
  else:
    barrier_ref = []
    arrive = False

  tcgen05_mma_p.bind(acc, a, b, accumulate, *barrier_ref, *scales, *sparse_metadata,
                     *acc_transforms_leaves, *a_transforms_leaves,
                     *b_transforms_leaves,
                     *barrier_transforms_leaves,
                     *a_scale_transforms_leaves, *b_scale_transforms_leaves,
                     *a_sparse_metadata_transforms_leaves,
                     acc_transforms_tree=acc_transforms_tree,
                     a_transforms_tree=a_transforms_tree,
                     b_transforms_tree=b_transforms_tree,
                     barrier_transforms_tree=barrier_transforms_tree,
                     a_scale_transforms_tree=a_scale_transforms_tree,
                     b_scale_transforms_tree=b_scale_transforms_tree,
                     a_sparse_metadata_transforms_tree=a_sparse_metadata_transforms_tree,
                     collective_axis=collective_axis,
                     arrive=arrive,
                     scaled=bool(scales),
                     sparse=is_sparse)


@tcgen05_mma_p.def_effectful_abstract_eval
def _tcgen05_mma_abstract_eval(acc, a, b, accumulate,
                               *barrier_scales_and_transforms_leaves,
                               acc_transforms_tree, a_transforms_tree,
                               b_transforms_tree,
                               barrier_transforms_tree,
                               a_scale_transforms_tree,
                               b_scale_transforms_tree,
                               a_sparse_metadata_transforms_tree,
                               collective_axis,
                               arrive,
                               scaled,
                               sparse):
  del accumulate, acc_transforms_tree, a_transforms_tree, b_transforms_tree, barrier_transforms_tree

  if acc.memory_space != gpu_core.TMEM:
    raise ValueError("Accumulator must be a TMEM Ref.")
  if a.memory_space not in (gpu_core.SMEM, gpu_core.TMEM):
    raise ValueError("LHS must be a TMEM/SMEM Ref.")
  if b.memory_space != gpu_core.SMEM:
    raise ValueError("RHS must be an SMEM Ref.")

  if collective_axis is not None:
    # TODO(justinfu): If under a core_map, the avals for acc/a
    # become normal MemRefs so we cannot check if they are collective.
    # Figure out a way to fix this.
    if isinstance(acc, gpu_core.AbstractTMEMRef) and not acc.collective:
      raise ValueError(
          "Accumulator Ref must be collective if collective_axis is set.")
    if isinstance(a, gpu_core.AbstractTMEMRef) and not a.collective:
      raise ValueError(
          "LHS Ref must be collective if collective_axis is set.")

  scales_and_transforms_leaves = barrier_scales_and_transforms_leaves
  if arrive:
    barrier, *scales_and_transforms_leaves = barrier_scales_and_transforms_leaves
    orders_tensor_core = getattr(
        barrier.inner_aval.dtype, "orders_tensor_core", False)
    if not orders_tensor_core:
      raise ValueError("MMA barrier must have orders_tensor_core set to True.")
  if scaled:
    a_scale, b_scale = scales_and_transforms_leaves[:2]
    if a_scale.memory_space != gpu_core.TMEM:
      raise ValueError("a_scale must be a TMEM Ref")
    if b_scale.memory_space != gpu_core.TMEM:
      raise ValueError("b_scale must be a TMEM Ref")

  return [], {gpu_core._memory_effect}


@lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWG_SEMANTICS)
@lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWARP_SEMANTICS)
def _tcgen05_mma_lowering(
    ctx: lowering.LoweringRuleContext,
    acc: tcgen05.TMEMRef,
    a_ref,
    b_ref,
    accumulate: bool | ir.Value,
    *barrier_scales_and_transforms_leaves,
    acc_transforms_tree,
    a_transforms_tree,
    b_transforms_tree,
    barrier_transforms_tree,
    a_scale_transforms_tree,
    b_scale_transforms_tree,
    a_sparse_metadata_transforms_tree,
    collective_axis,
    arrive,
    scaled: bool,
    sparse: bool,
):
  (
      acc_aval,
      a_aval,
      b_aval,
      accumulate_aval,
      *_,
  ) = ctx.avals_in
  assert isinstance(acc_aval, state_types.AbstractRef)
  assert isinstance(a_aval, state_types.AbstractRef)
  assert isinstance(b_aval, state_types.AbstractRef)
  del accumulate_aval
  lhs_swizzle: int | None = None
  lhs_transpose: bool = False
  leaves = list(barrier_scales_and_transforms_leaves)
  avals = list(ctx.avals_in[4:])
  if arrive:
    barrier_ref, leaves = leaves[0], leaves[1:]
    barrier_ref_aval, avals = avals[0], avals[1:]
  else:
    barrier_ref = None
    barrier_ref_aval = None
  if scaled:
    a_scale_ref, b_scale_ref, leaves = leaves[0], leaves[1], leaves[2:]
    a_scale_ref_aval, b_scale_ref_aval, avals = avals[0], avals[1], avals[2:]
  else:
    a_scale_ref = b_scale_ref = a_scale_ref_aval = b_scale_ref_aval = None
  if sparse:
    a_sparse_metadata_ref, leaves = leaves[0], leaves[1:]
    a_sparse_metadata_ref_aval, avals = avals[0], avals[1:]
  else:
    a_sparse_metadata_ref = a_sparse_metadata_ref_aval = None

  transforms_trees = (
      acc_transforms_tree,
      a_transforms_tree,
      b_transforms_tree,
      barrier_transforms_tree,
      a_scale_transforms_tree,
      b_scale_transforms_tree,
      a_sparse_metadata_transforms_tree,
  )
  ns = [getattr(tree, "num_leaves", 0) for tree in transforms_trees]
  transforms_leaves_lists = util.split_list_checked(leaves, ns)
  transforms_avals_lists = util.split_list_checked(avals, ns)

  (
      acc_transforms_leaves,
      a_transforms_leaves,
      b_transforms_leaves,
      barrier_transforms_leaves,
      a_scale_transforms_leaves,
      b_scale_transforms_leaves,
      a_sparse_metadata_transforms_leaves,
  ) = transforms_leaves_lists

  (
      acc_transforms_leaves_avals,
      a_transforms_leaves_avals,
      b_transforms_leaves_avals,
      barrier_transforms_leaves_avals,
      a_scale_transforms_leaves_avals,
      b_scale_transforms_leaves_avals,
      a_sparse_metadata_transforms_leaves_avals,
  ) = transforms_avals_lists

  if acc_transforms_tree is not None:
    acc_transforms = acc_transforms_tree.unflatten(acc_transforms_leaves)
    acc_transform_avals = acc_transforms_tree.unflatten(acc_transforms_leaves_avals)
  else:
    acc_transforms = ()
    acc_transform_avals = ()
  acc, _, acc_transforms = lowering._handle_transforms(
      ctx, acc_aval, acc, acc_transform_avals, acc_transforms,
      handle_transposes=False
  )
  if acc_transforms:
    raise NotImplementedError(
        f"Unsupported transforms for ACC: {acc_transforms}."
    )

  if a_transforms_tree is not None:
    a_transforms = a_transforms_tree.unflatten(a_transforms_leaves)
    a_transform_avals = a_transforms_tree.unflatten(a_transforms_leaves_avals)
  else:
    a_transforms = ()
    a_transform_avals = ()
  a_out_ty = state_types.transform_type(a_transforms, a_aval)
  assert isinstance(a_out_ty, state_types.AbstractRef)
  a_dtype = a_out_ty.dtype
  a_ref, _, a_transforms = lowering._handle_transforms(
      ctx, a_aval, a_ref, a_transform_avals, a_transforms,
      handle_transposes=False, handle_reshapes=True)
  match a_transforms:
    case (
        gpu_core.UnswizzleRef(lhs_swizzle),
        gpu_core.UntilingTransform(lhs_tiling),
    ):
      lhs_transpose = False
    case (
        gpu_core.UnswizzleRef(lhs_swizzle),
        gpu_core.UntilingTransform(lhs_tiling),
        state_types.TransposeTransform((1, 0)),
    ):
      lhs_transpose = True
    case () if isinstance(a_ref, tcgen05.TMEMRef):
      lhs_tiling = None
    case _:
      raise NotImplementedError(
          f"Unsupported transforms for LHS: {a_transforms}."
      )
  if not isinstance(a_ref, tcgen05.TMEMRef):
    assert lhs_swizzle is not None
    swizzle_elems = 8 * lhs_swizzle // dtypes.itemsize_bits(a_dtype)
    if lhs_tiling != (8, swizzle_elems):
      raise ValueError("MMA lhs tiling does not fit swizzle. "
                      f"{lhs_tiling=} expected={(8, swizzle_elems)}")

  if b_transforms_tree is not None:
    b_transforms = b_transforms_tree.unflatten(b_transforms_leaves)
    b_transform_avals = b_transforms_tree.unflatten(b_transforms_leaves_avals)
  else:
    b_transforms = ()
    b_transform_avals = ()
  b_out_ty = state_types.transform_type(b_transforms, b_aval)
  assert isinstance(b_out_ty, state_types.AbstractRef)
  b_dtype = b_out_ty.dtype
  b_ref, _, b_transforms = lowering._handle_transforms(
      ctx, b_aval, b_ref, b_transform_avals, b_transforms, handle_transposes=False,
      handle_reshapes=True)
  match b_transforms:
    case (
        gpu_core.UnswizzleRef(rhs_swizzle),
        gpu_core.UntilingTransform(rhs_tiling),
    ):
      rhs_transpose = False
    case (
        gpu_core.UnswizzleRef(rhs_swizzle),
        gpu_core.UntilingTransform(rhs_tiling),
        state_types.TransposeTransform((1, 0)),
    ):
      rhs_transpose = True
    case _:
      raise NotImplementedError(
          f"Unsupported transforms for RHS: {b_transforms}."
      )
  swizzle_elems = 8 * rhs_swizzle // dtypes.itemsize_bits(b_dtype)
  if rhs_tiling != (8, swizzle_elems):
    raise ValueError(
        "MMA rhs tiling does not fit swizzle"
        f" {rhs_tiling=} expected={(8, swizzle_elems)}"
    )

  if barrier_transforms_tree is not None and barrier_ref is not None:
    barrier_transforms = barrier_transforms_tree.unflatten(
        barrier_transforms_leaves
    )
    base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)
    if base_index is not None:
      barrier_ref = barrier_ref[base_index]

  if lhs_swizzle is None:
    lhs_swizzle = rhs_swizzle
  elif rhs_swizzle != lhs_swizzle:
    raise ValueError("MMA rhs swizzle must match lhs swizzle."
                      f" {lhs_swizzle=} {rhs_swizzle=}")
  if lhs_transpose:
    if isinstance(a_ref, tcgen05.TMEMRef):
      raise ValueError("TMEM transpose not allowed.")
    a_ref = mgpu.memref_transpose(a_ref, (1, 0, 3, 2))
  if rhs_transpose:
    b_ref = mgpu.memref_transpose(b_ref, (1, 0, 3, 2))
  if isinstance(accumulate, bool):
    accumulate = mgpu.c(accumulate, ir.IntegerType.get_signless(1))
  elif isinstance(accumulate, mgpu.FragmentedArray):
    accumulate = accumulate.registers.item()
    assert isinstance(accumulate, ir.Value)

  if a_scale_ref is not None:
    assert isinstance(a_scale_ref_aval, state.AbstractRef)
    if a_scale_transforms_tree is not None:
      a_scale_transforms = a_scale_transforms_tree.unflatten(
          a_scale_transforms_leaves
      )
      a_scale_transform_avals = a_scale_transforms_tree.unflatten(
          a_scale_transforms_leaves_avals
      )
    else:
      a_scale_transforms = ()
      a_scale_transform_avals = ()
    a_scale_ref, _, a_scale_transforms = lowering._handle_transforms(
        ctx, a_scale_ref_aval, a_scale_ref, a_scale_transform_avals,
        a_scale_transforms
    )
    if a_scale_transforms:
      raise NotImplementedError(
          f"Unsupported transforms: {a_scale_transforms}"
      )
  if b_scale_ref is not None:
    assert isinstance(b_scale_ref_aval, state.AbstractRef)
    if b_scale_transforms_tree is not None:
      b_scale_transforms = b_scale_transforms_tree.unflatten(
          b_scale_transforms_leaves
      )
      b_scale_transform_avals = b_scale_transforms_tree.unflatten(
          b_scale_transforms_leaves_avals
      )
    else:
      b_scale_transforms = ()
      b_scale_transform_avals = ()
    b_scale_ref, _, b_scale_transforms = lowering._handle_transforms(
        ctx, b_scale_ref_aval, b_scale_ref, b_scale_transform_avals,
        b_scale_transforms
    )
    if b_scale_transforms:
      raise NotImplementedError(f"Unsupported transforms: {b_scale_transforms}")
  if a_sparse_metadata_ref is not None:
    if a_sparse_metadata_transforms_tree is not None:
      a_sparse_metadata_transforms = a_sparse_metadata_transforms_tree.unflatten(
          a_sparse_metadata_transforms_leaves
      )
      a_sparse_metadata_transform_avals = (
          a_sparse_metadata_transforms_tree.unflatten(
              a_sparse_metadata_transforms_leaves_avals
          )
      )
    else:
      a_sparse_metadata_transforms = ()
      a_sparse_metadata_transform_avals = ()
    assert isinstance(a_sparse_metadata_ref_aval, state_types.AbstractRef)
    a_sparse_metadata_ref, _, a_sparse_metadata_transforms = (
        lowering._handle_transforms(  # pyrefly: ignore[bad-specialization]
            ctx, a_sparse_metadata_ref_aval, a_sparse_metadata_ref,
            a_sparse_metadata_transform_avals, a_sparse_metadata_transforms)
    )
    if a_sparse_metadata_transforms:
      raise NotImplementedError(
          f"Unsupported transforms: {a_sparse_metadata_transforms}"
      )

  predicate = ctx.module_ctx.single_lane_predicate
  if collective := collective_axis is not None:
    is_leader_block = _collective_mma_predicate(ctx, collective_axis)
    predicate = _andi_maybe_none(predicate, is_leader_block)

  with mgpu.when(predicate):
    tcgen05.mma(
        acc,
        a_ref,
        b_ref,
        a_swizzle=int(lhs_swizzle),
        b_swizzle=int(rhs_swizzle),
        a_scale=a_scale_ref,
        b_scale=b_scale_ref,
        a_sparse_metadata=a_sparse_metadata_ref,
        accumulate=accumulate,
        collective=collective,
    )
    if arrive:
      assert barrier_ref is not None
      tcgen05.commit_arrive(
          barrier_ref, collective=collective, ctx=ctx.launch_ctx
      )
  return []


@lowering.register_lowering_rule(tcgen05_mma_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.WGxWARP_SEMANTICS)
def _tcgen05_mma_lowering_wg(
    ctx: lowering.LoweringRuleContext,
    acc_ref,
    a_ref,
    b_ref,
    accumulate: bool | ir.Value,
    *barrier_scales_and_transforms_leaves,
    acc_transforms_tree,
    a_transforms_tree,
    b_transforms_tree,
    barrier_transforms_tree,
    a_scale_transforms_tree,
    b_scale_transforms_tree,
    a_sparse_metadata_transforms_tree,
    collective_axis,
    arrive,
    scaled: bool,
    sparse: bool,
):
  (
      acc_aval,
      a_aval,
      b_aval,
      accumulate_aval,
      *_,
  ) = ctx.avals_in
  assert isinstance(acc_aval, state_types.AbstractRef)
  assert isinstance(a_aval, state_types.AbstractRef)
  assert isinstance(b_aval, state_types.AbstractRef)
  del accumulate_aval

  leaves = list(barrier_scales_and_transforms_leaves)
  avals = list(ctx.avals_in[4:])
  if arrive:
    barrier_ref, leaves = leaves[0], leaves[1:]
    barrier_ref_aval, avals = avals[0], avals[1:]
  else:
    barrier_ref = None
    barrier_ref_aval = None

  if scaled:
    # Scales are not supported for WG semantics, but we still need to unpack them
    # if they are present in the leaves/avals to get to the transforms.
    a_scale_ref, b_scale_ref, leaves = leaves[0], leaves[1], leaves[2:]
    a_scale_ref_aval, b_scale_ref_aval, avals = avals[0], avals[1], avals[2:]
  else:
    a_scale_ref = b_scale_ref = a_scale_ref_aval = b_scale_ref_aval = None

  if sparse:
    a_sparse_metadata_ref, leaves = leaves[0], leaves[1:]
    a_sparse_metadata_ref_aval, avals = avals[0], avals[1:]
  else:
    a_sparse_metadata_ref = a_sparse_metadata_ref_aval = None

  transforms_trees = (
      acc_transforms_tree,
      a_transforms_tree,
      b_transforms_tree,
      barrier_transforms_tree,
      a_scale_transforms_tree,
      b_scale_transforms_tree,
      a_sparse_metadata_transforms_tree,
  )
  ns = [getattr(tree, "num_leaves", 0) for tree in transforms_trees]
  transforms_leaves_lists = util.split_list_checked(leaves, ns)
  transforms_avals_lists = util.split_list_checked(avals, ns)

  (
      acc_transforms_leaves,
      a_transforms_leaves,
      b_transforms_leaves,
      barrier_transforms_leaves,
      a_scale_transforms_leaves,
      b_scale_transforms_leaves,
      a_sparse_metadata_transforms_leaves,
  ) = transforms_leaves_lists

  (
      acc_transforms_leaves_avals,
      a_transforms_leaves_avals,
      b_transforms_leaves_avals,
      _,
      a_scale_transforms_leaves_avals,
      b_scale_transforms_leaves_avals,
      a_sparse_metadata_transforms_leaves_avals,
  ) = transforms_avals_lists

  def handle_transforms_and_get_ref(tree, leaves, leaves_avals, ref, ref_aval, handle_transposes=True):
    if ref is None:
      return None
    if tree is None:
      transforms = ()
      transform_avals = ()
    else:
      transforms = tree.unflatten(leaves)
      transform_avals = tree.unflatten(leaves_avals)
    ref, _, transforms = lowering._handle_transforms(
        ctx, ref_aval, ref, transform_avals, transforms, handle_transposes=handle_transposes
    )
    if transforms:
      raise NotImplementedError(
          f"Unsupported transforms for {ref}. Transforms {transforms}."
      )
    return ref

  acc_ref = handle_transforms_and_get_ref(
      acc_transforms_tree,
      acc_transforms_leaves,
      acc_transforms_leaves_avals,
      acc_ref,
      acc_aval,
      handle_transposes=False,
  )

  a_ref = handle_transforms_and_get_ref(
      a_transforms_tree,
      a_transforms_leaves,
      a_transforms_leaves_avals,
      a_ref,
      a_aval,
      handle_transposes=a_aval.memory_space == gpu_core.SMEM
  )

  b_ref = handle_transforms_and_get_ref(
      b_transforms_tree,
      b_transforms_leaves,
      b_transforms_leaves_avals,
      b_ref,
      b_aval,
  )

  a_sparse_metadata_ref = handle_transforms_and_get_ref(
      a_sparse_metadata_transforms_tree,
      a_sparse_metadata_transforms_leaves,
      a_sparse_metadata_transforms_leaves_avals,
      a_sparse_metadata_ref,
      a_sparse_metadata_ref_aval,
  )

  if barrier_transforms_tree is not None and barrier_ref is not None:
    barrier_transforms = barrier_transforms_tree.unflatten(
        barrier_transforms_leaves
    )
    base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)
    if base_index is not None:
      barrier_ref = barrier_ref[base_index]

  a_scale_ref = handle_transforms_and_get_ref(
      a_scale_transforms_tree,
      a_scale_transforms_leaves,
      a_scale_transforms_leaves_avals,
      a_scale_ref,
      a_scale_ref_aval,
  )

  b_scale_ref = handle_transforms_and_get_ref(
      b_scale_transforms_tree,
      b_scale_transforms_leaves,
      b_scale_transforms_leaves_avals,
      b_scale_ref,
      b_scale_ref_aval,
  )

  predicate_ctx: contextlib.AbstractContextManager[None]
  if collective_axis is not None:
    predicate_ctx = mgpu.when(_collective_mma_predicate(ctx, collective_axis))
    collective = True
  else:
    predicate_ctx = contextlib.nullcontext()
    collective = False

  if isinstance(accumulate, bool):
    i1 = ir.IntegerType.get_signless(1)
    accumulate = arith_dialect.constant(i1, accumulate)

  with predicate_ctx:
    mgpu.dialect.tcgen05_mma(
        acc_ref,
        a_ref,
        b_ref,
        accumulate=accumulate,
        collective=collective,
        a_scale=a_scale_ref,
        b_scale=b_scale_ref,
        a_sparse_metadata=a_sparse_metadata_ref,
    )
    if arrive:
      assert isinstance(barrier_ref, mgpu.DialectBarrierRef)
      mgpu.dialect.tcgen05_commit_arrive(
          barrier_ref.as_barrier_memref(), collective=collective
      )
  return []


tcgen05_commit_arrive_p = jax_core.Primitive("tcgen05_commit_arrive")
tcgen05_commit_arrive_p.multiple_results = True


def tcgen05_commit_arrive(
    barrier: _Ref,
    collective_axis: str | None = None,
    *,
    predicate: bool | jax.Array | None = None,
) -> None:
  """Tracks completion of all preceding ``tcgen05_mma`` and ``async_copy_smem_to_tmem`` calls.

  Args:
    barrier: Barrier Ref for synchronizing with the tensor core. Must have
      orders_tensor_core set to True.
    collective_axis: The name of the cluster axis along which the
      operations were performed if it was collective. The cluster axis should
      have a size of exactly 2, and must be on the minormost cluster axis.
    predicate: A boolean indicating whether the commit arrive should be
      performed. If ``None``, the commit arrive is always performed.

  See also:
    - :func:`jax.experimental.pallas.mosaic_gpu.tcgen05_mma`
    - :func:`jax.experimental.pallas.mosaic_gpu.async_copy_smem_to_tmem`
  """
  if isinstance(barrier, pallas_core.TransformedRef):
    barrier_transforms_leaves, barrier_transforms_tree = jax.tree.flatten(
        barrier.transforms
    )
    barrier = barrier.ref
  else:
    barrier_transforms_leaves, barrier_transforms_tree = [], None

  tcgen05_commit_arrive_p.bind(
      barrier,
      *barrier_transforms_leaves,
      *() if predicate is None else (predicate,),
      barrier_transforms_tree=barrier_transforms_tree,
      collective_axis=collective_axis,
      has_user_predicate=predicate is not None,
  )


@tcgen05_commit_arrive_p.def_effectful_abstract_eval
def _tcgen05_commit_arrive_abstract_eval(
    barrier,
    *args,
    barrier_transforms_tree,
    collective_axis,
    has_user_predicate: bool = False,
):
  del args, barrier_transforms_tree, collective_axis, has_user_predicate
  orders_tensor_core = getattr(
      barrier.inner_aval.dtype, "orders_tensor_core", False)
  if not orders_tensor_core:
    raise ValueError("MMA barrier must have orders_tensor_core set to True.")
  return [], {gpu_core._memory_effect}


@lowering.register_lowering_rule(
    tcgen05_commit_arrive_p, *gpu_core.LANExWG_SEMANTICS)
@lowering.register_lowering_rule(
    tcgen05_commit_arrive_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(
    tcgen05_commit_arrive_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(
    tcgen05_commit_arrive_p, *gpu_core.WGxWARP_SEMANTICS
)
def _tcgen05_commit_arrive_lowering(
    ctx: lowering.LoweringRuleContext,
    barrier_ref: mgpu.BarrierRef | mgpu.DialectBarrierRef,
    *flat_args,
    barrier_transforms_tree,
    collective_axis,
    has_user_predicate: bool = False,
):
  if has_user_predicate:
    *flat_args, user_predicate = flat_args
    predicate = lowering._ensure_ir_value(user_predicate, jnp.bool)  # pylint: disable=protected-access
  else:
    predicate = None

  barrier_ref_aval = ctx.avals_in[0]
  assert isinstance(barrier_ref_aval, state_types.AbstractRef)
  if barrier_transforms_tree is not None:
    barrier_transforms = barrier_transforms_tree.unflatten(flat_args)
    base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)
    if base_index is not None:
      barrier_ref = barrier_ref[base_index]

  if collective := collective_axis is not None:
    is_leader_block = _collective_mma_predicate(ctx, collective_axis)
    predicate = _andi_maybe_none(predicate, is_leader_block)

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    assert isinstance(barrier_ref, mgpu.DialectBarrierRef)
    memref = barrier_ref.as_barrier_memref()
    # TODO(cjfj): remove when minimum jaxlib version is 0.12.
    if hasattr(mgpu.dialect.TcGen05CommitArriveOp, "predicate"):
      mgpu.dialect.tcgen05_commit_arrive(
          memref, collective=collective, predicate=predicate  # pyrefly: ignore[unexpected-keyword]
      )
    elif predicate is None:
      mgpu.dialect.tcgen05_commit_arrive(memref, collective=collective)
    else:
      with mgpu.when(predicate):
        mgpu.dialect.tcgen05_commit_arrive(memref, collective=collective)
  else:
    assert isinstance(barrier_ref, mgpu.BarrierRef)
    single_lane = ctx.module_ctx.single_lane_predicate
    tcgen05.commit_arrive(
        barrier_ref,
        collective=collective,
        ctx=ctx.launch_ctx,
        predicate=_andi_maybe_none(predicate, single_lane),
    )
  return []


def _collective_mma_predicate(ctx: lowering.LoweringRuleContext,
                              collective_axis: str) -> ir.Value:
  """Computes a predicate to run only on the leader block."""
  cluster_axis = lowering._resolve_cluster_axis(
      ctx.module_ctx.axis_names, collective_axis)
  if cluster_axis != gpu_dialect.Dimension(0):
    # Note: resolve_cluster_axis checks if axis_names exists.
    assert ctx.module_ctx.axis_names is not None
    if len(ctx.module_ctx.axis_names.cluster) <= 1:
      raise ValueError("No cluster axes found.")
    minormost_cluster_axis = ctx.module_ctx.axis_names.cluster[0]
    raise ValueError(
        "Can only perform collective MMA along minormost cluster axis. "
        f"Got {collective_axis}, expected {minormost_cluster_axis}.")
  index = ir.IndexType.get()
  is_leader_block = arith_dialect.cmpi(
      arith_dialect.CmpIPredicate.eq,
      mgpu.utils.cluster_idx(cluster_axis), mgpu.c(0, index))
  return is_leader_block


commit_tmem_p = jax_core.Primitive("commit_tmem")
commit_tmem_p.multiple_results = True


@commit_tmem_p.def_effectful_abstract_eval
def _commit_tmem_abstract_eval():
  return (), {gpu_core._memory_effect}


@lowering.register_lowering_rule(commit_tmem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(
    commit_tmem_p, mgpu.LoweringSemantics.Warpgroup
)
def _commit_tmem_lowering(_):
  tcgen05.commit_tmem()
  return ()


def commit_tmem():
  """Commits all writes to TMEM issued by the current thread.

  Once this function returns, the effects of calling ``async_store_tmem`` from
  the current thread are visible to TMEM loads, MMA and barrier operations of
  ``Barrier``s with ``orders_tensor_core=True``.
  """
  commit_tmem_p.bind()


set_max_registers_p = jax_core.Primitive("set_max_registers")
set_max_registers_p.multiple_results = True


@set_max_registers_p.def_effectful_abstract_eval
def _set_max_registers_abstract_eval(*, n: int, action):
  # The register count of `setmaxnreg` must be a multiple of 8.
  if not isinstance(n, int) or n <= 0 or n % 8:
    raise ValueError(f"n must be a positive multiple of 8, got {n}")
  if action not in ("increase", "decrease"):
    raise ValueError(
        f"action must be 'increase' or 'decrease', got {action!r}"
    )
  return (), {gpu_core._memory_effect}


@lowering.register_lowering_rule(
    set_max_registers_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(
    set_max_registers_p, mgpu.LoweringSemantics.Warpgroup)
def _set_max_registers_lowering(
    ctx: lowering.LoweringRuleContext, *, n: int, action
):
  del ctx
  action_str = "inc" if action == "increase" else "dec"
  llvm_dialect.inline_asm(
      ir.Type.parse("!llvm.void"),
      [],
      f"setmaxnreg.{action_str}.sync.aligned.u32 {n};",
      "",
      has_side_effects=True,
  )
  return ()


def set_max_registers(n: int, *, action: Literal["increase", "decrease"]):
  """Sets the maximum number of per-lane registers in the thread."""
  set_max_registers_p.bind(n=n, action=action)


commit_smem_p = jax_core.Primitive("commit_smem")
commit_smem_p.multiple_results = True


@commit_smem_p.def_effectful_abstract_eval
def _commit_smem_abstract_eval():
  return (), {gpu_core._memory_effect}


@lowering.register_lowering_rule(commit_smem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(
    commit_smem_p, mgpu.LoweringSemantics.Warpgroup)
def _commit_smem_lowering(ctx: lowering.LoweringRuleContext):
  # TODO(bchetioui): add primitive for commit smem to mosaic_gpu dialect.
  mgpu.commit_shared()
  return ()


def commit_smem():
  """Commits all reads from/writes to SMEM, making them visible to TMA and MMA operations."""
  commit_smem_p.bind()


def broadcasted_iota(
    dtype: jax.typing.DTypeLike,
    shape: Sequence[int],
    dimension: int,
    *,
    layout: SomeLayout | None = None,
) -> jax.Array:
  result = jax.lax.broadcasted_iota(dtype, shape, dimension)
  if layout is not None:
    result = gpu_core.layout_cast(result, layout)
  return result


@dataclasses.dataclass(frozen=True)
class ShapeDtypeStruct:
  shape: tuple[int, ...]
  dtype: jnp.dtype
  layout: SomeLayout


griddepcontrol_wait_p = jax_core.Primitive("griddepcontrol_wait")
griddepcontrol_wait_p.multiple_results = True


@griddepcontrol_wait_p.def_effectful_abstract_eval
def _griddepcontrol_wait_abstract_eval():
  return (), {gpu_core._memory_effect, gpu_core._pdl_effect}


@lowering.register_lowering_rule(
    griddepcontrol_wait_p, *gpu_core.LANExWG_SEMANTICS
)
@lowering.register_lowering_rule(
    griddepcontrol_wait_p, *gpu_core.WGxWG_SEMANTICS
)
@lowering.register_lowering_rule(
    griddepcontrol_wait_p, *gpu_core.LANExWARP_SEMANTICS
)
@lowering.register_lowering_rule(
    griddepcontrol_wait_p, *gpu_core.WGxWARP_SEMANTICS
)
def _griddepcontrol_wait_lowering(ctx: lowering.LoweringRuleContext):
  void = ir.Type.parse("!llvm.void")
  with lowering._wrap_in_custom_primitive_if_wg(ctx, []):
    llvm_dialect.inline_asm(
        void, [], "griddepcontrol.wait;", "", has_side_effects=True
    )
  return ()


def griddepcontrol_wait():
  """Wait for dependent grids to finish."""
  griddepcontrol_wait_p.bind()


griddepcontrol_launch_dependents_p = jax_core.Primitive(
    "griddepcontrol_launch_dependents"
)
griddepcontrol_launch_dependents_p.multiple_results = True


@griddepcontrol_launch_dependents_p.def_effectful_abstract_eval
def _griddepcontrol_launch_dependents_abstract_eval():
  return (), {gpu_core._memory_effect}


@lowering.register_lowering_rule(
    griddepcontrol_launch_dependents_p, *gpu_core.LANExWG_SEMANTICS
)
@lowering.register_lowering_rule(
    griddepcontrol_launch_dependents_p, *gpu_core.WGxWG_SEMANTICS
)
@lowering.register_lowering_rule(
    griddepcontrol_launch_dependents_p, *gpu_core.LANExWARP_SEMANTICS
)
@lowering.register_lowering_rule(
    griddepcontrol_launch_dependents_p, *gpu_core.WGxWARP_SEMANTICS
)
def _griddepcontrol_launch_dependents_lowering(
    ctx: lowering.LoweringRuleContext,
):
  void = ir.Type.parse("!llvm.void")
  with lowering._wrap_in_custom_primitive_if_wg(ctx, []):
    llvm_dialect.inline_asm(
        void, [], "griddepcontrol.launch_dependents;", "", has_side_effects=True
    )
  return ()


def griddepcontrol_launch_dependents():
  """Signal that dependents can be launched."""
  griddepcontrol_launch_dependents_p.bind()


inline_mgpu_p = jax_core.Primitive("inline_mgpu")
inline_mgpu_p.multiple_results = True


@dataclasses.dataclass(frozen=True)
class RefType:
  transforms: tuple[state_types.Transform, ...] = ()


def inline_mgpu(*, arg_types=(), return_type=None):
  r"""Returns a decorator that inlines Mosaic GPU code.

  This allows using lower-level Mosaic GPU abstractions and operations, which
  are otherwise not directly exposed in Pallas.

  Example::

      layout = plgpu.Layout.WG_STRIDED(x_ref.shape, vec_size=4)

      @plgpu.inline_mgpu(
          arg_types=(plgpu.RefType(),),
          return_type=plgpu.ShapeDtypeStruct(
              (128, 128), dtype, layout=layout
          ),
      )
      def add_one(ctx, smem_ref):
        x = mgpu.FragmentedArray.load_tiled(smem_ref)
        y = mgpu.FragmentedArray.splat(
            mgpu.c(1, x.mlir_dtype), shape=x.shape, layout=x.layout
        )
        return x + y

  Args:
    arg_types: A sequence of pytrees where the leaves are
      :class:`~jax.experimental.pallas.mosaic_gpu.RefType`\s or
      :class:`~jax.experimental.pallas.mosaic_gpu.Layout`\s for reference or
      array arguments respectively.
    return_type: A pytree where the leaves are
      :class:`~jax.experimental.pallas.mosaic_gpu.ShapeDtypeStruct`\s
      representing the arrays returned by the decorated function.
  """
  flat_arg_types, treedef_ty = jax.tree.flatten(tuple(arg_types))
  flat_ret_ty, pytree_ret_ty = jax.tree.flatten(return_type)
  if return_type and not all(isinstance(r, ShapeDtypeStruct) for r in flat_ret_ty):
    raise ValueError(
        "inline_mgpu_p only supports plgpu.ShapeDtypeStruct return types."
    )
  if not all(isinstance(r, (SomeLayout, RefType)) for r in flat_arg_types):
    raise ValueError(
        "inline_mgpu_p only supports only SomeLayout and RefType arg types."
    )

  def inner(f):
    def wrapper(*args):
      flat_args, treedef = jax.tree.flatten(tuple(args))
      if treedef != treedef_ty:
        raise ValueError(f"Mismatched type shape: {treedef} != {treedef_ty}")

      # Strip the transforms from the refs since they will be recorded in
      # the types.
      ref_transforms: list[Any] = []
      raw_flat_args = []
      for a, t in zip(flat_args, flat_arg_types):
        if isinstance(a, state_types.TransformedRef) and isinstance(t, RefType):
          raw_flat_args.append(a.ref)
          ref_transforms.append(a.transforms)
        elif isinstance(aval := jax_core.typeof(a), jax_core.ShapedArray) and isinstance(t, SomeLayout):
          raw_flat_args.append(a)
          ref_transforms.append(None)
        elif isinstance(aval, state.AbstractRef) and isinstance(t, RefType):
          raw_flat_args.append(a)
          ref_transforms.append(())
        else:
          raise ValueError(f"Mismatched type: {a, t}")

      flat_ref_transforms, pytree_ref_transforms = jax.tree.flatten(ref_transforms)
      flat_ret = inline_mgpu_p.bind(
          *raw_flat_args,
          *flat_ref_transforms,
          flat_arg_types=tuple(flat_arg_types),
          flat_ret_ty=tuple(flat_ret_ty),
          pytree_ret_ty=pytree_ret_ty,
          pytree_args=treedef,
          pytree_ref_transforms=pytree_ref_transforms,
          mgpu_fn=f,
      )
      return jax.tree.unflatten(pytree_ret_ty, flat_ret)
    return wrapper

  return inner


@inline_mgpu_p.def_effectful_abstract_eval
def _inline_mgpu_abstract_eval(
    *flat_args_and_transforms,
    flat_arg_types,
    flat_ret_ty,
    pytree_args,
    pytree_ref_transforms,
    pytree_ret_ty,
    mgpu_fn,
):
  del flat_arg_types, pytree_ret_ty, pytree_ref_transforms, mgpu_fn  # Unused.
  aval_return = tuple(
      jax_core.ShapedArray(x.shape, x.dtype) for x in flat_ret_ty
  )
  # TODO(cperivol): Let the user set the effects.
  flat_args = flat_args_and_transforms[:pytree_args.num_leaves]
  return aval_return, {
      gpu_core._wgmma_pipeline_effect,
      gpu_core._memory_effect,
      *itertools.chain.from_iterable(
          (state.ReadEffect(i), state.WriteEffect(i))
          for i, r in enumerate(flat_args)
          if isinstance(r, state.AbstractRef)
      ),
  }


@discharge.register_discharge_rule(inline_mgpu_p)
def _inline_mgpu_discharge(*args, **kwargs):
  del args, kwargs
  raise NotImplementedError("inline_mgpu_p does not support discharge.")


def _type_check_mgpu_lane_semantics(v, ty):
  match (ty, v):
    case (RefType(), ir.Value()) if isinstance(v.type, ir.MemRefType):
      pass
    case (ShapeDtypeStruct(), mgpu.FragmentedArray()):
      mlir_dtype = mgpu_utils.dtype_to_ir_type(ty.dtype)
      if v.mlir_dtype != mlir_dtype:
        raise ValueError(
            f"Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtype}."
        )
      if ty.shape != v.shape:
        raise ValueError(
            f"Array shape mismatch: expected {ty.shape} got {v.shape}."
        )
      if v.layout != ty.layout.to_mgpu():
        raise ValueError(
            f"Array layout mismatch: expected {v.layout} got {ty.layout.to_mgpu()}."
        )
    case (SomeLayout(), mgpu.FragmentedArray()):
      if ty.to_mgpu() != v.layout:
        raise ValueError(f"Unexpected layout for {v} (expected: {ty})")
    case _:
      raise ValueError(f"Unexpected type {ty} for value {v}")


def _inline_mgpu_flat_transformed_args(
    ctx: lowering.LoweringRuleContext,
    flat_args_and_transforms,
    flat_arg_types,
    pytree_args,
    pytree_ref_transforms,
  ) -> Sequence[ir.Value | mgpu.FragmentedArray]:
  flat_args = flat_args_and_transforms[:pytree_args.num_leaves]
  flat_arg_avals = ctx.avals_in[:pytree_args.num_leaves]
  ref_transforms = pytree_ref_transforms.unflatten(flat_args_and_transforms[pytree_args.num_leaves:])
  ref_transform_avals = pytree_ref_transforms.unflatten(ctx.avals_in[pytree_args.num_leaves:])
  is_wg_semantics = (
      ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup
  )
  is_warp_semantics = (
      ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp
  )

  if is_wg_semantics:
    flat_args = [
        lowering._ensure_ir_value(a, aval.dtype) if not isinstance(t, RefType) else a
        for a, aval, t in zip(flat_args, flat_arg_avals, flat_arg_types)
    ]
  else:
    flat_args = [
        lowering._ensure_fa(a, aval.dtype) if not isinstance(t, RefType) else a
        for a, aval, t in zip(flat_args, flat_arg_avals, flat_arg_types)
    ]

  for a, aval, t in zip(flat_args, flat_arg_avals, flat_arg_types):
    if not is_wg_semantics:
      _type_check_mgpu_lane_semantics(a, t)
    if is_warp_semantics and not isinstance(t, RefType):
      if not isinstance(aval, jax_core.ShapedArray) or aval.shape:
        raise ValueError(
            "inline_mgpu in a single-warp context only supports scalar"
            f" arrays (and Refs). Got {aval}."
        )

  flat_transformed : list[ir.Value | mgpu.FragmentedArray] = []
  for a, aval, t, transforms, transform_avals in zip(
      flat_args,
      flat_arg_avals,
      flat_arg_types,
      ref_transforms,
      ref_transform_avals,
      strict=True,
  ):
    if not isinstance(t, RefType):
      flat_transformed.append(a)
      assert transforms is None
      continue
    assert isinstance(aval, state.AbstractRef)
    assert isinstance(a, ir.Value)
    a, aval, user_transforms = lowering._handle_transforms(
        ctx,
        aval,
        a,
        transform_avals,
        transforms,
        handle_transposes=is_wg_semantics,
        allow_peer_refs=True,
    )

    if is_wg_semantics:
      if user_transforms:
        raise NotImplementedError(
            "Not all transforms could be handled. Remaining transforms:"
            f" {user_transforms}."
        )
    else:
      # Transforms that do not originate from a MemoryRefTransform are
      # applied implicitly (eg by emit-pipeline) and therefore we do not
      # expect the user to pass them to the type. The transforms not
      # passed by the user here will be discharged.
      ty_transforms = tuple(pallas_core.undo_transforms(aval, t.transforms))
      if ty_transforms != tuple(user_transforms):
        raise ValueError(
            f"Transform mismatch: got {user_transforms}, expected"
            f" {ty_transforms}"
        )
    flat_transformed.append(a)

  return flat_transformed


@lowering.register_lowering_rule(inline_mgpu_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(inline_mgpu_p, *gpu_core.LANExWARP_SEMANTICS)
def _inline_mgpu_lowering_rule(
    ctx: lowering.LoweringRuleContext,
    *flat_args_and_transforms,
    mgpu_fn: Callable[..., Any],
    flat_arg_types,
    flat_ret_ty,
    pytree_args,
    pytree_ref_transforms,
    pytree_ret_ty,
):
  is_warp_semantics = (
      ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp
  )
  if is_warp_semantics:
    for r in flat_ret_ty:
      if isinstance(r, ShapeDtypeStruct) and r.shape:
        raise ValueError(
            "inline_mgpu in a single-warp context only supports scalar return"
            f" types. Got shape={r.shape}."
        )

  flat_transformed = _inline_mgpu_flat_transformed_args(
      ctx,
      flat_args_and_transforms,
      flat_arg_types,
      pytree_args,
      pytree_ref_transforms,
  )
  args = jax.tree.unflatten(pytree_args, flat_transformed)
  ret = mgpu_fn(ctx.launch_ctx, *args)
  ret_leaves, ret_tree = jax.tree.flatten(
      ret, lambda x: isinstance(x, mgpu.FragmentedArray)
  )

  if ret_tree != pytree_ret_ty:
    return_type = jax.tree.unflatten(pytree_ret_ty, flat_ret_ty)
    raise ValueError(
        f"inline_mgpu_p return type tree mismatch: {ret} != {return_type}"
    )

  for ty, r in zip(flat_ret_ty, ret_leaves):
    _type_check_mgpu_lane_semantics(r, ty)

  return ret_leaves


def _ref_type_to_transforms(ref_type: RefType) -> ir.ArrayAttr:
  """Returns the Mosaic GPU transforms for the given ref type."""
  transform_attrs = [gpu_core.to_transform_attr(t)
                     for t in ref_type.transforms]
  return ir.ArrayAttr.get(transform_attrs)


def _custom_primitive_in_specs(
    ctx: lowering.LoweringRuleContext,
    flat_arg_types,
    flat_transformed_args,
    pytree_args,
) -> tuple[Sequence[ir.Type], Sequence[ir.Attribute], Sequence[ir.ArrayAttr]]:
  """Returns a tuple containing the list of MLIR input types, layouts, and
  transforms for the given JAX array and ref arguments."""
  in_types : list[ir.Type] = []
  in_layouts : list[ir.Attribute] = []
  in_transforms : list[ir.ArrayAttr] = []
  flat_arg_avals = ctx.avals_in[:pytree_args.num_leaves]
  for aval, transformed, t in zip(
      flat_arg_avals, flat_transformed_args, flat_arg_types
  ):
    match aval:
      case state.AbstractRef():
        initial_ty = ir.MemRefType(transformed.type)
        in_types.append(initial_ty)
        if mgpu_utils.is_smem_ref(initial_ty):
          in_transforms.append(_ref_type_to_transforms(t))
      case jax_core.ShapedArray() if isinstance(t, SomeLayout):
        el_type = mgpu_utils.dtype_to_ir_type(aval.dtype)
        if len(aval.shape) == 0:
          in_types.append(el_type)
        else:
          vector_type = ir.VectorType.get(aval.shape, el_type)
          in_types.append(vector_type)
          in_layouts.append(mgpu_layouts.to_layout_attr(t.to_mgpu()))
      case _:
        raise NotImplementedError(
            f"Unsupported aval type: {aval}, {type(aval)}, {t}"
        )
  return in_types, in_layouts, in_transforms


def _custom_primitive_op_results(flat_ret_ty) -> tuple[
    Sequence[ir.Type],
    Sequence[ir.Attribute | None],
]:
  """Returns a tuple containing the list of output MLIR types, and layouts for
  the given JAX return types."""
  results_ty: list[ir.Type] = []
  out_layouts: list[ir.Attribute | None] = []
  for r in flat_ret_ty:
    if not isinstance(r, ShapeDtypeStruct):
      raise NotImplementedError(f"Expected a ShapeDtypeStruct, but got: {r}")
    el_type = mgpu_utils.dtype_to_ir_type(r.dtype)
    if not r.shape:  # scalar case.
      results_ty.append(el_type)
      out_layouts.append(None)
    else:
      results_ty.append(ir.VectorType.get(r.shape, el_type))
      layout = mgpu_layouts.to_layout_attr(r.layout.to_mgpu())
      out_layouts.append(layout)
  return results_ty, out_layouts


def _populate_custom_primitive_op_block(
    ctx: lowering.LoweringRuleContext,
    block: ir.Block,
    mgpu_fn: Callable[..., Any],
    pytree_args,
    in_layouts: Sequence[ir.Attribute],
    in_transforms: Sequence[ir.ArrayAttr],
    results_ty: Sequence[ir.Type],
    out_layouts: Sequence[ir.Attribute | None],
):
  """Calls the given mgpu_fn to populate the block, handling inputs and outputs.

  Block arguments that are references to SMEM or vectors are unwrapped to
  transformed references and fragmented arrays before they are passed to the
  python function mgpu_fn.

  The resulting fragmented arrays, if any, are wrapped as vectors before they
  are returned.
  """
  with ir.InsertionPoint(block):
    fn_inputs: list[ir.Value | mgpu.FragmentedArray] = []
    in_layouts_it = iter(in_layouts)
    in_transforms_it = iter(in_transforms)
    avals_in = ctx.avals_in[:pytree_args.num_leaves]
    for arg, aval in zip(block.arguments, avals_in, strict=True):
      if isinstance(arg.type, ir.MemRefType):
        memref_ty = ir.MemRefType(arg.type)
        if not mgpu_utils.is_smem_ref(memref_ty):
          fn_inputs.append(arg)
          continue

        transforms = ir.ArrayAttr(next(in_transforms_it))
        # The block arguments in the Mosaic GPU dialect are logical refs that
        # wrap the transfromed refs. Since the mgpu_fn works at the lowered
        # "lane" level, we need to transform (lower) the inputs before passing
        # them to the mgpu_fn.
        transformed_type = mgpu.dialect_lowering.transform_type(
            memref_ty, transforms
        )
        conversion_cast = builtin_dialect.UnrealizedConversionCastOp(
            [transformed_type], [arg]
        )
        fn_inputs.append(conversion_cast.result)
      elif isinstance(arg.type, ir.VectorType):
        layout_attr = next(in_layouts_it)
        layout = mgpu.layouts.from_layout_attr(layout_attr)

        vector_ty = ir.VectorType(arg.type)
        reg_shape = layout.registers_shape(tuple(vector_ty.shape))
        reg_ty = layout.registers_element_type(vector_ty.element_type)

        # The vector block arguments in the Mosaic GPU dialect are wrapped
        # Fragmented Arrays. Since the mgpu_fn works at the lowered
        # "lane" level, we need to unwrap (lower) the input vectors before
        # passing them to the mgpu_fn.
        conversion_cast = builtin_dialect.UnrealizedConversionCastOp(
            [reg_ty] * math.prod(reg_shape), [arg]
        )
        conversion_cast.attributes["registers_shape"] = ir.ArrayAttr.get([
            ir.IntegerAttr.get(ir.IntegerType.get_signless(64), s)
            for s in reg_shape
        ])
        conversion_cast.attributes["layout"] = layout_attr

        registers = np.array(list(conversion_cast.results)).reshape(reg_shape)
        is_signed = mgpu_utils.is_signed(aval.dtype)
        fa = mgpu.FragmentedArray(
            _registers=registers, _layout=layout, _is_signed=is_signed
        )
        fn_inputs.append(fa)
      else:  # scalar case.
        is_signed = mgpu_utils.is_signed(aval.dtype)
        fa = mgpu.FragmentedArray.splat(arg, (), is_signed=is_signed)
        fn_inputs.append(fa)

    args = jax.tree.unflatten(pytree_args, fn_inputs)
    inner_ret = mgpu_fn(ctx.launch_ctx, *args)
    if inner_ret is None:
      inner_ret = []
    elif not isinstance(inner_ret, tuple) and not isinstance(inner_ret, list):
      inner_ret = [inner_ret]
    ir_ret = []
    for fa, result_ty, out_layout in zip(
        inner_ret, results_ty, out_layouts, strict=True
    ):
      if not isinstance(fa, mgpu.FragmentedArray):
        raise ValueError(f"Expected a FragmentedArray, but got: {fa}")
      if isinstance(result_ty, ir.VectorType):
        result_shape = ir.VectorType(result_ty).shape
        if fa.shape != tuple(result_shape):
          raise ValueError(f"Expected {result_shape} but got {fa.shape}")
        if out_layout != mgpu.layouts.to_layout_attr(fa.layout):
          raise ValueError(
              f"Output layout {out_layout} does not match the layout of the"
              f" returned fragmented array {fa.layout}."
          )
        ir_ret.append(
            mgpu.dialect_lowering.fragmented_array_to_ir(fa, result_ty)
        )
      else:  # scalar case.
        assert out_layout is None
        if fa.shape:
          raise ValueError(f"Expected 0D shape, but got {fa.shape}")
        if not isinstance(fa.layout, mgpu.WGSplatFragLayout):
          raise ValueError(f"Expected WGSplatFragLayout, but got {fa.layout}")
        value = fa.registers.item()
        ir_ret.append(value)

    mgpu.dialect.return_(ir_ret)


@lowering.register_lowering_rule(inline_mgpu_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(inline_mgpu_p, *gpu_core.WGxWARP_SEMANTICS)
def _inline_mgpu_lowering_rule_wg_semantics(
    ctx: lowering.LoweringRuleContext,
    *flat_args_and_transforms,
    mgpu_fn: Callable[..., Any],
    flat_arg_types,
    flat_ret_ty,
    pytree_args,
    pytree_ref_transforms,
    pytree_ret_ty,
):
  del pytree_ret_ty
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    for r in flat_ret_ty:
      if isinstance(r, ShapeDtypeStruct) and r.shape:
        raise ValueError(
            "inline_mgpu in a single-warp context only supports scalar return"
            f" types. Got shape={r.shape}."
        )

  flat_transformed_args = _inline_mgpu_flat_transformed_args(
      ctx,
      flat_args_and_transforms,
      flat_arg_types,
      pytree_args,
      pytree_ref_transforms,
  )

  in_types, in_layouts, in_transforms = (
      _custom_primitive_in_specs(
          ctx, flat_arg_types, flat_transformed_args, pytree_args
      )
  )
  results_ty, out_layouts = _custom_primitive_op_results(flat_ret_ty)

  custom_op = mgpu.dialect.CustomPrimitiveOp(
      result=results_ty,
      operands_=flat_transformed_args,  # pyrefly: ignore[bad-argument-type]
      in_layouts=in_layouts,
      in_transforms=in_transforms,
      out_layouts=[l for l in out_layouts if l is not None],
  )
  block: ir.Block = custom_op.body.blocks.append(*in_types)
  _populate_custom_primitive_op_block(
      ctx,
      block,
      mgpu_fn,
      pytree_args,
      in_layouts,
      in_transforms,
      results_ty,
      out_layouts,
  )

  # We need to ensure that the block doesn't capture any values from the context
  # and uses args for everything instead. E.g. `LaunchContext.tma_descriptors`
  # will be captured when calling `ctx.async_copy`.
  custom_op = lowering._isolate_from_above(custom_op)

  return custom_op.results


load_p = jax_core.Primitive("load")


@load_p.def_effectful_abstract_eval
def _load_abstract_eval(src, *avals_flat, tree, optimized):
  del optimized  # Unused.
  transforms = list(tree.unflatten(avals_flat))
  if not transforms or not isinstance(transforms[-1], indexing.NDIndexer):
    tref_aval = state.transform_type(transforms, src)
    assert isinstance(tref_aval, state_types.AbstractRef)
    transforms.append(indexing.NDIndexer.make_trivial_indexer(tref_aval.shape))
  out_ty = state.transform_type(transforms, src)
  assert isinstance(out_ty, state_types.AbstractRef)
  return out_ty.inner_aval, {state.ReadEffect(0)}


lowering.register_lowering_rule(load_p, mgpu.LoweringSemantics.Lane)(
    lowering._get_lowering_rule
)
lowering.register_lowering_rule(load_p, *gpu_core.LANExWARP_SEMANTICS)(
    lowering._get_lowering_rule
)
lowering.register_lowering_rule(load_p, mgpu.LoweringSemantics.Warpgroup)(
    lowering._get_lowering_rule_wg
)
lowering.register_lowering_rule(load_p, *gpu_core.WGxWARP_SEMANTICS)(
    lowering._get_lowering_rule_wg
)


def load(
    src: _Ref, *, layout: SomeLayout | None = None, optimized: bool = True
) -> jax.Array:
  """Loads from a reference into an array with the specified layout.

  Args:
    src: The reference to load from. Can be either in SMEM or GMEM.
    layout: The optional layout to use for the resulting array.
    optimized: If True, a compilation error will be raised if no optimized
      implementation for the load is available.

  Returns:
    The loaded array.
  """
  src, src_transforms = state_primitives.get_ref_and_transforms(
      src, None, "load"
  )
  flat_src_transforms, src_transforms_treedef = tree_util.tree_flatten(
      src_transforms
  )
  result = load_p.bind(
      src,
      *flat_src_transforms,
      tree=src_transforms_treedef,
      optimized=optimized,
  )
  if layout is not None:
    result = gpu_core.layout_cast(result, layout)
  return result


store_p = jax_core.Primitive("store")
store_p.multiple_results = True


@store_p.def_effectful_abstract_eval
def _store_abstract_eval(dst, val, *avals_flat, tree, optimized):
  del optimized  # Unused.
  # `store` is `swap` without the result, so we can reuse its checks.
  _, effects = state_primitives._swap_abstract_eval(
      dst, val, *avals_flat, tree=tree
  )
  return (), effects


@lowering.register_lowering_rule(store_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(store_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(store_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(store_p, *gpu_core.WGxWARP_SEMANTICS)
def _store_lowering_rule(
    ctx: lowering.LoweringRuleContext, dst, value, *leaves, tree, optimized
):
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    rule = lowering._swap_lowering_rule_wg
  else:
    rule = lowering._swap_lowering_rule
  return rule(
      ctx, dst, value, *leaves, tree=tree, optimized=optimized, return_old=False
  )


def store(dst: _Ref, value: jax.Array, *, optimized: bool = True) -> None:
  """Stores an array to a reference in SMEM or GMEM.

  Currently, only unoptimized stores to GMEM are supported.

  Args:
    dst: The reference to store to. Can be either in SMEM or GMEM.
    value: The array to store.
    optimized: If True, a compilation error will be raised if no optimized
      implementation for the store is available.
  """
  dst, dst_transforms = state_primitives.get_ref_and_transforms(
      dst, None, "store"
  )
  flat_dst_transforms, dst_transforms_treedef = tree_util.tree_flatten(
      dst_transforms
  )
  store_p.bind(
      dst,
      value,
      *flat_dst_transforms,
      tree=dst_transforms_treedef,
      optimized=optimized,
  )


async_load_tmem_p = jax_core.Primitive("async_load")
async_load_tmem_p.multiple_results = True


@overload
def async_load_tmem(
    src: _Ref, *, layout: SomeLayout | None = ..., reduce: None = ...
) -> jax.Array:
  ...


@overload
def async_load_tmem(
    src: _Ref,
    *,
    layout: SomeLayout | None = ...,
    reduce: Literal["max", "min", "absmax", "absmin"] = ...,
) -> tuple[jax.Array, jax.Array]:
  ...


def async_load_tmem(
    src: _Ref,
    *,
    layout: SomeLayout | None = None,
    reduce: Literal["max", "min", "absmax", "absmin"] | None = None,
) -> jax.Array | tuple[jax.Array, jax.Array]:
  """Performs an asynchronous load from the TMEM array.

  The load operation is only partly asynchronous. The returned array can be used
  immediately, without any additional synchronization. However, it cannot be
  assumed that the read from TMEM has completed when the function returns. If
  you ever attempt to overwrite the read region, you should ensure that
  ``wait_load_tmem`` has been called before that happens. Failure to do so
  can result in nondeterministic data races.

  For example, the following sequence of operations at the end of the kernel is
  valid, even though the TMEM load is never awaited::

    smem_ref[...] = plgpu.async_load_tmem(tmem_ref)
    plgpu.commit_smem()
    plgpu.copy_smem_to_gmem(smem_ref, gmem_ref)
    plgpu.wait_smem_to_gmem(0)

  However, if the kernel was persistent and might reuse the TMEM again, the
  sequence should be extended with a call to ``wait_load_tmem``.

  Args:
    src: The TMEM reference to load from.
    layout: The optional layout hint to use for the resulting array.
    reduce: The optional reduction operation to perform on the loaded data. If
      specified, the function returns an extra array, obtained by reducing the
      loaded data along the last dimension.
  """
  src, src_transforms = state_primitives.get_ref_and_transforms(
      src, None, "async_load_tmem"
  )
  flat_src_transforms, src_transforms_treedef = tree_util.tree_flatten(
      src_transforms
  )
  results = async_load_tmem_p.bind(
      src, *flat_src_transforms, tree=src_transforms_treedef, reduce=reduce
  )
  if reduce is None:
    [result] = results
    if layout is not None:
      result = gpu_core.layout_cast(result, layout)
    return result
  else:
    result, reduced = results
    if layout is not None:
      result = gpu_core.layout_cast(result, layout)
    return result, reduced


@async_load_tmem_p.def_effectful_abstract_eval
def _async_load_tmem_abstract_eval(
    src, *avals_flat, tree, reduce: tcgen05.LoadReduceOp | None = None
):
  if src.memory_space != gpu_core.MemorySpace.TMEM:
    raise ValueError("Async load only supports TMEM refs")
  val_aval, effects = state_primitives._get_abstract_eval(
      src, *avals_flat, tree=tree
  )
  if reduce is None:
    return (val_aval,), effects
  if val_aval.dtype not in map(jnp.dtype, (jnp.float32, jnp.int32, jnp.uint32)):
    raise ValueError(
        f"Unsupported dtype for reduction: {val_aval.dtype}. Only float32, "
        " int32 and uint32 are supported."
    )
  reduced_aval = jax_core.ShapedArray(val_aval.shape[:-1], val_aval.dtype)
  return (val_aval, reduced_aval), effects


@lowering.register_lowering_rule(async_load_tmem_p, mgpu.LoweringSemantics.Lane)
def _async_load_tmem_lowering_rule(
    ctx: lowering.LoweringRuleContext,
    x_ref,
    *leaves,
    tree,
    reduce: Literal["max", "min", "absmax", "absmin"] | None = None,
):
  assert isinstance(x_ref, tcgen05.TMEMRef)
  x_aval = ctx.avals_in[0]
  assert isinstance(x_aval, state_types.AbstractRef)
  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = tree.unflatten(
      ctx.avals_in[1 : 1 + tree.num_leaves]
  )
  x_tmem, _, transforms = lowering._handle_transforms(
      ctx, x_aval, x_ref, transform_avals, transforms, handle_transposes=False,
      handle_reshapes=False)
  if transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for TMEM refs. {transforms=}"
    )
  layout_hint = None
  if isinstance(ctx.out_layout_hint, mgpu.TiledLayout):
    layout_hint = ctx.out_layout_hint
  is_signed = mgpu_utils.is_signed(ctx.avals_out[0].dtype)
  res = x_tmem.load(layout=layout_hint, is_signed=is_signed, reduce=reduce)
  return (res,) if reduce is None else res


@lowering.register_lowering_rule(
    async_load_tmem_p, mgpu.LoweringSemantics.Warpgroup
)
def _async_load_tmem_lowering_rule_wg(
    ctx: lowering.LoweringRuleContext,
    x_ref: ir.Value,
    *leaves,
    tree,
    reduce: tcgen05.LoadReduceOp | None = None,
):
  assert isinstance(x_ref, ir.Value)
  assert isinstance(x_ref.type, ir.MemRefType)
  x_aval = ctx.avals_in[0]
  assert isinstance(x_aval, state_types.AbstractRef)

  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = tree.unflatten(
      ctx.avals_in[1 : 1 + tree.num_leaves]
  )
  x_tmem, _, transforms = lowering._handle_transforms(
      ctx,
      x_aval,
      x_ref,
      transform_avals,
      transforms,
      handle_transposes=False,
      handle_reshapes=False,
  )
  if transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for TMEM refs. {transforms=}"
    )
  if reduce is not None:
    orig_reduce = reduce
    # Make sure we will interpret the reduction as unsigned, since we erase
    # signedness.
    if isinstance(x_aval.dtype, jnp.unsignedinteger) and "abs" not in reduce:
      reduce = "abs" + reduce  # type: ignore
    match reduce:
      case "min":
        reduce_attr = mgpu.dialect.TMEMLoadReduction.Min  # pyrefly: ignore[missing-attribute]
      case "max":
        reduce_attr = mgpu.dialect.TMEMLoadReduction.Max  # pyrefly: ignore[missing-attribute]
      case "absmin":
        reduce_attr = mgpu.dialect.TMEMLoadReduction.AbsMin  # pyrefly: ignore[missing-attribute]
      case "absmax":
        reduce_attr = mgpu.dialect.TMEMLoadReduction.AbsMax  # pyrefly: ignore[missing-attribute]
      case _:
        raise ValueError(f"Unsupported load reduce operation: {orig_reduce}")
    return tuple(mgpu.dialect.async_load_tmem(x_tmem, reduce=reduce_attr))  # type: ignore
  return (mgpu.dialect.async_load_tmem(x_tmem),)


wait_load_tmem_p = jax_core.Primitive("wait_load_tmem")
wait_load_tmem_p.multiple_results = True

def wait_load_tmem():
  """Awaits all previously asynchronous TMEM loads issued by the calling thread.

  Once this function returns, the TMEM loads issued by the calling thread are
  guaranteed to have completed. The read TMEM regions can be safely overwritten
  by the calling thread, or any threads signalled through ``Barrier``s with
  ``orders_tensor_core=True``.
  """
  wait_load_tmem_p.bind()


@wait_load_tmem_p.def_effectful_abstract_eval
def _wait_load_tmem_abstract_eval():
  return (), {gpu_core._memory_effect}


@lowering.register_lowering_rule(wait_load_tmem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(wait_load_tmem_p, mgpu.LoweringSemantics.Warpgroup)
def _wait_load_tmem_lowering(_):
  tcgen05.wait_load_tmem()
  return ()


async_store_tmem_p = jax_core.Primitive("async_store_tmem")
async_store_tmem_p.multiple_results = True

def async_store_tmem(ref: _Ref, value):
  """Stores the value to TMEM.

  The store is asynchronous and is not guaranteed to be visible (e.g. by reads
  or MMA operations) until ``commit_tmem`` has been called.

  Args:
    ref: The TMEM reference to store to.
    value: The value to store.
  """
  ref, ref_transforms = state_primitives.get_ref_and_transforms(
      ref, None, "async_store_tmem"
  )
  flat_ref_transforms, ref_transforms_treedef = tree_util.tree_flatten(
      ref_transforms
  )
  async_store_tmem_p.bind(
      ref, value, *flat_ref_transforms, tree=ref_transforms_treedef
  )

@async_store_tmem_p.def_effectful_abstract_eval
def _async_store_tmem_abstract_eval(ref, val, *avals_flat, tree):
  if ref.memory_space != gpu_core.MemorySpace.TMEM:
    raise ValueError("Async store only supports TMEM refs")
  _, effects = state_primitives._swap_abstract_eval(
      ref, val, *avals_flat, tree=tree
  )
  return (), effects

@lowering.register_lowering_rule(async_store_tmem_p, mgpu.LoweringSemantics.Lane)
def _async_store_tmem_lowering_rule(
    ctx: lowering.LoweringRuleContext, x_ref, value, *leaves, tree
):
  assert isinstance(x_ref, tcgen05.TMEMRef)
  x_aval = ctx.avals_in[0]
  assert isinstance(x_aval, state_types.AbstractRef)
  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = tree.unflatten(
      ctx.avals_in[2 : 2 + tree.num_leaves]
  )
  x_tmem, _, transforms = lowering._handle_transforms(
      ctx, x_aval, x_ref, transform_avals, transforms, handle_transposes=False,
      handle_reshapes=False)
  batch_shape = ()
  if transforms and isinstance(
      transforms[0], gpu_core.ExpandLeadingBatchDimensionsTransform
  ):
    batch_shape = transforms[0].batch_shape
    transforms = transforms[1:]
  if transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for TMEM refs. {transforms=}"
    )
  if batch_shape:
    m, n = value.shape[-2:]
    for batch_idx in np.ndindex(batch_shape):
      flat_batch_idx = int(np.ravel_multi_index(batch_idx, batch_shape))
      val_slice = value[batch_idx]
      col_start = flat_batch_idx * n
      tmem_slice = x_tmem.slice(slice(0, m), slice(col_start, col_start + n))
      tmem_slice.store(val_slice)
  else:
    x_tmem.store(value)
  return ()


@lowering.register_lowering_rule(
    async_store_tmem_p, mgpu.LoweringSemantics.Warpgroup
)
def _async_store_tmem_lowering_rule_wg(
    ctx: lowering.LoweringRuleContext,
    x_ref: ir.Value,
    value: ir.Value,
    *leaves,
    tree,
):
  assert isinstance(x_ref, ir.Value)
  assert isinstance(x_ref.type, ir.MemRefType)
  assert isinstance(value, ir.Value)
  assert isinstance(value.type, ir.VectorType)
  x_aval = ctx.avals_in[0]
  assert isinstance(x_aval, state_types.AbstractRef)

  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = tree.unflatten(
      ctx.avals_in[2 : 2 + tree.num_leaves]
  )
  x_tmem, _, transforms = lowering._handle_transforms(
      ctx,
      x_aval,
      x_ref,
      transform_avals,
      transforms,
      handle_transposes=False,
      handle_reshapes=False,
  )
  batch_shape = ()
  if transforms and isinstance(
      transforms[0], gpu_core.ExpandLeadingBatchDimensionsTransform
  ):
    batch_shape = transforms[0].batch_shape
    transforms = transforms[1:]
  if transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for TMEM refs. {transforms=}"
    )
  if batch_shape:
    m, n = ir.VectorType(value.type).shape[-2:]
    for batch_idx in np.ndindex(batch_shape):
      flat_batch_idx = int(np.ravel_multi_index(batch_idx, batch_shape))
      val_slice = vector_dialect.extract(
          value, dynamic_position=[], static_position=batch_idx
      )
      col_start = flat_batch_idx * n
      tmem_slice = mgpu_utils.memref_slice(
          x_tmem, (slice(0, m), slice(col_start, col_start + n))
      )
      mgpu.dialect.async_store_tmem(val_slice, tmem_slice)
  else:
    mgpu.dialect.async_store_tmem(value, x_tmem)
  return ()


async_copy_scales_to_tmem_p = jax_core.Primitive("async_copy_scales_to_tmem")
async_copy_scales_to_tmem_p.multiple_results = True


def async_copy_scales_to_tmem(
    smem_ref: _Ref, tmem_ref: _Ref, collective_axis: AxisName | None = None,
):
  """Copies the MMA scales from SMEM to TMEM.

  The copy is performed asynchronously and can be awaited by calling
  ``tcgen05_commit_arrive`` and waiting on the specified barrier. However, if
  the copy is consumed by an MMA operation issued in the same thread, no
  synchronization is necessary (except for eventually awaiting the MMA operation
  itself).
  """
  smem_ref, smem_transforms = state_primitives.get_ref_and_transforms(
      smem_ref, None, "async_copy_scales_to_tmem"
  )
  flat_smem_transforms, smem_transforms_treedef = tree_util.tree_flatten(
      smem_transforms
  )
  tmem_ref, tmem_transforms = state_primitives.get_ref_and_transforms(
      tmem_ref, None, "async_copy_scales_to_tmem"
  )
  flat_tmem_transforms, tmem_transforms_treedef = tree_util.tree_flatten(
      tmem_transforms
  )
  async_copy_scales_to_tmem_p.bind(
      smem_ref, tmem_ref, *flat_smem_transforms, *flat_tmem_transforms,
      smem_tree=smem_transforms_treedef, tmem_tree=tmem_transforms_treedef,
      collective_axis=collective_axis,
  )


async_copy_sparse_metadata_to_tmem_p = jax_core.Primitive("async_copy_sparse_metadata_to_tmem")
async_copy_sparse_metadata_to_tmem_p.multiple_results = True


def async_copy_sparse_metadata_to_tmem(
    smem_ref: _Ref, tmem_ref: _Ref, collective_axis: AxisName | None = None
):
  """Copies the MMA sparse metadata from SMEM to TMEM.

  The copy is performed asynchronously and can be awaited by calling
  ``tcgen05_commit_arrive`` and waiting on the specified barrier. However, if
  the copy is consumed by an MMA operation issued in the same thread, no
  synchronization is necessary (except for eventually awaiting the MMA operation
  itself).
  """
  smem_ref, smem_transforms = state_primitives.get_ref_and_transforms(
      smem_ref, None, "async_copy_sparse_metadata_to_tmem"
  )
  flat_smem_transforms, smem_transforms_treedef = tree_util.tree_flatten(
      smem_transforms
  )
  tmem_ref, tmem_transforms = state_primitives.get_ref_and_transforms(
      tmem_ref, None, "async_copy_sparse_metadata_to_tmem"
  )
  flat_tmem_transforms, tmem_transforms_treedef = tree_util.tree_flatten(
      tmem_transforms
  )
  async_copy_sparse_metadata_to_tmem_p.bind(
      smem_ref, tmem_ref, *flat_smem_transforms, *flat_tmem_transforms,
      smem_tree=smem_transforms_treedef, tmem_tree=tmem_transforms_treedef,
      collective_axis=collective_axis,
  )


@async_copy_scales_to_tmem_p.def_effectful_abstract_eval
@async_copy_sparse_metadata_to_tmem_p.def_effectful_abstract_eval
def _async_copy_to_tmem_abstract_eval(smem_ref, tmem_ref, *_args, **_kwargs):
  if smem_ref.memory_space != gpu_core.MemorySpace.SMEM:
    raise ValueError("async_copy_scales_to_tmem source must be an SMEM ref")
  if tmem_ref.memory_space != gpu_core.MemorySpace.TMEM:
    raise ValueError("async_copy_scales_to_tmem target must be a TMEM ref")
  return (), {state_types.ReadEffect(0), state_types.WriteEffect(1)}

def _async_copy_to_tmem_lowering_rule(
    impl, ctx: lowering.LoweringRuleContext, smem_ref, tmem_ref, *leaves, smem_tree, tmem_tree, collective_axis
):
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    assert isinstance(tmem_ref, tcgen05.TMEMRef)
  smem_leaves, tmem_leaves = util.split_list(leaves, [smem_tree.num_leaves])
  smem_transforms = jax.tree.unflatten(smem_tree, smem_leaves)
  tmem_transforms = jax.tree.unflatten(tmem_tree, tmem_leaves)
  smem_aval = ctx.avals_in[0]
  assert isinstance(smem_aval, state_types.AbstractRef)
  tmem_aval = ctx.avals_in[1]
  assert isinstance(tmem_aval, state_types.AbstractRef)
  transform_avals = util.split_list(
      ctx.avals_in[2:], [smem_tree.num_leaves]
  )
  smem_transform_avals = smem_tree.unflatten(transform_avals[0])
  tmem_transform_avals = tmem_tree.unflatten(transform_avals[1])
  smem_ref, _, smem_transforms = lowering._handle_transforms(
      ctx, smem_aval, smem_ref, smem_transform_avals, smem_transforms
  )
  tmem_ref, _, tmem_transforms = lowering._handle_transforms(
      ctx, tmem_aval, tmem_ref, tmem_transform_avals, tmem_transforms
  )
  if smem_transforms:
    raise NotImplementedError(f"Unimplemented transforms for SMEM refs: {smem_transforms}")
  if tmem_transforms:
    raise NotImplementedError(f"Unimplemented transforms for TMEM refs: {tmem_transforms}")

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    predicate_ctx: contextlib.AbstractContextManager[None]
    if collective_axis is not None:
      predicate_ctx = mgpu.when(_collective_mma_predicate(ctx, collective_axis))
      collective = True
    else:
      predicate_ctx = contextlib.nullcontext()
      collective = False
    with predicate_ctx:
      impl(smem_ref, tmem_ref, collective=collective)
    return ()

  predicate = ctx.module_ctx.single_lane_predicate
  if collective := collective_axis is not None:
    is_leader_block = _collective_mma_predicate(ctx, collective_axis)
    predicate = _andi_maybe_none(predicate, is_leader_block)

  with mgpu.when(predicate):
    impl(smem_ref, tmem_ref, collective=collective)
  return ()

@lowering.register_lowering_rule(
    async_copy_scales_to_tmem_p, mgpu.LoweringSemantics.Lane
)
@lowering.register_lowering_rule(
    async_copy_scales_to_tmem_p, *gpu_core.LANExWARP_SEMANTICS
)
def _async_copy_scales_to_tmem_lowering_rule(*args, **kwargs):
  return _async_copy_to_tmem_lowering_rule(
      tcgen05.async_copy_scales_smem_to_tmem, *args, **kwargs
  )


@lowering.register_lowering_rule(
    async_copy_scales_to_tmem_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(
    async_copy_scales_to_tmem_p, *gpu_core.WGxWARP_SEMANTICS
)
def _async_copy_scales_to_tmem_lowering_rule_wg(*args, **kwargs):
  return _async_copy_to_tmem_lowering_rule(
      mgpu.dialect.async_store_scales_smem_to_tmem,
      *args,
      **kwargs,
  )


@lowering.register_lowering_rule(
    async_copy_sparse_metadata_to_tmem_p, mgpu.LoweringSemantics.Lane
)
@lowering.register_lowering_rule(
    async_copy_sparse_metadata_to_tmem_p, *gpu_core.LANExWARP_SEMANTICS
)
def _async_copy_sparse_metadata_to_tmem_lowering_rule(*args, **kwargs):
  return _async_copy_to_tmem_lowering_rule(
      tcgen05.async_copy_sparse_metadata_smem_to_tmem, *args, **kwargs
  )


@lowering.register_lowering_rule(
    async_copy_sparse_metadata_to_tmem_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(
    async_copy_sparse_metadata_to_tmem_p, *gpu_core.WGxWARP_SEMANTICS
)
def _async_copy_sparse_metadata_to_tmem_lowering_rule_wg(*args, **kwargs):
  return _async_copy_to_tmem_lowering_rule(
      mgpu.dialect.async_store_sparse_metadata_smem_to_tmem,
      *args,
      **kwargs,
  )


async_copy_smem_to_tmem_p = jax_core.Primitive("async_copy_smem_to_tmem")
async_copy_smem_to_tmem_p.multiple_results = True


def async_copy_smem_to_tmem(
    smem_ref: _Ref,
    tmem_ref: _Ref,
    collective_axis: AxisName | None = None,
):
  """Copies data from SMEM to TMEM using the tcgen05.cp instruction.

  The source SMEM ref must have tiling and swizzle transforms applied. The
  destination TMEM ref must use packed layout (i.e. ``packed=True`` for sub-32b
  types).

  The copy is performed asynchronously. It can be awaited by calling
  ``tcgen05_commit_arrive`` and waiting on the specified barrier. No
  synchronization is necessary if the target of the copy is used by a
  tcgen05_mma operation.

  Args:
    smem_ref: The SMEM reference to copy from.
    tmem_ref: The TMEM reference to copy into.
    collective_axis: The name of the cluster axis along which the
      copy should be performed collectively. The cluster axis should have a
      size of exactly 2, and must be on the minormost cluster axis.
  """
  smem_ref, smem_transforms = state_primitives.get_ref_and_transforms(
      smem_ref, None, "async_copy_smem_to_tmem"
  )
  flat_smem_transforms, smem_transforms_treedef = tree_util.tree_flatten(
      smem_transforms
  )
  tmem_ref, tmem_transforms = state_primitives.get_ref_and_transforms(
      tmem_ref, None, "async_copy_smem_to_tmem"
  )
  flat_tmem_transforms, tmem_transforms_treedef = tree_util.tree_flatten(
      tmem_transforms
  )
  async_copy_smem_to_tmem_p.bind(
      smem_ref, tmem_ref, *flat_smem_transforms, *flat_tmem_transforms,
      smem_tree=smem_transforms_treedef, tmem_tree=tmem_transforms_treedef,
      collective_axis=collective_axis,
  )


@async_copy_smem_to_tmem_p.def_effectful_abstract_eval
def _async_copy_smem_to_tmem_abstract_eval(
    smem_ref, tmem_ref, *args, smem_tree, tmem_tree, **_kwargs
):
  if smem_ref.memory_space != gpu_core.MemorySpace.SMEM:
    raise ValueError("async_copy_smem_to_tmem source must be an SMEM ref")
  if tmem_ref.memory_space != gpu_core.MemorySpace.TMEM:
    raise ValueError("async_copy_smem_to_tmem target must be a TMEM ref")
  smem_transforms, tmem_transforms = util.split_list(
      args, [smem_tree.num_leaves]
  )
  smem_aval = smem_ref
  for t in jax.tree.unflatten(smem_tree, smem_transforms):
    smem_aval = t.transform_type(smem_aval)
  tmem_aval = tmem_ref
  for t in jax.tree.unflatten(tmem_tree, tmem_transforms):
    tmem_aval = t.transform_type(tmem_aval)
  if smem_aval.dtype != tmem_aval.dtype:
    raise ValueError(
        f"Expected SMEM element type ({smem_aval.dtype}) to equal the TMEM"
        f" element type ({tmem_aval.dtype})"
    )
  if smem_aval.shape != tmem_aval.shape:
    raise ValueError(
        f"Expected SMEM reference shape {smem_aval.shape} to equal the TMEM"
        f" reference shape {tmem_aval.shape}"
    )
  return (), {state_types.ReadEffect(0), state_types.WriteEffect(1)}


@lowering.register_lowering_rule(
    async_copy_smem_to_tmem_p, mgpu.LoweringSemantics.Lane
)
@lowering.register_lowering_rule(
    async_copy_smem_to_tmem_p, *gpu_core.LANExWARP_SEMANTICS
)
def _async_copy_smem_to_tmem_lowering_rule(
    ctx: lowering.LoweringRuleContext, smem_ref, tmem_ref, *leaves,
    smem_tree, tmem_tree, collective_axis,
):
  assert isinstance(tmem_ref, tcgen05.TMEMRef)
  smem_leaves, tmem_leaves = util.split_list(leaves, [smem_tree.num_leaves])
  smem_transforms = jax.tree.unflatten(smem_tree, smem_leaves)
  tmem_transforms = jax.tree.unflatten(tmem_tree, tmem_leaves)
  smem_aval = ctx.avals_in[0]
  assert isinstance(smem_aval, state_types.AbstractRef)
  tmem_aval = ctx.avals_in[1]
  assert isinstance(tmem_aval, state_types.AbstractRef)
  transform_avals = util.split_list(
      ctx.avals_in[2:], [smem_tree.num_leaves]
  )
  smem_transform_avals = smem_tree.unflatten(transform_avals[0])
  tmem_transform_avals = tmem_tree.unflatten(transform_avals[1])
  smem_ref, transformed_smem_aval, smem_transforms = lowering._handle_transforms(
      ctx, smem_aval, smem_ref, smem_transform_avals, smem_transforms,
      handle_transposes=False
  )
  tmem_ref, _, tmem_transforms = lowering._handle_transforms(
      ctx, tmem_aval, tmem_ref, tmem_transform_avals, tmem_transforms
  )
  match smem_transforms:
    case (
        gpu_core.UnswizzleRef(swizzle),
        gpu_core.UntilingTransform(tiling),
    ):
      pass
    case (gpu_core.UntilingTransform(tiling),):
      swizzle = 16  # swizzle=16 is equivalent to no swizzle.
    case _:
      raise NotImplementedError(
          f"Unsupported transforms for SMEM ref: {smem_transforms}"
      )
  swizzle_elems = 8 * swizzle // dtypes.itemsize_bits(transformed_smem_aval.dtype)
  if tiling != (8, swizzle_elems):
    raise ValueError(
        f"Tiling does not fit swizzle: expected (8, {swizzle_elems}), but got"
        f" {tiling}"
    )
  if tmem_transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for TMEM refs: {tmem_transforms}"
    )

  predicate = ctx.module_ctx.single_lane_predicate
  if collective := collective_axis is not None:
    is_leader_block = _collective_mma_predicate(ctx, collective_axis)
    predicate = _andi_maybe_none(predicate, is_leader_block)

  with mgpu.when(predicate):
    tcgen05.async_copy_smem_to_tmem(
        smem_ref, tmem_ref, swizzle, collective=collective
    )
  return ()


@lowering.register_lowering_rule(
    async_copy_smem_to_tmem_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(
    async_copy_smem_to_tmem_p, *gpu_core.WGxWARP_SEMANTICS
)
def _async_copy_smem_to_tmem_lowering_rule_wg(
    ctx: lowering.LoweringRuleContext, smem_ref, tmem_ref, *leaves,
    smem_tree, tmem_tree, collective_axis,
):
  smem_leaves, tmem_leaves = util.split_list(leaves, [smem_tree.num_leaves])
  smem_transforms = jax.tree.unflatten(smem_tree, smem_leaves)
  tmem_transforms = jax.tree.unflatten(tmem_tree, tmem_leaves)
  smem_aval = ctx.avals_in[0]
  assert isinstance(smem_aval, state_types.AbstractRef)
  tmem_aval = ctx.avals_in[1]
  assert isinstance(tmem_aval, state_types.AbstractRef)
  transform_avals = util.split_list(
      ctx.avals_in[2:], [smem_tree.num_leaves]
  )
  smem_transform_avals = smem_tree.unflatten(transform_avals[0])
  tmem_transform_avals = tmem_tree.unflatten(transform_avals[1])
  smem_ref, _, smem_transforms = lowering._handle_transforms(
      ctx, smem_aval, smem_ref, smem_transform_avals, smem_transforms,
      handle_transposes=False
  )
  if smem_transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for SMEM ref: {smem_transforms}"
    )
  tmem_ref, _, tmem_transforms = lowering._handle_transforms(
      ctx, tmem_aval, tmem_ref, tmem_transform_avals, tmem_transforms
  )
  if tmem_transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for TMEM refs: {tmem_transforms}"
    )

  predicate_ctx: contextlib.AbstractContextManager[None]
  if collective_axis is not None:
    predicate_ctx = mgpu.when(_collective_mma_predicate(ctx, collective_axis))
    collective = True
  else:
    predicate_ctx = contextlib.nullcontext()
    collective = False

  with predicate_ctx:
    mgpu.dialect.async_store_smem_to_tmem(
        smem_ref, tmem_ref, collective=collective
    )

  return []


semaphore_signal_parallel_p = jax_core.Primitive('semaphore_signal_parallel')
semaphore_signal_parallel_p.multiple_results = True


@dataclasses.dataclass(frozen=True)
class SemaphoreSignal:
  ref: _Ref
  _: dataclasses.KW_ONLY
  device_id: pallas_primitives.DeviceId | None
  inc: int | jax.Array = 1


def semaphore_signal_parallel(*signals: SemaphoreSignal):
  """Signals multiple semaphores without any guaranteed ordering of signal arrivals.

  This primitive is largely equivalent to::

    for sem in semaphores:
      pl.semaphore_signal(sem, inc, device_id=device_id)

  only unlike the loop above, it does not guarantee any ordering of signal
  arrivals. In particular, the target device might observe a signal on
  ``semaphores[1]`` before it observes a signal on ``semaphores[0]``.
  This operation still guarantees that any side effects performed before the
  signal will be fully performed and visible before any of the signals arrive.

  The relaxed requirements make the whole operation significantly cheaper on
  GPUs, as a single expensive memory fence can be used for all signals (instead
  of an expensive fence for each signal).
  """
  semaphores = [s.ref for s in signals]
  device_ids = [
      pallas_primitives.canonicalize_device_id(s.device_id) for s in signals
  ]
  incs = [jnp.asarray(s.inc, dtype=jnp.int32) for s in signals]
  refs, transforms = util.unzip2(
      map(pallas_primitives._get_ref_and_transforms, semaphores)
  )
  args = [refs, transforms, incs, device_ids]
  flat_args, args_tree = tree_util.tree_flatten(args)
  semaphore_signal_parallel_p.bind(
      *flat_args,
      args_tree=args_tree,
  )


@semaphore_signal_parallel_p.def_effectful_abstract_eval
def _semaphore_signal_parallel_abstract_eval(*avals, args_tree):
  (
      sem_avals,
      sem_transforms_avals,
      value_avals,
      device_id_avals,
  ) = tree_util.tree_unflatten(args_tree, avals)
  for sem_aval, sem_transform_avals in zip(sem_avals, sem_transforms_avals, strict=True):
    pallas_primitives.check_sem_avals(sem_aval, sem_transform_avals, "signal")
  if any(va.dtype != jnp.dtype("int32") for va in value_avals):
    raise ValueError(
        "Must signal int32 values, but got"
        f" {[aval.dtype for aval in value_avals]}"
    )
  effs = set()
  for device_id in device_id_avals:
    if device_id is not None:
      device_id_flat_avals = tree_util.tree_leaves(device_id)
      for aval in device_id_flat_avals:
        if aval.dtype != jnp.dtype("int32"):
          raise ValueError(
             f"`device_id`s must be int32 values, but got {aval.dtype}"
          )
      effs.add(pallas_core.comms_effect)
  return [], effs


@lowering.register_lowering_rule(semaphore_signal_parallel_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(semaphore_signal_parallel_p, mgpu.LoweringSemantics.Warpgroup)
def _semaphore_signal_lowering_rule(
    ctx: lowering.LoweringRuleContext, *args, args_tree,
):
  i32 = ir.IntegerType.get_signless(32)
  sems, transforms, values, device_ids = tree_util.tree_unflatten(
      args_tree, args
  )
  sem_avals, transform_avals, value_avals, device_id_avals = tree_util.tree_unflatten(args_tree, ctx.avals_in)
  transformed_sems = []
  for sem, sem_aval, sem_transform_avals, sem_transforms in zip(
      sems, sem_avals, transform_avals, transforms, strict=True
  ):
    assert isinstance(sem_aval, state_types.AbstractRef)
    sem, _, sem_transforms = lowering._handle_transforms(
        ctx, sem_aval, sem, sem_transform_avals, sem_transforms
    )
    if sem_transforms:
      raise NotImplementedError(f"Unhandled transforms for semaphore_signal_parallel: {sem_transforms}")
    transformed_sems.append(sem)
  del sems, transforms  # Use transformed_sems instead.
  for sem, value, device_id, device_id_aval in zip(
      transformed_sems, values, device_ids, device_id_avals, strict=True
  ):
    if device_id is not None:
      device_id = lowering._device_id_to_logical(
          ctx, device_id, pallas_primitives.DeviceIdType.MESH, device_id_aval
      )
      device_id = lowering._ensure_ir_value(device_id, jnp.int32)
      sem = ctx.launch_ctx.to_remote(sem, device_id)
    val = lowering._ir_constant(value, i32)
    with lowering._wrap_in_custom_primitive_if_wg(ctx, [sem, val]) as [sem, val]:
      sem_ptr = mgpu.utils.memref_ptr(sem)
      # TODO(apaszke): Narrow the scope from .sys to .gpu when the semaphore is local.
      # We only signal the semaphore from a single lane, which does not guarantee
      # anything about the state of the other three warps in the warpgroup (they
      # might still be e.g. reading memory that someone will overwrite once they
      # receive a signal).
      if ctx.module_ctx.auto_barriers:
        mgpu.utils.warpgroup_barrier()
      mgpu_utils.SemaphoreRef(sem_ptr).signal(
          val, predicate=ctx.module_ctx.single_wg_lane_predicate, relaxed=True,
      )
      mgpu_utils.fence_release_sys()
  return ()

try_cluster_cancel_p = jax_core.Primitive('try_cluster_cancel')
try_cluster_cancel_p.multiple_results = True

@try_cluster_cancel_p.def_effectful_abstract_eval
def _try_cluster_cancel_abstract_eval(*args, **params):
  del args, params

  return (), {gpu_core._memory_effect}

@lowering.register_lowering_rule(try_cluster_cancel_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(try_cluster_cancel_p, mgpu.LoweringSemantics.Warpgroup)
def try_cluster_cancel_lowering(
    ctx: lowering.LoweringRuleContext,
    result_ref,
    barrier,
    *transforms_leaves,
    result_transforms_tree,
    barrier_transforms_tree,
    collective_axes,
):
  i1 = ir.IntegerType.get_signless(1)
  i32 = ir.IntegerType.get_signless(32)

  result_aval = ctx.avals_in[0]
  assert isinstance(result_aval, state_types.AbstractRef)
  if result_transforms_tree is not None:
    res_transforms_leaves, barrier_transforms_leaves = util.split_list(
      transforms_leaves, [result_transforms_tree.num_leaves])
    res_transforms = result_transforms_tree.unflatten(res_transforms_leaves)
    res_transform_avals = result_transforms_tree.unflatten(
        ctx.avals_in[2 : 2 + result_transforms_tree.num_leaves]
    )
  else:
    barrier_transforms_leaves = transforms_leaves
    res_transforms = ()
    res_transform_avals = ()
  result_ref, _, res_transforms = lowering._handle_transforms(
      ctx, result_aval, result_ref, res_transform_avals, res_transforms)
  if res_transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for result ref: {res_transforms}"
    )

  if barrier_transforms_tree is not None:
    base_index = _get_barrier_base_index(
        ctx.avals_in[1],
        barrier_transforms_tree.unflatten(barrier_transforms_leaves),
    )
    if base_index is not None:
      barrier = barrier[base_index]

  result_ty = ir.MemRefType(result_ref.type)
  bits = math.prod(result_ty.shape) * mgpu.bitwidth(result_ty.element_type)
  if bits != 128:
    raise TypeError(
        f"Try cluster cancel response must be 128 bits, but is {bits} bits."
    )

  axis_names = ctx.module_ctx.axis_names
  eq = functools.partial(arith_dialect.cmpi, arith_dialect.CmpIPredicate.eq)
  zero = mgpu.c(0, ir.IndexType.get())
  is_first_wg = is_first_cta = mgpu.c(1, i1)

  for axis in collective_axes:
    if axis_names.wg is not None and axis == axis_names.wg:
      is_first_wg = eq(mgpu.warpgroup_idx(), mgpu.c(0, i32))
    else:
      cluster_dim = lowering._resolve_cluster_axis(axis_names, axis)
      cluster_idx = mgpu.utils.cluster_idx(cluster_dim)
      is_first_cta = arith_dialect.andi(is_first_cta, eq(cluster_idx, zero))

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:

    # TODO(b/415721295): Check whether this is slower than doing
    # arrive_expect_tx(select(is_first_wg, 16, 0)) and if so then implement
    # support for ir.Value in arrive_expect_tx in the MGPU dialect (right now
    # the tx_bytes is an attribute, not an operand).
    with mgpu.utils.when(is_first_wg):
      barrier.arrive_expect_tx(16)
    with mgpu.utils.when(arith_dialect.xori(is_first_wg, mgpu.c(1, i1))):
      barrier.arrive()
    mgpu.dialect.try_cluster_cancel(
        result_ref,
        barrier.as_barrier_memref(),
        predicate=arith_dialect.andi(is_first_cta, is_first_wg))
  else:
    assert (lane_pred := ctx.module_ctx.single_lane_predicate) is not None
    is_leader = arith_dialect.andi(lane_pred, is_first_wg)
    tx_count = arith_dialect.select(is_leader, mgpu.c(16, i32), mgpu.c(0, i32))
    barrier.arrive_expect_tx(tx_count)
    mgpu.try_cluster_cancel(
        result_ref,
        barrier,
        predicate=arith_dialect.andi(is_leader, is_first_cta),
    )

  return []


def try_cluster_cancel(
    result_ref: _Ref,
    barrier: _Ref,
    *,
    collective_axes: Hashable | tuple[Hashable, ...] = (),
) -> None:
  """Initiates an async request to claim a new work unit from the grid.

  It allows an SM to dynamically acquire work by atomically canceling the launch
  of a pending cluster from the grid and retrieving its indices.

  Args:
    result_ref: An SMEM ref where the 16-byte result will be stored.
    barrier: A barrier used to coordinate the completion of the query.
    collective_axes: The thread / cluster axis names across which the
      cancellation is collective. When specified, only the axis index 0 along
      these axes will issue the cancellation request. If more than one thread
      along an axis calls this function, the axis must be specified here to
      avoid inadvertently cancelling multiple clusters.

  See also:
    :func:`jax.experimental.pallas.mosaic_gpu.query_cluster_cancel`
  """
  if not isinstance(collective_axes, tuple):
    collective_axes = (collective_axes,)

  if isinstance(result_ref, pallas_core.TransformedRef):
    result_transforms_leaves, result_transforms_tree = jax.tree.flatten(
        result_ref.transforms
    )
    result_ref = result_ref.ref
  else:
    result_transforms_leaves, result_transforms_tree = [], None

  if isinstance(barrier, pallas_core.TransformedRef):
    barrier_transforms_leaves, barrier_transforms_tree = jax.tree.flatten(
        barrier.transforms
    )
    barrier = barrier.ref
  else:
    barrier_transforms_leaves, barrier_transforms_tree = [], None

  try_cluster_cancel_p.bind(
      result_ref,
      barrier,
      *result_transforms_leaves,
      *barrier_transforms_leaves,
      result_transforms_tree=result_transforms_tree,
      barrier_transforms_tree=barrier_transforms_tree,
      collective_axes=collective_axes,
  )


query_cluster_cancel_p = jax_core.Primitive("query_cluster_cancel")
query_cluster_cancel_p.multiple_results = True

@query_cluster_cancel_p.def_effectful_abstract_eval
def _query_cluster_cancel_abstract_eval(try_cancel_buffer,
                                        *transforms_leaves,
                                        grid_names,
                                        transforms_tree):
  del try_cancel_buffer, transforms_leaves, transforms_tree
  grid_idxs = (jax_core.ShapedArray((), jnp.int32),) * len(grid_names)
  return (
      (*grid_idxs, jax_core.ShapedArray((), jnp.bool_)),
      {gpu_core._memory_effect},
  )


@lowering.register_lowering_rule(query_cluster_cancel_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(query_cluster_cancel_p, mgpu.LoweringSemantics.Warpgroup)
def query_cluster_cancel_lowering(ctx: lowering.LoweringRuleContext,
                                  result_ref,
                                  *transforms_leaves,
                                  grid_names,
                                  transforms_tree):
  result_aval = ctx.avals_in[0]
  assert isinstance(result_aval, state_types.AbstractRef)
  if transforms_tree is not None:
    res_transforms = transforms_tree.unflatten(transforms_leaves)
    transform_avals = transforms_tree.unflatten(ctx.avals_in[1:])
  else:
    res_transforms = ()
    transform_avals = ()
  result_ref, _, res_transforms = lowering._handle_transforms(
      ctx, result_aval, result_ref, transform_avals, res_transforms)
  if res_transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for result ref: {res_transforms}"
    )

  result_ty = ir.MemRefType(result_ref.type)
  bits = math.prod(result_ty.shape) * mgpu.bitwidth(result_ty.element_type)
  if bits != 128:
    raise TypeError(f"Response to decode must be 128 bits, but is {bits} bits.")

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    x, y, z, success = mgpu.dialect.query_cluster_cancel(result_ref)
  else:
    x, y, z, success = mgpu.query_cluster_cancel(result_ref)

  cta_grid = [x, y, z]
  i32 = ir.IntegerType.get_signless(32)
  # Divide out the cluster dimensions.
  for axis in ctx.module_ctx.axis_names.cluster:
    dim = lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, axis)
    cta_grid[dim] = arith_dialect.divui(
        cta_grid[dim],
        mgpu.c(ctx.launch_ctx.cluster_size[dim], i32))
  # Convert to grid indices.
  requested_idxs = []
  for axis_name in grid_names:
    requested_idxs.append(lowering.block_id_to_grid_id(
        ctx, cta_grid, axis_name))
  return (*requested_idxs, success)


def query_cluster_cancel(
    result_ref: _Ref,
    grid_names: Sequence[Hashable]) -> tuple[tuple[jax.Array, ...], jax.Array]:
  """Decodes the result of a ``try_cluster_cancel`` operation.

  It interprets the 16-byte opaque response written to shared memory by a
  completed ``try_cluster_cancel`` call to determine if a new work unit was
  successfully claimed.

  Args:
    result_ref: The SMEM ref containing the query response.
    grid_names: A tuple of grid axis names to query for.

  Returns:
    A tuple containing the decoded response:
      - the grid indices for the requested axis names.
      - A boolean indicating if the cancellation was successful.

  See also:
    :func:`jax.experimental.pallas.mosaic_gpu.try_cluster_cancel`
  """
  if isinstance(result_ref, pallas_core.TransformedRef):
    result_transforms_leaves, result_transforms_tree = jax.tree.flatten(
        result_ref.transforms
    )
    result_ref = result_ref.ref
  else:
    result_transforms_leaves, result_transforms_tree = [], None
  result = query_cluster_cancel_p.bind(
      result_ref,
      *result_transforms_leaves,
      grid_names=grid_names,
      transforms_tree=result_transforms_tree)
  return tuple(result[:-1]), result[-1]


class AtomicOpType(enum.Enum):
  ADD = "add"
  MAX = "max"
  MIN = "min"
  AND = "and"
  OR = "or"
  XOR = "xor"


def _atomic_op_type_to_int(atomic_type: AtomicOpType) -> int:
  match atomic_type:
    case AtomicOpType.ADD:
      return 0
    case AtomicOpType.MIN:
      return 1
    case AtomicOpType.MAX:
      return 2
    case AtomicOpType.AND:
      return 3
    case AtomicOpType.OR:
      return 4
    case AtomicOpType.XOR:
      return 5
    case _:
      assert_never(atomic_type)


atomic_store_p = jax_core.Primitive("mgpu_atomic_store")
atomic_store_p.multiple_results = True


@atomic_store_p.def_effectful_abstract_eval
def _atomic_store_abstract_eval(*avals_flat, args_tree, atomic_type, optimized):
  del atomic_type, optimized
  ref, transforms, val = args_tree.unflatten(avals_flat)
  if transforms is not None:
    ref = pallas_core.TransformedRef(ref, transforms)
  if ref.shape != val.shape:
    raise ValueError(
        f"Invalid shape for `swap`. Ref shape: {ref.shape}. "
        f"Value shape: {val.shape}."
    )
  if ref.dtype != val.dtype:
    raise ValueError(
        f"Invalid dtype for `swap`. Ref dtype: {ref.dtype}. "
        f"Value dtype: {val.dtype}."
    )
  return (), {state.WriteEffect(0)}


@discharge.register_discharge_rule(atomic_store_p)
def _atomic_store_discharge_rule(
    ctx, *args_flat, args_tree, atomic_type: AtomicOpType, optimized: bool
):
  del optimized
  ref, transforms, val, mask = args_tree.unflatten(args_flat)
  *prev_transforms, idx = transforms
  ref = discharge.transform_array(ref, prev_transforms)

  if mask is not None:
    raise NotImplementedError

  if atomic_type == AtomicOpType.ADD:
    monoid = lambda x, y: x + y
  elif atomic_type == AtomicOpType.MAX:
    monoid = jnp.maximum
  elif atomic_type == AtomicOpType.MIN:
    monoid = jnp.minimum
  else:
    raise NotImplementedError(atomic_type)

  if all(
      (isinstance(s, indexing.Slice) or not s.shape) for s in idx.indices
  ):
    indices = idx.indices
    scalar_dims = [
        not isinstance(s, indexing.Slice) and not s.shape for s in indices
    ]
    slice_starts = [
        s.start if isinstance(s, indexing.Slice) else s for s in indices
    ]
    slice_sizes = tuple(
        s.size if isinstance(s, indexing.Slice) else 1 for s in indices
    )
    out_ones = lax.dynamic_slice(ref, slice_starts, slice_sizes=slice_sizes)
    val_indexer = tuple(
        None if scalar else slice(None) for scalar in scalar_dims
    )
    val = val[val_indexer]
    val = monoid(val, out_ones)
    x_new = lax.dynamic_update_slice(ref, val, start_indices=slice_starts)
  elif all(not isinstance(s, indexing.Slice) for s in idx.indices):
    out = ref[idx.indices]
    x_new = ref.at[idx.indices].set(monoid(out, val))
  else:
    raise NotImplementedError
  return (x_new,) + (None,) * (len(ctx.in_avals) - 1), ()


def _atomic_store(
    x_ref_or_view,
    val,
    *,
    atomic_type: AtomicOpType,
    optimized: bool = True,
):
  x_ref, transforms = state_primitives.get_ref_and_transforms(
      x_ref_or_view, None, "atomic_store"
  )
  args_flat, args_tree = tree_util.tree_flatten((x_ref, transforms, val))
  atomic_store_p.bind(
      *args_flat,
      args_tree=args_tree,
      atomic_type=atomic_type,
      optimized=optimized,
  )


@lowering.register_lowering_rule(
    atomic_store_p, mgpu.LoweringSemantics.Warpgroup
)
def _atomic_store_lowering_rule_wg(
    ctx: lowering.LoweringRuleContext,
    *args_flat,
    args_tree,
    atomic_type: AtomicOpType,
    optimized: bool,
):
  ref, transforms, value = args_tree.unflatten(args_flat)
  ref_aval, transforms_avals, value_aval = args_tree.unflatten(ctx.avals_in)
  value = lowering._ensure_ir_value(value, value_aval.dtype)
  assert isinstance(ref_aval, state_types.AbstractRef)
  ref, _, remaining_transforms = lowering._handle_transforms(
      ctx, ref_aval, ref, list(transforms_avals), list(transforms),
      allow_peer_refs=True,
  )
  if remaining_transforms:
    raise NotImplementedError(
        f"Unsupported transforms for atomic_store: {remaining_transforms}"
    )

  mgpu.dialect.vector_store(
      value,
      ref,
      optimized=optimized,
      atomic_type=_atomic_op_type_to_int(atomic_type),
  )
  return ()


@lowering.register_lowering_rule(atomic_store_p, mgpu.LoweringSemantics.Lane)
def _atomic_store_lowering_rule(
    ctx: lowering.LoweringRuleContext,
    *args_flat,
    args_tree,
    atomic_type: AtomicOpType,
    optimized: bool,
):
  ref, transforms, value = args_tree.unflatten(args_flat)
  ref_aval, transforms_avals, value_aval = args_tree.unflatten(ctx.avals_in)
  value = lowering._ensure_fa(value, value_aval.dtype)
  assert isinstance(ref_aval, state_types.AbstractRef)
  ref, _, remaining_transforms = lowering._handle_transforms(
      ctx,
      ref_aval,
      ref,
      list(transforms_avals),
      list(transforms),
      allow_peer_refs=True,
  )
  match remaining_transforms:
    case (
        gpu_core.UnswizzleRef(swizzle),
        gpu_core.UntilingTransform(tiling),
    ):
      if len(tiling) != 2:
        raise NotImplementedError(
            f"Only 2D tiling is supported, got: {tiling}"
        )
      value.store_tiled(
          ref,
          swizzle=swizzle,
          optimized=optimized,
          tiling_rank=len(tiling),
          atomic=atomic_type.value,  # pyrefly: ignore[bad-argument-type]
      )
    case ():
      value.store_untiled(ref, optimized=optimized, atomic=atomic_type.value)  # pyrefly: ignore[bad-argument-type]
    case _:
      raise NotImplementedError(
          f"Unsupported transforms for atomic_store: {remaining_transforms}"
      )
  return ()


def atomic_add(ref: _Ref, val, *, optimized: bool = True) -> None:
  """Performs an atomic store-add of the value to the reference.

  Note that atomicity is only guaranteed on the element-level
  and that floating-point addition is not associative, so this op
  can introduce non-determinism into the kernels.

  Args:
    ref: The reference to store the value to.
    val: The value to store.
    optimized: If True, a compilation error will be raised if no optimized
      implementation for the store is available.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.ADD, optimized=optimized)


def atomic_max(ref: _Ref, val, *, optimized: bool = True) -> None:
  """Performs an atomic store-max of the value to the reference.

  Note that atomicity is only guaranteed on the element-level.

  Args:
    ref: The reference to store the value to.
    val: The value to store.
    optimized: If True, a compilation error will be raised if no optimized
      implementation for the store is available.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.MAX, optimized=optimized)


def atomic_min(ref: _Ref, val, *, optimized: bool = True) -> None:
  """Performs an atomic store-min of the value to the reference.

  Note that atomicity is only guaranteed on the element-level.

  Args:
    ref: The reference to store the value to.
    val: The value to store.
    optimized: If True, a compilation error will be raised if no optimized
      implementation for the store is available.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.MIN, optimized=optimized)


def atomic_and(ref: _Ref, val, *, optimized: bool = True) -> None:
  """Performs an atomic store-and of the value to the reference.

  Note that atomicity is only guaranteed on the element-level.

  Args:
    ref: The reference to store the value to.
    val: The value to store.
    optimized: If True, a compilation error will be raised if no optimized
      implementation for the store is available.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.AND, optimized=optimized)


def atomic_or(ref: _Ref, val, *, optimized: bool = True) -> None:
  """Performs an atomic store-or of the value to the reference.

  Note that atomicity is only guaranteed on the element-level.

  Args:
    ref: The reference to store the value to.
    val: The value to store.
    optimized: If True, a compilation error will be raised if no optimized
      implementation for the store is available.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.OR, optimized=optimized)


def atomic_xor(ref: _Ref, val, *, optimized: bool = True) -> None:
  """Performs an atomic store-xor of the value to the reference.

  Note that atomicity is only guaranteed on the element-level.

  Args:
    ref: The reference to store the value to.
    val: The value to store.
    optimized: If True, a compilation error will be raised if no optimized
      implementation for the store is available.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.XOR, optimized=optimized)


multimem_store_p = jax_core.Primitive("multimem_store")
multimem_store_p.multiple_results = True


def multimem_store(source: jax.Array, ref: _Ref, collective_axes: Hashable | tuple[Hashable, ...]):
  """Stores the value to ref on all devices present in collective_axes.

  The stores is done using the multimem instructions, meaning that the data is
  only transferred to the switch once, and broadcasted to all other devices
  there.

  Args:
    source: The value to store.
    ref: The GMEM reference to store the value to.
    collective_axes: The JAX mesh axes indicating the devices to store to.
  """
  if isinstance(ref, pallas_core.TransformedRef):
    transforms_leaves, transforms_tree = jax.tree.flatten(
        ref.transforms
    )
    ref = ref.ref
  else:
    transforms_leaves, transforms_tree = [], None
  multimem_store_p.bind(
      source,
      ref,
      *transforms_leaves,
      collective_axes=collective_axes,
      transforms_tree=transforms_tree,
  )


@multimem_store_p.def_effectful_abstract_eval
def _multimem_store_abstract_eval(source, ref, *transforms_leaves, transforms_tree, **_):
  _check_ref(ref, "ref", gpu_core.GMEM)
  shape, dtype = ref.shape, ref.dtype
  if transforms_tree is not None:
    transforms = jax.tree.unflatten(transforms_tree, transforms_leaves)
    ty = state.transform_type(transforms, ref)
    assert isinstance(ty, state.AbstractRef)
    shape = ty.shape
    dtype = ty.dtype
  if source.dtype != dtype:
    raise ValueError(f"Value dtype {source.dtype} does not match ref dtype {dtype}")
  if source.shape != shape:
    raise ValueError(f"Value shape {source.shape} does not match ref shape {shape}")
  return [], {pallas_core.comms_effect, state.WriteEffect(1)}


@lowering.register_lowering_rule(multimem_store_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(multimem_store_p, mgpu.LoweringSemantics.Warpgroup)
def _multimem_store_lowering_rule(
    ctx: lowering.LoweringRuleContext, value, local_ref, *transforms_leaves, transforms_tree, collective_axes,
):
  if (mesh_info := ctx.module_ctx.mesh_info) is None:
    raise ValueError(
        "JAX device mesh is required by multimem_store, but not defined."
    )
  if set(collective_axes) != set(mesh_info.axis_names):
    raise NotImplementedError(
        "Only collective_axes that include all JAX device mesh"
        f" ({mesh_info.axis_names}) axes are supported, but got"
        f" {collective_axes}"
    )
  if transforms_tree is not None:
    transforms = tree_util.tree_unflatten(transforms_tree, transforms_leaves)
    local_ref_aval = ctx.avals_in[1]
    assert isinstance(local_ref_aval, state_types.AbstractRef)
    transform_avals = transforms_tree.unflatten(ctx.avals_in[2:])
    local_ref, _, transforms = lowering._handle_transforms(
        ctx, local_ref_aval, local_ref, transform_avals, transforms, allow_peer_refs=False
    )
    if transforms:
      raise NotImplementedError(
          f"Unhandled transforms for multimem_store: {transforms}"
      )
  multi_ref = ctx.launch_ctx.to_remote_multicast(local_ref)
  scalar = not ctx.avals_in[0].shape
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    val = lowering._ensure_ir_value(value, ctx.avals_in[0].dtype)
    if scalar:
      with lowering._wrap_in_custom_primitive_if_wg(ctx, [multi_ref.ref, val]) as [multi_ref, val]:
        mgpu_utils.MultimemRef(multi_ref).store(val, indices=[])
        if ctx.module_ctx.auto_barriers:
          mgpu.warpgroup_barrier()
    else:
      mgpu.dialect.vector_store(val, multi_ref.ref, optimized=False, multimem=True)
    return ()

  if scalar:
    multi_ref.store(lowering._ensure_ir_value(value, ctx.avals_in[0].dtype), [])
  else:
    value.store_untiled(multi_ref, optimized=False)
  if ctx.module_ctx.auto_barriers:
    mgpu.warpgroup_barrier()  # Make sure the writes have completed.
  return ()


multimem_load_reduce_p = jax_core.Primitive("multimem_load_reduce")

@multimem_load_reduce_p.def_effectful_abstract_eval
def _multimem_load_reduce_abstract_eval(ref, *avals_flat, tree, collective_axes, reduction_op):
  del collective_axes, reduction_op
  _check_ref(ref, "ref", gpu_core.GMEM)
  out_ref = ref
  if tree is not None:
    transforms = jax.tree.unflatten(tree, avals_flat)
    out_ref = state.transform_type(transforms, ref)
  assert isinstance(out_ref, state_types.AbstractRef)
  return out_ref.inner_aval, {pallas_core.comms_effect}

@lowering.register_lowering_rule(multimem_load_reduce_p, mgpu.LoweringSemantics.Lane)
def _multimem_load_reduce_lowering_rule(
    ctx: lowering.LoweringRuleContext, ref, *transforms_leaves, tree, collective_axes, reduction_op,
):
  if (mesh_info := ctx.module_ctx.mesh_info) is None:
    raise ValueError(
        "JAX device mesh is required by multimem_load_reduce, but not defined."
    )
  if set(collective_axes) != set(mesh_info.axis_names):
    raise NotImplementedError(
        "Only collective_axes that include all JAX device mesh"
        f" ({mesh_info.axis_names}) axes are supported, but got"
        f" {collective_axes}"
    )
  if (layout := ctx.out_layout_hint) is None:
    raise RuntimeError(
        "Failed to infer the output layout of multimem_load_reduce. Please apply"
        " plgpu.layout_cast to its output right after its creation."
    )
  if not isinstance(layout, (mgpu.TiledLayout, mgpu.WGStridedFragLayout)):
    raise ValueError(
        "Only tiled and WG strided layouts are supported by"
        f" multimem_load_reduce, but got {layout}"
    )
  dtype = ctx.avals_out[0].dtype
  transforms = tree.unflatten(transforms_leaves)
  transform_avals = tree.unflatten(ctx.avals_in[1:])
  ref_aval = ctx.avals_in[0]
  assert isinstance(ref_aval, state_types.AbstractRef)
  ref, _, transforms = lowering._handle_transforms(ctx, ref_aval, ref,
                                                   transform_avals, transforms,
                                                   allow_peer_refs=False)
  if transforms:
    raise NotImplementedError(
        f"Unhandled transforms for multimem_load_reduce: {transforms}"
    )
  multi_ref = ctx.launch_ctx.to_remote_multicast(ref)
  is_signed = mgpu_utils.is_signed(dtype)
  arr = mgpu.FragmentedArray.load_reduce_untiled(
      multi_ref, layout=layout, is_signed=is_signed, reduction=reduction_op
  )
  return arr


@lowering.register_lowering_rule(multimem_load_reduce_p, mgpu.LoweringSemantics.Warpgroup)
def _multimem_load_reduce_lowering_rule_wg(
    ctx: lowering.LoweringRuleContext, ref, *transforms_leaves, tree, collective_axes, reduction_op,
):
  if (mesh_info := ctx.module_ctx.mesh_info) is None:
    raise ValueError(
        "JAX device mesh is required by multimem_load_reduce, but not defined."
    )
  if set(collective_axes) != set(mesh_info.axis_names):
    raise NotImplementedError(
        "Only collective_axes that include all JAX device mesh"
        f" ({mesh_info.axis_names}) axes are supported, but got"
        f" {collective_axes}"
    )
  transforms = tree.unflatten(transforms_leaves)
  transform_avals = tree.unflatten(ctx.avals_in[1:])
  ref_aval = ctx.avals_in[0]
  assert isinstance(ref_aval, state_types.AbstractRef)
  ref, _, transforms = lowering._handle_transforms(ctx, ref_aval, ref,
                                                   transform_avals, transforms)
  if transforms:
    raise NotImplementedError(
        f"Unhandled transforms for multimem_load_reduce: {transforms}"
    )
  assert reduction_op is not None
  is_signed = mgpu_utils.is_signed(ref_aval.dtype)
  match reduction_op:
    case "min" if is_signed is None:
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.Min
    case "min" if is_signed is True:
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.Smin
    case "min" if is_signed is False:
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.Umin
    case "max" if is_signed is True:
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.Smax
    case "max" if is_signed is False:
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.Umax
    case "max" if is_signed is None:
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.Max
    case "add":
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.Add
    case "and":
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.And
    case "or":
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.Or
    case "xor":
      reduction_op_attr = mgpu.dialect.MultimemLoadReductionType.Xor
    case _:
      raise ValueError(f"Unsupported integer reduction op: {reduction_op}")

  multimem_ref = ctx.launch_ctx.to_remote_multicast(ref)
  return mgpu.dialect.multimem_load_reduce(
      source=multimem_ref.ref,
      reduction_type=reduction_op_attr,
  )


def multimem_load_reduce(
    ref: _Ref,
    *,
    collective_axes: Hashable | tuple[Hashable, ...],
    reduction_op: mgpu.MultimemReductionOp,
) -> jax.Array:
  """Loads from a GMEM reference on all devices present in collective_axes and reduces the loaded values.

  The supported dtypes are: ``jnp.float32``, ``jnp.float16``, ``jnp.bfloat16``,
  ``jnp.float8_e5m2``, ``jnp.float8_e4m3fn``, ``jnp.int32`` and ``jnp.int64``.

  8-bit floating point dtypes are only supported on Blackwell GPUs.

  Args:
    ref: The GMEM reference to load from.
    collective_axes: The JAX mesh axes indicating the devices to load from.
    reduction_op: The reduction operation to perform on the loaded values. The
      allowed values are add (all dtypes), min, max (all dtypes but f32), as
      well as and, or and xor (integer types only).
  """
  ref, ref_transforms = state_primitives.get_ref_and_transforms(
      ref, None, "multimem_load_reduce"
  )
  flat_ref_transforms, ref_transforms_treedef = tree_util.tree_flatten(
      ref_transforms
  )
  return multimem_load_reduce_p.bind(
      ref,
      *flat_ref_transforms,
      tree=ref_transforms_treedef,
      collective_axes=collective_axes,
      reduction_op=reduction_op,
  )

semaphore_signal_multicast_p = jax_core.Primitive("semaphore_signal_multicast")
semaphore_signal_multicast_p.multiple_results = True

def semaphore_signal_multicast(
    semaphore,
    value: int | jax.Array = 1,
    *,
    collective_axes: Hashable | tuple[Hashable, ...],
):
  """Signals a semaphore on all devices along collective_axes.

  At the moment only signals to all devices are supported.

  Args:
    semaphore: The semaphore reference to signal.
    value: The increment value for the semaphore.
    collective_axes: The mesh axes to multicast the signal across.
      Must contain all mesh axes.
  """
  if not isinstance(collective_axes, tuple):
    collective_axes = (collective_axes,)
  ref, transforms = pallas_primitives._get_ref_and_transforms(semaphore)
  value = jnp.asarray(value, dtype=jnp.int32)
  args = [ref, transforms, value]
  flat_args, args_tree = tree_util.tree_flatten(args)
  return semaphore_signal_multicast_p.bind(
      *flat_args,
      args_tree=args_tree,
      collective_axes=collective_axes,
  )


@semaphore_signal_multicast_p.def_effectful_abstract_eval
def _semaphore_signal_multicast_abstract_eval(
    *avals, args_tree, collective_axes
):
  del collective_axes  # Unused.
  sem_aval, transform_avals, _ = tree_util.tree_unflatten(args_tree, avals)
  pallas_primitives.check_sem_avals(
      sem_aval, transform_avals, "semaphore_signal_multicast"
  )
  return (), {pallas_core.comms_effect}


@lowering.register_lowering_rule(
    semaphore_signal_multicast_p, mgpu.LoweringSemantics.Lane
)
@lowering.register_lowering_rule(
    semaphore_signal_multicast_p, *gpu_core.LANExWARP_SEMANTICS
)
@lowering.register_lowering_rule(
    semaphore_signal_multicast_p, mgpu.LoweringSemantics.Warpgroup
)
@lowering.register_lowering_rule(
    semaphore_signal_multicast_p, *gpu_core.WGxWARP_SEMANTICS
)
def _semaphore_signal_multicast_lowering(
    ctx: lowering.LoweringRuleContext, *args, args_tree, collective_axes
):
  sem, transforms, value = tree_util.tree_unflatten(args_tree, args)
  sem_aval, transform_avals, _ = tree_util.tree_unflatten(args_tree, ctx.avals_in)
  assert isinstance(sem_aval, state_types.AbstractRef)
  sem, _, sem_transforms = lowering._handle_transforms(ctx, sem_aval, sem,
                                                       transform_avals,
                                                       transforms)
  if sem_transforms:
    raise NotImplementedError(
        f"Unhandled transforms for semaphore_signal_multicast: {sem_transforms}"
    )
  if not isinstance(collective_axes, (tuple, list)):
    collective_axes = (collective_axes,)
  if (mesh_info := ctx.module_ctx.mesh_info) is None:
    raise ValueError("collective_axes requires a mesh context")
  if set(collective_axes) != set(mesh_info.axis_names):
    raise ValueError(
        f"collective_axes {collective_axes} must equal entire mesh axes {mesh_info.axis_names}"
    )
  i32 = ir.IntegerType.get_signless(32)
  val = lowering._ir_constant(value, i32)
  multi_ref = ctx.launch_ctx.to_remote_multicast(sem).ref
  with lowering._wrap_in_custom_primitive_if_wg(ctx, [multi_ref, val]) as [multi_ref, val]:
    if ctx.module_ctx.auto_barriers:
      if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
        mgpu_utils.warp_barrier()
      else:
        mgpu_utils.warpgroup_barrier()

    assert ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane
    predicate = ctx.module_ctx.single_lane_predicate

    mgpu_utils.SemaphoreRef.signal_multimem(
        mgpu_utils.memref_ptr(multi_ref), val, predicate=predicate
    )
  return ()

semaphore_signal_p = jax_core.Primitive("semaphore_signal")
semaphore_signal_p.multiple_results = True


def semaphore_signal(
    semaphore,
    inc: int | jax.Array = 1,
    *,
    device_id: pallas_primitives.DeviceId | None = None,
    memory_scope: Literal["sys", "gpu"] = "sys",
):
  """Signals a semaphore, optionally on a remote device.

  This is the MGPU specific variant of :func:`pallas.semaphore_signal`,
  which additionally exposes the ``memory_scope`` of the underlying atomic.

  Args:
    semaphore: The semaphore reference to signal.
    inc: The increment value for the semaphore.
    device_id: Optional logical device id at which to signal the semaphore.
    memory_scope: The memory scope of the underlying atomic. Must be ``"sys"``
      or ``"gpu"``. Defaults to ``"sys"``.
  """
  ref, transforms = pallas_primitives._get_ref_and_transforms(semaphore)
  value = jnp.asarray(inc, dtype=jnp.int32)
  device_id = pallas_primitives.canonicalize_device_id(device_id)
  core_index = None
  args = [ref, transforms, value, device_id, core_index]
  flat_args, args_tree = tree_util.tree_flatten(args)
  semaphore_signal_p.bind(
      *flat_args,
      args_tree=args_tree,
      device_id_type=pallas_primitives.DeviceIdType.MESH,
      memory_scope=memory_scope,
  )


@semaphore_signal_p.def_effectful_abstract_eval
def _semaphore_signal_abstract_eval(
    *avals, args_tree, device_id_type, memory_scope
):
  del memory_scope  # Unused.
  return pallas_primitives.semaphore_signal_p.abstract_eval(
      *avals, args_tree=args_tree, device_id_type=device_id_type
  )


@lowering.register_lowering_rule(pallas_primitives.semaphore_signal_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(pallas_primitives.semaphore_signal_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(pallas_primitives.semaphore_signal_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(pallas_primitives.semaphore_signal_p, *gpu_core.WGxWARP_SEMANTICS)
@lowering.register_lowering_rule(semaphore_signal_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(semaphore_signal_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(semaphore_signal_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(semaphore_signal_p, *gpu_core.WGxWARP_SEMANTICS)
def _semaphore_signal_lowering_rule(
    ctx: lowering.LoweringRuleContext,
    *args,
    args_tree,
    device_id_type,
    memory_scope: Literal["sys", "gpu"] = "sys",
):
  i32 = ir.IntegerType.get_signless(32)
  sem, transforms, value, device_id, core_index = tree_util.tree_unflatten(
      args_tree, args
  )
  sem_aval, transform_avals, _, device_id_aval, _ = tree_util.tree_unflatten(
      args_tree, ctx.avals_in
  )
  if core_index is not None:
    raise NotImplementedError(
        "Mosaic GPU backend does not support the concept of cores, but"
        " core_index is specified"
    )
  assert isinstance(sem_aval, state_types.AbstractRef)
  sem, _, transforms = lowering._handle_transforms(
      ctx, sem_aval, sem, transform_avals, transforms
  )
  if transforms:
    raise NotImplementedError(f"Unhandled transforms for semaphore_signal: {transforms}")
  if device_id is not None:
    if memory_scope == "gpu":
      raise ValueError(
          "Cannot signal a GPU-local semaphore from a remote device. Please use"
          " `memory_scope='sys'` instead."
      )
    device_id = lowering._device_id_to_logical(
        ctx, device_id, device_id_type, device_id_aval
    )
    assert device_id is not None
    device_id = lowering._ensure_ir_value(device_id, jnp.int32)
    sem = ctx.launch_ctx.to_remote(sem, device_id)

  val = lowering._ir_constant(value, i32)
  with lowering._wrap_in_custom_primitive_if_wg(ctx, [sem, val]) as [sem, val]:
    sem_ptr = mgpu.utils.memref_ptr(sem)
    # We only signal the semaphore from a single lane, which does not guarantee
    # anything about the state of the other three warps in the warpgroup (they
    # might still be e.g. reading memory that someone will overwrite once they
    # receive a signal).
    if ctx.module_ctx.auto_barriers:
      if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
        mgpu_utils.warp_barrier()
      else:
        mgpu_utils.warpgroup_barrier()

    mgpu_utils.SemaphoreRef(sem_ptr).signal(
        val,
        predicate=ctx.module_ctx.single_lane_predicate,
        memory_scope=memory_scope,
    )
  return ()


semaphore_wait_p = jax_core.Primitive("semaphore_wait")
semaphore_wait_p.multiple_results = True


def semaphore_wait(
    semaphore,
    value: int | jax.Array = 1,
    *,
    decrement: bool = True,
    memory_scope: Literal["sys", "gpu"] = "sys",
):
  """Waits on a semaphore until it reaches at least ``value``.

  This is the MGPU specific variant of :func:`pallas.semaphore_wait`,
  which additionally exposes the ``memory_scope`` of the underlying atomic.

  Args:
    semaphore: The semaphore reference to wait on.
    value: The target value that the semaphore should reach before unblocking.
    decrement: Whether to decrement the semaphore by ``value`` once the wait
      succeeds.
    memory_scope: The memory scope of the underlying atomic. Must be ``"sys"``
      or ``"gpu"``. Defaults to ``"sys"``.
  """
  ref, transforms = pallas_primitives._get_ref_and_transforms(semaphore)
  value = jnp.asarray(value, dtype=jnp.int32)
  args = [ref, transforms, value, decrement]
  flat_args, args_tree = tree_util.tree_flatten(args)
  semaphore_wait_p.bind(
      *flat_args,
      args_tree=args_tree,
      memory_scope=memory_scope,
  )


@semaphore_wait_p.def_effectful_abstract_eval
def _semaphore_wait_abstract_eval(*avals, args_tree, memory_scope):
  del memory_scope  # Unused.
  return pallas_primitives.semaphore_wait_p.abstract_eval(
      *avals, args_tree=args_tree
  )


@lowering.register_lowering_rule(pallas_primitives.semaphore_wait_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(pallas_primitives.semaphore_wait_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(pallas_primitives.semaphore_wait_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(pallas_primitives.semaphore_wait_p, *gpu_core.WGxWARP_SEMANTICS)
@lowering.register_lowering_rule(semaphore_wait_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(semaphore_wait_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(semaphore_wait_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(semaphore_wait_p, *gpu_core.WGxWARP_SEMANTICS)
def _semaphore_wait_lowering_rule(
    ctx: lowering.LoweringRuleContext,
    *args,
    args_tree,
    memory_scope: Literal["sys", "gpu"] = "sys",
):
  sem, transforms, value, decrement = tree_util.tree_unflatten(args_tree, args)
  sem_aval, transform_avals, *_ = tree_util.tree_unflatten(
      args_tree, ctx.avals_in
  )
  assert isinstance(sem_aval, state_types.AbstractRef)
  sem, _, transforms = lowering._handle_transforms(ctx, sem_aval, sem, transform_avals, transforms)
  if transforms:
    raise NotImplementedError(
        f"Unhandled transforms for semaphore_wait: {transforms}"
    )
  val = lowering._ensure_ir_value(value, jnp.int32)

  scope = mgpu.ThreadSubset.WARPGROUP
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    scope = mgpu.ThreadSubset.WARP

  with lowering._wrap_in_custom_primitive_if_wg(ctx, [sem, val]) as [sem, val]:
    mgpu_utils.SemaphoreRef(mgpu.utils.memref_ptr(sem)).wait(
        val, decrement=decrement, scope=scope, memory_scope=memory_scope,
    )
  return ()


reduce_sum_p = jax_core.Primitive("mosaic_gpu_reduce_sum")
reduce_max_p = jax_core.Primitive("mosaic_gpu_reduce_max")
reduce_min_p = jax_core.Primitive("mosaic_gpu_reduce_min")
reduce_prod_p = jax_core.Primitive("mosaic_gpu_reduce_prod")


def _reduce_abstract_eval(
    x_aval: jax_core.ShapedArray,
    *,
    axes: tuple[int, ...],
    accumulator_ilp: int | None = None,
) -> jax_core.ShapedArray:
  if accumulator_ilp is not None and (
      not isinstance(accumulator_ilp, int)
      or isinstance(accumulator_ilp, bool)
      or accumulator_ilp <= 0
  ):
    raise ValueError(
        f"accumulator_ilp must be a positive integer, got: {accumulator_ilp}"
    )

  out_shape = tuple(s for i, s in enumerate(x_aval.shape) if i not in axes)
  return jax_core.ShapedArray(out_shape, x_aval.dtype)


for prim, op in (
    (reduce_sum_p, "add"),
    (reduce_max_p, "max"),
    (reduce_min_p, "min"),
    (reduce_prod_p, "prod"),
):
  prim.def_abstract_eval(_reduce_abstract_eval)
  batching.defreducer(prim)
  lowering._register_resource_estimator(prim)(
      lowering._reduce_resource_estimator
  )
  lowering.register_lowering_rule(prim, mgpu.LoweringSemantics.Lane)(
      functools.partial(lowering._reduce_lowering_rule, op)
  )

lowering.register_lowering_rule(reduce_sum_p, mgpu.LoweringSemantics.Warpgroup)(
    lowering._reduce_sum_lowering_rule_wg
)
lowering.register_lowering_rule(reduce_max_p, mgpu.LoweringSemantics.Warpgroup)(
    lowering._reduce_max_lowering_rule_wg
)
lowering.register_lowering_rule(reduce_min_p, mgpu.LoweringSemantics.Warpgroup)(
    lowering._reduce_min_lowering_rule_wg
)
lowering.register_lowering_rule(reduce_prod_p, mgpu.LoweringSemantics.Warpgroup)(
    lowering._reduce_prod_lowering_rule_wg
)


def _reduce(
    prim: jax_core.Primitive,
    x: Any,
    axis: int | Sequence[int] | None = None,
    keepdims: bool = False,
    *,
    accumulator_ilp: int | None = None,
) -> Any:
  if keepdims:
    raise NotImplementedError("keepdims=True is not yet supported.")
  axes = util.canonicalize_axis_tuple(axis, x.ndim)
  return prim.bind(x, axes=axes, accumulator_ilp=accumulator_ilp)


sum = functools.partial(_reduce, reduce_sum_p)
max = functools.partial(_reduce, reduce_max_p)
min = functools.partial(_reduce, reduce_min_p)
prod = functools.partial(_reduce, reduce_prod_p)