Hi community, recently I have been studying TVM’s auto-scheduler, both Ansor for TE and AutoTIR for TensorIR. I read the Ansor paper and basically understood its sketch generation algorithm. However, I had some confusion when trying to comprehend its AutoTIR counterpart, as specified in Section 3.4 of the RFC and the implementation of PostOrderVisit (line 96). For convenience, let me copy the implementation below:
Array<tir::Schedule> GenerateDesignSpace(const IRModule& mod_) final {
using ScheduleAndUnvisitedBlocks = std::pair<tir::Schedule, Array<tir::BlockRV>>;
tir::Schedule sch = tir::Schedule::Traced( //
/*mod=*/mod_, //
/*rand_state=*/ForkSeed(&this->rand_state_), //
/*debug_mode=*/tir::kVerifySRefTree | tir::kVerifyCachedFlags, //
/*error_render_level=*/tir::ScheduleErrorRenderLevel::kDetail);
std::vector<ScheduleAndUnvisitedBlocks> stack;
Array<tir::Schedule> result{sch};
// Enumerate the schedule rules first because you can
// always concat multiple schedule rules as one
Array<tir::BlockRV> all_blocks = BlockCollector::Collect(sch);
for (ScheduleRule sch_rule : sch_rules_) {
for (const tir::Schedule& sch : result) {
stack.emplace_back(sch, all_blocks);
}
result.clear();
while (!stack.empty()) {
// get the stack.top()
tir::Schedule sch;
Array<tir::BlockRV> blocks;
std::tie(sch, blocks) = stack.back();
stack.pop_back();
// if all blocks are visited
if (blocks.empty()) {
result.push_back(sch);
continue;
}
// otherwise, get the last block that is not visited
tir::BlockRV block_rv = blocks.back();
blocks.pop_back();
if (sch->HasBlock(block_rv)) {
Array<tir::Schedule> applied = sch_rule->Apply(sch, /*block=*/block_rv);
for (const tir::Schedule& sch : applied) {
stack.emplace_back(sch, blocks);
}
} else {
stack.emplace_back(sch, blocks);
}
}
}
return result;
}
My main confusion is that in this implementation, if all the sch_rules_ return only one schedule, the design space size would always be 1. However, in Ansor, for each iteration, multiple rules can be applied to one state to generate multiple succeeding states.