Skip to content

Commit

Permalink
[util] add tvm.support.dump_for_debug function
Browse files Browse the repository at this point in the history
- 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.
  • Loading branch information
Christian Convey committed Jul 26, 2022
1 parent 9ca6be8 commit c824720
Showing 1 changed file with 18 additions and 15 deletions.
33 changes: 18 additions & 15 deletions python/tvm/support.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)
Expand All @@ -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__)

0 comments on commit c824720

Please sign in to comment.