[Ansor][AutoTVM v2.0] Phase 1: Add annotation/compute_at/compute_root/compute_inline steps #6073

jcf94 · 2020-07-16T07:06:56Z

For full upstream plan, see Ansor RFC.

In this PR, we bring annotation/compute_at/compute_root/compute_inline steps for Ansor auto_scheduler.

To record the stage/iterator attach relations, we introduced a new structure AttachMap in State.

cc @tqchen @merrymercy @FrozenGene @minminsun @comaniac

junrushao · 2020-07-16T19:39:20Z

Just noticed that #6077 is also titled as "part 1", shall we make one of them part 2 instead?

python/tvm/auto_scheduler/loop_state.py

comaniac · 2020-07-16T23:40:55Z

src/auto_scheduler/measure_record.cc

@@ -169,6 +206,18 @@ struct Handler<::tvm::Array<::tvm::auto_scheduler::Step>> {
          fused_ids.push_back(i);
        }
        data->push_back(::tvm::auto_scheduler::FuseStep(stage_id, fused_ids));
+      } else if (name == "AN") {


I feel the current approach of (de)serializing records is not scalable and hard to be maintained. Specifically, we put all deserialization rules to Read and all serialization rules to Write. The key to connect the serialization logic to the corresponding deserialization logic is a two character short name (e.g., AN). It seems to me that it would be better to define a short name and (de)serialization logic in each step. In this case, we can use ps->serialize(writer) for serialization, and build a step factory to deserialize them.

Good point.
@merrymercy Currently to add a new transform steps, we have to modify 7 files ... Maybe it would better to move all specific implementation to transform_step.h/cc, and left the other file to just add some interface.

I agree with your points.

I'm working on this.

src/auto_scheduler/loop_state.h

src/auto_scheduler/loop_state.cc

tests/python/unittest/test_auto_scheduler_loop_state.py

tests/python/unittest/test_auto_scheduler_measure.py

comaniac · 2020-07-17T00:37:45Z

Just noticed that #6077 is also titled as "part 1", shall we make one of them part 2 instead?

Part 1 will have several PRs. "Phase 1" might be a better word to describe this process, but anyways.

jcf94 · 2020-07-17T01:32:05Z

Just noticed that #6077 is also titled as "part 1", shall we make one of them part 2 instead?

Part 1 will have several PRs. "Phase 1" might be a better word to describe this process, but anyways.

@junrushao1994 My bad, yes I was intended to express a meaning of stage/phase...
@comaniac Let's just use phase, I'll rename these PRs.

python/tvm/auto_scheduler/loop_state.py

tests/python/unittest/test_auto_scheduler_loop_state.py

src/auto_scheduler/transform_step.cc

junrushao

Thanks @jcf94! Overall it looks good. Just some nitpicks.

src/auto_scheduler/utils.h

src/auto_scheduler/loop_state.cc

src/auto_scheduler/loop_state.h

python/tvm/auto_scheduler/loop_state.py

src/auto_scheduler/compute_dag.cc

merrymercy · 2020-07-17T16:37:05Z

src/auto_scheduler/measure_record.cc

@@ -169,6 +206,18 @@ struct Handler<::tvm::Array<::tvm::auto_scheduler::Step>> {
          fused_ids.push_back(i);
        }
        data->push_back(::tvm::auto_scheduler::FuseStep(stage_id, fused_ids));
+      } else if (name == "AN") {


I agree with your points.

tests/python/unittest/test_auto_scheduler_loop_state.py

tests/python/unittest/test_auto_scheduler_measure.py

comaniac

Overall LGTM. What's the plan of fixing the step serialization issue? Will that be done in this PR?

src/auto_scheduler/loop_state.cc

tests/python/unittest/test_auto_scheduler_loop_state.py

junrushao · 2020-07-18T18:14:19Z

A minor typo we didn't catch before (github doesn't allow me to comment on that line because it is not related to this change): in file "loop_state.h", document for class StageNode, "Similar to te::Stage in `include/schedule.h`", the path is incorrect.

jcf94 · 2020-07-20T03:53:01Z

Moved all the implementation to transform_step.h/cc, and updated the order of steps.
In the current implementation, the next pr of steps adding will only need to change transform_step.h/cc two files.
cc @merrymercy @junrushao1994 @comaniac @FrozenGene
... I'm not sure why I failed to re-request review of @junrushao1994 and @comaniac

comaniac · 2020-07-20T16:44:35Z

src/auto_scheduler/loop_state.cc

