[Mosaic GPU][NFC] Make Equation a type of Constraint.

This change renames the `Equation` class to `Equals` and includes it in the `Constraint` type alias. The `EquationSystem` now stores all equality constraints within its `constraints` list, removing the separate `equations` list.

PiperOrigin-RevId: 836228575

Author: allanrenucci

jax/experimental/mosaic/gpu/equations.py

@@ -340,6 +340,21 @@ def reduce_expression(
     case _:
       assert_never(expr)
 
+@dataclasses.dataclass(frozen=True)
+class Equals:
+  """States that `lhs` and `rhs` are equal."""
+  lhs: Expression
+  rhs: Expression
+
+  def holds(self) -> bool | None:
+    if self.lhs == self.rhs:
+      return True
+    if isinstance(self.lhs, Constant) and isinstance(self.rhs, Constant):
+      return False
+    return None
+
+  def __str__(self):
+    return f"Equals({self.lhs} == {self.rhs})"
 
 _SUPPORTED_TILED_RELAYOUTS = frozenset([
     # Transposed layouts.
@@ -548,7 +563,7 @@ def __str__(self):
     return f"{self.tiling_multiple} % {self.expr} == 0"
 
 
-Constraint = Relayout | NotOfType | IsTransferable | Divides
+Constraint = Equals | Relayout | NotOfType | IsTransferable | Divides
 
 
 def reduce_constraint(
@@ -558,6 +573,14 @@ def reduce_constraint(
 
   new_constraint: 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()
+      new_constraint = Equals(lhs_red, rhs_red)
     case Relayout(source=source, target=target):
       source_red = reduce_expression(source, assignments)
       target_red = reduce_expression(target, assignments)
@@ -591,51 +614,6 @@ def reduce_constraint(
   return Tautological() if constraint_holds else Unsatisfiable()
 
 
-@dataclasses.dataclass(frozen=True)
-class Equation:
-  lhs: Expression
-  rhs: Expression
-
-  def __str__(self):
-    return f"{self.lhs} == {self.rhs}"
-
-
-def reduce_equation(
-    eq: Equation, assignments: dict[Variable, Constant]
-) -> Solution:
-  """Reduces an equation.
-
-  Args:
-    eq: the equation to reduce.
-    assignments: a set of known variable assignments.
-
-  Returns:
-    A Solution object representing the result of the evaluation. That is:
-      - Unsatisfiable(): if the equation is unsatisfiable.
-      - Tautological(): if the equation is tautological.
-      - Satisfiable(): if the equation is satisfiable by assigning a value to
-          a variable.
-      - Unknown(): if the equation contains remaining unknown variables.
-  """
-  lhs = reduce_expression(eq.lhs, assignments)
-  rhs = reduce_expression(eq.rhs, assignments)
-  match (lhs, rhs):
-    case (Variable(), Constant()):
-      return SatisfiedBy((lhs, rhs))
-    case (Constant(), Variable()):
-      return SatisfiedBy((rhs, lhs))
-    case (Constant(), Constant()) if lhs != rhs:
-      return Unsatisfiable()
-    case _ if isinstance(lhs, Unsatisfiable) or isinstance(rhs, Unsatisfiable):
-      return Unsatisfiable()
-    case _ if lhs == rhs:
-      return Tautological()
-    case _:
-      # This is covered above. Add a check here to appease the type checker.
-      assert not isinstance(lhs, Unsatisfiable) and not isinstance(rhs, Unsatisfiable)
-      return Unknown(Equation(lhs, rhs))
-
-
 @dataclasses.dataclass
 class EquationSystem:
   """An equation system contains a set of equations and assignments.
@@ -650,7 +628,6 @@ class EquationSystem:
   assignments: dict[Variable, Constant] = dataclasses.field(
       default_factory=dict
   )
-  equations: list[Equation] = dataclasses.field(default_factory=list)
   constraints: Sequence[Constraint] = dataclasses.field(default_factory=list)
 
   def unknowns(self) -> list[Variable]:
@@ -681,11 +658,11 @@ def extract_variables(expr: Expression) -> None:
           extract_variables(e)
         case _:
           assert_never(expr)
-    for equation in self.equations:
-      extract_variables(equation.lhs)
-      extract_variables(equation.rhs)
     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)
@@ -706,7 +683,6 @@ def __and__(self, other: EquationSystem) -> EquationSystem | Unsatisfiable:
         return Unsatisfiable()
     return EquationSystem(
         assignments=self.assignments | other.assignments,
-        equations=self.equations + other.equations,
         constraints=[*self.constraints, *other.constraints],
     )
 
@@ -715,9 +691,6 @@ def __str__(self):
     r += "  assignments:\n"
     for assignment, constant in self.assignments.items():
       r += f"    {assignment} ⟵ {constant}\n"
-    r += "  equations:\n"
-    for equation in self.equations:
-      r += f"    {equation}\n"
     r += "  constraints:\n"
     for constraint in self.constraints:
       r += f"    {constraint}\n"
@@ -729,16 +702,6 @@ def __and__(self, other: EquationSystem | Unsatisfiable) -> Unsatisfiable:
     return self
 
 
-@dataclasses.dataclass(frozen=True)
-class SatisfiedBy:
-  assignment: tuple[Variable, Constant]
-
-
-@dataclasses.dataclass(frozen=True)
-class Unknown:
-  equation: Equation
-
-
 class Tautological:
   ...
 
@@ -756,14 +719,6 @@ def non_splat_variables(
   return vars
 
 
-# The result of reducing an equation---and by extension, a system of
-# equations. An equation can either be unsatisfiable (i.e. there exists no
-# assignment for which it holds), satisfied by an assignment, unknown (i.e.
-# still undetermined), or tautological (i.e. the equation is guaranteed to
-# hold for any assignment).
-Solution = Unsatisfiable | SatisfiedBy | Unknown | Tautological
-
-
 def _has_relayout_of_non_splat_to_splat(constraints: Sequence[Constraint]) -> bool:
   """Returns whether the constraints imply a non-splat to splat relayout.
 
@@ -854,11 +809,14 @@ def union(v1: Variable, v2: Variable):
       parent[root2] = root1
 
   all_vars: set[Variable] = set()
-  for eq in system.equations:
-    if isinstance(eq.lhs, Variable) and isinstance(eq.rhs, Variable):
-      all_vars.add(eq.lhs)
-      all_vars.add(eq.rhs)
-      union(eq.lhs, eq.rhs)
+  for constraint in system.constraints:
+    match constraint:
+      case Equals(lhs=Variable() as lhs, rhs=Variable() as rhs):
+        assert isinstance(lhs, Variable)  # make pytype happy
+        assert isinstance(rhs, Variable)  # make pytype happy
+        all_vars.add(lhs)
+        all_vars.add(rhs)
+        union(lhs, rhs)
 
   # Group variables by their component representative.
   components: dict[Variable, list[Variable]] = {}
@@ -934,33 +892,28 @@ def _reduce_system_once(
     - None: if the equation system is not known unsatisfiable, but hasn't been
       reduced.
   """
+  assignments = equation_system.assignments
+  constraints: list[Constraint] = []
   changed = False
-  assignments: dict[Variable, Constant] = {}
-  equations: list[Equation] = []
-  for equation in equation_system.equations:
-    match reduce_equation(equation, equation_system.assignments):
-      case Unsatisfiable():
-        return Unsatisfiable()
-      case Tautological():
-        changed = True
-      case SatisfiedBy() as result:
-        variable, expression = result.assignment
-        if variable in assignments and assignments[variable] != expression:
-          return Unsatisfiable()
-        assignments[variable] = expression
-        changed = True
-      case Unknown(equation=reduced_equation):
-        equations.append(reduced_equation)
-        changed |= reduced_equation != equation
-      case _ as never:
-        assert_never(never)
 
-  assignments |= equation_system.assignments
-  constraints: list[Constraint] = []
+  def try_assign(var: Variable, cst: Constant) -> bool:
+    if var in assignments and assignments[var] != cst:
+      return False
+    assignments[var] = cst
+    return True
+
   for constraint in equation_system.constraints:
     match reduce_constraint(constraint, assignments):
       case Unsatisfiable():
         return Unsatisfiable()
+      case Equals(lhs=Variable() as var, rhs=Constant() as cst):
+        if not try_assign(var, cst):
+          return Unsatisfiable()
+        changed = True
+      case Equals(lhs=Constant() as cst, rhs=Variable() as var):
+        if not try_assign(var, cst):
+          return Unsatisfiable()
+        changed = True
       case Tautological():
         changed = True
       case _ as new_constraint:
@@ -979,7 +932,6 @@ def _reduce_system_once(
   if changed:
     return EquationSystem(
         assignments=assignments | equation_system.assignments,
-        equations=equations,
         constraints=constraints,
     )
   return None

jax/experimental/mosaic/gpu/layout_inference.py

@@ -1052,12 +1052,12 @@ def _vector_reduction_equation_system(
   return eqns.EquationSystem(), {in_variable: [in_variable.key]}, []
 
 
-def _reduction_equation_and_hint(
+def _reduction_constraint_and_hint(
     larger: eqns.Variable,
     smaller: eqns.Variable,
     larger_shape: tuple[int, ...],
-    reduction_dims: tuple[int, ...]
-) -> tuple[eqns.Equation, Hint]:
+    reduction_dims: tuple[int, ...],
+) -> tuple[eqns.Constraint, Hint]:
   reduce_expr = eqns.Reduce(larger, reduction_dims)
   # There are always many options for broadcasting a layout, so we can only
   # derive a broadcast hint in the out_variable -> source_variable direction.
@@ -1066,7 +1066,7 @@ def _reduction_equation_and_hint(
   )
   broadcast_expr = eqns.BroadcastInDim(smaller, broadcast_dims, larger_shape)
   broadcast_hint = Hint(variable=larger, expression=broadcast_expr)
-  return eqns.Equation(lhs=smaller, rhs=reduce_expr), broadcast_hint
+  return eqns.Equals(lhs=smaller, rhs=reduce_expr), broadcast_hint
 
 
 @_add_equation_system_derivation_rule(vector.MultiDimReductionOp)
@@ -1081,15 +1081,17 @@ def _multi_dim_reduction_equation_system(
   source_variable = eqns.Variable(source)
   out_variable = eqns.Variable(out)
 
-  reduction_equation, broadcast_hint = _reduction_equation_and_hint(
-      source_variable, out_variable,
-      tuple(ir.ShapedType(op.source.type).shape), tuple(op.reduction_dims)
+  reduction_constraint, broadcast_hint = _reduction_constraint_and_hint(
+      source_variable,
+      out_variable,
+      tuple(ir.ShapedType(op.source.type).shape),
+      tuple(op.reduction_dims),
   )
   # TODO(bchetioui): in the future, we may need to add rules that prevent
   # strided layouts from being chosen---since trying to reduce a strided layout
   # may cause us to raise an Exception at the moment.
   return (
-      eqns.EquationSystem(equations=[reduction_equation]),
+      eqns.EquationSystem(constraints=[reduction_constraint]),
       {source_variable: [source], out_variable: [acc, out]},
       [broadcast_hint],
   )
@@ -1108,13 +1110,16 @@ def _broadcast_in_dim_equation_system(
       i for i in range(len(out_shape)) if i not in op.broadcast_dimensions
   )
 
-  reduction_equation, broadcast_hint = _reduction_equation_and_hint(
+  reduction_constraint, broadcast_hint = _reduction_constraint_and_hint(
       out_variable, source_variable, out_shape, reduction_dims
   )
 
   return (
-      eqns.EquationSystem(equations=[reduction_equation]),
-      {source_variable: [source_variable.key], out_variable: [out_variable.key]},
+      eqns.EquationSystem(constraints=[reduction_constraint]),
+      {
+          source_variable: [source_variable.key],
+          out_variable: [out_variable.key],
+      },
       [broadcast_hint],
   )
 
@@ -1153,9 +1158,9 @@ def _shape_cast_equation_system(
 
   return (
       eqns.EquationSystem(
-          equations=[
-              eqns.Equation(lhs=out_variable, rhs=in_to_out),
-              eqns.Equation(lhs=in_variable, rhs=out_to_in),
+          constraints=[
+              eqns.Equals(lhs=out_variable, rhs=in_to_out),
+              eqns.Equals(lhs=in_variable, rhs=out_to_in),
           ],
       ),
       {in_variable: [in_variable.key], out_variable: [out_variable.key]},
@@ -1196,7 +1201,7 @@ def _custom_primitive_equation_system(
     op: mgpu.CustomPrimitiveOp,
 ) -> tuple[eqns.EquationSystem, ValueSitesForVariable, list[Hint]]:
   assignments: dict[eqns.Variable, eqns.Constant] = {}
-  equations: list[eqns.Equation] = []
+  constraints: list[eqns.Constraint] = []
   in_layouts = iter(op.in_layouts)
   in_transforms = iter(op.in_transforms)
   variables: list[eqns.Variable] = []
@@ -1220,7 +1225,7 @@ def _custom_primitive_equation_system(
       value_site = ValueSite(op, VariableType.OPERAND, i)
       source_var = ctx.producer_ref(value_site)
       v = eqns.Variable(value_site)
-      equations.append(eqns.Equation(lhs=source_var, rhs=v))
+      constraints.append(eqns.Equals(lhs=source_var, rhs=v))
       variables.append(v)
       transforms = next(in_transforms)
       ref_ty = value_site.value.type
@@ -1236,7 +1241,7 @@ def _custom_primitive_equation_system(
           layouts_lib.from_layout_attr(next(out_layouts))
       )
   return (
-      eqns.EquationSystem(equations=equations, assignments=assignments),
+      eqns.EquationSystem(assignments, constraints),
       {v: [v.key] for v in variables},
       [],
   )
@@ -1512,11 +1517,11 @@ def _memref_transpose_op_equation_system(
     return (eqns.EquationSystem(), {source_var: [source, dest]}, [])
 
   dest_var = eqns.Variable(dest)
-  equations = [
-      eqns.Equation(source_var, eqns.Transpose(dest_var)),
-      eqns.Equation(eqns.Transpose(source_var), dest_var),
+  constraints = [
+      eqns.Equals(eqns.Transpose(source_var), dest_var),
+      eqns.Equals(source_var, eqns.Transpose(dest_var)),
   ]
-  system = eqns.EquationSystem(equations=equations)
+  system = eqns.EquationSystem(constraints=constraints)
   return system, {source_var: [source], dest_var: [dest]}, []