[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)
ioana ghiban
llvmlistbot at llvm.org
Fri Jun 19 01:48:41 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/3] [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/3] 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/3] 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
More information about the Mlir-commits
mailing list