# 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 inspect
import itertools
import math
from typing import Any, Literal, assert_never

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.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 nvvm as nvvm_dialect
from jax._src.pallas import core as pallas_core
from jax._src.pallas import primitives as pallas_primitives
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 CopyPartition
from jax.experimental.mosaic.gpu.launch_context import OOBFillMode
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) or 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}"
    )


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, fmt, *_, **params):
  del aval_in, fmt, 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,
    fmt: str,
    *transforms_leaves,
    transforms_tree
):
  if transforms_leaves:
    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),
    )
    if remaining_transforms:
      raise NotImplementedError(
          f"Unsupported transforms {remaining_transforms}."
      )
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    print(fmt.format(mgpu.dialect_lowering.pprint_layout(x)))  # pyrefly: ignore[bad-argument-type]
  else:
    assert isinstance(x, ir.Value)
    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 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 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,
      fmt=fmt,
      *transforms_leaves,
      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)
  del args, params  # Unused.
  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"] = jax_core.pp_var(user_predicate, 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

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    if predicate is not None:
      assert ctx.module_ctx.single_lane_predicate is not None
      predicate = arith_dialect.andi(
          predicate, ctx.module_ctx.single_lane_predicate
      )
    else:
      predicate = ctx.module_ctx.single_lane_predicate

  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),
  }
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    ctx.launch_ctx.async_copy(
        src_ref=src,
        dst_ref=dst,
        predicate=predicate,
        arrive=commit_group,
        reduction_op=reduction_op,
        oob_mode=OOBFillMode.UNDEFINED,
        **copy_params,
    )
    return ()

  if "gmem_slice" not in copy_params:
    i32 = ir.IntegerType.get_signless(32)
    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
  if copy_params.get("gmem_peer_id", None) is not None:
    raise NotImplementedError(
        "GMEM refs with peer ids are not supported in warpgroup lowering."
    )
  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,
  )
  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: 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


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, barrier, *args, **params):
  del args, params  # Unused.
  _check_ref(src, "src", gpu_core.GMEM)
  _check_ref(dst, "dst", gpu_core.SMEM)
  _check_ref(barrier, "barrier", gpu_core.SMEM)
  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,
):
  src, dst, barrier, *flat_args = eqn.invars
  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
  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)
  barrier_transforms = barrier_transforms_treedef.unflatten(
      flat_barrier_transforms
  )
  return pp.concat([
      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),
      pp.text(" using "),
      state_primitives.pp_ref_transforms(context, barrier, barrier_transforms),
  ])


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,
    barrier,
    *flat_transforms,
    src_transforms_treedef,
    dst_transforms_treedef,
    barrier_transforms_treedef,
    collective_axes,
    leader_tracked,
    oob_mode,
):
  flat_src_transforms, flat_dst_transforms, flat_barrier_transforms = (
      util.split_list(
          flat_transforms,
          [
              src_transforms_treedef.num_leaves,
              dst_transforms_treedef.num_leaves,
          ],
      )
  )
  flat_src_transforms_avals, flat_dst_transforms_avals, _ = (
      util.split_list(
          ctx.avals_in[3:],
          [
              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]
  barrier_ref_aval = ctx.avals_in[2]
  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)
  assert isinstance(barrier_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),
  }
  base_index = _extract_barrier_slice_base(
      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 math.prod(ctx.launch_ctx.cluster_size) != 2:
      raise NotImplementedError(
          "Partitioned loads only supported for clusters of size 2. Got"
          f" cluster size {ctx.launch_ctx.cluster_size}."
      )

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    if 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 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
      if ctx.module_ctx.auto_barriers:
        mgpu.warpgroup_barrier()  # Make sure all reads have completed.
      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()),
        )
        barrier.arrive_expect_tx(bytes, predicate=first_block)
      else:
        barrier.arrive_expect_tx(bytes)
    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 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()),
        )
        with mgpu.when(first_block):
          barrier.arrive(arrival_count=3, can_complete=False)
          barrier.arrive_expect_tx(bytes)
      else:
        barrier.arrive(arrival_count=3, can_complete=False)
        barrier.arrive_expect_tx(bytes)

    # Gathers are a warpgroup-level collective and can't take a predicate.
    predicate_kwarg = dict(predicate=ctx.module_ctx.single_lane_predicate)
    if gmem_slice := copy_params.get("gmem_slice", ()):
      first_idx = gmem_slice[0]
      if isinstance(first_idx, mgpu.FragmentedArray) and first_idx.shape:
        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,
        **copy_params,
        **predicate_kwarg,  # pyrefly: ignore[bad-argument-type]
    )
    return ()
  i32 = ir.IntegerType.get_signless(32)
  if "gmem_slice" not in copy_params:
    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 copy_params.get("gmem_peer_id", None) is not None:
    raise NotImplementedError(
        "GMEM refs with peer ids are not supported in warpgroup lowering."
    )
  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

  barrier_ref = barrier.as_barrier_memref()

  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_ctx = mgpu.when(first_block)
  else:
    arrive_ctx = contextlib.nullcontext()
  with arrive_ctx:
    mgpu.dialect.arrive_expect_tx(barrier_ref, bytes)

  mgpu.dialect.async_load(
      src,
      dst,
      barrier_ref,
      indices,
      slice_lengths,
      collective=ir.ArrayAttr.get(
          [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)  # pyrefly: ignore[unexpected-keyword]
  )
  return ()

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

  If collective_axes is specified, this performs a multicast copy where
  all CUDA blocks that share the same index along the collective axis
  receive a copy of the same block of data loaded from `dst` to `src`.

  If both ``collective_axes`` and ``leader_tracked`` are specified as
  ``CopyPartition.PARTITIONED(axis)``, this will perform a partitioned
  collective copy where each block in the cluster will receive a tile of
  ``transfer_size // cluster_size`` data from the ``src`` Ref.
  For example, if ``src`` has a shape of (256, 256) and a partitioned
  copy is performed along axis 0 with cluster size 2, then the first block
  will receive ``src[0:128, :]`` and the second will receive
  ``src[128:256, :]``.

  If both ``collective_axes`` and ``leader_tracked`` are specified as
  ``CopyPartition.REPLICATED``, this will perform a replicated copy where
  all blocks load the same data but only the first block in the collective
  tracks progress via barrier arrivals.


  NOTE: Only the first block in the cluster will arrive on the barrier,
  and an additional cluster barrier is necessary to ensure that all blocks in
  the cluster 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.
    collective_axes: The collective axes to use for the copy.
    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.
    oob_mode: The optional out-of-bounds fill mode. Can be ``OOBFillMode.UNDEFINED``,
     ``OOBFillMode.PROMISE_IN_BOUNDS`` or ``OOBFillMode.ZEROS``.

  See also:
    :func:`jax.experimental.pallas.mosaic_gpu.barrier_arrive`
    :func:`jax.experimental.pallas.mosaic_gpu.barrier_wait`
  """
  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
  )
  barrier, barrier_transforms = state_primitives.get_ref_and_transforms(
      barrier, None, "copy_gmem_to_smem"
  )
  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,
      *flat_src_transforms,
      *flat_dst_transforms,
      *flat_barrier_transforms,
      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,
  )
  return None

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


@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_ref_transforms,
    ref_transforms_treedef,
    collective_axes,
    leader_tracked,
):
  ref_transforms = ref_transforms_treedef.unflatten(flat_ref_transforms)
  ref_transform_avals = ref_transforms_treedef.unflatten(ctx.avals_in[1:])
  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:
    predicate_kwarg: dict[str, Any] = dict(
        predicate=ctx.module_ctx.single_lane_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:
        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([])
  )
  return ()


def async_prefetch(
    ref: _Ref,
    *,
    collective_axes: str | tuple[str, ...] | None = None,
    leader_tracked: CopyPartition | 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.
  """
  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,
      ref_transforms_treedef=ref_transforms_treedef,
      collective_axes=collective_axes,
      leader_tracked=leader_tracked,
  )
  return None


