[test] add unified jit importer test config (#4083)

* make test config more clear

Author: qingyunqu

projects/pt1/e2e_testing/main.py

@@ -19,13 +19,11 @@
 # Available test configs.
 from torch_mlir_e2e_test.configs import (
     LazyTensorCoreTestConfig,
-    LinalgOnTensorsBackendTestConfig,
-    StablehloBackendTestConfig,
     NativeTorchTestConfig,
     OnnxBackendTestConfig,
     TorchScriptTestConfig,
-    TosaBackendTestConfig,
     TorchDynamoTestConfig,
+    JITImporterTestConfig,
     FxImporterTestConfig,
 )
 
@@ -152,15 +150,15 @@ def main():
 
     # Find the selected config.
     if args.config == "linalg":
-        config = LinalgOnTensorsBackendTestConfig(RefBackendLinalgOnTensorsBackend())
+        config = JITImporterTestConfig(RefBackendLinalgOnTensorsBackend())
         xfail_set = LINALG_XFAIL_SET
         crashing_set = LINALG_CRASHING_SET
     elif args.config == "stablehlo":
-        config = StablehloBackendTestConfig(LinalgOnTensorsStablehloBackend())
+        config = JITImporterTestConfig(LinalgOnTensorsStablehloBackend(), "stablehlo")
         xfail_set = all_test_unique_names - STABLEHLO_PASS_SET
         crashing_set = STABLEHLO_CRASHING_SET
     elif args.config == "tosa":
-        config = TosaBackendTestConfig(LinalgOnTensorsTosaBackend())
+        config = JITImporterTestConfig(LinalgOnTensorsTosaBackend(), "tosa")
         xfail_set = all_test_unique_names - TOSA_PASS_SET
         crashing_set = TOSA_CRASHING_SET
     elif args.config == "native_torch":

projects/pt1/python/torch_mlir_e2e_test/configs/__init__.py

@@ -4,11 +4,9 @@
 # Also available under a BSD-style license. See LICENSE.
 
 from .lazy_tensor_core import LazyTensorCoreTestConfig
-from .linalg_on_tensors_backend import LinalgOnTensorsBackendTestConfig
 from .native_torch import NativeTorchTestConfig
 from .onnx_backend import OnnxBackendTestConfig
 from .torchscript import TorchScriptTestConfig
-from .stablehlo_backend import StablehloBackendTestConfig
-from .tosa_backend import TosaBackendTestConfig
 from .torchdynamo import TorchDynamoTestConfig
+from .jit_importer_backend import JITImporterTestConfig
 from .fx_importer_backend import FxImporterTestConfig

projects/pt1/python/torch_mlir_e2e_test/configs/stablehlo_backend.py renamed to projects/pt1/python/torch_mlir_e2e_test/configs/jit_importer_backend.py

@@ -8,30 +8,27 @@
 import torch
 from torch_mlir import torchscript
 
-from torch_mlir_e2e_test.stablehlo_backends.abc import StablehloBackend
 from torch_mlir_e2e_test.framework import TestConfig, Trace, TraceItem
 from torch_mlir_e2e_test.utils import convert_annotations_to_placeholders
+
 from .utils import (
     recursively_convert_to_numpy,
     recursively_convert_from_numpy,
 )
 
 
-class StablehloBackendTestConfig(TestConfig):
-    """Base class for TestConfig's that are implemented with StableHLO.
-
-    This class handles all the common lowering that torch-mlir does before
-    reaching the StableHLO abstraction level.
-    """
+class JITImporterTestConfig(TestConfig):
+    """TestConfig that runs the torch.nn.Module with JIT Importer"""
 
-    def __init__(self, backend: StablehloBackend):
+    def __init__(self, backend, output_type="linalg-on-tensors"):
         super().__init__()
         self.backend = backend
+        self.output_type = output_type
 
     def compile(self, program: torch.nn.Module, verbose: bool = False) -> Any:
         example_args = convert_annotations_to_placeholders(program.forward)
         module = torchscript.compile(
-            program, example_args, output_type="stablehlo", verbose=verbose
+            program, example_args, output_type=self.output_type, verbose=verbose
         )
 
         return self.backend.compile(module)