[TVMC] Add tvmc flag to print ir before / after pass names

Author: Eirene Pandi

python/tvm/driver/tvmc/compiler.py

@@ -31,7 +31,7 @@
 from tvm import autotvm, auto_scheduler
 from tvm import relay
 from tvm.driver.tvmc.registry import generate_registry_args, reconstruct_registry_entity
-from tvm.ir.instrument import PassInstrument, PassTimingInstrument
+from tvm.ir.instrument import PassInstrument, PassTimingInstrument, PassPrintBefore, PassPrintAfter
 from tvm.ir.memory_pools import WorkspaceMemoryPools
 from tvm.target import Target
 from tvm.relay.backend import Executor, Runtime
@@ -162,6 +162,18 @@ def add_compile_parser(subparsers, _, json_params):
         action="store_true",
         help="print compilation time per pass",
     )
+    parser.add_argument(
+        "--print-ir-before",
+        help="print IR before each named pass of a comma-separated list of pass names."
+        "e.g. '--print-ir-before [tir.SplitHostDevice,tir.ConvertSSA]' ",
+        default="",
+    )
+    parser.add_argument(
+        "--print-ir-after",
+        help="print IR after each named pass of a comma-separated list of pass names."
+        "e.g. '--print-ir-after [tir.SplitHostDevice,tir.ConvertSSA]' ",
+        default="",
+    )
     for one_entry in json_params:
         parser.set_defaults(**one_entry)
 
@@ -220,6 +232,8 @@ def drive_compile(args):
             workspace_pools_recombobulate(args, [workspace_pools_target], extra_targets)
         ),
         print_pass_times=args.print_pass_times,
+        print_ir_before=args.print_ir_before,
+        print_ir_after=args.print_ir_after,
         **transform_args,
     )
 
@@ -247,6 +261,8 @@ def compile_model(
     mod_name: Optional[str] = "default",
     workspace_pools: Optional[WorkspaceMemoryPools] = None,
     print_pass_times: bool = False,
+    print_ir_before: List[str] = "",
+    print_ir_after: List[str] = "",
     instruments: Optional[Sequence[PassInstrument]] = None,
     desired_layout: Optional[str] = None,
     desired_layout_ops: Optional[List[str]] = None,
@@ -295,7 +311,7 @@ def compile_model(
         needs to be generated.
     disabled_pass: str, optional
         Comma-separated list of passes which needs to be disabled
-        during compilation
+        during compilation.
     pass_context_configs: list[str], optional
         List of strings containing a set of configurations to be passed to the
         PassContext.
@@ -310,6 +326,10 @@ def compile_model(
         compilation.
     print_pass_times: bool
         To enable printing a breakdown of compilation times by pass. Disabled by default.
+    print_ir_before: list[str]
+        To print ir before each named pass of a comma-separated list of passes.
+    print_ir_after: list[str]
+        To print ir after each named pass of a comma-separated list of passes.
     instruments: Optional[Sequence[PassInstrument]]
         The list of pass instrument implementations.
     desired_layout: str, optional
@@ -369,6 +389,14 @@ def compile_model(
         timing_inst = PassTimingInstrument()
         instruments = [timing_inst] if instruments is None else [timing_inst] + instruments
 
+    if print_ir_before:
+        print_ir_before_instr = PassPrintBefore(print_ir_before)
+        instruments = [print_ir_before_instr] if instruments is None else [print_ir_before_instr] + instruments
+    
+    if print_ir_after:
+        print_ir_after_instr = PassPrintAfter(print_ir_after)
+        instruments = [print_ir_after_instr] if instruments is None else [print_ir_after_instr] + instruments
+
     with tvm.transform.PassContext(
         opt_level=opt_level,
         config=config,

python/tvm/ir/instrument.py

@@ -216,6 +216,8 @@ def should_run(self, mod, pass_info)
             tvm.relay.build(mod, "llvm")
     """
 
+        
+            
     def create_pass_instrument(pi_cls):
         if not inspect.isclass(pi_cls):
             raise TypeError("pi_cls must be a class")
@@ -255,3 +257,21 @@ def render():
                 profiles = timing_inst.render()
         """
         return _ffi_instrument_api.RenderTimePassProfiles()
+    
+@pass_instrument
+class PassPrintBefore:
+            def __init__(self, print_pass_name):
+                self.print_pass_name = print_pass_name
+            def run_before_pass(self, mod, pass_info):
+                if (pass_info.name in self.print_pass_name):
+                    print("Print ir before:")
+                    print(str(pass_info.name) + '\n' + str(mod) + '\n\n')
+
+@pass_instrument
+class PassPrintAfter:
+            def __init__(self, print_pass_name):
+                self.print_pass_name = print_pass_name
+            def run_after_pass(self, mod, pass_info):
+                if (pass_info.name in self.print_pass_name ):
+                    print("Print ir after:")
+                    print(str(pass_info.name) + '\n' + str(mod) + '\n\n')

tests/python/driver/tvmc/test_command_line.py

@@ -289,3 +289,35 @@ def test_tvmc_print_pass_times(capsys, keras_simple, tmpdir_factory):
     captured_out = capsys.readouterr().out
     for exp_str in ("Compilation time breakdown by pass:", "sequential:", "us]"):
         assert exp_str in captured_out
+
+def test_tvmc_print_ir_before(capsys, keras_simple, tmpdir_factory):
+    pytest.importorskip("tensorflow")
+    tmpdir = tmpdir_factory.mktemp("out")
+    print_cmd = "--print-ir-before=[tir.SplitHostDevice]"
+
+    # Compile model
+    module_file = os.path.join(tmpdir, "keras-tvm.tar")
+    compile_cmd = f"tvmc compile --target 'llvm' {keras_simple} --output {module_file} {print_cmd}"
+    compile_args = compile_cmd.split(" ")[1:]
+    _main(compile_args)
+
+    # Check for timing results output
+    captured_out = capsys.readouterr().out
+    for exp_str in ("Print ir before:\n","tir.SplitHostDevice\n"):
+        assert exp_str in captured_out
+
+def test_tvmc_print_ir_after(capsys, keras_simple, tmpdir_factory):
+    pytest.importorskip("tensorflow")
+    tmpdir = tmpdir_factory.mktemp("out")
+    print_cmd = "--print-ir-after=[tir.SplitHostDevice]"
+
+    # Compile model
+    module_file = os.path.join(tmpdir, "keras-tvm.tar")
+    compile_cmd = f"tvmc compile --target 'llvm' {keras_simple} --output {module_file} {print_cmd}"
+    compile_args = compile_cmd.split(" ")[1:]
+    _main(compile_args)
+
+    # Check for timing results output
+    captured_out = capsys.readouterr().out
+    for exp_str in ("Print ir after:\n","tir.SplitHostDevice\n"):
+        assert exp_str in captured_out