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

"""Module for lowering JAX primitives to Mosaic GPU."""

from __future__ import annotations

import collections
from collections.abc import Callable, Generator, Hashable, Iterator, MutableMapping, MutableSequence, Sequence
import contextlib
import dataclasses
import functools
import itertools
import math
import operator
from typing import Any, Protocol, assert_never, cast, overload

import jax
from jax import api_util
from jax import lax
from jax._src import checkify
from jax._src import config
from jax._src import core as jax_core
from jax._src import custom_derivatives
from jax._src import debugging
from jax._src import deprecations
from jax._src import dtypes
from jax._src import flattree as ft
from jax._src import layout as jax_layout
from jax._src import literals as jax_literals
from jax._src import mesh as mesh_lib
from jax._src import pjit
from jax._src import source_info_util
from jax._src import tree_util
from jax._src import util
from jax._src.interpreters import mlir
from jax._src.interpreters import partial_eval as pe
from jax._src.lib import xla_client as xc
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 cf as cf_dialect
from jax._src.lib.mlir.dialects import gpu as gpu_dialect
from jax._src.lib.mlir.dialects import llvm as llvm_dialect
from jax._src.lib.mlir.dialects import math as math_dialect
from jax._src.lib.mlir.dialects import memref as memref_dialect
from jax._src.lib.mlir.dialects import scf as scf_dialect
from jax._src.lib.mlir.dialects import vector as vector_dialect
from jax._src.pallas import core as pallas_core
from jax._src.pallas import mpmd
from jax._src.pallas import primitives
from jax._src.pallas import utils as pallas_utils
from jax._src.pallas.mosaic_gpu import core as gpu_core
from jax._src.state import discharge
from jax._src.state import indexing
from jax._src.state import primitives as sp
from jax._src.state import types as state_types
from jax._src.state.types import ReshapeTransform
from jax._src.state.types import TransposeTransform
from jax._src.util import foreach
import jax.experimental.mosaic.gpu as mgpu
from jax.experimental.mosaic.gpu import core as mgpu_core
from jax.experimental.mosaic.gpu import fragmented_array as fa
from jax.experimental.mosaic.gpu import profiler as mgpu_profiler
from jax.experimental.mosaic.gpu import tcgen05
from jax.experimental.mosaic.gpu import utils as mgpu_utils
import jax.numpy as jnp
import numpy as np

map, unsafe_map = util.safe_map, map
zip, unsafe_zip = util.safe_zip, zip

partial = functools.partial
SMEM = gpu_core.SMEM
WARPGROUP_SIZE = 128

type ShapedAbstractValue = pallas_utils.ShapedAbstractValue

# TODO(slebedev): The type argument should also be comparable.
CollectiveAxesType = Sequence[Hashable]


@dataclasses.dataclass(frozen=True, kw_only=True)
class ResourceEstimatorContext:
  reduction_scratch_bytes: int
  axis_names: _AxisNames
  lowering_semantics: mgpu.LoweringSemantics

  @property
  def arrival_multiplier(self) -> int:
    return (
        WARPGROUP_SIZE
        if self.lowering_semantics == mgpu.LoweringSemantics.Lane
        else 1
    )


AnyBarrier = mgpu.Barrier | mgpu.ClusterBarrier


def _get_barrier(
    aval: ShapedAbstractValue, arrival_multiplier: int
) -> mgpu.Barrier:
  assert isinstance(aval.dtype, gpu_core.BarrierType)
  num_arrivals = aval.dtype.num_arrivals
  num_barriers = math.prod(aval.shape)
  if not (orders_tc := aval.dtype.orders_tensor_core):
    num_arrivals *= arrival_multiplier
  return mgpu.Barrier(num_arrivals, num_barriers, orders_tc)


def _get_cluster_barrier(
    aval: ShapedAbstractValue, axis_names: _AxisNames
) -> mgpu.ClusterBarrier:
  assert isinstance(aval.dtype, gpu_core.ClusterBarrierType)
  num_arrivals = aval.dtype.num_arrivals
  num_barriers = math.prod(aval.shape)
  resolve = functools.partial(_resolve_cluster_axis, axis_names)
  collective_dims = jax.tree.map(resolve, aval.dtype.collective_axes)
  return mgpu.ClusterBarrier(
      collective_dims, num_arrivals, num_barriers,
      orders_tensor_core=aval.dtype.orders_tensor_core,
      leader_tracked=aval.dtype.leader_tracked,
  )


@dataclasses.dataclass(kw_only=True, frozen=True)
class Resources:
  smem_scratch_bytes: int = 0
  tmem_scratch_cols: int = 0
  tmem_collective_scratch_cols: int = 0
  barrier_counts: collections.Counter[AnyBarrier] = dataclasses.field(
      default_factory=collections.Counter
  )
  # Maps from collective axes to number of semaphores.
  scoped_gmem_semaphores: dict[CollectiveAxesType, int] = dataclasses.field(
      default_factory=dict
  )

  def __post_init__(self):
    object.__setattr__(
        self,
        "smem_scratch_bytes",
        gpu_core.align_to(self.smem_scratch_bytes, gpu_core.SMEM_ALIGNMENT),
    )

    # TMEM must be allocated in 128x8 chunks.
    object.__setattr__(
        self,
        "tmem_scratch_cols",
        gpu_core.align_to(self.tmem_scratch_cols, 8),
    )
    object.__setattr__(
        self,
        "tmem_collective_scratch_cols",
        gpu_core.align_to(self.tmem_collective_scratch_cols, 8),
    )

  @property
  def barriers(self) -> Sequence[AnyBarrier]:
    return list(self.barrier_counts.elements())

  def __add__(self, other: Resources) -> Resources:
    # TODO(slebedev): Optimize this.
    #
    # At the moment, if we have run_scoped(b1) followed by run_scoped(b2)
    # we will allocate two barriers, even though one would be enough.
    sems = self.scoped_gmem_semaphores
    other_sems = other.scoped_gmem_semaphores
    scoped_gmem_semaphores = {key: sems.get(key, 0) + other_sems.get(key, 0)
                              for key in sems.keys() | other_sems.keys()}
    return Resources(
        smem_scratch_bytes=self.smem_scratch_bytes + other.smem_scratch_bytes,
        tmem_scratch_cols=self.tmem_scratch_cols + other.tmem_scratch_cols,
        tmem_collective_scratch_cols=self.tmem_collective_scratch_cols
        + other.tmem_collective_scratch_cols,
        barrier_counts=self.barrier_counts + other.barrier_counts,
        scoped_gmem_semaphores=scoped_gmem_semaphores,
    )

  def or_(self, other: Resources, axis_names: _AxisNames) -> Resources:
    sems = self.scoped_gmem_semaphores
    other_sems = other.scoped_gmem_semaphores
    scoped_gmem_semaphores = {}
    for sem_scope in sems.keys() | other_sems.keys():
      if _is_block_local_scope(sem_scope, axis_names):
        value = max(sems.get(sem_scope, 0), other_sems.get(sem_scope, 0))
      elif _is_global_scope(sem_scope, axis_names):
        value = sems.get(sem_scope, 0) + other_sems.get(sem_scope, 0)
      else:
        raise RuntimeError(f"Unrecognized semaphore scope: {sem_scope}")
      scoped_gmem_semaphores[sem_scope] = value
    return Resources(
        smem_scratch_bytes=max(
            self.smem_scratch_bytes, other.smem_scratch_bytes
        ),
        tmem_scratch_cols=max(self.tmem_scratch_cols, other.tmem_scratch_cols),
        tmem_collective_scratch_cols=max(
            self.tmem_collective_scratch_cols,
            other.tmem_collective_scratch_cols,
        ),
        barrier_counts=self.barrier_counts | other.barrier_counts,
        scoped_gmem_semaphores=scoped_gmem_semaphores,
    )


class ResourceEstimator(Protocol):

  def __call__(
      self, ctx: ResourceEstimatorContext, *args: Any, **params: Any
  ) -> Resources:
    ...


_resource_estimators: dict[jax_core.Primitive, ResourceEstimator] = {}


def _register_resource_estimator(primitive: jax_core.Primitive):
  def deco(fn):
    _resource_estimators[primitive] = fn
    return fn

  return deco


def _estimate_resources(
    ctx: ResourceEstimatorContext, jaxpr: jax_core.Jaxpr
) -> Resources:
  """Estimates the resources required by the kernel."""
  rs = Resources(smem_scratch_bytes=0)
  for eqn in jaxpr.eqns:
    # TODO(slebedev): Add support for other primitives, notably control flow.
    if rule := _resource_estimators.get(eqn.primitive):
      rs = rs.or_(
          rule(ctx, *(invar.aval for invar in eqn.invars), **eqn.params),
          ctx.axis_names,
      )
      continue
    # Assume that unsupported primitives are neutral wrt resource usage,
    # unless they have a jaxpr in their params.
    if any( isinstance(v, jax_core.Jaxpr) for v in eqn.params.values()):
      raise NotImplementedError(
          f"Resource estimation does not support {eqn.primitive}"
      )

  return rs


@_register_resource_estimator(lax.cond_p)
def _cond_resource_estimator(
    ctx: ResourceEstimatorContext, *args, branches
) -> Resources:
  del args  # Unused.
  return functools.reduce(
      lambda a, b: a.or_(b, ctx.axis_names),
      (_estimate_resources(ctx, branch) for branch in branches),
  )


@_register_resource_estimator(lax.scan_p)
def _scan_resource_estimator(
    ctx: ResourceEstimatorContext, *args, jaxpr: jax_core.Jaxpr, **params
) -> Resources:
  del args, params  # Unused.
  return _estimate_resources(ctx, jaxpr)


@_register_resource_estimator(lax.while_p)
def _while_resource_estimator(
    ctx: ResourceEstimatorContext,
    *args,
    cond_jaxpr: jax_core.Jaxpr,
    body_jaxpr: jax_core.Jaxpr,
    **params,
) -> Resources:
  del args, params  # Unused.
  return _estimate_resources(ctx, cond_jaxpr).or_(
      _estimate_resources(ctx, body_jaxpr), ctx.axis_names
  )


@_register_resource_estimator(pjit.jit_p)
def _pjit_resource_estimator(
    ctx: ResourceEstimatorContext,
    *args,
    jaxpr: jax_core.Jaxpr,
    **params,
) -> Resources:
  del args, params  # Unused.
  return _estimate_resources(ctx, jaxpr)


@_register_resource_estimator(pe.eval_jaxpr_p)
@_register_resource_estimator(custom_derivatives.custom_jvp_call_p)
@_register_resource_estimator(custom_derivatives.custom_vjp_call_p)
def _eval_jaxpr_resource_estimator(
    ctx: ResourceEstimatorContext,
    *args,
    call_jaxpr: jax_core.Jaxpr,
    **params,
) -> Resources:
  del args, params  # Unused.
  if call_jaxpr.consts:
    raise NotImplementedError
  return _estimate_resources(ctx, call_jaxpr)


@_register_resource_estimator(mpmd.mpmd_map_p)
def _mpmd_map_resource_estimator(
    ctx: ResourceEstimatorContext, *args, jaxprs: tuple[jax_core.Jaxpr, ...],
    **params
) -> Resources:
  del args, params  # Unused.
  if len(jaxprs) > 1:
    raise NotImplementedError(
        "MPMD map with multiple jaxprs not supported for resource estimation."
    )
  return _estimate_resources(ctx, jaxprs[0])


@_register_resource_estimator(discharge.run_state_p)
def _run_state_resource_estimator(
    ctx: ResourceEstimatorContext, *args, jaxpr: jax_core.Jaxpr, **params
) -> Resources:
  del args, params  # Unused.
  return _estimate_resources(ctx, jaxpr)


@_register_resource_estimator(primitives.run_scoped_p)
def _run_scoped_resource_estimator(
    ctx: ResourceEstimatorContext,
    *consts,
    jaxpr: jax_core.Jaxpr,
    collective_axes,
    **_,
) -> Resources:
  # NOTE: This rule assumes that the allocation happens collectively, although
  # it can't be checked here due to limited context. We check this in the actual
  # lowering rule.
  del consts  # Unused.
  rs = Resources()
  for v in jaxpr.invars:
    aval = cast(ShapedAbstractValue, v.aval)
    if isinstance(aval.dtype, gpu_core.BarrierType):
      barrier = _get_barrier(aval, ctx.arrival_multiplier)
      rs += Resources(barrier_counts=collections.Counter([barrier]))
      continue
    if isinstance(aval.dtype, gpu_core.ClusterBarrierType):
      barrier = _get_cluster_barrier(aval, ctx.axis_names)
      rs += Resources(barrier_counts=collections.Counter([barrier]))
      continue
    assert isinstance(aval, state_types.AbstractRef)
    if aval.memory_space == gpu_core.TMEM:
      if len(aval.shape) != 2:
        raise ValueError(f"TMEM allocations must be 2D. Got {aval.shape}")
      # Estimate columns used.
      if isinstance(aval, gpu_core.AbstractRefUnion):
        assert aval.shape[0] == 128
        cols_used = aval.shape[1]
      else:
        # pyrefly: ignore[missing-attribute]
        cols_used = aval.layout.cols_in_shape(
            aval.shape, dtypes.itemsize_bits(aval.dtype)
        )
      if aval.collective:  # pyrefly: ignore[missing-attribute]
        rs += Resources(tmem_collective_scratch_cols=cols_used)
      else:
        rs += Resources(tmem_scratch_cols=cols_used)
    elif aval.memory_space == gpu_core.SMEM:
      rs += Resources(
          smem_scratch_bytes=aval.size * dtypes.itemsize_bits(aval.dtype) // 8
      )
    elif aval.memory_space == gpu_core.REGS:
      # Don't need to allocate anything.
      pass
    elif aval.memory_space == gpu_core.GMEM and jnp.issubdtype(aval.dtype, pallas_core.semaphore):
      if _is_block_local_scope(collective_axes, ctx.axis_names):
        rs += Resources(scoped_gmem_semaphores={collective_axes: aval.size})
      else:
        raise ValueError(
            "Only thread-collective allocations are supported in run_scoped. To"
            " allocate global semaphores, use pl.get_global."
        )
    else:
      raise NotImplementedError(
          f"Unsupported memory space: {aval.memory_space}")
  return rs + _estimate_resources(ctx, jaxpr)


@_register_resource_estimator(lax.reduce_sum_p)
@_register_resource_estimator(lax.reduce_max_p)
@_register_resource_estimator(lax.reduce_min_p)
@_register_resource_estimator(lax.reduce_prod_p)
def _reduce_resource_estimator(
    ctx: ResourceEstimatorContext, x_aval: jax_core.ShapedArray, *, axes,
    **kwargs
) -> Resources:
  del x_aval, axes, kwargs  # Unused.
  # We don't need SMEM for some reductions, but it depends on the layout, so we
  # conservatively request the maximum scratch space we might need.
  return Resources(smem_scratch_bytes=ctx.reduction_scratch_bytes)


@dataclasses.dataclass(frozen=True)
class _AxisNames:
  grid: Sequence[Hashable]
  cluster: Sequence[Hashable] = ()
  wg: Hashable | None = None

  def __iter__(self) -> Iterator[Hashable]:
    return itertools.chain(
        self.grid, self.cluster, [self.wg] if self.wg is not None else []
    )

  def reverse(self) -> _AxisNames:
    return _AxisNames(self.grid[::-1], self.cluster[::-1], self.wg)


AnyBarrierRef = (
    mgpu.BarrierRef | mgpu.DialectBarrierRef | mgpu.CollectiveBarrierRef
)


@dataclasses.dataclass
class ModuleContext:
  name: str
  axis_names: _AxisNames
  program_ids: Sequence[ir.Value] | None
  approx_math: bool
  single_wg_lane_predicate: ir.Value
  single_warp_lane_predicate: ir.Value
  smem_requested_bytes: int
  smem_used_bytes: int
  tmem_requested_cols: int
  tmem_used_cols: int
  tmem_base: ir.Value | None
  scoped_gmem_used_semaphores: dict[CollectiveAxesType, int]
  scoped_gmem_semaphore_base_ptr: dict[CollectiveAxesType, ir.Value]
  runtime_barriers: MutableMapping[AnyBarrier, MutableSequence[AnyBarrierRef]]
  name_stack: source_info_util.NameStack
  traceback_caches: mlir.TracebackCaches
  squashed_dims: tuple[int, ...]
  lowering_semantics: mgpu.LoweringSemantics
  primitive_semantics: gpu_core.PrimitiveSemantics
  mesh_info: pallas_utils.MeshInfo | None
  # See the documentation of unsafe_no_auto_barriers in CompilerParams.
  auto_barriers: bool
  # See the documentation of reduction_scratch_bytes in CompilerParams.
  reduction_scratch_bytes: int
  warp_axis_name: jax_core.AxisName | None = None
  outer_traceback: xc.Traceback | None = None
  _smem_allocation_counter: int = dataclasses.field(default=0, init=False)
  _tmem_allocation_counter: int = dataclasses.field(default=0, init=False)

  def next_smem_allocation_id(self) -> int:
    assert self.lowering_semantics == mgpu.LoweringSemantics.Warpgroup
    self._smem_allocation_counter += 1
    return self._smem_allocation_counter

  def next_tmem_allocation_id(self) -> int:
    assert self.lowering_semantics == mgpu.LoweringSemantics.Warpgroup
    self._tmem_allocation_counter += 1
    return self._tmem_allocation_counter

  @property
  def single_lane_predicate(self) -> ir.Value:
    """Returns a predicate that is True for a single lane within the current
    thread semantics.
    """
    assert self.lowering_semantics == mgpu.LoweringSemantics.Lane
    match self.primitive_semantics:
      case gpu_core.PrimitiveSemantics.Warpgroup:
        return self.single_wg_lane_predicate
      case gpu_core.PrimitiveSemantics.Warp:
        return self.single_warp_lane_predicate
      case _:
        raise ValueError(f"Unknown semantics: {self.primitive_semantics}")

  @contextlib.contextmanager
  def reserve_barrier(
      self, barrier: mgpu.Barrier | mgpu.ClusterBarrier
  ) -> Generator[
      mgpu.BarrierRef | mgpu.DialectBarrierRef | mgpu.CollectiveBarrierRef,
  ]:
    """Reserves a barrier.

    Raises:
      RuntimeError: If the barrier is already reserved.
    """
    available = self.runtime_barriers.get(barrier, [])
    if not available:
      raise RuntimeError(f"Barrier {barrier} is already reserved")
    barrier_ref = available.pop()
    yield barrier_ref
    available.append(barrier_ref)

  @contextlib.contextmanager
  def reserve_semaphores(
      self, shape: tuple[int, ...], collective_axes: CollectiveAxesType
  ) -> Generator[ir.Value]:
    allocated_sems = math.prod(shape)
    ref = mgpu.memref_slice(
        self.scoped_gmem_semaphore_base_ptr[collective_axes],
        mgpu.ds(self.scoped_gmem_used_semaphores[collective_axes],
                allocated_sems),
    )
    ref = mgpu.memref_reshape(ref, shape)

    self.scoped_gmem_used_semaphores[collective_axes] += allocated_sems
    yield ref
    # TODO: In debug mode verify the values of all semaphores are again 0
    self.scoped_gmem_used_semaphores[collective_axes] -= allocated_sems

  @contextlib.contextmanager
  def alloc_tmem(
      self,
      struct: jax.ShapeDtypeStruct,
      *,
      layout: tcgen05.TMEMLayout,
  ) -> Generator[tcgen05.TMEMRef | ir.Value]:
    assert self.tmem_base is not None
    if self.lowering_semantics == mgpu.LoweringSemantics.Lane:
      off = arith_dialect.addi(
          self.tmem_base, _i32_constant(self.tmem_used_cols)
      )
      tmem_ref = tcgen05.TMEMRef(
          address=off,
          shape=struct.shape,
          dtype=mgpu_utils.dtype_to_ir_type(struct.dtype),
          layout=layout,
      )
    else:
      type = ir.MemRefType.get(
          struct.shape,
          mgpu_utils.dtype_to_ir_type(struct.dtype),
          memory_space=mgpu_utils.tmem(),
      )
      tmem_ref = mgpu.dialect.slice_tmem(
          type,
          self.tmem_base,
          self.tmem_used_cols,
          alias_id=self.next_tmem_allocation_id(),
      )
      layout_attr = mgpu.to_layout_attr(layout)
      tmem_ref = mgpu.dialect.tmem_layout_cast(tmem_ref, layout_attr)
    cols_used = layout.cols_in_shape(
        struct.shape, dtypes.itemsize_bits(struct.dtype)
    )
    cols_used = gpu_core.align_to(cols_used, gpu_core.TMEM_COL_ALIGNMENT)
    self.tmem_used_cols += cols_used
    yield tmem_ref
    self.tmem_used_cols -= cols_used

  @contextlib.contextmanager
  def scratch_view(
      self, struct: jax.ShapeDtypeStruct
  ) -> Generator[ir.Value]:
    """Creates a view into the runtime scratch buffer for the given struct.

    This is a low-level API. Use it only if you know what you are doing.

    After the context manager exits, the view is invalid and cannot be used.

    Args:
      struct: The shape and dtype of the view to create.

    Returns:
      A memref view into the runtime scratch buffer.
    """
    smem_base = None
    i8 = ir.IntegerType.get_signless(8)
    if self.lowering_semantics == mgpu.LoweringSemantics.Lane:
      smem_base = gpu_dialect.dynamic_shared_memory(
          ir.MemRefType.get(
              (mgpu_utils.DYNAMIC,), i8, memory_space=mgpu_utils.smem()
          )
      )
    off = initial_used_bytes = self.smem_used_bytes
    assert off % gpu_core.SMEM_ALIGNMENT == 0
    scratch_ty = ir.MemRefType.get(
        struct.shape,
        mgpu_utils.dtype_to_ir_type(struct.dtype),
        memory_space=mgpu_utils.smem(),
    )
    # The below code emission relies on the assumption that the first scratch
    # operand provided by Mosaic GPU always begins at the beginning of
    # dynamic SMEM. Mosaic GPU is expected to uphold that invariant.
    if self.lowering_semantics == mgpu.LoweringSemantics.Lane:
      assert smem_base is not None
      view = memref_dialect.view(scratch_ty, smem_base, _as_index(off), [])
    else:
      view = mgpu.dialect.slice_smem(
          scratch_ty, off, alias_id=self.next_smem_allocation_id()
      )

    off += gpu_core.align_to(
        math.prod(struct.shape)
        * dtypes.itemsize_bits(jnp.dtype(struct.dtype))
        // 8,
        gpu_core.SMEM_ALIGNMENT,
    )
    assert off <= self.smem_requested_bytes, "Ran out of scoped SMEM"
    assert off % gpu_core.SMEM_ALIGNMENT == 0

    self.smem_used_bytes = off
    yield view
    self.smem_used_bytes = initial_used_bytes


