[Metaschedule] New relay backend for meta schedule task extraction

[masahi]

Building on #10561, add a new interface for task extraction that doesn't require the whole VMCompiler.lower(). Extracted tasks are the same as ones from the existing flow. Currently only enabled for meta schedule, but I believe it is worth adapting for autotvm / ansor cc @comaniac @tkonolige

I'm currently working on a model where TE scheduling leads to an error during lowering, but TIR scheduling might work. I wanted to try the latter, but currently task extraction entails (for no reason) TE schedule application and lowering, which leads to the said error. This motivated this work, which aims to remove the unnecessary TE scheduling steps from the task extraction process.

Please review @junrushao1994 @mbs-octoml @jroesch

[tkonolige]

I like not having to go through VMCompiler to get tasks. And I'd like to see this adapted to autoscheduler/autotvm. Another improvement would be to remove all metascheduler/auotscheduler depenednt code from te_compiler. te_compiler could return a list of tasks and the metascheduler/autoscheduler could handle it as they please. Maybe Collage will do this.

[masahi]

Another improvement would be to remove all metascheduler/auotscheduler depenednt code from te_compiler. te_compiler could return a list of tasks and the metascheduler/autoscheduler could handle it as they please

This is basically my goal of this PR, and already achieved for meta scheduler. There is no longer meta scheduler-specific code path for task extraction in te_compiler_cache.cc. Most of the code in task_extraction.cc, in particular

Array<ExtractedTask> ExtractTask(IRModule mod, Target target, Map<String, Constant> params) {
  backend::BindParamsInModule(mod, params);

  // is_vm=true for backward compatibility
  Array<Pass> pass_seqs = relay::backend::GetPassPrefix(/*is_homogenous=*/true, /*is_vm=*/true);
  pass_seqs.push_back(transform::FuseOps());

  transform::Sequential seq(pass_seqs);
  auto opt_mod = seq(std::move(mod));

  Array<ExtractedTask> tasks;
  std::unordered_set<tec::CCacheKey> cache_;
  std::unordered_map<std::string, int> name_map;

  PostOrderVisit(opt_mod->Lookup("main"), [target, &tasks, &cache_, &name_map](const Expr& exp) {
    if (exp->IsInstance<FunctionNode>()) {
      Function relay_func = Downcast<Function>(exp);
      tec::CCacheKey cache_key(relay_func, target);
      if (relay_func->HasNonzeroAttr(attr::kPrimitive) && cache_.find(cache_key) == cache_.end()) {
        Array<te::Tensor> outputs;
        std::string fused_name;
        std::tie(outputs, fused_name) =
                  tec::LowerTECompute(relay_func, target, /*return_inputs*/ true);
       ...

are not specific to any tuners and thus can be shared. Different tuners can customize the behavior as they please, e.g. how to create "task" objects from lowered TE compute.

[junrushao]

Will review next week!

[junrushao]

Thanks for refactoring this interface! On the other hand, the interface itself might be unnecessary any more after the refactoring :-) Perhaps we would love to refactor this completely some time after this PR is merged!

[masahi]

yeah, based on the current use, I want to change the interface to be tir::Schedule -> tir::Schedule or PrimFunc -> PrimFunc, and rename Query to ApplySchedule or something like that.

When we do need IRModule -> IRModule interface later, I think it's better to introduce another function. Rather than one interface that does everything.

[junrushao]

I like the idea of renaming to something more meaningful :-)

On the interface side, I would prefer IRModule -> IRModule transformation instead of limiting it to something less general. Imagine the case where we do subgraph-level BYOC, which we already supported in Builder/Runner, and adding the database support might be less tricky if we use the most general IRModule -> IRModule transformation

[masahi]

As I commented above, for other use cases that may or may not happen in the future, I think it's better to introduce another interface dedicated for it. BYOC, as currently implemented today, doesn't go through TE Compiler.

[junrushao]

Overall looks good to me! Only a few nitpicks :-)