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

[Andrzej Warzyński]

================
@@ -29,128 +29,268 @@ 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").
+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;
+}
+
+static std::optional<unsigned>
+findBaseDimForViewStride(MemRefType baseType, ArrayRef<int64_t> baseStrides,
+                         ArrayRef<bool> usedBaseDims, int64_t viewStride,
+                         int64_t viewSize) {
+  std::optional<unsigned> fallback;
+  for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
+    if (usedBaseDims[idx] || stride != viewStride ||
+        baseType.getDimSize(idx) < viewSize)
+      continue;
+
+    // Prefer an exact shape match. Otherwise, use the first dimension large
+    // enough to contain the copied logical vector.
+    if (baseType.getDimSize(idx) == viewSize)
+      return idx;
+    if (!fallback)
+      fallback = idx;
+  }
+  return fallback;
+}
+
+static std::optional<SmallVector<int64_t>>
+delinearizeStaticOffset(int64_t offset, MemRefType baseType,
+                        ArrayRef<int64_t> baseStrides) {
+  if (offset < 0)
+    return std::nullopt;
+
+  SmallVector<int64_t> indices(baseType.getRank(), 0);
+  int64_t remainder = offset;
+  for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
+    indices[idx] = remainder / stride;
+    if (indices[idx] >= baseType.getDimSize(idx))
+      return std::nullopt;
+    remainder %= stride;
+  }
+
+  if (remainder != 0)
+    return std::nullopt;
+  return indices;
+}
+
+static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
+  if (!type.hasStaticShape() || type.getRank() == 0)
+    return std::nullopt;
+
+  std::optional<unsigned> nonUnitDim;
+  for (auto [idx, dim] : llvm::enumerate(type.getShape())) {
+    if (dim == 1)
+      continue;
+    if (nonUnitDim)
+      return std::nullopt;
+    nonUnitDim = idx;
+  }
+  return nonUnitDim;
+}
+
+struct CopyLoopDimInfo {
+  unsigned viewDim;
+  unsigned dstLoopDim;
+  int64_t loopSize;
+};
+
+struct CopyToLoadStoreInfo {
+  SmallVector<CopyLoopDimInfo> loopDims;
+  SmallVector<int64_t> staticOffsetIndices;
+  std::optional<unsigned> dynamicOffsetDim;
+};
+
+/// Builds the index mapping needed to replace a copy into a reinterpret_cast
+/// view with scalar stores into the reinterpret_cast base.
 ///
-/// Examples that return true:
+/// Checklist:
+/// - The copy destination must be a `memref.reinterpret_cast`.
+/// - The copy source, reinterpret_cast source, and reinterpret_cast result must
+///   be ranked memrefs with static shapes.
+/// - The reinterpret_cast source/result ranks must match.
+/// - The reinterpret_cast source must have static identity layout.
+/// - Each non-unit copied view dimension must have a static stride that maps to
+///   an identity-layout base dimension.
+/// - Static offsets, dynamic only in scalar or effectively-1D copies
+///   where the offset can be used directly as one base index.
 ///
-///   // 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>
+/// Examples that return true:
 ///
-///   // Column-major slice (first dim is non-unit)
-///   memref.reinterpret_cast %buff to offset: [%off],
-///     sizes: [1, 1], strides: [1, 1]
-///     : memref<2x1xf32> to memref<1x1xf32>
+///   // Scalar-shaped copy. There are no copied non-unit dimensions, so
+///   // dynamic strides in the scalar view do not affect index mapping.
+///   copy memref<1x...x1xf32>
+///     to reinterpret_cast memref<base-shape>
+///       to memref<1x...x1xf32, strided<[?, ..., ?], offset: ?>>
 ///
-///   // Random strides
-///   memref.reinterpret_cast %buff to offset: [%off],
-///     sizes: [1, 1], strides: [10, 100]
-///     : memref<2x1xf32, strided<[10, 100]>>
-///         to memref<1x1xf32>
+///   // Effectively-1D copy. The single non-unit view dimension is mapped to
+///   // an identity-layout base dimension by its static stride.
+///   copy memref<1x...xNx...x1xf32>
+///     to reinterpret_cast memref<base-shape>
+///       to memref<1x...xNx...x1xf32, strided<[..., S, ...]>>
 ///
-///   // Rank-1 case
-///   memref.reinterpret_cast %buf to offset: [%off],
-///     sizes: [1], strides: [1]
-///     : memref<8xi32> to memref<1xi32>
+///   // Multidimensional copy with static offset. Each non-unit view dimension
+///   // is mapped independently by its static stride.
+///   copy memref<1x...xNx...xKx...x1xf32>
+///     to reinterpret_cast memref<base-shape>
+///       to memref<1x...xNx...xKx...x1xf32,
+///                 strided<[..., S0, ..., S1, ...], offset: O>>
 ///
 /// Examples that return false:
 ///
-///   // More non-unit dims
-///   memref.reinterpret_cast %buff to offset: [%off],
-///     sizes: [1, 1, 1], strides: [1, 1, 1]
-///     : memref<1x2x8xi32> to memref<1x1x1xi32>
+///   // Dynamic stride on a copied view dimension.
+///   copy memref<1xNxf32>
+///     to reinterpret_cast memref<1xNxMxf32>
+///       to memref<1xNxf32, strided<[?, ?]>>
 ///
-///   // View is not scalar (size != 1)
-///   memref.reinterpret_cast %buff to offset: [%off],
-///     sizes: [2, 1], strides: [1, 1]
-///     : memref<1x2xf32> to memref<2x1xf32>
-///
-///   // Base has non-identity layout
-///   %buff = memref.alloc() : memref<1x2xf32, strided<[1, 3]>>
-///   memref.reinterpret_cast %buff to offset: [%off],
-///     sizes: [1, 1], strides: [1, 1]
-///     : memref<1x2xf32, strided<[1, 3]>> to memref<1x1xf32>
-static bool isScalarSlice(memref::ReinterpretCastOp rc) {
-  auto rcInputTy = dyn_cast<MemRefType>(rc.getSource().getType());
-  auto rcOutputTy = dyn_cast<MemRefType>(rc.getType());
-
-  // Reject strided base - logic for computing linear idx is TODO
-  if (!rcInputTy.getLayout().isIdentity())
-    return false;
+///   // Multidimensional copy with dynamic linear offset.
+///   copy memref<1xNxKxf32>
+///     to reinterpret_cast memref<1xNxMxf32>
+///       to memref<1xNxKxf32, strided<[N*M, M, 1], offset: ?>>
+static std::optional<CopyToLoadStoreInfo>
+getCopyToLoadStoreInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
----------------
banach-space wrote:

Shouldn't this be called `getCopyFromReinterCastInfo` istead?

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