Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[MetaSchedule][M4b] Task Extraction #9382

Merged
merged 3 commits into from
Nov 5, 2021
Merged

[MetaSchedule][M4b] Task Extraction #9382

merged 3 commits into from
Nov 5, 2021

Conversation

junrushao
Copy link
Member

@junrushao junrushao commented Oct 28, 2021

This PR enables meta schedule to extract tasks from the high-level framework, in our particular case, Relay.

A base class IntegrationContext is introduced as the context manager for both task extraction (observer mode) and tuned task injection (injector mode):

  • Derived class TaskExtraction serves as the observer mode to extract tasks from Relay
  • Derived class ApplyHistoryBest serves as the injector mode to inject the tuned task back to Relay

The primary API for task extraction is meta_schedule.integration.extract_task, which returns a list of tasks on both Relay and TIR-level.

The PR only handles task extraction, but doesn't include further work like:

  • Multi-dispatching from Relay to TE/TIR
  • Deduplication and similarity score calculation of extracted tasks
  • Relay TIR operator strategy
  • Layout rewriting

CC: @mbrookhart @mbs-octoml @sunggg @zxybazh @Hzfengsy @YuchenJin @altanh

python/tvm/meta_schedule/integration.py Outdated Show resolved Hide resolved
python/tvm/meta_schedule/integration.py Show resolved Hide resolved
python/tvm/meta_schedule/integration.py Outdated Show resolved Hide resolved
python/tvm/meta_schedule/integration.py Show resolved Hide resolved
python/tvm/meta_schedule/integration.py Outdated Show resolved Hide resolved
python/tvm/meta_schedule/integration.py Outdated Show resolved Hide resolved
python/tvm/te/operation.py Show resolved Hide resolved
src/meta_schedule/integration.cc Show resolved Hide resolved
python/tvm/meta_schedule/integration.py Show resolved Hide resolved
tests/python/unittest/test_meta_schedule_integration.py Outdated Show resolved Hide resolved
Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com>
Co-authored-by: Hongyi Jin <3231950289@qq.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Wuwei Lin <wuwei@apache.org>
Copy link
Contributor

@comaniac comaniac left a comment

Choose a reason for hiding this comment

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

LGTM. Thanks

include/tvm/meta_schedule/integration.h Outdated Show resolved Hide resolved
Copy link
Member

@zxybazh zxybazh left a comment

Choose a reason for hiding this comment

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

LGTM. Minor nit.

include/tvm/meta_schedule/integration.h Outdated Show resolved Hide resolved
@junrushao junrushao merged commit 85fd85f into apache:main Nov 5, 2021
mehrdadh pushed a commit to mehrdadh/tvm that referenced this pull request Dec 1, 2021
Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com>
Co-authored-by: Hongyi Jin <3231950289@qq.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Wuwei Lin <wuwei@apache.org>
mehrdadh pushed a commit to mehrdadh/tvm that referenced this pull request Dec 1, 2021
Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com>
Co-authored-by: Hongyi Jin <3231950289@qq.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Wuwei Lin <wuwei@apache.org>
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 7, 2022
Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com>
Co-authored-by: Hongyi Jin <3231950289@qq.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Wuwei Lin <wuwei@apache.org>
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022
Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com>
Co-authored-by: Hongyi Jin <3231950289@qq.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Wuwei Lin <wuwei@apache.org>
@junrushao junrushao changed the title [MetaSchedule] Task Extraction [MetaSchedule][M4b] Task Extraction Jan 26, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants