Uh oh!
There was an error while loading. Please reload this page.
[TECompiler] Decouple TE compute and schedule lowering in ScheduleBuilder - #10561
Conversation
f21f8f4 to
669a1f7Compare| auto outputs = lower_te_compute.Lower(prim_func, [&](std::string name) { return name; }); | ||
| return CachedFunc(tgt, GlobalVar(lower_te_compute.candidate_name_), lower_te_compute.fn_inputs_, | ||
| outputs, te::Schedule(), tir::PrimFunc(), {}, | ||
| IRModule(Map<GlobalVar, BaseFunc>({})), lower_te_compute.constant_tensors_); |
669a1f7 to
a167655Compare| int LowerToTECompute::const_index = 0; | ||
| // Construct a schedule for a given Relay primitive function and target. | ||
| class ScheduleBuilder : ExprVisitor { |
There was a problem hiding this comment.
Maybe we could give this class a better name while we are doing this cleanup?
There was a problem hiding this comment.
It's a bit hard to come up with a good name since this class does a lot more than "schedule building". See https://github.com/apache/tvm/pull/10561/files#r824096015
Over time we stuffed all lowering-dependent tasks onto this class...
| use_auto_scheduler_ = backend::IsAutoSchedulerEnabled(); | ||
| } | ||
| CachedFunc Create(const Function& relay_func, std::function<std::string(std::string)> renamer) { |
There was a problem hiding this comment.
While you are refactoring this code it might be good to better annotate it with comments describing the pieces here as now this code (due to me and Mark refactoring) has quite a few branches.
There was a problem hiding this comment.
Yeah, this function is not only used for ordinary lowering purposes, but also abused for ansor / meta schedule task extraction or "apply history best" depending on where it is called...
Optional<ObjectRef> opt_mod_or_base_func = meta_schedule::MetaScheduleContext::QueryInsideWithScope is very opaque since (1) we don't know if it returns anything and (2) we don't what concrete value it returns.
After my task extraction refactor, QueryInsideWithScope can always return a non-null, concrete Schedule object, since the None case is used only for task extraction currently. cc @junrushao1994
There was a problem hiding this comment.
Per discussion with Masa, we decided to just clean up meta_schedule/integration.cc and make it clearer
jroesch
commented
Mar 10, 2022
Overall looks great Masa thanks for doing this! |
a167655 to
8e4f69eCompare8e4f69e to
4cd3a16CompareUh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
| use_auto_scheduler_ = backend::IsAutoSchedulerEnabled(); | ||
| } | ||
| CachedFunc Create(const Function& relay_func, std::function<std::string(std::string)> renamer) { |
There was a problem hiding this comment.
Per discussion with Masa, we decided to just clean up meta_schedule/integration.cc and make it clearer
…lder (apache#10561) * Decouple TE compute and schedule lowering in ScheduleBuilder * fixed merge conflict * removed create_schedule stuff * add public, fix include path convention * Forgot visiting arg in ScheduleBuilder CallNode vsit * fixed anchor impl selection
Refactors
ScheduleBuilder, that currently does TE compute lowering and schedule application in one go, into two classes.The motivation is to use the TE compute lowering class,
LowerToTECompute, for meta schedule task extraction. For task extraction, there is no need to apply and lower schedules, but currently we doVMCompiler.lower()which is too hacky and overkill.@mbs-octoml@junrushao1994@jroesch@comaniac