diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index 64fabe5604e..f2bb7a45f3c 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -354,6 +354,14 @@ def call(self, graph_module: torch.fx.GraphModule): self.traverse(node, graph_module) self.insert_permute, self.transformed_tag = True, QCOM_AXIS_ORDER + for node in graph.nodes: + if hasattr(node, "meta"): + # Pop QCOM_AXIS_ORDER written by the to-edge LayoutTransform pass. + # Without this, the main for-loop below would see is_transformed_node=True + # for every sensitive node (deepcopy carries the tag from to-edge) and + # skip them entirely, so no permute nodes would ever be inserted. + node.meta.pop(QCOM_AXIS_ORDER, "") + for node in sensitive_nodes: if not self.is_transformed_node(node): self.mark_as_transformed(node) diff --git a/backends/qualcomm/qnn_preprocess.py b/backends/qualcomm/qnn_preprocess.py index 7b5aace51b8..94c884be2a2 100644 --- a/backends/qualcomm/qnn_preprocess.py +++ b/backends/qualcomm/qnn_preprocess.py @@ -23,10 +23,7 @@ from executorch.backends.qualcomm.serialization.qc_schema_serialize import ( flatbuffer_to_option, ) -from executorch.backends.qualcomm.utils.constants import ( - QCOM_AXIS_ORDER, - QCOM_TENSOR_NAME, -) +from executorch.backends.qualcomm.utils.constants import QCOM_TENSOR_NAME from executorch.backends.qualcomm.utils.qnn_manager_lifecycle import ( get_current_qnn_manager, ) @@ -100,10 +97,6 @@ def _build_op_wrappers( use_mha2sha: bool, backend_type: QnnExecuTorchBackendType, ): - for node in edge_program.graph_module.graph.nodes: - if hasattr(node, "meta"): - # pop certain keys in meta for not affecting the passes in compilation - node.meta.pop(QCOM_AXIS_ORDER, "") # QNN Delegate Specific Passes graph_module = get_qnn_pass_manager_cls( backend_type diff --git a/backends/qualcomm/tests/rework/passes/passes_helper.py b/backends/qualcomm/tests/rework/passes/passes_helper.py index 5b5a1eec6c6..07ac3b36a73 100644 --- a/backends/qualcomm/tests/rework/passes/passes_helper.py +++ b/backends/qualcomm/tests/rework/passes/passes_helper.py @@ -300,13 +300,6 @@ def lower_preprocess_gm( quantizer=quantizer, ) gm = edge_ep.graph_module - # Mirror qnn_preprocess.py: strip QCOM_AXIS_ORDER from all nodes - from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER - - for node in gm.graph.nodes: - if hasattr(node, "meta"): - node.meta.pop(QCOM_AXIS_ORDER, "") - pm_cls = get_qnn_pass_manager_cls(backend_type) pass_classes = PassPipeline._slice_to_target( pm_cls.get_preprocess_passes(use_mha2sha=use_mha2sha),