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

"""Layout and transform inference pass for the MLIR Mosaic GPU dialect."""

from __future__ import annotations

from collections.abc import Callable, Iterator, Sequence
import dataclasses
import enum
import itertools
import logging
import math
import re
from typing import Any, assert_never, cast

from jax._src import traceback_util
from jax._src.lib import mosaic_gpu_dialect as mgpu  # noqa: F401
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import arith
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._src.pallas.mosaic import error_handling as error
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

logger = logging.getLogger(__name__)

# This value was arrived at by looking at an existing kernel where layout
# inference would never be able to complete successfully, and kernels where it
# would, as well as existing tests as of 2025-11-03. We observed the following:
#
#   1. all tests would pass with a fuel that is at least ~15_000;
#   2. the kernel for which layout inference fails would fail in less than 12
#      seconds when using a fuel of 100_000.
#
# All in all, this seems like a reasonable compromise: the value is high
# enough that we can comfortably find a solution to even the most complicated
# layout inference problems that we have seen so far, but the runtime is fast
# enough that users will not waste much time waiting for a never-ending pass to
# complete when the system is unable to find a solution.
_DEFAULT_LAYOUT_INFERENCE_FUEL = 100_000


class VariableType(enum.IntEnum):
  """The type of a variable.

  Variables are operands, results, or arguments of MLIR operations.
  """
  OPERAND = 0
  RESULT = 1
  ARGUMENT = 2


_op_name_regex = re.compile(r"^(%\d+ = )?\S+")


@dataclasses.dataclass(frozen=True)
class ValueSite:
  """A unique identifier for a variable.

  This class describes a particular role of a Value, either as a result of an
  operation, an operand of an operation, or a block argument.
  """
  # A MLIR operation. If the type is `ARGUMENT`, this is the owner of the block
  # and region_index is the region that contains the block with the argument.
  # The block is always the first block of the region.
  operation: ir.OpView
  # Whether this represents an operand, a result, or an argument.
  type: VariableType
  # The index of the operand/result/argument within the op's
  # operands/results/arguments.
  index: int
  # The index of the region that contains the block with the argument.
  region_index: int | None = None

  def __post_init__(self):
    assert (self.type != VariableType.ARGUMENT) == (self.region_index is None)

  @property
  def value(self) -> ir.Value:
    """Returns the IR value corresponding to this value site."""
    if self.type == VariableType.OPERAND:
      return self.operation.operands[self.index]
    elif self.type == VariableType.RESULT:
      return self.operation.results[self.index]
    else:
      return self.operation.regions[self.region_index].blocks[0].arguments[self.index]

  @property
  def shape(self) -> tuple[int, ...]:
    """Returns the shape of the underlying value."""
    return tuple(self.value.type.shape)

  @property
  def memory_space(self) -> cs.MemorySpace:
    """Returns the memory space associated with this value."""
    ty = self.value.type
    if isinstance(ty, ir.VectorType):
      return cs.MemorySpace.REG
    assert isinstance(ty, ir.MemRefType)
    if utils.is_tmem_ref(ty):
      return cs.MemorySpace.TMEM
    elif utils.is_smem_ref(ty) or utils.is_cluster_smem_ref(ty):
      return cs.MemorySpace.SMEM
    raise ValueError(f"Unsupported memory space for: {ty}")

  def __str__(self):
    match = _op_name_regex.match(str(self.operation))
    assert match is not None
    if self.type == VariableType.OPERAND:
      return f"{match.group(0)}:o-{self.index}"
    elif self.type == VariableType.RESULT:
      return f"{match.group(0)}:r-{self.index}"
    else:
      return f"{match.group(0)}:a-{self.index}"


def is_hopper(arch: tuple[int, int]) -> bool:
  return arch == (9, 0)


def extract_assignment_candidates_from_reduce_equation(
    small: cs.RegisterLayout,
    large: cs.Variable,
    reduction_dims: tuple[int, ...],
    keep_dims: bool,
) -> Iterator[cs.RegisterLayout]:
  """Yields layout candidates for the reduce equation `small = reduce(large, reduction_dims)."""
  large_shape = large.shape

  if isinstance(small.value, fa.WGSplatFragLayout):
    yield cs.RegisterLayout(fa.WGSplatFragLayout(large_shape))
    return

  if isinstance(small.value, fa.WGStridedFragLayout):
    layout = fa.WGStridedFragLayout(large_shape, small.value.vec_size)
    yield cs.RegisterLayout(layout)
    return

  assert isinstance(small.value, fa.TiledLayout)
  # TODO(allanrenucci): Add support for reducing tiled layouts when keep_dims=True.
  if keep_dims:
    return

  candidates = [
      fa.WGMMA_LAYOUT,
      fa.WGMMA_TRANSPOSED_LAYOUT,
      fa.TCGEN05_LAYOUT,
      fa.TCGEN05_TRANSPOSED_LAYOUT,
      tcgen05.TMEM_NATIVE_LAYOUT,
  ]
  if large_shape[-1] % 16 == 0:
    candidates.append(tcgen05.fa_m64_collective_layout(large_shape[-1]))

  # In the case where we are not actually reducing any of the tiled dimensions,
  # we should consider the "small" layout as a candidate. We consider it last,
  # since it is always a valid assignment when reducing a leading tiled
  # dimension. Using it as a candidate early makes no theoretical difference,
  # but it practice could lead to longer-winded backtracking.
  candidates.append(small.value)

  for candidate in candidates:
    if len(candidate.base_tile_shape) > len(large_shape):
      continue
    num_untiled_dims = len(large_shape) - len(candidate.base_tile_shape)
    reduced_tiling_axes = tuple(
        a - num_untiled_dims for a in reduction_dims if a >= num_untiled_dims
    )
    if candidate.reduce(reduced_tiling_axes) == small.value:
      yield cs.RegisterLayout(candidate)


def _strided_layout_for_variable(
    variable: cs.Variable,
) -> fa.WGStridedFragLayout | None:
  """Returns a strided layout for the given variable.

  If the given variable cannot have a strided layout, returns `None`.
  """
  ty = variable.key.value.type
  assert isinstance(ty, ir.VectorType)
  return fa.WGStridedFragLayout.from_shaped_type(ty)


def _default_tmem_layout_for_variable(
    variable: cs.Variable,
) -> tcgen05.TMEMLayout | None:
  """Returns a default TMEM layout for the given variable, if one is defined."""
  value = variable.key.value
  parent = value.owner
  if isinstance(parent, mgpu.TmemAllocOp):
    return tcgen05._infer_tmem_layout(
        tuple(value.type.shape), bool(parent.collective), packing=1
    )
  return None


def _register_layouts_for_optimized_transfer_to_smem(
    shaped_type: ir.ShapedType,
    smem_layout: cs.SMEMTransforms,
    arch: tuple[int, int],
) -> Iterator[fa.FragmentedLayout]:
  """Yields register layout candidates for optimized transfers to SMEM."""
  if smem_layout.tiling is None:
    reg_layout = fa.WGStridedFragLayout.from_shaped_type(shaped_type)
    if reg_layout is not None:
      yield reg_layout
    return

  if is_hopper(arch):
    candidate_layouts = [
        fa.WGMMA_LAYOUT,
        fa.WGMMA_TRANSPOSED_LAYOUT,
    ]
  else:
    # For now, just assume that if it's not Hopper, it's Blackwell.
    candidate_layouts = [
        # Try the layouts with larger base tiles first.
        fa.TCGEN05_LAYOUT,
        fa.TCGEN05_TRANSPOSED_LAYOUT,
        # Keep using WGMMA and WGMMA_TRANSPOSED layouts here, simply because
        # they may apply to smaller shapes where TCGEN05 layouts do not apply.
        # This can be useful for kernels not involving MMAs that still need
        # optimized transfers to TiledLayouts, and actually shows up in some
        # tests.
        fa.WGMMA_LAYOUT,
        fa.WGMMA_TRANSPOSED_LAYOUT,
    ]

  yield from candidate_layouts


def _conjure_transforms_for_smem_ref(
    variable: cs.Variable,
    divide_constraints_per_var: dict[cs.Variable, cs.Divides]
) -> Iterator[tuple[tuple[int, ...], cs.Swizzle]]:
  ref_ty = variable.key.value.type
  if len(ref_ty.shape) < 2:
    return
  bitwidth = utils.bitwidth(ref_ty.element_type)
  strides, _ = ref_ty.get_strides_and_offset()
  rank = len(strides)
  dim_order = np.argsort(strides)
  divide_constraint = divide_constraints_per_var.get(variable)

  # We want to tile only the last two dimensions.
  if {dim_order[0], dim_order[1]} != {rank - 1, rank - 2}:
    return

  minor_dim = ref_ty.shape[dim_order[0]]
  second_to_minor_dim = ref_ty.shape[dim_order[1]]

  # The second to minor dimension must be tileable by 8.
  if second_to_minor_dim % 8 != 0:
    return

  transposed = dim_order[0] != rank - 1
  for swizzle in [128, 64, 32]:
    swizzle_elems = 8 * swizzle // bitwidth
    if minor_dim % swizzle_elems == 0:
      tiling = (swizzle_elems, 8) if transposed else (8, swizzle_elems)
      if divide_constraint is not None:
        tiling = cs.merge_divides_constraints(
            divide_constraint, cs.Divides(variable, tiling)
        ).tiling_multiple
        constrained_swizzle_elems = tiling[0] if transposed else tiling[1]
        constrained_swizzle = constrained_swizzle_elems * bitwidth // 8
        if constrained_swizzle not in {128, 64, 32}:
          continue
        yield tiling, constrained_swizzle
      else:
        yield tiling, swizzle  # pyrefly: ignore[invalid-yield]


