Update (base update)

[ghstack-poisoned]

Author: SS-JIA

backends/mlx/passes.py

@@ -99,49 +99,8 @@ def call(self, exported_program) -> ExportedProgramPassResult:
         }
         if ops_to_inplace:
             reinplace_pass(exported_program, ops_to_inplace=ops_to_inplace)
-            self._resync_output_specs(exported_program)
         return ExportedProgramPassResult(exported_program, True)
 
-    @staticmethod
-    def _resync_output_specs(exported_program) -> None:
-        """Re-sync graph-signature output names after reinplace.
-
-        ``reinplace_pass`` rewrites an output-producing node (e.g. the final
-        ``exp`` -> ``exp_``) via ``replace_all_uses_with`` + erase, but does not
-        update ``graph_signature.output_specs``. Output order is preserved, so we
-        positionally re-sync each spec's argument name to the current output node
-        arg; otherwise ``ExportedProgram.validate()`` (run by the pass manager)
-        raises a SpecViolationError.
-
-        This positional pairing is only valid because reinplace does a 1:1
-        ``replace_all_uses_with`` + erase and never drops, adds, or reorders
-        outputs. We assert ``len(output_specs) == len(out_args)`` so that a
-        future change violating that invariant fails loudly here instead of
-        silently mis-pairing names (``zip`` would otherwise truncate). Specs
-        whose ``arg`` is not a named tensor (e.g. ``ConstantArgument``) carry no
-        ``name`` and are skipped by the ``getattr`` guard below.
-        """
-        out_node = next(
-            n for n in reversed(exported_program.graph.nodes) if n.op == "output"
-        )
-        out_args = out_node.args[0]
-        if not isinstance(out_args, (tuple, list)):
-            out_args = (out_args,)
-        output_specs = exported_program.graph_signature.output_specs
-        assert len(output_specs) == len(out_args), (
-            "reinplace changed graph output count: "
-            f"{len(output_specs)} output_specs vs {len(out_args)} output args. "
-            "Positional output-spec re-sync assumes a 1:1, order-preserving "
-            "rewrite."
-        )
-        for spec, arg in zip(output_specs, out_args):
-            if (
-                isinstance(arg, torch.fx.Node)
-                and getattr(spec.arg, "name", None) is not None
-                and spec.arg.name != arg.name
-            ):
-                spec.arg.name = arg.name
-
 
 @dataclass
 class RMSNormMatch(PatternMatch):

backends/vulkan/test/op_tests/utils/aten_types.py

@@ -26,7 +26,6 @@
 OPT_MEMORY_FORMAT = "::std::optional<at::MemoryFormat>"
 OPT_SCALAR_TYPE = "::std::optional<at::ScalarType>"
 STRING = "std::string_view"
-OLD_STRING = "c10::string_view"
 TWO_TENSOR_TUPLE = "::std::tuple<at::Tensor,at::Tensor>"
 THREE_TENSOR_TUPLE = "::std::tuple<at::Tensor,at::Tensor,at::Tensor>"
 TENSOR_VECTOR = "::std::vector<at::Tensor>"

backends/vulkan/test/op_tests/utils/gen_computegraph.py

@@ -16,7 +16,6 @@
     BOOL,
     DOUBLE,
     INT,
-    OLD_STRING,
     OPT_AT_DOUBLE_ARRAY_REF,
     OPT_AT_INT_ARRAY_REF,
     OPT_AT_TENSOR,
