Qualcomm AI Engine Direct - accept an ExportedProgram in to_edge_tran… (#23025)

…sform_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.


cc @cbilgin

Author: psiddh

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()

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

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 (
@@ -783,6 +783,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.
 

backends/qualcomm/utils/utils.py

@@ -347,10 +347,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,
@@ -366,17 +368,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]]]):
@@ -473,12 +478,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