From 2e2e12cd6e5f7fbfbe9f4af37dc8d605cd9a88cf Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Tue, 29 Sep 2026 14:10:32 -0700 Subject: [PATCH] Qualcomm AI Engine Direct - accept an ExportedProgram in to_edge_transform_and_lower_to_qnn to_edge_transform_and_lower_to_qnn only took a module and always re-exported it internally, so a caller that has already captured a program - with its own dynamic shapes, fx fixes or graph surgery - had no way to hand that program in and had to reimplement the transform/partition/lower sequence against backends.qualcomm._passes directly. Accept ExportedProgram (and Dict[str, ExportedProgram]) and skip the internal re-export for those entries. Module inputs are unchanged. inputs is now Optional, matching the documented inputs=None usage, and supplying inputs or dynamic_shapes alongside an already-captured program warns rather than silently dropping them. A shape assertion alone cannot detect a re-export, since re-exporting exported.module() reproduces the original graph. The tests therefore spy on torch.export.export and require that lowering never calls it; both the unittest and the rework-framework test fail against a variant that accepts the program and re-exports it. --- backends/qualcomm/tests/rework/src/utils.py | 13 ++++ backends/qualcomm/tests/rework/utils/test.py | 4 ++ backends/qualcomm/tests/test_passes.py | 69 +++++++++++++++++++- backends/qualcomm/utils/utils.py | 48 ++++++++++---- 4 files changed, 121 insertions(+), 13 deletions(-) diff --git a/backends/qualcomm/tests/rework/src/utils.py b/backends/qualcomm/tests/rework/src/utils.py index 9202007b18a..148f78003d5 100644 --- a/backends/qualcomm/tests/rework/src/utils.py +++ b/backends/qualcomm/tests/rework/src/utils.py @@ -633,3 +633,16 @@ def test(subtests): assert len(weight_fqs) > 0, "no int8-range weight FQ found for fp16a8w" # should complete without error __class__._get_converted_module(module, prepared, inputs) + + +class LoweringWithExportedProgram: + @staticmethod + def test(compile_spec): + from unittest import mock + + module = _UtilsModel() + inputs = (torch.randn(1, 4, 8, 8), torch.randn(1, 4, 8, 8)) + exported = torch.export.export(module, inputs, strict=True) + with mock.patch("torch.export.export", wraps=torch.export.export) as export_spy: + to_edge_transform_and_lower_to_qnn(exported, None, compile_spec) + export_spy.assert_not_called() diff --git a/backends/qualcomm/tests/rework/utils/test.py b/backends/qualcomm/tests/rework/utils/test.py index 8330f79ab88..093c69ab210 100644 --- a/backends/qualcomm/tests/rework/utils/test.py +++ b/backends/qualcomm/tests/rework/utils/test.py @@ -39,3 +39,7 @@ def test_skip_node_quantizer(subtests, quantizer, compile_spec): def test_qat(subtests): QAT.test(subtests) # noqa: F405 + + +def test_lowering_with_exported_program(compile_spec): + LoweringWithExportedProgram.test(compile_spec) # noqa: F405 diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index abeb82f0e13..05ed2c5b7cd 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -1,5 +1,5 @@ import unittest -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import torch from executorch.backends.qualcomm._passes import ( @@ -760,6 +760,73 @@ def forward(self, x): if "QNN" in str(e) or "qnn" in str(e): self.skipTest(f"QNN SDK not available: {e}") + def test_lowering_accepts_exported_program_without_reexport(self): + """to_edge_transform_and_lower_to_qnn must lower a caller-supplied + ExportedProgram as captured, never re-exporting it. + + Callers that have already captured and patched a program - HF transformers' + ExecuTorch exporter, for instance, which owns its own dynamic shapes and fx + fixes - cannot express that through a module argument. + + A shape assertion alone does not discriminate here: re-exporting + exported.module() reproduces the original graph, so the contract is checked + by spying on torch.export.export and requiring that lowering never calls it. + """ + + class RetraceSensitive(torch.nn.Module): + def __init__(self): + super().__init__() + self.narrow = False + + def forward(self, x): + return torch.relu(x[:, :1] if self.narrow else x) + + module = RetraceSensitive().eval() + exported = torch.export.export(module, (torch.randn(1, 4),), strict=True) + + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + with patch("torch.export.export", wraps=torch.export.export) as export_spy: + try: + edge = to_edge_transform_and_lower_to_qnn( + exported, None, compiler_specs + ) + except RuntimeError as e: + if "QNN" in str(e) or "qnn" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + export_spy.assert_not_called() + + # The lowered graph is the one that was handed in, not a fresh capture. + output_node = edge.exported_program("forward").graph.output_node() + self.assertEqual(tuple(output_node.args[0][0].meta["val"].shape), (1, 4)) + + def test_lowering_warns_when_inputs_ignored_for_exported_program(self): + """inputs and dynamic_shapes cannot affect an already-captured program, so + supplying them alongside one must warn rather than silently drop them.""" + + class Tiny(torch.nn.Module): + def forward(self, x): + return torch.relu(x) + + sample_input = (torch.randn(1, 4),) + exported = torch.export.export(Tiny().eval(), sample_input, strict=True) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + try: + with self.assertWarnsRegex(UserWarning, "inputs ignored"): + to_edge_transform_and_lower_to_qnn( + exported, sample_input, compiler_specs + ) + except RuntimeError as e: + if "QNN" in str(e) or "qnn" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + def test_index_put_int64_value_not_quantized(self): """QNN's IndexPut annotator must skip a non-float (int64) value arg. diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index f78d47525ad..b0a06d79055 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -346,10 +346,12 @@ def to_edge_transform_and_lower_to_qnn( module: Union[ torch.nn.Module, torch.fx.GraphModule, + ExportedProgram, Dict[str, torch.nn.Module], Dict[str, torch.fx.GraphModule], + Dict[str, ExportedProgram], ], - inputs: Union[Tuple[torch.Tensor], Dict[str, Tuple[torch.Tensor]]], + inputs: Optional[Union[Tuple[torch.Tensor], Dict[str, Tuple[torch.Tensor]]]], compiler_specs: Union[List[Any], Dict[str, List[Any]]], constant_methods: Optional[Dict[str, Any]] = None, dynamic_shapes: Optional[Dict] = None, @@ -365,17 +367,20 @@ def to_edge_transform_and_lower_to_qnn( Transforms and lowers a given PyTorch module to the QNN backend. Args: - module (Union[torch.nn.Module, torch.fx.GraphModule,Dict[str, torch.nn.Module], Dict[str, torch.fx.GraphModule]]): - The PyTorch module or fx.GraphModule to be transformed. - inputs (Union[Tuple[torch.Tensor], Dict[str, Tuple[torch.Tensor]]]): - The input tensors for the module. + module (Union[torch.nn.Module, torch.fx.GraphModule, ExportedProgram, Dict[str, torch.nn.Module], Dict[str, torch.fx.GraphModule], Dict[str, ExportedProgram]]): + The PyTorch module or fx.GraphModule to be transformed, or an already captured + ExportedProgram to lower as-is. + inputs (Optional[Union[Tuple[torch.Tensor], Dict[str, Tuple[torch.Tensor]]]]): + The input tensors for the module. Ignored, and may be None, for graphs supplied + as an ExportedProgram: the program already carries its own example inputs. compiler_specs (Union[List[Any], Dict[str, List[Any]]]): Compiler specifications for Qualcomm AI Engine Direct. constant_methods (Optional[Dict[str, Any]]): An optional dictionary mapping method names to constant values returned by those methods in eager mode. Often used to store configuration information on Edge models. dynamic_shapes (Optional[Dict]): - Information about dynamic shapes. + Information about dynamic shapes. Ignored for graphs supplied as an + ExportedProgram, whose shape constraints were fixed at capture time. dep_table (Optional[Dict]): Dependency table for the transformation passes. passes_job (Optional[Union[OrderedDict, Dict[str, OrderedDict]]]): @@ -472,12 +477,31 @@ def ensure_graph_specific_dict(value, graph_names): } for graph_name, m in module.items(): - ep = torch.export.export( - m, - inputs[graph_name], - dynamic_shapes=dynamic_shapes[graph_name], - strict=True, - ) + if isinstance(m, ExportedProgram): + # Capture already happened; inputs and dynamic_shapes cannot influence a + # finished program, so say so rather than silently dropping them. + ignored = [ + name + for name, value in ( + ("inputs", inputs[graph_name]), + ("dynamic_shapes", dynamic_shapes[graph_name]), + ) + if value is not None + ] + if ignored: + warnings.warn( + f"{', '.join(ignored)} ignored for graph '{graph_name}': an " + "ExportedProgram is lowered as captured.", + stacklevel=2, + ) + ep = m + else: + ep = torch.export.export( + m, + inputs[graph_name], + dynamic_shapes=dynamic_shapes[graph_name], + strict=True, + ) option = generate_qnn_executorch_option(compiler_specs[graph_name]) python_options = flatbuffer_to_option(option) backend_type = python_options.backend_options.backend_type