Migrate executorch/ tests from exir.capture to torch.export + to_edge

Differential Revision: D95605454

Pull Request resolved: #18111

Author: JacobSzwejbka

exir/backend/test/BUCK

@@ -158,38 +158,6 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
-fbcode_target(_kind = runtime.python_test,
-    name = "test_backends",
-    srcs = [
-        "test_backends.py",
-    ],
-    preload_deps = [
-        "//executorch/configurations:optimized_native_cpu_ops",
-        "//executorch/kernels/quantized:custom_ops_generated_lib",
-        "//executorch/runtime/executor/test:test_backend_compiler_lib",
-    ],
-    deps = [
-        ":backend_with_compiler_demo",
-        ":hta_partitioner_demo",
-        ":op_partitioner_demo",
-        ":demo_backend",
-        "//caffe2:torch",
-        "//caffe2/functorch:functorch_src",
-        "//executorch/exir:delegate",
-        "//executorch/exir:graph_module",
-        "//executorch/exir:lib",
-        "//executorch/exir:lowered_backend_module",
-        "//executorch/exir:print_program",
-        "//executorch/exir:schema",
-        "//executorch/exir/backend:backend_api",
-        "//executorch/exir/backend:compile_spec_schema",
-        "//executorch/exir/backend:partitioner",
-        "//executorch/exir/dialects:lib",
-        "//executorch/extension/pybindings:portable_lib",  # @manual
-        "//executorch/extension/pytree:pylib",
-    ],
-)
-
 fbcode_target(_kind = runtime.python_test,
     name = "test_to_backend_multi_method",
     srcs = [

exir/backend/test/hta_partitioner_demo.py

@@ -9,6 +9,7 @@
 
 import torch
 from executorch import exir
+from executorch.exir import to_edge
 from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import (
     generate_pattern_op_partitions,
 )
@@ -20,7 +21,7 @@
 )
 from executorch.exir.backend.test.demo_backend import DemoBackend
 from executorch.exir.backend.utils import tag_constant_data
-from torch.export import ExportedProgram
+from torch.export import export, ExportedProgram
 from torch.fx.passes.infra.partitioner import Partition
 
 
@@ -63,56 +64,30 @@ def forward(self, x_raw, h, c):
         input_h = torch.ones([1, 32])
         input_c = torch.ones([1, 32])
 
-        pattern_lstm_conv_lifted = (
-            exir.capture(
-                LSTMConvPattern(),
-                (input_x, input_h, input_c),
-                exir.CaptureConfig(enable_aot=True),
-            )
-            .to_edge(
-                # torch._export.verifier.SpecViolationError: Operator torch._ops.aten.mkldnn_rnn_layer.default is not Aten Canonical.
-                exir.EdgeCompileConfig(_check_ir_validity=False)
-            )
-            .exported_program.graph_module
-        )
         pattern_lstm_conv = (
-            exir.capture(
-                LSTMConvPattern(),
-                (input_x, input_h, input_c),
-                exir.CaptureConfig(),
-            )
-            .to_edge(
+            to_edge(
+                export(LSTMConvPattern(), (input_x, input_h, input_c), strict=True),
                 # torch._export.verifier.SpecViolationError: Operator torch._ops.aten.mkldnn_rnn_layer.default is not Aten Canonical.
-                exir.EdgeCompileConfig(_check_ir_validity=False)
+                compile_config=exir.EdgeCompileConfig(_check_ir_validity=False),
             )
-            .exported_program.graph_module
+            .exported_program()
+            .graph_module
         )
 
-        def sub(x, y):
-            return torch.sub(x, y)
+        class SubModule(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.sub(x, y)
 
-        pattern_sub_lifted = (
-            exir.capture(
-                sub,
-                (input_x, input_h),
-                exir.CaptureConfig(enable_aot=True, _unlift=False),
-            )
-            .to_edge(exir.EdgeCompileConfig(_use_edge_ops=True))
-            .exported_program.graph_module
-        )
         pattern_sub = (
-            exir.capture(
-                sub,
-                (input_x, input_h),
-                exir.CaptureConfig(),
+            to_edge(
+                export(SubModule(), (input_x, input_h), strict=True),
+                compile_config=exir.EdgeCompileConfig(_use_edge_ops=True),
             )
-            .to_edge()
-            .exported_program.graph_module
+            .exported_program()
+            .graph_module
         )
         self.patterns = [
-            pattern_lstm_conv_lifted.graph,
             pattern_lstm_conv.graph,
-            pattern_sub_lifted.graph,
             pattern_sub.graph,
         ]
 
@@ -239,33 +214,17 @@ def forward(self, x_raw, h, c):
         input_h = torch.ones([1, 32])
         input_c = torch.ones([1, 32])
 
-        pattern_lstm_conv_lifted = (
-            exir.capture(
-                LSTMConvPattern(),
-                (input_x, input_h, input_c),
-                exir.CaptureConfig(enable_aot=True),
-            )
-            .to_edge(
-                # torch._export.verifier.SpecViolationError: Operator torch._ops.aten.mkldnn_rnn_layer.default is not Aten Canonical.
-                exir.EdgeCompileConfig(_check_ir_validity=False)
-            )
-            .exported_program.graph_module
-        )
-        pattern_lstm_conv_unlifted = (
-            exir.capture(
-                LSTMConvPattern(),
-                (input_x, input_h, input_c),
-                exir.CaptureConfig(),
-            )
-            .to_edge(
+        pattern_lstm_conv = (
+            to_edge(
+                export(LSTMConvPattern(), (input_x, input_h, input_c), strict=True),
                 # torch._export.verifier.SpecViolationError: Operator torch._ops.aten.mkldnn_rnn_layer.default is not Aten Canonical.
-                exir.EdgeCompileConfig(_check_ir_validity=False)
+                compile_config=exir.EdgeCompileConfig(_check_ir_validity=False),
             )
-            .exported_program.graph_module
+            .exported_program()
+            .graph_module
         )
         self.patterns = [
-            pattern_lstm_conv_lifted.graph,
-            pattern_lstm_conv_unlifted.graph,
+            pattern_lstm_conv.graph,
         ]
         # Only (lstm + conv) pattern is lowerable