Fix create_constant_placeholder crashing when torch.fx renames the requested placeholder (pytorch#21542)

Fixes pytorch#14055
Fixes pytorch#21541

## Problem

Found while verifying pytorch#21489: the textbook CNN used there, one top-level
`nn.Sequential` of conv, batchnorm, relu and maxpool, gets past the
pooling check it used to crash on and now dies one pass later, in
`FuseBatchNormPass`.

`create_constant_placeholder` records the name it requested, not the
name torch.fx assigned:

```python
node = graph.create_node(op="placeholder", name=name, target=name)
...
node_names = [n.name for n in graph.nodes if n.op == "placeholder"]
node_index = node_names.index(name)   # ValueError after any rename
...
arg_spec = TensorArgument(name)
```

torch.fx renames the node whenever the requested name is not a valid
identifier or collides with an existing node. Two plain models trigger
this today. A top-level `nn.Sequential` names its parameters `0.weight`,
so `FuseBatchNormPass` requests `0_weight_fused_bn`, fx assigns
`_0_weight_fused_bn`, and lowering to XNNPACK dies with `ValueError:
'0_weight_fused_bn' is not in list`. torchvision regnet under Vulkan
hits the same lines through its hyphenated module names (pytorch#14055). The
identical models behind a named attribute lower fine, which is why
existing coverage never reached this path.

Two further defects sit behind the crash. A placeholder's `target` is
emitted verbatim as a function parameter name on recompile, so fixing
only the lookup produces `def forward(self, 0_weight_fused_bn, ...)`, a
SyntaxError. And the shared-weight dedup from pytorch#18031 matched on `target
== name`, so after a rename it misses and creates a duplicate
placeholder.

## Fix

One name everywhere: `node.name`, `node.target`, the state_dict key and
the graph signature all follow the name fx assigned. The requested name
is kept in `node.meta` and used for the dedup, so the pytorch#18031 contract is
preserved: a second request for the same name returns the existing node,
now also when fx renamed it. Keying the state_dict by the assigned name
inherits the fx namespace uniqueness guarantee, so a colliding request
can never overwrite an existing parameter.

For any requested name that is a valid identifier with no collision, the
assigned name equals the requested name and behavior is unchanged.

## Effect

Same script on either side of the change, executed on executor_runner
built from this branch. maxdiff is deviation from eager.

```
                                     BEFORE                        AFTER
nn.Sequential(conv, bn, relu)        ValueError at utils.py:151    delegated, maxdiff 4.0e-07
two-block Sequential CNN             ValueError at utils.py:151    delegated, maxdiff 4.8e-07
same models behind an attribute      lower and run                 unchanged, maxdiff <= 5.5e-07
```

## Testing

Two helper tests in `test_create_delete_constant_placeholder.py`: a
digit-leading request checks node, signature, state_dict and a recompile
round trip; a dedup test checks that requesting the same name twice
returns the existing node for a clean, a digit-leading and a hyphenated
name. One pass test in `test_batch_norm_fusion.py` fuses a top-level
Sequential end to end. `test_conv_bn` in
`runtime/test/test_runtime_xnnpack.py`, skipped on this bug, is
re-enabled (linux-gated, runs in CI).

Verified failing-first by reverting only `backends/transforms/utils.py`:
exactly the three new tests fail, all with `ValueError` at
`utils.py:151`.

```
backends/transforms/test + backends/xnnpack/test/passes + runtime/test/test_runtime_xnnpack.py
  main    : 165 passed, 57 skipped
  this PR : 168 passed, 57 skipped
```

lintrunner reports no issues on all changed files.

The helper is shared by the Arm, Vulkan and XNNPACK backends. Callers
passing valid unique names get identical behavior; renamed cases
previously crashed, so no caller can depend on the old behavior.

cc @GregoryComer @digantdesai @cbilgin @JakeStevens

Author: slipstr34m

backends/transforms/test/test_create_delete_constant_placeholder.py

@@ -106,6 +106,90 @@ def _test_create_delete(kind: InputKind, persistent_buffer: bool = None):
     assert len(exported_program.constants) == 0
 
 
+def test_create_delete_sanitized_name():
+    """
+    torch.fx renames a placeholder whose requested name is not a valid
+    identifier, e.g. '0_weight_fused_bn' from a top-level nn.Sequential.
+    Target, graph signature and state_dict must follow the assigned name.
+    """
+    module = EmptyNetwork()
+    exported_program = export(module, args=module.test_data, strict=True)
+    exported_program = to_edge(exported_program).exported_program()
+    graph = exported_program.graph_module.graph
+
+    input_node = list(graph.nodes)[0]
+    with graph.inserting_before(input_node):
+        const_node = create_constant_placeholder(
+            exp_program=exported_program,
+            graph=graph,
+            kind=InputKind.PARAMETER,
+            name="0_test_node",
+            data=torch.ones(1),
+        )
+
+    assert const_node.name != "0_test_node"
+    assert const_node.name.isidentifier()
+    assert const_node.target == const_node.name
+    assert const_node.name in exported_program.graph_signature.inputs_to_parameters
+    target = exported_program.graph_signature.inputs_to_parameters[const_node.name]
+    assert target in exported_program.state_dict
+
+    with graph.inserting_after(input_node):
+        add_node = graph.create_node(
+            "call_function",
+            exir_ops.edge.aten.add.Tensor,
+            args=(input_node, const_node),
+            kwargs={},
+        )
+    output_node = graph.output_node()
+    output_node.replace_input_with(input_node, add_node)
+
+    # Recompiling emits the placeholder target as a function parameter name
+    assert exported_program.module()(torch.zeros(1)) == 1
+
+    output_node.replace_input_with(add_node, input_node)
+    graph.eliminate_dead_code()
+    delete_constant_placeholder(exported_program, const_node)
+    assert len(exported_program.state_dict) == 0
+    assert len(exported_program.graph_signature.input_specs) == 1
+
+
+def test_create_same_name_returns_existing_node():
+    """
+    Requesting the same name twice must return the existing node instead of
+    creating a duplicate, including when torch.fx renamed the placeholder.
+    """
+    # "block1-0_weight" mirrors torchvision regnet's hyphenated module names
+    for requested_name in ("test_node", "0_test_node", "block1-0_weight"):
+        module = EmptyNetwork()
+        exported_program = export(module, args=module.test_data, strict=True)
+        exported_program = to_edge(exported_program).exported_program()
+        graph = exported_program.graph_module.graph
+
+        input_node = list(graph.nodes)[0]
+        with graph.inserting_before(input_node):
+            first = create_constant_placeholder(
+                exp_program=exported_program,
+                graph=graph,
+                kind=InputKind.PARAMETER,
+                name=requested_name,
+                data=torch.ones(1),
+            )
+        with graph.inserting_before(input_node):
+            second = create_constant_placeholder(
+                exp_program=exported_program,
+                graph=graph,
+                kind=InputKind.PARAMETER,
+                name=requested_name,
+                data=torch.ones(1),
+            )
+
+        assert first is second
+        placeholders = [n for n in graph.nodes if n.op == "placeholder"]
+        assert len(placeholders) == 2  # the created node and the user input
+        assert len(exported_program.state_dict) == 1
+
+
 def test_create_delete_parameter():
     _test_create_delete(InputKind.PARAMETER)
 

backends/transforms/utils.py

@@ -109,16 +109,20 @@ def create_constant_placeholder(
     decide where to insert the node, at an insertion point before the first input node.
     """
 
-    target = name
+    # 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
 
-    # If a placeholder with this target already exists, return it to avoid
-    # duplicate parameter names in the generated function signature which would
-    # cause a SyntaxError on recompile. This can happen when multiple pattern
-    # replacements independently create placeholders for a shared weight.
-    if name in exp_program.state_dict or name in exp_program.constants:
-        for n in graph.nodes:
-            if n.op == "placeholder" and n.target == name:
-                return n
+    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.meta["val"] = fake_tensor
+    node.meta["requested_name"] = name
 
     # Add data to state_dict/ constants
     match kind:
@@ -140,15 +144,9 @@ def create_constant_placeholder(
         case _:
             raise RuntimeError("Can only create constant input nodes.")
 
-    fake_tensor = _get_fake_tensor_mode(graph, data)
-
-    # Create node
-    node = graph.create_node(op="placeholder", name=name, target=name)
-    node.meta["val"] = fake_tensor
-
     # Add tensor to graph_signature in the same order as nodes in the graph
     node_names = [n.name for n in graph.nodes if n.op == "placeholder"]
-    node_index = node_names.index(name)
+    node_index = node_names.index(node.name)
 
     input_specs = exp_program.graph_signature.input_specs
     user_input_indices = [
@@ -161,7 +159,7 @@ def create_constant_placeholder(
             f"Failed to insert {name}; Const placeholder nodes must be inserted before user input nodes in the graph."
         )
 
-    arg_spec = TensorArgument(name)
+    arg_spec = TensorArgument(node.name)
     input_spec = InputSpec(kind, arg_spec, target, persistent_buffer)
     input_specs.insert(node_index, input_spec)
 

backends/xnnpack/test/passes/test_batch_norm_fusion.py

@@ -89,6 +89,24 @@ def test_fp32_conv_batch_norm_fusion(self):
                 .run_method_and_compare_outputs()
             )
 
+    def test_fp32_conv_batch_norm_fusion_top_level_sequential(self):
+        """
+        A top-level nn.Sequential yields digit-leading fused placeholder names,
+        which torch.fx renames. The pass must survive the rename.
+        """
+        model = torch.nn.Sequential(
+            torch.nn.Conv2d(2, 2, (2, 2)),
+            torch.nn.BatchNorm2d(2),
+        ).eval()
+        (
+            Tester(model, (torch.randn(2, 2, 4, 4),))
+            .export()
+            .to_edge()
+            .run_passes(self.PassStage)
+            .check_not([self.bn_name])
+            .run_method_and_compare_outputs()
+        )
+
     def test_q8_conv_batch_norm_fusion(self):
         for transpose in [False, True]:
             (

runtime/test/test_runtime_xnnpack.py

@@ -116,10 +116,7 @@ def test_mlp(self):
 
     # ------------------------------------------------------------------
     # BatchNorm + Conv (common in MobileNet-style models)
-    # Skipped: FuseBatchNormPass crashes on Sequential(Conv2d, BN) export.
-    # TODO(#11225): re-enable once the XNNPACK pass is fixed.
     # ------------------------------------------------------------------
-    @unittest.skip("FuseBatchNormPass bug in XNNPACK backend")
     def test_conv_bn(self):
         model = torch.nn.Sequential(
             torch.nn.Conv2d(3, 16, 3, padding=1),