diff --git a/backends/transforms/test/test_create_mutable_buffer.py b/backends/transforms/test/test_create_mutable_buffer.py index 4eb4d9538f2..77917987a7c 100644 --- a/backends/transforms/test/test_create_mutable_buffer.py +++ b/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_renamed_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): """ diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index 941e07a3e95..51fc4684f58 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -95,6 +95,33 @@ def get_param_tensor( raise RuntimeError(f"unsupported param type, {node.op}.") +def _buffer_target(node_name: str) -> str: + """Map a placeholder name to its state_dict target per the export + convention: placeholder "b_foo" corresponds to buffer target "foo".""" + return node_name[2:] if node_name.startswith("b_") else node_name + + +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 +138,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: @@ -298,11 +322,7 @@ def create_mutable_buffer( if not isinstance(data, torch.Tensor): raise ValueError("Data must be a torch.Tensor") - # Extract target name (remove "b_" prefix if present, following export convention) - if name.startswith("b_"): - target = name[2:] - else: - target = name + target = _buffer_target(name) # Check if target already exists if target in exp_program.state_dict: @@ -311,10 +331,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 +362,30 @@ 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 fx renamed the node, the caller's name is stale; re-apply the target + # convention to the assigned name. + if node.name != name: + target = _buffer_target(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 +401,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)