From c14df75a9008ee91b564fa7b26cacec58f1b720a Mon Sep 17 00:00:00 2001 From: henry Date: Thu, 10 Sep 2026 00:45:08 -0700 Subject: [PATCH] Assign debug operation IDs in one pass per graph, not once per node `_DebugInfoRecorder._get_operations_for_node` ended with `_in_ir_order(added_operations, self._current_graph)`, which built a position map by walking every operation converted so far in order to sort the handful the current FX node had just added. That runs once per node, so conversion is O(nodes x ops) -- quadratic in graph size. The constant factor is large as well: each `for op in block` ends in a pybind11 stop_iteration, i.e. a C++ throw, and on macOS every throw calls `_dyld_find_unwind_sections`, which re-parses the Mach-O headers of loaded images. A `sample` of a real export shows `_Unwind_RaiseException`, `__gxx_personality_v0` and `mach_o::Header::parse_segment_command` on top. Operation IDs and debug locations are now assigned in a single IR-order pass once the graph body is complete (`finalize_node_operations`), called from `_get_graph_op` after the node loop. Measured on a chain of Linear+relu, `to_coreai()` alone: layers before after 128 1.30s 0.07s 256 4.91s 0.14s 512 19.98s 0.28s (~70x) i.e. quadratic becomes linear. A ten-procedure production model went from >45 min of conversion, unfinished, to 48 seconds. Two details worth review attention: * `update_output_maps` now guards on the target having debug info rather than on it having an ID. IDs are assigned later, so the old guard would have silently dropped every output map. * Operation IDs change. They are now genuinely IR-ordered; previously they followed lowering order, which is not the same thing, because a constant is inserted at the head of the block rather than appended -- a chain of linears numbered `[0, 3, 7, 10, ..., 1, 2, 4, 5, ...]` when read in IR order. Nothing pinned the old values. Adds three tests: IDs increase in IR order with none missing or repeated; output maps survive; and the graph is walked a fixed number of times per graph rather than once per lowered node. The last one is a counter rather than a timing assertion, so it is stable in CI. Both of the first two fail on the parent commit. --- coreai_torch/_debug_locations.py | 127 +++++++++++++------------------ coreai_torch/converter.py | 4 + tests/test_debug_locations.py | 100 +++++++++++++++++++++++- 3 files changed, 156 insertions(+), 75 deletions(-) diff --git a/coreai_torch/_debug_locations.py b/coreai_torch/_debug_locations.py index 16fe1c9..ffcd467 100644 --- a/coreai_torch/_debug_locations.py +++ b/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 diff --git a/coreai_torch/converter.py b/coreai_torch/converter.py index 8a44f8a..51d3a2e 100644 --- a/coreai_torch/converter.py +++ b/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]) diff --git a/tests/test_debug_locations.py b/tests/test_debug_locations.py index b85ea91..de0d85b 100644 --- a/tests/test_debug_locations.py +++ b/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"