[Ansor][AutoTVM v2.0] Phase 1: Add follow_split and follow_fused_split steps #6142

jiuqi-yang · 2020-07-27T07:00:44Z

For the full upstream plan, see Ansor RFC.

In this PR, we bring follow split and follow fused split steps for Ansor auto_scheduler.

cc @merrymercy @comaniac @junrushao1994 @FrozenGene @jroesch

Signed-off-by: jingbang.yjb <[email protected]> Conflicts: src/auto_scheduler/compute_dag.cc src/auto_scheduler/transform_step.cc src/auto_scheduler/transform_step.h tests/python/unittest/test_auto_scheduler_loop_state.py

Signed-off-by: jingbang.yjb <[email protected]>

1. delete a comment 2. add "fuse" between follow_split and follow_fused_split Signed-off-by: jingbang.yjb <[email protected]>

Signed-off-by: jingbang.yjb <[email protected]>

Signed-off-by: jingbang.yjb <[email protected]> Conflicts: include/tvm/auto_scheduler/loop_state.h include/tvm/auto_scheduler/transform_step.h src/auto_scheduler/compute_dag.cc src/auto_scheduler/compute_dag.h src/auto_scheduler/loop_state.cc src/auto_scheduler/transform_step.cc tests/python/unittest/test_auto_scheduler_loop_state.py tests/python/unittest/test_auto_scheduler_measure.py

include/tvm/auto_scheduler/loop_state.h

include/tvm/auto_scheduler/transform_step.h

src/auto_scheduler/loop_state.cc

src/auto_scheduler/transform_step.cc

tests/python/unittest/test_auto_scheduler_loop_state.py

tests/python/unittest/test_auto_scheduler_measure.py

jcf94 · 2020-07-27T07:21:35Z

Hi, all. This is an student intern of us, who is now helping us with the Ansor upstreaming. 😄

The follow_split & follow_fused_split are two steps extent to te.Stage.Split. Each of these will collect information from the former history and process the split.

FollowSplit

This is mainly used in stage fusion using compute at.
For example we have stages: Dense -> Relu:
We've already done some tiling on Relu, and we would like to compute the Dense at the Relu stage. FollowSplit step is used to keep the outer most few iterators of Dense the same as the Relu stage.
Since in Ansor, the split factor of Relu stage may be left as a None placeholder to be filled by search policy, by this way we can easily write a schedule with some kind of dynamic dependence.

FollowFusedSplit

This is mainly used in GPU cooperative fetching.
For example we have stages: Input -> Dense:

            for [email protected] = ... : Bind to blockIdx.x
                for [email protected] = ... : Bind to threadIdx.x
                    for [email protected] = ...
                        Input_shared = Input ...
                        for k = ...
                            Dense = ...

In Ansor's search policy, the outer stage has been tiled. The the threadIdx.x axis is binded to a iterator generated by split & fuse step.
We use FollowFusedSplit step to compute out the final extent of the threadIdx.x binded iterator, to make sure that Input_shared stage can split out a iterator with same extent.

yangjunpro · 2020-07-27T07:26:51Z

Thanks @jiuqi-yang for the nice work, @merrymercy @tqchen @FrozenGene @comaniac , would you please take a look at this PR? We are trying to accelerate the auto-schedule upstreaming process.

Signed-off-by: jingbang.yjb <[email protected]>

yangjunpro

Just some nitpicking comments.

src/auto_scheduler/transform_step.cc

yangjunpro · 2020-07-27T15:45:41Z

src/auto_scheduler/transform_step.cc

