Skip to content

reshard: fix reverse-iterator direction in getGatheringAxes - #1453

Open
EylonKrause wants to merge 1 commit into
openxla:mainfrom
EylonKrause:eylonk/reshard-gather-axes-iterator
Open

reshard: fix reverse-iterator direction in getGatheringAxes#1453
EylonKrause wants to merge 1 commit into
openxla:mainfrom
EylonKrause:eylonk/reshard-gather-axes-iterator

Conversation

@EylonKrause

Copy link
Copy Markdown

Problem

CollectiveInserter::getGatheringAxes (shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc) over-gathers a dimension during reshard lowering: it can return the entire axis list (including axes still present in the output sharding) instead of only the trailing axes absent from the output.

Root cause

The suffix scan starts at inAxes.rbegin() and does --axisRevIt, which advances the std::list reverse_iterator toward the back (off the end) rather than toward the front:

auto axisRevIt = inAxes.rbegin();
while (axisRevIt != inAxes.rend() &&
       !outAxisToDimAndIndex.contains(*axisRevIt)) {
  inAxisSet.erase(*axisRevIt);
  --axisRevIt;   // wrong direction
}

For the function's own documented example (inAxesPerDim = [[], ["x","y","z","w"]], outAxesPerDim = [["y"], []], expected ["z","w"]), -- stops after one step and axisRevIt.base() lands at the front, so the tail loop gathers ["x","y","z","w"] — wrongly gathering "y", which the output still needs — and leaves currentAxes / inAxisSet inconsistent.

The sibling reverse scan in getAllToAllInfo (line ~1134) and popBackFromCurrentAxes (line 118) both use ++.

Fix

Advance with ++axisRevIt.

Verification

Static reasoning + consistency with the sibling scans; reproduces the documented example. No bazel build available on my machine — happy to add a lit test for a reshard that lowers to an all-gather.

The suffix scan in CollectiveInserter::getGatheringAxes started at
inAxes.rbegin() and decremented the reverse_iterator (--axisRevIt),
which moves toward the back (off the end) instead of toward the front.
The loop therefore stopped after a single step and getGatheringAxes
over-gathered the entire dimension -- including axes still present in the
output sharding -- instead of only the trailing axes absent from the
output, leaving currentAxes/inAxisSet inconsistent. Advance with
++axisRevIt, matching the sibling reverse scan in getAllToAllInfo and
popBackFromCurrentAxes.
@EylonKrause

Copy link
Copy Markdown
Author

Disclosure: this contribution was authored with an AI coding assistant (Claude) and reviewed before submission.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant