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

"""Defines expressions and constraints over layouts."""

from __future__ import annotations

import abc
from collections.abc import Sequence
import dataclasses
import enum
import math
from typing import Any, TYPE_CHECKING, Literal, assert_never, cast, final

from jax._src.lib.mlir.dialects import builtin, func
from jax._src.lib.mlir import ir
import numpy as np

from . import dialect_lowering as lowering
from . import fragmented_array as fa
from . import launch_context as lc
from . import layouts as layouts_lib
from . import tcgen05
from . import utils


class MemorySpace(enum.Enum):
  """The memory space of a variable."""

  REG = enum.auto()
  SMEM = enum.auto()
  TMEM = enum.auto()


# TODO(bchetioui): consider defining an interface for variable keys that carry
# shape and memory space information.
VariableKey = Any


@dataclasses.dataclass(frozen=True)
class Variable:
  """A variable is an abstract identifier.

  `key` is supposed to be hashable.
  """
  key: VariableKey

  @property
  def shape(self) -> tuple[int, ...]:
    return self.key.shape

  @property
  def memory_space(self) -> MemorySpace:
    return self.key.memory_space

  def __str__(self):
    return f"V({self.key})"


@dataclasses.dataclass(frozen=True)
class RegisterLayout:
  """Wraps a known register layout."""

  value: fa.FragmentedLayout

  def __str__(self):
    return f"C({self.value})"


@dataclasses.dataclass(frozen=True)
class TMEMLayout:
  """Wraps a known TMEM layout."""

  value: tcgen05.TMEMLayout

  def __str__(self):
    return f"C({self.value})"


Swizzle = Literal[32, 64, 128] | None


@dataclasses.dataclass(frozen=True)
class SMEMTransforms:
  """Wraps known SMEM transforms.

  If an SMEM reference may, in principle, have transforms but should not be
  tiled, then `tiling` is `None`.

  Attributes:
    tiling: The tiling transform.
    swizzle: The swizzle pattern length in bytes. Note that it may not
      correspond to the trailing dimension of the provided `tiling`, since these
      expressions may be transposed (swizzle is not transposed along with
      tiling). A `None` swizzle means no swizzling.

  """

  tiling: lc.TileTransform | None
  swizzle: Swizzle

  def __post_init__(self):
    if self.swizzle and self.swizzle not in {32, 64, 128}:
      raise ValueError(
          f"Swizzle {self.swizzle} is not supported. Only 32, 64 and 128 are"
          " accepted."
      )

  def __str__(self):
    return f"C({self.tiling}, {self.swizzle})"


Constant = RegisterLayout | TMEMLayout | SMEMTransforms


@dataclasses.dataclass(frozen=True)
class Reduce:
  expression: Expression
  axes: tuple[int, ...]
  # The rank of the shape of the input to the reduction. It is necessary to
  # know this in order to reduce `TiledLayout`s correctly.
  rank: int
  # If `True`, the axes which are reduced are left in the result as dimensions
  # with size one.
  keep_dims: bool = False

  def __post_init__(self):
    assert self.axes
    assert self.axes == tuple(sorted(self.axes)), "axes must be sorted."
    assert self.axes[0] >= 0, "axes must not contain negative indices."

  def __str__(self):
    return (
        f"Reduce([{self.axes}], {self.expression}, rank={self.rank},"
        f" keep_dims={self.keep_dims})"
    )


@dataclasses.dataclass(frozen=True)
class Reshape:
  expression: Expression
  source_shape: tuple[int, ...]
  target_shape: tuple[int, ...]


@dataclasses.dataclass(frozen=True)
class Transpose:
  expression: Expression
  permutation: tuple[int, ...]

  def __post_init__(self):
    if sorted(self.permutation) != list(range(len(self.permutation))):
      raise ValueError(f"Invalid permutation {self.permutation}")

  def __str__(self):
    return f"T({self.expression}, permutation={self.permutation})"


@dataclasses.dataclass(frozen=True)
class CollapseShape:
  """Collapses a shape into a lower rank shape via a reassociation.

  `reassociation` is a tuple of integers, where each integer is the number of
  contiguous dimensions that must be collapsed together. Each group must
  consist of at least one dimension, and `sum(reassociation)` must be equal to
  the rank of the source shape.
  """
  expression: Expression
  source_shape: tuple[int, ...]
  reassociation: tuple[int, ...]

  def __post_init__(self):
    for num_collapsed_dims in self.reassociation:
      if num_collapsed_dims <= 0:
        raise ValueError(
            f"Invalid reassociation {self.reassociation}. Each group of "
            "collapsed dimensions must contain at least one dimension."
        )
    if sum(self.reassociation) != len(self.source_shape):
      raise ValueError(
          f"Invalid reassociation {self.reassociation}. The number of collapsed"
          f" dimensions must be equal to the rank of the source shape."
      )

  def __str__(self):
    return (f"CollapseShape({self.expression}, source_shape={self.source_shape}"
            f", reassociation={self.reassociation})")


Expression = (
    Variable
    | Constant
    | Reduce
    | Reshape
    | Transpose
    | CollapseShape
)


def reduce_reshape_expression(
    reshape: Reshape, assignments: dict[Variable, Constant]
) -> Expression | Unsatisfiable:
  reduced_expr = reduce_expression(reshape.expression, assignments)
  match reduced_expr:
    case Unsatisfiable():
      return Unsatisfiable()
    case RegisterLayout(value=layout):
      match layout:
        case fa.WGSplatFragLayout(shape=shape):
          assert math.prod(shape) == math.prod(reshape.target_shape)
          return RegisterLayout(
              fa.WGSplatFragLayout(shape=reshape.target_shape)
          )
        case fa.WGStridedFragLayout(shape=shape, vec_size=vec_size):
          assert math.prod(shape) == math.prod(reshape.target_shape)
          return RegisterLayout(
              fa.WGStridedFragLayout(
                  shape=reshape.target_shape, vec_size=vec_size
              )
          )
        case fa.TiledLayout() as tiled_layout:
          tile_shape = tiled_layout.base_tile_shape
          if len(reshape.target_shape) < len(tile_shape):
            return dataclasses.replace(reshape, expression=reduced_expr)
          # Even if the new shape is not perfectly tilable, it is possible that
          # we may be able to reshape the tiling itself in a way that is
          # compatible with the new shape. We do not handle this case at the
          # moment.
          for ts, s in zip(tile_shape, reshape.source_shape[-len(tile_shape):], strict=True):
            if s % ts != 0:
              return dataclasses.replace(reshape, expression=reduced_expr)

          # If minor tiled dimensions are modified, then reshaping is likely to
          # not be a no-op since the strides between tiles will change,
          # potentially mapping different elements to lanes and warps. We don't
          # attempt to handle this case at the moment.
          num_minor_tiled_dims = len(tile_shape) - 1
          source_minor_tiled_dims = reshape.source_shape[-num_minor_tiled_dims:]
          target_minor_tiled_dims = reshape.target_shape[-num_minor_tiled_dims:]
          major_tiled_dim = tile_shape[0]
          if (source_minor_tiled_dims != target_minor_tiled_dims or
              reshape.target_shape[-len(tile_shape)] % major_tiled_dim != 0):
            return dataclasses.replace(reshape, expression=reduced_expr)
          # At this point, we now that only non-tiled dimensions and/or the
          # majormost tiled dimensions may have changed. We also know that the
          # majormost tiled dimension is still tilable in the new shape.
          # Therefore, we can return the tiled layout as is.
          return RegisterLayout(tiled_layout)
    case _:
      return dataclasses.replace(reshape, expression=reduced_expr)