def _extract_barrier_slice_base(transforms) -> ir.Value | None:
  if not transforms:
    return None
  base_index: ir.Value | None = None
  while transforms:
    match transforms:
      case [indexing.NDIndexer(indices=[idx]) as indexer, *transforms]:
        if isinstance(idx, indexing.Slice):
          if indexing.Slice.from_slice(slice(None), *indexer.shape) == idx:
            # Special-case: the whole slice.
            continue
          idx = idx.start
        if isinstance(
            idx, (int, ir.Value, mgpu.FragmentedArray, literals.TypedNdArray)
        ):
          if base_index is None:
            base_index = lowering._as_index(idx)
          else:
            base_index = arith_dialect.addi(base_index, lowering._as_index(idx))
        else:
          raise ValueError(
              f"Barrier can only be indexed with integers or slices, got {idx}"
          )
      case [indexing.NDIndexer(), *_]:
        raise NotImplementedError("Barrier does not support multiple indices")
      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,
):
  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"),
      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_transforms,
    transforms_treedef,
):
  transforms = transforms_treedef.unflatten(flat_transforms)
  barrier_aval = ctx.avals_in[0]
  assert isinstance(barrier_aval, state_types.AbstractRef)
  base_index = _extract_barrier_slice_base(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)
  elif ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    barrier.arrive(orders_tensor_core)
  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=pred,
        scope=scope,
    )
  return ()


