Fix constant_prop_pass fqn collision with state_dict (pytorch#23266)

Summary:

`replace_with_constant_node` names propagated constants
`_prop_tensor_constant{len(exported_program.constants)}`, then checks for a
collision. Both the check and the max-suffix rescan look only at
`exported_program.constants`.

A program re-entering this pass can already carry a `_prop_tensor_constant*` in
`exported_program.state_dict` from an earlier run. That name is invisible to
both, so the pass reuses it and writes its tensor into `constants` alongside the
stale `state_dict` entry. Emission then resolves the placeholder against
`state_dict` before `constants` (`exir/emit/_emitter.py:2044`), picks up the
stale tensor, and fails the size check at `_emitter.py:513`:

  InternalError: Tensor spec has buffer of size 48, but expected nbytes of 4

The spec is correct at 4 bytes; the tensor bound to it is the wrong one. Note
the existing comment already reads "If prop_constant_tensor_fqn already exists
in the state dict" -- the implementation never matched that intent.

Widen the collision check and the suffix rescan to span `state_dict` as well as
`constants`.

Observed on a Turing/Coleman hand-tracking model compiled with
`fx_options.allow_quant_dequant_on_tce=True`, which folds a scalar `1` from a
`1 - valid_multi_view` expression whose name was already taken in `state_dict`.
Instrumenting the pass showed the collision directly:

  fqn=_prop_tensor_constant20 new_nbytes=4
  already_in_constants=False already_in_state_dict=True

With the fix the model compiles to a .pte (3 TCE delegates, 84.19% delegation);
without it emission aborts.

Differential Revision: D122298673

Pull Request resolved: pytorch#23266

Author: pcwu0329

exir/passes/constant_prop_pass.py

@@ -213,11 +213,17 @@ def replace_with_constant_node(
     # If prop_constant_tensor_fqn already exists in the state dict, we need
     # to create a new name. Find the largest suffix of "_prop_tensor_constant",
     # and increment it by 1 to form the new name.
-    if prop_constant_tensor_fqn in exported_program.constants:
+    # A program re-entering this pass can already carry a
+    # `_prop_tensor_constant*` in state_dict from an earlier run. Emission
+    # resolves placeholders against state_dict before constants, so reusing
+    # such a name lets the stale entry shadow the tensor written here. Both
+    # the check and the suffix scan therefore span state_dict too.
+    taken = exported_program.constants.keys() | exported_program.state_dict.keys()
+    if prop_constant_tensor_fqn in taken:
         suffix = 1 + max(
             (
                 int(name[len(prefix) :])
-                for name in exported_program.constants.keys()
+                for name in taken
                 if name.startswith(prefix) and name[len(prefix) :].isdigit()
             ),
             default=-1,

exir/tests/test_passes.py

@@ -1690,6 +1690,32 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             new_ep.graph_module.code
         )
 
+    def test_constant_prop_pass_avoids_state_dict_name_collision(self) -> None:
+        class Add(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return x + 3
+
+        edge = to_edge(
+            export(Add(), (torch.ones(1),), strict=True),
+            compile_config=EdgeCompileConfig(_skip_dim_order=False),
+        )
+        edge = edge.transform([ScalarToTensorPass(), RemoveMixedTypeOperators()])
+        exported_program = lift_constant_tensor_pass(edge.exported_program())
+
+        # A program re-entering this pass can already carry a
+        # `_prop_tensor_constant*` in state_dict from an earlier run.
+        # Emission resolves placeholders against state_dict before
+        # constants, so reusing the name lets the stale entry shadow the
+        # propagated tensor and the two disagree on size.
+        stale = torch.zeros(12)
+        exported_program.state_dict["_prop_tensor_constant0"] = stale
+
+        new_ep = constant_prop_pass(exported_program)
+
+        for name in new_ep.constants:
+            self.assertNotIn(name, new_ep.state_dict)
+        self.assertIs(new_ep.state_dict["_prop_tensor_constant0"], stale)
+
     def test_pass_no_user_inputs(self) -> None:
         class NoUserInputs(torch.nn.Module):
             def __init__(self):