def reduce_transpose_expression(
    transpose: Transpose, assignments: dict[Variable, Constant]
) -> Expression | Unsatisfiable:
  reduced_expr = reduce_expression(transpose.expression, assignments)
  match reduced_expr:
    case Unsatisfiable():
      return Unsatisfiable()
    case SMEMTransforms(tiling=tile_transform, swizzle=swizzle):
      if tile_transform is None:
        return SMEMTransforms(None, swizzle)
      tiling = tile_transform.tiling
      permutation = transpose.permutation
      tiling_offset = len(permutation) - len(tiling)
      # We reject if there's a swap between tiled dimensions and untiled dimensions.
      #
      # For example:
      #   A permutation (0, 3, 2, 1) and tiling of length <=2, we reject because tiling becomes non-contiguous.
      #   A permutation (0, 3, 2, 1) and tiling of length 3, we accept.
      if any(dim < tiling_offset for dim in permutation[-len(tiling) :]):
        return Unsatisfiable()
      new_tiling = tuple(tiling[dim - tiling_offset] for dim in permutation[-len(tiling):])
      return SMEMTransforms(lc.TileTransform(new_tiling), swizzle)
    case _:
      return Transpose(expression=reduced_expr, permutation=transpose.permutation)


def reduce_reduce_expression(
    expr: Reduce, assignments: dict[Variable, Constant]
) -> Expression | Unsatisfiable:
  reduced_expr = reduce_expression(expr.expression, assignments)

  def default():
    """We don't know how to reduce further."""
    assert not isinstance(reduced_expr, Unsatisfiable)
    return dataclasses.replace(expr, expression=reduced_expr)

  match reduced_expr:
    case Unsatisfiable():
      return Unsatisfiable()
    case RegisterLayout(value=fa.TiledLayout() as layout):
      # TODO(allanrenucci): Add support for reducing tiled layouts when keep_dims=True.
      if expr.keep_dims:
        return default()
      num_untiled_dims = expr.rank - len(layout.base_tile_shape)
      reduced_tiling_axes = [
          a - num_untiled_dims for a in expr.axes if a >= num_untiled_dims
      ]
      if reduced_tiling_axes:
        return RegisterLayout(layout.reduce(reduced_tiling_axes))
      return RegisterLayout(layout)
    case RegisterLayout(value=fa.WGStridedFragLayout() as layout):
      # We don't support cross-lane and non vec_size preserving reductions.
      trailing_dims = layout.shape[expr.axes[-1] + 1 :]
      if math.prod(trailing_dims) % (layout.vec_size * fa.WARPGROUP_SIZE):
        return default()
      shape = utils.reduce_shape(layout.shape, expr.axes, expr.keep_dims)
      return RegisterLayout(fa.WGStridedFragLayout(shape, layout.vec_size))
    case RegisterLayout(value=fa.WGSplatFragLayout() as layout):
      shape = utils.reduce_shape(layout.shape, expr.axes, expr.keep_dims)
      return RegisterLayout(fa.WGSplatFragLayout(shape))
    case _:
      return default()


def reduce_collapse_shape_expression(
    expr: CollapseShape, assignments: dict[Variable, Constant]
) -> Expression | Unsatisfiable:
  reduced_expr = reduce_expression(expr.expression, assignments)
  match reduced_expr:
    case Unsatisfiable():
      return Unsatisfiable()
    case SMEMTransforms(tiling=tile_transform, swizzle=swizzle):
      if tile_transform is None:
        return reduced_expr
      tiling = tile_transform.tiling
      rev_tiling_to_process = list(tiling)[::-1]
      rev_shape_to_process = expr.source_shape[-len(tiling):][::-1]
      # Ensure that the provided tiling applies to the shape. Otherwise, the
      # expression is unsatisfiable.
      for s, t in zip(rev_shape_to_process, rev_tiling_to_process):
        if s % t != 0:
          return Unsatisfiable()
      rev_new_tiling: list[int] = []
      for ndim in expr.reassociation[::-1]:
        # Collapsing tiled dimensions into untiled dimensions is not
        # supported. While there is a reasonable way of handling this case in
        # particular situations, we forbid it in the semantics of
        # `CollapseShape`.
        if len(rev_tiling_to_process) < ndim:
          return Unsatisfiable()

        rev_tiling_slice = rev_tiling_to_process[:ndim]
        rev_shape_slice = rev_shape_to_process[:ndim]
        new_tiling_dim = math.prod(rev_tiling_slice)
        num_elems = math.prod(rev_shape_slice)
        assert num_elems % new_tiling_dim == 0
        # We can collapse dimensions when the tiling is of the form
        # (1*, partial_dim?, full_dim*)---i.e. when it contains any number of
        # leading unit dimensions, followed by at most one arbitrary non-unit
        # dimension, and any number of trailing "full" dimensions (where the
        # tiling size equals the dimension size).
        #
        # Here, the tiling and shape are reversed, so we look for the pattern
        # (full_dim*, partial_dim?, 1*).
        suffix_length = 0
        for t, s in zip(rev_tiling_slice, rev_shape_slice):
          if t != s:
            break
          suffix_length += 1
        if (rev_unsuffixed_tiling := rev_tiling_slice[suffix_length:]):
          # Ignore the partial dimension, since it can be anything.
          _, *rev_prefix_tiling = rev_unsuffixed_tiling
          if any(t != 1 for t in rev_prefix_tiling):
            return Unsatisfiable()
        rev_new_tiling.append(new_tiling_dim)
        rev_tiling_to_process = rev_tiling_to_process[ndim:]
        rev_shape_to_process = rev_shape_to_process[ndim:]
        if not rev_tiling_to_process:
          break
      assert not rev_tiling_to_process
      assert not rev_shape_to_process
      new_tiling = tuple(rev_new_tiling[::-1])
      return SMEMTransforms(lc.TileTransform(tuple(new_tiling)), swizzle)
    case _ if isinstance(reduced_expr, Constant):
      raise NotImplementedError(
          "CollapseShape is only implemented for variables in SMEM")
    case _:
      return dataclasses.replace(expr, expression=reduced_expr)


