Open
Description
They currently assume DID is only in logical and will break when we switch to loop split.
The solution I have in mind and am open to discussion is:
- Change
ReorderShardedAxis
to move DID to the front in loop. We may even do this universally because many schedulers bear that assumption, but it can come separately. - Change
MakeReshardingContiguous
to set allocation around each resharding expression (which has been decomposed and therefore can be lowered to communication) to be the same as loop and be contiguous. - Change existing Communication IRs to respect the allocation order. I don't know exactly what needs to be fixed, but I think
Fuser/csrc/multidevice/communication.cpp
Line 313 in bd2c81d
That said, even without DID loop split, I think this is a good cleanup for ReorderShardedAxis. For example, it wouldn't need to insert Set.Permute, making fusion IR easier to follow.
For #2563
Activity