enhance verbose option in e2e_testing (llvm#3390)

so that `python3 e2e_testing/main.py -v` would print intermediate IR.

projects/pt1/python/torch_mlir/torchscript.py

@@ -382,6 +382,12 @@ def compile(
         ) from None
     finally:
         sys.stderr = original_stderr
+
+    if verbose:
+        print("\n====================")
+        print("TorchScript RAW IR")
+        print(mb.module)
+
     if output_type == OutputType.RAW:
         return mb.module
 

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

@@ -36,7 +36,9 @@ def __init__(self, backend, output_type="linalg-on-tensors"):
         self._backend = backend
         self._output_type = output_type
 
-    def compile(self, program: torch.nn.Module) -> torch.nn.Module:
+    def compile(
+        self, program: torch.nn.Module, verbose: bool = False
+    ) -> torch.nn.Module:
         return program
 
     def run(self, artifact: torch.nn.Module, trace: Trace) -> Trace:

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

@@ -22,7 +22,9 @@ def __init__(self):
         super().__init__()
         lazy_backend._initialize()
 
-    def compile(self, program: torch.nn.Module) -> torch.nn.Module:
+    def compile(
+        self, program: torch.nn.Module, verbose: bool = False
+    ) -> torch.nn.Module:
         return program.to("lazy")
 
     def run(self, artifact: torch.nn.Module, trace: Trace) -> Trace:

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

@@ -29,10 +29,10 @@ def __init__(self, backend: LinalgOnTensorsBackend):
         super().__init__()
         self.backend = backend
 
-    def compile(self, program: torch.nn.Module) -> Any:
+    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="linalg-on-tensors"
+            program, example_args, output_type="linalg-on-tensors", verbose=verbose
         )
 
         return self.backend.compile(module)

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

@@ -14,7 +14,9 @@ class NativeTorchTestConfig(TestConfig):
     def __init__(self):
         super().__init__()
 
-    def compile(self, program: torch.nn.Module) -> torch.nn.Module:
+    def compile(
+        self, program: torch.nn.Module, verbose: bool = False
+    ) -> torch.nn.Module:
         return program
 
     def run(self, artifact: torch.nn.Module, trace: Trace) -> Trace:

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

@@ -89,13 +89,23 @@ def _module_lowering(
     output_type,
     torch_mod,
 ):
+    if verbose:
+        print("\n====================")
+        print("ONNX RAW IR")
+        print(torch_mod)
+
     # Lower from ONNX to Torch
     run_pipeline_with_repro_report(
         torch_mod,
         f"builtin.module(func.func({ONNX_TO_TORCH_FUNC_PIPELINE}))",
         "Lowering Onnx backend contract to Linalg-on-Tensors backend contract",
     )
 
+    if verbose:
+        print("\n====================")
+        print("TorchFX IR")
+        print(torch_mod)
+
     backend_legal_ops = [
         "aten.flatten.using_ints",
         "aten.adaptive_avg_pool1d",

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

@@ -28,9 +28,11 @@ def __init__(self, backend: StablehloBackend):
         super().__init__()
         self.backend = backend
 
-    def compile(self, program: torch.nn.Module) -> Any:
+    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")
+        module = torchscript.compile(
+            program, example_args, output_type="stablehlo", verbose=verbose
+        )
 
         return self.backend.compile(module)
 

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

@@ -170,7 +170,9 @@ def __init__(self, backend):
         super().__init__()
         self.backend = backend
 
-    def compile(self, program: torch.nn.Module) -> torch.nn.Module:
+    def compile(
+        self, program: torch.nn.Module, verbose: bool = False
+    ) -> torch.nn.Module:
         return program
 
     def run(self, artifact: torch.nn.Module, trace: Trace) -> Trace:

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

@@ -17,7 +17,9 @@ class TorchScriptTestConfig(TestConfig):
     def __init__(self):
         super().__init__()
 
-    def compile(self, program: torch.nn.Module) -> torch.jit.ScriptModule:
+    def compile(
+        self, program: torch.nn.Module, verbose: bool = False
+    ) -> torch.jit.ScriptModule:
         return torch.jit.script(program)
 
     def run(self, artifact: torch.jit.ScriptModule, trace: Trace) -> Trace:

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

@@ -29,10 +29,14 @@ def __init__(self, backend: TosaBackend, use_make_fx: bool = False):
         self.backend = backend
         self.use_make_fx = use_make_fx
 
-    def compile(self, program: torch.nn.Module) -> Any:
+    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="tosa", use_make_fx=self.use_make_fx
+            program,
+            example_args,
+            output_type="tosa",
+            use_make_fx=self.use_make_fx,
+            verbose=verbose,
         )
 
         return self.backend.compile(module)

projects/pt1/python/torch_mlir_e2e_test/framework.py

@@ -310,7 +310,7 @@ def compile_and_run_test(test: Test, config: TestConfig, verbose=False) -> Any:
         golden_trace = generate_golden_trace(test)
         if verbose:
             print(f"Compiling {test.unique_name}...", file=sys.stderr)
-        compiled = config.compile(test.program_factory())
+        compiled = config.compile(test.program_factory(), verbose=verbose)
     except Exception as e:
         return TestResult(
             unique_name=test.unique_name,