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

"""Lowering rules and pass for the MLIR Mosaic GPU dialect."""

from __future__ import annotations

from collections.abc import Callable, Iterable, Sequence
import dataclasses
import functools
import math
import operator
from typing import Any, Literal, assert_never, cast

from jax._src.interpreters import mlir as mlir_interpreter
from jax._src.lib import mosaic_gpu_dialect as mgpu
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import arith
from jax._src.lib.mlir.dialects import builtin
from jax._src.lib.mlir.dialects import func
from jax._src.lib.mlir.dialects import gpu
from jax._src.lib.mlir.dialects import llvm
from jax._src.lib.mlir.dialects import math as mlir_math
from jax._src.lib.mlir.dialects import memref
from jax._src.lib.mlir.dialects import scf
from jax._src.lib.mlir.dialects import vector
from jax.experimental.mosaic.gpu.mma import mma as do_mma
from jax.experimental.mosaic.gpu.mma import MMALayouts
import numpy as np

from . import constraints as cs
from . import fragmented_array as fa
from . import inference_utils
from . import launch_context as lc
from . import layouts as layouts_lib
from . import tcgen05
from . import utils
from . import wgmma


@dataclasses.dataclass()
class LoweringContext:
  launch_context: lc.LaunchContext
  _single_thread_per_warp_predicate: ir.Value
  _single_thread_per_warpgroup_predicate: ir.Value
  single_thread_per_block_predicate: ir.Value
  single_warp_per_block_predicate: ir.Value
  auto_barriers: bool
  smem_requested_bytes: int
  is_collective_kernel: bool | None = dataclasses.field(
      init=False, default=None
  )
  thread_semantics: Literal[
      utils.ThreadSubset.WARPGROUP, utils.ThreadSubset.WARP
  ] = utils.ThreadSubset.WARPGROUP

  @property
  def single_lane_predicate(self) -> ir.Value:
    match self.thread_semantics:
      case utils.ThreadSubset.WARPGROUP:
        return self._single_thread_per_warpgroup_predicate
      case utils.ThreadSubset.WARP:
        return self._single_thread_per_warp_predicate
      case _:
        assert_never(self.thread_semantics)

  def check_collective(self, op: ir.OpView) -> None:
    """Checks that the collective attribute is consistent across operations.

    It is an error to mix collective and non-collective operations in the same
    kernel.
    """
    if "collective" not in op.attributes:
      return
    if self.is_collective_kernel is None:
      self.is_collective_kernel = bool(op.attributes["collective"])
    elif self.is_collective_kernel != bool(op.attributes["collective"]):
      raise ValueError(
          "Collective attributes are inconsistent across operations in the"
          " kernel."
      )

  def lower_op(self, op: ir.OpView):
    if not _should_lower(op):
      return

    # pyrefly: ignore[missing-attribute]
    if (name := op.OPERATION_NAME) not in _lowerings:
      raise NotImplementedError(f"Missing lowering rule for {op}")

    lowering_rule = _lowerings[name]

    if (
        self.thread_semantics == utils.ThreadSubset.WARP
        and name not in _supported_warp_lowerings
    ):
      raise NotImplementedError(f"Op {name} does not support warp semantics.")

    # TODO(bchetioui): make sure all layouts are set here.
    if inference_utils.should_have_layout(
        op
    ) and not inference_utils.has_any_layout_set(op):
      raise ValueError(f"{op} is missing a layout and can not be lowered.")

    with op.operation.location:
      new_results = lowering_rule(self, op)
    if not isinstance(new_results, Recursed):
      for old, new in zip(op.results, new_results):
        old.replace_all_uses_with(new)
      op.erase()


class Recursed:
  pass
RECURSED = Recursed()

MlirLoweringRuleResult = Sequence[ir.Value] | Recursed
MlirLoweringRule = Callable[
    [LoweringContext, Any], MlirLoweringRuleResult
]


_lowerings: dict[str, MlirLoweringRule] = {}


def _undo_conversion_cast(
    ir_value: ir.Value,
    expected_types: Sequence[ir.Type],
) -> tuple[builtin.UnrealizedConversionCastOp, Sequence[ir.Value]]:
  """Undoes the provided unrealized conversion cast.

  The `ir_value` must be an unrealized conversion cast. The returned tuple
  contains:
  - The original unrealized conversion cast (useful for extract attributes).
  - The list of operands of the original conversion cast.

  The function will verify that the returned values have types that match
  `expected_types`.
  """
  conversion_cast: Any = ir_value.owner

  if not isinstance(conversion_cast, builtin.UnrealizedConversionCastOp):
    raise ValueError(f"{conversion_cast} is not a conversion_cast")

  original_inputs: Sequence[ir.Value] = conversion_cast.operands

  for v, t in zip(original_inputs, expected_types, strict=True):
    if v.type != t:
      raise ValueError(f"Expected type {t} for value {v}")

  return conversion_cast, original_inputs


def fragmented_array_to_ir(
    fragmented_array: fa.FragmentedArray, ty: ir.Type
) -> ir.Value:
  """Converts a FragmentedArray to an IR value.

  The fragmented array's signedness is omitted from the IR representation.
  """
  conversion_cast = builtin.UnrealizedConversionCastOp(
      [ty], fragmented_array.registers.flatten().tolist()
  )

  conversion_cast.attributes["layout"] = layouts_lib.to_layout_attr(
      fragmented_array.layout
  )

  return conversion_cast.result


def _default_is_signed(dtype: ir.Type) -> bool | None:
  """Returns `False` for Integer types, `None` otherwise.

  When converting from Pallas dtype to IR type, we lose the `is_signed`
  information. We can default to `False` for most use cases.
  """
  return False if isinstance(dtype, ir.IntegerType) else None


def _fragmented_array_from_ir(
    fragmented_array_as_ir: ir.Value,
    layout: ir.Attribute,
    is_signed: bool | None = None,
) -> fa.FragmentedArray:
  producer_layout_attr = fragmented_array_as_ir.owner.attributes["layout"]  # pyrefly: ignore[missing-attribute]
  producer_layout = layouts_lib.from_layout_attr(producer_layout_attr)
  vector_ty = ir.VectorType(fragmented_array_as_ir.type)
  reg_shape = producer_layout.registers_shape(tuple(vector_ty.shape))
  reg_ty = producer_layout.registers_element_type(vector_ty.element_type)

  _, original_inputs = _undo_conversion_cast(
      fragmented_array_as_ir, [reg_ty] * math.prod(reg_shape)
  )

  registers = np.array(list(original_inputs)).reshape(reg_shape)

  if isinstance(vector_ty.element_type, ir.IntegerType):
    is_signed = False if is_signed is None else is_signed

  return fa.FragmentedArray(
      _registers=registers, _layout=producer_layout, _is_signed=is_signed
  ).to_layout(layouts_lib.from_layout_attr(layout))


def wrap_transformed_memref(
    transformed_memref: ir.Value,
    logical_type: ir.Type,
    transforms: ir.ArrayAttr,
) -> ir.Value:
  """Wraps a transformed memref to an unrealized cast with transforms.

  The return type of the cast is the untransformed logical type.
  """
  conversion_cast = builtin.UnrealizedConversionCastOp(
      [logical_type], [transformed_memref]
  )
  conversion_cast.attributes["transforms"] = transforms
  return conversion_cast.result


# TODO(bchetioui): Change ``expected_transforms`` to be a ``Sequence``.
def unwrap_transformed_memref(
    ref: ir.Value[ir.MemRefType], expected_transforms: ir.ArrayAttr
) -> ir.Value:
  """Uwraps a memref from an unrealized cast and verifies its transforms."""
  transformed_type = transform_type(ref.type, expected_transforms)
  conversion_cast, [result] = _undo_conversion_cast(ref, [transformed_type])

  # Check that the actual transforms match the expected ones.
  if expected_transforms != conversion_cast.attributes["transforms"]:
    raise ValueError(
        f"Expected transforms {expected_transforms} do not match actual"
        f" transforms {conversion_cast.attributes['transforms']}"
    )

  return result

_supported_warp_lowerings: set[str] = set()


def _register_lowering(
    op: str | type[ir.OpView] | None,
    support_warp_semantics: bool = False,
) -> Callable[[MlirLoweringRule], MlirLoweringRule]:
  def wrapper(f):
    if op is not None:
      op_name = op if isinstance(op, str) else op.OPERATION_NAME  # pyrefly: ignore[missing-attribute]
      _lowerings[op_name] = f
      if support_warp_semantics:
        _supported_warp_lowerings.add(op_name)
    return f

  return wrapper


def _lowered_barrier_type() -> ir.Type:
  return ir.IntegerType.get_signless(64)


@_register_lowering(mgpu.InitializeBarrierOp)
def _initialize_barrier_op_lowering_rule(
    ctx: LoweringContext,
    op: mgpu.InitializeBarrierOp,
) -> Sequence[ir.Value]:
  i32 = ir.IntegerType.get_signless(32)
  arrival_count = op.arrival_count.value * (
      utils.WARPGROUP_SIZE if not op.orders_tensor_core.value else 1
  )
  for i in range(op.num_barriers.value):
    bar_ptr = utils.getelementptr(op.base_pointer, [i], _lowered_barrier_type())
    llvm.inline_asm(
        ir.Type.parse("!llvm.void"),
        [bar_ptr, utils.c(arrival_count, i32), ctx.single_thread_per_block_predicate],
        "@$2 mbarrier.init.shared::cta.b64 [$0], $1;",
        "r,r,b",
        has_side_effects=True,
    )
  return []


@_register_lowering(mgpu.AssumeMultipleOp)
def _assume_multiple_op_lowering_rule(
    _: LoweringContext,
    op: mgpu.AssumeMultipleOp,
) -> Sequence[ir.Value]:
  return [op.value]


@_register_lowering(mgpu.OptimizationBarrierOp, support_warp_semantics=True)
def _optimization_barrier_op_lowering_rule(
    _: LoweringContext,
    op: mgpu.OptimizationBarrierOp,
) -> Sequence[ir.Value]:
  if inference_utils.has_in_layouts_set(op):
    in_layouts = inference_utils.in_layouts(op)
  else:
    in_layouts = []

  fragmented_arrays = []
  layout_idx = 0
  is_scalar = lambda dtype: isinstance(dtype, (ir.IntegerType, ir.IndexType, ir.FloatType))
  for operand in op.operands:
    if isinstance(operand.type, ir.VectorType):
      layout = in_layouts[layout_idx]
      layout_idx += 1
      fragmented_arrays.append(_fragmented_array_from_ir(operand, layout))
    else:
      if not is_scalar(operand.type):
        raise NotImplementedError(
            f"Unsupported optimization barrier for operand {operand.type}"
        )
      is_signed = _default_is_signed(operand.type)
      fragmented_arrays.append(
          fa.FragmentedArray.splat(operand, shape=(), is_signed=is_signed)
      )

  assert layout_idx == len(in_layouts)
  lowered_fragmented_arrays = fa.optimization_barrier(*fragmented_arrays)
  if isinstance(lowered_fragmented_arrays, fa.FragmentedArray):
    lowered_fragmented_arrays = [lowered_fragmented_arrays]

  result = []
  for arr, res in zip(lowered_fragmented_arrays, op.results, strict=True):
    if is_scalar(res.type):
      result.append(arr.registers.item())
    else:
      result.append(fragmented_array_to_ir(arr, res.type))
  return result


@_register_lowering(mgpu.GetClusterRefOp)
def _get_cluster_ref_op_lowering_rule(
    _: LoweringContext, op: mgpu.GetClusterRefOp,
) -> Sequence[ir.Value]:
  index = ir.IndexType.get()
  specified_idxs = [
      (d, dim) for d, dim in zip((op.x, op.y, op.z), gpu.Dimension)
      if d is not None
  ]
  if len(specified_idxs) != 1:
    raise ValueError(
        "Exactly one cluster dimension must be specified, got"
        f" {len(specified_idxs)}"
    )
  [(idx, dim)] = specified_idxs
  [in_transforms] = inference_utils.in_transforms(op)
  result = utils.get_cluster_ref(
      unwrap_transformed_memref(op.source, in_transforms),
      dim,
      arith.index_cast(index, idx),
      generic=False,
  )
  [out_transforms] = inference_utils.out_transforms(op)
  return [wrap_transformed_memref(result, op.result.type, out_transforms)]


@_register_lowering(arith.ConstantOp)
def _arith_constant_op_lowering_rule(
    _: LoweringContext, op: arith.ConstantOp
) -> Sequence[ir.Value]:
  if not isinstance(op.value, ir.DenseElementsAttr):
    raise NotImplementedError(f"Unsupported constant op: {op}")

  value = ir.DenseElementsAttr(op.value)
  if not value.is_splat:
    raise NotImplementedError(f"Unsupported constant op: {op}")

  ty = ir.VectorType(op.result.type)
  is_signed = _default_is_signed(ty.element_type)
  layout = layouts_lib.from_layout_attr(inference_utils.out_layouts(op)[0])
  return [
      fragmented_array_to_ir(
          fa.FragmentedArray.splat(
              arith.constant(ty.element_type, value.get_splat_value()),
              tuple(ty.shape),
              layout,
              is_signed=is_signed,
          ),
          op.result.type,
      )
  ]


@_register_lowering(mgpu.VectorLoadOp)
def _vector_load_op_lowering_rule(
    _: LoweringContext, op: mgpu.VectorLoadOp
) -> Sequence[ir.Value]:
  out_layout_attr, = inference_utils.out_layouts(op)
  out_layout = layouts_lib.from_layout_attr(out_layout_attr)

  element_type = ir.VectorType(op.result.type).element_type
  is_signed = _default_is_signed(element_type)

  orig_ref_ty = op.source.type
  if utils.is_smem_ref(orig_ref_ty) or utils.is_cluster_smem_ref(orig_ref_ty):
    [transforms_attr] = inference_utils.in_transforms(op)
    transformed_ref = unwrap_transformed_memref(op.source, transforms_attr)
  else:
    transforms_attr = None
    transformed_ref = op.source

  def _fragmented_array_to_ir(
      fragmented_array: fa.FragmentedArray,
  ) -> ir.Value:
    return fragmented_array_to_ir(fragmented_array, op.result.type)

  if isinstance(out_layout, fa.WGStridedFragLayout):
    # TODO(bchetioui): Process transforms.
    if transforms_attr is not None:
      swizzle = swizzle_from_transforms_attr(transforms_attr)
      transforms = memref_transforms_from_transforms_attr(transforms_attr)
      if swizzle != mgpu.SwizzlingMode.kNoSwizzle or transforms:
        raise NotImplementedError(
            "Transformed or swizzled strided loads are not supported"
        )

    fragmented_array = fa.FragmentedArray.load_strided(
        transformed_ref,
        is_signed=is_signed,
        vec_size=out_layout.vec_size,
    )
    return [_fragmented_array_to_ir(fragmented_array)]

  if not isinstance(out_layout, fa.TiledLayout):
    raise ValueError(f"{op} has an unsupported layout: {out_layout_attr}")

  optimized = op.optimized.value if op.optimized is not None else True
  if transformed_ref.type.memory_space is None:  # GMEM
    fragmented_array = fa.FragmentedArray.load_untiled(
        transformed_ref,
        layout=out_layout,
        is_signed=is_signed,
        optimized=optimized,
    )
    return [_fragmented_array_to_ir(fragmented_array)]

  if transforms_attr is None:
    raise ValueError(f"Unsupported memory space: {orig_ref_ty.memory_space}")

  swizzle = swizzle_from_transforms_attr(transforms_attr)
  transforms = memref_transforms_from_transforms_attr(transforms_attr)
  if transforms:
    [tiling_transform] = transforms
    assert isinstance(tiling_transform, lc.TileTransform)

    fragmented_array = fa.FragmentedArray.load_tiled(
        transformed_ref,
        swizzle,
        is_signed=is_signed,
        layout=out_layout,
        optimized=optimized,
        tiling_rank=len(tiling_transform.tiling),
    )
  else:
    fragmented_array = fa.FragmentedArray.load_untiled(
        transformed_ref,
        layout=out_layout,
        swizzle=swizzle,
        is_signed=is_signed,
        optimized=optimized,
    )

  return [_fragmented_array_to_ir(fragmented_array)]


