[flang-commits] [flang] [FIRToMemRef] Fix wrong indexing for converted array_coor over sliced fir.embox (PR #207749)

[Kareem Ergawy via flang-commits]

================
@@ -456,6 +460,100 @@ static Value castTypeToIndexType(Value originalValue,
                                     originalValue);
 }
 
+/// Fold the embox's slice per-dim (lb - 1) shift into `indices` after
+/// `getMemrefIndices` has produced them. Handles both range and rank-reducing
+/// (scalar-subscript) slice triples uniformly.
+///
+/// Necessary because `getMemrefIndices` only consumes the first `rank*3`
+/// entries of `sliceInfo.sliceVec` (the array_coor's own slice); the embox's
+/// triples sit at `sliceVec[rank*3 .. 2*rank*3-1]` and are otherwise dropped,
+/// leaving the memref access short by `(lb - 1)` per dim -- and, for
+/// collapsed dims, additionally polluted with the array_coor's own scalar
+/// subscript (typically a loop variable) that would drag the load off the
+/// target dim across iterations.
+///
+/// Behaviour, per Fortran dim `d` of the embox slice (memref position
+/// `m = parentRank - 1 - d`):
+///   - Range triple `(lb, ub, step)`:
+///       indices[m] += (lb - 1)
+///     Adds the (lb - 1) shift on top of whatever `getMemrefIndices` produced
+///     for that dim.
+///   - Scalar-subscript triple `(lb, undef, undef)`:
+///       indices[m]  = (lb - 1)
+///     Overwrites the memref position outright. The dim is collapsed in the
+///     box's rank, but the reinterpret_cast still keeps a memref index slot
+///     for it (with the correct parent stride). Placing `(lb - 1)` there is
+///     mathematically equivalent to shifting the flat offset by
+///     `(lb - 1) * parentStride[d]`, without needing any parent-shape
+///     lookups from the caller.
+///
+/// Example A -- non-rank-reducing (2-D):
+///
+///   `fir.array_coor %box(...)[%innerSlice] %i, %j` with
+///     - array_coor slice: identity in both dims
+///     - embox slice     : `fir.slice %c2,%c2,%c1, %c1,%c1,%c1`
+///                         (Fortran dim 0 lb = 2, dim 1 lb = 1)
+///
+///   Incoming `indices` (memref order): [ %j', %i' ]
+///   Fold:
+///     - dim 0 (range): delta = 1 -> indices[1] += 1
+///     - dim 1 (range): delta = 0 -> indices[0] += 0
+///   Result:  [ %j', %i' + 1 ]
+///
+/// Example B -- rank-reducing (3-D parent, dim 1 collapsed):
+///
+///   %parent : !fir.ref<!fir.array<3x2x2xi32>>
+///   %shape  = fir.shape %c3, %c2, %c2
+///   %eslc   = fir.slice %c1, %c3, %c1,  %c2, %undef, %undef,  %c1, %c2, %c1
+///   %box    = fir.embox %parent(%shape) [%eslc]
+///                                          : !fir.box<!fir.array<3x2xi32>>
+///   %bshape = fir.shape %c3, %c2                     ; box's rank-2 shape
+///   %addr   = fir.array_coor %box(%bshape) %ii, %jj
+///        : (!fir.box<!fir.array<3x2xi32>>, !fir.shape<2>, index, index)
+///        -> !fir.ref<i32>
+///
+///   parentRank = 3. `getRankFromEmbox` returns the parent's rank, so
+///   `getMemrefIndices` builds a rank-3 index vector (memref order):
+///     indices = [ %jj', <ac scalar sub for collapsed dim>, %ii' ]
+///                ^ dim 2   ^ dim 1 (collapsed)              ^ dim 0
+///
+///   Fold:
+///     - dim 0 (range, lb=1): indices[2] += 0
+///     - dim 1 (scalar, lb=2): indices[1]  = 1        (overwrite)
+///     - dim 2 (range, lb=1): indices[0] += 0
+///
+///   Final: `memref.load %reinterpret_cast[%jj', 1, %ii'] : memref<?x?x?xi32,
+///   strided<[?, ?, ?], offset: ?>>` -- stride 3 on the middle position
+///   lands on the correct plane, `offset` stays at 0.
+void FIRToMemRef::foldSliceLbIntoIndices(SmallVectorImpl<Value> &indices,
+                                         fir::EmboxOp embox,
+                                         PatternRewriter &rewriter,
+                                         Location loc) const {
+  auto sliceOp = getSliceOp(embox.getSlice());
+  if (!sliceOp)
+    return;
+  auto triples = sliceOp.getTriples();
+  unsigned parentRank = triples.size() / 3;
+  if (parentRank == 0)
+    return;
+
+  auto isUndef = [](Value v) {
+    return v && v.getDefiningOp<fir::UndefOp>() != nullptr;
+  };
+  Value cOne = arith::ConstantIndexOp::create(rewriter, loc, 1);
+  for (unsigned d = 0; d < parentRank; ++d) {
+    unsigned m = parentRank - 1 - d;
+    if (m >= indices.size())
+      continue;
+    Value lb = castTypeToIndexType(triples[d * 3], rewriter);
+    Value delta = arith::SubIOp::create(rewriter, loc, lb, cOne);
+    bool isScalar = isUndef(triples[d * 3 + 1]) || isUndef(triples[d * 3 + 2]);
+    indices[m] = isScalar
+                     ? delta
+                     : arith::AddIOp::create(rewriter, loc, indices[m], delta);
----------------
ergawy wrote:

Done. You uncovered a bug. The formula should be:
* 1-based indexing: `lb + (idx - 1) * stride`
* 0-based indexing: `(lb - 1) + idx * stride`.
Thanks a lot.

https://github.com/llvm/llvm-project/pull/207749