@@ -122,7 +121,9 @@ def __init__(
             ctype = cpp.argumenttype_type(
                 arg.type, mutable=arg.is_write, binds=arg.name
             )
-            cpp_type = ctype.cpp_type(strip_ref=True)
+            cpp_type = ctype.cpp_type(strip_ref=True).replace(
+                "c10::string_view", STRING
+            )
 
             self.args.append(
                 ATenArg(name=arg.name, cpp_type=cpp_type, default=arg.default)
@@ -494,7 +495,7 @@ def create_value_for(  # noqa: C901
             or ref.src_cpp_type == OPT_MEMORY_FORMAT
         ):
             ret_str += "add_none(); \n"
-        elif ref.src_cpp_type == STRING or ref.src_cpp_type == OLD_STRING:
+        elif ref.src_cpp_type == STRING:
             ret_str += f"add_string(std::string({ref.src_cpp_name})); \n"
         elif ref.src_cpp_type == TWO_TENSOR_TUPLE:
             ret_str += f"add_value_list({{{ref.name}_first, {ref.name}_second}}); \n"

backends/vulkan/test/op_tests/utils/gen_correctness_base.py

@@ -15,7 +15,6 @@
     BOOL,
     DOUBLE,
     INT,
-    OLD_STRING,
     OPT_AT_DOUBLE_ARRAY_REF,
     OPT_AT_INT_ARRAY_REF,
     OPT_AT_TENSOR,
@@ -156,7 +155,7 @@ def call_data_gen_fn(self, arg: Argument, data: Any, terminate: bool = True) ->
 
     def create_input_data(self, arg: Argument, data: Any) -> str:  # noqa: C901
         ctype = cpp.argumenttype_type(arg.type, mutable=arg.is_write, binds=arg.name)
-        cpp_type = ctype.cpp_type(strip_ref=True)
+        cpp_type = ctype.cpp_type(strip_ref=True).replace("c10::string_view", STRING)
 
         # Short cut exit for TENSORLIST, because it needs multiple lines of
         # construction, deviates from the rest.
@@ -218,7 +217,7 @@ def create_input_data(self, arg: Argument, data: Any) -> str:  # noqa: C901
                 ret_str += "std::nullopt;"
             else:
                 ret_str += f"{str(data)};"
-        elif cpp_type == STRING or cpp_type == OLD_STRING:
+        elif cpp_type == STRING:
             ret_str += f'std::string_view("{data}");'
         elif (
             cpp_type == OPT_SCALAR_TYPE

exir/passes/reinplace.py

@@ -507,4 +507,26 @@ def reinplace_pass(  # noqa: C901
             continue
 
         seen_nodes.update(node.all_input_nodes)
+
+    # Reinplace rewrites may rename output nodes (e.g. aten_relu_default →
+    # aten_relu__default) but replace_all_uses_with does not update the
+    # graph signature. Fix output_specs to match the actual output node args.
+    output_node = next((n for n in ep.graph.nodes if n.op == "output"), None)
+    if output_node is not None:
+        out_args = output_node.args[0]
+        if not isinstance(out_args, (tuple, list)):
+            out_args = (out_args,)
+        output_specs = ep.graph_signature.output_specs
+        assert len(output_specs) == len(out_args), (
+            f"reinplace: output spec count changed: "
+            f"{len(output_specs)} specs vs {len(out_args)} output args"
+        )
+        for spec, arg in zip(output_specs, out_args):
+            if (
+                isinstance(arg, torch.fx.Node)
+                and getattr(spec.arg, "name", None) is not None
+                and spec.arg.name != arg.name
+            ):
+                spec.arg.name = arg.name
+
     return ep

exir/tests/test_reinplace_pass.py

@@ -326,6 +326,25 @@ def forward(
         # kernels in the other tests in this file.
         edge.to_executorch()
 
+    def test_output_specs_updated_after_reinplace(self) -> None:
+        """reinplace_pass must update graph_signature.output_specs to match
+        the renamed output nodes (e.g. relu -> relu_)."""
+
+        class M(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return torch.relu(x + 1.0)
+
+        ep = export(M(), (torch.randn(4),), strict=True)
+        edge_program = to_edge(ep).exported_program()
+
+        custom_set = {edge_ops.edge.aten.relu.default}
+        ep = reinplace_pass(edge_program, ops_to_inplace=custom_set)
+
+        output_node = next(n for n in ep.graph.nodes if n.op == "output")
+        actual_names = [arg.name for arg in output_node.args[0] if hasattr(arg, "name")]
+        spec_names = [s.arg.name for s in ep.graph_signature.output_specs]
+        self.assertEqual(actual_names, spec_names)
+
     def test_broadcasting_self_not_reinplaced(self) -> None:
         """An op whose mutated arg (self) broadcasts up to a larger output
         must NOT be reinplaced: the in-place form cannot grow self

exir/verification/bindings.cpp

@@ -13,7 +13,6 @@
 #include <c10/core/ScalarType.h>
 #include <c10/macros/Macros.h>
 #include <c10/util/C++17.h>
-#include <c10/util/Optional.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 #include <torch/extension.h> // @manual=//caffe2:torch_extension