Remove deprecated XNNPACK capture utilities and migrate tests

backends/xnnpack/__init__.py

@@ -12,23 +12,17 @@
 
 # Exposed Configs in XNNPACK Package
 from .utils.configs import (
-    get_xnnpack_capture_config,
     get_xnnpack_edge_compile_config,
     get_xnnpack_executorch_backend_config,
 )
 
-# Easy util functions
-from .utils.utils import capture_graph_for_xnnpack
-
 # XNNPACK Backend
 from .xnnpack_preprocess import XnnpackBackend
 
 __all__ = [
     "XnnpackDynamicallyQuantizedPartitioner",
     "XnnpackPartitioner",
     "XnnpackBackend",
-    "capture_graph_for_xnnpack",
-    "get_xnnpack_capture_config",
     "get_xnnpack_edge_compile_config",
     "get_xnnpack_executorch_backend_config",
 ]

backends/xnnpack/test/test_xnnpack_utils.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -10,7 +10,7 @@
 
 import torch
 import torch.nn.functional as F
 from executorch import exir
 
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
     XnnpackDynamicallyQuantizedPartitioner,
@@ -25,7 +25,6 @@
     get_xnnpack_edge_compile_config,
     get_xnnpack_executorch_backend_config,
 )
-from executorch.backends.xnnpack.utils.utils import capture_graph_for_xnnpack
 
 # import the xnnpack backend implementation
 from executorch.backends.xnnpack.xnnpack_preprocess import XnnpackBackend
@@ -35,7 +34,7 @@
 from executorch.devtools.bundled_program.serialize import (
     serialize_from_bundled_program_to_flatbuffer,
 )
-from executorch.exir import ExecutorchProgram, ExirExportedProgram
+from executorch.exir import EdgeProgramManager, to_edge
 from executorch.exir.backend.backend_api import to_backend, validation_disabled
 
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
@@ -157,6 +156,14 @@
                 torch.allclose(model_output[0], ref_output, atol=1e-03, rtol=1e-03)
             )
 
+    def _capture_graph_for_xnnpack(
+        self, module: torch.nn.Module, sample_inputs: Tuple[torch.Tensor]
+    ) -> EdgeProgramManager:
+        return to_edge(
+            export(module, sample_inputs, strict=True),
+            compile_config=get_xnnpack_edge_compile_config(),
+        ).transform(*get_transform_passes())
+
     def lower_module_and_test_output(
         self,
         module: Any,
@@ -167,15 +174,15 @@
         # TODO: remove this after we migrate to use long term flow
         quantizer_api_test: bool = False,
         dump_bundled_program: bool = False,  # for debugging, dump the generated bundled program file
-    ) -> ExirExportedProgram:
+    ) -> EdgeProgramManager:
         """
         Helper testing function that takes a torch.nn.Module and lowers it to XNNPACK with
         the given sample inputs. It then runs the lowered module and compares its
         outputs with the outputs of the eager module.
         """
 
         if quantizer_api_test:
-            assert isinstance(module, ExirExportedProgram)
+            assert isinstance(module, EdgeProgramManager)
             edge_program = module
         else:
 
@@ -187,7 +194,7 @@
                 def forward(self, *args):
                     return self.one_module(*args)
 
-            edge_program = capture_graph_for_xnnpack(WrappedModule(), sample_inputs)
+            edge_program = self._capture_graph_for_xnnpack(WrappedModule(), sample_inputs)
 
         partitioner = None
         if quantized:
@@ -201,35 +208,32 @@
         if use_partitioner:
             with validation_disabled():
                 delegated_program = edge_program
