From 2f9dd1bff260784c370129a65c82f2fb55cf4f31 Mon Sep 17 00:00:00 2001 From: Shresth Rana <138112017+slipstr34m@users.noreply.github.com> Date: Mon, 3 Aug 2026 20:42:22 -0700 Subject: [PATCH 1/2] Fix create_mutable_buffer recording the requested name instead of the name torch.fx assigned --- .../test/test_create_mutable_buffer.py | 47 +++++++++++++++++++ backends/transforms/utils.py | 16 +++++-- 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/backends/transforms/test/test_create_mutable_buffer.py b/backends/transforms/test/test_create_mutable_buffer.py index 4eb4d9538f2..b4930883e8a 100644 --- a/backends/transforms/test/test_create_mutable_buffer.py +++ b/backends/transforms/test/test_create_mutable_buffer.py @@ -93,6 +93,53 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: == target_name ) + def test_create_mutable_buffer_sanitized_name(self): + """ + torch.fx renames a placeholder whose requested name is not a valid + identifier or collides with an existing node. Node target, graph + signature and state_dict must follow the assigned name. + """ + + class EmptyNetwork(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + test_data: torch.Tensor = (torch.zeros(1),) + + # "0_cache" is not a valid identifier; "x" collides with the user input + for requested_name in ("0_cache", "x"): + exported_program = export( + EmptyNetwork(), args=EmptyNetwork.test_data, strict=True + ) + exported_program = to_edge(exported_program).exported_program() + graph = exported_program.graph_module.graph + + buffer_node = create_mutable_buffer( + exp_program=exported_program, + name=requested_name, + data=torch.ones(1) * 2, + ) + + assert buffer_node.name != requested_name + assert buffer_node.target == buffer_node.name + signature = exported_program.graph_signature + assert buffer_node.name in signature.inputs_to_buffers + target = signature.inputs_to_buffers[buffer_node.name] + assert target in exported_program.state_dict + assert signature.buffers_to_mutate[buffer_node.name] == target + + input_node = list(graph.nodes)[1] + with graph.inserting_after(input_node): + graph.create_node( + "call_function", + exir_ops.edge.aten.add.Tensor, + args=(input_node, buffer_node), + kwargs={}, + ) + + # Recompiling emits the placeholder target as a parameter name + assert exported_program.module()(torch.zeros(1)) == 0 + class TestRegisterMutableBufferPass(unittest.TestCase): """ diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index 941e07a3e95..972f8101772 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -311,7 +311,6 @@ def create_mutable_buffer( _validate_graph_signature(exp_program) persistent_buffer = True - exp_program.state_dict[target] = data graph = exp_program.graph_module.graph @@ -349,9 +348,20 @@ def create_mutable_buffer( break assert node is not None, "node should be created at this point" + + # torch.fx may rename the node (invalid identifier or collision); the + # target, state_dict key and graph signature must follow the assigned name. + node.target = node.name + if node.name != name: + target = node.name[2:] if node.name.startswith("b_") else node.name + if target in exp_program.state_dict: + graph.erase_node(node) + raise RuntimeError(f"Buffer target '{target}' already exists in state_dict") + exp_program.state_dict[target] = data + node.meta["val"] = fake_tensor buffer_input_spec = InputSpec( - InputKind.BUFFER, TensorArgument(name), target, persistent_buffer + InputKind.BUFFER, TensorArgument(node.name), target, persistent_buffer ) input_specs.insert(node_index, buffer_input_spec) @@ -367,7 +377,7 @@ def create_mutable_buffer( output_specs = exp_program.graph_signature.output_specs mutation_output_spec = OutputSpec( - OutputKind.BUFFER_MUTATION, TensorArgument(name), target + OutputKind.BUFFER_MUTATION, TensorArgument(node.name), target ) output_specs.insert(output_index, mutation_output_spec) From 286b7bf12e8c942a717caa82c19f33a82c9b7f67 Mon Sep 17 00:00:00 2001 From: Shresth Rana <138112017+slipstr34m@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:37:42 -0700 Subject: [PATCH 2/2] Extract shared placeholder creation helper and reject duplicate mutable buffer requests --- .../test/test_create_mutable_buffer.py | 30 ++++++++++++ backends/transforms/utils.py | 46 +++++++++++++------ 2 files changed, 63 insertions(+), 13 deletions(-) diff --git a/backends/transforms/test/test_create_mutable_buffer.py b/backends/transforms/test/test_create_mutable_buffer.py index b4930883e8a..04105efac28 100644 --- a/backends/transforms/test/test_create_mutable_buffer.py +++ b/backends/transforms/test/test_create_mutable_buffer.py @@ -140,6 +140,36 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # Recompiling emits the placeholder target as a parameter name assert exported_program.module()(torch.zeros(1)) == 0 + def test_create_mutable_buffer_duplicate_name_raises(self): + """ + A mutable buffer cannot be silently shared, so repeating a request must + raise instead of deduplicating, even when torch.fx renamed the node. + """ + + class EmptyNetwork(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + test_data: torch.Tensor = (torch.zeros(1),) + + for requested_name in ("b_cache", "0_cache"): + exported_program = export( + EmptyNetwork(), args=EmptyNetwork.test_data, strict=True + ) + exported_program = to_edge(exported_program).exported_program() + + create_mutable_buffer( + exp_program=exported_program, + name=requested_name, + data=torch.ones(1), + ) + with self.assertRaises(RuntimeError): + create_mutable_buffer( + exp_program=exported_program, + name=requested_name, + data=torch.ones(1), + ) + class TestRegisterMutableBufferPass(unittest.TestCase): """ diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index 972f8101772..1813b4565f8 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -95,6 +95,27 @@ def get_param_tensor( raise RuntimeError(f"unsupported param type, {node.op}.") +def _find_placeholder(graph: torch.fx.Graph, name: str) -> Optional[torch.fx.Node]: + """Return the placeholder previously created for this requested name, if any.""" + for n in graph.nodes: + if n.op == "placeholder" and n.meta.get("requested_name") == name: + return n + return None + + +def _create_placeholder_node(graph: torch.fx.Graph, name: str) -> torch.fx.Node: + """Create a placeholder at the current insertion point. + + torch.fx may rename the node (invalid identifier or collision); the target, + state_dict key and graph signature must follow the assigned name, so the + target is set to node.name and the requested name is kept in node.meta. + """ + node = graph.create_node(op="placeholder", name=name, target=name) + node.target = node.name + node.meta["requested_name"] = name + return node + + def create_constant_placeholder( exp_program: ExportedProgram, graph: torch.fx.Graph, @@ -111,18 +132,15 @@ def create_constant_placeholder( # Multiple pattern replacements may request the same shared weight; return # the existing node to avoid duplicate parameter names on recompile. - for n in graph.nodes: - if n.op == "placeholder" and n.meta.get("requested_name") == name: - return n + existing = _find_placeholder(graph, name) + if existing is not None: + return existing fake_tensor = _get_fake_tensor_mode(graph, data) - # torch.fx may rename the node (invalid identifier or collision); the - # target, state_dict key and graph signature must follow the assigned name. - node = graph.create_node(op="placeholder", name=name, target=name) - target = node.target = node.name + node = _create_placeholder_node(graph, name) + target = node.name node.meta["val"] = fake_tensor - node.meta["requested_name"] = name # Add data to state_dict/ constants match kind: @@ -314,6 +332,11 @@ def create_mutable_buffer( graph = exp_program.graph_module.graph + # Unlike create_constant_placeholder, a mutable buffer cannot be silently + # shared, so a repeated request is an error rather than a dedup hit. + if _find_placeholder(graph, name) is not None: + raise RuntimeError(f"Placeholder for '{name}' already exists in the graph") + # Create fake tensor using helper function fake_tensor = _get_fake_tensor_mode(graph, data) @@ -337,21 +360,18 @@ def create_mutable_buffer( ): # No const or user input nodes node_index = len(input_specs) - node = graph.create_node(op="placeholder", name=name, target=name) + node = _create_placeholder_node(graph, name) else: # Find the first constant or user input node for i, spec in enumerate(input_specs): if spec.kind in [InputKind.CONSTANT_TENSOR, InputKind.USER_INPUT]: node_index = i with graph.inserting_before(_spec_to_node(exp_program, spec)): - node = graph.create_node(op="placeholder", name=name, target=name) + node = _create_placeholder_node(graph, name) break assert node is not None, "node should be created at this point" - # torch.fx may rename the node (invalid identifier or collision); the - # target, state_dict key and graph signature must follow the assigned name. - node.target = node.name if node.name != name: target = node.name[2:] if node.name.startswith("b_") else node.name if target in exp_program.state_dict: