Migrate devtools tests and clean up remaining executorch/ CaptureConfig refs (#18135)

Differential Revision: D95605485

Pull Request resolved: #18135

Author: JacobSzwejbka

backends/qualcomm/utils/utils.py

@@ -134,10 +134,6 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool:
         return True
 
 
-def qnn_capture_config():
-    return exir.CaptureConfig(enable_aot=True)
-
-
 def qnn_edge_config() -> exir.EdgeCompileConfig:
     return exir.EdgeCompileConfig(
         _check_ir_validity=False,

devtools/etrecord/tests/etrecord_test.py

@@ -618,12 +618,10 @@ def test_add_extra_export_modules(self):
 
         # Create additional module to add
         f2 = models.BasicSinMax()
-        captured_output2 = exir.capture(
-            f2, f2.get_random_inputs(), exir.CaptureConfig()
-        )
+        captured_output2 = export(f2, f2.get_random_inputs(), strict=True)
 
         extra_modules = {
-            "new_module": captured_output2.exported_program,
+            "new_module": captured_output2,
         }
 
         # Add extra export modules
@@ -640,7 +638,7 @@ def test_add_extra_export_modules(self):
         )
         self.check_graph_closeness(
             etrecord.graph_map["new_module/forward"],
-            captured_output2.exported_program.graph_module,
+            captured_output2.graph_module,
         )
 
     def test_add_extra_export_modules_reserved_name_validation(self):
@@ -1066,13 +1064,11 @@ def test_add_exported_program_already_exists_exception(self):
 
         # Create another exported program to try to add
         f2 = models.BasicSinMax()
-        captured_output2 = exir.capture(
-            f2, f2.get_random_inputs(), exir.CaptureConfig()
-        )
+        captured_output2 = export(f2, f2.get_random_inputs(), strict=True)
 
         # Verify that adding exported program raises RuntimeError
         with self.assertRaises(RuntimeError) as context:
-            etrecord.add_exported_program(captured_output2.exported_program)
+            etrecord.add_exported_program(captured_output2)
 
         self.assertIn(
             "Exported program already exists in the ETRecord",
@@ -1202,11 +1198,11 @@ def test_add_edge_dialect_program_already_exists_exception(self):
 
         # Create another edge program to try to add
         f2 = models.BasicSinMax()
-        captured_output2 = exir.capture(
-            f2, f2.get_random_inputs(), exir.CaptureConfig()
-        )
-        edge_output2 = captured_output2.to_edge(
-            exir.EdgeCompileConfig(_check_ir_validity=False, _use_edge_ops=False)
+        edge_output2 = to_edge(
+            export(f2, f2.get_random_inputs(), strict=True),
+            compile_config=exir.EdgeCompileConfig(
+                _check_ir_validity=False, _use_edge_ops=False
+            ),
         )
 
         # Verify that adding edge dialect program raises RuntimeError

devtools/size_analysis_tool/size_analysis_tool_test.py

@@ -11,6 +11,7 @@
     XnnpackFloatingPointPartitioner,
 )
 from executorch.backends.xnnpack.utils.configs import (
+    get_transform_passes,
     get_xnnpack_edge_compile_config,
     get_xnnpack_executorch_backend_config,
 )
@@ -19,6 +20,7 @@
     generate_model_size_information,
 )
 from executorch.exir import to_edge
+from executorch.exir.backend.backend_api import to_backend, validation_disabled
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
 from torch.export import export
 
@@ -56,10 +58,14 @@ def forward(self, x):
         edge_program = to_edge(
             export(mm, (test_input,), strict=True),
             compile_config=get_xnnpack_edge_compile_config(),
-        )
+        ).transform(get_transform_passes())
         partitioner = XnnpackFloatingPointPartitioner()
 