@dataclasses.dataclass(frozen=True)
class LoweringRuleContext:
  module_ctx: ModuleContext
  launch_ctx: mgpu.LaunchContext
  prim: jax_core.Primitive
  avals_in: Sequence[ShapedAbstractValue]
  avals_out: Sequence[ShapedAbstractValue]
  out_layout_hint: mgpu.FragmentedLayout | None

  replace = dataclasses.replace

  @property
  def estimator_ctx(self) -> ResourceEstimatorContext:
    return ResourceEstimatorContext(
        reduction_scratch_bytes=self.module_ctx.reduction_scratch_bytes,
        axis_names=self.module_ctx.axis_names,
        lowering_semantics=self.module_ctx.lowering_semantics
    )


@dataclasses.dataclass(frozen=True)
class LoweringResult:
  module: ir.Module
  grid: tuple[int, ...]
  block: tuple[int, ...]
  new_out_shapes: tuple[jax.ShapeDtypeStruct, ...]  # Does not include gmem scratch!
  profiler_spec: mgpu_profiler.ProfilerSpec | None
  gmem_scratch_shapes: tuple[jax.ShapeDtypeStruct, ...]


class LoweringError(Exception):
  pass


def _eval_index_map(
    module_ctx: ModuleContext,
    launch_ctx: mgpu.LaunchContext,
    idx: Sequence[ir.Value],
    block_mapping: pallas_core.BlockMapping,
) -> Sequence[ir.Value]:
  block_indices = lower_jaxpr_to_mosaic_gpu(
      module_ctx, launch_ctx, block_mapping.index_map_jaxpr, idx
  )
  result = []
  for i, b in zip(block_indices, block_mapping.block_shape):
    match b:
      case pallas_core.Squeezed() | pallas_core.Element():
        result.append(i)
      case pallas_core.Blocked():
        result.append(arith_dialect.muli(_as_index(i), _as_index(b)))
      case _:
        raise ValueError(f"Unsupported block dim type: {b}")
  return tuple(result)


def _check_block_mappings(
    block_mappings: Sequence[pallas_core.BlockMapping],
    debug_info: jax_core.DebugInfo,
) -> None:
  def err_details(bm: pallas_core.BlockMapping) -> str:
    return (
        f"Block spec for {bm.origin} in pallas_call {debug_info.func_src_info}"
        f" has block shape {bm.block_shape}, array shape"
        f" {bm.array_aval.shape},"
        # TODO(necula): add index_map source location info
        f" and index_map {bm.index_map_jaxpr} in"
        f" memory space {bm.transformed_block_aval.memory_space}."
        " See details at"
        " https://docs.jax.dev/en/latest/pallas/grid_blockspec.html#pallas-blockspec."
    )

  for bm in block_mappings:
    if (
        bm.transformed_block_aval.memory_space == gpu_core.GMEM
        and not bm.has_trivial_window()
    ):
      raise NotImplementedError(
          "Mosaic GPU lowering currently requires blocks in GMEM memory space "
          "to have same block shape as the array shape "
          "and a trivial index_map (returning all 0s).\n\n"
          + err_details(bm)
      )

    if any(isinstance(b, pallas_core.Element) for b in bm.block_shape):
      raise NotImplementedError(
          "Only Blocked indexing mode is supported in Mosaic GPU lowering.\n\n"
          + err_details(bm)
      )

    if bm.pipeline_mode is not None:
      raise NotImplementedError(
          "Pipeline mode is not supported in Mosaic GPU lowering.\n\n"
          + err_details(bm)
      )


def lower_unpipelined_jaxpr_to_module(
    gpu_mesh: gpu_core.Mesh,
    jax_mesh: mesh_lib.Mesh | None,
    jaxpr: jax_core.Jaxpr,
    in_shapes: Sequence[jax_core.ShapedArray],
    out_shapes: Sequence[jax_core.ShapedArray],
    params: gpu_core.CompilerParams,
    *,
    outer_traceback: xc.Traceback | None = None,
) -> LoweringResult:
  """Lowers a jaxpr that operates directly on GMEM refs, without pipelining.

  Unlike ``lower_pipelined_jaxpr_to_module``, the ``jaxpr`` is expected to take
  the GMEM refs for all inputs and outputs directly, without any block-level
  pipelining. This is used by ``mpmd_map``, whose kernels perform their own
  memory transfers.
  """
  assert len(jaxpr.outvars) == 0
  if params.dimension_semantics is not None:
    raise NotImplementedError(
        "dimension_semantics= is not supported by the unpipelined lowering"
    )
  block = (128 * (gpu_mesh.num_threads or 1), 1, 1)
  axis_names = _AxisNames(
      gpu_mesh.grid_names, gpu_mesh.cluster_names, gpu_mesh.thread_name
  )
  return lower_jaxpr_to_module(
      jax_mesh,
      axis_names,
      tuple(gpu_mesh.grid),
      block,
      tuple(gpu_mesh.cluster),
      in_shapes,
      out_shapes,
      jaxpr,
      params,
      (),
      outer_traceback=outer_traceback,
  )


def lower_jaxpr_to_module(
    jax_mesh: mesh_lib.Mesh | None,
    axis_names: _AxisNames,
    grid: tuple[int, ...],
    block: tuple[int, int, int],
    cluster: tuple[int, ...],
    in_shapes: Sequence[jax_core.ShapedArray],
    out_shapes: Sequence[jax_core.ShapedArray],
    jaxpr: jax_core.Jaxpr,
    params: gpu_core.CompilerParams,
    consts=(),
    outer_traceback: xc.Traceback | None = None,
) -> LoweringResult:
  debug_info = jaxpr.debug_info
  approx_math = params.approx_math
  lowering_semantics = params.lowering_semantics

  if len(cluster) < 3:
    cluster = (1,) * (3 - len(cluster)) + cluster
  else:
    assert len(cluster) == 3

  if len(grid) <= 3:
    squashed_dims = ()
    parallel_grid = (1,) * (3 - len(grid)) + grid
  else:
    # If we have >3 parallel dimensions, we flatten all but the minormost 2 dims.
    # Ex: (2, 3, 4, 5) -> (6, 4, 5)
    squashed_dims = grid[:-2]
    parallel_grid = (math.prod(grid[:-2]), *grid[-2:])

  # We reverse the order because Pallas prefers row-major iteration while the
  # CUDA runtime prefers column-major iteration.
  parallel_grid = parallel_grid[::-1]
  cluster = cluster[::-1]
  squashed_dims = squashed_dims[::-1]
  axis_names = axis_names.reverse()

  uses_pdl = gpu_core._pdl_effect in jaxpr.effects
  rs = _estimate_resources(
      ResourceEstimatorContext(
          reduction_scratch_bytes=params.reduction_scratch_bytes,
          axis_names=axis_names, lowering_semantics=lowering_semantics
      ),
      jaxpr,
  )

  def body(launch_ctx: mgpu.LaunchContext, *buffers: Any):
    *buffers_gmem, (
        runtime_smem,
        runtime_barriers,
        runtime_tmem,
    ) = buffers
    num_input_buffers = (len(in_shapes) +
                         len(rs.scoped_gmem_semaphores))
    input_buffers_gmem = buffers_gmem[:num_input_buffers]
    output_buffers_gmem = buffers_gmem[num_input_buffers:]

    scoped_gmem_semaphores = {}
    # pyrefly: ignore[no-matching-overload]
    for collective_axes in sorted(rs.scoped_gmem_semaphores, reverse=True):
      num_sems = rs.scoped_gmem_semaphores[collective_axes]
      # Extract the semaphores local to the current scope.
      index = ir.IndexType.get()
      # TODO(justinfu): Compute scope_idx for general collective_axes.
      # scope_idx computes axis_index(all_axes - collective_axes)
      if _is_block_local_scope(collective_axes, axis_names):
        scope_idx = arith_dialect.index_castui(index, mgpu_utils.block_idx())
      elif _is_global_scope(collective_axes, axis_names):
        scope_idx = _as_index(0)
      else:
        raise NotImplementedError(
            f"Unimplemented scope for semaphores: {collective_axes=}")
      scoped_gmem_semaphores[collective_axes] = mgpu.memref_slice(
          output_buffers_gmem[-1],
          mgpu.ds(
              arith_dialect.muli(
                  scope_idx, arith_dialect.constant(index, num_sems)
              ),
              num_sems,
          ),
      )
      # The semaphore buffer is an aliased input/output, so we need to skip it
      # in both the inputs and outputs.
      input_buffers_gmem = input_buffers_gmem[:-1]
      output_buffers_gmem = output_buffers_gmem[:-1]
    buffers_gmem = [*input_buffers_gmem, *output_buffers_gmem]

    grouped_barriers: MutableMapping[AnyBarrier, MutableSequence[AnyBarrierRef]]
    grouped_barriers = collections.defaultdict(list)
    for barrier, barrier_ref in zip(rs.barriers, runtime_barriers):
      grouped_barriers[barrier].append(barrier_ref)
    if runtime_tmem is not None:
      if lowering_semantics == mgpu.LoweringSemantics.Lane:
        tmem_cols = math.prod(runtime_tmem.shape) // tcgen05.TMEM_ROWS
        tmem_base = runtime_tmem.address
      else:
        tmem_cols = math.prod(runtime_tmem.type.shape) // tcgen05.TMEM_ROWS
        tmem_base = runtime_tmem
    else:
      tmem_cols = 0
      tmem_base = None

    single_wg_lane_predicate = mgpu.single_thread_predicate(
        scope=mgpu.ThreadSubset.WARPGROUP)
    single_warp_lane_predicate = mgpu.single_thread_predicate(
        scope=mgpu.ThreadSubset.WARP)

    module_ctx = ModuleContext(
        mlir.sanitize_name(debug_info.func_name),
        axis_names,
        [
            _program_id(axis, squashed_dims, len(grid))
            for axis in range(len(grid))
        ],
        approx_math,
        single_wg_lane_predicate,
        single_warp_lane_predicate,
        smem_requested_bytes=math.prod(ir.MemRefType(runtime_smem.type).shape),
        smem_used_bytes=0,
        tmem_requested_cols=tmem_cols,
        tmem_used_cols=0,
        tmem_base=tmem_base,
        scoped_gmem_used_semaphores={k: 0 for k in scoped_gmem_semaphores},
        scoped_gmem_semaphore_base_ptr=scoped_gmem_semaphores,
        runtime_barriers=grouped_barriers,
        name_stack=source_info_util.NameStack(),
        traceback_caches=mlir.TracebackCaches(),
        squashed_dims=squashed_dims,
        lowering_semantics=lowering_semantics,
        primitive_semantics=gpu_core.PrimitiveSemantics.Warpgroup,
        mesh_info=pallas_utils.MeshInfo.from_mesh(jax_mesh)
        if jax_mesh is not None
        else None,
        auto_barriers=not params.unsafe_no_auto_barriers,
        reduction_scratch_bytes=params.reduction_scratch_bytes,
        outer_traceback=outer_traceback,
    )
    del runtime_smem, grouped_barriers, runtime_barriers
    _ = lower_jaxpr_to_mosaic_gpu(
        module_ctx, launch_ctx, jaxpr, buffers_gmem, consts
    )

  scratch_buffers: list[Any] = [
      jax.ShapeDtypeStruct(shape=[rs.smem_scratch_bytes], dtype=np.int8),
      rs.barriers,
  ]
  if rs.tmem_scratch_cols > 0 and rs.tmem_collective_scratch_cols > 0:
    raise ValueError(
        "Can't mix collective and non-collective TMEM allocations within the"
        " same kernel."
    )
  tmem_scratch_cols = rs.tmem_scratch_cols + rs.tmem_collective_scratch_cols
  if tmem_scratch_cols > 0:
    scratch_buffers.append(
        mgpu.TMEM(
            shape=(tcgen05.TMEM_ROWS, tmem_scratch_cols),
            dtype=np.int32,
            collective=rs.tmem_collective_scratch_cols > 0,
        ),
    )
  else:
    scratch_buffers.append(None)

  prof_spec = None
  if params.profile_space:
    # Each range is 2 events, each event costs 2 entries.
    if params.profile_trace_scope == gpu_core.TraceScope.WARP:
      trace_scope = mgpu.ThreadSubset.WARP
    elif params.profile_trace_scope == gpu_core.TraceScope.WARPGROUP:
      trace_scope = mgpu.ThreadSubset.WARPGROUP
    else:
      raise ValueError(f"Unsupported trace scope: {params.profile_trace_scope}")
    prof_spec = mgpu_profiler.ProfilerSpec(
        params.profile_space * 2 * 2,
        dump_path=params.profile_dir,
        trace_scope=trace_scope,
        bounds_check=params.profile_bounds_check,
    )
  cuda_grid = tuple(map(operator.mul, parallel_grid, cluster))

  scoped_semaphores_shape = []
  for collective_axes in sorted(rs.scoped_gmem_semaphores):  # pyrefly: ignore[bad-specialization]
    num_sems = rs.scoped_gmem_semaphores[collective_axes]
    # TODO(justinfu): Compute axis_size for general collective_axes.
    # axis_size computes axis_size(all_axes - collective_axes)
    if _is_block_local_scope(collective_axes, axis_names):
      axis_size = math.prod(cuda_grid)
    elif _is_global_scope(collective_axes, axis_names):
      axis_size = 1
    else:
      raise NotImplementedError(
          f"Unimplemented scope for semaphores: {collective_axes=}")
    scoped_semaphores_shape.append(
        jax.ShapeDtypeStruct(
            shape=(axis_size * num_sems,), dtype=np.int32
        ),
    )
  scoped_semaphores_shape = tuple(scoped_semaphores_shape)

  if outer_traceback is not None:
    frame = source_info_util.user_frame(outer_traceback)
    if frame is not None:
      base_loc = ir.Location.file(
          mlir.get_canonical_source_file(
              frame.file_name, mlir.TracebackCaches()
          ),
          frame.start_line,
          frame.start_column,
      )
    else:
      base_loc = None
  else:
    base_loc = None

  def _to_physical_shape(
      s: jax.ShapeDtypeStruct | jax_core.ShapedArray,
  ) -> jax.ShapeDtypeStruct:
    dtype = s.dtype
    if isinstance(dtype, dtypes.ExtendedDType):
      dtype = jax_core.physical_element_aval(dtype).dtype
    return jax.ShapeDtypeStruct(s.shape, dtype)

  # NOTE: new_out_shapes has out_shapes, then semaphores_shape and
  # optionally the profiler buffer.
  module, new_out_shapes, _, launch_ctx = mgpu_core._lower_as_gpu_kernel(
      body,
      grid=cuda_grid,
      cluster=cast(tuple[int, int, int], cluster),
      block=block,
      in_shapes=tuple(_to_physical_shape(s) for s in (*in_shapes, *scoped_semaphores_shape)),
      out_shape=(*out_shapes, *scoped_semaphores_shape),
      inout_shape=(),
      smem_scratch_shape=scratch_buffers,
      lowering_semantics=lowering_semantics,
      module_name=mlir.sanitize_name(debug_info.func_name),
      kernel_name=mlir.sanitize_name(debug_info.func_name),
      prof_spec=prof_spec,
      jax_mesh=jax_mesh,
      base_loc=base_loc,
      uses_pdl=uses_pdl,
      is_multi_process=params.is_multi_process,
  )

  mgpu_core.lower_mgpu_module(
      module,
      launch_ctx,
      lowering_semantics,
      auto_barriers=not params.unsafe_no_auto_barriers,
  )

  dump_options = mgpu.dialect.get_or_set_dump_options(module)
  if dump_options.resources:
    if prof_spec is not None:
      smem_scratch_bytes = rs.smem_scratch_bytes + prof_spec.smem_bytes(block)
      rs = dataclasses.replace(rs, smem_scratch_bytes=smem_scratch_bytes)
    mgpu_utils.dump_to_file_or_stdout(
        str(rs),
        f"{dump_options.module_basename}.resources.txt",
        dump_options.dump_path,
    )

  return LoweringResult(
      module, cuda_grid, block, new_out_shapes, prof_spec,
      scoped_semaphores_shape,
  )


mosaic_lowering_rules = {
    # Lowering rules when using Mosaic GPU lane semantics.
    gpu_core.LANExWG_SEMANTICS: {} ,
    gpu_core.LANExWARP_SEMANTICS: {} ,
    # Lowering rules when using Mosaic GPU warpgroup semantics.
    gpu_core.WGxWG_SEMANTICS: {},
    gpu_core.WGxWARP_SEMANTICS: {},
}


def register_lowering_rule(
    primitive: jax_core.Primitive,
    lowering_semantics: mgpu.LoweringSemantics,
    primitive_semantics: gpu_core.PrimitiveSemantics = gpu_core.PrimitiveSemantics.Warpgroup,
):
  def deco(fn):
    mosaic_lowering_rules[
        (lowering_semantics, primitive_semantics)][primitive] = fn
    return fn

  return deco


def _compute_name_stack_updates(
    old_name_stack: list[str],
    new_name_stack: list[str]
) -> tuple[list[str], list[str]]:
  common_prefix_idx = 0
  for i, (old, new) in enumerate(unsafe_zip(old_name_stack, new_name_stack)):
    if old == new:
      common_prefix_idx = i+1
    else:
      break
  return old_name_stack[common_prefix_idx:], new_name_stack[common_prefix_idx:]


def lower_jaxpr_to_mosaic_gpu(
    module_ctx: ModuleContext,
    launch_ctx: mgpu.LaunchContext,
    jaxpr: jax_core.Jaxpr,
    args: Sequence[ir.Value],
    consts=(),
) -> Sequence[ir.Value]:
  env = {}

  def read_env(atom: jax_core.Atom):
    return atom.val if isinstance(atom, jax_core.Literal) else env[atom]

  def write_env(var: jax_core.Var, val, require_value: bool = True):
    env[var] = val
    # TODO(apaszke): Handle other avals (refs, etc.).
    if isinstance(aval := var.aval, jax_core.ShapedArray):
      # TODO(apaszke): Clarify the type invariants for lane semantics?
      if module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
        # Shaped arrays must be vectors if and only if their shape is non-empty.
        # Those with empty shapes should be represented by their scalar type.
        dtype = aval.dtype
        if isinstance(dtype, dtypes.ExtendedDType):
          dtype = jax_core.physical_element_aval(dtype).dtype
        mlir_dtype = mgpu_utils.dtype_to_ir_type(dtype)
        if not isinstance(val, ir.Value):
          if require_value:
            raise AssertionError(f"Shaped arrays must be represented by ir.Values, got: {val}")
          else:
            if aval.shape:
              raise AssertionError("Only scalars can be represented by non-ir.Values")
            return  # Skip following checks.
        if aval.shape:
          if not isinstance(val.type, ir.VectorType):
            raise AssertionError(f"Non-scalar arrays must be represented by vectors, got: {val.type}")
          vty = ir.VectorType(val.type)
          if vty.element_type != mlir_dtype:
            raise AssertionError(f"Vector element type must match ShapedArray dtype, got: {val.type} != {mlir_dtype}")
          if tuple(vty.shape) != aval.shape:
            raise AssertionError(f"Vector shape must match ShapedArray shape, got: {vty.shape} != {aval.shape}")
        else:
          if isinstance(val.type, ir.VectorType):
            raise AssertionError(f"Scalars must be represented by non-vector types, got: {val.type}")
          if val.type != mlir_dtype:
            raise AssertionError(f"Scalar type must match ShapedArray dtype, got: {val.type} != {mlir_dtype}")

  foreach(
      functools.partial(write_env, require_value=False), jaxpr.constvars, consts
  )
  foreach(functools.partial(write_env, require_value=False), jaxpr.invars, args)

  # TODO(justinfu): Handle transform scopes.
  last_local_name_stack: list[str] = []
  named_regions = []
  for i, eqn in enumerate(jaxpr.eqns):
    invals = map(read_env, eqn.invars)
    eqn_name_stack = module_ctx.name_stack + eqn.source_info.name_stack
    with config.include_full_tracebacks_in_locations(False):
      loc = mlir.source_info_to_location(
          module_ctx,
          eqn.primitive,
          eqn_name_stack,
          eqn.source_info.traceback or module_ctx.outer_traceback,
      )
    with source_info_util.user_context(eqn.source_info.traceback), loc:
      if eqn.primitive not in mosaic_lowering_rules[
          (module_ctx.lowering_semantics, module_ctx.primitive_semantics)]:
        raise NotImplementedError(
            "Unimplemented primitive in Pallas Mosaic GPU lowering:"
            f" {eqn.primitive.name} for lowering semantics"
            f" {module_ctx.lowering_semantics} and user thread semantics"
            f" {module_ctx.primitive_semantics}. Please file an issue at"
            " https://github.com/jax-ml/jax/issues/new/choose."
        )
      new_local_name_stack = [scope.name for scope in eqn.source_info.name_stack.stack]
      popped, pushed = _compute_name_stack_updates(last_local_name_stack, new_local_name_stack)
      last_local_name_stack = new_local_name_stack
      for _ in popped:
        named_regions.pop().close()
      for name in pushed:
        wrapper_stack = contextlib.ExitStack()
        wrapper_stack.enter_context(launch_ctx.named_region(name))
        named_regions.append(wrapper_stack)
      rule = mosaic_lowering_rules[
          (module_ctx.lowering_semantics, module_ctx.primitive_semantics)
          ][eqn.primitive]
      # If the equation is immediately followed by a layout cast on its output,
      # we provide the layout as a hint to the rule.
      out_layout_hint = None
      if i + 1 < len(jaxpr.eqns):
        lookahead_eqn = jaxpr.eqns[i + 1]
        is_layout_cast = lookahead_eqn.primitive == gpu_core.layout_cast_p
        # We provide the hint for the first output only.
        uses_eqn_output = lookahead_eqn.invars == eqn.outvars[:1]
        if is_layout_cast and uses_eqn_output:
          out_layout_hint = lookahead_eqn.params["new_layout"].to_mgpu()
      rule_ctx = LoweringRuleContext(
          module_ctx,
          launch_ctx,
          avals_in=[cast(ShapedAbstractValue, v.aval) for v in eqn.invars],
          avals_out=[cast(ShapedAbstractValue, v.aval) for v in eqn.outvars],
          prim=eqn.primitive,
          out_layout_hint=out_layout_hint,
      )
      try:
        outvals = rule(rule_ctx, *invals, **eqn.params)
      except LoweringError:
        raise  # We only add the extra info to the innermost exception.
      except Exception as e:
        if not config.jax_pallas_verbose_errors.value:
          raise
        inval_types = map(lambda t: getattr(t, "type", None), invals)
        raise LoweringError(
            f"Exception while lowering eqn:\n  {eqn}\nWith context:\n "
            f" {rule_ctx}\nWith inval types={inval_types}\nIn jaxpr:\n{jaxpr}"
        ) from e
      if eqn.primitive.multiple_results:
        foreach(write_env, eqn.outvars, outvals)
      else:
        write_env(eqn.outvars[0], outvals)
  while named_regions:  # Drain the name stack.
    named_regions.pop().close()
  return map(read_env, jaxpr.outvars)


