From ee37db45efac525237405c28cec5b4371b03976a Mon Sep 17 00:00:00 2001 From: Max Ren Date: Thu, 20 Mar 2025 13:44:10 -0700 Subject: [PATCH] [backend_api] Delete partitioner tags after lowering Pull Request resolved: https://github.com/pytorch/executorch/pull/9435 There was a slight bug i found when lowering with: XnnpackDynamicallyQuantizedPartitioner --> XnnpackPartitioner. The issue arrises because the second time we partition, delegation tags from the previous partitioner still exist. Specifically they exist on the getitem node because the metadata was propagated. See: https://github.com/pytorch/pytorch/blob/main/torch/fx/passes/utils/fuser_utils.py#L235 Since the getitem nodes created in by the XnnpackDynamicallyQuantizedPartitioner still have the delegation tags, they sneak into partitions from XnnpackPartitioner. This shouldn't happen because delegation tags from previous Partitioner should not exist after the previous lowering. We update the code in to_backend here to erase all the "delegation_tags" from node.meta after we've performed all `_partition_and_lower` ghstack-source-id: 273078717 @exported-using-ghexport Differential Revision: [D71518362](https://our.internmc.facebook.com/intern/diff/D71518362/) --- exir/backend/backend_api.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/exir/backend/backend_api.py b/exir/backend/backend_api.py index d5bd574ec5a..ab2e66f7885 100644 --- a/exir/backend/backend_api.py +++ b/exir/backend/backend_api.py @@ -401,6 +401,11 @@ def to_backend( tagged_exported_program, ) + # Partitioner added delegation tags to the graph module nodes, + # we make sure to remove them after we finished partition_and_lower + for node in tagged_graph_module.graph.nodes: + node.meta.pop("delegation_tag", None) + return ExportedProgram( root=tagged_graph_module, graph=tagged_graph_module.graph,