Allow context-binary lowering to use edge dialect ops (pytorch#20518)

Summary:

Unblocks the QNN context-binary path from lowering through `to_edge` with `_use_edge_ops=True` (the default). Previously it was pinned to `EdgeCompileConfig(_use_edge_ops=False)` purely to keep the `qaisw` context-loader custom op's original name, because loader detection was name-based.

We can simply unwrap the op overload and maintain the same check.

Differential Revision: D109598309

Author: JakeStevens

backends/qualcomm/qnn_preprocess.py

@@ -36,6 +36,7 @@
 )
 from executorch.exir.backend.utils import DelegateMappingBuilder
 from executorch.exir.debug_handle_utils import DEBUG_HANDLE_KEY
+from executorch.exir.operator.convert import unwrap_op_overload
 from torch.export.exported_program import ExportedProgram
 
 DEFAULT_DEBUG_HANDLE = 65535
@@ -90,11 +91,12 @@ def _build_op_wrappers(
                         "is not supported in Qnn Delegate"
                     )
                     try:
+                        op = unwrap_op_overload(node.target)
                         context_loader_target = eval(
-                            f"torch.ops.{OpContextLoader.namespace}.{node.target.__name__}",
+                            f"torch.ops.{OpContextLoader.namespace}.{op.__name__}",
                             globals().update(torch.__dict__),
                         )
-                        assert node.target == context_loader_target, err_msg
+                        assert op == context_loader_target, err_msg
                         # if graph has context binary loader node, return directly
                         return node.meta[OpContextLoader.meta_ctx_bin]
                     except:

backends/qualcomm/tests/test_passes.py

@@ -1,4 +1,5 @@
 import unittest
+from unittest.mock import MagicMock
 
 import torch
 from executorch.backends.qualcomm._passes import (
@@ -13,6 +14,9 @@
 from executorch.backends.qualcomm._passes.qnn_pass_manager import (
     get_qnn_pass_manager_cls,
 )
+from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
+from executorch.backends.qualcomm.partition.qnn_partitioner import QnnOperatorSupport
+from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
 from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer, QuantDtype
 from executorch.backends.qualcomm.serialization.qc_schema import (
     QcomChipset,
@@ -28,13 +32,101 @@
     generate_qnn_executorch_compiler_spec,
     to_edge_transform_and_lower_to_qnn,
 )
-from executorch.exir import to_edge
+from executorch.exir import EdgeCompileConfig, to_edge
 from executorch.exir.debug_handle_utils import DEBUG_HANDLE_KEY
 from executorch.exir.dialects._ops import ops as exir_ops
+from torch.library import Library
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 
 class TestPasses(unittest.TestCase):
+    def _build_context_loader_edge_program(self, op_name, check_ir_validity=True):
+        graph_name = "forward"
+        custom_op = Library(OpContextLoader.namespace, "FRAGMENT")
+        self.addCleanup(custom_op._destroy)
+        custom_op.define(f"{op_name}(Tensor[] inputs) -> Any")
+
+        @torch.library.impl(
+            custom_op, op_name, dispatch_key="CompositeExplicitAutograd"
+        )
+        def op_impl(inputs):
+            return (torch.zeros((1, 2), device="meta", dtype=inputs[0].dtype),)
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return getattr(
+                    getattr(torch.ops, OpContextLoader.namespace), op_name
+                ).default((x,))
+
+        exported_program = torch.export.export(
+            Model(), (torch.ones(1, 2),), strict=True
+        )
+        compile_config = (
+            None if check_ir_validity else EdgeCompileConfig(_check_ir_validity=False)
+        )
+        edge_program_manager = to_edge(
+            {graph_name: exported_program},
+            compile_config=compile_config,
+        )
+        edge_program = edge_program_manager._edge_programs[graph_name]
+        context_loader_nodes = [
+            node
+            for node in edge_program.graph.nodes
+            if node.op == "call_function"
+            and OpContextLoader.namespace in str(node.target)
+        ]
+        return edge_program, context_loader_nodes
+
+    def test_context_loader_edge_op_is_delegated(self):
+        op_name = "ctx_loader_delegation"
+        ctx_bin = b"qnn_context_binary"
+        _, context_loader_nodes = self._build_context_loader_edge_program(
+            op_name, check_ir_validity=False
+        )
+        self.assertEqual(1, len(context_loader_nodes))
+        context_loader_nodes[0].meta[OpContextLoader.meta_ctx_bin] = ctx_bin
+
+        # A fully constructed QnnOperatorSupport needs a live QNN manager, so
+        # drive is_node_supported with a mock self: the context-loader node is
+        # force-passed without touching instance state beyond the log label.
+        support = MagicMock()
+        support.phase = "QnnPartitioner"
+        self.assertTrue(
+            QnnOperatorSupport.is_node_supported(support, None, context_loader_nodes[0])
+        )
+
+    def test_build_op_wrappers_returns_context_binary(self):
+        op_name = "ctx_loader_build"
+        ctx_bin = b"qnn_context_binary"
+        edge_program, context_loader_nodes = self._build_context_loader_edge_program(
+            op_name, check_ir_validity=False
+        )
+        for node in context_loader_nodes:
+            node.meta[OpContextLoader.meta_ctx_bin] = ctx_bin
+
+        # For a graph whose only op is the context-binary loader, _build_op_wrappers
+        # returns the stamped context binary directly, before any QNN compilation.
+        result = QnnBackend._build_op_wrappers(
+            edge_program,
+            enable_tensor_dump=False,
+            op_package_infos=[],
+            use_mha2sha=False,
+            backend_type=QnnExecuTorchBackendType.kHtpBackend,
+        )
+        self.assertEqual(ctx_bin, result)
+
+    def test_context_loader_op_lowers_with_ir_validation(self):
+        op_name = "ctx_loader_validation"
+
+        # from_context_binary lowers with IR validity checks enabled (the
+        # default). The context-loader custom op survives the edge verifier
+        # because its namespace is not aten, so no validation is disabled and
+        # the loader node is still present for downstream stamping.
+        _, context_loader_nodes = self._build_context_loader_edge_program(
+            op_name, check_ir_validity=True
+        )
+        self.assertEqual(1, len(context_loader_nodes))
+
     def _build_quantized_graph(self):
         """Build a quantized graph through AnnotateQuantAttrs + FoldQDQ."""
 

backends/qualcomm/utils/utils.py

@@ -65,7 +65,7 @@
 )
 from executorch.backends.qualcomm.utils.qnn_manager_lifecycle import QnnManagerContext
 
-from executorch.exir import EdgeCompileConfig, ExirExportedProgram, to_edge
+from executorch.exir import ExirExportedProgram, to_edge
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.lowered_backend_module import LoweredBackendModule
 from executorch.exir.program._program import (
@@ -947,8 +947,6 @@ def preprocess_binary(ctx_bin, compiler_specs):
     # temporarily remove the first parameter name.
     edge_prog_mgr = to_edge(
         {graph_name: bundle_prog["exported_program"]},
-        # do not alter name for custom op
-        compile_config=EdgeCompileConfig(_use_edge_ops=False),
     )
 
     # update meta with context binary