diff --git a/coreai_torch/_utils.py b/coreai_torch/_utils.py index 436034a..a2211b0 100644 --- a/coreai_torch/_utils.py +++ b/coreai_torch/_utils.py @@ -1596,15 +1596,15 @@ def _find_program_for( name: str, op_name: str, programs: dict[str, ExportedProgram], -) -> ExportedProgram: - """Find the nearest ancestor program that contains *op_name*.""" +) -> ExportedProgram | None: + """Find the nearest ancestor program that contains *op_name*, or ``None``.""" for ancestor in _ancestor_paths(name): if ( ancestor in programs and _find_custom_op_node(programs[ancestor], op_name) is not None ): return programs[ancestor] - raise ValueError(f"Custom op for '{name}' not found in any ancestor program") + return None def _reverse_lookup(mapping: dict[str, str], value: str) -> str | None: diff --git a/coreai_torch/externalize.py b/coreai_torch/externalize.py index 03b201d..97eaf41 100644 --- a/coreai_torch/externalize.py +++ b/coreai_torch/externalize.py @@ -431,6 +431,18 @@ def __iter__(self) -> Iterator[_PreparedModule]: name: str = mod._externalize_name # type: ignore[attr-defined] op_name: str = mod._externalize_op_name # type: ignore[attr-defined] parent_ep = _find_program_for(name, op_name, self._programs) + if parent_ep is None: + # Marked submodule is not invoked in the exported graph. + # Skip it; its forward will be restored by _restore_externalized. + warnings.warn( + f"\n[WARN] coreai_torch.externalize: skipping unused submodule '{name}'.\n" + f" It matched an externalize_modules target class but is not " + f"reachable from the exported graph.\n" + f" Action: remove it from the model passed to add_pytorch_module, " + f"or ignore if intentional.\n", + stacklevel=2, + ) + continue preps = _prepare_module_export(mod, parent_ep) for prep in preps: prep._program_registry = self diff --git a/tests/test_externalize.py b/tests/test_externalize.py index fe8a74e..59661d5 100644 --- a/tests/test_externalize.py +++ b/tests/test_externalize.py @@ -3683,3 +3683,108 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: await _validate_numerics( coreai_program, model, sample, input_names=("x", "indices") ) + + +@pytest.mark.ir +def test_externalize_unused_submodule_ir() -> None: + """An externalizable submodule that the model's forward never calls is skipped. + + Previously, a registered submodule that did not appear in the exported graph + caused the externalize pipeline to raise ``ValueError: Custom op for ... + not found in any ancestor program``. The pipeline should instead warn and + proceed, lowering the rest of the model normally. + """ + + class InnerModule(nn.Module): + def __init__(self): + super().__init__() + self.fc = nn.Linear(4, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.fc(x)) + + class OuterModel(nn.Module): + def __init__(self): + super().__init__() + self.pre = nn.Linear(4, 4) + # Registered as a submodule but intentionally never invoked. + self.unused = InnerModule() + self.post = nn.Linear(4, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.post(self.pre(x)) + + torch.manual_seed(42) + model = OuterModel().eval() + sample = (torch.randn(2, 4),) + + with pytest.warns(UserWarning, match="skipping unused submodule"): + converter = TorchConverter().add_pytorch_module( + model, + export_fn=lambda m: torch.export.export(m, args=sample).run_decompositions( + get_decomp_table() + ), + externalize_modules=[InnerModule], + ) + coreai_program = converter.to_coreai() + + check_file = """ + // CHECK-LABEL: module { + // CHECK-NOT: coreai.graph noinline + // CHECK: coreai.graph @main( + // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x4xf32> {coreai.name = "x"} + // CHECK-SAME: ) -> (tensor<2x4xf32> + // CHECK: %[[MM0:[0-9a-z_]+]] = coreai.decomposable.broadcasting_batch_matmul %[[ARG0]], + // CHECK-SAME: : (tensor<2x4xf32>, tensor<4x4xf32>) -> tensor<2x4xf32> + // CHECK: %[[ADD0:[0-9a-z_]+]] = coreai.decomposable.broadcasting_add %[[MM0]], + // CHECK-SAME: : (tensor<2x4xf32>, tensor<4xf32>) -> tensor<2x4xf32> + // CHECK: %[[MM1:[0-9a-z_]+]] = coreai.decomposable.broadcasting_batch_matmul %[[ADD0]], + // CHECK-SAME: : (tensor<2x4xf32>, tensor<4x4xf32>) -> tensor<2x4xf32> + // CHECK: %[[ADD1:[0-9a-z_]+]] = coreai.decomposable.broadcasting_add %[[MM1]], + // CHECK-SAME: : (tensor<2x4xf32>, tensor<4xf32>) -> tensor<2x4xf32> + // CHECK-NOT: coreai.invoke + // CHECK-NOT: coreai.relu + // CHECK: coreai.output %[[ADD1]] : tensor<2x4xf32> + // CHECK: } + // CHECK-NOT: coreai.graph + // CHECK: } + """ + filecheck_pattern(str(coreai_program), check_file=check_file) + + +async def test_externalize_unused_submodule_numerics() -> None: + """Numerics: unused externalizable submodule does not affect output.""" + + class InnerModule(nn.Module): + def __init__(self): + super().__init__() + self.fc = nn.Linear(4, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.fc(x)) + + class OuterModel(nn.Module): + def __init__(self): + super().__init__() + self.pre = nn.Linear(4, 4) + self.unused = InnerModule() + self.post = nn.Linear(4, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.post(self.pre(x)) + + torch.manual_seed(42) + model = OuterModel().eval() + sample = (torch.randn(2, 4),) + + with pytest.warns(UserWarning, match="skipping unused submodule"): + converter = TorchConverter().add_pytorch_module( + model, + export_fn=lambda m: torch.export.export(m, args=sample).run_decompositions( + get_decomp_table() + ), + externalize_modules=[InnerModule], + ) + coreai_program = converter.to_coreai() + + await _validate_numerics(coreai_program, model, sample)