| CARVIEW |
Navigation Menu
-
Notifications
You must be signed in to change notification settings - Fork 5.8k
[CINN] Support grid reduce for multiple reduces #68750
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
Conversation
|
你的PR提交成功，感谢你对开源项目的贡献! |
lshpku
force-pushed
the
multiple_grid_reduce
branch
from
4a68544 to
ce38d2a
Compare
lshpku
force-pushed
the
multiple_grid_reduce
branch
from
ce38d2a to
b7aa0e4
Compare
| std::vector<ir::Expr*> upstream_bodies_; | ||
| std::vector<ir::Expr*> grid_reduce_bodies_; | ||
| std::vector<ir::Expr*> downstream_bodies_; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
这里的ir::Expr 可以直接用 ir::For 吗？
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
可以，但是没必要，后面CollectComputeBodies的时候还要转回Expr
| bool is_grid_reduce_downstream = IsGridReduceDownstream(*op); | ||
|
|
||
| if (is_grid_reduce || is_grid_reduce_downstream) { | ||
| downstream_names_.insert(schedule_block->name); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
is_grid_reduce_downstream的处理流程是什么样的？看上去是IsGridReduceDownstream里要判断downstream_names_，然后downstream_names_又依赖IsGridReduceDownstream的结果
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
IsGridReduceDownstream判断自己是不是依赖于grid reduce或者某个downstream，如果是的话把自己也加入downstream，是一个递推算法，只需从前到后遍历一次每个schedule block，就能判断每个schedule block是不是downstream
| std::vector<ir::Expr> CollectComputeBodiesInOrder() { | ||
| std::vector<ir::Expr> stmts; | ||
| for (auto* body : upstream_bodies_) { | ||
| stmts.push_back(*body); | ||
| } | ||
| for (auto* body : grid_reduce_bodies_) { | ||
| stmts.push_back(*body); | ||
| } | ||
| for (auto* body : downstream_bodies_) { | ||
| stmts.push_back(*body); | ||
| } | ||
| return stmts; | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
stmts的顺序调整可以确保不违反原先的依赖关系吗？
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
这个pass自己不做检查，由前面在判断can_apply_grid_reduce的时候就保证，如果can_apply_grid_reduce=true，则这张图里面所有的grid_reduce都是平行的，不存在一个grid_reduce依赖于另一个grid_reduce的情况，所以它们的upstream和downstream一定是严格分开的，不存在downstream反过来依赖upstream的情况
| * ... | ||
| * } | ||
| */ | ||
| struct CrossBlockReductionReorderer : public BaseMutator { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
这个reorder后续可以考虑换个名字比如聚类啥的会不会好点呢，后期如果要挪到schedule的话我的考虑也是把这个移动的语义限定为“分组，排序”之类的
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
可以改成CrossBlockReductionGroupingMutator
PR Category
CINN
PR Types
Improvements
Description
This PR mainly adds grid reduce support for fusion groups that have multiple reduces. To achieve this, this PR adds a
CrossBlockReductionReordererpass that separates the upstream and downstream schedule blocks of grid reduce. We can then insert a fence (update_semaphore) between the last upstream block and the first grid reduce block.After the
CrossBlockReductionReordererpass, the compute graph is like:This PR also adds tiling support for RS layout.
pcard-85711