[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)
ioana ghiban
llvmlistbot at llvm.org
Mon Jun 29 06:38:10 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 1/8] [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 2/8] 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 3/8] 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 4/8] 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 5/8] 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 6/8] 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 7/8] 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 ae4bd445afa3c6a31d7ed3468d05240209c5eed8 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 8/8] Address fourth round of comments
---
.../Transforms/ElideReinterpretCast.cpp | 23 +-
.../MemRef/elide-reinterpret-cast.mlir | 270 +++++++++++-------
2 files changed, 181 insertions(+), 112 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 435600f58f7fa..6cd56eefb9562 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -51,8 +51,8 @@ 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 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 +62,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 +77,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 +123,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..d622705a18025 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,72 @@ 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
+}
+
//===----------------------------------------------------------------------===//
// 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]>>) {
+ %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: [11, 80]
- : memref<8x1xf32, strided<[10, 2]>>
- to memref<1x1xf32, strided<[11, 80], offset: 6>>
+ 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<1x1xf32> to memref<1x1xf32, strided<[11, 80], offset: 6>>
+ : memref<1x1xf32> to memref<1x1xf32, strided<[10, 2], offset: 6>>
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>) {
- // 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]>>
-
- // 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]>>
- 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 +395,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 +477,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
}
@@ -451,19 +517,19 @@ 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>) {
+ %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
}
More information about the Mlir-commits
mailing list