diff --git a/docs/src/palettization/config.md b/docs/src/palettization/config.md index 5b07e4d..e83f777 100644 --- a/docs/src/palettization/config.md +++ b/docs/src/palettization/config.md @@ -116,7 +116,7 @@ op_config = OpKMeansPalettizerConfig( ## Examples -Several examples below configure specific module types or module names. To determine these for your model, see [How to get names + types](../quantization/config.md#how-to-get-names--types-for-modules-and-ops). Since palettization only supports eager execution mode, only the eager mode guidance in that section is relevant. +Several examples below configure specific module types or module names. To determine these for your model, use {class}`~coreai_opt.inspection.ModelInspector` with `execution_mode="eager"` — see [Inspecting Model Structure](../utils/model_inspection.md). Palettization supports eager mode only. ### Apply 4-bit palettization globally, 8-bit to linear layers diff --git a/docs/src/quantization/config.md b/docs/src/quantization/config.md index f1195e8..2adf49a 100644 --- a/docs/src/quantization/config.md +++ b/docs/src/quantization/config.md @@ -184,7 +184,7 @@ The defaults are: In [Quantization Overview](overview.md) we saw how to use the default `W_INT8_A_INT8` config. [Config classes and their defaults](#config-classes-and-their-defaults) described the default settings in `QuantizerConfig()`, `ModuleQuantizerConfig()`, and `OpQuantizerConfig()`. Let us now see how to configure quantization when non-default settings are desired. -Several examples below configure specific module names, module types, op names, or op types. To determine these for your model, see [How to get names + types for modules and ops](#how-to-get-names-types-for-modules-and-ops) (eager mode) or [Inspecting Model Structure](../utils/model_inspection.md) (graph mode). +Several examples below configure specific module names, module types, op names, or op types. To determine these for your model, see [Inspecting Model Structure](../utils/model_inspection.md). ### Example: `W_MXFP4_A_FP8` applied to all supported ops @@ -857,7 +857,7 @@ classDiagram ## How to get names + types for modules and ops -**Graph mode** (for `module_name_configs`, `module_type_configs`, `op_name_config`, `op_type_config`): use {class}`~coreai_opt.inspection.ModelInspector` to discover module names, module types, op names, and op types. +Use {class}`~coreai_opt.inspection.ModelInspector` to discover module names, module types, op names, and op types for both graph and eager execution modes. ```python import torch @@ -866,40 +866,12 @@ from coreai_opt.inspection import ModelInspector model = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5)) inspector = ModelInspector( - model, example_inputs=(torch.randn(1, 10),), execution_mode="graph" + # Use execution_mode="eager" for eager mode inspection. + model, + example_inputs=(torch.randn(1, 10),), + execution_mode="graph", ) print(inspector.format_summary()) ``` -See [Inspecting Model Structure](../utils/model_inspection.md) for full usage and examples. - -**Eager mode**: module names (for `module_name_configs`) can be obtained by inspecting `model.named_modules()`. This includes all modules in the model (nested and leaf). The names align with the structure of modules defined in code. - -Op names can be constructed by referring to the parent module and the op in it. Example: - -```python -class TwoAddModule(torch.nn.Module): - def forward(self, x): - a = x + x - b = a + a - return b - - -class Model(torch.nn.Module): - def __init__(self): - super().__init__() - self.submodule_a = TwoAddModule() - self.submodule_b = TwoAddModule() - - def forward(self, x): - x = self.submodule_a(x) - x = self.submodule_b(x) - return x - - -# the names of ops in the model will be : -# - submodule_a.add -# - submodule_a.add_1 -# - submodule_b.add -# - submodule_b.add_1 -``` +See [Inspecting Model Structure](../utils/model_inspection.md) for full usage, examples, and a comparison of graph and eager mode op naming. diff --git a/docs/src/utils/model_inspection.md b/docs/src/utils/model_inspection.md index d950c52..4b356e6 100644 --- a/docs/src/utils/model_inspection.md +++ b/docs/src/utils/model_inspection.md @@ -2,9 +2,16 @@ `coreai-opt` configs reference module names, module types, op names, and op types to target specific parts of a model. Before writing a config, you need to know exactly which strings your model exposes. {class}`~coreai_opt.inspection.ModelInspector` discovers these automatically and provides query methods corresponding to each config key type (`op_type_config`, `op_name_config`, `module_name_configs`, `module_type_configs`). -:::{note} -`ModelInspector` currently supports **graph execution mode only**. Eager mode support is planned. For eager mode op naming, see [How to get names + types](../quantization/config.md#how-to-get-names--types-for-modules-and-ops). -::: +## Execution Modes + +`ModelInspector` supports two execution modes, selected via the `execution_mode` argument: + +- **Graph mode** (`execution_mode="graph"`): Exports the model with `torch.export` and walks the resulting FX graph. Op names are global identifiers assigned during export (for example, `"linear"`, `"linear_1"`). The compressor must be `Quantizer` or `None`. +- **Eager mode** (`execution_mode="eager"`): Intercepts operations during a live forward pass. Op names are module-qualified identifiers that reflect the module hierarchy (for example, `"linear1.linear"`, `"linear2.linear"`). This mode supports both `Quantizer` and `KMeansPalettizer` as the compressor. + +If you plan to compress the model using one of `coreai-opt`'s compression techniques, choose the `execution_mode` you plan to use when compressing for inspection in order to identify the correct op and module names to use in the compression config. + +For more information on `graph` mode vs. `eager` mode, see [here](../quantization/overview.md#two-execution-modes-graph-and-eager). ## Basic Usage @@ -38,97 +45,125 @@ inspector = ModelInspector( compressor=Quantizer, ) -# Print a module-hierarchy tree showing ops, source locations, and connectivity +# Print a module-hierarchy tree showing ops, connectivity, and source locations print(inspector.format_summary()) ``` -Note the use of `compressor=Quantizer` in the list of arguments to `ModelInspector`. This filters the list of ops captured and displayed by `ModelInspector` to only those operations which are registered for compressibility by `Quantizer`. Omitting this argument allows for all ops to be captured and displayed. +Pass `colorize=False` to suppress ANSI color codes, for example when writing to a file. + +:::{note} +Note the use of `compressor=Quantizer`. This filters the captured and displayed ops to those registered as compressible by `Quantizer`. Omit this argument to capture and display all ops. +::: The above code produces output like the following (colors omitted for brevity): ```text -Legend: ■ module_name (module_type) ◆ op_name [op_type] - -(MyModel) - module inputs: linear - module outputs: linear_1 +Legend: + ■ module_name (module_type) ◆ op_name [op_type] + + op inputs: {I: producer[N]} — I = op_input_spec index; N = output slot of the producing op + op states: param_name — model parameter or buffer + op outputs: {N: [consumers]} — N = output slot index; consumers = ops receiving that output + untracked_N — input tensor whose producer was not intercepted (e.g. raw attribute or global tensor); still quantizable via op_input_spec + module inputs: {I: [op[N], ...]} — I = module_input_spec index; op[N] = op and its input slot receiving data from outside; absent keys = non-quantizable + module outputs: {I: op[N]} — I = module_output_spec index; op[N] = op and its output slot leaving the module; absent keys = non-quantizable + +(__main__.MyModel) + module inputs: {0: [linear[0]]} + module outputs: {0: linear_1[0]} ├── ■ linear1 (torch.nn.modules.linear.Linear) -│ module inputs: linear -│ module outputs: linear +│ module inputs: {0: [linear[0]]} +│ module outputs: {0: linear[0]} │ └── ◆ linear [linear] -│ op inputs: x, linear1_weight, linear1_bias -│ op outputs: relu -│ filepath: my_model.py:16 -│ code: x = self.linear1(x) +│ op inputs: {0: x[0]} +│ op states: weight, bias +│ op outputs: {0: [relu]} ├── ■ relu (torch.nn.modules.activation.ReLU) -│ module inputs: relu -│ module outputs: relu +│ module inputs: {0: [relu[0]]} +│ module outputs: {0: relu[0]} └── ■ linear2 (torch.nn.modules.linear.Linear) - module inputs: linear_1 - module outputs: linear_1 + module inputs: {0: [linear_1[0]]} + module outputs: {0: linear_1[0]} └── ◆ linear_1 [linear] - op inputs: relu, linear2_weight, linear2_bias - op outputs: output - filepath: my_model.py:18 - code: x = self.linear2(x) + op inputs: {0: relu[0]} + op states: weight, bias + op outputs: {0: [output]} ``` -The output shows the model's module hierarchy and the ops within each module. Note in particular that since `relu` is not a registered compressible op by `Quantizer`, it does not show up as an operation within the `ReLU` module. +Note that `relu` does not appear as an operation (`◆`) within the `relu` module, because `ReLU` is not a compressible op in `Quantizer`. It still appears as a module node (`■`) and in connectivity lines such as `op outputs: {0: [relu]}` and `op inputs: {0: relu[0]}`, because the relu tensor passes through and connects the two linear ops. -Reading the tree: +## Reading the Tree -- **Module name** and **module type** appear on module lines: `■ module_name (module_type)`. For example, `■ linear1 (torch.nn.modules.linear.Linear)` — `"linear1"` is the module name (usable in `module_name_configs`) and `"torch.nn.modules.linear.Linear"` is the module type (usable in `module_type_configs`). -- **Op name** and **op type** appear on operation lines: `◆ op_name [op_type]`. For example, `◆ linear_1 [linear]` — `"linear_1"` is the op name (usable in `op_name_config`) and `"linear"` is the op type (usable in `op_type_config`). -- **Op inputs/outputs** show connectivity between operations. -- **filepath** and **code** (shown for user-defined modules) show where in your source code the operation originates. +### Module lines -Using these strings directly in a config: +Module lines use the form `■ module_name (module_type)`. For example, `■ linear1 (torch.nn.modules.linear.Linear)`: -```python -config = QuantizerConfig( - # Target a specific module by name - module_name_configs={ - "linear1": ModuleQuantizerConfig(...), - }, - # Target all modules of a given type - module_type_configs={ - "torch.nn.modules.linear.Linear": ModuleQuantizerConfig(...), - }, -) +- `"linear1"` is the module name, usable in `module_name_configs`. +- `"torch.nn.modules.linear.Linear"` is the module type, usable in `module_type_configs`. -# Op-level targeting within a ModuleQuantizerConfig -config = QuantizerConfig( - global_config=ModuleQuantizerConfig( - # Target a specific op by name - op_name_config={ - "linear_1": OpQuantizerConfig(...), - }, - # Target all ops of a given type - op_type_config={ - "linear": OpQuantizerConfig(...), - }, - ), -) -``` +**Module boundaries** appear indented under the module header: + +- `module inputs: {I: [op[N], ...]}` — The activations entering this module from outside. `I` is the position in the module's input spec (matching `module_input_spec` in a config), `op` is the name of the first compressible op inside the module that receives data at that position, and `N` is the input slot on that op. A single external input can fan out to multiple ops. Keys absent from this dict correspond to non-quantizable positions (for example, state tensors or unused arguments). +- `module outputs: {I: op[N]}` — The activations leaving this module. `I` is the position in the module's output spec, `op` is the compressible op producing that output, and `N` is the op's output slot. Absent keys correspond to non-quantizable positions. + +### Op lines + +Op lines use the form `◆ op_name [op_type]`. For example, `◆ linear_1 [linear]`: + +- `"linear_1"` is the op name, usable in `op_name_config`. +- `"linear"` is the op type, usable in `op_type_config`. + +**Op connectivity** appears indented under the op header: + +- `op inputs: {I: producer[N]}` — Activation inputs only (parameters and buffers are on a separate line). `I` is the argument position (matching `op_input_spec` in a config), `producer` is the name of the op that produced this tensor, and `N` is the output slot of that producer. For example, `{0: relu[0]}` means argument 0 comes from output slot 0 of the `relu` op. +- `op states: param_name, ...` — Model parameters and buffers consumed by this op. This line is omitted if the op takes no states. +- `op outputs: {N: [consumer1, consumer2, ...]}` — `N` is the output slot index, and the list contains the names of all ops consuming that output. +- `untracked_N` — Appears in place of a producer name when the input tensor's origin was not intercepted (for example, a raw module attribute or global tensor). These tensors are still quantizable via `op_input_spec`. +- `filepath` and `code` — Source file and line of the call that produced the op, shown as dim text. + +## Eager Mode -Pass `colorize=False` to suppress ANSI color codes (e.g., when writing to a file). +To inspect using eager mode, pass `execution_mode="eager"`. The same `MyModel` example above yields: + +```text +(__main__.MyModel) + module inputs: {0: [linear1.linear[0]]} + module outputs: {0: linear2.linear[0]} +├── ■ linear1 (torch.nn.modules.linear.Linear) +│ module inputs: {0: [linear1.linear[0]]} +│ module outputs: {0: linear1.linear[0]} +│ └── ◆ linear1.linear [linear] +│ op inputs: {0: input_0} +│ op states: weight, bias +│ op outputs: {0: [relu.relu]} +│ filepath: my_model.py:16 +├── ■ relu (torch.nn.modules.activation.ReLU) +│ module inputs: {0: [relu.relu[0]]} +│ module outputs: {0: relu.relu[0]} +└── ■ linear2 (torch.nn.modules.linear.Linear) + module inputs: {0: [linear2.linear[0]]} + module outputs: {0: linear2.linear[0]} + └── ◆ linear2.linear [linear] + op inputs: {0: relu.relu[0]} + op states: weight, bias + op outputs: {0: [output_0]} + filepath: my_model.py:18 +``` ## Querying Operations by Config Key -Once you have reviewed the full summary to see what names and types are present, you can use query methods to check which operations would be matched by a specific name or type pattern. This is useful for verifying your config will target the intended ops before applying compression. +Once you have reviewed the full summary to see what names and types are present, use the query methods to check which operations a specific pattern matches. This is useful for verifying that a config targets the intended ops before applying compression. Each query method returns a tuple of {class}`~coreai_opt.inspection.OpInfo` objects matching the filter. The method names correspond directly to the config keys they help populate. -From the Basic Usage summary, this model exposes: +From the graph mode summary above, this model exposes: - **Op types**: `linear` - **Op names**: `linear`, `linear_1` - **Module types**: `torch.nn.modules.linear.Linear`, `torch.nn.modules.activation.ReLU` - **Module names**: `linear1`, `relu`, `linear2` -Op names and module names can be passed as a literal name or as a regex following [Python re syntax](https://docs.python.org/3/library/re.html) for wildcard matching; the pattern is matched against the entire string. The matching methodology is identical to how compression config entries match modules and ops in a model, allowing the user to see exactly which modules or ops would be matched given a particular string. - -Each query method returns a tuple of {class}`~coreai_opt.inspection.OpInfo` objects matching the filter: +Op names and module names can be passed as a literal string or as a regex following [Python re syntax](https://docs.python.org/3/library/re.html) for wildcard matching. The pattern is matched against the full string. The matching behavior is identical to how compression config entries match modules and ops, so you can see exactly which ops a given pattern would select. **By op type** — exact-string match against `op_type_config` keys: @@ -151,10 +186,10 @@ inspector.get_matched_ops_for_module_name( ) # matches the op in module "linear1" inspector.get_matched_ops_for_module_name( "linear[12]" -) # matches ops in modules "linear1" and "linear2" +) # matches ops in "linear1" and "linear2" ``` -Each returned {class}`~coreai_opt.inspection.OpInfo` provides `op_name`, `op_type`, and `module_stack` (the nesting of modules containing the op): +Each returned {class}`~coreai_opt.inspection.OpInfo` provides `op_name`, `op_type`, `module_stack`, `inputs`, `outputs`, and `is_state`. The `module_stack` is a tuple of {class}`~coreai_opt.inspection.ModuleContext` entries from outermost to innermost module: ```python >>> for op in inspector.get_matched_ops_for_op_type("linear"): @@ -166,6 +201,37 @@ Each returned {class}`~coreai_opt.inspection.OpInfo` provides `op_name`, `op_typ module: linear2 (torch.nn.modules.linear.Linear) ``` +`OpInfo.inputs` is a tuple of {class}`~coreai_opt.inspection.InputEdge` objects, one per input argument position. Each `InputEdge` carries the producing `OpInfo` and the output slot index (`output_idx`) of that producer. State inputs (parameters, buffers) are interleaved in the tuple at their actual argument positions, and their corresponding `InputEdge` objects have `is_state=True`. + +Using these strings directly in a config: + +```python +config = QuantizerConfig( + # Target a specific module by name + module_name_configs={ + "linear1": ModuleQuantizerConfig(...), + }, + # Target all modules of a given type + module_type_configs={ + "torch.nn.modules.linear.Linear": ModuleQuantizerConfig(...), + }, +) + +# Op-level targeting within a ModuleQuantizerConfig +config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + # Target a specific op by name + op_name_config={ + "linear_1": OpQuantizerConfig(...), + }, + # Target all ops of a given type + op_type_config={ + "linear": OpQuantizerConfig(...), + }, + ), +) +``` + ## Navigating the Module Hierarchy For programmatic access to the inspector's data structures, the {class}`~coreai_opt.inspection.ModelSummary` exposes a {class}`~coreai_opt.inspection.ModuleInfo` tree that mirrors the `nn.Module` hierarchy. These types are publicly exported from `coreai_opt.inspection` for use in custom analysis or tooling. @@ -180,11 +246,24 @@ linear2: torch.nn.modules.linear.Linear, 1 direct ops ``` ```python -# look up a specific submodule +# Look up a specific submodule linear2_module = root.get_submodule("linear2") -# get all ops under this subtree (depth-first) +# Get all ops under this subtree (depth-first) linear2_ops = linear2_module.all_ops() ``` `ModuleInfo` supports the same iteration patterns as `nn.Module`: `children()`, `named_children()`, `modules()`, `named_modules()`, and `get_submodule()`. + +`ModuleInfo` also exposes the module boundary connectivity described in the tree: + +- `input_ops` — dict mapping module input spec index to a list of {class}`~coreai_opt.inspection.BoundaryEdge` objects, each holding the op and input slot receiving data from outside the module. +- `output_ops` — dict mapping module output spec index to a single {class}`~coreai_opt.inspection.BoundaryEdge`, holding the op and output slot whose tensor leaves the module. + +```python +# Inspect boundary connectivity for a submodule +linear1_module = root.get_submodule("linear1") +for idx, edges in linear1_module.input_ops.items(): + for edge in edges: + print(f" module input {idx} -> {edge.op.op_name}[{edge.index}]") +``` diff --git a/src/coreai_opt/inspection/_common.py b/src/coreai_opt/inspection/_common.py index ec4dfba..c39e398 100644 --- a/src/coreai_opt/inspection/_common.py +++ b/src/coreai_opt/inspection/_common.py @@ -16,8 +16,9 @@ def _get_or_create_child(parent: ModuleInfo, module_name: str, module_type: str) -> ModuleInfo: """Get an existing child module or create a new one.""" - if module_name not in parent.child_modules: - parent.child_modules[module_name] = ModuleInfo( + child = parent.child_modules.get(module_name) + if child is None: + child = ModuleInfo( module_name=module_name, module_type=module_type, child_modules={}, @@ -25,7 +26,8 @@ def _get_or_create_child(parent: ModuleInfo, module_name: str, module_type: str) input_ops={}, output_ops={}, ) - return parent.child_modules[module_name] + parent.child_modules[module_name] = child + return child def build_module_tree( diff --git a/src/coreai_opt/inspection/_eager_mode.py b/src/coreai_opt/inspection/_eager_mode.py index 2e1c30d..ace497f 100644 --- a/src/coreai_opt/inspection/_eager_mode.py +++ b/src/coreai_opt/inspection/_eager_mode.py @@ -54,6 +54,22 @@ ) +class _TensorProducerMap: + """Maps a tensor's (id, version) to its producer edge, auto-evicting on GC.""" + + def __init__(self) -> None: + self._entries: dict[TensorIdVersion, InputEdge] = {} + + def register(self, tensor: torch.Tensor, edge: InputEdge) -> None: + """Associate ``edge`` with ``tensor``; remove the entry once ``tensor`` is GC'd.""" + key = TensorIdVersion(id(tensor), tensor._version) + self._entries[key] = edge + weakref.finalize(tensor, self._entries.pop, key, None) + + def get(self, key: TensorIdVersion) -> InputEdge | None: + return self._entries.get(key) + + class _EagerOpDiscoveryMode(TorchFunctionMode): """TorchFunctionMode that discovers ops during a forward pass. @@ -91,7 +107,7 @@ def __init__( # Tensor connectivity: (id, version) → InputEdge (producing op + output slot). # Tracks activation tensors only; state tensors are tracked separately below. - self._tensor_producers: dict[TensorIdVersion, InputEdge] = {} + self._tensor_producers = _TensorProducerMap() self._states_to_names = { id(state): name @@ -147,22 +163,26 @@ def _add_op(self, op_info: OpInfo) -> None: self._seen_op_names.add(op_info.op_name) self.all_ops.append(op_info) - def _get_or_create_ephemeral(self, tensor_id: int) -> InputEdge: - """Return the ephemeral InputEdge for this tensor id, creating one if needed.""" - op_info = self._ephemeral_op_infos.get(tensor_id) - if op_info is None: - op_info = OpInfo( - op_name=f"untracked_{self._ephemeral_counter}", - op_type=None, - module_stack=(), - source_frames=(), - inputs=(), - outputs={}, - is_state=False, - ) - self._ephemeral_op_infos[tensor_id] = op_info - self._ephemeral_counter += 1 - return InputEdge(op=op_info, output_idx=None) + def _create_and_register_ephemeral_tensor(self, tensor: torch.Tensor) -> InputEdge: + """Create and return the InputEdge for this ephemeral tensor, registering it in + self._tensor_producers. + + An ephemeral tensor is defined as a tensor that is not an input to the top level model + module, a state tensor, or a tensor produced by a previous operation. + """ + op_info = OpInfo( + op_name=f"untracked_{self._ephemeral_counter}", + op_type=None, + module_stack=(), + source_frames=(), + inputs=(), + outputs={}, + is_state=False, + ) + input_edge = InputEdge(op=op_info, output_idx=None) + self._tensor_producers.register(tensor, input_edge) + self._ephemeral_counter += 1 + return input_edge def _resolve_boundary_tensor(self, t: torch.Tensor) -> InputEdge | None: """Resolve a module-boundary tensor to its producer entry. @@ -182,7 +202,7 @@ def _resolve_boundary_tensor(self, t: torch.Tensor) -> InputEdge | None: return entry if id(t) in self._states_to_names: return None - return self._get_or_create_ephemeral(id(t)) + return self._create_and_register_ephemeral_tensor(t) def _enter_module(self, name: str) -> Callable: def hook(module: nn.Module, inputs: Any) -> None: @@ -225,9 +245,7 @@ def hook(module: nn.Module, inputs: Any, outputs: Any) -> None: strict=False, ): if entry is not None: - key = TensorIdVersion(id(tensor), tensor._version) - self._tensor_producers[key] = entry - weakref.finalize(tensor, self._tensor_producers.pop, key, None) + self._tensor_producers.register(tensor, entry) self.parents.pop() self.traversed_modules.add(module) @@ -246,9 +264,7 @@ def _capture_input_tensors(self, module: nn.Module, inputs: Any) -> None: is_state=False, ) self._add_op(op_info) - key = TensorIdVersion(id(tensor), tensor._version) - self._tensor_producers[key] = InputEdge(op=op_info, output_idx=None) - weakref.finalize(tensor, self._tensor_producers.pop, key, None) + self._tensor_producers.register(tensor, InputEdge(op=op_info, output_idx=None)) def _capture_output_tensors(self, module: nn.Module, inputs: Any, outputs: Any) -> None: """Create output-like ops for each module-level output tensor.""" @@ -296,36 +312,43 @@ def _extract_source_frames(self) -> tuple[SourceFrame, ...]: # Reverse: outermost forward first (matching graph mode order) return tuple(reversed(frames)) - def _resolve_inputs( - self, input_tensor_keys: tuple[TensorIdVersion, ...] - ) -> tuple[InputEdge, ...]: + def _get_or_create_state_edge(self, state_id: int) -> InputEdge: + """Return the edge for a registered state tensor, creating its OpInfo on first use. + + State ops are keyed by object id alone — version is irrelevant because states are not + produced by ops and their identity is stable for the model's lifetime. The created op + carries an empty ``module_stack``, keeping it out of the module tree and boundary lists. + """ + op_info = self._state_op_infos.get(state_id) + if op_info is None: + op_info = OpInfo( + op_name=self._states_to_names[state_id], + op_type=None, + module_stack=(), + source_frames=(), + inputs=(), + outputs={}, + is_state=True, + ) + self._add_op(op_info) + self._state_op_infos[state_id] = op_info + return InputEdge(op=op_info, output_idx=None) + + def _resolve_inputs(self, input_tensors: list[torch.Tensor]) -> tuple[InputEdge, ...]: """Look up which previously-recorded ops produced the input tensors. Returns an ordered tuple of :class:`InputEdge` objects (duplicates preserved), each carrying the producing op and its output slot. """ - input_edges: list[InputEdge] = [] - for key in input_tensor_keys: + + def resolve_input(tensor): + key = TensorIdVersion(id(tensor), tensor._version) if key.id in self._states_to_names: - # State tensor: look up or create by id only — version is irrelevant. - op_info = self._state_op_infos.get(key.id) - if op_info is None: - op_info = OpInfo( - op_name=self._states_to_names[key.id], - op_type=None, - module_stack=(), - source_frames=(), - inputs=(), - outputs={}, - is_state=True, - ) - self._add_op(op_info) - self._state_op_infos[key.id] = op_info - input_edges.append(InputEdge(op=op_info, output_idx=None)) + return self._get_or_create_state_edge(key.id) else: entry = self._tensor_producers.get(key) if entry is not None: - input_edges.append(entry) + return entry else: # Unknown tensor: not a registered state and not produced by any # intercepted op (e.g. a raw tensor attribute or a global tensor). @@ -333,16 +356,14 @@ def _resolve_inputs( # complete and the correct arg index appears in the formatted output. # Ephemeral ops are NOT added to all_ops — they never appear as their # own nodes in the summary tree. - input_edges.append(self._get_or_create_ephemeral(key.id)) + return self._create_and_register_ephemeral_tensor(tensor) - return tuple(input_edges) + return tuple(resolve_input(inp) for inp in input_tensors) def _record_outputs(self, out: Any, op_info: OpInfo) -> None: """Record that op_info produced these output tensors.""" for idx, tensor in enumerate(flatten_tensors_to_list(out)): - key = TensorIdVersion(id(tensor), tensor._version) - self._tensor_producers[key] = InputEdge(op=op_info, output_idx=idx) - weakref.finalize(tensor, self._tensor_producers.pop, key, None) + self._tensor_producers.register(tensor, InputEdge(op=op_info, output_idx=idx)) def _register_as_consumer(self, inputs: tuple[InputEdge, ...], consumer: OpInfo) -> None: """Append consumer to each producer's outputs dict at the given slot. @@ -372,13 +393,10 @@ def __torch_function__( if kwargs is None: kwargs = {} - # Snapshot input tensor keys BEFORE func() executes. + # Process input tensors before func() executes. # Critical for in-place ops: func() mutates _version, but # the producer was recorded at the pre-mutation version. - input_tensor_keys = tuple( - TensorIdVersion(id(t), t._version) - for t in flatten_tensors_to_list((*args, *kwargs.values())) - ) + input_edges = self._resolve_inputs(flatten_tensors_to_list((*args, *kwargs.values()))) out = func(*args, **kwargs) @@ -414,7 +432,6 @@ def __torch_function__( module_stack = self._get_module_stack() source_frames = self._extract_source_frames() - input_edges = self._resolve_inputs(input_tensor_keys) op_info = OpInfo( op_name=op_name, @@ -427,7 +444,6 @@ def __torch_function__( ) self._add_op(op_info) - self._register_as_consumer(input_edges, op_info) self._record_outputs(out, op_info) diff --git a/tests/quantization/test_is_state_node.py b/tests/quantization/test_is_coreai_compressed_state_node.py similarity index 91% rename from tests/quantization/test_is_state_node.py rename to tests/quantization/test_is_coreai_compressed_state_node.py index 674e28a..a9ee0e0 100644 --- a/tests/quantization/test_is_state_node.py +++ b/tests/quantization/test_is_coreai_compressed_state_node.py @@ -3,9 +3,9 @@ # 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 _is_state_node in PT2E annotation utils. +"""Tests for is_coreai_compressed_state_node in PT2E annotation utils. -Verifies that _is_state_node correctly identifies state nodes (weights, +Verifies that is_coreai_compressed_state_node correctly identifies state nodes (weights, parameters, compressed weight decompression ops) and rejects activation nodes. """ @@ -14,7 +14,7 @@ import pytest import torch -from coreai_opt._utils.fx_utils import is_coreai_compressed_state_node as is_state_node +from coreai_opt._utils.fx_utils import is_coreai_compressed_state_node from tests.test_utils.general import COREAI_AVAILABLE @@ -24,7 +24,7 @@ def _make_node( op_name: str | None = None, args: tuple = (), ) -> Mock: - """Create a mock FX node with the attributes needed by _is_state_node. + """Create a mock FX node with the attributes needed by is_coreai_compressed_state_node. Args: op (str): The FX node op type (e.g., "get_attr", "call_function"). @@ -34,7 +34,7 @@ def _make_node( args (tuple): Node args (other mock nodes or values). Returns: - Mock: A mock node suitable for passing to _is_state_node. + Mock: A mock node suitable for passing to is_coreai_compressed_state_node. """ node = Mock(spec=torch.fx.Node) node.op = op @@ -49,23 +49,23 @@ def _make_node( return node -class TestIsStateNode: +class TestIsCoreAICompressedStateNode: def test_get_attr_is_state(self): """get_attr nodes (direct parameter access) are state.""" node = _make_node("get_attr") - assert is_state_node(node) is True + assert is_coreai_compressed_state_node(node) is True def test_placeholder_is_not_state(self): """Placeholder nodes (model inputs) are not state.""" node = _make_node("placeholder") - assert is_state_node(node) is False + assert is_coreai_compressed_state_node(node) is False def test_lut_to_dense_is_state(self): """coreai.lut_to_dense call_function is state (palettized weights).""" indices = _make_node("get_attr") lut = _make_node("get_attr") node = _make_node("call_function", "coreai", "lut_to_dense", args=(indices, lut)) - assert is_state_node(node) is True + assert is_coreai_compressed_state_node(node) is True def test_shift_scale_with_lut_input_is_state(self): """constexpr_blockwise_shift_scale fed by lut_to_dense is state.""" @@ -76,7 +76,7 @@ def test_shift_scale_with_lut_input_is_state(self): node = _make_node( "call_function", "coreai", "constexpr_blockwise_shift_scale", args=(lut_node, scale) ) - assert is_state_node(node) is True + assert is_coreai_compressed_state_node(node) is True def test_shift_scale_is_state(self): """constexpr_blockwise_shift_scale is always state. This op is only @@ -87,7 +87,7 @@ def test_shift_scale_is_state(self): node = _make_node( "call_function", "coreai", "constexpr_blockwise_shift_scale", args=(data, scale) ) - assert is_state_node(node) is True + assert is_coreai_compressed_state_node(node) is True def test_aten_op_with_all_state_inputs_is_not_state(self): """An aten call_function whose inputs are all get_attr is NOT state. @@ -97,7 +97,7 @@ def test_aten_op_with_all_state_inputs_is_not_state(self): weight = _make_node("get_attr") bias = _make_node("get_attr") node = _make_node("call_function", "aten", "add", args=(weight, bias)) - assert is_state_node(node) is False + assert is_coreai_compressed_state_node(node) is False def _find_coreai_nodes(gm: torch.fx.GraphModule, op_name: str) -> list[torch.fx.Node]: @@ -115,7 +115,7 @@ def _find_coreai_nodes(gm: torch.fx.GraphModule, op_name: str) -> list[torch.fx. @pytest.mark.skipif(not COREAI_AVAILABLE, reason="Requires coreai") -class TestIsStateNodeIntegration: +class TestIsCoreAICompressedStateNodeIntegration: @pytest.mark.seed def test_joint_compression_lut_to_dense_not_quantized( self, simple_conv_linear_model, simple_model_input @@ -168,7 +168,7 @@ def test_joint_compression_lut_to_dense_not_quantized( # Each lut_to_dense must be recognized as state for node in lut_nodes: - assert is_state_node(node) is True, ( + assert is_coreai_compressed_state_node(node) is True, ( f"lut_to_dense node {node.name} not identified as state" ) @@ -255,7 +255,7 @@ def test_joint_compression_lut_quantized_shift_scale_not_quantized( # Both op types must be recognized as state for node in lut_nodes + shift_scale_nodes: - assert is_state_node(node) is True, ( + assert is_coreai_compressed_state_node(node) is True, ( f"{node.target._opname} node {node.name} not identified as state" ) diff --git a/tests/test_inspection.py b/tests/test_inspection.py index 483f3e4..ba73caa 100644 --- a/tests/test_inspection.py +++ b/tests/test_inspection.py @@ -1395,9 +1395,9 @@ def test_weakref_removes_producer_on_dealloc(self) -> None: ) mode._record_outputs(tensor, op) key = TensorIdVersion(id(tensor), tensor._version) - assert mode._tensor_producers[key] == InputEdge(op=op, output_idx=0) + assert mode._tensor_producers.get(key) == InputEdge(op=op, output_idx=0) del tensor - assert key not in mode._tensor_producers + assert mode._tensor_producers.get(key) is None def test_input_ops_disambiguates_multi_output_external_producer(self) -> None: """Verify input_ops correctly separates consumers of distinct outputs from the same op.