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

[Andrzej Warzyński]

================
@@ -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.
----------------
banach-space wrote:

I find this comment confusing.
* There was only _one_ offset that belonged to `memref.reinterpret_cast` (`offsets` --> `offset`)
* The fact that the offset has been de-linearized is IMO irrelevant here. IIUC, what you are trying to say is that for "static" offset, we re-use the indices used from de-linearization. For "dynamic" offset, we need to rely on a different mechanism.

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