[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)
ioana ghiban
llvmlistbot at llvm.org
Wed Jul 8 05:15:06 PDT 2026
https://github.com/ioghiban updated https://github.com/llvm/llvm-project/pull/203873
>From 5f862cac4e31c6a1481525fbf56f1abedfda06fe Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Fri, 12 Jun 2026 20:00:44 +0200
Subject: [PATCH 01/14] [memref] Support non-scalar copies in reinterpret_cast
elision
---
.../mlir/Dialect/MemRef/Transforms/Passes.td | 3 +
.../Dialect/MemRef/Transforms/CMakeLists.txt | 1 +
.../Transforms/ElideReinterpretCast.cpp | 455 +++++++++++++-----
.../MemRef/elide-reinterpret-cast.mlir | 258 +++++++++-
4 files changed, 577 insertions(+), 140 deletions(-)
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
index 915667190a3d3..7b54be85db340 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
@@ -19,6 +19,9 @@ def ElideReinterpretCastPass : Pass<"memref-elide-reinterpret-cast"> {
operations to obtain compatible shapes with equivalent ops that operate on
compatible shapes directly. This simplifies conversion to EmitC.
}];
+ let dependentDialects = [
+ "arith::ArithDialect", "memref::MemRefDialect", "scf::SCFDialect"
+ ];
}
def ExpandOpsPass : Pass<"memref-expand"> {
diff --git a/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt b/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
index 1c5e07f89b338..c2ead38d644da 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
@@ -38,6 +38,7 @@ add_mlir_dialect_library(MLIRMemRefTransforms
MLIRMemRefUtils
MLIRNVGPUDialect
MLIRPass
+ MLIRSCFDialect
MLIRTensorDialect
MLIRTransforms
MLIRValueBoundsOpInterface
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 41dad1384da75..47aac6b64d092 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -11,10 +11,10 @@
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"
-#include "llvm/ADT/Repeated.h"
#include <cassert>
#include <optional>
@@ -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) {
+ MemRefType srcType = dyn_cast<MemRefType>(op.getSource().getType());
+ MemRefType baseType = dyn_cast<MemRefType>(rc.getSource().getType());
+ MemRefType viewType = dyn_cast<MemRefType>(rc.getType());
+ // TODO: Support unranked copy sources or reinterpret_cast sources.
+ if (!srcType || !baseType || !viewType)
+ return std::nullopt;
- // Reject non-matching ranks
- unsigned srcRank = rcInputTy.getRank();
- if (srcRank != rcOutputTy.getRank())
- return false;
+ // TODO: Support rank-changing reinterpret_casts by converting the
+ // destination view indices to base indices. For example, a copy to a
+ // memref<2x3xf32> view of memref<6xf32> needs to linearize the view indices
+ // as `i * 3 + j`, then combine that with the reinterpret_cast offset before
+ // indexing the rank-1 base memref.
+ if (baseType.getRank() != viewType.getRank())
+ return std::nullopt;
- ArrayRef<int64_t> sizes = rc.getStaticSizes();
+ // TODO: Support dynamic shapes with mixed size operands as loop bounds.
+ if (!(srcType.hasStaticShape() && baseType.hasStaticShape() &&
+ viewType.hasStaticShape()))
+ return std::nullopt;
- // View must be scalar: memref<1x...x1>
- if (!llvm::all_of(rcOutputTy.getShape(),
- [](int64_t dim) { return dim == 1; }))
- return false;
+ assert(srcType.getShape() == viewType.getShape() &&
+ "copy source and destination are expected to have the same shape");
- // Sizes must all be statically 1
- if (!llvm::all_of(sizes, [](int64_t size) {
- return !ShapedType::isDynamic(size) && size == 1;
- }))
- return false;
+ // Store indices are formed in the reinterpret_cast source layout.
+ std::optional<SmallVector<int64_t>> baseStrides =
+ getIdentityStrides(baseType);
+ // TODO: Support non-identity reinterpret_cast source layouts by using the
+ // source layout strides as base strides.
+ if (!baseStrides)
+ return std::nullopt;
- // Rank-1 special case
- if (srcRank == 1) {
- // Reject non-scalar output
- if (rcOutputTy.getDimSize(0) > 1)
- return false;
+ CopyToLoadStoreInfo info;
+ SmallVector<bool> usedBaseDims(baseType.getRank(), false);
+
+ // Non-unit view dimensions become loop dimensions in the scalar rewrite.
+ for (auto [viewDim, viewSize] : llvm::enumerate(viewType.getShape())) {
+ if (viewSize == 1)
+ continue;
+
+ // TODO: Support dynamic strides on copied view dimensions.
+ if (ShapedType::isDynamic(rc.getStaticStrides()[viewDim]))
+ return std::nullopt;
+
+ std::optional<unsigned> dstLoopDim =
+ findBaseDimForViewStride(baseType, *baseStrides, usedBaseDims,
+ rc.getStaticStrides()[viewDim], viewSize);
+ assert(dstLoopDim &&
+ "static reinterpret_cast stride must map to an identity base "
+ "dimension");
+
+ usedBaseDims[*dstLoopDim] = true;
+ info.loopDims.push_back(
+ CopyLoopDimInfo{static_cast<unsigned>(viewDim), *dstLoopDim, viewSize});
+ }
+
+ ArrayRef<int64_t> staticOffsets = rc.getStaticOffsets();
+ assert(staticOffsets.size() == 1 && "Expecting single offset");
+ if (!ShapedType::isDynamic(staticOffsets[0])) {
+ // Static offsets are converted to base indices.
+ std::optional<SmallVector<int64_t>> offsetIndices =
+ delinearizeStaticOffset(staticOffsets[0], baseType, *baseStrides);
+ assert(offsetIndices &&
+ "static reinterpret_cast offset must delinearize to in-bounds base "
+ "indices");
+
+ for (const CopyLoopDimInfo &loopDim : info.loopDims) {
+ assert((*offsetIndices)[loopDim.dstLoopDim] + loopDim.loopSize <=
+ baseType.getDimSize(loopDim.dstLoopDim) &&
+ "reinterpret_cast metadata describes an invalid accessible "
+ "region");
+ }
+ info.staticOffsetIndices = std::move(*offsetIndices);
+ return info;
}
- int nonUnitCount =
- std::count_if(rcInputTy.getShape().begin(), rcInputTy.getShape().end(),
- [](int dim) { return dim != 1; });
- return nonUnitCount == 1;
+ // Dynamic offsets are kept only when they can be used as a single base index.
+ // TODO: Support multidimensional dynamic offsets with div/mod
+ // delinearization.
+ if (info.loopDims.size() > 1)
+ return std::nullopt;
+
+ if (info.loopDims.empty()) {
+ // TODO: Support scalar dynamic offsets into bases with multiple non-unit
+ // dimensions, and all-unit bases with a provably zero offset.
+ std::optional<unsigned> nonUnitDim = getSingleNonUnitDim(baseType);
+ if (!nonUnitDim)
+ return std::nullopt;
+
+ info.dynamicOffsetDim = *nonUnitDim;
+ return info;
+ }
+
+ unsigned dstLoopDim = info.loopDims.front().dstLoopDim;
+ info.dynamicOffsetDim =
+ (*baseStrides)[dstLoopDim] == 1 ? dstLoopDim : baseStrides->size() - 1;
+ return info;
}
-/// Rewrites `memref.copy` of a 1-element MemRef as a scalar load-store pair
+/// Rewrites supported copy operations through `memref.reinterpret_cast` to
+/// scalar load/store operations.
///
-/// The pattern matches a reinterpret_cast that creates a scalar view
-/// (`sizes = [1, ..., 1]`) into a memref with a single non-unit dimension.
-/// Since the view contains only one element, the accessed address is
-/// determined solely by the base pointer and the offset.
+/// // BEFORE (scalar copy)
+/// %view = memref.reinterpret_cast %dst
+/// to offset: [O], sizes: [1, ..., 1], strides: [...]
+/// memref.copy %src, %view
///
-/// Two layouts are supported:
-/// * row-major slice (stride pattern [N, ..., 1])
-/// * column-major slice (stride pattern [1, ..., N])
+/// // AFTER
+/// %v = memref.load %src[0, ..., 0]
+/// memref.store %v, %dst[delinearized(O)]
///
-/// BEFORE (row-major slice)
-/// %view = memref.reinterpret_cast %base
-/// to offset: [%off], sizes: [1, ..., 1], strides: [N, ..., 1]
-/// : memref<1x...xNxf32>
-/// to memref<1x...x1xf32, strided<[N, ..., 1], offset: ?>>
+/// // BEFORE (effectively-1D copy)
+/// %view = memref.reinterpret_cast %dst
+/// to offset: [O], sizes: [1, N, 1], strides: [...]
/// memref.copy %src, %view
-/// : memref<1x...x1xf32>
-/// to memref<1x...x1xf32, strided<[N, ..., 1], offset: ?>>
///
-/// AFTER
-/// %c0 = arith.constant 0 : index
-/// %v = memref.load %src[%c0, ..., %c0] : memref<1x...x1xf32>
-/// memref.store %v, %base[%c0, ..., %off] : memref<1x...xNxf32>
+/// // AFTER
+/// scf.for %i = 0 to N step 1 {
+/// %v = memref.load %src[0, %i, 0]
+/// memref.store %v, %dst[delinearized(O) + mapped(%i)]
+/// }
///
-/// BEFORE (column-major slice)
-/// %view = memref.reinterpret_cast %base
-/// to offset: [%off], sizes: [1, ..., 1], strides: [1, ..., N]
-/// : memref<Nx...x1xf32>
-/// to memref<1x...x1xf32, strided<[1, ..., N], offset: ?>>
+/// // BEFORE (multidimensional copy with static offset)
+/// %view = memref.reinterpret_cast %dst
+/// to offset: [O], sizes: [1, N, K], strides: [...]
/// memref.copy %src, %view
-/// : memref<1x...x1xf32>
-/// to memref<1x...x1xf32, strided<[1, ..., N], offset: ?>>
///
-/// AFTER
-/// %c0 = arith.constant 0 : index
-/// %v = memref.load %src[%c0, ..., %c0] : memref<1x...x1xf32>
-/// memref.store %v, %base[%off, ..., %c0] : memref<Nx...x1xf32>
-struct CopyToScalarLoadAndStore : public OpRewritePattern<memref::CopyOp> {
+/// // AFTER
+/// scf.for %i = 0 to N step 1 {
+/// scf.for %j = 0 to K step 1 {
+/// %v = memref.load %src[0, %i, %j]
+/// memref.store %v, %dst[delinearized(O) + mapped(%i, %j)]
+/// }
+/// }
+struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
public:
using OpRewritePattern::OpRewritePattern;
@@ -162,38 +302,112 @@ struct CopyToScalarLoadAndStore : public OpRewritePattern<memref::CopyOp> {
return rewriter.notifyMatchFailure(
op, "target is not a memref.reinterpret_cast");
- if (!isScalarSlice(rc))
+ std::optional<CopyToLoadStoreInfo> copyInfo =
+ getCopyToLoadStoreInfo(op, rc);
+ if (!copyInfo)
return rewriter.notifyMatchFailure(
- op, "reinterpret_cast does not match scalar slice");
+ op, "reinterpret_cast does not match scalar or loop copy region");
Location loc = op.getLoc();
-
Value src = op.getSource();
Value dst = rc.getSource();
- auto dstType = cast<MemRefType>(dst.getType());
- unsigned dstRank = dstType.getRank();
+ MemRefType srcType = cast<MemRefType>(src.getType());
+ MemRefType dstType = cast<MemRefType>(dst.getType());
Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ Value one;
+ // Reuse common index constants across bounds, steps, and static offsets.
+ // Keep `%c1` lazy so scalar copies without loops do not create an unused
+ // loop-step constant.
+ auto getOrCreateIndexConstant = [&](int64_t value) -> Value {
+ if (value == 0)
+ return zero;
+ if (value == 1) {
+ if (!one)
+ one = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ return one;
+ }
+ return arith::ConstantIndexOp::create(rewriter, loc, value);
+ };
+
+ // Materialize all loop bounds before building the loop nest. Otherwise an
+ // inner-loop bound may be created inside an outer loop body.
+ SmallVector<Value> upperBounds;
+ upperBounds.reserve(copyInfo->loopDims.size());
+ for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims)
+ upperBounds.push_back(getOrCreateIndexConstant(loopDim.loopSize));
+
+ SmallVector<Value> baseStoreIndices(dstType.getRank(), zero);
+ // Static offsets were already delinearized into base indices. Materialize
+ // the non-zero starting indices before creating loop bodies.
+ if (!copyInfo->staticOffsetIndices.empty()) {
+ for (auto [idx, offset] :
+ llvm::enumerate(copyInfo->staticOffsetIndices)) {
+ if (offset == 0)
+ continue;
+ baseStoreIndices[idx] = getOrCreateIndexConstant(offset);
+ }
+ } else if (copyInfo->dynamicOffsetDim) {
+ // Supported dynamic offsets are used directly in exactly one base
+ // dimension selected by getCopyToLoadStoreInfo.
+ SmallVector<OpFoldResult> offsets = rc.getMixedOffsets();
+ assert(offsets.size() == 1 && "Expecting single offset");
+ baseStoreIndices[*copyInfo->dynamicOffsetDim] =
+ getValueOrCreateConstantIndexOp(rewriter, loc, offsets[0]);
+ }
- auto srcType = cast<MemRefType>(src.getType());
- Repeated<Value> loadIndices(srcType.getRank(), zero);
- auto offsets = rc.getMixedOffsets();
- assert(offsets.size() == 1 && "Expecting single offset");
- OpFoldResult offset = offsets[0];
- Value storeOffset = getValueOrCreateConstantIndexOp(rewriter, loc, offset);
- unsigned offsetDim = dstType.getDimSize(0) == 1 ? dstRank - 1 : 0;
- SmallVector<Value> storeIndices(dstRank, zero);
- storeIndices[offsetDim] = storeOffset;
- // If the only user of `rc` is the current Op (which is about to be erased),
- // we can safely erase it.
- if (rcOutput.hasOneUse())
- rewriter.eraseOp(rc);
+ // Scope for OpBuilder::InsertionGuard.
+ {
+ OpBuilder::InsertionGuard guard(rewriter);
+ Value step;
+ if (!upperBounds.empty())
+ step = getOrCreateIndexConstant(1);
+
+ SmallVector<Value> loopIvs;
+ loopIvs.reserve(copyInfo->loopDims.size());
+
+ // Build one nested loop per non-unit copied view dimension.
+ for (Value upperBound : upperBounds) {
+ scf::ForOp loop =
+ scf::ForOp::create(rewriter, loc, zero, upperBound, step);
+ loopIvs.push_back(loop.getInductionVar());
+ rewriter.setInsertionPointToStart(loop.getBody());
+ }
+
+ // Load indices are zero except for copied view dimensions, which use the
+ // corresponding loop induction variables.
+ SmallVector<Value> loadIndices(srcType.getRank(), zero);
+ unsigned loopIndex = 0;
+ for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims)
+ loadIndices[loopDim.viewDim] = loopIvs[loopIndex++];
+
+ // Store indices start from the offset-derived base indices. Add each loop
+ // IV to the mapped base dimension.
+ SmallVector<Value> storeIndices(baseStoreIndices);
+ loopIndex = 0;
+ for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims) {
+ Value iv = loopIvs[loopIndex++];
+ if (storeIndices[loopDim.dstLoopDim] == zero) {
+ storeIndices[loopDim.dstLoopDim] = iv;
+ } else {
+ storeIndices[loopDim.dstLoopDim] = arith::AddIOp::create(
+ rewriter, loc, storeIndices[loopDim.dstLoopDim], iv);
+ }
+ }
- Value val = memref::LoadOp::create(rewriter, loc, src, loadIndices);
- memref::StoreOp::create(rewriter, loc, val, dst, storeIndices);
+ // Emit the scalar load/store at the innermost loop body, or directly at
+ // the original copy location for scalar copies.
+ Value val = memref::LoadOp::create(rewriter, loc, src, loadIndices);
+ memref::StoreOp::create(rewriter, loc, val, dst, storeIndices);
+ }
+ // If the only user of `rc` is the current Op (which is about to be erased),
+ // we can safely erase it.
+ bool eraseRc = rcOutput.hasOneUse();
rewriter.eraseOp(op);
+ if (eraseRc)
+ rewriter.eraseOp(rc);
return success();
}
};
@@ -466,7 +680,7 @@ struct ElideReinterpretCastPass
auto rc = op.getTarget().getDefiningOp<memref::ReinterpretCastOp>();
if (!rc)
return true;
- return !isScalarSlice(rc);
+ return !getCopyToLoadStoreInfo(op, rc);
});
target.addDynamicallyLegalOp<memref::LoadOp>([](memref::LoadOp op) {
auto rc = op.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
@@ -474,7 +688,8 @@ struct ElideReinterpretCastPass
return true;
return !isPureRankExpansionOrCollapsingRC(rc);
});
- target.addLegalDialect<arith::ArithDialect, memref::MemRefDialect>();
+ target.addLegalDialect<arith::ArithDialect, memref::MemRefDialect,
+ scf::SCFDialect>();
if (failed(applyPartialConversion(getOperation(), target,
std::move(patterns))))
signalPassFailure();
@@ -485,6 +700,6 @@ struct ElideReinterpretCastPass
void mlir::memref::populateElideReinterpretCastPatterns(
RewritePatternSet &patterns) {
- patterns.add<CopyToScalarLoadAndStore, RewriteLoadFromReinterpretCast>(
+ patterns.add<CopyToLoadAndStore, RewriteLoadFromReinterpretCast>(
patterns.getContext());
}
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 61b6d480ce7a0..90431fb507c1e 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -19,9 +19,8 @@ func.func private @concat_zero_offset(%src : memref<1x1xf32>,
/// Ensure copy was replaced
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0_0]]] : memref<1x108xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
memref.copy %src, %reinterpret_cast
: memref<1x1xf32> to memref<1x1xf32>
return
@@ -85,14 +84,14 @@ func.func private @concat_strided(%src : memref<1x1xf32>,
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0_0]]] : memref<1x108xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
memref.copy %src, %reinterpret_cast
: memref<1x1xf32> to memref<1x1xf32, strided<[107, 2]>>
return
}
+// Dynamic strides are irrelevant because all view indices are zero.
// CHECK-LABEL: func.func private @concat_dynamic_stride(
// CHECK-SAME: %[[STR0:[A-Za-z][A-Za-z0-9-]*]]: index
// CHECK-SAME: %[[STR1:[A-Za-z][A-Za-z0-9-]*]]: index
@@ -108,10 +107,9 @@ func.func private @concat_dynamic_stride(%stride0: index,
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
/// Dynamic offset used in store
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0_0]]] : memref<1x108xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
memref.copy %src, %reinterpret_cast
: memref<1x1xf32>
to memref<1x1xf32, strided<[?, ?]>>
@@ -129,9 +127,8 @@ func.func private @concat_rank1(%src : memref<1xf32>, %dst : memref<108xf32>) {
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]]] : memref<1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0_0]]] : memref<108xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]]] : memref<108xf32>
memref.copy %src, %reinterpret_cast
: memref<1xf32> to memref<1xf32>
return
@@ -149,14 +146,165 @@ func.func private @concat_rank3(%src : memref<1x1x1xf32>,
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0_0]]] : memref<1x1x108xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x108xf32>
memref.copy %src, %reinterpret_cast
: memref<1x1x1xf32> to memref<1x1x1xf32>
return
}
+// CHECK-LABEL: func.func private @concat_0d(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
+func.func private @concat_0d(
+ %src : memref<1x1x1xf32>, %dst : memref<1x33x42xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 1, 1], strides: [1, 1, 1]
+ : memref<1x33x42xf32>
+ to memref<1x1x1xf32>
+ // CHECK-NOT: memref.copy
+ // CHECK: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x33x42xf32>
+ memref.copy %src, %reinterpret_cast
+ : memref<1x1x1xf32> to memref<1x1x1xf32>
+ return
+}
+
+// CHECK-LABEL: func.func private @concat_1d_vector_zero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x33x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
+func.func private @concat_1d_vector_zero_offset(
+ %src : memref<1x33x1xf32>, %dst : memref<1x33x42xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 33, 1], strides: [1386, 42, 1]
+ : memref<1x33x42xf32>
+ to memref<1x33x1xf32, strided<[1386, 42, 1]>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C33:.*]] = arith.constant 33 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x42xf32>
+ // CHECK: }
+ memref.copy %src, %reinterpret_cast
+ : memref<1x33x1xf32>
+ to memref<1x33x1xf32, strided<[1386, 42, 1]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @concat_1d_vector_nonzero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x33x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
+func.func private @concat_1d_vector_nonzero_offset(
+ %src : memref<1x33x1xf32>, %dst : memref<1x33x42xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [41], sizes: [1, 33, 1], strides: [1386, 42, 1]
+ : memref<1x33x42xf32>
+ to memref<1x33x1xf32, strided<[1386, 42, 1], offset: 41>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C33:.*]] = arith.constant 33 : index
+ // CHECK-DAG: %[[C41:.*]] = arith.constant 41 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C41]]] : memref<1x33x42xf32>
+ // CHECK: }
+ memref.copy %src, %reinterpret_cast
+ : memref<1x33x1xf32>
+ to memref<1x33x1xf32, strided<[1386, 42, 1], offset: 41>>
+ return
+}
+
+// CHECK-LABEL: func.func private @concat_1d_vector_dynamic_offset_same_dim(
+// CHECK-SAME: %[[OFF:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<42xf32>
+func.func private @concat_1d_vector_dynamic_offset_same_dim(
+ %offset : index, %src : memref<4xf32>, %dst : memref<42xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [%offset], sizes: [4], strides: [1]
+ : memref<42xf32> to memref<4xf32, strided<[1], offset: ?>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX]] : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[DST_IDX]]] : memref<42xf32>
+ // CHECK: }
+ memref.copy %src, %reinterpret_cast
+ : memref<4xf32> to memref<4xf32, strided<[1], offset: ?>>
+ return
+}
+
+// CHECK-LABEL: func.func private @concat_1d_vector_dynamic_offset_separate_dim(
+// CHECK-SAME: %[[OFF:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<1x33x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
+func.func private @concat_1d_vector_dynamic_offset_separate_dim(
+ %offset : index, %src : memref<1x33x1xf32>,
+ %dst : memref<1x33x42xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [%offset], sizes: [1, 33, 1], strides: [1386, 42, 1]
+ : memref<1x33x42xf32>
+ to memref<1x33x1xf32, strided<[1386, 42, 1], offset: ?>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C33:.*]] = arith.constant 33 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[OFF]]] : memref<1x33x42xf32>
+ // CHECK: }
+ memref.copy %src, %reinterpret_cast
+ : memref<1x33x1xf32>
+ to memref<1x33x1xf32, strided<[1386, 42, 1], offset: ?>>
+ return
+}
+
+// CHECK-LABEL: func.func private @concat_2d_vector_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x33x4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
+func.func private @concat_2d_vector_offset(
+ %src : memref<1x33x4xf32>, %dst : memref<1x33x42xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [16], sizes: [1, 33, 4], strides: [1386, 42, 1]
+ : memref<1x33x42xf32>
+ to memref<1x33x4xf32, strided<[1386, 42, 1], offset: 16>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C33:.*]] = arith.constant 33 : index
+ // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
+ // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
+ // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[C16]], %[[IDX1]] : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x33x4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[DST_IDX]]] : memref<1x33x42xf32>
+ // CHECK: }
+ // CHECK: }
+ memref.copy %src, %reinterpret_cast
+ : memref<1x33x4xf32>
+ to memref<1x33x4xf32, strided<[1386, 42, 1], offset: 16>>
+ return
+}
+
//===----------------------------------------------------------------------===//
// Negative tests (must NOT rewrite)
//===----------------------------------------------------------------------===//
@@ -179,8 +327,8 @@ func.func private @negative_concat_strided_base(%src: memref<1x1xf32>,
return
}
-// CHECK-LABEL: func.func private @negative_rank_change(
-func.func private @negative_rank_change(%src : memref<2x3xf32>,
+// CHECK-LABEL: func.func private @negative_concat_rank_change(
+func.func private @negative_concat_rank_change(%src : memref<2x3xf32>,
%dst : memref<6xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -195,19 +343,90 @@ func.func private @negative_rank_change(%src : memref<2x3xf32>,
return
}
-// CHECK-LABEL: func.func private @negative_concat_multiple_non_unit_dims(
-func.func private @negative_concat_multiple_non_unit_dims(
- %src : memref<1x1xf32>, %dst : memref<2x108xf32>) {
+// CHECK-LABEL: func.func private @negative_concat_dynamic_copy_source_shape(
+func.func private @negative_concat_dynamic_copy_source_shape(%src : memref<?xf32>,
+ %dst : memref<4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 1], strides: [1, 1]
- : memref<2x108xf32>
- to memref<1x1xf32>
+ to offset: [0], sizes: [4], strides: [1]
+ : memref<4xf32> to memref<4xf32>
+
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %reinterpret_cast
- : memref<1x1xf32> to memref<1x1xf32>
+ : memref<?xf32> to memref<4xf32>
+ return
+}
+
+// CHECK-LABEL: func.func private @negative_concat_dynamic_rc_shapes(
+func.func private @negative_concat_dynamic_rc_shapes(%dim : index,
+ %src : memref<4xf32>, %dst : memref<?xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [%dim], strides: [1]
+ : memref<?xf32> to memref<?xf32, strided<[1]>>
+
+ // CHECK: memref.copy %arg1, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %reinterpret_cast
+ : memref<4xf32> to memref<?xf32, strided<[1]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @negative_concat_dynamic_offset_multi_dim_base(
+func.func private @negative_concat_dynamic_offset_multi_dim_base(
+ %offset : index, %src : memref<1x1xf32>, %dst : memref<4x8xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [%offset], sizes: [1, 1], strides: [8, 1]
+ : memref<4x8xf32> to memref<1x1xf32, strided<[8, 1], offset: ?>>
+
+ // CHECK: memref.copy %arg1, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %reinterpret_cast
+ : memref<1x1xf32>
+ to memref<1x1xf32, strided<[8, 1], offset: ?>>
+ return
+}
+
+// CHECK-LABEL: func.func private @negative_concat_2d_dynamic_offset(
+func.func private @negative_concat_2d_dynamic_offset(
+ %offset : index, %src : memref<1x33x4xf32>,
+ %dst : memref<1x33x42xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [%offset], sizes: [1, 33, 4], strides: [1386, 42, 1]
+ : memref<1x33x42xf32>
+ to memref<1x33x4xf32, strided<[1386, 42, 1], offset: ?>>
+
+ // CHECK: memref.copy %arg1, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %reinterpret_cast
+ : memref<1x33x4xf32>
+ to memref<1x33x4xf32, strided<[1386, 42, 1], offset: ?>>
+ return
+}
+
+/// Non-unit copied dimension needs stride-based address computation.
+// CHECK-LABEL: func.func private @negative_concat_dynamic_rc_stride(
+func.func private @negative_concat_dynamic_rc_stride(%stride : index,
+ %src : memref<1x33x1xf32>, %dst : memref<1x33x42xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 33, 1], strides: [1386, %stride, 1]
+ : memref<1x33x42xf32>
+ to memref<1x33x1xf32, strided<[1386, ?, 1]>>
+
+ // CHECK: memref.copy %arg1, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %reinterpret_cast
+ : memref<1x33x1xf32>
+ to memref<1x33x1xf32, strided<[1386, ?, 1]>>
return
}
@@ -222,7 +441,6 @@ func.func private @negative_plain_copy(%src : memref<1x1xf32>,
return
}
-
// -----
//===----------------------------------------------------------------------===//
>From d355a9b2dcfedb56f62a12eadf20de1d9c8b79a9 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Wed, 17 Jun 2026 17:10:59 +0200
Subject: [PATCH 02/14] Address first round of comments
---
.../Transforms/ElideReinterpretCast.cpp | 268 ++++++++++--------
.../MemRef/elide-reinterpret-cast.mlir | 264 ++++++++---------
2 files changed, 264 insertions(+), 268 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 47aac6b64d092..e510ef9057b3f 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -15,6 +15,7 @@
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"
+#include <array>
#include <cassert>
#include <optional>
@@ -29,6 +30,10 @@ using namespace mlir;
namespace {
+/// Copy rewrite helpers
+
+/// The copy rewrite stores directly into the reinterpret_cast source, so it
+/// needs concrete base strides instead of layout-map reasoning.
static std::optional<SmallVector<int64_t>> getIdentityStrides(MemRefType type) {
if (!type.getLayout().isIdentity() || !type.hasStaticShape())
return std::nullopt;
@@ -42,19 +47,22 @@ static std::optional<SmallVector<int64_t>> getIdentityStrides(MemRefType type) {
return strides;
}
+/// Non-unit strided memref dimensions are lowered to loops over base memref
+/// dimensions. Static reinterpret_cast strides connect those dimension spaces
+/// in the currently supported cases.
static std::optional<unsigned>
-findBaseDimForViewStride(MemRefType baseType, ArrayRef<int64_t> baseStrides,
- ArrayRef<bool> usedBaseDims, int64_t viewStride,
- int64_t viewSize) {
+findBaseDimForResultStride(MemRefType baseType, ArrayRef<int64_t> baseStrides,
+ ArrayRef<bool> usedBaseDims, int64_t resultStride,
+ int64_t resultSize) {
std::optional<unsigned> fallback;
for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
- if (usedBaseDims[idx] || stride != viewStride ||
- baseType.getDimSize(idx) < viewSize)
+ if (usedBaseDims[idx] || stride != resultStride ||
+ baseType.getDimSize(idx) < resultSize)
continue;
// Prefer an exact shape match. Otherwise, use the first dimension large
// enough to contain the copied logical vector.
- if (baseType.getDimSize(idx) == viewSize)
+ if (baseType.getDimSize(idx) == resultSize)
return idx;
if (!fallback)
fallback = idx;
@@ -62,6 +70,8 @@ findBaseDimForViewStride(MemRefType baseType, ArrayRef<int64_t> baseStrides,
return fallback;
}
+/// reinterpret_cast offsets are linear element offsets, while `memref.store`
+/// needs one index per base dimension.
static std::optional<SmallVector<int64_t>>
delinearizeStaticOffset(int64_t offset, MemRefType baseType,
ArrayRef<int64_t> baseStrides) {
@@ -82,6 +92,8 @@ delinearizeStaticOffset(int64_t offset, MemRefType baseType,
return indices;
}
+/// Dynamic scalar offsets cannot be delinearized statically. They can be used
+/// directly only when the base has a single non-unit dimension to receive them.
static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
if (!type.hasStaticShape() || type.getRank() == 0)
return std::nullopt;
@@ -97,56 +109,48 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
return nonUnitDim;
}
+/// Per-dimension loop nest info.
struct CopyLoopDimInfo {
- unsigned viewDim;
- unsigned dstLoopDim;
- int64_t loopSize;
+ unsigned copyDim;
+ unsigned baseDim;
+ int64_t size;
};
-struct CopyToLoadStoreInfo {
+/// Rewrite info from reinterpret_cast layout, captured after passing legality
+/// checks.
+struct CopyFromReinterCastInfo {
SmallVector<CopyLoopDimInfo> loopDims;
- SmallVector<int64_t> staticOffsetIndices;
+ std::optional<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.
-///
-/// 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.
+/// strided memref with scalar stores into the reinterpret_cast base.
///
-/// Examples that return true:
+/// Examples that return rewrite info:
///
-/// // 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>
+/// // 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<1x...x1xf32, strided<[?, ..., ?], offset: ?>>
+/// to memref<1 x ... x 1 x f32, strided<[?, ..., ?], offset: ?>>
///
-/// // 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>
+/// // 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<1x...xNx...x1xf32, strided<[..., S, ...]>>
+/// to memref<1 x ... x N x ... x 1 x f32, strided<[..., S, ...]>>
///
-/// // Multidimensional copy with static offset. Each non-unit view dimension
-/// // is mapped independently by its static stride.
-/// copy memref<1x...xNx...xKx...x1xf32>
+/// // 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<1x...xNx...xKx...x1xf32,
-/// strided<[..., S0, ..., S1, ...], offset: O>>
+/// to memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32,
+/// strided<[..., S_0, ..., S_1, ...], offset: O>>
///
-/// Examples that return false:
+/// Examples that return no info:
///
-/// // Dynamic stride on a copied view dimension.
+/// // Dynamic stride on a copied strided memref dimension.
/// copy memref<1xNxf32>
/// to reinterpret_cast memref<1xNxMxf32>
/// to memref<1xNxf32, strided<[?, ?]>>
@@ -155,32 +159,35 @@ struct CopyToLoadStoreInfo {
/// 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) {
+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 viewType = dyn_cast<MemRefType>(rc.getType());
- // TODO: Support unranked copy sources or reinterpret_cast sources.
- if (!srcType || !baseType || !viewType)
+ 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;
+
+ if (srcType.getShape() != resultType.getShape())
return std::nullopt;
// TODO: Support rank-changing reinterpret_casts by converting the
- // destination view indices to base indices. For example, a copy to a
- // memref<2x3xf32> view of memref<6xf32> needs to linearize the view indices
- // as `i * 3 + j`, then combine that with the reinterpret_cast offset before
- // indexing the rank-1 base memref.
- if (baseType.getRank() != viewType.getRank())
+ // 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;
// TODO: Support dynamic shapes with mixed size operands as loop bounds.
if (!(srcType.hasStaticShape() && baseType.hasStaticShape() &&
- viewType.hasStaticShape()))
+ resultType.hasStaticShape()))
return std::nullopt;
- assert(srcType.getShape() == viewType.getShape() &&
- "copy source and destination are expected to have the same shape");
-
- // Store indices are formed in the reinterpret_cast source layout.
+ // Map reinterpret_cast strided memref dimensions to source dimensions by
+ // stride.
std::optional<SmallVector<int64_t>> baseStrides =
getIdentityStrides(baseType);
// TODO: Support non-identity reinterpret_cast source layouts by using the
@@ -188,60 +195,68 @@ getCopyToLoadStoreInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
if (!baseStrides)
return std::nullopt;
- CopyToLoadStoreInfo info;
+ CopyFromReinterCastInfo info;
SmallVector<bool> usedBaseDims(baseType.getRank(), false);
- // Non-unit view dimensions become loop dimensions in the scalar rewrite.
- for (auto [viewDim, viewSize] : llvm::enumerate(viewType.getShape())) {
- if (viewSize == 1)
+ for (auto [resultDim, resultSize] : llvm::enumerate(resultType.getShape())) {
+ if (resultSize == 1)
continue;
- // TODO: Support dynamic strides on copied view dimensions.
- if (ShapedType::isDynamic(rc.getStaticStrides()[viewDim]))
+ // TODO: Support dynamic strides on copied dimensions.
+ if (ShapedType::isDynamic(rc.getStaticStrides()[resultDim]))
return std::nullopt;
- std::optional<unsigned> dstLoopDim =
- findBaseDimForViewStride(baseType, *baseStrides, usedBaseDims,
- rc.getStaticStrides()[viewDim], viewSize);
- assert(dstLoopDim &&
+ // Non-unit strided memref dimensions become loop dimensions in the scalar
+ // rewrite.
+ std::optional<unsigned> baseDim = findBaseDimForResultStride(
+ baseType, *baseStrides, usedBaseDims, rc.getStaticStrides()[resultDim],
+ resultSize);
+ assert(baseDim &&
"static reinterpret_cast stride must map to an identity base "
"dimension");
- usedBaseDims[*dstLoopDim] = true;
- info.loopDims.push_back(
- CopyLoopDimInfo{static_cast<unsigned>(viewDim), *dstLoopDim, viewSize});
+ usedBaseDims[*baseDim] = true;
+ info.loopDims.push_back(CopyLoopDimInfo{static_cast<unsigned>(resultDim),
+ *baseDim, resultSize});
}
- ArrayRef<int64_t> staticOffsets = rc.getStaticOffsets();
- assert(staticOffsets.size() == 1 && "Expecting single offset");
- if (!ShapedType::isDynamic(staticOffsets[0])) {
+ 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");
+ if (!ShapedType::isDynamic(offsets[0])) {
// Static offsets are converted to base indices.
std::optional<SmallVector<int64_t>> offsetIndices =
- delinearizeStaticOffset(staticOffsets[0], baseType, *baseStrides);
+ delinearizeStaticOffset(offsets[0], baseType, *baseStrides);
assert(offsetIndices &&
"static reinterpret_cast offset must delinearize to in-bounds base "
"indices");
for (const CopyLoopDimInfo &loopDim : info.loopDims) {
- assert((*offsetIndices)[loopDim.dstLoopDim] + loopDim.loopSize <=
- baseType.getDimSize(loopDim.dstLoopDim) &&
+ assert((*offsetIndices)[loopDim.baseDim] + loopDim.size <=
+ baseType.getDimSize(loopDim.baseDim) &&
"reinterpret_cast metadata describes an invalid accessible "
"region");
}
- info.staticOffsetIndices = std::move(*offsetIndices);
+ info.staticOffsetIndices = std::move(offsetIndices);
return info;
}
// Dynamic offsets are kept only when they can be used as a single base index.
- // TODO: Support multidimensional dynamic offsets with div/mod
- // delinearization.
+ // TODO: Support dynamic offsets for copies with multiple loop dimensions by
+ // delinearizing the offset into base start indices at runtime before adding
+ // loop IVs.
if (info.loopDims.size() > 1)
return std::nullopt;
if (info.loopDims.empty()) {
- // TODO: Support scalar dynamic offsets into bases with multiple non-unit
- // dimensions, and all-unit bases with a provably zero offset.
+ // Scalar-shaped strided memrefs copy one element, so a dynamic linear
+ // offset can be used directly only when the base has exactly one non-unit
+ // dimension.
std::optional<unsigned> nonUnitDim = getSingleNonUnitDim(baseType);
+ // TODO: Support scalar dynamic offsets into bases with multiple non-unit
+ // dimensions by delinearizing the single accessed element offset at
+ // runtime.
if (!nonUnitDim)
return std::nullopt;
@@ -249,9 +264,9 @@ getCopyToLoadStoreInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
return info;
}
- unsigned dstLoopDim = info.loopDims.front().dstLoopDim;
+ unsigned baseDim = info.loopDims.front().baseDim;
info.dynamicOffsetDim =
- (*baseStrides)[dstLoopDim] == 1 ? dstLoopDim : baseStrides->size() - 1;
+ (*baseStrides)[baseDim] == 1 ? baseDim : baseStrides->size() - 1;
return info;
}
@@ -259,18 +274,18 @@ getCopyToLoadStoreInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
/// scalar load/store operations.
///
/// // BEFORE (scalar copy)
-/// %view = memref.reinterpret_cast %dst
+/// %strided = memref.reinterpret_cast %dst
/// to offset: [O], sizes: [1, ..., 1], strides: [...]
-/// memref.copy %src, %view
+/// memref.copy %src, %strided
///
/// // AFTER
/// %v = memref.load %src[0, ..., 0]
/// memref.store %v, %dst[delinearized(O)]
///
/// // BEFORE (effectively-1D copy)
-/// %view = memref.reinterpret_cast %dst
+/// %strided = memref.reinterpret_cast %dst
/// to offset: [O], sizes: [1, N, 1], strides: [...]
-/// memref.copy %src, %view
+/// memref.copy %src, %strided
///
/// // AFTER
/// scf.for %i = 0 to N step 1 {
@@ -279,9 +294,9 @@ getCopyToLoadStoreInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
/// }
///
/// // BEFORE (multidimensional copy with static offset)
-/// %view = memref.reinterpret_cast %dst
+/// %strided = memref.reinterpret_cast %dst
/// to offset: [O], sizes: [1, N, K], strides: [...]
-/// memref.copy %src, %view
+/// memref.copy %src, %strided
///
/// // AFTER
/// scf.for %i = 0 to N step 1 {
@@ -302,8 +317,8 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
return rewriter.notifyMatchFailure(
op, "target is not a memref.reinterpret_cast");
- std::optional<CopyToLoadStoreInfo> copyInfo =
- getCopyToLoadStoreInfo(op, rc);
+ std::optional<CopyFromReinterCastInfo> copyInfo =
+ getCopyFromReinterCastInfo(op, rc);
if (!copyInfo)
return rewriter.notifyMatchFailure(
op, "reinterpret_cast does not match scalar or loop copy region");
@@ -315,43 +330,51 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
MemRefType srcType = cast<MemRefType>(src.getType());
MemRefType dstType = cast<MemRefType>(dst.getType());
- Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
- Value one;
- // Reuse common index constants across bounds, steps, and static offsets.
- // Keep `%c1` lazy so scalar copies without loops do not create an unused
- // loop-step constant.
+ // Reuse common index constants across bounds, steps, and static offsets,
+ // but avoid creating them for rank-0 copies.
+ std::array<Value, 2> cachedIndexConstants;
auto getOrCreateIndexConstant = [&](int64_t value) -> Value {
- if (value == 0)
- return zero;
- if (value == 1) {
- if (!one)
- one = arith::ConstantIndexOp::create(rewriter, loc, 1);
- return one;
+ if (value == 0 || value == 1) {
+ Value &cached = cachedIndexConstants[value];
+ if (!cached)
+ cached = arith::ConstantIndexOp::create(rewriter, loc, value);
+ return cached;
}
return arith::ConstantIndexOp::create(rewriter, loc, value);
};
+ auto getZeroIndices = [&](int64_t rank) {
+ SmallVector<Value> indices;
+ indices.reserve(rank);
+ if (rank != 0)
+ indices.append(rank, getOrCreateIndexConstant(0));
+ return indices;
+ };
- // Materialize all loop bounds before building the loop nest. Otherwise an
- // inner-loop bound may be created inside an outer loop body.
+ // Create all loop bounds before building the loop nest. Otherwise an
+ // inner-loop bound can be inserted inside an outer loop body.
SmallVector<Value> upperBounds;
upperBounds.reserve(copyInfo->loopDims.size());
for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims)
- upperBounds.push_back(getOrCreateIndexConstant(loopDim.loopSize));
+ upperBounds.push_back(getOrCreateIndexConstant(loopDim.size));
- SmallVector<Value> baseStoreIndices(dstType.getRank(), zero);
- // Static offsets were already delinearized into base indices. Materialize
- // the non-zero starting indices before creating loop bodies.
- if (!copyInfo->staticOffsetIndices.empty()) {
+ SmallVector<Value> baseStoreIndices = getZeroIndices(dstType.getRank());
+ // Static offsets were already delinearized into base indices. Fill the
+ // non-zero starting indices before creating loop bodies.
+ if (copyInfo->staticOffsetIndices) {
for (auto [idx, offset] :
- llvm::enumerate(copyInfo->staticOffsetIndices)) {
+ llvm::enumerate(*copyInfo->staticOffsetIndices)) {
if (offset == 0)
continue;
baseStoreIndices[idx] = getOrCreateIndexConstant(offset);
}
- } else if (copyInfo->dynamicOffsetDim) {
+ } else {
// Supported dynamic offsets are used directly in exactly one base
- // dimension selected by getCopyToLoadStoreInfo.
+ // dimension selected by getCopyFromReinterCastInfo.
+ assert(copyInfo->dynamicOffsetDim &&
+ "expected dynamic offset dimension for dynamic offset");
SmallVector<OpFoldResult> offsets = rc.getMixedOffsets();
+ // FIXME: Despite what `getMixedOffsets` implies, `reinterpret_cast` takes
+ // only a single offset. That should be fixed at the op definition level.
assert(offsets.size() == 1 && "Expecting single offset");
baseStoreIndices[*copyInfo->dynamicOffsetDim] =
getValueOrCreateConstantIndexOp(rewriter, loc, offsets[0]);
@@ -360,27 +383,30 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
// Scope for OpBuilder::InsertionGuard.
{
OpBuilder::InsertionGuard guard(rewriter);
+ Value lowerBound;
Value step;
- if (!upperBounds.empty())
+ if (!upperBounds.empty()) {
+ lowerBound = getOrCreateIndexConstant(0);
step = getOrCreateIndexConstant(1);
+ }
SmallVector<Value> loopIvs;
loopIvs.reserve(copyInfo->loopDims.size());
- // Build one nested loop per non-unit copied view dimension.
+ // Build one nested loop per non-unit copied strided memref dimension.
for (Value upperBound : upperBounds) {
scf::ForOp loop =
- scf::ForOp::create(rewriter, loc, zero, upperBound, step);
+ scf::ForOp::create(rewriter, loc, lowerBound, upperBound, step);
loopIvs.push_back(loop.getInductionVar());
rewriter.setInsertionPointToStart(loop.getBody());
}
- // Load indices are zero except for copied view dimensions, which use the
- // corresponding loop induction variables.
- SmallVector<Value> loadIndices(srcType.getRank(), zero);
+ // Load indices are zero except for copied strided memref dimensions,
+ // which use the corresponding loop induction variables.
+ SmallVector<Value> loadIndices = getZeroIndices(srcType.getRank());
unsigned loopIndex = 0;
for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims)
- loadIndices[loopDim.viewDim] = loopIvs[loopIndex++];
+ loadIndices[loopDim.copyDim] = loopIvs[loopIndex++];
// Store indices start from the offset-derived base indices. Add each loop
// IV to the mapped base dimension.
@@ -388,11 +414,11 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
loopIndex = 0;
for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims) {
Value iv = loopIvs[loopIndex++];
- if (storeIndices[loopDim.dstLoopDim] == zero) {
- storeIndices[loopDim.dstLoopDim] = iv;
+ if (storeIndices[loopDim.baseDim] == getOrCreateIndexConstant(0)) {
+ storeIndices[loopDim.baseDim] = iv;
} else {
- storeIndices[loopDim.dstLoopDim] = arith::AddIOp::create(
- rewriter, loc, storeIndices[loopDim.dstLoopDim], iv);
+ storeIndices[loopDim.baseDim] = arith::AddIOp::create(
+ rewriter, loc, storeIndices[loopDim.baseDim], iv);
}
}
@@ -680,7 +706,7 @@ struct ElideReinterpretCastPass
auto rc = op.getTarget().getDefiningOp<memref::ReinterpretCastOp>();
if (!rc)
return true;
- return !getCopyToLoadStoreInfo(op, rc);
+ return !getCopyFromReinterCastInfo(op, rc);
});
target.addDynamicallyLegalOp<memref::LoadOp>([](memref::LoadOp op) {
auto rc = op.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 90431fb507c1e..044e12cec42af 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -5,10 +5,10 @@
// Positive tests
//===----------------------------------------------------------------------===//
-// CHECK-LABEL: func.func private @concat_zero_offset(
+// CHECK-LABEL: func.func private @copy_to_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
-func.func private @concat_zero_offset(%src : memref<1x1xf32>,
+func.func private @copy_to_strided_zero_offset(%src : memref<1x1xf32>,
%dst : memref<1x108xf32>) {
/// reinterpret_cast removed
// CHECK-NOT: memref.reinterpret_cast
@@ -23,13 +23,14 @@ func.func private @concat_zero_offset(%src : memref<1x1xf32>,
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
memref.copy %src, %reinterpret_cast
: memref<1x1xf32> to memref<1x1xf32>
+ // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @concat_nonzero_offset(
+// CHECK-LABEL: func.func private @copy_to_strided_nonzero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
-func.func private @concat_nonzero_offset(%src : memref<1x1xf32>,
+func.func private @copy_to_strided_nonzero_offset(%src : memref<1x1xf32>,
%dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -45,14 +46,15 @@ func.func private @concat_nonzero_offset(%src : memref<1x1xf32>,
memref.copy %src, %reinterpret_cast
: memref<1x1xf32>
to memref<1x1xf32, strided<[1, 1], offset: 1>>
+ // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @concat_dynamic_offset(
+// CHECK-LABEL: func.func private @copy_to_strided_dynamic_offset(
// CHECK-SAME: %[[OFF:.*]]: index
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
-func.func private @concat_dynamic_offset(%offset: index, %src : memref<1x1xf32>,
+func.func private @copy_to_strided_dynamic_offset(%offset: index, %src : memref<1x1xf32>,
%dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -69,35 +71,17 @@ func.func private @concat_dynamic_offset(%offset: index, %src : memref<1x1xf32>,
memref.copy %src, %reinterpret_cast
: memref<1x1xf32>
to memref<1x1xf32, strided<[1, 1], offset: ?>>
- return
-}
-
-// CHECK-LABEL: func.func private @concat_strided(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
-func.func private @concat_strided(%src : memref<1x1xf32>,
- %dst : memref<1x108xf32>) {
- // CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 1], strides: [107, 2]
- : memref<1x108xf32> to memref<1x1xf32, strided<[107, 2]>>
-
// CHECK-NOT: memref.copy
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
- memref.copy %src, %reinterpret_cast
- : memref<1x1xf32> to memref<1x1xf32, strided<[107, 2]>>
return
}
-// Dynamic strides are irrelevant because all view indices are zero.
-// CHECK-LABEL: func.func private @concat_dynamic_stride(
+// Dynamic strides are irrelevant because all strided memref indices are zero.
+// CHECK-LABEL: func.func private @copy_to_strided_dynamic_stride(
// CHECK-SAME: %[[STR0:[A-Za-z][A-Za-z0-9-]*]]: index
// CHECK-SAME: %[[STR1:[A-Za-z][A-Za-z0-9-]*]]: index
// CHECK-SAME: %[[SRC:[A-Za-z][A-Za-z0-9-]*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:[A-Za-z][A-Za-z0-9-]*]]: memref<1x108xf32>
-func.func private @concat_dynamic_stride(%stride0: index,
+func.func private @copy_to_strided_dynamic_stride(%stride0: index,
%stride1: index, %src : memref<1x1xf32>, %dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -113,50 +97,30 @@ func.func private @concat_dynamic_stride(%stride0: index,
memref.copy %src, %reinterpret_cast
: memref<1x1xf32>
to memref<1x1xf32, strided<[?, ?]>>
- return
-}
-
-// CHECK-LABEL: func.func private @concat_rank1(
-// CHECK-SAME: %[[SRC:.*]]: memref<1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<108xf32>
-func.func private @concat_rank1(%src : memref<1xf32>, %dst : memref<108xf32>) {
- // CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1], strides: [1]
- : memref<108xf32> to memref<1xf32>
-
// CHECK-NOT: memref.copy
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]]] : memref<1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]]] : memref<108xf32>
- memref.copy %src, %reinterpret_cast
- : memref<1xf32> to memref<1xf32>
return
}
-// CHECK-LABEL: func.func private @concat_rank3(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x1x108xf32>
-func.func private @concat_rank3(%src : memref<1x1x1xf32>,
- %dst : memref<1x1x108xf32>) {
+// CHECK-LABEL: func.func private @copy_to_strided_rank0(
+// CHECK-SAME: %[[SRC:.*]]: memref<f32>, %[[DST:.*]]: memref<f32>
+func.func private @copy_to_strided_rank0(%src : memref<f32>, %dst : memref<f32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 1, 1], strides: [1, 1, 1]
- : memref<1x1x108xf32> to memref<1x1x1xf32>
+ to offset: [0], sizes: [], strides: []
+ : memref<f32> to memref<f32>
// CHECK-NOT: memref.copy
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x108xf32>
- memref.copy %src, %reinterpret_cast
- : memref<1x1x1xf32> to memref<1x1x1xf32>
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][] : memref<f32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][] : memref<f32>
+ memref.copy %src, %reinterpret_cast : memref<f32> to memref<f32>
+ // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @concat_0d(
+// CHECK-LABEL: func.func private @copy_to_strided_0d_2d_base(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
-func.func private @concat_0d(
+func.func private @copy_to_strided_0d_2d_base(
%src : memref<1x1x1xf32>, %dst : memref<1x33x42xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -169,65 +133,68 @@ func.func private @concat_0d(
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x33x42xf32>
memref.copy %src, %reinterpret_cast
: memref<1x1x1xf32> to memref<1x1x1xf32>
+ // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @concat_1d_vector_zero_offset(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x33x1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
-func.func private @concat_1d_vector_zero_offset(
- %src : memref<1x33x1xf32>, %dst : memref<1x33x42xf32>) {
+// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_zero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_to_strided_1d_vector_zero_offset(
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 33, 1], strides: [1386, 42, 1]
- : memref<1x33x42xf32>
- to memref<1x33x1xf32, strided<[1386, 42, 1]>>
+ to offset: [0], sizes: [1, 3, 1], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1]>>
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C33:.*]] = arith.constant 33 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x42xf32>
+ // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x11xf32>
// CHECK: }
memref.copy %src, %reinterpret_cast
- : memref<1x33x1xf32>
- to memref<1x33x1xf32, strided<[1386, 42, 1]>>
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1]>>
+ // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @concat_1d_vector_nonzero_offset(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x33x1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
-func.func private @concat_1d_vector_nonzero_offset(
- %src : memref<1x33x1xf32>, %dst : memref<1x33x42xf32>) {
+// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_nonzero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_to_strided_1d_vector_nonzero_offset(
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [41], sizes: [1, 33, 1], strides: [1386, 42, 1]
- : memref<1x33x42xf32>
- to memref<1x33x1xf32, strided<[1386, 42, 1], offset: 41>>
+ to offset: [10], sizes: [1, 3, 1], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: 10>>
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C33:.*]] = arith.constant 33 : index
- // CHECK-DAG: %[[C41:.*]] = arith.constant 41 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C41]]] : memref<1x33x42xf32>
+ // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C10]]] : memref<1x3x11xf32>
// CHECK: }
memref.copy %src, %reinterpret_cast
- : memref<1x33x1xf32>
- to memref<1x33x1xf32, strided<[1386, 42, 1], offset: 41>>
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: 10>>
+ // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @concat_1d_vector_dynamic_offset_same_dim(
+// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_dynamic_offset_in_loop_dim(
// CHECK-SAME: %[[OFF:.*]]: index
// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
// CHECK-SAME: %[[DST:.*]]: memref<42xf32>
-func.func private @concat_1d_vector_dynamic_offset_same_dim(
+func.func private @copy_to_strided_1d_vector_dynamic_offset_in_loop_dim(
%offset : index, %src : memref<4xf32>, %dst : memref<42xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -245,63 +212,66 @@ func.func private @concat_1d_vector_dynamic_offset_same_dim(
// CHECK: }
memref.copy %src, %reinterpret_cast
: memref<4xf32> to memref<4xf32, strided<[1], offset: ?>>
+ // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @concat_1d_vector_dynamic_offset_separate_dim(
+// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_dynamic_offset_not_in_loop_dim(
// CHECK-SAME: %[[OFF:.*]]: index
-// CHECK-SAME: %[[SRC:.*]]: memref<1x33x1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
-func.func private @concat_1d_vector_dynamic_offset_separate_dim(
- %offset : index, %src : memref<1x33x1xf32>,
- %dst : memref<1x33x42xf32>) {
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_to_strided_1d_vector_dynamic_offset_not_in_loop_dim(
+ %offset : index, %src : memref<1x3x1xf32>,
+ %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [1, 33, 1], strides: [1386, 42, 1]
- : memref<1x33x42xf32>
- to memref<1x33x1xf32, strided<[1386, 42, 1], offset: ?>>
+ to offset: [%offset], sizes: [1, 3, 1], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C33:.*]] = arith.constant 33 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[OFF]]] : memref<1x33x42xf32>
+ // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[OFF]]] : memref<1x3x11xf32>
// CHECK: }
memref.copy %src, %reinterpret_cast
- : memref<1x33x1xf32>
- to memref<1x33x1xf32, strided<[1386, 42, 1], offset: ?>>
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
+ // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @concat_2d_vector_offset(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x33x4xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
-func.func private @concat_2d_vector_offset(
- %src : memref<1x33x4xf32>, %dst : memref<1x33x42xf32>) {
+// CHECK-LABEL: func.func private @copy_to_strided_2d_vector_nonzero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_to_strided_2d_vector_nonzero_offset(
+ %src : memref<1x3x4xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [16], sizes: [1, 33, 4], strides: [1386, 42, 1]
- : memref<1x33x42xf32>
- to memref<1x33x4xf32, strided<[1386, 42, 1], offset: 16>>
+ to offset: [7], sizes: [1, 3, 4], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1], offset: 7>>
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C33:.*]] = arith.constant 33 : index
+ // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
- // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
- // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
+ // CHECK-DAG: %[[C7:.*]] = arith.constant 7 : index
+ // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
// CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
- // CHECK: %[[DST_IDX:.*]] = arith.addi %[[C16]], %[[IDX1]] : index
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x33x4xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[DST_IDX]]] : memref<1x33x42xf32>
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[C7]], %[[IDX1]] : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[DST_IDX]]] : memref<1x3x11xf32>
// CHECK: }
// CHECK: }
memref.copy %src, %reinterpret_cast
- : memref<1x33x4xf32>
- to memref<1x33x4xf32, strided<[1386, 42, 1], offset: 16>>
+ : memref<1x3x4xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1], offset: 7>>
+ // CHECK-NOT: memref.copy
return
}
@@ -309,8 +279,8 @@ func.func private @concat_2d_vector_offset(
// Negative tests (must NOT rewrite)
//===----------------------------------------------------------------------===//
-// CHECK-LABEL: func.func private @negative_concat_strided_base(
-func.func private @negative_concat_strided_base(%src: memref<1x1xf32>,
+// CHECK-LABEL: func.func private @negative_copy_to_strided_non_identity_base(
+func.func private @negative_copy_to_strided_non_identity_base(%src: memref<1x1xf32>,
%dst: memref<8x1xf32, strided<[10, 2]>>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -327,8 +297,8 @@ func.func private @negative_concat_strided_base(%src: memref<1x1xf32>,
return
}
-// CHECK-LABEL: func.func private @negative_concat_rank_change(
-func.func private @negative_concat_rank_change(%src : memref<2x3xf32>,
+// CHECK-LABEL: func.func private @negative_copy_to_strided_rank_change(
+func.func private @negative_copy_to_strided_rank_change(%src : memref<2x3xf32>,
%dst : memref<6xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -343,8 +313,8 @@ func.func private @negative_concat_rank_change(%src : memref<2x3xf32>,
return
}
-// CHECK-LABEL: func.func private @negative_concat_dynamic_copy_source_shape(
-func.func private @negative_concat_dynamic_copy_source_shape(%src : memref<?xf32>,
+// CHECK-LABEL: func.func private @negative_copy_to_strided_dynamic_copy_source_shape(
+func.func private @negative_copy_to_strided_dynamic_copy_source_shape(%src : memref<?xf32>,
%dst : memref<4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -359,8 +329,8 @@ func.func private @negative_concat_dynamic_copy_source_shape(%src : memref<?xf32
return
}
-// CHECK-LABEL: func.func private @negative_concat_dynamic_rc_shapes(
-func.func private @negative_concat_dynamic_rc_shapes(%dim : index,
+// CHECK-LABEL: func.func private @negative_copy_to_strided_dynamic_rc_shapes(
+func.func private @negative_copy_to_strided_dynamic_rc_shapes(%dim : index,
%src : memref<4xf32>, %dst : memref<?xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -375,8 +345,8 @@ func.func private @negative_concat_dynamic_rc_shapes(%dim : index,
return
}
-// CHECK-LABEL: func.func private @negative_concat_dynamic_offset_multi_dim_base(
-func.func private @negative_concat_dynamic_offset_multi_dim_base(
+// CHECK-LABEL: func.func private @negative_copy_to_strided_dynamic_offset_multi_dim_base(
+func.func private @negative_copy_to_strided_dynamic_offset_multi_dim_base(
%offset : index, %src : memref<1x1xf32>, %dst : memref<4x8xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -392,41 +362,41 @@ func.func private @negative_concat_dynamic_offset_multi_dim_base(
return
}
-// CHECK-LABEL: func.func private @negative_concat_2d_dynamic_offset(
-func.func private @negative_concat_2d_dynamic_offset(
- %offset : index, %src : memref<1x33x4xf32>,
- %dst : memref<1x33x42xf32>) {
+// CHECK-LABEL: func.func private @negative_copy_to_strided_2d_dynamic_offset(
+func.func private @negative_copy_to_strided_2d_dynamic_offset(
+ %offset : index, %src : memref<1x3x4xf32>,
+ %dst : memref<1x3x11xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [1, 33, 4], strides: [1386, 42, 1]
- : memref<1x33x42xf32>
- to memref<1x33x4xf32, strided<[1386, 42, 1], offset: ?>>
+ to offset: [%offset], sizes: [1, 3, 4], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1], offset: ?>>
// CHECK: memref.copy %arg1, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %reinterpret_cast
- : memref<1x33x4xf32>
- to memref<1x33x4xf32, strided<[1386, 42, 1], offset: ?>>
+ : memref<1x3x4xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1], offset: ?>>
return
}
/// Non-unit copied dimension needs stride-based address computation.
-// CHECK-LABEL: func.func private @negative_concat_dynamic_rc_stride(
-func.func private @negative_concat_dynamic_rc_stride(%stride : index,
- %src : memref<1x33x1xf32>, %dst : memref<1x33x42xf32>) {
+// CHECK-LABEL: func.func private @negative_copy_to_strided_dynamic_rc_stride(
+func.func private @negative_copy_to_strided_dynamic_rc_stride(%stride : index,
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 33, 1], strides: [1386, %stride, 1]
- : memref<1x33x42xf32>
- to memref<1x33x1xf32, strided<[1386, ?, 1]>>
+ to offset: [0], sizes: [1, 3, 1], strides: [33, %stride, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, ?, 1]>>
// CHECK: memref.copy %arg1, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %reinterpret_cast
- : memref<1x33x1xf32>
- to memref<1x33x1xf32, strided<[1386, ?, 1]>>
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[33, ?, 1]>>
return
}
>From 7399855abf4eaf455ae06282947c2345d41d5a41 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Fri, 19 Jun 2026 10:45:19 +0200
Subject: [PATCH 03/14] Refine helper function comments
---
.../Transforms/ElideReinterpretCast.cpp | 108 ++++++++++--------
1 file changed, 60 insertions(+), 48 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index e510ef9057b3f..95d7b1f1dc286 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -30,10 +30,11 @@ using namespace mlir;
namespace {
-/// Copy rewrite helpers
+//===----------------------------------------------------------------------===//
+// Copy Rewrite Helpers
+//===----------------------------------------------------------------------===//
-/// The copy rewrite stores directly into the reinterpret_cast source, so it
-/// needs concrete base strides instead of layout-map reasoning.
+/// 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;
@@ -47,42 +48,54 @@ static std::optional<SmallVector<int64_t>> getIdentityStrides(MemRefType type) {
return strides;
}
-/// Non-unit strided memref dimensions are lowered to loops over base memref
-/// dimensions. Static reinterpret_cast strides connect those dimension spaces
-/// in the currently supported cases.
-static std::optional<unsigned>
-findBaseDimForResultStride(MemRefType baseType, ArrayRef<int64_t> baseStrides,
- ArrayRef<bool> usedBaseDims, int64_t resultStride,
- int64_t resultSize) {
- std::optional<unsigned> fallback;
- for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
- if (usedBaseDims[idx] || stride != resultStride ||
- baseType.getDimSize(idx) < resultSize)
+/// 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;
- // Prefer an exact shape match. Otherwise, use the first dimension large
- // enough to contain the copied logical vector.
- if (baseType.getDimSize(idx) == resultSize)
- return idx;
- if (!fallback)
- fallback = idx;
+ if (!sourceDim ||
+ sourceType.getDimSize(idx) < sourceType.getDimSize(*sourceDim))
+ sourceDim = idx;
}
- return fallback;
+ return sourceDim;
}
-/// reinterpret_cast offsets are linear element offsets, while `memref.store`
-/// needs one index per base dimension.
+/// Returns source indices for a static reinterpret_cast offset.
static std::optional<SmallVector<int64_t>>
-delinearizeStaticOffset(int64_t offset, MemRefType baseType,
- ArrayRef<int64_t> baseStrides) {
- if (offset < 0)
+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(baseType.getRank(), 0);
- int64_t remainder = offset;
- for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
+ 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] >= baseType.getDimSize(idx))
+ if (indices[idx] >= sourceType.getDimSize(idx))
return std::nullopt;
remainder %= stride;
}
@@ -92,14 +105,15 @@ delinearizeStaticOffset(int64_t offset, MemRefType baseType,
return indices;
}
-/// Dynamic scalar offsets cannot be delinearized statically. They can be used
-/// directly only when the base has a single non-unit dimension to receive them.
+/// Returns the dimension whose static size is not one if it is unique.
static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
- if (!type.hasStaticShape() || type.getRank() == 0)
+ 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(type.getShape())) {
+ for (auto [idx, dim] : llvm::enumerate(shape)) {
if (dim == 1)
continue;
if (nonUnitDim)
@@ -186,12 +200,10 @@ getCopyFromReinterCastInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
resultType.hasStaticShape()))
return std::nullopt;
- // Map reinterpret_cast strided memref dimensions to source dimensions by
- // stride.
std::optional<SmallVector<int64_t>> baseStrides =
getIdentityStrides(baseType);
- // TODO: Support non-identity reinterpret_cast source layouts by using the
- // source layout strides as base strides.
+ // TODO: Support non-identity source layouts by computing source strides from
+ // the layout map.
if (!baseStrides)
return std::nullopt;
@@ -206,11 +218,11 @@ getCopyFromReinterCastInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
if (ShapedType::isDynamic(rc.getStaticStrides()[resultDim]))
return std::nullopt;
- // Non-unit strided memref dimensions become loop dimensions in the scalar
- // rewrite.
- std::optional<unsigned> baseDim = findBaseDimForResultStride(
- baseType, *baseStrides, usedBaseDims, rc.getStaticStrides()[resultDim],
- resultSize);
+ // Each copied result dimension must map by stride to a source dimension
+ // whose static size >= result dimension size.
+ std::optional<unsigned> baseDim =
+ findSourceDimForResultDim(rc, static_cast<unsigned>(resultDim),
+ baseType, *baseStrides, usedBaseDims);
assert(baseDim &&
"static reinterpret_cast stride must map to an identity base "
"dimension");
@@ -225,9 +237,9 @@ getCopyFromReinterCastInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
// only a single offset. That should be fixed at the op definition level.
assert(offsets.size() == 1 && "Expecting single offset");
if (!ShapedType::isDynamic(offsets[0])) {
- // Static offsets are converted to base indices.
+ // Static offset must delinearize to in-bounds source indices.
std::optional<SmallVector<int64_t>> offsetIndices =
- delinearizeStaticOffset(offsets[0], baseType, *baseStrides);
+ delinearizeStaticOffset(rc, baseType, *baseStrides);
assert(offsetIndices &&
"static reinterpret_cast offset must delinearize to in-bounds base "
"indices");
@@ -250,9 +262,9 @@ getCopyFromReinterCastInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
return std::nullopt;
if (info.loopDims.empty()) {
- // Scalar-shaped strided memrefs copy one element, so a dynamic linear
- // offset can be used directly only when the base has exactly one non-unit
- // dimension.
+ // Dynamic scalar offsets cannot be delinearized statically. They can be
+ // used directly only when the base has a single non-unit dimension to
+ // receive them.
std::optional<unsigned> nonUnitDim = getSingleNonUnitDim(baseType);
// TODO: Support scalar dynamic offsets into bases with multiple non-unit
// dimensions by delinearizing the single accessed element offset at
>From ec9d6848bd0b349da2436011bf7c9c6b1e15478f Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Fri, 19 Jun 2026 13:05:33 +0200
Subject: [PATCH 04/14] Remove redundant dialect dependency
---
mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
index 7b54be85db340..c6be600247696 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
@@ -20,7 +20,7 @@ def ElideReinterpretCastPass : Pass<"memref-elide-reinterpret-cast"> {
compatible shapes directly. This simplifies conversion to EmitC.
}];
let dependentDialects = [
- "arith::ArithDialect", "memref::MemRefDialect", "scf::SCFDialect"
+ "arith::ArithDialect", "scf::SCFDialect"
];
}
>From bb787ee565ccd81cdd5e29a086e2b26f80665431 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Mon, 22 Jun 2026 18:01:38 +0200
Subject: [PATCH 05/14] Address second round of comments
---
.../Transforms/ElideReinterpretCast.cpp | 247 +++++++++---------
.../MemRef/elide-reinterpret-cast.mlir | 16 ++
2 files changed, 135 insertions(+), 128 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 95d7b1f1dc286..6d806b087bdb4 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -12,6 +12,7 @@
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"
@@ -34,68 +35,87 @@ namespace {
// 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;
+/// Per-dimension loop nest info.
+struct CopyLoopDimInfo {
+ unsigned copyDim;
+ unsigned baseDim;
+ int64_t size;
+};
- 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;
-}
+/// 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());
-/// 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)
+ // 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;
- if (!sourceDim ||
- sourceType.getDimSize(idx) < sourceType.getDimSize(*sourceDim))
- sourceDim = idx;
+ // 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 sourceDim;
+ return true;
}
-/// Returns source indices for a static reinterpret_cast offset.
+/// Returns base 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();
+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(offsets.size() == 1 && "Expecting single offset");
- assert(!ShapedType::isDynamic(offsets[0]) && "expected static offset");
+ assert(rcOffsets.size() == 1 && "Expecting single offset");
+ assert(!ShapedType::isDynamic(rcOffsets[0]) && "expected static offset");
- if (offsets[0] < 0)
+ if (rcOffsets[0] < 0)
return std::nullopt;
- SmallVector<int64_t> indices(sourceType.getRank(), 0);
- int64_t remainder = offsets[0];
- for (auto [idx, stride] : llvm::enumerate(sourceStrides)) {
+ MemRefType baseType = dyn_cast<MemRefType>(rc.getSource().getType());
+ SmallVector<int64_t> indices(baseType.getRank(), 0);
+ int64_t remainder = rcOffsets[0];
+ SmallVector<int64_t> baseStrides = computeStrides(baseType.getShape());
+ for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
indices[idx] = remainder / stride;
- if (indices[idx] >= sourceType.getDimSize(idx))
+ if (indices[idx] >= baseType.getDimSize(idx))
return std::nullopt;
remainder %= stride;
}
@@ -123,21 +143,6 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
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.
///
@@ -174,86 +179,70 @@ struct CopyFromReinterCastInfo {
/// 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());
+getCopyFromReinterCastInfo(memref::CopyOp copy, memref::ReinterpretCastOp rc) {
+ MemRefType cpSrcType = dyn_cast<MemRefType>(copy.getSource().getType());
+ MemRefType rcBaseType = dyn_cast<MemRefType>(rc.getSource().getType());
+ MemRefType rcResType = dyn_cast<MemRefType>(rc.getType());
// Ranked memref types are required to statically build load/store index
// lists.
- if (!srcType || !baseType || !resultType)
+ if (!cpSrcType || !rcBaseType || !rcResType)
return std::nullopt;
- if (srcType.getShape() != resultType.getShape())
+ if (cpSrcType.getShape() != rcResType.getShape())
return std::nullopt;
- // TODO: Support rank-changing reinterpret_casts by converting the
+ // TODO: Support rank-modifying 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())
+ if (rcBaseType.getRank() != rcResType.getRank())
return std::nullopt;
// TODO: Support dynamic shapes with mixed size operands as loop bounds.
- if (!(srcType.hasStaticShape() && baseType.hasStaticShape() &&
- resultType.hasStaticShape()))
+ if (!(cpSrcType.hasStaticShape() && rcBaseType.hasStaticShape() &&
+ rcResType.hasStaticShape()))
return std::nullopt;
- std::optional<SmallVector<int64_t>> baseStrides =
- getIdentityStrides(baseType);
- // TODO: Support non-identity source layouts by computing source strides from
+ // TODO: Support non-identity base layouts by computing base strides from
// the layout map.
- if (!baseStrides)
+ if (!rcBaseType.getLayout().isIdentity())
return std::nullopt;
CopyFromReinterCastInfo info;
- SmallVector<bool> usedBaseDims(baseType.getRank(), false);
- for (auto [resultDim, resultSize] : llvm::enumerate(resultType.getShape())) {
- if (resultSize == 1)
- continue;
-
- // TODO: Support dynamic strides on copied dimensions.
- if (ShapedType::isDynamic(rc.getStaticStrides()[resultDim]))
- return std::nullopt;
-
- // Each copied result dimension must map by stride to a source dimension
- // whose static size >= result dimension size.
- std::optional<unsigned> baseDim =
- findSourceDimForResultDim(rc, static_cast<unsigned>(resultDim),
- baseType, *baseStrides, usedBaseDims);
- assert(baseDim &&
- "static reinterpret_cast stride must map to an identity base "
- "dimension");
-
- usedBaseDims[*baseDim] = true;
- info.loopDims.push_back(CopyLoopDimInfo{static_cast<unsigned>(resultDim),
- *baseDim, resultSize});
- }
+ // reinterpret_cast result dimensions must map to distinct base dimensions.
+ // The rewrite emits one loop per copied dimension and adds each IV to one
+ // base index.
+ if (!findBaseDimForResultDim(rc, info))
+ return std::nullopt;
- ArrayRef<int64_t> offsets = rc.getStaticOffsets();
+ 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(offsets.size() == 1 && "Expecting single offset");
- if (!ShapedType::isDynamic(offsets[0])) {
- // Static offset must delinearize to in-bounds source indices.
- std::optional<SmallVector<int64_t>> offsetIndices =
- delinearizeStaticOffset(rc, baseType, *baseStrides);
- assert(offsetIndices &&
+ assert(rcOffsets.size() == 1 && "Expecting single offset");
+ // CASE 1: Static ReinterpretCast offset
+ if (!ShapedType::isDynamic(rcOffsets[0])) {
+ // Delinearize static ReinterpretCast offset as in-bounds indices (one for
+ // every base dimension).
+ std::optional<SmallVector<int64_t>> offsetIdxs =
+ delinearizeStaticRCOffset(rc);
+ assert(offsetIdxs &&
"static reinterpret_cast offset must delinearize to in-bounds base "
"indices");
- for (const CopyLoopDimInfo &loopDim : info.loopDims) {
- assert((*offsetIndices)[loopDim.baseDim] + loopDim.size <=
- baseType.getDimSize(loopDim.baseDim) &&
- "reinterpret_cast metadata describes an invalid accessible "
- "region");
- }
- info.staticOffsetIndices = std::move(offsetIndices);
+ assert(llvm::all_of(info.loopDims,
+ [&](const CopyLoopDimInfo &loopDim) {
+ return (*offsetIdxs)[loopDim.baseDim] +
+ loopDim.size <=
+ rcBaseType.getDimSize(loopDim.baseDim);
+ }) &&
+ "reinterpret_cast metadata describes an invalid accessible region");
+ info.staticOffsetIdxs = std::move(offsetIdxs);
return info;
}
-
+ // CASE 2: Dynamic ReinterpretCast offset
// Dynamic offsets are kept only when they can be used as a single base index.
// TODO: Support dynamic offsets for copies with multiple loop dimensions by
// delinearizing the offset into base start indices at runtime before adding
@@ -265,7 +254,7 @@ getCopyFromReinterCastInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
// Dynamic scalar offsets cannot be delinearized statically. They can be
// used directly only when the base has a single non-unit dimension to
// receive them.
- std::optional<unsigned> nonUnitDim = getSingleNonUnitDim(baseType);
+ std::optional<unsigned> nonUnitDim = getSingleNonUnitDim(rcBaseType);
// TODO: Support scalar dynamic offsets into bases with multiple non-unit
// dimensions by delinearizing the single accessed element offset at
// runtime.
@@ -276,9 +265,12 @@ getCopyFromReinterCastInfo(memref::CopyOp op, memref::ReinterpretCastOp rc) {
return info;
}
- unsigned baseDim = info.loopDims.front().baseDim;
- info.dynamicOffsetDim =
- (*baseStrides)[baseDim] == 1 ? baseDim : baseStrides->size() - 1;
+ unsigned rcBaseDim = info.loopDims.front().baseDim;
+ SmallVector<int64_t> rcBaseIdentityStrides =
+ computeStrides(rcBaseType.getShape());
+ info.dynamicOffsetDim = rcBaseIdentityStrides[rcBaseDim] == 1
+ ? rcBaseDim
+ : rcBaseIdentityStrides.size() - 1;
return info;
}
@@ -339,7 +331,7 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
Value src = op.getSource();
Value dst = rc.getSource();
- MemRefType srcType = cast<MemRefType>(src.getType());
+ MemRefType cpSrcType = cast<MemRefType>(src.getType());
MemRefType dstType = cast<MemRefType>(dst.getType());
// Reuse common index constants across bounds, steps, and static offsets,
@@ -372,9 +364,8 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
SmallVector<Value> baseStoreIndices = getZeroIndices(dstType.getRank());
// Static offsets were already delinearized into base indices. Fill the
// non-zero starting indices before creating loop bodies.
- if (copyInfo->staticOffsetIndices) {
- for (auto [idx, offset] :
- llvm::enumerate(*copyInfo->staticOffsetIndices)) {
+ if (copyInfo->staticOffsetIdxs) {
+ for (auto [idx, offset] : llvm::enumerate(*copyInfo->staticOffsetIdxs)) {
if (offset == 0)
continue;
baseStoreIndices[idx] = getOrCreateIndexConstant(offset);
@@ -384,12 +375,12 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
// dimension selected by getCopyFromReinterCastInfo.
assert(copyInfo->dynamicOffsetDim &&
"expected dynamic offset dimension for dynamic offset");
- SmallVector<OpFoldResult> offsets = rc.getMixedOffsets();
+ SmallVector<OpFoldResult> rcOffsets = rc.getMixedOffsets();
// FIXME: Despite what `getMixedOffsets` 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(rcOffsets.size() == 1 && "Expecting single offset");
baseStoreIndices[*copyInfo->dynamicOffsetDim] =
- getValueOrCreateConstantIndexOp(rewriter, loc, offsets[0]);
+ getValueOrCreateConstantIndexOp(rewriter, loc, rcOffsets[0]);
}
// Scope for OpBuilder::InsertionGuard.
@@ -415,7 +406,7 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
// Load indices are zero except for copied strided memref dimensions,
// which use the corresponding loop induction variables.
- SmallVector<Value> loadIndices = getZeroIndices(srcType.getRank());
+ SmallVector<Value> loadIndices = getZeroIndices(cpSrcType.getRank());
unsigned loopIndex = 0;
for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims)
loadIndices[loopDim.copyDim] = loopIvs[loopIndex++];
@@ -548,7 +539,7 @@ static bool isPureRankExpansionOrCollapsingRC(memref::ReinterpretCastOp rc) {
if (!inputNonUnitDim || !outputNonUnitDim)
return false;
- // The source and result must either both have a single non-unit dimension
+ // The base and result must either both have a single non-unit dimension
// or both be all-ones.
if (inputNonUnitDim->allOnes != outputNonUnitDim->allOnes)
return false;
@@ -592,7 +583,7 @@ static bool isPureRankExpansionOrCollapsingRC(memref::ReinterpretCastOp rc) {
}
/// Rewrites `memref.load` through a pure rank-only `reinterpret_cast` by
-/// mapping the load indices directly onto the source MemRef.
+/// mapping the load indices directly onto the base MemRef.
/// Shape restriction gated by isPureRankExpansionOrCollapsingRC().
///
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 044e12cec42af..34678871e8cfe 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -297,6 +297,22 @@ func.func private @negative_copy_to_strided_non_identity_base(%src: memref<1x1xf
return
}
+// CHECK-LABEL: func.func private @negative_copy_to_strided_overlapping_result_dims(
+func.func private @negative_copy_to_strided_overlapping_result_dims(%src : memref<3x3xf32>,
+ %dst : memref<4x4xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [3, 3], strides: [1, 1]
+ : memref<4x4xf32> to memref<3x3xf32, strided<[1, 1]>>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %reinterpret_cast
+ : memref<3x3xf32> to memref<3x3xf32, strided<[1, 1]>>
+ return
+}
+
// CHECK-LABEL: func.func private @negative_copy_to_strided_rank_change(
func.func private @negative_copy_to_strided_rank_change(%src : memref<2x3xf32>,
%dst : memref<6xf32>) {
>From c11bb46b611739656cf1229dfc5191b416a53a96 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Thu, 25 Jun 2026 09:41:58 +0200
Subject: [PATCH 06/14] Address third round of comments
---
.../Transforms/ElideReinterpretCast.cpp | 378 +++++++++---------
.../MemRef/elide-reinterpret-cast.mlir | 76 ++--
2 files changed, 230 insertions(+), 224 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 6d806b087bdb4..435600f58f7fa 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -35,141 +35,154 @@ namespace {
// Copy Rewrite Helpers
//===----------------------------------------------------------------------===//
-/// Per-dimension loop nest info.
-struct CopyLoopDimInfo {
- unsigned copyDim;
- unsigned baseDim;
- int64_t size;
+/// Non-unit reinterpret_cast result dimension and the source dimension it
+/// advances through.
+struct NonUnitDimAssocMapForRC {
+ unsigned resultDimPos;
+ unsigned sourceDimPos;
};
-/// 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;
+/// 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;
};
-/// Maps non-unit reinterpret_cast result dimensions to distinct base
-/// dimensions.
-static bool findBaseDimForResultDim(memref::ReinterpretCastOp rc,
- CopyFromReinterCastInfo &info) {
+/// 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 baseType = dyn_cast<MemRefType>(rc.getSource().getType());
- SmallVector<int64_t> baseIdentityStrides =
- computeStrides(baseType.getShape());
+ MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
+ assert(srcType.getLayout().isIdentity() &&
+ "Expecting identity source layout.");
- // 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);
+ 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);
- // 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.
+ // 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> 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)
+ 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 (!baseDim || baseType.getDimSize(idx) < baseType.getDimSize(*baseDim))
- baseDim = idx;
+ if (!srcDim || srcType.getDimSize(idx) < srcType.getDimSize(*srcDim))
+ srcDim = idx;
}
- if (!baseDim)
+ if (!srcDim)
return false;
- usedBaseDims[*baseDim] = true;
- info.loopDims.push_back(CopyLoopDimInfo{static_cast<unsigned>(resultDim),
- *baseDim, resultSize});
+ usedSrcDims[*srcDim] = true;
+ mapAndOffs.assocMap.push_back(
+ NonUnitDimAssocMapForRC{static_cast<unsigned>(resultDim), *srcDim});
}
return true;
}
-/// Returns base indices for a static reinterpret_cast offset.
+/// 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::isDynamic(rcOffsets[0]) && "expected static offset");
-
- if (rcOffsets[0] < 0)
- return std::nullopt;
+ 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>{};
+ }
- MemRefType baseType = dyn_cast<MemRefType>(rc.getSource().getType());
- SmallVector<int64_t> indices(baseType.getRank(), 0);
+ SmallVector<int64_t> offsetIdxs(srcType.getRank(), 0);
int64_t remainder = rcOffsets[0];
- SmallVector<int64_t> baseStrides = computeStrides(baseType.getShape());
- for (auto [idx, stride] : llvm::enumerate(baseStrides)) {
- indices[idx] = remainder / stride;
- if (indices[idx] >= baseType.getDimSize(idx))
- return std::nullopt;
+ SmallVector<int64_t> srcStrides = computeStrides(srcType.getShape());
+ // Convert the linear reinterpret_cast offset to per-dimension source starting
+ // indices.
+ for (auto [dim, stride] : llvm::enumerate(srcStrides)) {
+ offsetIdxs[dim] = remainder / stride;
+ assert(offsetIdxs[dim] < srcType.getDimSize(dim) &&
+ "static reinterpret_cast offset must delinearize to in-bounds "
+ "source indices");
remainder %= stride;
}
- if (remainder != 0)
- return std::nullopt;
- return indices;
+ assert(remainder == 0 &&
+ "Assuming identity source layout, the trailing stride == 1 "
+ "so, the remainder should be 0 at the end of index calculation.");
+ return offsetIdxs;
}
-/// Returns the dimension whose static size is not one if it is unique.
+/// Returns the unique non-unit dim or nullopt of # non-unit-dims != 1.
static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
assert(type.hasStaticShape() && "expected static shape");
ArrayRef<int64_t> shape = type.getShape();
- if (shape.empty())
+
+ // Find all non-unit dims
+ auto nonUnitDims = llvm::make_filter_range(
+ llvm::enumerate(shape), [](auto it) { return it.value() != 1; });
+
+ // Expect single non-unit dims
+ if (llvm::range_size(nonUnitDims) != 1)
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;
+ // Return the index of the unique non-unit dim.
+ return (*nonUnitDims.begin()).index();
}
-/// Builds the index mapping needed to replace a copy into a reinterpret_cast
-/// strided memref with scalar stores into the reinterpret_cast base.
+/// Returns the copy-relevant reinterpret_cast information: non-unit result
+/// dimensions, their source-dimension mapping, and optional source starting
+/// indices for a static offset.
///
/// Examples that return rewrite info:
///
-/// // Scalar-shaped copy. There are no copied non-unit dimensions, so dynamic
-/// // strides in the strided memref do not affect index mapping.
+/// // Scalar-shaped copy into a source with at most one non-unit dimension.
+/// There are
+/// // no non-unit result 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 reinterpret_cast memref<source-shape>
/// to memref<1 x ... x 1 x f32, strided<[?, ..., ?], offset: ?>>
///
/// // Effectively-1D copy. The single non-unit strided memref dimension is
-/// // mapped to an identity-layout base dimension by its static stride.
+/// // mapped to an identity-layout source dimension by its static stride.
/// copy memref<1 x ... x N x ... x 1 x f32>
-/// to reinterpret_cast memref<base-shape>
+/// to reinterpret_cast memref<source-shape>
/// to memref<1 x ... x N x ... x 1 x f32, strided<[..., S, ...]>>
///
/// // 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 reinterpret_cast memref<source-shape>
/// to memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32,
/// strided<[..., S_0, ..., S_1, ...], offset: O>>
///
/// Examples that return no info:
///
-/// // Dynamic stride on a copied strided memref dimension.
+/// // Dynamic stride on a non-unit strided memref dimension.
/// copy memref<1xNxf32>
/// to reinterpret_cast memref<1xNxMxf32>
/// to memref<1xNxf32, strided<[?, ?]>>
@@ -178,100 +191,73 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
/// copy memref<1xNxKxf32>
/// to reinterpret_cast memref<1xNxMxf32>
/// to memref<1xNxKxf32, strided<[N*M, M, 1], offset: ?>>
-static std::optional<CopyFromReinterCastInfo>
-getCopyFromReinterCastInfo(memref::CopyOp copy, memref::ReinterpretCastOp rc) {
- MemRefType cpSrcType = dyn_cast<MemRefType>(copy.getSource().getType());
- MemRefType rcBaseType = dyn_cast<MemRefType>(rc.getSource().getType());
- MemRefType rcResType = dyn_cast<MemRefType>(rc.getType());
+static std::optional<AssocMapAndOffsetsForRC>
+getAssocMapAndOffsetsForRC(memref::ReinterpretCastOp rc) {
+ MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
+ MemRefType resType = dyn_cast<MemRefType>(rc.getType());
// Ranked memref types are required to statically build load/store index
// lists.
- if (!cpSrcType || !rcBaseType || !rcResType)
- return std::nullopt;
-
- if (cpSrcType.getShape() != rcResType.getShape())
+ if (!srcType || !resType)
return std::nullopt;
- // TODO: Support rank-modifying 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 (rcBaseType.getRank() != rcResType.getRank())
+ // TODO: Support rank-modifying reinterpret_casts
+ if (srcType.getRank() != resType.getRank())
return std::nullopt;
// TODO: Support dynamic shapes with mixed size operands as loop bounds.
- if (!(cpSrcType.hasStaticShape() && rcBaseType.hasStaticShape() &&
- rcResType.hasStaticShape()))
+ if (!(srcType.hasStaticShape() && resType.hasStaticShape()))
return std::nullopt;
- // TODO: Support non-identity base layouts by computing base strides from
+ // TODO: Support non-identity source layouts by computing source strides from
// the layout map.
- if (!rcBaseType.getLayout().isIdentity())
+ if (!srcType.getLayout().isIdentity())
return std::nullopt;
- CopyFromReinterCastInfo info;
+ AssocMapAndOffsetsForRC mapAndOffs;
- // reinterpret_cast result dimensions must map to distinct base dimensions.
- // The rewrite emits one loop per copied dimension and adds each IV to one
- // base index.
- if (!findBaseDimForResultDim(rc, info))
+ // reinterpret_cast result dimensions must map to distinct source dimensions.
+ if (!findSourceDimForResultDim(rc, mapAndOffs))
return std::nullopt;
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");
- // CASE 1: Static ReinterpretCast offset
- if (!ShapedType::isDynamic(rcOffsets[0])) {
+ // Static ReinterpretCast offset
+ if (ShapedType::isStatic(rcOffsets[0])) {
// Delinearize static ReinterpretCast offset as in-bounds indices (one for
- // every base dimension).
- std::optional<SmallVector<int64_t>> offsetIdxs =
- delinearizeStaticRCOffset(rc);
- assert(offsetIdxs &&
- "static reinterpret_cast offset must delinearize to in-bounds base "
- "indices");
-
- assert(llvm::all_of(info.loopDims,
- [&](const CopyLoopDimInfo &loopDim) {
- return (*offsetIdxs)[loopDim.baseDim] +
- loopDim.size <=
- rcBaseType.getDimSize(loopDim.baseDim);
- }) &&
- "reinterpret_cast metadata describes an invalid accessible region");
- info.staticOffsetIdxs = std::move(offsetIdxs);
- return info;
+ // every source dimension).
+ mapAndOffs.delinearizedOffsets = delinearizeStaticRCOffset(rc);
+ assert(
+ mapAndOffs.delinearizedOffsets &&
+ "static reinterpret_cast offset must delinearize to in-bounds source "
+ "indices");
+
+ assert(
+ llvm::all_of(
+ mapAndOffs.assocMap,
+ [&](const NonUnitDimAssocMapForRC &assocMap) {
+ return (*mapAndOffs.delinearizedOffsets)[assocMap.sourceDimPos] +
+ resType.getDimSize(assocMap.resultDimPos) <=
+ srcType.getDimSize(assocMap.sourceDimPos);
+ }) &&
+ "reinterpret_cast metadata describes an invalid accessible region");
+ return mapAndOffs;
}
- // CASE 2: Dynamic ReinterpretCast offset
- // Dynamic offsets are kept only when they can be used as a single base index.
- // TODO: Support dynamic offsets for copies with multiple loop dimensions by
- // delinearizing the offset into base start indices at runtime before adding
- // loop IVs.
- if (info.loopDims.size() > 1)
- return std::nullopt;
- if (info.loopDims.empty()) {
- // Dynamic scalar offsets cannot be delinearized statically. They can be
- // used directly only when the base has a single non-unit dimension to
- // receive them.
- std::optional<unsigned> nonUnitDim = getSingleNonUnitDim(rcBaseType);
- // TODO: Support scalar dynamic offsets into bases with multiple non-unit
- // dimensions by delinearizing the single accessed element offset at
- // runtime.
- if (!nonUnitDim)
- return std::nullopt;
-
- info.dynamicOffsetDim = *nonUnitDim;
- return info;
- }
+ // Dynamic ReinterpretCast offset.
+ // TODO: Support dynamic offsets into sources with multiple non-unit
+ // dimensions by delinearizing the offset into source start indices at runtime
+ // before adding loop IVs.
+ if (mapAndOffs.assocMap.size() > 1)
+ return std::nullopt;
+ // Only sources with a single non-unit dimension can receive a dynamic offset
+ // directly.
+ if (!getSingleNonUnitDim(srcType))
+ return std::nullopt;
- unsigned rcBaseDim = info.loopDims.front().baseDim;
- SmallVector<int64_t> rcBaseIdentityStrides =
- computeStrides(rcBaseType.getShape());
- info.dynamicOffsetDim = rcBaseIdentityStrides[rcBaseDim] == 1
- ? rcBaseDim
- : rcBaseIdentityStrides.size() - 1;
- return info;
+ return mapAndOffs;
}
/// Rewrites supported copy operations through `memref.reinterpret_cast` to
@@ -315,23 +301,25 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
LogicalResult matchAndRewrite(memref::CopyOp op,
PatternRewriter &rewriter) const final {
+ Value src = op.getSource();
+ MemRefType cpSrcType = cast<MemRefType>(src.getType());
+ if (!cpSrcType || !cpSrcType.hasStaticShape())
+ return rewriter.notifyMatchFailure(
+ op, "only ranked, static copy sources are supported.");
Value rcOutput = op.getTarget();
auto rc = rcOutput.getDefiningOp<memref::ReinterpretCastOp>();
if (!rc)
return rewriter.notifyMatchFailure(
op, "target is not a memref.reinterpret_cast");
- std::optional<CopyFromReinterCastInfo> copyInfo =
- getCopyFromReinterCastInfo(op, rc);
- if (!copyInfo)
+ std::optional<AssocMapAndOffsetsForRC> mapAndOffs =
+ getAssocMapAndOffsetsForRC(rc);
+ if (!mapAndOffs)
return rewriter.notifyMatchFailure(
op, "reinterpret_cast does not match scalar or loop copy region");
Location loc = op.getLoc();
- Value src = op.getSource();
Value dst = rc.getSource();
-
- MemRefType cpSrcType = cast<MemRefType>(src.getType());
MemRefType dstType = cast<MemRefType>(dst.getType());
// Reuse common index constants across bounds, steps, and static offsets,
@@ -346,40 +334,48 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
}
return arith::ConstantIndexOp::create(rewriter, loc, value);
};
- auto getZeroIndices = [&](int64_t rank) {
- SmallVector<Value> indices;
- indices.reserve(rank);
+ auto getZeroIdxs = [&](int64_t rank) {
+ SmallVector<Value> idxs;
+ idxs.reserve(rank);
if (rank != 0)
- indices.append(rank, getOrCreateIndexConstant(0));
- return indices;
+ idxs.append(rank, getOrCreateIndexConstant(0));
+ return idxs;
};
- // Create all loop bounds before building the loop nest. Otherwise an
- // inner-loop bound can be inserted inside an outer loop body.
+ // Create loop bounds before moving the insertion point into the loop nest,
+ // so loop-invariant constants are emitted outside the generated loops.
SmallVector<Value> upperBounds;
- upperBounds.reserve(copyInfo->loopDims.size());
- for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims)
- upperBounds.push_back(getOrCreateIndexConstant(loopDim.size));
-
- SmallVector<Value> baseStoreIndices = getZeroIndices(dstType.getRank());
- // Static offsets were already delinearized into base indices. Fill the
- // non-zero starting indices before creating loop bodies.
- if (copyInfo->staticOffsetIdxs) {
- for (auto [idx, offset] : llvm::enumerate(*copyInfo->staticOffsetIdxs)) {
+ upperBounds.reserve(mapAndOffs->assocMap.size());
+ MemRefType rcResType = dyn_cast<MemRefType>(rc.getType());
+ for (const NonUnitDimAssocMapForRC &assocMap : mapAndOffs->assocMap)
+ upperBounds.push_back(getOrCreateIndexConstant(
+ rcResType.getDimSize(assocMap.resultDimPos)));
+
+ SmallVector<Value> rcSrcStoreIdxs = getZeroIdxs(dstType.getRank());
+ std::optional<unsigned> srcNonUnitDimPos;
+ // Static offset has been delinearized in function gating rewrite.
+ if (mapAndOffs->delinearizedOffsets) {
+ for (auto [idx, offset] :
+ llvm::enumerate(*mapAndOffs->delinearizedOffsets)) {
if (offset == 0)
continue;
- baseStoreIndices[idx] = getOrCreateIndexConstant(offset);
+ rcSrcStoreIdxs[idx] = getOrCreateIndexConstant(offset);
}
} else {
- // Supported dynamic offsets are used directly in exactly one base
- // dimension selected by getCopyFromReinterCastInfo.
- assert(copyInfo->dynamicOffsetDim &&
- "expected dynamic offset dimension for dynamic offset");
+ // Without runtime delinearization, use the dynamic offset directly only
+ // when the source has a single non-unit dimension.
+ assert(mapAndOffs->assocMap.size() <= 1 &&
+ "Expecting single non-unit dimension mapping.");
+ srcNonUnitDimPos = getSingleNonUnitDim(dstType);
+ assert(srcNonUnitDimPos &&
+ "Expecting single non-unit dimension source to receive the "
+ "dynamic offset.");
+
SmallVector<OpFoldResult> rcOffsets = rc.getMixedOffsets();
// FIXME: Despite what `getMixedOffsets` implies, `reinterpret_cast` takes
// only a single offset. That should be fixed at the op definition level.
assert(rcOffsets.size() == 1 && "Expecting single offset");
- baseStoreIndices[*copyInfo->dynamicOffsetDim] =
+ rcSrcStoreIdxs[*srcNonUnitDimPos] =
getValueOrCreateConstantIndexOp(rewriter, loc, rcOffsets[0]);
}
@@ -392,11 +388,10 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
lowerBound = getOrCreateIndexConstant(0);
step = getOrCreateIndexConstant(1);
}
-
SmallVector<Value> loopIvs;
- loopIvs.reserve(copyInfo->loopDims.size());
+ loopIvs.reserve(mapAndOffs->assocMap.size());
- // Build one nested loop per non-unit copied strided memref dimension.
+ // Build one nested loop per non-unit strided memref dimension.
for (Value upperBound : upperBounds) {
scf::ForOp loop =
scf::ForOp::create(rewriter, loc, lowerBound, upperBound, step);
@@ -404,31 +399,32 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
rewriter.setInsertionPointToStart(loop.getBody());
}
- // Load indices are zero except for copied strided memref dimensions,
+ // Load indices are zero except for non-unit strided memref dimensions,
// which use the corresponding loop induction variables.
- SmallVector<Value> loadIndices = getZeroIndices(cpSrcType.getRank());
+ SmallVector<Value> loadIdxs = getZeroIdxs(cpSrcType.getRank());
unsigned loopIndex = 0;
- for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims)
- loadIndices[loopDim.copyDim] = loopIvs[loopIndex++];
+ for (const NonUnitDimAssocMapForRC &assocMap : mapAndOffs->assocMap)
+ loadIdxs[assocMap.resultDimPos] = loopIvs[loopIndex++];
- // Store indices start from the offset-derived base indices. Add each loop
- // IV to the mapped base dimension.
- SmallVector<Value> storeIndices(baseStoreIndices);
+ // Store indices start from the offset-derived source indices. Add each
+ // loop IV to the mapped source dimension.
+ SmallVector<Value> storeIdxs(rcSrcStoreIdxs);
loopIndex = 0;
- for (const CopyLoopDimInfo &loopDim : copyInfo->loopDims) {
+ for (const NonUnitDimAssocMapForRC &assocMap : mapAndOffs->assocMap) {
Value iv = loopIvs[loopIndex++];
- if (storeIndices[loopDim.baseDim] == getOrCreateIndexConstant(0)) {
- storeIndices[loopDim.baseDim] = iv;
+ // Add each IV to one source index.
+ if (storeIdxs[assocMap.sourceDimPos] == getOrCreateIndexConstant(0)) {
+ storeIdxs[assocMap.sourceDimPos] = iv;
} else {
- storeIndices[loopDim.baseDim] = arith::AddIOp::create(
- rewriter, loc, storeIndices[loopDim.baseDim], iv);
+ storeIdxs[assocMap.sourceDimPos] = arith::AddIOp::create(
+ rewriter, loc, storeIdxs[assocMap.sourceDimPos], iv);
}
}
// Emit the scalar load/store at the innermost loop body, or directly at
// the original copy location for scalar copies.
- Value val = memref::LoadOp::create(rewriter, loc, src, loadIndices);
- memref::StoreOp::create(rewriter, loc, val, dst, storeIndices);
+ Value val = memref::LoadOp::create(rewriter, loc, src, loadIdxs);
+ memref::StoreOp::create(rewriter, loc, val, dst, storeIdxs);
}
// If the only user of `rc` is the current Op (which is about to be erased),
@@ -709,7 +705,11 @@ struct ElideReinterpretCastPass
auto rc = op.getTarget().getDefiningOp<memref::ReinterpretCastOp>();
if (!rc)
return true;
- return !getCopyFromReinterCastInfo(op, rc);
+ // Pattern applies only when the copy source shape is static and the
+ // reinterpret_cast result can be mapped back to base memref indices.
+ MemRefType cpSrcType = dyn_cast<MemRefType>(op.getSource().getType());
+ return !(cpSrcType && cpSrcType.hasStaticShape() &&
+ getAssocMapAndOffsetsForRC(rc));
});
target.addDynamicallyLegalOp<memref::LoadOp>([](memref::LoadOp op) {
auto rc = op.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 34678871e8cfe..4a879a916c625 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -92,7 +92,6 @@ func.func private @copy_to_strided_dynamic_stride(%stride0: index,
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
- /// Dynamic offset used in store
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
memref.copy %src, %reinterpret_cast
: memref<1x1xf32>
@@ -117,10 +116,10 @@ func.func private @copy_to_strided_rank0(%src : memref<f32>, %dst : memref<f32>)
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_0d_2d_base(
+// CHECK-LABEL: func.func private @copy_to_strided_0d_base_2d(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
-func.func private @copy_to_strided_0d_2d_base(
+func.func private @copy_to_strided_0d_base_2d(
%src : memref<1x1x1xf32>, %dst : memref<1x33x42xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -190,11 +189,11 @@ func.func private @copy_to_strided_1d_vector_nonzero_offset(
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_dynamic_offset_in_loop_dim(
+// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_1d_base_dynamic_offset(
// CHECK-SAME: %[[OFF:.*]]: index
// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
// CHECK-SAME: %[[DST:.*]]: memref<42xf32>
-func.func private @copy_to_strided_1d_vector_dynamic_offset_in_loop_dim(
+func.func private @copy_to_strided_1d_vector_1d_base_dynamic_offset(
%offset : index, %src : memref<4xf32>, %dst : memref<42xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -216,34 +215,6 @@ func.func private @copy_to_strided_1d_vector_dynamic_offset_in_loop_dim(
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_dynamic_offset_not_in_loop_dim(
-// CHECK-SAME: %[[OFF:.*]]: index
-// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_to_strided_1d_vector_dynamic_offset_not_in_loop_dim(
- %offset : index, %src : memref<1x3x1xf32>,
- %dst : memref<1x3x11xf32>) {
- // CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [1, 3, 1], strides: [33, 11, 1]
- : memref<1x3x11xf32>
- to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
-
- // CHECK-NOT: memref.copy
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[OFF]]] : memref<1x3x11xf32>
- // CHECK: }
- memref.copy %src, %reinterpret_cast
- : memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
- // CHECK-NOT: memref.copy
- return
-}
-
// CHECK-LABEL: func.func private @copy_to_strided_2d_vector_nonzero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
@@ -361,8 +332,24 @@ func.func private @negative_copy_to_strided_dynamic_rc_shapes(%dim : index,
return
}
-// CHECK-LABEL: func.func private @negative_copy_to_strided_dynamic_offset_multi_dim_base(
-func.func private @negative_copy_to_strided_dynamic_offset_multi_dim_base(
+// CHECK-LABEL: func.func private @negative_copy_to_strided_unranked_rc_base(
+func.func private @negative_copy_to_strided_unranked_rc_base(
+ %src : memref<4xf32>, %dst : memref<*xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [4], strides: [1]
+ : memref<*xf32> to memref<4xf32>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %reinterpret_cast
+ : memref<4xf32> to memref<4xf32>
+ return
+}
+
+// CHECK-LABEL: func.func private @negative_copy_to_strided_0d_base_2d_dynamic_offset(
+func.func private @negative_copy_to_strided_0d_base_2d_dynamic_offset(
%offset : index, %src : memref<1x1xf32>, %dst : memref<4x8xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -378,6 +365,25 @@ func.func private @negative_copy_to_strided_dynamic_offset_multi_dim_base(
return
}
+// CHECK-LABEL: func.func private @negative_copy_to_strided_1d_base_2d_dynamic_offset(
+func.func private @negative_copy_to_strided_1d_base_2d_dynamic_offset(
+ %offset : index, %src : memref<1x3x1xf32>,
+ %dst : memref<1x3x11xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [%offset], sizes: [1, 3, 1], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
+
+ // CHECK: memref.copy %arg1, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %reinterpret_cast
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
+ return
+}
+
// CHECK-LABEL: func.func private @negative_copy_to_strided_2d_dynamic_offset(
func.func private @negative_copy_to_strided_2d_dynamic_offset(
%offset : index, %src : memref<1x3x4xf32>,
>From 8c90ec715552457fc2676af41000a7b008719197 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Fri, 26 Jun 2026 17:50:06 +0200
Subject: [PATCH 07/14] Fixup tests
---
.../MemRef/elide-reinterpret-cast.mlir | 245 +++++++++++-------
1 file changed, 145 insertions(+), 100 deletions(-)
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 4a879a916c625..eaa823e4afd3d 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -5,14 +5,16 @@
// Positive tests
//===----------------------------------------------------------------------===//
-// CHECK-LABEL: func.func private @copy_to_strided_zero_offset(
+/// Effectively copied scalar within a MemRef with rank >= 0
+/// to practically a 1D array within a MemRef with rank >= 1
+// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
-func.func private @copy_to_strided_zero_offset(%src : memref<1x1xf32>,
+func.func private @copy_scalar_into_1D_strided_zero_offset(%src : memref<1x1xf32>,
%dst : memref<1x108xf32>) {
/// reinterpret_cast removed
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [1, 1], strides: [1, 1]
: memref<1x108xf32> to memref<1x1xf32>
@@ -21,19 +23,19 @@ func.func private @copy_to_strided_zero_offset(%src : memref<1x1xf32>,
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x1xf32> to memref<1x1xf32>
// CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_nonzero_offset(
+// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_nonzero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
-func.func private @copy_to_strided_nonzero_offset(%src : memref<1x1xf32>,
+func.func private @copy_scalar_into_1D_strided_nonzero_offset(%src : memref<1x1xf32>,
%dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [1], sizes: [1, 1], strides: [1, 1]
: memref<1x108xf32>
to memref<1x1xf32, strided<[1, 1], offset: 1>>
@@ -43,21 +45,21 @@ func.func private @copy_to_strided_nonzero_offset(%src : memref<1x1xf32>,
// CHECK: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C1]]] : memref<1x108xf32>
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x1xf32>
to memref<1x1xf32, strided<[1, 1], offset: 1>>
// CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_dynamic_offset(
+// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_dynamic_offset(
// CHECK-SAME: %[[OFF:.*]]: index
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
-func.func private @copy_to_strided_dynamic_offset(%offset: index, %src : memref<1x1xf32>,
+func.func private @copy_scalar_into_1D_strided_dynamic_offset(%offset: index, %src : memref<1x1xf32>,
%dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [%offset], sizes: [1, 1], strides: [1, 1]
: memref<1x108xf32>
to memref<1x1xf32, strided<[1, 1], offset: ?>>
@@ -68,7 +70,7 @@ func.func private @copy_to_strided_dynamic_offset(%offset: index, %src : memref<
// CHECK-SAME: : memref<1x1xf32>
/// Dynamic offset used in store
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[OFF]]] : memref<1x108xf32>
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x1xf32>
to memref<1x1xf32, strided<[1, 1], offset: ?>>
// CHECK-NOT: memref.copy
@@ -76,15 +78,15 @@ func.func private @copy_to_strided_dynamic_offset(%offset: index, %src : memref<
}
// Dynamic strides are irrelevant because all strided memref indices are zero.
-// CHECK-LABEL: func.func private @copy_to_strided_dynamic_stride(
+// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_dynamic_stride(
// CHECK-SAME: %[[STR0:[A-Za-z][A-Za-z0-9-]*]]: index
// CHECK-SAME: %[[STR1:[A-Za-z][A-Za-z0-9-]*]]: index
// CHECK-SAME: %[[SRC:[A-Za-z][A-Za-z0-9-]*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:[A-Za-z][A-Za-z0-9-]*]]: memref<1x108xf32>
-func.func private @copy_to_strided_dynamic_stride(%stride0: index,
+func.func private @copy_scalar_into_1D_strided_dynamic_stride(%stride0: index,
%stride1: index, %src : memref<1x1xf32>, %dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [1, 1], strides: [%stride0, %stride1]
: memref<1x108xf32>
to memref<1x1xf32, strided<[?, ?]>>
@@ -93,36 +95,37 @@ func.func private @copy_to_strided_dynamic_stride(%stride0: index,
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x1xf32>
to memref<1x1xf32, strided<[?, ?]>>
// CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_rank0(
+// CHECK-LABEL: func.func private @copy_0D_into_0D_strided(
// CHECK-SAME: %[[SRC:.*]]: memref<f32>, %[[DST:.*]]: memref<f32>
-func.func private @copy_to_strided_rank0(%src : memref<f32>, %dst : memref<f32>) {
+func.func private @copy_0D_into_0D_strided(%src : memref<f32>, %dst : memref<f32>) {
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [], strides: []
: memref<f32> to memref<f32>
// CHECK-NOT: memref.copy
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][] : memref<f32>
// CHECK: memref.store %[[VAL]], %[[DST]][] : memref<f32>
- memref.copy %src, %reinterpret_cast : memref<f32> to memref<f32>
+ memref.copy %src, %rc : memref<f32> to memref<f32>
// CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_0d_base_2d(
+/// Effectively a 2D array within a MemRef with rank >= 2
+// CHECK-LABEL: func.func private @copy_scalar_into_2D_strided(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
-func.func private @copy_to_strided_0d_base_2d(
+func.func private @copy_scalar_into_2D_strided(
%src : memref<1x1x1xf32>, %dst : memref<1x33x42xf32>) {
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [1, 1, 1], strides: [1, 1, 1]
: memref<1x33x42xf32>
to memref<1x1x1xf32>
@@ -130,19 +133,19 @@ func.func private @copy_to_strided_0d_base_2d(
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x33x42xf32>
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x1x1xf32> to memref<1x1x1xf32>
// CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_zero_offset(
+// CHECK-LABEL: func.func private @copy_1D_into_1D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_to_strided_1d_vector_zero_offset(
+func.func private @copy_1D_into_1D_strided_zero_offset(
%src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [1, 3, 1], strides: [33, 11, 1]
: memref<1x3x11xf32>
to memref<1x3x1xf32, strided<[33, 11, 1]>>
@@ -155,20 +158,20 @@ func.func private @copy_to_strided_1d_vector_zero_offset(
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x11xf32>
// CHECK: }
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x3x1xf32>
to memref<1x3x1xf32, strided<[33, 11, 1]>>
// CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_nonzero_offset(
+// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_nonzero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_to_strided_1d_vector_nonzero_offset(
+func.func private @copy_1D_into_2D_strided_nonzero_offset(
%src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [10], sizes: [1, 3, 1], strides: [33, 11, 1]
: memref<1x3x11xf32>
to memref<1x3x1xf32, strided<[33, 11, 1], offset: 10>>
@@ -182,21 +185,21 @@ func.func private @copy_to_strided_1d_vector_nonzero_offset(
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C10]]] : memref<1x3x11xf32>
// CHECK: }
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x3x1xf32>
to memref<1x3x1xf32, strided<[33, 11, 1], offset: 10>>
// CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_1d_base_dynamic_offset(
+// CHECK-LABEL: func.func private @copy_1D_into_1D_strided_dynamic_offset(
// CHECK-SAME: %[[OFF:.*]]: index
// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
// CHECK-SAME: %[[DST:.*]]: memref<42xf32>
-func.func private @copy_to_strided_1d_vector_1d_base_dynamic_offset(
+func.func private @copy_1D_into_1D_strided_dynamic_offset(
%offset : index, %src : memref<4xf32>, %dst : memref<42xf32>) {
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [%offset], sizes: [4], strides: [1]
: memref<42xf32> to memref<4xf32, strided<[1], offset: ?>>
@@ -209,19 +212,19 @@ func.func private @copy_to_strided_1d_vector_1d_base_dynamic_offset(
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<4xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[DST_IDX]]] : memref<42xf32>
// CHECK: }
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<4xf32> to memref<4xf32, strided<[1], offset: ?>>
// CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_to_strided_2d_vector_nonzero_offset(
+// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_nonzero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_to_strided_2d_vector_nonzero_offset(
+func.func private @copy_2D_into_2D_strided_nonzero_offset(
%src : memref<1x3x4xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [7], sizes: [1, 3, 4], strides: [33, 11, 1]
: memref<1x3x11xf32>
to memref<1x3x4xf32, strided<[33, 11, 1], offset: 7>>
@@ -239,7 +242,7 @@ func.func private @copy_to_strided_2d_vector_nonzero_offset(
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[DST_IDX]]] : memref<1x3x11xf32>
// CHECK: }
// CHECK: }
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x3x4xf32>
to memref<1x3x4xf32, strided<[33, 11, 1], offset: 7>>
// CHECK-NOT: memref.copy
@@ -250,11 +253,11 @@ func.func private @copy_to_strided_2d_vector_nonzero_offset(
// Negative tests (must NOT rewrite)
//===----------------------------------------------------------------------===//
-// CHECK-LABEL: func.func private @negative_copy_to_strided_non_identity_base(
-func.func private @negative_copy_to_strided_non_identity_base(%src: memref<1x1xf32>,
+// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity(
+func.func private @negative_copy_into_strided_non_identity(%src: memref<1x1xf32>,
%dst: memref<8x1xf32, strided<[10, 2]>>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [6], sizes: [1, 1], strides: [11, 80]
: memref<8x1xf32, strided<[10, 2]>>
to memref<1x1xf32, strided<[11, 80], offset: 6>>
@@ -262,153 +265,195 @@ func.func private @negative_copy_to_strided_non_identity_base(%src: memref<1x1xf
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x1xf32> to memref<1x1xf32, strided<[11, 80], offset: 6>>
return
}
-// CHECK-LABEL: func.func private @negative_copy_to_strided_overlapping_result_dims(
-func.func private @negative_copy_to_strided_overlapping_result_dims(%src : memref<3x3xf32>,
+// CHECK-LABEL: func.func private @negative_copy_into_strided_overlapping_result_dims(
+func.func private @negative_copy_into_strided_overlapping_result_dims(%src : memref<3x3xf32>,
%dst : memref<4x4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [3, 3], strides: [1, 1]
: memref<4x4xf32> to memref<3x3xf32, strided<[1, 1]>>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<3x3xf32> to memref<3x3xf32, strided<[1, 1]>>
return
}
-// CHECK-LABEL: func.func private @negative_copy_to_strided_rank_change(
-func.func private @negative_copy_to_strided_rank_change(%src : memref<2x3xf32>,
+// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_strided_unmatched_strides(
+func.func private @negative_copy_1D_into_2D_strided_unmatched_strides(
+ %src : memref<1x3x1xf32>, %dst : memref<1x4x8xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 3, 1], strides: [32, 5, 1]
+ : memref<1x4x8xf32>
+ to memref<1x3x1xf32, strided<[32, 5, 1]>>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[32, 5, 1]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @negative_copy_2D_into_2D_strided_unmatched_strides(
+func.func private @negative_copy_2D_into_2D_strided_unmatched_strides(
+ %src : memref<1x3x4xf32>, %dst : memref<1x4x3xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 3, 4], strides: [12, 4, 1]
+ : memref<1x4x3xf32>
+ to memref<1x3x4xf32, strided<[12, 4, 1]>>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<1x3x4xf32>
+ to memref<1x3x4xf32, strided<[12, 4, 1]>>
+ return
+}
+
+/// The linear accessed region is in-bounds, but result dim 1 has size 5 while
+/// the matching source dim has size 4. Directly using the loop IV as the source
+/// dim index would be out-of-bounds. Each loop IV value would need to be
+/// linearized with the result stride and then delinearized into source indices.
+// CHECK-LABEL: func.func private @negative_copy_into_strided_crosses_source_dim_boundary(
+func.func private @negative_copy_into_strided_crosses_source_dim_boundary(
+ %src : memref<1x5xf32>, %dst : memref<2x4xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 5], strides: [4, 1]
+ : memref<2x4xf32>
+ to memref<1x5xf32, strided<[4, 1]>>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<1x5xf32>
+ to memref<1x5xf32, strided<[4, 1]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @negative_copy_into_strided_rank_change(
+func.func private @negative_copy_into_strided_rank_change(%src : memref<2x3xf32>,
%dst : memref<6xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [2, 3], strides: [3, 1]
: memref<6xf32> to memref<2x3xf32>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<2x3xf32> to memref<2x3xf32>
return
}
-// CHECK-LABEL: func.func private @negative_copy_to_strided_dynamic_copy_source_shape(
-func.func private @negative_copy_to_strided_dynamic_copy_source_shape(%src : memref<?xf32>,
+// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_copy_source_shape(
+func.func private @negative_copy_into_strided_dynamic_copy_source_shape(%src : memref<?xf32>,
%dst : memref<4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [4], strides: [1]
: memref<4xf32> to memref<4xf32>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<?xf32> to memref<4xf32>
return
}
-// CHECK-LABEL: func.func private @negative_copy_to_strided_dynamic_rc_shapes(
-func.func private @negative_copy_to_strided_dynamic_rc_shapes(%dim : index,
+// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_rc_shapes(
+func.func private @negative_copy_into_strided_dynamic_rc_shapes(%dim : index,
%src : memref<4xf32>, %dst : memref<?xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [%dim], strides: [1]
: memref<?xf32> to memref<?xf32, strided<[1]>>
// CHECK: memref.copy %arg1, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<4xf32> to memref<?xf32, strided<[1]>>
return
}
-// CHECK-LABEL: func.func private @negative_copy_to_strided_unranked_rc_base(
-func.func private @negative_copy_to_strided_unranked_rc_base(
+// CHECK-LABEL: func.func private @negative_copy_into_strided_unranked_rc_base(
+func.func private @negative_copy_into_strided_unranked_rc_base(
%src : memref<4xf32>, %dst : memref<*xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [4], strides: [1]
: memref<*xf32> to memref<4xf32>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<4xf32> to memref<4xf32>
return
}
-// CHECK-LABEL: func.func private @negative_copy_to_strided_0d_base_2d_dynamic_offset(
-func.func private @negative_copy_to_strided_0d_base_2d_dynamic_offset(
+// CHECK-LABEL: func.func private @negative_copy_into_multidim_strided_dynamic_offset(
+func.func private @negative_copy_into_multidim_strided_dynamic_offset(
%offset : index, %src : memref<1x1xf32>, %dst : memref<4x8xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [%offset], sizes: [1, 1], strides: [8, 1]
: memref<4x8xf32> to memref<1x1xf32, strided<[8, 1], offset: ?>>
// CHECK: memref.copy %arg1, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x1xf32>
to memref<1x1xf32, strided<[8, 1], offset: ?>>
return
}
-// CHECK-LABEL: func.func private @negative_copy_to_strided_1d_base_2d_dynamic_offset(
-func.func private @negative_copy_to_strided_1d_base_2d_dynamic_offset(
- %offset : index, %src : memref<1x3x1xf32>,
- %dst : memref<1x3x11xf32>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
- %reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [1, 3, 1], strides: [33, 11, 1]
- : memref<1x3x11xf32>
- to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
-
- // CHECK: memref.copy %arg1, %reinterpret_cast
- // CHECK-NOT: memref.load
- // CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
- : memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
- return
-}
-
-// CHECK-LABEL: func.func private @negative_copy_to_strided_2d_dynamic_offset(
-func.func private @negative_copy_to_strided_2d_dynamic_offset(
+/// For effectively-1D copies with dynamic offsets, the single non-unit result
+/// dimension must map to the single non-unit source dimension. Otherwise the
+/// dynamic linear offset would need runtime delinearization before adding IVs.
+// CHECK-LABEL: func.func private @negative_copy_multidim_into_1D_strided_dynamic_offset(
+func.func private @negative_copy_multidim_into_1D_strided_dynamic_offset(
%offset : index, %src : memref<1x3x4xf32>,
- %dst : memref<1x3x11xf32>) {
+ %dst : memref<1x12x1xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
- %reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [1, 3, 4], strides: [33, 11, 1]
- : memref<1x3x11xf32>
- to memref<1x3x4xf32, strided<[33, 11, 1], offset: ?>>
+ %rc = memref.reinterpret_cast %dst
+ to offset: [%offset], sizes: [1, 3, 4], strides: [12, 4, 1]
+ : memref<1x12x1xf32>
+ to memref<1x3x4xf32, strided<[12, 4, 1], offset: ?>>
// CHECK: memref.copy %arg1, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x3x4xf32>
- to memref<1x3x4xf32, strided<[33, 11, 1], offset: ?>>
+ to memref<1x3x4xf32, strided<[12, 4, 1], offset: ?>>
return
}
/// Non-unit copied dimension needs stride-based address computation.
-// CHECK-LABEL: func.func private @negative_copy_to_strided_dynamic_rc_stride(
-func.func private @negative_copy_to_strided_dynamic_rc_stride(%stride : index,
+// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_rc_stride(
+func.func private @negative_copy_into_strided_dynamic_rc_stride(%stride : index,
%src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
- %reinterpret_cast = memref.reinterpret_cast %dst
+ %rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [1, 3, 1], strides: [33, %stride, 1]
: memref<1x3x11xf32>
to memref<1x3x1xf32, strided<[33, ?, 1]>>
@@ -416,7 +461,7 @@ func.func private @negative_copy_to_strided_dynamic_rc_stride(%stride : index,
// CHECK: memref.copy %arg1, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %reinterpret_cast
+ memref.copy %src, %rc
: memref<1x3x1xf32>
to memref<1x3x1xf32, strided<[33, ?, 1]>>
return
>From 96a6bb69f529416e1bda4e1bd4461933b6e311ea Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Mon, 29 Jun 2026 15:04:44 +0200
Subject: [PATCH 08/14] Address fourth round of comments
---
.../Transforms/ElideReinterpretCast.cpp | 24 +-
.../MemRef/elide-reinterpret-cast.mlir | 304 ++++++++++++------
2 files changed, 214 insertions(+), 114 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 435600f58f7fa..c79e522661350 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -51,8 +51,9 @@ struct AssocMapAndOffsetsForRC {
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.
+/// Maps each non-unit result dimension to a source dimension. Returns false if
+/// a stride of a non-unit dimension in the rc result is dynamic or no distinct
+/// source dimension matches.
static bool findSourceDimForResultDim(memref::ReinterpretCastOp rc,
AssocMapAndOffsetsForRC &mapAndOffs) {
MemRefType resType = dyn_cast<MemRefType>(rc.getType());
@@ -62,8 +63,9 @@ static bool findSourceDimForResultDim(memref::ReinterpretCastOp rc,
SmallVector<int64_t> srcIdentityStrides = computeStrides(srcType.getShape());
- // Reusing a source dimension would require delinearizing the combined linear
- // offset, which is TODO.
+ // Each non-unit result dimension needs one source dimension to receive its
+ // loop IV. Combining multiple IVs into one source index would require
+ // linearizing the result position which is TODO.
SmallVector<bool> usedSrcDims(srcType.getRank(), false);
for (auto [resultDim, resultSize] : llvm::enumerate(resType.getShape())) {
@@ -76,14 +78,16 @@ static bool findSourceDimForResultDim(memref::ReinterpretCastOp rc,
int64_t resultStride = rc.getStaticStrides()[resultDim];
std::optional<unsigned> srcDim;
- // Find an unused source dimension with matching stride and enough elements.
+ // Pick the first unused source dimension with the same stride.
for (auto [idx, stride] : llvm::enumerate(srcIdentityStrides)) {
- if (usedSrcDims[idx] || stride != resultStride ||
- srcType.getDimSize(idx) < resultSize)
+ if (usedSrcDims[idx] || stride != resultStride)
continue;
- if (!srcDim || srcType.getDimSize(idx) < srcType.getDimSize(*srcDim))
- srcDim = idx;
+ assert(srcType.getDimSize(idx) >= resultSize &&
+ "reinterpret_cast result dimension does not fit in the matching "
+ "source dimension");
+ srcDim = idx;
+ break;
}
if (!srcDim)
return false;
@@ -120,7 +124,7 @@ delinearizeStaticRCOffset(memref::ReinterpretCastOp rc) {
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
+ // Convert the scalar reinterpret_cast offset to per-dimension source starting
// indices.
for (auto [dim, stride] : llvm::enumerate(srcStrides)) {
offsetIdxs[dim] = remainder / stride;
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index eaa823e4afd3d..f0b449d59a1ae 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -3,10 +3,10 @@
//===----------------------------------------------------------------------===//
// Positive tests
+//
+// The destination is effectively a 1D array within a MemRef with rank >= 1
//===----------------------------------------------------------------------===//
-/// Effectively copied scalar within a MemRef with rank >= 0
-/// to practically a 1D array within a MemRef with rank >= 1
// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
@@ -42,9 +42,9 @@ func.func private @copy_scalar_into_1D_strided_nonzero_offset(%src : memref<1x1x
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK: %[[OFF:.*]] = arith.constant 1 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C1]]] : memref<1x108xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[OFF]]] : memref<1x108xf32>
memref.copy %src, %rc
: memref<1x1xf32>
to memref<1x1xf32, strided<[1, 1], offset: 1>>
@@ -77,7 +77,26 @@ func.func private @copy_scalar_into_1D_strided_dynamic_offset(%offset: index, %s
return
}
-// Dynamic strides are irrelevant because all strided memref indices are zero.
+// Strides are irrelevant because all strided memref indices are zero.
+// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_static_stride(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
+func.func private @copy_scalar_into_1D_strided_static_stride(%src : memref<1x1xf32>,
+ %dst : memref<1x108xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 1], strides: [107, 2]
+ : memref<1x108xf32> to memref<1x1xf32, strided<[107, 2]>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
+ memref.copy %src, %reinterpret_cast
+ : memref<1x1xf32> to memref<1x1xf32, strided<[107, 2]>>
+ return
+}
+
// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_dynamic_stride(
// CHECK-SAME: %[[STR0:[A-Za-z][A-Za-z0-9-]*]]: index
// CHECK-SAME: %[[STR1:[A-Za-z][A-Za-z0-9-]*]]: index
@@ -102,9 +121,41 @@ func.func private @copy_scalar_into_1D_strided_dynamic_stride(%stride0: index,
return
}
-// CHECK-LABEL: func.func private @copy_0D_into_0D_strided(
+// CHECK-LABEL: func.func private @copy_1D_into_1D_strided_dynamic_offset(
+// CHECK-SAME: %[[OFF:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<108xf32>
+func.func private @copy_1D_into_1D_strided_dynamic_offset(
+ %offset : index, %src : memref<4xf32>, %dst : memref<108xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [%offset], sizes: [4], strides: [1]
+ : memref<108xf32> to memref<4xf32, strided<[1], offset: ?>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[UB:.*]] = arith.constant 4 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX]] : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[DST_IDX]]] : memref<108xf32>
+ // CHECK: }
+ memref.copy %src, %rc
+ : memref<4xf32> to memref<4xf32, strided<[1], offset: ?>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// Positive tests
+//
+// The destination is effectively a scalar within a MemRef with rank == 0
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func private @copy_scalar_into_0D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<f32>, %[[DST:.*]]: memref<f32>
-func.func private @copy_0D_into_0D_strided(%src : memref<f32>, %dst : memref<f32>) {
+func.func private @copy_scalar_into_0D_strided_zero_offset(%src : memref<f32>, %dst : memref<f32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [], strides: []
@@ -118,31 +169,36 @@ func.func private @copy_0D_into_0D_strided(%src : memref<f32>, %dst : memref<f32
return
}
-/// Effectively a 2D array within a MemRef with rank >= 2
+//===----------------------------------------------------------------------===//
+// Positive tests
+//
+// The destination is effectively a 2D array within a MemRef with rank >= 2
+//===----------------------------------------------------------------------===//
+
// CHECK-LABEL: func.func private @copy_scalar_into_2D_strided(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x33x42xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
func.func private @copy_scalar_into_2D_strided(
- %src : memref<1x1x1xf32>, %dst : memref<1x33x42xf32>) {
+ %src : memref<1x1x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [1, 1, 1], strides: [1, 1, 1]
- : memref<1x33x42xf32>
+ : memref<1x3x11xf32>
to memref<1x1x1xf32>
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x33x42xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x3x11xf32>
memref.copy %src, %rc
: memref<1x1x1xf32> to memref<1x1x1xf32>
// CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_1D_into_1D_strided_zero_offset(
+// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_1D_into_1D_strided_zero_offset(
+func.func private @copy_1D_into_2D_strided_zero_offset(
%src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
@@ -153,8 +209,8 @@ func.func private @copy_1D_into_1D_strided_zero_offset(
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
+ // CHECK-DAG: %[[UB:.*]] = arith.constant 3 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x11xf32>
// CHECK: }
@@ -179,11 +235,11 @@ func.func private @copy_1D_into_2D_strided_nonzero_offset(
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
- // CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
+ // CHECK-DAG: %[[UB:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[OFF:.*]] = arith.constant 10 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C10]]] : memref<1x3x11xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[OFF]]] : memref<1x3x11xf32>
// CHECK: }
memref.copy %src, %rc
: memref<1x3x1xf32>
@@ -192,28 +248,31 @@ func.func private @copy_1D_into_2D_strided_nonzero_offset(
return
}
-// CHECK-LABEL: func.func private @copy_1D_into_1D_strided_dynamic_offset(
-// CHECK-SAME: %[[OFF:.*]]: index
-// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<42xf32>
-func.func private @copy_1D_into_1D_strided_dynamic_offset(
- %offset : index, %src : memref<4xf32>, %dst : memref<42xf32>) {
+// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_zero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_2D_into_2D_strided_zero_offset(
+ %src : memref<1x3x4xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [4], strides: [1]
- : memref<42xf32> to memref<4xf32, strided<[1], offset: ?>>
+ to offset: [0], sizes: [1, 3, 4], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1]>>
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
- // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX]] : index
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<4xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[DST_IDX]]] : memref<42xf32>
+ // CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[UB1:.*]] = arith.constant 4 : index
+ // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
+ // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x11xf32>
+ // CHECK: }
// CHECK: }
memref.copy %src, %rc
- : memref<4xf32> to memref<4xf32, strided<[1], offset: ?>>
+ : memref<1x3x4xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1]>>
// CHECK-NOT: memref.copy
return
}
@@ -232,12 +291,12 @@ func.func private @copy_2D_into_2D_strided_nonzero_offset(
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
- // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
- // CHECK-DAG: %[[C7:.*]] = arith.constant 7 : index
- // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
- // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
- // CHECK: %[[DST_IDX:.*]] = arith.addi %[[C7]], %[[IDX1]] : index
+ // CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[UB1:.*]] = arith.constant 4 : index
+ // CHECK-DAG: %[[OFF:.*]] = arith.constant 7 : index
+ // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
+ // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX1]] : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x4xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[DST_IDX]]] : memref<1x3x11xf32>
// CHECK: }
@@ -249,59 +308,102 @@ func.func private @copy_2D_into_2D_strided_nonzero_offset(
return
}
-//===----------------------------------------------------------------------===//
-// Negative tests (must NOT rewrite)
-//===----------------------------------------------------------------------===//
-
-// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity(
-func.func private @negative_copy_into_strided_non_identity(%src: memref<1x1xf32>,
- %dst: memref<8x1xf32, strided<[10, 2]>>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_rc_non_identity_strides(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x4x3xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x4xf32>
+func.func private @copy_2D_into_2D_strided_rc_non_identity_strides(
+ %src : memref<1x4x3xf32>, %dst : memref<1x3x4xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
- to offset: [6], sizes: [1, 1], strides: [11, 80]
- : memref<8x1xf32, strided<[10, 2]>>
- to memref<1x1xf32, strided<[11, 80], offset: 6>>
+ to offset: [0], sizes: [1, 4, 3], strides: [12, 1, 4]
+ : memref<1x3x4xf32>
+ to memref<1x4x3xf32, strided<[12, 1, 4]>>
- // CHECK: memref.copy %arg0, %reinterpret_cast
- // CHECK-NOT: memref.load
- // CHECK-NOT: memref.store
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[UB1:.*]] = arith.constant 4 : index
+ // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
+ // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x4x3xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX1]], %[[IDX0]]] : memref<1x3x4xf32>
+ // CHECK: }
+ // CHECK: }
memref.copy %src, %rc
- : memref<1x1xf32> to memref<1x1xf32, strided<[11, 80], offset: 6>>
+ : memref<1x4x3xf32>
+ to memref<1x4x3xf32, strided<[12, 1, 4]>>
+ // CHECK-NOT: memref.copy
+ return
+}
+/// Only strides in non-unit dimensions of reinterpret_cast result are checked.
+// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_rc_unit_dim_dynamic_stride(
+// CHECK-SAME: %{{[^:]+}}: index
+// CHECK-SAME: %{{[^:]+}}: index
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_1D_into_2D_strided_rc_unit_dim_dynamic_stride(
+ %stride0 : index, %stride2 : index,
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 3, 1], strides: [%stride0, 11, %stride2]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[?, 11, ?]>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[UB:.*]] = arith.constant 3 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x11xf32>
+ // CHECK: }
+ memref.copy %src, %rc
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[?, 11, ?]>>
+ // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @negative_copy_into_strided_overlapping_result_dims(
-func.func private @negative_copy_into_strided_overlapping_result_dims(%src : memref<3x3xf32>,
- %dst : memref<4x4xf32>) {
+//===----------------------------------------------------------------------===//
+// Negative tests (must NOT rewrite)
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity(
+func.func private @negative_copy_into_strided_non_identity(%src: memref<1x1xf32>,
+ %dst: memref<108x1xf32, strided<[10, 2]>>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [3, 3], strides: [1, 1]
- : memref<4x4xf32> to memref<3x3xf32, strided<[1, 1]>>
+ to offset: [6], sizes: [1, 1], strides: [10, 2]
+ : memref<108x1xf32, strided<[10, 2]>>
+ to memref<1x1xf32, strided<[10, 2], offset: 6>>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
- : memref<3x3xf32> to memref<3x3xf32, strided<[1, 1]>>
+ : memref<1x1xf32> to memref<1x1xf32, strided<[10, 2], offset: 6>>
+
return
}
// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_strided_unmatched_strides(
func.func private @negative_copy_1D_into_2D_strided_unmatched_strides(
- %src : memref<1x3x1xf32>, %dst : memref<1x4x8xf32>) {
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 3, 1], strides: [32, 5, 1]
- : memref<1x4x8xf32>
- to memref<1x3x1xf32, strided<[32, 5, 1]>>
+ to offset: [0], sizes: [1, 3, 1], strides: [33, 10, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, 10, 1]>>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
: memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[32, 5, 1]>>
+ to memref<1x3x1xf32, strided<[33, 10, 1]>>
return
}
@@ -323,67 +425,61 @@ func.func private @negative_copy_2D_into_2D_strided_unmatched_strides(
return
}
-/// The linear accessed region is in-bounds, but result dim 1 has size 5 while
-/// the matching source dim has size 4. Directly using the loop IV as the source
-/// dim index would be out-of-bounds. Each loop IV value would need to be
-/// linearized with the result stride and then delinearized into source indices.
-// CHECK-LABEL: func.func private @negative_copy_into_strided_crosses_source_dim_boundary(
-func.func private @negative_copy_into_strided_crosses_source_dim_boundary(
- %src : memref<1x5xf32>, %dst : memref<2x4xf32>) {
+// CHECK-LABEL: func.func private @negative_copy_into_strided_rank_change(
+func.func private @negative_copy_into_strided_rank_change(%src : memref<3x4xf32>,
+ %dst : memref<12xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 5], strides: [4, 1]
- : memref<2x4xf32>
- to memref<1x5xf32, strided<[4, 1]>>
+ to offset: [0], sizes: [3, 4], strides: [1, 1]
+ : memref<12xf32> to memref<3x4xf32, strided<[1, 1]>>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
- : memref<1x5xf32>
- to memref<1x5xf32, strided<[4, 1]>>
+ : memref<3x4xf32> to memref<3x4xf32, strided<[1, 1]>>
return
}
-// CHECK-LABEL: func.func private @negative_copy_into_strided_rank_change(
-func.func private @negative_copy_into_strided_rank_change(%src : memref<2x3xf32>,
- %dst : memref<6xf32>) {
+// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_copy_source_shape(
+func.func private @negative_copy_into_strided_dynamic_copy_source_shape(%src : memref<?xf32>,
+ %dst : memref<4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [2, 3], strides: [3, 1]
- : memref<6xf32> to memref<2x3xf32>
+ to offset: [0], sizes: [4], strides: [1]
+ : memref<4xf32> to memref<4xf32>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
- : memref<2x3xf32> to memref<2x3xf32>
+ : memref<?xf32> to memref<4xf32>
return
}
-// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_copy_source_shape(
-func.func private @negative_copy_into_strided_dynamic_copy_source_shape(%src : memref<?xf32>,
- %dst : memref<4xf32>) {
+// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_rc_source_shape(
+func.func private @negative_copy_into_strided_dynamic_rc_source_shape(
+ %src : memref<4xf32>, %dst : memref<?xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [4], strides: [1]
- : memref<4xf32> to memref<4xf32>
+ : memref<?xf32> to memref<4xf32, strided<[1]>>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
- : memref<?xf32> to memref<4xf32>
+ : memref<4xf32> to memref<4xf32, strided<[1]>>
return
}
-// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_rc_shapes(
-func.func private @negative_copy_into_strided_dynamic_rc_shapes(%dim : index,
- %src : memref<4xf32>, %dst : memref<?xf32>) {
+// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_rc_result_shape(
+func.func private @negative_copy_into_strided_dynamic_rc_result_shape(%dim : index,
+ %src : memref<4xf32>, %dst : memref<12xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [%dim], strides: [1]
- : memref<?xf32> to memref<?xf32, strided<[1]>>
+ : memref<12xf32> to memref<?xf32, strided<[1]>>
// CHECK: memref.copy %arg1, %reinterpret_cast
// CHECK-NOT: memref.load
@@ -411,18 +507,18 @@ func.func private @negative_copy_into_strided_unranked_rc_base(
// CHECK-LABEL: func.func private @negative_copy_into_multidim_strided_dynamic_offset(
func.func private @negative_copy_into_multidim_strided_dynamic_offset(
- %offset : index, %src : memref<1x1xf32>, %dst : memref<4x8xf32>) {
+ %offset : index, %src : memref<1x1xf32>, %dst : memref<3x4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%rc = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [1, 1], strides: [8, 1]
- : memref<4x8xf32> to memref<1x1xf32, strided<[8, 1], offset: ?>>
+ to offset: [%offset], sizes: [1, 1], strides: [12, 1]
+ : memref<3x4xf32> to memref<1x1xf32, strided<[12, 1], offset: ?>>
// CHECK: memref.copy %arg1, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
: memref<1x1xf32>
- to memref<1x1xf32, strided<[8, 1], offset: ?>>
+ to memref<1x1xf32, strided<[12, 1], offset: ?>>
return
}
@@ -449,21 +545,21 @@ func.func private @negative_copy_multidim_into_1D_strided_dynamic_offset(
}
/// Non-unit copied dimension needs stride-based address computation.
-// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_rc_stride(
-func.func private @negative_copy_into_strided_dynamic_rc_stride(%stride : index,
- %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
+// CHECK-LABEL: func.func private @negative_copy_into_strided_rc_non_unit_dim_dynamic_stride(
+func.func private @negative_copy_into_strided_rc_non_unit_dim_dynamic_stride(%stride : index,
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 3, 1], strides: [33, %stride, 1]
- : memref<1x3x11xf32>
- to memref<1x3x1xf32, strided<[33, ?, 1]>>
+ to offset: [0], sizes: [1, 3, 1], strides: [12, %stride, 1]
+ : memref<1x3x4xf32>
+ to memref<1x3x1xf32, strided<[12, ?, 1]>>
// CHECK: memref.copy %arg1, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
: memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[33, ?, 1]>>
+ to memref<1x3x1xf32, strided<[12, ?, 1]>>
return
}
>From d0c18498d7a5748cf17d8ac19694b92b77290c88 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Wed, 1 Jul 2026 14:13:18 +0200
Subject: [PATCH 09/14] Drop support for non-identical rc source and result
strides
---
.../Transforms/ElideReinterpretCast.cpp | 94 +++---
.../MemRef/elide-reinterpret-cast.mlir | 295 +++++++-----------
2 files changed, 156 insertions(+), 233 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index c79e522661350..57a10b0f24ce9 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -51,9 +51,9 @@ struct AssocMapAndOffsetsForRC {
std::optional<SmallVector<int64_t>> delinearizedOffsets;
};
-/// Maps each non-unit result dimension to a source dimension. Returns false if
-/// a stride of a non-unit dimension in the rc result is dynamic or no distinct
-/// source dimension matches.
+/// Maps each non-unit result dimension to a source dimension by stride. Returns
+/// false if a stride of a non-unit dimension in the rc result is dynamic or rc
+/// result strides are not equivalent to rc source strides.
static bool findSourceDimForResultDim(memref::ReinterpretCastOp rc,
AssocMapAndOffsetsForRC &mapAndOffs) {
MemRefType resType = dyn_cast<MemRefType>(rc.getType());
@@ -62,40 +62,34 @@ static bool findSourceDimForResultDim(memref::ReinterpretCastOp rc,
"Expecting identity source layout.");
SmallVector<int64_t> srcIdentityStrides = computeStrides(srcType.getShape());
+ ArrayRef<int64_t> rcResultStrides = rc.getStaticStrides();
- // Each non-unit result dimension needs one source dimension to receive its
- // loop IV. Combining multiple IVs into one source index would require
- // linearizing the result position which is TODO.
- SmallVector<bool> usedSrcDims(srcType.getRank(), false);
+ assert(srcType.getRank() == resType.getRank() &&
+ "Expecting rank-preserving reinterpret_casts");
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]))
+ int64_t resultStride = rcResultStrides[resultDim];
+
+ if (ShapedType::isDynamic(resultStride))
return false;
- int64_t resultStride = rc.getStaticStrides()[resultDim];
- std::optional<unsigned> srcDim;
- // Pick the first unused source dimension with the same stride.
- for (auto [idx, stride] : llvm::enumerate(srcIdentityStrides)) {
- if (usedSrcDims[idx] || stride != resultStride)
- continue;
-
- assert(srcType.getDimSize(idx) >= resultSize &&
- "reinterpret_cast result dimension does not fit in the matching "
- "source dimension");
- srcDim = idx;
- break;
- }
- if (!srcDim)
+ // For non-scalar strided memrefs, only support result strides identical to
+ // the identity strides of the source. This enables direct indexing into the
+ // same source dimensions, without linearization.
+ if (resultStride != srcIdentityStrides[resultDim])
return false;
- usedSrcDims[*srcDim] = true;
- mapAndOffs.assocMap.push_back(
- NonUnitDimAssocMapForRC{static_cast<unsigned>(resultDim), *srcDim});
+ assert(srcType.getDimSize(resultDim) >= resultSize &&
+ "reinterpret_cast result dimension does not fit in the matching "
+ "source dimension");
+
+ mapAndOffs.assocMap.push_back(NonUnitDimAssocMapForRC{
+ static_cast<unsigned>(resultDim), static_cast<unsigned>(resultDim)});
}
+
return true;
}
@@ -164,29 +158,28 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
/// Examples that return rewrite info:
///
/// // Scalar-shaped copy into a source with at most one non-unit dimension.
-/// There are
-/// // no non-unit result dimensions, so dynamic strides in the strided memref
-/// // do not affect index mapping.
+/// // There are no non-unit result dimensions, so strides in the strided
+/// // memref do not affect index mapping.
/// copy memref<1 x ... x 1 x f32>
/// to reinterpret_cast memref<source-shape>
/// to memref<1 x ... x 1 x f32, strided<[?, ..., ?], offset: ?>>
///
/// // Effectively-1D copy. The single non-unit strided memref dimension is
-/// // mapped to an identity-layout source dimension by its static stride.
+/// // mapped to its matching identity-layout stride source dimension.
/// copy memref<1 x ... x N x ... x 1 x f32>
/// to reinterpret_cast memref<source-shape>
/// to memref<1 x ... x N x ... x 1 x f32, strided<[..., S, ...]>>
///
/// // 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>
+/// // dimension is mapped in order, checking for matching source static
+/// stride. copy memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32>
/// to reinterpret_cast memref<source-shape>
/// to memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32,
/// strided<[..., S_0, ..., S_1, ...], offset: O>>
///
/// Examples that return no info:
///
-/// // Dynamic stride on a non-unit strided memref dimension.
+/// // Non-scalar copies: dynamic strides in strided memref.
/// copy memref<1xNxf32>
/// to reinterpret_cast memref<1xNxMxf32>
/// to memref<1xNxf32, strided<[?, ?]>>
@@ -220,24 +213,31 @@ getAssocMapAndOffsetsForRC(memref::ReinterpretCastOp rc) {
AssocMapAndOffsetsForRC mapAndOffs;
- // reinterpret_cast result dimensions must map to distinct source dimensions.
- if (!findSourceDimForResultDim(rc, mapAndOffs))
+ assert(resType.hasStaticShape() && "expected static shape");
+ // For scalar copies, result strides are irrelevant, including dynamic ones.
+ // For non-scalar copies, require static result strides identical to the
+ // identity strides of the reinterpret_cast source.
+ if (!llvm::all_of(resType.getShape(),
+ [](int64_t size) { return size == 1; }) &&
+ !findSourceDimForResultDim(rc, mapAndOffs))
return std::nullopt;
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");
- // Static ReinterpretCast offset
+
+ // CASE 1: Static ReinterpretCast offset
if (ShapedType::isStatic(rcOffsets[0])) {
// Delinearize static ReinterpretCast offset as in-bounds indices (one for
// every source dimension).
mapAndOffs.delinearizedOffsets = delinearizeStaticRCOffset(rc);
- assert(
- mapAndOffs.delinearizedOffsets &&
- "static reinterpret_cast offset must delinearize to in-bounds source "
- "indices");
+ assert(mapAndOffs.delinearizedOffsets &&
+ "static reinterpret_cast offset must delinearize to in-bounds "
+ "reinterpret_cast source indices");
+ // Relevant for non-scalar copies: assert that the rectangular
+ // copied slice is in bounds.
assert(
llvm::all_of(
mapAndOffs.assocMap,
@@ -250,17 +250,21 @@ getAssocMapAndOffsetsForRC(memref::ReinterpretCastOp rc) {
return mapAndOffs;
}
- // Dynamic ReinterpretCast offset.
+ // CASE 2: Dynamic ReinterpretCast offset.
// TODO: Support dynamic offsets into sources with multiple non-unit
// dimensions by delinearizing the offset into source start indices at runtime
// before adding loop IVs.
- if (mapAndOffs.assocMap.size() > 1)
- return std::nullopt;
- // Only sources with a single non-unit dimension can receive a dynamic offset
- // directly.
+
+ // With an effectively-1D source, a dynamic linear offset can be used directly
+ // as the index of the unique non-unit source dimension.
if (!getSingleNonUnitDim(srcType))
return std::nullopt;
+ // Non-scalar copies require identical strides and no rank-changing,
+ // so there can be at most one non-unit result dimension in this case.
+ assert(mapAndOffs.assocMap.size() <= 1 &&
+ "effectively-1D source cannot have multiple mapped non-unit dims");
+
return mapAndOffs;
}
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index f0b449d59a1ae..908cfdcfbcc7f 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -3,10 +3,28 @@
//===----------------------------------------------------------------------===//
// Positive tests
-//
-// The destination is effectively a 1D array within a MemRef with rank >= 1
+//
+// Scalar (0D) copy
//===----------------------------------------------------------------------===//
+// The destination is effectively a scalar within a MemRef with rank == 0
+// CHECK-LABEL: func.func private @copy_scalar_into_0D_strided(
+// CHECK-SAME: %[[SRC:.*]]: memref<f32>, %[[DST:.*]]: memref<f32>
+func.func private @copy_scalar_into_0D_strided(%src : memref<f32>, %dst : memref<f32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [], strides: []
+ : memref<f32> to memref<f32>
+
+ // CHECK-NOT: memref.copy
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][] : memref<f32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][] : memref<f32>
+ memref.copy %src, %rc : memref<f32> to memref<f32>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+/// The destination is effectively a 1D array within a MemRef with rank >= 1
// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
@@ -77,11 +95,12 @@ func.func private @copy_scalar_into_1D_strided_dynamic_offset(%offset: index, %s
return
}
-// Strides are irrelevant because all strided memref indices are zero.
-// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_static_stride(
+// Scalar copies have no varying result dimensions, so rc result strides do
+// not affect the copy destination address and are ignored.
+// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_non_identity_stride(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
-func.func private @copy_scalar_into_1D_strided_static_stride(%src : memref<1x1xf32>,
+func.func private @copy_scalar_into_1D_strided_non_identity_stride(%src : memref<1x1xf32>,
%dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
@@ -121,64 +140,11 @@ func.func private @copy_scalar_into_1D_strided_dynamic_stride(%stride0: index,
return
}
-// CHECK-LABEL: func.func private @copy_1D_into_1D_strided_dynamic_offset(
-// CHECK-SAME: %[[OFF:.*]]: index
-// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<108xf32>
-func.func private @copy_1D_into_1D_strided_dynamic_offset(
- %offset : index, %src : memref<4xf32>, %dst : memref<108xf32>) {
- // CHECK-NOT: memref.reinterpret_cast
- %rc = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [4], strides: [1]
- : memref<108xf32> to memref<4xf32, strided<[1], offset: ?>>
-
- // CHECK-NOT: memref.copy
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[UB:.*]] = arith.constant 4 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
- // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX]] : index
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<4xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[DST_IDX]]] : memref<108xf32>
- // CHECK: }
- memref.copy %src, %rc
- : memref<4xf32> to memref<4xf32, strided<[1], offset: ?>>
- // CHECK-NOT: memref.copy
- return
-}
-
-//===----------------------------------------------------------------------===//
-// Positive tests
-//
-// The destination is effectively a scalar within a MemRef with rank == 0
-//===----------------------------------------------------------------------===//
-
-// CHECK-LABEL: func.func private @copy_scalar_into_0D_strided_zero_offset(
-// CHECK-SAME: %[[SRC:.*]]: memref<f32>, %[[DST:.*]]: memref<f32>
-func.func private @copy_scalar_into_0D_strided_zero_offset(%src : memref<f32>, %dst : memref<f32>) {
- // CHECK-NOT: memref.reinterpret_cast
- %rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [], strides: []
- : memref<f32> to memref<f32>
-
- // CHECK-NOT: memref.copy
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][] : memref<f32>
- // CHECK: memref.store %[[VAL]], %[[DST]][] : memref<f32>
- memref.copy %src, %rc : memref<f32> to memref<f32>
- // CHECK-NOT: memref.copy
- return
-}
-
-//===----------------------------------------------------------------------===//
-// Positive tests
-//
-// The destination is effectively a 2D array within a MemRef with rank >= 2
-//===----------------------------------------------------------------------===//
-
-// CHECK-LABEL: func.func private @copy_scalar_into_2D_strided(
+/// The destination is effectively a 2D array within a MemRef with rank >= 2
+// CHECK-LABEL: func.func private @copy_scalar_into_2D_strided_non_identity_stride(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_scalar_into_2D_strided(
+func.func private @copy_scalar_into_2D_strided_non_identity_stride(
%src : memref<1x1x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
@@ -195,6 +161,12 @@ func.func private @copy_scalar_into_2D_strided(
return
}
+//===----------------------------------------------------------------------===//
+// Positive tests
+//
+// Non-scalar (ND) copy
+//===----------------------------------------------------------------------===//
+
// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
@@ -308,139 +280,107 @@ func.func private @copy_2D_into_2D_strided_nonzero_offset(
return
}
-// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_rc_non_identity_strides(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x4x3xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x3x4xf32>
-func.func private @copy_2D_into_2D_strided_rc_non_identity_strides(
- %src : memref<1x4x3xf32>, %dst : memref<1x3x4xf32>) {
- // CHECK-NOT: memref.reinterpret_cast
- %rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 4, 3], strides: [12, 1, 4]
- : memref<1x3x4xf32>
- to memref<1x4x3xf32, strided<[12, 1, 4]>>
-
- // CHECK-NOT: memref.copy
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
- // CHECK-DAG: %[[UB1:.*]] = arith.constant 4 : index
- // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
- // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x4x3xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX1]], %[[IDX0]]] : memref<1x3x4xf32>
- // CHECK: }
- // CHECK: }
- memref.copy %src, %rc
- : memref<1x4x3xf32>
- to memref<1x4x3xf32, strided<[12, 1, 4]>>
- // CHECK-NOT: memref.copy
- return
-}
-
-/// Only strides in non-unit dimensions of reinterpret_cast result are checked.
-// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_rc_unit_dim_dynamic_stride(
-// CHECK-SAME: %{{[^:]+}}: index
-// CHECK-SAME: %{{[^:]+}}: index
-// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_1D_into_2D_strided_rc_unit_dim_dynamic_stride(
- %stride0 : index, %stride2 : index,
- %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
+/// rc result dynamic offset:
+/// supported only for effectively-1D rc source
+/// (runtime delinearization not implemented)
+// CHECK-LABEL: func.func private @copy_1D_into_1D_strided_dynamic_offset(
+// CHECK-SAME: %[[OFF:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<108xf32>
+func.func private @copy_1D_into_1D_strided_dynamic_offset(
+ %offset : index, %src : memref<4xf32>, %dst : memref<108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 3, 1], strides: [%stride0, 11, %stride2]
- : memref<1x3x11xf32>
- to memref<1x3x1xf32, strided<[?, 11, ?]>>
+ to offset: [%offset], sizes: [4], strides: [1]
+ : memref<108xf32> to memref<4xf32, strided<[1], offset: ?>>
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[UB:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[UB:.*]] = arith.constant 4 : index
// CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x11xf32>
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX]] : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[DST_IDX]]] : memref<108xf32>
// CHECK: }
memref.copy %src, %rc
- : memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[?, 11, ?]>>
+ : memref<4xf32> to memref<4xf32, strided<[1], offset: ?>>
// CHECK-NOT: memref.copy
return
}
//===----------------------------------------------------------------------===//
// Negative tests (must NOT rewrite)
+//
+// Either scalar (0D) OR non-scalar (ND) copy
//===----------------------------------------------------------------------===//
-// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity(
-func.func private @negative_copy_into_strided_non_identity(%src: memref<1x1xf32>,
- %dst: memref<108x1xf32, strided<[10, 2]>>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
- %rc = memref.reinterpret_cast %dst
- to offset: [6], sizes: [1, 1], strides: [10, 2]
- : memref<108x1xf32, strided<[10, 2]>>
- to memref<1x1xf32, strided<[10, 2], offset: 6>>
-
- // CHECK: memref.copy %arg0, %reinterpret_cast
+/// Reject copies that don't target a strided memref
+// CHECK-LABEL: func.func private @negative_plain_copy(
+func.func private @negative_plain_copy(%src : memref<1x1xf32>,
+ %dst : memref<1x1xf32>) {
+ // CHECK: memref.copy %arg0, %arg1
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
- memref.copy %src, %rc
- : memref<1x1xf32> to memref<1x1xf32, strided<[10, 2], offset: 6>>
-
+ memref.copy %src, %dst
+ : memref<1x1xf32> to memref<1x1xf32>
return
}
-// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_strided_unmatched_strides(
-func.func private @negative_copy_1D_into_2D_strided_unmatched_strides(
- %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
+/// Reject unranked memref operands
+// CHECK-LABEL: func.func private @negative_copy_into_strided_unranked_rc_base(
+func.func private @negative_copy_into_strided_unranked_rc_base(
+ %src : memref<4xf32>, %dst : memref<*xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 3, 1], strides: [33, 10, 1]
- : memref<1x3x11xf32>
- to memref<1x3x1xf32, strided<[33, 10, 1]>>
+ to offset: [0], sizes: [4], strides: [1]
+ : memref<*xf32> to memref<4xf32>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
- : memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[33, 10, 1]>>
+ : memref<4xf32> to memref<4xf32>
return
}
-// CHECK-LABEL: func.func private @negative_copy_2D_into_2D_strided_unmatched_strides(
-func.func private @negative_copy_2D_into_2D_strided_unmatched_strides(
- %src : memref<1x3x4xf32>, %dst : memref<1x4x3xf32>) {
+/// Reject rank-changing reinterpet_casts
+// CHECK-LABEL: func.func private @negative_copy_into_strided_rank_change(
+func.func private @negative_copy_into_strided_rank_change(%src : memref<3x4xf32>,
+ %dst : memref<12xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 3, 4], strides: [12, 4, 1]
- : memref<1x4x3xf32>
- to memref<1x3x4xf32, strided<[12, 4, 1]>>
+ to offset: [0], sizes: [3, 4], strides: [1, 1]
+ : memref<12xf32> to memref<3x4xf32, strided<[1, 1]>>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
- : memref<1x3x4xf32>
- to memref<1x3x4xf32, strided<[12, 4, 1]>>
+ : memref<3x4xf32> to memref<3x4xf32, strided<[1, 1]>>
return
}
-// CHECK-LABEL: func.func private @negative_copy_into_strided_rank_change(
-func.func private @negative_copy_into_strided_rank_change(%src : memref<3x4xf32>,
- %dst : memref<12xf32>) {
+/// Reject non-identity layout rc source strides
+// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity(
+func.func private @negative_copy_into_strided_non_identity(%src: memref<1x1xf32>,
+ %dst: memref<108x1xf32, strided<[10, 2]>>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [3, 4], strides: [1, 1]
- : memref<12xf32> to memref<3x4xf32, strided<[1, 1]>>
+ to offset: [6], sizes: [1, 1], strides: [10, 2]
+ : memref<108x1xf32, strided<[10, 2]>>
+ to memref<1x1xf32, strided<[10, 2], offset: 6>>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
- : memref<3x4xf32> to memref<3x4xf32, strided<[1, 1]>>
+ : memref<1x1xf32> to memref<1x1xf32, strided<[10, 2], offset: 6>>
+
return
}
+/// Reject dynamic shapes
// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_copy_source_shape(
func.func private @negative_copy_into_strided_dynamic_copy_source_shape(%src : memref<?xf32>,
%dst : memref<4xf32>) {
@@ -489,24 +429,10 @@ func.func private @negative_copy_into_strided_dynamic_rc_result_shape(%dim : ind
return
}
-// CHECK-LABEL: func.func private @negative_copy_into_strided_unranked_rc_base(
-func.func private @negative_copy_into_strided_unranked_rc_base(
- %src : memref<4xf32>, %dst : memref<*xf32>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
- %rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [4], strides: [1]
- : memref<*xf32> to memref<4xf32>
-
- // CHECK: memref.copy %arg0, %reinterpret_cast
- // CHECK-NOT: memref.load
- // CHECK-NOT: memref.store
- memref.copy %src, %rc
- : memref<4xf32> to memref<4xf32>
- return
-}
-
-// CHECK-LABEL: func.func private @negative_copy_into_multidim_strided_dynamic_offset(
-func.func private @negative_copy_into_multidim_strided_dynamic_offset(
+/// Reject dynamic offsets for rc sources with > 1 non-unit dimension -
+/// runtime delinearization of these offsets is TODO.
+// CHECK-LABEL: func.func private @negative_copy_into_ND_strided_dynamic_offset(
+func.func private @negative_copy_into_ND_strided_dynamic_offset(
%offset : index, %src : memref<1x1xf32>, %dst : memref<3x4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%rc = memref.reinterpret_cast %dst
@@ -522,31 +448,35 @@ func.func private @negative_copy_into_multidim_strided_dynamic_offset(
return
}
-/// For effectively-1D copies with dynamic offsets, the single non-unit result
-/// dimension must map to the single non-unit source dimension. Otherwise the
-/// dynamic linear offset would need runtime delinearization before adding IVs.
-// CHECK-LABEL: func.func private @negative_copy_multidim_into_1D_strided_dynamic_offset(
-func.func private @negative_copy_multidim_into_1D_strided_dynamic_offset(
- %offset : index, %src : memref<1x3x4xf32>,
- %dst : memref<1x12x1xf32>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
+//===----------------------------------------------------------------------===//
+// Negative tests
+//
+// Non-scalar (ND) copy
+//===----------------------------------------------------------------------===//
+
+/// Reject rc result strides that not equal to rc source identity strides.
+/// (non-unit copied dimension needs stride-based address computation)
+// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity_strides(
+func.func private @negative_copy_into_strided_non_identity_strides(
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [1, 3, 4], strides: [12, 4, 1]
- : memref<1x12x1xf32>
- to memref<1x3x4xf32, strided<[12, 4, 1], offset: ?>>
+ to offset: [0], sizes: [1, 3, 1], strides: [33, 10, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, 10, 1]>>
- // CHECK: memref.copy %arg1, %reinterpret_cast
+ // CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
- : memref<1x3x4xf32>
- to memref<1x3x4xf32, strided<[12, 4, 1], offset: ?>>
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[33, 10, 1]>>
return
}
-/// Non-unit copied dimension needs stride-based address computation.
-// CHECK-LABEL: func.func private @negative_copy_into_strided_rc_non_unit_dim_dynamic_stride(
-func.func private @negative_copy_into_strided_rc_non_unit_dim_dynamic_stride(%stride : index,
+/// Reject dynamic rc result strides.
+// CHECK-LABEL: func.func private @negative_copy_ND_into_strided_dynamic_stride(
+func.func private @negative_copy_ND_into_strided_dynamic_stride(%stride : index,
%src : memref<1x3x1xf32>, %dst : memref<1x3x4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%rc = memref.reinterpret_cast %dst
@@ -563,17 +493,6 @@ func.func private @negative_copy_into_strided_rc_non_unit_dim_dynamic_stride(%st
return
}
-// CHECK-LABEL: func.func private @negative_plain_copy(
-func.func private @negative_plain_copy(%src : memref<1x1xf32>,
- %dst : memref<1x1xf32>) {
- // CHECK: memref.copy %arg0, %arg1
- // CHECK-NOT: memref.load
- // CHECK-NOT: memref.store
- memref.copy %src, %dst
- : memref<1x1xf32> to memref<1x1xf32>
- return
-}
-
// -----
//===----------------------------------------------------------------------===//
>From c4437cd9d1dde7e1bb8447e075ca7b53edcfa6d1 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Thu, 2 Jul 2026 09:47:56 +0200
Subject: [PATCH 10/14] Fixup
Remove reduntant dimension mapping, refine comments and tests
---
.../Transforms/ElideReinterpretCast.cpp | 304 ++++++++++--------
.../MemRef/elide-reinterpret-cast.mlir | 32 +-
2 files changed, 178 insertions(+), 158 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 57a10b0f24ce9..483d5c94ea97e 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -35,64 +35,15 @@ namespace {
// 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 {
+struct ResultNonUnitDimsAndOffsetsForRC {
// Non-unit dimensions of the reinterpret_cast result.
- SmallVector<NonUnitDimAssocMapForRC> assocMap;
+ SmallVector<unsigned> nonUnitDimsPos;
// Delinearized offsets to in-bounds reinterpret_cast source indices.
// Optional since it is only supported for static offsets.
std::optional<SmallVector<int64_t>> delinearizedOffsets;
};
-/// Maps each non-unit result dimension to a source dimension by stride. Returns
-/// false if a stride of a non-unit dimension in the rc result is dynamic or rc
-/// result strides are not equivalent to rc source strides.
-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());
- ArrayRef<int64_t> rcResultStrides = rc.getStaticStrides();
-
- assert(srcType.getRank() == resType.getRank() &&
- "Expecting rank-preserving reinterpret_casts");
-
- for (auto [resultDim, resultSize] : llvm::enumerate(resType.getShape())) {
- if (resultSize == 1)
- continue;
-
- int64_t resultStride = rcResultStrides[resultDim];
-
- if (ShapedType::isDynamic(resultStride))
- return false;
-
- // For non-scalar strided memrefs, only support result strides identical to
- // the identity strides of the source. This enables direct indexing into the
- // same source dimensions, without linearization.
- if (resultStride != srcIdentityStrides[resultDim])
- return false;
-
- assert(srcType.getDimSize(resultDim) >= resultSize &&
- "reinterpret_cast result dimension does not fit in the matching "
- "source dimension");
-
- mapAndOffs.assocMap.push_back(NonUnitDimAssocMapForRC{
- static_cast<unsigned>(resultDim), static_cast<unsigned>(resultDim)});
- }
-
- return true;
-}
-
/// Returns source indices for a static reinterpret_cast offset of an
/// identity-layout source.
static std::optional<SmallVector<int64_t>>
@@ -157,39 +108,93 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
///
/// Examples that return rewrite info:
///
-/// // Scalar-shaped copy into a source with at most one non-unit dimension.
-/// // There are no non-unit result dimensions, so strides in the strided
-/// // memref do not affect index mapping.
+/// // Scalar-shaped copy. There are no non-unit result dimensions, so result
+/// // strides do not affect index mapping - may be static or
+/// // dynamic.
+/// copy memref<1 x ... x 1 x f32>
+/// to reinterpret_cast memref<source-shape, identity-layout>
+/// to memref<1 x ... x 1 x f32,
+/// strided<[?, ..., ?], offset: O>>
+///
+/// // Scalar-shaped copy with dynamic offset into an effectively-1D source.
+/// // The dynamic offset can be used directly as the index of the source's
+/// // unique non-unit dimension.
/// copy memref<1 x ... x 1 x f32>
-/// to reinterpret_cast memref<source-shape>
-/// to memref<1 x ... x 1 x f32, strided<[?, ..., ?], offset: ?>>
+/// to reinterpret_cast memref<1 x ... x M x ... x 1 x f32,
+/// identity-layout>
+/// to memref<1 x ... x 1 x f32,
+/// strided<[?, ..., ?], offset: ?>>
+///
+/// // Non-scalar effectively-1D copy with static offset. Result strides must
+/// // be static and identical to the identity strides of the source.
+/// copy memref<1 x ... x N x ... x 1 x f32>
+/// to reinterpret_cast memref<1 x ... x M x ... x 1 x f32,
+/// identity-layout>
+/// to memref<1 x ... x N x ... x 1 x f32,
+/// strided<[source-identity-strides], offset: O>>
///
-/// // Effectively-1D copy. The single non-unit strided memref dimension is
-/// // mapped to its matching identity-layout stride source dimension.
+/// // Non-scalar effectively-1D copy with dynamic offset into an
+/// // effectively-1D source. Runtime delinearization is not needed because
+/// the
+/// // source has a unique non-unit dimension.
/// copy memref<1 x ... x N x ... x 1 x f32>
-/// to reinterpret_cast memref<source-shape>
-/// to memref<1 x ... x N x ... x 1 x f32, strided<[..., S, ...]>>
+/// to reinterpret_cast memref<1 x ... x M x ... x 1 x f32,
+/// identity-layout>
+/// to memref<1 x ... x N x ... x 1 x f32,
+/// strided<[source-identity-strides], offset: ?>>
///
-/// // Multidimensional copy with static offset. Each non-unit strided memref
-/// // dimension is mapped in order, checking for matching source static
-/// stride. copy memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32>
-/// to reinterpret_cast memref<source-shape>
+/// // Non-scalar multidimensional copy with static offset. Result strides
+/// must
+/// // be static and identical to the identity strides of the source.
+/// copy memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32>
+/// to reinterpret_cast memref<source-shape, identity-layout>
/// to memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32,
-/// strided<[..., S_0, ..., S_1, ...], offset: O>>
+/// strided<[source-identity-strides], offset: O>>
///
/// Examples that return no info:
///
-/// // Non-scalar copies: dynamic strides in strided memref.
+/// // Rank-changing reinterpret_casts are not supported.
+/// copy memref<1xNxf32>
+/// to reinterpret_cast memref<Mxf32, identity-layout>
+/// to memref<1xNxf32, strided<[N, 1]>>
+///
+/// // Dynamic shapes are not supported.
+/// copy memref<?xNxf32>
+/// to reinterpret_cast memref<?xMxf32, identity-layout>
+/// to memref<?xNxf32, strided<[M, 1]>>
+///
+/// // Non-identity source layouts are not supported.
+/// copy memref<1xNxf32>
+/// to reinterpret_cast memref<1xMxf32, strided<[S, 1]>>
+/// to memref<1xNxf32, strided<[M, 1]>>
+///
+/// // Dynamic offset into a source with more than one non-unit dimension is
+/// // not supported because runtime delinearization is not implemented.
+/// copy memref<1x1xf32>
+/// to reinterpret_cast memref<1xNxMxf32, identity-layout>
+/// to memref<1x1xf32, strided<[?, ?], offset: ?>>
+///
+/// // Non-scalar copies with dynamic result strides are not supported.
+/// copy memref<1xNxf32>
+/// to reinterpret_cast memref<1xMxf32, identity-layout>
+/// to memref<1xNxf32, strided<[?, 1]>>
+///
+/// // Non-scalar copies with result strides different from the source
+/// identity
+/// // strides are not supported.
/// copy memref<1xNxf32>
-/// to reinterpret_cast memref<1xNxMxf32>
-/// to memref<1xNxf32, strided<[?, ?]>>
+/// to reinterpret_cast memref<1xMxf32, identity-layout>
+/// to memref<1xNxf32, strided<[S, 1]>>
///
-/// // Multidimensional copy with dynamic linear offset.
+/// // Multidimensional non-scalar copies with dynamic offset are not
+/// supported
+/// // unless the source is effectively 1D.
/// copy memref<1xNxKxf32>
-/// to reinterpret_cast memref<1xNxMxf32>
-/// to memref<1xNxKxf32, strided<[N*M, M, 1], offset: ?>>
-static std::optional<AssocMapAndOffsetsForRC>
-getAssocMapAndOffsetsForRC(memref::ReinterpretCastOp rc) {
+/// to reinterpret_cast memref<1xNxMxf32, identity-layout>
+/// to memref<1xNxKxf32,
+/// strided<[N*M, M, 1], offset: ?>>
+static std::optional<ResultNonUnitDimsAndOffsetsForRC>
+getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
MemRefType resType = dyn_cast<MemRefType>(rc.getType());
@@ -211,16 +216,34 @@ getAssocMapAndOffsetsForRC(memref::ReinterpretCastOp rc) {
if (!srcType.getLayout().isIdentity())
return std::nullopt;
- AssocMapAndOffsetsForRC mapAndOffs;
+ ResultNonUnitDimsAndOffsetsForRC dimsAndOffs;
assert(resType.hasStaticShape() && "expected static shape");
// For scalar copies, result strides are irrelevant, including dynamic ones.
// For non-scalar copies, require static result strides identical to the
// identity strides of the reinterpret_cast source.
if (!llvm::all_of(resType.getShape(),
- [](int64_t size) { return size == 1; }) &&
- !findSourceDimForResultDim(rc, mapAndOffs))
- return std::nullopt;
+ [](int64_t size) { return size == 1; })) {
+ SmallVector<int64_t> srcIdentityStrides =
+ computeStrides(srcType.getShape());
+ ArrayRef<int64_t> rcResultStrides = rc.getStaticStrides();
+
+ assert((srcIdentityStrides.size() == rcResultStrides.size()) &&
+ "Expecting same number of strides for rank-preserving "
+ "reinterpret_casts.");
+
+ if (!llvm::all_of(llvm::zip_equal(srcIdentityStrides, rcResultStrides),
+ [](auto pair) {
+ auto [srcStride, resultStride] = pair;
+ return !ShapedType::isDynamic(resultStride) &&
+ srcStride == resultStride;
+ }))
+ return std::nullopt;
+ for (auto [dim, resultSize] : llvm::enumerate(resType.getShape())) {
+ if (resultSize != 1)
+ dimsAndOffs.nonUnitDimsPos.push_back(static_cast<unsigned>(dim));
+ }
+ }
ArrayRef<int64_t> rcOffsets = rc.getStaticOffsets();
// FIXME: Despite what `getStaticOffsets` implies, `reinterpret_cast` takes
@@ -231,23 +254,23 @@ getAssocMapAndOffsetsForRC(memref::ReinterpretCastOp rc) {
if (ShapedType::isStatic(rcOffsets[0])) {
// Delinearize static ReinterpretCast offset as in-bounds indices (one for
// every source dimension).
- mapAndOffs.delinearizedOffsets = delinearizeStaticRCOffset(rc);
- assert(mapAndOffs.delinearizedOffsets &&
+ dimsAndOffs.delinearizedOffsets = delinearizeStaticRCOffset(rc);
+ assert(dimsAndOffs.delinearizedOffsets &&
"static reinterpret_cast offset must delinearize to in-bounds "
"reinterpret_cast source indices");
// Relevant for non-scalar copies: assert that the rectangular
// copied slice is in bounds.
- assert(
- llvm::all_of(
- mapAndOffs.assocMap,
- [&](const NonUnitDimAssocMapForRC &assocMap) {
- return (*mapAndOffs.delinearizedOffsets)[assocMap.sourceDimPos] +
- resType.getDimSize(assocMap.resultDimPos) <=
- srcType.getDimSize(assocMap.sourceDimPos);
- }) &&
- "reinterpret_cast metadata describes an invalid accessible region");
- return mapAndOffs;
+ assert(llvm::all_of(llvm::enumerate(resType.getShape()),
+ [&](auto it) {
+ unsigned dim = it.index();
+ int64_t resultSize = it.value();
+ return (*dimsAndOffs.delinearizedOffsets)[dim] +
+ resultSize <=
+ srcType.getDimSize(dim);
+ }) &&
+ "reinterpret_cast metadata describes an invalid accessible region");
+ return dimsAndOffs;
}
// CASE 2: Dynamic ReinterpretCast offset.
@@ -262,10 +285,10 @@ getAssocMapAndOffsetsForRC(memref::ReinterpretCastOp rc) {
// Non-scalar copies require identical strides and no rank-changing,
// so there can be at most one non-unit result dimension in this case.
- assert(mapAndOffs.assocMap.size() <= 1 &&
- "effectively-1D source cannot have multiple mapped non-unit dims");
+ assert(dimsAndOffs.nonUnitDimsPos.size() <= 1 &&
+ "effectively-1D source cannot have multiple non-unit result dims");
- return mapAndOffs;
+ return dimsAndOffs;
}
/// Rewrites supported copy operations through `memref.reinterpret_cast` to
@@ -314,21 +337,23 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
if (!cpSrcType || !cpSrcType.hasStaticShape())
return rewriter.notifyMatchFailure(
op, "only ranked, static copy sources are supported.");
+
Value rcOutput = op.getTarget();
auto rc = rcOutput.getDefiningOp<memref::ReinterpretCastOp>();
if (!rc)
return rewriter.notifyMatchFailure(
op, "target is not a memref.reinterpret_cast");
- std::optional<AssocMapAndOffsetsForRC> mapAndOffs =
- getAssocMapAndOffsetsForRC(rc);
- if (!mapAndOffs)
+ std::optional<ResultNonUnitDimsAndOffsetsForRC> dimsAndOffs =
+ getResultNonUnitDimsAndOffsetsForRC(rc);
+ if (!dimsAndOffs)
return rewriter.notifyMatchFailure(
op, "reinterpret_cast does not match scalar or loop copy region");
Location loc = op.getLoc();
Value dst = rc.getSource();
MemRefType dstType = cast<MemRefType>(dst.getType());
+ MemRefType rcResType = cast<MemRefType>(rc.getType());
// Reuse common index constants across bounds, steps, and static offsets,
// but avoid creating them for rank-0 copies.
@@ -342,6 +367,7 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
}
return arith::ConstantIndexOp::create(rewriter, loc, value);
};
+
auto getZeroIdxs = [&](int64_t rank) {
SmallVector<Value> idxs;
idxs.reserve(rank);
@@ -353,27 +379,28 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
// Create loop bounds before moving the insertion point into the loop nest,
// so loop-invariant constants are emitted outside the generated loops.
SmallVector<Value> upperBounds;
- upperBounds.reserve(mapAndOffs->assocMap.size());
- MemRefType rcResType = dyn_cast<MemRefType>(rc.getType());
- for (const NonUnitDimAssocMapForRC &assocMap : mapAndOffs->assocMap)
- upperBounds.push_back(getOrCreateIndexConstant(
- rcResType.getDimSize(assocMap.resultDimPos)));
+ upperBounds.reserve(dimsAndOffs->nonUnitDimsPos.size());
+ for (unsigned dim : dimsAndOffs->nonUnitDimsPos) {
+ upperBounds.push_back(
+ getOrCreateIndexConstant(rcResType.getDimSize(dim)));
+ }
SmallVector<Value> rcSrcStoreIdxs = getZeroIdxs(dstType.getRank());
std::optional<unsigned> srcNonUnitDimPos;
- // Static offset has been delinearized in function gating rewrite.
- if (mapAndOffs->delinearizedOffsets) {
+ if (dimsAndOffs->delinearizedOffsets) {
+ // Initialize store indices from the static reinterpret_cast offset,
+ // delinearized in function gating rewrite.
for (auto [idx, offset] :
- llvm::enumerate(*mapAndOffs->delinearizedOffsets)) {
+ llvm::enumerate(*dimsAndOffs->delinearizedOffsets)) {
if (offset == 0)
continue;
rcSrcStoreIdxs[idx] = getOrCreateIndexConstant(offset);
}
} else {
- // Without runtime delinearization, use the dynamic offset directly only
- // when the source has a single non-unit dimension.
- assert(mapAndOffs->assocMap.size() <= 1 &&
- "Expecting single non-unit dimension mapping.");
+ // Dynamic offsets are used directly only for effectively-1D sources.
+ assert(dimsAndOffs->nonUnitDimsPos.size() <= 1 &&
+ "Expecting at most one non-unit result dimension.");
+
srcNonUnitDimPos = getSingleNonUnitDim(dstType);
assert(srcNonUnitDimPos &&
"Expecting single non-unit dimension source to receive the "
@@ -387,45 +414,38 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
getValueOrCreateConstantIndexOp(rewriter, loc, rcOffsets[0]);
}
- // Scope for OpBuilder::InsertionGuard.
+ // Create the loop nest and emit the load/store at the innermost insertion
+ // point.
{
OpBuilder::InsertionGuard guard(rewriter);
- Value lowerBound;
- Value step;
- if (!upperBounds.empty()) {
- lowerBound = getOrCreateIndexConstant(0);
- step = getOrCreateIndexConstant(1);
- }
- SmallVector<Value> loopIvs;
- loopIvs.reserve(mapAndOffs->assocMap.size());
-
- // Build one nested loop per non-unit strided memref dimension.
- for (Value upperBound : upperBounds) {
- scf::ForOp loop =
- scf::ForOp::create(rewriter, loc, lowerBound, upperBound, step);
- loopIvs.push_back(loop.getInductionVar());
- rewriter.setInsertionPointToStart(loop.getBody());
- }
- // Load indices are zero except for non-unit strided memref dimensions,
- // which use the corresponding loop induction variables.
SmallVector<Value> loadIdxs = getZeroIdxs(cpSrcType.getRank());
- unsigned loopIndex = 0;
- for (const NonUnitDimAssocMapForRC &assocMap : mapAndOffs->assocMap)
- loadIdxs[assocMap.resultDimPos] = loopIvs[loopIndex++];
-
- // Store indices start from the offset-derived source indices. Add each
- // loop IV to the mapped source dimension.
SmallVector<Value> storeIdxs(rcSrcStoreIdxs);
- loopIndex = 0;
- for (const NonUnitDimAssocMapForRC &assocMap : mapAndOffs->assocMap) {
- Value iv = loopIvs[loopIndex++];
- // Add each IV to one source index.
- if (storeIdxs[assocMap.sourceDimPos] == getOrCreateIndexConstant(0)) {
- storeIdxs[assocMap.sourceDimPos] = iv;
- } else {
- storeIdxs[assocMap.sourceDimPos] = arith::AddIOp::create(
- rewriter, loc, storeIdxs[assocMap.sourceDimPos], iv);
+
+ if (!dimsAndOffs->nonUnitDimsPos.empty()) {
+ Value lowerBound = getOrCreateIndexConstant(0);
+ Value step = getOrCreateIndexConstant(1);
+
+ // Build one nested loop per non-unit reinterpret_cast result dimension.
+ for (auto [loopIndex, dim] :
+ llvm::enumerate(dimsAndOffs->nonUnitDimsPos)) {
+ scf::ForOp loop = scf::ForOp::create(rewriter, loc, lowerBound,
+ upperBounds[loopIndex], step);
+
+ rewriter.setInsertionPointToStart(loop.getBody());
+
+ Value iv = loop.getInductionVar();
+ // Since result strides match source identity strides dimension-wise,
+ // each IV indexes the same dimension in both the copy source and rc
+ // source.
+ loadIdxs[dim] = iv;
+
+ if (storeIdxs[dim] == getOrCreateIndexConstant(0)) {
+ storeIdxs[dim] = iv;
+ } else {
+ storeIdxs[dim] =
+ arith::AddIOp::create(rewriter, loc, storeIdxs[dim], iv);
+ }
}
}
@@ -717,7 +737,7 @@ struct ElideReinterpretCastPass
// reinterpret_cast result can be mapped back to base memref indices.
MemRefType cpSrcType = dyn_cast<MemRefType>(op.getSource().getType());
return !(cpSrcType && cpSrcType.hasStaticShape() &&
- getAssocMapAndOffsetsForRC(rc));
+ getResultNonUnitDimsAndOffsetsForRC(rc));
});
target.addDynamicallyLegalOp<memref::LoadOp>([](memref::LoadOp op) {
auto rc = op.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 908cfdcfbcc7f..be0fca979f764 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -104,15 +104,15 @@ func.func private @copy_scalar_into_1D_strided_non_identity_stride(%src : memref
%dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 1], strides: [107, 2]
- : memref<1x108xf32> to memref<1x1xf32, strided<[107, 2]>>
+ to offset: [0], sizes: [1, 1], strides: [2, 54]
+ : memref<1x108xf32> to memref<1x1xf32, strided<[2, 54]>>
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
memref.copy %src, %reinterpret_cast
- : memref<1x1xf32> to memref<1x1xf32, strided<[107, 2]>>
+ : memref<1x1xf32> to memref<1x1xf32, strided<[2, 54]>>
return
}
@@ -315,9 +315,9 @@ func.func private @copy_1D_into_1D_strided_dynamic_offset(
// Either scalar (0D) OR non-scalar (ND) copy
//===----------------------------------------------------------------------===//
-/// Reject copies that don't target a strided memref
-// CHECK-LABEL: func.func private @negative_plain_copy(
-func.func private @negative_plain_copy(%src : memref<1x1xf32>,
+/// Reject copies that don't target a reinterpret_cast result
+// CHECK-LABEL: func.func private @negative_no_rc(
+func.func private @negative_no_rc(%src : memref<1x1xf32>,
%dst : memref<1x1xf32>) {
// CHECK: memref.copy %arg0, %arg1
// CHECK-NOT: memref.load
@@ -364,18 +364,18 @@ func.func private @negative_copy_into_strided_rank_change(%src : memref<3x4xf32>
/// Reject non-identity layout rc source strides
// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity(
func.func private @negative_copy_into_strided_non_identity(%src: memref<1x1xf32>,
- %dst: memref<108x1xf32, strided<[10, 2]>>) {
+ %dst: memref<12x1xf32, strided<[10, 2]>>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [6], sizes: [1, 1], strides: [10, 2]
- : memref<108x1xf32, strided<[10, 2]>>
- to memref<1x1xf32, strided<[10, 2], offset: 6>>
+ to offset: [0], sizes: [1, 1], strides: [10, 2]
+ : memref<12x1xf32, strided<[10, 2]>>
+ to memref<1x1xf32, strided<[10, 2]>>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
- : memref<1x1xf32> to memref<1x1xf32, strided<[10, 2], offset: 6>>
+ : memref<1x1xf32> to memref<1x1xf32, strided<[10, 2]>>
return
}
@@ -458,19 +458,19 @@ func.func private @negative_copy_into_ND_strided_dynamic_offset(
/// (non-unit copied dimension needs stride-based address computation)
// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity_strides(
func.func private @negative_copy_into_strided_non_identity_strides(
- %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x4xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 3, 1], strides: [33, 10, 1]
- : memref<1x3x11xf32>
- to memref<1x3x1xf32, strided<[33, 10, 1]>>
+ to offset: [0], sizes: [1, 3, 1], strides: [12, 4, 4]
+ : memref<1x3x4xf32>
+ to memref<1x3x1xf32, strided<[12, 4, 4]>>
// CHECK: memref.copy %arg0, %reinterpret_cast
// CHECK-NOT: memref.load
// CHECK-NOT: memref.store
memref.copy %src, %rc
: memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[33, 10, 1]>>
+ to memref<1x3x1xf32, strided<[12, 4, 4]>>
return
}
>From 4c67fa2f9ad493d9da385e1fc48577f0e97a50d1 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Thu, 2 Jul 2026 15:54:30 +0200
Subject: [PATCH 11/14] Fixup
---
.../Transforms/ElideReinterpretCast.cpp | 117 ++-------
.../MemRef/elide-reinterpret-cast.mlir | 224 ++++++++++--------
2 files changed, 146 insertions(+), 195 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 483d5c94ea97e..5048a6507c08e 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -49,15 +49,14 @@ struct ResultNonUnitDimsAndOffsetsForRC {
static std::optional<SmallVector<int64_t>>
delinearizeStaticRCOffset(memref::ReinterpretCastOp rc) {
ArrayRef<int64_t> rcOffsets = rc.getStaticOffsets();
+ MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
// 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(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) {
@@ -102,97 +101,29 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
return (*nonUnitDims.begin()).index();
}
-/// Returns the copy-relevant reinterpret_cast information: non-unit result
-/// dimensions, their source-dimension mapping, and optional source starting
-/// indices for a static offset.
-///
-/// Examples that return rewrite info:
+/// Returns reinterpret_cast result non-unit dimensions and, for static offsets,
+/// the corresponding source indices.
///
-/// // Scalar-shaped copy. There are no non-unit result dimensions, so result
-/// // strides do not affect index mapping - may be static or
-/// // dynamic.
-/// copy memref<1 x ... x 1 x f32>
-/// to reinterpret_cast memref<source-shape, identity-layout>
-/// to memref<1 x ... x 1 x f32,
-/// strided<[?, ..., ?], offset: O>>
+/// Supports ranked, static-shape, rank-preserving reinterpret_casts from
+/// identity-layout sources. Non-scalar results must have static strides
+/// identical to the source identity strides. Dynamic offsets are supported only
+/// for effectively-1D sources.
///
-/// // Scalar-shaped copy with dynamic offset into an effectively-1D source.
-/// // The dynamic offset can be used directly as the index of the source's
-/// // unique non-unit dimension.
-/// copy memref<1 x ... x 1 x f32>
-/// to reinterpret_cast memref<1 x ... x M x ... x 1 x f32,
-/// identity-layout>
-/// to memref<1 x ... x 1 x f32,
-/// strided<[?, ..., ?], offset: ?>>
+/// Examples that return info:
///
-/// // Non-scalar effectively-1D copy with static offset. Result strides must
-/// // be static and identical to the identity strides of the source.
-/// copy memref<1 x ... x N x ... x 1 x f32>
-/// to reinterpret_cast memref<1 x ... x M x ... x 1 x f32,
-/// identity-layout>
-/// to memref<1 x ... x N x ... x 1 x f32,
-/// strided<[source-identity-strides], offset: O>>
+/// reinterpret_cast memref<1xNxMxf32, identity-layout>
+/// to memref<1xNxKxf32, strided<[N*M, M, 1], offset: O>>
///
-/// // Non-scalar effectively-1D copy with dynamic offset into an
-/// // effectively-1D source. Runtime delinearization is not needed because
-/// the
-/// // source has a unique non-unit dimension.
-/// copy memref<1 x ... x N x ... x 1 x f32>
-/// to reinterpret_cast memref<1 x ... x M x ... x 1 x f32,
-/// identity-layout>
-/// to memref<1 x ... x N x ... x 1 x f32,
-/// strided<[source-identity-strides], offset: ?>>
-///
-/// // Non-scalar multidimensional copy with static offset. Result strides
-/// must
-/// // be static and identical to the identity strides of the source.
-/// copy memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32>
-/// to reinterpret_cast memref<source-shape, identity-layout>
-/// to memref<1 x ... x N_0 x ... x N_K x ... x 1 x f32,
-/// strided<[source-identity-strides], offset: O>>
+/// reinterpret_cast memref<1xMxf32, identity-layout>
+/// to memref<1x1xf32, strided<[?, ?], offset: ?>>
///
/// Examples that return no info:
///
-/// // Rank-changing reinterpret_casts are not supported.
-/// copy memref<1xNxf32>
-/// to reinterpret_cast memref<Mxf32, identity-layout>
-/// to memref<1xNxf32, strided<[N, 1]>>
-///
-/// // Dynamic shapes are not supported.
-/// copy memref<?xNxf32>
-/// to reinterpret_cast memref<?xMxf32, identity-layout>
-/// to memref<?xNxf32, strided<[M, 1]>>
+/// reinterpret_cast memref<1xNxMxf32, identity-layout>
+/// to memref<1xNxKxf32, strided<[?, M, 1]>>
///
-/// // Non-identity source layouts are not supported.
-/// copy memref<1xNxf32>
-/// to reinterpret_cast memref<1xMxf32, strided<[S, 1]>>
-/// to memref<1xNxf32, strided<[M, 1]>>
-///
-/// // Dynamic offset into a source with more than one non-unit dimension is
-/// // not supported because runtime delinearization is not implemented.
-/// copy memref<1x1xf32>
-/// to reinterpret_cast memref<1xNxMxf32, identity-layout>
-/// to memref<1x1xf32, strided<[?, ?], offset: ?>>
-///
-/// // Non-scalar copies with dynamic result strides are not supported.
-/// copy memref<1xNxf32>
-/// to reinterpret_cast memref<1xMxf32, identity-layout>
-/// to memref<1xNxf32, strided<[?, 1]>>
-///
-/// // Non-scalar copies with result strides different from the source
-/// identity
-/// // strides are not supported.
-/// copy memref<1xNxf32>
-/// to reinterpret_cast memref<1xMxf32, identity-layout>
-/// to memref<1xNxf32, strided<[S, 1]>>
-///
-/// // Multidimensional non-scalar copies with dynamic offset are not
-/// supported
-/// // unless the source is effectively 1D.
-/// copy memref<1xNxKxf32>
-/// to reinterpret_cast memref<1xNxMxf32, identity-layout>
-/// to memref<1xNxKxf32,
-/// strided<[N*M, M, 1], offset: ?>>
+/// reinterpret_cast memref<1xNxMxf32, identity-layout>
+/// to memref<1xNx1xf32, strided<[N*M, M, K]>>
static std::optional<ResultNonUnitDimsAndOffsetsForRC>
getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
@@ -219,11 +150,12 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
ResultNonUnitDimsAndOffsetsForRC dimsAndOffs;
assert(resType.hasStaticShape() && "expected static shape");
+
+ bool isScalarCopy =
+ llvm::all_of(resType.getShape(), [](int64_t size) { return size == 1; });
+
// For scalar copies, result strides are irrelevant, including dynamic ones.
- // For non-scalar copies, require static result strides identical to the
- // identity strides of the reinterpret_cast source.
- if (!llvm::all_of(resType.getShape(),
- [](int64_t size) { return size == 1; })) {
+ if (!isScalarCopy) {
SmallVector<int64_t> srcIdentityStrides =
computeStrides(srcType.getShape());
ArrayRef<int64_t> rcResultStrides = rc.getStaticStrides();
@@ -231,7 +163,8 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
assert((srcIdentityStrides.size() == rcResultStrides.size()) &&
"Expecting same number of strides for rank-preserving "
"reinterpret_casts.");
-
+ // For non-scalar copies, require static result strides identical to the
+ // identity strides of the reinterpret_cast source.
if (!llvm::all_of(llvm::zip_equal(srcIdentityStrides, rcResultStrides),
[](auto pair) {
auto [srcStride, resultStride] = pair;
@@ -239,6 +172,8 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
srcStride == resultStride;
}))
return std::nullopt;
+ // Track result dimensions that produce varying indices; unit dimensions are
+ // always indexed at 0.
for (auto [dim, resultSize] : llvm::enumerate(resType.getShape())) {
if (resultSize != 1)
dimsAndOffs.nonUnitDimsPos.push_back(static_cast<unsigned>(dim));
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index be0fca979f764..f3c47aeab7f58 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -2,15 +2,16 @@
// RUN: | FileCheck %s
//===----------------------------------------------------------------------===//
-// Positive tests
-//
// Scalar (0D) copy
+//
+// No varying RC result dimensions =>
+// RC result strides do not affect copy destination address and are ignored.
//===----------------------------------------------------------------------===//
// The destination is effectively a scalar within a MemRef with rank == 0
-// CHECK-LABEL: func.func private @copy_scalar_into_0D_strided(
+// CHECK-LABEL: func.func private @copy_scalar_into_0D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<f32>, %[[DST:.*]]: memref<f32>
-func.func private @copy_scalar_into_0D_strided(%src : memref<f32>, %dst : memref<f32>) {
+func.func private @copy_scalar_into_0D_strided_zero_offset(%src : memref<f32>, %dst : memref<f32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [], strides: []
@@ -47,6 +48,25 @@ func.func private @copy_scalar_into_1D_strided_zero_offset(%src : memref<1x1xf32
return
}
+/// Reject non-identity layout rc source strides
+// CHECK-LABEL: func.func private @negative_copy_scalar_into_1D_strided_zero_offset_non_identity_layout(
+func.func private @negative_copy_scalar_into_1D_strided_zero_offset_non_identity_layout(
+ %src: memref<1x1xf32>, %dst: memref<1x108xf32, strided<[54, 2]>>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 1], strides: [54, 2]
+ : memref<1x108xf32, strided<[54, 2]>>
+ to memref<1x1xf32, strided<[54, 2]>>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<1x1xf32> to memref<1x1xf32, strided<[54, 2]>>
+
+ return
+}
+
// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_nonzero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
@@ -95,33 +115,31 @@ func.func private @copy_scalar_into_1D_strided_dynamic_offset(%offset: index, %s
return
}
-// Scalar copies have no varying result dimensions, so rc result strides do
-// not affect the copy destination address and are ignored.
-// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_non_identity_stride(
+// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_zero_offset_non_identity_stride(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x108xf32>
-func.func private @copy_scalar_into_1D_strided_non_identity_stride(%src : memref<1x1xf32>,
- %dst : memref<1x108xf32>) {
+func.func private @copy_scalar_into_1D_strided_zero_offset_non_identity_stride(
+ %src : memref<1x1xf32>, %dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 1], strides: [2, 54]
- : memref<1x108xf32> to memref<1x1xf32, strided<[2, 54]>>
+ to offset: [0], sizes: [1, 1], strides: [54, 2]
+ : memref<1x108xf32> to memref<1x1xf32, strided<[54, 2]>>
// CHECK-NOT: memref.copy
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
// CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
memref.copy %src, %reinterpret_cast
- : memref<1x1xf32> to memref<1x1xf32, strided<[2, 54]>>
+ : memref<1x1xf32> to memref<1x1xf32, strided<[54, 2]>>
return
}
-// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_dynamic_stride(
+// CHECK-LABEL: func.func private @copy_scalar_into_1D_strided_zero_offset_dynamic_stride(
// CHECK-SAME: %[[STR0:[A-Za-z][A-Za-z0-9-]*]]: index
// CHECK-SAME: %[[STR1:[A-Za-z][A-Za-z0-9-]*]]: index
// CHECK-SAME: %[[SRC:[A-Za-z][A-Za-z0-9-]*]]: memref<1x1xf32>
// CHECK-SAME: %[[DST:[A-Za-z][A-Za-z0-9-]*]]: memref<1x108xf32>
-func.func private @copy_scalar_into_1D_strided_dynamic_stride(%stride0: index,
+func.func private @copy_scalar_into_1D_strided_zero_offset_dynamic_stride(%stride0: index,
%stride1: index, %src : memref<1x1xf32>, %dst : memref<1x108xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
@@ -141,10 +159,10 @@ func.func private @copy_scalar_into_1D_strided_dynamic_stride(%stride0: index,
}
/// The destination is effectively a 2D array within a MemRef with rank >= 2
-// CHECK-LABEL: func.func private @copy_scalar_into_2D_strided_non_identity_stride(
+// CHECK-LABEL: func.func private @copy_scalar_into_2D_strided_zero_offset_non_identity_stride(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_scalar_into_2D_strided_non_identity_stride(
+func.func private @copy_scalar_into_2D_strided_zero_offset_non_identity_stride(
%src : memref<1x1x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
@@ -161,9 +179,26 @@ func.func private @copy_scalar_into_2D_strided_non_identity_stride(
return
}
+/// Reject dynamic offsets for rc sources with > 1 non-unit dimension -
+/// runtime delinearization of these offsets is TODO.
+// CHECK-LABEL: func.func private @negative_copy_scalar_into_2D_strided_dynamic_offset(
+func.func private @negative_copy_scalar_into_2D_strided_dynamic_offset(
+ %offset : index, %src : memref<1x1x1xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
+ %rc = memref.reinterpret_cast %dst
+ to offset: [%offset], sizes: [1, 1, 1], strides: [33, 11, 1]
+ : memref<1x3x11xf32> to memref<1x1x1xf32, strided<[33, 11, 1], offset: ?>>
+
+ // CHECK: memref.copy %arg1, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<1x1x1xf32>
+ to memref<1x1x1xf32, strided<[33, 11, 1], offset: ?>>
+ return
+}
+
//===----------------------------------------------------------------------===//
-// Positive tests
-//
// Non-scalar (ND) copy
//===----------------------------------------------------------------------===//
@@ -220,6 +255,72 @@ func.func private @copy_1D_into_2D_strided_nonzero_offset(
return
}
+// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_delinearized_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x4x11xf32>
+func.func private @copy_1D_into_2D_strided_delinearized_offset(
+ %src : memref<1x3x1xf32>, %dst : memref<1x4x11xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [12], sizes: [1, 3, 1], strides: [44, 11, 1]
+ : memref<1x4x11xf32>
+ to memref<1x3x1xf32, strided<[44, 11, 1], offset: 12>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[UB:.*]] = arith.constant 3 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[C1]], %[[IDX]] : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[DST_IDX]], %[[C1]]] : memref<1x4x11xf32>
+ // CHECK: }
+ memref.copy %src, %rc
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[44, 11, 1], offset: 12>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+/// Reject rc result strides that not equal to rc source identity strides.
+/// (non-unit copied dimension needs stride-based address computation)
+// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_strided_zero_offset_non_identity_strides(
+func.func private @negative_copy_1D_into_2D_strided_zero_offset_non_identity_strides(
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x4xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 3, 1], strides: [12, 4, 4]
+ : memref<1x3x4xf32>
+ to memref<1x3x1xf32, strided<[12, 4, 4]>>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[12, 4, 4]>>
+ return
+}
+
+/// Reject dynamic rc result strides.
+// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_strided_zero_offset_dynamic_stride(
+func.func private @negative_copy_1D_into_2D_strided_zero_offset_dynamic_stride(%stride : index,
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x4xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 3, 1], strides: [12, %stride, 1]
+ : memref<1x3x4xf32>
+ to memref<1x3x1xf32, strided<[12, ?, 1]>>
+
+ // CHECK: memref.copy %arg1, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[12, ?, 1]>>
+ return
+}
+
// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
@@ -309,9 +410,7 @@ func.func private @copy_1D_into_1D_strided_dynamic_offset(
return
}
-//===----------------------------------------------------------------------===//
-// Negative tests (must NOT rewrite)
-//
+//===----------------------------------------------------------------------===//
// Either scalar (0D) OR non-scalar (ND) copy
//===----------------------------------------------------------------------===//
@@ -361,25 +460,6 @@ func.func private @negative_copy_into_strided_rank_change(%src : memref<3x4xf32>
return
}
-/// Reject non-identity layout rc source strides
-// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity(
-func.func private @negative_copy_into_strided_non_identity(%src: memref<1x1xf32>,
- %dst: memref<12x1xf32, strided<[10, 2]>>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
- %rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 1], strides: [10, 2]
- : memref<12x1xf32, strided<[10, 2]>>
- to memref<1x1xf32, strided<[10, 2]>>
-
- // CHECK: memref.copy %arg0, %reinterpret_cast
- // CHECK-NOT: memref.load
- // CHECK-NOT: memref.store
- memref.copy %src, %rc
- : memref<1x1xf32> to memref<1x1xf32, strided<[10, 2]>>
-
- return
-}
-
/// Reject dynamic shapes
// CHECK-LABEL: func.func private @negative_copy_into_strided_dynamic_copy_source_shape(
func.func private @negative_copy_into_strided_dynamic_copy_source_shape(%src : memref<?xf32>,
@@ -429,70 +509,6 @@ func.func private @negative_copy_into_strided_dynamic_rc_result_shape(%dim : ind
return
}
-/// Reject dynamic offsets for rc sources with > 1 non-unit dimension -
-/// runtime delinearization of these offsets is TODO.
-// CHECK-LABEL: func.func private @negative_copy_into_ND_strided_dynamic_offset(
-func.func private @negative_copy_into_ND_strided_dynamic_offset(
- %offset : index, %src : memref<1x1xf32>, %dst : memref<3x4xf32>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
- %rc = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [1, 1], strides: [12, 1]
- : memref<3x4xf32> to memref<1x1xf32, strided<[12, 1], offset: ?>>
-
- // CHECK: memref.copy %arg1, %reinterpret_cast
- // CHECK-NOT: memref.load
- // CHECK-NOT: memref.store
- memref.copy %src, %rc
- : memref<1x1xf32>
- to memref<1x1xf32, strided<[12, 1], offset: ?>>
- return
-}
-
-//===----------------------------------------------------------------------===//
-// Negative tests
-//
-// Non-scalar (ND) copy
-//===----------------------------------------------------------------------===//
-
-/// Reject rc result strides that not equal to rc source identity strides.
-/// (non-unit copied dimension needs stride-based address computation)
-// CHECK-LABEL: func.func private @negative_copy_into_strided_non_identity_strides(
-func.func private @negative_copy_into_strided_non_identity_strides(
- %src : memref<1x3x1xf32>, %dst : memref<1x3x4xf32>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
- %rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 3, 1], strides: [12, 4, 4]
- : memref<1x3x4xf32>
- to memref<1x3x1xf32, strided<[12, 4, 4]>>
-
- // CHECK: memref.copy %arg0, %reinterpret_cast
- // CHECK-NOT: memref.load
- // CHECK-NOT: memref.store
- memref.copy %src, %rc
- : memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[12, 4, 4]>>
- return
-}
-
-/// Reject dynamic rc result strides.
-// CHECK-LABEL: func.func private @negative_copy_ND_into_strided_dynamic_stride(
-func.func private @negative_copy_ND_into_strided_dynamic_stride(%stride : index,
- %src : memref<1x3x1xf32>, %dst : memref<1x3x4xf32>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
- %rc = memref.reinterpret_cast %dst
- to offset: [0], sizes: [1, 3, 1], strides: [12, %stride, 1]
- : memref<1x3x4xf32>
- to memref<1x3x1xf32, strided<[12, ?, 1]>>
-
- // CHECK: memref.copy %arg1, %reinterpret_cast
- // CHECK-NOT: memref.load
- // CHECK-NOT: memref.store
- memref.copy %src, %rc
- : memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[12, ?, 1]>>
- return
-}
-
// -----
//===----------------------------------------------------------------------===//
>From 53a57bba5c89d408c64fec9187c8c81b7593ca1c Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Fri, 3 Jul 2026 12:19:31 +0200
Subject: [PATCH 12/14] Fixup
---
.../Transforms/ElideReinterpretCast.cpp | 54 ++++++++++---------
.../MemRef/elide-reinterpret-cast.mlir | 12 +++--
2 files changed, 36 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 5048a6507c08e..38c1b7d9a8d35 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -107,7 +107,8 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
/// Supports ranked, static-shape, rank-preserving reinterpret_casts from
/// identity-layout sources. Non-scalar results must have static strides
/// identical to the source identity strides. Dynamic offsets are supported only
-/// for effectively-1D sources.
+/// for effectively-1D sources. Returns nullopt for unsupported
+/// reinterpret_casts.
///
/// Examples that return info:
///
@@ -123,7 +124,8 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
/// to memref<1xNxKxf32, strided<[?, M, 1]>>
///
/// reinterpret_cast memref<1xNxMxf32, identity-layout>
-/// to memref<1xNx1xf32, strided<[N*M, M, K]>>
+/// to memref<1xNx1xf32, strided<[K, M, L]>>
+/// ( identity-layout != [K, M, L] )
static std::optional<ResultNonUnitDimsAndOffsetsForRC>
getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
@@ -151,20 +153,23 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
assert(resType.hasStaticShape() && "expected static shape");
- bool isScalarCopy =
+ // Track result dimensions that produce varying indices; unit dimensions are
+ // always indexed at 0.
+ for (auto [dim, resultSize] : llvm::enumerate(resType.getShape())) {
+ if (resultSize != 1)
+ dimsAndOffs.nonUnitDimsPos.push_back(static_cast<unsigned>(dim));
+ }
+
+ bool isScalarRes =
llvm::all_of(resType.getShape(), [](int64_t size) { return size == 1; });
- // For scalar copies, result strides are irrelevant, including dynamic ones.
- if (!isScalarCopy) {
+ // For non-scalar results, verify that strides do not introduce
+ // non-contiguity that would require extra logic.
+ if (!isScalarRes) {
SmallVector<int64_t> srcIdentityStrides =
computeStrides(srcType.getShape());
ArrayRef<int64_t> rcResultStrides = rc.getStaticStrides();
- assert((srcIdentityStrides.size() == rcResultStrides.size()) &&
- "Expecting same number of strides for rank-preserving "
- "reinterpret_casts.");
- // For non-scalar copies, require static result strides identical to the
- // identity strides of the reinterpret_cast source.
if (!llvm::all_of(llvm::zip_equal(srcIdentityStrides, rcResultStrides),
[](auto pair) {
auto [srcStride, resultStride] = pair;
@@ -172,14 +177,15 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
srcStride == resultStride;
}))
return std::nullopt;
- // Track result dimensions that produce varying indices; unit dimensions are
- // always indexed at 0.
- for (auto [dim, resultSize] : llvm::enumerate(resType.getShape())) {
- if (resultSize != 1)
- dimsAndOffs.nonUnitDimsPos.push_back(static_cast<unsigned>(dim));
- }
}
+ std::optional<unsigned> srcNonUnitDim = getSingleNonUnitDim(srcType);
+
+ // A source with exactly one non-unit dimension cannot be indexed directly by
+ // multiple non-unit result dimensions.
+ if (srcNonUnitDim && dimsAndOffs.nonUnitDimsPos.size() > 1)
+ return std::nullopt;
+
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.
@@ -194,8 +200,9 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
"static reinterpret_cast offset must delinearize to in-bounds "
"reinterpret_cast source indices");
- // Relevant for non-scalar copies: assert that the rectangular
- // copied slice is in bounds.
+ // Sanity check that the reinterpret_cast doesn't create an out-of-bounds
+ // MemRef. Such cases should probably be rejected by Op verifier.
+ // FIXME: Add run-time verification for cases like this.
assert(llvm::all_of(llvm::enumerate(resType.getShape()),
[&](auto it) {
unsigned dim = it.index();
@@ -213,16 +220,11 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
// dimensions by delinearizing the offset into source start indices at runtime
// before adding loop IVs.
- // With an effectively-1D source, a dynamic linear offset can be used directly
- // as the index of the unique non-unit source dimension.
- if (!getSingleNonUnitDim(srcType))
+ // With an effectively-1D source, a dynamic offset can be mapped to the unique
+ // non-unit dim. For other cases, bail out.
+ if (!srcNonUnitDim)
return std::nullopt;
- // Non-scalar copies require identical strides and no rank-changing,
- // so there can be at most one non-unit result dimension in this case.
- assert(dimsAndOffs.nonUnitDimsPos.size() <= 1 &&
- "effectively-1D source cannot have multiple non-unit result dims");
-
return dimsAndOffs;
}
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index f3c47aeab7f58..a1181b3d54a29 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -228,10 +228,12 @@ func.func private @copy_1D_into_2D_strided_zero_offset(
return
}
-// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_nonzero_offset(
+/// Offset 10 delinearizes to [0, 0, 10], therefore is only
+/// added to the trailing source dimension.
+// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_nonzero_offset_delinierized_v1(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_1D_into_2D_strided_nonzero_offset(
+func.func private @copy_1D_into_2D_strided_nonzero_offset_delinierized_v1(
%src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
@@ -255,10 +257,12 @@ func.func private @copy_1D_into_2D_strided_nonzero_offset(
return
}
-// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_delinearized_offset(
+/// Offset 12 delinearizes to [0, 1, 1], therefore is split across
+/// both the looped source dimension and the trailing source dimension.
+// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_nonzero_offset_delinearized_v2(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x4x11xf32>
-func.func private @copy_1D_into_2D_strided_delinearized_offset(
+func.func private @copy_1D_into_2D_strided_nonzero_offset_delinearized_v2(
%src : memref<1x3x1xf32>, %dst : memref<1x4x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
>From e89c1d53a2027757778ea25c7e0db3cf1197381d Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Tue, 7 Jul 2026 17:26:24 +0200
Subject: [PATCH 13/14] Constrain RC to either scalar OR single non-unit dim
collapsed result
---
.../Transforms/ElideReinterpretCast.cpp | 179 +++++++------
.../MemRef/elide-reinterpret-cast.mlir | 242 ++++++++++++------
2 files changed, 268 insertions(+), 153 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 38c1b7d9a8d35..e3a2bb6457b69 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -84,7 +84,35 @@ delinearizeStaticRCOffset(memref::ReinterpretCastOp rc) {
return offsetIdxs;
}
-/// Returns the unique non-unit dim or nullopt of # non-unit-dims != 1.
+static bool hasExactlyOneCollapsedNonUnitDim(MemRefType srcType,
+ MemRefType resType) {
+ assert(srcType.hasStaticShape() && resType.hasStaticShape() &&
+ "expected static shapes");
+ assert(srcType.getRank() == resType.getRank() &&
+ "expected rank-preserving reinterpret_cast");
+
+ unsigned collapsedDims = 0;
+
+ for (auto [srcSize, resSize] :
+ llvm::zip_equal(srcType.getShape(), resType.getShape())) {
+ if (srcSize == resSize)
+ continue;
+
+ // Only allow collapsing one non-unit source dim to a unit result dim.
+ if (srcSize != 1 && resSize == 1) {
+ ++collapsedDims;
+ continue;
+ }
+
+ // The sizes differ and both of them are non-unit - ATM not supported.
+ return false;
+ }
+
+ // Make sure there is only one collapsed dimension.
+ return collapsedDims == 1;
+}
+
+/// Returns the unique non-unit dim or nullopt if # non-unit-dims != 1.
static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
assert(type.hasStaticShape() && "expected static shape");
ArrayRef<int64_t> shape = type.getShape();
@@ -105,27 +133,27 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
/// the corresponding source indices.
///
/// Supports ranked, static-shape, rank-preserving reinterpret_casts from
-/// identity-layout sources. Non-scalar results must have static strides
-/// identical to the source identity strides. Dynamic offsets are supported only
-/// for effectively-1D sources. Returns nullopt for unsupported
+/// identity-layout sources. Scalar-shaped results may have arbitrary result
+/// strides. Non-scalar results must have static offsets, static result strides
+/// identical to the source identity strides, and exactly one non-unit source
+/// dimension collapsed to unit size. Returns nullopt for unsupported
/// reinterpret_casts.
///
/// Examples that return info:
///
-/// reinterpret_cast memref<1xNxMxf32, identity-layout>
-/// to memref<1xNxKxf32, strided<[N*M, M, 1], offset: O>>
+/// reinterpret_cast memref<1xMxNxf32, identity-layout>
+/// to memref<1xMx1xf32, strided<[M*N, N, 1], offset: OFF>>
///
/// reinterpret_cast memref<1xMxf32, identity-layout>
/// to memref<1x1xf32, strided<[?, ?], offset: ?>>
///
/// Examples that return no info:
///
-/// reinterpret_cast memref<1xNxMxf32, identity-layout>
-/// to memref<1xNxKxf32, strided<[?, M, 1]>>
+/// reinterpret_cast memref<1xMxNxf32, identity-layout>
+/// to memref<1xMx1xf32, strided<[?, N, 1]>>
///
-/// reinterpret_cast memref<1xNxMxf32, identity-layout>
-/// to memref<1xNx1xf32, strided<[K, M, L]>>
-/// ( identity-layout != [K, M, L] )
+/// reinterpret_cast memref<1xMxNxf32, identity-layout>
+/// to memref<1xKx1xf32, strided<[M*N, N, 1], offset: OFF>>
static std::optional<ResultNonUnitDimsAndOffsetsForRC>
getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
@@ -136,7 +164,7 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
if (!srcType || !resType)
return std::nullopt;
- // TODO: Support rank-modifying reinterpret_casts
+ // TODO: Support rank-modifying reinterpret_casts.
if (srcType.getRank() != resType.getRank())
return std::nullopt;
@@ -151,8 +179,6 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
ResultNonUnitDimsAndOffsetsForRC dimsAndOffs;
- assert(resType.hasStaticShape() && "expected static shape");
-
// Track result dimensions that produce varying indices; unit dimensions are
// always indexed at 0.
for (auto [dim, resultSize] : llvm::enumerate(resType.getShape())) {
@@ -160,12 +186,20 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
dimsAndOffs.nonUnitDimsPos.push_back(static_cast<unsigned>(dim));
}
+ 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");
+
bool isScalarRes =
llvm::all_of(resType.getShape(), [](int64_t size) { return size == 1; });
- // For non-scalar results, verify that strides do not introduce
- // non-contiguity that would require extra logic.
if (!isScalarRes) {
+ // Non-scalar cases are restricted to same-dimension, one-collapsed-dim
+ // views with static metadata.
+ if (ShapedType::isDynamic(rcOffsets[0]))
+ return std::nullopt;
+
SmallVector<int64_t> srcIdentityStrides =
computeStrides(srcType.getShape());
ArrayRef<int64_t> rcResultStrides = rc.getStaticStrides();
@@ -177,53 +211,32 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
srcStride == resultStride;
}))
return std::nullopt;
- }
- std::optional<unsigned> srcNonUnitDim = getSingleNonUnitDim(srcType);
-
- // A source with exactly one non-unit dimension cannot be indexed directly by
- // multiple non-unit result dimensions.
- if (srcNonUnitDim && dimsAndOffs.nonUnitDimsPos.size() > 1)
- return std::nullopt;
-
- 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");
+ if (!hasExactlyOneCollapsedNonUnitDim(srcType, resType))
+ return std::nullopt;
+ }
- // CASE 1: Static ReinterpretCast offset
- if (ShapedType::isStatic(rcOffsets[0])) {
- // Delinearize static ReinterpretCast offset as in-bounds indices (one for
- // every source dimension).
- dimsAndOffs.delinearizedOffsets = delinearizeStaticRCOffset(rc);
- assert(dimsAndOffs.delinearizedOffsets &&
- "static reinterpret_cast offset must delinearize to in-bounds "
- "reinterpret_cast source indices");
+ // CASE 1: Dynamic ReinterpretCast offset.
+ //
+ // Dynamic offsets are supported only for scalar-shaped results, under the
+ // previous effectively-1D source restriction.
+ if (ShapedType::isDynamic(rcOffsets[0])) {
+ // With an effectively-1D source, a dynamic offset can be mapped to its
+ // unique non-unit dim. For other cases, bail out.
+ if (llvm::count_if(srcType.getShape(),
+ [](int64_t size) { return size != 1; }) != 1)
+ return std::nullopt;
- // Sanity check that the reinterpret_cast doesn't create an out-of-bounds
- // MemRef. Such cases should probably be rejected by Op verifier.
- // FIXME: Add run-time verification for cases like this.
- assert(llvm::all_of(llvm::enumerate(resType.getShape()),
- [&](auto it) {
- unsigned dim = it.index();
- int64_t resultSize = it.value();
- return (*dimsAndOffs.delinearizedOffsets)[dim] +
- resultSize <=
- srcType.getDimSize(dim);
- }) &&
- "reinterpret_cast metadata describes an invalid accessible region");
return dimsAndOffs;
}
- // CASE 2: Dynamic ReinterpretCast offset.
- // TODO: Support dynamic offsets into sources with multiple non-unit
- // dimensions by delinearizing the offset into source start indices at runtime
- // before adding loop IVs.
-
- // With an effectively-1D source, a dynamic offset can be mapped to the unique
- // non-unit dim. For other cases, bail out.
- if (!srcNonUnitDim)
- return std::nullopt;
+ // CASE 2: Static ReinterpretCast offset
+ // Delinearize static ReinterpretCast offset as in-bounds indices (one for
+ // every source dimension).
+ dimsAndOffs.delinearizedOffsets = delinearizeStaticRCOffset(rc);
+ assert(dimsAndOffs.delinearizedOffsets &&
+ "static reinterpret_cast offset must delinearize to in-bounds "
+ "reinterpret_cast source indices");
return dimsAndOffs;
}
@@ -231,37 +244,36 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
/// Rewrites supported copy operations through `memref.reinterpret_cast` to
/// scalar load/store operations.
///
-/// // BEFORE (scalar copy)
-/// %strided = memref.reinterpret_cast %dst
-/// to offset: [O], sizes: [1, ..., 1], strides: [...]
-/// memref.copy %src, %strided
+/// Supported cases:
+/// 1. Scalar-shaped reinterpret_cast results. Result strides are ignored;
+/// the store index is derived from the reinterpret_cast offset.
///
-/// // AFTER
-/// %v = memref.load %src[0, ..., 0]
-/// memref.store %v, %dst[delinearized(O)]
+/// 2. Non-scalar reinterpret_cast results that preserve all non-unit source
+/// dimensions except one collapsed-to-unit dimension. Result strides must
+/// be static and identical to the identity strides of the source, and the
+/// static offset selects the collapsed dimension.
///
-/// // BEFORE (effectively-1D copy)
+/// // BEFORE (scalar-shaped result)
/// %strided = memref.reinterpret_cast %dst
-/// to offset: [O], sizes: [1, N, 1], strides: [...]
+/// to offset: [OFF], sizes: [1, ..., 1], strides: [...]
/// memref.copy %src, %strided
///
/// // AFTER
-/// scf.for %i = 0 to N step 1 {
-/// %v = memref.load %src[0, %i, 0]
-/// memref.store %v, %dst[delinearized(O) + mapped(%i)]
-/// }
+/// %v = memref.load %src[0, ..., 0]
+/// memref.store %v, %dst[delinearized(OFF)]
///
-/// // BEFORE (multidimensional copy with static offset)
+/// // BEFORE (one collapsed non-unit dimension)
/// %strided = memref.reinterpret_cast %dst
-/// to offset: [O], sizes: [1, N, K], strides: [...]
+/// to offset: [OFF], sizes: [1, M, 1], strides: [M*N, N, 1]
+/// : memref<1xMxNxf32>
+/// to memref<1xMx1xf32, strided<[M*N, N, 1], offset: OFF>>
/// memref.copy %src, %strided
///
/// // AFTER
-/// scf.for %i = 0 to N step 1 {
-/// scf.for %j = 0 to K step 1 {
-/// %v = memref.load %src[0, %i, %j]
-/// memref.store %v, %dst[delinearized(O) + mapped(%i, %j)]
-/// }
+/// // Assuming OFF delinearizes to [0, 0, OFF]:
+/// scf.for %i = 0 to M step 1 {
+/// %v = memref.load %src[0, %i, 0]
+/// memref.store %v, %dst[0, %i, OFF]
/// }
struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
public:
@@ -292,6 +304,19 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
MemRefType dstType = cast<MemRefType>(dst.getType());
MemRefType rcResType = cast<MemRefType>(rc.getType());
+ // Sanity check that the copy doesn't access strided MemRef out-of-bounds.
+ // Such cases should probably be rejected by Op verifier.
+ // FIXME: Add run-time verification for cases like this.
+ if (ShapedType::isStatic(rc.getStaticOffsets()[0]) &&
+ llvm::any_of(llvm::enumerate(rcResType.getShape()), [&](auto it) {
+ unsigned dim = it.index();
+ int64_t rcResultSize = it.value();
+ return (*dimsAndOffs->delinearizedOffsets)[dim] + rcResultSize >
+ dstType.getDimSize(dim);
+ }))
+ return rewriter.notifyMatchFailure(
+ op, "copy accesses invalid accessible region");
+
// Reuse common index constants across bounds, steps, and static offsets,
// but avoid creating them for rank-0 copies.
std::array<Value, 2> cachedIndexConstants;
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index a1181b3d54a29..afa6833598c02 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -179,8 +179,59 @@ func.func private @copy_scalar_into_2D_strided_zero_offset_non_identity_stride(
return
}
-/// Reject dynamic offsets for rc sources with > 1 non-unit dimension -
-/// runtime delinearization of these offsets is TODO.
+/// Offset delinearized to [0, 0, 10], therefore is only
+/// added to the trailing source dimension.
+// CHECK-LABEL: func.func private @copy_scalar_into_2D_scalar_strided_nonzero_offset_delinearized_v1(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_scalar_into_2D_scalar_strided_nonzero_offset_delinearized_v1(
+ %src : memref<1x1x1xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [10], sizes: [1, 1, 1], strides: [1, 1, 1]
+ : memref<1x3x11xf32>
+ to memref<1x1x1xf32, strided<[1, 1, 1], offset: 10>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[OFF:.*]] = arith.constant 10 : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[OFF]]] : memref<1x3x11xf32>
+ memref.copy %src, %rc
+ : memref<1x1x1xf32>
+ to memref<1x1x1xf32, strided<[1, 1, 1], offset: 10>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+// Offset delinearized into more dimensions: [0, 2, 1]
+// CHECK-LABEL: func.func private @copy_scalar_into_2D_scalar_strided_nonzero_offset_delinearized_v2(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_scalar_into_2D_scalar_strided_nonzero_offset_delinearized_v2(
+ %src : memref<1x1x1xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [23], sizes: [1, 1, 1], strides: [1, 1, 1]
+ : memref<1x3x11xf32>
+ to memref<1x1x1xf32, strided<[1, 1, 1], offset: 23>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C2]], %[[C1]]] : memref<1x3x11xf32>
+ memref.copy %src, %rc
+ : memref<1x1x1xf32>
+ to memref<1x1x1xf32, strided<[1, 1, 1], offset: 23>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+/// rc result dynamic offset:
+/// supported only for effectively-1D rc source
+/// (runtime delinearization not implemented)
// CHECK-LABEL: func.func private @negative_copy_scalar_into_2D_strided_dynamic_offset(
func.func private @negative_copy_scalar_into_2D_strided_dynamic_offset(
%offset : index, %src : memref<1x1x1xf32>, %dst : memref<1x3x11xf32>) {
@@ -202,6 +253,23 @@ func.func private @negative_copy_scalar_into_2D_strided_dynamic_offset(
// Non-scalar (ND) copy
//===----------------------------------------------------------------------===//
+/// No non-unit dimension collapsed
+// CHECK-LABEL: func.func private @negative_copy_1D_into_1D_strided(
+func.func private @negative_copy_1D_into_1D_strided(
+ %src : memref<4xf32>, %dst : memref<108xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [4], strides: [1]
+ : memref<108xf32> to memref<4xf32, strided<[1]>>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<4xf32> to memref<4xf32, strided<[1]>>
+ return
+}
+
// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
@@ -228,12 +296,38 @@ func.func private @copy_1D_into_2D_strided_zero_offset(
return
}
-/// Offset 10 delinearizes to [0, 0, 10], therefore is only
+// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_zero_offset_loop_trailing_dim(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x1x11xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_1D_into_2D_strided_zero_offset_loop_trailing_dim(
+ %src : memref<1x1x11xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 1, 11], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x1x11xf32, strided<[33, 11, 1]>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[UB:.*]] = arith.constant 11 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[IDX]]] : memref<1x1x11xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[IDX]]] : memref<1x3x11xf32>
+ // CHECK: }
+ memref.copy %src, %rc
+ : memref<1x1x11xf32>
+ to memref<1x1x11xf32, strided<[33, 11, 1]>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+/// Offset delinearized to [0, 0, 10], therefore is only
/// added to the trailing source dimension.
-// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_nonzero_offset_delinierized_v1(
+// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_nonzero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_1D_into_2D_strided_nonzero_offset_delinierized_v1(
+func.func private @copy_1D_into_2D_strided_nonzero_offset(
%src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
@@ -257,36 +351,25 @@ func.func private @copy_1D_into_2D_strided_nonzero_offset_delinierized_v1(
return
}
-/// Offset 12 delinearizes to [0, 1, 1], therefore is split across
-/// both the looped source dimension and the trailing source dimension.
-// CHECK-LABEL: func.func private @copy_1D_into_2D_strided_nonzero_offset_delinearized_v2(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x4x11xf32>
-func.func private @copy_1D_into_2D_strided_nonzero_offset_delinearized_v2(
- %src : memref<1x3x1xf32>, %dst : memref<1x4x11xf32>) {
- // CHECK-NOT: memref.reinterpret_cast
+// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_strided_dynamic_offset(
+func.func private @negative_copy_1D_into_2D_strided_dynamic_offset(
+ %offset : index, %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg2
%rc = memref.reinterpret_cast %dst
- to offset: [12], sizes: [1, 3, 1], strides: [44, 11, 1]
- : memref<1x4x11xf32>
- to memref<1x3x1xf32, strided<[44, 11, 1], offset: 12>>
+ to offset: [%offset], sizes: [1, 3, 1], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
- // CHECK-NOT: memref.copy
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[UB:.*]] = arith.constant 3 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
- // CHECK: %[[DST_IDX:.*]] = arith.addi %[[C1]], %[[IDX]] : index
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[DST_IDX]], %[[C1]]] : memref<1x4x11xf32>
- // CHECK: }
+ // CHECK: memref.copy %arg1, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
memref.copy %src, %rc
: memref<1x3x1xf32>
- to memref<1x3x1xf32, strided<[44, 11, 1], offset: 12>>
- // CHECK-NOT: memref.copy
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: ?>>
return
}
-/// Reject rc result strides that not equal to rc source identity strides.
+/// Reject rc result strides that are not equal to rc source identity strides.
/// (non-unit copied dimension needs stride-based address computation)
// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_strided_zero_offset_non_identity_strides(
func.func private @negative_copy_1D_into_2D_strided_zero_offset_non_identity_strides(
@@ -325,91 +408,98 @@ func.func private @negative_copy_1D_into_2D_strided_zero_offset_dynamic_stride(%
return
}
-// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_zero_offset(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_2D_into_2D_strided_zero_offset(
+// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_strided_diff_dim_sizes(
+func.func private @negative_copy_1D_into_2D_strided_diff_dim_sizes(
+ %src : memref<1x3x1xf32>, %dst : memref<1x4x11xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 3, 1], strides: [44, 11, 1]
+ : memref<1x4x11xf32>
+ to memref<1x3x1xf32, strided<[44, 11, 1]>>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[44, 11, 1]>>
+ return
+}
+
+/// No non-unit dimension collapsed
+// CHECK-LABEL: func.func private @negative_copy_2D_into_2D_strided(
+func.func private @negative_copy_2D_into_2D_strided(
%src : memref<1x3x4xf32>, %dst : memref<1x3x11xf32>) {
- // CHECK-NOT: memref.reinterpret_cast
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [1, 3, 4], strides: [33, 11, 1]
: memref<1x3x11xf32>
to memref<1x3x4xf32, strided<[33, 11, 1]>>
- // CHECK-NOT: memref.copy
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
- // CHECK-DAG: %[[UB1:.*]] = arith.constant 4 : index
- // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
- // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x4xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x11xf32>
- // CHECK: }
- // CHECK: }
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
memref.copy %src, %rc
: memref<1x3x4xf32>
to memref<1x3x4xf32, strided<[33, 11, 1]>>
- // CHECK-NOT: memref.copy
return
}
-// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_nonzero_offset(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
-func.func private @copy_2D_into_2D_strided_nonzero_offset(
- %src : memref<1x3x4xf32>, %dst : memref<1x3x11xf32>) {
+// CHECK-LABEL: func.func private @copy_2D_into_3D_strided_zero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<3x1x4x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<3x1x4x11xf32>
+func.func private @copy_2D_into_3D_strided_zero_offset(
+ %src : memref<3x1x4x1xf32>, %dst : memref<3x1x4x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
- to offset: [7], sizes: [1, 3, 4], strides: [33, 11, 1]
- : memref<1x3x11xf32>
- to memref<1x3x4xf32, strided<[33, 11, 1], offset: 7>>
+ to offset: [0], sizes: [3, 1, 4, 1], strides: [44, 44, 11, 1]
+ : memref<3x1x4x11xf32>
+ to memref<3x1x4x1xf32, strided<[44, 44, 11, 1]>>
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[UB1:.*]] = arith.constant 4 : index
- // CHECK-DAG: %[[OFF:.*]] = arith.constant 7 : index
// CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
// CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
- // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX1]] : index
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x4xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[DST_IDX]]] : memref<1x3x11xf32>
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX0]], %[[C0]], %[[IDX1]], %[[C0]]] : memref<3x1x4x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[IDX0]], %[[C0]], %[[IDX1]], %[[C0]]] : memref<3x1x4x11xf32>
// CHECK: }
// CHECK: }
memref.copy %src, %rc
- : memref<1x3x4xf32>
- to memref<1x3x4xf32, strided<[33, 11, 1], offset: 7>>
+ : memref<3x1x4x1xf32>
+ to memref<3x1x4x1xf32, strided<[44, 44, 11, 1]>>
// CHECK-NOT: memref.copy
return
}
-/// rc result dynamic offset:
-/// supported only for effectively-1D rc source
-/// (runtime delinearization not implemented)
-// CHECK-LABEL: func.func private @copy_1D_into_1D_strided_dynamic_offset(
-// CHECK-SAME: %[[OFF:.*]]: index
-// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
-// CHECK-SAME: %[[DST:.*]]: memref<108xf32>
-func.func private @copy_1D_into_1D_strided_dynamic_offset(
- %offset : index, %src : memref<4xf32>, %dst : memref<108xf32>) {
+// CHECK-LABEL: func.func private @copy_2D_into_3D_strided_nonzero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<3x1x4x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<3x1x4x11xf32>
+func.func private @copy_2D_into_3D_strided_nonzero_offset(
+ %src : memref<3x1x4x1xf32>, %dst : memref<3x1x4x11xf32>) {
// CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
- to offset: [%offset], sizes: [4], strides: [1]
- : memref<108xf32> to memref<4xf32, strided<[1], offset: ?>>
+ to offset: [10], sizes: [3, 1, 4, 1], strides: [44, 44, 11, 1]
+ : memref<3x1x4x11xf32>
+ to memref<3x1x4x1xf32, strided<[44, 44, 11, 1], offset: 10>>
// CHECK-NOT: memref.copy
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[UB:.*]] = arith.constant 4 : index
- // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
- // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX]] : index
- // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<4xf32>
- // CHECK: memref.store %[[VAL]], %[[DST]][%[[DST_IDX]]] : memref<108xf32>
+ // CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[UB1:.*]] = arith.constant 4 : index
+ // CHECK-DAG: %[[OFF:.*]] = arith.constant 10 : index
+ // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
+ // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX0]], %[[C0]], %[[IDX1]], %[[C0]]] : memref<3x1x4x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[IDX0]], %[[C0]], %[[IDX1]], %[[OFF]]] : memref<3x1x4x11xf32>
+ // CHECK: }
// CHECK: }
memref.copy %src, %rc
- : memref<4xf32> to memref<4xf32, strided<[1], offset: ?>>
+ : memref<3x1x4x1xf32>
+ to memref<3x1x4x1xf32, strided<[44, 44, 11, 1], offset: 10>>
// CHECK-NOT: memref.copy
return
}
>From 439436a5b975771987f7d519d98df388e9246227 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Wed, 8 Jul 2026 14:13:32 +0200
Subject: [PATCH 14/14] Fixup
---
.../Transforms/ElideReinterpretCast.cpp | 54 ++++++++++---------
1 file changed, 30 insertions(+), 24 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index e3a2bb6457b69..0285620e22524 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -84,8 +84,9 @@ delinearizeStaticRCOffset(memref::ReinterpretCastOp rc) {
return offsetIdxs;
}
-static bool hasExactlyOneCollapsedNonUnitDim(MemRefType srcType,
- MemRefType resType) {
+static bool hasExactlyOneCollapsedNonUnitDim(memref::ReinterpretCastOp rc) {
+ MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
+ MemRefType resType = dyn_cast<MemRefType>(rc.getType());
assert(srcType.hasStaticShape() && resType.hasStaticShape() &&
"expected static shapes");
assert(srcType.getRank() == resType.getRank() &&
@@ -129,14 +130,19 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
return (*nonUnitDims.begin()).index();
}
-/// Returns reinterpret_cast result non-unit dimensions and, for static offsets,
-/// the corresponding source indices.
+/// Returns reinterpret_cast's result non-unit dimensions and, for static
+/// offsets, delinearized offset.
///
/// Supports ranked, static-shape, rank-preserving reinterpret_casts from
-/// identity-layout sources. Scalar-shaped results may have arbitrary result
-/// strides. Non-scalar results must have static offsets, static result strides
-/// identical to the source identity strides, and exactly one non-unit source
-/// dimension collapsed to unit size. Returns nullopt for unsupported
+/// identity-layout sources. In addition:
+/// identical to the source identity strides, and exactly one non-unit
+/// source
+/// * Non-scalar results must have static offsets, static result strides
+/// dimension collapsed to unit size
+/// Scalar-shaped results may have arbitrary result strides (i.e. for scalars,
+/// strides are effectively irrelevant).
+///
+/// Returns nullopt for unsupported
/// reinterpret_casts.
///
/// Examples that return info:
@@ -179,7 +185,7 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
ResultNonUnitDimsAndOffsetsForRC dimsAndOffs;
- // Track result dimensions that produce varying indices; unit dimensions are
+ // Track non-unit result dimensions; unit dimensions are
// always indexed at 0.
for (auto [dim, resultSize] : llvm::enumerate(resType.getShape())) {
if (resultSize != 1)
@@ -194,10 +200,13 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
bool isScalarRes =
llvm::all_of(resType.getShape(), [](int64_t size) { return size == 1; });
+ bool isOffsetDynamic = ShapedType::isDynamic(rcOffsets[0]);
+
+ // Cases with at least one non-unit dimension in reinterpret_cast's result are
+ // restricted to preserving all but one dimension from the source, which
+ // collapsed to `1` in the result, and fully static metadata.
if (!isScalarRes) {
- // Non-scalar cases are restricted to same-dimension, one-collapsed-dim
- // views with static metadata.
- if (ShapedType::isDynamic(rcOffsets[0]))
+ if (isOffsetDynamic)
return std::nullopt;
SmallVector<int64_t> srcIdentityStrides =
@@ -212,15 +221,15 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
}))
return std::nullopt;
- if (!hasExactlyOneCollapsedNonUnitDim(srcType, resType))
+ if (!hasExactlyOneCollapsedNonUnitDim(rc))
return std::nullopt;
}
// CASE 1: Dynamic ReinterpretCast offset.
//
- // Dynamic offsets are supported only for scalar-shaped results, under the
- // previous effectively-1D source restriction.
- if (ShapedType::isDynamic(rcOffsets[0])) {
+ // Dynamic offsets are supported only for effectively-1D to scalar
+ // reinterpret_casts.
+ if (isOffsetDynamic) {
// With an effectively-1D source, a dynamic offset can be mapped to its
// unique non-unit dim. For other cases, bail out.
if (llvm::count_if(srcType.getShape(),
@@ -234,9 +243,6 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
// Delinearize static ReinterpretCast offset as in-bounds indices (one for
// every source dimension).
dimsAndOffs.delinearizedOffsets = delinearizeStaticRCOffset(rc);
- assert(dimsAndOffs.delinearizedOffsets &&
- "static reinterpret_cast offset must delinearize to in-bounds "
- "reinterpret_cast source indices");
return dimsAndOffs;
}
@@ -297,7 +303,8 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
getResultNonUnitDimsAndOffsetsForRC(rc);
if (!dimsAndOffs)
return rewriter.notifyMatchFailure(
- op, "reinterpret_cast does not match scalar or loop copy region");
+ op,
+ "unsupported reinterpret_cast result dimensions, strides, or offset");
Location loc = op.getLoc();
Value dst = rc.getSource();
@@ -314,11 +321,10 @@ struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
return (*dimsAndOffs->delinearizedOffsets)[dim] + rcResultSize >
dstType.getDimSize(dim);
}))
- return rewriter.notifyMatchFailure(
- op, "copy accesses invalid accessible region");
+ return rewriter.notifyMatchFailure(op, "copy accesses are OOB");
- // Reuse common index constants across bounds, steps, and static offsets,
- // but avoid creating them for rank-0 copies.
+ // Constant Op cache to reuse common index constants across bounds, steps,
+ // and static offsets: 0 is stored at index 0 and 1 is stored at index 1.
std::array<Value, 2> cachedIndexConstants;
auto getOrCreateIndexConstant = [&](int64_t value) -> Value {
if (value == 0 || value == 1) {
More information about the Mlir-commits
mailing list