Fix out-of-bounds read for unscattered dims in multi-dim reduce-scatter lowering

[EylonKrause]

Problem

-sdy-convert-global-to-local='combine-multi-dimension-reduce-scatter=true' reads past the end of a SmallVector (and can divide by a garbage factor) when lowering a reduce_scatter that scatters two or more dimensions but leaves at least one dimension unscattered.

Root cause

In ReduceScatterOpPattern::rewriteReduceScatterCombiningMultipleDims (convert_global_to_local.cc), the loop that splits each scattered dimension into (factor, quotient) handles a non-scattered dimension but forgets to continue:

auto* it = llvm::find(scatteredDims, i);
if (it == scatteredDims.end()) {
  splitShape.push_back(inputShape[i]);
}                                              // <-- missing continue
int64_t factor =
    scatteredFactors[std::distance(scatteredDims.begin(), it)];
...
if (inputShape[i] == ShapedType::kDynamic || inputShape[i] % factor != 0) { ... }

When dimension i is not scattered, it == scatteredDims.end(), so std::distance(begin, end) == scatteredDims.size(), and scatteredFactors[scatteredDims.size()] is one past the end (scatteredFactors.size() == scatteredDims.size()). Execution then pushes a bogus factor sub-dimension into splitShape, records it in factorDimIndices (corrupting the subsequent transpose permutation), and computes inputShape[i] % factor on the garbage value (a SIGFPE if it happens to be 0). Under an assertions-enabled build (standard for MLIR) the SmallVector::operator[] bounds assert aborts.

The sibling all-to-all combining routine has the continue; this is an omission in the reduce-scatter adaptation. The only existing multi-dimension test happens to scatter all dimensions, so the unscattered-dimension path is untested.

Fix

Add the missing continue so a non-scattered dimension keeps its original size and does not emit a factor sub-dimension. This is a no-op for the all-dimensions-scattered case.

Test

sdy_reduce_scatter_partial_dims.mlir: a rank-3 reduce_scatter that scatters dimensions 0 and 1 but leaves dimension 2 unscattered now lowers to a single stablehlo.reduce_scatter instead of reading out of bounds.

---

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