def barrier_arrive(barrier: state.AbstractRef) -> 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, 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 = _extract_barrier_slice_base(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, *, wait_read_only):
  del n, wait_read_only  # Unused.
  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, *, 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, 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.
  nvvm_dialect.cp_async_bulk_commit_group()
  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}"
    )

  if a.dtype != b.dtype:
    raise ValueError(f"Mixed input dtypes for matrix multiplication unsupported: 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(in_avals, out_avals, *args, **kwargs):
  del in_avals, out_avals
  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]
  )
  if a_transforms is not None:
    a_aval = ctx.avals_in[1]
    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
    assert transform_treedefs[1] is not None
    a_transform_avals = transform_treedefs[1].unflatten(transform_avals_list[1])
    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(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
    if not isinstance(a, mgpu.FragmentedArray):
      raise ValueError(
          "When WGMMA lhs is passed in as a ref, it must be transformed by"
          " swizzling and tiling appropriately."
      )

  assert b_transforms is not None
  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
  assert transform_treedefs[2] is not None
  b_transform_avals = transform_treedefs[2].unflatten(transform_avals_list[2])
  b, _, b_transforms = lowering._handle_transforms(
      ctx, b_ref_aval, b, b_transform_avals, b_transforms,
      handle_transposes=False, handle_reshapes=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 (
        gpu_core.UnswizzleRef(rhs_swizzle),
        state_types.TransposeTransform((1, 0, 2, 3, 4)),
        gpu_core.UntilingTransform(rhs_tiling),
        state_types.TransposeTransform(permutation=(1, 0, 2)),
        state.types.ReshapeTransform(shape=new_shape),
    ):
      if len(rhs_tiling) != 2 or len(new_shape) != 2:
        raise ValueError("WGMMA expects shapes 2D tiled into 2D tiles.")

      if any(d % t != 0 for d, t in util.safe_zip(new_shape, rhs_tiling)):
        raise ValueError(
            f"The last reshape {new_shape} is not divisible by the tiling"
            f" {rhs_tiling}."
        )

      high_dims = [d // t for d, t in util.safe_zip(new_shape, rhs_tiling)]
      b = mgpu.memref_reshape(b, (*high_dims, *rhs_tiling))
      rhs_transpose = False
    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
    )
  nvvm_dialect.wgmma_commit_group_sync_aligned()
  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,
):
  if acc_transforms_tree is not None:
    raise NotImplementedError
  if a_transforms_tree is not None:
    a_transforms_leaves, b_transforms_leaves = util.split_list(
        transforms_leaves, [a_transforms_tree.num_leaves]
    )
    a_transforms = a_transforms_tree.unflatten(a_transforms_leaves)
    a_aval = ctx.avals_in[1]
    assert isinstance(a_aval, state_types.AbstractRef)
    a_transform_avals = a_transforms_tree.unflatten(
        ctx.avals_in[3 : 3 + a_transforms_tree.num_leaves]
    )
    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}."
      )
    a_transforms_num_leaves = a_transforms_tree.num_leaves
  else:
    a_transforms_num_leaves = 0
    b_transforms_leaves = transforms_leaves

  if b_transforms_tree is not None:
    b_transforms = b_transforms_tree.unflatten(b_transforms_leaves)
    b_aval = ctx.avals_in[2]
    assert isinstance(b_aval, state_types.AbstractRef)
    b_transform_avals = b_transforms_tree.unflatten(
        ctx.avals_in[3 + a_transforms_num_leaves:]
    )
    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}."
      )

  new_acc = mgpu.dialect.wgmma(acc, a, b)
  nvvm_dialect.wgmma_commit_group_sync_aligned()
  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 [])
  }

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)


@wgmma_wait_p.def_effectful_abstract_eval
def wgmma_wait_effectful_abstract_eval(_):
  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, allow_groups):
  del ctx
  nvvm_dialect.wgmma_wait_group_sync_aligned(allow_groups)
  return ()


wgmma_accumulator_deref_p = jax_core.Primitive("wgmma_accumulator_deref_p")


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(in_avals, out_avals, acc, *, wait_n):
  del in_avals, out_avals
  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:
    nvvm_dialect.wgmma_wait_group_sync_aligned(wait_n)
  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_p")


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(in_avals, out_avals, acc, val):
  del in_avals, out_avals
  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])
  nvvm_dialect.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
  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)
    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_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_transform_avals = a_transforms_tree.unflatten(a_transforms_leaves_avals)
    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)}")

  assert b_transforms_tree is not None
  b_transforms = b_transforms_tree.unflatten(b_transforms_leaves)
  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_transform_avals = b_transforms_tree.unflatten(b_transforms_leaves_avals)
  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 = _extract_barrier_slice_base(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 and a_scale_transforms_tree is not None:
    assert isinstance(a_scale_ref_aval, state.AbstractRef)
    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
    )
    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 and b_scale_transforms_tree is not None:
    assert isinstance(b_scale_ref_aval, state.AbstractRef)
    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
    )
    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_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
        )
    )
    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_axis is not None:
    assert predicate is not None
    is_leader_block = _collective_mma_predicate(ctx, collective_axis)
    predicate = arith_dialect.andi(predicate, is_leader_block)
    collective = True
  else:
    collective = False

  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,  # pyrefly: ignore[bad-argument-type]
              b_scale=b_scale_ref,  # pyrefly: ignore[bad-argument-type]
              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

  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 tree is None:
      return ref
    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 = _extract_barrier_slice_base(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):
  """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.

  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,
      barrier_transforms_tree=barrier_transforms_tree,
      collective_axis=collective_axis)


@tcgen05_commit_arrive_p.def_effectful_abstract_eval
def _tcgen05_commit_arrive_abstract_eval(barrier,
                               *barrier_transforms_leaves,
                               barrier_transforms_tree,
                               collective_axis):
  del barrier_transforms_leaves, barrier_transforms_tree, collective_axis
  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)
def _tcgen05_commit_arrive_lowering(
    ctx: lowering.LoweringRuleContext,
    barrier_ref: mgpu.BarrierRef,
    *barrier_transforms_leaves,
    barrier_transforms_tree,
    collective_axis,
):
  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(
        barrier_transforms_leaves
    )
    base_index = _extract_barrier_slice_base(barrier_transforms)
    if base_index is not None:
      barrier_ref = barrier_ref[base_index]

  predicate = ctx.module_ctx.single_lane_predicate
  if collective_axis is not None:
    assert predicate is not None
    is_leader_block = _collective_mma_predicate(ctx, collective_axis)
    predicate = arith_dialect.andi(predicate, is_leader_block)
    collective = True
  else:
    collective = False

  with mgpu.when(predicate):
    tcgen05.commit_arrive(barrier_ref,
                          collective=collective,
                          ctx=ctx.launch_ctx)
  return []


@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_wg(
    ctx: lowering.LoweringRuleContext,
    barrier_ref: mgpu.DialectBarrierRef,
    *barrier_transforms_leaves,
    barrier_transforms_tree,
    collective_axis,
):
  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(
        barrier_transforms_leaves
    )
    base_index = _extract_barrier_slice_base(barrier_transforms)
    if base_index is not None:
      barrier_ref = barrier_ref[base_index]

  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.tcgen05_commit_arrive(
        barrier_ref.as_barrier_memref(), collective=collective
    )
  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_p")
set_max_registers_p.multiple_results = True


@set_max_registers_p.def_effectful_abstract_eval
def _set_max_registers_abstract_eval(n, *, action):
  del n, action  # Unused.
  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, *, action
):
  del ctx
  nvvm_dialect.setmaxregister(
      n,
      nvvm_dialect.SetMaxRegisterAction.increase
      if action == "increase"
      else nvvm_dialect.SetMaxRegisterAction.decrease,
  )
  return ()


def set_max_registers(n: int, *, action: Literal["increase", "decrease"]):
  """Sets the maximum number of registers owned by a warp."""
  set_max_registers_p.bind(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


@lowering.register_lowering_rule(jax_core.closed_call_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(jax_core.closed_call_p, mgpu.LoweringSemantics.Warpgroup)
def _closed_call_lowering_rule(ctx, *args, call_jaxpr: jax_core.ClosedJaxpr):
  if call_jaxpr.consts: raise NotImplementedError
  return lowering.lower_jaxpr_to_mosaic_gpu(
      ctx.module_ctx, ctx.launch_ctx, call_jaxpr.jaxpr, args)


@lowering._register_resource_estimator(jax_core.closed_call_p)
def _closed_call_resource_estimator(ctx, *args, call_jaxpr):
  del args  # Unused.
  if call_jaxpr.consts: raise NotImplementedError
  return lowering._estimate_resources(ctx, call_jaxpr.jaxpr)


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


inline_mgpu_p = jax_core.Primitive("inline_mgpu_p")
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_partial_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,
    )

    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 = (
            mgpu.dialect_lowering.swizzle_and_transforms_from_transforms_attr(
                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.ReturnOp(operands_=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
)


def load(
    src: _Ref,
    idx,
    *,
    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.
    idx: The index to load from.
    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, idx, "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


async_load_tmem_p = jax_core.Primitive("async_load")

def async_load_tmem(src: _Ref, *, layout: SomeLayout | None = None) -> 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.
  """
  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
  )
  result = async_load_tmem_p.bind(
      src, *flat_src_transforms, tree=src_transforms_treedef
  )
  if layout is not None:
    result = gpu_core.layout_cast(result, layout)
  return result

@async_load_tmem_p.def_effectful_abstract_eval
def _async_load_tmem_abstract_eval(src, *avals_flat, tree):
  if src.memory_space != gpu_core.MemorySpace.TMEM:
    raise ValueError("Async load only supports TMEM refs")
  return state_primitives._get_abstract_eval(src, *avals_flat, tree=tree)

@lowering.register_lowering_rule(async_load_tmem_p, mgpu.LoweringSemantics.Lane)
def _async_load_tmem_lowering_rule(
    ctx: lowering.LoweringRuleContext, x_ref, *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[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)
  return x_tmem.load(layout=layout_hint, is_signed=is_signed)


@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
):
  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=}"
    )
  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)
  if transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for TMEM refs. {transforms=}"
    )
  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,
  )
  if transforms:
    raise NotImplementedError(
        f"Unimplemented transforms for TMEM refs. {transforms=}"
    )
  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_axis is not None:
    assert predicate is not None
    is_leader_block = _collective_mma_predicate(ctx, collective_axis)
    predicate = arith_dialect.andi(predicate, is_leader_block)
    collective = True
  else:
    collective = False

  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
)
def _async_copy_scales_to_tmem_lowering_rule_wg(*args, **kwargs):
  # TODO(olechwierowicz): remove this check once minimum jaxlib version is 0.10.0.
  if not hasattr(mgpu.dialect, "async_store_scales_smem_to_tmem"):
    raise NotImplementedError(
        "async_copy_scales_to_tmem_p WG lowering is not implemented."
    )
  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
)
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, **_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 = jax.tree.unflatten(smem_tree, args[:smem_tree.num_leaves])
  smem_aval = smem_ref
  for t in smem_transforms:
    smem_aval = t.transform_type(smem_aval)
  if smem_aval.dtype != tmem_ref.dtype:
    raise ValueError(
        f"Expected SMEM element type ({smem_aval.dtype}) to equal the TMEM"
        f" element type ({tmem_ref.dtype})"
    )
  if smem_aval.shape != tmem_ref.shape:
    raise ValueError(
        f"Expected SMEM reference shape {smem_aval.shape} to equal the TMEM"
        f" reference shape {tmem_ref.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_axis is not None:
    assert predicate is not None
    is_leader_block = _collective_mma_predicate(ctx, collective_axis)
    predicate = arith_dialect.andi(predicate, is_leader_block)
    collective = True
  else:
    collective = False

  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 = [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,
):
  i1 = ir.IntegerType.get_signless(1)
  i32 = ir.IntegerType.get_signless(32)

  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]
    )
    result_aval = ctx.avals_in[0]
    assert isinstance(result_aval, state_types.AbstractRef)
    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}"
      )
  else:
    barrier_transforms_leaves = transforms_leaves

  if barrier_transforms_tree is not None:
    base_index = _extract_barrier_slice_base(
        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."
    )

  is_first_wg = arith_dialect.cmpi(
      arith_dialect.CmpIPredicate.eq, mgpu.warpgroup_idx(), mgpu.c(0, i32)
  )

  is_first_cta = mgpu.c(1, i1)
  for dim in gpu_dialect.Dimension:
    is_first_cta = arith_dialect.andi(
        is_first_cta,
        arith_dialect.cmpi(
            arith_dialect.CmpIPredicate.eq,
            mgpu.utils.cluster_idx(dim),
            mgpu.c(0, ir.IndexType.get()),
        ),
    )

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    # TODO(b/415721295): remove this check once minimum jaxlib version is 0.10.0
    if not hasattr(mgpu.dialect, "TryClusterCancelOp"):
      raise NotImplementedError("TryClusterCancelOp is not available.")

    # 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 ctx.module_ctx.single_lane_predicate is not None
    is_leader_thread = arith_dialect.andi(
        ctx.module_ctx.single_lane_predicate, is_first_wg
    )
    bytes = arith_dialect.select(is_leader_thread, mgpu.c(16, i32), mgpu.c(0, i32))
    barrier.arrive_expect_tx(bytes)
    mgpu.try_cluster_cancel(
        result_ref,
        barrier,
        predicate=arith_dialect.andi(is_leader_thread, is_first_cta),
    )

  return []