def reduce_expression(
    expr: Expression, assignments: dict[Variable, Constant]
) -> Expression | Unsatisfiable:
  """Reduces an expression as much as is possible given a set of known variable assignments."""
  match expr:
    case RegisterLayout() | TMEMLayout() | SMEMTransforms():
      return expr
    case Variable():
      return assignments.get(expr, expr)
    case Reduce():
      return reduce_reduce_expression(expr, assignments)
    case Reshape():
      return reduce_reshape_expression(expr, assignments)
    case Transpose():
      return reduce_transpose_expression(expr, assignments)
    case CollapseShape():
      return reduce_collapse_shape_expression(expr, assignments)
    case _:
      assert_never(expr)


class _BaseConstraint(abc.ABC):

  @property
  @abc.abstractmethod
  def _is_constant(self) -> bool:
    """Returns true all expressions this constraint depends on are `Constant`, false otherwise."""

  def holds(self) -> bool | None:
    if self._is_constant:
      return self._constant_holds()
    else:
      return None

  @abc.abstractmethod
  def _constant_holds(self) -> bool:
    """Evaluates the constraint when all underlying expressions are constants."""


@dataclasses.dataclass(frozen=True)
class AlwaysTrue(_BaseConstraint):
  def holds(self) -> bool | None:
    return True

  def _constant_holds(self) -> bool:
    raise NotImplementedError

  @property
  def _is_constant(self):
    raise NotImplementedError


@dataclasses.dataclass(frozen=True)
class Equals(_BaseConstraint):
  """States that `lhs` and `rhs` are equal."""
  lhs: Expression
  rhs: Expression

  @property
  def _is_constant(self) -> bool:
    return isinstance(self.lhs, Constant) and isinstance(self.rhs, Constant)

  def canonicalize(self) -> Constraint:
    if self.lhs == self.rhs:
      return AlwaysTrue()
    return self

  def _constant_holds(self) -> bool:
    return self.lhs == self.rhs

  def __str__(self):
    return f"Equals({self.lhs} == {self.rhs})"


def _is_supported_tiled_relayout(
    src: fa.TiledLayout, dst: fa.TiledLayout, bitwidth: int
) -> bool:
  """Returns whether the source->target relayout is supported for values of types with the given bitwidth."""
  match src, dst:
    # Transposed layouts.
    case fa.WGMMA_LAYOUT, fa.WGMMA_TRANSPOSED_LAYOUT:
      return True
    case fa.WGMMA_TRANSPOSED_LAYOUT, fa.WGMMA_LAYOUT:
      return True
    case fa.TCGEN05_LAYOUT, fa.TCGEN05_TRANSPOSED_LAYOUT:
      return True
    case fa.TCGEN05_TRANSPOSED_LAYOUT, fa.TCGEN05_LAYOUT:
      return True
    # "Conversion-optimized" layouts.
    case (fa.WGMMA_LAYOUT_UPCAST_2X, fa.WGMMA_LAYOUT) | (
        fa.WGMMA_LAYOUT,
        fa.WGMMA_LAYOUT_UPCAST_2X,
    ):
      return fa.can_relayout_wgmma_2x_to_wgmma(bitwidth)
    case fa.WGMMA_LAYOUT_UPCAST_4X, fa.WGMMA_LAYOUT_UPCAST_2X:
      return fa.can_relayout_wgmma_4x_to_wgmma_2x(bitwidth)
    case fa.WGMMA_LAYOUT_UPCAST_4X, fa.WGMMA_LAYOUT:
      return fa.can_relayout_wgmma_4x_to_wgmma_2x(
          bitwidth
      ) and fa.can_relayout_wgmma_2x_to_wgmma(bitwidth)
  if src == fa.tmem_native_layout(
      src.vector_length
  ) and dst == fa.tmem_native_layout(dst.vector_length):
    return True
  return False


@dataclasses.dataclass(frozen=True)
class Relayout(_BaseConstraint):
  """States that `source` must be relayout-able to `target`.

  Relayout-ability here is not defined as a fundamental property of layouts, but
  rather a reflection of our implementation. For instance, when evaluating this
  constraint, we will return `False` systematically if a relayout exists but we
  do not ever plan to support it.

  Modeling this constraint this way is helpful, in order to allow pruning
  inefficient solutions when attempting to solve a constraint system.

  We include here the bitwidth of the element type we want to associate with
  this constraint, as certain relayouts are only supported for specific
  bitwidths.

  If `strict` is `True`, only allows relayout from splat layouts and force
  layout equality otherwise.
  """

  source: Expression
  target: Expression
  bitwidth: int
  strict: bool = False

  def canonicalize(self) -> Constraint:
    if self.source == self.target:
      return AlwaysTrue()
    match self:
      # The only valid strict tiled and strided relayout is the identity.
      case Relayout(
          source=RegisterLayout(
              value=fa.TiledLayout() | fa.WGStridedFragLayout()
          ) as cst,
          target=target,
          strict=True,
      ):
        return Equals(lhs=cst, rhs=target)
      case _:
        return self

  @property
  def _is_constant(self) -> bool:
    return isinstance(self.source, Constant) and isinstance(self.target, Constant)

  def _constant_holds(self) -> bool:
    source = self.source
    target = self.target

    if not isinstance(source, RegisterLayout) or not isinstance(
        target, RegisterLayout
    ):
      raise ValueError(
          f"Relayout can only be applied to registers, got source {self.source}"
          f" to {self.target}"
      )

    if source == target:
      return True

    source_layout, target_layout = source.value, target.value
    match source_layout, target_layout:
      case fa.WGSplatFragLayout() as splat, fa.WGStridedFragLayout() as strided:
        return splat.shape == strided.shape
      case fa.WGSplatFragLayout(), fa.TiledLayout():
        return layouts_lib.splat_is_compatible_with_tiled(
            source_layout, target_layout
        )
      case fa.TiledLayout(), fa.TiledLayout() if not self.strict:
        return _is_supported_tiled_relayout(
            source_layout, target_layout, self.bitwidth
        )
      case _:
        return False

  def __str__(self):
    return f"Relayout({self.source}  ⟶ {self.target})"


@dataclasses.dataclass(frozen=True)
class IsTransferable(_BaseConstraint, abc.ABC):
  """States that `source` layout must be transferable across memory spaces to `target` layout."""
  source: Expression
  target: Expression
  shape: tuple[int, ...]

  @property
  def _is_constant(self) -> bool:
    return isinstance(self.source, Constant) and isinstance(self.target, Constant)


