[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)
ioana ghiban
llvmlistbot at llvm.org
Mon Jun 15 08:54:56 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] [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
}
-
// -----
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list