[rfc] Use logging.getLogger for projects/pt1/e2e_testing

projects/pt1/e2e_testing/main.py

@@ -4,6 +4,7 @@
 # Also available under a BSD-style license. See LICENSE.
 
 import argparse
+import logging
 import re
 import sys
 
@@ -75,6 +76,10 @@ def _get_argparse():
                         default=False,
                         action="store_true",
                         help="report test results with additional detail")
+    parser.add_argument("--print-ir",
+                        default=False,
+                        action="store_true",
+                        help="Set logging level to DEBUG and causes the IR to be printed for each test.")
     parser.add_argument("-s", "--sequential",
                         default=False,
                         action="store_true",
@@ -93,6 +98,18 @@ def _get_argparse():
 def main():
     args = _get_argparse().parse_args()
 
+    ir_printer = logging.getLogger("ir_printer")
+    if args.print_ir:
+        print("WARNING: --print-ir is a work in progress feature.")
+        print("print-ir: Setting logging level to DEBUG and enabling IR printing.")
+        print("print-ir: This currently only affects the Linalg-on-Tensors and onnx configs.")
+        print("print-ir: Work in progress. See https://github.com/llvm/torch-mlir/issues/3172")
+        ir_printer.setLevel(logging.DEBUG)
+    else:
+        # disable logging
+        ir_printer.setLevel(logging.CRITICAL+1)
+
+
     all_test_unique_names = set(
         test.unique_name for test in GLOBAL_TEST_REGISTRY)
 

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

@@ -2,6 +2,8 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 # Also available under a BSD-style license. See LICENSE.
+import logging
+ir_printer = logging.getLogger("ir_printer")
 
 from typing import Any
 
@@ -32,7 +34,9 @@ def compile(self, program: torch.nn.Module) -> Any:
         example_args = convert_annotations_to_placeholders(program.forward)
         module = torchscript.compile(
             program, example_args, output_type="linalg-on-tensors")
-
+        ir_printer.debug("LinalgOnTensorsBackendTestConfig compiled module:")
+        ir_printer.debug(module)
+        ir_printer.debug("End LinalgOnTensorsBackendTestConfig compiled module")
         return self.backend.compile(module)
 
 

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

@@ -22,7 +22,8 @@
 from torch_mlir.extras import onnx_importer
 from torch_mlir.dialects import torch as torch_d
 from torch_mlir.ir import Context, Module
-
+import logging
+ir_printer = logging.getLogger("ir_printer")
 
 def import_onnx(contents):
     # Import the ONNX model proto from the file contents:
@@ -39,7 +40,7 @@ def import_onnx(contents):
     return m
 
 
-def convert_onnx(model, inputs):
+def convert_onnx(model: torch.nn.Module, inputs):
     buffer = io.BytesIO()
 
     # Process the type information so we export with the dynamic shape information
@@ -82,6 +83,9 @@ def __init__(self, backend: OnnxBackend, use_make_fx: bool = False):
     def compile(self, program: torch.nn.Module) -> Any:
         example_args = convert_annotations_to_placeholders(program.forward)
         onnx_module = convert_onnx(program, example_args)
+        ir_printer.debug("OnnxBackendTestConfig imported module:")
+        ir_printer.debug(onnx_module)
+        ir_printer.debug("End OnnxBackendTestConfig imported module")
         compiled_module = self.backend.compile(onnx_module)
         return compiled_module
 

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

@@ -2,6 +2,8 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 # Also available under a BSD-style license. See LICENSE.
+import logging
+ir_printer = logging.getLogger("ir_printer")
 
 from typing import Any
 
@@ -31,7 +33,9 @@ def __init__(self, backend: StablehloBackend):
     def compile(self, program: torch.nn.Module) -> Any:
         example_args = convert_annotations_to_placeholders(program.forward)
         module = torchscript.compile(program, example_args, output_type="stablehlo")
-
+        ir_printer.debug("StablehloBackendTestConfig compiled module:")
+        ir_printer.debug(module)
+        ir_printer.debug("End StablehloBackendTestConfig compiled module")
         return self.backend.compile(module)
 
     def run(self, artifact: Any, trace: Trace) -> Trace:

projects/pt1/python/torch_mlir_e2e_test/framework.py

@@ -143,6 +143,7 @@ def compile(self, program: torch.nn.Module) -> CompiledArtifact:
         """Compile the provided torch.nn.Module into a compiled artifact"""
         pass
 
+
     # Any should match result of `compile`.
 
     @abc.abstractmethod