@_register_lowering(mgpu.MultimemLoadReduceOp)
def _multimem_load_reduce_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.MultimemLoadReduceOp
) -> Sequence[ir.Value]:
  [out_layout_attr] = inference_utils.out_layouts(op)
  out_layout = layouts_lib.from_layout_attr(out_layout_attr)
  # TODO(apaszke): DO NOT IGNORE TRANSFORMS.

  # pyrefly: ignore[missing-attribute]
  reduction = str(mgpu.MultimemLoadReductionType(op.reduction_type.value))
  if isinstance(op.source.type.element_type, ir.IntegerType):
    is_signed = reduction in ("smax", "smin")
    if reduction in ("smax", "smin", "umax", "umin"):
      reduction = reduction[1:]
  else:
    is_signed = None

  multimem_ref = utils.MultimemRef(op.source)
  fa_res = fa.FragmentedArray.load_reduce_untiled(
      multimem_ref,
      out_layout,  # pyrefly: ignore[bad-argument-type]
      reduction,  # pyrefly: ignore[bad-argument-type]
      is_signed=is_signed
  )

  return [fragmented_array_to_ir(fa_res, op.result.type)]


@_register_lowering(mgpu.VectorStoreOp)
def _vector_store_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.VectorStoreOp
) -> Sequence[ir.Value]:
  if op.atomic_type is not None:
    # pyrefly: ignore[missing-attribute]
    atomic: Any = str(mgpu.AtomicOpType(op.atomic_type.value))
  else:
    atomic = None

  [to_store_layout] = inference_utils.in_layouts(op)
  fragmented_array = _fragmented_array_from_ir(op.valueToStore, to_store_layout)

  if ctx.auto_barriers:
    utils.warpgroup_barrier()  # Make sure the reads have completed.

  ref = op.destination
  ref_type = ir.MemRefType(ref.type)
  optimized = op.optimized.value if op.optimized is not None else True

  if ref_type.memory_space is None:  # GMEM
    ref = utils.MultimemRef(ref) if op.multimem.value else ref
    fragmented_array.store_untiled(ref, optimized=optimized, atomic=atomic)
  elif ref_type.memory_space == utils.smem():
    transforms_attr = inference_utils.in_transforms(op)[0]
    swizzle = swizzle_from_transforms_attr(transforms_attr)
    transforms = memref_transforms_from_transforms_attr(transforms_attr)
    if transforms:
      unwrapped_ref = unwrap_transformed_memref(ref, transforms_attr)
      [tiling_transform] = transforms
      assert isinstance(tiling_transform, lc.TileTransform)
      fragmented_array.store_tiled(
          unwrapped_ref,
          swizzle,
          optimized,
          tiling_rank=len(tiling_transform.tiling),
          atomic=atomic,
      )
    else:
      fragmented_array.store_untiled(
          ref, optimized=optimized, swizzle=swizzle, atomic=atomic
      )
  else:
    raise ValueError(f"Unsupported memory space: {ref_type.memory_space}")

  if ctx.auto_barriers:
    utils.warpgroup_barrier()  # Make sure the writes have completed.

  return []


@_register_lowering(mgpu.AsyncStoreSmemOp)
def _async_store_smem_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.AsyncStoreSmemOp
) -> Sequence[ir.Value]:
  index = ir.IndexType.get()

  [to_store_layout] = inference_utils.in_layouts(op)
  value = _fragmented_array_from_ir(op.valueToStore, to_store_layout)
  layout = layouts_lib.from_layout_attr(to_store_layout)
  if not isinstance(layout, fa.TiledLayout):
    raise NotImplementedError(f"Expected TiledLayout, got {type(layout)}")

  ref = op.destination
  transforms_attr = inference_utils.in_transforms(op)[0]
  swizzle = swizzle_from_transforms_attr(transforms_attr)
  unwrapped_ref = unwrap_transformed_memref(ref, transforms_attr)
  tiling_transform, = memref_transforms_from_transforms_attr(transforms_attr)
  assert isinstance(tiling_transform, lc.TileTransform)

  dialect_barrier = utils.DialectBarrierRef.from_barrier_memref(op.barrier)
  barrier_ref = dialect_barrier.barrier_ref

  cluster_dim = gpu.Dimension(op.cluster_dim.value)  # pyrefly: ignore[missing-attribute]
  cluster_idx = arith.index_cast(index, op.cluster_idx)
  cluster_barrier_ref = barrier_ref.remap_to_cluster(cluster_dim, cluster_idx)

  transfer_bytes = math.prod(value.shape) * utils.bitwidth(value.mlir_dtype)
  assert transfer_bytes % 8 == 0
  transfer_bytes //= 8
  if transfer_bytes % utils.WARPGROUP_SIZE:
    raise NotImplementedError(
        f"Transfer of {transfer_bytes} bytes is not divisible by "
        f"{utils.WARPGROUP_SIZE}"
    )

  if (orders_tc := dialect_barrier.orders_tensor_core):
    tx_count = transfer_bytes
    predicate = ctx.single_lane_predicate
  else:
    tx_count = transfer_bytes // utils.WARPGROUP_SIZE
    predicate = None

  cluster_barrier_ref.arrive_expect_tx(
      tx_count, predicate=predicate,
      tensor_core_order_scope=ctx.thread_semantics if orders_tc else None,
  )

  atomic = None
  if op.atomic_type is not None:
    atomic = str(mgpu.AtomicOpType(op.atomic_type.value))  # pyrefly: ignore[missing-attribute]

  optimized = op.optimized.value if op.optimized is not None else True
  value.store_tiled_async(
      unwrapped_ref,
      barrier_ref,
      cluster_dim=cluster_dim,
      cluster_idx=cluster_idx,
      swizzle=swizzle.value if swizzle != mgpu.SwizzlingMode.kNoSwizzle else None,
      optimized=optimized,
      tiling_rank=len(tiling_transform.tiling),
      atomic=atomic,  # pyrefly: ignore[bad-argument-type]
  )
  return []


@_register_lowering(mgpu.DebugPrintOp)
def _debug_print_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.DebugPrintOp
) -> Sequence[ir.Value]:
  del ctx
  [layout] = inference_utils.in_layouts(op)
  a = _fragmented_array_from_ir(op.value, layout)
  a.debug_print(op.format.value)
  return []


def pprint_layout(v: fa.FragmentedArray | tcgen05.TMEMRef) -> str:
  if isinstance(v, fa.FragmentedArray):
    return layouts_lib.pprint_layout(v.layout)
  assert isinstance(v, tcgen05.TMEMRef), v
  if v.layout == tcgen05.tmem_default_layout(packing=v.packing):
    return f"TMEM_DEFAULT(packing={v.packing})"
  return str(v.layout)


def pprint_transforms(
    transforms: Sequence[lc.MemRefTransform],
    swizzle: int | mgpu.SwizzlingMode | None = None,
) -> str:
  parts = []
  for t in transforms:
    if isinstance(t, lc.TileTransform):
      parts.append(f"TilingTransform(tiling={t.tiling})")  # skip rounding
    else:
      parts.append(str(t))
  if swizzle is not None and swizzle != mgpu.SwizzlingMode.kNoSwizzle:
    if isinstance(swizzle, mgpu.SwizzlingMode):
      swizzle = swizzle.value
    parts.append(f"SwizzleTransform(swizzle={swizzle})")
  return f"({', '.join(parts)})"


@_register_lowering(mgpu.PrintLayoutOp)
def _print_layout_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.PrintLayoutOp
) -> Sequence[ir.Value]:
  del ctx
  if isinstance(op.value.type, ir.VectorType):
    (layout,) = inference_utils.in_layouts(op)
    a = _fragmented_array_from_ir(op.value, layout)
    print(op.format.value.format(pprint_layout(a)))
  elif utils.is_smem_ref(op.value):
    [transforms_attr] = inference_utils.in_transforms(op)
    swizzle = swizzle_from_transforms_attr(transforms_attr)
    transforms = memref_transforms_from_transforms_attr(transforms_attr)
    print(op.format.value.format(pprint_transforms(transforms, swizzle)))
  else:
    (layout,) = inference_utils.in_tmem_layouts(op)
    ref = _tmem_ref_from_ir(op.value, layout)
    print(op.format.value.format(pprint_layout(ref)))
  return []


@_register_lowering(mgpu.BroadcastedIotaOp)
def _broadcasted_iota_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.BroadcastedIotaOp
) -> Sequence[ir.Value]:
  del ctx
  [layout] = inference_utils.out_layouts(op)
  result_type = ir.VectorType(op.result.type)
  a = fa.FragmentedArray.broadcasted_iota(
      result_type.element_type,
      tuple(result_type.shape),
      op.dimension.value,
      layouts_lib.from_layout_attr(layout),
      is_signed=_default_is_signed(result_type.element_type),
  )
  return [fragmented_array_to_ir(a, result_type)]


@_register_lowering(vector.BroadcastOp)
def _vector_broadcast_op_lowering_rule(
    _: LoweringContext, op: vector.BroadcastOp
) -> Sequence[ir.Value]:
  out_vec_ty = ir.VectorType(op.vector.type)
  fragmented_array = fa.FragmentedArray.splat(
      op.source,
      tuple(out_vec_ty.shape),
      layouts_lib.from_layout_attr(inference_utils.out_layouts(op)[0]),
      is_signed=_default_is_signed(out_vec_ty.element_type),
  )
  return [fragmented_array_to_ir(fragmented_array, out_vec_ty)]


@_register_lowering(vector.ShapeCastOp)
def _vector_shape_cast_op_lowering_rule(
    _: LoweringContext, op: vector.ShapeCastOp
) -> Sequence[ir.Value]:
  [layout] = inference_utils.in_layouts(op)
  out_vec_ty = ir.VectorType(op.result.type)
  assert out_vec_ty.has_static_shape
  a = _fragmented_array_from_ir(op.source, layout)
  return [
      fragmented_array_to_ir(a.reshape(tuple(out_vec_ty.shape)), out_vec_ty)
  ]


@_register_lowering(vector.ExtractStridedSliceOp)
def _vector_extract_strided_slice_op_lowering_rule(
    ctx: LoweringContext, op: vector.ExtractStridedSliceOp
) -> Sequence[ir.Value]:
  del ctx
  if any(ir.IntegerAttr(s).value != 1 for s in op.strides):
    raise NotImplementedError("`strides` must contain only 1s.")
  [in_layout] = inference_utils.in_layouts(op)
  [out_layout] = inference_utils.out_layouts(op)
  assert in_layout == out_layout
  out_vec_ty = ir.VectorType(op.result.type)
  assert out_vec_ty.has_static_shape
  a = _fragmented_array_from_ir(op.source, in_layout)
  indices = tuple(
      utils.DynamicSlice(
          ir.IntegerAttr(offset).value, ir.IntegerAttr(length).value
      )
      for offset, length in zip(op.offsets, op.sizes, strict=True)
  )
  result = a[indices]
  assert result.layout == layouts_lib.from_layout_attr(out_layout)
  return [fragmented_array_to_ir(result, out_vec_ty)]


@_register_lowering(vector.ExtractOp)
def _vector_extract_op_lowering_rule(
    ctx: LoweringContext, op: vector.ExtractOp
) -> Sequence[ir.Value]:
  del ctx
  if op.dynamic_position:
    raise NotImplementedError("Only slicing with static indices allowed.")

  [in_layout] = inference_utils.in_layouts(op)
  a = _fragmented_array_from_ir(op.source, in_layout)

  if not isinstance(op.result.type, ir.VectorType):  # scalar result
    result = a[tuple(op.static_position)]
    assert isinstance(result.layout, fa.WGSplatFragLayout)
    return [result.registers.item()]

  [out_layout] = inference_utils.out_layouts(op)
  assert in_layout == out_layout
  a = _fragmented_array_from_ir(op.source, in_layout)
  result_type = ir.VectorType(op.result.type)
  result = a[tuple(op.static_position)]
  assert result.layout == layouts_lib.from_layout_attr(out_layout)
  return [fragmented_array_to_ir(result, result_type)]


@_register_lowering(vector.InsertStridedSliceOp)
def _vector_insert_strided_slice_op_lowering_rule(
    ctx: LoweringContext, op: vector.InsertStridedSliceOp
) -> Sequence[ir.Value]:
  del ctx
  if any(ir.IntegerAttr(s).value != 1 for s in op.strides):
    raise NotImplementedError("`strides` must contain only 1s.")
  [src_layout, dst_layout] = inference_utils.in_layouts(op)
  [out_layout] = inference_utils.out_layouts(op)
  assert src_layout == dst_layout == out_layout
  src = _fragmented_array_from_ir(op.valueToStore, src_layout)
  indices = tuple(
      slice(ir.IntegerAttr(offset).value, ir.IntegerAttr(offset).value + size)
      for offset, size in zip(op.offsets, src.shape, strict=True)
  )
  result = _fragmented_array_from_ir(op.dest, dst_layout)
  result[indices] = src
  assert result.layout == layouts_lib.from_layout_attr(out_layout)
  return [fragmented_array_to_ir(result, op.result.type)]


def _combining_kind(attr: ir.Attribute) -> vector.CombiningKind:
  return vector.CombiningKind[
      str(attr).removeprefix("#vector.kind<").removesuffix(">").upper()
  ]


def _is_reduction_signed(kind: vector.CombiningKind) -> bool | None:
  if kind in (vector.CombiningKind.MAXSI, vector.CombiningKind.MINSI):
    return True
  if kind in (vector.CombiningKind.MAXUI, vector.CombiningKind.MINUI):
    return False
  return None


