diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index ac146956611..447f80fec28 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -351,6 +351,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 a267dc2f763..dae209ac170 100644 --- a/backends/qualcomm/qnn_preprocess.py +++ b/backends/qualcomm/qnn_preprocess.py @@ -22,10 +22,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, ) @@ -56,10 +53,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 834e1571907..acea8013b66 100644 --- a/backends/qualcomm/tests/rework/passes/passes_helper.py +++ b/backends/qualcomm/tests/rework/passes/passes_helper.py @@ -287,13 +287,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),