-        delegated_program = edge_program.to_backend(partitioner)
+        with validation_disabled():
+            delegated_program = edge_program
+            delegated_program._edge_programs["forward"] = to_backend(
+                edge_program.exported_program(), partitioner
+            )
 
         program = delegated_program.to_executorch(
             get_xnnpack_executorch_backend_config([SpecPropPass()]),

test/end2end/exported_module.py

@@ -67,7 +67,6 @@ def export(
         methods: Sequence[str] = ("forward",),
         ignore_to_out_var_failure: bool = False,
         dynamic_memory_planning_mode: DynamicMemoryPlanningMode = DynamicMemoryPlanningMode.UPPER_BOUND,
-        capture_config=None,
         export_joint_graph: bool = False,
         external_constants: bool = False,
         export_state_names: bool = False,
@@ -146,8 +145,6 @@ def return_wrapper():
 
         method_name_to_dynamic_shapes = None
         if hasattr(eager_module, "get_dynamic_shapes"):
-            assert capture_config is not None
-            assert capture_config.enable_aot is True
             trace_dynamic_shapes = eager_module.get_dynamic_shapes()  # type: ignore[operator]
             method_name_to_dynamic_shapes = {}
             for method in methods:

test/end2end/test_end2end.py

@@ -21,12 +21,7 @@
 import executorch.extension.pytree as pytree
 import torch
 
-from executorch.exir import (
-    CaptureConfig,
-    EdgeCompileConfig,
-    ExecutorchBackendConfig,
-    memory,
-)
+from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, memory
 from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
 from executorch.exir.emit import emit_program
 from executorch.exir.pass_manager import PassManager
@@ -471,7 +466,6 @@ def maketest(
     allow_non_contiguous_tensor: bool = False,
     method: str = "forward",
     dynamic_memory_planning_mode: DynamicMemoryPlanningMode = DynamicMemoryPlanningMode.UPPER_BOUND,
-    capture_config=None,
     verify_graph: Optional[Callable] = None,
 ) -> Callable[[unittest.TestCase], None]:
     r"""Returns a TestCase method to test the provided module class and method.
@@ -507,7 +501,6 @@ def wrapper(self: unittest.TestCase) -> None:
             methods=(method,),
             ignore_to_out_var_failure=ignore_to_out_var_failure,
             dynamic_memory_planning_mode=dynamic_memory_planning_mode,
-            capture_config=capture_config,
         )
         if verify_graph:
             verify_graph(self, module.exported_program.graph_module)
@@ -599,9 +592,6 @@ def test_ops_return_multi(self):
     def test_mem_planning_toy_model(self):
         maketest(
             ToyModelForMemPlanning,
-            capture_config=exir.CaptureConfig(
-                enable_dynamic_shape=True,
-            ),
         )(self)
 
     # TODO: add ops implementations and turn on 'run_executor'
@@ -621,9 +611,6 @@ def test_containers(self):
         maketest(
             ModuleContainers,
             do_tree_flatten=True,
-            capture_config=exir.CaptureConfig(
-                enable_dynamic_shape=True,
-            ),
         )(self)
 
     # can not run the graph module since the out variance with tensor list out
@@ -675,9 +662,6 @@ def test_intermediate_dynamic_shape(self):
             ModuleIntermediateDynamicShape,
             run_graph_module=False,
             allow_non_contiguous_tensor=True,
-            capture_config=exir.CaptureConfig(
-                enable_dynamic_shape=True,
-            ),
         )(self)
 
     # TODO(shunting): some non constant tensors for transformer are non-contiguous.
@@ -697,10 +681,6 @@ def test_transformer_encode(self):
     def test_ft_cond_basic(self):
         maketest(
             FTCondBasic,
-            capture_config=exir.CaptureConfig(
-                enable_dynamic_shape=True,
-                enable_functionalization=False,  # TODO enable functionalization
-            ),
         )(self)
 
     def test_ft_map_basic(self):
@@ -746,10 +726,6 @@ def test_ft_map_basic(self):
     def test_ft_cond_dynshape(self):
         maketest(
             FTCondDynShape,
-            capture_config=exir.CaptureConfig(
-                enable_dynamic_shape=True,
-                enable_functionalization=False,  # TODO enable functionalization
-            ),
         )(self)
 
     def test_ft_map_dynshape(self):
@@ -802,9 +778,6 @@ def test_ft_map_dynshape(self):
     def test_batch_norm(self):
         maketest(
             BatchNormModel,
-            capture_config=exir.CaptureConfig(
-                enable_dynamic_shape=True,
-            ),
             verify_graph=BatchNormModel.verify_graph,
             # TODO: lean mode does not have native_batch_norm.out implemented
             # run this on aten mode.

test/models/export_program.py

@@ -13,7 +13,6 @@
 from typing import Any, Dict, List, Type
 
 import torch
-from executorch.exir import CaptureConfig
 from executorch.exir.passes import MemoryPlanningPass
 from executorch.exir.program._program import ExecutorchProgramManager
 from torch import nn
@@ -140,7 +139,7 @@ def get_memory_planning_pass(self):
 
     @staticmethod
     def get_export_kwargs():
-        return {"capture_config": CaptureConfig(pt2_mode=True, enable_aot=True)}
+        return {}
 
 
 class ModuleAddMul(torch.nn.Module):