[TECompiler] Decouple TE compute and schedule lowering in ScheduleBuilder #10561
Conversation
masahi
force-pushed
the
schedule-builder-refactor
branch
from
f21f8f4
to
669a1f7
Compare
masahi
requested review from
jroesch,
slyubomirsky,
icemelon,
MarisaKirisame,
ZihengJiang,
yzhliu,
vinx13,
mbrookhart,
jwfromm,
zhiics,
anijain2305,
wweic and
junrushao
as code owners
masahi
changed the title
| 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_); |
There was a problem hiding this comment.
cc @mbaret for this change
masahi
force-pushed
the
schedule-builder-refactor
branch
from
669a1f7
to
a167655
Compare
| 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
|
Overall looks great Masa thanks for doing this! |
masahi
force-pushed
the
schedule-builder-refactor
branch
from
a167655
to
8e4f69e
Compare
masahi
force-pushed
the
schedule-builder-refactor
branch
from
8e4f69e
to
4cd3a16
Compare
| 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