Enable grid reductions within loops. #1681
Conversation
…oops without turning the kernel cooperative.
| const nvfuser_index_t entrance_ind, | ||
| const nvfuser_index_t n_entrances) { |
There was a problem hiding this comment.
nit: can't imagine these would need 64 bits...
| if (fl->isTrivial()) { | ||
| continue; | ||
| } | ||
| if (fl->iter_domain()->isThread()) { |
There was a problem hiding this comment.
fl->isTrivial() should include this condition.
| (grouped_rop->isAllreduce() && is_within_a_loop ? 2 : 1)), | ||
| output->dtype(), | ||
| false); | ||
| }); | ||
|
|
||
| const auto sync_buffer = ir_utils::allocGlobalBufferForGridComm( | ||
| getGridSyncBufferSize(out_domain), DataType::Int, true); | ||
| getGridSyncBufferSize(out_domain, for_loops_), DataType::Int, true); |
There was a problem hiding this comment.
GroupedReduction doesn't support reentrance, so this is not necessary right now.
| @@ -271,12 +347,25 @@ void IndexLowering::handleGridReduction( | |||
|
|
|||
| const auto reduce_buffer = ir_utils::allocGlobalBufferForGridComm( | |||
| getGridCommWorkBufferSize( | |||
| out_domain, rop->isAllreduce() && is_within_a_loop ? 2 : 1), | |||
| out_domain, | |||
| rop->isAllreduce() ? std::vector<kir::ForLoop*>() : for_loops_, | |||
There was a problem hiding this comment.
Would be nice if we could refactor the code on the conditional processing of when to expand the buffer.
| return buffer_size; | ||
| } | ||
|
|
||
| Val* getGridSyncBufferSize(const TensorDomain* td) { | ||
| Val* getGridSyncBufferSize( |
There was a problem hiding this comment.
Do we need to expand the sync buffer?
There was a problem hiding this comment.
No, we shouldn't need to. Because we wait until all iterations are done to start cleaning any of them up. Maybe that's a reason it's slow, I think we should use multiple sync buffers for each reduction!
There was a problem hiding this comment.
I don't know if it would make a difference, but seems like it's low risk.
|
Closing in favor of #1698 |
Allow re-entrant grid reductions, allowing them to be placed inside loops without turning the kernel cooperative.