[Mosaic GPU] Remove TraversalOrder from traverse_op.

We only use it in one place now, and therefore no longer need this
parametrization.

PiperOrigin-RevId: 826556517

Author: bchetioui

jax/experimental/mosaic/gpu/inference_utils.py

@@ -14,12 +14,10 @@
 
 """Layout & transform inference convenience utils."""
 
-from collections.abc import Callable, Sequence
-import enum
+from collections.abc import Sequence
 from functools import partial
 from typing import cast, Union
 
-from jax._src.lib import mosaic_gpu_dialect as mgpu
 from jax._src.lib.mlir import ir
 
 from . import fragmented_array as fa
@@ -303,37 +301,3 @@ def is_mma_layout(layout: fa.FragmentedLayout) -> bool:
   return columns % 16 == 0 and (
       layout == tcgen05.fa_m64_collective_layout(columns)
   )
-
-
-class TraversalOrder(enum.Enum):
-  """Traversal orders with respect to the data flow for IR."""
-
-  FORWARD = 1
-  BACKWARDS = 2
-
-
-def traverse_op(
-    op: ir.OpView,
-    callback: Callable[[ir.OpView], None],
-    traversal_order: TraversalOrder = TraversalOrder.FORWARD,
-    do_not_recurse_into_ops: tuple[type, ...] = (mgpu.CustomPrimitiveOp,),
-):
-  """Traverses the operation and applies the callback in the given order.
-
-  If do_not_recurse_into_ops is provided, the callback will be executed on these
-  ops, but any regions they might have will not be traversed.
-  """
-  callback(op)
-  if not isinstance(op, do_not_recurse_into_ops):
-    # The block of a mosaic_gpu.custom_primitive op is already lowered so it
-    # should not be traversed.
-    for region in op.operation.regions:
-      for block in region:
-        if traversal_order == TraversalOrder.FORWARD:
-          ops_to_traverse = list(block)
-        else:
-          ops_to_traverse = reversed(list(block))  # type: ignore
-        for block_op in ops_to_traverse:
-          traverse_op(
-              block_op, callback, traversal_order, do_not_recurse_into_ops
-          )

jax/experimental/mosaic/gpu/layout_inference.py

@@ -1850,6 +1850,25 @@ 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 infer_layout(module: ir.Module):
   """Infers layouts for the given module.
 
@@ -1891,7 +1910,7 @@ def gather_equations(op: ir.Operation):
     hints.extend(op_hints)
 
   for op in module.body:
-    inference_utils.traverse_op(op, gather_equations)
+    traverse_op(op, gather_equations)
 
   if isinstance(global_equation_system, eqns.Unsatisfiable):
     raise ValueError(
@@ -1933,4 +1952,4 @@ def gather_equations(op: ir.Operation):
 
   # Sanity check: ensure that all ops have the right number of in/out layouts.
   for op in module.body:
-    inference_utils.traverse_op(op, _ensure_all_layouts_are_set)
+    traverse_op(op, _ensure_all_layouts_are_set)