[AutoScheduler] Relay integration : Task extraction - #6710
Conversation
|
cc @icemelon9 @tqchen @minminsun @jcf94 @comaniac @FrozenGene |
7758f24 to
5d93c61
Compare
comaniac
left a comment
There was a problem hiding this comment.
Theoretically we need to add add_auto_schedular to all ops but we only have 2 right now. Is there any concern or plan?
| auto_schedule_impl_suffix = ".auto_scheduler" | ||
|
|
||
|
|
||
| def auto_schedule_topi(outs): |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
comaniac
left a comment
There was a problem hiding this comment.
LGTM. Please fix the CI.
|
@comaniac I only added conv2d_nhwc and dense in this PR to test the functionality. |
df161b9 to
ba6d2aa
Compare
| try: | ||
| import psutil | ||
| except ImportError: | ||
| raise ImportError("psutil not found, try `pip install psutil` to fix this") |
There was a problem hiding this comment.
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") |
There was a problem hiding this comment.
Ah...you move it here. But i am curious why we should need it here. For testing?
There was a problem hiding this comment.
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 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 |
|
Fixed by #6807 |
* 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
* 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
* 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
This pr implements the basic relay integration for auto-scheduler.
Approach
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
dispather.py