@_register_lowering(vector.ReductionOp)
def _vector_reduction_op_lowering_rule(
    ctx: LoweringContext, op: vector.ReductionOp
) -> Sequence[ir.Value]:
  [layout] = inference_utils.in_layouts(op)
  element_type = op.vector.type.element_type
  scratch = _slice_smem(
      ir.MemRefType.get([4], element_type, memory_space=utils.smem()),
      ir.IntegerAttr(op.attributes["offset"]).value,
      ctx.smem_requested_bytes,
  )
  axes = range(op.vector.type.rank)
  op_kind = _combining_kind(op.kind)
  is_signed = _is_reduction_signed(op_kind)
  a = _fragmented_array_from_ir(op.vector, layout, is_signed)
  acc_ilp_attr = op.attributes.get("acc_ilp")
  acc_ilp = None if acc_ilp_attr is None else ir.IntegerAttr(acc_ilp_attr).value
  match op_kind:
    case vector.CombiningKind.ADD:
      result = a.reduce("add", axes, scratch, acc_ilp=acc_ilp)
    case vector.CombiningKind.MUL:
      result = a.reduce("prod", axes, scratch, acc_ilp=acc_ilp)
    case vector.CombiningKind.MAXSI | vector.CombiningKind.MAXUI | vector.CombiningKind.MAXIMUMF:
      result = a.reduce("max", axes, scratch, acc_ilp=acc_ilp)
    case vector.CombiningKind.MINUI | vector.CombiningKind.MINSI | vector.CombiningKind.MINIMUMF:
      result = a.reduce("min", axes, scratch, acc_ilp=acc_ilp)
    case _:
      raise NotImplementedError(f"Unsupported reduction kind: {op.kind}")
  assert isinstance(result.layout, fa.WGSplatFragLayout)
  return [result.registers.item()]


@_register_lowering(vector.MultiDimReductionOp)
def _vector_multi_dim_reduction_op_lowering_rule(
    ctx: LoweringContext, op: vector.MultiDimReductionOp
) -> Sequence[ir.Value]:
  [in_layout, acc_layout] = inference_utils.in_layouts(op)
  [out_layout] = inference_utils.out_layouts(op)
  if out_layout != acc_layout:
    raise ValueError(
        f"Output layout {out_layout} must match the accumulator layout"
        f" {acc_layout}"
    )

  if len(op.reduction_dims) != 1:
    raise NotImplementedError("Only 1 reduction dimension is supported.")

  op_kind = _combining_kind(op.kind)
  is_signed = _is_reduction_signed(op_kind)
  src = _fragmented_array_from_ir(op.source, in_layout, is_signed)
  acc = _fragmented_array_from_ir(op.acc, acc_layout, is_signed)

  if not isinstance(src.layout, fa.TiledLayout):
    raise NotImplementedError(f"Unsupported layout: {src.layout}")
  reduced_dim = src.layout.tiling.tile_dimension(op.reduction_dims[0])
  if any(reduced_dim[d] for d in src.layout.partitioned_warp_dims):
    # cross-warp reductions require scratch space.
    dtype = op.source.type.element_type
    allocation_size = ir.IntegerAttr(op.attributes["scratch_size"]).value * 8 // utils.bitwidth(dtype)
    scratch = _slice_smem(
        ir.MemRefType.get([allocation_size], dtype, memory_space=utils.smem()),
        ir.IntegerAttr(op.attributes["offset"]).value,
        ctx.smem_requested_bytes,
    )
  else:
    scratch = None

  acc_ilp_attr = op.attributes.get("acc_ilp")
  acc_ilp = None if acc_ilp_attr is None else ir.IntegerAttr(acc_ilp_attr).value
  reduce_src = functools.partial(
      src.reduce,
      axis=op.reduction_dims[0],
      scratch=scratch,
      acc_ilp=acc_ilp,
  )
  match op_kind:
    case vector.CombiningKind.ADD:
      result = reduce_src("add") + acc
    case vector.CombiningKind.MUL:
      result = reduce_src("prod") * acc
    case vector.CombiningKind.MAXSI | vector.CombiningKind.MAXUI | vector.CombiningKind.MAXIMUMF:
      result = reduce_src("max").max(acc)
    case vector.CombiningKind.MINUI | vector.CombiningKind.MINSI | vector.CombiningKind.MINIMUMF:
      result = reduce_src("min").min(acc)
    case _:
      raise NotImplementedError(f"Unsupported reduction kind: {op.kind}")
  assert result.layout == layouts_lib.from_layout_attr(out_layout)
  return [fragmented_array_to_ir(result, op.result.type)]


@_register_lowering(mgpu.LayoutCastOp)
def _mgpu_layout_cast_op_lowering_rule(
    _: LoweringContext, op: mgpu.LayoutCastOp
) -> Sequence[ir.Value]:
  [in_layout] = inference_utils.in_layouts(op)
  [out_layout] = inference_utils.out_layouts(op)
  in_array = _fragmented_array_from_ir(op.x, in_layout)
  out_array = in_array.to_layout(layouts_lib.from_layout_attr(out_layout))
  return [fragmented_array_to_ir(out_array, op.result.type)]


@_register_lowering(mgpu.BroadcastInDimOp)
def _mgpu_broadcast_in_dim_op_lowering_rule(
    _: LoweringContext, op: mgpu.BroadcastInDimOp
) -> Sequence[ir.Value]:
  out_ty = ir.VectorType(op.result.type)
  broadcast_dims = tuple(op.broadcast_dimensions)
  in_layout_attr = inference_utils.in_layouts(op)[0]
  operand_fa = _fragmented_array_from_ir(op.operand, in_layout_attr)
  out_layout = layouts_lib.from_layout_attr(inference_utils.out_layouts(op)[0])
  out = operand_fa.broadcast_in_dim(
      tuple(out_ty.shape), broadcast_dims, out_layout
  )
  return [fragmented_array_to_ir(out, out_ty)]


@_register_lowering(mgpu.VectorConcatOp)
def _mgpu_vector_concat_op_lowering_rule(
    _: LoweringContext, op: mgpu.VectorConcatOp
) -> Sequence[ir.Value]:
  in_layouts = inference_utils.in_layouts(op)
  out_layout, = inference_utils.out_layouts(op)
  operands_fa = [
      _fragmented_array_from_ir(opr, l)
      for opr, l in zip(op.operands, in_layouts, strict=True)
  ]
  out = fa.concatenate(operands_fa, axis=op.dimension.value)
  assert out.layout == layouts_lib.from_layout_attr(out_layout)
  return [fragmented_array_to_ir(out, op.result.type)]


def swizzle_from_transforms_attr(attr: ir.ArrayAttr) -> mgpu.SwizzlingMode:
  swizzle = None
  for transform in attr:
    if isinstance(transform, mgpu.SwizzleTransformAttr):
      if swizzle is not None:
        raise ValueError("Found multiple SwizzleTransformAttr")
      swizzle = mgpu.SwizzlingMode(mgpu.SwizzleTransformAttr(transform).swizzle)
  return swizzle or mgpu.SwizzlingMode.kNoSwizzle


def memref_transforms_from_transforms_attr(
    attr: ir.ArrayAttr,
) -> tuple[lc.MemRefTransform, ...]:
  gmem_transforms: list[lc.MemRefTransform] = []
  for transform in attr:
    if isinstance(transform, mgpu.TileTransformAttr):
      tile_transform = lc.TileTransform(tuple(transform.tiling))
      gmem_transforms.append(tile_transform)
    elif not isinstance(transform, mgpu.SwizzleTransformAttr):
      raise NotImplementedError(f"Unsupported transform: {transform}")
  return tuple(gmem_transforms)


def tile_offset(
    offsets: tuple[int, ...], tiling: tuple[int, ...]
) -> tuple[int, ...]:
  """Tiles the trailing offsets in `offsets` according to `tiling`.

  Raises if the offsets are not aligned with the start of a tile.
  """
  if len(offsets) < len(tiling):
    raise ValueError(f"Offsets {offsets} have lower rank than tiling {tiling}")
  untiled_offsets, tiled_offsets = (
      offsets[: -len(tiling)],
      offsets[-len(tiling) :],
  )
  for i, t in zip(tiled_offsets, tiling, strict=True):
    if i % t != 0:
      raise ValueError(f"Offset {i} is not divisible by tile size {t}")
  return (
      *untiled_offsets,
      *[i // t for i, t in zip(tiled_offsets, tiling, strict=True)],
      *[0] * len(tiling),
  )


def tile_strides(
    strides: tuple[int, ...], tiling: tuple[int, ...]
) -> tuple[int, ...]:
  """Tiles the trailing strides in `strides` according to `tiling`.

  The `len(tiling)` trailing strides in `strides` must be the `len(tiling)`
  smallest strides in `strides`. The same property holds in the result, i.e.,
  given two tiles with indices i and j (i < j) with strides tiled according to
  this function, then all the elements in tile i are physically ordered before
  all the elements in tile j.

  E.g., tile_strides((2048, 32, 1), (8, 4)) = (2048, 256, 32, 4, 1)
  """
  if len(strides) < len(tiling):
    raise ValueError(f"Strides {strides} have lower rank than tiling {tiling}")
  ordered_strides = sorted(strides, reverse=True)
  if set(ordered_strides[-len(tiling):]) != set(strides[-len(tiling):]):
    raise ValueError(
        "Can not tile strides when tiled dimensions have been transposed with "
        f"untiled dimensions. Strides: {strides}, tiling: {tiling}"
    )
  untiled_strides, tiled_strides = strides[:-len(tiling)], strides[-len(tiling):]

  # Zip the strides and tiling together, in order to sort them together. This
  # allows handling cases where multiple tiling dimensions have the same stride,
  # which can occur with size-1 dimensions.
  tiled_strides_and_tiling: list[tuple[int, int]] = list(
      zip(tiled_strides, tiling, strict=True))
  tiled_ordered_strides_and_tiling = sorted(
      tiled_strides_and_tiling, reverse=True)

  to_ordered = lambda i: tiled_ordered_strides_and_tiling.index(tiled_strides_and_tiling[i])
  from_ordered = lambda i: tiled_strides_and_tiling.index(tiled_ordered_strides_and_tiling[i])

  ordered_tiling = [tiling[from_ordered(i)] for i in range(len(tiling))]
  ordered_tiled_strides = [tiled_strides[from_ordered(i)] for i in range(len(tiling))]

  ordered_tiled_tiling_strides = [1]
  for t in reversed(ordered_tiling):
    ordered_tiled_tiling_strides.append(ordered_tiled_tiling_strides[-1] * t)

  prev_s = ordered_tiled_strides[-1]
  for s, t in zip(ordered_tiled_strides[:-1][::-1], ordered_tiling[1:][::-1], strict=True):
    d = prev_s * t
    prev_s = s
    if s % d != 0:
      raise ValueError(
          f"Stride {s} is not divisible by {d} (tile size = {t}). "
          f"Strides: {strides}, tiling: {tiling}"
      )
    ordered_tiled_tiling_strides.append(s // d * ordered_tiled_tiling_strides[-1])

  ordered_tiled_tiling_strides.reverse()

  return (
      *untiled_strides,
      *[ordered_tiled_tiling_strides[to_ordered(i)] for i in range(len(tiling))],
      *[ordered_tiled_tiling_strides[len(tiling) + to_ordered(i)] for i in range(len(tiling))]
  )


def transform_type(
    ref_ty: ir.MemRefType,
    transforms: tuple[lc.MemRefTransform, ...] | ir.ArrayAttr,
) -> ir.MemRefType:
  if isinstance(transforms, ir.ArrayAttr):
    transforms = memref_transforms_from_transforms_attr(transforms)
  if not (utils.is_smem_ref(ref_ty) or utils.is_cluster_smem_ref(ref_ty)):
    raise ValueError(f"Only workgroup memory is supported but got {ref_ty}.")
  if utils.is_cluster_smem_ref(ref_ty):
    i32 = ir.IntegerType.get_signless(32)
    ref_ty = ir.MemRefType.get(
        ref_ty.shape,
        ref_ty.element_type,
        ref_ty.layout,
        memory_space=ir.IntegerAttr.get(i32, 7),
    )

  if not transforms:
    return ref_ty

  # TODO(bchetioui): this should be trivial to relax if ever necessary.
  if len(transforms) > 1 or not isinstance(transforms[0], lc.TileTransform):
    raise NotImplementedError(f"Unsupported transforms: {transforms}")
  tile_transform: lc.TileTransform = transforms[0]

  strides, offset = ref_ty.get_strides_and_offset()
  tiled_shape = tile_transform.transform_shape(ref_ty.shape)
  tiled_strides = tile_strides(tuple(strides), tile_transform.tiling)

  if offset == ir.ShapedType.get_dynamic_stride_or_offset():
    tiled_offset = offset
  else:
    delinearized_offset = [0] * len(strides)
    for i, stride in sorted(enumerate(strides), key=lambda es: es[1], reverse=True):
      delinearized_offset[i] = offset // stride
      offset %= stride
    tiled_delinearized_offset = tile_offset(
        tuple(delinearized_offset), tile_transform.tiling
    )
    tiled_offset = sum(o * s for o, s in zip(tiled_delinearized_offset, tiled_strides, strict=True))

  if isinstance(ref_ty.layout, ir.StridedLayoutAttr):
    layout = ir.StridedLayoutAttr.get(tiled_offset, tiled_strides)
  else:
    layout = None

  return ir.MemRefType.get(
      tiled_shape,
      ref_ty.element_type,
      memory_space=ref_ty.memory_space,
      layout=layout
  )


def _gmem_slice_and_predicate(
    ctx: LoweringContext,
    op: mgpu.AsyncLoadOp | mgpu.AsyncPrefetchOp | mgpu.AsyncStoreOp,
) -> tuple[
    tuple[ir.Value | fa.FragmentedArray | utils.DynamicSlice, ...],
    dict[str, ir.Value],
]:
  """Returns the GMEM slice and predicate for the given async op."""
  gmem_slice: list[ir.Value | fa.FragmentedArray | utils.DynamicSlice]
  gmem_slice = []
  predicate = ctx.single_lane_predicate
  if op.predicate is not None:
    predicate = arith.andi(predicate, op.predicate)
  predicate = dict(predicate=predicate)
  for idx, size in zip(op.indices, op.slice_lengths, strict=True):
    if isinstance(idx.type, ir.IntegerType):
      idx_int = arith.index_cast(ir.IndexType.get(), idx)
      v = idx_int if size < 0 else utils.DynamicSlice(idx_int, size)
      gmem_slice.append(v)
    elif isinstance(idx.type, ir.VectorType):
      layout = inference_utils.in_layouts(op)[0]
      assert layouts_lib.from_layout_attr(layout) in (fa.TMA_INDICES_LAYOUT, fa.TMA_INDICES_4_LAYOUT), layout
      idx_fa = _fragmented_array_from_ir(idx, layout)
      gmem_slice.append(idx_fa)
      predicate = dict()
    else:
      raise TypeError(f"Unsupported index type: {idx.type}")
  return tuple(gmem_slice), predicate


@_register_lowering(mgpu.AsyncLoadOp, support_warp_semantics=True)
def _mgpu_async_load_op_lowering_rule(
    ctx: LoweringContext, load_op: mgpu.AsyncLoadOp
) -> Sequence[ir.Value]:
  if is_cp_async := load_op.barrier is None:
    barrier = None
  else:
    barrier = utils.DialectBarrierRef.from_barrier_memref(load_op.barrier)

  [transforms_attr] = inference_utils.in_transforms(load_op)
  swizzle = swizzle_from_transforms_attr(transforms_attr)
  transforms = memref_transforms_from_transforms_attr(transforms_attr)

  unwrapped_dst = unwrap_transformed_memref(
      load_op.destination, transforms_attr
  )
  assert isinstance(unwrapped_dst.type, ir.MemRefType)
  if utils.is_memref_transposed(unwrapped_dst.type):
    strides, _ = ir.MemRefType(unwrapped_dst.type).get_strides_and_offset()
    permutation = tuple(
        sorted(range(len(strides)), key=lambda i: strides[i], reverse=True)
    )
    # We undo the tranpose and apply it as a transform.
    unwrapped_dst = utils.memref_transpose(
        unwrapped_dst, permutation
    )
    transforms = (*transforms, lc.TransposeTransform(permutation))

  gmem_slice, predicate = _gmem_slice_and_predicate(ctx, load_op)

  collective = tuple(
      gpu.Dimension(ir.IntegerAttr(axis).value)
      for axis in load_op.collective or []
  )

  match load_op.leader_tracked:
    case mgpu.CopyReplicatedAttr():
      leader_tracked = lc.CopyPartition.REPLICATED
    case mgpu.CopyPartitionedAttr() as attr:
      leader_tracked = lc.CopyPartition.PARTITIONED(attr.axis)
    case _:
      leader_tracked = None

  # TODO(dasenov): async_copy requires all GMEM strides except the last one
  # to be a multiple of 16 bytes. This restriction could be loosned with
  # strided layouts when they are contiguous in GMEM. In that case, we could do:
  # flatten -> async_copy -> unflatted here, as long as flattened size is a
  # multiple of 16.

  if ctx.auto_barriers and ctx.thread_semantics == utils.ThreadSubset.WARPGROUP:
    utils.warpgroup_barrier()  # Make sure the writes have completed.

  # TODO(dasenov): Add support for the remaining op properties.
  oob_mode = lc.OOBFillMode(ir.IntegerAttr(load_op.oob_fill_mode).value)
  ctx.launch_context.async_copy(
      src_ref=load_op.source,
      dst_ref=unwrapped_dst,
      gmem_slice=gmem_slice,
      barrier=barrier.barrier_ref if barrier is not None else None,
      collective=collective,
      arrive=False,
      swizzle=swizzle,
      gmem_transform=transforms,
      leader_tracked=leader_tracked,
      oob_mode=oob_mode,
      implementation=(
          lc.AsyncCopyImplementation.CP_ASYNC
          if is_cp_async
          else lc.AsyncCopyImplementation.TMA
      ),
      gmem_peer_id=load_op.gmem_peer_id,
      **{} if is_cp_async else predicate,  # pyrefly: ignore[bad-argument-type]
  )
  return []


@_register_lowering(mgpu.AsyncPrefetchOp, support_warp_semantics=True)
def _mgpu_async_prefetch_op_lowering_rule(
    ctx: LoweringContext, load_op: mgpu.AsyncPrefetchOp
) -> Sequence[ir.Value]:
  gmem_slice, predicate = _gmem_slice_and_predicate(ctx, load_op)

  if load_op.collective:
    raise NotImplementedError("Collective prefetches are not supported yet.")

  ctx.launch_context.async_prefetch(
      gmem_ref=load_op.source,
      gmem_slice=gmem_slice,
      swizzle=None,
      gmem_transform=(),
      **predicate,  # pyrefly: ignore[bad-argument-type]
  )
  return []


@_register_lowering(mgpu.AsyncStoreOp, support_warp_semantics=True)
def _mgpu_async_store_op_lowering_rule(
    ctx: LoweringContext, store_op: mgpu.AsyncStoreOp
) -> Sequence[ir.Value]:
  [transforms_attr] = inference_utils.in_transforms(store_op)
  swizzle = swizzle_from_transforms_attr(transforms_attr)
  transforms = memref_transforms_from_transforms_attr(transforms_attr)
  unwrapped_source = unwrap_transformed_memref(store_op.source, transforms_attr)
  assert isinstance(unwrapped_source.type, ir.MemRefType)
  if utils.is_memref_transposed(unwrapped_source.type):
    strides, _ = ir.MemRefType(unwrapped_source.type).get_strides_and_offset()
    permutation = tuple(
        sorted(range(len(strides)), key=lambda i: strides[i], reverse=True)
    )
    # We undo the tranpose and apply it as a transform.
    unwrapped_source = utils.memref_transpose(
        unwrapped_source, permutation
    )
    transforms = (*transforms, lc.TransposeTransform(permutation))

  gmem_slice, predicate = _gmem_slice_and_predicate(ctx, store_op)

  arrive = (
      None if store_op.commit_group is None else bool(store_op.commit_group)
  )

  # TODO(dasenov): async_copy requires all GMEM strides except the last one
  # to be a multiple of 16 bytes. This restriction could be loosned with
  # strided layouts when they are contiguous in GMEM. In that case, we could do:
  # flatten -> async_copy -> unflatted here, as long as flattened size is a
  # multiple of 16.
  if store_op.reduction_op is not None:
    # pyrefly: ignore[missing-attribute]
    reduction_op = mgpu.TMAReduction(store_op.reduction_op.value).name.lower()
  else:
    reduction_op = None

  peer_id: ir.Value | lc.GlobalBroadcast | None = None
  if store_op.is_global_broadcast.value:
    peer_id = lc.GLOBAL_BROADCAST
  elif store_op.gmem_peer_id is not None:
    peer_id = store_op.gmem_peer_id

  # TODO(dasenov): Add support for the remaining op properties.
  ctx.launch_context.async_copy(
      src_ref=unwrapped_source,
      dst_ref=store_op.destination,
      gmem_slice=gmem_slice,
      swizzle=swizzle,
      gmem_transform=transforms,
      **predicate,  # pyrefly: ignore[bad-argument-type]
      arrive=arrive,
      reduction_op=reduction_op,  # pyrefly: ignore[bad-argument-type]
      gmem_peer_id=peer_id,
      oob_mode=lc.OOBFillMode.UNDEFINED,
  )
  return []


@_register_lowering(
    mgpu.AsyncStoreScalesSmemToTmemOp, support_warp_semantics=True
)
def _async_copy_scales_smem_to_tmem_lowering_rule(
    ctx: LoweringContext, op: mgpu.AsyncStoreScalesSmemToTmemOp
) -> Sequence[ir.Value]:
  ctx.check_collective(op)
  [in_layout_attr] = inference_utils.in_tmem_layouts(op)
  tmem_ref = _tmem_ref_from_ir(op.destination, in_layout_attr)
  smem_ref = unwrap_transformed_memref(op.source, ir.ArrayAttr.get([]))
  with utils.when(ctx.single_lane_predicate):
    tcgen05.async_copy_scales_smem_to_tmem(
      smem_ref, tmem_ref, collective=bool(op.collective)
    )
  return []


@_register_lowering(
    mgpu.AsyncStoreSparseMetadataSmemToTmemOp, support_warp_semantics=True
)
def _async_copy_sparse_metadata_smem_to_tmem_lowering_rule(
    ctx: LoweringContext, op: mgpu.AsyncStoreSparseMetadataSmemToTmemOp
) -> Sequence[ir.Value]:
  ctx.check_collective(op)
  [in_layout_attr] = inference_utils.in_tmem_layouts(op)
  tmem_ref = _tmem_ref_from_ir(op.destination, in_layout_attr)
  smem_ref = unwrap_transformed_memref(op.source, ir.ArrayAttr.get([]))
  with utils.when(ctx.single_lane_predicate):
    tcgen05.async_copy_sparse_metadata_smem_to_tmem(
      smem_ref, tmem_ref, collective=bool(op.collective)
    )
  return []


@_register_lowering(mgpu.AsyncStoreSmemToTmemOp, support_warp_semantics=True)
def _async_store_smem_to_tmem_lowering_rule(
    ctx: LoweringContext, op: mgpu.AsyncStoreSmemToTmemOp
) -> Sequence[ir.Value]:
  ctx.check_collective(op)
  [transforms_attr] = inference_utils.in_transforms(op)
  swizzle = swizzle_from_transforms_attr(transforms_attr)
  smem_ref = unwrap_transformed_memref(op.source, transforms_attr)
  smem_ref_ty = op.source.type
  assert isinstance(smem_ref_ty, ir.MemRefType)

  [in_layout_attr] = inference_utils.in_tmem_layouts(op)
  tmem_ref = _tmem_ref_from_ir(op.destination, in_layout_attr)
  with utils.when(ctx.single_lane_predicate):
    tcgen05.async_copy_smem_to_tmem(
        smem_ref, tmem_ref, swizzle, collective=bool(op.collective)
    )
  return []


@_register_lowering(mgpu.TmemLayoutCastOp, support_warp_semantics=True)
def _tmem_layout_cast_lowering_rule(
    ctx: LoweringContext,
    op: mgpu.TmemLayoutCastOp,
) -> Sequence[ir.Value]:
  del ctx
  in_layout = inference_utils.in_tmem_layouts(op)[0]
  tmem_ref = _tmem_ref_from_ir(op.ref, in_layout)
  # We can't relayout TMEM.
  assert layouts_lib.to_layout_attr(tmem_ref.layout) == op.new_layout
  return [op.ref]


@_register_lowering(mgpu.SliceTmemOp, support_warp_semantics=True)
def _slice_tmem_lowering_rule(
    ctx: LoweringContext, op: mgpu.SliceTmemOp
) -> Sequence[ir.Value]:
  del ctx
  in_layout_attr = inference_utils.in_tmem_layouts(op)[0]
  out_layout_attr = inference_utils.out_tmem_layouts(op)[0]
  source = _tmem_ref_from_ir(op.source, in_layout_attr)
  i32 = ir.IntegerType.get_signless(32)
  offset = arith.constant(i32, op.offset)
  dest_addr = arith.addi(source.address, offset)
  conversion_cast = builtin.UnrealizedConversionCastOp([op.result.type], [dest_addr])
  conversion_cast.attributes["layout"] = out_layout_attr
  return [conversion_cast.result]


def _conversion_op_lowering_rule(
    _: LoweringContext,
    op: ir.OpView,
    source_is_signed: bool | None,
    target_is_signed: bool | None,
) -> Sequence[ir.Value]:
  [in_layout] = inference_utils.in_layouts(op)
  [layout] = inference_utils.out_layouts(op)
  if in_layout != layout:
    raise ValueError("Layout mismatch")

  target_ty = op.result.type.element_type
  operand = _fragmented_array_from_ir(op.operands[0], layout, source_is_signed)
  converted = operand.astype(target_ty, is_signed=target_is_signed)
  return [fragmented_array_to_ir(converted, op.result.type)]


for _op, _source_is_signed, _target_is_signed in [
    (arith.ExtFOp, None, None),
    (arith.ExtSIOp, True, True),
    (arith.ExtUIOp, False, False),
    (arith.FPToSIOp, None, True),
    (arith.FPToUIOp, None, False),
    (arith.SIToFPOp, True, None),
    (arith.TruncFOp, None, None),
    (arith.TruncIOp, False, False),
    (arith.UIToFPOp, False, None),
]:
  _lowerings[_op.OPERATION_NAME] = functools.partial(
      _conversion_op_lowering_rule,
      source_is_signed=_source_is_signed,
      target_is_signed=_target_is_signed,
  )


def _unary_op_lowering_rule(
    _: LoweringContext,
    op: Any,
    impl: Callable[..., fa.FragmentedArray],
    is_signed: bool | None = None,
) -> Sequence[ir.Value]:
  in_layouts = inference_utils.in_layouts(op)
  [layout] = inference_utils.out_layouts(op)
  if any(in_layout != layout for in_layout in in_layouts):
    raise ValueError("Layout mismatch")
  a = _fragmented_array_from_ir(op.operand, layout, is_signed)
  if hasattr(op, "fastmath"):
    if op.fastmath == ir.Attribute.parse("#arith.fastmath<afn>"):
      result_fa = impl(a, approx=True)
    else:
      result_fa = impl(a)
  else:
    result_fa = impl(a)

  return [fragmented_array_to_ir(result_fa, op.result.type)]


for _op, _unary_impl, _is_signed in [
    (mlir_math.RsqrtOp, fa.FragmentedArray.rsqrt, None),
    (mlir_math.SqrtOp, fa.FragmentedArray.sqrt, None),
    (mlir_math.ExpOp, fa.FragmentedArray.exp, None),
    (mlir_math.Exp2Op, fa.FragmentedArray.exp2, None),
    (mlir_math.SinOp, fa.FragmentedArray.sin, None),
    (mlir_math.CosOp, fa.FragmentedArray.cos, None),
    (mlir_math.LogOp, fa.FragmentedArray.log, None),
    (mlir_math.Log2Op, fa.FragmentedArray.log2, None),
    (mlir_math.TanhOp, fa.FragmentedArray.tanh, None),
    (mlir_math.AbsFOp, fa.FragmentedArray.abs, None),
    (mlir_math.AbsIOp, fa.FragmentedArray.abs, True),
    (mlir_math.RoundOp, fa.FragmentedArray.round, None),
    (mlir_math.RoundEvenOp, fa.FragmentedArray.round_even, None),
    (mlir_math.ErfOp, fa.FragmentedArray.erf, None),
    (
        mlir_math.CountLeadingZerosOp,
        lambda x: x._pointwise(mlir_math.ctlz, restrict_bitwidth=False),
        None,
    ),
    (arith.NegFOp, operator.neg, None),
]:
  _lowerings[_op.OPERATION_NAME] = functools.partial(
      _unary_op_lowering_rule, impl=_unary_impl, is_signed=_is_signed
  )


def _binary_op_lowering_rule(
    _: LoweringContext,
    op: Any,
    is_signed: bool | None,
    impl: Callable[
        [fa.FragmentedArray, fa.FragmentedArray], fa.FragmentedArray
    ],
) -> Sequence[ir.Value]:
  in_layouts = inference_utils.in_layouts(op)
  [layout] = inference_utils.out_layouts(op)
  if any(in_layout != layout for in_layout in in_layouts):
    raise ValueError("Layout mismatch")
  lhs = _fragmented_array_from_ir(op.lhs, layout, is_signed)
  rhs = _fragmented_array_from_ir(op.rhs, layout, is_signed)
  return [fragmented_array_to_ir(impl(lhs, rhs), op.result.type)]


for _op, _binary_impl, _is_signed in [
    (arith.AddIOp, operator.add, False),
    (arith.AddFOp, operator.add, None),
    (arith.SubIOp, operator.sub, False),
    (arith.SubFOp, operator.sub, None),
    (arith.MulIOp, operator.mul, False),
    (arith.MulFOp, operator.mul, None),
    (arith.FloorDivSIOp, operator.floordiv, True),
    (arith.DivSIOp, lambda x, y: x.trunc_div(y), True),
    (arith.DivUIOp, lambda x, y: x.trunc_div(y), False),
    (arith.RemSIOp, operator.mod, True),
    (arith.RemUIOp, operator.mod, False),
    (arith.RemFOp, operator.mod, None),
    (arith.AndIOp, operator.and_, False),
    (arith.OrIOp, operator.or_, False),
    (arith.XOrIOp, operator.xor, False),
    (arith.MaxSIOp, fa.FragmentedArray.max, True),
    (arith.MaxUIOp, fa.FragmentedArray.max, False),
    (arith.MaximumFOp, fa.FragmentedArray.max, None),
    (arith.MinSIOp, fa.FragmentedArray.min, True),
    (arith.MinUIOp, fa.FragmentedArray.min, False),
    (arith.MinimumFOp, fa.FragmentedArray.min, None),
    (mlir_math.Atan2Op, fa.FragmentedArray.atan2, None),
    (mlir_math.CopySignOp, fa.FragmentedArray.copysign, None),
    (arith.ShLIOp, operator.lshift, False),
    (arith.ShRUIOp, operator.rshift, False),
    (arith.ShRSIOp, operator.rshift, True),
]:
  _lowerings[_op.OPERATION_NAME] = functools.partial(
      _binary_op_lowering_rule, impl=_binary_impl, is_signed=_is_signed
  )


def _divf_lowering_rule(
    ctx: LoweringContext, op: Any
) -> Sequence[ir.Value]:
  [layout] = inference_utils.out_layouts(op)
  lhs_ty = ir.VectorType(op.lhs.type)
  if isinstance(lhs_ty.element_type, ir.Float8E8M0FNUType):
    lhs_layout = layouts_lib.from_layout_attr(
        op.lhs.owner.attributes["layout"]
    )
    if isinstance(lhs_layout, fa.WGSplatFragLayout):
      [source] = op.lhs.owner.opview.operands
      if (
          isinstance(source.owner.opview, arith.ConstantOp)
          and float(ir.FloatAttr(source.owner.opview.value)) == 1.0
      ):
        rhs = _fragmented_array_from_ir(op.rhs, layout)
        return [fragmented_array_to_ir(1 / rhs, op.result.type)]
  return _binary_op_lowering_rule(
      ctx, op, is_signed=None, impl=operator.truediv
  )

_lowerings[arith.DivFOp.OPERATION_NAME] = _divf_lowering_rule


CMPI_IMPLS = {
    arith.CmpIPredicate.eq: (operator.eq, False),
    arith.CmpIPredicate.ne: (operator.ne, False),
    arith.CmpIPredicate.slt: (operator.lt, True),
    arith.CmpIPredicate.sle: (operator.le, True),
    arith.CmpIPredicate.sgt: (operator.gt, True),
    arith.CmpIPredicate.sge: (operator.ge, True),
    arith.CmpIPredicate.ult: (operator.lt, False),
    arith.CmpIPredicate.ule: (operator.le, False),
    arith.CmpIPredicate.ugt: (operator.gt, False),
    arith.CmpIPredicate.uge: (operator.ge, False),
}


@_register_lowering(arith.CmpIOp)
def _cmpi_op_lowering_rule(
    _: LoweringContext, op: arith.CmpIOp
) -> Sequence[ir.Value]:
  in_layouts = inference_utils.in_layouts(op)
  [layout] = inference_utils.out_layouts(op)
  if any(in_layout != layout for in_layout in in_layouts):
    raise ValueError("Layout mismatch")
  # pyrefly: ignore[missing-attribute]
  impl, is_signed = CMPI_IMPLS[op.predicate.value]
  lhs = _fragmented_array_from_ir(op.lhs, layout, is_signed)
  rhs = _fragmented_array_from_ir(op.rhs, layout, is_signed)
  return [fragmented_array_to_ir(impl(lhs, rhs), op.result.type)]


CMPF_IMPLS = {
    arith.CmpFPredicate.OEQ: operator.eq,
    arith.CmpFPredicate.UNE: operator.ne,
    arith.CmpFPredicate.OLT: operator.lt,
    arith.CmpFPredicate.OLE: operator.le,
    arith.CmpFPredicate.OGT: operator.gt,
    arith.CmpFPredicate.OGE: operator.ge,
}


@_register_lowering(arith.CmpFOp)
def _cmpf_op_lowering_rule(
    _: LoweringContext, op: arith.CmpFOp
) -> Sequence[ir.Value]:
  in_layouts = inference_utils.in_layouts(op)
  [layout] = inference_utils.out_layouts(op)
  if any(in_layout != layout for in_layout in in_layouts):
    raise ValueError("Layout mismatch")
  # pyrefly: ignore[missing-attribute]
  impl = CMPF_IMPLS[op.predicate.value]
  lhs = _fragmented_array_from_ir(op.lhs, layout)
  rhs = _fragmented_array_from_ir(op.rhs, layout)
  return [fragmented_array_to_ir(impl(lhs, rhs), op.result.type)]


@_register_lowering(arith.BitcastOp)
def _bitcast_op_lowering_rule(
    _: LoweringContext, op: arith.BitcastOp
) -> Sequence[ir.Value]:
  in_layouts = inference_utils.in_layouts(op)
  [layout] = inference_utils.out_layouts(op)
  if any(in_layout != layout for in_layout in in_layouts):
    raise ValueError("Layout mismatch")
  in_ = _fragmented_array_from_ir(op.in_, layout)
  out_element_type = ir.VectorType(op.result.type).element_type
  out = in_.bitcast(
      out_element_type,
      output_is_signed=_default_is_signed(out_element_type),
  )
  return [fragmented_array_to_ir(out, op.result.type)]


@_register_lowering(arith.SelectOp)
def _select_op_lowering_rule(
    ctx: LoweringContext, op: arith.SelectOp
) -> Sequence[ir.Value]:
  del ctx
  in_layouts = inference_utils.in_layouts(op)
  [layout] = inference_utils.out_layouts(op)
  if any(in_layout != layout for in_layout in in_layouts):
    raise ValueError("Layout mismatch")
  pred = _fragmented_array_from_ir(op.condition, layout)
  true_value = _fragmented_array_from_ir(op.true_value, layout)
  false_value = _fragmented_array_from_ir(op.false_value, layout)
  result = pred.select(true_value, false_value)
  return [fragmented_array_to_ir(result, op.result.type)]


@_register_lowering(mgpu.WGMMAOp)
def _mgpu_wgmma_op_lowering_rule(
    ctx: LoweringContext, wgmma_op: mgpu.WGMMAOp
) -> Sequence[ir.Value]:
  del ctx
  in_layouts = inference_utils.in_layouts(wgmma_op)
  assert in_layouts[0] == layouts_lib.to_layout_attr(fa.WGMMA_LAYOUT)
  [out_layout] = inference_utils.out_layouts(wgmma_op)
  assert out_layout == layouts_lib.to_layout_attr(fa.WGMMA_LAYOUT)

  # s8/i8 WGMMA expects signed integer accumulator.
  element_type = wgmma_op.a.type.element_type
  is_signed = True if isinstance(element_type, ir.IntegerType) else None
  regs = _fragmented_array_from_ir(
      wgmma_op.accumulator, in_layouts[0], is_signed
  )
  acc = wgmma.WGMMAAccumulator.from_registers(regs, sync=False)  # pyrefly: ignore[missing-attribute]

  if isinstance(wgmma_op.a.type, ir.VectorType):
    a_transforms = None
    b_transforms = inference_utils.in_transforms(wgmma_op)[0]
    unwrapped_a_ref = None
    unwrapped_b_ref = unwrap_transformed_memref(wgmma_op.b, b_transforms)
  else:
    a_transforms, b_transforms = inference_utils.in_transforms(wgmma_op)
    unwrapped_a_ref = unwrap_transformed_memref(wgmma_op.a, a_transforms)
    unwrapped_b_ref = unwrap_transformed_memref(wgmma_op.b, b_transforms)

  b_swizzle = swizzle_from_transforms_attr(b_transforms)

  if isinstance(wgmma_op.a.type, ir.VectorType):
    expected_a_layout = (
        fa.WGMMA_LAYOUT_8BIT
        if utils.bitwidth(element_type) == 8
        else fa.WGMMA_LAYOUT
    )
    assert in_layouts[1] == layouts_lib.to_layout_attr(expected_a_layout)
    a_operand = _fragmented_array_from_ir(wgmma_op.a, in_layouts[1], is_signed)
  else:
    assert a_transforms is not None
    a_swizzle = swizzle_from_transforms_attr(a_transforms)
    if a_swizzle != b_swizzle:
      raise ValueError(
          f"Non-matching swizzles of operands a and b in WGMMA: {a_swizzle} !="
          f" {b_swizzle}"
      )
    assert unwrapped_a_ref is not None
    a_operand = unwrapped_a_ref

  # pyrefly: ignore[missing-attribute]
  new_acc = wgmma.wgmma(acc, a_operand, unwrapped_b_ref, swizzle=b_swizzle)
  return [
      fragmented_array_to_ir(
          new_acc.value.to_layout(fa.WGMMA_LAYOUT),
          wgmma_op.accumulator.type,
      )
  ]


@_register_lowering(mgpu.MMAOp)
def _mgpu_mma_op_lowering_rule(
    ctx: LoweringContext, mma_op: mgpu.MMAOp
) -> Sequence[ir.Value]:
  del ctx
  acc_layout, a_layout, b_layout = inference_utils.in_layouts(mma_op)
  [out_layout] = inference_utils.out_layouts(mma_op)

  a_element_type = mma_op.a.type.element_type
  m, n = mma_op.accumulator.type.shape
  mma_layouts = MMALayouts.for_shape(a_element_type, m, n)
  expected_acc_layout = layouts_lib.to_layout_attr(mma_layouts.acc)
  assert acc_layout == expected_acc_layout
  assert out_layout == expected_acc_layout
  is_signed = True if isinstance(a_element_type, ir.IntegerType) else None

  acc = _fragmented_array_from_ir(
      mma_op.accumulator, acc_layout, is_signed=is_signed
  )
  a = _fragmented_array_from_ir(mma_op.a, a_layout, is_signed=is_signed)
  b = _fragmented_array_from_ir(mma_op.b, b_layout, is_signed=is_signed)

  new_acc = do_mma(acc, a, b)
  return [fragmented_array_to_ir(new_acc, mma_op.result.type)]


@_register_lowering(mgpu.ArriveOp, support_warp_semantics=True)
def _mgpu_arrive_op_lowering_rule(
    ctx: LoweringContext, arrive_op: mgpu.ArriveOp
) -> Sequence[ir.Value]:
  barrier = utils.DialectBarrierRef.from_barrier_memref(arrive_op.barrier)
  orders_tc = arrive_op.orders_tensor_core.value
  predicate = arrive_op.predicate
  if orders_tc:
    # Barrier expects a single thread arrival.
    pred = ctx.single_lane_predicate
    predicate = pred if predicate is None else arith.andi(predicate, pred)
    arrival_count = 1
  elif ctx.thread_semantics == utils.ThreadSubset.WARP:
    # In warp-level lowering, we arrive on each CUDA thread in a warp, but the
    # barrier still expects a full 128 arrivals so we arrive 4 times on each
    # CUDA thread instead.
    arrival_count = 4
  else:
    # Barrier expects each thread arrives once.
    arrival_count = 1

  barrier.barrier_ref.arrive(
      arrival_count,
      orders_tensor_core=orders_tc,
      predicate=predicate,
      scope=ctx.thread_semantics,
  )
  return []


@_register_lowering(mgpu.ArriveExpectTxOp, support_warp_semantics=True)
def _mgpu_arrive_expect_tx_op_lowering_rule(
    ctx: LoweringContext, arrive_expect_tx_op: mgpu.ArriveExpectTxOp
) -> Sequence[ir.Value]:
  i32 = ir.IntegerType.get_signless(32)
  num_lanes = (
      utils.WARPGROUP_SIZE
      if ctx.thread_semantics == utils.ThreadSubset.WARPGROUP
      else utils.WARP_SIZE
  )
  num_bytes = arrive_expect_tx_op.expect_tx
  barrier = utils.DialectBarrierRef.from_barrier_memref(
      arrive_expect_tx_op.barrier
  )
  orders_tc = barrier.orders_tensor_core

  if orders_tc:
    tx_bytes = num_bytes
  elif isinstance(num_bytes.owner, arith.ConstantOp):
    num_bytes_int = int(num_bytes.owner.value)
    if num_bytes_int % num_lanes == 0:
      # Prefer uniform arrival whenever possible because it's more efficient.
      # We arrive uniformly from each lane in the WG/Warp, so we need to divide
      # the number of bytes by the number of lanes in the WG/Warp.
      tx_bytes = utils.c(num_bytes_int // num_lanes, i32)
    else:
      tx_bytes = arith.select(
          ctx.single_lane_predicate,
          num_bytes,
          utils.c(0, i32),
      )
  else:
    tx_bytes = arith.select(
        ctx.single_lane_predicate,
        num_bytes,
        utils.c(0, i32),
    )

  # In Warp-level lowering, we arrive on each CUDA thread in a warp, but the
  # barrier still expects a full 128 arrivals so we arrive 4 times on each CUDA
  # thread instead.
  if ctx.thread_semantics == utils.ThreadSubset.WARP and not orders_tc:
    barrier.barrier_ref.arrive(arrival_count=3, can_complete=False)
  barrier.barrier_ref.arrive_expect_tx(
      tx_bytes,
      predicate=ctx.single_lane_predicate if orders_tc else None,
      tensor_core_order_scope=ctx.thread_semantics if orders_tc else None,
  )

  return []


@_register_lowering(mgpu.SliceSMEMOp, support_warp_semantics=True)
def _mgpu_slice_smem_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.SliceSMEMOp
) -> Sequence[ir.Value]:
  ref_ty = ir.MemRefType(op.result.type)
  offset = op.offset.value
  if isinstance(ref_ty.element_type, mgpu.BarrierType):
    # Barrier memrefs are not transformed and must not be wrapped.
    assert not inference_utils.has_out_transforms_set(op)
    return [_slice_smem(ref_ty, offset, ctx.smem_requested_bytes)]

  [out_transforms] = inference_utils.out_transforms(op)
  transformed_ref_ty = transform_type(ref_ty, out_transforms)
  transformed_ref = _slice_smem(transformed_ref_ty, offset, ctx.smem_requested_bytes)
  return [wrap_transformed_memref(transformed_ref, op.result.type, out_transforms)]


def _slice_smem(result: ir.MemRefType, offset: int, smem_size: int):
  size = math.prod(result.shape) * utils.bitwidth(result.element_type) // 8
  if offset + size > smem_size:
    raise ValueError("Ran out of shared memory.")

  i8 = ir.IntegerType.get_signless(8)
  smem_base = gpu.dynamic_shared_memory(
      ir.MemRefType.get((utils.DYNAMIC,), i8, memory_space=utils.smem())
  )
  ir_offset = arith.constant(ir.IndexType.get(), offset)
  lowered_result_type = result
  if isinstance(result.element_type, mgpu.BarrierType):
    lowered_result_type = ir.MemRefType.get(
        result.shape, _lowered_barrier_type(), memory_space=utils.smem()
    )
  view = memref.view(lowered_result_type, smem_base, ir_offset, [])
  if result == lowered_result_type:
    return view
  return builtin.unrealized_conversion_cast([result], [view])


@_register_lowering(mgpu.WithTransformsOp, support_warp_semantics=True)
def _mgpu_with_transforms_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.WithTransformsOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.WithTransformsOp.
  This is a noop that simply returns its input.
  """
  del ctx

  [in_transforms] = inference_utils.in_transforms(op)
  unwrapped_source_ref = unwrap_transformed_memref(op.ref, in_transforms)
  out_transforms = inference_utils.out_transforms(op)[0]
  wrapped_ref = wrap_transformed_memref(
      unwrapped_source_ref, op.result.type, out_transforms
  )
  return [wrapped_ref]


def _tile_transform_offsets(
    tiling: Sequence[int],
    static_offsets: Sequence[int],
    dynamic_offsets: Sequence[ir.Value],
) -> tuple[Sequence[int], Sequence[ir.Value]]:
  """Computes the static and dynamic offsets after the given tiling is applied.

  Conceptually, this function is analogous to
  tile.transform_shape(static_offsets), except that it also handles dynamic offsets.
  """
  dynamic_offset_index = 0
  new_static_offsets = []
  new_dynamic_offsets = []

  # Preserve all offsets in non-tiled dimensions.
  for offset in static_offsets[: -len(tiling)]:
    new_static_offsets.append(offset)
    if offset == ir.ShapedType.get_dynamic_stride_or_offset():
      new_dynamic_offsets.append(dynamic_offsets[dynamic_offset_index])
      dynamic_offset_index += 1

  # Compute static and dynamic offsets of tiled dimensions.
  for tile_size, offset in zip(
      tiling, static_offsets[-len(tiling) :], strict=True
  ):
    if offset == ir.ShapedType.get_dynamic_stride_or_offset():
      # Here we assume that the offset is divisble by the tile size, but we
      # don't check it. This has been established at the time the tiling was
      # inferred.
      dyn_offset = arith.divui(
          dynamic_offsets[dynamic_offset_index],
          utils.c(tile_size, ir.IndexType.get()),
      )
      new_dynamic_offsets.append(dyn_offset)
      new_static_offsets.append(ir.ShapedType.get_dynamic_stride_or_offset())
      dynamic_offset_index += 1
    else:
      assert offset % tile_size == 0
      new_static_offsets.append(offset // tile_size)

  # Add 0 offsets for the newly created dimension of the tile.
  new_static_offsets += [0] * len(tiling)

  return new_static_offsets, new_dynamic_offsets


@_register_lowering(memref.SubViewOp, support_warp_semantics=True)
def _memref_subview_op_lowering_rule(
    ctx: LoweringContext, op: memref.SubViewOp
) -> Sequence[ir.Value]:
  del ctx

  if any(s != 1 for s in op.static_strides):
    raise NotImplementedError("SubViewOp only supports static strides of 1.")
  if op.sizes:
    raise NotImplementedError("SubViewOp only supports static sizes.")
  src_ty = ir.MemRefType(op.source.type)


  if utils.is_tmem_ref(src_ty):
    [in_tmem_layout] = inference_utils.in_tmem_layouts(op)
    [out_tmem_layout] = inference_utils.out_tmem_layouts(op)
    assert in_tmem_layout == out_tmem_layout
    ref = _tmem_ref_from_ir(op.source, in_tmem_layout)
    indices = []
    dynamic_offset_index = 0
    for offset, size in zip(op.static_offsets, op.static_sizes, strict=True):
      if ir.ShapedType.is_dynamic_size(offset):
        offset = op.offsets[dynamic_offset_index]
        dynamic_offset_index += 1
      indices.append(utils.DynamicSlice(offset, size))
    return [_tmem_ref_to_ir(ref.slice(*indices), op.result.type)]

  in_transforms = inference_utils.in_transforms(op)[0]
  out_transforms = inference_utils.out_transforms(op)[0]

  if in_transforms != out_transforms:
    raise NotImplementedError(
        "SubViewOp transforms for the input and output refs must be identical."
    )

  swizzle = swizzle_from_transforms_attr(out_transforms)
  if swizzle != mgpu.SwizzlingMode.kNoSwizzle:
    swizzle_elems = swizzle * 8 // utils.bitwidth(src_ty.element_type)
    source_strides, _ = src_ty.get_strides_and_offset()
    dyn_offset_index = 0
    for stride, static_offset, size in zip(
        source_strides, op.static_offsets, op.static_sizes, strict=True
    ):
      offset: int | ir.Value
      if ir.ShapedType.is_dynamic_size(static_offset):
        offset = op.offsets[dyn_offset_index]
        dyn_offset_index += 1
      else:
        offset = static_offset
      if stride != 1:
        continue
      # A dimension with stride 1 is a minor dimension and is swizzled.
      if size % swizzle_elems != 0:
        raise ValueError(
            f"Swizzled dimension of {size=} is not a multiple of"
            f" {swizzle_elems=}."
        )
      if isinstance(offset, ir.Value):
        if not utils.is_known_divisible(offset, swizzle_elems):
          raise ValueError(
              "subview dynamic offset is not a known multiple of"
              f" {swizzle_elems=}."
          )
      elif offset % swizzle_elems != 0:
        raise ValueError(
            f"subview {offset=} is not a multiple of {swizzle_elems=}."
        )

  unwrapped_source_ref = unwrap_transformed_memref(op.source, in_transforms)
  transforms = memref_transforms_from_transforms_attr(out_transforms)
  match transforms:
    case ():
      new_subview_op = memref.SubViewOp(
          op.result.type,
          unwrapped_source_ref,
          op.offsets,
          sizes=[],
          strides=[],
          static_offsets=op.static_offsets,
          static_sizes=op.static_sizes,
          static_strides=op.static_strides,
      )
    case (lc.TileTransform() as tile_transform, ):
      in_transformed_ty = ir.MemRefType(unwrapped_source_ref.type)
      tiling = tile_transform.tiling
      if any(
          ir.ShapedType.is_dynamic_size(s)
          for s in list(op.static_sizes)[-len(tiling) :]
      ):
        raise NotImplementedError(
            "SubViewOp only supports static sizes for the tiled dimensions."
        )
      new_sizes = tile_transform.transform_shape(list(op.static_sizes))
      # TODO(bchetioui): support transposed offsets.
      new_static_offsets, new_dynamic_offsets = _tile_transform_offsets(
          tiling, list(op.static_offsets), list(op.offsets)
      )

      new_subview_op = memref.SubViewOp(
          transform_type(ir.MemRefType(op.result.type), transforms),
          unwrapped_source_ref,
          new_dynamic_offsets,
          sizes=[],
          strides=[],
          static_offsets=new_static_offsets,
          static_sizes=new_sizes,
          static_strides=[1] * len(in_transformed_ty.shape),
      )
    case _:
      raise NotImplementedError(
          "SubViewOp only supports a single tile transform."
      )

  wrapped_ref = wrap_transformed_memref(
      new_subview_op.result, op.result.type, out_transforms
  )
  return [wrapped_ref]


# memref.cast shows up when we slice a ref with a dynamic index, that later gets
# folded into a constant. At that time, the sliced ref type is simplified from
# having a dynamic offset to a constant offset. However, downstream consumer ops
# still expect the offset to be dynamic, forcing the insertion of a memref.cast
# op to reintroduce the dynamic offset.
@_register_lowering(memref.CastOp, support_warp_semantics=True)
def _memref_cast_op_lowering_rule(
    ctx: LoweringContext, op: memref.CastOp
) -> Sequence[ir.Value]:
  del ctx
  in_ty = ir.MemRefType(op.source.type)
  out_ty = ir.MemRefType(op.result.type)
  in_strides, _ = in_ty.get_strides_and_offset()
  out_strides, _ = out_ty.get_strides_and_offset()

  unoffseted_in_ty = ir.MemRefType.get(
      in_ty.shape,
      in_ty.element_type,
      memory_space=in_ty.memory_space,
      layout=ir.StridedLayoutAttr.get(0, in_strides),
  )
  unoffseted_out_ty = ir.MemRefType.get(
      out_ty.shape,
      out_ty.element_type,
      memory_space=out_ty.memory_space,
      layout=ir.StridedLayoutAttr.get(0, out_strides),
  )

  if unoffseted_in_ty != unoffseted_out_ty:
    raise NotImplementedError(
        "Only support memref.cast where the input and output types are the "
        f"same up to offset, but got {in_ty=} and {out_ty=}."
    )

  memory_space = ir.MemRefType(op.result.type).memory_space
  if memory_space == utils.smem():
    [in_transforms] = inference_utils.in_transforms(op)
    [out_transforms] = inference_utils.out_transforms(op)
    if in_transforms != out_transforms:
      raise NotImplementedError(
          "memref.cast transforms must have identical transforms for both "
          f"input and output but got {in_transforms=} and {out_transforms=}"
      )
    result = memref.cast(
        transform_type(ir.MemRefType(op.result.type), out_transforms),
        unwrap_transformed_memref(op.source, in_transforms),
    )
    return [wrap_transformed_memref(result, op.result.type, out_transforms)]

  if memory_space == utils.tmem():
    [in_tmem_layout] = inference_utils.in_tmem_layouts(op)
    [out_tmem_layout] = inference_utils.out_tmem_layouts(op)
    if in_tmem_layout != out_tmem_layout:
      raise NotImplementedError(
          "memref.cast tmem layouts must be identical for both input and"
          f" output but got {in_tmem_layout=} and {out_tmem_layout=}"
      )
    return [_tmem_ref_to_ir(_tmem_ref_from_ir(op.source, in_tmem_layout),
                            op.result.type)]

  raise NotImplementedError(
      f"Unsupported memory space when lowering memref.cast: {memory_space}"
  )


def _permutation_to_affine_map_attr(
    permutation: Sequence[int],
) -> ir.AffineMapAttr:
  return ir.AffineMapAttr.get(ir.AffineMap.get_permutation(permutation))


@_register_lowering(memref.TransposeOp, support_warp_semantics=True)
def _memref_transpose_op_lowering_rule(
    ctx: LoweringContext, op: memref.TransposeOp
) -> Sequence[ir.Value]:
  del ctx

  in_transforms_attr = inference_utils.in_transforms(op)[0]
  unwrapped_in_ref = unwrap_transformed_memref(op.in_, in_transforms_attr)
  in_swizzle = swizzle_from_transforms_attr(in_transforms_attr)
  in_transforms = memref_transforms_from_transforms_attr(in_transforms_attr)
  out_transforms_attr = inference_utils.out_transforms(op)[0]
  out_swizzle = swizzle_from_transforms_attr(out_transforms_attr)
  out_transforms = memref_transforms_from_transforms_attr(out_transforms_attr)

  if in_swizzle != out_swizzle:
    raise ValueError(
        f"Swizzle mismatch. In transforms swizzle: {in_swizzle}, out transforms"
        f" swizzle {out_swizzle}."
    )
  if len(out_transforms) != len(in_transforms):
    raise ValueError(
        f"Size mismatch for in/out transforms. In transforms: {in_transforms},"
        f" out transforms: {out_transforms}."
    )
  if not out_transforms:
    new_permutation = op.permutation
  else:
    permutation = [
        ir.AffineDimExpr(e).position
        for e in op.permutation.value.results
    ]
    # We expect to have the same transforms on in/out, up to permutation of the
    # out transforms.
    # For example, for 3D input:
    #   permutation: (0, 2, 1)
    #   in_transforms: TilingTransform((32, 8))
    # We expect:
    #   out_transforms: TilingTransform((8, 32))
    # TODO(olechwierowicz): Support multiple transforms.
    [transform] = out_transforms
    [in_transform] = in_transforms
    if not isinstance(transform, lc.TileTransform) or not isinstance(
        in_transform, lc.TileTransform
    ):
      raise NotImplementedError(
          f"Unsupported in/out transforms. In transform: {in_transform}, out"
          f" transform: {transform}"
      )
    tiling_len = len(in_transform.tiling)
    tiling_offset = len(permutation) - tiling_len
    if any(dim < tiling_offset for dim in permutation[-tiling_len :]):
      raise ValueError(
          f"Cannot tile a transpose ({permutation}). Tiling dims"
          f" ({permutation[-tiling_len:]}) cannot contain non-tiled dims."
          f" All of them must be >= {tiling_offset}."
      )
    dims = [-1] * len(permutation)
    dims = dims[:-tiling_len] + list(in_transform.tiling)
    permuted_dims = tuple(dims[permutation[i]] for i in range(len(dims)))
    if permuted_dims[-tiling_len:] != transform.tiling:
      raise ValueError(
          f"Invalid in/out transforms. In transform: {in_transform}, out"
          f" transform: {transform}"
      )
    new_permutation = permutation + [
        x + tiling_len for x in permutation[-tiling_len:]
    ]
    new_permutation = _permutation_to_affine_map_attr(new_permutation)

  new_transpose_op = memref.TransposeOp(
      transform_type(ir.MemRefType(op.result.type), out_transforms),
      unwrapped_in_ref,
      new_permutation,
  )

  out_transforms = inference_utils.out_transforms(op)[0]
  wrapped_ref = wrap_transformed_memref(
      new_transpose_op.result, op.result.type, out_transforms
  )
  return [wrapped_ref]


@_register_lowering(memref.ExpandShapeOp)
def _memref_expand_shape_op_lowering_rule(
    ctx: LoweringContext, op: memref.ExpandShapeOp
) -> Sequence[ir.Value]:
  del ctx

  in_transforms = inference_utils.in_transforms(op)[0]
  unwrapped_in_ref = unwrap_transformed_memref(op.src, in_transforms)
  in_transformed_ty = ir.MemRefType(unwrapped_in_ref.type)

  out_transforms = inference_utils.out_transforms(op)[0]
  out_transformed_ty = transform_type(op.result.type, out_transforms)

  reassociation = cast(list[ir.ArrayAttr], list(op.reassociation))
  num_tiling_dims = len(in_transformed_ty.shape) - len(op.src.type.shape)

  # We don't currently allow expanding tiled dimensions. So to compute the
  # reassociation on the lowered types, we just need to backfill the original
  # one with the number of missing dimensions.
  if num_tiling_dims > 0 and any(
      len(x) > 1 for x in reassociation[-num_tiling_dims:]
  ):
    # If we ever remove this restriction, we will need to ensure this is
    # compatible with `transform_type`.
    raise NotImplementedError("Expanding tiled dimensions is not supported.")

  start_index = len(op.static_output_shape)
  for i in range(start_index, start_index + num_tiling_dims):
    reassociation.append([i])  # pyrefly: ignore[bad-argument-type]

  new_expand_shape_op = memref.ExpandShapeOp(
      out_transformed_ty,
      unwrapped_in_ref,
      reassociation,
      output_shape=op.output_shape,
      static_output_shape=out_transformed_ty.shape,
  )

  wrapped_ref = wrap_transformed_memref(
      new_expand_shape_op.result, op.result.type, out_transforms
  )
  return [wrapped_ref]


# TODO(bchetioui): find a way to consolidate the logic that is shared logic with
# layout inference. It is not entirely clear what the best approach is.
def _check_collapse_shape(
    op: memref.CollapseShapeOp,
    in_transforms: Sequence[lc.MemRefTransform],
    out_transforms: Sequence[lc.MemRefTransform],
):
  if len(in_transforms) != len(out_transforms):
    raise ValueError(
        "Expected the same number of in/out transforms, but got "
        f"{in_transforms=} and {out_transforms=}"
    )
  if not in_transforms:
    return
  t_in, *in_transforms = in_transforms
  t_out, *_ = out_transforms
  if (in_transforms or
      not isinstance(t_in, lc.TileTransform) or
      not isinstance(t_out, lc.TileTransform)):
    raise NotImplementedError(
        "Only a single tiling transform is supported when collapsing a shape, "
        f"but got {in_transforms=} and {out_transforms=}"
    )
  src_ty = ir.MemRefType(op.src.type)
  strides, _ = src_ty.get_strides_and_offset()
  if strides != utils.get_contiguous_strides(src_ty.shape):
    raise NotImplementedError(
        "Collapsing the shape of a memref with non-contiguous strides is not "
        "supported"
    )
  reassociation = tuple(len(ir.ArrayAttr(idx)) for idx in op.reassociation)

  collapsed_tiling = cs.reduce_expression(
      cs.CollapseShape(cs.SMEMTransforms(t_in, None), tuple(src_ty.shape),
                       reassociation),
      {},
  )

  if isinstance(collapsed_tiling, cs.Unsatisfiable):
    raise ValueError(f"Input tiling {t_in.tiling} is not compatible with {op}")

  assert isinstance(collapsed_tiling, cs.SMEMTransforms)
  expected_t_out = collapsed_tiling.tiling
  assert expected_t_out is not None
  if expected_t_out != t_out:
    raise ValueError(
        "Input/output tiling mismatch when attempting to collapse a shape. "
        f"Expected output tiling to be {expected_t_out.tiling} for input "
        f"tiling {t_in.tiling}, but got {t_out.tiling}"
    )


@_register_lowering(memref.CollapseShapeOp, support_warp_semantics=True)
def _memref_collapse_shape_op_lowering_rule(
    ctx: LoweringContext, op: memref.CollapseShapeOp
) -> Sequence[ir.Value]:
  del ctx

  [in_transforms_attr] = inference_utils.in_transforms(op)
  [out_transforms_attr] = inference_utils.out_transforms(op)

  in_swizzle = swizzle_from_transforms_attr(in_transforms_attr)
  in_transforms = memref_transforms_from_transforms_attr(in_transforms_attr)
  out_swizzle = swizzle_from_transforms_attr(out_transforms_attr)
  out_transforms = memref_transforms_from_transforms_attr(out_transforms_attr)

  if in_swizzle != out_swizzle:
    raise ValueError(
        f"Swizzle mismatch. In transforms swizzle: {in_swizzle}, out transforms"
        f" swizzle {out_swizzle}."
    )
  _check_collapse_shape(op, in_transforms, out_transforms)
  reassociation = [
      [ir.IntegerAttr(i).value for i in ir.ArrayAttr(dims)]
      for dims in op.reassociation
  ]
  new_reassociation = reassociation.copy()
  if in_transforms:
    [t_in] = in_transforms
    assert isinstance(t_in, lc.TileTransform)
    tiling_rank = to_process =  len(t_in.tiling)
    for index_from_end, dims in enumerate(reassociation[::-1]):
      to_process -= len(dims)
      if to_process < 0:
        # This should be caught by `_check_collapse_shape` today, but we check
        # it here as well in case `cs.CollapseShape` ever changes to allow this.
        raise ValueError(
            f"Reassociation {reassociation} is not compatible with tiling "
            f"{t_in.tiling}, as it causes tiled and untiled dimensions to "
            "be collapsed together"
        )
      if to_process == 0:
        for t_dims in reassociation[-index_from_end - 1:]:
          new_reassociation.append([dim + tiling_rank for dim in t_dims])
        break
    assert to_process == 0

  result = memref.collapse_shape(
      transform_type(op.result.type, out_transforms),
      unwrap_transformed_memref(op.src, in_transforms_attr),
      new_reassociation,
  )
  return [wrap_transformed_memref(result, op.result.type, out_transforms_attr)]


@_register_lowering(memref.LoadOp)
def _memref_load_op_lowering_rule(
    ctx: LoweringContext, op: memref.LoadOp
) -> Sequence[ir.Value]:
  """Lowering rule for memref.LoadOp.

  Loads are never transformed so this rule is mostly just a pass-through.
  """
  del ctx

  in_transforms = inference_utils.in_transforms(op)[0]
  if in_transforms:
    raise NotImplementedError(f"memref.LoadOp does not support transforms: {op}")

  new_load_op = memref.LoadOp(
      memref=unwrap_transformed_memref(op.memref, in_transforms),
      indices=op.indices,
      nontemporal=op.nontemporal,
  )
  return [new_load_op.result]


@_register_lowering(memref.StoreOp, support_warp_semantics=True)
def _memref_store_op_lowering_rule(
    ctx: LoweringContext, op: memref.StoreOp
) -> Sequence[ir.Value]:
  """Lowering rule for memref.StoreOp.

  Stores are never transformed so this rule is mostly just a pass-through.
  """
  del ctx

  in_transforms = inference_utils.in_transforms(op)[0]
  if in_transforms:
    raise NotImplementedError(f"memref.StoreOp does not support transforms: {op}")

  memref.StoreOp(
      value=op.value,
      memref=unwrap_transformed_memref(op.memref, in_transforms),
      indices=op.indices,
      nontemporal=op.nontemporal,
  )
  return []


@_register_lowering(mgpu.TmemAllocOp)
def _tmem_alloc_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.TmemAllocOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.TmemAllocOp."""
  ctx.check_collective(op)

  output_shape = ir.MemRefType(op.result.type).shape
  ncols = output_shape[1] // op.packing.value

  with utils.when(ctx.single_warp_per_block_predicate):
    tcgen05.tmem_alloc(op.smem_ptr, ncols, bool(op.collective), exact=False)
  gpu.barrier()
  tmem_addr = memref.load(op.smem_ptr, [])

  cast_op = builtin.UnrealizedConversionCastOp(
      [op.result.type], [tmem_addr]
  )
  cast_op.attributes["collective"] = op.collective
  cast_op.attributes["packing"] = op.packing
  cast_op.attributes["layout"] = inference_utils.out_tmem_layouts(op)[0]

  return [cast_op.result]


@_register_lowering(mgpu.TmemRelinquishAllocPermitOp)
def _tmem_relinquish_alloc_permit_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.TmemRelinquishAllocPermitOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.TmemRelinquishAllocPermitOp."""
  ctx.check_collective(op)
  with utils.when(ctx.single_warp_per_block_predicate):
    tcgen05.tmem_relinquish_alloc_permit(bool(op.collective))
  return []


@_register_lowering(mgpu.TmemDeallocOp)
def _tmem_dealloc_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.TmemDeallocOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.TmemDeallocOp."""
  i32 = ir.IntegerType.get_signless(32)
  conversion_cast, [tmem_addr] = _undo_conversion_cast(op.tmem_ref, [i32])
  collective = ir.BoolAttr(conversion_cast.attributes["collective"]).value
  packing: int = ir.IntegerAttr(conversion_cast.attributes["packing"]).value

  output_shape = ir.MemRefType(op.tmem_ref.type).shape
  ncols = output_shape[1] // packing

  with utils.when(ctx.single_warp_per_block_predicate):
    tcgen05.tmem_dealloc(tmem_addr, ncols, collective, exact=False)

  return []


def _tmem_ref_from_ir(
    ref: ir.Value, expected_layout: ir.Attribute
) -> tcgen05.TMEMRef:
  """Returns a TMEMRef from an IR value.

  Throws an error if the annotated layout does not match the expected layout.
  """
  if not isinstance(ref.type, ir.MemRefType):
    raise ValueError(f"{ref} is not a memref.")
  mem_ref_ty = ir.MemRefType(ref.type)

  if mem_ref_ty.memory_space != utils.tmem():
    raise ValueError(
        f"{ref} has a memory space {mem_ref_ty.memory_space} that is not TMEM."
    )

  i32 = ir.IntegerType.get_signless(32)
  conversion_cast, [tmem_addr] = _undo_conversion_cast(ref, [i32])

  assert mem_ref_ty.rank == 2
  shape = cast(tuple[int, int], tuple(mem_ref_ty.shape))
  el_ty = mem_ref_ty.element_type
  layout_attr = conversion_cast.attributes["layout"]
  if layout_attr != expected_layout:
    raise ValueError(
        f"{ref} has a layout {layout_attr} that does not match the expected"
        f" layout {expected_layout}."
    )
  layout = layouts_lib.from_layout_attr(layout_attr)
  assert isinstance(layout, fa.TiledLayout)
  tmem_layout = tcgen05.TMEMLayout(
      layout.tiling, layout.warp_dims, layout.lane_dims, layout.vector_dim
  )
  return tcgen05.TMEMRef(tmem_addr, shape, el_ty, tmem_layout)


def _tmem_ref_to_ir(ref: tcgen05.TMEMRef, ty: ir.MemRefType) -> ir.Value:
  """Returns an IR value from a TMEMRef."""
  conversion_cast = builtin.UnrealizedConversionCastOp([ty], [ref.address])
  conversion_cast.attributes["layout"] = layouts_lib.to_layout_attr(ref.layout)
  return conversion_cast.result


@_register_lowering(mgpu.TcGen05MMAOp, support_warp_semantics=True)
def _tcgen05_mma_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.TcGen05MMAOp
) -> Sequence[ir.Value]:
  ctx.check_collective(op)

  def tmem_layout(operand):
    result = inference_utils.in_tmem_layout_for_operand(op, operand)
    # satisfy the type checker
    assert result is not None
    return result

  acc_ref = _tmem_ref_from_ir(op.accumulator, tmem_layout(op.accumulator))

  if op.a_sparse_metadata is not None:
    a_sparse_metadata = _tmem_ref_from_ir(
        op.a_sparse_metadata, tmem_layout(op.a_sparse_metadata)
    )
  else:
    a_sparse_metadata = None
  if (scaled := op.a_scale is not None) != (op.b_scale is not None):
    raise ValueError("Expected both or neither of scales to be specified.")
  if scaled:
    a_scale = _tmem_ref_from_ir(op.a_scale, tmem_layout(op.a_scale))  # pyrefly: ignore[bad-argument-type]
    b_scale = _tmem_ref_from_ir(op.b_scale, tmem_layout(op.b_scale))  # pyrefly: ignore[bad-argument-type]
  else:
    a_scale = None
    b_scale = None

  if utils.is_smem_ref(op.a):
    a_transforms, b_transforms = inference_utils.in_transforms(op)
    a_swizzle = swizzle_from_transforms_attr(a_transforms)
    b_swizzle = swizzle_from_transforms_attr(b_transforms)
    a_ref = unwrap_transformed_memref(op.a, a_transforms)
    b_ref = unwrap_transformed_memref(op.b, b_transforms)
  else:
    a_ref = _tmem_ref_from_ir(op.a, tmem_layout(op.a))
    [b_transforms] = inference_utils.in_transforms(op)
    b_swizzle = swizzle_from_transforms_attr(b_transforms)
    a_swizzle = b_swizzle
    b_ref = unwrap_transformed_memref(op.b, b_transforms)

  with utils.when(ctx.single_lane_predicate):
    tcgen05.mma(
        acc_ref,
        a_ref,
        b_ref,
        a_swizzle=a_swizzle,
        b_swizzle=b_swizzle,
        a_scale=a_scale,
        b_scale=b_scale,
        accumulate=op.accumulate,
        collective=op.collective.value,
        a_sparse_metadata=a_sparse_metadata,
    )

  return []


@_register_lowering(mgpu.AsyncLoadTmemOp)
def _async_load_tmem_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.AsyncLoadTmemOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.AsyncLoadTmemOp."""
  del ctx
  in_layout_attr = inference_utils.in_tmem_layouts(op)[0]
  tmem_ref = _tmem_ref_from_ir(op.source, in_layout_attr)
  out_layout_attr = inference_utils.out_layouts(op)[0]
  out_layout = layouts_lib.from_layout_attr(out_layout_attr)
  assert isinstance(out_layout, fa.TiledLayout)
  element_type = ir.MemRefType(op.source.type).element_type
  if op.reduce is not None:
    reduce_str = cast(
        tcgen05.LoadReduceOp, str(mgpu.TMEMLoadReduction(op.reduce.value))  # pyrefly: ignore[missing-attribute]
    )
    if isinstance(element_type, ir.IntegerType):
      is_signed = not reduce_str.startswith("abs")
      reduce_str = reduce_str[-3:]
    else:
      is_signed = None
    loaded, reduced = tmem_ref.load(out_layout, is_signed, reduce=reduce_str)
    return [
        fragmented_array_to_ir(loaded, op.results[0].type),
        fragmented_array_to_ir(reduced, op.results[1].type),
    ]
  else:
    is_signed = _default_is_signed(element_type)
    arr = tmem_ref.load(out_layout, is_signed)
    return [fragmented_array_to_ir(arr, op.results[0].type)]


@_register_lowering(mgpu.AsyncStoreTmemOp)
def _async_store_tmem_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.AsyncStoreTmemOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.AsyncStoreTmemOp."""
  del ctx
  in_layout_attr = inference_utils.in_tmem_layouts(op)[0]
  tmem_ref = _tmem_ref_from_ir(op.destination, in_layout_attr)
  in_layout_attr = inference_utils.in_layouts(op)[0]
  arr = _fragmented_array_from_ir(op.source, in_layout_attr)
  tmem_ref.store(arr)

  return []


@_register_lowering(mgpu.TcGen05CommitArriveOp, support_warp_semantics=True)
def _tcgen05_commit_arrive_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.TcGen05CommitArriveOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.TcGen05CommitArriveOp."""
  ctx.check_collective(op)
  barrier = utils.DialectBarrierRef.from_barrier_memref(op.barrier)
  predicate = ctx.single_lane_predicate
  # TODO(cjfj): simplify when minimum jaxlib version is 0.12.
  if (pred := getattr(op, "predicate", None)) is not None:
    predicate = arith.andi(predicate, pred)
  tcgen05.commit_arrive(
      barrier.barrier_ref,
      op.collective.value,
      ctx.launch_context,
      predicate=predicate,
  )
  return []


@_register_lowering(mgpu.CustomPrimitiveOp, support_warp_semantics=True)
def _mgpu_custom_primitive_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.CustomPrimitiveOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.CustomPrimitiveOp."""
  del ctx
  block = op.body.blocks[0]
  for arg, operand in zip(block.arguments, op.operands, strict=True):
    arg.replace_all_uses_with(operand)

  return_op = None
  ip = ir.InsertionPoint.current
  for block_op in block.operations:
    if isinstance(block_op.opview, mgpu.ReturnOp):
      assert return_op is None
      return_op = block_op.opview
      continue
    block_op.detach_from_parent()
    ip.insert(block_op)

  if return_op is None:
    raise ValueError("A custom return op must terminate the block.")

  return return_op.operands


@_register_lowering(mgpu.WarpMapOp)
def _mgpu_warp_map_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.WarpMapOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.WarpMapOp."""
  for a, o in zip(op.body.arguments, op.operands, strict=True):
    a.replace_all_uses_with(o)
  warp_ctx = dataclasses.replace(ctx, thread_semantics=utils.ThreadSubset.WARP)
  # 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.auto_barriers:
    utils.warpgroup_barrier()
  ip = ir.InsertionPoint.current
  for op in op.body.operations:
    op.detach_from_parent()
    ip.insert(op)
    warp_ctx.lower_op(op)
  # We need to ensure that any effects produced by one warp (e.g. async copies)
  # are observable by all other warps.
  if ctx.auto_barriers:
    utils.warpgroup_barrier()
  return []


# The metadata needed to recostruct a vector from its flattened representation.
_VectorTemplate = tuple[Sequence[int], fa.FragmentedLayout, ir.VectorType]


def _flatten_ir_values(
    values: Sequence[ir.Value], fa_layouts: Iterable[ir.Attribute]
) -> tuple[Sequence[ir.Value], Sequence[_VectorTemplate | None]]:
  """Flattens a sequence of values.

  Non-vector values are preserved as is. Vectors are mapped to fragmented
  arrays and then flattened into per-register values.

  Args:
    values: The sequence of values to flatten.
    fa_layouts: The layouts of vectors in ``values``.

  Returns:
    A tuple of (flattened values, templates). The templates are used to
    reconstruct the vectors from the per-register  values.
  """
  fa_layouts_it = iter(fa_layouts)
  result: list[ir.Value] = []
  templates: list[_VectorTemplate | None] = []
  for v in values:
    if isinstance(v.type, ir.VectorType):
      arr = _fragmented_array_from_ir(v, next(fa_layouts_it))
      result.extend(arr.registers.flat)
      templates.append((arr.registers.shape, arr.layout, ir.VectorType(v.type)))
    else:
      result.append(v)
      templates.append(None)
  return result, templates


def _unflatten_ir_values(
    flat_values: Sequence[ir.Value], templates: Sequence[_VectorTemplate | None]
) -> Sequence[ir.Value]:
  """The inverse of ``_flatten_ir_values``."""
  result = []
  flat_values_it = iter(flat_values)
  for template in templates:
    if template is None:
      result.append(next(flat_values_it))
      continue
    registers_shape, layout, vec_type = template
    value_registers = np.asarray(
        [next(flat_values_it) for _ in range(math.prod(registers_shape))],
        dtype=object,
    )
    value = fa.FragmentedArray(
        _registers=value_registers.reshape(registers_shape),
        _layout=layout,
        _is_signed=_default_is_signed(vec_type.element_type),
    )
    result.append(fragmented_array_to_ir(value, vec_type))
  return result


def _move_scf_block_to_block_with_flattened_arguments(
    ctx: LoweringContext,
    old_block: ir.Block,
    new_block: ir.Block,
    last_op_type: type[ir.OpView],
    args_template: Sequence[_VectorTemplate | None],
    *new_leading_args: ir.Value,
) -> Sequence[_VectorTemplate | None]:
  """Moves the operations from `old_block` to `new_block`.

  The input arguments to the block, if any, are flattened using the provided
  `args_template`, except for any new_leading_args which are simply prepended
  to the flattened arguments and must be part of the template.

  The last operation of the old block must be of type `last_op_type` which
  is expected to be either a `scf.YieldOp` or a `scf.ConditionOp`. This
  operation is recreated with flattened output arguments.
  """
  out_template = None
  with ir.InsertionPoint(new_block):
    new_carry = _unflatten_ir_values(new_block.arguments[len(new_leading_args):], args_template)
    new_args = new_leading_args + tuple(new_carry)
    for old_arg, new_arg in zip(old_block.arguments, new_args, strict=True):
      old_arg.replace_all_uses_with(new_arg)
    for op in [*old_block]:
      if not isinstance(op, last_op_type):
        # `append` moves the operation.
        new_block.append(op)
        ctx.lower_op(op)
      else:
        assert out_template is None
        layouts = (
            inference_utils.in_layouts(op)
            if inference_utils.has_in_layouts_set(op)
            else []
        )
        if isinstance(op, scf.YieldOp):
          flat_operands, out_template = _flatten_ir_values(op.operands, layouts)
          scf.yield_(flat_operands)
        elif isinstance(op, scf.ConditionOp):
          flat_carry, out_template = _flatten_ir_values(op.args, layouts)
          scf.condition(op.condition, flat_carry)
        else:
          raise NotImplementedError(f"Unsupported op type: {op}")
        op.erase()
  assert out_template is not None
  return out_template


@_register_lowering(scf.ForOp, support_warp_semantics=True)
def _for_op_lowering_rule(
    ctx: LoweringContext, for_op: scf.ForOp
) -> MlirLoweringRuleResult:
  if not inference_utils.should_have_layout(for_op):
    return _traverse_op_lowering_rule(ctx, for_op)
  in_layouts = inference_utils.in_layouts(for_op)
  out_layouts = inference_utils.out_layouts(for_op)
  yield_op = for_op.body.operations[len(for_op.body.operations) - 1]
  yield_layouts = inference_utils.in_layouts(yield_op)
  if in_layouts != out_layouts or in_layouts != yield_layouts:
    raise ValueError("Layout mismatch")

  flat_init_args, args_template = _flatten_ir_values(
      for_op.initArgs, in_layouts
  )
  new_for_op = scf.ForOp(
      for_op.lowerBound,
      for_op.upperBound,
      for_op.step,
      flat_init_args,
  )

  _move_scf_block_to_block_with_flattened_arguments(
      ctx,
      for_op.body,
      new_for_op.body,
      scf.YieldOp,
      args_template,
      new_for_op.induction_variable,
  )

  return _unflatten_ir_values(new_for_op.results, args_template)


@_register_lowering(scf.WhileOp)
def _while_op_lowering_rule(
    ctx: LoweringContext, while_op: scf.WhileOp
) -> MlirLoweringRuleResult:
  if not inference_utils.should_have_layout(while_op):
    return _traverse_op_lowering_rule(ctx, while_op)

  before_block = while_op.before.blocks[0]
  after_block = while_op.after.blocks[0]
  condition_op = before_block.operations[len(before_block.operations) - 1]
  yield_op = after_block.operations[len(after_block.operations) - 1]

  in_layouts = (
      inference_utils.in_layouts(while_op)
      if inference_utils.should_have_in_layout(while_op)
      else []
  )
  out_layouts = (
      inference_utils.out_layouts(while_op)
      if inference_utils.should_have_out_layout(while_op)
      else []
  )

  if in_layouts:
    yield_layouts = inference_utils.in_layouts(yield_op)
    if in_layouts != yield_layouts:
      raise ValueError(
          f"Input layouts {in_layouts} do not match yield layouts"
          f" {yield_layouts}"
      )

  if out_layouts:
    condition_layouts = inference_utils.in_layouts(condition_op)
    if out_layouts != condition_layouts:
      raise ValueError(
          f"Output layouts {out_layouts} do not match condition layouts"
          f" {condition_layouts}"
      )

  flat_inits, inits_template = _flatten_ir_values(while_op.inits, in_layouts)
  result_types = _infer_flat_result_types(while_op, out_layouts)
  new_while_op = scf.WhileOp(result_types, flat_inits)

  # Before block
  init_types = [v.type for v in flat_inits]
  new_before_block = new_while_op.before.blocks.append(*init_types)
  results_template = _move_scf_block_to_block_with_flattened_arguments(
      ctx,
      before_block,
      new_before_block,
      scf.ConditionOp,
      inits_template,
  )

  # After block
  new_after_block = new_while_op.after.blocks.append(*result_types)
  _move_scf_block_to_block_with_flattened_arguments(
      ctx,
      after_block,
      new_after_block,
      scf.YieldOp,
      results_template,
  )

  return _unflatten_ir_values(new_while_op.results, results_template)


def _infer_flat_result_types(
    op: ir.OpView, out_layouts: Sequence[ir.Attribute]
) -> Sequence[ir.Type]:
  result_types: list[ir.Type] = []
  out_layouts_it = iter(out_layouts)
  for r in op.results:
    if not isinstance(r.type, ir.VectorType):
      result_types.append(r.type)
      continue
    vec_type = ir.VectorType(r.type)
    layout = layouts_lib.from_layout_attr(next(out_layouts_it))
    reg_type: ir.Type = layout.registers_element_type(vec_type.element_type)
    result_types.extend(
        [reg_type] * math.prod(layout.registers_shape(tuple(vec_type.shape)))
    )
  return result_types


@_register_lowering(scf.IfOp, support_warp_semantics=True)
def _if_op_lowering_rule(
    ctx: LoweringContext, if_op: scf.IfOp
) -> MlirLoweringRuleResult:
  if not inference_utils.should_have_layout(if_op):
    return _traverse_op_lowering_rule(ctx, if_op)

  raise NotImplementedError


@_register_lowering(scf.IndexSwitchOp, support_warp_semantics=True)
def _index_switch_op_lowering_rule(
    ctx: LoweringContext, switch_op: scf.IndexSwitchOp
) -> MlirLoweringRuleResult:
  if not inference_utils.should_have_layout(switch_op):
    return _traverse_op_lowering_rule(ctx, switch_op)

  out_layouts = inference_utils.out_layouts(switch_op)
  new_switch_op = scf.IndexSwitchOp(
      _infer_flat_result_types(switch_op, out_layouts),
      switch_op.arg,
      switch_op.cases,
  )

  results_template: Sequence[_VectorTemplate | None] = []
  for region, new_region in zip(
      switch_op.regions, new_switch_op.regions, strict=True
  ):
    [block] = region.blocks
    new_block = new_region.blocks[0]
    results_template = _move_scf_block_to_block_with_flattened_arguments(
        ctx, block, new_block, scf.YieldOp, []
    )
  return _unflatten_ir_values(new_switch_op.results, results_template)

@_register_lowering(mgpu.TryClusterCancelOp)
def _try_cluster_cancel_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.TryClusterCancelOp
) -> Sequence[ir.Value]:
  barrier = utils.DialectBarrierRef.from_barrier_memref(op.barrier)
  predicate = ctx.single_lane_predicate
  if op.predicate is not None:
    predicate = arith.andi(predicate, op.predicate)
  utils.try_cluster_cancel(op.cancellation_result, barrier.barrier_ref, predicate)
  return []


@_register_lowering(mgpu.QueryClusterCancelOp)
def _query_cluster_cancel_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.QueryClusterCancelOp
) -> Sequence[ir.Value]:
  del ctx
  return utils.query_cluster_cancel(op.cancellation_result)


@_register_lowering(func.FuncOp)
@_register_lowering(gpu.LaunchOp)
def _traverse_op_lowering_rule(
    ctx: LoweringContext, op: ir.OpView
) -> MlirLoweringRuleResult:
  if inference_utils.should_have_layout(op):
    raise ValueError(
        f"Rule cannot handle an op with vector operands or results: {op}"
    )
  for region in op.operation.regions:
    for block in region:
      for block_op in list(block):
        with ir.InsertionPoint(block_op):
          ctx.lower_op(block_op)
  return RECURSED


def _should_lower(op: ir.OpView) -> bool:
  """Returns 'true' if the operation should be lowered."""
  return (
      # pyrefly: ignore[missing-attribute]
      op.OPERATION_NAME.startswith("mosaic_gpu.")
      or inference_utils.should_have_layout(op)
      or inference_utils.should_have_transforms(op)
      or inference_utils.should_have_tmem_layout(op)
      # Does it have subblocks?
      or any(bool(b) for r in op.regions for b in r)
  )


def _gpu_launch_op(module: ir.Module) -> gpu.LaunchOp:
  for op in module.body.operations:
    for region in op.operation.regions:
      for block in region.blocks:
        for sub_op in block.operations:
          if isinstance(sub_op, gpu.LaunchOp):
            return sub_op
  raise ValueError("gpu.launch op not found.")


def _lowering_context(
    module: ir.Module,
    launch_context: lc.LaunchContext,
    auto_barriers: bool,
) -> LoweringContext:
  """Returns a `LoweringContext` for the given `LaunchContext`."""
  gpu_launch_op = _gpu_launch_op(module)
  with ir.InsertionPoint.at_block_begin(gpu_launch_op.regions[0].blocks[0]):
    eq = arith.CmpIPredicate.eq
    i32 = ir.IntegerType.get_signless(32)
    single_warp_per_block_predicate = arith.cmpi(
        eq, utils.warp_idx(sync=False), utils.c(0, i32)
    )
    smem_size = gpu_launch_op.dynamicSharedMemorySize
    assert smem_size is not None
    assert isinstance(smem_size.owner, arith.ConstantOp)
    smem_size = ir.IntegerAttr(smem_size.owner.value).value
    return LoweringContext(
        launch_context,
        utils.single_thread_predicate(scope=utils.ThreadSubset.WARP),
        utils.single_thread_predicate(scope=utils.ThreadSubset.WARPGROUP),
        utils.single_thread_predicate(scope=utils.ThreadSubset.BLOCK),
        single_warp_per_block_predicate,
        auto_barriers,
        smem_size,
    )


def lower_mgpu_dialect(
    module: ir.Module,
    launch_context: lc.LaunchContext,
    auto_barriers: bool = True,
):
  # TODO(apaszke,bchetioui): Make sure the layouts match.
  # TODO(bchetioui): rethink this API. It doesn't make sense to pass in a full
  # module and to traverse all `gpu.LaunchOp`s if we have a `LaunchContext` that
  # references a single `gpu.LaunchOp`.
  #
  # A `LaunchContext` should have all the information needed to lower a single
  # kernel.
  module.context.append_dialect_registry(mlir_interpreter.upstream_dialects)
  module.context.load_all_available_dialects()
  ctx = _lowering_context(module, launch_context, auto_barriers)
  with ir.InsertionPoint(module.body):
    for op in list(module.body):
      ctx.lower_op(op)