Skip to content

[TensorIR] Cross-Thread Reduction#9360

Merged
junrushao merged 6 commits into
apache:mainfrom
MasterJH5574:m2/cross-thread-reduction
Nov 14, 2021
Merged

[TensorIR] Cross-Thread Reduction#9360
junrushao merged 6 commits into
apache:mainfrom
MasterJH5574:m2/cross-thread-reduction

Conversation

@MasterJH5574

Copy link
Copy Markdown
Contributor

Hi community! This PR adds cross-thread reduction support for TensorIR. After this PR, cross-thread reduction patterns in TIR can be successfully lowered.

cc @Hzfengsy @vinx13 @comaniac @junrushao1994 @jcf94 @jinhongyii @spectrometerHBH @tqchen

Co-authored-by: Wuwei Lin wuwei@apache.org
Co-authored-by: Junru Shao junrushao1994@gmail.com
Co-authored-by: Siyuan Feng Hzfengsy@sjtu.edu.cn
Co-authored-by: Hongyi Jin 3231950289@qq.com
Co-authored-by: Bohan Hou 32121147+spectrometerHBH@users.noreply.github.com

Comment thread src/tir/schedule/analysis.h Outdated
Comment thread src/tir/schedule/analysis/analysis.cc
Comment thread src/tir/transforms/lower_cross_thread_reduction.cc Outdated
@junrushao

Copy link
Copy Markdown
Member

I will do another round of review next week!

@junrushao

Copy link
Copy Markdown
Member

Will do the review tomorrow

Comment thread src/tir/transforms/lower_cross_thread_reduction.cc Outdated
Comment thread src/tir/transforms/lower_cross_thread_reduction.cc Outdated
@junrushao

Copy link
Copy Markdown
Member

Finally got some time for a detailed code review! Will take over this PR and try to get it merged!

@junrushao junrushao force-pushed the m2/cross-thread-reduction branch from 3ca0e56 to a560a9b Compare November 12, 2021 03:24
@junrushao

Copy link
Copy Markdown
Member

Did a pass over analysis and misc changes

@junrushao

Copy link
Copy Markdown
Member

Love the PR and very comprehensively tested implementation ❤️

@junrushao junrushao force-pushed the m2/cross-thread-reduction branch from 4ac6c65 to 6a162d7 Compare November 12, 2021 20:29
@junrushao

Copy link
Copy Markdown
Member

@Hzfengsy @MasterJH5574 Should be good to go. Please take another look :-)

Comment thread src/tir/transforms/lower_cross_thread_reduction.cc Outdated
@MasterJH5574

Copy link
Copy Markdown
Contributor Author

@Hzfengsy Could you take another look? Junru's polishing looks very good, but I myself as the author cannot approve this PR 😅.

@Hzfengsy Hzfengsy left a comment

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.

LGTM. Thanks @MasterJH5574 for such a great effort on this PR.

@junrushao junrushao merged commit 08898e1 into apache:main Nov 14, 2021
mehrdadh pushed a commit to mehrdadh/tvm that referenced this pull request Dec 1, 2021
* [TensorIR] Cross-Thread Reduction

* Code revision on analysis and misc

* Refactor TransformReductionBlock

* Refactor code organization

* Address comment

* Use `std::make_tuple`

Co-authored-by: Junru Shao <junrushao1994@gmail.com>
mehrdadh pushed a commit to mehrdadh/tvm that referenced this pull request Dec 1, 2021
* [TensorIR] Cross-Thread Reduction

* Code revision on analysis and misc

* Refactor TransformReductionBlock

* Refactor code organization

* Address comment

* Use `std::make_tuple`

Co-authored-by: Junru Shao <junrushao1994@gmail.com>
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 7, 2022
* [TensorIR] Cross-Thread Reduction

* Code revision on analysis and misc

* Refactor TransformReductionBlock

* Refactor code organization

* Address comment

* Use `std::make_tuple`

Co-authored-by: Junru Shao <junrushao1994@gmail.com>
yangulei pushed a commit to yangulei/tvm that referenced this pull request Jan 11, 2022
* [TensorIR] Cross-Thread Reduction

* Code revision on analysis and misc

* Refactor TransformReductionBlock

* Refactor code organization

* Address comment

* Use `std::make_tuple`

Co-authored-by: Junru Shao <junrushao1994@gmail.com>
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022
* [TensorIR] Cross-Thread Reduction

* Code revision on analysis and misc

* Refactor TransformReductionBlock

* Refactor code organization

* Address comment

* Use `std::make_tuple`

Co-authored-by: Junru Shao <junrushao1994@gmail.com>
@LeiWang1999

Copy link
Copy Markdown
Contributor

hi @MasterJH5574 , do we currently only support block reduction with warp level __shfl_down_sync ? or the reduction can also be lowered into block reduce within shared memory.

@LeiWang1999

Copy link
Copy Markdown
Contributor

I've found the related implementation, thanks.

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.

5 participants