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

[Andrzej Warzyński]

================
@@ -42,26 +47,31 @@ static std::optional<SmallVector<int64_t>> getIdentityStrides(MemRefType type) {
   return strides;
 }
 
+/// Non-unit strided memref dimensions are lowered to loops over base memref
+/// dimensions. Static reinterpret_cast strides connect those dimension spaces
+/// in the currently supported cases.
 static std::optional<unsigned>
-findBaseDimForViewStride(MemRefType baseType, ArrayRef<int64_t> baseStrides,
-                         ArrayRef<bool> usedBaseDims, int64_t viewStride,
-                         int64_t viewSize) {
+findBaseDimForResultStride(MemRefType baseType, ArrayRef<int64_t> baseStrides,
+                           ArrayRef<bool> usedBaseDims, int64_t resultStride,
+                           int64_t resultSize) {
   std::optional<unsigned> fallback;
   for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
-    if (usedBaseDims[idx] || stride != viewStride ||
-        baseType.getDimSize(idx) < viewSize)
+    if (usedBaseDims[idx] || stride != resultStride ||
+        baseType.getDimSize(idx) < resultSize)
       continue;
 
     // Prefer an exact shape match. Otherwise, use the first dimension large
     // enough to contain the copied logical vector.
-    if (baseType.getDimSize(idx) == viewSize)
+    if (baseType.getDimSize(idx) == resultSize)
       return idx;
     if (!fallback)
       fallback = idx;
   }
   return fallback;
 }
 
+/// reinterpret_cast offsets are linear element offsets, while `memref.store`
+/// needs one index per base dimension.
----------------
banach-space wrote:

Similar comment as above :)

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