Fix debug info attribution and module instance numbering (#61)

Operations and module instance names were attributed to the wrong source
of origin in four ways: instance counts collided across separately
converted hierarchies, counts were never reset between them, the
new-operation walk escaped the graph and handed one node's debug info to
the entire graph body, and operations reachable only through operand
edges received no debug info at all.

Author: cymbalrush

coreai_torch/_debug_locations.py

@@ -1071,6 +1071,22 @@ def update_output_maps(
             )
             target_debug_info.output_maps.append(output_map)
 
+    def reset_module_registry(self: Self) -> None:
+        """Start instance numbering again from one.
+
+        Instance counts are meant to number the instances of a class *within one
+        module hierarchy* -- ``Block$1``, ``Block$2``, ``Block$3`` for a model with
+        three of them. A submodule converted standalone is a different hierarchy:
+        its root is ``L__self__`` of that submodule's own type, not a path within
+        the model. Counting both in one sequence consumed numbers that never
+        appear in the emitted asset, so a three-block model reported
+        ``Block$2..Block$4`` and had no ``Block$1`` at all.
+
+        Called between hierarchies rather than reaching into the registry, so what
+        the numbering promises stays stated in one place.
+        """
+        self.module_registry = _ModuleInstanceRegistry()
+
     @contextmanager
     def record_module(self: Self, module: Module):
         """Context manager for module-level debug recording.
@@ -1152,7 +1168,9 @@ def _create_location_for_node(
     def _find_new_operations(self: Self) -> list[Operation]:
         """Find operations that were added during the current operation context.
 
-        Uses op_results to find candidate operations, then checks all nested operations.
+        Uses op_results to find candidate operations, then walks operands to
+        reach the rest of the ops the lowering emitted, then checks all nested
+        operations.
 
         Returns:
             List of newly added operations that don't have debug info yet
@@ -1167,6 +1185,18 @@ def should_process_op(op) -> bool:
                 and isinstance(op, Operation)
                 and op not in seen_ops
                 and op not in self._debug_info_map
+                # Never step out of the graph. The walk below goes *up* through
+                # parents, so without this it reaches the graph op itself, and
+                # the nested-operation scan a few lines down then yields every
+                # operation in the graph -- handing the whole body the debug info
+                # of whichever node was being lowered at the time.
+                #
+                # The first node lowered is the one that pays: it collects every
+                # parameter constant materialized during graph setup. In a
+                # 100-layer model that was 701 of 712 constants, all reporting
+                # the module of the first node (`Model$1/Block$1/RMSNorm$1`), so
+                # every Linear's weight claimed to belong to the first norm.
+                and op != self._current_graph
             )
 
         operations = []
@@ -1178,6 +1208,25 @@ def should_process_op(op) -> bool:
                 operations.append(op)
                 op = op.parent
 
+        # Lowering one FX node can emit a chain of operations, of which only the
+        # last produces a returned result. The rest are reachable only through
+        # operand edges, so without this they never receive the node's debug info
+        # and fall through to _ensure_all_operations_have_debug_locations, which
+        # can give them nothing but an operation ID. aten.addmm is one such case:
+        # it lowers to a transpose feeding a batch matmul feeding an add, and
+        # only the add would keep its file, line and module hierarchy.
+        #
+        # Ops already carrying debug info belong to an earlier node and end the
+        # walk there, so attribution is never reassigned from one node to another.
+        queue = list(operations)
+        while queue:
+            for operand in queue.pop().operands:
+                producer = operand.owner  # a Block, for a block argument
+                if should_process_op(producer):
+                    seen_ops.add(producer)
+                    operations.append(producer)
+                    queue.append(producer)
+
         added_operations = []
         for op in operations:
             added_operations.append(op)

coreai_torch/_utils.py

@@ -1122,7 +1122,7 @@ class _ModuleInstanceRegistry:
     def __init__(self) -> None:
         """Initialize empty module bookkeeping state."""
         self.module_type_to_next_count: dict[str, int] = {}