@register_lowering_rule(primitives.program_id_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(
    primitives.program_id_p, mgpu.LoweringSemantics.Warpgroup)
def _program_id_lowering_rule(ctx: LoweringRuleContext, axis):
  if ctx.module_ctx.program_ids is None:
    raise NotImplementedError("pl.program_id() is not supported in this context")
  deprecations.warn(
      "jax-pallas-mgpu-program-id",
      "Using pl.program_id in Pallas MGPU is deprecated. Use jax.lax.axis_index"
      " instead.",
      stacklevel=2,
  )
  return ctx.module_ctx.program_ids[axis]


@register_lowering_rule(primitives.multiple_of_p, mgpu.LoweringSemantics.Lane)
def _multiple_of_lane_lowering_rule(ctx: LoweringRuleContext, val, *, values):
  del ctx, values
  # Under Lane lowering semantics, we currently don't do anything with the
  # annotation.
  return val


@register_lowering_rule(
    primitives.multiple_of_p, mgpu.LoweringSemantics.Warpgroup
)
def _multiple_of_wg_lowering_rule(ctx: LoweringRuleContext, val, *, values):
  [aval] = ctx.avals_in
  if aval.shape:
    raise NotImplementedError("multiple_of only supports scalar inputs.")
  for multiple in values:
    val = mgpu.dialect.assume_multiple(val, multiple)
  return val


def _unravel_program_id(
    block_id: ir.Value,
    axis: int,
    dimensions: tuple[int, ...],
    row_major: bool = False
) -> ir.Value:
  """Computes the program ID for axes compressed into one block dimension."""
  if row_major:
    div_value = math.prod(dimensions[axis+1:])
  else:
    div_value = math.prod(dimensions[:axis])
  div_value = _as_index(_i32_constant(div_value))
  pid = arith_dialect.divui(block_id, div_value)
  axis_size = _as_index(_i32_constant(dimensions[axis]))
  pid = arith_dialect.remui(pid, axis_size)
  return arith_dialect.index_cast(ir.IntegerType.get_signless(32), pid)


def _program_id(
    parallel_axis: int, squashed_dims: tuple[int, ...], grid_size: int
) -> ir.Value:
  """Returns the id of the current kernel instance along the given axis in the original Pallas grid."""
  if parallel_axis < len(squashed_dims):
    # All squashed dimensions are mapped to Dimension.z.
    block_id = gpu_dialect.block_id(gpu_dialect.Dimension.z)
    idx = len(squashed_dims) - 1 - parallel_axis
    return _unravel_program_id(block_id, idx, squashed_dims)
  else:
    idx = grid_size - 1 - parallel_axis
    assert idx in (0, 1, 2)
    return arith_dialect.index_cast(
        ir.IntegerType.get_signless(32),
        gpu_dialect.block_id(gpu_dialect.Dimension(idx)))


def _lower_fun(
    fun: Callable,
    *,
    in_avals: Any | None = None,
) -> Callable:

  def f_lowered(ctx: LoweringRuleContext, *args, **params):
    is_leaf = lambda x: isinstance(
        x, (mgpu.FragmentedArray, mgpu.WGMMAAccumulator)
    )
    flat_args, in_tree = tree_util.tree_flatten(args, is_leaf=is_leaf)
    if in_avals is None:
      unflat_avals = tree_util.tree_unflatten(in_tree, ctx.avals_in)
    else:
      _, aval_tree = tree_util.tree_flatten(in_avals)
      if in_tree != aval_tree:
        raise ValueError(
            "args and in_avals pytrees mismatch:\nargs tree:"
            f" {in_tree}\navals tree: {aval_tree}\nargs: {args}\navals:"
            f" {in_avals}"
        )
      unflat_avals = in_avals

    in_avals_ft = ft.flatten_static_argnums_argnames(
        unflat_avals,
        params,
        static_argnums=(),
        static_argnames=tuple(params.keys()),
    )
    debug_info = api_util.debug_info(
        "mosaic_gpu lower_fun", fun, unflat_avals, params,
        static_argnames=tuple(params.keys()),
    )
    closed_jaxpr, out_avals_ft = pe.trace_to_jaxpr(
        fun, in_avals_ft, debug_info=debug_info, requires_low=True
    )
    if closed_jaxpr.consts:
      raise NotImplementedError("lower_fun should not capture constvars")
    jaxpr = pe.convert_constvars_jaxpr(closed_jaxpr)
    out = lower_jaxpr_to_mosaic_gpu(
        ctx.module_ctx, ctx.launch_ctx, jaxpr, flat_args, closed_jaxpr.consts
    )
    out_tree = out_avals_ft.tree
    return tree_util.tree_unflatten(out_tree, out)

  return f_lowered


@overload
def _extract_aliased_ref(
  ref: ir.Value,
  ref_aval: state_types.AbstractRef,
  transform_avals: Sequence[state_types.Transform],
  transforms: Sequence[state_types.Transform],
  lowering_semantics: mgpu.LoweringSemantics,
) -> tuple[
    ir.Value,
    state_types.AbstractRef,
    Sequence[state_types.Transform],
    Sequence[state_types.Transform]
]:
  ...

@overload
def _extract_aliased_ref(
  ref: tcgen05.TMEMRef,
  ref_aval: state_types.AbstractRef,
  transform_avals: Sequence[state_types.Transform],
  transforms: Sequence[state_types.Transform],
  lowering_semantics: mgpu.LoweringSemantics,
) -> tuple[
    tcgen05.TMEMRef,
    state_types.AbstractRef,
    Sequence[state_types.Transform],
    Sequence[state_types.Transform]
]:
  ...

def _extract_aliased_ref(
    ref: ir.Value | tcgen05.TMEMRef,
    ref_aval: state_types.AbstractRef,
    transform_avals: Sequence[state_types.Transform],
    transforms: Sequence[state_types.Transform],
    lowering_semantics: mgpu.LoweringSemantics,
) -> tuple[
    ir.Value | tcgen05.TMEMRef,
    state_types.AbstractRef,
    Sequence[state_types.Transform],
    Sequence[state_types.Transform],
]:
  # Looks for the first transform being an ExtractAliasedRef and pulls out the
  # Ref there, updating the transforms.
  match transforms:
    case (
        gpu_core.ExtractAliasedRef(dtype, transformed_shape, offset, alias_group_idx, layout) as t,
        *other_transforms,
    ):
      ref_aval = t.transform_type(ref_aval)
      mlir_dtype = mgpu_utils.dtype_to_ir_type(dtype)
      if isinstance(ref, tcgen05.TMEMRef):
        assert layout is not None
        if ref.shape[0] != transformed_shape[0]:
          raise ValueError(
              "TMEM aliasing only supported for Refs with the same first"
              f" dimension, got {ref.shape[0]} != {transformed_shape[0]}."
          )
        address = arith_dialect.addi(ref.address, _i32_constant(offset))
        ref = tcgen05.TMEMRef(
            address=address,
            shape=cast(tuple[int, int], transformed_shape),
            dtype=mlir_dtype,
            layout=layout,
        )
      else:
        assert isinstance(ref, ir.Value)  # make pyrefly happy
        input_ref_ty = ir.MemRefType(ref.type)
        if input_ref_ty.memory_space == mgpu_utils.smem():
          assert layout is None
          ref_bits = math.prod(transformed_shape) * mgpu_utils.bitwidth(
              mlir_dtype
          )
          if ref_bits % 8:
            raise NotImplementedError("Only byte-aligned bitcasts are supported.")
          assert offset % gpu_core.SMEM_ALIGNMENT == 0

          ref_ty = ir.MemRefType.get(
              transformed_shape, mlir_dtype, memory_space=mgpu_utils.smem()
          )
          if lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
            if not isinstance(ref.owner, mgpu.dialect.SliceSMEMOp):
              # This restriction can be lifted by:
              # - Using memref ops to get the pointer and offset of the base ref.
              # - Subtracting gpu_dialect.dynamic_shared_memory() from those to
              #   get the base offset relative to the beginning of SMEM.
              # - Implementing layout and lowering rules for all ops above.
              raise NotImplementedError(
                  "The base ref for aliases must come from a slice_smem op."
              )

            base_offset = ref.owner.offset.value
            total_offset = base_offset + offset

            assert ref.owner.alias_id is not None
            alloc_id = ref.owner.alias_id.value
            # TODO(bchetioui): Use a scheme resilient to hash collisions.
            alias_id = hash((offset, alloc_id, alias_group_idx))
            # The composite key formed of `(offset, alloc_id, alias_group_idx)`
            # is a unique identifier across:
            #   - different RefUnions (different `alloc_id`, since two
            #     distinct RefUnions represent two SMEM allocations);
            #   - different ref_groups within a RefUnion (different
            #     `alias_group_idx`);
            #   - different elements within a ref_group (different
            #     `offset` which is the offset of the particular element to
            #     the beginning of the RefUnion). This only holds in the absence
            #     of 0-sized refs, which don't serve a practical purpose anyway.
            ref = mgpu.dialect.slice_smem(
                ref_ty, total_offset, alias_id=alias_id
            )
          else:
            ref = memref_dialect.view(ref_ty, ref, _as_index(offset), [])
        elif input_ref_ty.memory_space == mgpu_utils.tmem():

          if isinstance(ref.owner, mgpu.dialect.SliceTmemOp):
            source_slice_op = ref.owner
          elif isinstance(
              ref.owner, mgpu.dialect.TmemLayoutCastOp
          ) and isinstance(ref.owner.operands[0].owner, mgpu.dialect.SliceTmemOp):
            source_slice_op = ref.owner.operands[0].owner
          else:
            raise NotImplementedError(f"Unsupported TMEM ref {ref}.")

          base_offset = source_slice_op.offset.value
          assert isinstance(base_offset, int)  # make pyrefly happy
          total_offset = base_offset + offset
          ref_ty = ir.MemRefType.get(
              transformed_shape, mlir_dtype, memory_space=mgpu_utils.tmem()
          )
          alloc_id = source_slice_op.alias_id
          assert alloc_id is not None
          # TODO(bchetioui): Use a scheme resilient to hash collisions.
          alias_id = hash((offset, alloc_id.value, alias_group_idx))
          slice_op = mgpu.dialect.SliceTmemOp(
              ref_ty, source_slice_op.source, total_offset
          )
          i64 = ir.IntegerType.get_signless(64)
          slice_op.attributes["alias_id"] = ir.IntegerAttr.get(i64, alias_id)
          ref = slice_op.result
          assert layout is not None
          layout_attr = mgpu.layouts.to_layout_attr(layout)
          ref = mgpu.dialect.tmem_layout_cast(ref, layout_attr)
        else:
          raise NotImplementedError("Unsupported memory space.")
      return (
          ref,
          ref_aval,
          transform_avals[1:],
          tuple(other_transforms),
      )
    case _:
      # No ExtractAliasedRef found, don't do anything.
      return ref, ref_aval, transform_avals, transforms


def _commute_transform(
    aval: jax_core.AbstractValue,
    t1: state_types.Transform,
    t2: state_types.Transform,
) -> tuple[state_types.Transform, state_types.Transform]:
  """Commutes two transforms.

  We pass in an `aval` to handle cases where the transforms, by themselves, do
  not provide enough information to determine how to commute.

  Args:
    aval: An abstract value.
    t1: A transform.
    t2: Another transform.

  Returns:
    Returns a tuple of transforms (t2', t1') such that
    t2'(t1'(aval)) == t1(t2(aval)).
  """
  match t1, t2:
    case (
        (
            gpu_core.UntilingTransform()
            | gpu_core.UnswizzleRef()
            | gpu_core.ExpandLeadingBatchDimensionsTransform()
        ) as t1,
        indexing.NDIndexer() as t2,
    ):
      new_indexer, new_t1 = t1.commute_ndindexer(aval, t2)
      return new_indexer, new_t1
    case (
        state_types.TransposeTransform() as t1,
        indexing.NDIndexer() as t2,
    ):
      return gpu_core.commute_transpose_indexer(aval, t1, t2)
    case (
        gpu_core.UnswizzleRef() as t1,
        state_types.ReshapeTransform() as t2,
    ):
      # pyrefly: ignore[bad-argument-type]
      new_reshape, new_unswizzle = t1.commute_reshape(aval, t2)
      return new_reshape, new_unswizzle
    case (
        gpu_core.UntilingTransform() | gpu_core.UnswizzleRef() as t1,
        gpu_core.TransposeTransform() as t2,
    ):
      if isinstance(aval, state_types.AbstractRef):
        aval = aval.inner_aval
      assert isinstance(aval, jax_core.ShapedArray)
      new_reshape, new_unswizzle = t1.commute_transpose(aval, t2)
      return new_reshape, new_unswizzle
    case (
        gpu_core.UntilingTransform() as t1,
        state_types.ReshapeTransform() as t2,
    ):
      if isinstance(aval, state_types.AbstractRef):
        aval = aval.inner_aval
      assert isinstance(aval, jax_core.ShapedArray)
      new_reshape, new_untile = t1.commute_reshape(aval, t2)
      return new_reshape, new_untile
    case (
        gpu_core.UntilingTransform() | gpu_core.UnswizzleRef() as t1,
        gpu_core.ClusterRefTransform()
        | gpu_core.MulticastRef()
        | gpu_core.PeerMemRef() as t2,
    ):
      return t2, t1
    case _:
      raise NotImplementedError(t1, t2)


def _lower_fn_with_avals(f, avals_in):
  def inner(ctx, *args):
    flat_args = tree_util.tree_leaves(args)
    in_avals_ft = ft.flatten_args(*avals_in)
    debug_info = api_util.debug_info(
        "Pallas Mosaic GPU lower_fn_with_avals", f, ("",) * len(args), {}
    ).with_unknown_names()
    jaxpr, out_avals_ft = pe.trace_to_jaxpr(
        f, in_avals_ft, debug_info=debug_info, requires_low=True
    )
    out_flat = lower_jaxpr_to_mosaic_gpu(
        ctx.module_ctx, ctx.launch_ctx, jaxpr, flat_args, jaxpr.consts
    )
    out_tree = out_avals_ft.tree
    return out_tree.unflatten(out_flat), out_avals_ft.unflatten()
  return inner


def _bubble_up_transform[T: state_types.Transform](
    ctx: LoweringRuleContext,
    aval: jax_core.AbstractValue,
    transforms: Sequence[state_types.Transform],
    transforms_avals: Sequence[state_types.Transform],
    t: T,
    t_aval: T,
) -> tuple[T, T, list[state_types.Transform], list[state_types.Transform]]:
  new_transforms_rev = []
  new_transforms_avals_rev = []
  for transform, transform_aval in reversed(zip(transforms, transforms_avals)):
    avals = (transform_aval, t_aval)
    (t, new_transform), (t_aval, new_transform_aval) = _lower_fn_with_avals(
        functools.partial(_commute_transform, aval), avals
    )(ctx, transform, t)
    if not isinstance(new_transform, gpu_core.IdentityTransform):
      new_transforms_rev.append(new_transform)
      new_transforms_avals_rev.append(new_transform_aval)
  new_transforms_rev.reverse()
  new_transforms_avals_rev.reverse()
  return (
      t,
      t_aval,
      new_transforms_rev,
      new_transforms_avals_rev,
  )


def _reinterpret_cast(ref: ir.Value, new_ref_aval: state_types.AbstractRef) -> ir.Value:
  ref_ty = ir.MemRefType(ref.type)
  strides, offset = ref_ty.get_strides_and_offset()
  # A sanity check. It doesn't do much to check that an offset is dynamic, but
  # if we ever get here through an unexpected path that slices with a non-zero
  # static offset, we'll at least catch it.
  assert offset == 0  or offset == ir.ShapedType.get_dynamic_stride_or_offset()
  expected_strides = mgpu_utils.get_contiguous_strides(ref_ty.shape)
  if expected_strides != strides:
    raise NotImplementedError(
        f"Expected contiguous strides {expected_strides} when applying "
        f"reinterpret_cast to {ref_ty} but got {strides}"
    )
  if offset == 0:
    layout = None
  else:
    layout = ir.StridedLayoutAttr.get(
        offset, mgpu_utils.get_contiguous_strides(new_ref_aval.shape)
    )
  new_ty = ir.MemRefType.get(
      new_ref_aval.shape, mgpu_utils.dtype_to_ir_type(new_ref_aval.dtype),
      memory_space=ref_ty.memory_space,
      layout=layout
  )
  if new_ty == ref_ty:
    return ref
  return mgpu.dialect.reinterpret_cast(new_ty, ref)


def _bubble_up_transforms_for_lowering(
    ctx: LoweringRuleContext,
    aval: jax_core.AbstractValue,
    transforms: Sequence[state_types.Transform],
    transform_avals: Sequence[state_types.Transform],
    *,
    handle_transposes: bool = True,
    handle_reshapes: bool = True,
) -> tuple[
    list[state_types.Transform],
    list[state_types.Transform],
    list[state_types.Transform],
    list[state_types.Transform],
]:
  """Bubbles up eligible `transforms` to the head of the sequence.

  The transforms to lower are commuted to the head of the sequence of
  transforms, such that the (unmaterialized) new order of transform
  application is:

    (*transforms_to_lower, *remaining_transforms)

  Returns a tuple where:
    * first element is a list of bubbled up transforms (i.e.,
    `transforms_to_lower`)
    * second element is a list of avals corresponding to bubbled up transforms
    * third element is a list of remaining transforms (i.e.,
    `remaining_transforms`)
    * fourth element is a list of avals corresponding to remaining transforms
  """
  bubbled_up_transforms = []
  bubbled_up_transform_avals = []
  remaining_transforms = []
  remaining_transform_avals = []

  for t_aval, t in zip(transform_avals, transforms):
    match t:
      case TransposeTransform():
        should_bubble_up = handle_transposes
      case ReshapeTransform():
        should_bubble_up = handle_reshapes
      case (
          indexing.NDIndexer()
          | gpu_core.PeerMemRef()
          | gpu_core.MulticastRef()
          | gpu_core.ClusterRefTransform()
      ):
        should_bubble_up = True
      case _:
        should_bubble_up = False

    if should_bubble_up:
      (
          t,
          t_aval,
          remaining_transforms,
          remaining_transform_avals,
      ) = _bubble_up_transform(
          ctx,
          aval,
          remaining_transforms,
          remaining_transform_avals,
          t,
          t_aval,
      )
      bubbled_up_transforms.append(t)
      bubbled_up_transform_avals.append(t_aval)
      aval = t_aval.transform_type(aval)
    else:
      remaining_transforms.append(t)
      remaining_transform_avals.append(t_aval)

  assert len(bubbled_up_transforms) == len(bubbled_up_transform_avals)
  assert len(remaining_transforms) == len(remaining_transform_avals)
  return (
      bubbled_up_transforms,
      bubbled_up_transform_avals,
      remaining_transforms,
      remaining_transform_avals,
  )


def _handle_transforms[T: (ir.Value, tcgen05.TMEMRef)](
    ctx: LoweringRuleContext,
    ref_aval: state_types.AbstractRef,
    ref: T,
    transform_avals: Sequence[state_types.Transform],
    transforms: Sequence[state_types.Transform],
    *,
    handle_transposes=True,
    handle_reshapes=True,
    allow_peer_refs=False,
    allow_multicast_refs=False,
) -> tuple[T, state_types.AbstractRef, Sequence[state_types.Transform]]:
  # Before we handle other transforms, we resolve any possible leading
  # aliasing transform.
  ref, ref_aval, transform_avals, transforms = _extract_aliased_ref(  # pyrefly: ignore[bad-assignment,no-matching-overload]
      ref,
      ref_aval,
      transform_avals,
      transforms,
      ctx.module_ctx.lowering_semantics,
  )

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    # We only bubble up transforms here to verify that all the specified
    # transforms can be commuted correctly with the BlockSpec transforms.
    _bubble_up_transforms_for_lowering(
        ctx,
        ref_aval,
        transforms,
        transform_avals,
        handle_transposes=False,
        handle_reshapes=handle_reshapes,
    )

    spec_transforms = []
    num_block_spec_transforms = 0
    for t in transforms:
      if isinstance(t, (gpu_core.UntilingTransform, gpu_core.UnswizzleRef)):
        spec_transforms = [t.undo(ref_aval)] + spec_transforms
        ref_aval = cast(state_types.AbstractRef, t.transform_type(ref_aval))
        num_block_spec_transforms += 1
      else:
        break
    assert isinstance(ref, ir.Value)
    if spec_transforms:
      transforms_attr = ir.ArrayAttr.get([
          gpu_core.to_transform_attr(t) for t in spec_transforms
      ])
      ref = cast(
          T,
          mgpu.dialect.with_transforms(_reinterpret_cast(ref, ref_aval), transforms_attr)
      )
      transforms = transforms[num_block_spec_transforms:]
      transform_avals = transform_avals[num_block_spec_transforms:]
      if any(isinstance(t, (gpu_core.UntilingTransform, gpu_core.UnswizzleRef)) for t in transforms):
        raise ValueError("Unexpected untiling or unswizzle transform found in "
                         f"remaining transforms: {transforms}.")

  (
      bubbled_up_transforms,
      bubbled_up_transform_avals,
      remaining_transforms,
      _,
  ) = _bubble_up_transforms_for_lowering(
      ctx,
      ref_aval,
      transforms,
      transform_avals,
      handle_transposes=handle_transposes,
      handle_reshapes=handle_reshapes,
  )

  transformed_ref: Any = ref
  peer_device_id = None
  is_multicast = False
  cluster_dim = None
  cluster_idx = None

  for t_aval, t in zip(bubbled_up_transform_avals, bubbled_up_transforms):
    match t:
      case indexing.NDIndexer() as indexer:
        assert isinstance(t_aval, indexing.NDIndexer)
        if t_aval.int_indexer_shape:
          raise NotImplementedError("int_indexer_shape non-empty")
        indices = _ndindexer_indices(indexer)
        if (
            isinstance(transformed_ref, tcgen05.TMEMRef)
            and ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane
        ):
          transformed_ref = transformed_ref.slice(*indices)
        else:
          transformed_ref = mgpu_utils.memref_slice(transformed_ref, indices)
        ref_aval = t_aval.transform_type(ref_aval)
      case TransposeTransform() as t:
        assert handle_transposes
        if isinstance(transformed_ref, tcgen05.TMEMRef):
          raise ValueError("TMEM transpose not allowed.")
        transformed_ref = mgpu.memref_transpose(
            transformed_ref, t.permutation
        )
        ref_aval = t_aval.transform_type(ref_aval)  # pyrefly: ignore [bad-assignment]
      case ReshapeTransform() as t:
        assert handle_reshapes
        if isinstance(transformed_ref, tcgen05.TMEMRef):
          raise ValueError("TMEM reshape not allowed.")
        transformed_ref = mgpu.memref_reshape(transformed_ref, t.shape)
        ref_aval = t_aval.transform_type(ref_aval)  # pyrefly: ignore [bad-assignment]
      case gpu_core.PeerMemRef(device_id, device_id_type):
        assert isinstance(t_aval, gpu_core.PeerMemRef)
        peer_device_id = _device_id_to_logical(
            ctx, device_id, device_id_type, t_aval.device_id
        )
      case gpu_core.MulticastRef(_):
        if not allow_multicast_refs:
          raise NotImplementedError(
              "Multicast references are not allowed in the lowering of this"
              " primitive."
          )
        is_multicast = True
      case gpu_core.ClusterRefTransform(dims, idxs):
        if len(dims) != 1:
          raise NotImplementedError(
              "Only overriding a single cluster axis is supported for now."
          )
        cluster_dim = _resolve_cluster_axis(ctx.module_ctx.axis_names, dims[0])
        cluster_idx = _as_index(idxs[0])
      case _:
        raise AssertionError(
            f"Transform {t} has no defined lowering rule."
        )

  if cluster_dim is not None:
    assert cluster_idx is not None
    if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
      i32 = ir.IntegerType.get_signless(32)
      kwargs = dict(x=None, y=None, z=None)
      kwargs[cluster_dim.name] = arith_dialect.index_cast(i32, cluster_idx)
      transformed_ref = mgpu.dialect.get_cluster_ref(transformed_ref, **kwargs)
    else:
      transformed_ref = mgpu.get_cluster_ref(
          transformed_ref, cluster_dim, cluster_idx, generic=False
      )
  if peer_device_id is not None:
    assert not is_multicast
    if not allow_peer_refs:
      raise NotImplementedError(
          "Peer device references are not allowed in the lowering of this"
          " primitive."
      )
    transformed_ref = ctx.launch_ctx.to_remote(
        transformed_ref, _ensure_ir_value(peer_device_id, jnp.int32)
    )
  if is_multicast:
    transformed_ref = ctx.launch_ctx.to_remote_multicast(transformed_ref)
  assert isinstance(ref_aval, state_types.AbstractRef)
  return transformed_ref, ref_aval, remaining_transforms


def _ndindexer_indices(
    indexer: indexing.NDIndexer, allow_arrays: bool = False
) -> tuple[Any, ...]:
  indices: list[Any] = []
  for idx in indexer.indices:
    if (isinstance(idx, mgpu.FragmentedArray) and idx.shape) or (
        isinstance(idx, ir.Value) and isinstance(idx.type, ir.VectorType)
    ):
      if not allow_arrays:
        raise ValueError("Arrays are not supported as indices.")
      indices.append(idx)
    elif not isinstance(idx, indexing.Slice):
      indices.append(_as_index(idx))
    elif not idx.is_dynamic_start and not idx.is_dynamic_size:
      indices.append(slice(idx.start, idx.start + idx.size, idx.stride))
    elif idx.stride == 1:
      if idx.is_dynamic_size:
        raise NotImplementedError("Dynamic slice size not supported.")
      indices.append(
          mgpu.DynamicSlice(
              _as_index(idx.start) if idx.is_dynamic_start else idx.start,  # pyrefly: ignore[bad-argument-type]
              int(idx.size),
          )
      )
    else:
      raise NotImplementedError(f"Unsupported slice: {idx}")
  return tuple(indices)


@register_lowering_rule(sp.get_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(sp.get_p, *gpu_core.LANExWARP_SEMANTICS)
def _get_lowering_rule(
    ctx: LoweringRuleContext, x_ref, *leaves, tree, optimized=True
):
  if isinstance(x_ref, tcgen05.TMEMRef):
    raise RuntimeError(
        "Loads from TMEM are asynchronous operations and cannot be performed"
        " using the usual syntax. Please use plgpu.async_load_tmem instead."
    )
  if (
      ctx.avals_out[0].shape
      and ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp
  ):
    raise ValueError("Can only load scalars in warp-level code.")
  if not isinstance(x_ref, ir.Value) and isinstance(x_ref, ir.MemRefType):
    raise TypeError(f"Can only load from references (got {x_ref}).")
  dtype = ctx.avals_out[0].dtype

  transforms = jax.tree.unflatten(tree, leaves)
  assert isinstance(ctx.avals_in[0], state_types.AbstractRef)
  transform_avals = tree.unflatten(ctx.avals_in[1:])

  # Swizzle always applies first, to the raw addresses, so we pop it immediately.
  if transforms and isinstance(transforms[0], gpu_core.UnswizzleRef):
    swizzle = transforms[0].swizzle
    transforms = transforms[1:]
    transform_avals = transform_avals[1:]
  elif (
      len(transforms) > 1
      and isinstance(transforms[0], gpu_core.ExtractAliasedRef)
      and isinstance(transforms[1], gpu_core.UnswizzleRef)
  ):
    swizzle = transforms[1].swizzle
    transforms = [transforms[0], *transforms[2:]]
    transform_avals = [transform_avals[0], *transform_avals[2:]]
  else:
    swizzle = None

  # We verify tiling against swizzling in memory before transposes flip it.
  # Note: we do not check len(tiling) == 2 here because subsequent transforms
  # (e.g. ReshapeTransform) may commute with it to produce 2D tiling.
  tiling_transform = None
  if transforms and isinstance(transforms[0], gpu_core.UntilingTransform):
    tiling_transform = transforms[0]
  elif (
      len(transforms) > 1
      and isinstance(transforms[0], gpu_core.ExtractAliasedRef)
      and isinstance(transforms[1], gpu_core.UntilingTransform)
  ):
    tiling_transform = transforms[1]
  if tiling_transform is not None:
    if swizzle is None:
      raise NotImplementedError("Tiling without swizzle is not supported.")
    bw = dtypes.itemsize_bits(ctx.avals_out[0].dtype)
    expected_minor_tiling = swizzle * 8 // bw
    if tiling_transform.tiling[-1] != expected_minor_tiling:
      raise NotImplementedError(
          "Minor tiling dimension does not fit swizzle: "
          f" expected {expected_minor_tiling}, got {tiling_transform.tiling[-1]}"
      )

  x_smem, _, transforms = _handle_transforms(
      ctx, ctx.avals_in[0], x_ref, transform_avals, transforms,
      allow_peer_refs=True
  )
  del x_ref  # Don't use x_ref anymore. Use x_smem instead!

  is_signed = mgpu_utils.is_signed(dtype)

  if not ctx.avals_out[0].shape:  # The scalar case is simple.
    val = memref_dialect.load(x_smem, [])
    return mgpu.FragmentedArray.splat(val, shape=(), is_signed=is_signed)

  match transforms:
    case (gpu_core.UntilingTransform(tiling),):
      if len(tiling) != 2:
        raise NotImplementedError(f"Only 2D tiling is supported, got: {tiling}")
      return mgpu.FragmentedArray.load_tiled(
          x_smem,
          is_signed=is_signed,
          swizzle=swizzle,
          layout=ctx.out_layout_hint or mgpu.WGMMA_LAYOUT,
          optimized=optimized,
          tiling_rank=len(tiling),
      )
    case ():
      match ctx.out_layout_hint:
        case mgpu.WGStridedFragLayout(shape=shape, vec_size=vec_size):
          ref_ty = ir.MemRefType(x_smem.type)
          if shape != tuple(ref_ty.shape):
            raise ValueError(
                f"Unsupported shape {shape}, (expected {tuple(ref_ty.shape)})"
            )
          if swizzle is not None:
            raise NotImplementedError(
                "Unsupported swizzle transform with strided layout"
            )
          return mgpu.FragmentedArray.load_strided(
              x_smem,
              is_signed=is_signed,
              vec_size=vec_size,
          )
        case None:
          if swizzle is not None:
            raise NotImplementedError(
                "Unsupported swizzle transform with strided layout"
            )
          return mgpu.FragmentedArray.load_strided(x_smem, is_signed=is_signed)
        case _:
          assert isinstance(ctx.out_layout_hint, mgpu.TiledLayout)
          return mgpu.FragmentedArray.load_untiled(
              x_smem,
              is_signed=is_signed,
              layout=ctx.out_layout_hint,
              swizzle=swizzle or 16,
              optimized=optimized,
          )
    case _:
      raise NotImplementedError(f"Unsupported transforms: {transforms}")


@register_lowering_rule(sp.get_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(sp.get_p, *gpu_core.WGxWARP_SEMANTICS)
def _get_lowering_rule_wg(
    ctx: LoweringRuleContext, x_ref, *leaves, tree, optimized=True
):
  if not isinstance(x_ref, ir.Value) and isinstance(x_ref, ir.MemRefType):
    raise TypeError(f"Can only load from references (got {x_ref}).")
  shape = ctx.avals_out[0].shape
  if shape and ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    raise ValueError("Can only load scalars in warp-level code.")

  transforms = jax.tree.unflatten(tree, leaves)
  assert isinstance(ctx.avals_in[0], state_types.AbstractRef)
  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[1:])
  x_ref, _, transforms = _handle_transforms(
      ctx, ctx.avals_in[0], x_ref, transform_avals, transforms,
      allow_peer_refs=True
  )

  if transforms:
    raise NotImplementedError(
        "Transforms are not yet implemented for warpgroup semantics"
    )

  assert isinstance(x_ref, ir.Value)
  shape = ctx.avals_out[0].shape
  if shape:
    return mgpu.dialect.vector_load(x_ref, optimized=optimized)
  else:
    return memref_dialect.load(x_ref, [])


@register_lowering_rule(sp.swap_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(sp.swap_p, *gpu_core.LANExWARP_SEMANTICS)
def _swap_lowering_rule(
    ctx: LoweringRuleContext,
    x_ref,
    value,
    *leaves,
    tree,
    optimized: bool = True,
    return_old: bool = True,
):
  if isinstance(x_ref, tcgen05.TMEMRef):
    raise RuntimeError(
        "Stores to TMEM are asynchronous operations and cannot be performed"
        " using the usual syntax. Please use plgpu.async_store_tmem instead."
    )
  v_aval = ctx.avals_in[1]
  barrier = mgpu.warpgroup_barrier
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if v_aval.shape:
      raise NotImplementedError("Can only store scalars in warp-level lowering.")
    barrier = mgpu_utils.warp_barrier
  value = _ensure_fa(value, v_aval.dtype)

  if not isinstance(x_ref, ir.Value) and isinstance(x_ref, ir.MemRefType):
    raise TypeError(f"Can only store to references (got {x_ref}).")
  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[2:])

  if ctx.module_ctx.auto_barriers:
    barrier()  # Make sure reads have completed before we write.

  if transforms and isinstance(transforms[0], gpu_core.UnswizzleRef):
    swizzle = transforms[0].swizzle
    transforms = transforms[1:]
    transform_avals = transform_avals[1:]
  elif (
      len(transforms) > 1
      and isinstance(transforms[0], gpu_core.ExtractAliasedRef)
      and isinstance(transforms[1], gpu_core.UnswizzleRef)
  ):
    swizzle = transforms[1].swizzle
    transforms = [transforms[0], *transforms[2:]]
    transform_avals = [transform_avals[0], *transform_avals[2:]]
  else:
    swizzle = None

  # We verify tiling against swizzling in memory before transposes flip it.
  tiling_transform = None
  if transforms and isinstance(transforms[0], gpu_core.UntilingTransform):
    tiling_transform = transforms[0]
  elif (
      len(transforms) > 1
      and isinstance(transforms[0], gpu_core.ExtractAliasedRef)
      and isinstance(transforms[1], gpu_core.UntilingTransform)
  ):
    tiling_transform = transforms[1]
  if tiling_transform is not None:
    if swizzle is None:
      raise NotImplementedError("Tiling without swizzle is not supported.")
    bw = dtypes.itemsize_bits(v_aval.dtype)
    expected_minor_tiling = swizzle * 8 // bw
    if tiling_transform.tiling[-1] != expected_minor_tiling:
      raise NotImplementedError(
          "Minor tiling dimension does not fit swizzle: "
          f" expected {expected_minor_tiling}, got {tiling_transform.tiling[-1]}"
      )

  assert isinstance(ctx.avals_in[0], state_types.AbstractRef)
  x_smem, _, transforms = _handle_transforms(
      ctx, ctx.avals_in[0], x_ref, transform_avals, transforms,
      allow_peer_refs=True
  )
  del x_ref  # Don't use x_ref anymore. Use x_smem instead!

  old_value = ()

  match transforms:
    case _ if math.prod(v_aval.shape) == 1:  # Scalar case.
      zero_idx = _ir_constant(0, ir.IndexType.get())
      indices = [zero_idx] * len(v_aval.shape)
      if return_old:
        old_value = mgpu.FragmentedArray.splat(
            memref_dialect.load(x_smem, indices),
            shape=(),
            is_signed=mgpu_utils.is_signed(v_aval.dtype),
        )
      value.store_untiled(x_smem)
    case (gpu_core.UntilingTransform(tiling),):
      if len(tiling) != 2:
        raise NotImplementedError(f"Only 2D tiling is supported, got: {tiling}")
      if return_old:
        old_value = mgpu.FragmentedArray.load_tiled(
            x_smem,
            is_signed=mgpu_utils.is_signed(v_aval.dtype),
            swizzle=swizzle,
            layout=value.layout,
            optimized=optimized,
            tiling_rank=len(tiling),
        )
      value.store_tiled(
          x_smem,
          swizzle=swizzle,
          optimized=optimized,
          tiling_rank=len(tiling),
      )
    case ():
      match value.layout:
        case mgpu.TiledLayout():
          if return_old:
            old_value = mgpu.FragmentedArray.load_untiled(
                x_smem,
                layout=value.layout,
                is_signed=mgpu_utils.is_signed(v_aval.dtype),
                swizzle=swizzle or 16,
                optimized=optimized,
            )
          value.store_untiled(
              x_smem, swizzle=swizzle or 16, optimized=optimized
          )
        case _:
          if swizzle is not None:
            raise NotImplementedError(
                "Unsupported swizzle transform with strided layout"
            )
          if return_old:
            old_value = mgpu.FragmentedArray.load_strided(
                x_smem, is_signed=mgpu_utils.is_signed(v_aval.dtype)
            )
          value.store_untiled(x_smem)
    case _:
      raise NotImplementedError(f"Unsupported transforms: {transforms}")
  if ctx.module_ctx.auto_barriers:
    barrier()  # Make sure the writes have completed.
  return old_value


@register_lowering_rule(sp.swap_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(sp.swap_p, *gpu_core.WGxWARP_SEMANTICS)
def _swap_lowering_rule_wg(
    ctx: LoweringRuleContext,
    x_smem,
    value,
    *leaves,
    tree,
    optimized: bool = True,
    return_old: bool = True,
):
  v_aval = ctx.avals_in[1]
  shape = v_aval.shape
  if shape and not isinstance(value.type, ir.VectorType):
    raise TypeError(f"Can only store scalars or vectors (got {value}).")
  if not (
      isinstance(x_smem, ir.Value) and isinstance(x_smem.type, ir.MemRefType)
  ):
    raise TypeError(f"Can only store to references (got {x_smem}).")
  if shape and ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    raise NotImplementedError("Can only store scalars in warp-level lowering.")
  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[2:])
  assert isinstance(ctx.avals_in[0], state_types.AbstractRef)
  x_smem, _, transforms = _handle_transforms(
      ctx, ctx.avals_in[0], x_smem, transform_avals, transforms,
      allow_peer_refs=True
  )
  if transforms:
    raise NotImplementedError(
        "Transforms are not yet implemented for warpgroup semantics"
    )
  assert isinstance(x_smem, ir.Value)
  value = _ensure_ir_value(value, v_aval.dtype)
  old_value = ()
  if shape:
    if return_old:
      old_value = mgpu.dialect.vector_load(x_smem, optimized=optimized)
    mgpu.dialect.vector_store(value, x_smem, optimized=optimized)
  else:
    if return_old:
      old_value = memref_dialect.load(x_smem, [])
    memref_dialect.store(value, x_smem, [])
  return old_value


@register_lowering_rule(pjit.jit_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(pjit.jit_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(pjit.jit_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(pjit.jit_p, *gpu_core.WGxWARP_SEMANTICS)
def _pjit_lowering_rule(ctx: LoweringRuleContext, *args, jaxpr, **kwargs):
  if jaxpr.consts:
    raise NotImplementedError
  return lower_jaxpr_to_mosaic_gpu(
      ctx.module_ctx, ctx.launch_ctx, jaxpr, args,
  )


@register_lowering_rule(pe.eval_jaxpr_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(pe.eval_jaxpr_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(pe.eval_jaxpr_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(pe.eval_jaxpr_p, *gpu_core.WGxWARP_SEMANTICS)
@register_lowering_rule(
    custom_derivatives.custom_jvp_call_p, mgpu.LoweringSemantics.Lane
)
@register_lowering_rule(
    custom_derivatives.custom_jvp_call_p, mgpu.LoweringSemantics.Warpgroup
)
@register_lowering_rule(
    custom_derivatives.custom_jvp_call_p, *gpu_core.LANExWARP_SEMANTICS
)
@register_lowering_rule(
    custom_derivatives.custom_jvp_call_p, *gpu_core.WGxWARP_SEMANTICS
)
@register_lowering_rule(
    custom_derivatives.custom_vjp_call_p, mgpu.LoweringSemantics.Lane
)
@register_lowering_rule(
    custom_derivatives.custom_vjp_call_p, mgpu.LoweringSemantics.Warpgroup
)
@register_lowering_rule(
    custom_derivatives.custom_vjp_call_p, *gpu_core.LANExWARP_SEMANTICS
)
@register_lowering_rule(
    custom_derivatives.custom_vjp_call_p, *gpu_core.WGxWARP_SEMANTICS
)
def _eval_jaxpr_lowering_rule(
    ctx: LoweringRuleContext, *args, call_jaxpr: jax_core.Jaxpr, **kwargs
):
  del kwargs  # Unused.
  if call_jaxpr.consts:
    raise NotImplementedError
  return lower_jaxpr_to_mosaic_gpu(
      ctx.module_ctx, ctx.launch_ctx, call_jaxpr, args
  )


@register_lowering_rule(lax.slice_p, mgpu.LoweringSemantics.Lane)
def _slice_lowering_rule(
    ctx: LoweringRuleContext, x, limit_indices, start_indices, strides
):
  if strides is not None:
    raise NotImplementedError("Strides are not supported.")

  return x[tuple(slice(b, e) for b, e in zip(start_indices, limit_indices))]


@register_lowering_rule(lax.slice_p, mgpu.LoweringSemantics.Warpgroup)
def _slice_lowering_rule_wg(
    ctx: LoweringRuleContext, x, limit_indices, start_indices, strides
):
  del limit_indices
  assert isinstance(x.type, ir.VectorType)
  if strides is not None:
    raise NotImplementedError("Strides are not supported.")
  out_ty = ir.VectorType.get(
      ctx.avals_out[0].shape, ir.VectorType(x.type).element_type
  )
  sizes = ctx.avals_out[0].shape
  strides = [1] * len(start_indices)
  return vector_dialect.extract_strided_slice(
      out_ty, x, start_indices, sizes, strides
  )


@register_lowering_rule(lax.concatenate_p, mgpu.LoweringSemantics.Lane)
def _concatenate_lowering_rule(ctx: LoweringRuleContext, *args, dimension):
  arrays = [_ensure_fa(x, aval.dtype) for x, aval in zip(args, ctx.avals_in)]
  return mgpu.concatenate(arrays, axis=dimension)


@register_lowering_rule(lax.concatenate_p, mgpu.LoweringSemantics.Warpgroup)
def _concatenate_lowering_rule_wg(ctx: LoweringRuleContext, *args, dimension):
  operands = [_ensure_ir_value(x, a.dtype) for x, a in zip(args, ctx.avals_in)]
  return mgpu.dialect.vector_concat(operands, dimension)


@register_lowering_rule(lax.select_n_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.select_n_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(lax.select_n_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(lax.select_n_p, *gpu_core.WGxWARP_SEMANTICS)
def _select_n_lowering_rule(ctx: LoweringRuleContext, pred, *cases):
  if len(cases) != 2:
    raise NotImplementedError(
        "Mosaic GPU lowering only supports select_n with 2 cases, got"
        f" {len(cases)}"
    )
  pred_aval, *cases_avals = ctx.avals_in
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if not all(aval.shape == () for aval in ctx.avals_in):
      raise NotImplementedError(
          "Can only select on scalars in warp-level lowering.")
  [out_aval] = ctx.avals_out
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    pred = _ensure_fa(pred, pred_aval.dtype)
    cases = _bcast(*cases, *cases_avals, out_aval=out_aval)
    # ``select`` expects the first case to be the true branch, but ``select_n``
    # orders the cases in reverse.
    return pred.select(*reversed(cases))
  else:
    pred = _ensure_ir_value(pred, pred_aval.dtype)
    cases = [_ensure_ir_value(c, c_aval.dtype) for c, c_aval in zip(cases, cases_avals)]
    # TODO(bchetioui): support implicit broadcast.
    if any(a.shape != out_aval.shape for a in ctx.avals_in):
      raise NotImplementedError(
          "Implicit broadcast not implemented with warpgroup semantics")
    # ``select`` expects the first case to be the true branch, but ``select_n``
    # orders the cases in reverse.
    return arith_dialect.select(pred, *reversed(cases))


@register_lowering_rule(lax.broadcast_in_dim_p, mgpu.LoweringSemantics.Lane)
def _broadcast_in_dim_lowering_rule(
    ctx: LoweringRuleContext,
    x: mgpu.FragmentedArray,
    *,
    broadcast_dimensions,
    shape,
    sharding,
):
  del sharding
  [x_aval] = ctx.avals_in
  [y_aval] = ctx.avals_out
  x = _ensure_fa(x, x_aval.dtype)
  rank_diff = y_aval.ndim - x_aval.ndim
  if (isinstance(x.layout, mgpu.WGSplatFragLayout) and
      broadcast_dimensions == tuple(range(rank_diff, rank_diff + x_aval.ndim))):
    return x.broadcast(shape)
  new_layout = None
  if (
      isinstance(x.layout, mgpu.WGStridedFragLayout)
      and broadcast_dimensions == tuple(range(rank_diff, y_aval.ndim))
  ):
    new_layout = mgpu.WGStridedFragLayout(
        shape=y_aval.shape, vec_size=x.layout.vec_size
    )
    return x.broadcast_in_dim(y_aval.shape, broadcast_dimensions, new_layout)
  if not isinstance(layout := x.layout, mgpu.TiledLayout):
    raise NotImplementedError(f"Unsupported layout: {x.layout}")
  if any(d1 >= d2 for d1, d2 in zip(broadcast_dimensions[:-1], broadcast_dimensions[1:])):
    raise NotImplementedError("broadcast_dimensions must be strictly increasing")
  new_dims = [d for d in range(y_aval.ndim) if d not in broadcast_dimensions]
  if (new_layout := ctx.out_layout_hint) is None:
    candidates = [
      mgpu.WGMMA_LAYOUT,
      mgpu.WGMMA_TRANSPOSED_LAYOUT,
      mgpu.TCGEN05_LAYOUT,
      mgpu.TCGEN05_TRANSPOSED_LAYOUT,
      tcgen05.TMEM_NATIVE_LAYOUT,
    ]
    if y_aval.shape[-1] % 16 == 0:
      candidates.append(tcgen05.fa_m64_collective_layout(y_aval.shape[-1]))
    for candidate in candidates:
      if len(candidate.base_tile_shape) != len(shape):
        continue
      # check if base_tile_shape is compatible with shape
      if any(s % t != 0 for s, t in zip(shape, candidate.base_tile_shape)):
        continue
      if candidate.reduce(new_dims) == layout:
        if new_layout is None:
          new_layout = candidate
        elif candidate == mgpu.TCGEN05_LAYOUT and new_layout == mgpu.WGMMA_LAYOUT:
          continue  # Choosing WGMMA_LAYOUT for backwards compatibility.
        else:
          raise NotImplementedError(
              "Multiple options for the layout of the broadcast result (found"
              f" at least {new_layout} and {candidate}). Use plgpu.layout_cast"
              " on the output to suggest the desired output layout."
          )
  if new_layout is None:
    raise NotImplementedError(
        "No compatible layout found for the broadcast result. Use"
        " plgpu.layout_cast on the output to suggest the desired output layout."
    )
  return x.broadcast_in_dim(y_aval.shape, broadcast_dimensions, new_layout)


@register_lowering_rule(
    lax.broadcast_in_dim_p, mgpu.LoweringSemantics.Warpgroup)
def _broadcast_in_dim_lowering_rule_wg(
    ctx: LoweringRuleContext,
    x,
    *,
    broadcast_dimensions,
    shape,
    sharding,
):
  del sharding
  [x_aval] = ctx.avals_in
  mlir_type = mgpu_utils.dtype_to_ir_type(x_aval.dtype)
  result_ty = ir.VectorType.get(shape, mlir_type)
  if not broadcast_dimensions:
    # Even though we could implement this case by passing a 0D vector as input
    # to mgpu.dialect.BroadcastInDimOp we don't want that. 0D vectors are
    # generally problematic and so we avoid them by specializing that case
    # directly here.
    x = _ensure_ir_value(x, x_aval.dtype)
    return vector_dialect.broadcast(result_ty, x)
  return mgpu.dialect.broadcast_in_dim(result_ty, x, broadcast_dimensions)


@register_lowering_rule(lax.convert_element_type_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.convert_element_type_p, *gpu_core.LANExWARP_SEMANTICS)
def _convert_element_type_lowering_rule(
    ctx: LoweringRuleContext, x, *, new_dtype, weak_type, sharding
):
  del weak_type, sharding
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if x_aval.shape != ():
      raise NotImplementedError(
          "Non-scalar arithmetic is not supported in warp-level lowering.")
  return _ensure_fa(x, x_aval.dtype).astype(
      mgpu_utils.dtype_to_ir_type(new_dtype), is_signed=mgpu_utils.is_signed(new_dtype)
  )


@register_lowering_rule(
    lax.convert_element_type_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(
    lax.convert_element_type_p, *gpu_core.WGxWARP_SEMANTICS
)
def _convert_element_type_lowering_rule_wg(
    ctx: LoweringRuleContext, x, *, new_dtype, weak_type, sharding
):
  del weak_type, sharding
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if x_aval.shape:
      raise NotImplementedError(
          "Non-scalar arithmetic is not supported in warp-level lowering.")
  [y_aval] = ctx.avals_out
  x = _ensure_ir_value(x, x_aval.dtype)

  cur_dtype = mgpu_utils.dtype_to_ir_type(x_aval.dtype)
  new_dtype = mgpu_utils.dtype_to_ir_type(new_dtype)

  if cur_dtype == new_dtype:
    return x

  from_float = isinstance(cur_dtype, ir.FloatType)
  to_float = isinstance(new_dtype, ir.FloatType)
  from_integer = isinstance(cur_dtype, ir.IntegerType)
  to_integer = isinstance(new_dtype, ir.IntegerType)
  if from_float and to_float:
    cur_ty_width = ir.FloatType(cur_dtype).width
    new_ty_width = ir.FloatType(new_dtype).width
    if cur_ty_width == new_ty_width:
      # There is no instruction to perform conversions between two float types
      # of the same width. Go through the next-larger standard type.
      # TODO(bchetioui): support conversions between float types of width 8.
      # Which larger type to pick will depend on the number of bits in the
      # smallest exponent.
      if cur_ty_width != 16:
        raise NotImplementedError(
            "Conversion between float types of width other than 16 not"
            " supported"
        )
      larger_ty = ir.F32Type.get()
      if x_aval.shape:
        upcast_ty = ir.VectorType.get(x_aval.shape, larger_ty)
      else:
        upcast_ty = larger_ty

      def convert(ty, x):
        return arith_dialect.truncf(ty, arith_dialect.extf(upcast_ty, x))

    elif ir.FloatType(cur_dtype).width > ir.FloatType(new_dtype).width:
      convert = arith_dialect.truncf
    else:
      convert = arith_dialect.extf
  elif from_integer and to_integer:
    if ir.IntegerType(cur_dtype).width > ir.IntegerType(new_dtype).width:
      convert = arith_dialect.trunci
    elif ir.IntegerType(cur_dtype).width < ir.IntegerType(new_dtype).width:
      if mgpu_utils.is_signed(x_aval.dtype):
        convert = arith_dialect.extsi
      else:
        convert = arith_dialect.extui
    else:
      convert = lambda _, x: x  # signed <-> unsigned conversions
  elif from_integer and to_float:
    if mgpu_utils.is_signed(x_aval.dtype):
      convert = arith_dialect.sitofp
    else:
      convert = arith_dialect.uitofp
  elif from_float and to_integer:
    dst_width = mgpu_utils.bitwidth(new_dtype)
    # We clamp the float value to the min/max integer destination value
    # in order to match JAX/XLA casting behavior. Note that this differs
    # from numpy casting behavior.
    if mgpu_utils.is_signed(y_aval.dtype):
      maxint = 2 ** (dst_width - 1) - 1
      minint = -(2 ** (dst_width - 1))
      convert = arith_dialect.fptosi
    else:
      maxint = 2**dst_width - 1
      minint = 0
      convert = arith_dialect.fptoui

    maxint = _ir_constant(maxint, cur_dtype)
    minint = _ir_constant(minint, cur_dtype)
    if x_aval.shape:
      maxint = vector_dialect.broadcast(x.type, maxint)
      minint = vector_dialect.broadcast(x.type, minint)
    x = arith_dialect.minimumf(x, maxint)
    x = arith_dialect.maximumf(x, minint)
  else:
    raise NotImplementedError(f"Unsupported conversion {cur_dtype} -> {new_dtype}")

  ty = ir.VectorType.get(x_aval.shape, new_dtype) if x_aval.shape else new_dtype
  return convert(ty, x)


mosaic_lowering_rules[gpu_core.LANExWG_SEMANTICS].update({
    lax.neg_p: lambda ctx, x: -x,
    lax.not_p: lambda ctx, x: ~x,
})

def _unary_warp_lowering_rule(impl):
  def _lowering_rule(ctx: LoweringRuleContext, x):
    if not all(aval_in.shape == () for aval_in in ctx.avals_in):
      raise NotImplementedError(
          "Non-scalar arithmetic is not supported in warp-level lowering.")
    return impl(x)
  return _lowering_rule

mosaic_lowering_rules[gpu_core.LANExWARP_SEMANTICS].update({
    lax.neg_p: _unary_warp_lowering_rule(lambda x: -x),
    lax.not_p: _unary_warp_lowering_rule(lambda x: ~x)
})


def _neg_lowering_rule_wg(ctx: LoweringRuleContext, x):
  [x_aval] = ctx.avals_in
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    return arith_dialect.negf(_ensure_ir_value(x, x_aval.dtype))
  return _lower_fun(lambda x: jnp.subtract(0, x))(ctx, x)


mosaic_lowering_rules[gpu_core.WGxWG_SEMANTICS].update({
    lax.neg_p: _neg_lowering_rule_wg,
    lax.not_p: _lower_fun(
        lambda x: jnp.astype(
            jnp.bitwise_xor(jnp.astype(x, int), -1), jnp.dtype(x)
        )
    ),
})


def _binary_op_lowering_rule(ctx: LoweringRuleContext, x, y, *, impl, out_dtype=None):
  del out_dtype  # unused here.
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if not all(aval_in.shape == () for aval_in in ctx.avals_in):
      raise NotImplementedError(
          "Non-scalar arithmetic is not supported in warp-level lowering.")
  x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)
  return impl(x, y)


def _div(x, y):
  return x / y if isinstance(x.mlir_dtype, ir.FloatType) else x.trunc_div(y)


for semantics in [gpu_core.LANExWG_SEMANTICS, gpu_core.LANExWARP_SEMANTICS]:
  mosaic_lowering_rules[semantics].update({
    lax.add_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x + y),
    lax.sub_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x - y),
    lax.mul_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x * y),
    lax.div_p: partial(_binary_op_lowering_rule, impl=_div),
    lax.rem_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x % y),
    lax.and_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x & y),
    lax.or_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x | y),
    lax.xor_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x ^ y),
    lax.gt_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x > y),
    lax.lt_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x < y),
    lax.ge_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x >= y),
    lax.le_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x <= y),
    lax.eq_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x == y),
    lax.ne_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x != y),
    lax.max_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x.max(y)),
    lax.min_p: partial(_binary_op_lowering_rule, impl=lambda x, y: x.min(y)),
    lax.shift_left_p: partial(
        _binary_op_lowering_rule,
        impl=lambda x, y: x._pointwise(arith_dialect.shli, y),
    ),
    lax.shift_right_logical_p: partial(
        _binary_op_lowering_rule,
        impl=lambda x, y: x._pointwise(arith_dialect.shrui, y),
    ),
    lax.shift_right_arithmetic_p: partial(
        _binary_op_lowering_rule,
        impl=lambda x, y: x._pointwise(arith_dialect.shrsi, y),
    ),
  })

def _binary_op_lowering_rule_wg(
    ctx: LoweringRuleContext, x, y, *, ui_impl, si_impl, f_impl=None, **kwargs,
):
  if kwargs.get('out_dtype') is not None:
    raise NotImplementedError("out_dtype argument in binary_op_lowering_rule_wg")
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if any(aval_in.shape for aval_in in ctx.avals_in):
      raise NotImplementedError(
          "Non-scalar arithmetic is not supported in warp-level lowering.")
  x_aval, y_aval = ctx.avals_in
  [out_aval] = ctx.avals_out
  x, y = _bcast_wg(x, y, *ctx.avals_in, *ctx.avals_out)
  if jnp.issubdtype(out_aval, jnp.signedinteger):
    return si_impl(x, y)
  elif jnp.issubdtype(out_aval, jnp.integer):
    return ui_impl(x, y)
  elif f_impl is not None and jnp.issubdtype(out_aval, jnp.floating):
    return f_impl(x, y)
  else:
    raise NotImplementedError(
        f"{ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype}"
    )


for op, si_impl, ui_impl, f_impl in [
    (lax.add_p, arith_dialect.addi, arith_dialect.addi, arith_dialect.addf),
    (lax.sub_p, arith_dialect.subi, arith_dialect.subi, arith_dialect.subf),
    (lax.mul_p, arith_dialect.muli, arith_dialect.muli, arith_dialect.mulf),
    (
        lax.div_p,
        arith_dialect.divsi,
        arith_dialect.divui,
        arith_dialect.divf,
    ),
    (lax.rem_p, arith_dialect.remsi, arith_dialect.remui, arith_dialect.remf),
    (
        lax.max_p,
        arith_dialect.maxsi,
        arith_dialect.maxui,
        arith_dialect.maximumf,
    ),
    (
        lax.min_p,
        arith_dialect.minsi,
        arith_dialect.minui,
        arith_dialect.minimumf,
    ),
    (lax.shift_left_p, arith_dialect.shli, arith_dialect.shli, None),
    (lax.shift_right_logical_p, arith_dialect.shrui, arith_dialect.shrui, None),
    (lax.shift_right_arithmetic_p, arith_dialect.shrsi, arith_dialect.shrsi, None),
]:
  rule = partial(
      _binary_op_lowering_rule_wg,
      si_impl=si_impl,
      ui_impl=ui_impl,
      f_impl=f_impl,
  )
  mosaic_lowering_rules[gpu_core.WGxWG_SEMANTICS][op] = rule
  mosaic_lowering_rules[gpu_core.WGxWARP_SEMANTICS][op] = rule


