From 786d78666f8a3d48818bb563f7ee7fd6dbcda53e Mon Sep 17 00:00:00 2001 From: Max Ren Date: Wed, 19 Mar 2025 18:08:26 -0700 Subject: [PATCH] [backend_api] Delete partitioner tags after lowering 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` Differential Revision: [D71518362](https://our.internmc.facebook.com/intern/diff/D71518362/) [ghstack-poisoned] --- exir/backend/backend_api.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/exir/backend/backend_api.py b/exir/backend/backend_api.py index d5bd574ec5a..6345a455caf 100644 --- a/exir/backend/backend_api.py +++ b/exir/backend/backend_api.py @@ -401,6 +401,12 @@ 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: + if "delegation_tag" in node.meta: + del node.meta["delegation_tag"] + return ExportedProgram( root=tagged_graph_module, graph=tagged_graph_module.graph,