[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)

[ioana ghiban]

================
@@ -162,38 +321,122 @@ struct CopyToScalarLoadAndStore : public OpRewritePattern<memref::CopyOp> {
       return rewriter.notifyMatchFailure(
           op, "target is not a memref.reinterpret_cast");
 
-    if (!isScalarSlice(rc))
+    std::optional<CopyFromReinterCastInfo> copyInfo =
+        getCopyFromReinterCastInfo(op, rc);
+    if (!copyInfo)
       return rewriter.notifyMatchFailure(
-          op, "reinterpret_cast does not match scalar slice");
+          op, "reinterpret_cast does not match scalar or loop copy region");
 
     Location loc = op.getLoc();
-
     Value src = op.getSource();
     Value dst = rc.getSource();
 
-    auto dstType = cast<MemRefType>(dst.getType());
-    unsigned dstRank = dstType.getRank();
+    MemRefType cpSrcType = cast<MemRefType>(src.getType());
+    MemRefType dstType = cast<MemRefType>(dst.getType());
+
+    // Reuse common index constants across bounds, steps, and static offsets,
+    // but avoid creating them for rank-0 copies.
+    std::array<Value, 2> cachedIndexConstants;
+    auto getOrCreateIndexConstant = [&](int64_t value) -> Value {
+      if (value == 0 || value == 1) {
+        Value &cached = cachedIndexConstants[value];
+        if (!cached)
+          cached = arith::ConstantIndexOp::create(rewriter, loc, value);
+        return cached;
+      }
+      return arith::ConstantIndexOp::create(rewriter, loc, value);
+    };
+    auto getZeroIndices = [&](int64_t rank) {
+      SmallVector<Value> indices;
+      indices.reserve(rank);
+      if (rank != 0)
+        indices.append(rank, getOrCreateIndexConstant(0));
+      return indices;
+    };
+
+    // Create all loop bounds before building the loop nest. Otherwise an
+    // inner-loop bound can be inserted inside an outer loop body.
+    SmallVector<Value> upperBounds;
+    upperBounds.reserve(copyInfo->loopDims.size());
+    for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims)
+      upperBounds.push_back(getOrCreateIndexConstant(loopDim.size));
+
+    SmallVector<Value> baseStoreIndices = getZeroIndices(dstType.getRank());
+    // Static offsets were already delinearized into base indices. Fill the
+    // non-zero starting indices before creating loop bodies.
+    if (copyInfo->staticOffsetIdxs) {
+      for (auto [idx, offset] : llvm::enumerate(*copyInfo->staticOffsetIdxs)) {
+        if (offset == 0)
+          continue;
+        baseStoreIndices[idx] = getOrCreateIndexConstant(offset);
+      }
+    } else {
+      // Supported dynamic offsets are used directly in exactly one base
+      // dimension selected by getCopyFromReinterCastInfo.
+      assert(copyInfo->dynamicOffsetDim &&
+             "expected dynamic offset dimension for dynamic offset");
+      SmallVector<OpFoldResult> rcOffsets = rc.getMixedOffsets();
+      // FIXME: Despite what `getMixedOffsets` implies, `reinterpret_cast` takes
+      // only a single offset. That should be fixed at the op definition level.
+      assert(rcOffsets.size() == 1 && "Expecting single offset");
+      baseStoreIndices[*copyInfo->dynamicOffsetDim] =
+          getValueOrCreateConstantIndexOp(rewriter, loc, rcOffsets[0]);
+    }
 
-    Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+    // Scope for OpBuilder::InsertionGuard.
+    {
+      OpBuilder::InsertionGuard guard(rewriter);
+      Value lowerBound;
+      Value step;
+      if (!upperBounds.empty()) {
+        lowerBound = getOrCreateIndexConstant(0);
+        step = getOrCreateIndexConstant(1);
+      }
 
-    auto srcType = cast<MemRefType>(src.getType());
-    Repeated<Value> loadIndices(srcType.getRank(), zero);
-    auto offsets = rc.getMixedOffsets();
-    assert(offsets.size() == 1 && "Expecting single offset");
-    OpFoldResult offset = offsets[0];
-    Value storeOffset = getValueOrCreateConstantIndexOp(rewriter, loc, offset);
-    unsigned offsetDim = dstType.getDimSize(0) == 1 ? dstRank - 1 : 0;
-    SmallVector<Value> storeIndices(dstRank, zero);
-    storeIndices[offsetDim] = storeOffset;
-    // If the only user of `rc` is the current Op (which is about to be erased),
-    // we can safely erase it.
-    if (rcOutput.hasOneUse())
-      rewriter.eraseOp(rc);
+      SmallVector<Value> loopIvs;
+      loopIvs.reserve(copyInfo->loopDims.size());
 
-    Value val = memref::LoadOp::create(rewriter, loc, src, loadIndices);
-    memref::StoreOp::create(rewriter, loc, val, dst, storeIndices);
+      // Build one nested loop per non-unit copied strided memref dimension.
+      for (Value upperBound : upperBounds) {
+        scf::ForOp loop =
+            scf::ForOp::create(rewriter, loc, lowerBound, upperBound, step);
+        loopIvs.push_back(loop.getInductionVar());
+        rewriter.setInsertionPointToStart(loop.getBody());
+      }
 
+      // Load indices are zero except for copied strided memref dimensions,
+      // which use the corresponding loop induction variables.
+      SmallVector<Value> loadIndices = getZeroIndices(cpSrcType.getRank());
----------------
ioghiban wrote:

TODO: `loadIndices` -> `loadIdxs`

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