def _binary_boolean_op_lowering_rule_wg(
    ctx: LoweringRuleContext, x, y, *, impl
):
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if any(aval_in.shape for aval_in in ctx.avals_in):
      raise NotImplementedError(
          "Non-scalar arithmetic is not supported in warp-level lowering.")
  x, y = _bcast_wg(x, y, *ctx.avals_in, *ctx.avals_out)
  return impl(x, y)

for op, impl in [
    (lax.and_p, arith_dialect.andi),
    (lax.or_p, arith_dialect.ori),
    (lax.xor_p, arith_dialect.xori),
]:
  rule = partial(
      _binary_boolean_op_lowering_rule_wg,
      impl=impl,
  )
  mosaic_lowering_rules[gpu_core.WGxWG_SEMANTICS][op] = rule
  mosaic_lowering_rules[gpu_core.WGxWARP_SEMANTICS][op] = rule

CmpIPred = arith_dialect.CmpIPredicate
CmpFPred = arith_dialect.CmpFPredicate

def _comparison_lowering_rule_wg(
    ctx: LoweringRuleContext, x, y, *, si_pred, ui_pred, f_pred
):
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if any(aval_in.shape for aval_in in ctx.avals_in):
      raise NotImplementedError(
          "Non-scalar arithmetic is not supported in warp-level lowering.")
  x_aval, y_aval = ctx.avals_in
  x, y = _bcast_wg(x, y, *ctx.avals_in, *ctx.avals_out)
  if jnp.issubdtype(x_aval, jnp.signedinteger):
    return arith_dialect.cmpi(si_pred, x, y)
  elif jnp.issubdtype(x_aval, jnp.unsignedinteger) or jnp.issubdtype(x_aval, jnp.bool):
    return arith_dialect.cmpi(ui_pred, x, y)
  elif jnp.issubdtype(x_aval, jnp.floating):
    return arith_dialect.cmpf(f_pred, x, y)
  else:
    raise NotImplementedError(
        f"{ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype}"
    )


for op, si_pred, ui_pred, f_pred in [
    (lax.eq_p, CmpIPred.eq, CmpIPred.eq, CmpFPred.OEQ),
    (lax.ne_p, CmpIPred.ne, CmpIPred.ne, CmpFPred.UNE),
    (lax.lt_p, CmpIPred.slt, CmpIPred.ult, CmpFPred.OLT),
    (lax.le_p, CmpIPred.sle, CmpIPred.ule, CmpFPred.OLE),
    (lax.gt_p, CmpIPred.sgt, CmpIPred.ugt, CmpFPred.OGT),
    (lax.ge_p, CmpIPred.sge, CmpIPred.uge, CmpFPred.OGE),
]:
  rule = partial(
      _comparison_lowering_rule_wg,
      si_pred=si_pred,
      ui_pred=ui_pred,
      f_pred=f_pred,
  )
  mosaic_lowering_rules[gpu_core.WGxWG_SEMANTICS][op] = rule
  mosaic_lowering_rules[gpu_core.WGxWARP_SEMANTICS][op] = rule

@register_lowering_rule(lax.integer_pow_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.integer_pow_p, mgpu.LoweringSemantics.Warpgroup)
def _integer_pow_lowering_rule(ctx: LoweringRuleContext, x, y):
  [x_aval] = ctx.avals_in
  if y == -1:
    return _lower_fun(lambda x: 1 / x)(ctx, x)
  if y <= 1:
    raise NotImplementedError

  mul_op: Callable[[Any, Any], Any]
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    mul_op = operator.mul
  elif jnp.issubdtype(x_aval.dtype, jnp.integer):
    mul_op = arith_dialect.muli
  elif jnp.issubdtype(x_aval.dtype, jnp.floating):
    mul_op = arith_dialect.mulf
  else:
    raise NotImplementedError(f"Unsupported dtype {x_aval.dtype}")

  # Y is an integer. Here we start with res = x so the range is y-1
  res = x
  # Repeated doubling algorithm.
  for i in reversed(range(y.bit_length() - 1)):
    res = mul_op(res, res)  # pyrefly: ignore[no-matching-overload]
    if (y >> i) & 1:
      res = mul_op(res, x)
  return res


@register_lowering_rule(lax.clamp_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.clamp_p, mgpu.LoweringSemantics.Warpgroup)
def _clamp_lowering_rule(ctx: LoweringRuleContext, l, x, u):
  return _lower_fun(lambda l, x, u: lax.min(lax.max(x, l), u))(ctx, l, x, u)


@register_lowering_rule(lax.square_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.square_p, mgpu.LoweringSemantics.Warpgroup)
def _square_lowering_rule(ctx: LoweringRuleContext, x):
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    x = _ensure_fa(x, x_aval.dtype)
    return x * x
  if jnp.issubdtype(x_aval.dtype, jnp.integer):
    return arith_dialect.muli(x, x)
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    return arith_dialect.mulf(x, x)
  raise NotImplementedError(f"Unsupported dtype {x_aval.dtype}")


@register_lowering_rule(lax.one_minus_square_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.one_minus_square_p, mgpu.LoweringSemantics.Warpgroup)
def _one_minus_square_lowering_rule(ctx: LoweringRuleContext, x):
  return _lower_fun(lambda x: (1 + x) * (1 - x))(ctx, x)


@register_lowering_rule(lax.clz_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.clz_p, mgpu.LoweringSemantics.Warpgroup)
def _clz_lowering_rule(ctx: LoweringRuleContext, x):
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype)._pointwise(math_dialect.ctlz, restrict_bitwidth=False)
  x = _ensure_ir_value(x, x_aval.dtype)
  return math_dialect.ctlz(x)


@register_lowering_rule(lax.rsqrt_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.rsqrt_p, mgpu.LoweringSemantics.Warpgroup)
def _rsqrt_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype).rsqrt(approx=ctx.module_ctx.approx_math)
  fastmath = (
      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
  )
  return math_dialect.rsqrt(
      _ensure_ir_value(x, x_aval.dtype), fastmath=fastmath
  )


@register_lowering_rule(lax.sqrt_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.sqrt_p, mgpu.LoweringSemantics.Warpgroup)
def _sqrt_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype).sqrt(approx=ctx.module_ctx.approx_math)
  fastmath = (
      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
  )
  return math_dialect.sqrt(
      _ensure_ir_value(x, x_aval.dtype), fastmath=fastmath
  )


@register_lowering_rule(lax.tanh_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.tanh_p, mgpu.LoweringSemantics.Warpgroup)
def _tanh_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype).tanh(approx=ctx.module_ctx.approx_math)
  fastmath = (
      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
  )
  return math_dialect.tanh(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)


def _logistic(x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  return 1.0 / (1 + lax.exp(-x))


mosaic_lowering_rules[gpu_core.LANExWG_SEMANTICS][lax.logistic_p] = _lower_fun(
    _logistic
)
mosaic_lowering_rules[gpu_core.WGxWG_SEMANTICS][lax.logistic_p] = _lower_fun(
    _logistic
)


@register_lowering_rule(lax.exp_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.exp_p, mgpu.LoweringSemantics.Warpgroup)
def _exp_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype).exp(approx=ctx.module_ctx.approx_math)
  fastmath = (
      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
  )
  return math_dialect.exp(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)


@register_lowering_rule(lax.exp2_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.exp2_p, mgpu.LoweringSemantics.Warpgroup)
def _exp2_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype).exp2(approx=ctx.module_ctx.approx_math)
  fastmath = (
      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
  )
  return math_dialect.exp2(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)

@register_lowering_rule(lax.sin_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.sin_p, mgpu.LoweringSemantics.Warpgroup)
def _sin_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype).sin(approx=ctx.module_ctx.approx_math)
  fastmath = (
      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
  )
  return math_dialect.sin(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)

@register_lowering_rule(lax.cos_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.cos_p, mgpu.LoweringSemantics.Warpgroup)
def _cos_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype).cos(approx=ctx.module_ctx.approx_math)
  fastmath = (
      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
  )
  return math_dialect.cos(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)

@register_lowering_rule(lax.log_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.log_p, mgpu.LoweringSemantics.Warpgroup)
def _log_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype).log(approx=ctx.module_ctx.approx_math)
  fastmath = (
      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
  )
  return math_dialect.log(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)


@register_lowering_rule(lax.log2_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.log2_p, mgpu.LoweringSemantics.Warpgroup)
def _log2_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
  if accuracy is not None:
    raise NotImplementedError("Not implemented: accuracy")
  [x_aval] = ctx.avals_in
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(x, x_aval.dtype).log2(approx=ctx.module_ctx.approx_math)
  fastmath = (
      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
  )
  return math_dialect.log2(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)


@register_lowering_rule(lax.abs_p, mgpu.LoweringSemantics.Lane)
def _abs_lowering_rule(ctx: LoweringRuleContext, x):
  [x_aval] = ctx.avals_in
  return _ensure_fa(x, x_aval.dtype).abs()


@register_lowering_rule(lax.abs_p, mgpu.LoweringSemantics.Warpgroup)
def _abs_lowering_rule_wg(ctx: LoweringRuleContext, x):
  [x_aval] = ctx.avals_in
  x = _ensure_ir_value(x, x_aval.dtype)
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    return math_dialect.absf(x)
  if jnp.issubdtype(x_aval.dtype, jnp.integer):
    return math_dialect.absi(x)
  raise NotImplementedError(f"Unsupported dtype for abs: {x_aval.dtype}")


@register_lowering_rule(lax.round_p, mgpu.LoweringSemantics.Lane)
def _round_lowering_rule(ctx: LoweringRuleContext, x, rounding_method):
  [x_aval] = ctx.avals_in
  x = _ensure_fa(x, x_aval.dtype)
  match rounding_method:
    case lax.RoundingMethod.AWAY_FROM_ZERO:
      return x.round()
    case lax.RoundingMethod.TO_NEAREST_EVEN:
      return x.round_even()
    case _:
      assert_never(rounding_method)


@register_lowering_rule(lax.round_p, mgpu.LoweringSemantics.Warpgroup)
def _round_lowering_rule_wg(ctx: LoweringRuleContext, x, rounding_method):
  [x_aval] = ctx.avals_in
  x = _ensure_ir_value(x, x_aval.dtype)
  if not jnp.issubdtype(x_aval.dtype, jnp.floating):
    raise NotImplementedError(f"Unsupported dtype for round: {x_aval.dtype}")
  match rounding_method:
    case lax.RoundingMethod.AWAY_FROM_ZERO:
      return math_dialect.round(x)
    case lax.RoundingMethod.TO_NEAREST_EVEN:
      return math_dialect.roundeven(x)
    case _:
      assert_never(rounding_method)

_copysign_p = jax_core.Primitive("_copysign")


def _copysign(x1: jax.typing.ArrayLike, x2: jax.typing.ArrayLike) -> jax.Array:
  return _copysign_p.bind(x1, x2)


@_copysign_p.def_abstract_eval
def _copysign_abstract_eval(x1, x2):
  return jax_core.ShapedArray(x2.shape, x2.dtype)


@register_lowering_rule(_copysign_p, mgpu.LoweringSemantics.Lane)
def _copysign_lowering_rule(ctx: LoweringRuleContext, x1, x2):
  [x1_aval, x2_aval] = ctx.avals_in
  x1 = _ensure_fa(x1, x1_aval.dtype)
  x2 = _ensure_fa(x2, x2_aval.dtype)
  return x1.copysign(x2)


@register_lowering_rule(_copysign_p, mgpu.LoweringSemantics.Warpgroup)
def _copysign_lowering_rule(ctx: LoweringRuleContext, x1, x2):
  [x1_aval, x2_aval] = ctx.avals_in
  x1 = _ensure_ir_value(x1, x1_aval.dtype)
  x2 = _ensure_ir_value(x2, x2_aval.dtype)
  return math_dialect.copysign(x1, x2)


