Assign debug operation IDs in one pass per graph, not once per node

coreai_torch/_debug_locations.py

@@ -578,53 +578,6 @@ def _get_compile_unit_attr(
 # =============================================================================
 
 
-def _in_ir_order(
-    operations: list[Operation], graph_operation: Operation | None
-) -> list[Operation]:
-    """Sort *operations* into IR order, so ids are assigned in dataflow order.
-
-    IR order is used rather than reversing the walk. Reversal is only correct when every
-    path from a result to a producer has the same length; where an operation is reachable
-    by two paths of different lengths the backwards walk records it at the shorter one,
-    and reversing then places it after something it feeds. MLIR block order is already a
-    valid topological order -- a value is defined before it is used -- so sorting by it
-    needs no such assumption.
-
-    Args:
-        operations: Operations to order; may contain duplicates and nested operations.
-        graph_operation: The graph whose IR order to use. When ``None`` the input is
-            returned unchanged, since there is nothing to sort against.
-
-    Returns:
-        The operations, without duplicates, in IR order. Any operation not found in
-        *graph_operation* keeps its original relative position at the end.
-    """
-    if graph_operation is None:
-        return operations
-
-    position = {
-        _operation_key(op): index
-        for index, op in enumerate(_get_nested_operations(graph_operation))
-    }
-    # Stable sort, so operations absent from the graph keep their relative order rather
-    # than being reshuffled among themselves.
-    unseen = len(position)
-    ordered = sorted(
-        dict.fromkeys(operations, None),
-        key=lambda op: position.get(_operation_key(op), unseen),
-    )
-    return list(ordered)
-
-
-def _operation_key(operation: Operation) -> int:
-    """A stable identity for an MLIR operation.
-
-    ``id()`` is not usable: the bindings hand back a fresh Python wrapper on each access,
-    so two wrappers for one operation compare unequal and hash differently.
-    """
-    return hash(operation)
-
-
 def _get_nested_operations(operation: Operation) -> Iterator[Operation]:
     """Iteratively yield all nested operations from regions/blocks.
 
@@ -1101,11 +1054,9 @@ def update_output_maps(
                 continue
 
             target_debug_info = self._debug_info_map.get(target_op)
-            target_op_id = (
-                target_debug_info.operation_id.value if target_debug_info else None
-            )
-
-            if target_op_id is None:
+            # Presence of debug info, not of an ID: IDs are assigned later, in IR order, by
+            # `finalize_node_operations`.
+            if target_debug_info is None:
                 continue
 
             output_map = OutputMap.create_torch_mapping(
@@ -1280,7 +1231,10 @@ def should_process_op(op) -> bool:
             for nested_op in _get_nested_operations(op):
                 added_operations.append(nested_op)
 
-        return _in_ir_order(added_operations, self._current_graph)
+        # Deduplicated but NOT sorted into IR order: doing that here meant walking every
+        # operation converted so far, once per FX node, which made conversion quadratic in
+        # graph size. `finalize_node_operations` orders them in a single pass instead.
+        return list(dict.fromkeys(added_operations))
 
     def _assign_debug_info_to_operations(
         self: Self,
@@ -1289,21 +1243,22 @@ def _assign_debug_info_to_operations(
     ) -> None:
         """Assign debug information to a list of operations.
 
+        The operation ID is left unset: it is assigned in IR order by
+        :meth:`finalize_node_operations` once the graph body is complete.
+
         Args:
             operations: List of operations to assign debug info to
             base_debug_info: Base debug information to copy and modify
         """
         for operation in operations:
-            operation_id = self._operation_id
             op_debug_info = DebugInfo(
-                operation_id=OperationID(type="coreai", value=operation_id),
+                operation_id=None,
                 source=base_debug_info.source,
                 file_locations=base_debug_info.file_locations,
                 output_maps=[],
                 call_stack=base_debug_info.call_stack,
             )
             self._debug_info_map[operation] = op_debug_info
-            self._operation_id += 1
 
     def _process_terminating_operations(self: Self, source_operation_id: int) -> None:
         """Process terminating operations and update output maps.
@@ -1363,28 +1318,52 @@ def record_operation(self: Self, node: fx.Node):
             # Process terminating operations and update output maps
             self._process_terminating_operations(source_op_id)
 
-            for operation in added_operations:
-                debug_info = self._debug_info_map[operation]
-                context = operation.context
+        self._op_results = None
+
+    def finalize_node_operations(self: Self) -> None:
+        """Assign operation IDs and materialize locations for the current graph, in IR order.
 
-                if not self.config.include_stack_trace:
-                    # Create unknown location with metadata if locations are disabled
-                    location = self._get_unknown_location_with_operation_id(
-                        debug_info, context
-                    )
-                else:
-                    # Create location based on the specified mode
-                    scope = _get_parent_scope(operation)
-                    location = self._create_operation_location(
-                        debug_info, context, scope
-                    )
+        Done once per graph rather than once per lowered node. Resolving IR order per node
+        meant walking every operation converted so far, making conversion quadratic in graph
+        size -- and each block walk ends in a binding-level C++ exception, so the constant
+        factor is large.
 
-                set_op_location(operation, location)
+        One pass over the finished graph also makes the IDs match IR order, which per-node
+        ordering did not: a constant is inserted at the top of the block rather than
+        appended, so numbering as nodes were lowered left the IDs out of order in the IR.
 
-                # Set block argument locations using the operation's location
-                self._set_block_argument_locations(operation)
+        Call this after the graph body is complete but *before* any pass that moves
+        operations out of the graph: an operation that has been outlined elsewhere is no
+        longer reachable from this graph and would never get its location.
+        """
+        if self._current_graph is None:
+            return
 
-        self._op_results = None
+        for operation in _get_nested_operations(self._current_graph):
+            debug_info = self._debug_info_map.get(operation)
+            if debug_info is None or debug_info.operation_id is not None:
+                continue
+
+            debug_info.operation_id = OperationID(
+                type="coreai", value=self._operation_id
+            )
+            self._operation_id += 1
+
+            context = operation.context
+            if not self.config.include_stack_trace:
+                # Create unknown location with metadata if locations are disabled
+                location = self._get_unknown_location_with_operation_id(
+                    debug_info, context
+                )
+            else:
+                # Create location based on the specified mode
+                scope = _get_parent_scope(operation)
+                location = self._create_operation_location(debug_info, context, scope)
+
+            set_op_location(operation, location)
+
+            # Set block argument locations using the operation's location
+            self._set_block_argument_locations(operation)
 
     def _ensure_all_operations_have_debug_locations(
         self: Self, graph_operation: Operation

coreai_torch/converter.py

@@ -836,6 +836,10 @@ def _get_graph_op(
                 with graph_op.block:
                     self._get_operation(node)
 
+            # Operation IDs and debug locations for everything just lowered, in one
+            # IR-order pass, now that the graph body is complete.
+            self._debug_info_recorder.finalize_node_operations()
+
             # Assemble outputs with resolved names
             outputs_name_value: list[tuple[str, Value]] = [
                 (resolved_name, self._values_map[fx_name])

tests/test_debug_locations.py

@@ -11,9 +11,10 @@
 from coreai._compiler.ir import Location
 from torch.export.exported_program import ExportedProgram
 
-from coreai_torch import get_decomp_table
+from coreai_torch import _debug_locations, get_decomp_table
 from coreai_torch._debug_locations import _DebugInfoRecorder, _get_nested_operations
 from coreai_torch.converter import TorchConverter
+from coreai_torch.debugging.debug_info import get_operation_id
 
 from .debugging.test_model import HierarchicalModel
 
@@ -173,3 +174,100 @@ def test_intermediate_ops_of_a_lowering_keep_their_attribution() -> None:
             location.filename.endswith(".py") and location.line >= 1
             for location in locations
         ), locations
+
+
+class DeepChainModel(nn.Module):
+    """A chain long enough that per-node whole-graph work is visible."""
+
+    def __init__(self, depth: int) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([nn.Linear(8, 8) for _ in range(depth)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            x = torch.relu(layer(x))
+        return x
+
+
+def _convert(model: nn.Module, depth_input: torch.Tensor) -> object:
+    exported_program: ExportedProgram = torch.export.export(
+        model.eval(), (depth_input,)
+    )
+    exported_program = exported_program.run_decompositions(get_decomp_table())
+    converter: TorchConverter = TorchConverter()
+    converter.add_exported_program(
+        exported_program, entrypoint_name="f", input_names=["x"], output_names=["y"]
+    )
+    return converter.to_coreai()
+
+
+def test_operation_ids_increase_in_ir_order() -> None:
+    """Operation IDs follow IR order, with none missing and none repeated.
+
+    This is a *stronger* guarantee than before, not a preserved one. IDs used to be
+    assigned as each node was lowered, which is not the same as IR order: constants are
+    inserted at the top of the block rather than appended, so a chain of linears produced
+    ``[0, 3, 7, 10, ..., 1, 2, 4, 5, ...]`` when read in IR order. Assigning them in one
+    pass over the finished graph makes the numbering match the IR, which is what
+    per-node ordering was reaching for.
+    """
+    program = _convert(DeepChainModel(6), torch.randn(1, 8))
+
+    body_ids = [
+        get_operation_id(operation)
+        for operation in _get_nested_operations(program._mlir_module.operation)
+        if operation.name != "coreai.graph"
+    ]
+
+    assert body_ids, "expected a graph body"
+    assert None not in body_ids, "every operation gets an ID"
+    assert len(body_ids) == len(set(body_ids)), f"repeated IDs: {body_ids}"
+    assert all(later > earlier for earlier, later in zip(body_ids, body_ids[1:])), (
+        f"IDs are out of IR order: {body_ids}"
+    )
+
+
+def test_the_graph_is_not_rewalked_for_every_lowered_node(monkeypatch) -> None:  # type: ignore[no-untyped-def]
+    """The graph is walked a fixed number of times, not once per lowered node.
+
+    This is the shape of the performance fix rather than a timing assertion: walking the
+    graph per node made conversion quadratic in graph size (~70x on a 512-layer chain),
+    and each walk ends in a binding-level C++ exception, so the constant factor is large
+    too. Counting walks is stable in CI in a way that wall-clock is not.
+    """
+    walks: list[str] = []
+    original = _debug_locations._get_nested_operations
+
+    def counting_get_nested_operations(operation):  # type: ignore[no-untyped-def]
+        if operation.name == "coreai.graph":
+            walks.append(operation.name)
+        return original(operation)
+
+    monkeypatch.setattr(
+        _debug_locations, "_get_nested_operations", counting_get_nested_operations
+    )
+
+    _convert(DeepChainModel(4), torch.randn(1, 8))
+    shallow = len(walks)
+    walks.clear()
+    _convert(DeepChainModel(16), torch.randn(1, 8))
+    deep = len(walks)
+
+    assert shallow > 0, "the graph is still walked once per graph"
+    assert deep == shallow, (
+        f"graph walks scale with node count: {shallow} for 4 layers, {deep} for 16"
+    )
+
+
+def test_output_maps_survive_deferred_operation_ids() -> None:
+    """Torch-to-Core AI output maps are still recorded.
+
+    They are attached while a node is lowered, but keyed on the operation having debug
+    info -- not on it having an ID, which is only assigned later. Keying on the ID drops
+    every output map silently.
+    """
+    program = _convert(DeepChainModel(3), torch.randn(1, 8))
+
+    asm = program._mlir_module.operation.get_asm(enable_debug_info=True)
+
+    assert "output_maps" in asm, "no output maps were recorded"