From a031f63d3a03fc630d4c24055f65c43f492159e5 Mon Sep 17 00:00:00 2001 From: john-rocky Date: Fri, 21 Aug 2026 01:55:37 +0900 Subject: [PATCH 1/2] Recognise a buffer as mutated when the write is not its direct user MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tag_constant_data decided whether a buffer is mutated by looking through its users for the node that performs the write. A cache that is first concatenated with its new values and written back further down the chain has the concat as its user, not the write, so it was taken for constant data and tagged into the delegate. A backend that compiles mutable buffers into state then produced a model that asks the runtime for state — after the partitioner had been told, by take_over_mutable_buffer=False, that this runtime has none. On the Core ML backend that is a load that succeeds and an execute that fails with 'The input feature for layers_N_conv_conv_state must be an MLState, but it was not'. The signature already records which buffers are mutated, so match on the target instead of on the mutating node's name. The direct-user check stays for parameters and lifted constants, which have no target to match. --- exir/backend/test/test_partitioner.py | 50 +++++++++++++++++++++++++++ exir/backend/utils.py | 16 +++++++++ 2 files changed, 66 insertions(+) diff --git a/exir/backend/test/test_partitioner.py b/exir/backend/test/test_partitioner.py index f2ce8e14988..7422b136f84 100644 --- a/exir/backend/test/test_partitioner.py +++ b/exir/backend/test/test_partitioner.py @@ -620,6 +620,56 @@ def partition( ] self.assertEqual(len(copy_node), 1) + def test_buffer_mutated_indirectly_is_not_taken_for_constant(self): + """A buffer written a step after it is read still counts as mutated. + + `tag_constant_data` used to decide by looking at a buffer's direct users for the + node that performs the mutation. A cache that is concatenated with its new values + and written back further down has something else as its user, so it was taken for + a constant and handed to the delegate as a buffer. A backend that compiles mutable + buffers into state then produced a model asking for state that the partitioner had + been told not to take over. + """ + + class IndirectlyMutated(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cache", torch.zeros(1, 4)) + + def forward(self, x): + # The buffer's user is the cat, not the write. + joined = torch.cat([self.cache, x], dim=1) + self.cache.copy_(joined[:, -4:]) + return joined.sum(1) + + edge = exir.to_edge( + torch.export.export(IndirectlyMutated(), (torch.zeros(1, 2),), strict=True) + ) + program = edge.exported_program() + signature = program.graph_signature + self.assertIn("cache", signature.buffers_to_mutate.values()) + + placeholder = next( + node + for node in program.graph.nodes + if node.op == "placeholder" and signature.inputs_to_buffers.get(node.name) + ) + self.assertNotIn( + placeholder.name, + {user.name for user in placeholder.users} & set(signature.buffers_to_mutate), + "this test is only meaningful while the mutation is not a direct user", + ) + + for node in program.graph.nodes: + if node.op == "call_function": + node.meta["delegation_tag"] = "tag0" + tag_constant_data(program) + + self.assertIsNone( + placeholder.meta.get("delegation_tag"), + "a mutated buffer must not be tagged as constant data", + ) + def test_buffer_mutation1(self): class TestModule(torch.nn.Module): def __init__(self): diff --git a/exir/backend/utils.py b/exir/backend/utils.py index 9bdba810138..a77f687dcb6 100644 --- a/exir/backend/utils.py +++ b/exir/backend/utils.py @@ -357,6 +357,16 @@ def tag_constant_data(edge_program: ExportedProgram) -> None: constants_map = sig.inputs_to_lifted_tensor_constants buffers_to_mutate = sig.buffers_to_mutate + # Which buffers the program says are mutated. Reading the targets rather than + # matching the mutating node's name against a buffer's direct users is what makes + # this hold for a buffer that is not mutated in one step: a cache that is first + # concatenated with the new values and written back further down the chain has + # something other than the mutation as its user, and used to be taken for a + # constant. It would then be handed to a delegate as a buffer, and a backend that + # compiles mutable buffers into state — Core ML does — produced a model asking for + # state that the runtime had been told, by take_over_mutable_buffer=False, not to + # provide. + mutated_targets = set(buffers_to_mutate.values()) mutated_buffer = set() for node in edge_program.graph.nodes: if node.op == "placeholder" and ( @@ -364,6 +374,12 @@ def tag_constant_data(edge_program: ExportedProgram) -> None: or node.name in buffers_map or node.name in constants_map ): + if buffers_map.get(node.name) in mutated_targets: + logging.info( + "The buffer node is a mutated buffer node, which is not constant." + ) + mutated_buffer.add(node) + continue for node_user in node.users: if node_user.name in buffers_to_mutate: logging.info( From d3d6fcf551d08f02ad9c4b4fbdcc0ae5c433f87d Mon Sep 17 00:00:00 2001 From: john-rocky Date: Thu, 27 Aug 2026 10:14:12 +0900 Subject: [PATCH 2/2] Extract mutated-buffer detection into a helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review asked for two things here: the comment explaining the target-matching should not name a specific backend, and flake8 C901 flagged tag_constant_data at complexity 13 against the limit of 12. Moving the detection loop into _mutated_buffer_placeholders answers both — the explanation moves into the helper's docstring with the backend-specific wording dropped, and both functions sit below the complexity limit. Tagging behaviour is unchanged. --- exir/backend/utils.py | 45 +++++++++++++++++++++++++++---------------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/exir/backend/utils.py b/exir/backend/utils.py index a77f687dcb6..b2d80357072 100644 --- a/exir/backend/utils.py +++ b/exir/backend/utils.py @@ -342,31 +342,24 @@ def format_delegated_graph(graph_module: torch.fx.GraphModule) -> str: return graph_format_str -def tag_constant_data(edge_program: ExportedProgram) -> None: +def _mutated_buffer_placeholders(edge_program: ExportedProgram) -> Set[torch.fx.Node]: """ - Util function for partitioners. This function tags the const/param/buffers nodes - whose users all belong within the same partition. This should be called after tagging all other nodes. - Any const/param/buffer which is used as input to a subgraph, will be tagged with the same tag as that - subgraph. Throw error when const/param/buffers is used across different partitions. That is the - underlying data will be owned by multiple delegates. + Placeholders whose data the program mutates, which must not be treated as + constant. A buffer is matched against the signature's mutation targets rather + than by looking through its users for the mutating node: a buffer that is not + written in one step — a cache first concatenated with new values and written + back further down the chain — has something other than the mutation as its + user, and would otherwise be taken for a constant and folded into a delegate + instead of remaining a mutable buffer input. Parameters and lifted constants + have no mutation target, so for them the direct-user check remains. """ - # Cache signature lookups to avoid rebuilding dicts on every access. sig = edge_program.graph_signature params_map = sig.inputs_to_parameters buffers_map = sig.inputs_to_buffers constants_map = sig.inputs_to_lifted_tensor_constants buffers_to_mutate = sig.buffers_to_mutate - - # Which buffers the program says are mutated. Reading the targets rather than - # matching the mutating node's name against a buffer's direct users is what makes - # this hold for a buffer that is not mutated in one step: a cache that is first - # concatenated with the new values and written back further down the chain has - # something other than the mutation as its user, and used to be taken for a - # constant. It would then be handed to a delegate as a buffer, and a backend that - # compiles mutable buffers into state — Core ML does — produced a model asking for - # state that the runtime had been told, by take_over_mutable_buffer=False, not to - # provide. mutated_targets = set(buffers_to_mutate.values()) + mutated_buffer = set() for node in edge_program.graph.nodes: if node.op == "placeholder" and ( @@ -386,6 +379,24 @@ def tag_constant_data(edge_program: ExportedProgram) -> None: "The buffer node is a mutated buffer node, which is not constant." ) mutated_buffer.add(node) + return mutated_buffer + + +def tag_constant_data(edge_program: ExportedProgram) -> None: + """ + Util function for partitioners. This function tags the const/param/buffers nodes + whose users all belong within the same partition. This should be called after tagging all other nodes. + Any const/param/buffer which is used as input to a subgraph, will be tagged with the same tag as that + subgraph. Throw error when const/param/buffers is used across different partitions. That is the + underlying data will be owned by multiple delegates. + """ + # Cache signature lookups to avoid rebuilding dicts on every access. + sig = edge_program.graph_signature + params_map = sig.inputs_to_parameters + buffers_map = sig.inputs_to_buffers + constants_map = sig.inputs_to_lifted_tensor_constants + + mutated_buffer = _mutated_buffer_placeholders(edge_program) for node in edge_program.graph.nodes: # go through const/param/buffer nodes, if all users of const/param/buffer nodes are partitioned then partition