@dataclasses.dataclass(frozen=True)
class IsTransferableTmemRegisters(IsTransferable):
  """States that `source` layout must be transferable across memory spaces to `target` layout.

  In this case, one of `source` and `target` must be in TMEM, and the other must
  be in registers. `bitwidth` is the bitwidth of the element type.
  """
  bitwidth: int

  def __post_init__(self):
    assert len(self.shape) == 2
    assert 0 < self.bitwidth <= 32

  def is_valid_tmem_transfer(
      self, tmem_layout: tcgen05.TMEMLayout, reg_layout: fa.FragmentedLayout
  ) -> bool:
    if not isinstance(reg_layout, fa.TiledLayout):
      return False
    packing = tmem_layout.vector_length
    columns = self.shape[1]
    if (
        reg_layout == fa.TCGEN05_LAYOUT
        and tmem_layout == tcgen05.tmem_default_layout(packing)
    ):
      return True
    if (
        reg_layout == tmem_layout.as_tiled_layout()
        and packing * self.bitwidth == 32
    ):
      return True
    if (
        reg_layout == fa.TMEM_NATIVE_LAYOUT
        and tmem_layout == tcgen05.tmem_default_layout(packing)
        and ((self.bitwidth == 16 and packing == 1) or self.bitwidth == 32)
    ):
      return True
    if columns % 16 == 0 and packing <= columns // 2 and (
        reg_layout == fa.WGMMA_LAYOUT
        and tmem_layout == tcgen05.tmem_half_lane_layout(columns, packing)
    ):
      return True
    if columns % 16 == 0 and packing <= 8 and (
        reg_layout == tcgen05.fa_m64_collective_layout(columns)
        and tmem_layout == tcgen05.tmem_m64_collective_layout(columns, packing)
    ):
      return True
    return False

  def _constant_holds(self) -> bool:
    match self.source, self.target:
      case RegisterLayout(value=src), TMEMLayout(value=dst):
        return self.is_valid_tmem_transfer(dst, src)
      case TMEMLayout(value=src), RegisterLayout(value=dst):
        return self.is_valid_tmem_transfer(src, dst)
      case _:
        raise ValueError(
            f"{self.source} -> {self.target} is not a TMEM <-> Registers"
            " transfer."
        )

  def __str__(self):
    return f"IsTransferableTmemRegisters({self.source} ⟶ {self.target})"


@dataclasses.dataclass(frozen=True)
class IsTransferableSmemRegisters(IsTransferable):
  """States that `source` layout must be transferable across memory spaces to `target` layout.

  In this case, one of `source` and `target` must be in SMEM, and the other must
  be in registers.
  """
  strides: tuple[int, ...]
  bitwidth: int
  optimized: bool

  def _is_supported_smem_transfer(
      self,
      smem_layout: SMEMTransforms,
      reg_layout: fa.FragmentedLayout,
  ) -> bool:
    tiling_transform = smem_layout.tiling
    swizzle = smem_layout.swizzle

    if not isinstance(reg_layout, fa.TiledLayout):
      return tiling_transform is None and swizzle is None
    tiling = tiling_transform.tiling if tiling_transform is not None else ()
    tiling_rank = len(tiling)

    # If `tiling_rank` is 0, then we tile by the shape. This is the logic that
    # is implemented in `load_untiled` and `store_untiled`.
    if tiling_rank == 0:
      tiling = self.shape
      tiling_rank = len(tiling)

    int_ty = ir.IntegerType.get_signless(self.bitwidth)
    layout = ir.StridedLayoutAttr.get(0, lowering.tile_strides(self.strides, tiling))
    ref_shape = utils.tile_shape(self.shape, tiling)
    assert len(layout.strides) == len(ref_shape), (len(layout.strides), len(ref_shape))
    memref_ty = ir.MemRefType.get(ref_shape, int_ty, layout, utils.smem())
    dummy_func = func.FuncOp("dummy_transfer_test", ir.FunctionType.get([], []))
    dummy_block = dummy_func.add_entry_block()
    try:
      with ir.InsertionPoint(dummy_block):
        fake_ref_op = builtin.UnrealizedConversionCastOp([memref_ty], [])
        fake_ref = fake_ref_op.results[0]
        for use_txmatrix in [True, False]:
          try:
            next(
                fa.FragmentedArray.transfer_tiled(
                    fake_ref,
                    swizzle or 16,
                    reg_layout,
                    self.shape,
                    optimized=self.optimized,
                    ref_tiling_rank=tiling_rank,
                    use_txmatrix=use_txmatrix,
                )
            )
            return True
          except (fa.TxMatrixIneligible, fa.TransferPlanDerivationError, fa.UnsupportedTransferError):
            continue
        return False
    finally:
      dummy_func.erase()

  def _constant_holds(self) -> bool:
    match self.source, self.target:
      case SMEMTransforms() as src, RegisterLayout(value=dst):
        return self._is_supported_smem_transfer(src, dst)
      case RegisterLayout(value=src), SMEMTransforms() as dst:
        return self._is_supported_smem_transfer(dst, src)
      case _:
        raise ValueError(
            f"{self.source} -> {self.target} is not a SMEM <-> Registers transfer."
        )

  def __str__(self):
    return f"IsTransferableSmemRegisters({self.source} ⟶ {self.target})"


@dataclasses.dataclass(frozen=True)
class NotOfType(_BaseConstraint):
  """States that `expr` is not an instance of `type`."""

  expr: Expression
  type: type[fa.FragmentedLayout]

  @property
  def _is_constant(self) -> bool:
    return isinstance(self.expr, Constant)

  def _constant_holds(self) -> bool:
    # TODO(olechwierowicz): We should raise here because this constraint is
    # supposed to be used only for register expressions.
    if not isinstance(self.expr, RegisterLayout):
      return True
    return not isinstance(self.expr.value, self.type)

  def __str__(self):
    return f"type({self.expr}) ≠ {self.type.__name__}"


