[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 15 04:43:03 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: ioana ghiban (ioghiban)
<details>
<summary>Changes</summary>
Extends `memref-elide-reinterpret-cast` copy rewriting from scalar-only `memref.copy` into `memref.reinterpret_cast` to also handle copies of effectively 1D and static multidimensional regions.
The rewrite now lowers supported copies to scalar `memref.load`/`memref.store`, using nested `scf.for` loops for non-unit copied dimensions. Store indices are computed directly into the reinterpret_cast source memref by mapping view strides to identity-layout base dimensions and delinearizing static offsets.
## Examples
Scalar copy, unchanged behavior:
```mlir
// BEFORE
%view = memref.reinterpret_cast %dst
to offset: [1], sizes: [1, 1], strides: [1, 1]
: memref<1x108xf32> to memref<1x1xf32, strided<[1, 1], offset: 1>>
memref.copy %src, %view : memref<1x1xf32> to memref<1x1xf32, strided<[1, 1], offset: 1>>
// AFTER
%v = memref.load %src[%c0, %c0] : memref<1x1xf32>
memref.store %v, %dst[%c0, %c1] : memref<1x108xf32>
```
Effectively 1D copy:
```mlir
// BEFORE
%view = memref.reinterpret_cast %dst
to offset: [0], sizes: [1, 33, 1], strides: [1386, 42, 1]
: memref<1x33x42xf32> to memref<1x33x1xf32, strided<[1386, 42, 1]>>
memref.copy %src, %view : memref<1x33x1xf32> to memref<1x33x1xf32, strided<[1386, 42, 1]>>
// AFTER
scf.for %i = %c0 to %c33 step %c1 {
%v = memref.load %src[%c0, %i, %c0] : memref<1x33x1xf32>
memref.store %v, %dst[%c0, %i, %c0] : memref<1x33x42xf32>
}
```
Static 2D copy with offset becomes a nested loop with the static offset folded into the base store indices.
```mlir
// BEFORE
%view = 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>>
memref.copy %src, %view
// AFTER
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
%c16 = arith.constant 16 : index
%c33 = arith.constant 33 : index
scf.for %i = %c0 to %c33 step %c1 {
scf.for %j = %c0 to %c4 step %c1 {
%v = memref.load %src[%c0, %i, %j]
memref.store %v, %dst[%c0, %i, %c16 + %j]
}
}
```
## Scope
Supported:
- scalar copies, including dynamic offsets and dynamic view strides where all copied indices are zero
- effectively 1D copies with static strides
- static multidimensional copies with static offsets
- dynamic offsets where the offset can be used directly in one base dimension
Not addressed:
- rank-changing copy destinations
- dynamic copy/source shapes
- dynamic strides on non-unit copied dimensions
- multidimensional copies with dynamic offsets requiring div/mod delinearization
- non-identity-layout reinterpret_cast sources
## Correctness
The rewrite is restricted to static identity-layout base memrefs so view strides and static offsets can be mapped to concrete base indices. Cases assumed invalid by reinterpret_cast/copy semantics are guarded with assertions.
---
Patch is 51.50 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/203873.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td (+3)
- (modified) mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt (+1)
- (modified) mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp (+337-126)
- (modified) mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir (+229-323)
``````````diff
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..10b6b4e5412e7 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,264 @@ using namespace mlir;
namespace {
-/// Returns true if `rc` represents a scalar view (all sizes == 1)
-/// into a memref that has exactly one non-unit dimension located at
-/// either the first or last position (i.e. a "row" or "column").
-///
-/// Examples that return true:
-///
-/// // Row-major slice (last dim is non-unit)
-/// memref.reinterpret_cast %buff to offset: [%off],
-/// sizes: [1, 1, 1], strides: [1, 1, 1]
-/// : memref<1x1x8xi32> to memref<1x1x1xi32>
-///
-/// // 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>
+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 CopiedDimInfo {
+ unsigned viewDim;
+ unsigned baseDim;
+ int64_t size;
+};
+
+struct CopyToLoadStoreInfo {
+ SmallVector<CopiedDimInfo> copiedDims;
+ SmallVector<int64_t> staticOffsetIndices;
+ std::optional<unsigned> dynamicOffsetBaseDim;
+};
+
+static std::optional<SmallVector<CopiedDimInfo>>
+getCopiedDims(MemRefType baseType, MemRefType viewType,
+ ArrayRef<int64_t> baseStrides, ArrayRef<int64_t> viewStrides) {
+ SmallVector<CopiedDimInfo> copiedDims;
+ SmallVector<bool> usedBaseDims(baseType.getRank(), false);
+
+ for (auto [viewDim, viewSize] : llvm::enumerate(viewType.getShape())) {
+ if (viewSize == 1)
+ continue;
+
+ // TODO: Support dynamic strides on copied view dimensions.
+ if (ShapedType::isDynamic(viewStrides[viewDim]))
+ return std::nullopt;
+
+ std::optional<unsigned> baseDim = findBaseDimForViewStride(
+ baseType, baseStrides, usedBaseDims, viewStrides[viewDim], viewSize);
+ assert(baseDim &&
+ "static reinterpret_cast stride must map to an identity base "
+ "dimension");
+
+ usedBaseDims[*baseDim] = true;
+ copiedDims.push_back(
+ CopiedDimInfo{static_cast<unsigned>(viewDim), *baseDim, viewSize});
+ }
+
+ return copiedDims;
+}
+
+static bool copiedRegionFitsInBase(MemRefType baseType,
+ ArrayRef<int64_t> offsetIndices,
+ ArrayRef<CopiedDimInfo> copiedDims) {
+ for (const CopiedDimInfo &copiedDim : copiedDims) {
+ if (offsetIndices[copiedDim.baseDim] + copiedDim.size >
+ baseType.getDimSize(copiedDim.baseDim))
+ return false;
+ }
+ return true;
+}
+
+static std::optional<unsigned>
+getDynamicOffsetBaseDim(MemRefType baseType, ArrayRef<int64_t> baseStrides,
+ ArrayRef<CopiedDimInfo> copiedDims) {
+ // TODO: Support multidimensional dynamic offsets with div/mod
+ // delinearization.
+ if (copiedDims.size() > 1)
+ return std::nullopt;
+
+ if (copiedDims.empty()) {
+ // TODO: Support scalar dynamic offsets into bases with multiple non-unit
+ // dimensions, and all-unit bases with a provably zero offset.
+ return getSingleNonUnitDim(baseType);
+ }
+
+ unsigned baseDim = copiedDims.front().baseDim;
+ if (baseStrides[baseDim] == 1)
+ return baseDim;
+ return baseStrides.size() - 1;
+}
+
+/// Builds the index mapping needed to replace a copy into a reinterpret_cast
+/// view with scalar stores into the reinterpret_cast base.
///
-/// // Random strides
-/// memref.reinterpret_cast %buff to offset: [%off],
-/// sizes: [1, 1], strides: [10, 100]
-/// : memref<2x1xf32, strided<[10, 100]>>
-/// to memref<1x1xf32>
+/// 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.
///
-/// // Rank-1 case
-/// memref.reinterpret_cast %buf to offset: [%off],
-/// sizes: [1], strides: [1]
-/// : memref<8xi32> to memref<1xi32>
+/// Supported shape examples:
+/// copy memref<1x1xf32>
+/// to reinterpret_cast memref<1x108xf32>
+/// to memref<1x1xf32, strided<[?, ?], offset: ?>>
///
-/// Examples that return false:
+/// copy memref<1xNxf32>
+/// to reinterpret_cast memref<1xNxMxf32>
+/// to memref<1xNxf32, strided<[N*M, M]>>
///
-/// // More non-unit dims
-/// memref.reinterpret_cast %buff to offset: [%off],
-/// sizes: [1, 1, 1], strides: [1, 1, 1]
-/// : memref<1x2x8xi32> to memref<1x1x1xi32>
+/// copy memref<1xNxKxf32>
+/// to reinterpret_cast memref<1xNxMxf32>
+/// to memref<1xNxKxf32, strided<[N*M, M, 1], offset: O>>
///
-/// // View is not scalar (size != 1)
-/// memref.reinterpret_cast %buff to offset: [%off],
-/// sizes: [2, 1], strides: [1, 1]
-/// : memref<1x2xf32> to memref<2x1xf32>
+/// TODO examples:
+/// // Dynamic stride on a copied view dimension.
+/// copy memref<1xNxf32>
+/// to reinterpret_cast memref<1xNxMxf32>
+/// to memref<1xNxf32, strided<[?, ?]>>
///
-/// // Base has non-identity layout
-/// %buff = memref.alloc() : memref<1x2xf32, strided<[1, 3]>>
-/// memref.reinterpret_cast %buff to offset: [%off],
-/// sizes: [1, 1], strides: [1, 1]
-/// : memref<1x2xf32, strided<[1, 3]>> to memref<1x1xf32>
-static bool isScalarSlice(memref::ReinterpretCastOp rc) {
- auto rcInputTy = dyn_cast<MemRefType>(rc.getSource().getType());
- auto rcOutputTy = dyn_cast<MemRefType>(rc.getType());
-
- // Reject strided base - logic for computing linear idx is TODO
- if (!rcInputTy.getLayout().isIdentity())
- return false;
+/// // Multidimensional copy with dynamic linear offset.
+/// copy memref<1xNxKxf32>
+/// to reinterpret_cast memref<1xNxMxf32>
+/// to memref<1xNxKxf32, strided<[N*M, M, 1], offset: ?>>
+static std::optional<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;
+ // Non-unit view dimensions become copied dimensions in the scalar rewrite.
+ std::optional<SmallVector<CopiedDimInfo>> copiedDims =
+ getCopiedDims(baseType, viewType, *baseStrides, rc.getStaticStrides());
+ if (!copiedDims)
+ return std::nullopt;
+ info.copiedDims = std::move(*copiedDims);
+
+ 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");
+
+ assert(copiedRegionFitsInBase(baseType, *offsetIndices, info.copiedDims) &&
+ "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.
+ info.dynamicOffsetBaseDim =
+ getDynamicOffsetBaseDim(baseType, *baseStrides, info.copiedDims);
+ if (!info.dynamicOffsetBaseDim)
+ return std::nullopt;
+ 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.
-///
-/// Two layouts are supported:
-/// * row-major slice (stride pattern [N, ..., 1])
-/// * column-major slice (stride pattern [1, ..., N])
-///
-/// BEFORE (row-major slice)
+/// Before:
/// %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: ?>>
+/// to offset: [O], sizes: [...], 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>
///
-/// 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: ?>>
-/// memref.copy %src, %view
-/// : memref<1x...x1xf32>
-/// to memref<1x...x1xf32, strided<[1, ..., N], offset: ?>>
+/// After scalar copy:
+/// %v = memref.load %src[0, ..., 0]
+/// memref.store %v, %base[delinearized(O)]
///
-/// 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 copied dimensions:
+/// for each index tuple in the copied shape:
+/// %v = memref.load %src[index tuple]
+/// memref.store %v, %base[delinearized(O) + mapped index tuple]
+struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
public:
using OpRewritePattern::OpRewritePattern;
@@ -162,38 +298,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->copiedDims.size());
+ for (const CopiedDimInfo &copiedDim : copyInfo->copiedDims)
+ upperBounds.push_back(getOrCreateIndexConstant(copiedDim.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()) {
+ for (auto [idx, offset] :
+ llvm::enumerate(copyInfo->staticOffsetIndices)) {
+ if (offset == 0)
+ continue;
+ baseStoreIndices[idx] = getOrCreateIndexConstant(offset);
+ }
+ } else if (copyInfo->dynamicOffsetBaseDim) {
+ // 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->dynamicOffsetBaseDim] =
+ 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->copiedDims.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());
+ }
- Value val = memref::LoadOp::create(rewriter, loc, src, loadIndices);
- memref::StoreOp::create(rewriter, loc, val, dst, storeIndices);
+ // Load indices are zero except for copied view dimensions, which use the
+ // corresponding loop induction variables.
+ SmallVector<Value> loadIndices(srcType.getRank(), zero);
+ unsign...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/203873
More information about the Mlir-commits
mailing list