def try_cluster_cancel(result_ref: _Ref, barrier: _Ref) -> 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 claiming its CTA ID as the next unit
  of work.

  Note that this operation must be called collectively by all Pallas threads.

  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.

  See also:
    :func:`jax.experimental.pallas.mosaic_gpu.query_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

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


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):
  if transforms_tree is not None:
    res_transforms = transforms_tree.unflatten(transforms_leaves)
    result_aval = ctx.avals_in[0]
    assert isinstance(result_aval, state_types.AbstractRef)
    transform_avals = transforms_tree.unflatten(ctx.avals_in[1:])
    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:
    # TODO(b/415721295): remove this check once minimum jaxlib version is 0.10.0
    if not hasattr(mgpu.dialect, "QueryClusterCancelOp"):
      raise NotImplementedError("QueryClusterCancelOp is not available.")
    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"


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):
  del atomic_type
  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(
    in_avals, out_avals, *args_flat, args_tree, atomic_type: AtomicOpType
):
  del out_avals
  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(in_avals) - 1), ()


def _atomic_store(
    x_ref_or_view,
    val,
    *,
    atomic_type: AtomicOpType,
):
  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
  )


@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,
):
  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)
  )
  if remaining_transforms:
    raise NotImplementedError(
        f"Unsupported transforms for atomic_store: {remaining_transforms}"
    )

  def to_int_attr(atomic_type: AtomicOpType):
    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)

  mgpu.dialect.vector_store(value, ref, atomic_type=to_int_attr(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,
):
  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)
  )
  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, tiling_rank=len(tiling),
          atomic=atomic_type.value,  # pyrefly: ignore[bad-argument-type]
      )
    case ():
      value.store_untiled(ref, optimized=False, 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) -> 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.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.ADD)


def atomic_max(ref: _Ref, val) -> 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.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.MAX)


def atomic_min(ref: _Ref, val) -> 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.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.MIN)


def atomic_and(ref: _Ref, val) -> 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.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.AND)


def atomic_or(ref: _Ref, val) -> 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.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.OR)


def atomic_xor(ref: _Ref, val) -> 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.
  """
  _atomic_store(ref, val, atomic_type=AtomicOpType.XOR)


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,
):
  # TODO(olechwierowicz): Remove once min version of jaxlib is 0.10.0.
  multimem_available = "multimem" in inspect.signature(mgpu.dialect.vector_store).parameters
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup and not multimem_available:
    raise NotImplementedError("_multimem_store_lowering_rule WG lowering not implemented")

  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}"
      )
  multicast_ref = lambda r: ctx.launch_ctx.to_remote_multicast(r)
  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, [local_ref, val]) as [local_ref, val]:
        multicast_ref(local_ref).store(val, indices=[])
        if ctx.module_ctx.auto_barriers:
          mgpu.warpgroup_barrier()
    else:
      mc_ref = multicast_ref(local_ref).ref
      mgpu.dialect.vector_store(val, mc_ref, optimized=False, multimem=True)
    return ()

  multi_ref = multicast_ref(local_ref)
  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,
):
  # TODO(olechwierowicz): Remove this check once min jaxlib version is 0.10.0
  if not hasattr(mgpu.dialect, "multimem_load_reduce"):
    raise NotImplementedError("multimem_load_reduce_p is unsupported for WG semantics.")
  assert hasattr(mgpu.dialect, "MultimemLoadReductionType")
  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
)
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)
  with lowering._wrap_in_custom_primitive_if_wg(ctx, [sem, val]) as [sem, val]:
    multi_ref = ctx.launch_ctx.to_remote_multicast(sem)
    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()

    predicate = None
    if 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.ref), val, predicate=predicate
    )
  return ()
