From efc6e3cd50049504e9b6013f3536fbba3e656390 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Fri, 17 Jul 2026 07:54:47 -0700 Subject: [PATCH] Allow context-binary lowering to use edge dialect ops (#20518) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/20518 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. Reviewed By: rascani Differential Revision: D109598309 --- backends/qualcomm/qnn_preprocess.py | 6 +- backends/qualcomm/tests/test_passes.py | 94 +++++++++++++++++++++++++- backends/qualcomm/utils/utils.py | 4 +- 3 files changed, 98 insertions(+), 6 deletions(-) diff --git a/backends/qualcomm/qnn_preprocess.py b/backends/qualcomm/qnn_preprocess.py index cbe96b5954a..a267dc2f763 100644 --- a/backends/qualcomm/qnn_preprocess.py +++ b/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: diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 1124b01d613..7ba0e4f9de1 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/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.""" diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index b1c236647f9..18de0a85ed5 100644 --- a/backends/qualcomm/utils/utils.py +++ b/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