Skip to content

Commit

Permalink
[TE][BuildModule] Fix import in dump pass ir (apache#5327)
Browse files Browse the repository at this point in the history
  • Loading branch information
vinx13 authored and dpankratz committed Apr 24, 2020
1 parent b4a649e commit 6843828
Showing 1 changed file with 5 additions and 5 deletions.
10 changes: 5 additions & 5 deletions python/tvm/target/build_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,15 +59,15 @@ def dump(*args, **kwargs):

def decorate_irpass(self):
"""decorate ir_pass and ScheduleOps"""
self._old_sgpass = schedule.ScheduleOps
schedule.ScheduleOps = self.decorate(schedule.ScheduleOps)
vset = vars(ir_pass)
self._old_sgpass = tvm.te.schedule.ScheduleOps
tvm.te.schedule.ScheduleOps = self.decorate(tvm.te.schedule.ScheduleOps)
vset = vars(tvm.tir.ir_pass)
k = v = 0
def recover():
vset[k] = v
for k, v in vset.items():
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):
"""decorate given list of custom passes, and return decorated passes"""
Expand All @@ -93,7 +93,7 @@ def exit(self):
# recover decorated functions
for f in self._recover_list:
f()
schedule.ScheduleOps = self._old_sgpass
tvm.te.schedule.ScheduleOps = self._old_sgpass
DumpIR.scope_level -= 1


Expand Down

0 comments on commit 6843828

Please sign in to comment.