-                delegated_program.exported_program = to_backend(
-                    edge_program.exported_program, partitioner
+                delegated_program._edge_programs["forward"] = to_backend(
+                    edge_program.exported_program(), partitioner
                 )
 
-            executorch_program: ExecutorchProgram = delegated_program.to_executorch(
+            executorch_program = delegated_program.to_executorch(
                 get_xnnpack_executorch_backend_config([SpecPropPass()]),
             )
         else:
-            delegated_program = to_backend(
-                "XnnpackBackend", edge_program.exported_program, []
+            delegated_module = to_backend(
+                "XnnpackBackend", edge_program.exported_program(), []
             )
 
-            exported_program: ExirExportedProgram = capture_graph_for_xnnpack(
-                delegated_program, sample_inputs
+            exported_program = self._capture_graph_for_xnnpack(
+                delegated_module, sample_inputs
             )
-            executorch_program: ExecutorchProgram = exported_program.to_executorch(
+            executorch_program = exported_program.to_executorch(
                 get_xnnpack_executorch_backend_config(),
             )
 
-        # print("Graph Module with delegate:")
-        # delegated_module.print_readable()
-
         # Assert the backend name is xnnpack
         self.assertEqual(
-            executorch_program.program.execution_plan[0].delegates[0].id,
+            executorch_program.executorch_program.execution_plan[0].delegates[0].id,
             XnnpackBackend.__name__,
         )
 
-        ref_output = delegated_program(*sample_inputs)
+        ref_output = delegated_program.exported_program().module()(*sample_inputs)
         if dump_bundled_program:
             save_bundled_program(
                 representative_inputs=sample_inputs,
@@ -325,14 +329,9 @@
         prepared = prepare_pt2e(m, quantizer)
         converted = convert_pt2e(prepared)
 
-        captured_program = exir.capture(
-            converted,
-            example_inputs,
-            config=exir.CaptureConfig(enable_aot=True, _unlift=True),
-        )
-
-        edge_program = captured_program.to_edge(
-            get_xnnpack_edge_compile_config()
+        edge_program = to_edge(
+            export(converted, example_inputs, strict=True),
+            compile_config=get_xnnpack_edge_compile_config(),
         ).transform(*get_transform_passes())
         delegated_module = self.lower_module_and_test_output(
             module=edge_program,
@@ -350,7 +349,7 @@
         }
         for op in supported_ops:
             FileCheck().check_count(op, 0, exactly=True).run(
-                delegated_module.exported_program.graph_module.code
+                delegated_module.exported_program().graph_module.code
             )
 
     def _test_xnnpack_dqlinear(
@@ -398,12 +397,12 @@
             prepared_linear,
         )
 
-        captured_dqlinear = capture_graph_for_xnnpack(converted_linear, example_inputs)
+        captured_dqlinear = self._capture_graph_for_xnnpack(converted_linear, example_inputs)
 
-        captured_dqlinear.exported_program.graph_module.graph.print_tabular()
+        captured_dqlinear.exported_program().graph_module.graph.print_tabular()
 
         lowered_module = to_backend(
-            "XnnpackBackend", captured_dqlinear.exported_program, []
+            "XnnpackBackend", captured_dqlinear.exported_program(), []
         )
 
         class CompositeModule(torch.nn.Module):
@@ -417,19 +416,19 @@
         composite_model = CompositeModule()
         composite_model(*example_inputs)
 
-        exported_program: ExirExportedProgram = capture_graph_for_xnnpack(
+        exported_program = self._capture_graph_for_xnnpack(
             composite_model, example_inputs
         )
-        executorch_program: ExecutorchProgram = exported_program.to_executorch(
+        executorch_program = exported_program.to_executorch(
             get_xnnpack_executorch_backend_config(),
         )
 
         self.assertEqual(
-            executorch_program.program.execution_plan[0].delegates[0].id,
+            executorch_program.executorch_program.execution_plan[0].delegates[0].id,
             XnnpackBackend.__name__,
         )
 
-        ref_output = captured_dqlinear(*example_inputs)
+        ref_output = captured_dqlinear.exported_program().module()(*example_inputs)
         ref_output = composite_model(*example_inputs)
         print("ref_output:", ref_output)
 

backends/xnnpack/utils/configs.py

@@ -4,10 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import List, Optional
+from typing import List
 
 import executorch.exir as exir
-from executorch.exir import CaptureConfig
 from executorch.exir.pass_manager import PassType
 
 
@@ -33,17 +32,3 @@ def get_xnnpack_executorch_backend_config(
         passes=additional_passes,
         extract_delegate_segments=True,
     )
-
-
-def get_xnnpack_capture_config(
-    dynamic_shape=False,
-    enable_aot: Optional[bool] = None,
-    unlift: Optional[bool] = None,
-):
-    if enable_aot is None:
-        return CaptureConfig(enable_dynamic_shape=dynamic_shape)
-    else:
-        unlift = unlift if unlift is not None else enable_aot
-        return CaptureConfig(
-            enable_dynamic_shape=dynamic_shape, enable_aot=enable_aot, _unlift=unlift
-        )

backends/xnnpack/utils/utils.py

@@ -6,14 +6,8 @@
 
 from typing import Any, cast, Optional, Tuple
 
-import executorch.exir as exir
 import torch
 
-from executorch.backends.xnnpack.utils.configs import (
-    get_transform_passes,
-    get_xnnpack_capture_config,
-    get_xnnpack_edge_compile_config,
-)
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
 
@@ -28,24 +22,6 @@
 from torchao.quantization.pt2e.utils import _is_conv_node, _is_conv_transpose_node
 
 
-### XNNPACK Capture ###
-def capture_graph_for_xnnpack(
-    module: torch.nn.Module,
-    inputs: Tuple[torch.Tensor],
-    enable_aot: Optional[bool] = None,
-    unlift: Optional[bool] = None,
-) -> exir.ExirExportedProgram:
-    return (
-        exir.capture(
-            module,
-            inputs,
-            get_xnnpack_capture_config(enable_aot=enable_aot, unlift=unlift),
-        )
-        .to_edge(get_xnnpack_edge_compile_config())
-        .transform(*get_transform_passes())
-    )
-
-
 ### XNNPACK Utils ###
 PERM_NCHW_TO_NHWC = [0, 2, 3, 1]
 PERM_NHWC_TO_NCHW = [0, 3, 1, 2]

devtools/size_analysis_tool/size_analysis_tool_test.py

@@ -11,15 +11,16 @@
     XnnpackFloatingPointPartitioner,
 )
 from executorch.backends.xnnpack.utils.configs import (
+    get_xnnpack_edge_compile_config,
     get_xnnpack_executorch_backend_config,
 )
-from executorch.backends.xnnpack.utils.utils import capture_graph_for_xnnpack
 
 from executorch.devtools.size_analysis_tool.size_analysis_tool import (
     generate_model_size_information,
 )
-from executorch.exir.backend.backend_api import to_backend, validation_disabled
+from executorch.exir import to_edge
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
+from torch.export import export
 
 
 class SizeAnalysisToolTest(unittest.TestCase):
@@ -52,14 +53,13 @@ def forward(self, x):
 
         test_input = torch.ones(size=(4, 7, 5, 6), dtype=torch.float)
 
-        edge_program = capture_graph_for_xnnpack(mm, (test_input,))
+        edge_program = to_edge(
+            export(mm, (test_input,), strict=True),
+            compile_config=get_xnnpack_edge_compile_config(),
+        )
         partitioner = XnnpackFloatingPointPartitioner()
 
-        with validation_disabled():
-            delegated_program = edge_program
-            delegated_program.exported_program = to_backend(
-                edge_program.exported_program, partitioner
-            )
+        delegated_program = edge_program.to_backend(partitioner)
 
         program = delegated_program.to_executorch(
             get_xnnpack_executorch_backend_config([SpecPropPass()]),

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 = [