-        self.module_instance_to_count: dict[str, int] = {}
+        self.module_instance_to_count: dict[tuple[str, str], int] = {}
 
     def get_instance_count(
         self,
@@ -1134,20 +1134,27 @@ def get_instance_count(
         If the instance was already seen, return its existing count. Otherwise,
         assign the next count for the given module type, store it, and return it.
 
+        Keyed on the instance name *and* its type. Keying on the name alone meant
+        one name reused for a different type inherited the other type's count:
+        when a submodule is converted standalone its root is ``L__self__`` of that
+        submodule's type, and the whole model's root is ``L__self__`` too, so the
+        model's root silently took the number assigned to the submodule.
+
         Args:
             module_instance_name: Unique module instance identifier.
             module_type: Module type name, for example "Linear".
 
         Returns:
             The stable per-type instance count for this module instance.
         """
-        existing_count = self.module_instance_to_count.get(module_instance_name)
+        key = (module_instance_name, module_type)
+        existing_count = self.module_instance_to_count.get(key)
         if existing_count is not None:
             return existing_count
 
         next_count = self.module_type_to_next_count.get(module_type, 0) + 1
         self.module_type_to_next_count[module_type] = next_count
-        self.module_instance_to_count[module_instance_name] = next_count
+        self.module_instance_to_count[key] = next_count
         return next_count
 
 

coreai_torch/converter.py

@@ -409,6 +409,13 @@ def _perform_externalization(self, context) -> None:
 
         self.exported_program = whole_program
 
+        # Each submodule above was converted as its own hierarchy, rooted at the
+        # submodule, so it consumed instance numbers from a different namespace
+        # than the whole model's -- and those numbers reach nothing in the emitted
+        # asset. Left in, a model with three blocks reported Block$2, Block$3 and
+        # Block$4, with no Block$1 anywhere.
+        self._debug_info_recorder.reset_module_registry()
+
     def _clean(self) -> None:
         """Reset all internal state dictionaries to empty.
 

tests/test_debug_locations.py

@@ -7,12 +7,16 @@
 
 import torch
 import torch.nn as nn
+from coreai._compiler._mlir_libs._coreaiIR._bindings import mlir as _mlir
 from coreai._compiler.ir import Location
 from torch.export.exported_program import ExportedProgram
 
-from coreai_torch._debug_locations import _DebugInfoRecorder
+from coreai_torch import get_decomp_table
+from coreai_torch._debug_locations import _DebugInfoRecorder, _get_nested_operations
 from coreai_torch.converter import TorchConverter
 
+from .debugging.test_model import HierarchicalModel
+
 
 class SimpleModel(nn.Module):
     """Simple test model."""
@@ -124,3 +128,48 @@ def test_debug_locations_multiple_programs() -> None:
     converter.add_exported_program(exported_program2, entrypoint_name="model_2")
     # Verification happens automatically during conversion via _verify_debuginfo_locations
     _ = converter.to_coreai()
+
+
+def test_intermediate_ops_of_a_lowering_keep_their_attribution() -> None:
+    """Test that every op a multi-op lowering emits keeps its debug info.
+
+    Lowering one FX node can emit a chain of operations, of which only the last
+    produces a returned result. aten.addmm is such a case: it becomes a transpose
+    feeding a batch matmul feeding an add. The ops that only feed another op are
+    reachable from the returned result solely through operand edges, and when
+    those were not followed they received no file, no line and no module
+    hierarchy -- just an operation ID.
+
+    Uses HierarchicalModel because its Linear layers sit at two different depths,
+    so the recovered hierarchy has to be the right one rather than merely present.
+    """
+    model: HierarchicalModel = HierarchicalModel()
+    example_input: torch.Tensor = torch.randn(2, 4)
+
+    exported_program: ExportedProgram = torch.export.export(model, (example_input,))
+    exported_program = exported_program.run_decompositions(get_decomp_table())
+
+    converter: TorchConverter = TorchConverter()
+    converter.add_exported_program(exported_program)
+    program = converter.to_coreai()
+
+    matmuls = [
+        operation
+        for operation in _get_nested_operations(program._mlir_module.operation)
+        if "batch_matmul" in operation.name
+    ]
+    assert matmuls, "expected the linear layers to lower to batch matmuls"
+
+    for operation in matmuls:
+        stack_trace = _mlir.get_stack_trace(operation.location)  # type: ignore[attr-defined]
+        assert stack_trace, f"{operation.name} has no module hierarchy"
+        # The matmul belongs to the Linear that produced its weight, not to an
+        # enclosing module and not to nothing at all.
+        assert stack_trace[-1].startswith("Linear"), stack_trace
+
+        locations = _mlir.get_file_line_col_locations(operation.location)  # type: ignore[attr-defined]
+        assert locations, f"{operation.name} has no source location"
+        assert any(
+            location.filename.endswith(".py") and location.line >= 1
+            for location in locations
+        ), locations