Unverified Commit 1df6bb6d by Wuwei Lin Committed by GitHub

[TE][BuildModule] Fix import in dump pass ir (#5327)

parent b7545eb5
...@@ -59,15 +59,15 @@ class DumpIR(object): ...@@ -59,15 +59,15 @@ class DumpIR(object):
def decorate_irpass(self): def decorate_irpass(self):
"""decorate ir_pass and ScheduleOps""" """decorate ir_pass and ScheduleOps"""
self._old_sgpass = schedule.ScheduleOps self._old_sgpass = tvm.te.schedule.ScheduleOps
schedule.ScheduleOps = self.decorate(schedule.ScheduleOps) tvm.te.schedule.ScheduleOps = self.decorate(tvm.te.schedule.ScheduleOps)
vset = vars(ir_pass) vset = vars(tvm.tir.ir_pass)
k = v = 0 k = v = 0
def recover(): def recover():
vset[k] = v vset[k] = v
for k, v in vset.items(): for k, v in vset.items():
self._recover_list.append(recover) self._recover_list.append(recover)
vset[k] = self.decorate(v) if isinstance(v, Function) else v vset[k] = self.decorate(v) if isinstance(v, tvm.runtime.PackedFunc) else v
def decorate_custompass(self, custom_pass): def decorate_custompass(self, custom_pass):
"""decorate given list of custom passes, and return decorated passes""" """decorate given list of custom passes, and return decorated passes"""
...@@ -93,7 +93,7 @@ class DumpIR(object): ...@@ -93,7 +93,7 @@ class DumpIR(object):
# recover decorated functions # recover decorated functions
for f in self._recover_list: for f in self._recover_list:
f() f()
schedule.ScheduleOps = self._old_sgpass tvm.te.schedule.ScheduleOps = self._old_sgpass
DumpIR.scope_level -= 1 DumpIR.scope_level -= 1
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment