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

[Andrzej Warzyński]

================
@@ -29,128 +29,268 @@ using namespace mlir;
 
 namespace {
 
-/// Returns true if `rc` represents a scalar view (all sizes == 1)
-/// into a memref that has exactly one non-unit dimension located at
-/// either the first or last position (i.e. a "row" or "column").
+static std::optional<SmallVector<int64_t>> getIdentityStrides(MemRefType type) {
+  if (!type.getLayout().isIdentity() || !type.hasStaticShape())
+    return std::nullopt;
+
+  SmallVector<int64_t> strides(type.getRank(), 1);
+  int64_t stride = 1;
+  for (int64_t dim = type.getRank() - 1; dim >= 0; --dim) {
+    strides[dim] = stride;
+    stride *= type.getDimSize(dim);
+  }
+  return strides;
+}
+
+static std::optional<unsigned>
+findBaseDimForViewStride(MemRefType baseType, ArrayRef<int64_t> baseStrides,
+                         ArrayRef<bool> usedBaseDims, int64_t viewStride,
+                         int64_t viewSize) {
+  std::optional<unsigned> fallback;
+  for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
+    if (usedBaseDims[idx] || stride != viewStride ||
+        baseType.getDimSize(idx) < viewSize)
+      continue;
+
+    // Prefer an exact shape match. Otherwise, use the first dimension large
+    // enough to contain the copied logical vector.
+    if (baseType.getDimSize(idx) == viewSize)
+      return idx;
+    if (!fallback)
+      fallback = idx;
+  }
+  return fallback;
+}
+
+static std::optional<SmallVector<int64_t>>
+delinearizeStaticOffset(int64_t offset, MemRefType baseType,
+                        ArrayRef<int64_t> baseStrides) {
+  if (offset < 0)
+    return std::nullopt;
+
+  SmallVector<int64_t> indices(baseType.getRank(), 0);
+  int64_t remainder = offset;
+  for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
+    indices[idx] = remainder / stride;
+    if (indices[idx] >= baseType.getDimSize(idx))
+      return std::nullopt;
+    remainder %= stride;
+  }
+
+  if (remainder != 0)
+    return std::nullopt;
+  return indices;
+}
+
+static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
----------------
banach-space wrote:

DOCUMENTME

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