def _extract_layout_candidates_from_tmem_registers_transfer(
    constraint: cs.IsTransferableTmemRegisters,
) -> Iterator[tuple[cs.Variable, cs.Constant]]:
  match constraint.source, constraint.target:
    case cs.Variable() as src, tgt if isinstance(tgt, cs.Constant):
      variable, constant = src, tgt
    case src, cs.Variable() as tgt if isinstance(src, cs.Constant):
      variable, constant = tgt, src
    case _:
      return

  columns = constraint.shape[1]

  if isinstance(constant, cs.TMEMLayout):
    candidates = [
      fa.TCGEN05_LAYOUT,
      constant.value.as_tiled_layout(),
      fa.TMEM_NATIVE_LAYOUT,
      fa.WGMMA_LAYOUT,
    ]
    if columns % 16 == 0:
      candidates.append(tcgen05.fa_m64_collective_layout(columns))
    for candidate in candidates:
     if constraint.is_valid_tmem_transfer(constant.value, candidate):
       yield variable, cs.RegisterLayout(candidate)
    return

  assert isinstance(constant, cs.RegisterLayout)
  if not isinstance(constant.value, fa.TiledLayout):
    return
  for packing in dict.fromkeys([32 // constraint.bitwidth, 1]):
    candidates = [tcgen05.tmem_default_layout(packing)]
    if columns % 16 == 0 and packing <= columns // 2:
      candidates.append(tcgen05.tmem_half_lane_layout(columns, packing))
    if columns % 16 == 0 and packing <= 8:
      candidates.append(tcgen05.tmem_m64_collective_layout(columns, packing))
    for candidate in candidates:
      if constraint.is_valid_tmem_transfer(candidate, constant.value):
        yield variable, cs.TMEMLayout(candidate)


def _extract_layout_candidates_from_smem_registers_transfer(
    constraint: cs.IsTransferableSmemRegisters,
    division_constraint_per_var: dict[cs.Variable, cs.Divides],
    arch: tuple[int, int]
) -> Iterator[tuple[cs.Variable, cs.Constant]]:
  src, tgt = constraint.source, constraint.target
  match src, tgt:
    case cs.Variable(), cst if isinstance(cst, cs.Constant):
      variable, constant = src, tgt
    case cst, cs.Variable() if isinstance(cst, cs.Constant):
      variable, constant = tgt, src
    case _:
      return

  assert isinstance(variable, cs.Variable)  # Satisfy type checkers.
  if isinstance(constant, cs.RegisterLayout):
    layout = constant.value
    assert variable.memory_space == cs.MemorySpace.SMEM
    if isinstance(layout, fa.TiledLayout) and len(variable.shape) >= 2:
      # Maintain a set of yielded tilings to avoid duplicates caused by existing
      # divides constraints.
      yielded = set()
      for transform in _conjure_transforms_for_smem_ref(
          variable, division_constraint_per_var
      ):
        tiling, swizzle = transform
        if transform in yielded:
          continue
        yielded.add(transform)
        yield variable, cs.SMEMTransforms(lc.TileTransform(tiling), swizzle)
    return

  assert isinstance(constant, cs.SMEMTransforms)
  assert variable.memory_space == cs.MemorySpace.REG
  for layout in _register_layouts_for_optimized_transfer_to_smem(
      variable.key.value.type, constant, arch
  ):
    yield variable, cs.RegisterLayout(layout)


def mma_tiling_constraint(
    var: cs.Variable,
    op: ir.OpView,
    allow_unswizzled: bool = False,
) -> cs.Constraint:
  """Returns a constraint stating that `var` has a valid MMA SMEM tiling.

  For both tcgen05.mma and wgmma, tiling is valid if it is of the form
  (8, swizzle_elems), with
      swizzle_elems in {s * 8 // dtype_bitwidth for s in [32, 64, 128]},
  as support for unswizzled tilings is not yet supported.

  If `allow_unswizzled` is True, then we additionally accept
  (8, 16 * 8 // dtype_bitwidth) as a valid tiling.

  If the reference is transposed, the tilings are transposed to
  (swizzle_elems, 8).
  """
  ref_type = ir.MemRefType(var.key.value.type)
  assert len(ref_type.shape) == 2
  bitwidth = utils.bitwidth(ref_type.element_type)
  is_transposed = utils.is_memref_transposed(ref_type)
  allowed: list[cs.Constant] = []
  swizzles = (128, 64, 32, None) if allow_unswizzled else (128, 64, 32)
  tiling = None
  for swizzle in swizzles:
    swizzle_elems = (swizzle or 16) * 8 // bitwidth
    tiling = (swizzle_elems, 8) if is_transposed else (8, swizzle_elems)
    if all(s % t == 0 for s, t in zip(var.shape, tiling)):
      allowed.append(cs.SMEMTransforms(lc.TileTransform(tiling), swizzle))
  if not allowed:
    assert tiling is not None
    raise _construct_value_error_with_op_stacktrace(
        f"No valid tiling found for SMEM reference. {var.shape} not"
        f" divisible by candidate tiling {tiling}",
        op,
    )
  return cs.OneOf(var, tuple(allowed))


def _divides_per_var(
    constraints: Sequence[cs.Constraint],
) -> dict[cs.Variable, cs.Divides]:
  result: dict[cs.Variable, cs.Divides] = {}
  for constraint in constraints:
    if isinstance(constraint, cs.Divides) and isinstance(
        constraint.expr, cs.Variable
    ):
      assert constraint.expr not in result
      result[constraint.expr] = constraint
  return result


def _extract_layout_candidates_from_broadcast(
    src: cs.RegisterLayout,
    dst: cs.Variable,
    dims: tuple[int, ...],
) -> Iterator[tuple[cs.Variable, cs.Constant]]:
  """Yields layout candidates for a broadcast equation."""
  match src.value:
    case fa.WGSplatFragLayout():
      yield dst, cs.RegisterLayout(fa.WGSplatFragLayout(dst.shape))
    case fa.WGStridedFragLayout() as src:
      dst_layout = fa.WGStridedFragLayout(dst.shape, src.vec_size)
      if fa.is_supported_strided_layout_broadcast(src, dst_layout, dims):
        yield dst, cs.RegisterLayout(dst_layout)


# TODO(bchetioui): flatten this call hierarchy.
def _extract_variable_assignments_from_constraints(
    constraints: Sequence[cs.Constraint], arch: tuple[int, int],
) -> Iterator[tuple[cs.Variable, cs.Constant]]:
  """Attempts to extract variable assignments from all constraints."""
  dpv = _divides_per_var(constraints)
  def priority(constraint: cs.Constraint) -> int:
    match constraint:
      # We want to minimize the number of relayouts in the kernel, so we first
      # try to satisfy relayout constraints via identity relayouts.
      case cs.Relayout():
        return 0  # Highest priority
      case _:
        return 1
  for c in sorted(constraints, key=priority):
    match c:
      case cs.IsTransferableTmemRegisters():
        yield from _extract_layout_candidates_from_tmem_registers_transfer(c)
      case cs.IsTransferableSmemRegisters():
        yield from _extract_layout_candidates_from_smem_registers_transfer(c, dpv, arch)
      case cs.Equals(cs.Reduce(cs.Variable() as large, axes=axes, keep_dims=keep_dims), cs.RegisterLayout() as small):
        for layout in extract_assignment_candidates_from_reduce_equation(small, large, axes, keep_dims):
          yield large, layout
      case cs.Equals(cs.RegisterLayout() as small, cs.Reduce(cs.Variable() as large, axes=axes, keep_dims=keep_dims)):
        for layout in extract_assignment_candidates_from_reduce_equation(small, large, axes, keep_dims):
          yield large, layout
      case cs.Relayout(cs.Variable() as var, cs.RegisterLayout() as layout):
        yield var, layout
      case cs.Relayout(cs.RegisterLayout() as layout, cs.Variable() as var):
        yield var, layout
      case cs.IsSupportedBroadcast(cs.RegisterLayout() as src, cs.Variable() as dst, dims=dims):
        yield from _extract_layout_candidates_from_broadcast(src, dst, dims)
      case cs.OneOf(expr=cs.Variable() as var, allowed=allowed_constants):
        for const in allowed_constants:
          yield var, const


def conjure_assignment(
    unknowns: Sequence[cs.Variable],
    constraint_system: cs.ConstraintSystem,
    arch: tuple[int, int],
) -> Iterator[tuple[cs.Variable, cs.Constant]]:
  """Attempts to conjure an assignment for an unknown variable."""
  # TODO(allanrenucci): We should be able to short-circuit the search here if
  # the constraint is not satisfiable.

  # As we extract assignment candidates from constraints, we prioritize
  # candidates that are more "interesting"; e.g., in the case of registers,
  # introducing splat layout candidate assignments often leads to a dead end in
  # practice---as opposed to tiled layouts, which are more likely to yield
  # solutions to the constraint system.
  low_priority_assignments: list[tuple[cs.Variable, cs.Constant]] = []
  for variable, constant in _extract_variable_assignments_from_constraints(
      constraint_system.constraints, arch
  ):
    match constant:
      case cs.RegisterLayout(value=value) if not isinstance(
          value, fa.TiledLayout
      ):
        low_priority_assignments.append((variable, constant))
      case _:
        yield variable, constant

  # After all high-priority assignments have been attempted, switch to using
  # low-priority assignments.
  for variable, constant in low_priority_assignments:
    yield variable, constant

  # Here, we have not managed to find an assignment for all the unknown
  # variables. We now try to introduce new arbitrary (valid) assignments into
  # the system, and hope that they turn out to be compatible with the constraint
  # system.
  for variable in unknowns:
    if variable in constraint_system.assignments:
      continue
    # Try to instantiate a single variable to a default layout and see if it
    # reduces the system.
    match variable.memory_space:
      case cs.MemorySpace.REG:
        layout = _strided_layout_for_variable(variable)
        if layout is not None:
          yield variable, cs.RegisterLayout(layout)
      case cs.MemorySpace.SMEM:
        # TODO(olechwierowicz): Consider inferring a standalone swizzle here.
        yield variable, cs.SMEMTransforms(None, None)
      case cs.MemorySpace.TMEM:
        layout = _default_tmem_layout_for_variable(variable)
        if layout is not None:
          yield variable, cs.TMEMLayout(layout)
      case never:
        assert_never(never)


def find_assignments_for(
    unknowns: Sequence[cs.Variable],
    constraint_system: cs.ConstraintSystem,
    *,
    fuel: int,
    arch: tuple[int, int],
) -> tuple[dict[cs.Variable, cs.Constant] | cs.Unsatisfiable, int]:
  """Attempts to find assignments that satisfy `constraint_system` for `unknowns`.

  Args:
    unknowns: The set of variables that are unknown. Represented as a sequence
      of `Variable`s for determinism purposes.
    constraint_system: the constraint system to satisfy.
    fuel: The fuel to use for the search. Once the fuel is exhausted, we raise
      an error.
    arch: The architecture to target in the search.

  Returns:
    A tuple where the first element is the solution, and the second element is
    the fuel remaining after the search. The solution is either:
      - Unsatisfiable() if the constraint system has unsatisfiable constraints.
      - A dictionary assigning all the unknown variables to
        `ConstantExpression`s such that the assignment satisfies the constraint
        system otherwise.
  """
  if isinstance(result := cs.reduce(constraint_system), cs.Unsatisfiable):
    return cs.Unsatisfiable(), fuel

  constraint_system = result
  remaining_unknowns = [
      u for u in unknowns if u not in constraint_system.assignments.keys()
  ]

  # In this case, we have determined an assignment for all the unknown
  # variables. Return their respective assignment.
  if not remaining_unknowns:
    assert not constraint_system.constraints, (
        "A satisfiable system should not have remaining unsatisfied"
        " constraints. This is a bug."
    )
    return {u: constraint_system.assignments[u] for u in unknowns}, fuel

  # If unknowns remain and we have fully reduced the system, we may still
  # be able to make progress by trying out potential assignments. These
  # new assignments could make the system unsatisfiable, so we use a recursive
  # call to be able to backtrack if necessary.
  for assignment in conjure_assignment(
      remaining_unknowns, constraint_system, arch
  ):
    if fuel <= 0:
      raise ValueError(
          "Layout inference failed to find a solution. Consider adding layout "
          "annotations to your program to guide the search."
      )
    variable, expr = assignment
    assert isinstance(expr, cs.Constant)
    if not cs.is_valid_assignment(variable, expr):
      continue
    # Trying one valid assignment consumes fuel.
    fuel -= 1
    new_constraint_system = (
        cs.ConstraintSystem(assignments={variable: expr}) & constraint_system
    )
    if isinstance(new_constraint_system, cs.Unsatisfiable):
      # This assignment is not compatible with the constraint system.
      continue
    solution, fuel = find_assignments_for(
        unknowns, new_constraint_system, fuel=fuel, arch=arch
    )
    if not isinstance(solution, cs.Unsatisfiable):
      return solution, fuel

  # TODO(bchetioui): should we have a way to give a useful dump to the user
  # here, perhaps indicating what to layout cast.
  return cs.Unsatisfiable(), fuel


@dataclasses.dataclass()
class DerivationContext:
  """Holds context information used for deriving an constraint system."""
  # A map of `ValueSite` to the variable that it is associated with.
  variable_for_value_site: dict[ValueSite, cs.Variable] = dataclasses.field(
      default_factory=dict, init=False
  )
  # A map of `cs.Variable` to all the `ValueSite`s that it is associated with.
  value_sites_for_variable: ValueSitesForVariable = (
      dataclasses.field(default_factory=dict, init=False)
  )

  # A `slice_smem` or `slice_tmem` op may carry an optional `alias_id`. It is
  # used to disambiguate whether overlapping memory slices are independent or
  # related for the purposes of layout inference.
  #
  # Concretely, if a memory slice op shares the same `alias_id` with another
  # memory slice op, then they are considered to be the same ref and must have
  # the same layout. Otherwise, they are independent.
  #
  # It is incorrect to annotate non-overlapping refs or refs with incompatible
  # result types with the same `alias_id`.
  slice_smem_aliases: dict[int, cs.Variable] = dataclasses.field(
      default_factory=dict, init=False
  )
  slice_tmem_aliases: dict[int, cs.Variable] = dataclasses.field(
      default_factory=dict, init=False
  )

  # For each constraint, records the operation it was derived from.
  op_for_constraint: dict[cs.Constraint, ir.OpView] = dataclasses.field(
      default_factory=dict, init=False
  )

  def update(self, mapping: ValueSitesForVariable) -> None:
    for variable, value_sites in mapping.items():
      if variable in self.value_sites_for_variable:
        self.value_sites_for_variable[variable].extend(value_sites)
      else:
        self.value_sites_for_variable[variable] = value_sites
      for value_site in value_sites:
        assert value_site not in self.variable_for_value_site
        self.variable_for_value_site[value_site] = variable

  def producer_ref(self, operand: ValueSite) -> cs.Variable:
    """Returns the producer reference variable for the given operand."""
    return self.variable_for_value_site[producer_result(operand)]


ValueSitesForVariable = dict[cs.Variable, list[ValueSite]]

# A constraint system derivation rule is a function that takes an MLIR operation
# and returns a constraint system, and a mapping from variables to value site
# identifiers.
#
# The intended meaning of the mapping is that, for each identifier in the list
# keyed by a given variable, the MLIR operand/result/argument corresponding to
# that identifier has the same layout as the variable.
#
# A `ConstraintSystemDerivationRule` must return a mapping such that the
# identifier corresponding to each value site must appear in the mapping,
# and each identifier in the mapping must be keyed by exactly one variable.
# Lastly, the mapping must only refer to variables and
# operands/results/arguments that correspond to the given operation.
ConstraintSystemDerivationRuleResult = tuple[
    cs.ConstraintSystem, ValueSitesForVariable
]
ConstraintSystemDerivationRule = Callable[
    [DerivationContext, ir.OpView],
    ConstraintSystemDerivationRuleResult,
]
_constraint_system_derivation_rules: dict[
    str, ConstraintSystemDerivationRule
] = {}


def _add_constraint_system_derivation_rule(op: type[ir.OpView]):
  def wrapper(rule: ConstraintSystemDerivationRule):
    if op is not None:
      assert hasattr(op, "OPERATION_NAME")
      _constraint_system_derivation_rules[op.OPERATION_NAME] = rule
    return rule

  return wrapper


def is_vector(v: ir.Value) -> bool:
  return isinstance(v.type, ir.VectorType)


def _is_smem_ref(v: ir.Value) -> bool:
  return isinstance(v.type, ir.MemRefType) and (
      utils.is_smem_ref(v) or utils.is_cluster_smem_ref(v)
  )


def _is_tmem_ref(v: ir.Value) -> bool:
  return isinstance(v.type, ir.MemRefType) and utils.is_tmem_ref(v)


def _pointwise_op_constraint_system(
    ctx: DerivationContext,
    op: ir.OpView,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  all_value_sites = vector_value_sites(op)
  variable = cs.Variable(all_value_sites[-1])
  return cs.ConstraintSystem(), {variable: all_value_sites}


for _op in [
    arith.AddIOp,
    arith.AddFOp,
    arith.AndIOp,
    arith.BitcastOp,
    arith.CmpFOp,
    arith.CmpIOp,
    arith.ExtFOp,
    arith.ExtSIOp,
    arith.ExtUIOp,
    arith.FPToSIOp,
    arith.FPToUIOp,
    arith.MaximumFOp,
    arith.MaxUIOp,
    arith.MaxSIOp,
    arith.MinimumFOp,
    arith.MinUIOp,
    arith.MinSIOp,
    arith.MulIOp,
    arith.MulFOp,
    arith.NegFOp,
    arith.OrIOp,
    arith.FloorDivSIOp,
    arith.DivSIOp,
    arith.DivUIOp,
    arith.DivFOp,
    arith.RemUIOp,
    arith.RemSIOp,
    arith.RemFOp,
    arith.SIToFPOp,
    arith.UIToFPOp,
    arith.SubIOp,
    arith.SubFOp,
    arith.TruncFOp,
    arith.TruncIOp,
    arith.XOrIOp,
    arith.SelectOp,
    mlir_math.ExpOp,
    mlir_math.Exp2Op,
    mlir_math.SinOp,
    mlir_math.CosOp,
    mlir_math.LogOp,
    mlir_math.Log2Op,
    mlir_math.RsqrtOp,
    mlir_math.SqrtOp,
    mlir_math.TanhOp,
    mlir_math.AbsFOp,
    mlir_math.AbsIOp,
    mlir_math.RoundOp,
    mlir_math.RoundEvenOp,
    mlir_math.CopySignOp,
    mlir_math.CountLeadingZerosOp,
    arith.ShLIOp,
    arith.ShRUIOp,
    arith.ShRSIOp,
]:
  _add_constraint_system_derivation_rule(_op)(_pointwise_op_constraint_system)


@_add_constraint_system_derivation_rule(mgpu.VectorLoadOp)
def _vector_load_constraint_system(
    ctx: DerivationContext,
    op: mgpu.VectorLoadOp,
) -> ConstraintSystemDerivationRuleResult:
  # Registers
  dest = ValueSite(op, VariableType.RESULT, 0)
  dest_var = cs.Variable(dest)
  value_sites_for_variable = {dest_var: [dest]}
  constraints: list[cs.Constraint]
  constraints = [cs.NotOfType(dest_var, fa.WGSplatFragLayout)]

  # An unset `optimized` attribute means that we require an optimized transfer.
  optimized = op.optimized is None or bool(op.optimized)

  # SMEM
  if _is_smem_ref(op.source):
    source = ValueSite(op, VariableType.OPERAND, 0)
    source_var = ctx.producer_ref(source)
    value_sites_for_variable[source_var] = [source]
    ref_ty = ir.MemRefType(op.source.type)
    shape = tuple(ref_ty.shape)
    strides, _ = ref_ty.get_strides_and_offset()
    constraints.append(
        cs.IsTransferableSmemRegisters(
            source_var, dest_var, shape, tuple(strides),
            bitwidth=utils.bitwidth(ref_ty.element_type),
            optimized=optimized
        )
    )

  system = cs.ConstraintSystem(constraints=constraints)
  return system, value_sites_for_variable


@_add_constraint_system_derivation_rule(mgpu.MultimemLoadReduceOp)
def _multimem_load_reduce_constraint_system(
    _: DerivationContext,
    op: mgpu.MultimemLoadReduceOp,
) -> ConstraintSystemDerivationRuleResult:
  dest = ValueSite(op, VariableType.RESULT, 0)
  dest_var = cs.Variable(dest)
  system = cs.ConstraintSystem(
      constraints=[cs.NotOfType(dest_var, fa.WGSplatFragLayout)]
  )
  return system, {dest_var: [dest]}


@_add_constraint_system_derivation_rule(mgpu.VectorStoreOp)
def _vector_store_constraint_system(
    ctx: DerivationContext,
    op: Any,  # This is mgpu.VectorStoreOp | mgpu.AsyncStoreSmemOp
) -> ConstraintSystemDerivationRuleResult:
  # Registers
  value = ValueSite(op, VariableType.OPERAND, 0)
  value_var = cs.Variable(value)
  value_sites_for_variable = {value_var: [value]}

  # An unset `optimized` attribute means that we require an optimized transfer.
  optimized = op.optimized is None or bool(op.optimized)

  # SMEM
  constraints = []
  if _is_smem_ref(op.destination):
    dest = ValueSite(op, VariableType.OPERAND, 1)
    dest_var = ctx.producer_ref(dest)
    value_sites_for_variable[dest_var] = [dest]
    ref_ty = ir.MemRefType(op.destination.type)
    shape = tuple(ref_ty.shape)
    strides, _ = ref_ty.get_strides_and_offset()
    constraints.append(
        cs.IsTransferableSmemRegisters(
            value_var, dest_var, shape, tuple(strides),
            bitwidth=utils.bitwidth(ref_ty.element_type),
            optimized=optimized
        )
    )

  system = cs.ConstraintSystem(constraints=constraints)
  return system, value_sites_for_variable


@_add_constraint_system_derivation_rule(mgpu.AsyncStoreSmemOp)
def _async_store_smem_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncStoreSmemOp,
) -> ConstraintSystemDerivationRuleResult:
  system, value_sites_for_variable = _vector_store_constraint_system(ctx, op)
  var = cs.Variable(ValueSite(op, VariableType.OPERAND, 0))
  extra_constraints = cs.ConstraintSystem(
      constraints=[
          cs.NotOfType(var, fa.WGStridedFragLayout),
          cs.NotOfType(var, fa.WGSplatFragLayout),
      ]
  )
  return system & extra_constraints, value_sites_for_variable


@_add_constraint_system_derivation_rule(mgpu.DebugPrintOp)
def _debug_print_constraint_system(
    ctx: DerivationContext,
    op: mgpu.DebugPrintOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  value = ValueSite(op, VariableType.OPERAND, 0)
  return cs.ConstraintSystem(), {cs.Variable(value): [value]}


@_add_constraint_system_derivation_rule(mgpu.PrintLayoutOp)
def _print_layout_constraint_system(
    ctx: DerivationContext,
    op: mgpu.PrintLayoutOp,
) -> ConstraintSystemDerivationRuleResult:
  value = ValueSite(op, VariableType.OPERAND, 0)
  var = cs.Variable(value) if is_vector(op.value) else ctx.producer_ref(value)
  return cs.ConstraintSystem(), {var: [value]}


@_add_constraint_system_derivation_rule(mgpu.BroadcastedIotaOp)
def _broadcasted_iota_constraint_system(
    ctx: DerivationContext,
    op: mgpu.BroadcastedIotaOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  value = ValueSite(op, VariableType.RESULT, 0)
  var = cs.Variable(value)
  constraints = [cs.NotOfType(var, fa.WGSplatFragLayout)]
  return cs.ConstraintSystem(constraints=constraints), {var: [value]}


@_add_constraint_system_derivation_rule(mgpu.OptimizationBarrierOp)
def _optimization_barrier_constraint_system(
    ctx: DerivationContext,
    op: mgpu.OptimizationBarrierOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  value_sites_for_variable: ValueSitesForVariable = {}

  for i, operand in enumerate(op.operands):
    if not is_vector(operand):
      continue
    variable = cs.Variable(ValueSite(op, VariableType.OPERAND, i))
    value_sites_for_variable[variable] = [
        ValueSite(op, VariableType.OPERAND, i),
        ValueSite(op, VariableType.RESULT, i)
    ]

  return cs.ConstraintSystem(), value_sites_for_variable


@_add_constraint_system_derivation_rule(arith.ConstantOp)
def _constant_constraint_system(
    ctx: DerivationContext,
    constant_op: arith.ConstantOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  value = constant_op.value
  result = ValueSite(constant_op, VariableType.RESULT, 0)
  variable = cs.Variable(result)
  shape = tuple(ir.ShapedType(constant_op.result.type).shape)
  if (
      isinstance(value, ir.DenseElementsAttr)
      and ir.DenseElementsAttr(value).is_splat
  ):
    layout = fa.WGSplatFragLayout(shape=shape)
    system = cs.ConstraintSystem(
        assignments={variable: cs.RegisterLayout(layout)}
    )
  else:
    constant_is_not_splat = cs.NotOfType(variable, fa.WGSplatFragLayout)
    system = cs.ConstraintSystem(constraints=[constant_is_not_splat])

  return system, {variable: [result]}


def _terminator(
    block: ir.Block, expected_terminator: type[ir.OpView]
) -> ir.OpView:
  """Returns the terminator of the given block.

  Checks that the terminator is of the expected type.
  """
  terminator = block.operations[len(block.operations) - 1]
  assert isinstance(terminator, expected_terminator)
  return terminator.opview


@_add_constraint_system_derivation_rule(scf.ForOp)
def _for_constraint_system(
    ctx: DerivationContext,
    op: scf.ForOp,
) -> ConstraintSystemDerivationRuleResult:
  [block] = op.region.blocks
  yield_op = _terminator(block, scf.YieldOp)
  value_sites_for_variable: ValueSitesForVariable = {}

  # Account for the lower bound, upper bound, and step of the loop, which appear
  # in the operands but not in the results.
  num_leading_args = 3
  for index, o in enumerate(op.operands):
    if not is_vector(o) and not _is_smem_ref(o):
      continue
    result_index = index - num_leading_args
    arg_index = index - num_leading_args + 1  # Account for the induction var.
    operand = ValueSite(op, VariableType.OPERAND, index)
    arg = ValueSite(op, VariableType.ARGUMENT, arg_index, region_index=0)
    result = ValueSite(op, VariableType.RESULT, result_index)
    yield_operand = ValueSite(
        yield_op, VariableType.OPERAND, result_index
    )
    var = cs.Variable(operand) if is_vector(o) else ctx.producer_ref(operand)
    value_sites_for_variable[var] = [operand, arg, result, yield_operand]

  return cs.ConstraintSystem(), value_sites_for_variable


def prime_decomposition(n: int) -> list[int]:
  """Returns the prime decomposition of the given number `n` as a list of ints.

  A factor appears as many times in the list as the power up to which it divides
  `n`.
  """
  # This implementation should be sufficiently efficient for small `n`, which
  # should always be the case for us.
  prime_factors = []
  divisor = 2
  while divisor * divisor <= n:
    while n % divisor == 0:
      n //= divisor
      prime_factors.append(divisor)
    divisor += 1
  if n != 1:
    prime_factors.append(n)
  return prime_factors


# TODO(bchetioui): let's see if we need to parametrize this by depth.
def dynamic_gcd(a: int, b: ir.Value) -> int:
  if a <= 0:
    raise ValueError("a must be strictly positive")
  if isinstance(b.type, ir.VectorType):
    # We don't actually know the values of the vector elements, so we pick 1
    # as the only safe value.
    return 1
  if not isinstance(b.type, (ir.IntegerType, ir.IndexType)):
    raise ValueError(f"Expected an integer dynamic value, got a {b.type}")
  if isinstance(b.owner, arith.ConstantOp):
    assert isinstance(b.owner.literal_value, int)
    return math.gcd(a, b.owner.literal_value)
  running_gcd = 1
  for factor in prime_decomposition(a):
    if utils.is_known_divisible(b, running_gcd * factor):
      running_gcd *= factor
  return running_gcd


@_add_constraint_system_derivation_rule(scf.WhileOp)
def _while_constraint_system(
    ctx: DerivationContext,
    op: scf.WhileOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  [before_block] = op.before.blocks
  [after_block] = op.after.blocks
  cond_op = _terminator(before_block, scf.ConditionOp)
  yield_op = _terminator(after_block, scf.YieldOp)

  value_sites_for_variable: ValueSitesForVariable = {}

  for value_site in vector_value_sites(op):
    idx = value_site.index
    match value_site.type:
      case VariableType.OPERAND:
        arg = ValueSite(op, VariableType.ARGUMENT, idx, region_index=0)
        yield_operand = ValueSite(yield_op, VariableType.OPERAND, idx)
        value_sites_for_variable[cs.Variable(value_site)] = [
            value_site,
            arg,
            yield_operand,
        ]
      case VariableType.RESULT:
        # Increment by 1 to account for the conditional.
        cond_operand = ValueSite(cond_op, VariableType.OPERAND, idx + 1)
        arg = ValueSite(op, VariableType.ARGUMENT, idx, region_index=1)
        value_sites_for_variable[cs.Variable(value_site)] = [
            value_site,
            arg,
            cond_operand,
        ]
      case _:
        raise RuntimeError(f"Unexpected value site type: {value_site.type}")

  return cs.ConstraintSystem(), value_sites_for_variable


@_add_constraint_system_derivation_rule(scf.IndexSwitchOp)
def _index_switch_constraint_system(
    ctx: DerivationContext,
    op: scf.IndexSwitchOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  value_sites_for_variable: ValueSitesForVariable = {
      cs.Variable(o): [o] for o in vector_value_sites(op)
  }
  for region in op.regions:
    [block] = region.blocks
    yield_op = _terminator(block, scf.YieldOp)
    for var in value_sites_for_variable.keys():
      assert var.key.type == VariableType.RESULT
      yield_operand = ValueSite(yield_op, VariableType.OPERAND, var.key.index)
      value_sites_for_variable[var].append(yield_operand)

  return cs.ConstraintSystem(), value_sites_for_variable


@_add_constraint_system_derivation_rule(mgpu.LayoutCastOp)
def _layout_cast_constraint_system(
    ctx: DerivationContext,
    op: mgpu.LayoutCastOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  operand = ValueSite(op, VariableType.OPERAND, 0)
  operand_var = cs.Variable(operand)
  result = ValueSite(op, VariableType.RESULT, 0)
  result_var = cs.Variable(result)
  out_layout = cs.RegisterLayout(layouts_lib.from_layout_attr(op.new_layout))
  if not cs.is_valid_assignment(result_var, out_layout):
    raise ValueError(
        f"Cannot cast to layout {out_layout.value}: the layout is not"
        f"compatible with the operand shape {operand.shape} in {op}."
    )
  bitwidth = utils.bitwidth(op.x.type.element_type)
  # TODO(bchetioui): remove `getattr` once minimum jaxlib version is 0.12.0.
  strict = bool(getattr(op, "strict", False))
  return (
      cs.ConstraintSystem(
          assignments={result_var: out_layout},
          constraints=[
              cs.Relayout(operand_var, result_var, bitwidth, strict=strict),
          ],
      ),
      {operand_var: [operand], result_var: [result]},
  )


@_add_constraint_system_derivation_rule(mgpu.WGMMAOp)
def _wgmma_constraint_system(
    ctx: DerivationContext,
    op: mgpu.WGMMAOp,
) -> ConstraintSystemDerivationRuleResult:
  assignments: dict[cs.Variable, cs.Constant] = {}
  constraints: list[cs.Constraint] = []
  value_sites_for_variable: ValueSitesForVariable = {}

  acc_out = ValueSite(op, VariableType.RESULT, 0)
  acc_in = ValueSite(op, VariableType.OPERAND, 0)
  acc_var = cs.Variable(acc_out)
  acc_layout = cs.RegisterLayout(fa.WGMMA_LAYOUT)
  assignments[acc_var] = acc_layout
  if not cs.is_valid_assignment(acc_var, acc_layout):
    raise ValueError(
        f"Cannot assign layout {acc_layout.value} to the accumulator of a wgmma"
        f" op: the layout is not compatible with the accumulator shape"
        f" {acc_out.shape}."
    )
  value_sites_for_variable[acc_var] = [acc_in, acc_out]

  b = ValueSite(op, VariableType.OPERAND, 2)
  b_var = ctx.producer_ref(b)
  constraints.append(mma_tiling_constraint(b_var, op))
  value_sites_for_variable[b_var] = [b]

  a = ValueSite(op, VariableType.OPERAND, 1)
  if _is_smem_ref(op.a):
    a_var = ctx.producer_ref(a)
    # We expect the tiling transform to be physically the same on both sides.
    # However, the constraint system assigns tiling transforms based on the
    # logical shape. In the case the tiled dimensions of exactly one of the
    # operands are transposed, we need to transpose the transform as well.
    a_is_transposed = utils.is_memref_transposed(ir.MemRefType(op.a.type))
    b_is_transposed = utils.is_memref_transposed(ir.MemRefType(op.b.type))
    if a_is_transposed != b_is_transposed:
      constraints.append(cs.Equals(lhs=a_var, rhs=cs.Transpose(b_var, (1, 0))))
    else:
      constraints.append(cs.Equals(lhs=a_var, rhs=b_var))
  else:
    a_var = cs.Variable(a)
    if utils.bitwidth(op.a.type.element_type) == 8:
      layout = fa.WGMMA_LAYOUT_8BIT
    else:
      layout = fa.WGMMA_LAYOUT
    layout = cs.RegisterLayout(layout)
    assignments[a_var] = layout
    if not cs.is_valid_assignment(a_var, layout):
      raise ValueError(
          f"Cannot assign layout {layout.value} to the 'a' operand of WGMMAOp: "
          f"the layout is not compatible with the operand shape {a.shape}."
      )

  value_sites_for_variable[a_var] = [a]
  return cs.ConstraintSystem(assignments, constraints), value_sites_for_variable


@_add_constraint_system_derivation_rule(mgpu.MMAOp)
def _mma_constraint_system(
    ctx: DerivationContext, op: mgpu.MMAOp
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  element_type = op.a.type.element_type
  m, n = op.accumulator.type.shape
  layouts = MMALayouts.for_shape(element_type, m, n)

  assignments: dict[cs.Variable, cs.Constant] = {}
  value_sites_for_variable: ValueSitesForVariable = {}

  acc_out = ValueSite(op, VariableType.RESULT, 0)
  acc_in = ValueSite(op, VariableType.OPERAND, 0)
  acc_var = cs.Variable(acc_out)
  acc_layout = cs.RegisterLayout(layouts.acc)
  if not cs.is_valid_assignment(acc_var, acc_layout):
    raise ValueError(
        f"Cannot assign layout {acc_layout.value} to the accumulator of an mma"
        " op: the layout is not compatible with the accumulator shape"
        f" {acc_out.shape}."
    )
  assignments[acc_var] = acc_layout
  value_sites_for_variable[acc_var] = [acc_in, acc_out]

  a_site = ValueSite(op, VariableType.OPERAND, 1)
  a_var = cs.Variable(a_site)
  a_layout = cs.RegisterLayout(layouts.lhs)
  if not cs.is_valid_assignment(a_var, a_layout):
    raise ValueError(
        f"Cannot assign layout {a_layout.value} to the 'a' operand of MMAOp: "
        f"the layout is not compatible with the operand shape {a_site.shape}."
    )
  assignments[a_var] = a_layout
  value_sites_for_variable[a_var] = [a_site]

  b_site = ValueSite(op, VariableType.OPERAND, 2)
  b_var = cs.Variable(b_site)
  b_layout = cs.RegisterLayout(layouts.rhs)
  if not cs.is_valid_assignment(b_var, b_layout):
    raise ValueError(
        f"Cannot assign layout {b_layout.value} to the 'b' operand of MMAOp: "
        f"the layout is not compatible with the operand shape {b_site.shape}."
    )
  assignments[b_var] = b_layout
  value_sites_for_variable[b_var] = [b_site]

  return cs.ConstraintSystem(assignments, []), value_sites_for_variable


@_add_constraint_system_derivation_rule(vector.BroadcastOp)
def _vector_broadcast_constraint_system(
    ctx: DerivationContext,
    op: vector.BroadcastOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  # This is not expected to be necessary at the moment. We should be using
  # mgpu.BroadcastInDimOp instead when dealing with broadcasting vectors.
  if isinstance(op.source.type, ir.ShapedType):
    raise NotImplementedError("Only vector broadcasts from scalars are supported.")
  out_variable = cs.Variable(ValueSite(op, VariableType.RESULT, 0))
  layout = cs.RegisterLayout(fa.WGSplatFragLayout(tuple(op.result.type.shape)))
  return (
      cs.ConstraintSystem(assignments={out_variable: layout}),
      {out_variable: [out_variable.key]},
  )


@_add_constraint_system_derivation_rule(vector.ReductionOp)
def _vector_reduction_constraint_system(
    ctx: DerivationContext,
    op: vector.ReductionOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  in_variable = cs.Variable(ValueSite(op, VariableType.OPERAND, 0))
  return cs.ConstraintSystem(), {in_variable: [in_variable.key]}


@_add_constraint_system_derivation_rule(vector.MultiDimReductionOp)
def _multi_dim_reduction_constraint_system(
    ctx: DerivationContext,
    op: vector.MultiDimReductionOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  source = ValueSite(op, VariableType.OPERAND, 0)
  acc = ValueSite(op, VariableType.OPERAND, 1)
  out = ValueSite(op, VariableType.RESULT, 0)
  source_variable = cs.Variable(source)
  out_variable = cs.Variable(out)
  constraints = [
      cs.NotOfType(source_variable, fa.WGStridedFragLayout),
      cs.Equals(
          out_variable,
          cs.Reduce(
              source_variable, tuple(op.reduction_dims), rank=len(source.shape)
          ),
      ),
  ]
  return (
      cs.ConstraintSystem(constraints=constraints),
      {source_variable: [source], out_variable: [acc, out]},
  )


@_add_constraint_system_derivation_rule(mgpu.BroadcastInDimOp)
def _broadcast_in_dim_constraint_system(
    ctx: DerivationContext,
    op: mgpu.BroadcastInDimOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  src_variable = cs.Variable(ValueSite(op, VariableType.OPERAND, 0))
  dst_variable = cs.Variable(ValueSite(op, VariableType.RESULT, 0))
  src_shape = tuple(op.operand.type.shape)
  dst_shape = tuple(op.result.type.shape)

  # Map destination index -> source index
  dst_to_src = {dst: src for src, dst in enumerate(op.broadcast_dimensions)}

  kept_dims = []
  collapsed_dims = []
  for dim in range(len(dst_shape)):
    if dim in dst_to_src:
      s_idx = dst_to_src[dim]
      # If the source was 1 but destination is > 1, we reduce but keep the dim.
      if src_shape[s_idx] == 1 and dst_shape[dim] > 1:
        kept_dims.append(dim)
    else:
      # If the dimension didn't exist in src_shape at all, we remove it.
      collapsed_dims.append(dim)

  assert kept_dims or collapsed_dims

  reduce_expr = dst_variable

  # 1. Apply keep_dims=True first. This keeps the rank the same,
  # so collapsed_dims indices remain valid.
  if kept_dims:
    reduce_expr = cs.Reduce(
        reduce_expr, axes=tuple(kept_dims), rank=len(dst_shape), keep_dims=True
    )

  # 2. Apply keep_dims=False to remove the added dimensions.
  if collapsed_dims:
    reduce_expr = cs.Reduce(
        reduce_expr, axes=tuple(collapsed_dims), rank=len(dst_shape), keep_dims=False
    )

  constraints = [
      # We need the `src = Reduce(...)` constraint to enforce correctness.
      # Alternatively, we could enforce it via the `IsSupportedBroadcast`
      # constraint but currently it doesn't do the necessary checks.
      cs.Equals(src_variable, reduce_expr),
      cs.IsSupportedBroadcast(
          src_variable, dst_variable, tuple(op.broadcast_dimensions)
      ),
  ]
  return (
      cs.ConstraintSystem(constraints=constraints),
      {
          src_variable: [src_variable.key],
          dst_variable: [dst_variable.key],
      },
  )


@_add_constraint_system_derivation_rule(mgpu.VectorConcatOp)
def _vector_concat_constraint_system(
    ctx: DerivationContext, op: mgpu.VectorConcatOp
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  result = ValueSite(op, VariableType.RESULT, 0)
  result_var = cs.Variable(result)
  value_sites_for_variable: ValueSitesForVariable = {result_var: [result]}
  constraints: list[cs.Constraint] = [
      # TODO(allanrenucci): Add support for strided layout.
      cs.NotOfType(result_var, fa.WGSplatFragLayout),
      cs.NotOfType(result_var, fa.WGStridedFragLayout),
  ]
  for i, _ in enumerate(op.operands):
    operand_site = ValueSite(op, VariableType.OPERAND, i)
    operand_var = cs.Variable(operand_site)
    value_sites_for_variable[operand_var] = [operand_site]
    constraints.append(cs.Equals(operand_var, result_var))
  return cs.ConstraintSystem(constraints=constraints), value_sites_for_variable


@_add_constraint_system_derivation_rule(vector.ShapeCastOp)
def _shape_cast_constraint_system(
    ctx: DerivationContext, op: vector.ShapeCastOp
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  in_shape = tuple(cast(ir.ShapedType, op.source.type).shape)
  out_shape = tuple(cast(ir.ShapedType, op.result.type).shape)

  in_variable = cs.Variable(ValueSite(op, VariableType.OPERAND, 0))
  out_variable = cs.Variable(ValueSite(op, VariableType.RESULT, 0))

  # Here, we are in a case where we are stating
  #
  #   out_variable = reshape(in_variable, in_shape, out_shape).
  #
  # Thanks to the symmetric property of reshape, we can also issue a constraint
  # in the other direction, i.e.
  #
  #   in_variable = reshape(out_variable, out_shape, in_shape)
  #
  # in order to be able to figure out an assignment for `in_variable`. if we
  # happen to know `out_variable`. If we only issue the first constraint, then
  # we will not be able to figure out an assignment for `in_variable` if we
  # only know `out_variable`, even though their relationship is fully
  # determined.
  in_to_out = cs.Reshape(
      in_variable, source_shape=in_shape, target_shape=out_shape
  )
  out_to_in = cs.Reshape(
      out_variable, source_shape=out_shape, target_shape=in_shape
  )

  return (
      cs.ConstraintSystem(
          constraints=[
              cs.Equals(lhs=out_variable, rhs=in_to_out),
              cs.Equals(lhs=in_variable, rhs=out_to_in),
          ],
      ),
      {in_variable: [in_variable.key], out_variable: [out_variable.key]},
  )


@_add_constraint_system_derivation_rule(vector.ExtractStridedSliceOp)
def _extract_strided_slice_constraint_system(
    ctx: DerivationContext, op: vector.ExtractStridedSliceOp
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  if any(ir.IntegerAttr(s).value != 1 for s in op.strides):
    raise NotImplementedError("`strides` must contain only 1s.")
  operand = cs.Variable(ValueSite(op, VariableType.OPERAND, 0))
  result = cs.Variable(ValueSite(op, VariableType.RESULT, 0))
  offsets = tuple(ir.IntegerAttr(o).value for o in op.offsets)
  constraints = [
      cs.Divides(operand, offsets),
      cs.Equals(operand, result),
      # TODO(allanrenucci): Remove once vectors with splat and strided layouts
      # can be sliced.
      cs.NotOfType(result, fa.WGSplatFragLayout),
      cs.NotOfType(result, fa.WGStridedFragLayout),
  ]
  return (
      cs.ConstraintSystem(constraints=constraints),
      {operand: [operand.key], result: [result.key]},
  )


@_add_constraint_system_derivation_rule(vector.ExtractOp)
def _vector_extract_constraint_system(
    ctx: DerivationContext, op: vector.ExtractOp
) -> tuple[cs.ConstraintSystem, ValueSitesForVariable]:
  del ctx
  operand = cs.Variable(ValueSite(op, VariableType.OPERAND, 0))
  if not isinstance(op.result.type, ir.VectorType):  # scalar result
    layout = fa.WGSplatFragLayout(tuple(op.source.type.shape))
    # We only support indexing for splat layout.
    assignments: dict[cs.Variable, cs.Constant] = {
        operand: cs.RegisterLayout(layout)
    }
    return cs.ConstraintSystem(assignments), {operand: [operand.key]}

  if op.dynamic_position:
    raise NotImplementedError("Only slicing with static indices allowed.")
  result = cs.Variable(ValueSite(op, VariableType.RESULT, 0))
  constraints = [
      cs.Equals(operand, result),
      cs.Divides(operand, tuple(op.source.type.shape)),
      cs.Divides(result, tuple(op.result.type.shape)),
      # TODO(allanrenucci): Remove once vectors with splat and strided layouts
      # can be sliced.
      cs.NotOfType(result, fa.WGSplatFragLayout),
      cs.NotOfType(result, fa.WGStridedFragLayout),
  ]
  return (
      cs.ConstraintSystem(constraints=constraints),
      {operand: [operand.key], result: [result.key]},
  )


@_add_constraint_system_derivation_rule(vector.InsertStridedSliceOp)
def _insert_strided_slice_constraint_system(
    ctx: DerivationContext, op: vector.InsertStridedSliceOp
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  if any(ir.IntegerAttr(s).value != 1 for s in op.strides):
    raise NotImplementedError("`strides` must contain only 1s.")
  src = cs.Variable(ValueSite(op, VariableType.OPERAND, 0))
  dst = cs.Variable(ValueSite(op, VariableType.OPERAND, 1))
  result = cs.Variable(ValueSite(op, VariableType.RESULT, 0))
  offsets = tuple(ir.IntegerAttr(o).value for o in op.offsets)
  constraints = [
      cs.Divides(dst, offsets),
      cs.Equals(dst, src),
      cs.Equals(dst, result),
      cs.Equals(src, result),
      # TODO(allanrenucci): Remove once vectors with splat and strided layouts
      # can be sliced.
      cs.NotOfType(dst, fa.WGSplatFragLayout),
      cs.NotOfType(dst, fa.WGStridedFragLayout),
  ]
  return (
      cs.ConstraintSystem(constraints=constraints),
      {src: [src.key], dst: [dst.key], result: [result.key]},
  )


@_add_constraint_system_derivation_rule(mgpu.CustomPrimitiveOp)
def _custom_primitive_constraint_system(
    ctx: DerivationContext,
    op: mgpu.CustomPrimitiveOp,
) -> ConstraintSystemDerivationRuleResult:
  assignments: dict[cs.Variable, cs.Constant] = {}
  constraints: list[cs.Constraint] = []
  in_layouts = iter(op.in_layouts)
  in_transforms = iter(op.in_transforms)
  variables: list[cs.Variable] = []
  for i, operand in enumerate(op.operands):
    if is_vector(operand):
      v = cs.Variable(ValueSite(op, VariableType.OPERAND, i))
      variables.append(v)
      assignments[v] = cs.RegisterLayout(
          layouts_lib.from_layout_attr(next(in_layouts))
      )
    elif _is_smem_ref(operand):
      # Here we need to create a new variable, even though it is equal to the
      # source operand. This is because we directly assign the new variable and
      # if we did that to the source there could be conflicting assignments.
      # For example, the same ref could be passed into the custom op twice with
      # different transforms, which needs to yield an unsatisfiable system.
      #
      # TODO(b/447079781): Consider creating the final constraint system using
      # __and__ and potentially returning Unsatisfiable() directly if there is
      # a conflict between the assignments.
      value_site = ValueSite(op, VariableType.OPERAND, i)
      source_var = ctx.producer_ref(value_site)
      v = cs.Variable(value_site)
      constraints.append(cs.Equals(lhs=source_var, rhs=v))
      variables.append(v)
      transforms = next(in_transforms)
      assert isinstance(transforms, ir.ArrayAttr)
      tiling = _extract_smem_transforms_from_custom_transform_attrs(transforms)
      assignments[v] = tiling

  out_layouts = iter(op.out_layouts)
  for i, result in enumerate(op.results):
    if isinstance(result.type, ir.VectorType):
      v = cs.Variable(ValueSite(op, VariableType.RESULT, i))
      variables.append(v)
      assignments[v] = cs.RegisterLayout(
          layouts_lib.from_layout_attr(next(out_layouts))
      )
  return (
      cs.ConstraintSystem(assignments, constraints),
      {v: [v.key] for v in variables},
  )


def _tmem_layout_from_layout_attr(
    layout_attr: ir.Attribute,
) -> tcgen05.TMEMLayout:
  layout = layouts_lib.from_layout_attr(layout_attr)
  assert isinstance(layout, fa.TiledLayout)
  return tcgen05.TMEMLayout(
      layout.tiling, layout.warp_dims, layout.lane_dims, layout.vector_dim
  )


@_add_constraint_system_derivation_rule(mgpu.TmemLayoutCastOp)
def _tmem_layout_cast_constraint_system(
    ctx: DerivationContext,
    op: mgpu.TmemLayoutCastOp,
) -> ConstraintSystemDerivationRuleResult:
  operand = ValueSite(op, VariableType.OPERAND, 0)
  variable = ctx.producer_ref(operand)
  result = ValueSite(op, VariableType.RESULT, 0)
  tmem_layout = cs.TMEMLayout(_tmem_layout_from_layout_attr(op.new_layout))
  if not cs.is_valid_assignment(variable, tmem_layout):
    raise ValueError(
        f"Cannot cast to TMEM layout {tmem_layout.value} in {op}: the layout is"
        f" not compatible with the operand shape {operand.shape}."
    )
  return (
      cs.ConstraintSystem(assignments={variable: tmem_layout}),
      {variable: [operand, result]},
  )


@_add_constraint_system_derivation_rule(mgpu.TmemAllocOp)
def _tmem_alloc_constraint_system(
    ctx: DerivationContext,
    op: mgpu.TmemAllocOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  result = ValueSite(op, VariableType.RESULT, 0)
  result_var = cs.Variable(result)
  in_smem = ValueSite(op, VariableType.OPERAND, 0)
  in_smem_var = cs.Variable(in_smem)
  assignments: dict[cs.Variable, cs.Constant] = {
      in_smem_var: cs.SMEMTransforms(None, None)
  }
  operands_for_variable = {result_var: [result], in_smem_var: [in_smem]}
  return cs.ConstraintSystem(assignments=assignments), operands_for_variable


@_add_constraint_system_derivation_rule(mgpu.TmemDeallocOp)
def _tmem_dealloc_constraint_system(
    ctx: DerivationContext,
    op: mgpu.TmemDeallocOp,
) -> ConstraintSystemDerivationRuleResult:
  operand = ValueSite(op, VariableType.OPERAND, 0)
  variable = ctx.producer_ref(operand)
  return cs.ConstraintSystem(), {variable: [operand]}


@_add_constraint_system_derivation_rule(mgpu.TcGen05MMAOp)
def _tcgen05_mma_constraint_system(
    ctx: DerivationContext,
    op: mgpu.TcGen05MMAOp,
) -> ConstraintSystemDerivationRuleResult:
  assignments: dict[cs.Variable, cs.Constant] = {}
  constraints: list[cs.Constraint] = []
  operands_for_variable: ValueSitesForVariable = {}

  # TMEM
  acc = ValueSite(op, VariableType.OPERAND, 0)
  acc_variable = ctx.producer_ref(acc)
  acc_type = ir.ShapedType(op.accumulator.type)
  acc_layout = tcgen05._infer_tmem_layout(
      tuple(acc_type.shape), bool(op.collective), packing=1
  )
  acc_layout = cs.TMEMLayout(acc_layout)
  assignments[acc_variable] = acc_layout
  acc_is_valid = cs.is_valid_assignment(acc_variable, acc_layout)
  if not acc_is_valid:
    raise ValueError(
        f"Cannot assign TMEM layout {acc_layout.value} to the accumulator of"
        f" a tcgen05 MMA op: the layout is not compatible with the accumulator"
        f" shape {acc.shape}."
    )
  operands_for_variable[acc_variable] = [acc]

  element_type_bitwidth = utils.bitwidth(op.b.type.element_type)
  b = ValueSite(op, VariableType.OPERAND, 2)
  b_var = ctx.producer_ref(b)
  operands_for_variable[b_var] = [b]
  constraints.append(mma_tiling_constraint(b_var, op))

  # SMEM
  M = op.accumulator.type.shape[0]
  if M == 64 and not op.collective.value:
    # We can't split N into groups if we would partition it below the tile size.
    N = op.b.type.shape[1]
    n_lane_groups = 2
    max_swizzle_elems = next(
        8 * s // element_type_bitwidth
        for s in reversed(mgpu.SwizzlingMode)
        if 8 * s // element_type_bitwidth <= N // n_lane_groups
    )
    if utils.is_memref_transposed(op.b.type):
      constraints.append(cs.Divides(b_var, (max_swizzle_elems, 8)))
    else:
      constraints.append(cs.Divides(b_var, (8, max_swizzle_elems)))

  if _is_tmem_ref(op.a):
    a = ValueSite(op, VariableType.OPERAND, 1)
    a_type = ir.ShapedType(op.a.type)
    a_var = ctx.producer_ref(a)
    packing = 32 // utils.bitwidth(a_type.element_type)
    a_layout = tcgen05._infer_tmem_layout(
        tuple(a_type.shape), bool(op.collective), packing
    )
    a_layout = cs.TMEMLayout(a_layout)
    assignments[a_var] = a_layout
    operands_for_variable[a_var] = [a]
    a_is_valid = cs.is_valid_assignment(a_var, a_layout)
    if not a_is_valid:
      raise ValueError(
          f"Cannot assign TMEM layout {a_layout.value} to the 'a' operand of"
          f" a tcgen05 MMA op: the layout is not compatible with the operand"
          f" shape {a.shape}."
      )
  else:
    assert _is_smem_ref(op.a)
    a = ValueSite(op, VariableType.OPERAND, 1)
    a_var = ctx.producer_ref(a)
    operands_for_variable[a_var] = [a]
    constraints.append(mma_tiling_constraint(a_var, op))

  if (sparse_meta_operand := getattr(op, "a_sparse_metadata")) is not None:
    sparse_meta = ValueSite(
        op, VariableType.OPERAND, list(op.operands).index(sparse_meta_operand)
    )
    sparse_meta_var = ctx.producer_ref(sparse_meta)
    sparse_meta_layout = cs.TMEMLayout(tcgen05.sparse_meta_layout())
    assignments[sparse_meta_var] = sparse_meta_layout
    if not cs.is_valid_assignment(sparse_meta_var, sparse_meta_layout):
      raise ValueError(
          f"Cannot assign TMEM layout {sparse_meta_layout.value} to the"
          f" 'a_sparse_metadata' operand of a tcgen05 MMA op: the layout is not"
          f" compatible with the operand shape {sparse_meta.shape}."
      )
    operands_for_variable[sparse_meta_var] = [sparse_meta]

  if (scaled := op.a_scale is not None) != (op.b_scale is not None):
    raise ValueError(
        f"Expecting neither or both scales to be present. Got {op.a_scale=},"
        f" {op.b_scale=}"
    )

  def assign_scaled_layout(scale_operand):
    scale_index = list(op.operands).index(scale_operand)
    scale = ValueSite(op, VariableType.OPERAND, scale_index)
    scale_var = ctx.producer_ref(scale)
    if op.collective and scale_operand == op.b_scale and M == 64:
      layout = tcgen05.b_scales_m64_collective_layout()
    else:
      layout = tcgen05.scales_layout()
    layout = cs.TMEMLayout(layout)
    assignments[scale_var] = layout
    if not cs.is_valid_assignment(scale_var, layout):
      raise ValueError(
          f"Cannot assign {layout} to {scale_operand=} with"
          f" shape {scale.shape}."
      )
    operands_for_variable[scale_var] = [scale]

  if scaled:
    assign_scaled_layout(op.a_scale)
    assign_scaled_layout(op.b_scale)

  return cs.ConstraintSystem(assignments=assignments, constraints=constraints), operands_for_variable


@_add_constraint_system_derivation_rule(mgpu.AsyncLoadTmemOp)
def _async_load_tmem_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncLoadTmemOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  source_variable = ctx.producer_ref(source)
  destination = ValueSite(op, VariableType.RESULT, 0)
  destination_variable = cs.Variable(destination)
  constraint = cs.IsTransferableTmemRegisters(
      source_variable,
      destination_variable,
      tuple(ir.ShapedType(op.source.type).shape),
      bitwidth=utils.bitwidth(op.source.type.element_type),
  )
  constraints: list[cs.Constraint] = [
      constraint,
      # The following `NotOfType` constraints are shortcuts, helping
      # the layout inference system converge faster.
      # They are implicitly rejected by `IsTransferableTmemRegisters`.
      cs.NotOfType(destination_variable, fa.WGSplatFragLayout),
      cs.NotOfType(destination_variable, fa.WGStridedFragLayout),
  ]
  operands_for_variable = {
      source_variable: [source],
      destination_variable: [destination],
  }
  if op.reduce is not None:
    reduced = ValueSite(op, VariableType.RESULT, 1)
    reduced_variable = cs.Variable(reduced)
    operands_for_variable[reduced_variable] = [reduced]
    source_rank = len(ir.MemRefType(op.source.type).shape)
    assert source_rank == 2
    constraints.append(
        cs.Equals(
            reduced_variable,
            cs.Reduce(
                destination_variable, axes=(1,), rank=source_rank, keep_dims=False
            ),
        )
    )
  return (
      cs.ConstraintSystem(constraints=constraints),
      operands_for_variable,
  )


@_add_constraint_system_derivation_rule(mgpu.AsyncStoreSmemToTmemOp)
def _async_store_smem_to_tmem_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncStoreSmemToTmemOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  source_variable = ctx.producer_ref(source)
  destination = ValueSite(op, VariableType.OPERAND, 1)
  destination_variable = ctx.producer_ref(destination)
  bitwidth = utils.bitwidth(op.destination.type.element_type)
  packing = 32 // bitwidth
  tmem_layout = cs.TMEMLayout(tcgen05.tmem_default_layout(packing))
  if not cs.is_valid_assignment(destination_variable, tmem_layout):
    raise ValueError(
        f"Cannot assign TMEM layout {tmem_layout.value} to a TMEM ref "
        f"with shape {destination.shape}"
    )
  constraint = mma_tiling_constraint(source_variable, op, allow_unswizzled=True)
  return (
      cs.ConstraintSystem(
          assignments={destination_variable: tmem_layout},
          constraints=[constraint],
      ),
      {source_variable: [source], destination_variable: [destination]},
  )


@_add_constraint_system_derivation_rule(
    mgpu.AsyncStoreSparseMetadataSmemToTmemOp
)
def _async_store_sparse_metadata_smem_to_tmem_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncStoreSparseMetadataSmemToTmemOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  source_variable = ctx.producer_ref(source)
  destination = ValueSite(op, VariableType.OPERAND, 1)
  destination_variable = ctx.producer_ref(destination)
  return (
      cs.ConstraintSystem(
          assignments={
              destination_variable: cs.TMEMLayout(
                  tcgen05.sparse_meta_layout()
              ),
              source_variable: cs.SMEMTransforms(None, None),
          },
      ),
      {source_variable: [source], destination_variable: [destination]},
  )

@_add_constraint_system_derivation_rule(mgpu.AsyncStoreScalesSmemToTmemOp)
def _async_store_scales_smem_to_tmem_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncStoreScalesSmemToTmemOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  source_variable = ctx.producer_ref(source)
  destination = ValueSite(op, VariableType.OPERAND, 1)
  destination_variable = ctx.producer_ref(destination)

  assignments: dict[cs.Variable, cs.Constant] = {
      source_variable: cs.SMEMTransforms(None, None)
  }
  k_tiles = destination.shape[1] // 4
  if source.shape == (1, k_tiles, 64, 16):
    assignments[destination_variable] = cs.TMEMLayout(
        tcgen05.b_scales_m64_collective_layout()
    )
  else:
    assignments[destination_variable] = cs.TMEMLayout(tcgen05.scales_layout())
  return (
      cs.ConstraintSystem(assignments),
      {source_variable: [source], destination_variable: [destination]},
  )


@_add_constraint_system_derivation_rule(mgpu.SliceTmemOp)
def _slice_tmem_constraint_system(
    ctx: DerivationContext,
    op: mgpu.SliceTmemOp,
) -> ConstraintSystemDerivationRuleResult:
  operand = ValueSite(op, VariableType.OPERAND, 0)
  operand_variable = ctx.producer_ref(operand)
  result = ValueSite(op, VariableType.RESULT, 0)
  # TODO(bchetioui): enforce that the parent is a TmemAllocOp.
  if op.alias_id is not None:
    alias_id = op.alias_id.value
    if (cached_variable := ctx.slice_tmem_aliases.get(alias_id)) is not None:
      result_variable = cached_variable
    else:
      result_variable = cs.Variable(result)
      ctx.slice_tmem_aliases[alias_id] = result_variable
  else:
    result_variable = cs.Variable(result)
  return (
      cs.ConstraintSystem(),
      {operand_variable: [operand], result_variable: [result]},
  )


@_add_constraint_system_derivation_rule(mgpu.AsyncStoreTmemOp)
def _async_store_tmem_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncStoreTmemOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  source_variable = cs.Variable(source)
  destination = ValueSite(op, VariableType.OPERAND, 1)
  destination_variable = ctx.producer_ref(destination)
  constraint = cs.IsTransferableTmemRegisters(
      source_variable,
      destination_variable,
      tuple(ir.ShapedType(op.source.type).shape),
      bitwidth=utils.bitwidth(op.source.type.element_type),
  )
  return (
      cs.ConstraintSystem(constraints=[constraint]),
      {source_variable: [source], destination_variable: [destination]},
  )


@_add_constraint_system_derivation_rule(mgpu.SliceSMEMOp)
def _slice_smem_constraint_system(
    ctx: DerivationContext,
    op: mgpu.SliceSMEMOp,
) -> ConstraintSystemDerivationRuleResult:
  result = ValueSite(op, VariableType.RESULT, 0)
  if op.alias_id is not None:
    alias_id = op.alias_id.value
    if (cached_variable := ctx.slice_smem_aliases.get(alias_id)) is not None:
      result_variable = cached_variable
    else:
      result_variable = cs.Variable(result)
      ctx.slice_smem_aliases[alias_id] = result_variable
  else:
    result_variable = cs.Variable(result)
  return cs.ConstraintSystem(), {result_variable: [result]}


@_add_constraint_system_derivation_rule(mgpu.GetClusterRefOp)
def _get_cluster_ref_constraint_system(
    ctx: DerivationContext,
    op: mgpu.GetClusterRefOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  var_source_dest = ctx.producer_ref(source)
  dest = ValueSite(op, VariableType.RESULT, 0)
  return cs.ConstraintSystem(), {var_source_dest: [source, dest]}


@_add_constraint_system_derivation_rule(memref.SubViewOp)
def _memref_subview_constraint_system(
    ctx: DerivationContext,
    op: memref.SubViewOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  source_var = ctx.producer_ref(source)
  result = ValueSite(op, VariableType.RESULT, 0)
  result_var = cs.Variable(result)

  if any(s != 1 for s in op.static_strides):
    raise NotImplementedError(
        f"Only unit strides are supported but got {op.static_strides}."
    )

  # Collect all the constraints from all dimensions.
  tiling_multiple = []
  dynamic_offset_index = 0
  for i, size in enumerate(op.static_sizes):
    offset = op.static_offsets[i]
    if offset == ir.ShapedType.get_dynamic_size():
      offset = op.offsets[dynamic_offset_index]
      dynamic_offset_index += 1

    # Drop all dimensions up to and including the last dynamic size. Dynamic
    # sizes are not supported yet.
    #
    # Supporting dynamic sizes here can be done analogously to how dynamic
    # offsets are supported. The reason we don't support dynamic sizes now is
    # because the lowering does not yet support them.
    if ir.ShapedType.is_dynamic_size(size):
      tiling_multiple = []
    else:
      src_type = ir.MemRefType(op.source.type)
      divisibility_constraint = math.gcd(size, src_type.shape[i])
      if isinstance(offset, int):
        divisibility_constraint = math.gcd(divisibility_constraint, offset)
      else:
        divisibility_constraint = dynamic_gcd(divisibility_constraint, offset)
      tiling_multiple.append(divisibility_constraint)

  constraints = [
      cs.Divides(source_var, tuple(tiling_multiple)),
      cs.Equals(source_var, result_var),
  ]
  system = cs.ConstraintSystem(constraints=constraints)
  return system, {source_var: [source], result_var: [result]}


@_add_constraint_system_derivation_rule(memref.CastOp)
def _memref_cast_op_constraint_system(
    ctx: DerivationContext,
    op: memref.CastOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  var_source_dest = ctx.producer_ref(source)
  dest = ValueSite(op, VariableType.RESULT, 0)
  return cs.ConstraintSystem(), {var_source_dest: [source, dest]}


@_add_constraint_system_derivation_rule(memref.TransposeOp)
def _memref_transpose_op_constraint_system(
    ctx: DerivationContext,
    op: memref.TransposeOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  dest = ValueSite(op, VariableType.RESULT, 0)
  source_var = ctx.producer_ref(source)

  dest_var = cs.Variable(dest)

  permutation = tuple(
      ir.AffineDimExpr(e).position for e in op.permutation.value.results
  )
  inv_permutation = tuple(permutation.index(i) for i in range(len(permutation)))

  constraints = [
      cs.Equals(cs.Transpose(source_var, permutation=permutation), dest_var),
      cs.Equals(source_var, cs.Transpose(dest_var, permutation=inv_permutation)),
  ]
  system = cs.ConstraintSystem(constraints=constraints)
  return system, {source_var: [source], dest_var: [dest]}


@_add_constraint_system_derivation_rule(memref.ExpandShapeOp)
def _memref_expand_shape_op_equation_system(
    ctx: DerivationContext,
    op: memref.ExpandShapeOp,
) -> ConstraintSystemDerivationRuleResult:
  if utils.is_memref_transposed(ir.MemRefType(op.src.type)):
    raise NotImplementedError(
        "Transposed memrefs are not supported in ExpandShapeOp."
    )

  source = ValueSite(op, VariableType.OPERAND, 0)
  source_var = ctx.producer_ref(source)
  dest = ValueSite(op, VariableType.RESULT, 0)
  dest_var = cs.Variable(dest)

  reverse_tiling_multiple = []
  for dim, idx in zip(
      reversed(op.static_output_shape), reversed(op.reassociation)
  ):
    # pyrefly: ignore[bad-argument-type]
    if ir.ShapedType.is_dynamic_size(dim) or len(idx) > 1:
      # For simplicity, we only support tiling non-expanded static dimensions.
      # These limitations could be lifted later if needed.
      break
    reverse_tiling_multiple.append(dim)

  constraints = [
      cs.Divides(source_var, tuple(reversed(reverse_tiling_multiple))),
      cs.Equals(source_var, dest_var),
  ]
  return cs.ConstraintSystem(constraints=constraints), {
      source_var: [source],
      dest_var: [dest],
  }


@_add_constraint_system_derivation_rule(memref.CollapseShapeOp)
def _memref_collapse_shape_op_constraint_system(
    ctx: DerivationContext,
    op: memref.CollapseShapeOp,
) -> ConstraintSystemDerivationRuleResult:
  reassociation = tuple(len(ir.ArrayAttr(idx)) for idx in op.reassociation)
  # This should only occur when going from a (1, ...) shape to an empty shape.
  # We can handle it if needed, but right now `CollapseShape` will not deal with
  # this case.
  if not reassociation:
    raise NotImplementedError(
        "CollapseShapeOp with empty reassociation is not supported."
    )

  source = ValueSite(op, VariableType.OPERAND, 0)
  source_var = ctx.producer_ref(source)
  dest = ValueSite(op, VariableType.RESULT, 0)
  dest_var = cs.Variable(dest)

  strides, _ = ir.MemRefType(source.value.type).get_strides_and_offset()
  # In this case, we'd need additional checks to produce a correct constraint.
  if strides != utils.get_contiguous_strides(source.shape):
    raise NotImplementedError(
        "CollapseShapeOp with non-contiguous strides is not supported."
    )

  # TODO(bchetioui): We could generate an inverse expression `ExpandShape` in
  # order to allow inferring layouts bidirectionally. This would allow removing
  # transforms from some kernels' BlockSpecs, but is not necessary at this time.
  collapse_expr = cs.CollapseShape(source_var, source.shape, reassociation)
  return cs.ConstraintSystem(constraints=[cs.Equals(dest_var, collapse_expr)]), {
      source_var: [source],
      dest_var: [dest],
  }


# `memref.load` and `memref.store` are used to load barrier phases which are
# scalars---the rule needn't do anything interesting, but we need to have it.
@_add_constraint_system_derivation_rule(memref.LoadOp)
@_add_constraint_system_derivation_rule(memref.StoreOp)
def _memref_load_store_op_constraint_system(
    ctx: DerivationContext,
    op: memref.LoadOp | memref.StoreOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx

  ref_shape = ir.MemRefType(op.memref.type).shape
  if ref_shape and ref_shape != [1]:
    raise NotImplementedError(
        f"Only scalar memrefs are supported, got {ref_shape}"
    )

  ref_op_index = 0 if isinstance(op, memref.LoadOp) else 1
  ref = ValueSite(op, VariableType.OPERAND, ref_op_index)
  var = cs.Variable(ref)
  assignments: dict[cs.Variable, cs.Constant] = {var: cs.SMEMTransforms(None, None)}
  return cs.ConstraintSystem(assignments=assignments), {var: [ref]}


@_add_constraint_system_derivation_rule(mgpu.TryClusterCancelOp)
@_add_constraint_system_derivation_rule(mgpu.QueryClusterCancelOp)
def _cluster_launch_control_ops_constraint_system(
    ctx: DerivationContext,
    op: mgpu.TryClusterCancelOp | mgpu.QueryClusterCancelOp,
) -> ConstraintSystemDerivationRuleResult:
  ref = ValueSite(op, VariableType.OPERAND, 0)
  var = ctx.producer_ref(ref)
  assignments: dict[cs.Variable, cs.Constant] = {var: cs.SMEMTransforms(None, None)}
  return cs.ConstraintSystem(assignments=assignments), {var: [ref]}


def _extract_smem_transforms_from_custom_transform_attrs(
    transform_attrs: ir.ArrayAttr,
) -> cs.SMEMTransforms:
  transforms = [layouts_lib.from_transform_attr(x) for x in transform_attrs]
  match transforms:
    case []:
      tile_transform = None
      swizzle = None
    case [lc.TileTransform() as t]:
      tile_transform = t
      swizzle = None
    case [lc.TileTransform() as t, mgpu.SwizzlingMode() as s]:
      tile_transform = t
      # TODO(olechwierowicz): We should eliminate `kNoSwizzle`, representing
      # this state as None is enough.
      swizzle = s.value if s != mgpu.SwizzlingMode.kNoSwizzle else None
    case [mgpu.SwizzlingMode() as s]:
      tile_transform = None
      swizzle = s.value if s != mgpu.SwizzlingMode.kNoSwizzle else None
    case _:
      raise NotImplementedError(f"Unsupported transforms {transforms}")

  return cs.SMEMTransforms(tile_transform, swizzle)


@_add_constraint_system_derivation_rule(mgpu.WithTransformsOp)
def _with_transforms_constraint_system(
    ctx: DerivationContext,
    op: mgpu.WithTransformsOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  dest = ValueSite(op, VariableType.RESULT, 0)
  var = ctx.producer_ref(source)
  smem_transforms = _extract_smem_transforms_from_custom_transform_attrs(
      op.transforms
  )

  if not cs.is_valid_assignment(var, smem_transforms):
    tiling_transform_str = (
        f"tiling {smem_transforms.tiling}"
        if smem_transforms.tiling
        else "empty tiling"
    )
    swizzle_str = (
        f"{smem_transforms.swizzle} swizzle"
        if smem_transforms.swizzle
        else "no swizzle"
    )
    raise ValueError(
        f"Cannot apply {tiling_transform_str} with {swizzle_str} to memref with"
        f" shape {source.shape}."
    )
  assignments: dict[cs.Variable, cs.Constant] = {var: smem_transforms}
  return cs.ConstraintSystem(assignments=assignments), {var: [source, dest]}


def _vector_value_sites_and_constraints_for_async_ops(
    op: mgpu.AsyncLoadOp | mgpu.AsyncStoreOp | mgpu.AsyncPrefetchOp,
) -> tuple[ValueSitesForVariable, list[cs.Constraint]]:
  values_sites: ValueSitesForVariable = dict()
  constraints: list[cs.Constraint] = []

  match op:
    case mgpu.AsyncLoadOp():
      base_operand_index = 3
    case mgpu.AsyncStoreOp():
      base_operand_index = 2
    case mgpu.AsyncPrefetchOp():
      base_operand_index = 1

  for i, idx in enumerate(op.indices):
    if isinstance(idx.type, ir.VectorType):
      value_site = ValueSite(op, VariableType.OPERAND, base_operand_index + i)
      value_site_var = cs.Variable(value_site)
      shape = tuple(idx.type.shape)

      allowed_layouts = []
      if shape[0] % 16 == 0:
        allowed_layouts.append(cs.RegisterLayout(value=fa.TMA_INDICES_LAYOUT))
      if shape[0] % 4 == 0:
        allowed_layouts.append(cs.RegisterLayout(value=fa.TMA_INDICES_4_LAYOUT))
      if not allowed_layouts:
        raise ValueError(f"Unsupported TMA index shape {shape}")
      values_sites[value_site_var] = [value_site]
      constraints.append(cs.OneOf(value_site_var, tuple(allowed_layouts)))
  return values_sites, constraints


@_add_constraint_system_derivation_rule(mgpu.AsyncLoadOp)
@_add_constraint_system_derivation_rule(mgpu.AsyncStoreOp)
def _async_load_store_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncLoadOp | mgpu.AsyncStoreOp,
) -> ConstraintSystemDerivationRuleResult:
  # We only support 2D gathers/scatters along the leading dimension. Tiling
  # either keeps the gather/scatter dimension leading or allows
  # collapsing leading dimensions to maintain contiguity without
  # transforming global memory.
  tiling_multiple = []
  for i, (size, index) in enumerate(zip(op.slice_lengths, op.indices, strict=True)):
    if size == -1:
      # This dimension does not appear in the final smem memref shape.
      continue
    if isinstance(index.type, ir.VectorType):
      if i != 0:
        raise NotImplementedError("Only leading gather dimensions allowed.")
      if isinstance(op, mgpu.AsyncStoreOp):
        gmem_shape = ir.MemRefType(op.destination.type).shape
      else:
        gmem_shape = ir.MemRefType(op.source.type).shape
      if len(gmem_shape) != 2:
        raise NotImplementedError("Only 2D gathers/scatters for async load/store are supported.")
      tiling_multiple.append(size)
      continue
    tiling_multiple.append(dynamic_gcd(size, index))

  operand_index = 1 if isinstance(op, mgpu.AsyncLoadOp) else 0
  operand = ValueSite(op, VariableType.OPERAND, operand_index)
  var = ctx.producer_ref(operand)
  constraints: list[cs.Constraint] = [
      cs.Divides(expr=var, tiling_multiple=tuple(tiling_multiple))
  ]
  if any(isinstance(idx.type, ir.VectorType) for idx in op.indices):
    element_bitwidth = utils.bitwidth(op.source.type.element_type)
    # This constraint enforces sufficient SMEM-alignment.
    # The transfer chunk needs to be 1024 bit-aligned. For each write in the
    # lowering we transfer 4 rows, so each row must be 256 bit-aligned.
    divisor = (1024 // 4) // element_bitwidth
    slice_lengths = [s for s in op.slice_lengths if s != -1]
    if slice_lengths and (slice_lengths[-1] % divisor):
      raise ValueError(
          "Cannot assign layout to async load with gather indices since"
          f" minor dim={slice_lengths[-1]} is not divisible by {divisor=}"
          " bits."
      )
    constraints.append(cs.MinorDimDivisibleBy(expr=var, divisor=divisor))

  value_sites_for_variable = {var: [operand]}
  value_sites, extra_constraints = _vector_value_sites_and_constraints_for_async_ops(op)
  value_sites_for_variable.update(value_sites)
  constraints.extend(extra_constraints)
  return cs.ConstraintSystem(constraints=constraints), value_sites_for_variable


@_add_constraint_system_derivation_rule(mgpu.AsyncPrefetchOp)
def _async_prefetch_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncPrefetchOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx
  value_sites, constraints = _vector_value_sites_and_constraints_for_async_ops(op)
  return cs.ConstraintSystem(constraints=constraints), value_sites


@_add_constraint_system_derivation_rule(mgpu.WarpMapOp)
def _warp_map_constraint_system(
    ctx: DerivationContext,
    op: mgpu.WarpMapOp,
) -> ConstraintSystemDerivationRuleResult:
  value_sites_for_variable: ValueSitesForVariable = dict()
  for i, o in enumerate(op.operands):
    if _is_tmem_ref(o) or _is_smem_ref(o):
      operand = ValueSite(op, VariableType.OPERAND, i)
      arg = ValueSite(op, VariableType.ARGUMENT, i, region_index=0)
      var = ctx.producer_ref(operand)
      value_sites_for_variable.setdefault(var, []).extend([operand, arg])
  return cs.ConstraintSystem(), value_sites_for_variable


def _ensure_all_layouts_are_set(op: ir.OpView) -> None:
  if inference_utils.should_have_layout(op):
    _ensure_right_number_of_layouts(is_vector, "layouts", "vector", op)
  if inference_utils.should_have_tmem_layout(op):
    _ensure_right_number_of_layouts(_is_tmem_ref, "tmem_layouts", "TMEM ref", op)
  if inference_utils.should_have_transforms(op):
    _ensure_right_number_of_layouts(
        inference_utils.is_transformable_smem_memref, "transforms", "SMEM ref", op,
    )


def _ensure_right_number_of_layouts(
    filter_fn: Callable[[ir.Value], bool],
    attr_suffix: str,
    value_type: str,
    op: ir.OpView,
) -> None:
  """Ensures that the right number of in/out layouts are provided for an op.

  Layouts here are can be vector layouts, TMEM layouts, or SMEM transforms.
  """
  layouts = lambda attr: op.attributes[attr] if attr in op.attributes else []
  in_layouts = layouts(f"in_{attr_suffix}")
  out_layouts = layouts(f"out_{attr_suffix}")

  num_matching_operands = sum(map(filter_fn, op.operands))
  if len(in_layouts) != num_matching_operands:
    raise ValueError(
        f"Expected the same number of in_{attr_suffix} ({len(in_layouts)}) as "
        f"{value_type} operands ({num_matching_operands}). op=\n  {op}"
    )
  num_matching_results = sum(map(filter_fn, op.results))
  if len(out_layouts) != num_matching_results:
    raise ValueError(
        f"Expected the same number of out_{attr_suffix} ({len(out_layouts)}) "
        f"as {value_type} results ({num_matching_results}). op=\n  {op}"
    )


@dataclasses.dataclass(frozen=True)
class _TypeAndLayout:
  type: ir.Type
  layout: cs.Constant


def assign_layouts(solution: dict[ValueSite, cs.Constant]) -> None:
  """Assigns the layouts in `solution` to the MLIR ops they belong to.

  This function requires that, for each MLIR op that appears in `solution`,
  `solution` contains a layout assignment for all of its `vector`, TMEM, and
  SMEM operands and results. Block arguments are ignored.
  """
  solution_sorted_by_op = sorted(
      solution.items(), key=lambda kv: id(kv[0].operation)
  )
  solution_per_op = itertools.groupby(
      solution_sorted_by_op, key=lambda kv: kv[0].operation
  )

  for op, assignments in solution_per_op:
    assignments_sorted_by_type = sorted(assignments, key=lambda kv: kv[0].type)
    assignments_by_type = {
        ty: list(group)
        for ty, group in itertools.groupby(
            assignments_sorted_by_type, key=lambda kv: kv[0].type
        )
    }

    in_assignments = assignments_by_type.get(VariableType.OPERAND, [])
    out_assignments = assignments_by_type.get(VariableType.RESULT, [])

    index = lambda kv: kv[0].index
    in_tls = [
        _TypeAndLayout(v.value.type, ce)
        for v, ce in sorted(in_assignments, key=index)
    ]
    out_tls = [
        _TypeAndLayout(v.value.type, ce)
        for v, ce in sorted(out_assignments, key=index)
    ]

    in_layouts = [
        tl.layout.value
        for tl in in_tls
        if isinstance(tl.layout, cs.RegisterLayout)
    ]
    out_layouts = [
        tl.layout.value
        for tl in out_tls
        if isinstance(tl.layout, cs.RegisterLayout)
    ]
    in_tmem_layouts = [
        tl.layout.value for tl in in_tls if isinstance(tl.layout, cs.TMEMLayout)
    ]
    out_tmem_layouts = [
        tl.layout.value
        for tl in out_tls
        if isinstance(tl.layout, cs.TMEMLayout)
    ]
    in_transforms = [
        tl for tl in in_tls if isinstance(tl.layout, cs.SMEMTransforms)
    ]
    out_transforms = [
        tl for tl in out_tls if isinstance(tl.layout, cs.SMEMTransforms)
    ]

    if inference_utils.should_have_in_layout(op):
      attrs = [layouts_lib.to_layout_attr(l) for l in in_layouts]
      op.attributes["in_layouts"] = ir.ArrayAttr.get(attrs)
    if inference_utils.should_have_out_layout(op):
      attrs = [layouts_lib.to_layout_attr(l) for l in out_layouts]
      op.attributes["out_layouts"] = ir.ArrayAttr.get(attrs)
    if inference_utils.should_have_in_tmem_layout(op):
      attrs = [layouts_lib.to_layout_attr(l) for l in in_tmem_layouts]
      op.attributes["in_tmem_layouts"] = ir.ArrayAttr.get(attrs)
    if inference_utils.should_have_out_tmem_layout(op):
      attrs = [layouts_lib.to_layout_attr(l) for l in out_tmem_layouts]
      op.attributes["out_tmem_layouts"] = ir.ArrayAttr.get(attrs)

    def _to_transform_attrs(
        transforms: list[_TypeAndLayout],
    ) -> list[ir.ArrayAttr]:
      all_attrs: list[ir.ArrayAttr] = []
      for tl in transforms:
        assert isinstance(tl.layout, cs.SMEMTransforms)
        attrs = []
        tiling_transform = tl.layout.tiling
        swizzle_v = tl.layout.swizzle
        if tiling_transform is not None:
          attrs.append(layouts_lib.to_transform_attr(tiling_transform))
        if swizzle_v is not None:
          swizzle = mgpu.SwizzlingMode(swizzle_v)
          attrs.append(layouts_lib.to_transform_attr(swizzle))
        all_attrs.append(ir.ArrayAttr.get(attrs))
      return all_attrs

    if inference_utils.should_have_in_transforms(op):
      attrs = _to_transform_attrs(in_transforms)
      op.attributes["in_transforms"] = ir.ArrayAttr.get(attrs)
    if inference_utils.should_have_out_transforms(op):
      attrs = _to_transform_attrs(out_transforms)
      op.attributes["out_transforms"] = ir.ArrayAttr.get(attrs)

    _ensure_all_layouts_are_set(op)


def vector_value_sites(op: ir.OpView) -> list[ValueSite]:
  """Returns all the vector operands and results for the given op."""
  value_sites = [
      ValueSite(op, VariableType.OPERAND, i)
      for i, o in enumerate(op.operands)
      if is_vector(o)
  ]
  value_sites.extend([
      ValueSite(op, VariableType.RESULT, i)
      for i, o in enumerate(op.results)
      if is_vector(o)
  ])
  return value_sites


def producer_result(operand: ValueSite) -> ValueSite:
  """Given an operand, returns the corresponding result in its producer.

  When the producer is a block, we return the corresponding operand in the
  operation that owns the block.
  """
  assert operand.type == VariableType.OPERAND
  value = operand.value
  producer = value.owner
  if isinstance(producer, ir.OpView):
    index = list[ir.Value](producer.results).index(value)
    return ValueSite(producer, VariableType.RESULT, index)

  if isinstance(producer, ir.Block):
    index = list[ir.Value](producer.arguments).index(value)
    region_index = list(producer.owner.regions).index(producer.region)
    return ValueSite(producer.owner, VariableType.ARGUMENT, index, region_index)

  raise TypeError(
      f"Producer {producer} is not an operation nor a block: {type(producer)}."
  )


def consumer_operands(result: ValueSite) -> Sequence[ValueSite]:
  """Given a result or an argument, returns the corresponding operands in its consumers."""
  assert result.type in (VariableType.RESULT, VariableType.ARGUMENT)
  results: list[ValueSite] = []
  # The layout can also be chosen from the layout of the consumers of the
  # results.
  for use in result.value.uses:
    consumer = use.owner
    index = use.operand_number
    results.append(ValueSite(consumer, VariableType.OPERAND, index))
  return results


def derive_relayout_constraints(
    value_sites_for_variable: ValueSitesForVariable,
) -> list[cs.Relayout]:
  """Derives relayout constraints from the given variable mapping."""
  constraints: list[cs.Relayout] = []
  variable_for_value_site: dict[ValueSite, cs.Variable] = {}
  for variable, value_sites in value_sites_for_variable.items():
    for value_site in value_sites:
      if value_site in variable_for_value_site:
        raise ValueError(
            f"{value_site} is mapped to both {variable} and "
            f"{variable_for_value_site[value_site]}"
        )
    variable_for_value_site |= {k: variable for k in value_sites}

  visited: set[cs.Variable] = set()
  for variable, value_sites in value_sites_for_variable.items():
    for value_site in value_sites:
      # We can only relayout variables that are in registers.
      if value_site.memory_space != cs.MemorySpace.REG:
        continue

      elt_bitwidth = utils.bitwidth(value_site.value.type.element_type)
      if value_site.type == VariableType.OPERAND:
        pr = producer_result(value_site)
        producer_variable = variable_for_value_site[pr]
        # Only add the constraint if we haven't already created that constraint
        # when processing this variable as one of the producer's consumers.
        if producer_variable not in visited:
          # The producer of a variable must be relayout-able to the variable.
          constraints.append(
              cs.Relayout(
                  producer_variable, variable, elt_bitwidth, strict=True
              )
          )
      elif value_site.type in (VariableType.RESULT, VariableType.ARGUMENT):
        for co in consumer_operands(value_site):
          consumer_variable = variable_for_value_site[co]
          # Only add the constraint if we haven't already created that
          # constraint when processing this variable as the consumer's producer.
          if consumer_variable not in visited:
            # A variable must be relayout-able to its consumers.
            constraints.append(
                cs.Relayout(
                    variable, consumer_variable, elt_bitwidth, strict=True
                )
            )
    visited.add(variable)
  return constraints


def is_terminator(op: ir.OpView) -> bool:
  return isinstance(op, (scf.YieldOp, scf.ConditionOp))


def traverse_op(
    op: ir.OpView,
    callback: Callable[[ir.OpView], None],
):
  """Traverses the operation and applies the callback in pre-order fashion.

  Skips recursing into `mgpu.CustomPrimitiveOp`s, and assumes that the values
  iterated on are not being modified.
  """
  callback(op)
  # The block of a mosaic_gpu.custom_primitive op is already lowered so it
  # should not be traversed.
  if not isinstance(op, mgpu.CustomPrimitiveOp):
    for region in op.operation.regions:
      for block in region:
        for block_op in block.operations:
          traverse_op(block_op, callback)


def check_layout_assignment(var: cs.Variable, layout: cs.Constant) -> None:
  """Raises if the given layout can not be assigned to the given `Variable`."""
  if not cs.is_valid_assignment(var, layout):
    raise ValueError(
        f"Variable {var} in memory space {var.memory_space} should not be "
        f"assigned a layout of type {type(layout)}. This is a bug."
    )


def _construct_value_error_with_op_stacktrace(
    msg: str, culprit_op: ir.OpView
) -> ValueError:
  tb = None
  try:
    tb = error.traceback_from_op(culprit_op.operation)
  except Exception:
    pass
  ve = ValueError(msg)
  if tb is not None:
    ve.__traceback__ = traceback_util.filter_traceback(tb)
  return ve


def _check_unsatisfiable_divisibility_constraints(
    ctx: DerivationContext,
    system: cs.ConstraintSystem,
) -> None:
  """Given `system` attempts to find unsatisfiable `cs.Divides` constraints.

  If at least one such a constraint is found, a `ValueError` is raised, pointing
  to the problematic value site.
  """
  divides_per_var = _divides_per_var(system.constraints)
  if not divides_per_var:
    return

  # Maps each variable with a divides constraint to its candidate assignments.
  # This list is intended to contain assignments only when we can list them
  # exhaustively.
  candidates_per_var: dict[cs.Variable, list[cs.SMEMTransforms]] = {
      v: [] for v in divides_per_var
  }

  for var, cst in system.assignments.items():
    if var in candidates_per_var:
      if isinstance(cst, cs.SMEMTransforms) and cst.tiling is not None:
        candidates_per_var[var].append(cst)

  # We only look at candidates from `OneOf` because candidates extracted from it
  # are exhaustive. Semantically it allows us to raise an error message saying
  # there's a problem with the kernel, and not the layout inference system
  # itself.
  for constraint in system.constraints:
    match constraint:
      case cs.OneOf(expr=cs.Variable() as var, allowed=allowed):
        if var not in candidates_per_var or var in system.assignments:
          continue
        for cand in allowed:
          if isinstance(cand, cs.SMEMTransforms) and cand.tiling is not None:
            candidates_per_var[var].append(cand)

  def _culprit_indices(tiling: tuple[int, ...], multiple: tuple[int, ...]):
    """Returns negative indices into the tiling for which the divisibility

    constraint is not satisfied.
    """
    assert len(tiling) <= len(multiple)
    for i, (t, m) in enumerate(zip(reversed(tiling), reversed(multiple))):
      if m % t != 0:
        yield -(i + 1)

  for var, constraint in divides_per_var.items():
    candidates = candidates_per_var[var]
    multiple = constraint.tiling_multiple
    if any(cs.Divides(c, multiple).holds() for c in candidates):
      continue

    # At this point we know that the contradiction for `var` within the `system`
    # exist. We attempt to fetch the troublesome value site to build a localized
    # error message.
    def transfer_ops():
      for vs in [var.key] + ctx.value_sites_for_variable.get(var, []):
        if isinstance(vs, ValueSite) and isinstance(
            vs.operation, (mgpu.AsyncLoadOp, mgpu.AsyncStoreOp)
        ):
          yield vs.operation

    for op in transfer_ops():
      for cand in candidates:
        # TODO(olechwierowicz): If only len(multiple) < len(tiling) candidates
        # are present we should raise too.
        assert cand.tiling is not None
        tiling = cand.tiling.tiling
        if len(tiling) > len(multiple):
          continue
        # None of the tiling candidates satisfy the constraint. Here
        # we attempt to find a case of a dynamic index of which we do not
        # know the value at compile time.
        for bad_idx in _culprit_indices(tiling, multiple):
          assert len(op.indices) == len(multiple)
          idx = op.indices[bad_idx]
          # TODO(olechwierowicz): We check for dynamic_gcd == 1, but a
          # dynamic_gcd > 1 could still be insufficient to satisfy the tiling
          # constraint. Relax this check to compare dynamic_gcd against the
          # required multiple.
          slice_length = list(op.slice_lengths)[bad_idx]
          if dynamic_gcd(slice_length, idx) != 1:
            continue
          msg = (
              "Failed to infer a possible set of layouts. You need to prove"
              f" the divisibility of index {bad_idx + len(op.indices)} using"
              " pl.multiple_of."
          )
          raise _construct_value_error_with_op_stacktrace(msg, idx)


def try_raise_constraint_specific_error(op: ir.OpView, constraint: cs.Constraint):
  # TODO(bchetioui): handle other constraint types.
  msg: str | None = None
  match constraint:
    case cs.IsTransferableSmemRegisters(
        source=cs.RegisterLayout(value=reg_layout),
        target=cs.SMEMTransforms() as smem_transforms,
        optimized=optimized
    ) | cs.IsTransferableSmemRegisters(
        source=cs.SMEMTransforms() as smem_transforms,
        target=cs.RegisterLayout(value=reg_layout),
        optimized=optimized
    ):
      match smem_transforms:
        case cs.SMEMTransforms(tiling=None, swizzle=None):
          ref_str = "untransformed ref"
        case cs.SMEMTransforms(tiling=lc.TileTransform() as t, swizzle=None):
          ref_str = f"ref with transforms (TilingTransform({t.tiling}),)"
        case cs.SMEMTransforms(tiling=None, swizzle=s) if s is not None:
          ref_str = f"ref with transforms (SwizzleTransform({s}),)"
        case cs.SMEMTransforms(tiling=t, swizzle=s):
          assert t is not None and s is not None
          ref_str = (
              "ref with transforms "
              f"(TilingTransform({t.tiling}), SwizzleTransform({s}))"
          )
      opt_str = "optimized " if optimized else ""
      msg = (
          f"no {opt_str}SMEM <-> registers transfer plan could be synthesized "
          f"for register layout {layouts_lib.pprint_layout(reg_layout)} and "
          f"{ref_str}"
      )
    case cs.Relayout(
        source=cs.RegisterLayout(value=source_layout),
        target=cs.RegisterLayout(value=target_layout),
    ):
      msg = (
          "cannot relayout from register layout "
          f"{layouts_lib.pprint_layout(source_layout)} to "
          f"{layouts_lib.pprint_layout(target_layout)}"
      )
    case cs.Equals(
        lhs=cs.RegisterLayout(value=lhs_layout),
        rhs=cs.RegisterLayout(value=rhs_layout),
    ):
      msg = (
          "expected register layouts "
          f"{layouts_lib.pprint_layout(lhs_layout)} and "
          f"{layouts_lib.pprint_layout(rhs_layout)} to be equal"
      )

  if msg is not None:
    raise _construct_value_error_with_op_stacktrace(
        f"Failed to infer a possible set of layouts: {msg}", op
    )


def diagnose_unsatisfiable_system(
    ctx: DerivationContext, initial_system: cs.ConstraintSystem
):
  """Raises the most specific error message possible for an unsatisfiable system."""
  # First, we deduce as many assignments as possible, until we reach the
  # expected unsatisfiable state, or can not reduce further.
  system = initial_system
  unsatisfied_assignments: dict[cs.Variable, cs.Constant] | None = None
  while True:
    match cs._reduce_system_once(system):
      case cs.Unsatisfiable():
        unsatisfied_assignments = system.assignments
        break
      case cs.ConstraintSystem() as reduced:
        system = reduced
      case None:
        break

  # If the constraint system is already known to be unsatisfiable here, we can
  # attempt to raise a specialized error message.
  if unsatisfied_assignments is not None:
    for constraint, op in ctx.op_for_constraint.items():
      reduced = cs.reduce_constraint(constraint, unsatisfied_assignments)
      if not isinstance(reduced, cs.Unsatisfiable) and reduced.holds() is False:
        try_raise_constraint_specific_error(op, reduced)

  # In the case below, we have not yet been able to find a single contradiction
  # to point to as the culprit. We reuse the initial, unreduced constraint
  # system to check for divisibility constraints.
  #
  # We saturate `cs.OneOf` constraints across equal variables to propagate
  # layout candidates (e.g. from ops like WGMMA) to equal variables/subviews
  # (e.g. created by slicing), allowing
  # `_check_unsatisfiable_divisibility_constraints` to inspect all relevant
  # candidate layouts for each variable when constructing error messages.
  saturated_system = cs.saturate_one_of_constraints_for_equal_vars(
      initial_system
  )
  if not isinstance(saturated_system, cs.Unsatisfiable):
    _check_unsatisfiable_divisibility_constraints(ctx, saturated_system)

  # TODO(bchetioui): add more error diagnostics.
  raise ValueError(
      "Failed to infer a possible set of layouts. This should only happen if "
      "user-provided layout casts are unsatisfiable. To dump the constraint "
      "system for debugging, use `MOSAIC_GPU_DUMP_CONSTRAINT_SYSTEM=1`."
  )


def _is_expensive_relayout(
    layout1: fa.FragmentedLayout, layout2: fa.FragmentedLayout
) -> bool:
  # TODO(bchetioui): handle TMEM_NATIVE relayouts, which can be free.
  return layout1 != layout2 and not isinstance(layout1, fa.WGSplatFragLayout)


def check_for_expensive_relayout(module: ir.Module):
  """Returns whether the given module has an expensive relayout."""
  layout_for_variable: dict[ir.Value, fa.FragmentedLayout] = {}

  def _check_for_expensive_relayout(op: ir.OpView):
    if not inference_utils.should_have_layout(op):
      return
    if inference_utils.has_in_layouts_set(op):
      in_layouts = iter(layouts_lib.from_layout_attr(l)
                        for l in cast(ir.ArrayAttr, op.attributes["in_layouts"]))
      for operand in op.operands:
        assert isinstance(operand, ir.Value)
        if not isinstance(operand.type, ir.VectorType):
          continue
        if operand not in layout_for_variable:
          raise ValueError(
              f"Inferred layout not found for operand {operand}."
          )
        operand_layout = next(in_layouts)
        # TODO(bchetioui): refine to figure out whether it's a cheap relayout.
        if _is_expensive_relayout(layout_for_variable[operand], operand_layout):
          raise ValueError(
              f"Inferred layout {operand_layout} for operand {operand} does "
              f"not match the layout in layout_for_variable "
              f"{layout_for_variable[operand]}."
          )
    if inference_utils.has_out_layouts_set(op):
      out_layouts = iter(layouts_lib.from_layout_attr(l)
                         for l in cast(ir.ArrayAttr, op.attributes["out_layouts"]))
      for result in op.results:
        assert isinstance(result, ir.Value)
        if not isinstance(result.type, ir.VectorType):
          continue
        assert result not in layout_for_variable
        layout_for_variable[result] = next(out_layouts)

  for op in module.body:
    traverse_op(op, _check_for_expensive_relayout)


def infer_layout(
    module: ir.Module, *, fuel: int = _DEFAULT_LAYOUT_INFERENCE_FUEL,
    arch: tuple[int, int] = (9, 0)
):
  """Infers layouts for the given module.

  * If there are vector (respectively SMEM refs, TMEM refs) operands,
  `in_layouts` (respectively `in_transforms`, `in_tmem_layouts`) will be set and
  contain one element per relevant argument in the memory space.
  * If there are vector (respectively SMEM refs, TMEM refs) outputs,
  `out_layouts` (respectively `out_transforms`, `out_tmem_layouts`) will be set
  and contain one element per relevant argument in the memory space.
  * Any of these attributes is guaranteed to not be set if there is no relevant
  input/output in the corresponding memory space.

  Args:
    module: The module to infer layouts for.
    fuel: The fuel is provided in order to limit the number of attempts made by
      the solver.
    arch: The architecture to infer layouts for.
  """
  global_constraint_system: cs.ConstraintSystem | cs.Unsatisfiable
  global_constraint_system = cs.ConstraintSystem()
  ctx = DerivationContext()

  def gather_constraints(op: Any):
    # Terminator ops are handled directly by the op whose region they belong to.
    # This is because they need to be in sync with their parent op's inputs and
    # outputs---and the parent op's constraints therefore need to take them into
    # account.
    if is_terminator(op):
      return
    should_have_layout = (
        inference_utils.should_have_layout(op)
        or inference_utils.should_have_tmem_layout(op)
        or inference_utils.should_have_transforms(op)
    )
    if not should_have_layout:
      return
    rule = _constraint_system_derivation_rules.get(op.OPERATION_NAME, None)
    if rule is None:
      raise NotImplementedError(f"No layout inference rule defined for {op}")
    rule_result = rule(ctx, op)
    nonlocal global_constraint_system
    constraint_system, mapping = rule_result
    for var, sites in mapping.items():
      assert isinstance(var.key, ValueSite)
      for site in sites:
        if site.memory_space != var.memory_space:
          raise ValueError(
              f"Memory space mismatch between variable and {site}:"
              f" {var.memory_space} != {site.memory_space}."
          )
        if site.shape != var.shape:
          raise ValueError(
              f"Shape mismatch between variable and {site}:"
              f" {var.shape} != {site.shape}."
          )

    for constraint in constraint_system.constraints:
      ctx.op_for_constraint.setdefault(constraint, op)

    global_constraint_system &= constraint_system
    ctx.update(mapping)

  for op in module.body:
    traverse_op(op, gather_constraints)
    # Short-circuit if we have an unsatisfiable constraint system, we won't
    # construct anything useful anymore.
    if isinstance(global_constraint_system, cs.Unsatisfiable):
      break

  if isinstance(global_constraint_system, cs.Unsatisfiable):
    # TODO(bchetioui): rework this to raise a nicer error message here too. This
    # case is only hit if there are directly conflicting assignments (not even
    # constraints), so it should be easy.
    raise ValueError(
        "Failed to infer a possible set of layouts. This should only happen if "
        "user-provided layout casts are unsatisfiable. To dump the constraint "
        "system for debugging, use `MOSAIC_GPU_DUMP_CONSTRAINT_SYSTEM=1`."
    )

  constraints = derive_relayout_constraints(ctx.value_sites_for_variable)
  for constraint in constraints:
    assert isinstance(constraint.target, cs.Variable)
    ctx.op_for_constraint.setdefault(
        constraint, constraint.target.key.operation
    )
  global_constraint_system &= cs.ConstraintSystem(constraints=constraints)
  assert not isinstance(global_constraint_system, cs.Unsatisfiable)

  # Add additional (redundant) constraints which helps the search converge
  # faster.
  global_constraint_system = cs.saturate_distinct_from_splat(
      global_constraint_system
  )
  assert not isinstance(global_constraint_system, cs.Unsatisfiable)
  global_constraint_system = (
      cs.canonicalize_strict_non_splat_relayouts_to_equals(
          global_constraint_system
      )
  )
  global_constraint_system = cs.saturate_divides_constraints_for_equal_vars(
      global_constraint_system
  )

  dump_options = mgpu.get_or_set_dump_options(module)  # pyrefly: ignore[missing-attribute]
  if dump_options.constraint_system:
    utils.dump_to_file_or_stdout(
        str(global_constraint_system),
        f"{dump_options.module_basename}.constraint_system.txt",
        dump_options.dump_path
    )

  # Attempt to find assignments that satisfy the constraint system.
  solution, remaining_fuel = find_assignments_for(
      list(ctx.value_sites_for_variable.keys()),
      global_constraint_system,
      fuel=fuel,
      arch=arch,
  )

  if logger.isEnabledFor(logging.DEBUG):
    print("Finding a solution (or exhausting the entire search space) "
          f"consumed {fuel - remaining_fuel}/{fuel} fuel.")

  if isinstance(solution, cs.Unsatisfiable):
    # At this point we know the system is unsatisfiable. We attempt to find
    # contradictions in the `cs.ConstraintSystem`, and raise a meaningful error
    # message.
    diagnose_unsatisfiable_system(ctx, global_constraint_system)
    raise RuntimeError("Unreachable")

  layout_for_value_site: dict[ValueSite, cs.Constant] = {}
  for variable, value_sites in ctx.value_sites_for_variable.items():
    layout = solution[variable]
    # Ensure that the layout assignment is valid for the variable. This should
    # only ever fail if our implementation is buggy.
    check_layout_assignment(variable, layout)
    for value_site in value_sites:
      layout_for_value_site[value_site] = layout

  # Assigns the layouts that we found to the ops.
  assign_layouts(layout_for_value_site)

  # Sanity check: ensure that all ops have the right number of in/out layouts.
  for op in module.body:
    traverse_op(op, _ensure_all_layouts_are_set)