[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)
Andrzej Warzyński
llvmlistbot at llvm.org
Fri Jun 19 07:53:03 PDT 2026
================
@@ -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;
+ std::optional<SmallVector<int64_t>> staticOffsetIndices;
+ std::optional<unsigned> dynamicOffsetDim;
+};
+
+/// Builds the index mapping needed to replace a copy into a reinterpret_cast
+/// strided memref with scalar stores into the reinterpret_cast base.
///
-/// // 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>
+/// Examples that return rewrite info:
///
-/// // Random strides
-/// memref.reinterpret_cast %buff to offset: [%off],
-/// sizes: [1, 1], strides: [10, 100]
-/// : memref<2x1xf32, strided<[10, 100]>>
-/// to memref<1x1xf32>
+/// // Scalar-shaped copy. There are no copied non-unit dimensions, so dynamic
+/// // strides in the strided memref do not affect index mapping.
+/// copy memref<1 x ... x 1 x f32>
+/// to reinterpret_cast memref<base-shape>
+/// to memref<1 x ... x 1 x f32, strided<[?, ..., ?], offset: ?>>
///
-/// // Rank-1 case
-/// memref.reinterpret_cast %buf to offset: [%off],
-/// sizes: [1], strides: [1]
-/// : memref<8xi32> to memref<1xi32>
+/// // Effectively-1D copy. The single non-unit strided memref dimension is
+/// // mapped to an identity-layout base dimension by its static stride.
+/// copy memref<1 x ... x N x ... x 1 x f32>
+/// to reinterpret_cast memref<base-shape>
+/// to memref<1 x ... x N x ... x 1 x f32, strided<[..., S, ...]>>
///
-/// Examples that return false:
+/// // Multidimensional copy with static offset. Each non-unit strided memref
+/// // dimension is mapped independently by its static stride.
+/// copy memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32>
+/// to reinterpret_cast memref<base-shape>
+/// to memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32,
+/// strided<[..., S_0, ..., S_1, ...], offset: O>>
///
-/// // More non-unit dims
-/// memref.reinterpret_cast %buff to offset: [%off],
-/// sizes: [1, 1, 1], strides: [1, 1, 1]
-/// : memref<1x2x8xi32> to memref<1x1x1xi32>
+/// Examples that return no info:
///
-/// // View is not scalar (size != 1)
-/// memref.reinterpret_cast %buff to offset: [%off],
-/// sizes: [2, 1], strides: [1, 1]
-/// : memref<1x2xf32> to memref<2x1xf32>
+/// // Dynamic stride on a copied strided memref dimension.
+/// copy memref<1xNxf32>
+/// to reinterpret_cast memref<1xNxMxf32>
+/// to memref<1xNxf32, strided<[?, ?]>>
///
-/// // 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<CopyFromReinterCastInfo>
+getCopyFromReinterCastInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
+ MemRefType srcType = dyn_cast<MemRefType>(op.getSource().getType());
+ MemRefType baseType = dyn_cast<MemRefType>(rc.getSource().getType());
+ MemRefType resultType = dyn_cast<MemRefType>(rc.getType());
+
+ // Ranked memref types are required to statically build load/store index
+ // lists.
+ if (!srcType || !baseType || !resultType)
+ return std::nullopt;
- // Reject non-matching ranks
- unsigned srcRank = rcInputTy.getRank();
- if (srcRank != rcOutputTy.getRank())
- return false;
+ if (srcType.getShape() != resultType.getShape())
+ return std::nullopt;
- ArrayRef<int64_t> sizes = rc.getStaticSizes();
+ // TODO: Support rank-changing reinterpret_casts by converting the
+ // strided memref indices to base indices. For example, a copy to
+ // a strided memref<2x3xf32> of base memref<6xf32> needs to linearize the
+ // strided memref indices as `i * 3 + j`, then combine that with the
+ // reinterpret_cast offset before indexing the rank-1 base memref.
+ if (baseType.getRank() != resultType.getRank())
+ return std::nullopt;
- // View must be scalar: memref<1x...x1>
- if (!llvm::all_of(rcOutputTy.getShape(),
- [](int64_t dim) { return dim == 1; }))
- return false;
+ // TODO: Support dynamic shapes with mixed size operands as loop bounds.
+ if (!(srcType.hasStaticShape() && baseType.hasStaticShape() &&
+ resultType.hasStaticShape()))
+ return std::nullopt;
- // Sizes must all be statically 1
- if (!llvm::all_of(sizes, [](int64_t size) {
- return !ShapedType::isDynamic(size) && size == 1;
- }))
- return false;
+ std::optional<SmallVector<int64_t>> baseStrides =
+ getIdentityStrides(baseType);
+ // TODO: Support non-identity source layouts by computing source strides from
+ // the layout map.
+ if (!baseStrides)
+ return std::nullopt;
- // Rank-1 special case
- if (srcRank == 1) {
- // Reject non-scalar output
- if (rcOutputTy.getDimSize(0) > 1)
- return false;
+ CopyFromReinterCastInfo info;
+ SmallVector<bool> usedBaseDims(baseType.getRank(), false);
+
+ for (auto [resultDim, resultSize] : llvm::enumerate(resultType.getShape())) {
----------------
banach-space wrote:
What does this loop compute? Could you document?
https://github.com/llvm/llvm-project/pull/203873
More information about the Mlir-commits
mailing list