-    } else {
-      LOG(FATAL) << "Invalid step: " << step;
-    }
+    StepApplyToState(step, this, dag);


Looks like we can remove StepApplyToState and simply use step->ApplyToState(this)?

dag seems unused in both this function and StepApplyToState.

We can remove StepApplyToState

dag will be used in other steps such as cache_write and rfactor.

The step->ApplyToState() is actually not a virtual function(the return value of different steps may be different), so I add the StepApplyToState to pack all those step type check to transform_steps.cc.

comaniac · 2020-07-20T16:52:16Z

src/auto_scheduler/transform_step.cc

-    if (i != after_ids.size() - 1) {
-      ss << ", ";
-    }
+void StepApplyToSchedule(const Step& step, Array<te::Stage>* stages,


For 3 functions in this file (StepApplyToState, StepApplyToSchedule, and StepApplyToPythonAPI), it should be fine and clearer to directly call step->XXX(...);incompute_dag.cc` so that we can avoid unnecessary step node dispatching.

Emm, step->ApplyToState, step->ApplyToSchedule & step->PrintAsPythonAPI all have problems to be written to virtual functions, we can later try to see if there's any better way to deal with this.

comaniac · 2020-07-20T17:02:42Z

src/auto_scheduler/transform_step.cc

-  node->stage_id = stage_id;
-  for (const auto& x : after_ids) {
-    CHECK(x->IsInstance<IntImmNode>());
+Step StepReadFromRecord(dmlc::JSONReader* reader) {


It would be better to have a step registration mechanism to establish a step lookup table from record prefix to the step node class, so that we can totally hide step node dispatching from developers. The potential issues in the current implementation is that 1) people may forget to update this function when adding a new step, and 2) there's no mechanism to check duplicated record prefix of two steps.

In addition, I feel this issue is important but I'm not sure if it's proper to address it in this PR. We may need versioning for records.

Yes, we can add a registration mechanism. But I think it can be done in future PRs.

The current records do have a version field
https://github.com/apache/incubator-tvm/blob/7ac78d96611ab3a069fae9eb5210e433eb240cc5/src/auto_scheduler/measure_record.cc#L314

merrymercy · 2020-07-21T00:53:03Z

src/auto_scheduler/loop_state.cc

-    } else {
-      LOG(FATAL) << "Invalid step: " << step;
-    }
+    StepApplyToState(step, this, dag);


We can remove StepApplyToState

dag will be used in other steps such as cache_write and rfactor.

src/auto_scheduler/measure_record.cc

merrymercy

Overall the re-organization looks good to me. We can merge it after fixing the several remaining comments from Cody and me.

src/auto_scheduler/compute_dag.cc

FrozenGene

Overall LGTM

src/auto_scheduler/compute_dag.cc

junrushao

LGTM. Thanks for the contribution!

comaniac

Given that the issue of 3 dispatching functions will be addressed later, this PR LGTM. Thanks for the efforts!.

jcf94 · 2020-07-22T01:53:01Z

Given that the issue of 3 dispatching functions will be addressed later, this PR LGTM. Thanks for the efforts!.

Thanks, updated our merge plan to track this issue: https://github.com/merrymercy/Ansor/issues/65

…/compute_inline steps (apache#6073) * Add annotation step * Add compute_at/compute_root/compute_inline * Doc update * Update * Update * Update measure record UT * Update * Update * Update * Move state implementation to step * Move measure_record implementation to step * Order update & API update * Update the order of state api * Update

jcf94 added 6 commits July 15, 2020 16:21

Add annotation step

603d3d2

Add compute_at/compute_root/compute_inline

d81553c

Doc update

78db4c9

Update

308f141

Update

35eb81d

Update measure record UT

a4589a5

comaniac requested changes Jul 17, 2020

View reviewed changes

jcf94 changed the title ~~[Ansor][AutoTVM v2.0] Part 1: Add annotation/compute_at/compute_root/compute_inline steps~~ [Ansor][AutoTVM v2.0] Phase 1: Add annotation/compute_at/compute_root/compute_inline steps Jul 17, 2020

FrozenGene reviewed Jul 17, 2020

View reviewed changes

python/tvm/auto_scheduler/loop_state.py Outdated Show resolved Hide resolved

junrushao reviewed Jul 17, 2020

View reviewed changes

src/auto_scheduler/utils.h Outdated Show resolved Hide resolved

src/auto_scheduler/loop_state.cc Outdated Show resolved Hide resolved

src/auto_scheduler/loop_state.h Outdated Show resolved Hide resolved

jcf94 added 2 commits July 17, 2020 17:05

Update

8b05133

Update

8fb87ae

tqchen added the status: need review label Jul 17, 2020

merrymercy requested changes Jul 17, 2020

View reviewed changes

comaniac reviewed Jul 17, 2020

View reviewed changes

src/auto_scheduler/loop_state.cc Outdated Show resolved Hide resolved

tests/python/unittest/test_auto_scheduler_loop_state.py Outdated Show resolved Hide resolved

Update

cc0ad1d

jcf94 added 3 commits July 19, 2020 22:36

Move state implementation to step

13809c2

Move measure_record implementation to step

30828a8

Order update & API update

4bb58ce

jcf94 requested review from merrymercy and FrozenGene July 20, 2020 03:49

Update the order of state api

0cb3dbd

comaniac reviewed Jul 20, 2020

View reviewed changes

merrymercy requested changes Jul 21, 2020

View reviewed changes

merrymercy reviewed Jul 21, 2020

View reviewed changes

src/auto_scheduler/compute_dag.cc Outdated Show resolved Hide resolved

jcf94 added 2 commits July 21, 2020 10:03

Update

d217fd4

Merge branch 'upstream_master' into compute_at_annotation

5e983e1

FrozenGene approved these changes Jul 21, 2020

View reviewed changes

src/auto_scheduler/compute_dag.cc Show resolved Hide resolved

src/auto_scheduler/compute_dag.cc Show resolved Hide resolved

jcf94 requested a review from merrymercy July 21, 2020 08:09

merrymercy approved these changes Jul 21, 2020

View reviewed changes

merrymercy mentioned this pull request Jul 21, 2020

[Ansor][AutoTVM v2.0] Phase 1: Access Analyzer #6103

Merged

junrushao approved these changes Jul 21, 2020

View reviewed changes

comaniac approved these changes Jul 21, 2020

View reviewed changes

merrymercy merged commit 71533a5 into apache:master Jul 21, 2020

jcf94 deleted the compute_at_annotation branch July 29, 2020 07:29

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 annotation/compute_at/compute_root/compute_inline steps #6073

[Ansor][AutoTVM v2.0] Phase 1: Add annotation/compute_at/compute_root/compute_inline steps #6073

jcf94 commented Jul 16, 2020 •

edited

Loading

junrushao commented Jul 16, 2020

comaniac Jul 16, 2020

jcf94 Jul 17, 2020

merrymercy Jul 17, 2020

jcf94 Jul 18, 2020

comaniac commented Jul 17, 2020 •

edited

Loading

jcf94 commented Jul 17, 2020

junrushao left a comment

merrymercy Jul 17, 2020

comaniac left a comment

junrushao commented Jul 18, 2020

jcf94 commented Jul 20, 2020 •

edited

Loading

comaniac Jul 20, 2020

merrymercy Jul 21, 2020

jcf94 Jul 21, 2020

comaniac Jul 20, 2020

jcf94 Jul 21, 2020

comaniac Jul 20, 2020

merrymercy Jul 21, 2020 •

edited

Loading

merrymercy Jul 21, 2020

merrymercy left a comment •

edited

Loading

FrozenGene left a comment

junrushao left a comment

comaniac left a comment

jcf94 commented Jul 22, 2020

[Ansor][AutoTVM v2.0] Phase 1: Add annotation/compute_at/compute_root/compute_inline steps #6073

[Ansor][AutoTVM v2.0] Phase 1: Add annotation/compute_at/compute_root/compute_inline steps #6073

Conversation

jcf94 commented Jul 16, 2020 • edited Loading

junrushao commented Jul 16, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

comaniac commented Jul 17, 2020 • edited Loading

jcf94 commented Jul 17, 2020

junrushao left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

comaniac left a comment

Choose a reason for hiding this comment

junrushao commented Jul 18, 2020

jcf94 commented Jul 20, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

merrymercy Jul 21, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

merrymercy left a comment • edited Loading

Choose a reason for hiding this comment

FrozenGene left a comment

Choose a reason for hiding this comment

junrushao left a comment

Choose a reason for hiding this comment

comaniac left a comment

Choose a reason for hiding this comment

jcf94 commented Jul 22, 2020

jcf94 commented Jul 16, 2020 •

edited

Loading

comaniac commented Jul 17, 2020 •

edited

Loading

jcf94 commented Jul 20, 2020 •

edited

Loading

merrymercy Jul 21, 2020 •

edited

Loading

merrymercy left a comment •

edited

Loading