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

[Andrzej Warzyński]

================
@@ -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) {
+  ArrayRef<int64_t> offsets = rc.getStaticOffsets();
+  // FIXME: Despite what `getStaticOffsets` implies, `reinterpret_cast` takes
+  // only a single offset. That should be fixed at the op definition level.
+  assert(offsets.size() == 1 && "Expecting single offset");
+  assert(!ShapedType::isDynamic(offsets[0]) && "expected static offset");
+
+  if (offsets[0] < 0)
+    return std::nullopt;
+
+  SmallVector<int64_t> indices(sourceType.getRank(), 0);
+  int64_t remainder = offsets[0];
+  for (auto [idx, stride] : llvm::enumerate(sourceStrides)) {
+    indices[idx] = remainder / stride;
+    if (indices[idx] >= sourceType.getDimSize(idx))
+      return std::nullopt;
+    remainder %= stride;
+  }
+
+  if (remainder != 0)
+    return std::nullopt;
+  return indices;
+}
+
+/// Returns the dimension whose static size is not one if it is unique.
+static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
+  assert(type.hasStaticShape() && "expected static shape");
+  ArrayRef<int64_t> shape = type.getShape();
+  if (shape.empty())
+    return std::nullopt;
+
+  std::optional<unsigned> nonUnitDim;
+  for (auto [idx, dim] : llvm::enumerate(shape)) {
+    if (dim == 1)
+      continue;
+    if (nonUnitDim)
+      return std::nullopt;
+    nonUnitDim = idx;
+  }
+  return nonUnitDim;
+}
+
+/// Per-dimension loop nest info.
+struct CopyLoopDimInfo {
+  unsigned copyDim;
+  unsigned baseDim;
+  int64_t size;
+};
+
+/// Rewrite info from reinterpret_cast layout, captured after passing legality
+/// checks.
+struct CopyFromReinterCastInfo {
+  SmallVector<CopyLoopDimInfo> loopDims;
----------------
banach-space wrote:

IIUC, `loopBound` are required for non-scalar loads which are "lowered" to loops, right? Could you document that?

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