[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;
----------------
ioghiban wrote:

TODO: `indices` -> `idxs`

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