Copy returned views out of AOTI methods as contiguous tensors

ReplaceViewCopyWithViewPass turns a method's slice_copy, select_copy and
split_copy into views so that AOTInductor compiles them natively. When
the view is what the method returns, AOTInductor returns it with its
base's strides, and the CUDA backend fails to copy it into the output
tensor ("SlimTensor must be contiguous"); a policy that returns
sample[..., :action_dim] cannot run.

Clone a returned view with contiguous memory format. A strided view is
then copied out by a kernel. Inductor still drops the clone of a view
that is already contiguous, so a returned slice that starts at an
offset, such as x[1:], still fails the runtime's storage offset check.

Author: edmundshieh

backends/aoti/passes/replace_view_copy_with_view.py

@@ -9,7 +9,8 @@
 # because this should only be used in the AOTInductor backend, as it
 # has less restrictions on whether the tensor memory is densely packed,
 
-from typing import Dict, Iterable
+import operator
+from typing import Dict, Iterable, List
 
 import torch
 from executorch.exir.dialects._ops import ops
@@ -40,22 +41,47 @@ class ReplaceViewCopyWithViewPass(ExportPass):
     def call(self, graph_module: fx.GraphModule) -> PassResult:
         graph_changed = False
 
-        for node in graph_module.graph.nodes:
+        for node in list(graph_module.graph.nodes):
             if node.op != "call_function" or node.target not in _VIEW_TARGETS:
                 continue
 
             if self._has_blocking_user(node, node.users.keys()):
                 continue
 
+            edge = isinstance(node.target, EdgeOpOverload)
             node.target = _VIEW_TARGETS[node.target]
             graph_changed = True
+            # Method outputs are contiguous tensors, so a returned view is
+            # copied out as one.
+            for returned in self._returned_values(node):
+                self._clone_contiguous(graph_module.graph, returned, edge)
 
         if graph_changed:
             graph_module.graph.lint()
             graph_module.recompile()
 
         return PassResult(graph_module, graph_changed)
 
+    def _returned_values(self, node: fx.Node) -> List[fx.Node]:
+        values = [node] if any(user.op == "output" for user in node.users) else []
+        for user in node.users:
+            if user.target is operator.getitem and any(
+                getitem_user.op == "output" for getitem_user in user.users
+            ):
+                values.append(user)
+        return values
+
+    def _clone_contiguous(self, graph: fx.Graph, node: fx.Node, edge: bool) -> None:
+        # The returned node becomes the clone, so the outputs keep their names.
+        with graph.inserting_before(node):
+            view = graph.call_function(node.target, node.args, node.kwargs)
+        view.meta = dict(node.meta)
+        node.target = (
+            ops.edge.aten.clone.default if edge else torch.ops.aten.clone.default
+        )
+        node.args = (view,)
+        node.kwargs = {"memory_format": torch.contiguous_format}
+
     def _has_blocking_user(self, node: fx.Node, users: Iterable[fx.Node]) -> bool:
         for user in users:
             if self._is_mutating_user(node, user) or self._is_view_user(node, user):

backends/cuda/passes/tests/test_replace_view_copy_with_view.py

@@ -0,0 +1,60 @@
+# 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.
+
+import unittest
+
+import torch
+from executorch.backends.aoti.passes.replace_view_copy_with_view import (
+    ReplaceViewCopyWithViewPass,
+)
+from executorch.exir import to_edge
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export import export
+
+
+class TestReplaceViewCopyWithViewPass(unittest.TestCase):
+    def _targets(self, module, inputs):
+        graph_module = (
+            to_edge(export(module, inputs, strict=True)).exported_program().graph_module
+        )
+        ReplaceViewCopyWithViewPass()(graph_module)
+        return [
+            node.target
+            for node in graph_module.graph.nodes
+            if node.op == "call_function"
+        ]
+
+    def test_intermediate_slice_becomes_view(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x[:, :6] * 2
+
+        targets = self._targets(M(), (torch.randn(4, 32),))
+        self.assertIn(exir_ops.edge.aten.slice.Tensor, targets)
+        self.assertNotIn(exir_ops.edge.aten.slice_copy.Tensor, targets)
+        self.assertNotIn(exir_ops.edge.aten.clone.default, targets)
+
+    def test_returned_slice_is_cloned_contiguous(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return (x * 2)[:, :6]
+
+        targets = self._targets(M(), (torch.randn(4, 32),))
+        self.assertIn(exir_ops.edge.aten.slice.Tensor, targets)
+        self.assertEqual(targets[-1], exir_ops.edge.aten.clone.default)
+
+    def test_returned_select_is_cloned_contiguous(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return (x * 2)[:, 2]
+
+        targets = self._targets(M(), (torch.randn(4, 32),))
+        self.assertIn(exir_ops.edge.aten.select.int, targets)
+        self.assertEqual(targets[-1], exir_ops.edge.aten.clone.default)
+
+
+if __name__ == "__main__":
+    unittest.main()