Fix create_mutable_buffer recording the requested name instead of the name torch.fx assigned

backends/transforms/test/test_create_mutable_buffer.py

@@ -93,6 +93,83 @@ 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
+
+    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):
     """

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:
@@ -311,10 +329,14 @@ def create_mutable_buffer(
     _validate_graph_signature(exp_program)
 
     persistent_buffer = True
-    exp_program.state_dict[target] = data
 
     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)
 
@@ -338,20 +360,28 @@ 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"
+
+    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 +397,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)