[util] add tvm.support.dump_for_debug function

- add a utility function to help with debugging / investigation

  The function is particularly geared to help programmers examine
  changes to TVM modules and functions as they progress through
  the scheduling and lowering.

python/tvm/support.py

@@ -89,7 +89,7 @@ def __setitem__(self, key, value):
         self.add_function(key, value)
 
 
-def dump_for_debug(obj, obj_name:Optional[str]='obj', io_tensors: Optional = None) -> str:
+def dump_for_debug(obj, obj_name: Optional[str] = "obj", io_tensors: Optional = None) -> str:
     """
     Return a multi-line, human-readable dump of various details about the TVM
     object 'obj'.  Intended for debugging and/or exposition.
@@ -103,29 +103,31 @@ def dump_for_debug(obj, obj_name:Optional[str]='obj', io_tensors: Optional = Non
     """
     output = io.StringIO()
 
-    def handle_primfunc(name_prefix:str, func: tvm.tir.function.PrimFunc) -> None:
-        output.write(f'type({name_prefix}): {type(func)}\n')
-        output.write(f'{name_prefix}:\n{func}\n')
+    def handle_primfunc(name_prefix: str, func: tvm.tir.function.PrimFunc) -> None:
+        output.write(f"type({name_prefix}): {type(func)}\n")
+        output.write(f"{name_prefix}:\n{func}\n")
 
         script = str(func.script())
-        output.write(f'{name_prefix}.script():\n{script}\n')
+        output.write(f"{name_prefix}.script():\n{script}\n")
 
-    def handle_irmodule(name_prefix:str, mod: tvm.ir.module.IRModule, already_lowered:bool) -> None:
-        output.write(f'type({name_prefix}): {type(mod)}\n')
-        output.write(f'{name_prefix}:\n{mod}\n')
+    def handle_irmodule(
+        name_prefix: str, mod: tvm.ir.module.IRModule, already_lowered: bool
+    ) -> None:
+        output.write(f"type({name_prefix}): {type(mod)}\n")
+        output.write(f"{name_prefix}:\n{mod}\n")
 
         if (not already_lowered) and (io_tensors is not None):
             lmod = tvm.lower(mod, io_tensors)
-            lname_prefix = f'tvm.lower({name_prefix}, io_tensors)'
+            lname_prefix = f"tvm.lower({name_prefix}, io_tensors)"
             handle_irmodule(lname_prefix, lmod, True)
 
-    def handle_schedule(name_prefix:str, sched: tvm.tir.schedule.schedule.Schedule) -> None:
-        output.write(f'type({name_prefix}): {type(sched)}\n')
-        output.write(f'{name_prefix}:\n{sched}\n')
+    def handle_schedule(name_prefix: str, sched: tvm.tir.schedule.schedule.Schedule) -> None:
+        output.write(f"type({name_prefix}): {type(sched)}\n")
+        output.write(f"{name_prefix}:\n{sched}\n")
 
         # Note: We're assuming that this module hasn't yet been lowered.
         mod = sched.mod
-        handle_irmodule(f'{name_prefix}.mod', mod, False)
+        handle_irmodule(f"{name_prefix}.mod", mod, False)
 
     if isinstance(obj, tvm.tir.function.PrimFunc):
         handle_primfunc(obj_name, obj)
@@ -136,9 +138,10 @@ def handle_schedule(name_prefix:str, sched: tvm.tir.schedule.schedule.Schedule)
         handle_irmodule(obj_name, obj, False)
     else:
         # catch-all logic...
-        output.write(f'type({obj_name}): {type(obj)}\n')
-        output.write(f'{obj_name}:\n{obj}\n')
+        output.write(f"type({obj_name}): {type(obj)}\n")
+        output.write(f"{obj_name}:\n{obj}\n")
 
     return output.getvalue()
 
+
 tvm._ffi._init_api("support", __name__)