Skip to content

[AutoScheduler] Relay integration : Task extraction - #6710

Merged
merrymercy merged 12 commits into
apache:mainfrom
merrymercy:master
Oct 30, 2020
Merged

[AutoScheduler] Relay integration : Task extraction#6710
merrymercy merged 12 commits into
apache:mainfrom
merrymercy:master

Conversation

@merrymercy

@merrymercy merrymercy commented Oct 19, 2020

Copy link
Copy Markdown
Member

This pr implements the basic relay integration for auto-scheduler.

Approach

  1. register auto-scheduler as an implementation in the OpStrategy. We have a universal schedule function for all topi compute functions.
  2. Use tracing to extract all tasks similar to autotvm.

If auto-scheduler is enabled by auto_scheduler.enable_relay_integration, then auto-scheduler is always preferred to any other registered implementations.

Limitation

Ideally, we should be able to mix autotvm, auto-scheduler, and static libraries, then we can select the best implementation according to profiling records.
However, this is not trivial to implement, because auto-scheduler works on a subgraph level while the current OpStrategy is designed for single operator level.
The limitation of OpStrategy also makes auto-scheduler unable to support multiple compute declarations. For example, auto-scheduler cannot select direct vs. winograd based on profiling records, while autotvm can easily do this.

Minor todo

  • Add an API to disable the cache in relay compilation, so we can extract the weights (i.e., the number of appearances) of each task for the task-scheduler.
  • Clean up dispather.py

@merrymercy

Copy link
Copy Markdown
Member Author

cc @icemelon9 @tqchen @minminsun @jcf94 @comaniac @FrozenGene

@merrymercy
merrymercy force-pushed the master branch 3 times, most recently from 7758f24 to 5d93c61 Compare October 19, 2020 09:44

@comaniac comaniac left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Theoretically we need to add add_auto_schedular to all ops but we only have 2 right now. Is there any concern or plan?

Comment thread python/tvm/auto_scheduler/relay_integration.py
Comment thread python/tvm/auto_scheduler/dispatcher.py Outdated
Comment thread python/tvm/auto_scheduler/dispatcher.py Outdated
Comment thread python/tvm/auto_scheduler/dispatcher.py Outdated
Comment thread python/tvm/auto_scheduler/dispatcher.py
Comment thread python/tvm/auto_scheduler/dispatcher.py
Comment thread python/tvm/auto_scheduler/dispatcher.py Outdated
Comment thread tests/python/relay/test_auto_scheduler_task_extraction.py
auto_schedule_impl_suffix = ".auto_scheduler"


def auto_schedule_topi(outs):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it better to call it auto_schedule_te (and replace all topi in this file with te)? I imagine there are cases where inputs to this function don't come from topi, such as TE from gradient ops generated by Relay.

@merrymercy merrymercy Oct 26, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We cannot call it auto_schedule_te. We have another function auto_schedule (https://github.com/apache/incubator-tvm/blob/e59c603515befb02035e237794aa0645dbfbaf09/python/tvm/auto_scheduler/auto_schedule.py#L161) for your use case.
But this function is designed to be used as a TOPI schedule function for Relay, because it does a lot of other things.

Comment thread python/tvm/auto_scheduler/dispatcher.py Outdated
Comment thread tests/python/relay/test_auto_scheduler_task_extraction.py Outdated

@comaniac comaniac left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Please fix the CI.

@merrymercy

merrymercy commented Oct 27, 2020

Copy link
Copy Markdown
Member Author

@comaniac I only added conv2d_nhwc and dense in this PR to test the functionality.
I will send another PR with more ops and a tutorial on tuning networks for CUDA.

@merrymercy
merrymercy force-pushed the master branch 2 times, most recently from df161b9 to ba6d2aa Compare October 27, 2020 02:44
try:
import psutil
except ImportError:
raise ImportError("psutil not found, try `pip install psutil` to fix this")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I prefer keeping this ImportError. During internal development, we have met several times our environment doesn't import this package, we could tune but meet process error and find this root cause is hard.

def kill_child_processes(parent_pid, sig=signal.SIGTERM):
"""kill all child processes recursively"""
if not psutil:
raise ImportError("psutil not found, try `pip install psutil` to fix this")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah...you move it here. But i am curious why we should need it here. For testing?

@merrymercy merrymercy Oct 30, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, it is for testing. The goal is to minimize the dependency to run tvm.
If users do not use auto-scheduler, then they do not need to install this package.

@merrymercy
merrymercy merged commit a261454 into apache:main Oct 30, 2020
@tqchen

tqchen commented Oct 31, 2020

Copy link
Copy Markdown
Member

@merrymercy the code seems to cause test problem on MacOS, please keep tornado (rpc tracker) as an optional dep https://github.com/apache/incubator-tvm/actions/runs/338532340

@merrymercy

Copy link
Copy Markdown
Member Author

Fixed by #6807

trevor-m pushed a commit to trevor-m/tvm that referenced this pull request Dec 2, 2020
* add task extraction

* fix evo search

* fix tests

* fix test

* fix docstring

* fix docstring

* update workload registry

* fix warning

* fix test

* fix fallback

* fix lint

* fix tests
trevor-m pushed a commit to trevor-m/tvm that referenced this pull request Dec 4, 2020
* add task extraction

* fix evo search

* fix tests

* fix test

* fix docstring

* fix docstring

* update workload registry

* fix warning

* fix test

* fix fallback

* fix lint

* fix tests
trevor-m pushed a commit to neo-ai/tvm that referenced this pull request Dec 4, 2020
* add task extraction

* fix evo search

* fix tests

* fix test

* fix docstring

* fix docstring

* update workload registry

* fix warning

* fix test

* fix fallback

* fix lint

* fix tests
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants