From 507b9b8fe1ddc681112397fb784ee1341b8cc480 Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Mon, 19 Oct 2026 16:22:55 -0700 Subject: [PATCH 01/11] Add CoreAIProbedInspector: read intermediates as extra model outputs CoreAIInspector relies on the runtime's intermediate logger, which does not fire for work placed on the neural engine. This copies the program, appends the requested operations' values to the graph's outputs, compiles the copy and runs it, so an intermediate comes back like any other output on any compute unit. The rewrite lives in probes.py and does not need the Python runtime. The inspector keeps compiled copies per set of probed operations. --- coreai_torch/debugging/inspector.py | 204 +++++++++++++++++++ coreai_torch/debugging/probes.py | 222 +++++++++++++++++++++ tests/debugging/test_probed_inspector.py | 241 +++++++++++++++++++++++ 3 files changed, 667 insertions(+) create mode 100644 coreai_torch/debugging/probes.py create mode 100644 tests/debugging/test_probed_inspector.py diff --git a/coreai_torch/debugging/inspector.py b/coreai_torch/debugging/inspector.py index 377f779..d87a595 100644 --- a/coreai_torch/debugging/inspector.py +++ b/coreai_torch/debugging/inspector.py @@ -13,11 +13,14 @@ from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from enum import Enum, auto +from pathlib import Path +from tempfile import TemporaryDirectory from typing import Any, Union, cast import coreai import numpy as np import torch +from coreai.authoring import AIProgram from numpy.typing import NDArray from .debug_info import ( @@ -26,6 +29,7 @@ _build_compile_id_to_coreai_map, parse_debug_infos, ) +from .probes import ProbedProgram, add_probe_outputs from .torch_utils import _TorchFXNodeValueInterpreter logger = logging.getLogger(__name__) @@ -42,6 +46,11 @@ def _uses_os_framework() -> bool: return "USE_OS_COREAI" in os.environ +def _as_numpy(value: Any) -> NDArray[Any]: + """A numpy view of a Core AI NDArray, or the array itself if it is not one.""" + return value.numpy() if isinstance(value, coreai.runtime.NDArray) else value + + async def _wait_for_async_callbacks() -> None: """ Wait for asynchronously-invoked intermediate capture callbacks to complete. @@ -1142,3 +1151,198 @@ def get_all_compile_identifiers( ) for (source_op_id, _output_idx), mk in compile_map.items() } + + +@dataclass +class _CompiledProbe: + """A probed program compiled and loaded once, with the files it was loaded from.""" + + probed: ProbedProgram + function: Any + temp_dir: TemporaryDirectory[str] + + +class CoreAIProbedInspector(Inspector): + """ + Inspector that reads intermediates by adding them to the model's outputs. + + :class:`CoreAIInspector` relies on the runtime's intermediate logger, which does + not fire for work placed on the neural engine. This inspector needs no logger: + each call copies the program, appends the requested operations' values to the + graph's outputs, compiles the copy and runs it, so a value comes back the way any + model output does on any compute unit. + + The price is a compile, and that a probed value is an output the compiler may not + fuse away, so placement of the probed copy can differ from the original's. A + compiled copy is kept per set of probed operations, so reading the same set again + with new inputs does not compile again; see ``max_cached_models``. Wrap in + :class:`CachingInspector` to avoid re-reading values. + + Operation IDs are Core AI operation IDs of the program passed in. + """ + + def __init__( + self, + program: AIProgram, + function_name: str = "main", + specialization_options: coreai.runtime.SpecializationOptions | None = None, + max_cached_models: int = 4, + ): + """ + Initialize the probed inspector. + + Args: + program: The program to inspect. Not modified. + function_name: Graph to run and probe (default: "main") + specialization_options: Options the probed copy is compiled with, for + example a preferred compute unit. Runtime defaults when None. + max_cached_models: How many compiled copies to keep, one per distinct + set of probed operations, least recently used dropped first. A + compiled copy of a large model is large, so lower this to save + memory; 0 compiles on every call. + + Raises: + ValueError: If max_cached_models is negative. + + """ + if max_cached_models < 0: + msg = f"max_cached_models must not be negative, got {max_cached_models}" + raise ValueError(msg) + self._max_cached_models = max_cached_models + self._compiled: OrderedDict[frozenset[int], _CompiledProbe] = OrderedDict() + self._program = program + self._function_name = function_name + self._specialization_options = specialization_options + self._last_outputs: dict[str, NDArray[Any]] | None = None + self._last_state: dict[str, NDArray[Any]] | None = None + + async def get_intermediates_for_ops( + self, + op_ids: list[Inspector.OpID], + inputs: Any, + state: Mapping[str, Any] | None = None, + ) -> dict[Inspector.OpID, list[NDArray[Any] | None] | None]: + """ + Run the model with *op_ids* added as outputs and return their values. + + Args: + op_ids: Core AI operation IDs to read. + inputs: Model inputs, mapping input names to NDArray or numpy arrays. + state: Starting value of each state the model reads or writes, by name. + The runtime updates state in place as the model runs, so this copies + it: every call starts from the same values and the caller's arrays + are left as they were. The values the run left behind are in + :attr:`last_state`. + + Returns: + Each requested ID mapped to its output arrays, one per result. None + means the operation cannot be probed: it has no tensor result, is not in + the function's top-level graph, or has no ID in this program. + + Raises: + TypeError: If inputs is not a dictionary. + + """ + if not isinstance(inputs, dict): + msg = "inputs must be a dictionary mapping input names to NDArray" + raise TypeError(msg) + + int_op_ids = cast(list[int], op_ids) + compiled = await self._compiled_for(int_op_ids) + probed = compiled.probed + + ndarray_inputs = { + name: value + if isinstance(value, coreai.runtime.NDArray) + else coreai.runtime.NDArray(value) + for name, value in inputs.items() + } + + ndarray_state = { + name: coreai.runtime.NDArray(np.array(_as_numpy(value))) + for name, value in (state or {}).items() + } + + try: + outputs = await compiled.function( + inputs=ndarray_inputs, state=ndarray_state + ) + arrays = {name: array.numpy() for name, array in outputs.items()} + self._last_state = { + name: np.array(array.numpy()) for name, array in ndarray_state.items() + } + finally: + if self._max_cached_models == 0: + compiled.temp_dir.cleanup() + + self._last_outputs, probed_arrays = probed.split_outputs(arrays) + + results: dict[int, list[NDArray[Any] | None] | None] = dict.fromkeys(int_op_ids) + for op_id, values in probed_arrays.items(): + results[op_id] = [np.asarray(value) for value in values] + return results # type: ignore[return-value] + + async def _compiled_for(self, op_ids: list[int]) -> _CompiledProbe: + """The compiled copy for this set of operations, compiling it if not kept.""" + key = frozenset(op_ids) + if (kept := self._compiled.get(key)) is not None: + self._compiled.move_to_end(key) + return kept + + probed = add_probe_outputs(self._program, op_ids, self._function_name) + if probed.unprobed: + logger.warning("Cannot probe operations: %s", probed.unprobed) + + temp_dir = TemporaryDirectory() + try: + asset = probed.program.save_asset(Path(temp_dir.name) / "probed.aimodel") + model = await coreai.runtime.AIModel.load( + asset.path, self._specialization_options + ) + function = model.load_function(self._function_name) + except BaseException: + temp_dir.cleanup() + raise + + compiled = _CompiledProbe(probed, function, temp_dir) + if self._max_cached_models > 0: + self._compiled[key] = compiled + while len(self._compiled) > self._max_cached_models: + self._compiled.popitem(last=False)[1].temp_dir.cleanup() + return compiled + + def clear_cache(self) -> None: + """Drop every compiled copy and delete the files they were loaded from.""" + while self._compiled: + self._compiled.popitem(last=False)[1].temp_dir.cleanup() + + @property + def cached_models(self) -> int: + """How many compiled copies are currently kept.""" + return len(self._compiled) + + @property + def last_outputs(self) -> dict[str, NDArray[Any]] | None: + """ + Get the model's own outputs from the most recent execution. + + The probe outputs are not included. + + Returns: + Dictionary mapping output names to NumPy arrays, or None if the model + hasn't been executed yet. + + """ + return self._last_outputs + + @property + def last_state(self) -> dict[str, NDArray[Any]] | None: + """ + Get the state the most recent execution left behind. + + Returns: + Dictionary mapping state names to NumPy arrays, or None if the model + hasn't been executed yet. + + """ + return self._last_state diff --git a/coreai_torch/debugging/probes.py b/coreai_torch/debugging/probes.py new file mode 100644 index 0000000..955c6fc --- /dev/null +++ b/coreai_torch/debugging/probes.py @@ -0,0 +1,222 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Read intermediate values by making them outputs of the program.""" + +from __future__ import annotations + +import copy +import json +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from coreai._compiler.dialects import coreai as coreai_dialect +from coreai._compiler.ir import ( + ArrayAttr, + DictAttr, + FunctionType, + InsertionPoint, + RankedTensorType, + StringAttr, + TypeAttr, +) +from coreai.authoring import AIProgram + +from .debug_info import _build_coreai_op_map + +_PROBE_NAME_PREFIX = "__probe" + + +@dataclass(frozen=True) +class ProbedProgram: + """ + A copy of a program with extra outputs that carry intermediate values. + + Nothing here needs the Python runtime. The rewritten program can be saved and run by + any runtime, and :meth:`manifest` says which of its outputs are probes, so a caller + that gets the arrays back some other way can still tell them apart. + + Attributes: + program: The rewritten copy. The original is not modified. + function_name: The graph the outputs were added to. + outputs: Output names, one per result, by the operation ID probed. + unprobed: Requested operation IDs that could not be probed. + + """ + + program: AIProgram + function_name: str + outputs: dict[int, list[str]] + unprobed: list[int] + + @property + def probe_names(self) -> set[str]: + """Every output name that carries a probe rather than a model output.""" + return {name for names in self.outputs.values() for name in names} + + def split_outputs( + self, + arrays: Mapping[str, Any], + ) -> tuple[dict[str, Any], dict[int, list[Any]]]: + """ + Separate a run's outputs into the model's own and the probed values. + + Args: + arrays: Every output of the run, by output name. + + Returns: + The model's own outputs by name, and the probed arrays by operation ID, one + per result. + + """ + probes = self.probe_names + model_outputs = {k: v for k, v in arrays.items() if k not in probes} + probed = { + op_id: [arrays[name] for name in names] + for op_id, names in self.outputs.items() + } + return model_outputs, probed + + def manifest(self) -> dict[str, Any]: + """ + Describe the probes as plain values. + + Returns: + The function name, the output names by operation ID (keys are strings, as + JSON requires), and the operation IDs that could not be probed. + + """ + return { + "function": self.function_name, + "probes": {str(i): names for i, names in self.outputs.items()}, + "unprobed": list(self.unprobed), + } + + def save(self, directory: Path) -> Path: + """ + Save the rewritten program and its manifest for another runtime. + + Args: + directory: Where to write ``probed.aimodel`` and ``probes.json``. + + Returns: + The path of the saved asset. + + """ + directory.mkdir(parents=True, exist_ok=True) + asset = self.program.save_asset(directory / "probed.aimodel") + (directory / "probes.json").write_text(json.dumps(self.manifest(), indent=2)) + return Path(asset.path) + + +def _probe_name(op_id: int, result_idx: int) -> str: + return f"{_PROBE_NAME_PREFIX}_{op_id}_{result_idx}" + + +def add_probe_outputs( + program: AIProgram, + op_ids: Iterable[int], + function_name: str = "main", +) -> ProbedProgram: + """ + Copy *program* and return the values of the given operations as extra outputs. + + This reads an intermediate without the runtime's intermediate logger, which does + not fire on the neural engine: the value leaves the graph like any other output. + The extra outputs are appended after the existing ones, so the original outputs + keep their names and positions. + + Only tensor results of operations in the entry block of *function_name* can be + probed. An operation elsewhere (a nested graph, say), one with no results, or one + with no operation ID is reported in :attr:`ProbedProgram.unprobed`. + + Probing changes what the compiler sees: a value that is an output cannot be + fused away, so the compiled program can place work differently from the + original. Compare placement before trusting a probed value as a stand-in. + + Args: + program: The program to read values from. Not modified. + op_ids: Core AI operation IDs to read. + function_name: The graph to add outputs to. + + Returns: + The rewritten copy, the output names by operation ID, and the IDs that could + not be probed. + + Raises: + ValueError: If *program* has no graph named *function_name* or the graph has + no ``coreai.output``. + + """ + requested = list(dict.fromkeys(op_ids)) + probed = copy.deepcopy(program) + graph = probed._get_graph(function_name) # noqa: SLF001 + if graph is None: + msg = f"program has no graph named {function_name!r}" + raise ValueError(msg) + + block = graph.entry_block + terminator = next( + (op for op in block.operations if op.name == "coreai.output"), None + ) + if terminator is None: + msg = f"graph {function_name!r} has no coreai.output" + raise ValueError(msg) + + all_ops = _build_coreai_op_map(probed) + + values = [] + names: list[str] = [] + outputs: dict[int, list[str]] = {} + unprobed: list[int] = [] + for op_id in requested: + op = all_ops.get(op_id) + results = ( + [r for r in op.results if isinstance(r.type, RankedTensorType)] + if op is not None and op.operation.parent == graph.operation + else [] + ) + if not results: + unprobed.append(op_id) + continue + outputs[op_id] = [] + for result_idx, result in enumerate(results): + name = _probe_name(op_id, result_idx) + values.append(result) + names.append(name) + outputs[op_id].append(name) + + if values: + with probed._module._mlir_module.context, terminator.location: # noqa: SLF001 + with InsertionPoint(terminator): + coreai_dialect.OutputOp([*terminator.operands, *values]) + terminator.erase() + + signature = graph.function_type.value + attributes = graph.operation.attributes + attributes["function_type"] = TypeAttr.get( + FunctionType.get( + list(signature.inputs), + [*signature.results, *(v.type for v in values)], + ) + ) + attributes["res_attrs"] = ArrayAttr.get( + [ + *attributes["res_attrs"], + *( + DictAttr.get({"coreai.name": StringAttr.get(name)}) + for name in names + ), + ] + ) + + return ProbedProgram( + program=probed, + function_name=function_name, + outputs=outputs, + unprobed=unprobed, + ) diff --git a/tests/debugging/test_probed_inspector.py b/tests/debugging/test_probed_inspector.py new file mode 100644 index 0000000..5b5b4cd --- /dev/null +++ b/tests/debugging/test_probed_inspector.py @@ -0,0 +1,241 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for the inspector that reads intermediates as extra model outputs.""" + +import json +from pathlib import Path + +import numpy as np +import pytest +import torch +from coreai.authoring import AIProgram +from coreai.runtime import SpecializationOptions + +from coreai_torch.converter import TorchConverter +from coreai_torch.debugging.debug_info import _build_coreai_op_map +from coreai_torch.debugging.inspector import CoreAIProbedInspector +from coreai_torch.debugging.probes import add_probe_outputs + + +class _Net(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.conv(x)) + + +@pytest.fixture +def program_and_input() -> tuple[AIProgram, torch.Tensor]: + torch.manual_seed(0) + model = _Net().eval().half() + x = torch.randn(1, 4, 8, 8).half() + exported = torch.export.export(model, (x,)).run_decompositions() + converter = TorchConverter() + converter.add_exported_program(exported, entrypoint_name="main") + return converter.to_coreai(), x + + +def _ids_named(program: AIProgram, name: str) -> list[int]: + return [i for i, op in _build_coreai_op_map(program).items() if op.name == name] + + +def test_probing_leaves_the_original_program_and_outputs_alone( + program_and_input: tuple[AIProgram, torch.Tensor], +) -> None: + program, _ = program_and_input + before = str(program) + (relu_id,) = _ids_named(program, "coreai.relu") + + probed = add_probe_outputs(program, [relu_id]) + + assert str(program) == before + assert probed.outputs == {relu_id: [f"__probe_{relu_id}_0"]} + graph = probed.program._get_graph("main") + assert [r["coreai.name"].value for r in graph.res_attrs] == [ + "relu", + f"__probe_{relu_id}_0", + ] + + +def test_ids_that_cannot_be_probed_are_reported( + program_and_input: tuple[AIProgram, torch.Tensor], +) -> None: + program, _ = program_and_input + (output_id,) = _ids_named(program, "coreai.output") + + probed = add_probe_outputs(program, [output_id, 10_000]) + + assert probed.outputs == {} + assert probed.unprobed == [output_id, 10_000] + + +async def test_probed_inspector_returns_intermediates( + program_and_input: tuple[AIProgram, torch.Tensor], + specialization_options: SpecializationOptions | None, +) -> None: + program, x = program_and_input + (conv_id,) = _ids_named(program, "coreai.conv2d") + (relu_id,) = _ids_named(program, "coreai.relu") + inspector = CoreAIProbedInspector( + program, specialization_options=specialization_options + ) + + result = await inspector.get_intermediates_for_ops( + [conv_id, relu_id, 10_000], {"x": x.numpy()} + ) + + assert result[10_000] is None + (relu,) = result[relu_id] + (conv,) = result[conv_id] + assert conv.shape == relu.shape == (1, 4, 8, 8) + np.testing.assert_array_equal(relu, inspector.last_outputs["relu"]) + assert np.all(relu >= 0) + # The value is the convolution before the bias add, so relu(conv) is not it. + assert not np.array_equal(conv, relu) + + +class _Accumulator(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("acc", torch.zeros(1, 4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = self.acc + x + self.acc.copy_(y * 2) + return torch.relu(y) + + +async def test_probed_inspector_reads_stateful_models_without_advancing_state( + specialization_options: SpecializationOptions | None, +) -> None: + x = torch.randn(1, 4) + exported = torch.export.export(_Accumulator().eval(), (x,)).run_decompositions() + converter = TorchConverter() + converter.add_exported_program(exported, entrypoint_name="main") + program = converter.to_coreai() + (read_id,) = _ids_named(program, "coreai.read_handle") + (written_id,) = _ids_named(program, "coreai.decomposable.broadcasting_mul") + inspector = CoreAIProbedInspector( + program, specialization_options=specialization_options + ) + start = np.ones((1, 4), dtype=np.float32) + caller_state = {"b_acc": start} + + result = await inspector.get_intermediates_for_ops( + [read_id, written_id], {"x": x.numpy()}, state=caller_state + ) + + ((read,),) = (result[read_id],) + ((written,),) = (result[written_id],) + np.testing.assert_array_equal(read, start) + np.testing.assert_allclose(written, (start + x.numpy()) * 2, rtol=1e-5) + # What was written is what the model left in its state, and the caller's own + # array is untouched so a second call starts from the same place. + np.testing.assert_array_equal(inspector.last_state["b_acc"], written) + np.testing.assert_array_equal(caller_state["b_acc"], np.ones((1, 4))) + assert set(inspector.last_outputs) == {"relu"} + + +def test_probed_program_describes_itself_without_a_runtime( + program_and_input: tuple[AIProgram, torch.Tensor], + tmp_path: Path, +) -> None: + program, _ = program_and_input + (relu_id,) = _ids_named(program, "coreai.relu") + probed = add_probe_outputs(program, [relu_id, 10_000]) + + # A run's outputs split by name alone, so a caller using another runtime can do it. + own, by_op = probed.split_outputs({"relu": 1, f"__probe_{relu_id}_0": 2}) + assert own == {"relu": 1} + assert by_op == {relu_id: [2]} + + manifest = probed.manifest() + assert manifest == { + "function": "main", + "probes": {str(relu_id): [f"__probe_{relu_id}_0"]}, + "unprobed": [10_000], + } + + asset = probed.save(tmp_path / "out") + assert asset.exists() + assert json.loads((tmp_path / "out" / "probes.json").read_text()) == manifest + + +class _CountLoads: + """Count compiles by wrapping AIModel.load.""" + + def __init__(self, monkeypatch: pytest.MonkeyPatch) -> None: + import coreai.runtime + + self.count = 0 + real = coreai.runtime.AIModel.load + + async def counting(*args, **kwargs): # noqa: ANN002, ANN003, ANN202 + self.count += 1 + return await real(*args, **kwargs) + + monkeypatch.setattr(coreai.runtime.AIModel, "load", counting) + + +async def test_compiled_copy_is_reused_for_the_same_operations( + program_and_input: tuple[AIProgram, torch.Tensor], + specialization_options: SpecializationOptions | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + program, x = program_and_input + (conv_id,) = _ids_named(program, "coreai.conv2d") + (relu_id,) = _ids_named(program, "coreai.relu") + loads = _CountLoads(monkeypatch) + inspector = CoreAIProbedInspector( + program, specialization_options=specialization_options + ) + + first = await inspector.get_intermediates_for_ops([relu_id], {"x": x.numpy()}) + # Same set in another order, other inputs: no compile, and the answer follows the input. + again = await inspector.get_intermediates_for_ops([relu_id], {"x": x.numpy() * 2}) + assert loads.count == 1 + assert not np.array_equal(first[relu_id][0], again[relu_id][0]) + + await inspector.get_intermediates_for_ops([conv_id, relu_id], {"x": x.numpy()}) + assert loads.count == 2 + assert inspector.cached_models == 2 + + +async def test_cache_is_bounded_and_can_be_switched_off( + program_and_input: tuple[AIProgram, torch.Tensor], + specialization_options: SpecializationOptions | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + program, x = program_and_input + (conv_id,) = _ids_named(program, "coreai.conv2d") + (relu_id,) = _ids_named(program, "coreai.relu") + loads = _CountLoads(monkeypatch) + inputs = {"x": x.numpy()} + + bounded = CoreAIProbedInspector( + program, specialization_options=specialization_options, max_cached_models=1 + ) + await bounded.get_intermediates_for_ops([conv_id], inputs) + await bounded.get_intermediates_for_ops([relu_id], inputs) # evicts conv + assert bounded.cached_models == 1 + await bounded.get_intermediates_for_ops([conv_id], inputs) # compiles again + assert loads.count == 3 + bounded.clear_cache() + assert bounded.cached_models == 0 + + loads.count = 0 + uncached = CoreAIProbedInspector( + program, specialization_options=specialization_options, max_cached_models=0 + ) + await uncached.get_intermediates_for_ops([relu_id], inputs) + await uncached.get_intermediates_for_ops([relu_id], inputs) + assert loads.count == 2 + assert uncached.cached_models == 0 + + with pytest.raises(ValueError, match="must not be negative"): + CoreAIProbedInspector(program, max_cached_models=-1) From 2ca3e34edf5559f3c87e930d2a138a82de36eee2 Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Mon, 19 Oct 2026 17:34:18 -0700 Subject: [PATCH 02/11] Add state, perturbation checks and comparator support to probed inspection Thread an optional state argument through the inspectors, validator and comparator so stateful models can be debugged. CoreAIProbedInspector now compares each probe set's outputs and state with an unprobed run, within rtol and atol, and reduces the set when they move or when the probed copy fails to compile or run. Operations that cannot be read this way are reported in perturbing_ops and failed_ops. create_comparator_for_programs can read target values with the probed inspector via target_inspector='probed'. --- coreai_torch/debugging/comparator.py | 75 +++-- coreai_torch/debugging/inspector.py | 341 +++++++++++++++++------ coreai_torch/debugging/intermediates.py | 6 + coreai_torch/debugging/probes.py | 94 ++++++- coreai_torch/debugging/utils.py | 5 + coreai_torch/debugging/validator.py | 13 +- tests/debugging/test_comparator.py | 48 ++++ tests/debugging/test_probed_inspector.py | 189 ++++++++++++- 8 files changed, 645 insertions(+), 126 deletions(-) diff --git a/coreai_torch/debugging/comparator.py b/coreai_torch/debugging/comparator.py index 6263ef1..5aeff6c 100644 --- a/coreai_torch/debugging/comparator.py +++ b/coreai_torch/debugging/comparator.py @@ -20,12 +20,12 @@ import logging from collections import Counter -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass, field from enum import Enum, auto from pathlib import Path from tempfile import TemporaryDirectory -from typing import Any, Generic, TypeVar +from typing import Any, Generic, Literal, TypeVar import numpy as np import torch @@ -43,6 +43,7 @@ from .inspector import ( CachingInspector, CoreAIInspector, + CoreAIProbedInspector, Inspector, TorchFXInspector, ) @@ -52,7 +53,7 @@ ) from .table_writer import _Column, _Row, _TableSpec, _write_table from .torch_utils import get_torch_to_coreai_output_mapping -from .utils import _plain, _with_debug +from .utils import _plain, _state_kwargs, _with_debug logger = logging.getLogger(__name__) @@ -865,6 +866,7 @@ async def _fetch_target_intermediates( self, batch: list[ComputationGraph.Node], inputs: Any, + state: Mapping[str, Any] | None = None, ) -> dict[ComputationGraph.OpID, list[NDArray[Any] | None] | None]: """Fetch intermediates for target nodes corresponding to source batch.""" target_batch_op_ids: list[ComputationGraph.OpID] = [] @@ -879,6 +881,7 @@ async def _fetch_target_intermediates( return await self.target.inspector.get_intermediates_for_ops( target_batch_op_ids, inputs, + **_state_kwargs(state), ) @staticmethod @@ -1045,6 +1048,7 @@ async def _process_batch( ], Status, ], + state: Mapping[str, Any] | None = None, ) -> BatchResult: """ Process a batch of nodes for comparison. @@ -1069,8 +1073,9 @@ async def _process_batch( source_results = await self.source.inspector.get_intermediates_for_ops( source_batch_op_ids, inputs, + **_state_kwargs(state), ) - target_results = await self._fetch_target_intermediates(batch, inputs) + target_results = await self._fetch_target_intermediates(batch, inputs, state) # Check each node pair in the batch batch_results = [] @@ -1130,6 +1135,7 @@ async def compare( Status, ], inputs: Any, + state: Mapping[str, Any] | None = None, ) -> Result: """ Compare operations between source and target graphs. @@ -1138,6 +1144,8 @@ async def compare( check_fn: Function that takes (source_node, target_node, source_outputs, target_outputs) and returns Status inputs: Model inputs to use during execution + state: Starting value of each stateful buffer, by name. Both sides start + from it, so the buffer names must match in source and target. Returns: Result containing failed and unknown comparisons sorted by source topological order. @@ -1165,6 +1173,7 @@ async def compare( batch, inputs, check_fn, + state, ) pass_count += batch_result.pass_count @@ -1362,6 +1371,7 @@ async def compare_with_tolerance( rtol: float = _DEFAULT_RTOL, atol: float = _DEFAULT_ATOL, exclude_multi_output: bool = False, + state: Mapping[str, Any] | None = None, ) -> Result: """ Compare outputs between graphs with numerical tolerance. @@ -1375,6 +1385,7 @@ async def compare_with_tolerance( up. exclude_multi_output: Skip any operation whose source has several outputs, reporting it as EXCLUDED rather than comparing. + state: Starting value of each stateful buffer, by name. Returns: Result containing operations where outputs differ beyond tolerance @@ -1452,7 +1463,7 @@ def check_fn( return Comparator.Status.PASS - return await self.compare(check_fn, inputs) + return await self.compare(check_fn, inputs, state) def _create_id_map_from_coreai_program( @@ -1516,6 +1527,8 @@ async def create_comparator_for_programs( use_caching: bool = True, exclude_ops: frozenset[str] = _DEFAULT_EXCLUDED_OPS, specialization_options: SpecializationOptions | None = None, + target_inspector: Literal["logger", "probed"] = "logger", + probed_options: Mapping[str, Any] | None = None, ) -> Comparator[torch.fx.Node, torch.fx.Graph, Operation, Module]: """ Create a comparator between PyTorch ExportedProgram and AIProgram. @@ -1538,11 +1551,29 @@ async def create_comparator_for_programs( Defaults to _DEFAULT_EXCLUDED_OPS which includes view/reshape operations. Pass frozenset() to disable exclusions. specialization_options: Options for configuring model specialization + target_inspector: How target intermediates are read. ``"logger"`` uses the + runtime's intermediate logger. ``"probed"`` adds the values to the model's + outputs (:class:`CoreAIProbedInspector`), so the values come from a copy of + the model with extra outputs; check its ``perturbations``. + probed_options: Extra keyword arguments for :class:`CoreAIProbedInspector` + (``rtol``, ``atol``, ``on_perturbation``, ``max_cached_models``). Only + valid with ``target_inspector="probed"``. Returns: Comparator instance configured for comparing the two programs + Raises: + ValueError: If target_inspector is unknown, or probed_options is given + without ``"probed"``. + """ + if target_inspector not in ("logger", "probed"): + msg = f"target_inspector must be logger or probed, got {target_inspector}" + raise ValueError(msg) + if probed_options and target_inspector != "probed": + msg = "probed_options is only valid with target_inspector='probed'" + raise ValueError(msg) + # Create ID mapping from torch to coreai operations id_map = _create_id_map_from_coreai_program( target_program, @@ -1558,22 +1589,26 @@ async def create_comparator_for_programs( # Create source graph source_graph = create_graph_from_exported_program(source_program) - # Create target (AIProgram) inspector based on inspector type - temp_dir = TemporaryDirectory() - asset_path = Path(temp_dir.name) / "model.aimodel" - - # Create asset from AIProgram and load model from asset - asset = target_program.save_asset(asset_path) - specialization_options = _with_debug(specialization_options) - model = await AIModel.load(asset.path, specialization_options) - target_inspector = CoreAIInspector( - model=model, - function_name=target_entry_point, - temp_dir=temp_dir, - ) + target: Inspector + if target_inspector == "probed": + target = CoreAIProbedInspector( + target_program, + function_name=target_entry_point, + specialization_options=specialization_options, + **(probed_options or {}), + ) + else: + temp_dir = TemporaryDirectory() + asset = target_program.save_asset(Path(temp_dir.name) / "model.aimodel") + model = await AIModel.load(asset.path, _with_debug(specialization_options)) + target = CoreAIInspector( + model=model, + function_name=target_entry_point, + temp_dir=temp_dir, + ) if use_caching: - target_inspector = CachingInspector(target_inspector) + target = CachingInspector(target) # Create target graph target_graph = create_graph_from_coreai_program( @@ -1588,7 +1623,7 @@ async def create_comparator_for_programs( ) target_debug_graph = DebugGraph( graph=target_graph, - inspector=target_inspector, + inspector=target, ) return Comparator( diff --git a/coreai_torch/debugging/inspector.py b/coreai_torch/debugging/inspector.py index d87a595..aa188f1 100644 --- a/coreai_torch/debugging/inspector.py +++ b/coreai_torch/debugging/inspector.py @@ -15,7 +15,7 @@ from enum import Enum, auto from pathlib import Path from tempfile import TemporaryDirectory -from typing import Any, Union, cast +from typing import Any, Literal, Union, cast import coreai import numpy as np @@ -29,8 +29,14 @@ _build_compile_id_to_coreai_map, parse_debug_infos, ) -from .probes import ProbedProgram, add_probe_outputs +from .probes import ( + Perturbation, + ProbedProgram, + add_probe_outputs, + compare_to_reference, +) from .torch_utils import _TorchFXNodeValueInterpreter +from .utils import _state_kwargs logger = logging.getLogger(__name__) @@ -463,6 +469,7 @@ async def get_intermediates_for_ops( self, op_ids: list[OpID], inputs: Any, + state: Mapping[str, Any] | None = None, ) -> dict[OpID, list[NDArray[Any] | None] | None]: """ Execute model and capture intermediate outputs for specified operations. @@ -470,6 +477,8 @@ async def get_intermediates_for_ops( Args: op_ids: List of operation identifiers to capture outputs for inputs: Model inputs (format varies by implementation) + state: Starting value of each stateful buffer, by name. Copied, so the + caller's arrays are never updated. Returns: Dictionary mapping each operation ID to a list of output arrays. @@ -525,6 +534,7 @@ def __init__(self, inspector: Inspector, max_cache_size: int | None = None): OrderedDict() ) self._current_inputs: Any = None + self._current_state: Any = None self._max_cache_size = max_cache_size @staticmethod @@ -595,25 +605,31 @@ async def get_intermediates_for_ops( self, op_ids: list[Inspector.OpID], inputs: Any, + state: Mapping[str, Any] | None = None, ) -> dict[Inspector.OpID, list[NDArray[Any] | None] | None]: """ Retrieve intermediate outputs with automatic caching. Returns cached values when available, only executing the model for operations - not yet in the cache. Automatically clears the cache when inputs change. - Implements LRU eviction when max_cache_size is set. + not yet in the cache. Automatically clears the cache when inputs or state + change. Implements LRU eviction when max_cache_size is set. Args: op_ids: List of operation identifiers to capture outputs for inputs: Model inputs (cache is invalidated if these change) + state: Starting state (cache is invalidated if this changes) Returns: Dictionary mapping operation IDs to output arrays (cached or freshly captured) """ - if not self._inputs_equal(self._current_inputs, inputs): + if not ( + self._inputs_equal(self._current_inputs, inputs) + and self._inputs_equal(self._current_state, state) + ): self._cache.clear() self._current_inputs = self._copy_inputs(inputs) + self._current_state = self._copy_inputs(state) uncached_ops = [op_id for op_id in op_ids if op_id not in self._cache] @@ -621,6 +637,7 @@ async def get_intermediates_for_ops( results = await self._inspector.get_intermediates_for_ops( uncached_ops, inputs, + **_state_kwargs(state), ) if results is not None: for op_id, value in results.items(): @@ -645,6 +662,7 @@ def clear_cache(self) -> None: """Clear all cached intermediate values and reset input tracking.""" self._cache.clear() self._current_inputs = None + self._current_state = None class TorchFXInspector(Inspector): @@ -672,6 +690,7 @@ async def get_intermediates_for_ops( self, op_ids: list[Inspector.OpID], inputs: Any, + state: Mapping[str, Any] | None = None, ) -> dict[Inspector.OpID, list[NDArray[Any] | None] | None]: """ Capture intermediate values at specified FX graph nodes. @@ -682,6 +701,11 @@ async def get_intermediates_for_ops( Args: op_ids: List of FX node names (operation IDs) to capture inputs: Tuple of input tensors matching the model's expected signature + state: Starting value of each buffer, by its placeholder name (for + example ``"b_cache"``). Buffers not named keep the exported value. + + Raises: + KeyError: If a state name is not a buffer of the exported program. Returns: Dictionary mapping node names to lists of output arrays (as NumPy). @@ -701,10 +725,16 @@ def capture_callback(node: torch.fx.Node, result: Any) -> None: else: results[node.name] = [self.__class__.convert_to_numpy(result)] - interpreter = _TorchFXNodeValueInterpreter( - self.exported_program.module(), - callback=capture_callback, - ) + module = self.exported_program.module() + buffers = self.exported_program.graph_signature.inputs_to_buffers + for name, value in (state or {}).items(): + if name not in buffers: + msg = f"{name!r} is not a buffer; buffers: {sorted(buffers)}" + raise KeyError(msg) + buffer = module.get_buffer(buffers[name]) + buffer.copy_(torch.as_tensor(np.array(_as_numpy(value))).to(buffer.dtype)) + + interpreter = _TorchFXNodeValueInterpreter(module, callback=capture_callback) # Accept the same shape of inputs as CoreAIInspector, which requires a # name -> array dict. Without this the two inspectors in a Comparator need @@ -791,6 +821,7 @@ def __init__( self._model = model self._function_name = function_name self._last_outputs: dict[str, NDArray[Any]] | None = None + self._last_state: dict[str, NDArray[Any]] | None = None self._temp_dir = temp_dir # Keep temp directory alive # Load debug info from model @@ -984,6 +1015,7 @@ async def get_intermediates_for_ops( self, op_ids: list[Inspector.OpID], inputs: Any, + state: Mapping[str, Any] | None = None, ) -> dict[Inspector.OpID, list[NDArray[Any] | None] | None]: """ Capture intermediate outputs from Core AI Runtime model execution. @@ -994,6 +1026,8 @@ async def get_intermediates_for_ops( Args: op_ids: List of source operation debug IDs to capture inputs: Model inputs (dictionary mapping input names to NDArray or numpy arrays) + state: Starting value of each state, by name. Copied, because the runtime + updates state in place. Returns: Dictionary mapping operation IDs to lists of output arrays. @@ -1054,10 +1088,17 @@ async def get_intermediates_for_ops( intermediate_logger=intermediate_logger, ) - outputs = await inference_function(inputs=ndarray_inputs) + ndarray_state = { + name: coreai.runtime.NDArray(np.array(_as_numpy(value))) + for name, value in (state or {}).items() + } + outputs = await inference_function(inputs=ndarray_inputs, state=ndarray_state) await _wait_for_async_callbacks() self._last_outputs = {name: array.numpy() for name, array in outputs.items()} + self._last_state = { + name: np.array(array.numpy()) for name, array in ndarray_state.items() + } logger.debug("Successfully captured %d source operations", len(results)) @@ -1090,6 +1131,11 @@ def last_outputs(self) -> dict[str, NDArray[Any]] | None: """ return self._last_outputs + @property + def last_state(self) -> dict[str, NDArray[Any]] | None: + """State left by the most recent execution, or None if never executed.""" + return self._last_state + @staticmethod def get_compile_identifiers_for_op( source_level: str, @@ -1153,6 +1199,10 @@ def get_all_compile_identifiers( } +class PerturbationError(RuntimeError): + """Probing moved the model's own outputs beyond the tolerance.""" + + @dataclass class _CompiledProbe: """A probed program compiled and loaded once, with the files it was loaded from.""" @@ -1166,19 +1216,22 @@ class CoreAIProbedInspector(Inspector): """ Inspector that reads intermediates by adding them to the model's outputs. - :class:`CoreAIInspector` relies on the runtime's intermediate logger, which does - not fire for work placed on the neural engine. This inspector needs no logger: - each call copies the program, appends the requested operations' values to the - graph's outputs, compiles the copy and runs it, so a value comes back the way any - model output does on any compute unit. + Needs no intermediate logger: it copies the program, appends the requested + operations' values to the graph's outputs, compiles the copy and runs it, so a + value comes back like any model output on any compute unit. + + A probed value is an output the compiler may not fuse away, so the copy can be + placed or fused differently from the original. Each probe set is checked: the + copy's outputs and state are compared with an unprobed run on the same inputs, + within ``rtol`` and ``atol``. A set that moves them, or fails to compile or run, + is reduced (see ``on_perturbation``) so every value returned comes from a copy + that left the model's outputs in place. - The price is a compile, and that a probed value is an output the compiler may not - fuse away, so placement of the probed copy can differ from the original's. A - compiled copy is kept per set of probed operations, so reading the same set again - with new inputs does not compile again; see ``max_cached_models``. Wrap in - :class:`CachingInspector` to avoid re-reading values. + Passing the check means no evidence the probes changed the model, not that a + probed value equals the unprobed model's internal value. - Operation IDs are Core AI operation IDs of the program passed in. + Compiled copies are kept per probe set; see ``max_cached_models``. Operation IDs + are Core AI operation IDs of the program passed in. """ def __init__( @@ -1187,34 +1240,58 @@ def __init__( function_name: str = "main", specialization_options: coreai.runtime.SpecializationOptions | None = None, max_cached_models: int = 4, + rtol: float = 1e-3, + atol: float = 1e-3, + on_perturbation: Literal["split", "warn", "raise"] = "split", ): """ Initialize the probed inspector. Args: program: The program to inspect. Not modified. - function_name: Graph to run and probe (default: "main") - specialization_options: Options the probed copy is compiled with, for - example a preferred compute unit. Runtime defaults when None. - max_cached_models: How many compiled copies to keep, one per distinct - set of probed operations, least recently used dropped first. A - compiled copy of a large model is large, so lower this to save - memory; 0 compiles on every call. + function_name: Graph to run and probe. + specialization_options: Options the copies are compiled with. Runtime + defaults when None. + max_cached_models: Probed copies to keep, least recently used dropped + first. 0 compiles on every call. The unprobed reference is kept + separately. + rtol: Relative tolerance of the perturbation check. + atol: Absolute tolerance of the perturbation check. Neither default is + tuned to a model; set them for the precision you run. + on_perturbation: For a set that moves the outputs. ``"split"`` retries + each half until every returned value comes from a clean set; + ``"warn"`` logs and returns the values; ``"raise"`` raises + :class:`PerturbationError`. A set that fails to compile or run is + always split, or raises under ``"raise"``. Raises: - ValueError: If max_cached_models is negative. + ValueError: If max_cached_models, rtol or atol is negative, or + on_perturbation is not one of the three. """ if max_cached_models < 0: msg = f"max_cached_models must not be negative, got {max_cached_models}" raise ValueError(msg) + if rtol < 0 or atol < 0: + msg = f"rtol and atol must not be negative, got {rtol} and {atol}" + raise ValueError(msg) + if on_perturbation not in ("split", "warn", "raise"): + msg = f"on_perturbation must be split, warn or raise, got {on_perturbation}" + raise ValueError(msg) self._max_cached_models = max_cached_models + self._rtol = rtol + self._atol = atol + self._on_perturbation = on_perturbation self._compiled: OrderedDict[frozenset[int], _CompiledProbe] = OrderedDict() + self._reference: _CompiledProbe | None = None self._program = program self._function_name = function_name self._specialization_options = specialization_options self._last_outputs: dict[str, NDArray[Any]] | None = None self._last_state: dict[str, NDArray[Any]] | None = None + self._perturbations: list[Perturbation] = [] + self._perturbing_ops: set[int] = set() + self._failed_ops: set[int] = set() async def get_intermediates_for_ops( self, @@ -1228,67 +1305,150 @@ async def get_intermediates_for_ops( Args: op_ids: Core AI operation IDs to read. inputs: Model inputs, mapping input names to NDArray or numpy arrays. - state: Starting value of each state the model reads or writes, by name. - The runtime updates state in place as the model runs, so this copies - it: every call starts from the same values and the caller's arrays - are left as they were. The values the run left behind are in - :attr:`last_state`. + state: Starting value of each state, by name. Copied, so every run starts + from the same values and the caller's arrays are left alone. Returns: - Each requested ID mapped to its output arrays, one per result. None - means the operation cannot be probed: it has no tensor result, is not in - the function's top-level graph, or has no ID in this program. + Each requested ID mapped to its output arrays, one per result. None means + the operation cannot be probed (no tensor result, not in the function's + top-level graph, or no ID in this program), or it is in + :attr:`perturbing_ops` or :attr:`failed_ops`. Raises: TypeError: If inputs is not a dictionary. + PerturbationError: If a set moved the outputs, or failed to compile or + run, and ``on_perturbation="raise"``. """ if not isinstance(inputs, dict): msg = "inputs must be a dictionary mapping input names to NDArray" raise TypeError(msg) - int_op_ids = cast(list[int], op_ids) - compiled = await self._compiled_for(int_op_ids) - probed = compiled.probed - + int_op_ids = list(dict.fromkeys(cast(list[int], op_ids))) ndarray_inputs = { name: value if isinstance(value, coreai.runtime.NDArray) else coreai.runtime.NDArray(value) for name, value in inputs.items() } - - ndarray_state = { - name: coreai.runtime.NDArray(np.array(_as_numpy(value))) - for name, value in (state or {}).items() + start_state = { + name: np.array(_as_numpy(value)) for name, value in (state or {}).items() } + self._perturbations = [] + self._perturbing_ops = set() + self._failed_ops = set() + reference_compiled = await self._reference_copy() try: - outputs = await compiled.function( - inputs=ndarray_inputs, state=ndarray_state + reference_arrays, reference_state = await self._run( + reference_compiled, ndarray_inputs, start_state ) - arrays = {name: array.numpy() for name, array in outputs.items()} - self._last_state = { - name: np.array(array.numpy()) for name, array in ndarray_state.items() - } finally: if self._max_cached_models == 0: - compiled.temp_dir.cleanup() - - self._last_outputs, probed_arrays = probed.split_outputs(arrays) + self._release_reference() + self._last_outputs = reference_arrays + self._last_state = reference_state + reference = {**reference_arrays, **reference_state} results: dict[int, list[NDArray[Any] | None] | None] = dict.fromkeys(int_op_ids) - for op_id, values in probed_arrays.items(): - results[op_id] = [np.asarray(value) for value in values] + pending = [int_op_ids] + while pending: + batch = pending.pop(0) + try: + compiled = await self._compiled_for(batch) + try: + arrays, left_state = await self._run( + compiled, ndarray_inputs, start_state + ) + finally: + if self._max_cached_models == 0: + compiled.temp_dir.cleanup() + except Exception as error: + if self._on_perturbation == "raise": + msg = f"probing {batch} failed to compile or run: {error}" + raise PerturbationError(msg) from error + self._split_or_give_up( + batch, + pending, + self._failed_ops, + f"fails to compile or run: {error}", + ) + continue + + own, probed_arrays = compiled.probed.split_outputs(arrays) + probeable = list(compiled.probed.outputs) + report = compare_to_reference( + reference, {**own, **left_state}, probeable, self._rtol, self._atol + ) + if probeable: + self._perturbations.append(report) + + if report.within_tolerance or self._on_perturbation == "warn": + if not report.within_tolerance: + logger.warning( + "Probing %s moved the model's outputs (max abs %.3g, " + "max rel %.3g)", + probeable, + report.max_abs_diff, + report.max_rel_diff, + ) + for op_id, values in probed_arrays.items(): + results[op_id] = [np.asarray(value) for value in values] + elif self._on_perturbation == "raise": + msg = ( + f"probing {probeable} moved the model's outputs " + f"(max abs {report.max_abs_diff:.3g}, " + f"max rel {report.max_rel_diff:.3g})" + ) + raise PerturbationError(msg) + else: + self._split_or_give_up( + probeable, + pending, + self._perturbing_ops, + "moves the model's outputs", + ) return results # type: ignore[return-value] - async def _compiled_for(self, op_ids: list[int]) -> _CompiledProbe: - """The compiled copy for this set of operations, compiling it if not kept.""" - key = frozenset(op_ids) - if (kept := self._compiled.get(key)) is not None: - self._compiled.move_to_end(key) - return kept + @staticmethod + def _split_or_give_up( + batch: list[int], + pending: list[list[int]], + culprits: set[int], + why: str, + ) -> None: + """Queue the halves of a failing set, or record a failing single op.""" + if len(batch) == 1: + culprits.add(batch[0]) + logger.warning( + "Probing operation %d alone %s; no value returned", *batch, why + ) + return + middle = len(batch) // 2 + pending.extend([batch[:middle], batch[middle:]]) + + async def _run( + self, + compiled: _CompiledProbe, + inputs: Mapping[str, Any], + state: Mapping[str, Any], + ) -> tuple[dict[str, NDArray[Any]], dict[str, NDArray[Any]]]: + """Run a compiled copy on fresh copies of the state. + + Returns: + Every output by name, and the state the run left behind. + """ + ndarray_state = { + name: coreai.runtime.NDArray(np.array(value)) + for name, value in state.items() + } + outputs = await compiled.function(inputs=dict(inputs), state=ndarray_state) + arrays = {name: np.array(array.numpy()) for name, array in outputs.items()} + left = {name: np.array(array.numpy()) for name, array in ndarray_state.items()} + return arrays, left + async def _compile(self, op_ids: list[int]) -> _CompiledProbe: + """Rewrite, save, compile and load a copy that probes *op_ids*.""" probed = add_probe_outputs(self._program, op_ids, self._function_name) if probed.unprobed: logger.warning("Cannot probe operations: %s", probed.unprobed) @@ -1303,8 +1463,27 @@ async def _compiled_for(self, op_ids: list[int]) -> _CompiledProbe: except BaseException: temp_dir.cleanup() raise + return _CompiledProbe(probed, function, temp_dir) + + async def _reference_copy(self) -> _CompiledProbe: + """The program compiled with nothing probed.""" + if self._reference is None: + self._reference = await self._compile([]) + return self._reference + + def _release_reference(self) -> None: + if self._reference is not None: + self._reference.temp_dir.cleanup() + self._reference = None - compiled = _CompiledProbe(probed, function, temp_dir) + async def _compiled_for(self, op_ids: list[int]) -> _CompiledProbe: + """The compiled copy for this set of operations, compiling it if not kept.""" + key = frozenset(op_ids) + if (kept := self._compiled.get(key)) is not None: + self._compiled.move_to_end(key) + return kept + + compiled = await self._compile(op_ids) if self._max_cached_models > 0: self._compiled[key] = compiled while len(self._compiled) > self._max_cached_models: @@ -1312,37 +1491,37 @@ async def _compiled_for(self, op_ids: list[int]) -> _CompiledProbe: return compiled def clear_cache(self) -> None: - """Drop every compiled copy and delete the files they were loaded from.""" + """Drop every compiled copy, the reference included.""" while self._compiled: self._compiled.popitem(last=False)[1].temp_dir.cleanup() + self._release_reference() @property def cached_models(self) -> int: - """How many compiled copies are currently kept.""" + """Probed copies currently kept, the reference not counted.""" return len(self._compiled) @property - def last_outputs(self) -> dict[str, NDArray[Any]] | None: - """ - Get the model's own outputs from the most recent execution. + def perturbations(self) -> list[Perturbation]: + """The check result for each probe set the latest call ran, in order.""" + return list(self._perturbations) - The probe outputs are not included. + @property + def perturbing_ops(self) -> set[int]: + """Operations that moved the model's outputs even when probed alone.""" + return set(self._perturbing_ops) - Returns: - Dictionary mapping output names to NumPy arrays, or None if the model - hasn't been executed yet. + @property + def failed_ops(self) -> set[int]: + """Operations whose copy failed to compile or run even when probed alone.""" + return set(self._failed_ops) - """ + @property + def last_outputs(self) -> dict[str, NDArray[Any]] | None: + """The model's outputs from the latest unprobed run, or None before one.""" return self._last_outputs @property def last_state(self) -> dict[str, NDArray[Any]] | None: - """ - Get the state the most recent execution left behind. - - Returns: - Dictionary mapping state names to NumPy arrays, or None if the model - hasn't been executed yet. - - """ + """The state the latest unprobed run left behind, or None before one.""" return self._last_state diff --git a/coreai_torch/debugging/intermediates.py b/coreai_torch/debugging/intermediates.py index 768b639..c4c7086 100644 --- a/coreai_torch/debugging/intermediates.py +++ b/coreai_torch/debugging/intermediates.py @@ -20,6 +20,7 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass, field from pathlib import Path from tempfile import TemporaryDirectory @@ -436,6 +437,7 @@ async def compare_program_intermediates( rtol: float = _DEFAULT_RTOL, atol: float = _DEFAULT_ATOL, exclude_multi_output: bool = True, + state: Mapping[str, Any] | None = None, ) -> IntermediatesReport: """ Run both programs on the same inputs and compare every mapped operation. @@ -464,6 +466,8 @@ async def compare_program_intermediates( rtol: Relative tolerance. atol: Absolute tolerance. exclude_multi_output: As :func:`compare_intermediates`. + state: Starting value of each stateful buffer, by its placeholder name. Both + sides start from it. Returns: One entry per mapping, classified. @@ -492,6 +496,7 @@ async def compare_program_intermediates( if node.op == "call_function" ], torch_arguments, + state, ) mappings = get_torch_to_coreai_output_mapping(program) @@ -503,6 +508,7 @@ async def compare_program_intermediates( coreai_intermediates = await coreai_inspector.get_intermediates_for_ops( [mapping.target_op_id for mapping in mappings.values()], numpy_inputs, + state, ) return compare_intermediates( diff --git a/coreai_torch/debugging/probes.py b/coreai_torch/debugging/probes.py index 955c6fc..3e54f20 100644 --- a/coreai_torch/debugging/probes.py +++ b/coreai_torch/debugging/probes.py @@ -14,6 +14,7 @@ from pathlib import Path from typing import Any +import numpy as np from coreai._compiler.dialects import coreai as coreai_dialect from coreai._compiler.ir import ( ArrayAttr, @@ -36,9 +37,8 @@ class ProbedProgram: """ A copy of a program with extra outputs that carry intermediate values. - Nothing here needs the Python runtime. The rewritten program can be saved and run by - any runtime, and :meth:`manifest` says which of its outputs are probes, so a caller - that gets the arrays back some other way can still tell them apart. + Independent of the Python runtime: the saved program runs anywhere, and + :meth:`manifest` says which outputs are probes. Attributes: program: The rewritten copy. The original is not modified. @@ -125,18 +125,13 @@ def add_probe_outputs( """ Copy *program* and return the values of the given operations as extra outputs. - This reads an intermediate without the runtime's intermediate logger, which does - not fire on the neural engine: the value leaves the graph like any other output. - The extra outputs are appended after the existing ones, so the original outputs - keep their names and positions. + The extra outputs follow the existing ones, which keep their names and positions. Only tensor results of operations in the entry block of *function_name* can be - probed. An operation elsewhere (a nested graph, say), one with no results, or one - with no operation ID is reported in :attr:`ProbedProgram.unprobed`. + probed. Others (nested graphs, no results, no operation ID) are reported in + :attr:`ProbedProgram.unprobed`. - Probing changes what the compiler sees: a value that is an output cannot be - fused away, so the compiled program can place work differently from the - original. Compare placement before trusting a probed value as a stand-in. + An output cannot be fused away, so the compiled copy may differ from the original. Args: program: The program to read values from. Not modified. @@ -220,3 +215,78 @@ def add_probe_outputs( outputs=outputs, unprobed=unprobed, ) + + +@dataclass(frozen=True) +class Perturbation: + """ + How far probing moved the model's own outputs. + + Attributes: + op_ids: The operations that were probed together. + max_abs_diff: Largest absolute difference from the unprobed run, over every + output and state. + max_rel_diff: Largest difference relative to the unprobed value's magnitude. + within_tolerance: Whether every output and state stayed within the tolerance. + + """ + + op_ids: tuple[int, ...] + max_abs_diff: float + max_rel_diff: float + within_tolerance: bool + + +def compare_to_reference( + reference: Mapping[str, Any], + probed: Mapping[str, Any], + op_ids: Iterable[int], + rtol: float, + atol: float, +) -> Perturbation: + """ + Compare a probed run's own outputs with an unprobed run of the same program. + + Needs only numpy, so it is usable wherever the arrays came from. + + A value counts as moved when ``|probed - reference| > atol + rtol * |reference|``, + the rule ``numpy.allclose`` uses. A NaN or infinity that appears, disappears or + changes sign is always a move. + + Args: + reference: Outputs and state of the unprobed run, by name. + probed: The same names from the probed run. + op_ids: The operations that were probed, recorded in the result. + rtol: Relative tolerance. + atol: Absolute tolerance. + + Returns: + The differences found and whether they were within tolerance. + + Raises: + KeyError: If the probed run lacks a name the reference has. + + """ + max_abs = 0.0 + max_rel = 0.0 + within = True + for name, expected in reference.items(): + got = np.asarray(probed[name], dtype=np.float64) + want = np.asarray(expected, dtype=np.float64) + if got.shape != want.shape or not np.array_equal( + np.isfinite(got), np.isfinite(want) + ): + max_abs = max_rel = float("inf") + within = False + continue + finite = np.isfinite(want) + diff = np.abs(got[finite] - want[finite]) + if diff.size == 0: + continue + scale = np.abs(want[finite]) + max_abs = max(max_abs, float(diff.max())) + max_rel = max( + max_rel, float((diff / np.maximum(scale, np.finfo(float).tiny)).max()) + ) + within = within and bool(np.all(diff <= atol + rtol * scale)) + return Perturbation(tuple(op_ids), max_abs, max_rel, within) diff --git a/coreai_torch/debugging/utils.py b/coreai_torch/debugging/utils.py index 869b104..78b3cb4 100644 --- a/coreai_torch/debugging/utils.py +++ b/coreai_torch/debugging/utils.py @@ -45,6 +45,11 @@ def _with_debug( return specialization_options.with_debug(enabled=True) +def _state_kwargs(state: Any) -> dict[str, Any]: + """``state`` as a keyword only when given, so inspectors without it still work.""" + return {} if state is None else {"state": state} + + def _plain(value: Any) -> Any: """ Normalise one leaf value for a report's ``to_dict``. diff --git a/coreai_torch/debugging/validator.py b/coreai_torch/debugging/validator.py index 1580439..30da71d 100644 --- a/coreai_torch/debugging/validator.py +++ b/coreai_torch/debugging/validator.py @@ -18,7 +18,7 @@ - Concrete implementations for Core AI AIProgram and PyTorch ExportedProgram """ -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass, field from pathlib import Path from tempfile import TemporaryDirectory @@ -47,7 +47,7 @@ ExhaustiveStrategy, SearchStrategy, ) -from .utils import _with_debug +from .utils import _state_kwargs, _with_debug TNode = TypeVar("TNode") TGraph = TypeVar("TGraph") @@ -247,6 +247,7 @@ async def check( self, check_fn: Callable[[list[NDArray[Any] | None]], bool], inputs: Any, + state: Mapping[str, Any] | None = None, ) -> Result: """ Find operations where check_fn returns True. @@ -254,6 +255,7 @@ async def check( Args: check_fn: Function that returns True if outputs have issues inputs: Model inputs to use during execution + state: Starting value of each stateful buffer, by name Returns: ValidationResult containing failed and unknown operations sorted by topological order. @@ -279,6 +281,7 @@ async def check( results = await self.inspector.get_intermediates_for_ops( batch_op_ids, inputs, + **_state_kwargs(state), ) # Check each node in the batch and collect results @@ -346,12 +349,14 @@ async def check( async def check_for_nans( self, inputs: Any, + state: Mapping[str, Any] | None = None, ) -> Result: """ Find operations that produce NaN values. Args: inputs: Model inputs to use during execution + state: Starting value of each stateful buffer, by name Returns: ValidationResult containing operations that produced NaN and unknown operations, @@ -366,17 +371,20 @@ async def check_for_nans( for arr in outputs ), inputs, + state, ) async def check_for_infs( self, inputs: Any, + state: Mapping[str, Any] | None = None, ) -> Result: """ Find operations that produce infinite values. Args: inputs: Model inputs to use during execution + state: Starting value of each stateful buffer, by name Returns: ValidationResult containing operations that produced inf and unknown operations, @@ -391,6 +399,7 @@ async def check_for_infs( for arr in outputs ), inputs, + state, ) diff --git a/tests/debugging/test_comparator.py b/tests/debugging/test_comparator.py index 9174797..d5b14f3 100644 --- a/tests/debugging/test_comparator.py +++ b/tests/debugging/test_comparator.py @@ -651,3 +651,51 @@ async def test_comparator_catches_modified_ops_at_different_positions( assert len(result_2.failed_nodes) > 0, ( "Should detect failure from changing second mul to div" ) + + +class _Accumulator(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("acc", torch.zeros(1, 4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = self.acc + x + self.acc.copy_(y * 2) + return torch.relu(y) + + +@pytest.mark.parametrize("target_inspector", ["logger", "probed"]) +async def test_comparator_starts_both_sides_from_the_given_state( + specialization_options: SpecializationOptions | None, + target_inspector: str, +) -> None: + x = torch.randn(1, 4) + exported = torch.export.export(_Accumulator().eval(), (x,)).run_decompositions() + program = await _create_coreai_program_from_model(exported) + comparator = await create_comparator_for_programs( + source_program=exported, + target_program=program, + target_entry_point="main", + specialization_options=specialization_options, + target_inspector=target_inspector, # type: ignore[arg-type] + ) + state = {"b_acc": np.full((1, 4), 3.0, dtype=np.float32)} + + result = await comparator.compare_with_tolerance( + inputs={"x": x.numpy()}, state=state + ) + + assert len(result.compared) > 0 + assert result.failed_nodes == [] + # The caller's array is not updated by either run. + np.testing.assert_array_equal(state["b_acc"], np.full((1, 4), 3.0)) + + +async def test_probed_options_need_the_probed_inspector() -> None: + with pytest.raises(ValueError, match="only valid"): + await create_comparator_for_programs( + source_program=None, # type: ignore[arg-type] + target_program=None, # type: ignore[arg-type] + target_entry_point="main", + probed_options={"rtol": 1e-2}, + ) diff --git a/tests/debugging/test_probed_inspector.py b/tests/debugging/test_probed_inspector.py index 5b5b4cd..3136e2a 100644 --- a/tests/debugging/test_probed_inspector.py +++ b/tests/debugging/test_probed_inspector.py @@ -15,9 +15,17 @@ from coreai.runtime import SpecializationOptions from coreai_torch.converter import TorchConverter +from coreai_torch.debugging import inspector as inspector_module from coreai_torch.debugging.debug_info import _build_coreai_op_map -from coreai_torch.debugging.inspector import CoreAIProbedInspector -from coreai_torch.debugging.probes import add_probe_outputs +from coreai_torch.debugging.inspector import ( + CoreAIProbedInspector, + PerturbationError, +) +from coreai_torch.debugging.probes import ( + Perturbation, + add_probe_outputs, + compare_to_reference, +) class _Net(torch.nn.Module): @@ -93,7 +101,10 @@ async def test_probed_inspector_returns_intermediates( (relu,) = result[relu_id] (conv,) = result[conv_id] assert conv.shape == relu.shape == (1, 4, 8, 8) - np.testing.assert_array_equal(relu, inspector.last_outputs["relu"]) + # last_outputs is the unprobed run's; the probed copy is within an fp16 step. + np.testing.assert_allclose( + relu, inspector.last_outputs["relu"], rtol=1e-2, atol=2e-3 + ) assert np.all(relu >= 0) # The value is the convolution before the bias add, so relu(conv) is not it. assert not np.array_equal(conv, relu) @@ -134,8 +145,7 @@ async def test_probed_inspector_reads_stateful_models_without_advancing_state( ((written,),) = (result[written_id],) np.testing.assert_array_equal(read, start) np.testing.assert_allclose(written, (start + x.numpy()) * 2, rtol=1e-5) - # What was written is what the model left in its state, and the caller's own - # array is untouched so a second call starts from the same place. + # The written value is the state left behind; the caller's array is untouched. np.testing.assert_array_equal(inspector.last_state["b_acc"], written) np.testing.assert_array_equal(caller_state["b_acc"], np.ones((1, 4))) assert set(inspector.last_outputs) == {"relu"} @@ -149,7 +159,7 @@ def test_probed_program_describes_itself_without_a_runtime( (relu_id,) = _ids_named(program, "coreai.relu") probed = add_probe_outputs(program, [relu_id, 10_000]) - # A run's outputs split by name alone, so a caller using another runtime can do it. + # Splitting needs only output names. own, by_op = probed.split_outputs({"relu": 1, f"__probe_{relu_id}_0": 2}) assert own == {"relu": 1} assert by_op == {relu_id: [2]} @@ -196,13 +206,15 @@ async def test_compiled_copy_is_reused_for_the_same_operations( ) first = await inspector.get_intermediates_for_ops([relu_id], {"x": x.numpy()}) - # Same set in another order, other inputs: no compile, and the answer follows the input. + # Reference plus probed copy. + assert loads.count == 2 + # Same set, new inputs: no compile. again = await inspector.get_intermediates_for_ops([relu_id], {"x": x.numpy() * 2}) - assert loads.count == 1 + assert loads.count == 2 assert not np.array_equal(first[relu_id][0], again[relu_id][0]) await inspector.get_intermediates_for_ops([conv_id, relu_id], {"x": x.numpy()}) - assert loads.count == 2 + assert loads.count == 3 assert inspector.cached_models == 2 @@ -224,7 +236,8 @@ async def test_cache_is_bounded_and_can_be_switched_off( await bounded.get_intermediates_for_ops([relu_id], inputs) # evicts conv assert bounded.cached_models == 1 await bounded.get_intermediates_for_ops([conv_id], inputs) # compiles again - assert loads.count == 3 + # Reference, conv, relu, conv again. + assert loads.count == 4 bounded.clear_cache() assert bounded.cached_models == 0 @@ -234,8 +247,162 @@ async def test_cache_is_bounded_and_can_be_switched_off( ) await uncached.get_intermediates_for_ops([relu_id], inputs) await uncached.get_intermediates_for_ops([relu_id], inputs) - assert loads.count == 2 + # Nothing kept: two compiles per call. + assert loads.count == 4 assert uncached.cached_models == 0 with pytest.raises(ValueError, match="must not be negative"): CoreAIProbedInspector(program, max_cached_models=-1) + + +def test_compare_to_reference_applies_the_allclose_rule() -> None: + reference = {"y": np.array([1.0, 100.0], dtype=np.float32)} + + near = compare_to_reference( + reference, {"y": np.array([1.0005, 100.05])}, [7], rtol=1e-3, atol=1e-3 + ) + assert near.within_tolerance + assert near.op_ids == (7,) + assert near.max_abs_diff == pytest.approx(0.05) + + far = compare_to_reference( + reference, {"y": np.array([1.0, 100.5])}, [7], rtol=1e-3, atol=1e-3 + ) + assert not far.within_tolerance + + # A new NaN or a changed shape always fails. + assert not compare_to_reference( + reference, {"y": np.array([1.0, np.nan])}, [7], rtol=1.0, atol=1.0 + ).within_tolerance + assert not compare_to_reference( + reference, {"y": np.zeros(3)}, [7], rtol=1.0, atol=1.0 + ).within_tolerance + + +def _perturbed_by(culprit: int): # noqa: ANN202 + """A stand-in comparison that finds any set containing *culprit* perturbed.""" + + def compare(reference, probed, op_ids, rtol, atol): # noqa: ANN001, ANN202, ARG001 + ids = tuple(op_ids) + return Perturbation( + ids, 1.0 if culprit in ids else 0.0, 0.0, culprit not in ids + ) + + return compare + + +async def test_a_perturbing_set_is_split_until_every_value_is_clean( + program_and_input: tuple[AIProgram, torch.Tensor], + specialization_options: SpecializationOptions | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + program, x = program_and_input + (conv_id,) = _ids_named(program, "coreai.conv2d") + (relu_id,) = _ids_named(program, "coreai.relu") + (add_id,) = _ids_named(program, "coreai.decomposable.broadcasting_add") + monkeypatch.setattr( + inspector_module, "compare_to_reference", _perturbed_by(relu_id) + ) + inspector = CoreAIProbedInspector( + program, specialization_options=specialization_options + ) + + result = await inspector.get_intermediates_for_ops( + [conv_id, add_id, relu_id], {"x": x.numpy()} + ) + + assert result[relu_id] is None + assert result[conv_id] is not None + assert result[add_id] is not None + assert inspector.perturbing_ops == {relu_id} + # The whole set fails, then its halves; the half without the culprit passes. + sets = [(p.op_ids, p.within_tolerance) for p in inspector.perturbations] + assert sets[0] == ((conv_id, add_id, relu_id), False) + assert any(ids == (relu_id,) and not ok for ids, ok in sets) + assert any(relu_id not in ids and ok for ids, ok in sets) + + +async def test_perturbation_can_warn_or_raise_instead_of_splitting( + program_and_input: tuple[AIProgram, torch.Tensor], + specialization_options: SpecializationOptions | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + program, x = program_and_input + (conv_id,) = _ids_named(program, "coreai.conv2d") + (relu_id,) = _ids_named(program, "coreai.relu") + monkeypatch.setattr( + inspector_module, "compare_to_reference", _perturbed_by(relu_id) + ) + inputs = {"x": x.numpy()} + + warn = CoreAIProbedInspector( + program, specialization_options=specialization_options, on_perturbation="warn" + ) + result = await warn.get_intermediates_for_ops([conv_id, relu_id], inputs) + assert result[relu_id] is not None # returned anyway, and logged + assert not warn.perturbations[0].within_tolerance + + strict = CoreAIProbedInspector( + program, specialization_options=specialization_options, on_perturbation="raise" + ) + with pytest.raises(PerturbationError): + await strict.get_intermediates_for_ops([conv_id, relu_id], inputs) + + with pytest.raises(ValueError, match="on_perturbation"): + CoreAIProbedInspector(program, on_perturbation="ignore") # type: ignore[arg-type] + with pytest.raises(ValueError, match="rtol and atol"): + CoreAIProbedInspector(program, rtol=-1.0) + + +async def test_a_clean_probe_reports_no_perturbation( + program_and_input: tuple[AIProgram, torch.Tensor], + specialization_options: SpecializationOptions | None, +) -> None: + program, x = program_and_input + (relu_id,) = _ids_named(program, "coreai.relu") + inspector = CoreAIProbedInspector( + program, specialization_options=specialization_options + ) + + await inspector.get_intermediates_for_ops([relu_id], {"x": x.numpy()}) + + (report,) = inspector.perturbations + assert report.within_tolerance + assert inspector.perturbing_ops == set() + + +async def test_a_set_that_fails_to_compile_is_reduced( + program_and_input: tuple[AIProgram, torch.Tensor], + specialization_options: SpecializationOptions | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + program, x = program_and_input + (conv_id,) = _ids_named(program, "coreai.conv2d") + (relu_id,) = _ids_named(program, "coreai.relu") + real = CoreAIProbedInspector._compile + + async def failing(self, op_ids): # noqa: ANN001, ANN202 + if relu_id in op_ids: + msg = "cannot load" + raise RuntimeError(msg) + return await real(self, op_ids) + + monkeypatch.setattr(CoreAIProbedInspector, "_compile", failing) + inspector = CoreAIProbedInspector( + program, specialization_options=specialization_options + ) + + result = await inspector.get_intermediates_for_ops( + [conv_id, relu_id], {"x": x.numpy()} + ) + + assert result[relu_id] is None + assert result[conv_id] is not None + assert inspector.failed_ops == {relu_id} + assert inspector.perturbing_ops == set() + + strict = CoreAIProbedInspector( + program, specialization_options=specialization_options, on_perturbation="raise" + ) + with pytest.raises(PerturbationError, match="cannot load"): + await strict.get_intermediates_for_ops([conv_id, relu_id], {"x": x.numpy()}) From 6f48ff2c809ad71e99e7f60938bbdfd1f02fb342 Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Mon, 19 Oct 2026 18:13:50 -0700 Subject: [PATCH 03/11] Use enums for probe policy and comparator target; trim probed tests PerturbationPolicy and TargetInspector replace string options and keep the string values working. The probed inspector tests now cover its main behaviors. --- coreai_torch/debugging/comparator.py | 41 ++- coreai_torch/debugging/inspector.py | 43 ++- tests/debugging/test_comparator.py | 17 +- tests/debugging/test_probed_inspector.py | 385 ++++++----------------- 4 files changed, 159 insertions(+), 327 deletions(-) diff --git a/coreai_torch/debugging/comparator.py b/coreai_torch/debugging/comparator.py index 5aeff6c..bdd9c91 100644 --- a/coreai_torch/debugging/comparator.py +++ b/coreai_torch/debugging/comparator.py @@ -25,7 +25,7 @@ from enum import Enum, auto from pathlib import Path from tempfile import TemporaryDirectory -from typing import Any, Generic, Literal, TypeVar +from typing import Any, Generic, TypeVar import numpy as np import torch @@ -57,6 +57,17 @@ logger = logging.getLogger(__name__) + +class TargetInspector(str, Enum): + """How a comparator reads the target program's intermediates.""" + + LOGGER = "logger" + """The runtime's intermediate logger. Leaves the model as it is.""" + + PROBED = "probed" + """Extra model outputs (:class:`CoreAIProbedInspector`), for work the logger misses.""" + + _DEFAULT_RTOL: float = 1e-5 _DEFAULT_ATOL: float = 1e-3 @@ -1527,7 +1538,7 @@ async def create_comparator_for_programs( use_caching: bool = True, exclude_ops: frozenset[str] = _DEFAULT_EXCLUDED_OPS, specialization_options: SpecializationOptions | None = None, - target_inspector: Literal["logger", "probed"] = "logger", + target_inspector: TargetInspector | str = TargetInspector.LOGGER, probed_options: Mapping[str, Any] | None = None, ) -> Comparator[torch.fx.Node, torch.fx.Graph, Operation, Module]: """ @@ -1551,27 +1562,23 @@ async def create_comparator_for_programs( Defaults to _DEFAULT_EXCLUDED_OPS which includes view/reshape operations. Pass frozenset() to disable exclusions. specialization_options: Options for configuring model specialization - target_inspector: How target intermediates are read. ``"logger"`` uses the - runtime's intermediate logger. ``"probed"`` adds the values to the model's - outputs (:class:`CoreAIProbedInspector`), so the values come from a copy of - the model with extra outputs; check its ``perturbations``. - probed_options: Extra keyword arguments for :class:`CoreAIProbedInspector` - (``rtol``, ``atol``, ``on_perturbation``, ``max_cached_models``). Only - valid with ``target_inspector="probed"``. + target_inspector: A :class:`TargetInspector` (or its string value). With + ``PROBED`` the values come from a copy of the model with extra outputs; + check the inspector's ``perturbations``. + probed_options: Keyword arguments for :class:`CoreAIProbedInspector`. Only + valid with ``PROBED``. Returns: Comparator instance configured for comparing the two programs Raises: - ValueError: If target_inspector is unknown, or probed_options is given - without ``"probed"``. + ValueError: If target_inspector is not a :class:`TargetInspector`, or + probed_options is given without ``PROBED``. """ - if target_inspector not in ("logger", "probed"): - msg = f"target_inspector must be logger or probed, got {target_inspector}" - raise ValueError(msg) - if probed_options and target_inspector != "probed": - msg = "probed_options is only valid with target_inspector='probed'" + target_inspector = TargetInspector(target_inspector) + if probed_options and target_inspector is not TargetInspector.PROBED: + msg = "probed_options is only valid with TargetInspector.PROBED" raise ValueError(msg) # Create ID mapping from torch to coreai operations @@ -1590,7 +1597,7 @@ async def create_comparator_for_programs( source_graph = create_graph_from_exported_program(source_program) target: Inspector - if target_inspector == "probed": + if target_inspector is TargetInspector.PROBED: target = CoreAIProbedInspector( target_program, function_name=target_entry_point, diff --git a/coreai_torch/debugging/inspector.py b/coreai_torch/debugging/inspector.py index aa188f1..291a67d 100644 --- a/coreai_torch/debugging/inspector.py +++ b/coreai_torch/debugging/inspector.py @@ -15,7 +15,7 @@ from enum import Enum, auto from pathlib import Path from tempfile import TemporaryDirectory -from typing import Any, Literal, Union, cast +from typing import Any, Union, cast import coreai import numpy as np @@ -1199,6 +1199,19 @@ def get_all_compile_identifiers( } +class PerturbationPolicy(str, Enum): + """What a probed inspector does with a probe set that moves the model's outputs.""" + + SPLIT = "split" + """Retry each half until every returned value comes from a clean set.""" + + WARN = "warn" + """Log and return the values anyway.""" + + RAISE = "raise" + """Raise :class:`PerturbationError`.""" + + class PerturbationError(RuntimeError): """Probing moved the model's own outputs beyond the tolerance.""" @@ -1242,7 +1255,7 @@ def __init__( max_cached_models: int = 4, rtol: float = 1e-3, atol: float = 1e-3, - on_perturbation: Literal["split", "warn", "raise"] = "split", + on_perturbation: PerturbationPolicy | str = PerturbationPolicy.SPLIT, ): """ Initialize the probed inspector. @@ -1258,15 +1271,13 @@ def __init__( rtol: Relative tolerance of the perturbation check. atol: Absolute tolerance of the perturbation check. Neither default is tuned to a model; set them for the precision you run. - on_perturbation: For a set that moves the outputs. ``"split"`` retries - each half until every returned value comes from a clean set; - ``"warn"`` logs and returns the values; ``"raise"`` raises - :class:`PerturbationError`. A set that fails to compile or run is - always split, or raises under ``"raise"``. + on_perturbation: A :class:`PerturbationPolicy` (or its string value) for + a set that moves the outputs. A set that fails to compile or run is + split, or raises under ``RAISE``. Raises: ValueError: If max_cached_models, rtol or atol is negative, or - on_perturbation is not one of the three. + on_perturbation is not a policy. """ if max_cached_models < 0: @@ -1275,13 +1286,10 @@ def __init__( if rtol < 0 or atol < 0: msg = f"rtol and atol must not be negative, got {rtol} and {atol}" raise ValueError(msg) - if on_perturbation not in ("split", "warn", "raise"): - msg = f"on_perturbation must be split, warn or raise, got {on_perturbation}" - raise ValueError(msg) self._max_cached_models = max_cached_models self._rtol = rtol self._atol = atol - self._on_perturbation = on_perturbation + self._on_perturbation = PerturbationPolicy(on_perturbation) self._compiled: OrderedDict[frozenset[int], _CompiledProbe] = OrderedDict() self._reference: _CompiledProbe | None = None self._program = program @@ -1317,7 +1325,7 @@ async def get_intermediates_for_ops( Raises: TypeError: If inputs is not a dictionary. PerturbationError: If a set moved the outputs, or failed to compile or - run, and ``on_perturbation="raise"``. + run, and the policy is ``RAISE``. """ if not isinstance(inputs, dict): @@ -1364,7 +1372,7 @@ async def get_intermediates_for_ops( if self._max_cached_models == 0: compiled.temp_dir.cleanup() except Exception as error: - if self._on_perturbation == "raise": + if self._on_perturbation is PerturbationPolicy.RAISE: msg = f"probing {batch} failed to compile or run: {error}" raise PerturbationError(msg) from error self._split_or_give_up( @@ -1383,7 +1391,10 @@ async def get_intermediates_for_ops( if probeable: self._perturbations.append(report) - if report.within_tolerance or self._on_perturbation == "warn": + if ( + report.within_tolerance + or self._on_perturbation is PerturbationPolicy.WARN + ): if not report.within_tolerance: logger.warning( "Probing %s moved the model's outputs (max abs %.3g, " @@ -1394,7 +1405,7 @@ async def get_intermediates_for_ops( ) for op_id, values in probed_arrays.items(): results[op_id] = [np.asarray(value) for value in values] - elif self._on_perturbation == "raise": + elif self._on_perturbation is PerturbationPolicy.RAISE: msg = ( f"probing {probeable} moved the model's outputs " f"(max abs {report.max_abs_diff:.3g}, " diff --git a/tests/debugging/test_comparator.py b/tests/debugging/test_comparator.py index d5b14f3..04f6cb1 100644 --- a/tests/debugging/test_comparator.py +++ b/tests/debugging/test_comparator.py @@ -20,6 +20,7 @@ from coreai_torch.debugging.comparator import ( Comparator, DebugGraph, + TargetInspector, create_comparator_for_programs, ) from coreai_torch.debugging.graph import ComputationGraph @@ -664,10 +665,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.relu(y) -@pytest.mark.parametrize("target_inspector", ["logger", "probed"]) +@pytest.mark.parametrize("target_inspector", list(TargetInspector)) async def test_comparator_starts_both_sides_from_the_given_state( specialization_options: SpecializationOptions | None, - target_inspector: str, + target_inspector: TargetInspector, ) -> None: x = torch.randn(1, 4) exported = torch.export.export(_Accumulator().eval(), (x,)).run_decompositions() @@ -677,7 +678,7 @@ async def test_comparator_starts_both_sides_from_the_given_state( target_program=program, target_entry_point="main", specialization_options=specialization_options, - target_inspector=target_inspector, # type: ignore[arg-type] + target_inspector=target_inspector, ) state = {"b_acc": np.full((1, 4), 3.0, dtype=np.float32)} @@ -689,13 +690,3 @@ async def test_comparator_starts_both_sides_from_the_given_state( assert result.failed_nodes == [] # The caller's array is not updated by either run. np.testing.assert_array_equal(state["b_acc"], np.full((1, 4), 3.0)) - - -async def test_probed_options_need_the_probed_inspector() -> None: - with pytest.raises(ValueError, match="only valid"): - await create_comparator_for_programs( - source_program=None, # type: ignore[arg-type] - target_program=None, # type: ignore[arg-type] - target_entry_point="main", - probed_options={"rtol": 1e-2}, - ) diff --git a/tests/debugging/test_probed_inspector.py b/tests/debugging/test_probed_inspector.py index 3136e2a..161838e 100644 --- a/tests/debugging/test_probed_inspector.py +++ b/tests/debugging/test_probed_inspector.py @@ -5,9 +5,6 @@ """Tests for the inspector that reads intermediates as extra model outputs.""" -import json -from pathlib import Path - import numpy as np import pytest import torch @@ -20,6 +17,7 @@ from coreai_torch.debugging.inspector import ( CoreAIProbedInspector, PerturbationError, + PerturbationPolicy, ) from coreai_torch.debugging.probes import ( Perturbation, @@ -37,58 +35,63 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.relu(self.conv(x)) +class _Accumulator(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("acc", torch.zeros(1, 4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = self.acc + x + self.acc.copy_(y * 2) + return torch.relu(y) + + +def _convert(model: torch.nn.Module, x: torch.Tensor) -> AIProgram: + exported = torch.export.export(model.eval(), (x,)).run_decompositions() + converter = TorchConverter() + converter.add_exported_program(exported, entrypoint_name="main") + return converter.to_coreai() + + @pytest.fixture def program_and_input() -> tuple[AIProgram, torch.Tensor]: torch.manual_seed(0) - model = _Net().eval().half() x = torch.randn(1, 4, 8, 8).half() - exported = torch.export.export(model, (x,)).run_decompositions() - converter = TorchConverter() - converter.add_exported_program(exported, entrypoint_name="main") - return converter.to_coreai(), x + return _convert(_Net().half(), x), x -def _ids_named(program: AIProgram, name: str) -> list[int]: +def _ids(program: AIProgram, name: str) -> list[int]: return [i for i, op in _build_coreai_op_map(program).items() if op.name == name] -def test_probing_leaves_the_original_program_and_outputs_alone( +def test_probe_outputs_are_appended_without_touching_the_program( program_and_input: tuple[AIProgram, torch.Tensor], ) -> None: program, _ = program_and_input before = str(program) - (relu_id,) = _ids_named(program, "coreai.relu") + (relu_id,) = _ids(program, "coreai.relu") + (output_id,) = _ids(program, "coreai.output") - probed = add_probe_outputs(program, [relu_id]) + probed = add_probe_outputs(program, [relu_id, output_id]) assert str(program) == before assert probed.outputs == {relu_id: [f"__probe_{relu_id}_0"]} - graph = probed.program._get_graph("main") - assert [r["coreai.name"].value for r in graph.res_attrs] == [ - "relu", - f"__probe_{relu_id}_0", + assert probed.unprobed == [output_id] + names = [ + r["coreai.name"].value for r in probed.program._get_graph("main").res_attrs ] + assert names == ["relu", f"__probe_{relu_id}_0"] + own, by_op = probed.split_outputs({"relu": 1, f"__probe_{relu_id}_0": 2}) + assert (own, by_op) == ({"relu": 1}, {relu_id: [2]}) -def test_ids_that_cannot_be_probed_are_reported( - program_and_input: tuple[AIProgram, torch.Tensor], -) -> None: - program, _ = program_and_input - (output_id,) = _ids_named(program, "coreai.output") - - probed = add_probe_outputs(program, [output_id, 10_000]) - - assert probed.outputs == {} - assert probed.unprobed == [output_id, 10_000] - - -async def test_probed_inspector_returns_intermediates( +async def test_intermediates_come_back_by_op_id( program_and_input: tuple[AIProgram, torch.Tensor], specialization_options: SpecializationOptions | None, ) -> None: program, x = program_and_input - (conv_id,) = _ids_named(program, "coreai.conv2d") - (relu_id,) = _ids_named(program, "coreai.relu") + (conv_id,) = _ids(program, "coreai.conv2d") + (relu_id,) = _ids(program, "coreai.relu") inspector = CoreAIProbedInspector( program, specialization_options=specialization_options ) @@ -98,311 +101,131 @@ async def test_probed_inspector_returns_intermediates( ) assert result[10_000] is None - (relu,) = result[relu_id] - (conv,) = result[conv_id] + ((conv,), (relu,)) = result[conv_id], result[relu_id] assert conv.shape == relu.shape == (1, 4, 8, 8) + assert np.all(relu >= 0) # last_outputs is the unprobed run's; the probed copy is within an fp16 step. np.testing.assert_allclose( relu, inspector.last_outputs["relu"], rtol=1e-2, atol=2e-3 ) - assert np.all(relu >= 0) - # The value is the convolution before the bias add, so relu(conv) is not it. - assert not np.array_equal(conv, relu) + assert all(p.within_tolerance for p in inspector.perturbations) -class _Accumulator(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.register_buffer("acc", torch.zeros(1, 4)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - y = self.acc + x - self.acc.copy_(y * 2) - return torch.relu(y) - - -async def test_probed_inspector_reads_stateful_models_without_advancing_state( +async def test_state_is_read_and_the_callers_copy_is_left_alone( specialization_options: SpecializationOptions | None, ) -> None: x = torch.randn(1, 4) - exported = torch.export.export(_Accumulator().eval(), (x,)).run_decompositions() - converter = TorchConverter() - converter.add_exported_program(exported, entrypoint_name="main") - program = converter.to_coreai() - (read_id,) = _ids_named(program, "coreai.read_handle") - (written_id,) = _ids_named(program, "coreai.decomposable.broadcasting_mul") + program = _convert(_Accumulator(), x) + (read_id,) = _ids(program, "coreai.read_handle") + (written_id,) = _ids(program, "coreai.decomposable.broadcasting_mul") inspector = CoreAIProbedInspector( program, specialization_options=specialization_options ) start = np.ones((1, 4), dtype=np.float32) - caller_state = {"b_acc": start} result = await inspector.get_intermediates_for_ops( - [read_id, written_id], {"x": x.numpy()}, state=caller_state + [read_id, written_id], {"x": x.numpy()}, state={"b_acc": start} ) - ((read,),) = (result[read_id],) - ((written,),) = (result[written_id],) - np.testing.assert_array_equal(read, start) - np.testing.assert_allclose(written, (start + x.numpy()) * 2, rtol=1e-5) - # The written value is the state left behind; the caller's array is untouched. - np.testing.assert_array_equal(inspector.last_state["b_acc"], written) - np.testing.assert_array_equal(caller_state["b_acc"], np.ones((1, 4))) - assert set(inspector.last_outputs) == {"relu"} - - -def test_probed_program_describes_itself_without_a_runtime( - program_and_input: tuple[AIProgram, torch.Tensor], - tmp_path: Path, -) -> None: - program, _ = program_and_input - (relu_id,) = _ids_named(program, "coreai.relu") - probed = add_probe_outputs(program, [relu_id, 10_000]) - - # Splitting needs only output names. - own, by_op = probed.split_outputs({"relu": 1, f"__probe_{relu_id}_0": 2}) - assert own == {"relu": 1} - assert by_op == {relu_id: [2]} - - manifest = probed.manifest() - assert manifest == { - "function": "main", - "probes": {str(relu_id): [f"__probe_{relu_id}_0"]}, - "unprobed": [10_000], - } - - asset = probed.save(tmp_path / "out") - assert asset.exists() - assert json.loads((tmp_path / "out" / "probes.json").read_text()) == manifest - - -class _CountLoads: - """Count compiles by wrapping AIModel.load.""" - - def __init__(self, monkeypatch: pytest.MonkeyPatch) -> None: - import coreai.runtime - - self.count = 0 - real = coreai.runtime.AIModel.load - - async def counting(*args, **kwargs): # noqa: ANN002, ANN003, ANN202 - self.count += 1 - return await real(*args, **kwargs) - - monkeypatch.setattr(coreai.runtime.AIModel, "load", counting) - - -async def test_compiled_copy_is_reused_for_the_same_operations( - program_and_input: tuple[AIProgram, torch.Tensor], - specialization_options: SpecializationOptions | None, - monkeypatch: pytest.MonkeyPatch, -) -> None: - program, x = program_and_input - (conv_id,) = _ids_named(program, "coreai.conv2d") - (relu_id,) = _ids_named(program, "coreai.relu") - loads = _CountLoads(monkeypatch) - inspector = CoreAIProbedInspector( - program, specialization_options=specialization_options + np.testing.assert_array_equal(result[read_id][0], start) + np.testing.assert_allclose( + result[written_id][0], (start + x.numpy()) * 2, rtol=1e-5 ) + np.testing.assert_array_equal(inspector.last_state["b_acc"], result[written_id][0]) + np.testing.assert_array_equal(start, np.ones((1, 4))) - first = await inspector.get_intermediates_for_ops([relu_id], {"x": x.numpy()}) - # Reference plus probed copy. - assert loads.count == 2 - # Same set, new inputs: no compile. - again = await inspector.get_intermediates_for_ops([relu_id], {"x": x.numpy() * 2}) - assert loads.count == 2 - assert not np.array_equal(first[relu_id][0], again[relu_id][0]) - - await inspector.get_intermediates_for_ops([conv_id, relu_id], {"x": x.numpy()}) - assert loads.count == 3 - assert inspector.cached_models == 2 - -async def test_cache_is_bounded_and_can_be_switched_off( +async def test_compiled_copies_are_reused_and_bounded( program_and_input: tuple[AIProgram, torch.Tensor], specialization_options: SpecializationOptions | None, monkeypatch: pytest.MonkeyPatch, ) -> None: program, x = program_and_input - (conv_id,) = _ids_named(program, "coreai.conv2d") - (relu_id,) = _ids_named(program, "coreai.relu") - loads = _CountLoads(monkeypatch) - inputs = {"x": x.numpy()} + (conv_id,) = _ids(program, "coreai.conv2d") + (relu_id,) = _ids(program, "coreai.relu") + loads = [] + real = CoreAIProbedInspector._compile + + async def counting(self, op_ids): # noqa: ANN001, ANN202 + loads.append(tuple(op_ids)) + return await real(self, op_ids) - bounded = CoreAIProbedInspector( + monkeypatch.setattr(CoreAIProbedInspector, "_compile", counting) + inspector = CoreAIProbedInspector( program, specialization_options=specialization_options, max_cached_models=1 ) - await bounded.get_intermediates_for_ops([conv_id], inputs) - await bounded.get_intermediates_for_ops([relu_id], inputs) # evicts conv - assert bounded.cached_models == 1 - await bounded.get_intermediates_for_ops([conv_id], inputs) # compiles again - # Reference, conv, relu, conv again. - assert loads.count == 4 - bounded.clear_cache() - assert bounded.cached_models == 0 - - loads.count = 0 - uncached = CoreAIProbedInspector( - program, specialization_options=specialization_options, max_cached_models=0 - ) - await uncached.get_intermediates_for_ops([relu_id], inputs) - await uncached.get_intermediates_for_ops([relu_id], inputs) - # Nothing kept: two compiles per call. - assert loads.count == 4 - assert uncached.cached_models == 0 + inputs = {"x": x.numpy()} + + await inspector.get_intermediates_for_ops([relu_id], inputs) + await inspector.get_intermediates_for_ops([relu_id], {"x": x.numpy() * 2}) + assert loads == [(), (relu_id,)] # the reference once, the probed copy once - with pytest.raises(ValueError, match="must not be negative"): - CoreAIProbedInspector(program, max_cached_models=-1) + await inspector.get_intermediates_for_ops([conv_id], inputs) # evicts relu + await inspector.get_intermediates_for_ops([relu_id], inputs) + assert loads[2:] == [(conv_id,), (relu_id,)] + assert inspector.cached_models == 1 -def test_compare_to_reference_applies_the_allclose_rule() -> None: +def test_comparison_follows_the_allclose_rule() -> None: reference = {"y": np.array([1.0, 100.0], dtype=np.float32)} - near = compare_to_reference( - reference, {"y": np.array([1.0005, 100.05])}, [7], rtol=1e-3, atol=1e-3 - ) - assert near.within_tolerance - assert near.op_ids == (7,) - assert near.max_abs_diff == pytest.approx(0.05) + near = compare_to_reference(reference, {"y": [1.0005, 100.05]}, [7], 1e-3, 1e-3) + far = compare_to_reference(reference, {"y": [1.0, 100.5]}, [7], 1e-3, 1e-3) + nan = compare_to_reference(reference, {"y": [1.0, np.nan]}, [7], 1.0, 1.0) - far = compare_to_reference( - reference, {"y": np.array([1.0, 100.5])}, [7], rtol=1e-3, atol=1e-3 - ) + assert near.within_tolerance assert not far.within_tolerance + assert not nan.within_tolerance - # A new NaN or a changed shape always fails. - assert not compare_to_reference( - reference, {"y": np.array([1.0, np.nan])}, [7], rtol=1.0, atol=1.0 - ).within_tolerance - assert not compare_to_reference( - reference, {"y": np.zeros(3)}, [7], rtol=1.0, atol=1.0 - ).within_tolerance - - -def _perturbed_by(culprit: int): # noqa: ANN202 - """A stand-in comparison that finds any set containing *culprit* perturbed.""" - - def compare(reference, probed, op_ids, rtol, atol): # noqa: ANN001, ANN202, ARG001 - ids = tuple(op_ids) - return Perturbation( - ids, 1.0 if culprit in ids else 0.0, 0.0, culprit not in ids - ) - - return compare - -async def test_a_perturbing_set_is_split_until_every_value_is_clean( +@pytest.mark.parametrize("failure", ["moves_outputs", "fails_to_compile"]) +async def test_a_bad_probe_set_is_reduced( program_and_input: tuple[AIProgram, torch.Tensor], specialization_options: SpecializationOptions | None, monkeypatch: pytest.MonkeyPatch, + failure: str, ) -> None: program, x = program_and_input - (conv_id,) = _ids_named(program, "coreai.conv2d") - (relu_id,) = _ids_named(program, "coreai.relu") - (add_id,) = _ids_named(program, "coreai.decomposable.broadcasting_add") - monkeypatch.setattr( - inspector_module, "compare_to_reference", _perturbed_by(relu_id) - ) - inspector = CoreAIProbedInspector( - program, specialization_options=specialization_options - ) + (conv_id,) = _ids(program, "coreai.conv2d") + (relu_id,) = _ids(program, "coreai.relu") + (add_id,) = _ids(program, "coreai.decomposable.broadcasting_add") - result = await inspector.get_intermediates_for_ops( - [conv_id, add_id, relu_id], {"x": x.numpy()} - ) + if failure == "moves_outputs": - assert result[relu_id] is None - assert result[conv_id] is not None - assert result[add_id] is not None - assert inspector.perturbing_ops == {relu_id} - # The whole set fails, then its halves; the half without the culprit passes. - sets = [(p.op_ids, p.within_tolerance) for p in inspector.perturbations] - assert sets[0] == ((conv_id, add_id, relu_id), False) - assert any(ids == (relu_id,) and not ok for ids, ok in sets) - assert any(relu_id not in ids and ok for ids, ok in sets) + def compare(reference, probed, op_ids, rtol, atol): # noqa: ANN001, ANN202, ARG001 + ids = tuple(op_ids) + return Perturbation(ids, 0.0, 0.0, relu_id not in ids) + monkeypatch.setattr(inspector_module, "compare_to_reference", compare) + else: + real = CoreAIProbedInspector._compile -async def test_perturbation_can_warn_or_raise_instead_of_splitting( - program_and_input: tuple[AIProgram, torch.Tensor], - specialization_options: SpecializationOptions | None, - monkeypatch: pytest.MonkeyPatch, -) -> None: - program, x = program_and_input - (conv_id,) = _ids_named(program, "coreai.conv2d") - (relu_id,) = _ids_named(program, "coreai.relu") - monkeypatch.setattr( - inspector_module, "compare_to_reference", _perturbed_by(relu_id) - ) - inputs = {"x": x.numpy()} + async def compile_(self, op_ids): # noqa: ANN001, ANN202 + if relu_id in op_ids: + msg = "cannot load" + raise RuntimeError(msg) + return await real(self, op_ids) - warn = CoreAIProbedInspector( - program, specialization_options=specialization_options, on_perturbation="warn" - ) - result = await warn.get_intermediates_for_ops([conv_id, relu_id], inputs) - assert result[relu_id] is not None # returned anyway, and logged - assert not warn.perturbations[0].within_tolerance - - strict = CoreAIProbedInspector( - program, specialization_options=specialization_options, on_perturbation="raise" - ) - with pytest.raises(PerturbationError): - await strict.get_intermediates_for_ops([conv_id, relu_id], inputs) + monkeypatch.setattr(CoreAIProbedInspector, "_compile", compile_) - with pytest.raises(ValueError, match="on_perturbation"): - CoreAIProbedInspector(program, on_perturbation="ignore") # type: ignore[arg-type] - with pytest.raises(ValueError, match="rtol and atol"): - CoreAIProbedInspector(program, rtol=-1.0) - - -async def test_a_clean_probe_reports_no_perturbation( - program_and_input: tuple[AIProgram, torch.Tensor], - specialization_options: SpecializationOptions | None, -) -> None: - program, x = program_and_input - (relu_id,) = _ids_named(program, "coreai.relu") inspector = CoreAIProbedInspector( program, specialization_options=specialization_options ) - - await inspector.get_intermediates_for_ops([relu_id], {"x": x.numpy()}) - - (report,) = inspector.perturbations - assert report.within_tolerance - assert inspector.perturbing_ops == set() - - -async def test_a_set_that_fails_to_compile_is_reduced( - program_and_input: tuple[AIProgram, torch.Tensor], - specialization_options: SpecializationOptions | None, - monkeypatch: pytest.MonkeyPatch, -) -> None: - program, x = program_and_input - (conv_id,) = _ids_named(program, "coreai.conv2d") - (relu_id,) = _ids_named(program, "coreai.relu") - real = CoreAIProbedInspector._compile - - async def failing(self, op_ids): # noqa: ANN001, ANN202 - if relu_id in op_ids: - msg = "cannot load" - raise RuntimeError(msg) - return await real(self, op_ids) - - monkeypatch.setattr(CoreAIProbedInspector, "_compile", failing) - inspector = CoreAIProbedInspector( - program, specialization_options=specialization_options - ) - - result = await inspector.get_intermediates_for_ops( - [conv_id, relu_id], {"x": x.numpy()} - ) + ids = [conv_id, add_id, relu_id] + result = await inspector.get_intermediates_for_ops(ids, {"x": x.numpy()}) assert result[relu_id] is None assert result[conv_id] is not None - assert inspector.failed_ops == {relu_id} - assert inspector.perturbing_ops == set() + assert result[add_id] is not None + culprits = ( + inspector.perturbing_ops if failure == "moves_outputs" else inspector.failed_ops + ) + assert culprits == {relu_id} strict = CoreAIProbedInspector( - program, specialization_options=specialization_options, on_perturbation="raise" + program, + specialization_options=specialization_options, + on_perturbation=PerturbationPolicy.RAISE, ) - with pytest.raises(PerturbationError, match="cannot load"): - await strict.get_intermediates_for_ops([conv_id, relu_id], {"x": x.numpy()}) + with pytest.raises(PerturbationError): + await strict.get_intermediates_for_ops(ids, {"x": x.numpy()}) From 82402186b99e15db3cde7c6bfc4ad4ca289a946e Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Tue, 20 Oct 2026 11:36:47 -0700 Subject: [PATCH 04/11] Add state to the benchmarker benchmark, benchmark_coreai_program and compare_runs take an optional state. Every run, warmup included, starts from a fresh copy, so no run is timed on the state a previous one left. --- coreai_torch/debugging/benchmarker.py | 19 +++++++++++---- coreai_torch/debugging/timing_diff.py | 5 +++- coreai_torch/debugging/utils.py | 11 ++++++++- tests/debugging/test_benchmarker.py | 34 +++++++++++++++++++++++++++ 4 files changed, 63 insertions(+), 6 deletions(-) diff --git a/coreai_torch/debugging/benchmarker.py b/coreai_torch/debugging/benchmarker.py index 6fa903a..c871db7 100644 --- a/coreai_torch/debugging/benchmarker.py +++ b/coreai_torch/debugging/benchmarker.py @@ -22,7 +22,7 @@ import threading from abc import ABC, abstractmethod from collections import defaultdict -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass, field from enum import Enum from pathlib import Path @@ -66,6 +66,7 @@ ) from .utils import ( LocationInfo, + _copy_state, _plain, _with_debug, get_operation_locations, @@ -1727,6 +1728,7 @@ async def benchmark( self: Self, inputs: dict[str, Any], num_runs: int = 1, + state: Mapping[str, Any] | None = None, ) -> BenchmarkResult: """ Benchmark the coreai_program program with the given inputs. @@ -1736,6 +1738,9 @@ async def benchmark( num_runs: Number of timed iterations (default: 1). Always preceded by an untimed warmup run, which is not one of them, so the samples describe the steady state rather than first-inference costs. + state: Starting value of each stateful buffer, by name. Every run, + warmup included, starts from a fresh copy, so no run is timed on the + state a previous one left. Returns: BenchmarkResult containing timing information for all operations @@ -1751,6 +1756,7 @@ async def benchmark( self: Self, inputs: dict[str, Any], num_runs: int = 1, + state: Mapping[str, Any] | None = None, ) -> BenchmarkResult: """ Benchmark using Core AI Runtime. @@ -1760,6 +1766,9 @@ async def benchmark( num_runs: Number of timed iterations (default: 1). Always preceded by an untimed warmup run, which is not one of them, so the samples describe the steady state rather than first-inference costs. + state: Starting value of each stateful buffer, by name. Every run, + warmup included, starts from a fresh copy, so no run is timed on the + state a previous one left. Returns: BenchmarkResult containing timing information @@ -1823,7 +1832,7 @@ async def benchmark( # cold one, so a single-run number was not comparable with any other -- # including the other side of a `timing_diff` comparison. logger.debug("Warmup run (untimed)") - await function(nd_inputs) + await function(nd_inputs, state=_copy_state(state)) # Transition to RUNNING state to start collecting timing data. Written # under the lock the callbacks take, so the transition is ordered against @@ -1833,7 +1842,7 @@ async def benchmark( for i in range(num_runs): logger.debug("Benchmark run %d/%d", i + 1, num_runs) - await function(nd_inputs) + await function(nd_inputs, state=_copy_state(state)) # Callbacks arrive asynchronously, so intervals from the final runs may # still be open. Block (off the event loop) until they close: flipping to @@ -1861,6 +1870,7 @@ async def benchmark_coreai_program( # noqa: PLR0913 num_runs: int = 1, excluded_operations: tuple[str, ...] | None = None, specialization_options: SpecializationOptions | None = None, + state: Mapping[str, Any] | None = None, ) -> BenchmarkResult: """ Benchmark an AIProgram with profiling. @@ -1873,6 +1883,7 @@ async def benchmark_coreai_program( # noqa: PLR0913 excluded_operations: Tuple of operation names to exclude from timing (default: ("coreai.graph", "coreai.constant")) specialization_options: Options for configuring model specialization + state: Starting value of each stateful buffer, by name Returns: BenchmarkResult containing timing information for all operations @@ -1886,4 +1897,4 @@ async def benchmark_coreai_program( # noqa: PLR0913 benchmarker = CoreAIBenchmarker( coreai_program, entry_point, excluded_operations, specialization_options ) - return await benchmarker.benchmark(inputs, num_runs) + return await benchmarker.benchmark(inputs, num_runs, state) diff --git a/coreai_torch/debugging/timing_diff.py b/coreai_torch/debugging/timing_diff.py index 7af0053..949102d 100644 --- a/coreai_torch/debugging/timing_diff.py +++ b/coreai_torch/debugging/timing_diff.py @@ -11,6 +11,7 @@ import logging from collections import defaultdict +from collections.abc import Mapping from dataclasses import dataclass, field from enum import Enum from typing import Any, TextIO @@ -623,6 +624,7 @@ async def compare_runs( excluded_operations: tuple[str, ...] | None = None, specialization_options: SpecializationOptions | None = None, weights: WeightPolicy = WeightPolicy.IGNORE, + state: Mapping[str, Any] | None = None, ) -> TimingDiff: """ Benchmark both programs under identical settings and compare them. @@ -650,6 +652,7 @@ async def compare_runs( otherwise unchanged dispatch on one side. specialization_options: Specialization options. Applied to both runs. weights: Whether parameter values count towards an operation's identity. + state: Starting state for both, by name. Returns: What changed between the two runs. @@ -672,7 +675,7 @@ async def compare_runs( benchmarker = CoreAIBenchmarker( program, entry_point, excluded_operations, specialization_options ) - results.append(await benchmarker.benchmark(inputs, num_runs)) + results.append(await benchmarker.benchmark(inputs, num_runs, state)) for label, result in zip(("before", "after"), results, strict=True): if not result.operation_timings: diff --git a/coreai_torch/debugging/utils.py b/coreai_torch/debugging/utils.py index 78b3cb4..6b34ad8 100644 --- a/coreai_torch/debugging/utils.py +++ b/coreai_torch/debugging/utils.py @@ -22,9 +22,10 @@ from typing import Any import coreai._compiler._mlir_libs._coreaiIR._bindings.mlir as _mlir +import numpy as np from coreai._compiler.ir import Operation from coreai.authoring import AIProgram -from coreai.runtime import SpecializationOptions +from coreai.runtime import NDArray, SpecializationOptions def _with_debug( @@ -45,6 +46,14 @@ def _with_debug( return specialization_options.with_debug(enabled=True) +def _copy_state(state: Mapping[str, Any] | None) -> dict[str, NDArray]: + """Fresh runtime arrays holding *state*, since a run updates its state in place.""" + return { + name: NDArray(np.array(value.numpy() if isinstance(value, NDArray) else value)) + for name, value in (state or {}).items() + } + + def _state_kwargs(state: Any) -> dict[str, Any]: """``state`` as a keyword only when given, so inspectors without it still work.""" return {} if state is None else {"state": state} diff --git a/tests/debugging/test_benchmarker.py b/tests/debugging/test_benchmarker.py index b8933e3..b691a52 100644 --- a/tests/debugging/test_benchmarker.py +++ b/tests/debugging/test_benchmarker.py @@ -334,3 +334,37 @@ async def test_a_benchmark_that_times_nothing_says_what_it_measured( f"unattributed_samples={result.unattributed_samples}. An empty benchmark with " "no explanation cannot be distinguished from a model that cost nothing." ) + + +class _Accumulator(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("acc", torch.zeros(1, 4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = self.acc + x + self.acc.copy_(y * 2) + return torch.relu(y) + + +@pytest.mark.skipif(sys.platform != "darwin", reason="Test only runs on macOS") +async def test_a_stateful_program_benchmarks_from_the_given_state() -> None: + """Every run starts from a fresh copy of the state, and the caller's is untouched.""" + import numpy as np + + x = torch.randn(1, 4) + exported = torch.export.export(_Accumulator().eval(), (x,)).run_decompositions() + converter = TorchConverter() + converter.add_exported_program(exported, entrypoint_name="main") + start = np.ones((1, 4), dtype=np.float32) + + with pytest.raises(RuntimeError, match="Missing state view for b_acc"): + await benchmark_coreai_program( + converter.to_coreai(), {"x": x.numpy()}, num_runs=2 + ) + + await benchmark_coreai_program( + converter.to_coreai(), {"x": x.numpy()}, num_runs=2, state={"b_acc": start} + ) + + np.testing.assert_array_equal(start, np.ones((1, 4))) From 0d28d614c3537d43c2d35af11dc574526e54ba8f Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Wed, 21 Oct 2026 18:03:04 -0700 Subject: [PATCH 05/11] Read module outputs from the torch graph and fused locations Add module_output_op_ids: which Core AI operations hold each module instance's output. A module's output is a torch node that something outside the module reads; the operation holding it is the one realizing that node (the largest torch id in its location). The module is left out rather than guessed when the value is not certain: no or several realizing operations, a multi-output node, or a reader that left no operation of its own (batch norm folded into a conv). Read stack traces through fused locations. get_stack_trace returns nothing for a fused location, which core-elide-broadcast builds for an nn.Linear's matmul, so those operations had no module. operation_stack_traces returns every stack on an operation, and operation_stack_trace the one it belongs to; module paths and the benchmarker use the latter. Record the FX graph while converting (TorchConverter.torch_graph) so a module's output node can be found, which the Core AI program alone cannot say. Normalise OpResult.owner to a plain Operation before the isinstance checks. Some Core AI builds return the typed op view, which failed them, so no location or output map was recorded at all. --- coreai_torch/_debug_locations.py | 60 +++++++++- coreai_torch/converter.py | 11 +- coreai_torch/debugging/benchmarker.py | 7 +- coreai_torch/debugging/modules.py | 127 +++++++++++++++++++- coreai_torch/debugging/source_annotator.py | 116 +++++++++++++++++- tests/debugging/test_modules.py | 129 +++++++++++++++++++++ tests/test_debug_locations.py | 3 +- 7 files changed, 437 insertions(+), 16 deletions(-) diff --git a/coreai_torch/_debug_locations.py b/coreai_torch/_debug_locations.py index 7e4f5dc..2a7bcdc 100644 --- a/coreai_torch/_debug_locations.py +++ b/coreai_torch/_debug_locations.py @@ -126,6 +126,19 @@ def create_torch_mapping( return cls(source=source_value, target=target_value) +@dataclass(frozen=True) +class TorchNode: + """One FX node as the converter saw it, for reasoning about module boundaries.""" + + id: int # Torch operation ID, as in a Core AI operation's `torch` sources + op: str # FX node kind: "call_function", "placeholder", "output", ... + target: str # Called function; empty unless `op` is "call_function" + stack: tuple[str, ...] # Module instance path, outermost first + inputs: tuple[int, ...] # IDs of the nodes it reads + users: tuple[int, ...] # IDs of the nodes that read it + multi_output: bool # Produces several values, so one ID is not one value + + @dataclass class DebugInfo: """Comprehensive debug information for an operation. @@ -610,6 +623,27 @@ def _get_scope(operation: Operation) -> Optional[Attribute]: return None +def _as_operation(owner: object) -> Operation | None: + """The plain ``Operation`` behind a value's owner. + + ``OpResult.owner`` is a plain ``Operation`` in some Core AI builds and the typed + op view (``ReshapeOp``, ...) in others, and the view fails ``isinstance(..., + Operation)``. Attribution then silently found nothing: no operation reached the + debug info map, every location fell back to the empty one, and no output maps were + recorded. ``.operation`` is the same object on a plain ``Operation``. + + Args: + owner: An ``OpResult.owner`` or an operand's owner. + + Returns: + The operation, or None when *owner* is not one (a ``Block``, for a block + argument). + + """ + operation = getattr(owner, "operation", None) + return operation if isinstance(operation, Operation) else None + + def _get_parent_scope(operation: Operation) -> Optional[Attribute]: """Get SymbolRefAttr from parent operation with symbol table trait or return None.""" current = operation.parent @@ -752,6 +786,7 @@ def __init__( self.config = config self.module_registry = _ModuleInstanceRegistry() self._source_operation_id = 0 + self._torch_nodes: dict[int, tuple[fx.Node, tuple[str, ...]]] = {} self._debug_info_map = _DebugInfoMap() self._current_module: Module | None = None self._current_graph: Operation | None = None @@ -1015,8 +1050,8 @@ def update_output_maps( count, logs the mismatch but continues without adding output maps. """ for i, target_result in enumerate(op_results): - target_op = target_result.owner - if not isinstance(target_op, Operation): + target_op = _as_operation(target_result.owner) + if target_op is None: continue target_debug_info = self._debug_info_map.get(target_op) @@ -1166,7 +1201,7 @@ def should_process_op(op) -> bool: operations = [] seen_ops = set() for result in self._op_results: - op = result.owner + op = _as_operation(result.owner) while should_process_op(op): seen_ops.add(op) operations.append(op) @@ -1185,7 +1220,7 @@ def should_process_op(op) -> bool: queue = list(operations) while queue: for operand in queue.pop().operands: - producer = operand.owner # a Block, for a block argument + producer = _as_operation(operand.owner) # None for a block argument if should_process_op(producer): seen_ops.add(producer) operations.append(producer) @@ -1268,6 +1303,7 @@ def record_operation(self: Self, node: fx.Node): source_op_id = self._source_operation_id self._source_operation_id += 1 + self._torch_nodes[source_op_id] = (node, tuple(reversed(debug_info.call_stack))) # Create location from debug info location = self._create_location_for_node(debug_info) @@ -1284,6 +1320,22 @@ def record_operation(self: Self, node: fx.Node): self._op_results = None + def torch_graph(self: Self) -> dict[int, TorchNode]: + """The FX nodes recorded so far, by torch operation ID.""" + ids = {node: node_id for node_id, (node, _) in self._torch_nodes.items()} + return { + node_id: TorchNode( + id=node_id, + op=node.op, + target=str(node.target) if node.op == "call_function" else "", + stack=stack, + inputs=tuple(ids[n] for n in node.all_input_nodes if n in ids), + users=tuple(ids[n] for n in node.users if n in ids), + multi_output=isinstance(node.meta.get("val"), (tuple, list)), + ) + for node_id, (node, stack) in self._torch_nodes.items() + } + def finalize_node_operations(self: Self) -> None: """Materialize locations for the operations lowered into the current graph. diff --git a/coreai_torch/converter.py b/coreai_torch/converter.py index a45fabd..fc7e7b1 100644 --- a/coreai_torch/converter.py +++ b/coreai_torch/converter.py @@ -39,7 +39,7 @@ from ._composite_declaration import generate_composite_decl from ._compression.utils import inject_subbyte_tensors from ._custom_to_core import _custom_to_core_resolver -from ._debug_locations import _DebugInfoRecorder +from ._debug_locations import TorchNode, _DebugInfoRecorder from ._torch_metal_kernel import TorchMetalKernel from ._utils import ( _EXTERNALIZE_NAMESPACE, @@ -930,6 +930,15 @@ def _get_graph_op( return graph_op + @property + def torch_graph(self) -> dict[int, TorchNode]: + """The FX nodes converted so far, by the torch operation ID Core AI reports. + + Needed to say which node is a module's output, which the Core AI program alone + cannot: nodes the compiler folds away leave no location behind. + """ + return self._debug_info_recorder.torch_graph() + def to_coreai(self, *, entrypoints: Sequence[str] | None = None) -> AIProgram: """Convert staged programs to a Core AI AIProgram. diff --git a/coreai_torch/debugging/benchmarker.py b/coreai_torch/debugging/benchmarker.py index c871db7..54a3213 100644 --- a/coreai_torch/debugging/benchmarker.py +++ b/coreai_torch/debugging/benchmarker.py @@ -55,7 +55,7 @@ get_operation_id, parse_debug_infos, ) -from .source_annotator import _annotate_operations +from .source_annotator import _annotate_operations, operation_stack_trace from .table_writer import ( _Column, _Row, @@ -908,10 +908,7 @@ def module_at(path: tuple[str, ...]) -> ModuleTiming | None: for timing in self.operation_timings: paths = [ - tuple( - _mlir.get_stack_trace(operation.location) # type: ignore[attr-defined] - or ("",) - ) + operation_stack_trace(operation) or ("",) for operation in timing.operations ] if not paths: diff --git a/coreai_torch/debugging/modules.py b/coreai_torch/debugging/modules.py index 8fbb4c4..e201e9b 100644 --- a/coreai_torch/debugging/modules.py +++ b/coreai_torch/debugging/modules.py @@ -9,19 +9,22 @@ from __future__ import annotations -from collections.abc import Callable, Iterable, Iterator +from collections.abc import Callable, Iterable, Iterator, Mapping from dataclasses import dataclass, field from typing import Any, Generic, TypeVar from coreai.authoring import AIProgram from typing_extensions import Self +from .._debug_locations import TorchNode +from .debug_info import _build_coreai_op_map, get_operation_id from .source_annotator import ( ModulePath, module_paths_by_op_id, + operation_torch_ids, source_locations_by_op_id, ) -from .utils import LocationInfo, split_module_frame +from .utils import STRUCTURAL_OPS, LocationInfo, split_module_frame T = TypeVar("T") @@ -269,3 +272,123 @@ def attributions_by_op_id(program: AIProgram) -> dict[int, Attribution]: ) for op_id in sorted({*paths, *lines}) } + + +_SHAPE_ONLY = frozenset( + { + "view", + "reshape", + "_unsafe_view", + "squeeze", + "unsqueeze", + "flatten", + "contiguous", + "clone", + "alias", + "detach", + } +) +"""ATen operations that return their input's values unchanged, in another shape.""" + + +def _is_shape_only(node: TorchNode) -> bool: + parts = node.target.split(".") + return node.op == "call_function" and len(parts) > 1 and parts[1] in _SHAPE_ONLY + + +def module_output_op_ids( + program: AIProgram, torch_graph: Mapping[int, TorchNode] +) -> dict[ModulePath, list[int]]: + """ + The Core AI operations that hold each module instance's output. + + A module's output is a torch node inside it that a node outside it reads, or the + graph returns. An operation realizes the torch node with the largest ID in its + location; the smaller ones are producers it absorbed, and say where it came from, + not what it computes. A node's value is therefore the operation realizing it whose + value something realizing another node reads. + + The module is left out, rather than guessed, when that value is not certain: the + node has no such operation, or several, or produces several values, or a node that + reads it left no operation of its own. The last means the reader was folded into + this operation, which then holds the reader's value instead, as batch norm folded + into a conv does. A shape-only node (``view``) with no operation of its own is + looked through to the node it reshapes, so the value returned has the producer's + shape. + + The program's output maps are not used: they list every operation that carries a + node's location and leaves its cluster, which holds for the downstream operation + that absorbed the node as much as for the one that computes it. + + Args: + program: Program to read. + torch_graph: The FX nodes behind it, from `TorchConverter.torch_graph`. + + Returns: + Module path to the sorted ids of the operations holding its output, for every + module whose output is certain. + + """ + realizing: dict[int, list[int]] = {} + ops = {} + for op_id, operation in _build_coreai_op_map(program).items(): + if operation.name in STRUCTURAL_OPS or operation.name == "coreai.constant": + continue + ids = operation_torch_ids(operation) + if ids: + ops[op_id] = operation + realizing.setdefault(max(ids), []).append(op_id) + realized_by = {op_id: node for node, held in realizing.items() for op_id in held} + + def read_outside(op_id: int) -> bool: + """Whether something realizing another node, or the program, reads the op.""" + return any( + user.name == "coreai.output" + or realized_by.get(get_operation_id(user)) != realized_by[op_id] + for result in ops[op_id].results + for user in (use.owner for use in result.uses) + ) + + def accounted(node_id: int) -> bool: + """Whether a reader of a value has its own operation, so it did not fold in.""" + node = torch_graph[node_id] + if node.op == "output" or realizing.get(node_id): + return True + return _is_shape_only(node) and all(accounted(user) for user in node.users) + + def value_of(node: TorchNode) -> int | None: + """The Core AI operation holding *node*'s value, None when not certain.""" + while True: + if node.multi_output: + return None + holders = [op for op in realizing.get(node.id, []) if read_outside(op)] + if holders: + if len(holders) != 1 or not all(map(accounted, node.users)): + return None + return holders[0] + # No operation of its own: only a shape-only node still has its input's values. + if not _is_shape_only(node) or len(node.inputs) != 1: + return None + node = torch_graph[node.inputs[0]] + + def inside(node: TorchNode, module: ModulePath) -> bool: + return node.op == "call_function" and node.stack[: len(module)] == module + + nodes = [node for node in torch_graph.values() if node.op == "call_function"] + modules = { + node.stack[:depth] + for node in nodes + for depth in range(1, len(node.stack) + 1) + } + outputs: dict[ModulePath, list[int]] = {} + for module in sorted(modules): + boundary = [ + node + for node in nodes + if inside(node, module) + and any(not inside(torch_graph[user], module) for user in node.users) + ] + values = [value_of(node) for node in boundary] + if boundary and None not in values: + outputs[module] = sorted(set(values)) # type: ignore[arg-type] + return outputs diff --git a/coreai_torch/debugging/source_annotator.py b/coreai_torch/debugging/source_annotator.py index ab718ef..54b5d61 100644 --- a/coreai_torch/debugging/source_annotator.py +++ b/coreai_torch/debugging/source_annotator.py @@ -248,6 +248,117 @@ def _trim_excluded_leaf_modules( return module[: max(0, len(module) - excluded_leaves)] +def _leaf_locations(location: Any) -> Iterable[Any]: + """ + Flatten a (possibly nested) fused location into its non-fused locations. + + Args: + location: Location to flatten. + + Yields: + Each non-fused location; *location* itself when it is not fused. + + """ + # Some builds answer ``is_a_fused()`` on every Location; others return a + # ``FusedLoc`` that alone has ``locations``. + is_a_fused = getattr(location, "is_a_fused", None) + fused = is_a_fused() if is_a_fused else hasattr(location, "locations") + if fused: + for child in location.locations: + yield from _leaf_locations(child) + else: + yield location + + +def _leaf_stack_traces(operation: Operation) -> list[tuple[tuple[str, ...], int]]: + """ + Read each stack trace on an operation's location, with the torch node it is for. + + Args: + operation: Operation to read. + + Returns: + ``(stack trace, torch operation ID)`` per location that has a stack trace, in + location order. The ID is -1 when the location names no torch node. + + """ + found = [] + for leaf in _leaf_locations(operation.location): + stack = tuple(_mlir.get_stack_trace(leaf) or ()) # type: ignore[attr-defined] + if not stack: + continue + torch_id = max( + ( + source.id + for source in _mlir.get_source_info(leaf) # type: ignore[attr-defined] + if source.name == "torch" and source.id is not None + ), + default=-1, + ) + found.append((stack, torch_id)) + return found + + +def operation_torch_ids(operation: Operation) -> set[int]: + """ + The torch operation IDs recorded on an operation's location. + + Args: + operation: Operation to read. + + Returns: + The ID of every torch node the operation came from, reading through fused + locations. The largest is the node the operation itself realizes; the rest are + producers it absorbed. + + """ + return { + source.id + for leaf in _leaf_locations(operation.location) + for source in _mlir.get_source_info(leaf) # type: ignore[attr-defined] + if source.name == "torch" and source.id is not None + } + + +def operation_stack_traces(operation: Operation) -> list[tuple[str, ...]]: + """ + Every stack trace on an operation, outermost frame first. + + Reads through fused locations, so an operation that replaced several has the stack + of each. + + Args: + operation: Operation to read. + + Returns: + The distinct stack traces, in location order. Empty when none is recorded. + + """ + return list(dict.fromkeys(stack for stack, _ in _leaf_stack_traces(operation))) + + +def operation_stack_trace(operation: Operation) -> tuple[str, ...]: + """ + The stack trace of the module an operation belongs to, outermost frame first. + + Of an operation's stack traces, that of the latest torch node (highest operation + ID) is its own; the earlier ones are the producers it absorbed. + + Args: + operation: Operation to read. + + Returns: + The stack trace, empty when none is recorded. + + """ + best: tuple[str, ...] = () + best_id = -1 + for stack, torch_id in _leaf_stack_traces(operation): + if not best or torch_id > best_id: + best, best_id = stack, torch_id + return best + + def _get_module_path( operation: Operation, exclude: Callable[[LocationInfo], bool], @@ -269,7 +380,7 @@ def _get_module_path( their own annotatable file. Empty when there is no stack trace. """ - stack_trace = _mlir.get_stack_trace(operation.location) # type: ignore[attr-defined] + stack_trace = operation_stack_trace(operation) if not stack_trace: return () @@ -398,8 +509,7 @@ def _operation_in_module(operation: Operation, module: ModulePath) -> bool: True if the operation belongs to ``module`` (or one of its children). """ - stack_trace = tuple(_mlir.get_stack_trace(operation.location) or ()) # type: ignore[attr-defined] - return stack_trace[: len(module)] == module + return operation_stack_trace(operation)[: len(module)] == module def _should_exclude_location(location: LocationInfo) -> bool: diff --git a/tests/debugging/test_modules.py b/tests/debugging/test_modules.py index 97cf619..bc2fd1a 100644 --- a/tests/debugging/test_modules.py +++ b/tests/debugging/test_modules.py @@ -15,11 +15,25 @@ import json +import pytest +import torch +from coreai.authoring import AIProgram + +from coreai_torch.converter import TorchConverter +from coreai_torch.debugging.debug_info import _build_coreai_op_map from coreai_torch.debugging.modules import ( UNATTRIBUTED, ModuleNode, build_module_tree, + module_output_op_ids, +) +from coreai_torch.debugging.source_annotator import ( + module_paths_by_op_id, + operation_stack_traces, ) +from coreai_torch.debugging.utils import STRUCTURAL_OPS + +from .test_model import HierarchicalModel, get_example_inputs def test_nests_by_path() -> None: @@ -134,3 +148,118 @@ def test_to_dict_is_json_serialisable() -> None: assert plain["count"] == 2 assert plain["items"] == [] assert plain["children"][0]["items"] == [1, 2] + + +@pytest.fixture +def hierarchical_converter() -> TorchConverter: + """A converter that has converted the hierarchical model, with stack traces.""" + exported = torch.export.export( + HierarchicalModel().eval(), + args=tuple(get_example_inputs(HierarchicalModel).values()), + ).run_decompositions() + converter = TorchConverter() + converter.add_exported_program(exported, entrypoint_name="main") + converter.program = converter.to_coreai() + return converter + + +@pytest.fixture +def hierarchical_program(hierarchical_converter: TorchConverter) -> AIProgram: + """The Core AI program of the hierarchical model.""" + return hierarchical_converter.program + + +def test_module_output_ops_are_where_values_leave( + hierarchical_converter: TorchConverter, +) -> None: + """A module's output is the operation realizing a node a node outside reads.""" + root = ("HierarchicalModel$1",) + outputs = module_output_op_ids( + hierarchical_converter.program, hierarchical_converter.torch_graph + ) + assert outputs == { + root: [24], + (*root, "Linear$1"): [11], + (*root, "SubModel$1"): [15], + (*root, "SubModel$1", "Linear$2"): [14], + (*root, "SubModel$2"): [20], + (*root, "SubModel$2", "Linear$3"): [19], + (*root, "Linear$4"): [23], + } + + +def test_output_is_the_operation_realizing_the_node( + hierarchical_converter: TorchConverter, +) -> None: + """An operation that absorbed a producer does not hold that producer's value.""" + outputs = module_output_op_ids( + hierarchical_converter.program, hierarchical_converter.torch_graph + ) + + # Op 13 is Linear$2's matmul and also carries Linear$1's location, and the last + # sigmoid carries Linear$4's. Neither holds the value of the layer it absorbed. + assert outputs[("HierarchicalModel$1", "Linear$1")] == [11] + assert outputs[("HierarchicalModel$1", "Linear$4")] == [23] + + +def test_torch_graph_records_modules_and_edges( + hierarchical_converter: TorchConverter, +) -> None: + """Each FX node is recorded with its module path and what reads it.""" + graph = hierarchical_converter.torch_graph + addmm = next(n for n in graph.values() if n.target == "aten.addmm.default") + + assert addmm.stack == ("HierarchicalModel$1", "Linear$1") + assert addmm.users + assert all(user in graph for user in addmm.users) + + +def test_every_operation_has_a_stack_trace(hierarchical_program: AIProgram) -> None: + """No operation of a converted model is left without a module.""" + paths = module_paths_by_op_id(hierarchical_program) + ops = _build_coreai_op_map(hierarchical_program) + missing = [ + op_id + for op_id, op in ops.items() + if op.name not in STRUCTURAL_OPS and op.name != "coreai.constant" + if op_id not in paths + ] + assert missing == [] + + +def test_fused_location_gives_every_stack(hierarchical_program: AIProgram) -> None: + """A matmul rebuilt from the previous layer's output is in both layers' stacks.""" + ops = _build_coreai_op_map(hierarchical_program) + stacks = operation_stack_traces(ops[13]) + + assert ("HierarchicalModel$1", "SubModel$1", "Linear$2") in stacks + assert ("HierarchicalModel$1", "Linear$1") in stacks + assert len(stacks) == len(set(stacks)) + + +class _ConvBatchNorm(torch.nn.Module): + """A conv whose batch norm the compiler folds into its weights.""" + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3) + self.bn = torch.nn.BatchNorm2d(4) + self.head = torch.nn.Conv2d(4, 4, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.head(torch.relu(self.bn(self.conv(x)))) + + +def test_module_whose_reader_was_folded_in_is_left_out() -> None: + """A conv's op holds the batch norm's value once folded, not the conv's.""" + exported = torch.export.export( + _ConvBatchNorm().eval(), (torch.randn(1, 3, 8, 8),) + ).run_decompositions() + converter = TorchConverter() + converter.add_exported_program(exported, entrypoint_name="main") + outputs = module_output_op_ids(converter.to_coreai(), converter.torch_graph) + + root = ("_ConvBatchNorm$1",) + assert (*root, "Conv2d$1") not in outputs + assert (*root, "BatchNorm2d$1") not in outputs + assert (*root, "Conv2d$2") in outputs diff --git a/tests/test_debug_locations.py b/tests/test_debug_locations.py index 132f678..ae9a592 100644 --- a/tests/test_debug_locations.py +++ b/tests/test_debug_locations.py @@ -15,6 +15,7 @@ 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 coreai_torch.debugging.source_annotator import operation_stack_trace from .debugging.test_model import HierarchicalModel @@ -162,7 +163,7 @@ def test_intermediate_ops_of_a_lowering_keep_their_attribution() -> None: 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] + stack_trace = operation_stack_trace(operation) 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. From b872ec4cf7502839209bde1be92bab71c9543c1e Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Wed, 21 Oct 2026 18:07:29 -0700 Subject: [PATCH 06/11] Add get_module_outputs: read one module's output through an inspector Take a module path, resolve the Core AI operations that certainly hold its output with module_output_op_ids, and return their values from a CoreAIProbedInspector or a CoreAIInspector. Raise ModuleOutputUnknownError when the module is absent or its output is not certain, rather than guess. --- coreai_torch/debugging/inspector.py | 59 +++++++++++++++++++++ tests/debugging/test_modules.py | 81 +++++++++++++++++++++++++++++ 2 files changed, 140 insertions(+) diff --git a/coreai_torch/debugging/inspector.py b/coreai_torch/debugging/inspector.py index 291a67d..c9caa09 100644 --- a/coreai_torch/debugging/inspector.py +++ b/coreai_torch/debugging/inspector.py @@ -23,18 +23,22 @@ from coreai.authoring import AIProgram from numpy.typing import NDArray +from .._debug_locations import TorchNode + from .debug_info import ( DebugInfo, DebugInfoRecord, _build_compile_id_to_coreai_map, parse_debug_infos, ) +from .modules import module_output_op_ids from .probes import ( Perturbation, ProbedProgram, add_probe_outputs, compare_to_reference, ) +from .source_annotator import ModulePath from .torch_utils import _TorchFXNodeValueInterpreter from .utils import _state_kwargs @@ -1536,3 +1540,58 @@ def last_outputs(self) -> dict[str, NDArray[Any]] | None: def last_state(self) -> dict[str, NDArray[Any]] | None: """The state the latest unprobed run left behind, or None before one.""" return self._last_state + + +class ModuleOutputUnknownError(LookupError): + """A module's output cannot be read from the program with certainty.""" + + +async def get_module_outputs( + inspector: Inspector, + program: AIProgram, + torch_graph: Mapping[int, TorchNode], + module: ModulePath, + inputs: Any, + state: Mapping[str, Any] | None = None, +) -> list[NDArray[Any]]: + """ + Run a model and return the output of one module instance. + + Args: + inspector: Inspector over *program*, a `CoreAIInspector` or a + `CoreAIProbedInspector`. + program: The program the inspector runs. + torch_graph: The FX nodes behind it, from `TorchConverter.torch_graph`. + module: Module instance path, outermost first, as in ``("Net$1", "Block$2")``. + inputs: Model inputs, as the inspector takes them. + state: Starting value of each state, by name. + + Returns: + One array per value the module returns, ordered by Core AI operation ID. A + value has its producer's shape when the module reshapes it on the way out. + + Raises: + ModuleOutputUnknownError: If the module is not in the graph, or the compiler + left no operation that certainly holds its output (see + `module_output_op_ids`). + RuntimeError: If the inspector returns nothing for one of its operations. + + """ + if not any(node.stack[: len(module)] == module for node in torch_graph.values()): + msg = f"no module {'/'.join(module)!r} in the graph" + raise ModuleOutputUnknownError(msg) + op_ids = module_output_op_ids(program, torch_graph).get(module) + if not op_ids: + msg = f"the output of {'/'.join(module)!r} is not certain in this program" + raise ModuleOutputUnknownError(msg) + + values = await inspector.get_intermediates_for_ops(op_ids, inputs, state) + + arrays: list[NDArray[Any]] = [] + for op_id in op_ids: + found = values.get(op_id) + if not found or any(array is None for array in found): + msg = f"the inspector returned no value for operation {op_id}" + raise RuntimeError(msg) + arrays.extend(np.asarray(array) for array in found) + return arrays diff --git a/tests/debugging/test_modules.py b/tests/debugging/test_modules.py index bc2fd1a..105a2f4 100644 --- a/tests/debugging/test_modules.py +++ b/tests/debugging/test_modules.py @@ -15,12 +15,19 @@ import json +import numpy as np import pytest import torch from coreai.authoring import AIProgram +from coreai.runtime import SpecializationOptions from coreai_torch.converter import TorchConverter from coreai_torch.debugging.debug_info import _build_coreai_op_map +from coreai_torch.debugging.inspector import ( + CoreAIProbedInspector, + ModuleOutputUnknownError, + get_module_outputs, +) from coreai_torch.debugging.modules import ( UNATTRIBUTED, ModuleNode, @@ -263,3 +270,77 @@ def test_module_whose_reader_was_folded_in_is_left_out() -> None: assert (*root, "Conv2d$1") not in outputs assert (*root, "BatchNorm2d$1") not in outputs assert (*root, "Conv2d$2") in outputs + + +async def test_get_module_outputs_matches_torch( + hierarchical_converter: TorchConverter, + specialization_options: SpecializationOptions | None, +) -> None: + """Probing a module's output gives the value torch's own module returns.""" + torch.manual_seed(0) + model = HierarchicalModel().eval() + x = get_example_inputs(HierarchicalModel)["x"] + exported = torch.export.export(model, (x,)).run_decompositions() + converter = TorchConverter() + converter.add_exported_program(exported, entrypoint_name="main") + program = converter.to_coreai() + + seen: list[torch.Tensor] = [] + handle = model.sub_model1.register_forward_hook( + lambda _m, _i, out: seen.append(out.detach()) + ) + model(x) + handle.remove() + + inspector = CoreAIProbedInspector( + program, specialization_options=specialization_options + ) + (value,) = await get_module_outputs( + inspector, + program, + converter.torch_graph, + ("HierarchicalModel$1", "SubModel$1"), + {"x": x.numpy()}, + ) + + np.testing.assert_allclose( + value.reshape(-1), seen[0].numpy().reshape(-1), rtol=1e-3, atol=1e-3 + ) + + +async def test_get_module_outputs_refuses_what_it_cannot_place( + hierarchical_converter: TorchConverter, + specialization_options: SpecializationOptions | None, +) -> None: + """A module that is absent or folded away is an error, never a guess.""" + inspector = CoreAIProbedInspector( + hierarchical_converter.program, specialization_options=specialization_options + ) + graph = hierarchical_converter.torch_graph + x = get_example_inputs(HierarchicalModel)["x"].numpy() + + with pytest.raises(ModuleOutputUnknownError, match="no module"): + await get_module_outputs( + inspector, + hierarchical_converter.program, + graph, + ("HierarchicalModel$1", "Nope$1"), + {"x": x}, + ) + + exported = torch.export.export( + _ConvBatchNorm().eval(), (torch.randn(1, 3, 8, 8),) + ).run_decompositions() + folded = TorchConverter() + folded.add_exported_program(exported, entrypoint_name="main") + program = folded.to_coreai() + with pytest.raises(ModuleOutputUnknownError, match="not certain"): + await get_module_outputs( + CoreAIProbedInspector( + program, specialization_options=specialization_options + ), + program, + folded.torch_graph, + ("_ConvBatchNorm$1", "Conv2d$1"), + {"x": np.zeros((1, 3, 8, 8), dtype=np.float32)}, + ) From 084fcef46d032c8d9dcd7d4451a5e349591511be Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Thu, 22 Oct 2026 13:28:11 -0700 Subject: [PATCH 07/11] Read module outputs from the program alone module_output_op_ids and get_module_outputs no longer need the converter or the FX graph. The converter now records, on each graph's location, a `source_nodes` metadata: the ATen target of every node converted into the graph and the id of its output node. It survives compilation, is saved with the asset and comes back out of a loaded model. An operation is owned by the largest torch id in its location that is not shape-only. A module's output is an owned operation that an operation of another module reads, or that returns the program's output. The module is left out when an id between the operation's owner and its reader's owner shows up in no operation and is not known to be shape-only: a node was folded away without a trace (batch norm into a conv), or its kind is unknown. Without the table no id is known to be shape-only and only the root may take the program's output, so fewer modules are reported; with no debug info none. --- coreai_torch/_debug_locations.py | 84 ++++++---- coreai_torch/converter.py | 11 +- coreai_torch/debugging/inspector.py | 27 ++-- coreai_torch/debugging/modules.py | 233 +++++++++++++++++++--------- tests/debugging/test_modules.py | 105 ++++++------- 5 files changed, 278 insertions(+), 182 deletions(-) diff --git a/coreai_torch/_debug_locations.py b/coreai_torch/_debug_locations.py index 2a7bcdc..25d36c6 100644 --- a/coreai_torch/_debug_locations.py +++ b/coreai_torch/_debug_locations.py @@ -126,17 +126,29 @@ def create_torch_mapping( return cls(source=source_value, target=target_value) -@dataclass(frozen=True) -class TorchNode: - """One FX node as the converter saw it, for reasoning about module boundaries.""" +NODE_TABLE_METADATA = "source_nodes" +"""Name of the metadata on a graph's location that lists the source nodes converted into it.""" + +NODE_TABLE_VERSION = 1 + + +class _NodeTable: + """The torch nodes converted into one graph: what each was, and where the graph ends. + + Nodes the compiler folds away or reshapes leave nothing in the program, so their + kind cannot be read from an operation. This is what says that an id with no + operation was a ``view`` and not, say, a batch norm. + """ + + def __init__(self) -> None: + self.targets: dict[int, str] = {} + self.output_id: int | None = None - id: int # Torch operation ID, as in a Core AI operation's `torch` sources - op: str # FX node kind: "call_function", "placeholder", "output", ... - target: str # Called function; empty unless `op` is "call_function" - stack: tuple[str, ...] # Module instance path, outermost first - inputs: tuple[int, ...] # IDs of the nodes it reads - users: tuple[int, ...] # IDs of the nodes that read it - multi_output: bool # Produces several values, so one ID is not one value + def add(self, node_id: int, node: fx.Node) -> None: + if node.op == "call_function": + self.targets[node_id] = str(node.target) + elif node.op == "output": + self.output_id = node_id @dataclass @@ -786,7 +798,7 @@ def __init__( self.config = config self.module_registry = _ModuleInstanceRegistry() self._source_operation_id = 0 - self._torch_nodes: dict[int, tuple[fx.Node, tuple[str, ...]]] = {} + self._graph_tables: list[_NodeTable] = [] self._debug_info_map = _DebugInfoMap() self._current_module: Module | None = None self._current_graph: Operation | None = None @@ -927,11 +939,14 @@ def _get_files_for_graph( return graph_files - def _set_graph_location(self: Self, graph_operation: Operation) -> None: + def _set_graph_location( + self: Self, graph_operation: Operation, table: _NodeTable + ) -> None: """Set debug location for a graph operation. Args: graph_operation: The graph operation to set location for + table: The torch nodes converted into the graph, recorded on its location """ context = graph_operation.context @@ -980,6 +995,7 @@ def _set_graph_location(self: Self, graph_operation: Operation) -> None: debug_location = location_attr( src=src_location, scope=subprogram, + metadata=self._node_table_metadata(table, context), context=context, ) @@ -1125,10 +1141,13 @@ def record_graph(self, graph_operation: Operation): # Store current graph operation prev_graph_operation = self._current_graph self._current_graph = graph_operation + table = _NodeTable() + self._graph_tables.append(table) yield graph_operation + self._graph_tables.pop() # Set graph location first so it can be reused by block arguments - self._set_graph_location(graph_operation) + self._set_graph_location(graph_operation, table) # Ensure ALL operations in the graph have debug locations self._ensure_all_operations_have_debug_locations(graph_operation) @@ -1303,7 +1322,8 @@ def record_operation(self: Self, node: fx.Node): source_op_id = self._source_operation_id self._source_operation_id += 1 - self._torch_nodes[source_op_id] = (node, tuple(reversed(debug_info.call_stack))) + if self._graph_tables: + self._graph_tables[-1].add(source_op_id, node) # Create location from debug info location = self._create_location_for_node(debug_info) @@ -1320,21 +1340,29 @@ def record_operation(self: Self, node: fx.Node): self._op_results = None - def torch_graph(self: Self) -> dict[int, TorchNode]: - """The FX nodes recorded so far, by torch operation ID.""" - ids = {node: node_id for node_id, (node, _) in self._torch_nodes.items()} - return { - node_id: TorchNode( - id=node_id, - op=node.op, - target=str(node.target) if node.op == "call_function" else "", - stack=stack, - inputs=tuple(ids[n] for n in node.all_input_nodes if n in ids), - users=tuple(ids[n] for n in node.users if n in ids), - multi_output=isinstance(node.meta.get("val"), (tuple, list)), - ) - for node_id, (node, stack) in self._torch_nodes.items() + def _node_table_metadata( + self: Self, table: _NodeTable, context: Context + ) -> list[metadata_attr]: + """The graph location's metadata listing the torch nodes, ``[]`` when none.""" + if not table.targets: + return [] + entries = { + str(node_id): StringAttr.get(target, context=context) + for node_id, target in table.targets.items() + } + data = { + "version": _create_uint64_attr(NODE_TABLE_VERSION, context), + "nodes": DictAttr.get(entries, context=context), } + if table.output_id is not None: + data["output_id"] = _create_uint64_attr(table.output_id, context) + return [ + metadata_attr( + name=NODE_TABLE_METADATA, + data=DictAttr.get(data, context=context), + context=context, + ) + ] def finalize_node_operations(self: Self) -> None: """Materialize locations for the operations lowered into the current graph. diff --git a/coreai_torch/converter.py b/coreai_torch/converter.py index fc7e7b1..a45fabd 100644 --- a/coreai_torch/converter.py +++ b/coreai_torch/converter.py @@ -39,7 +39,7 @@ from ._composite_declaration import generate_composite_decl from ._compression.utils import inject_subbyte_tensors from ._custom_to_core import _custom_to_core_resolver -from ._debug_locations import TorchNode, _DebugInfoRecorder +from ._debug_locations import _DebugInfoRecorder from ._torch_metal_kernel import TorchMetalKernel from ._utils import ( _EXTERNALIZE_NAMESPACE, @@ -930,15 +930,6 @@ def _get_graph_op( return graph_op - @property - def torch_graph(self) -> dict[int, TorchNode]: - """The FX nodes converted so far, by the torch operation ID Core AI reports. - - Needed to say which node is a module's output, which the Core AI program alone - cannot: nodes the compiler folds away leave no location behind. - """ - return self._debug_info_recorder.torch_graph() - def to_coreai(self, *, entrypoints: Sequence[str] | None = None) -> AIProgram: """Convert staged programs to a Core AI AIProgram. diff --git a/coreai_torch/debugging/inspector.py b/coreai_torch/debugging/inspector.py index c9caa09..5396062 100644 --- a/coreai_torch/debugging/inspector.py +++ b/coreai_torch/debugging/inspector.py @@ -23,8 +23,6 @@ from coreai.authoring import AIProgram from numpy.typing import NDArray -from .._debug_locations import TorchNode - from .debug_info import ( DebugInfo, DebugInfoRecord, @@ -38,7 +36,7 @@ add_probe_outputs, compare_to_reference, ) -from .source_annotator import ModulePath +from .source_annotator import ModulePath, module_paths_by_op_id from .torch_utils import _TorchFXNodeValueInterpreter from .utils import _state_kwargs @@ -1549,7 +1547,6 @@ class ModuleOutputUnknownError(LookupError): async def get_module_outputs( inspector: Inspector, program: AIProgram, - torch_graph: Mapping[int, TorchNode], module: ModulePath, inputs: Any, state: Mapping[str, Any] | None = None, @@ -1560,8 +1557,7 @@ async def get_module_outputs( Args: inspector: Inspector over *program*, a `CoreAIInspector` or a `CoreAIProbedInspector`. - program: The program the inspector runs. - torch_graph: The FX nodes behind it, from `TorchConverter.torch_graph`. + program: The program the inspector runs, converted with debug info. module: Module instance path, outermost first, as in ``("Net$1", "Block$2")``. inputs: Model inputs, as the inspector takes them. state: Starting value of each state, by name. @@ -1571,18 +1567,23 @@ async def get_module_outputs( value has its producer's shape when the module reshapes it on the way out. Raises: - ModuleOutputUnknownError: If the module is not in the graph, or the compiler - left no operation that certainly holds its output (see + ModuleOutputUnknownError: If no operation belongs to the module, or the + compiler left no operation that certainly holds its output (see `module_output_op_ids`). RuntimeError: If the inspector returns nothing for one of its operations. """ - if not any(node.stack[: len(module)] == module for node in torch_graph.values()): - msg = f"no module {'/'.join(module)!r} in the graph" - raise ModuleOutputUnknownError(msg) - op_ids = module_output_op_ids(program, torch_graph).get(module) + op_ids = module_output_op_ids(program).get(module) if not op_ids: - msg = f"the output of {'/'.join(module)!r} is not certain in this program" + known = any( + path[: len(module)] == module + for path in module_paths_by_op_id(program).values() + ) + msg = ( + f"the output of {'/'.join(module)!r} is not certain in this program" + if known + else f"no operation belongs to module {'/'.join(module)!r}" + ) raise ModuleOutputUnknownError(msg) values = await inspector.get_intermediates_for_ops(op_ids, inputs, state) diff --git a/coreai_torch/debugging/modules.py b/coreai_torch/debugging/modules.py index e201e9b..e400c74 100644 --- a/coreai_torch/debugging/modules.py +++ b/coreai_torch/debugging/modules.py @@ -9,18 +9,21 @@ from __future__ import annotations -from collections.abc import Callable, Iterable, Iterator, Mapping +from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field from typing import Any, Generic, TypeVar +from coreai._compiler.ir import Attribute, DictAttr, IntegerAttr, StringAttr from coreai.authoring import AIProgram from typing_extensions import Self -from .._debug_locations import TorchNode +from .._debug_locations import NODE_TABLE_METADATA, NODE_TABLE_VERSION from .debug_info import _build_coreai_op_map, get_operation_id from .source_annotator import ( ModulePath, + _leaf_stack_traces, module_paths_by_op_id, + operation_stack_trace, operation_torch_ids, source_locations_by_op_id, ) @@ -291,104 +294,180 @@ def attributions_by_op_id(program: AIProgram) -> dict[int, Attribution]: """ATen operations that return their input's values unchanged, in another shape.""" -def _is_shape_only(node: TorchNode) -> bool: - parts = node.target.split(".") - return node.op == "call_function" and len(parts) > 1 and parts[1] in _SHAPE_ONLY +@dataclass(frozen=True) +class _NodeTable: + """What a graph's location says about the torch nodes converted into it.""" + + targets: dict[int, str] + """ATen target of each ``call_function`` node, by torch operation ID.""" + + output_id: int | None + """ID of the node that returns the graph's values, None when not recorded.""" + + def shape_only(self, node_id: int) -> bool: + """Whether the node returns its input's values unchanged, in another shape.""" + parts = self.targets.get(node_id, "").split(".") + return len(parts) > 1 and parts[1] in _SHAPE_ONLY + + +def _metadata_data(printed: str, name: str) -> str | None: + """ + The printed data of the metadata called *name* in a printed location, if any. + + A location's metadata is read from its printed form because builds differ in + what they expose of it: some give the list of metadata, others only the + attribute. The data is a brace-delimited dictionary, which MLIR parses back. + """ + marker = f'<"{name}": ' + start = printed.find(marker) + if start < 0: + return None + start += len(marker) + depth, quoted = 0, False + for index in range(start, len(printed)): + char = printed[index] + if char == '"' and printed[index - 1] != "\\": + quoted = not quoted + elif not quoted and char == "{": + depth += 1 + elif not quoted and char == "}": + depth -= 1 + if depth == 0: + return printed[start : index + 1] + return None + + +def _read_node_tables(program: AIProgram) -> dict[str, _NodeTable]: + """ + The node table recorded on each graph's location, by graph name. + A graph with no table, or one this version cannot read, is absent: nothing is + assumed about the nodes of a program that was not converted with one. + """ + tables: dict[str, _NodeTable] = {} + for operation in _build_coreai_op_map(program).values(): + if operation.name != "coreai.graph": + continue + printed = _metadata_data(str(operation.location), NODE_TABLE_METADATA) + if printed is None: + continue + with operation.context: + data = DictAttr(Attribute.parse(printed)) + if IntegerAttr(data["version"]).value != NODE_TABLE_VERSION: + continue + nodes = DictAttr(data["nodes"]) + name = StringAttr(operation.attributes["sym_name"]).value + tables[name] = _NodeTable( + targets={ + int(entry.name): StringAttr(entry.attr).value for entry in nodes + }, + output_id=( + IntegerAttr(data["output_id"]).value if "output_id" in data else None + ), + ) + return tables -def module_output_op_ids( - program: AIProgram, torch_graph: Mapping[int, TorchNode] -) -> dict[ModulePath, list[int]]: + +def module_output_op_ids(program: AIProgram) -> dict[ModulePath, list[int]]: """ The Core AI operations that hold each module instance's output. - A module's output is a torch node inside it that a node outside it reads, or the - graph returns. An operation realizes the torch node with the largest ID in its - location; the smaller ones are producers it absorbed, and say where it came from, - not what it computes. A node's value is therefore the operation realizing it whose - value something realizing another node reads. + A module's output is a value that leaves it: an operation it owns that an operation + of another module reads, or that returns the program's own output. An operation is + owned by the torch node with the largest ID in its location that is not shape-only; + the smaller ones are producers it absorbed, and say where it came from, not what it + computes. - The module is left out, rather than guessed, when that value is not certain: the - node has no such operation, or several, or produces several values, or a node that - reads it left no operation of its own. The last means the reader was folded into - this operation, which then holds the reader's value instead, as batch norm folded - into a conv does. A shape-only node (``view``) with no operation of its own is - looked through to the node it reshapes, so the value returned has the producer's - shape. + A module is left out, rather than guessed, whenever a node could have changed the + value on its way out. The torch IDs between the owner of the operation and the + owner of its reader are the nodes on the path between them, and every one must show + up in some operation's location or be shape-only. An ID that appears nowhere was + folded away without a trace, as batch norm is into a conv, and a reshape is + harmless only if the program's node table says it was one. A reader with no torch + ID, or one whose ID is not larger, leaves the module out too. - The program's output maps are not used: they list every operation that carries a - node's location and leaves its cluster, which holds for the downstream operation - that absorbed the node as much as for the one that computes it. + Without the node table, written when `coreai_torch` converts the program, no ID is + known to be shape-only and the program's output is accepted only for the root + module, so fewer modules are reported. A program with no debug info reports none. Args: - program: Program to read. - torch_graph: The FX nodes behind it, from `TorchConverter.torch_graph`. + program: Program to read, converted with debug info. Returns: Module path to the sorted ids of the operations holding its output, for every module whose output is certain. """ - realizing: dict[int, list[int]] = {} - ops = {} - for op_id, operation in _build_coreai_op_map(program).items(): - if operation.name in STRUCTURAL_OPS or operation.name == "coreai.constant": + tables = _read_node_tables(program) + targets: dict[int, str] = {} + for table in tables.values(): + targets.update(table.targets) + shape_only = _NodeTable(targets, None).shape_only + + ops = _build_coreai_op_map(program) + visible: set[int] = set() + owner_id: dict[int, int] = {} + owner_stack: dict[int, ModulePath] = {} + for op_id, operation in ops.items(): + if operation.name in STRUCTURAL_OPS: continue ids = operation_torch_ids(operation) - if ids: - ops[op_id] = operation - realizing.setdefault(max(ids), []).append(op_id) - realized_by = {op_id: node for node, held in realizing.items() for op_id in held} + visible |= ids + if operation.name == "coreai.constant" or not ids: + continue + owner_id[op_id] = max({i for i in ids if not shape_only(i)} or ids) + stacks = {node_id: stack for stack, node_id in _leaf_stack_traces(operation)} + owner_stack[op_id] = stacks.get(owner_id[op_id]) or operation_stack_trace( + operation + ) - def read_outside(op_id: int) -> bool: - """Whether something realizing another node, or the program, reads the op.""" + def hidden_between(low: int, high: int) -> bool: return any( - user.name == "coreai.output" - or realized_by.get(get_operation_id(user)) != realized_by[op_id] - for result in ops[op_id].results - for user in (use.owner for use in result.uses) + node_id not in visible and not shape_only(node_id) + for node_id in range(low + 1, high) ) - def accounted(node_id: int) -> bool: - """Whether a reader of a value has its own operation, so it did not fold in.""" - node = torch_graph[node_id] - if node.op == "output" or realizing.get(node_id): - return True - return _is_shape_only(node) and all(accounted(user) for user in node.users) - - def value_of(node: TorchNode) -> int | None: - """The Core AI operation holding *node*'s value, None when not certain.""" - while True: - if node.multi_output: - return None - holders = [op for op in realizing.get(node.id, []) if read_outside(op)] - if holders: - if len(holders) != 1 or not all(map(accounted, node.users)): - return None - return holders[0] - # No operation of its own: only a shape-only node still has its input's values. - if not _is_shape_only(node) or len(node.inputs) != 1: - return None - node = torch_graph[node.inputs[0]] - - def inside(node: TorchNode, module: ModulePath) -> bool: - return node.op == "call_function" and node.stack[: len(module)] == module - - nodes = [node for node in torch_graph.values() if node.op == "call_function"] + def in_module(stack: ModulePath, module: ModulePath) -> bool: + return stack[: len(module)] == module + modules = { - node.stack[:depth] - for node in nodes - for depth in range(1, len(node.stack) + 1) + stack[:depth] for stack in owner_stack.values() for depth in range(1, len(stack) + 1) } outputs: dict[ModulePath, list[int]] = {} for module in sorted(modules): - boundary = [ - node - for node in nodes - if inside(node, module) - and any(not inside(torch_graph[user], module) for user in node.users) - ] - values = [value_of(node) for node in boundary] - if boundary and None not in values: - outputs[module] = sorted(set(values)) # type: ignore[arg-type] + found: set[int] = set() + certain = True + for op_id, stack in owner_stack.items(): + if not in_module(stack, module): + continue + for result in ops[op_id].results: + for use in result.uses: + user = use.owner + user_id = get_operation_id(user) + if user.name == "coreai.output": + graph = StringAttr(user.parent.attributes["sym_name"]).value + end = tables[graph].output_id if graph in tables else None + # The program's own output, so nothing can have changed it unless a + # node after this module was folded in unseen. + if len(module) == 1 or ( + end is not None + and end > owner_id[op_id] + and not hidden_between(owner_id[op_id], end) + ): + found.add(op_id) + else: + certain = False + continue + if user_id not in owner_stack: + certain = False + continue + if in_module(owner_stack[user_id], module): + continue + low, high = owner_id[op_id], owner_id[user_id] + if high <= low or hidden_between(low, high): + certain = False + else: + found.add(op_id) + if certain and found: + outputs[module] = sorted(found) return outputs diff --git a/tests/debugging/test_modules.py b/tests/debugging/test_modules.py index 105a2f4..57f318b 100644 --- a/tests/debugging/test_modules.py +++ b/tests/debugging/test_modules.py @@ -21,6 +21,7 @@ from coreai.authoring import AIProgram from coreai.runtime import SpecializationOptions +from coreai_torch._debug_locations import _DebugInfoRecorder from coreai_torch.converter import TorchConverter from coreai_torch.debugging.debug_info import _build_coreai_op_map from coreai_torch.debugging.inspector import ( @@ -32,6 +33,7 @@ UNATTRIBUTED, ModuleNode, build_module_tree, + _read_node_tables, module_output_op_ids, ) from coreai_torch.debugging.source_annotator import ( @@ -157,34 +159,27 @@ def test_to_dict_is_json_serialisable() -> None: assert plain["children"][0]["items"] == [1, 2] -@pytest.fixture -def hierarchical_converter() -> TorchConverter: - """A converter that has converted the hierarchical model, with stack traces.""" - exported = torch.export.export( - HierarchicalModel().eval(), - args=tuple(get_example_inputs(HierarchicalModel).values()), - ).run_decompositions() +def _convert(model: torch.nn.Module, *inputs: torch.Tensor) -> AIProgram: + exported = torch.export.export(model.eval(), inputs).run_decompositions() converter = TorchConverter() converter.add_exported_program(exported, entrypoint_name="main") - converter.program = converter.to_coreai() - return converter + return converter.to_coreai() @pytest.fixture -def hierarchical_program(hierarchical_converter: TorchConverter) -> AIProgram: - """The Core AI program of the hierarchical model.""" - return hierarchical_converter.program +def hierarchical_program() -> AIProgram: + """The Core AI program of the hierarchical model, with stack traces.""" + return _convert( + HierarchicalModel(), *get_example_inputs(HierarchicalModel).values() + ) def test_module_output_ops_are_where_values_leave( - hierarchical_converter: TorchConverter, + hierarchical_program: AIProgram, ) -> None: """A module's output is the operation realizing a node a node outside reads.""" root = ("HierarchicalModel$1",) - outputs = module_output_op_ids( - hierarchical_converter.program, hierarchical_converter.torch_graph - ) - assert outputs == { + assert module_output_op_ids(hierarchical_program) == { root: [24], (*root, "Linear$1"): [11], (*root, "SubModel$1"): [15], @@ -196,12 +191,10 @@ def test_module_output_ops_are_where_values_leave( def test_output_is_the_operation_realizing_the_node( - hierarchical_converter: TorchConverter, + hierarchical_program: AIProgram, ) -> None: """An operation that absorbed a producer does not hold that producer's value.""" - outputs = module_output_op_ids( - hierarchical_converter.program, hierarchical_converter.torch_graph - ) + outputs = module_output_op_ids(hierarchical_program) # Op 13 is Linear$2's matmul and also carries Linear$1's location, and the last # sigmoid carries Linear$4's. Neither holds the value of the layer it absorbed. @@ -209,16 +202,36 @@ def test_output_is_the_operation_realizing_the_node( assert outputs[("HierarchicalModel$1", "Linear$4")] == [23] -def test_torch_graph_records_modules_and_edges( - hierarchical_converter: TorchConverter, +def test_without_the_node_table_fewer_modules_are_certain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """No id is known to be a reshape without the table, so only some modules remain.""" + monkeypatch.setattr( + _DebugInfoRecorder, "_node_table_metadata", lambda *_args, **_kw: [] + ) + program = _convert( + HierarchicalModel(), *get_example_inputs(HierarchicalModel).values() + ) + + root = ("HierarchicalModel$1",) + assert module_output_op_ids(program) == { + root: [24], + (*root, "Linear$4"): [23], + (*root, "SubModel$1"): [15], + (*root, "SubModel$2"): [20], + } + + +def test_node_table_is_read_back_from_the_program( + hierarchical_program: AIProgram, ) -> None: - """Each FX node is recorded with its module path and what reads it.""" - graph = hierarchical_converter.torch_graph - addmm = next(n for n in graph.values() if n.target == "aten.addmm.default") + """The graph's location carries what each torch node was, and where the graph ends.""" + (table,) = _read_node_tables(hierarchical_program).values() - assert addmm.stack == ("HierarchicalModel$1", "Linear$1") - assert addmm.users - assert all(user in graph for user in addmm.users) + assert table.targets[11] == "aten.addmm.default" + assert table.shape_only(9) + assert not table.shape_only(11) + assert table.output_id == 29 def test_every_operation_has_a_stack_trace(hierarchical_program: AIProgram) -> None: @@ -259,12 +272,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def test_module_whose_reader_was_folded_in_is_left_out() -> None: """A conv's op holds the batch norm's value once folded, not the conv's.""" - exported = torch.export.export( - _ConvBatchNorm().eval(), (torch.randn(1, 3, 8, 8),) - ).run_decompositions() - converter = TorchConverter() - converter.add_exported_program(exported, entrypoint_name="main") - outputs = module_output_op_ids(converter.to_coreai(), converter.torch_graph) + outputs = module_output_op_ids( + _convert(_ConvBatchNorm(), torch.randn(1, 3, 8, 8)) + ) root = ("_ConvBatchNorm$1",) assert (*root, "Conv2d$1") not in outputs @@ -273,17 +283,13 @@ def test_module_whose_reader_was_folded_in_is_left_out() -> None: async def test_get_module_outputs_matches_torch( - hierarchical_converter: TorchConverter, specialization_options: SpecializationOptions | None, ) -> None: """Probing a module's output gives the value torch's own module returns.""" torch.manual_seed(0) model = HierarchicalModel().eval() x = get_example_inputs(HierarchicalModel)["x"] - exported = torch.export.export(model, (x,)).run_decompositions() - converter = TorchConverter() - converter.add_exported_program(exported, entrypoint_name="main") - program = converter.to_coreai() + program = _convert(model, x) seen: list[torch.Tensor] = [] handle = model.sub_model1.register_forward_hook( @@ -298,7 +304,6 @@ async def test_get_module_outputs_matches_torch( (value,) = await get_module_outputs( inspector, program, - converter.torch_graph, ("HierarchicalModel$1", "SubModel$1"), {"x": x.numpy()}, ) @@ -309,38 +314,30 @@ async def test_get_module_outputs_matches_torch( async def test_get_module_outputs_refuses_what_it_cannot_place( - hierarchical_converter: TorchConverter, + hierarchical_program: AIProgram, specialization_options: SpecializationOptions | None, ) -> None: """A module that is absent or folded away is an error, never a guess.""" inspector = CoreAIProbedInspector( - hierarchical_converter.program, specialization_options=specialization_options + hierarchical_program, specialization_options=specialization_options ) - graph = hierarchical_converter.torch_graph x = get_example_inputs(HierarchicalModel)["x"].numpy() - with pytest.raises(ModuleOutputUnknownError, match="no module"): + with pytest.raises(ModuleOutputUnknownError, match="no operation belongs"): await get_module_outputs( inspector, - hierarchical_converter.program, - graph, + hierarchical_program, ("HierarchicalModel$1", "Nope$1"), {"x": x}, ) - exported = torch.export.export( - _ConvBatchNorm().eval(), (torch.randn(1, 3, 8, 8),) - ).run_decompositions() - folded = TorchConverter() - folded.add_exported_program(exported, entrypoint_name="main") - program = folded.to_coreai() + program = _convert(_ConvBatchNorm(), torch.randn(1, 3, 8, 8)) with pytest.raises(ModuleOutputUnknownError, match="not certain"): await get_module_outputs( CoreAIProbedInspector( program, specialization_options=specialization_options ), program, - folded.torch_graph, ("_ConvBatchNorm$1", "Conv2d$1"), {"x": np.zeros((1, 3, 8, 8), dtype=np.float32)}, ) From 7dcadfb3c1622a4d43b85d171a4ebdcd59d5bb69 Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Thu, 22 Oct 2026 15:45:55 -0700 Subject: [PATCH 08/11] Sort imports and format the module output code with ruff --- coreai_torch/debugging/modules.py | 8 ++++---- tests/debugging/test_modules.py | 6 ++---- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/coreai_torch/debugging/modules.py b/coreai_torch/debugging/modules.py index e400c74..7dd3a2f 100644 --- a/coreai_torch/debugging/modules.py +++ b/coreai_torch/debugging/modules.py @@ -358,9 +358,7 @@ def _read_node_tables(program: AIProgram) -> dict[str, _NodeTable]: nodes = DictAttr(data["nodes"]) name = StringAttr(operation.attributes["sym_name"]).value tables[name] = _NodeTable( - targets={ - int(entry.name): StringAttr(entry.attr).value for entry in nodes - }, + targets={int(entry.name): StringAttr(entry.attr).value for entry in nodes}, output_id=( IntegerAttr(data["output_id"]).value if "output_id" in data else None ), @@ -431,7 +429,9 @@ def in_module(stack: ModulePath, module: ModulePath) -> bool: return stack[: len(module)] == module modules = { - stack[:depth] for stack in owner_stack.values() for depth in range(1, len(stack) + 1) + stack[:depth] + for stack in owner_stack.values() + for depth in range(1, len(stack) + 1) } outputs: dict[ModulePath, list[int]] = {} for module in sorted(modules): diff --git a/tests/debugging/test_modules.py b/tests/debugging/test_modules.py index 57f318b..bff9fe2 100644 --- a/tests/debugging/test_modules.py +++ b/tests/debugging/test_modules.py @@ -32,8 +32,8 @@ from coreai_torch.debugging.modules import ( UNATTRIBUTED, ModuleNode, - build_module_tree, _read_node_tables, + build_module_tree, module_output_op_ids, ) from coreai_torch.debugging.source_annotator import ( @@ -272,9 +272,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def test_module_whose_reader_was_folded_in_is_left_out() -> None: """A conv's op holds the batch norm's value once folded, not the conv's.""" - outputs = module_output_op_ids( - _convert(_ConvBatchNorm(), torch.randn(1, 3, 8, 8)) - ) + outputs = module_output_op_ids(_convert(_ConvBatchNorm(), torch.randn(1, 3, 8, 8))) root = ("_ConvBatchNorm$1",) assert (*root, "Conv2d$1") not in outputs From b415c4cb5e2a3669b07bfd8aa44b5e4f0c37bb83 Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Thu, 22 Oct 2026 17:56:55 -0700 Subject: [PATCH 09/11] Move module outputs to their own file and recover more modules Split module_output_op_ids out of modules.py into module_outputs.py, built from a _Boundaries class and a three-way _Verdict (the value stays, leaves, or is not certain) instead of one long function. The behaviour is unchanged by the move. Two changes recover modules that were left out: - a value only written to state (create_token, read_handle, write_handle) is not a module output and no longer makes a module uncertain, so a stateful model with several outputs keeps its root; - expand returns its input's values, so it is shape-only like view and no longer blocks a boundary it sits on. Move the module-output tests to test_module_outputs.py and the stack-trace tests to test_source_annotator.py. Add a test that a stateful two-head model returns its two heads, matching torch. --- coreai_torch/debugging/inspector.py | 2 +- coreai_torch/debugging/module_outputs.py | 285 +++++++++++++++++++++++ coreai_torch/debugging/modules.py | 204 +--------------- tests/debugging/test_module_outputs.py | 235 +++++++++++++++++++ tests/debugging/test_modules.py | 205 ---------------- tests/debugging/test_source_annotator.py | 31 +++ 6 files changed, 553 insertions(+), 409 deletions(-) create mode 100644 coreai_torch/debugging/module_outputs.py create mode 100644 tests/debugging/test_module_outputs.py diff --git a/coreai_torch/debugging/inspector.py b/coreai_torch/debugging/inspector.py index 5396062..6dfa1d6 100644 --- a/coreai_torch/debugging/inspector.py +++ b/coreai_torch/debugging/inspector.py @@ -29,7 +29,7 @@ _build_compile_id_to_coreai_map, parse_debug_infos, ) -from .modules import module_output_op_ids +from .module_outputs import module_output_op_ids from .probes import ( Perturbation, ProbedProgram, diff --git a/coreai_torch/debugging/module_outputs.py b/coreai_torch/debugging/module_outputs.py new file mode 100644 index 0000000..0b3cac4 --- /dev/null +++ b/coreai_torch/debugging/module_outputs.py @@ -0,0 +1,285 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +""" +Which Core AI operations hold a module's output. + +The compiler fuses, folds and reshapes, so a module's output is not simply the last +operation carrying its name. This reads each operation's torch ids and module stacks, +and the node table `coreai_torch` records on the graph, and reports a module only when +nothing could have changed its value on the way out. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass +from enum import Enum, auto + +from coreai._compiler.ir import Attribute, DictAttr, IntegerAttr, Operation, StringAttr +from coreai.authoring import AIProgram + +from .._debug_locations import NODE_TABLE_METADATA, NODE_TABLE_VERSION +from .debug_info import _build_coreai_op_map, get_operation_id +from .source_annotator import ( + ModulePath, + _leaf_stack_traces, + operation_stack_trace, + operation_torch_ids, +) +from .utils import STRUCTURAL_OPS + +_SHAPE_ONLY = frozenset( + { + "view", + "reshape", + "_unsafe_view", + "squeeze", + "unsqueeze", + "flatten", + "contiguous", + "clone", + "alias", + "detach", + "expand", + } +) +"""ATen operations that return their input's values unchanged, in another shape or broadcast.""" + +_STATE_OPS = frozenset( + {"coreai.create_token", "coreai.read_handle", "coreai.write_handle"} +) +"""Operations that move a value into or out of program state.""" + + +@dataclass(frozen=True) +class _SourceNodes: + """What a graph's location says about the torch nodes converted into it.""" + + targets: dict[int, str] + """ATen target of each ``call_function`` node, by torch operation ID.""" + + output_id: int | None + """ID of the node that returns the graph's values, None when not recorded.""" + + def shape_only(self, node_id: int) -> bool: + """Whether the node returns its input's values unchanged, in another shape.""" + parts = self.targets.get(node_id, "").split(".") + return len(parts) > 1 and parts[1] in _SHAPE_ONLY + + +def _metadata_data(printed: str, name: str) -> str | None: + """ + The printed data of the metadata called *name* in a printed location, if any. + + A location's metadata is read from its printed form because builds differ in + what they expose of it: some give the list of metadata, others only the + attribute. The data is a brace-delimited dictionary, which MLIR parses back. + """ + marker = f'<"{name}": ' + start = printed.find(marker) + if start < 0: + return None + start += len(marker) + depth, quoted = 0, False + for index in range(start, len(printed)): + char = printed[index] + if char == '"' and printed[index - 1] != "\\": + quoted = not quoted + elif not quoted and char == "{": + depth += 1 + elif not quoted and char == "}": + depth -= 1 + if depth == 0: + return printed[start : index + 1] + return None + + +def _read_source_nodes(program: AIProgram) -> dict[str, _SourceNodes]: + """ + The node table recorded on each graph's location, by graph name. + + A graph with no table, or one this version cannot read, is absent: nothing is + assumed about the nodes of a program that was not converted with one. + """ + graphs: dict[str, _SourceNodes] = {} + for operation in _build_coreai_op_map(program).values(): + if operation.name != "coreai.graph": + continue + printed = _metadata_data(str(operation.location), NODE_TABLE_METADATA) + if printed is None: + continue + with operation.context: + data = DictAttr(Attribute.parse(printed)) + if IntegerAttr(data["version"]).value != NODE_TABLE_VERSION: + continue + name = StringAttr(operation.attributes["sym_name"]).value + graphs[name] = _SourceNodes( + targets={ + int(entry.name): StringAttr(entry.attr).value + for entry in DictAttr(data["nodes"]) + }, + output_id=( + IntegerAttr(data["output_id"]).value if "output_id" in data else None + ), + ) + return graphs + + +class _Verdict(Enum): + """What one reader of an operation's value says about its module.""" + + STAYS = auto() + """The reader is in the same module, so the value has not left it.""" + + LEAVES = auto() + """The reader is in another module, with every node on the way accounted for.""" + + UNKNOWN = auto() + """A node on the way could have changed the value, so the boundary is not certain.""" + + +class _Boundaries: + """ + The module each operation of a program belongs to, and what lies between them. + + An operation is owned by the torch node with the largest ID in its location that + is not shape-only. The smaller IDs belong to producers it absorbed: they say where + it came from, not what it computes. + """ + + def __init__(self, program: AIProgram) -> None: + self._graphs = _read_source_nodes(program) + targets = { + i: t for graph in self._graphs.values() for i, t in graph.targets.items() + } + self._shape_only = _SourceNodes(targets, None).shape_only + self._ops = _build_coreai_op_map(program) + self._node: dict[int, int] = {} + self._stack: dict[int, ModulePath] = {} + self._seen: set[int] = set() + for op_id, operation in self._ops.items(): + if operation.name not in STRUCTURAL_OPS: + self._own(op_id, operation) + + def _own(self, op_id: int, operation: Operation) -> None: + """Record which torch node and module an operation belongs to.""" + ids = operation_torch_ids(operation) + self._seen |= ids + if operation.name == "coreai.constant" or not ids: + return + node = max({i for i in ids if not self._shape_only(i)} or ids) + stacks = {node_id: stack for stack, node_id in _leaf_stack_traces(operation)} + self._node[op_id] = node + self._stack[op_id] = stacks.get(node) or operation_stack_trace(operation) + + def modules(self) -> list[ModulePath]: + """Every module that owns an operation, parents included.""" + return sorted( + { + stack[:depth] + for stack in self._stack.values() + for depth in range(1, len(stack) + 1) + } + ) + + def members(self, module: ModulePath) -> list[int]: + """The operations owned by *module* or by a module inside it.""" + return [ + op_id + for op_id, stack in self._stack.items() + if stack[: len(module)] == module + ] + + def readers(self, op_id: int) -> Iterator[Operation]: + """ + The operations that read any value of operation *op_id*. + + A state operation is not a reader: writing a value to state is not it leaving + the module, and a module's output is what its caller receives, not what it + stores. + """ + for result in self._ops[op_id].results: + for use in result.uses: + if use.owner.name not in _STATE_OPS: + yield use.owner + + def verdict(self, op_id: int, reader: Operation, module: ModulePath) -> _Verdict: + """What *reader* says about whether *op_id*'s value has left *module*.""" + if reader.name == "coreai.output": + return self._verdict_for_output(op_id, reader, module) + reader_id = get_operation_id(reader) + if reader_id not in self._stack: + return _Verdict.UNKNOWN + if self._stack[reader_id][: len(module)] == module: + return _Verdict.STAYS + return self._verdict_between(self._node[op_id], self._node[reader_id]) + + def _verdict_for_output( + self, op_id: int, output: Operation, module: ModulePath + ) -> _Verdict: + """The program's own output: the root's, or a module's if nothing follows it.""" + if len(module) == 1: + return _Verdict.LEAVES + graph = StringAttr(output.parent.attributes["sym_name"]).value + end = self._graphs[graph].output_id if graph in self._graphs else None + if end is None: + return _Verdict.UNKNOWN + return self._verdict_between(self._node[op_id], end) + + def _verdict_between(self, low: int, high: int) -> _Verdict: + """LEAVES when every torch node between the two is accounted for.""" + if high <= low: + return _Verdict.UNKNOWN + # An id on no operation was folded away without a trace, as batch norm is into a + # conv. A reshape is harmless only if the node table says it was one. + hidden = any( + node_id not in self._seen and not self._shape_only(node_id) + for node_id in range(low + 1, high) + ) + return _Verdict.UNKNOWN if hidden else _Verdict.LEAVES + + +def module_output_op_ids(program: AIProgram) -> dict[ModulePath, list[int]]: + """ + The Core AI operations that hold each module instance's output. + + A module's output is a value that leaves it: an operation it owns that an operation + of another module reads, or that returns the program's own output. The module is + left out, rather than guessed, whenever a node could have changed the value on its + way out: an id between the producer's torch node and its reader's that shows up in + no operation and is not known to be shape-only, a reader with no torch id, or one + whose id is not larger. A value that is only written to state is not an output. + + Without the node table, written when `coreai_torch` converts the program, no id is + known to be shape-only and the program's output is accepted only for the root + module, so fewer modules are reported. A program with no debug info reports none. + + Args: + program: Program to read, converted with debug info. + + Returns: + Module path to the sorted ids of the operations holding its output, for every + module whose output is certain. + + """ + boundaries = _Boundaries(program) + outputs: dict[ModulePath, list[int]] = {} + for module in boundaries.modules(): + verdicts = { + op_id: [ + boundaries.verdict(op_id, reader, module) + for reader in boundaries.readers(op_id) + ] + for op_id in boundaries.members(module) + } + if any(_Verdict.UNKNOWN in found for found in verdicts.values()): + continue + leaving = sorted( + op_id for op_id, found in verdicts.items() if _Verdict.LEAVES in found + ) + if leaving: + outputs[module] = leaving + return outputs diff --git a/coreai_torch/debugging/modules.py b/coreai_torch/debugging/modules.py index 7dd3a2f..8fbb4c4 100644 --- a/coreai_torch/debugging/modules.py +++ b/coreai_torch/debugging/modules.py @@ -13,21 +13,15 @@ from dataclasses import dataclass, field from typing import Any, Generic, TypeVar -from coreai._compiler.ir import Attribute, DictAttr, IntegerAttr, StringAttr from coreai.authoring import AIProgram from typing_extensions import Self -from .._debug_locations import NODE_TABLE_METADATA, NODE_TABLE_VERSION -from .debug_info import _build_coreai_op_map, get_operation_id from .source_annotator import ( ModulePath, - _leaf_stack_traces, module_paths_by_op_id, - operation_stack_trace, - operation_torch_ids, source_locations_by_op_id, ) -from .utils import STRUCTURAL_OPS, LocationInfo, split_module_frame +from .utils import LocationInfo, split_module_frame T = TypeVar("T") @@ -275,199 +269,3 @@ def attributions_by_op_id(program: AIProgram) -> dict[int, Attribution]: ) for op_id in sorted({*paths, *lines}) } - - -_SHAPE_ONLY = frozenset( - { - "view", - "reshape", - "_unsafe_view", - "squeeze", - "unsqueeze", - "flatten", - "contiguous", - "clone", - "alias", - "detach", - } -) -"""ATen operations that return their input's values unchanged, in another shape.""" - - -@dataclass(frozen=True) -class _NodeTable: - """What a graph's location says about the torch nodes converted into it.""" - - targets: dict[int, str] - """ATen target of each ``call_function`` node, by torch operation ID.""" - - output_id: int | None - """ID of the node that returns the graph's values, None when not recorded.""" - - def shape_only(self, node_id: int) -> bool: - """Whether the node returns its input's values unchanged, in another shape.""" - parts = self.targets.get(node_id, "").split(".") - return len(parts) > 1 and parts[1] in _SHAPE_ONLY - - -def _metadata_data(printed: str, name: str) -> str | None: - """ - The printed data of the metadata called *name* in a printed location, if any. - - A location's metadata is read from its printed form because builds differ in - what they expose of it: some give the list of metadata, others only the - attribute. The data is a brace-delimited dictionary, which MLIR parses back. - """ - marker = f'<"{name}": ' - start = printed.find(marker) - if start < 0: - return None - start += len(marker) - depth, quoted = 0, False - for index in range(start, len(printed)): - char = printed[index] - if char == '"' and printed[index - 1] != "\\": - quoted = not quoted - elif not quoted and char == "{": - depth += 1 - elif not quoted and char == "}": - depth -= 1 - if depth == 0: - return printed[start : index + 1] - return None - - -def _read_node_tables(program: AIProgram) -> dict[str, _NodeTable]: - """ - The node table recorded on each graph's location, by graph name. - - A graph with no table, or one this version cannot read, is absent: nothing is - assumed about the nodes of a program that was not converted with one. - """ - tables: dict[str, _NodeTable] = {} - for operation in _build_coreai_op_map(program).values(): - if operation.name != "coreai.graph": - continue - printed = _metadata_data(str(operation.location), NODE_TABLE_METADATA) - if printed is None: - continue - with operation.context: - data = DictAttr(Attribute.parse(printed)) - if IntegerAttr(data["version"]).value != NODE_TABLE_VERSION: - continue - nodes = DictAttr(data["nodes"]) - name = StringAttr(operation.attributes["sym_name"]).value - tables[name] = _NodeTable( - targets={int(entry.name): StringAttr(entry.attr).value for entry in nodes}, - output_id=( - IntegerAttr(data["output_id"]).value if "output_id" in data else None - ), - ) - return tables - - -def module_output_op_ids(program: AIProgram) -> dict[ModulePath, list[int]]: - """ - The Core AI operations that hold each module instance's output. - - A module's output is a value that leaves it: an operation it owns that an operation - of another module reads, or that returns the program's own output. An operation is - owned by the torch node with the largest ID in its location that is not shape-only; - the smaller ones are producers it absorbed, and say where it came from, not what it - computes. - - A module is left out, rather than guessed, whenever a node could have changed the - value on its way out. The torch IDs between the owner of the operation and the - owner of its reader are the nodes on the path between them, and every one must show - up in some operation's location or be shape-only. An ID that appears nowhere was - folded away without a trace, as batch norm is into a conv, and a reshape is - harmless only if the program's node table says it was one. A reader with no torch - ID, or one whose ID is not larger, leaves the module out too. - - Without the node table, written when `coreai_torch` converts the program, no ID is - known to be shape-only and the program's output is accepted only for the root - module, so fewer modules are reported. A program with no debug info reports none. - - Args: - program: Program to read, converted with debug info. - - Returns: - Module path to the sorted ids of the operations holding its output, for every - module whose output is certain. - - """ - tables = _read_node_tables(program) - targets: dict[int, str] = {} - for table in tables.values(): - targets.update(table.targets) - shape_only = _NodeTable(targets, None).shape_only - - ops = _build_coreai_op_map(program) - visible: set[int] = set() - owner_id: dict[int, int] = {} - owner_stack: dict[int, ModulePath] = {} - for op_id, operation in ops.items(): - if operation.name in STRUCTURAL_OPS: - continue - ids = operation_torch_ids(operation) - visible |= ids - if operation.name == "coreai.constant" or not ids: - continue - owner_id[op_id] = max({i for i in ids if not shape_only(i)} or ids) - stacks = {node_id: stack for stack, node_id in _leaf_stack_traces(operation)} - owner_stack[op_id] = stacks.get(owner_id[op_id]) or operation_stack_trace( - operation - ) - - def hidden_between(low: int, high: int) -> bool: - return any( - node_id not in visible and not shape_only(node_id) - for node_id in range(low + 1, high) - ) - - def in_module(stack: ModulePath, module: ModulePath) -> bool: - return stack[: len(module)] == module - - modules = { - stack[:depth] - for stack in owner_stack.values() - for depth in range(1, len(stack) + 1) - } - outputs: dict[ModulePath, list[int]] = {} - for module in sorted(modules): - found: set[int] = set() - certain = True - for op_id, stack in owner_stack.items(): - if not in_module(stack, module): - continue - for result in ops[op_id].results: - for use in result.uses: - user = use.owner - user_id = get_operation_id(user) - if user.name == "coreai.output": - graph = StringAttr(user.parent.attributes["sym_name"]).value - end = tables[graph].output_id if graph in tables else None - # The program's own output, so nothing can have changed it unless a - # node after this module was folded in unseen. - if len(module) == 1 or ( - end is not None - and end > owner_id[op_id] - and not hidden_between(owner_id[op_id], end) - ): - found.add(op_id) - else: - certain = False - continue - if user_id not in owner_stack: - certain = False - continue - if in_module(owner_stack[user_id], module): - continue - low, high = owner_id[op_id], owner_id[user_id] - if high <= low or hidden_between(low, high): - certain = False - else: - found.add(op_id) - if certain and found: - outputs[module] = sorted(found) - return outputs diff --git a/tests/debugging/test_module_outputs.py b/tests/debugging/test_module_outputs.py new file mode 100644 index 0000000..aa172d0 --- /dev/null +++ b/tests/debugging/test_module_outputs.py @@ -0,0 +1,235 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +""" +Which Core AI operations hold a module's output. + +Checked on a model whose layers fuse (each `nn.Linear` is a matmul that absorbs the +previous layer's reshape), and on one whose batch norm the compiler folds into a conv. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from coreai.authoring import AIProgram +from coreai.runtime import SpecializationOptions + +from coreai_torch._debug_locations import _DebugInfoRecorder +from coreai_torch.converter import TorchConverter +from coreai_torch.debugging.inspector import ( + CoreAIProbedInspector, + ModuleOutputUnknownError, + get_module_outputs, +) +from coreai_torch.debugging.module_outputs import ( + _read_source_nodes, + module_output_op_ids, +) + +from .test_model import HierarchicalModel, get_example_inputs + + +def _convert(model: torch.nn.Module, *inputs: torch.Tensor) -> AIProgram: + exported = torch.export.export(model.eval(), inputs).run_decompositions() + converter = TorchConverter() + converter.add_exported_program(exported, entrypoint_name="main") + return converter.to_coreai() + + +@pytest.fixture +def hierarchical_program() -> AIProgram: + """The Core AI program of the hierarchical model, with stack traces.""" + return _convert( + HierarchicalModel(), *get_example_inputs(HierarchicalModel).values() + ) + + +def test_module_output_ops_are_where_values_leave( + hierarchical_program: AIProgram, +) -> None: + """A module's output is the operation realizing a node a node outside reads.""" + root = ("HierarchicalModel$1",) + assert module_output_op_ids(hierarchical_program) == { + root: [24], + (*root, "Linear$1"): [11], + (*root, "SubModel$1"): [15], + (*root, "SubModel$1", "Linear$2"): [14], + (*root, "SubModel$2"): [20], + (*root, "SubModel$2", "Linear$3"): [19], + (*root, "Linear$4"): [23], + } + + +def test_output_is_the_operation_realizing_the_node( + hierarchical_program: AIProgram, +) -> None: + """An operation that absorbed a producer does not hold that producer's value.""" + outputs = module_output_op_ids(hierarchical_program) + + # Op 13 is Linear$2's matmul and also carries Linear$1's location, and the last + # sigmoid carries Linear$4's. Neither holds the value of the layer it absorbed. + assert outputs[("HierarchicalModel$1", "Linear$1")] == [11] + assert outputs[("HierarchicalModel$1", "Linear$4")] == [23] + + +def test_without_the_node_table_fewer_modules_are_certain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """No id is known to be a reshape without the table, so only some modules remain.""" + monkeypatch.setattr( + _DebugInfoRecorder, "_node_table_metadata", lambda *_args, **_kw: [] + ) + program = _convert( + HierarchicalModel(), *get_example_inputs(HierarchicalModel).values() + ) + + root = ("HierarchicalModel$1",) + assert module_output_op_ids(program) == { + root: [24], + (*root, "Linear$4"): [23], + (*root, "SubModel$1"): [15], + (*root, "SubModel$2"): [20], + } + + +def test_node_table_is_read_back_from_the_program( + hierarchical_program: AIProgram, +) -> None: + """The graph's location carries what each torch node was, and where the graph ends.""" + (table,) = _read_source_nodes(hierarchical_program).values() + + assert table.targets[11] == "aten.addmm.default" + assert table.shape_only(9) + assert not table.shape_only(11) + assert table.output_id == 29 + + +class _ConvBatchNorm(torch.nn.Module): + """A conv whose batch norm the compiler folds into its weights.""" + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3) + self.bn = torch.nn.BatchNorm2d(4) + self.head = torch.nn.Conv2d(4, 4, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.head(torch.relu(self.bn(self.conv(x)))) + + +def test_module_whose_reader_was_folded_in_is_left_out() -> None: + """A conv's op holds the batch norm's value once folded, not the conv's.""" + outputs = module_output_op_ids(_convert(_ConvBatchNorm(), torch.randn(1, 3, 8, 8))) + + root = ("_ConvBatchNorm$1",) + assert (*root, "Conv2d$1") not in outputs + assert (*root, "BatchNorm2d$1") not in outputs + assert (*root, "Conv2d$2") in outputs + + +async def test_get_module_outputs_matches_torch( + specialization_options: SpecializationOptions | None, +) -> None: + """Probing a module's output gives the value torch's own module returns.""" + torch.manual_seed(0) + model = HierarchicalModel().eval() + x = get_example_inputs(HierarchicalModel)["x"] + program = _convert(model, x) + + seen: list[torch.Tensor] = [] + handle = model.sub_model1.register_forward_hook( + lambda _m, _i, out: seen.append(out.detach()) + ) + model(x) + handle.remove() + + inspector = CoreAIProbedInspector( + program, specialization_options=specialization_options + ) + (value,) = await get_module_outputs( + inspector, + program, + ("HierarchicalModel$1", "SubModel$1"), + {"x": x.numpy()}, + ) + + np.testing.assert_allclose( + value.reshape(-1), seen[0].numpy().reshape(-1), rtol=1e-3, atol=1e-3 + ) + + +async def test_get_module_outputs_refuses_what_it_cannot_place( + hierarchical_program: AIProgram, + specialization_options: SpecializationOptions | None, +) -> None: + """A module that is absent or folded away is an error, never a guess.""" + inspector = CoreAIProbedInspector( + hierarchical_program, specialization_options=specialization_options + ) + x = get_example_inputs(HierarchicalModel)["x"].numpy() + + with pytest.raises(ModuleOutputUnknownError, match="no operation belongs"): + await get_module_outputs( + inspector, + hierarchical_program, + ("HierarchicalModel$1", "Nope$1"), + {"x": x}, + ) + + program = _convert(_ConvBatchNorm(), torch.randn(1, 3, 8, 8)) + with pytest.raises(ModuleOutputUnknownError, match="not certain"): + await get_module_outputs( + CoreAIProbedInspector( + program, specialization_options=specialization_options + ), + program, + ("_ConvBatchNorm$1", "Conv2d$1"), + {"x": np.zeros((1, 3, 8, 8), dtype=np.float32)}, + ) + + +class _StatefulHeads(torch.nn.Module): + """Two heads over a shared body, and a buffer that is updated in place.""" + + def __init__(self) -> None: + super().__init__() + self.body = torch.nn.Linear(6, 6) + self.a = torch.nn.Linear(6, 3) + self.b = torch.nn.Linear(6, 2) + self.register_buffer("acc", torch.zeros(1, 6)) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + h = torch.relu(self.body(x)) + self.acc.copy_(self.acc + h.mean(0, keepdim=True)) + return self.a(h), self.b(h) + + +async def test_a_value_written_to_state_is_not_a_module_output( + specialization_options: SpecializationOptions | None, +) -> None: + """The root returns its two heads; the buffer it updates is not one of them.""" + torch.manual_seed(0) + model = _StatefulHeads().eval() + x = torch.randn(2, 6) + program = _convert(model, x) + with torch.no_grad(): + expected = [out.numpy() for out in model(x)] + + inspector = CoreAIProbedInspector( + program, specialization_options=specialization_options + ) + arrays = await get_module_outputs( + inspector, + program, + ("_StatefulHeads$1",), + {"x": x.numpy()}, + state={"b_acc": np.zeros((1, 6), dtype=np.float32)}, + ) + + assert len(arrays) == len(expected) + for array, want in zip(arrays, expected, strict=True): + np.testing.assert_allclose(array, want, rtol=1e-3, atol=1e-3) diff --git a/tests/debugging/test_modules.py b/tests/debugging/test_modules.py index bff9fe2..97cf619 100644 --- a/tests/debugging/test_modules.py +++ b/tests/debugging/test_modules.py @@ -15,34 +15,11 @@ import json -import numpy as np -import pytest -import torch -from coreai.authoring import AIProgram -from coreai.runtime import SpecializationOptions - -from coreai_torch._debug_locations import _DebugInfoRecorder -from coreai_torch.converter import TorchConverter -from coreai_torch.debugging.debug_info import _build_coreai_op_map -from coreai_torch.debugging.inspector import ( - CoreAIProbedInspector, - ModuleOutputUnknownError, - get_module_outputs, -) from coreai_torch.debugging.modules import ( UNATTRIBUTED, ModuleNode, - _read_node_tables, build_module_tree, - module_output_op_ids, -) -from coreai_torch.debugging.source_annotator import ( - module_paths_by_op_id, - operation_stack_traces, ) -from coreai_torch.debugging.utils import STRUCTURAL_OPS - -from .test_model import HierarchicalModel, get_example_inputs def test_nests_by_path() -> None: @@ -157,185 +134,3 @@ def test_to_dict_is_json_serialisable() -> None: assert plain["count"] == 2 assert plain["items"] == [] assert plain["children"][0]["items"] == [1, 2] - - -def _convert(model: torch.nn.Module, *inputs: torch.Tensor) -> AIProgram: - exported = torch.export.export(model.eval(), inputs).run_decompositions() - converter = TorchConverter() - converter.add_exported_program(exported, entrypoint_name="main") - return converter.to_coreai() - - -@pytest.fixture -def hierarchical_program() -> AIProgram: - """The Core AI program of the hierarchical model, with stack traces.""" - return _convert( - HierarchicalModel(), *get_example_inputs(HierarchicalModel).values() - ) - - -def test_module_output_ops_are_where_values_leave( - hierarchical_program: AIProgram, -) -> None: - """A module's output is the operation realizing a node a node outside reads.""" - root = ("HierarchicalModel$1",) - assert module_output_op_ids(hierarchical_program) == { - root: [24], - (*root, "Linear$1"): [11], - (*root, "SubModel$1"): [15], - (*root, "SubModel$1", "Linear$2"): [14], - (*root, "SubModel$2"): [20], - (*root, "SubModel$2", "Linear$3"): [19], - (*root, "Linear$4"): [23], - } - - -def test_output_is_the_operation_realizing_the_node( - hierarchical_program: AIProgram, -) -> None: - """An operation that absorbed a producer does not hold that producer's value.""" - outputs = module_output_op_ids(hierarchical_program) - - # Op 13 is Linear$2's matmul and also carries Linear$1's location, and the last - # sigmoid carries Linear$4's. Neither holds the value of the layer it absorbed. - assert outputs[("HierarchicalModel$1", "Linear$1")] == [11] - assert outputs[("HierarchicalModel$1", "Linear$4")] == [23] - - -def test_without_the_node_table_fewer_modules_are_certain( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """No id is known to be a reshape without the table, so only some modules remain.""" - monkeypatch.setattr( - _DebugInfoRecorder, "_node_table_metadata", lambda *_args, **_kw: [] - ) - program = _convert( - HierarchicalModel(), *get_example_inputs(HierarchicalModel).values() - ) - - root = ("HierarchicalModel$1",) - assert module_output_op_ids(program) == { - root: [24], - (*root, "Linear$4"): [23], - (*root, "SubModel$1"): [15], - (*root, "SubModel$2"): [20], - } - - -def test_node_table_is_read_back_from_the_program( - hierarchical_program: AIProgram, -) -> None: - """The graph's location carries what each torch node was, and where the graph ends.""" - (table,) = _read_node_tables(hierarchical_program).values() - - assert table.targets[11] == "aten.addmm.default" - assert table.shape_only(9) - assert not table.shape_only(11) - assert table.output_id == 29 - - -def test_every_operation_has_a_stack_trace(hierarchical_program: AIProgram) -> None: - """No operation of a converted model is left without a module.""" - paths = module_paths_by_op_id(hierarchical_program) - ops = _build_coreai_op_map(hierarchical_program) - missing = [ - op_id - for op_id, op in ops.items() - if op.name not in STRUCTURAL_OPS and op.name != "coreai.constant" - if op_id not in paths - ] - assert missing == [] - - -def test_fused_location_gives_every_stack(hierarchical_program: AIProgram) -> None: - """A matmul rebuilt from the previous layer's output is in both layers' stacks.""" - ops = _build_coreai_op_map(hierarchical_program) - stacks = operation_stack_traces(ops[13]) - - assert ("HierarchicalModel$1", "SubModel$1", "Linear$2") in stacks - assert ("HierarchicalModel$1", "Linear$1") in stacks - assert len(stacks) == len(set(stacks)) - - -class _ConvBatchNorm(torch.nn.Module): - """A conv whose batch norm the compiler folds into its weights.""" - - def __init__(self) -> None: - super().__init__() - self.conv = torch.nn.Conv2d(3, 4, 3) - self.bn = torch.nn.BatchNorm2d(4) - self.head = torch.nn.Conv2d(4, 4, 1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.head(torch.relu(self.bn(self.conv(x)))) - - -def test_module_whose_reader_was_folded_in_is_left_out() -> None: - """A conv's op holds the batch norm's value once folded, not the conv's.""" - outputs = module_output_op_ids(_convert(_ConvBatchNorm(), torch.randn(1, 3, 8, 8))) - - root = ("_ConvBatchNorm$1",) - assert (*root, "Conv2d$1") not in outputs - assert (*root, "BatchNorm2d$1") not in outputs - assert (*root, "Conv2d$2") in outputs - - -async def test_get_module_outputs_matches_torch( - specialization_options: SpecializationOptions | None, -) -> None: - """Probing a module's output gives the value torch's own module returns.""" - torch.manual_seed(0) - model = HierarchicalModel().eval() - x = get_example_inputs(HierarchicalModel)["x"] - program = _convert(model, x) - - seen: list[torch.Tensor] = [] - handle = model.sub_model1.register_forward_hook( - lambda _m, _i, out: seen.append(out.detach()) - ) - model(x) - handle.remove() - - inspector = CoreAIProbedInspector( - program, specialization_options=specialization_options - ) - (value,) = await get_module_outputs( - inspector, - program, - ("HierarchicalModel$1", "SubModel$1"), - {"x": x.numpy()}, - ) - - np.testing.assert_allclose( - value.reshape(-1), seen[0].numpy().reshape(-1), rtol=1e-3, atol=1e-3 - ) - - -async def test_get_module_outputs_refuses_what_it_cannot_place( - hierarchical_program: AIProgram, - specialization_options: SpecializationOptions | None, -) -> None: - """A module that is absent or folded away is an error, never a guess.""" - inspector = CoreAIProbedInspector( - hierarchical_program, specialization_options=specialization_options - ) - x = get_example_inputs(HierarchicalModel)["x"].numpy() - - with pytest.raises(ModuleOutputUnknownError, match="no operation belongs"): - await get_module_outputs( - inspector, - hierarchical_program, - ("HierarchicalModel$1", "Nope$1"), - {"x": x}, - ) - - program = _convert(_ConvBatchNorm(), torch.randn(1, 3, 8, 8)) - with pytest.raises(ModuleOutputUnknownError, match="not certain"): - await get_module_outputs( - CoreAIProbedInspector( - program, specialization_options=specialization_options - ), - program, - ("_ConvBatchNorm$1", "Conv2d$1"), - {"x": np.zeros((1, 3, 8, 8), dtype=np.float32)}, - ) diff --git a/tests/debugging/test_source_annotator.py b/tests/debugging/test_source_annotator.py index e85cfcd..0565618 100644 --- a/tests/debugging/test_source_annotator.py +++ b/tests/debugging/test_source_annotator.py @@ -13,13 +13,17 @@ from coreai.authoring import AIProgram from coreai_torch.converter import TorchConverter, _DebugInfoRecorder +from coreai_torch.debugging.debug_info import _build_coreai_op_map from coreai_torch.debugging.source_annotator import ( _annotate_source, _Annotation, _SourceAnnotator, _TextAnnotation, _walk_operations, + module_paths_by_op_id, + operation_stack_traces, ) +from coreai_torch.debugging.utils import STRUCTURAL_OPS from .test_model import HierarchicalModel, get_example_inputs @@ -127,3 +131,30 @@ async def test_source_annotator_get_annotation_unknown_location( ) assert annotator.get_annotation("/does/not/exist.py", 1) == [] + + +async def test_every_operation_has_a_stack_trace( + hierarchical_coreai_program: AIProgram, +) -> None: + """No operation of a converted model is left without a module.""" + paths = module_paths_by_op_id(hierarchical_coreai_program) + ops = _build_coreai_op_map(hierarchical_coreai_program) + missing = [ + op_id + for op_id, op in ops.items() + if op.name not in STRUCTURAL_OPS and op.name != "coreai.constant" + if op_id not in paths + ] + assert missing == [] + + +async def test_fused_location_gives_every_stack( + hierarchical_coreai_program: AIProgram, +) -> None: + """A matmul rebuilt from the previous layer's output is in both layers' stacks.""" + ops = _build_coreai_op_map(hierarchical_coreai_program) + stacks = operation_stack_traces(ops[13]) + + assert ("HierarchicalModel$1", "SubModel$1", "Linear$2") in stacks + assert ("HierarchicalModel$1", "Linear$1") in stacks + assert len(stacks) == len(set(stacks)) From 2753b35be4663f29b65e54518b83f5fff2f85481 Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Thu, 22 Oct 2026 18:11:39 -0700 Subject: [PATCH 10/11] Reset the search strategy at the start of every validation and comparison A search consumes its strategy: ExhaustiveStrategy yields its single batch once and keeps what it was told, LevelOrderStrategy narrows its depth range and keeps its results. Validator.check and Comparator.compare reused the strategy as it was, so a second call on the same object yielded nothing (the comparator compared no operations and reported no failures) or returned the first call's findings (the validator reported no infinities for an input that has them). Add SearchStrategy.reset, implemented for both strategies and used by their constructors, and call it at the start of each search. Add tests that call a validator and a comparator repeatedly with other inputs and state, which fail without the reset, and tests for the strategies' reset. --- coreai_torch/debugging/comparator.py | 3 ++ coreai_torch/debugging/search_strategy.py | 18 ++++++++- coreai_torch/debugging/validator.py | 3 ++ tests/debugging/test_comparator.py | 28 ++++++++++++++ tests/debugging/test_search_strategy.py | 47 ++++++++++++++++++++++- tests/debugging/test_validator.py | 27 +++++++++++++ 6 files changed, 124 insertions(+), 2 deletions(-) diff --git a/coreai_torch/debugging/comparator.py b/coreai_torch/debugging/comparator.py index bdd9c91..705edc1 100644 --- a/coreai_torch/debugging/comparator.py +++ b/coreai_torch/debugging/comparator.py @@ -1166,6 +1166,9 @@ async def compare( op_ids = self.source.graph.get_op_ids() total_ops = len(op_ids) + # A search consumes the strategy, so another call starts it afresh. + self.strategy.reset() + # Initialize progress tracking and result tracking pass_count = 0 fail_count = 0 diff --git a/coreai_torch/debugging/search_strategy.py b/coreai_torch/debugging/search_strategy.py index c0d1da9..685e6f1 100644 --- a/coreai_torch/debugging/search_strategy.py +++ b/coreai_torch/debugging/search_strategy.py @@ -72,6 +72,16 @@ async def __anext__(self) -> list[ComputationGraph.Node]: """Return the next batch of nodes to check, or raise StopAsyncIteration.""" raise StopAsyncIteration + @abstractmethod + def reset(self) -> None: + """ + Forget everything an earlier search left, so the strategy can search again. + + A search consumes the strategy: it yields its batches once and keeps the results + it was given. A caller that searches again, with other inputs or state, resets it + first, or the second search yields nothing and reports the first one's results. + """ + @abstractmethod async def update( self, @@ -126,6 +136,10 @@ def __init__(self, graph: ComputationGraph[TNode, TGraph]): """ self._graph = graph + self.reset() + + def reset(self) -> None: + """Yield every node again, and forget the results.""" self._yielded = False self._failed: list[ComputationGraph.Node] = [] self._unknown: list[ComputationGraph.Node] = [] @@ -325,8 +339,10 @@ def __init__( self.batch_size = batch_size self.level_selector = level_selector self._initial_scope_id = initial_scope_id + self.reset() - # Track search state + def reset(self) -> None: + """Start the search over from the initial scope, with no results.""" self._scope_nodes: list[ComputationGraph.Node] = [] self._depth_range: tuple[int, int] | None = None self._node_results: dict[ diff --git a/coreai_torch/debugging/validator.py b/coreai_torch/debugging/validator.py index 30da71d..3fa5aef 100644 --- a/coreai_torch/debugging/validator.py +++ b/coreai_torch/debugging/validator.py @@ -265,6 +265,9 @@ async def check( op_ids = self.graph.get_op_ids() total_ops = len(op_ids) + # A search consumes the strategy, so another call starts it afresh. + self.strategy.reset() + # Initialize progress tracking pass_count = 0 excluded: list[ComputationGraph.Node] = [] diff --git a/tests/debugging/test_comparator.py b/tests/debugging/test_comparator.py index 04f6cb1..402adcf 100644 --- a/tests/debugging/test_comparator.py +++ b/tests/debugging/test_comparator.py @@ -690,3 +690,31 @@ async def test_comparator_starts_both_sides_from_the_given_state( assert result.failed_nodes == [] # The caller's array is not updated by either run. np.testing.assert_array_equal(state["b_acc"], np.full((1, 4), 3.0)) + + +async def test_a_comparator_compares_again_on_every_call( + specialization_options: SpecializationOptions | None, +) -> None: + """A second comparison, from another state, covers every operation as the first did.""" + x = torch.randn(1, 4) + exported = torch.export.export(_Accumulator().eval(), (x,)).run_decompositions() + program = await _create_coreai_program_from_model(exported) + comparator = await create_comparator_for_programs( + source_program=exported, + target_program=program, + target_entry_point="main", + specialization_options=specialization_options, + ) + + compared = [] + for value in (3.0, 3.0, -7.0): + state = {"b_acc": np.full((1, 4), value, dtype=np.float32)} + result = await comparator.compare_with_tolerance( + inputs={"x": x.numpy()}, state=state + ) + assert result.failed_nodes == [] + compared.append(len(result.compared)) + np.testing.assert_array_equal(state["b_acc"], np.full((1, 4), value)) + + assert compared[0] > 0 + assert compared == [compared[0]] * 3 diff --git a/tests/debugging/test_search_strategy.py b/tests/debugging/test_search_strategy.py index f2d1ccc..fc26051 100644 --- a/tests/debugging/test_search_strategy.py +++ b/tests/debugging/test_search_strategy.py @@ -6,7 +6,11 @@ """Tests for search strategy with hierarchical graphs.""" from coreai_torch.debugging.graph import ComputationGraph -from coreai_torch.debugging.search_strategy import LevelOrderStrategy, SearchStrategy +from coreai_torch.debugging.search_strategy import ( + ExhaustiveStrategy, + LevelOrderStrategy, + SearchStrategy, +) def create_hierarchical_graph_with_dependencies() -> ComputationGraph: @@ -452,3 +456,44 @@ async def test_skipped_nodes_do_not_narrow_the_search() -> None: assert not strategy.get_unknown_operations(), ( "A skipped node is not an unknown one: it was never a candidate" ) + + +async def test_a_reset_level_order_strategy_searches_again_from_scratch() -> None: + """A finished search leaves the strategy spent; reset makes it a fresh one.""" + graph = create_hierarchical_graph_with_dependencies() + strategy = LevelOrderStrategy.top_down(graph) + + # Failing everything narrows the search, so the first run is cut short. + await _drain(strategy, SearchStrategy.ValidationResult.FAIL) + assert strategy.get_problematic_operations() + assert await _drain(strategy, SearchStrategy.ValidationResult.PASS) == [] + + strategy.reset() + + assert not strategy.get_problematic_operations() + again = await _drain(strategy, SearchStrategy.ValidationResult.PASS) + fresh = await _drain( + LevelOrderStrategy.top_down(graph), SearchStrategy.ValidationResult.PASS + ) + assert [n.op_id for n in again] == [n.op_id for n in fresh] + assert len(again) > 2 + assert not strategy.get_problematic_operations() + + +async def test_a_reset_exhaustive_strategy_offers_every_node_again() -> None: + """It yields everything once, and forgets what it was told when reset.""" + graph = create_hierarchical_graph_with_dependencies() + strategy = ExhaustiveStrategy(graph=graph) + + (batch,) = [b async for b in strategy] + await strategy.update( + [(node, SearchStrategy.ValidationResult.FAIL) for node in batch] + ) + assert strategy.get_problematic_operations() + assert [b async for b in strategy] == [] + + strategy.reset() + + assert not strategy.get_problematic_operations() + (again,) = [b async for b in strategy] + assert [n.op_id for n in again] == [n.op_id for n in batch] diff --git a/tests/debugging/test_validator.py b/tests/debugging/test_validator.py index 5e92b6d..1a21eb0 100644 --- a/tests/debugging/test_validator.py +++ b/tests/debugging/test_validator.py @@ -462,3 +462,30 @@ async def test_coreai_program_with_inspector( ) # Failed nodes can be inspected for debugging _ = [str(n) for n in result.failed_nodes[:3]] + + +class _Reciprocal(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(1.0 / x) + 1.0 + + +async def test_a_validator_checks_again_on_every_call( + specialization_options: SpecializationOptions | None, +) -> None: + """A second check, on other inputs, reports what it finds there, not the first's.""" + x = torch.ones(1, 4) + exported = torch.export.export(_Reciprocal().eval(), (x,)).run_decompositions() + converter = TorchConverter() + converter.add_exported_program(exported, entrypoint_name="main") + validator = await create_validator_for_coreai_program( + converter.to_coreai(), "main", specialization_options=specialization_options + ) + clean = {"x": np.ones((1, 4), dtype=np.float32)} + infinite = {"x": np.zeros((1, 4), dtype=np.float32)} # 1 / 0 + + assert (await validator.check_for_infs(clean)).failed_nodes == [] + first = await validator.check_for_infs(infinite) + again = await validator.check_for_infs(infinite) + assert len(first.failed_nodes) > 0 + assert len(again.failed_nodes) == len(first.failed_nodes) + assert (await validator.check_for_infs(clean)).failed_nodes == [] From 185750e7f710dbbee9ce6f49aed9876856e29a1b Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Mon, 26 Oct 2026 12:36:05 -0700 Subject: [PATCH 11/11] Read the node table from its metadata entry, not the printed location The node table on a graph's location was found by scanning the printed location for its brace-delimited data. Find the entry by name in the location's metadata list and print that attribute alone instead, so MLIR parses exactly the table. A build whose locations do not list their metadata has no table to read, so the reduced behaviour applies there. Drop the scanner and its tests, and test that a program converted without debug info reports no modules. --- coreai_torch/debugging/module_outputs.py | 43 +++++++++++------------- tests/debugging/test_module_outputs.py | 13 +++++++ 2 files changed, 33 insertions(+), 23 deletions(-) diff --git a/coreai_torch/debugging/module_outputs.py b/coreai_torch/debugging/module_outputs.py index 0b3cac4..69e21b6 100644 --- a/coreai_torch/debugging/module_outputs.py +++ b/coreai_torch/debugging/module_outputs.py @@ -18,7 +18,14 @@ from dataclasses import dataclass from enum import Enum, auto -from coreai._compiler.ir import Attribute, DictAttr, IntegerAttr, Operation, StringAttr +from coreai._compiler.ir import ( + Attribute, + DictAttr, + IntegerAttr, + Location, + Operation, + StringAttr, +) from coreai.authoring import AIProgram from .._debug_locations import NODE_TABLE_METADATA, NODE_TABLE_VERSION @@ -70,30 +77,20 @@ def shape_only(self, node_id: int) -> bool: return len(parts) > 1 and parts[1] in _SHAPE_ONLY -def _metadata_data(printed: str, name: str) -> str | None: +def _node_table_data(location: Location) -> str | None: """ - The printed data of the metadata called *name* in a printed location, if any. + The printed data of the node table on a graph's location, if it carries one. - A location's metadata is read from its printed form because builds differ in - what they expose of it: some give the list of metadata, others only the - attribute. The data is a brace-delimited dictionary, which MLIR parses back. + The entry is found by name in the location's metadata and printed on its own, which + is the attribute alone, and MLIR parses its dictionary back. A build whose locations + do not list their metadata has no table to read. """ - marker = f'<"{name}": ' - start = printed.find(marker) - if start < 0: - return None - start += len(marker) - depth, quoted = 0, False - for index in range(start, len(printed)): - char = printed[index] - if char == '"' and printed[index - 1] != "\\": - quoted = not quoted - elif not quoted and char == "{": - depth += 1 - elif not quoted and char == "}": - depth -= 1 - if depth == 0: - return printed[start : index + 1] + entries = getattr(location.attr.maybe_downcast(), "metadata", None) or [] + prefix = f'#debuginfo.metadata<"{NODE_TABLE_METADATA}": ' + for entry in entries: + printed = str(entry) + if entry.name == NODE_TABLE_METADATA and printed.startswith(prefix): + return printed[len(prefix) : -1] return None @@ -108,7 +105,7 @@ def _read_source_nodes(program: AIProgram) -> dict[str, _SourceNodes]: for operation in _build_coreai_op_map(program).values(): if operation.name != "coreai.graph": continue - printed = _metadata_data(str(operation.location), NODE_TABLE_METADATA) + printed = _node_table_data(operation.location) if printed is None: continue with operation.context: diff --git a/tests/debugging/test_module_outputs.py b/tests/debugging/test_module_outputs.py index aa172d0..9ee895f 100644 --- a/tests/debugging/test_module_outputs.py +++ b/tests/debugging/test_module_outputs.py @@ -96,6 +96,19 @@ def test_without_the_node_table_fewer_modules_are_certain( } +def test_a_program_without_debug_info_reports_no_modules() -> None: + """With no locations there is nothing to read, and nothing is guessed.""" + model = HierarchicalModel() + x = get_example_inputs(HierarchicalModel)["x"] + exported = torch.export.export(model.eval(), (x,)).run_decompositions() + converter = TorchConverter(mode=TorchConverter.Mode.RELEASE) + converter.add_exported_program(exported, entrypoint_name="main") + program = converter.to_coreai() + + assert _read_source_nodes(program) == {} + assert module_output_op_ids(program) == {} + + def test_node_table_is_read_back_from_the_program( hierarchical_program: AIProgram, ) -> None: