Add to_out_var pass variant that support cadence custom inplace ops.

Differential Revision: D93004811

Pull Request resolved: #17406

Author: hsharma35

backends/cadence/aot/BUCK

@@ -187,6 +187,20 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "to_out_var_pass",
+    srcs = [
+        "to_out_var_pass.py",
+    ],
+    typing = True,
+    deps = [
+        ":ops_registrations",
+        "//caffe2:torch",
+        "//executorch/exir/dialects:lib",
+        "//executorch/exir/passes:lib",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name = "graph_builder",
     srcs = [
@@ -649,3 +663,20 @@ fbcode_target(_kind = python_unittest,
         "//pytorch/ao:torchao",
     ],
 )
+
+fbcode_target(_kind = python_unittest,
+    name = "test_to_out_var_pass",
+    srcs = [
+        "tests/test_to_out_var_pass.py",
+    ],
+    typing = True,
+    deps = [
+        ":ops_registrations",
+        ":program_builder",
+        ":to_out_var_pass",
+        "//caffe2:torch",
+        "//executorch/exir:lib",
+        "//executorch/exir/dialects:lib",
+        "//later:lib",
+    ],
+)

backends/cadence/aot/tests/test_to_out_var_pass.py

@@ -0,0 +1,98 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import executorch.backends.cadence.aot.ops_registrations  # noqa
+import torch
+from executorch.backends.cadence.aot.program_builder import ProgramBuilder
+from executorch.backends.cadence.aot.to_out_var_pass import CadenceToOutVarPass
+from executorch.exir import ExecutorchBackendConfig
+from executorch.exir.dialects._ops import ops as exir_ops
+from later.unittest import TestCase
+
+
+class TestCadenceToOutVarPass(TestCase):
+    def test_serialize_with_slice_scatter_inplace(self) -> None:
+        """Test that a graph with slice_scatter_ can be serialized after CadenceToOutVarPass."""
+        builder = ProgramBuilder()
+        # Create input tensor placeholder
+        x = builder.placeholder("x", torch.randn(10, dtype=torch.float32))
+        # Create source tensor placeholder
+        src = builder.placeholder("src", torch.randn(3, dtype=torch.float32))
+
+        # Call slice_scatter_ inplace op
+        result = builder.call_operator(
+            op=exir_ops.edge.cadence.slice_scatter_.default,
+            args=(x, src, 0, 2, 5, 1),
+        )
+        builder.output([result])
+
+        # Get the edge program
+        edge_program = builder.get_edge_program()
+
+        # Apply CadenceToOutVarPass and serialize
+        exec_program = edge_program.to_executorch(
+            ExecutorchBackendConfig(
+                to_out_var_pass=CadenceToOutVarPass(),
+            )
+        )
+
+        # Verify serialization succeeded by checking the buffer is non-empty
+        buffer = exec_program.buffer
+        self.assertIsNotNone(buffer)
+        self.assertGreater(len(buffer), 0)
+
+    def test_serialize_with_mixed_ops(self) -> None:
+        """Test that a graph with mixed ops including slice_scatter_ can be serialized."""
+        builder = ProgramBuilder()
+        x = builder.placeholder("x", torch.randn(10, dtype=torch.float32))
+        y = builder.placeholder("y", torch.randn(10, dtype=torch.float32))
+        src = builder.placeholder("src", torch.randn(3, dtype=torch.float32))
+
+        # Add operation
+        add_result = builder.call_operator(exir_ops.edge.aten.add.Tensor, (x, y))
+
+        # Slice scatter inplace operation
+        scatter_result = builder.call_operator(
+            op=exir_ops.edge.cadence.slice_scatter_.default,
+            args=(add_result, src, 0, 2, 5, 1),
+        )
+
+        builder.output([scatter_result])
+
+        edge_program = builder.get_edge_program()
+
+        exec_program = edge_program.to_executorch(
+            ExecutorchBackendConfig(
+                to_out_var_pass=CadenceToOutVarPass(),
+            )
+        )
+
+        buffer = exec_program.buffer
+        self.assertIsNotNone(buffer)
+        self.assertGreater(len(buffer), 0)
+
+    def test_serialize_with_add_tensor(self) -> None:
+        """Test that a simple add graph can be serialized with CadenceToOutVarPass."""
+        builder = ProgramBuilder()
+        x = builder.placeholder("x", torch.randn(3, 5, dtype=torch.float32))
+        y = builder.placeholder("y", torch.randn(3, 5, dtype=torch.float32))
+
+        add = builder.call_operator(exir_ops.edge.aten.add.Tensor, (x, y))
+        builder.output([add])
+
+        edge_program = builder.get_edge_program()
+
+        exec_program = edge_program.to_executorch(
+            ExecutorchBackendConfig(
+                to_out_var_pass=CadenceToOutVarPass(),
+            )
+        )
+
+        buffer = exec_program.buffer
+        self.assertIsNotNone(buffer)
+        self.assertGreater(len(buffer), 0)

backends/cadence/aot/to_out_var_pass.py

@@ -0,0 +1,24 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import executorch.backends.cadence.aot.ops_registrations  # noqa
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.passes import ToOutVarPass
+from torch.fx.passes.infra.pass_base import PassResult
+
+
+class CadenceToOutVarPass(ToOutVarPass):
+    """Adds support for custom cadence inplace ops."""
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        for slice_scatter_inplace in graph_module.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.cadence.slice_scatter_.default
+        ):
+            slice_scatter_inplace.target = slice_scatter_inplace.target._op
+        return super().call(graph_module)