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"