[Ansor][AutoTVM v2.0] Phase 1: Add pragma/storage_align/rfactor steps #6141
Conversation
jcf94
force-pushed
the
pragma_rfactor_storage_align
branch
from
e4408da
to
082df7f
Compare
| // add FuseStepNode required by rfactor | ||
| if (i >= 2 && transform_steps[i - 2]->IsInstance<FuseStepNode>()) { | ||
| const Step& fuse_step = transform_steps[i - 2]; | ||
| if (fuse_step->stage_id == step->stage_id) { | ||
| ret_steps.push_back(fuse_step); | ||
| } | ||
| } | ||
| // add SplitStepNode required by rfactor | ||
| CHECK_GE(i, 1); | ||
| CHECK(transform_steps[i - 1]->IsInstance<SplitStepNode>()); | ||
| const Step& split_step = transform_steps[i - 1]; | ||
| CHECK_EQ(split_step->stage_id, step->stage_id); | ||
| ret_steps.push_back(split_step); | ||
| // add RfactorStepNode | ||
| ret_steps.push_back(step); |
There was a problem hiding this comment.
These checks here are specialized for the Rfactor in our search policy, will the rules continue to be kept like this? @merrymercy
There was a problem hiding this comment.
yes. Otherwise, there will be errors when creating the new dag
jcf94
force-pushed
the
pragma_rfactor_storage_align
branch
from
dab6ae8
to
62a51f2
Compare
jcf94
force-pushed
the
pragma_rfactor_storage_align
branch
from
49acf58
to
f05ec96
Compare
|
Thanks to all! It's really great to work with you. Next, @merrymercy will continue to add the cost model while at the same time I'll start to bring our Sketch Search Policy(proposed in our paper) to TVM, and tutorials about how to tune a op/subgraph will together to be added. Hope the PR of Sketch Search Policy comes soon. The API design and our custom sketch rule support (which works the same as the current AutoTVM template) are ongoing to be refined. We would be happy if anyone has suggestions on this project, and you're welcomed to join us to improve Ansor. cc @minminsun @comaniac @FrozenGene @yangjunpro @tqchen @junrushao1994 @jroesch |
…apache#6141) * Add pragma/storage_align/rfactor step * Update * Update * Update UT * Update
…apache#6141) * Add pragma/storage_align/rfactor step * Update * Update * Update UT * Update
…apache#6141) * Add pragma/storage_align/rfactor step * Update * Update * Update UT * Update
…apache#6141) * Add pragma/storage_align/rfactor step * Update * Update * Update UT * Update
…apache#6141) * Add pragma/storage_align/rfactor step * Update * Update * Update UT * Update
For the full upstream plan, see Ansor RFC.
In this PR, we bring pragma/storage_align/rfactor steps for Ansor auto_scheduler.
The Pragma step and StorageAlign step are two simple steps that add attributes to schedule.
The Rfactor step adds extra stage like CacheRead/CacheWrite do.
cc @merrymercy @comaniac @junrushao1994 @FrozenGene @jroesch