[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;
----------------
banach-space wrote:

We would only evaluate this condition if `srcDim` is already set: `srcType.getDimSize(idx) < srcType.getDimSize(*srcDim)` . That would mean that this condition was met for two different `resultStride`s: `resultStride == stride`. Is it possible?

I just find counter-intuitive and fragile if we update `srcDim` like this. It means that there can be more than one candidate. If that's the case, we need more determinism (as opposed to simply selecting the largest size). Do you have a test that would demonstrate this case?

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