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

[ioana ghiban]

================
@@ -29,128 +31,285 @@ 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
+//===----------------------------------------------------------------------===//
+
+/// 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 {
+  // Loop bounds that non-scalar loads "lower" to
+  SmallVector<CopyLoopDimInfo> loopDims;
+  // Deinearized offsets to in-bounds base indices.
+  std::optional<SmallVector<int64_t>> staticOffsetIdxs;
+  // reinterpret_cast dynamic offsets only supported for single non-unit
+  // dimension base, stored here to receive them.
+  std::optional<unsigned> dynamicOffsetDim;
+};
+
+/// Maps non-unit reinterpret_cast result dimensions to distinct base
+/// dimensions.
+static bool findBaseDimForResultDim(memref::ReinterpretCastOp rc,
+                                    CopyFromReinterCastInfo &info) {
+  MemRefType resType = dyn_cast<MemRefType>(rc.getType());
+  MemRefType baseType = dyn_cast<MemRefType>(rc.getSource().getType());
+  SmallVector<int64_t> baseIdentityStrides =
+      computeStrides(baseType.getShape());
+
+  // Each result loop IV is added directly to one base index. Reusing a base
+  // dimension would require delinearizing the combined linear offset.
+  SmallVector<bool> usedBaseDims(baseType.getRank(), false);
+
+  // Populate one loop-dimension entry for each non-unit result dimension.
+  for (auto [resultDim, resultSize] : llvm::enumerate(resType.getShape())) {
+    if (resultSize == 1)
+      continue;
+
+    // TODO: Support dynamic strides on copied dimensions.
+    if (ShapedType::isDynamic(rc.getStaticStrides()[resultDim]))
+      return false;
+
+    int64_t resultStride = rc.getStaticStrides()[resultDim];
+    std::optional<unsigned> baseDim;
+    // Find an unused base dimension with matching stride and enough elements.
+    for (auto [idx, stride] : llvm::enumerate(baseIdentityStrides)) {
+      if (usedBaseDims[idx] || stride != resultStride ||
+          baseType.getDimSize(idx) < resultSize)
+        continue;
+
+      if (!baseDim || baseType.getDimSize(idx) < baseType.getDimSize(*baseDim))
+        baseDim = idx;
+    }
+    if (!baseDim)
+      return false;
+
+    usedBaseDims[*baseDim] = true;
+    info.loopDims.push_back(CopyLoopDimInfo{static_cast<unsigned>(resultDim),
+                                            *baseDim, resultSize});
+  }
+  return true;
+}
+
+/// Returns base indices for a static reinterpret_cast offset.
+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::isDynamic(rcOffsets[0]) && "expected static offset");
+
+  if (rcOffsets[0] < 0)
+    return std::nullopt;
+
+  MemRefType baseType = dyn_cast<MemRefType>(rc.getSource().getType());
+  SmallVector<int64_t> indices(baseType.getRank(), 0);
----------------
ioghiban wrote:

TODO: `indices` -> `idxs`

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