Elide static view copies

Differential Revision: D68984189

Pull Request resolved: #8197

Author: metascroy

exir/capture/_config.py

@@ -87,6 +87,7 @@ class ExecutorchBackendConfig:
 
     # If set to true, view_copy operations will be converted to lightweight
     # view operations in the ET runtime
+    # Moreover, static views will be elided from the ExecuTorch graph
     remove_view_copy: bool = True
 
     # If set to true, all constant tensors will be stored in a separate file,

exir/emit/_emitter.py

@@ -943,6 +943,16 @@ def _emit_control_flow(
     def _emit_view(self, args: Tuple[_Argument, ...]) -> _EmitterValue:
         assert len(args) == 2
 
+        # Elide the view if it is static and memory planned
+        spec = self.node.meta["spec"]
+        is_static = spec.is_static_shape_tensor
+        is_memory_planned = (spec.mem_id is not None) and (spec.mem_offset is not None)
+        is_memory_planned = is_memory_planned or (
+            spec.const and spec.storage is not None
+        )
+        if is_static and is_memory_planned:
+            return self._emit_spec(spec)
+
         self_arg = self._emit_argument(args[0], torch.TensorType)  # pyre-ignore[6]
         size_arg = self._emit_argument(args[1], torch.ListType.ofInts())
         out_arg = self._emit_argument(

exir/emit/test/test_emit.py

@@ -331,29 +331,38 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             "aten::sin",
             "aten::relu",
             "aten::max",
-            "executorch_prim::et_view",  # aten::view_copy if ExecutorchBackendConfig.remove_view_copy = False
         ]
 
+        def expected_view_ops(config):
+            if config.remove_view_copy:
+                return []
+            else:
+                return ["aten::view_copy"]
+
         for opname in removed_ops:
             self.assertEqual(
                 self.count_node(edge.exported_program().graph_module, opname), 0
             )
         for opname in expected_ops:
-            if (
-                opname != "executorch_prim::et_view"
-            ):  # et_view appears as call_function with target = memory.view in graph
-                self.assertTrue(
-                    self.count_node(edge.exported_program().graph_module, opname) >= 1
-                )
-
-        program = edge.to_executorch().executorch_program
-        for opname in removed_ops:
             self.assertTrue(
-                all(op.name != opname for op in program.execution_plan[0].operators)
+                self.count_node(edge.exported_program().graph_module, opname) >= 1
             )
-        for opname in expected_ops:
+
+        for remove_view_copy in [True, False]:
+            config = exir.ExecutorchBackendConfig(remove_view_copy=remove_view_copy)
+            edge_copy = deepcopy(edge)
+            program = edge_copy.to_executorch(config=config).executorch_program
+            for opname in removed_ops:
+                self.assertTrue(
+                    all(op.name != opname for op in program.execution_plan[0].operators)
+                )
+            for opname in expected_ops + expected_view_ops(config):
+                self.assertTrue(
+                    any(op.name == opname for op in program.execution_plan[0].operators)
+                )
             self.assertTrue(
-                any(op.name == opname for op in program.execution_plan[0].operators)
+                len(program.execution_plan[0].operators)
+                == len(expected_ops + expected_view_ops(config))
             )
 
     def test_operators_unique(self) -> None:

exir/tests/test_remove_view_copy.py

@@ -194,16 +194,41 @@ def test_spec(self) -> None:
         self.assertEqual(plan.operators[2].name, "aten::view_copy")
 
         instructions = plan.chains[0].instructions
-        self.assertEqual(len(instructions), 7)
+        self.assertEqual(len(instructions), 5)
 
         self.assertEqual(instructions[0].instr_args.op_index, 0)  # view @ idx2
-        self.assertEqual(instructions[1].instr_args.op_index, 0)  # view @ idx3
-        self.assertEqual(instructions[2].instr_args.op_index, 1)  # aten:mul @ idx6
-        self.assertEqual(instructions[3].instr_args.op_index, 0)  # view @ idx7
-        self.assertEqual(instructions[4].instr_args.op_index, 1)  # aten:mul @ idx9
+        self.assertEqual(instructions[1].instr_args.op_index, 1)  # aten:mul @ idx6
+        self.assertEqual(instructions[2].instr_args.op_index, 1)  # aten:mul @ idx9
         self.assertEqual(
-            instructions[5].instr_args.op_index, 2
+            instructions[3].instr_args.op_index, 2
         )  # aten:view_copy @ idx11
         self.assertEqual(
-            instructions[6].instr_args.op_index, 2
+            instructions[4].instr_args.op_index, 2
         )  # aten:view_copy @ idx11
+
+    def test_elide_static_views_does_not_remove_dynamic_views(self) -> None:
+        class TestModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                x = x + x
+                x = x.view(-1, 1)
+                return 2 * x
+
+        model = TestModel()
+        model.eval()
+        example_inputs = (torch.rand(5, 6),)
+        dynamic_shapes = {"x": {0: torch.export.Dim("dim0", min=1, max=10)}}
+        ep = torch.export.export(
+            model, example_inputs, strict=True, dynamic_shapes=dynamic_shapes
+        )
+        etpm = to_edge(ep).to_executorch(
+            config=ExecutorchBackendConfig(
+                remove_view_copy=True,
+                memory_planning_pass=MemoryPlanningPass(alloc_graph_input=True),
+            ),
+        )
+        plan = etpm.executorch_program.execution_plan[0]
+        op_names = [op.name for op in plan.operators]
+        self.assertTrue("executorch_prim::et_view" in op_names)