@dataclasses.dataclass(frozen=True)
class Divides(_BaseConstraint):
  """States that the `expr` tiling is a divisor of `tiling_multiple`.

  That is to say that, for each tiled dimension in `expr`, the dimension must
  divide its corresponding dimension in `tiling_multiple` starting from the
  tail.

  If `tiling_multiple` contains more dimensions than `expr`, then
  the extra dimensions in `tiling_multiple` are ignored for the purposes of the
  check.

  `expr` is not allowed to contain more dimensions than `tiling_multiple`, and
  this constraint therefore also constrains the rank of `expr`.

  If `expr` is a `RegisterLayout` of `WGStridedFragLayout` or
  `WGSplatFragLayout` we return True.
  """
  expr: Expression
  tiling_multiple: tuple[int, ...]

  @property
  def _is_constant(self) -> bool:
    return isinstance(self.expr, Constant)

  def _constant_holds(self) -> bool:
    assert isinstance(self.expr, Constant)
    match self.expr:
      case SMEMTransforms(tiling=None):
        # If there is no tiling, then this holds trivially.
        return True
      case SMEMTransforms(tiling=lc.TileTransform(tiling=t)):
        tiling = t
      case RegisterLayout(
          value=fa.WGStridedFragLayout() | fa.WGSplatFragLayout()
      ):
        return True
      case RegisterLayout(value=fa.TiledLayout() as layout):
        tiling = layout.base_tile_shape
      case TMEMLayout(value):
        tiling = value.base_tile_shape
      # We check the case below only because type checker does not recognize
      # we've covered all `RegisterLayout` or `SMEMTransforms` cases.
      case SMEMTransforms() | RegisterLayout():
        raise ValueError(f"Unhandled expression: {self.expr}")
      case _ as never:
        assert_never(never)

    if len(tiling) > len(self.tiling_multiple):
      # The rank of the tiling is larger than the rank of the constraint. This
      # is not allowed.
      return False

    for size, multiple in zip(reversed(tiling), reversed(self.tiling_multiple)):
      if multiple % size:
        return False
    return True

  def __str__(self):
    return f"{self.tiling_multiple} % {self.expr} == 0"


@dataclasses.dataclass(frozen=True)
class MinorDimDivisibleBy(_BaseConstraint):
  """States that the minor dimension of the `expr` tiling is divisible by `divisor`.

  If the last dimension is untiled, then `true` is returned.

  If `expr` is not `SMEMTransforms` but any other constant `ValueError` is raised.
  """
  expr: Expression
  divisor: int

  @property
  def _is_constant(self) -> bool:
    return isinstance(self.expr, Constant)

  def _constant_holds(self) -> bool:
    assert isinstance(self.expr, Constant)
    match self.expr:
      case SMEMTransforms(tiling=None):
        return True
      case SMEMTransforms(tiling=lc.TileTransform(tiling=t)):
        tiling = t
      case RegisterLayout() | SMEMTransforms() | TMEMLayout():
        raise ValueError(
            f"Unexpected expression {self.expr} in MinorDimDivisibleBy constraint"
        )
      case _ as never:
        assert_never(never)

    if not tiling:
      return True

    return tiling[-1] % self.divisor == 0

  def __str__(self):
    return f"{self.expr}.tiling[-1] % {self.divisor} == 0"


@dataclasses.dataclass(frozen=True)
class IsSupportedBroadcast(_BaseConstraint):
  """States that `src` can be broadcasted to `dst`.

  See `FragmentedArray.broadcast_in_dim` for more details.
  """

  src: Expression
  dst: Expression
  dims: tuple[int, ...]

  @property
  def _is_constant(self) -> bool:
    return isinstance(self.src, Constant) and isinstance(self.dst, Constant)

  def _constant_holds(self) -> bool:
    assert isinstance(self.src, Constant) and isinstance(self.dst, Constant)
    match self.src, self.dst:
      case RegisterLayout(
          value=fa.WGStridedFragLayout() as src_layout
      ), RegisterLayout(value=fa.WGStridedFragLayout() as dst_layout):
        return fa.is_supported_strided_layout_broadcast(src_layout, dst_layout, self.dims)
      case RegisterLayout(value=src_layout), RegisterLayout(value=dst_layout):
        # This is an intentionally loose check. We rely on the presence of a
        # `src = Reduce(dst)` constraint to enforce correctness.
        return type(src_layout) == type(dst_layout)
      case _:
        raise ValueError(
            f"Unexpected values {self.src=} {self.dst=} in IsSupportedBroadcast"
            " constraint"
        )

  def __str__(self):
    return (
        f"IsSupportedBroadcast(src={self.src}, dst={self.dst},"
        f" dims={self.dims})"
    )


@dataclasses.dataclass(frozen=True)
class OneOf(_BaseConstraint):
  """States that `expr` evaluates to one of the given constants."""

  expr: Expression
  allowed: tuple[Constant, ...]

  @property
  def _is_constant(self) -> bool:
    return isinstance(self.expr, Constant)

  def _constant_holds(self) -> bool:
    return self.expr in self.allowed

  def canonicalize(self) -> Constraint:
    # Deduplicate while preserving insertion order.
    allowed = tuple(dict.fromkeys(self.allowed))
    if len(allowed) == 1:
      return Equals(self.expr, allowed[0]).canonicalize()
    # Short-circuit to avoid reconstructing the constraint.
    if len(allowed) == len(self.allowed):
      return self
    return OneOf(self.expr, allowed)

  def __str__(self):
    return f"OneOf({self.expr}, {self.allowed})"


Constraint = (
    Equals
    | Relayout
    | NotOfType
    | IsTransferable
    | Divides
    | IsSupportedBroadcast
    | MinorDimDivisibleBy
    | AlwaysTrue
    | OneOf
)

if TYPE_CHECKING:
  _: _BaseConstraint = cast(Constraint, None)


