[Mosaic GPU] Remove {Least,Most}ReplicatedExpression constructors.

bchetioui · Google-ML-Automation · commit 3bec82db421b · 2025-12-09T01:00:40.000-08:00
After getting rid of hints, these constructors are no longer necessary.

PiperOrigin-RevId: 842112040
diff --git a/jax/experimental/mosaic/gpu/constraints.py b/jax/experimental/mosaic/gpu/constraints.py
@@ -24,7 +24,7 @@
 from collections.abc import Sequence
 import dataclasses
 import math
-from typing import Any, Callable, assert_never, final
+from typing import Any, assert_never, final
 
 from . import fragmented_array as fa
 from . import launch_context as lc
@@ -86,22 +86,6 @@ def __str__(self):
     return f"C({self.value})"
 
 
-@dataclasses.dataclass(frozen=True)
-class LeastReplicated:
-  expressions: tuple[Expression, ...]
-
-  def __post_init__(self):
-    assert len(self.expressions) >= 1
-
-
-@dataclasses.dataclass(frozen=True)
-class MostReplicated:
-  expressions: tuple[Expression, ...]
-
-  def __post_init__(self):
-    assert len(self.expressions) >= 1
-
-
 @dataclasses.dataclass(frozen=True)
 class Reduce:
   expression: Expression
@@ -136,71 +120,13 @@ def __str__(self):
 Expression = (
     Variable
     | Constant
-    | LeastReplicated
-    | MostReplicated
     | Reduce
     | BroadcastInDim
     | Reshape
     | Transpose
 )
 
 
-def reduce_replicated_expression(
-    input_expr: LeastReplicated | MostReplicated,
-    assignments: dict[Variable, Constant],
-    reducer: Callable[[fa.FragmentedLayout, fa.FragmentedLayout], fa.FragmentedLayout | None]
-) -> Expression | Unsatisfiable:
-  assert input_expr.expressions
-
-  new_expressions: list[Expression] = []
-  # Use a set to eliminate duplicates, but preserve the order.
-  seen: set[Expression] = set()
-  for expr in input_expr.expressions:
-    reduced_expr = reduce_expression(expr, assignments)
-    if isinstance(reduced_expr, Unsatisfiable):
-      return Unsatisfiable()
-    if reduced_expr in seen:
-      continue
-    new_expressions.append(reduced_expr)
-    seen.add(reduced_expr)
-
-  if len(new_expressions) == 1:
-    return new_expressions[0]
-
-  consts = []
-  unknowns = []
-  for e in new_expressions:
-    if not isinstance(e, Constant):
-      unknowns.append(e)
-      continue
-    if not isinstance(e, RegisterLayout):
-      raise ValueError(
-          f"Reduction of non-register layout constant is not supported: {e}"
-      )
-    consts.append(e)
-
-  if consts:
-    const_red, *consts = consts
-    red = const_red
-    for cst in consts:
-      red_value = reducer(red.value, cst.value)
-      if red_value is None:
-        # The layouts are not compatible up to replication, this expression
-        # cannot be simplified.
-        return Unsatisfiable()
-      red = RegisterLayout(red_value)
-  else:
-    red = None
-
-  constructor = type(input_expr)
-  if red is not None:
-    if unknowns:
-      return constructor((red, *unknowns))
-    return red
-
-  return constructor(tuple(unknowns))
-
-
 def reduce_broadcast_expression(
     broadcast: BroadcastInDim, assignments: dict[Variable, Constant]
 ) -> Expression | Unsatisfiable:
