[Apache TVM Discuss] [Questions] [TensorIR] Question on the AutoScheduler-style design space generation

[Yuxiang Wei via Apache TVM Discuss <no...@discuss.tvm.ai>]

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).