@register_lowering_rule(lax.sign_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.sign_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(lax.sign_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(lax.sign_p, *gpu_core.WGxWARP_SEMANTICS)
def _sign_lowering_rule(ctx: LoweringRuleContext, x):
  if (ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp and
      ctx.avals_in[0].shape):
    raise ValueError("Warp-level sign is not supported for non-scalar inputs.")
  def sign(x):
    if jnp.issubdtype(x.dtype, jnp.floating):
      ones = lax.full(x.shape, 1.0, dtype=x.dtype)
      zeros = lax.full(x.shape, 0.0, dtype=x.dtype)
      return lax.select(x != 0, _copysign(ones, x), zeros)
    if jnp.issubdtype(x.dtype, jnp.signedinteger):
      return (x > 0).astype(x.dtype) - (x < 0).astype(x.dtype)
    if jnp.issubdtype(x.dtype, jnp.unsignedinteger):
      return (x != 0).astype(x.dtype)
    raise ValueError(f"Unsupported dtype for sign: {x.dtype}")

  return _lower_fun(sign)(ctx, x)


@register_lowering_rule(lax.erf_p, mgpu.LoweringSemantics.Lane)
def _erf_lowering_rule(ctx: LoweringRuleContext, x):
  [x_aval] = ctx.avals_in
  return _ensure_fa(x, x_aval.dtype).erf()


@register_lowering_rule(lax.erf_p, mgpu.LoweringSemantics.Warpgroup)
def _erf_lowering_rule_wg(ctx: LoweringRuleContext, x):
  [x_aval] = ctx.avals_in
  return math_dialect.erf(_ensure_ir_value(x, x_aval.dtype))


@register_lowering_rule(lax.atan2_p, mgpu.LoweringSemantics.Lane)
def _atan2_lowering_rule(ctx: LoweringRuleContext, y, x):
  y, x = _bcast(y, x, *ctx.avals_in, *ctx.avals_out)
  return y.atan2(x)


@register_lowering_rule(lax.atan2_p, mgpu.LoweringSemantics.Warpgroup)
def _atan2_lowering_rule_wg(ctx: LoweringRuleContext, y, x):
  y, x = _bcast_wg(y, x, *ctx.avals_in, *ctx.avals_out)
  return math_dialect.atan2(y, x)


@register_lowering_rule(lax.reshape_p, mgpu.LoweringSemantics.Lane)
def _reshape_lowering_rule(
    ctx: LoweringRuleContext, x, new_sizes, dimensions, sharding
):
  if dimensions is not None:
    raise NotImplementedError("Not implemented: dimensions")
  if sharding is not None:
    raise NotImplementedError("Not implemented: sharding")
  [x_aval] = ctx.avals_in
  return _ensure_fa(x, x_aval.dtype).reshape(new_sizes)


@register_lowering_rule(lax.reshape_p, mgpu.LoweringSemantics.Warpgroup)
def _reshape_lowering_rule_wg(
    ctx: LoweringRuleContext, x, new_sizes, dimensions, sharding
):
  if dimensions is not None:
    raise NotImplementedError("Not implemented: dimensions")
  if sharding is not None:
    raise NotImplementedError("Not implemented: sharding")
  [x_aval] = ctx.avals_in
  x = _ensure_ir_value(x, x_aval.dtype)
  if x_aval.ndim == 0:  # scalar
    res_ty = ir.VectorType.get(new_sizes, x.type)
    return vector_dialect.broadcast(res_ty, x)
  else:
    res_ty = ir.VectorType.get(new_sizes, ir.VectorType(x.type).element_type)
    return vector_dialect.shape_cast(res_ty, x)


@register_lowering_rule(lax.squeeze_p, mgpu.LoweringSemantics.Lane)
def _squeeze_lowering_rule(ctx: LoweringRuleContext, x, dimensions):
  [x_aval] = ctx.avals_in
  [y_aval] = ctx.avals_out
  return _ensure_fa(x, x_aval.dtype).reshape(y_aval.shape)


@register_lowering_rule(lax.squeeze_p, mgpu.LoweringSemantics.Warpgroup)
def _squeeze_lowering_rule_wg(ctx: LoweringRuleContext, x, dimensions):
  [x_aval] = ctx.avals_in
  [y_aval] = ctx.avals_out
  x = _ensure_ir_value(x, x_aval.dtype)
  if y_aval.ndim == 0:  # scalar
    return vector_dialect.extract(
        x, dynamic_position=[], static_position=[0] * x_aval.ndim
    )
  else:
    res_ty = ir.VectorType.get(y_aval.shape, ir.VectorType(x.type).element_type)
    return vector_dialect.shape_cast(res_ty, x)


def _reduce_lowering_rule(
    op,
    ctx: LoweringRuleContext,
    x,
    *,
    axes,
    accumulator_ilp: int | None = None,
    **kwargs,
):
  del kwargs
  [x_aval] = ctx.avals_in
  match x.layout:
    case mgpu.WGStridedFragLayout():
      if set(axes) != set(range(x_aval.ndim)):
        raise NotImplementedError("No support for axes yet")
      # To relax the restriction below, you need to ensure sufficient
      # synchronization with other places that use `scratch_view` (which at the
      # time of writing is only `run_scoped`).
      if ctx.module_ctx.axis_names.wg is not None:
        raise NotImplementedError(
            "No support for reduce_sum over all axes and multiple Pallas"
            " threads"
        )
      scratch_ty = jax.ShapeDtypeStruct(shape=(4,), dtype=x_aval.dtype)
      with ctx.module_ctx.scratch_view(scratch_ty) as scratch:
        return x.reduce(op, axes, scratch, acc_ilp=accumulator_ilp)
    case mgpu.TiledLayout():
      if len(axes) != 1:
        raise NotImplementedError("Multi-axis reductions not supported")
      reduced_dim = x.layout.tiling.tile_dimension(axes[0])
      if any(reduced_dim[d] for d in x.layout.partitioned_warp_dims):
        dtype_bitwidth = dtypes.itemsize_bits(x_aval.dtype)
        if dtype_bitwidth % 8:
          raise NotImplementedError("Sub-byte dtypes not supported")
        scratch_elems = ctx.module_ctx.reduction_scratch_bytes * 8 // dtype_bitwidth
        scratch_ty = jax.ShapeDtypeStruct(shape=(scratch_elems,), dtype=x_aval.dtype)
        scratch_ctx = ctx.module_ctx.scratch_view(scratch_ty)
      else:
        scratch_ctx = contextlib.nullcontext(None)
      with scratch_ctx as scratch:
        return x.reduce(op, axes[0], scratch=scratch, acc_ilp=accumulator_ilp)
    case _:
      raise NotImplementedError(f"Unsupported layout {x.layout}")

register_lowering_rule(lax.reduce_sum_p, mgpu.LoweringSemantics.Lane)(
    functools.partial(_reduce_lowering_rule, "add")
)
register_lowering_rule(lax.reduce_max_p, mgpu.LoweringSemantics.Lane)(
    functools.partial(_reduce_lowering_rule, "max")
)
register_lowering_rule(lax.reduce_min_p, mgpu.LoweringSemantics.Lane)(
    functools.partial(_reduce_lowering_rule, "min")
)
register_lowering_rule(lax.reduce_prod_p, mgpu.LoweringSemantics.Lane)(
    functools.partial(_reduce_lowering_rule, "prod")
)


def _reduce_lowering_rule_wg(
    ctx: LoweringRuleContext,
    kind: vector_dialect.CombiningKind,
    acc: int | float,
    x,
    axes,
    accumulator_ilp: int | None = None,
) -> ir.Value:
  [x_aval] = ctx.avals_in
  [out_aval] = ctx.avals_out
  x = _ensure_ir_value(x, x_aval.dtype)
  out_type = mgpu_utils.dtype_to_ir_type(out_aval.dtype)
  if not out_aval.shape:
    # Special-case: reducing to a scalar.
    if x_aval.ndim != 1:
      # Flatten to 1D, since vector.reduction only supports 1D inputs.
      x = vector_dialect.shape_cast(
          ir.VectorType.get([x_aval.size], out_type), x
      )
    reduction = vector_dialect.ReductionOp(out_type, kind, x)
  else:
    acc_vec = vector_dialect.broadcast(
        ir.VectorType.get(out_aval.shape, out_type),
        _ensure_ir_value(acc, out_aval.dtype),
    )
    reduction = vector_dialect.MultiDimReductionOp(kind, x, acc_vec, axes)
  def i32_attr(value: int) -> ir.IntegerAttr:
    return ir.IntegerAttr.get(ir.IntegerType.get_signless(32), value)
  reduction.attributes["offset"] = i32_attr(ctx.module_ctx.smem_used_bytes)
  # TODO(bchetioui): here, we could just donate all the remaining free SMEM that
  # we have at this point in time.
  reduction.attributes["scratch_size"] = i32_attr(ctx.module_ctx.reduction_scratch_bytes)
  if accumulator_ilp is not None:
    reduction.attributes["acc_ilp"] = i32_attr(accumulator_ilp)
  return reduction.result


@register_lowering_rule(lax.reduce_sum_p, mgpu.LoweringSemantics.Warpgroup)
def _reduce_sum_lowering_rule_wg(
    ctx: LoweringRuleContext,
    x,
    *,
    axes,
    out_sharding=None,
    accumulator_ilp: int | None = None,
):
  del out_sharding
  kind = vector_dialect.CombiningKind.ADD
  return _reduce_lowering_rule_wg(
      ctx, kind, 0, x, axes, accumulator_ilp=accumulator_ilp
  )


@register_lowering_rule(lax.reduce_max_p, mgpu.LoweringSemantics.Warpgroup)
def _reduce_max_lowering_rule_wg(
    ctx: LoweringRuleContext,
    x,
    *,
    axes,
    out_sharding=None,
    accumulator_ilp: int | None = None,
):
  del out_sharding
  [x_aval] = ctx.avals_in
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    kind = vector_dialect.CombiningKind.MAXIMUMF
    acc = float("-inf")
  elif jnp.issubdtype(x_aval.dtype, jnp.integer):
    if jnp.issubdtype(x_aval.dtype, jnp.signedinteger):
      kind = vector_dialect.CombiningKind.MAXSI
    else:
      kind = vector_dialect.CombiningKind.MAXUI
    acc = np.iinfo(x_aval.dtype).min
  else:
    raise NotImplementedError(f"Unsupported dtype {x_aval.dtype}")
  return _reduce_lowering_rule_wg(
      ctx, kind, acc, x, axes, accumulator_ilp=accumulator_ilp
  )


@register_lowering_rule(lax.reduce_min_p, mgpu.LoweringSemantics.Warpgroup)
def _reduce_min_lowering_rule_wg(
    ctx: LoweringRuleContext,
    x,
    *,
    axes,
    out_sharding=None,
    accumulator_ilp: int | None = None,
):
  del out_sharding
  [x_aval] = ctx.avals_in
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    kind = vector_dialect.CombiningKind.MINIMUMF
    acc = float("inf")
  elif jnp.issubdtype(x_aval.dtype, jnp.integer):
    if jnp.issubdtype(x_aval.dtype, jnp.signedinteger):
      kind = vector_dialect.CombiningKind.MINSI
    else:
      kind = vector_dialect.CombiningKind.MINUI
    acc = np.iinfo(x_aval.dtype).max
  else:
    raise NotImplementedError(f"Unsupported dtype {x_aval.dtype}")
  return _reduce_lowering_rule_wg(
      ctx, kind, acc, x, axes, accumulator_ilp=accumulator_ilp
  )


@register_lowering_rule(lax.reduce_prod_p, mgpu.LoweringSemantics.Warpgroup)
def _reduce_prod_lowering_rule_wg(
    ctx: LoweringRuleContext,
    x,
    *,
    axes,
    accumulator_ilp: int | None = None,
):
  [x_aval] = ctx.avals_in
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    acc = 1.0
  elif jnp.issubdtype(x_aval.dtype, jnp.integer):
    acc = 1
  else:
    raise NotImplementedError(f"Unsupported dtype {x_aval.dtype}")
  kind = vector_dialect.CombiningKind.MUL
  return _reduce_lowering_rule_wg(
      ctx, kind, acc, x, axes, accumulator_ilp=accumulator_ilp
  )


def _block_id(ctx: LoweringRuleContext, dim: gpu_dialect.Dimension) -> ir.Value:
  result = gpu_dialect.block_id(dim)
  cluster_size = ctx.launch_ctx.cluster_size
  if math.prod(cluster_size) == 1 or cluster_size[dim.value] == 1:
    return result
  # We scale the grid in the presence of clusters, so we need to scale the
  # block ID back here.
  return arith_dialect.divui(result, _as_index(cluster_size[dim.value]))


def _resolve_cluster_axis(axis_names: _AxisNames | None, axis_name: Hashable):
  if not axis_names:
    raise LookupError(
        "No axis names are available. Make sure you are using `pl.core_map`"
        " with a `plgpu.Mesh`."
    )
  if not axis_names or axis_name not in axis_names.cluster:
    raise LookupError(
        f"Unknown cluster axis {axis_name}, available axes:"
        f" {[*axis_names.cluster]}"
    )
  return gpu_dialect.Dimension(axis_names.cluster.index(axis_name))


def _is_block_local_scope(collective_axes: CollectiveAxesType,
                          axis_names: _AxisNames):
  """Returns whether the collective axes represents a block scope."""
  if axis_names.wg is None:
    return not collective_axes
  else:
    return collective_axes == (axis_names.wg,)


def _is_global_scope(collective_axes: CollectiveAxesType,
                     axis_names: _AxisNames):
  """Returns whether the collective axes represents a GPU global scope."""
  return set(collective_axes) == set(axis_names)

def block_id_to_grid_id(ctx: LoweringRuleContext,
                        block_ids: Sequence[ir.Value],
                        axis_name: Hashable):
  squashed_dims = ctx.module_ctx.squashed_dims
  axis_names = ctx.module_ctx.axis_names
  if squashed_dims:
    unsquashed_names = axis_names.grid[:2]
    squashed_names = axis_names.grid[2:]
  else:
    # These are unused but initialized for type checkers.
    unsquashed_names = squashed_names = ()

  if squashed_dims:
    if axis_name in unsquashed_names:
      # We reversed the grid and cluster axes.
      # e.g. for the grid (a, b, c, d, wg)
      # squashed = (a, b)  Mapped to Dimension.z (2)
      # unsquashed = (c, d)  Mapped to Dimension.y (1) and Dimension.x (0)
      idx = unsquashed_names.index(axis_name)
      return block_ids[gpu_dialect.Dimension(idx)]
    else:
      assert axis_name in squashed_names
      # All squashed dimensions are mapped to Dimension.z.
      axis = squashed_names.index(axis_name)
      return _unravel_program_id(
          _as_index(block_ids[gpu_dialect.Dimension.z]), axis, squashed_dims
      )
  else:
    assert axis_name in axis_names.grid
    idx = axis_names.grid.index(axis_name)
    return block_ids[gpu_dialect.Dimension(idx)]


@register_lowering_rule(lax.axis_index_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.axis_index_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(lax.axis_index_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(lax.axis_index_p, *gpu_core.WGxWARP_SEMANTICS)
def _axis_index_rule(ctx: LoweringRuleContext, *, axis_name: Hashable):
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if axis_name == ctx.module_ctx.warp_axis_name:
      w_idx = mgpu.warp_idx(sync=True)
      i32 = ir.IntegerType.get_signless(32)
      return arith_dialect.remui(w_idx, _ir_constant(4, i32))
  gpu_axis_names = ctx.module_ctx.axis_names
  jax_axis_names = getattr(ctx.module_ctx.mesh_info, "axis_names", ())
  if gpu_axis_names is None and not jax_axis_names:
    raise LookupError(
        "No axis names are available. Make sure you are using `pl.core_map`"
        " with a `plgpu.Mesh` or an appropriate JAX device mesh."
    )
  if axis_name not in itertools.chain(gpu_axis_names or (), jax_axis_names):
    raise LookupError(
        f"Axis {axis_name} does not refer to a GPU mesh axis (available axes:"
        f" {[*gpu_axis_names]}) or a JAX mesh axis (available axes:"
        f" {[*jax_axis_names]})"
    )
  if axis_name in jax_axis_names:
    jax_mesh = ctx.module_ctx.mesh_info
    assert jax_mesh is not None
    device_id = ctx.launch_ctx.device_id()
    jax_mesh_shape = jax_mesh.mesh_shape
    axis_index = jax_axis_names.index(axis_name)
    i32 = ir.IntegerType.get_signless(32)
    axis_size = _ir_constant(jax_mesh_shape[axis_index], i32)
    minor_divisor = _ir_constant(
        np.prod(jax_mesh_shape[axis_index + 1 :], dtype=np.int32), i32
    )
    return arith_dialect.remsi(arith_dialect.divsi(device_id, minor_divisor), axis_size)

  # We already checked that the axis is in scope and it wasn't a JAX mesh axis.
  assert gpu_axis_names is not None

  # We only deal with GPU axes from now on.
  axis_names = gpu_axis_names
  if axis_names.wg is not None and axis_name == axis_names.wg:
    return mgpu.warpgroup_idx(sync=True)

  if axis_name in axis_names.cluster:
    return arith_dialect.index_cast(
        ir.IntegerType.get_signless(32),
        gpu_dialect.cluster_block_id(
            gpu_dialect.Dimension(axis_names.cluster.index(axis_name))
        ),
    )
  block_ids = tuple(arith_dialect.index_cast(
          ir.IntegerType.get_signless(32),
          _block_id(ctx, dimension),
  ) for dimension in gpu_dialect.Dimension)
  return block_id_to_grid_id(ctx, block_ids, axis_name)


@register_lowering_rule(debugging.debug_print_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(debugging.debug_print_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(
    debugging.debug_print_p, mgpu.LoweringSemantics.Warpgroup
)
@register_lowering_rule(debugging.debug_print_p, *gpu_core.WGxWARP_SEMANTICS)
def _debug_print_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    fmt,
    ordered,
    partitioned,
    in_tree,
    static_args,
    np_printoptions,
    has_placeholders,
    logging_record,
):
  del partitioned, np_printoptions, has_placeholders
  if ordered:
    raise NotImplementedError("Ordered debug_print is not supported on Pallas.")
  args, kwargs = debugging.merge_callback_args(in_tree, args, static_args)
  if kwargs:
    raise ValueError(
        "Only positional arguments are supported by debug_print on Pallas."
    )
  primitives.check_debug_print_format(fmt, *args)
  if not any(aval.shape for aval in ctx.avals_in):
    scope = mgpu.ThreadSubset.WARPGROUP
    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
      scope = mgpu.ThreadSubset.WARP
    mgpu.debug_print(
        fmt,
        *(
            _ensure_ir_value(arg, aval.dtype)
            for arg, aval in zip(args, ctx.avals_in)
        ),
        scope=scope
    )
  elif len(ctx.avals_in) == 1:
    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
      raise NotImplementedError("Can only print scalars in warp-level code.")
    [arg] = args
    if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
      mgpu.dialect.debug_print(fmt, arg)
    else:
      arg.debug_print(fmt)

  else:
    raise NotImplementedError(
        "debug_print only supports printing of scalar values, or a single array"
        " value when using the Mosaic GPU backend."
    )

  return ()


def _wgmma_fence_aligned():
  void = ir.Type.parse("!llvm.void")
  llvm_dialect.inline_asm(
      void, [], "wgmma.fence.sync.aligned;", "", has_side_effects=True
  )


@register_lowering_rule(primitives.run_scoped_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(primitives.run_scoped_p, mgpu.LoweringSemantics.Warpgroup)
def _run_scoped_lowering_rule(
    ctx: LoweringRuleContext,
    *consts,
    jaxpr: jax_core.Jaxpr,
    collective_axes,
    **_,
):
  if pallas_core.poison_buffers_enabled():
    raise NotImplementedError("Buffer poisoning is not supported on GPU yet.")
  input_refs = []
  should_discharge = []
  wg_axis = ctx.module_ctx.axis_names.wg
  is_multithreaded = wg_axis is not None
  is_thread_collective = is_multithreaded and collective_axes == (wg_axis,)
  # Make sure everyone has exited previous scoped allocations. Note that we
  # don't synchronize when we exit the allocation, but only when we might want
  # to reuse its memory again.
  if collective_axes and collective_axes != (wg_axis,):
    raise ValueError(
        "Only thread-collective allocations are supported in run_scoped."
    )
  if is_multithreaded and is_thread_collective:
    gpu_dialect.barrier()
  with contextlib.ExitStack() as alloc_stack:
    for v in jaxpr.invars:
      aval = cast(ShapedAbstractValue, v.aval)
      if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):
        if collective_axes:
          raise ValueError(
              "WGMMA accumulators can only be allocated non-collectively. Hint:"
              " remove collective_axes from run_scoped. If other allocations"
              " are performed as well, split the run_scoped into two."
          )
        is_signed = mgpu_utils.is_signed(aval.dtype)
        if is_signed is not None and not is_signed:
          raise ValueError(
              "Invalid WGMMA accumulator dtype for s8/i8 WGMMA. "
              f"Expected signed integer, but got {aval.dtype}."
          )

        dtype = mlir.dtype_to_ir_type(aval.dtype)
        if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
          input_refs.append(
              mgpu.WGMMAAccumulator.zero(
                  *aval.shape, dtype=dtype, is_signed=is_signed
              )
          )
        else:
          zero = _ir_constant(0, dtype)
          acc_type = ir.VectorType.get(aval.shape, dtype)
          acc = vector_dialect.broadcast(acc_type, zero)
          acc = mgpu.dialect.optimization_barrier([acc])
          _wgmma_fence_aligned()
          input_refs.append(acc)
        should_discharge.append(True)
        continue
      if (
          isinstance(aval, state_types.AbstractRef)
          and aval.memory_space == gpu_core.GMEM
          and jnp.issubdtype(aval.dtype, pallas_core.semaphore)
      ):
        input_ref = alloc_stack.enter_context(
            ctx.module_ctx.reserve_semaphores(
                aval.shape, collective_axes=collective_axes
            )
        )
        input_refs.append(input_ref)
        should_discharge.append(False)
        continue

      # All other allocations must be made collectively across all threads.
      if is_multithreaded and not is_thread_collective:
        raise NotImplementedError(
            "Only thread-collective allocations are supported in multithreaded"
            " kernels. Hint: add"
            f" collective_axes={ctx.module_ctx.axis_names.wg} to your"
            " run_scoped if you intend all threads to share the same"
            f" allocation (currently collective_axes={collective_axes})."
        )
      if isinstance(aval.dtype, gpu_core.BarrierType):
        barrier = _get_barrier(aval, ctx.estimator_ctx.arrival_multiplier)
        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)
        input_refs.append(alloc_stack.enter_context(barrier_ctx))
        should_discharge.append(False)
        continue
      if isinstance(aval.dtype, gpu_core.ClusterBarrierType):
        barrier = _get_cluster_barrier(aval, ctx.module_ctx.axis_names)
        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)
        input_refs.append(alloc_stack.enter_context(barrier_ctx))
        should_discharge.append(False)
        continue

      if not isinstance(aval, state_types.AbstractRef):
        raise ValueError(f"Can't convert to ref: {aval}")
      if aval.memory_space == gpu_core.SMEM:
        input_ref = alloc_stack.enter_context(
            ctx.module_ctx.scratch_view(
                jax.ShapeDtypeStruct(shape=aval.shape, dtype=aval.dtype)
            )
        )
        input_refs.append(input_ref)
        should_discharge.append(False)
      elif aval.memory_space == gpu_core.TMEM:
        input_ref = alloc_stack.enter_context(
            ctx.module_ctx.alloc_tmem(
                jax.ShapeDtypeStruct(shape=aval.shape, dtype=aval.dtype),
                layout=aval.layout,  # pyrefly: ignore[missing-attribute]
            )
        )
        input_refs.append(input_ref)
        should_discharge.append(False)

    if any(should_discharge):
      # We convert consts to args, because we only have ir.Values and
      # not JAX values during lowering. discharge_state() produces JAX
      # valiues for the arguments but expects them to be provided for the
      # consts. We also don't want to wrap the values in refs.
      no_const_jaxpr = pe.convert_constvars_jaxpr(jaxpr)
      should_discharge = [False] * len(consts) + should_discharge
      with config._check_vma(False):
        discharged_closed_jaxpr = discharge.discharge_state(
            no_const_jaxpr,
            should_discharge=should_discharge,
            strip_memory_space=True,
        )
        discharged_jaxpr, _ = discharged_closed_jaxpr, discharged_closed_jaxpr.consts
      new_input_vals = (*consts, *input_refs)
      outs = lower_jaxpr_to_mosaic_gpu(
          ctx.module_ctx,
          ctx.launch_ctx,
          discharged_jaxpr,
          new_input_vals,
          (),
      )
      # Discharge appends to the output the refs that got discharged.
      outs = outs[:-sum(should_discharge)]
    else:
      outs = lower_jaxpr_to_mosaic_gpu(
          ctx.module_ctx,
          ctx.launch_ctx,
          jaxpr,
          input_refs,
          consts,
      )

  assert len(outs) == len(jaxpr.outvars), (jaxpr, outs)
  return outs


@_register_resource_estimator(primitives.get_global_p)
def _get_global_resource_estimator(
    ctx: ResourceEstimatorContext, *, what
) -> Resources:
  if what.memory_space == gpu_core.GMEM and jnp.issubdtype(
      what.dtype, pallas_core.semaphore
  ):
    collective_axes = tuple(ctx.axis_names)
    return Resources(scoped_gmem_semaphores={collective_axes: what.size})
  raise NotImplementedError(f"get_global only supports semaphores, got {what}")


@register_lowering_rule(primitives.get_global_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(
    primitives.get_global_p, mgpu.LoweringSemantics.Warpgroup
)
def _get_global_lowering_rule(ctx: LoweringRuleContext, *, what):
  if what.memory_space == gpu_core.GMEM and jnp.issubdtype(
      what.dtype, pallas_core.semaphore
  ):
    collective_axes = tuple(ctx.module_ctx.axis_names)
    return ctx.module_ctx.reserve_semaphores(
        what.shape, collective_axes=collective_axes
    ).__enter__()
  raise NotImplementedError(f"get_global only supports semaphores, got {what}")


@register_lowering_rule(discharge.run_state_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(discharge.run_state_p, mgpu.LoweringSemantics.Warpgroup)
def _run_state_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    jaxpr: jax_core.Jaxpr,
    which_linear: tuple[bool, ...],
    is_initialized: tuple[bool, ...],
):
  del which_linear
  # TODO(apaszke): This should be unified with run_scoped.
  if not all(is_initialized):
    raise NotImplementedError("Uninitialized Refs are not supported in lowering of run_state.")

  should_discharge = []
  new_input_vals = []
  # `should_deref_acc` is used under lane lowering semantics, to figure out
  # whether we need to return a `WGMMAAccumulator` or a `FragmentedArray` when
  # encountering a `WGMMAAbstractAccumulatorRef` as input.
  #
  # We can't tell the difference under warpgroup lowering semantics, but we do
  # not need to since we always return a `vector` anyway.
  should_deref_acc = []
  for arg, v, out_aval in zip(args, jaxpr.invars, ctx.avals_out):
    aval = v.aval
    if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):
      if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
        arg = mgpu.dialect.optimization_barrier([arg])
        _wgmma_fence_aligned()
        new_input_vals.append(arg)
      else:
        if isinstance(arg, mgpu.WGMMAAccumulator):
          should_deref_acc.append(False)
          new_input_vals.append(arg)  # pyrefly: ignore[bad-argument-type]
        else:
          should_deref_acc.append(True)
          new_input_vals.append(mgpu.WGMMAAccumulator.from_registers(arg))
      should_discharge.append(True)
      assert isinstance(out_aval, jax_core.ShapedArray)
    else:
      new_input_vals.append(arg)
      should_discharge.append(not isinstance(out_aval, state_types.AbstractRef))
  if not any(should_discharge):
    raise NotImplementedError(
        "Expected at least one accumulator to in run_state."
    )

  with config._check_vma(False):
    discharged_closed_jaxpr = discharge.discharge_state(
        jaxpr, should_discharge=should_discharge,
        strip_memory_space=True
    )
    discharged_jaxpr, new_consts = discharged_closed_jaxpr, discharged_closed_jaxpr.consts
  assert not new_consts
  outs = lower_jaxpr_to_mosaic_gpu(
      ctx.module_ctx, ctx.launch_ctx, discharged_jaxpr, new_input_vals, ()  # pyrefly: ignore[bad-argument-type]
  )
  # Await the accumulators and extract their final values.
  #
  # Note: there may be accumulators that we have closed over and that should not
  # technically not be dereferenced. That's not a big deal: since `run_state` is
  # always used with at least one accumulator though, we need to await here
  # anyway.
  deref_acc = iter(should_deref_acc)
  void = ir.Type.parse("!llvm.void")
  llvm_dialect.inline_asm(
      void, [], "wgmma.wait_group.sync.aligned 0;", "", has_side_effects=True,
  )
  outs = [
      out.value if isinstance(out, mgpu.WGMMAAccumulator) and next(deref_acc) else out
      for out in outs
  ]
  # Blend the discharge results with refs we closed over. I don't fully
  # understand the reasons behind this calling convention, but sharadmv@ has
  # assured me that this is ok.
  outs_it = iter(outs)
  return [next(outs_it) if d else a for d, a in zip(should_discharge, args)]


def _lower_jaxpr_to_for_loop(
    ctx: LoweringRuleContext,
    jaxpr: jax_core.Jaxpr,
    start: ir.Value,
    length: int | ir.Value,
    consts,
    *args,
    has_loop_index: bool,
    unroll: int | None = None,
):
  _consts_avals, arg_avals = util.split_list(ctx.avals_in, [len(consts)])
  arg_avals = arg_avals[has_loop_index:]
  out_avals = []
  if arg_avals:
    out_avals = ctx.avals_out[-len(arg_avals):]

  is_acc = [isinstance(v, mgpu.WGMMAAccumulator) for v in args]
  def as_values(vals, avals):
    if is_acc != [isinstance(v, mgpu.WGMMAAccumulator) for v in vals]:
      raise ValueError("Unexpected loop carry w.r.t. accumulators.")

    _ensure = (
        _ensure_fa
        if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane
        else _ensure_ir_value
    )
    return [
        v if a else _ensure(v, av.dtype)
        for a, v, av in zip(is_acc, vals, avals)
    ]

  def loop(base_loop_index, body_args):
    outs = body_args
    if unroll is not None:
      base_loop_index = arith_dialect.muli(
          base_loop_index, _ir_constant(unroll, start.type)
      )
    base_loop_index = arith_dialect.addi(base_loop_index, start)
    for step in range(unroll or 1):
      if has_loop_index:
        loop_index = arith_dialect.addi(
            base_loop_index, _ir_constant(step, start.type)
        )
        jaxpr_args = [*consts, loop_index, *outs]
      else:
        jaxpr_args = [*consts, *outs]
      outs = lower_jaxpr_to_mosaic_gpu(
          ctx.module_ctx, ctx.launch_ctx, jaxpr, jaxpr_args
      )
    return as_values(outs, out_avals)

  if unroll is not None:
    if not isinstance(length, int):
      raise NotImplementedError(
          "``length`` must be an integer when ``unroll` is specified, got"
          f" {length}"
      )
    if length % unroll:
      # TODO(slebedev): Emit an epilogue taking care of the remaining steps.
      raise NotImplementedError(
          f"``unroll`` must divide ``length``, got {unroll=} and {length=}"
      )
    if unroll == length:
      # Special-case: the loop is fully unrolled.
      return loop(_ir_constant(0, start.type), as_values(args, arg_avals))
    return mgpu.fori(
        _ir_constant(length // unroll, start.type), as_values(args, arg_avals)
    )(loop).results
  else:
    if not isinstance(length, ir.Value):
      length = _ir_constant(length, start.type)
    return mgpu.fori(length, as_values(args, arg_avals))(loop).results


@register_lowering_rule(lax.scan_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.scan_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(lax.scan_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(lax.scan_p, *gpu_core.WGxWARP_SEMANTICS)
def _scan_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    jaxpr: jax_core.Jaxpr,
    length: int,
    reverse: bool,
    unroll: int,
    ft_in,
    ft_out,
):
  num_consts, num_carry, _ = (len(g) for g in ft_in.unpack())
  # Can only handle fori_loop-like scans.
  if reverse:
    raise NotImplementedError
  del reverse

  jaxpr, jaxpr_consts = jaxpr, jaxpr.consts
  if jaxpr_consts:
    raise NotImplementedError
  del jaxpr_consts

  body_jaxpr, has_loop_index = pallas_utils.pattern_match_scan_to_fori_loop(
      jaxpr, num_consts, num_carry
  )
  consts, args = util.split_list(args, [num_consts])
  _consts_avals, arg_avals = util.split_list(ctx.avals_in, [num_consts])
  if has_loop_index:
    start, *args = args
    index_aval, *_ = arg_avals
    start: ir.Value = _ensure_ir_value(start, index_aval.dtype)
  else:
    start = _i32_constant(0)

  for_out = _lower_jaxpr_to_for_loop(
      ctx,
      body_jaxpr,
      start,
      length,
      consts,
      *args,
      has_loop_index=has_loop_index,
      unroll=unroll,
  )
  if has_loop_index:
    # Need to return the final loop index value if the outer scan expects
    # it as an output.
    loop_index = arith_dialect.addi(start, _ir_constant(length, start.type))
    return [loop_index, *for_out]
  return for_out


def _lower_while_via_fori(
    ctx: LoweringRuleContext,
    *args,
    fori_jaxpr,
    cond_nconsts,
    body_nconsts,
):
  assert not fori_jaxpr.constvars
  # The pattern matcher looks for conditions with no constants.
  assert cond_nconsts == 0

  # Reflect the changes of the pattern matcher to the context.
  lb_aval, ub_aval, *_ = ctx.avals_in[cond_nconsts + body_nconsts:]
  ctx = ctx.replace(
      avals_in=(
          *ctx.avals_in[cond_nconsts:body_nconsts],
          ctx.avals_in[body_nconsts],  # the index
          *ctx.avals_in[body_nconsts + 2 :],
      ),
      avals_out=tuple(ctx.avals_out[2:]),
  )
  _, consts, (lb, ub, *args) = util.split_list(
      args, [cond_nconsts, body_nconsts]
  )
  lb = _ensure_ir_value(lb, lb_aval.dtype)
  ub = _ensure_ir_value(ub, ub_aval.dtype)
  for_out = _lower_jaxpr_to_for_loop(
      ctx,
      fori_jaxpr,
      lb,
      arith_dialect.subi(ub, lb),
      consts,
      *args,
      has_loop_index=True,
  )
  return ub, ub, *for_out


@register_lowering_rule(lax.while_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.while_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(lax.while_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(lax.while_p, *gpu_core.WGxWARP_SEMANTICS)
def _while_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    cond_jaxpr,
    body_jaxpr,
    cond_nconsts,
    body_nconsts,
):
  # First try to lower via a simpler fori loop, which may optimize better.
  fori_jaxpr, _ = pallas_utils.pattern_match_while_to_fori_loop(
      cond_jaxpr, cond_nconsts, body_jaxpr, body_nconsts
  )
  if fori_jaxpr is not None:
    return _lower_while_via_fori(
        ctx,
        *args,
        fori_jaxpr=fori_jaxpr,
        cond_nconsts=cond_nconsts,
        body_nconsts=body_nconsts,
    )

  _is_acc = lambda x: isinstance(x, mgpu.WGMMAAccumulator)
  _ensure = _ensure_ir_value
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    _ensure = lambda v, aval: v if _is_acc(v) else _ensure_fa(v, aval.dtype)

  # If we fail conversion to fori, fallback to an ordinary while loop.
  cond_consts, body_consts, carry = util.split_list(
      args, [cond_nconsts, body_nconsts]
  )
  _cond_avals, _body_avals, carry_avals = util.split_list(
      ctx.avals_in, [cond_nconsts, body_nconsts]
  )
  carry = [*map(_ensure, carry, carry_avals)]
  # Flatten the carry to get a concatenated list of registers from each FA.
  # Note that the treedef is also used below to unflatten the body results.
  flat_carry, carry_treedef = jax.tree.flatten(carry)
  flat_carry_types = [a.type for a in flat_carry]
  while_op = scf_dialect.WhileOp(flat_carry_types, flat_carry)

  before_block = while_op.before.blocks.append(*flat_carry_types)
  with ir.InsertionPoint.at_block_begin(before_block):
    cond_args = [*cond_consts, *carry_treedef.unflatten(before_block.arguments)]
    [cond] = lower_jaxpr_to_mosaic_gpu(
        ctx.module_ctx, ctx.launch_ctx, cond_jaxpr, cond_args
    )
    scf_dialect.condition(
        _ensure_ir_value(cond, *cond_jaxpr.out_avals), before_block.arguments
    )

  after_block = while_op.after.blocks.append(*flat_carry_types)
  with ir.InsertionPoint.at_block_begin(after_block):
    body_args = [*body_consts, *carry_treedef.unflatten(after_block.arguments)]
    loop_out = lower_jaxpr_to_mosaic_gpu(
        ctx.module_ctx, ctx.launch_ctx, body_jaxpr, body_args
    )
    loop_out = [*map(_ensure, loop_out, carry_avals)]
    if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
      for idx, (carry_fa, out_fa) in enumerate(zip(carry, loop_out)):
        if _is_acc(carry_fa) != _is_acc(out_fa):
          raise ValueError(
              f"The loop body output has unexpected accumulator type:"
              f" output[{idx}] is {out_fa}, when it should be {carry_fa}."
          )

        if not _is_acc(out_fa) and carry_fa.layout != out_fa.layout:
          raise ValueError(
              f"The loop body output has unexpected layout: output[{idx}] has"
              f" layout {out_fa.layout}, when it should be {carry_fa.layout}."
          )
    scf_dialect.yield_(
        carry_treedef.flatten_up_to(loop_out) if loop_out else []
    )
  return carry_treedef.unflatten(list(while_op.results))


@register_lowering_rule(lax.cond_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.cond_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(lax.cond_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(lax.cond_p, *gpu_core.WGxWARP_SEMANTICS)
def _cond_lowering_rule(ctx: LoweringRuleContext, index, *args, branches,
                        **params):
  if params:
    raise NotImplementedError("platform_dependent cond")
  index_aval, *_arg_avals = ctx.avals_in

  _is_acc = lambda x: isinstance(x, mgpu.WGMMAAccumulator)
  _ensure = _ensure_ir_value
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    _ensure = lambda v, aval: v if _is_acc(v) else _ensure_fa(v, aval.dtype)

  def _yielded_values(outs, avals):
    return [*map(_ensure, outs, avals)]

  # We need to know the result types ahead of time to construct the switch
  # operation. Below we lower the first branch in a throw-away module to
  # extract them.
  tmp_module = ir.Module.create()
  # The module attributes are needed for `get_arch` to work correctly.
  for k, v in ctx.launch_ctx.module.operation.attributes.items():
    tmp_module.operation.attributes[k] = v

  with ir.InsertionPoint(tmp_module.body):
    outs = lower_jaxpr_to_mosaic_gpu(
        ctx.module_ctx, ctx.launch_ctx, branches[0], args
    )
    yielded_types = [
        v.type for v in jax.tree.leaves(_yielded_values(outs, ctx.avals_out))
    ]
    del outs

  switch_op = scf_dialect.IndexSwitchOp(
      yielded_types,
      _as_index(_ensure_ir_value(index, index_aval.dtype)),
      range(len(branches) - 1),
  )

  # ``RegionSequence`` in MLIR does not support slicing, so the
  # auto-generated Python bindings for ``caseRegions`` fail at runtime!
  # We convert it to a list to work around that.
  regions = list(switch_op.regions)
  # Move the default region to the back.
  regions = regions[1:] + regions[:1]
  treedef: tree_util.PyTreeDef | None = None
  for branch, region in zip(branches, regions):
    block = region.blocks[0]
    with ir.InsertionPoint(block):
      outs = lower_jaxpr_to_mosaic_gpu(
          ctx.module_ctx, ctx.launch_ctx, branch, args, consts=branch.consts
      )

      yielded_leaves, yielded_treedef = jax.tree.flatten(_yielded_values(outs, ctx.avals_out))
      if treedef is None:
        treedef = yielded_treedef
      else:
        assert treedef == yielded_treedef

      scf_dialect.yield_(yielded_leaves)

  assert treedef is not None
  return treedef.unflatten(list(switch_op.results))


@register_lowering_rule(lax.bitcast_convert_type_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(
    lax.bitcast_convert_type_p, mgpu.LoweringSemantics.Warpgroup
)
def _bitcast_convert_type_lowering_rule(
    ctx: LoweringRuleContext, x, *, new_dtype
):
  [x_aval] = ctx.avals_in
  src_elem_type = mgpu_utils.dtype_to_ir_type(x_aval.dtype)
  dst_elem_type = mgpu_utils.dtype_to_ir_type(new_dtype)
  assert isinstance(src_elem_type, (ir.IntegerType, ir.FloatType))
  assert isinstance(dst_elem_type, (ir.IntegerType, ir.FloatType))
  if src_elem_type.width != dst_elem_type.width:
    raise NotImplementedError(
        f"Cannot bitcast from {x_aval.dtype} to {new_dtype} because they"
        " have different widths"
    )

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    x = _ensure_ir_value(x, x_aval.dtype)
    return arith_dialect.bitcast(
        ir.VectorType.get(x_aval.shape, dst_elem_type), x
    )

  x = _ensure_fa(x, x_aval.dtype)
  output_is_signed = mgpu_utils.is_signed(new_dtype)
  return mgpu.FragmentedArray.bitcast(
      x, dst_elem_type, output_is_signed=output_is_signed
  )


@register_lowering_rule(lax.optimization_barrier_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(
    lax.optimization_barrier_p, *gpu_core.LANExWARP_SEMANTICS
)
def _optimization_barrier_lowering(ctx: LoweringRuleContext, *args):
  result = mgpu.optimization_barrier(
      *(_ensure_fa(arg, aval.dtype) for arg, aval in zip(args, ctx.avals_in))
  )
  return (result,) if len(ctx.avals_in) == 1 else result


@register_lowering_rule(
    lax.optimization_barrier_p, mgpu.LoweringSemantics.Warpgroup
)
@register_lowering_rule(
    lax.optimization_barrier_p, *gpu_core.WGxWARP_SEMANTICS
)
def _optimization_barrier_lowering_wg(ctx: LoweringRuleContext, *args):
  result = mgpu.dialect.optimization_barrier([
      _ensure_ir_value(arg, aval.dtype) for arg, aval in zip(args, ctx.avals_in)
  ])
  return (result,) if len(ctx.avals_in) == 1 else result


def _closed_over_values(block: ir.Block) -> list[ir.Value]:
  """Returns a list of unique values used within `block` that are defined outside of `block`."""
  def _closed_over_values_inner(
      block: ir.Block, vals_in_block: set[ir.Value]
  ) -> list[ir.Value]:
    closed_over_values = []
    for arg in block.arguments:
      vals_in_block.add(arg)
    for op in block.operations:
      for o in op.operands:
        if o not in vals_in_block:
          closed_over_values.append(o)
      for r in op.regions:
        for b in r.blocks:
          closed_over_values.extend(_closed_over_values_inner(b, vals_in_block))
      for r in op.results:
        vals_in_block.add(r)
    return closed_over_values
  result = _closed_over_values_inner(block, set())
  # Remove duplicates while preserving order.
  return list(dict.fromkeys(result))


def _replace_uses_in_block(old: ir.Value, new: ir.Value, block: ir.Block):
  """Replaces all uses of the `old` value with the `new` value in `block`."""

  def is_contained_within_block(operand: ir.OpOperand, block: ir.Block) -> bool:
    current_op = operand.owner.operation
    while (parent := current_op.parent) is not None:
      if current_op.block == block:
        return True
      current_op = parent
    return False

  for use in old.uses:
    if is_contained_within_block(use, block):
      use.owner.operands[use.operand_number] = new


def _isolate_from_above(
    op: ir.Operation | ir.OpView
) -> ir.Operation | ir.OpView:
  """Makes `op` conform to the `IsolatedFromAbove` trait.

  This replaces all captured values with new op operands.
  """
  if len(op.regions) != 1:
    raise NotImplementedError("Only support ops with one region.")
  if len(op.regions[0].blocks) != 1:
    raise NotImplementedError("Only support ops with one block.")

  block = op.regions[0].blocks[0]
  captures = _closed_over_values(block)
  if not captures:
    return op

  # 1. Create the new operation shell with the expanded operands.
  new_op = ir.Operation.create(
      name=op.name,
      results=[res.type for res in op.results],
      operands=list(op.operands) + captures,
      attributes=dict(op.attributes),
      regions=1,
  )

  # 2. Move the block from the old op to the new op.
  block.append_to(new_op.regions[0])

  # 3. Create new block arguments for the captured values.
  new_args = []
  for capture in captures:
    new_args.append(block.add_argument(capture.type, capture.location))

  # 4. Redirect all references from captured values to the newly created
  # block's arguments.
  for capture, new_arg in zip(captures, block.arguments[-len(captures) :]):
    _replace_uses_in_block(capture, new_arg, block)

  # 5. Clean up the now-empty old operation.
  op.erase()

  return new_op


@register_lowering_rule(mpmd.mpmd_map_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(mpmd.mpmd_map_p, mgpu.LoweringSemantics.Warpgroup)
def _mpmd_map_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    jaxprs,
    meshes,
    **_,
):
  if len(jaxprs) > 1:
    raise NotImplementedError(
        "MPMD map with multiple jaxprs not implemented."
    )
  mesh = meshes[0]
  jaxpr = jaxprs[0]
  if not isinstance(mesh, gpu_core.WarpMesh):
    raise NotImplementedError(f"Unsupported mesh: {mesh}")
  # A mpmd_map over a WarpMesh represents a fork/join over individual
  # warps in a warpgroup.
  if (ctx.module_ctx.warp_axis_name or
      ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp):
    raise LoweringError(
        "Cannot nest mpmd_maps. Already under mpmd_map with warp_axis_name "
        f"{ctx.module_ctx.warp_axis_name}.")
  module_ctx = dataclasses.replace(
      ctx.module_ctx,
      warp_axis_name=mesh.axis_name,
      primitive_semantics=gpu_core.PrimitiveSemantics.Warp,
  )
  for aval_in in ctx.avals_in:
    if isinstance(aval_in, jax_core.ShapedArray) and aval_in.shape:
      raise LoweringError(
        "Can only close over scalars and Refs when using mpmd_map with "
        f"WarpMesh. Found array of shape {aval_in}."
      )
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    # We allow the warps to schedule async copies without synchronizing with
    # other warps, so we need to add a barrier here to make sure all reads and
    # writes have completed.
    if ctx.module_ctx.auto_barriers:
      mgpu.warpgroup_barrier()
    _ = lower_jaxpr_to_mosaic_gpu(
        module_ctx,
        ctx.launch_ctx,
        jaxpr,
        args=args,
        consts=(),
    )
    if ctx.module_ctx.auto_barriers:
      # We need to ensure that any effects produced by one warp
      # (e.g. async copies) are observable by all other warps.
      mgpu.warpgroup_barrier()
  else:
    warp_map_op = mgpu.dialect.WarpMapOp(operands=[])
    with ir.InsertionPoint(warp_map_op.body):
      _ = lower_jaxpr_to_mosaic_gpu(
          module_ctx,
          ctx.launch_ctx,
          jaxpr,
          args=args,
          consts=(),
      )
    _isolate_from_above(warp_map_op)
  return []


def _bcast(
    x: Any,
    y: Any,
    x_aval: ShapedAbstractValue,
    y_aval: ShapedAbstractValue,
    out_aval: ShapedAbstractValue,
) -> tuple[mgpu.FragmentedArray, mgpu.FragmentedArray]:
  if not isinstance(x, mgpu.FragmentedArray):
    x_dtype = x_aval.dtype
    if x_aval.weak_type:
      x_dtype = y_aval.dtype
    x = _ensure_fa(x, x_dtype)
  if not isinstance(y, mgpu.FragmentedArray):
    y_dtype = y_aval.dtype
    if y_aval.weak_type:
      y_dtype = x_aval.dtype
    y = _ensure_fa(y, y_dtype)
  if x_aval.shape != out_aval.shape:
    x = x.broadcast(out_aval.shape)
  if y_aval.shape != out_aval.shape:
    y = y.broadcast(out_aval.shape)
  return x, y


def _ensure_fa(x: object, dtype: jax.typing.DTypeLike) -> mgpu.FragmentedArray:
  if isinstance(x, mgpu.FragmentedArray):
    assert x.mlir_dtype == mgpu_utils.dtype_to_ir_type(dtype)
    return x
  return mgpu.FragmentedArray.splat(
      _ensure_ir_value(x, dtype), (), is_signed=mgpu_utils.is_signed(dtype)
  )


def _bcast_wg(
    x: Any,
    y: Any,
    x_aval: ShapedAbstractValue,
    y_aval: ShapedAbstractValue,
    out_aval: ShapedAbstractValue,
) -> tuple[ir.Value, ir.Value]:
  """Ensures that ``x`` and ``y`` have the expected shapes and dtypes.

  More specifically, the inputs are converted to vectors of the same dtype
  as ``x_aval`` and ``y_aval``, and broadcasted to the output shape
  if necessary.
  """
  x_dtype = x_aval.dtype
  if not isinstance(x, ir.Value):
    if x_aval.weak_type:
      x_dtype = y_aval.dtype
    x = _ensure_ir_value(x, x_dtype)
  y_dtype = y_aval.dtype
  if not isinstance(y, ir.Value):
    if y_aval.weak_type:
      y_dtype = x_aval.dtype
    y = _ensure_ir_value(y, y_dtype)
  if not out_aval.shape:
    return x, y

  def bcast(value, dtype):
    ty = ir.VectorType.get(out_aval.shape, mgpu_utils.dtype_to_ir_type(dtype))
    if isinstance(value.type, ir.VectorType):
      assert value.type.rank == len(out_aval.shape), "broadcast ranks mismatch"
      dims = range(len(out_aval.shape))
      return mgpu.dialect.broadcast_in_dim(ty, value, dims)
    else:  # scalar broadcast
      return vector_dialect.broadcast(ty, value)

  if x_aval.shape != out_aval.shape:
    x = bcast(x, x_dtype)
  if y_aval.shape != out_aval.shape:
    y = bcast(y, y_dtype)
  return x, y


def _ensure_ir_value(x: Any, dtype: jax.typing.DTypeLike) -> ir.Value:
  mlir_dtype = mgpu_utils.dtype_to_ir_type(dtype)
  if isinstance(x, ir.Value):
    if isinstance(x.type, ir.VectorType):
      assert ir.VectorType(x.type).element_type == mlir_dtype
    else:
      assert x.type == mlir_dtype, (x.type, mlir_dtype)
    return x
  elif isinstance(x, mgpu.FragmentedArray):
    assert x.mlir_dtype == mlir_dtype
    if isinstance(x.layout, mgpu.WGSplatFragLayout):
      return x.registers.item()
    raise NotImplementedError(f"Unsupported layout: {x.layout}")
  return _ir_constant(x, mlir_dtype)


def _device_id_to_logical(
    ctx: LoweringRuleContext, device_id,
    device_id_type: primitives.DeviceIdType,
    device_id_aval: Any):
  def jax_fn(device_id_val):
    logical_device_id, non_mesh_axes = primitives.device_id_to_logical(
        ctx.module_ctx.mesh_info,
        device_id_val,
        device_id_type,
        lax.axis_index
    )
    if non_mesh_axes:
        raise ValueError(f"Unrecognized axes in device_id: {non_mesh_axes}")
    return logical_device_id

  return _lower_fun(jax_fn, in_avals=(device_id_aval,))(ctx, device_id)


def _ir_constant(v: object, t: ir.Type) -> ir.Value:
  if isinstance(
      v, (np.number, np.ndarray, int, float, jax_literals.TypedNdArray)
  ):
    if isinstance(t, (ir.IntegerType, ir.IndexType)):
      v = int(v)
    else:
      assert isinstance(t, ir.FloatType)
      v = float(v)
    return arith_dialect.constant(t, v)
  raise NotImplementedError(f"Unsupported constant: {v!r}")


def _i32_constant(v: int) -> ir.Value:
  if v < jnp.iinfo(jnp.int32).min or v > jnp.iinfo(jnp.int32).max:
    raise ValueError(f"Integer constant out of range for i32: {v}")
  return arith_dialect.constant(ir.IntegerType.get_signless(32), v)


def _i64_constant(v: int) -> ir.Value:
  if v < jnp.iinfo(jnp.int64).min or v > jnp.iinfo(jnp.int64).max:
    raise ValueError(f"Integer constant out of range for i64: {v}")
  return arith_dialect.constant(ir.IntegerType.get_signless(64), v)


def _as_index(v: object) -> ir.Value:
  match v:
    case int():
      return arith_dialect.constant(ir.IndexType.get(), v)
    case ir.Value() if isinstance(v.type, ir.IndexType):
      return v
    case ir.Value() if isinstance(v.type, ir.IntegerType):
      return arith_dialect.index_cast(ir.IndexType.get(), v)
    case mgpu.FragmentedArray(layout=mgpu.WGSplatFragLayout()):
      return _as_index(v.registers.item())
    case jax_literals.TypedNdArray() if (
        np.issubdtype(v.dtype, np.integer) and v.ndim == 0
    ):
      return arith_dialect.constant(ir.IndexType.get(), int(v))
    case _:
      raise ValueError(f"Unsupported index: {v} of type {type(v)}")


def merge_indexers(
    indexers: Sequence[indexing.NDIndexer]) -> indexing.NDIndexer:
  """Merges multiple indexers into a single indexer.

  This function computes a new indexer such that applying the
  new indexer produces the same result as applying the sequence
  of input indexers in order from first-to-last.
  """
  if len(indexers) == 0:
    raise ValueError("Cannot merge empty list of indexers")
  if len(indexers) == 1:
    return indexers[0]
  root_shape = indexers[0].shape
  current_indices = [indexing.Slice(0, size, 1) for size in root_shape]
  removed_dimensions = set()
  for indexer in indexers:
    if indexer.int_indexer_shape:
      raise NotImplementedError()

    def _ensure_idx_fa(x: Any) -> mgpu.FragmentedArray:
      i32 = ir.IntegerType.get_signless(32)
      if isinstance(x, ir.Value):
        # TODO(cperivol): We assume all indices are signed. We should
        # look at the JAX avals to see if the integers are signed or
        # not to figure out is_signed.
        is_signed = False if isinstance(x.type, ir.IntegerType) else None
        return mgpu.FragmentedArray.splat(x, (), is_signed=is_signed).astype(
            i32, is_signed=False
        )
      if isinstance(x, mgpu.FragmentedArray):
        return x.astype(i32, is_signed=False)
      if isinstance(x, int):
        return mgpu.FragmentedArray.splat(mgpu.c(x, i32), (), is_signed=False)
      if (
          isinstance(x, jax_literals.TypedNdArray)
          and x.ndim == 0
          and np.issubdtype(x.dtype, np.signedinteger)
      ):
        return mgpu.FragmentedArray.splat(
            mgpu.c(int(x), i32), (), is_signed=False
        )
      raise NotImplementedError(x)

    num_skipped = 0
    for i in range(len(current_indices)):
      # Integer indexers remove dimensions which should be
      # skipped by following indexers.
      if i in removed_dimensions:
        num_skipped += 1
        continue
      dim_indexer = indexer.indices[i - num_skipped]
      current_index = current_indices[i]
      assert isinstance(current_index, indexing.Slice)

      current_start_index = _ensure_idx_fa(current_index.start)
      if isinstance(dim_indexer, indexing.Slice):
        if dim_indexer.stride != 1:
          raise NotImplementedError("Non-unit strides not implemented.")
        current_indices[i] = indexing.Slice(
            current_start_index + _ensure_idx_fa(dim_indexer.start),
            dim_indexer.size,
            1,
        )
      else:
        current_indices[i] = current_start_index + _ensure_idx_fa(dim_indexer)
        removed_dimensions.add(i)
  return indexing.NDIndexer(
      indices=tuple(current_indices),
      shape=root_shape,
      int_indexer_shape=(),
  )


@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Warpgroup)
def _semaphore_read_lowering_rule(ctx: LoweringRuleContext, *args, args_tree):
  sem, transforms = tree_util.tree_unflatten(args_tree, args)
  sem_aval, transform_avals = tree_util.tree_unflatten(args_tree, ctx.avals_in)
  assert isinstance(sem_aval, state_types.AbstractRef)
  sem, _, transforms = _handle_transforms(ctx, sem_aval, sem, transform_avals, transforms)
  if transforms:
    raise NotImplementedError(f"Unhandled transforms for semaphore_read: {transforms}")
  sem_ptr = mgpu.utils.memref_ptr(sem)
  i32_ty = ir.IntegerType.get_signless(32)
  result = llvm_dialect.inline_asm(
      i32_ty,
      [sem_ptr],
      "ld.acquire.sys.u32 $0,[$1];",
      "=r,l",
      has_side_effects=True,
  )
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(result, jnp.int32)
  return result


@contextlib.contextmanager
def _wrap_in_custom_primitive_if_wg(
    ctx: LoweringRuleContext, operands: Sequence[ir.Value]
):
  """Wraps the body in a CustomPrimitiveOp for warpgroup semantics.

  For warpgroup lowering semantics, yields remapped block arguments that
  should be used instead of the original operands, and temporarily switches
  the lowering semantics to Lane for the duration of the context. For lane
  semantics, yields the original operands unchanged.
  """
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    custom_op = mgpu.dialect.CustomPrimitiveOp(
        result=[],
        operands_=list(operands),
        in_layouts=[],
        in_transforms=[],
        out_layouts=[],
    )
    block = custom_op.body.blocks.append(*[o.type for o in operands])
    with ir.InsertionPoint(block):
      ctx.module_ctx.lowering_semantics = mgpu.LoweringSemantics.Lane
      try:
        yield list(block.arguments)
      finally:
        ctx.module_ctx.lowering_semantics = mgpu.LoweringSemantics.Warpgroup
      mgpu.dialect.ReturnOp(operands_=[])
    _isolate_from_above(custom_op)
  else:
    yield list(operands)


@register_lowering_rule(checkify.check_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(checkify.check_p, mgpu.LoweringSemantics.Warpgroup)
def _check_lowering_rule(ctx: LoweringRuleContext, *err_args, err_tree, debug):
  if not debug:
    raise NotImplementedError(
        "Non-debug checks are not supported by the Mosaic GPU backend."
        " Functionalize them via `jax.experimental.checkify`."
    )
  if not pallas_core.debug_checks_enabled():
    return []

  error = jax.tree.unflatten(err_tree, err_args)
  [pred] = error._pred.values()
  [exception_tree] = error._metadata.values()
  [payload] = error._payload.values()
  exception = jax.tree.unflatten(exception_tree, payload)
  assert isinstance(exception, checkify.FailedCheckError)

  # check_p has an inverted predicate compared to assert, so we need to compute
  # ``not pred`` here.
  minus_one = _ir_constant(-1, mgpu_utils.dtype_to_ir_type(jnp.bool))
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    pred = pred.registers.item()
  not_pred = arith_dialect.xori(pred, minus_one)
  cf_dialect.assert_(not_pred, exception.fmt_string)
  return []


@register_lowering_rule(pjit.relayout_p, mgpu.LoweringSemantics.Lane)
def _relayout_lowering_lane(
    ctx: LoweringRuleContext, x, *, dst_layout
):
  del ctx, x, dst_layout
  raise NotImplementedError(
      "relayout_p is not supported with Lane semantics."
  )


@register_lowering_rule(pjit.relayout_p, mgpu.LoweringSemantics.Warpgroup)
def _relayout_lowering_wg(
    ctx: LoweringRuleContext, x, *, dst_layout
):
  if dst_layout is jax_layout.AutoLayout:
    return x
  layout = fa.TiledLayout(
      dst_layout.tiling,
      dst_layout.warp_dims,
      dst_layout.lane_dims,
      dst_layout.vector_dim,
  )
  if ctx.avals_in[0].ndim == 0:  # scalar case
    if layout != mgpu.WGSplatFragLayout():
      raise ValueError(
          "Only plgpu.Layout.WG_SPLAT is supported for scalar values."
      )
    return x
  return mgpu.dialect.layout_cast(x, mgpu.to_layout_attr(layout))


@register_lowering_rule(gpu_core.layout_cast_p, mgpu.LoweringSemantics.Lane)
def _layout_cast_lowering(ctx: LoweringRuleContext, x, *, new_layout):
  del ctx  # Unused.
  return x.to_layout(new_layout.to_mgpu())


@register_lowering_rule(gpu_core.layout_cast_p, mgpu.LoweringSemantics.Warpgroup)
def _layout_cast_lowering_wg(
    ctx: LoweringRuleContext, x, *, new_layout
):
  layout = new_layout.to_mgpu()
  if ctx.avals_in[0].ndim == 0:  # scalar case
    if layout != mgpu.WGSplatFragLayout():
      raise ValueError(
          "Only plgpu.Layout.WG_SPLAT is supported for scalar values."
      )
    return x
  return mgpu.dialect.layout_cast(x, mgpu.to_layout_attr(layout))


@register_lowering_rule(lax.iota_p, mgpu.LoweringSemantics.Lane)
def _iota_lowering(
    ctx: LoweringRuleContext, dtype, shape, dimension, sharding
):
  del sharding  # Unused.
  if ctx.out_layout_hint is None:
    raise RuntimeError(
        "Failed to infer the output layout of the iota. Please apply"
        " plgpu.layout_cast to its output right after its creation."
    )
  mlir_dtype = mgpu_utils.dtype_to_ir_type(dtype)
  is_signed = mgpu_utils.is_signed(dtype)
  return mgpu.FragmentedArray.broadcasted_iota(
      mlir_dtype, shape, dimension, ctx.out_layout_hint, is_signed=is_signed
  )


@register_lowering_rule(lax.iota_p, mgpu.LoweringSemantics.Warpgroup)
def _iota_lowering_wg(
    ctx: LoweringRuleContext, dtype, shape, dimension, sharding
):
  del ctx, sharding
  result_type = ir.VectorType.get(shape, mgpu_utils.dtype_to_ir_type(dtype))
  return mgpu.dialect.broadcasted_iota(result_type, dimension)


@register_lowering_rule(primitives.delay_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(primitives.delay_p, mgpu.LoweringSemantics.Warpgroup)
def _delay_lowering(ctx: LoweringRuleContext, nanos):
  del ctx  # Unused.
  if not isinstance(nanos, ir.Value):
    nanos = _i32_constant(nanos)
  mgpu.nanosleep(nanos)
  return []


@register_lowering_rule(pjit.reshard_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(pjit.reshard_p, mgpu.LoweringSemantics.Warpgroup)
def _reshard_lowering_rule(ctx, x, dst_sharding, concrete_mesh):
  del ctx, dst_sharding, concrete_mesh
  return x