@@ -314,14 +240,6 @@ def reduce_expression(
       return expr
     case Variable():
       return assignments.get(expr, expr)
-    case MostReplicated():
-      return reduce_replicated_expression(
-          expr, assignments, layouts_lib.join_layouts
-      )
-    case LeastReplicated():
-      return reduce_replicated_expression(
-          expr, assignments, layouts_lib.meet_layouts
-      )
     case Reduce(expression=expr, axes=axes):
       reduced_expr = reduce_expression(expr, assignments)
       match reduced_expr:
@@ -640,12 +558,6 @@ def extract_variables(expr: Expression) -> None:
             free_variables.append(expr)
         case Constant():
           ...
-        case MostReplicated(expressions=expressions):
-          for e in expressions:
-            extract_variables(e)
-        case LeastReplicated(expressions=expressions):
-          for e in expressions:
-            extract_variables(e)
         case Reduce(expression=e):
           extract_variables(e)
         case BroadcastInDim(expression=e):
diff --git a/jax/experimental/mosaic/gpu/layouts.py b/jax/experimental/mosaic/gpu/layouts.py
@@ -15,7 +15,6 @@
 """Layout utilities."""
 
 import re
-from typing import assert_never
 
 from jax._src.lib import mosaic_gpu_dialect as mgpu
 from jax._src.lib.mlir import ir
@@ -224,118 +223,6 @@ def splat_is_compatible_with_tiled(
   return all(d1 % d2 == 0 for d1, d2 in zip(s1, s2))
 
 
-def meet_layouts(
-    layout1: fa.FragmentedLayout, layout2: fa.FragmentedLayout
-) -> fa.FragmentedLayout | None:
-  """Returns the "meet" of two layouts that are compatible up to replication.
-
-  The "meet" of the two layouts is the most replicated layout that is still
-  less replicated than the arguments.
-
-  This is the dual of `join_layouts`.
-
-  Returns:
-    The "meet" of the two layouts if both layouts are compatible up to
-    replication.
-
-  Raises:
-    ValueError: if the two layouts are not compatible up to replication.
-  """
-  if layout1 == layout2:
-    return layout1
-
-  match (layout1, layout2):
-    case (fa.WGSplatFragLayout(), _):
-      if isinstance(layout2, fa.TiledLayout):
-        if splat_is_compatible_with_tiled(layout1, layout2):
-          return layout2
-      elif layout1.shape == layout2.shape:
-        return layout2
-    case (_, fa.WGSplatFragLayout()):
-      if isinstance(layout1, fa.TiledLayout):
-        if splat_is_compatible_with_tiled(layout2, layout1):
-          return layout1
-      elif layout1.shape == layout2.shape:
-        return layout1
-    case (fa.TiledLayout(), fa.TiledLayout()):
-      # TODO(bchetioui): handle `TiledLayout` replication.
-      raise NotImplementedError("TiledLayout replication not supported yet")
-
-  # Layouts are not compatible up to replication.
-  return None
-
-# NOTE: We say that two layouts are compatible up to replication if the two
-# layouts satisfy at least one of the following conditions together:
-#
-# - The two layouts are equal;
-# - One of the layouts is a `WGSplatFragLayout`, and
-#   * The other layout is a `WGStridedFragLayout` with the same shape;
-#   * The other layout is a `TiledLayout` that can be used to tile the shape
-#     embedded in the `WGSplatFragLayout`.
-#
-# If any of these conditions hold, then we are always able to substitute one
-# layout with the other without having to reorder any data in the underlying
-# array---i.e. a relayout is free.
-#
-# Note that there are other combinations of layouts for which relayout is free,
-# but we voluntarily narrowed down our definition to span a small, useful
-# subset.
-
-def join_layouts(
-    layout1: fa.FragmentedLayout, layout2: fa.FragmentedLayout
-) -> fa.FragmentedLayout | None:
-  """Returns the "join" of two layouts that are compatible up to replication.
-
-  The "join" of the two layouts is the least replicated layout that is still
-  more replicated than the arguments.
-
-  This is the dual of `meet_layouts`.
-
-  Returns:
-    The "join" of the two layouts if both layouts are compatible up to
-    replication.
-
-  Raises:
-    ValueError: if the two layouts are not compatible up to replication.
-  """
-  if layout1 == layout2:
-    return layout1
-
-  match (layout1, layout2):
-    case (fa.WGSplatFragLayout(), _):
-      if isinstance(layout2, fa.TiledLayout):
-        if splat_is_compatible_with_tiled(layout1, layout2):
-          return layout1
-      elif layout1.shape == layout2.shape:
-        return layout1
-    case (_, fa.WGSplatFragLayout()):
-      if isinstance(layout1, fa.TiledLayout):
-        if splat_is_compatible_with_tiled(layout2, layout1):
-          return layout2
-      elif layout1.shape == layout2.shape:
-        return layout2
-    case (fa.TiledLayout(), fa.TiledLayout()):
-      # TODO(bchetioui): handle `TiledLayout` replication.
-      raise NotImplementedError("TiledLayout replication not supported yet")
-
-  # Layouts are not compatible up to replication.
-  return None
-
-
-def has_any_replication(layout: fa.FragmentedLayout) -> bool:
-  match layout:
-    case fa.WGSplatFragLayout():
-      return True
-    case fa.WGStridedFragLayout():
-      return False
-    case fa.TiledLayout():
-      is_warp_replicated = any(isinstance(d, fa.Replicated) for d in layout.warp_dims)
-      is_lane_replicated = any(isinstance(d, fa.Replicated) for d in layout.lane_dims)
-      return is_warp_replicated or is_lane_replicated
-    case _ as unreachable:
-      return assert_never(unreachable)  # pytype: disable=wrong-arg-types
-
-
 _tile_transform_attr_pattern = re.compile(
     r"^#mosaic_gpu.tile<[^>]+>$"
 )
diff --git a/tests/mosaic/gpu_constraints_test.py b/tests/mosaic/gpu_constraints_test.py
@@ -127,11 +127,9 @@ def test_constraint_system_unknowns_are_all_the_variables_without_assignment(
   ):
     v0, v1, v2, v3 = V(0), V(1), V(2), V(3)
     layout = RL(mgpu.WGSplatFragLayout((1, 1)))
-    least_replicated = cs.LeastReplicated((v2, v3))
-    most_replicated = cs.MostReplicated((least_replicated,))
     system = cs.ConstraintSystem(
         assignments={v0: layout},
-        constraints=[Eq(v1, most_replicated)],
+        constraints=[Eq(v1, v2), cs.Relayout(v2, v3)],
     )
     self.assertSequenceEqual(system.unknowns(), [v1, v2, v3])
 
@@ -164,112 +162,6 @@ def test_intersection_of_compatible_systems_is_union_of_fields(self):
     self.assertSequenceEqual(system1.unknowns(), [v1])
     self.assertSequenceEqual(system_intersection.unknowns(), [v0, v1])
 
-  def test_reduce_extracts_most_replicated_expression_correctly(self):
-    v0 = V(0)
-    shape = (1, 128)
-    layout0 = RL(mgpu.WGSplatFragLayout(shape))
-    layout1 = RL(mgpu.WGStridedFragLayout(shape, vec_size=1))
-    with self.subTest("most-replicated-expression-exists"):
-      system = cs.ConstraintSystem(
-          constraints=[Eq(v0, cs.MostReplicated((layout0, layout1)))],
-      )
-      self.assertEqual(
-          cs.reduce(system),
-          cs.ConstraintSystem(assignments={v0: layout0}),
-      )
-
-    with self.subTest("most-replicated-expression-is-unique-expression"):
-      system = cs.ConstraintSystem(
-          constraints=[Eq(v0, cs.MostReplicated((layout0,)))],
-      )
-      self.assertEqual(
-          cs.reduce(system),
-          cs.ConstraintSystem(assignments={v0: layout0}),
-      )
-
-    with self.subTest("most-replicated-expression-does-not-exist"):
-      system = cs.ConstraintSystem(
-          constraints=[Eq(v0, cs.MostReplicated((layout1, v0)))],
-      )
-      self.assertEqual(cs.reduce(system), system)
-
-  def test_reduce_extracts_least_replicated_expression_correctly(self):
-    v0 = V(0)
-    shape = (1, 128)
-    layout0 = RL(mgpu.WGSplatFragLayout(shape))
-    layout1 = RL(mgpu.WGStridedFragLayout(shape, vec_size=1))
-    with self.subTest("least-replicated-expression-exists"):
-      system = cs.ConstraintSystem(
-          constraints=[Eq(v0, cs.LeastReplicated([layout0, layout1]))],
-      )
-      self.assertEqual(
-          cs.reduce(system),
-          cs.ConstraintSystem(assignments={v0: layout1}),
-      )
-
-    with self.subTest("least-replicated-expression-is-unique-expression"):
-      system = cs.ConstraintSystem(
-          constraints=[Eq(v0, cs.LeastReplicated((layout0,)))],
-      )
-      self.assertEqual(
-          cs.reduce(system),
-          cs.ConstraintSystem(assignments={v0: layout0}),
-      )
-
-    with self.subTest("least-replicated-expression-does-not-exist"):
-      system = cs.ConstraintSystem(
-          constraints=[Eq(v0, cs.LeastReplicated((layout0, v0)))],
-      )
-      self.assertEqual(cs.reduce(system), system)
-
-  def test_reduce_most_replicated_expression_reduces_compatible_layouts(self):
-    splat_layout = RL(mgpu.WGSplatFragLayout((128, 64)))
-    tiled_layout = RL(mgpu.WGMMA_LAYOUT)
-    self.assertEqual(
-        cs.reduce_expression(
-            cs.MostReplicated((splat_layout, tiled_layout)),
-            {},
-        ),
-        splat_layout,
-    )
-
-  def test_reduce_most_replicated_expression_is_unsatisfiable_for_incompatible_layouts(
-      self,
-  ):
-    splat_layout = RL(mgpu.WGSplatFragLayout((1, 2)))
-    tiled_layout = RL(mgpu.WGMMA_LAYOUT)
-    self.assertIsInstance(
-        cs.reduce_expression(
-            cs.MostReplicated((splat_layout, tiled_layout)),
-            {},
-        ),
-        cs.Unsatisfiable,
-    )
-
-  def test_reduce_least_replicated_expression_reduces_compatible_layouts(self):
-    splat_layout = RL(mgpu.WGSplatFragLayout((128, 64)))
-    tiled_layout = RL(mgpu.WGMMA_LAYOUT)
-    self.assertEqual(
-        cs.reduce_expression(
-            cs.LeastReplicated((splat_layout, tiled_layout)),
-            {},
-        ),
-        tiled_layout,
-    )
-
-  def test_reduce_least_replicated_expression_is_unsatisfiable_for_incompatible_layouts(
-      self,
-  ):
-    splat_layout = RL(mgpu.WGSplatFragLayout((1, 2)))
-    tiled_layout = RL(mgpu.WGMMA_LAYOUT)
-    self.assertIsInstance(
-        cs.reduce_expression(
-            cs.LeastReplicated((splat_layout, tiled_layout)),
-            {},
-        ),
-        cs.Unsatisfiable,
-    )
-
   @parameterized.named_parameters(
       ("reduce_to_row_layout", (1,), mgpu.WGMMA_ROW_LAYOUT),
       ("reduce_to_col_layout", (0,), mgpu.WGMMA_COL_LAYOUT),
diff --git a/tests/pallas/mosaic_gpu_test.py b/tests/pallas/mosaic_gpu_test.py