def reduce_constraint(
    constraint: Constraint, assignments: dict[Variable, Constant]
) -> Constraint | Unsatisfiable:
  """Reduces a constraint."""

  match constraint:
    case Equals(lhs=lhs, rhs=rhs):
      lhs_red = reduce_expression(lhs, assignments)
      if isinstance(lhs_red, Unsatisfiable):
        return Unsatisfiable()
      rhs_red = reduce_expression(rhs, assignments)
      if isinstance(rhs_red, Unsatisfiable):
        return Unsatisfiable()
      # Short-circuit to avoid reconstructing the constraint.
      if lhs_red is lhs and rhs_red is rhs:
        return constraint.canonicalize()
      return Equals(lhs_red, rhs_red).canonicalize()
    case Relayout(source=source, target=target) as relayout:
      source_red = reduce_expression(source, assignments)
      target_red = reduce_expression(target, assignments)
      if isinstance(source_red, Unsatisfiable) or isinstance(
          target_red, Unsatisfiable
      ):
        return Unsatisfiable()
      # Short-circuit to avoid reconstructing the constraint.
      if source_red is source and target_red is target:
        return relayout.canonicalize()
      reduced = dataclasses.replace(
          relayout, source=source_red, target=target_red
      )
      return reduced.canonicalize()
    case NotOfType(expr=expr, type=ty):
      expr_red = reduce_expression(expr, assignments)
      if isinstance(expr_red, Unsatisfiable):
        return Unsatisfiable()
      # Short-circuit to avoid reconstructing the constraint.
      if expr_red is expr:
        return constraint
      return NotOfType(expr_red, ty)
    case IsTransferable(source=source, target=target) as transfer:
      source_red = reduce_expression(source, assignments)
      target_red = reduce_expression(target, assignments)
      if isinstance(source_red, Unsatisfiable) or isinstance(target_red, Unsatisfiable):
        return Unsatisfiable()
      # Short-circuit to avoid reconstructing the constraint.
      if source_red is source and target_red is target:
        return transfer
      return dataclasses.replace(transfer, source=source_red, target=target_red)
    case Divides(expr=expr, tiling_multiple=tiling_multiple):
      expr_red = reduce_expression(expr, assignments)
      if isinstance(expr_red, Unsatisfiable):
        return Unsatisfiable()
      # Short-circuit to avoid reconstructing the constraint.
      if expr_red is expr:
        return constraint
      return Divides(expr_red, tiling_multiple)
    case MinorDimDivisibleBy(expr=expr, divisor=divisor):
      expr_red = reduce_expression(expr, assignments)
      if isinstance(expr_red, Unsatisfiable):
        return Unsatisfiable()
      # Short-circuit to avoid reconstructing the constraint.
      if expr_red is expr:
        return constraint
      return MinorDimDivisibleBy(expr_red, divisor)
    case IsSupportedBroadcast(src=src, dst=dst, dims=dims):
      src_red = reduce_expression(src, assignments)
      dst_red = reduce_expression(dst, assignments)
      if isinstance(src_red, Unsatisfiable) or isinstance(
          dst_red, Unsatisfiable
      ):
        return Unsatisfiable()
      # Short-circuit to avoid reconstructing the constraint.
      if src_red is src and dst_red is dst:
        return constraint
      return IsSupportedBroadcast(src_red, dst_red, dims)
    case AlwaysTrue():
      return constraint
    case OneOf(expr=expr) as oneof:
      expr_red = reduce_expression(expr, assignments)
      if isinstance(expr_red, Unsatisfiable):
        return Unsatisfiable()
      # Short-circuit to avoid reconstructing the constraint.
      if expr_red is expr:
        return oneof.canonicalize()
      return OneOf(expr_red, oneof.allowed).canonicalize()
    case _ as never:
      assert_never(never)


@dataclasses.dataclass
class ConstraintSystem:
  """A constraint system contains a set of constraints and assignments.

  Assignments assign constant values to variables in the system (bound
  variables). Constraints describe relationships between variables that must be
  upheld, and can be used to determine assignments for unknown (free) variables.
  """

  assignments: dict[Variable, Constant] = dataclasses.field(
      default_factory=dict
  )
  constraints: Sequence[Constraint] = dataclasses.field(default_factory=list)

  def unknowns(self) -> list[Variable]:
    """Returns the list of free variables in the system."""
    seen_variables: set[Variable] = set()
    free_variables: list[Variable] = []
    def extract_variables(expr: Expression) -> None:
      match expr:
        case Variable():
          if expr not in seen_variables and expr not in self.assignments:
            seen_variables.add(expr)
            free_variables.append(expr)
        case RegisterLayout() | SMEMTransforms() | TMEMLayout():
          ...
        case Reduce(expression=e):
          extract_variables(e)
        case Reshape(expression=e):
          extract_variables(e)
        case Transpose(expression=e):
          extract_variables(e)
        case CollapseShape(expression=e):
          extract_variables(e)
        case _ as never:
          assert_never(never)
    for constraint in self.constraints:
      match constraint:
        case Equals(lhs=lhs, rhs=rhs):
          extract_variables(lhs)
          extract_variables(rhs)
        case Relayout(source=source, target=target):
          extract_variables(source)
          extract_variables(target)
        case NotOfType(expr=expr):
          extract_variables(expr)
        case IsTransferable(source=source, target=target):
          extract_variables(source)
          extract_variables(target)
        case Divides(expr=expr):
          extract_variables(expr)
        case MinorDimDivisibleBy(expr=expr):
          extract_variables(expr)
        case IsSupportedBroadcast(src=src, dst=dst):
          extract_variables(src)
          extract_variables(dst)
        case AlwaysTrue():
          ...
        case OneOf(expr=expr):
          extract_variables(expr)
        case _ as never:
          assert_never(never)
    return free_variables

  def __and__(
      self, other: ConstraintSystem | Unsatisfiable
  ) -> ConstraintSystem | Unsatisfiable:
    if isinstance(other, Unsatisfiable):
      return Unsatisfiable()
    smaller, larger = (
        (self.assignments, other.assignments)
        if len(self.assignments) <= len(other.assignments)
        else (other.assignments, self.assignments)
    )
    for variable, assignment in smaller.items():
      if variable in larger and assignment != larger[variable]:
        return Unsatisfiable()
    return ConstraintSystem(
        assignments=self.assignments | other.assignments,
        constraints=[*self.constraints, *other.constraints],
    )

  def __str__(self):
    r = "ConstraintSystem\n"
    r += "  assignments:\n"
    for assignment, constant in self.assignments.items():
      r += f"    {assignment} ⟵ {constant}\n"
    r += "  constraints:\n"
    for constraint in self.constraints:
      r += f"    {constraint}\n"
    return r


@final
class Unsatisfiable:

  def __and__(self, other: ConstraintSystem | Unsatisfiable) -> Unsatisfiable:
    return self


def non_splat_variables(
    constraints: Sequence[Constraint],
) -> set[Variable]:
  """Returns a all vars distinct from a splat."""
  vs: set[Variable] = set()
  for constraint in constraints:
    match constraint:
      case NotOfType(expr=Variable() as v, type=fa.WGSplatFragLayout):
        vs.add(v)
  return vs


def _has_relayout_of_non_splat_to_splat(constraints: Sequence[Constraint]) -> bool:
  """Returns whether the constraints imply a non-splat to splat relayout.

  Such relayouts are impossible and this helps shortcut the search.

  If this function returns False, this doesn't necessarily mean that there are
  no non-splat to splat relayouts, just that this is not known yet.
  """
  non_splat = non_splat_variables(constraints)
  if not non_splat:
    return False

  def is_constant_splat(e) -> bool:
    return isinstance(e, RegisterLayout) and isinstance(
        e.value, fa.WGSplatFragLayout
    )

  for constraint in constraints:
    match constraint:
      case Relayout(source=source, target=target):
        if source in non_splat and is_constant_splat(target):
          return True
      case _:
        pass
  return False


