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

[ioana ghiban]

================
@@ -29,128 +30,294 @@ 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").
-///
-/// Examples that return true:
-///
-///   // Row-major slice (last dim is non-unit)
-///   memref.reinterpret_cast %buff to offset: [%off],
-///     sizes: [1, 1, 1], strides: [1, 1, 1]
-///     : memref<1x1x8xi32> to memref<1x1x1xi32>
+//===----------------------------------------------------------------------===//
+// Copy Rewrite Helpers
+//===----------------------------------------------------------------------===//
+
+/// Returns row-major strides for static identity-layout memref type.
+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;
+}
+
+/// Finds the source dimension for a static reinterpret_cast result dimension.
+/// Dimensions marked in `usedSourceDims` are skipped. Returns the smallest
+/// source dimension whose size is at least the result dimension size, with the
+/// same stride.
+static std::optional<unsigned> findSourceDimForResultDim(
+    memref::ReinterpretCastOp rc, unsigned resultDim, MemRefType sourceType,
+    ArrayRef<int64_t> sourceStrides, ArrayRef<bool> usedSourceDims) {
+  MemRefType resultType = cast<MemRefType>(rc.getType());
+  assert(resultDim < resultType.getRank() && "result dimension out of range");
+  assert(sourceType.getRank() == static_cast<int64_t>(sourceStrides.size()) &&
+         sourceStrides.size() == usedSourceDims.size() &&
+         "expected same-rank source type, strides, and used-dimension mask");
+  assert(!ShapedType::isDynamic(rc.getStaticStrides()[resultDim]) &&
+         "expected static result stride");
+
+  int64_t resultStride = rc.getStaticStrides()[resultDim];
+  int64_t resultSize = resultType.getDimSize(resultDim);
+  std::optional<unsigned> sourceDim;
+  for (auto [idx, stride] : llvm::enumerate(sourceStrides)) {
+    if (usedSourceDims[idx] || stride != resultStride ||
+        sourceType.getDimSize(idx) < resultSize)
+      continue;
+
+    if (!sourceDim ||
+        sourceType.getDimSize(idx) < sourceType.getDimSize(*sourceDim))
+      sourceDim = idx;
+  }
+  return sourceDim;
+}
+
+/// Returns source indices for a static reinterpret_cast offset.
+static std::optional<SmallVector<int64_t>>
+delinearizeStaticOffset(memref::ReinterpretCastOp rc, MemRefType sourceType,
+                        ArrayRef<int64_t> sourceStrides) {
----------------
ioghiban wrote:

yes, they are taken from `rc`. so we can reduce the `delinearizeStaticOffset()` args to `rc` only. I will also re-compute these within `findBaseDimForResultDim()`, and move the logic for getting `rcResultDim` and `usedBaseDims` there, since those also depend only on `rc`.

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