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.

Sign up for GitHub

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

Jump to bottom

[TE] Support schedulable TIR compute definitions in TOPI #11589

Merged
merged 7 commits into from
Jun 13, 2022
Merged

[TE] Support schedulable TIR compute definitions in TOPI #11589

merged 7 commits into from
Jun 13, 2022

Conversation

csullivan
Copy link
Contributor

This PR adds te.extern_primfunc which provides the interface around TE ExternOp that allows a TVMScript defined schedulable TIR PrimFunc to be inlined into a TE compute graph. The result is that TIR can be used for compute definitions in Relay OpStrategies and, paired with meta-scheduler support in relay as introduced in #10578, these compute definitions can be scheduled and tuned as demonstrated in the attached tests.

Prior to this, compute definitions were limited to those definable in TE only. As a consequence of this patch and ongoing improvements to TVMScript meta-programming (#11097), TOPI can be extended to include compute and scheduling functions targeting schedulable TIR uniformly.

@csullivan
[TIR] Update domain touched to track touched domains for all buffers
eb3feba 
for loads, stores and both in one pass. Expose an FFI to return the
buffer domain access map.
@csullivan
[TE] Refactor logical steps of CreatePrimFunc into separate functions.
fa888bf
@csullivan
[TE] Update CreatePrimFunc to pre-populate buffer binds map with
f75d82d 
user defined buffers expressed in te.extern ops.
@csullivan
[TE] Introduce te.extern_primfunc which allows for inlining a
5efedea 
TVMScript defined schedulable TIR primfunc into TE. In this way
TOPI compute definitions can be defined via TIR and scheduled
using TIR scheduling primitives via meta-scheduler.
@csullivan
Add test showing the use of te.extern on a TVMScript PrimFunc
659dce6 
Add test that uses TVMScript defined TOPI compute definition for relay meta-schedule execution.
@csullivan csullivan requested a review from masahi June 6, 2022 17:10
@masahi
Copy link
Member

masahi commented Jun 6, 2022

cc @Hzfengsy

Hzfengsy
Hzfengsy approved these changes Jun 7, 2022
View reviewed changes
Copy link
Member

@Hzfengsy Hzfengsy left a comment

Choose a reason for hiding this comment

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

Overall LGTM. Thanks @csullivan for such great contribution!!!

src/arith/domain_touched.cc Show resolved Hide resolved
src/te/operation/create_primfunc.cc Outdated Show resolved Hide resolved
tests/python/unittest/test_meta_schedule_relay_tir_compute.py Show resolved Hide resolved
csullivan
csullivan commented Jun 8, 2022
View reviewed changes
Copy link
Contributor Author

@csullivan csullivan left a comment •
edited
Loading

Choose a reason for hiding this comment

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

Many thanks @Hzfengsy! PTAL again if/when you have the cycles.

src/arith/domain_touched.cc Show resolved Hide resolved
Hzfengsy
Hzfengsy requested changes Jun 9, 2022
View reviewed changes
tests/python/unittest/test_tir_te_extern_primfunc.py
Comment on lines +142 to +145
for i in T.serial(
0,
16,
):
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
for i in T.serial(
0,
16,
):
for i in range(16):

tests/python/unittest/test_tir_te_extern_primfunc.py
Comment on lines +100 to +103
for i in T.serial(
0,
16,
):
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
for i in T.serial(
0,
16,
):
for i in range(16):

tests/python/unittest/test_tir_te_extern_primfunc.py
Comment on lines +66 to +69
for i in T.serial(
0,
16,
):
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
for i in T.serial(
0,
16,
):
for i in range(16):

tests/python/unittest/test_tir_te_extern_primfunc.py
Comment on lines +37 to +40
for i in T.serial(
0,
16,
):
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
for i in T.serial(
0,
16,
):
for i in range(16):

@junrushao
Copy link
Member

@Hzfengsy for context, as a rule of thumb, usually we don't use "requested changes" if the comment is not emotional :-) in our case, it's just some syntactic nitpicking, so let's just use "comment" as review type

@Hzfengsy
Copy link
Member

Thanks for the reminder of @junrushao1994. And I'm sorry for bringing ambiguity to @csullivan. I will change the Changes requested into Approve and we can merge it if we want.

@csullivan
Copy link
Contributor Author

Thanks @Hzfengsy. I think I prefer T.serial personally but I continue to be impressed by the TVMScript parser. I agree the linter made this a bit ugly. I want to add some additional tests in a follow up so I'll happily address the formatting then and consider the change to range.

@Hzfengsy
Copy link
Member

Thanks @Hzfengsy. I think I prefer T.serial personally but I continue to be impressed by the TVMScript parser. I agree the linter made this a bit ugly. I want to add some additional tests in a follow up so I'll happily address the formatting then and consider the change to range.

T.serial is good too! Just make it in one line, e.g. for i in T.serial(0, 16): :)

@junrushao junrushao merged commit 1420df7 into apache:main Jun 13, 2022
@csullivan csullivan deleted the te/extern_primfunc branch June 13, 2022 15:51
@areusch areusch mentioned this pull request Jun 14, 2022
Closed
juda pushed a commit to juda/tvm that referenced this pull request Jun 21, 2022
@csullivan
[TE] Support schedulable TIR compute definitions in TOPI (apache#11589)
2fb7d59 
This PR adds `te.extern_primfunc` which provides the interface around TE ExternOp that allows a TVMScript defined schedulable TIR PrimFunc to be inlined into a TE compute graph. The result is that TIR can be used for compute definitions in Relay OpStrategies and, paired with meta-scheduler support in relay as introduced in apache#10578, these compute definitions can be scheduled and tuned as demonstrated in the attached tests.  

Prior to this, compute definitions were limited to those definable in TE only. As a consequence of this patch and ongoing improvements to TVMScript meta-programming (apache#11097), TOPI can be extended to include compute and scheduling functions targeting schedulable TIR uniformly.
@Lunderberg Lunderberg mentioned this pull request Jun 30, 2022
Merged
@driazati driazati mentioned this pull request Jul 14, 2022
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Hzfengsy Hzfengsy Hzfengsy approved these changes

@masahi masahi Awaiting requested review from masahi

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@csullivan @masahi @junrushao @Hzfengsy