def saturate_distinct_from_splat(
    constraint_system: ConstraintSystem,
) -> ConstraintSystem | Unsatisfiable:
  """Adds transitive NotOfType constraints for all non-splat variables.

  Given `n` variables `l0`, ... `l{n-1}`, and a set of relayouts
  `{ Relayout(l{i}, l{i+1}) : 0 <= i < n }`, if we also know that
  `l{0}` is not splat, then we can automatically deduce that none of
  `l0`, ..., `l{n-1}` are splat either.

  This helps us quickly conclude that a system is unsatisfiable in cases where
  a non-splat variable is transitively relaid out into a splat layout.
  """
  non_splat = non_splat_variables(constraint_system.constraints)
  new_constraints: list[Constraint] = []
  new_non_splat_found = bool(non_splat)

  while new_non_splat_found:
    new_non_splat_found = False
    for constraint in constraint_system.constraints:
      match constraint:
        case Relayout(source=source, target=target):
          if (
              isinstance(target, Variable)
              and source in non_splat
              and target not in non_splat
          ):
            new_non_splat_found = True
            non_splat.add(target)
            new_constraints.append(NotOfType(target, fa.WGSplatFragLayout))
        case _:
          pass
  return constraint_system & ConstraintSystem(constraints=new_constraints)


def canonicalize_strict_non_splat_relayouts_to_equals(
    constraint_system: ConstraintSystem,
) -> ConstraintSystem:
  """Replaces strict relayouts from non-splat sources with Equals constraints.

  If we have `Relayout(source, target, strict=True)`, and `source` is known to
  be non-splat, then the only valid strict relayout is the identity. Thus,
  `source` must be equal to `target`.
  """
  non_splat = non_splat_variables(constraint_system.constraints)
  new_constraints: list[Constraint] = []
  for constraint in constraint_system.constraints:
    match constraint:
      case Relayout(
          source=Variable() as source, target=target, strict=True
      ) if source in non_splat:
        new_constraints.append(Equals(source, target))
      case _:
        new_constraints.append(constraint)
  return ConstraintSystem(
      assignments=constraint_system.assignments, constraints=new_constraints
  )


def compute_transitively_equal_vars(
    system: ConstraintSystem,
) -> dict[Variable, list[Variable]]:
  """Computes all transitively equal variables in a constraint system.

  The output dictionary maps each variable that appears in constraints in the
  constraint system to all the variables it is transitively equal to.
  """
  # The equality relations between variables form a graph where variables are
  # nodes and a constraint `v1 == v2` forms an edge. All variables in a
  # connected component are transitively equal. We use a Union-Find data
  # structure with path compression to efficiently find these connected
  # components (i.e., equivalence classes).
  parent: dict[Variable, Variable] = {}
  def find(v: Variable) -> Variable:
    if v not in parent:
      parent[v] = v
    if parent[v] != v:
      parent[v] = find(parent[v])
    return parent[v]

  def union(v1: Variable, v2: Variable):
    root1 = find(v1)
    root2 = find(v2)
    if root1 != root2:
      parent[root2] = root1

  all_vars: dict[Variable, None] = {}
  for constraint in system.constraints:
    match constraint:
      case Equals(lhs=Variable() as lhs, rhs=Variable() as rhs):
        all_vars[lhs] = None
        all_vars[rhs] = None
        union(lhs, rhs)

  # Group variables by their component representative.
  components: dict[Variable, list[Variable]] = {}
  for v in all_vars:
    root = find(v)
    components.setdefault(root, []).append(v)

  equal_vars: dict[Variable, list[Variable]] = {}
  for component_vars in components.values():
    for v in component_vars:
      equal_vars[v] = [other for other in component_vars if other != v]

  return equal_vars


def saturate_divides_constraints_for_equal_vars(
    system: ConstraintSystem,
) -> ConstraintSystem:
  """Saturates Divides constraints between all transitively equal vars."""
  equal_vars = compute_transitively_equal_vars(system)
  new_constraints: list[Constraint] = []
  for constraint in system.constraints:
    new_constraints.append(constraint)
    match constraint:
      case Divides(expr=Variable() as expr, tiling_multiple=tiling_multiple):
        for equal_var in equal_vars.get(expr, []):
          new_constraints.append(Divides(equal_var, tiling_multiple))
      case _:
        pass
  new_constraints = _merge_all_divides_constraints(new_constraints)
  return dataclasses.replace(system, constraints=new_constraints)


def saturate_one_of_constraints_for_equal_vars(
    system: ConstraintSystem,
) -> ConstraintSystem | Unsatisfiable:
  """Saturates OneOf constraints between all transitively equal vars."""
  equal_vars = compute_transitively_equal_vars(system)
  new_constraints: list[Constraint] = []
  for constraint in system.constraints:
    new_constraints.append(constraint)
    match constraint:
      case OneOf(expr=Variable() as expr, allowed=allowed):
        for equal_var in equal_vars.get(expr, []):
          new_constraints.append(OneOf(equal_var, allowed))
      case _:
        pass
  merged_constraints = _merge_all_one_of_constraints(new_constraints)
  if isinstance(merged_constraints, Unsatisfiable):
    return Unsatisfiable()
  return dataclasses.replace(system, constraints=merged_constraints)


def _merge_all_one_of_constraints(
    constraints: Sequence[Constraint],
) -> list[Constraint] | Unsatisfiable:
  """Merges OneOf constraints associated to the same variable."""
  result: list[Constraint] = []
  var_to_one_of: dict[Variable, OneOf] = {}
  for constraint in constraints:
    match constraint:
      case OneOf(expr=Variable() as v) as o1:
        if (o0 := var_to_one_of.get(v)) is None:
          var_to_one_of[v] = o1
          continue
        merged = _merge_one_of_constraints(o0, o1)
        if isinstance(merged, Unsatisfiable):
          return Unsatisfiable()
        var_to_one_of[v] = merged
      case _:
        result.append(constraint)
  result.extend(var_to_one_of.values())
  return result


def _merge_one_of_constraints(o0: OneOf, o1: OneOf) -> OneOf | Unsatisfiable:
  if o0.expr != o1.expr:
    raise ValueError("OneOf constraints must apply to the same expression.")
  allowed_set = set(o1.allowed)
  # Preserve order from o0 while intersecting with o1.
  merged_allowed = tuple(c for c in o0.allowed if c in allowed_set)
  if not merged_allowed:
    return Unsatisfiable()
  return OneOf(o0.expr, merged_allowed)


def _merge_all_divides_constraints(constraints: Sequence[Constraint]) -> list[Constraint]:
  """Merges Divides constraints that can be merged."""
  result: list[Constraint] = []
  var_to_divides : dict[Variable, Divides] = {}
  for constraint in constraints:
    match constraint:
      case Divides(expr=Variable() as v) as d1:
        if (d0 := var_to_divides.get(v)) is None:
          var_to_divides[v] = d1
          continue
        var_to_divides[v] = merge_divides_constraints(d0, d1)
      case _:
        result.append(constraint)
  result.extend(var_to_divides.values())
  return result


def merge_divides_constraints(d0: Divides, d1: Divides) -> Divides:
  if d0.expr != d1.expr:
    raise ValueError("Divides constraints must apply to the same expression.")
  # If the two tuples are of different lengths, the larger tuple will be
  # truncated to the length of the smaller tuple. This preserves the semantics
  # of the Divides constraints where a tiling's rank cannot exceed the size of
  # tiling_multiple.
  min_len = min(len(d0.tiling_multiple), len(d1.tiling_multiple))
  if min_len == 0:
    return Divides(d0.expr, ())
  tiling_multiple = []
  for t0, t1 in zip(d0.tiling_multiple[-min_len:], d1.tiling_multiple[-min_len:], strict=True):
    tiling_multiple.append(math.gcd(t0, t1))
  return Divides(d0.expr, tuple(tiling_multiple))


def _is_valid_register_layout_assignment(
    shape: tuple[int, ...], layout: fa.FragmentedLayout
) -> bool:
  match layout:
    case fa.WGStridedFragLayout() as strided_layout:
      return strided_layout.shape == shape
    case fa.WGSplatFragLayout() as splat_layout:
      return splat_layout.shape == shape
    case fa.TiledLayout(tiling=tiling):
      try:
        # `tiling.tile_shape` will raise if the shape is not tileable.
        _ = tiling.tile_shape(shape)
      except ValueError:
        return False
      return True
    case _:
      assert_never(layout)


def _is_valid_smem_layout_assignment(
    var: Variable, swizzle: int | None, tiling: lc.TileTransform | None
) -> bool:
  assert var.memory_space == MemorySpace.SMEM
  if tiling and not tiling.tiling:
    raise NotImplementedError("Empty tiling unsupported")
  # Scalar in SMEM is only valid if it's unswizzled and untiled.
  if not var.shape:
    return tiling is None and swizzle is None

  ref_ty = var.key.value.type
  strides, _ = ref_ty.get_strides_and_offset()

  min_stride = np.min(strides)
  if min_stride != 1:
    raise NotImplementedError("We cannot apply swizzle to non-contiguous refs")

  bitwidth = utils.bitwidth(ref_ty.element_type)
  swizzle_elems = 8 * swizzle // bitwidth if swizzle is not None else None
  duplicated_strides = strides.count(min_stride) > 1
  if tiling is None:
    # No swizzle and no tiling means it's always a valid assignment.
    if swizzle_elems is None:
      return True
    if duplicated_strides:
      return swizzle_elems is None
    # Otherwise, swizzle exist and we're dealing with non-scalar.
    # We check the divisibility of the minor most dim.
    minor_dim_index = np.argmin(strides)
    assert var.shape
    return var.shape[minor_dim_index] % swizzle_elems == 0
  tiling_v = tiling.tiling
  try:
    # `tiling.transform_shape` will raise if the shape is not tileable.
    _ = tiling.transform_shape(var.shape)
    if duplicated_strides:
      return swizzle_elems is None
    tiled_strides = lowering.tile_strides(strides, tiling_v)
    minor_tiling_dim_index = np.argmin(tiled_strides[-len(tiling_v):])
    if swizzle_elems:
      return tiling.tiling[minor_tiling_dim_index] % swizzle_elems == 0
  except ValueError:
    return False
  return True


def _is_valid_tmem_layout_assignment(
    shape: tuple[int, ...], layout: tcgen05.TMEMLayout
) -> bool:
  try:
    # `layout.tiling.tile_shape` will raise if the shape is not tileable.
    _ = layout.tiling.tile_shape(shape)
  except ValueError:
    return False
  return True


def is_valid_assignment(var: Variable, layout: Constant) -> bool:
  match layout:
    case RegisterLayout(value=reg_layout):
      assert var.memory_space == MemorySpace.REG
      return _is_valid_register_layout_assignment(var.shape, reg_layout)
    case TMEMLayout(value=tmem_layout):
      assert var.memory_space == MemorySpace.TMEM
      return _is_valid_tmem_layout_assignment(var.shape, tmem_layout)
    case SMEMTransforms(tiling=tiling, swizzle=swizzle):
      assert var.memory_space == MemorySpace.SMEM
      return _is_valid_smem_layout_assignment(var, swizzle, tiling)
    case _:
      raise ValueError(f"Unsupported layout type: {type(layout)}")


def _reduce_system_once(
    constraint_system: ConstraintSystem,
) -> ConstraintSystem | Unsatisfiable | None:
  """Performs one reduction step over each constraint in a constraint system.

  Returns:
    - Unsatisfiable(): if the constraint system is unsatisfiable.
    - A new constraint system if any constraint was reduced.
    - None: if the constraint system is not known unsatisfiable, but hasn't been
      reduced.
  """
  assignments = dict(constraint_system.assignments)
  constraints: list[Constraint] = []
  changed = False

  def try_assign(var: Variable, cst: Constant) -> bool:
    if var in assignments and assignments[var] != cst:
      return False
    if not is_valid_assignment(var, cst):
      return False
    assignments[var] = cst
    return True

  for constraint in constraint_system.constraints:
    match reduce_constraint(constraint, assignments):
      case Unsatisfiable():
        return Unsatisfiable()
      case Equals(lhs=Variable() as var, rhs=cst) if isinstance(cst, Constant):
        if not try_assign(var, cst):
          return Unsatisfiable()
        changed = True
      case Equals(lhs=cst, rhs=Variable() as var) if isinstance(cst, Constant):
        if not try_assign(var, cst):
          return Unsatisfiable()
        changed = True
      case new_constraint:
        match new_constraint.holds():  # pyrefly: ignore[missing-attribute]
          case None:
            constraints.append(new_constraint)  # pyrefly: ignore[bad-argument-type]
            changed |= new_constraint != constraint
          case False:
            return Unsatisfiable()
          case True:
            changed = True

  new_constraints = _merge_all_divides_constraints(constraints)
  changed |= len(new_constraints) != len(constraints)
  constraints = new_constraints

  # Shortcut for a specific case of unsatisfiability. This shortcut
  # drastically reduces the size of the search space.
  if _has_relayout_of_non_splat_to_splat(constraints):
    return Unsatisfiable()

  if changed:
    return ConstraintSystem(
        assignments=assignments,
        constraints=constraints,
    )
  return None


def reduce(
    constraint_system: ConstraintSystem,
) -> ConstraintSystem | Unsatisfiable:
  """Reduces a constraint system until it can no longer be reduced.

  Returns:
    - Unsatisfiable(): if the constraint system is unsatisfiable.
    - The maximally reduced constraint system otherwise.
  """
  while True:
    match _reduce_system_once(constraint_system):
      case None:
        break
      case Unsatisfiable():
        return Unsatisfiable()
      case ConstraintSystem() as new_system:
        constraint_system = new_system
      case _ as never:
        assert_never(never)  # pyrefly: ignore[bad-argument-type]  # pyrefly#2858

  return constraint_system