[AutoSchedule] Extract tasks via compile engine

[comaniac]

See RFC: https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444 for details.

Change highlight:

• auto_scheduler enabling.
  • Implicitly enable when calling auto_scheduler.extract_tasks.
  • Explicitly enable using PassContext(config={"relay.backend.use_auto_schedule": True}) when compilation.
  • Remove auto_scheduler.enable_relay_integration.
• Task extraction.
  • Add a flag include_simple_task=False. When setting it to True, the tasks with only injective ops will be extracted as well.
  • Extract tasks with Relay op strategy, but only rely on the TOPI computes and ignore the TOPI schedules, meaning that the selection of TOPI compute is purely based on plevel.
  • Skip the compute with unsupported ops (e.g., extern, hybrid ops).
• Schedule querying.
  • Change auto_scheduler.FallbackContext.silent to verbose. It can be 0 (silent), 1 (only warn the missing configs for workload with complex ops, default), and 2 (warn all missing configs).

This PR is ready for review. cc @merrymercy @tqchen @zhiics @jcf94 @icemelon9.

TODO

• Collect feedback from RFC.
• Avoid collecting tasks with only simple ops per users' choice.
• Deal with fallback schedules for non-collected tasks on GPU.
• Add unit tests.
• Specify the use of auto_scheduler in PassContext.
• Surpass the fallback warnings for simple tasks.

[mbaret]

At a high level, it seems to me this may benefit from the refactor I proposed in #6888. In particular, rather than switching the behaviour of ScheduleGetter with a flag it may make more sense to simply provide an alternative ScheduleGetter for the auto-scheduler. With a 'TETranslator' de-coupled, we should be able to avoid duplicating lots of code. What are your thoughts?

[comaniac]

At a high level, it seems to me this may benefit from the refactor I proposed in #6888. In particular, rather than switching the behaviour of ScheduleGetter with a flag it may make more sense to simply provide an alternative ScheduleGetter for the auto-scheduler. With a 'TETranslator' de-coupled, we should be able to avoid duplicating lots of code. What are your thoughts?

This PR can definitely work with #6888 to have a more concise implementation. However, this PR needs more than a TE translator, and that's why I exposed it first to have more discussions instead of waiting for #6888. Although we might be able to simplify the changes of compile_engine.cc in this PR with the TE translator, other changes in this PR (e.g., add a flag to graph runtime codegen and many other places) are still required. Thus, we can still work in parallel, and modify the later merged one accordingly to achieve the same goal.

[merrymercy]

The auto-scheduler part looks good to me.
Other people can comment more on whether adding use_auto_schedule to all compile engine functions is a good design.

[comaniac]

@merrymercy @zhiics thanks for the review and the comments are addressed. PTAL.

[zhiics]

please rebase

[comaniac]

Comment addressed. Now we could also use the flag in PassContext to add the TOPI compute for auto_scheduler with higher plevel (e.g., 15). Since this changes one task from direct to Winograd in the tutorial, I'm re-generating CI logs for the tutorial and will update it tomorrow.

[jcf94]

This feature is awesome! Thanks!

[zhiics]

Thanks everyone.