Skip to content

Commit e49d500

Browse files
committed
return reversed list
1 parent 74636be commit e49d500

File tree

2 files changed

+6
-1
lines changed

2 files changed

+6
-1
lines changed

python/tvm/meta_schedule/integration.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -247,4 +247,5 @@ def extract_task_from_relay(
247247
disabled_pass=disabled_pass,
248248
):
249249
tasks = extract_task_func(mod, target, relay_params)
250-
return tasks
250+
# Tasks are extracted via post order visit, return the reversed list.
251+
return list(reversed(tasks))

src/meta_schedule/integration.cc

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,10 @@ IRModule ApplyHistoryBestNode::Query(runtime::String task_name, IRModule mod, Ta
120120
ICHECK(HasOnlyOneFunction<relay::Function>(mod)) << mod;
121121
IRModule prim_mod = dispatched.value()[0];
122122
ICHECK(HasOnlyOneFunction<tir::PrimFunc>(prim_mod)) << prim_mod;
123+
// TODO(masahi): parse_mod below replaces the orginal function key with "main".
124+
// This is necessary because some scheduling primitives requires the PrimFunc key be "main".
125+
// If we can remove this restriction, there would no need for GetOnlyOneFunction* calls below
126+
// and we can directly return sch->mod().
123127
auto gv = GetOnlyOneFunctionKey<tir::PrimFunc>(prim_mod).value();
124128
// Unify func name to make sure it can be found in database
125129
const auto* parse_mod_func = runtime::Registry::Get("tvm.meta_schedule.tune.parse_mod");

0 commit comments

Comments
 (0)