@@ -136,6 +144,10 @@ void StepApplyToSchedule(const Step& step, Array<te::Stage>* stages, StageToAxes
    ps->ApplyToSchedule(stages, stage_to_axes);
  } else if (auto ps = step.as<SplitStepNode>()) {
    ps->ApplyToSchedule(stages, stage_to_axes);
+  } else if (auto ps = step.as<FollowSplitStepNode>()) {


yangjunpro · 2020-07-27T15:49:10Z

src/auto_scheduler/transform_step.cc

+void FollowSplitStepNode::WriteToRecord(dmlc::JSONWriter* writer) const {
+  writer->WriteArraySeperator();
+  writer->WriteString(record_prefix_str);
+  writer->WriteArrayItem(stage_id);


What will happen if the order of writing is changed?

The order here corresponds to the read order defined in the constructor of this step.

src/auto_scheduler/transform_step.cc

python/tvm/auto_scheduler/loop_state.py

include/tvm/auto_scheduler/transform_step.h

src/auto_scheduler/transform_step.cc

include/tvm/auto_scheduler/loop_state.h

python/tvm/auto_scheduler/loop_state.py

src/auto_scheduler/transform_step.cc

Signed-off-by: jingbang.yjb <[email protected]>

jcf94

Others look good to me. Just update these descriptions to follow the other functions.

include/tvm/auto_scheduler/transform_step.h

Signed-off-by: jingbang.yjb <[email protected]>

jcf94

Fine to me, I'll rebase the #6141 after this has been merged.

…t steps (apache#6142) * Add cache_read/cache_write step * Update * Add follow split and follow fused split Signed-off-by: jingbang.yjb <[email protected]> Conflicts: src/auto_scheduler/compute_dag.cc src/auto_scheduler/transform_step.cc src/auto_scheduler/transform_step.h tests/python/unittest/test_auto_scheduler_loop_state.py * add loop_state.py Signed-off-by: jingbang.yjb <[email protected]> * Update * Update * Update state->current_compute_dag to Optional * Add some doc strings for Follow_Split and Follow_fused_split Signed-off-by: jingbang.yjb <[email protected]> * Check code using c-lint Signed-off-by: jingbang.yjb <[email protected]> * Add more doc strings and change the order for follow split. Signed-off-by: jingbang.yjb <[email protected]> * Add record test for follow_split and follow_fused_split Signed-off-by: jingbang.yjb <[email protected]> * Add record test for follow_split Signed-off-by: jingbang.yjb <[email protected]> * Add record test for follow_fused_split. Signed-off-by: jingbang.yjb <[email protected]> * Add test record for follow_fused_split 1. delete a comment 2. add "fuse" between follow_split and follow_fused_split Signed-off-by: jingbang.yjb <[email protected]> * Add doc strings for some functions and variables Signed-off-by: jingbang.yjb <[email protected]> * Fix the code format in src/auto_scheduler/transform_step.h Signed-off-by: jingbang.yjb <[email protected]> * Update * Update doc * Update * Update * Fix follow_split and follow_fused_split record test. Signed-off-by: jingbang.yjb <[email protected]> * Doc update * Update some doc strings Signed-off-by: jingbang.yjb <[email protected]> * Fix code style and some function definitions. Signed-off-by: jingbang.yjb <[email protected]> * Update Signed-off-by: jingbang.yjb <[email protected]> * Add comments on parameters. Signed-off-by: jingbang.yjb <[email protected]> * Add more doc strings and fix some. Signed-off-by: jingbang.yjb <[email protected]> * Update Signed-off-by: jingbang.yjb <[email protected]> * Update Signed-off-by: jingbang.yjb <[email protected]> * Update Signed-off-by: jingbang.yjb <[email protected]> * Update. Signed-off-by: jingbang.yjb <[email protected]> Co-authored-by: chengfan.jcf <[email protected]> Co-authored-by: jingbang.yjb <[email protected]>

jcf94 and others added 26 commits July 22, 2020 09:55

Add cache_read/cache_write step

d12465d

Update

920f4b1

Add follow split and follow fused split

90e6391

Signed-off-by: jingbang.yjb <[email protected]> Conflicts: src/auto_scheduler/compute_dag.cc src/auto_scheduler/transform_step.cc src/auto_scheduler/transform_step.h tests/python/unittest/test_auto_scheduler_loop_state.py

add loop_state.py

e144082

Signed-off-by: jingbang.yjb <[email protected]>

Update

86c3670

Update

abfb150

Update state->current_compute_dag to Optional

3c1da64

Add some doc strings for Follow_Split and Follow_fused_split

c4a344c

Signed-off-by: jingbang.yjb <[email protected]>

Check code using c-lint

d3969b8

Signed-off-by: jingbang.yjb <[email protected]>

Add more doc strings and change the order for follow split.

f209525

Signed-off-by: jingbang.yjb <[email protected]>

Add record test for follow_split and follow_fused_split

50f7c4a

Signed-off-by: jingbang.yjb <[email protected]>

Add record test for follow_split

7bf8dd5

Signed-off-by: jingbang.yjb <[email protected]>

Add record test for follow_fused_split.

98d943b

Signed-off-by: jingbang.yjb <[email protected]>

Add test record for follow_fused_split

296cb36

1. delete a comment 2. add "fuse" between follow_split and follow_fused_split Signed-off-by: jingbang.yjb <[email protected]>

Add doc strings for some functions and variables

a7b1294

Signed-off-by: jingbang.yjb <[email protected]>

Fix the code format in src/auto_scheduler/transform_step.h

5220a68

Signed-off-by: jingbang.yjb <[email protected]>

Update

2a113d3

Update doc

bf660a8

Update

3649e26

Update

85da7e0

Fix follow_split and follow_fused_split record test.

1a87244

Signed-off-by: jingbang.yjb <[email protected]>

Merge branch 'follow_split' into merge_follow_split_cache

4b09317

Doc update

87e703a

Merge branch 'cache_read_write' into merge_follow_split_cache

3eeaff0

Update some doc strings

704c2cc

Signed-off-by: jingbang.yjb <[email protected]>

jcf94 requested changes Jul 27, 2020

View reviewed changes

Fix code style and some function definitions.

40f9638

Signed-off-by: jingbang.yjb <[email protected]>

jingbang.yjb added 2 commits July 27, 2020 15:56

Update

22b1f3b

Signed-off-by: jingbang.yjb <[email protected]>

Add comments on parameters.

d2e3da6

Signed-off-by: jingbang.yjb <[email protected]>

yangjunpro reviewed Jul 27, 2020

View reviewed changes

comaniac requested changes Jul 27, 2020

View reviewed changes

merrymercy requested changes Jul 27, 2020

View reviewed changes

jingbang.yjb added 3 commits July 28, 2020 10:13

Add more doc strings and fix some.

072c868

Signed-off-by: jingbang.yjb <[email protected]>

Update

be75123

Signed-off-by: jingbang.yjb <[email protected]>

Update

f4771f9

Signed-off-by: jingbang.yjb <[email protected]>

jcf94 mentioned this pull request Jul 28, 2020

[Ansor][AutoTVM v2.0] Phase 1: Add pragma/storage_align/rfactor steps #6141

Merged

jcf94 requested changes Jul 28, 2020

View reviewed changes

jingbang.yjb added 3 commits July 28, 2020 14:20

Fix some function parameters description.

ad94db3

Update

0780757

Signed-off-by: jingbang.yjb <[email protected]>

Update.

4d3a426

Signed-off-by: jingbang.yjb <[email protected]>

jcf94 approved these changes Jul 28, 2020

View reviewed changes

merrymercy approved these changes Jul 28, 2020

View reviewed changes

merrymercy merged commit bbc2dbf into apache:master Jul 28, 2020

ZihengJiang mentioned this pull request Sep 25, 2020

TVM v0.7 Release Note Candidate #6486

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Ansor][AutoTVM v2.0] Phase 1: Add follow_split and follow_fused_split steps #6142

[Ansor][AutoTVM v2.0] Phase 1: Add follow_split and follow_fused_split steps #6142

jiuqi-yang commented Jul 27, 2020

jcf94 commented Jul 27, 2020 •

edited

Loading

yangjunpro commented Jul 27, 2020

yangjunpro left a comment

yangjunpro Jul 27, 2020

yangjunpro Jul 27, 2020

comaniac Jul 27, 2020

jcf94 left a comment

jcf94 left a comment

[Ansor][AutoTVM v2.0] Phase 1: Add follow_split and follow_fused_split steps #6142

[Ansor][AutoTVM v2.0] Phase 1: Add follow_split and follow_fused_split steps #6142

Conversation

jiuqi-yang commented Jul 27, 2020

jcf94 commented Jul 27, 2020 • edited Loading

FollowSplit

FollowFusedSplit

yangjunpro commented Jul 27, 2020

yangjunpro left a comment

Choose a reason for hiding this comment

yangjunpro Jul 27, 2020

Choose a reason for hiding this comment

yangjunpro Jul 27, 2020

Choose a reason for hiding this comment

comaniac Jul 27, 2020

Choose a reason for hiding this comment

jcf94 left a comment

Choose a reason for hiding this comment

jcf94 left a comment

Choose a reason for hiding this comment

jcf94 commented Jul 27, 2020 •

edited

Loading