Hi community, recently I have been studying TVM's auto-scheduler, both
[Ansor](https://www.usenix.org/system/files/osdi20-zheng.pdf) for TE and
[AutoTIR](https://github.com/apache/tvm-rfcs/blob/main/rfcs/0005-meta-schedule-autotensorir.md)
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](https://github.com/apache/tvm-rfcs/blob/main/rfcs/0005-meta-schedule-autotensorir.md#34-autoscheduler-style-design-space-generation)
and the implementation of `PostOrderVisit` ([line
96](https://github.com/apache/tvm/blob/9ee0007c20c1a3701ddd0685b3d914fd2d2c76d6/src/meta_schedule/space_generator/post_order_apply.cc#L96)).
For convenience, let me copy the implementation below:
```c++
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.
---
[Visit
Topic](https://discuss.tvm.apache.org/t/tensorir-question-on-the-autoscheduler-style-design-space-generation/11794/1)
to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.apache.org/email/unsubscribe/dfe58f6c640e9f8d57e104edefde66629ef6a9377e0e60fdc358cc9f915aa7a0).
