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

[Andrzej Warzyński]

================
@@ -29,171 +31,408 @@ 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
+//===----------------------------------------------------------------------===//
+
+/// Non-unit reinterpret_cast result dimension and the source dimension it
+/// advances through.
+struct NonUnitDimAssocMapForRC {
+  unsigned resultDimPos;
+  unsigned sourceDimPos;
+};
+
+/// Copy-relevant information derived from a reinterpret_cast.
+struct AssocMapAndOffsetsForRC {
+  // Non-unit dimensions of the reinterpret_cast result.
+  SmallVector<NonUnitDimAssocMapForRC> assocMap;
+  // Delinearized offsets to in-bounds reinterpret_cast source indices.
+  // Optional since it is only supported for static offsets.
+  std::optional<SmallVector<int64_t>> delinearizedOffsets;
+};
+
+/// Records the reinterpret_cast result dimensions that span more than one
+/// element and maps each one to its corresponding source dimension.
+static bool findSourceDimForResultDim(memref::ReinterpretCastOp rc,
+                                      AssocMapAndOffsetsForRC &mapAndOffs) {
+  MemRefType resType = dyn_cast<MemRefType>(rc.getType());
+  MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
+  assert(srcType.getLayout().isIdentity() &&
+         "Expecting identity source layout.");
+
+  SmallVector<int64_t> srcIdentityStrides = computeStrides(srcType.getShape());
+
+  // Reusing a source dimension would require delinearizing the combined linear
+  // offset, which is TODO.
+  SmallVector<bool> usedSrcDims(srcType.getRank(), false);
+
+  for (auto [resultDim, resultSize] : llvm::enumerate(resType.getShape())) {
+    if (resultSize == 1)
+      continue;
+
+    // TODO: Support dynamic strides on non-unit result dimensions.
+    if (ShapedType::isDynamic(rc.getStaticStrides()[resultDim]))
+      return false;
+
+    int64_t resultStride = rc.getStaticStrides()[resultDim];
+    std::optional<unsigned> srcDim;
+    // Find an unused source dimension with matching stride and enough elements.
+    for (auto [idx, stride] : llvm::enumerate(srcIdentityStrides)) {
+      if (usedSrcDims[idx] || stride != resultStride ||
+          srcType.getDimSize(idx) < resultSize)
+        continue;
+
+      if (!srcDim || srcType.getDimSize(idx) < srcType.getDimSize(*srcDim))
+        srcDim = idx;
+    }
+    if (!srcDim)
+      return false;
+
+    usedSrcDims[*srcDim] = true;
+    mapAndOffs.assocMap.push_back(
+        NonUnitDimAssocMapForRC{static_cast<unsigned>(resultDim), *srcDim});
+  }
+  return true;
+}
+
+/// Returns source indices for a static reinterpret_cast offset of an
+/// identity-layout source.
+static std::optional<SmallVector<int64_t>>
+delinearizeStaticRCOffset(memref::ReinterpretCastOp rc) {
+  ArrayRef<int64_t> rcOffsets = rc.getStaticOffsets();
+  // FIXME: Despite what `getStaticOffsets` implies, `reinterpret_cast` takes
+  // only a single offset. That should be fixed at the op definition level.
+  assert(rcOffsets.size() == 1 && "Expecting single offset");
+  assert(ShapedType::isStatic(rcOffsets[0]) && "expected static offset");
+
+  assert(rcOffsets[0] >= 0 &&
+         "static reinterpret_cast offset must be non-negative");
+
+  MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
+  assert(srcType.getLayout().isIdentity() &&
+         "Expecting identity source layout.");
+  if (srcType.getRank() == 0) {
+    assert(rcOffsets[0] == 0 &&
+           "non-zero static offset is invalid for rank-0 source memref");
+    return SmallVector<int64_t>{};
+  }
+
+  SmallVector<int64_t> offsetIdxs(srcType.getRank(), 0);
+  int64_t remainder = rcOffsets[0];
+  SmallVector<int64_t> srcStrides = computeStrides(srcType.getShape());
+  // Convert the linear reinterpret_cast offset to per-dimension source starting
----------------
banach-space wrote:

[nit]
```suggestion
  // Convert the scalar reinterpret_cast offset to per-dimension source starting
```

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