[Ansor][AutoTVM v2.0] Phase 2: Basic CPU Sketch Search Policy #6184

jcf94 · 2020-08-01T10:56:08Z

For full upstream plan, see Ansor RFC.

In this PR, we bring the basic Sketch search policy which is proposed in our paper on CPU ops/subgraphs.
The complete search policy still waits for the feature extraction & cost model support, so this will be a basic search policy using random cost model.

We'll continue to work on the upstreaming process to bring the complete auto schedule support for various workloads on CPU/GPU with some reproducible performance results, a custom sketch rule support to work like the AutoTVM, as well as the tutorials of how to play with AutoScheduler.

Review guide

Main changes:

Update the API of auto_scheduler::AutoSchedule, TuningOptions and SearchPolicy::Search;
Some API name refine for compute_dag.h/cc, cost_model.h/cc;
Add search_policy/sketch_search_policy.h/cc, search_policy/utils.h/cc, these're the key parts of this PR;

In the SketchSearchPolicy, we can start the review from SketchSearchPolicyNode::Search and SketchSearchPolicyNode::SearchOneRound, which are the main entrance.

SketchGenerationRule and InitPopulationRule are used to generate the schedule automatically.

Some methods used by the SketchSearchPolicy are put on the search_policy/utils.h/cc.

jcf94 · 2020-08-01T11:11:03Z

@merrymercy I realized that a base class of CostModel is also needed in this PR. We can sync that part first when you start to work on the cost model PR(Or have a separate PR for the basic random model support?).

Another problem is that I found the setting of random seed does not work on the current implementation of random number generation in RandomModel(neither the python np.random, nor the C++ rand_r).

merrymercy · 2020-08-02T10:54:25Z

@jcf94 I sent a separate PR for the base class for cost models (#6187).
Why do we need to rely on the random seed here?

jcf94 · 2020-08-03T01:27:34Z

@jcf94 I sent a separate PR for the base class for cost models (#6187).
Why do we need to rely on the random seed here?

Thanks!
The random seed issue does not block anything. I addressed it here is just to make sure we're able to fully reproduce a search process. The current UT generates different states in each run.

jcf94 · 2020-08-05T12:46:55Z

~~This PR shares the base class with cost models, will rebase the code after #6187 has been merged.~~
~~The other parts of code is ready for review.~~
Rebased with #6187, ready for review.
cc @merrymercy @comaniac @FrozenGene @junrushao1994 @tqchen

include/tvm/auto_scheduler/auto_schedule.h

src/auto_scheduler/compute_dag.cc

include/tvm/auto_scheduler/search_policy.h

src/auto_scheduler/compute_dag.cc

src/auto_scheduler/search_policy/utils.cc

merrymercy

LGTM. Only some minor doc improvements.

include/tvm/auto_scheduler/compute_dag.h

include/tvm/auto_scheduler/search_policy.h

python/tvm/auto_scheduler/auto_schedule.py

src/auto_scheduler/compute_dag.cc

src/auto_scheduler/search_policy/sketch_search_policy.cc

comaniac

Overall LGTM, but I'd like to raise the discussion about the file organization for search policy. Now sketch_search_policy.cc has about one thousand line and it might continue to grow in the future. Here is the organization in my mind:

auto_scheduler
|- serach_policy
  |- sketch_search
    |- sketch_search_policy.{h,cc}
    |- sketch_rules.{h,cc}
    |- utils.{h,cc}
  |- empty_search
  |- utils.{h,cc}

Have auto_scheduler/search_policy/{sketch_policy, empty_policy}.
Separate all SketchGenerationRule and InitPopulationRule to search_policy/sketch_policy/sketch_rules.
Rename src/auto_scheduler/search_policy/utils.{h,cc} to 'src/auto_scheduler/search_policy/utils.{h,cc}' (still under search_policy), and move all sketch search specific functions such as Mutation (not included in this PR) to auto_scheduler/search_policy/sketch_policy/utils.{h,cc}.

python/tvm/auto_scheduler/auto_schedule.py

comaniac · 2020-08-10T21:02:43Z

python/tvm/auto_scheduler/auto_schedule.py

-    sch, tensors = _ffi_api.AutoSchedule(task, search_policy,
-                                         tuning_options if tuning_options else TuningOptions())
+    # TODO(jcf94): Remove EmptyPolicy after finish the porting of SketchSearchPolicy
+    sch, tensors = _ffi_api.AutoSchedule(search_policy or EmptyPolicy(task), tuning_options)


Should we make sketch search policy as the default one? Otherwise I can imagine lots of people will ask in the discuss forum about why auto scheduler doesn't do any search...

tests/python/unittest/test_auto_scheduler_sketch_generation.py

jcf94 · 2020-08-11T06:49:57Z

Overall LGTM, but I'd like to raise the discussion about the file organization for search policy. Now sketch_search_policy.cc has about one thousand line and it might continue to grow in the future. Here is the organization in my mind:
auto_scheduler
|- serach_policy
  |- sketch_search
    |- sketch_search_policy.{h,cc}
    |- sketch_rules.{h,cc}
    |- utils.{h,cc}
  |- empty_search
  |- utils.{h,cc}
Have auto_scheduler/search_policy/{sketch_policy, empty_policy}.

Separate all SketchGenerationRule and InitPopulationRule to search_policy/sketch_policy/sketch_rules.

Rename src/auto_scheduler/search_policy/utils.{h,cc} to 'src/auto_scheduler/search_policy/utils.{h,cc}' (still under search_policy), and move all sketch search specific functions such as Mutation (not included in this PR) to auto_scheduler/search_policy/sketch_policy/utils.{h,cc}.

Currently I just split out the sketch_policy_rules.h/cc, we can continue to consider about the directory structure.
We can also put the Evolutionary Search to a separate file.

merrymercy · 2020-08-11T20:45:06Z

@comaniac @junrushao1994 Wait for your approval and I will merge this

junrushao

LGTM

comaniac

LGTM. The file organization is much clearer now. Thanks :)

…#6184) * Init commit to pass the compile * First commit to pass the test * Update * Add UTs for sketch generation * Update * Add ASF to new UT file. * Update rule for winograd * Update * File renamed * Lint fix

jcf94 force-pushed the sketch_search_policy branch 3 times, most recently from d1a5777 to 510339c Compare August 5, 2020 12:37

jcf94 marked this pull request as ready for review August 5, 2020 12:42

jcf94 changed the title ~~[Ansor][AutoTVM v2.0] Phase 2: Basic CPU Sketch Search Policy [WIP]~~ [Ansor][AutoTVM v2.0] Phase 2: Basic CPU Sketch Search Policy Aug 6, 2020

Init commit to pass the compile

00acaba

junrushao reviewed Aug 6, 2020

View reviewed changes

include/tvm/auto_scheduler/auto_schedule.h Show resolved Hide resolved

jcf94 added 2 commits August 7, 2020 09:21

First commit to pass the test

37a3ded

Update

e625554

jcf94 force-pushed the sketch_search_policy branch from a232b40 to e625554 Compare August 7, 2020 02:41

jcf94 added 3 commits August 7, 2020 15:22

Add UTs for sketch generation

74e2a37

Update

395b963

Add ASF to new UT file.

76aa618

merrymercy requested review from merrymercy, FrozenGene, tqchen and kevinthesun August 8, 2020 10:09

Update rule for winograd

dd54d22

jcf94 commented Aug 10, 2020

View reviewed changes

src/auto_scheduler/compute_dag.cc Show resolved Hide resolved

FrozenGene reviewed Aug 10, 2020

View reviewed changes

include/tvm/auto_scheduler/search_policy.h Outdated Show resolved Hide resolved

src/auto_scheduler/compute_dag.cc Show resolved Hide resolved

src/auto_scheduler/search_policy/utils.cc Outdated Show resolved Hide resolved

merrymercy approved these changes Aug 10, 2020

View reviewed changes

comaniac reviewed Aug 10, 2020

View reviewed changes

jcf94 added 3 commits August 11, 2020 10:40

Update

4c510cb

File renamed

0f96fd0

Lint fix

5facafd

jcf94 requested review from FrozenGene and merrymercy August 11, 2020 10:38

junrushao approved these changes Aug 11, 2020

View reviewed changes

comaniac approved these changes Aug 11, 2020

View reviewed changes

merrymercy merged commit 75b8318 into apache:master Aug 11, 2020

jcf94 deleted the sketch_search_policy branch August 12, 2020 01:53

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 2: Basic CPU Sketch Search Policy #6184

[Ansor][AutoTVM v2.0] Phase 2: Basic CPU Sketch Search Policy #6184

jcf94 commented Aug 1, 2020 •

edited

Loading

jcf94 commented Aug 1, 2020 •

edited

Loading

merrymercy commented Aug 2, 2020

jcf94 commented Aug 3, 2020

jcf94 commented Aug 5, 2020 •

edited

Loading

merrymercy left a comment •

edited

Loading

comaniac left a comment

comaniac Aug 10, 2020

jcf94 commented Aug 11, 2020

merrymercy commented Aug 11, 2020

junrushao left a comment

comaniac left a comment

[Ansor][AutoTVM v2.0] Phase 2: Basic CPU Sketch Search Policy #6184

[Ansor][AutoTVM v2.0] Phase 2: Basic CPU Sketch Search Policy #6184

Conversation

jcf94 commented Aug 1, 2020 • edited Loading

Review guide

jcf94 commented Aug 1, 2020 • edited Loading

merrymercy commented Aug 2, 2020

jcf94 commented Aug 3, 2020

jcf94 commented Aug 5, 2020 • edited Loading

merrymercy left a comment • edited Loading

Choose a reason for hiding this comment

comaniac left a comment

Choose a reason for hiding this comment

comaniac Aug 10, 2020

Choose a reason for hiding this comment

jcf94 commented Aug 11, 2020

merrymercy commented Aug 11, 2020

junrushao left a comment

Choose a reason for hiding this comment

comaniac left a comment

Choose a reason for hiding this comment

jcf94 commented Aug 1, 2020 •

edited

Loading

jcf94 commented Aug 1, 2020 •

edited

Loading

jcf94 commented Aug 5, 2020 •

edited

Loading

merrymercy left a comment •

edited

Loading