[Mlir-commits] [mlir] 985bb1a - [memref] Simplify loads from `reinterpret_casts` preserving non-unit dims (#202683)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 29 03:12:46 PDT 2026
Author: ioana ghiban
Date: 2026-06-29T12:12:41+02:00
New Revision: 985bb1a0597959bdc82aae3e16fc8812159ab5e0
URL: https://github.com/llvm/llvm-project/commit/985bb1a0597959bdc82aae3e16fc8812159ab5e0
DIFF: https://github.com/llvm/llvm-project/commit/985bb1a0597959bdc82aae3e16fc8812159ab5e0.diff
LOG: [memref] Simplify loads from `reinterpret_casts` preserving non-unit dims (#202683)
Extend `memref.load` rewriting from `memref.reinterpret_cast` introduced
in [#188459](https://github.com/llvm/llvm-project/pull/188459) to handle
unit-dimension insertion/removal with multiple preserved non-unit
dimensions. This generalization addresses patterns in the MLGO regalloc
model.
## Example
Before:
```mlir
%view = memref.reinterpret_cast %src
to offset: [0], sizes: [33, 1, 1, 40], strides: [40, 40, 40, 1]
: memref<1x33x40xf32> to memref<33x1x1x40xf32,
strided<[40, 40, 40, 1]>>
%v = memref.load %view[%i, %c0, %c0, %j]
: memref<33x1x1x40xf32, strided<[40, 40, 40, 1]>>
```
After:
```mlir
%v = memref.load %src[%c0, %i, %j] : memref<1x33x40xf32>
```
## Scope
- Requires zero reinterpret-cast offset.
- Requires fully static sizes and strides.
- Preserved non-unit dimensions must have the same static sizes and
appear in the same order.
- Dropped/inserted dimensions must be unit dimensions, so their indices
are known to be zero for in-bounds loads.
This does not try to handle general `memref.reinterpret_cast` operations
with arbitrary offsets, arbitrary stride remapping, or cases where
dropping an index could change semantics.
Correctness is preserved by matching non-unit dimensions between source
and result in order. The rewritten load maps each preserved non-unit
result index back to its corresponding source dimension, and fills
source unit dimensions with zero indices. Constant out-of-bounds indices
are still rejected by the existing debug assertion path.
This PR aims to address part of the patterns raised in
https://github.com/llvm/llvm-project/issues/202536
Assisted-by: Codex (refine implementation + tests). I reviewed all code
and tests before submission.
Added:
Modified:
mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 41dad1384da75..0fd1ff83f15b5 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -198,43 +198,9 @@ struct CopyToScalarLoadAndStore : public OpRewritePattern<memref::CopyOp> {
}
};
-/// Captures info about MemRefs that are effectively 1D (the leading or trailing
-/// dims are all 1). The only accepted non-unit dim is either the leading of the
-/// trailing dim.
-///
-/// Examples:
-/// memref<1x1x4xf32>, memref<4x1x1xf32>, memref<1x1x1xf32>
-///
-struct ShapeInfoFor1DMemRef {
- // Are all dims == 1? `false` means that there is exactly one dim != 1.
- bool allOnes = true;
- // If there is a non-unit boundary dim, is it the leading or the trailing dim?
- bool isLeadingDimNonUnit = false;
-};
-
-/// Returns information about a MemRef if it contains at most one non-unit
-/// dimension.
-///
-/// The single non-unit dimension, if present, must be on the left or right
-/// boundary. Rank-1 non-unit MemRefs are treated as being on both boundaries.
-static std::optional<ShapeInfoFor1DMemRef>
-getShapeInfoFor1DMemRef(MemRefType type) {
- ArrayRef<int64_t> shape = type.getShape();
- int64_t nonUnitCount =
- llvm::count_if(shape, [](int64_t dim) { return dim != 1; });
- // Return default values if missing non-unit dimension (all-ones MemRef).
- if (nonUnitCount == 0)
- return ShapeInfoFor1DMemRef{};
- // Return no info if MemRef has more non-unit dimensions.
- if (nonUnitCount > 1)
- return std::nullopt;
- // Return no info if MemRef has non-unit dimension in non-boundary positions.
- if (shape.front() == 1 && shape.back() == 1)
- return std::nullopt;
-
- return ShapeInfoFor1DMemRef{/*allOnes=*/false,
- /*isLeadingDimNonUnit=*/shape.front() != 1};
-}
+//===----------------------------------------------------------------------===//
+// Load Rewrite Helpers
+//===----------------------------------------------------------------------===//
static bool hasStaticZeroOffset(memref::ReinterpretCastOp rc) {
ArrayRef<int64_t> offsets = rc.getStaticOffsets();
@@ -261,14 +227,22 @@ static bool isConstantIndexExplicitlyOutOfBounds(Value idx,
return idxVal && (*idxVal < 0 || *idxVal >= upperBound);
}
-/// Examples accepted by this shape restriction:
-/// memref<999xf32> <-> memref<1x1x999xf32>
-/// memref<1x108xf32> <-> memref<1x1x1x108xf32>
-/// memref<100x1xf32> <-> memref<100x1x1xf32>
-/// memref<1> <-> memref<1x1x1>
+using NonUnitDimMapping = SmallVector<std::pair<int64_t, int64_t>>;
+
+/// Shape restriction accepting only unit-dim insertion/removal
+/// reinterpret_casts.
///
-/// General reinterpret_casts are intentionally rejected.
-static bool isPureRankExpansionOrCollapsingRC(memref::ReinterpretCastOp rc) {
+/// Examples accepted:
+/// memref<1x1x1x108xf32> <-> memref<1x108xf32>
+/// memref<100x1xf32> <-> memref<100x1x1xf32>
+/// memref<1x33x40xf32> <-> memref<33x1x1x40xf32>
+/// memref<1> <-> memref<1x1x1>
+///
+/// Returns the mapping of non-unit dimensions from the source
+/// to the result MemRef if the reinterpret_cast preserved sizes and order (no
+/// transposition) of these dimensions.
+static std::optional<NonUnitDimMapping>
+getNonUnitDimMapping(memref::ReinterpretCastOp rc) {
auto inputTy = cast<MemRefType>(rc.getSource().getType());
auto outputTy = cast<MemRefType>(rc.getResult().getType());
@@ -276,53 +250,52 @@ static bool isPureRankExpansionOrCollapsingRC(memref::ReinterpretCastOp rc) {
// offsets would require reasoning about storage shifts in the underlying
// reinterpret_cast, which this helper does not model.
if (!hasStaticZeroOffset(rc))
- return false;
+ return std::nullopt;
// Dynamic sizes/strides prevent precise reasoning about the underlying
// reinterpret_cast, so only fully static shape metadata is accepted.
if (llvm::any_of(rc.getStaticSizes(), ShapedType::isDynamic) ||
llvm::any_of(rc.getStaticStrides(), ShapedType::isDynamic))
- return false;
+ return std::nullopt;
- // Only shapes with at most one non-unit dimension are accepted. This rules
- // out more general multi-dimensional reinterpret_casts and restricts the
- // helper to unit-dim insertion/removal around a single logical dimension.
- std::optional<ShapeInfoFor1DMemRef> inputNonUnitDim =
- getShapeInfoFor1DMemRef(inputTy);
- std::optional<ShapeInfoFor1DMemRef> outputNonUnitDim =
- getShapeInfoFor1DMemRef(outputTy);
- // Bail out if either type does not satisfy the single-boundary-non-unit-dim
- // restriction described above.
- if (!inputNonUnitDim || !outputNonUnitDim)
- return false;
+ ArrayRef<int64_t> inputShape = inputTy.getShape();
+ ArrayRef<int64_t> outputShape = outputTy.getShape();
+ int64_t inputDim = 0;
+ int64_t outputDim = 0;
+ int64_t inputRank = inputTy.getRank();
+ int64_t outputRank = outputTy.getRank();
+ NonUnitDimMapping mapping;
+
+ // The preserved non-unit dimensions must have the same static sizes and
+ // appear in the same order.
+ while (inputDim < inputRank || outputDim < outputRank) {
+ if (inputDim < inputRank && inputShape[inputDim] == 1) {
+ ++inputDim;
+ continue;
+ }
+ if (outputDim < outputRank && outputShape[outputDim] == 1) {
+ ++outputDim;
+ continue;
+ }
- // The source and result must either both have a single non-unit dimension
- // or both be all-ones.
- if (inputNonUnitDim->allOnes != outputNonUnitDim->allOnes)
- return false;
- if (inputNonUnitDim->allOnes)
- return true;
-
- // The preserved non-unit dimension must have the same size.
- if (inputTy.getDimSize(
- inputNonUnitDim->isLeadingDimNonUnit ? 0 : inputTy.getRank() - 1) !=
- outputTy.getDimSize(
- outputNonUnitDim->isLeadingDimNonUnit ? 0 : outputTy.getRank() - 1))
- return false;
+ if (inputDim == inputRank || outputDim == outputRank)
+ return std::nullopt;
- // If both sides have rank > 1, the non-unit dimension must be on the same
- // boundary. Rank-1 MemRefs are accepted against either boundary.
- if (inputTy.getRank() != 1 && outputTy.getRank() != 1 &&
- inputNonUnitDim->isLeadingDimNonUnit !=
- outputNonUnitDim->isLeadingDimNonUnit)
- return false;
+ if (ShapedType::isDynamic(inputShape[inputDim]) ||
+ ShapedType::isDynamic(outputShape[outputDim]) ||
+ inputShape[inputDim] != outputShape[outputDim])
+ return std::nullopt;
- return true;
+ mapping.push_back({inputDim, outputDim});
+ ++inputDim;
+ ++outputDim;
+ }
+ return mapping;
}
-/// Checks statically known and constant indices accessed by a load from a pure
-/// rank expansion/collapsing to ensure in-bounds only access. Fully dynamic
-/// indices are skipped (there is no way to verify them).
+/// Checks statically known and constant indices accessed by a load from a
+/// unit-dim insertion/removal reinterpret_cast to ensure in-bounds only access.
+/// Fully dynamic indices are skipped (there is no way to verify them).
[[maybe_unused]] static bool areIndicesInBounds(memref::LoadOp load) {
auto rc = load.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
auto rcOutputTy = cast<MemRefType>(rc.getResult().getType());
@@ -339,27 +312,26 @@ static bool isPureRankExpansionOrCollapsingRC(memref::ReinterpretCastOp rc) {
return true;
}
-/// Rewrites `memref.load` through a pure rank-only `reinterpret_cast` by
-/// mapping the load indices directly onto the source MemRef.
-
-/// Shape restriction gated by isPureRankExpansionOrCollapsingRC().
+/// Rewrites `memref.load` through a reinterpret_cast that only inserts/removes
+/// unit dimensions by mapping the load indices directly onto the source MemRef.
+///
+/// Shape restriction gated by getNonUnitDimMapping().
///
/// BEFORE (rank expansion)
/// %view = memref.reinterpret_cast %src
-/// : memref<Nxf32> to memref<1x1xNxf32>
-/// %v = memref.load %view[%c0, %c0, %i] : memref<1x1xNxf32>
+/// : memref<1xNxMxf32> to memref<Nx1x1xMxf32>
+/// %v = memref.load %view[%i, %c0, %c0, %j] : memref<Nx1x1xMxf32>
///
/// AFTER
-/// %v = memref.load %src[%i] : memref<Nxf32>
+/// %v = memref.load %src[%c0, %i, %j] : memref<1xNxMxf32>
///
/// BEFORE (rank collapsing)
/// %view = memref.reinterpret_cast %src
-/// : memref<1x1xNxf32> to memref<Nxf32>
-/// %v = memref.load %view[%i] : memref<Nxf32>
+/// : memref<Nx1x1xMxf32> to memref<1xNxMxf32>
+/// %v = memref.load %view[%c0, %i, %j] : memref<1xNxMxf32>
///
/// AFTER
-/// %c0 = arith.constant 0 : index
-/// %v = memref.load %src[%c0, %c0, %i] : memref<1x1xNxf32>
+/// %v = memref.load %src[%i, %c0, %c0, %j] : memref<Nx1x1xMxf32>
struct RewriteLoadFromReinterpretCast
: public OpRewritePattern<memref::LoadOp> {
public:
@@ -371,77 +343,37 @@ struct RewriteLoadFromReinterpretCast
if (!rc)
return rewriter.notifyMatchFailure(
op, "target is not a memref.reinterpret_cast");
- if (!isPureRankExpansionOrCollapsingRC(rc))
+ std::optional<NonUnitDimMapping> dimMapping = getNonUnitDimMapping(rc);
+ if (!dimMapping)
return rewriter.notifyMatchFailure(
- op, "reinterpret_cast is not a pure rank expansion or collapsing of "
- "a single dimension");
+ op, "reinterpret_cast is not a unit-dim insertion/removal preserving "
+ "non-unit dimensions");
assert(areIndicesInBounds(op) &&
"load from reinterpret_cast indexes out of bounds!");
- auto rcOutputTy = cast<MemRefType>(rc.getResult().getType());
auto rcInputTy = cast<MemRefType>(rc.getSource().getType());
- int64_t rcOutputRank = rcOutputTy.getRank();
int64_t rcInputRank = rcInputTy.getRank();
- SmallVector<Value> idxs(op.getIndices().begin(), op.getIndices().end());
- SmallVector<Value> rcInputIdxs;
- rcInputIdxs.reserve(rcInputRank);
-
- // The rewrite only supports reinterpret_casts with at most one non-unit
- // dimension, located at the left or right boundary.
- //
- // The higher-rank side tells which side the reinterpret_cast has
- // expanded/collapsed.
- //
- // expansion: rcOutput has the higher rank
- // collapsing : rcInput has the higher rank
- //
- // Example:
- // memref<999> -> memref<1x1x999> : leading extra dims
- // memref<999x1x1> -> memref<999> : trailing extra dims
- MemRefType expandedTy =
- rcOutputRank >= rcInputRank ? rcOutputTy : rcInputTy;
- std::optional<ShapeInfoFor1DMemRef> expandedNonUnitDim =
- getShapeInfoFor1DMemRef(expandedTy);
- assert(expandedNonUnitDim && "expected a single boundary non-unit dim");
- bool keepLeadingIndices = expandedNonUnitDim->isLeadingDimNonUnit;
-
- if (rcOutputRank >= rcInputRank) {
- // Rank expansion:
- // memref<N> -> memref<1x1xN> : keep the last rcInputRank indices
- // memref<N> -> memref<Nx1x1> : keep the first rcInputRank indices
- // memref<1> -> memref<1x1x1> : all indices are zero
- //
- // Any discarded indices are known to be zero from
- // areIndicesInBounds().
- int64_t firstKeptPos =
- keepLeadingIndices ? 0 : rcOutputRank - rcInputRank;
- rcInputIdxs.append(idxs.begin() + firstKeptPos,
- idxs.begin() + firstKeptPos + rcInputRank);
- } else {
- // Rank collapsing:
- // memref<1x1xN> -> memref<N> : reinsert leading zeros
- // memref<Nx1x1> -> memref<N> : reinsert trailing zeros
- // memref<1x1x1> -> memref<1> : all indices are zero
- //
- // The collapsed-away dimensions are unit dims, so re-adding them with
- // zero indices preserves semantics.
- Value c0 = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
- int64_t rankDiff = rcInputRank - rcOutputRank;
-
- if (keepLeadingIndices) {
- rcInputIdxs.append(idxs.begin(), idxs.end());
- rcInputIdxs.append(rankDiff, c0);
- } else {
- rcInputIdxs.append(rankDiff, c0);
- rcInputIdxs.append(idxs.begin(), idxs.end());
+ SmallVector<Value> oldIdxs(op.getIndices().begin(), op.getIndices().end());
+
+ // Prefer reusing an explicit constant-zero index from the old load.
+ Value zeroIndex;
+ for (Value idx : oldIdxs) {
+ std::optional<int64_t> idxVal = getConstantIndex(idx);
+ if (idxVal && *idxVal == 0) {
+ zeroIndex = idx;
+ break;
}
}
+ if (!zeroIndex)
+ zeroIndex = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
- assert(rcInputIdxs.size() == static_cast<size_t>(rcInputRank) &&
- "Incorrect number of indices!");
+ // Initialize new load indices to all 0s.
+ SmallVector<Value> rcInputIdxs(rcInputRank, zeroIndex);
+ for (auto [inputDim, outputDim] : *dimMapping)
+ rcInputIdxs[inputDim] = oldIdxs[outputDim];
auto rcInput = rc.getSource();
// If the only user of rc is the current Op (which is about to be erased),
@@ -472,7 +404,7 @@ struct ElideReinterpretCastPass
auto rc = op.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
if (!rc)
return true;
- return !isPureRankExpansionOrCollapsingRC(rc);
+ return !getNonUnitDimMapping(rc);
});
target.addLegalDialect<arith::ArithDialect, memref::MemRefDialect>();
if (failed(applyPartialConversion(getOperation(), target,
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 61b6d480ce7a0..0e23372d754a5 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -229,48 +229,47 @@ func.func private @negative_plain_copy(%src : memref<1x1xf32>,
// Positive tests
//===----------------------------------------------------------------------===//
-/// For rank-1 MemRefs, expansion/collapsing may be considered on either side.
-
// CHECK-LABEL: func.func private @expand_scalar(
// CHECK-SAME: %[[SRC:.*]]: memref<1xi64>) {
func.func private @expand_scalar(%src : memref<1xi64>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- %c0 = arith.constant 0 : index
+ // CHECK: %[[IDX:.*]] = arith.constant 0 : index
+ %idx = arith.constant 0 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [1, 1, 1], strides: [1, 1, 1]
: memref<1xi64> to memref<1x1x1xi64>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]]] : memref<1xi64>
- %0 = memref.load %reinterpret_cast[%c0, %c0, %c0] : memref<1x1x1xi64>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<1xi64>
+ %0 = memref.load %reinterpret_cast[%idx, %idx, %idx] : memref<1x1x1xi64>
return
}
// CHECK-LABEL: func.func private @collapse_scalar(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1xi64>) {
func.func private @collapse_scalar(%src : memref<1x1x1xi64>) {
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C0_0:.*]] = arith.constant 0 : index
- %c0 = arith.constant 0 : index
+ // CHECK: %[[IDX:.*]] = arith.constant 0 : index
+ %idx = arith.constant 0 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [1, 1], strides: [1, 1]
: memref<1x1x1xi64> to memref<1x1xi64>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0_0]], %[[C0]], %[[C0]]] : memref<1x1x1xi64>
- %0 = memref.load %reinterpret_cast[%c0, %c0] : memref<1x1xi64>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX]], %[[IDX]], %[[IDX]]] : memref<1x1x1xi64>
+ %0 = memref.load %reinterpret_cast[%idx, %idx] : memref<1x1xi64>
return
}
// CHECK-LABEL: func.func private @expand_left_vector(
// CHECK-SAME: %[[SRC:.*]]: memref<999xi64>) {
func.func private @expand_left_vector(%src : memref<999xi64>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- %c0 = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [1, 1, 999], strides: [999, 999, 1]
: memref<999xi64> to memref<1x1x999xi64>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]]] : memref<999xi64>
- %0 = memref.load %reinterpret_cast[%c0, %c0, %c0] : memref<1x1x999xi64>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_2]]] : memref<999xi64>
+ %0 = memref.load %reinterpret_cast[%idx_1, %idx_1, %idx_2] : memref<1x1x999xi64>
return
}
@@ -279,29 +278,29 @@ func.func private @expand_left_vector(%src : memref<999xi64>) {
// CHECK-SAME: %[[SRC:.*]]: memref<999xi64>) {
func.func private @expand_left_vector_dynamic_index(%i : index,
%src : memref<999xi64>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- %c0 = arith.constant 0 : index
+ // CHECK: %[[IDX:.*]] = arith.constant 0 : index
+ %idx = arith.constant 0 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [1, 1, 999], strides: [999, 999, 1]
: memref<999xi64> to memref<1x1x999xi64>
// CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[I]]] : memref<999xi64>
- %0 = memref.load %reinterpret_cast[%c0, %c0, %i] : memref<1x1x999xi64>
+ %0 = memref.load %reinterpret_cast[%idx, %idx, %i] : memref<1x1x999xi64>
return
}
// CHECK-LABEL: func.func private @collapse_left_vector(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x999xi64>) {
func.func private @collapse_left_vector(%src : memref<1x1x999xi64>) {
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- %c1 = arith.constant 1 : index
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx = arith.constant 13 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [999], strides: [1]
: memref<1x1x999xi64> to memref<999xi64>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C1]]] : memref<1x1x999xi64>
- %0 = memref.load %reinterpret_cast[%c1] : memref<999xi64>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_1]], %[[IDX_1]], %[[IDX_2]]] : memref<1x1x999xi64>
+ %0 = memref.load %reinterpret_cast[%idx] : memref<999xi64>
return
}
@@ -309,16 +308,16 @@ func.func private @collapse_left_vector(%src : memref<1x1x999xi64>) {
// CHECK-SAME: %[[SRC:.*]]: memref<1x999xf32>) {
func.func private @partial_expand_left_vector(
%src : memref<1x999xf32>) {
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [1, 1, 999], strides: [999, 999, 1]
: memref<1x999xf32> to memref<1x1x999xf32>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C1]]] : memref<1x999xf32>
- %0 = memref.load %reinterpret_cast[%c0, %c0, %c1]
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_1]], %[[IDX_2]]] : memref<1x999xf32>
+ %0 = memref.load %reinterpret_cast[%idx_1, %idx_1, %idx_2]
: memref<1x1x999xf32>
return
}
@@ -327,31 +326,32 @@ func.func private @partial_expand_left_vector(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x999xf32>) {
func.func private @partial_collapse_left_vector(
%src : memref<1x1x999xf32>) {
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C0_0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [1, 999], strides: [999, 1]
: memref<1x1x999xf32> to memref<1x999xf32>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0_0]], %[[C0]], %[[C1]]] : memref<1x1x999xf32>
- %0 = memref.load %reinterpret_cast[%c0, %c1] : memref<1x999xf32>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_1]], %[[IDX_1]], %[[IDX_2]]] : memref<1x1x999xf32>
+ %0 = memref.load %reinterpret_cast[%idx_1, %idx_2] : memref<1x999xf32>
return
}
// CHECK-LABEL: func.func private @expand_right_vector(
// CHECK-SAME: %[[SRC:.*]]: memref<999xi64>) {
func.func private @expand_right_vector(%src : memref<999xi64>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- %c0 = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [999, 1, 1], strides: [1, 999, 999]
: memref<999xi64> to memref<999x1x1xi64, strided<[1, 999, 999]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]]] : memref<999xi64>
- %0 = memref.load %reinterpret_cast[%c0, %c0, %c0] : memref<999x1x1xi64,
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_2]]] : memref<999xi64>
+ %0 = memref.load %reinterpret_cast[%idx_2, %idx_1, %idx_1] : memref<999x1x1xi64,
strided<[1, 999, 999]>>
return
}
@@ -359,15 +359,15 @@ func.func private @expand_right_vector(%src : memref<999xi64>) {
// CHECK-LABEL: func.func private @collapse_right_vector(
// CHECK-SAME: %[[SRC:.*]]: memref<999x1x1xi64>) {
func.func private @collapse_right_vector(%src : memref<999x1x1xi64>) {
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- %c1 = arith.constant 1 : index
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx = arith.constant 13 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [999], strides: [1]
: memref<999x1x1xi64> to memref<999xi64>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C1]], %[[C0]], %[[C0]]] : memref<999x1x1xi64>
- %0 = memref.load %reinterpret_cast[%c1] : memref<999xi64>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_2]], %[[IDX_1]], %[[IDX_1]]] : memref<999x1x1xi64>
+ %0 = memref.load %reinterpret_cast[%idx] : memref<999xi64>
return
}
@@ -376,12 +376,12 @@ func.func private @collapse_right_vector(%src : memref<999x1x1xi64>) {
// CHECK-SAME: %[[SRC:.*]]: memref<999x1x1xi64>) {
func.func private @collapse_right_vector_dynamic_index(%i : index,
%src : memref<999x1x1xi64>) {
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX:.*]] = arith.constant 0 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [999], strides: [1]
: memref<999x1x1xi64> to memref<999xi64>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[I]], %[[C0]], %[[C0]]] : memref<999x1x1xi64>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[I]], %[[IDX]], %[[IDX]]] : memref<999x1x1xi64>
%0 = memref.load %reinterpret_cast[%i] : memref<999xi64>
return
}
@@ -390,16 +390,16 @@ func.func private @collapse_right_vector_dynamic_index(%i : index,
// CHECK-SAME: %[[SRC:.*]]: memref<999x1xf32>) {
func.func private @partial_expand_right_vector(
%src : memref<999x1xf32>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[C1:.*]] = arith.constant 1 : index
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [999, 1, 1], strides: [1, 999, 999]
: memref<999x1xf32> to memref<999x1x1xf32, strided<[1, 999, 999]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C1]], %[[C0]]] : memref<999x1xf32>
- %0 = memref.load %reinterpret_cast[%c1, %c0, %c0]
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_2]], %[[IDX_1]]] : memref<999x1xf32>
+ %0 = memref.load %reinterpret_cast[%idx_2, %idx_1, %idx_1]
: memref<999x1x1xf32, strided<[1, 999, 999]>>
return
}
@@ -408,115 +408,232 @@ func.func private @partial_expand_right_vector(
// CHECK-SAME: %[[SRC:.*]]: memref<999x1x1xf32>) {
func.func private @partial_collapse_right_vector(
%src : memref<999x1x1xf32>) {
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C0_0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [999, 1], strides: [1, 999]
: memref<999x1x1xf32> to memref<999x1xf32, strided<[1, 999]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C1]], %[[C0]], %[[C0_0]]] : memref<999x1x1xf32>
- %0 = memref.load %reinterpret_cast[%c1, %c0] : memref<999x1xf32,
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_2]], %[[IDX_1]], %[[IDX_1]]] : memref<999x1x1xf32>
+ %0 = memref.load %reinterpret_cast[%idx_2, %idx_1] : memref<999x1xf32,
strided<[1, 999]>>
return
}
-//===----------------------------------------------------------------------===//
-// Negative tests (must NOT rewrite)
-//===----------------------------------------------------------------------===//
+// CHECK-LABEL: func.func private @expand_multiple_non_unit_dims(
+// CHECK-SAME: %[[SRC:.*]]: memref<17x100xf32>) {
+func.func private @expand_multiple_non_unit_dims(
+ %src : memref<17x100xf32>) {
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %src
+ to offset: [0], sizes: [17, 1, 1, 100], strides: [100, 100, 100, 1]
+ : memref<17x100xf32> to memref<17x1x1x100xf32,
+ strided<[100, 100, 100, 1]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_2]], %[[IDX_2]]] : memref<17x100xf32>
+ %0 = memref.load %reinterpret_cast[%idx_2, %idx_1, %idx_1, %idx_2]
+ : memref<17x1x1x100xf32, strided<[100, 100, 100, 1]>>
+ return
+}
-// CHECK-LABEL: func.func private @negative_nonzero_offset(
-// CHECK-SAME: %[[SRC:.*]]: memref<1xi64>) {
-func.func private @negative_nonzero_offset(
- %src : memref<1xi64>) {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
- // CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[SRC]]
+// CHECK-LABEL: func.func private @collapse_multiple_non_unit_dims(
+// CHECK-SAME: %[[SRC:.*]]: memref<17x1x1x100xf32>) {
+func.func private @collapse_multiple_non_unit_dims(
+ %src : memref<17x1x1x100xf32>) {
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx = arith.constant 13 : index
+ // CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
- to offset: [1], sizes: [1, 1, 1], strides: [1, 1, 1]
- : memref<1xi64> to memref<1x1x1xi64, strided<[1, 1, 1], offset: 1>>
- // CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %c0, %c1]
- : memref<1x1x1xi64, strided<[1, 1, 1], offset: 1>>
+ to offset: [0], sizes: [17, 100], strides: [100, 1]
+ : memref<17x1x1x100xf32> to memref<17x100xf32>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_2]], %[[IDX_1]], %[[IDX_1]], %[[IDX_2]]] : memref<17x1x1x100xf32>
+ %0 = memref.load %reinterpret_cast[%idx, %idx] : memref<17x100xf32>
return
}
-// CHECK-LABEL: func.func private @negative_dynamic_shape(
-// CHECK-SAME: %[[SRC:[A-Za-z][A-Za-z0-9-]*]]: memref<?xi64>
-func.func private @negative_dynamic_shape(%dim : index, %i : index,
- %src : memref<?xi64>) {
- %c0 = arith.constant 0 : index
- // CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[SRC]]
+// CHECK-LABEL: func.func private @expand_inner_non_unit_dims(
+// CHECK-SAME: %[[I:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<1x33xf32>) {
+func.func private @expand_inner_non_unit_dims(%i : index,
+ %src : memref<1x33xf32>) {
+ // CHECK: %[[IDX:.*]] = arith.constant 0 : index
+ %idx = arith.constant 0 : index
+ // CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
- to offset: [0], sizes: [1, %dim], strides: [1, 1]
- : memref<?xi64> to memref<1x?xi64>
- // CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %i] : memref<1x?xi64>
+ to offset: [0], sizes: [1, 33, 1, 1], strides: [33, 1, 1, 1]
+ : memref<1x33xf32> to memref<1x33x1x1xf32>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX]], %[[I]]] : memref<1x33xf32>
+ %0 = memref.load %reinterpret_cast[%idx, %i, %idx, %idx]
+ : memref<1x33x1x1xf32>
return
}
-// CHECK-LABEL: func.func private @negative_dynamic_stride(
-// CHECK-SAME: %[[SRC:[A-Za-z][A-Za-z0-9-]*]]: memref<1x108xi64>
-func.func private @negative_dynamic_stride(%stride0: index,
- %stride1: index, %src : memref<1x108xi64>) {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+// CHECK-LABEL: func.func private @collapse_inner_non_unit_dims(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1x100xf32>) {
+func.func private @collapse_inner_non_unit_dims(
+ %src : memref<1x1x1x100xf32>) {
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %src
+ to offset: [0], sizes: [1, 100, 1], strides: [100, 1, 100]
+ : memref<1x1x1x100xf32> to memref<1x100x1xf32, strided<[100, 1, 100]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_1]], %[[IDX_1]], %[[IDX_1]], %[[IDX_2]]] : memref<1x1x1x100xf32>
+ %0 = memref.load %reinterpret_cast[%idx_1, %idx_2, %idx_1] : memref<1x100x1xf32,
+ strided<[100, 1, 100]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @expand_
diff _non_unit_boundary(
+// CHECK-SAME: %[[I:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<1x33xf32>) {
+func.func private @expand_
diff _non_unit_boundary(%i : index,
+ %src : memref<1x33xf32>) {
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %src
+ to offset: [0], sizes: [33, 1, 1], strides: [1, 33, 33]
+ : memref<1x33xf32> to memref<33x1x1xf32, strided<[1, 33, 33]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_1]], %[[IDX_2]]] : memref<1x33xf32>
+ %0 = memref.load %reinterpret_cast[%idx_2, %idx_1, %idx_1]
+ : memref<33x1x1xf32, strided<[1, 33, 33]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @collapse_
diff _non_unit_boundary(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1x100xf32>) {
+func.func private @collapse_
diff _non_unit_boundary(
+ %src : memref<1x1x1x100xf32>) {
+ // CHECK-DAG: %[[IDX_1:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[IDX_2:.*]] = arith.constant 13 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %src
+ to offset: [0], sizes: [100, 1, 1], strides: [1, 100, 100]
+ : memref<1x1x1x100xf32> to memref<100x1x1xf32, strided<[1, 100, 100]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_1]], %[[IDX_1]], %[[IDX_1]], %[[IDX_2]]] : memref<1x1x1x100xf32>
+ %0 = memref.load %reinterpret_cast[%idx_2, %idx_1, %idx_1] : memref<100x1x1xf32,
+ strided<[1, 100, 100]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @expand_3d_moved_unit_dims(
+// CHECK-SAME: %[[I:[A-Za-z0-9_]+]]: index
+// CHECK-SAME: %[[J:[A-Za-z0-9_]+]]: index
+// CHECK-SAME: %[[K:[A-Za-z0-9_]+]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x22x3xf32>) {
+func.func private @expand_3d_moved_unit_dims(%i : index, %j : index,
+ %k : index, %src : memref<1x3x22x3xf32>) {
+ // CHECK: %[[IDX:.*]] = arith.constant 0 : index
+ %idx = arith.constant 0 : index
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %src
+ to offset: [0], sizes: [3, 1, 1, 22, 1, 3],
+ strides: [66, 66, 66, 3, 3, 1]
+ : memref<1x3x22x3xf32> to memref<3x1x1x22x1x3xf32,
+ strided<[66, 66, 66, 3, 3, 1]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX]], %[[I]], %[[J]], %[[K]]] : memref<1x3x22x3xf32>
+ %0 = memref.load %reinterpret_cast[%i, %idx, %idx, %j, %idx, %k]
+ : memref<3x1x1x22x1x3xf32, strided<[66, 66, 66, 3, 3, 1]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @collapse_3d_moved_unit_dims(
+// CHECK-SAME: %[[I:[A-Za-z0-9_]+]]: index
+// CHECK-SAME: %[[J:[A-Za-z0-9_]+]]: index
+// CHECK-SAME: %[[K:[A-Za-z0-9_]+]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1x1x22x1x3xf32>) {
+func.func private @collapse_3d_moved_unit_dims(%i : index, %j : index,
+ %k : index, %src : memref<1x3x1x1x22x1x3xf32>) {
+ // CHECK: %[[IDX:.*]] = arith.constant 0 : index
+ %idx_1 = arith.constant 0 : index
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %src
+ to offset: [0], sizes: [3, 1, 22, 3, 1, 1],
+ strides: [66, 66, 3, 1, 1, 1]
+ : memref<1x3x1x1x22x1x3xf32> to memref<3x1x22x3x1x1xf32,
+ strided<[66, 66, 3, 1, 1, 1]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX]], %[[I]], %[[IDX]], %[[IDX]], %[[J]], %[[IDX]], %[[K]]] : memref<1x3x1x1x22x1x3xf32>
+ %0 = memref.load %reinterpret_cast[%i, %idx_1, %j, %k, %idx_1, %idx_1]
+ : memref<3x1x22x3x1x1xf32, strided<[66, 66, 3, 1, 1, 1]>>
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// Negative tests (must NOT rewrite)
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func private @negative_nonzero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x100xf32>) {
+func.func private @negative_nonzero_offset(
+ %src : memref<1x100xf32>) {
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[SRC]]
%reinterpret_cast = memref.reinterpret_cast %src
- to offset: [0], sizes: [1, 1], strides: [%stride0, %stride1]
- : memref<1x108xi64> to memref<1x1xi64, strided<[?, ?]>>
+ to offset: [1], sizes: [1, 1, 100], strides: [1, 1, 1]
+ : memref<1x100xf32> to memref<1x1x100xf32, strided<[1, 1, 1], offset: 1>>
// CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %c1]
- : memref<1x1xi64, strided<[?, ?]>>
+ %0 = memref.load %reinterpret_cast[%idx_1, %idx_1, %idx_2]
+ : memref<1x1x100xf32, strided<[1, 1, 1], offset: 1>>
return
}
-// CHECK-LABEL: func.func private @negative_multiple_non_unit_dims(
-// CHECK-SAME: %[[SRC:.*]]: memref<2x1x1x100xf32>) {
-func.func private @negative_multiple_non_unit_dims(
- %src : memref<2x1x1x100xf32>) {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+// CHECK-LABEL: func.func private @negative_dynamic_shape(
+// CHECK-SAME: %[[SRC:[A-Za-z][A-Za-z0-9-]*]]: memref<?xf32>
+func.func private @negative_dynamic_shape(%dim : index,
+ %src : memref<?xf32>) {
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[SRC]]
%reinterpret_cast = memref.reinterpret_cast %src
- to offset: [0], sizes: [2, 100], strides: [100, 1]
- : memref<2x1x1x100xf32> to memref<2x100xf32>
+ to offset: [0], sizes: [1, %dim], strides: [1, 1]
+ : memref<?xf32> to memref<1x?xf32>
// CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %c1] : memref<2x100xf32>
+ %0 = memref.load %reinterpret_cast[%idx_1, %idx_2] : memref<1x?xf32>
return
}
-// CHECK-LABEL: func.func private @negative_inner_non_unit_dims(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1x100xf32>) {
-func.func private @negative_inner_non_unit_dims(
- %src : memref<1x1x1x100xf32>) {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+// CHECK-LABEL: func.func private @negative_dynamic_stride(
+// CHECK-SAME: %[[SRC:[A-Za-z][A-Za-z0-9-]*]]: memref<1x108xf32>
+func.func private @negative_dynamic_stride(%stride: index,
+ %src : memref<1x108xf32>) {
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[SRC]]
%reinterpret_cast = memref.reinterpret_cast %src
- to offset: [0], sizes: [1, 100, 1], strides: [100, 1, 100]
- : memref<1x1x1x100xf32> to memref<1x100x1xf32, strided<[100, 1, 100]>>
+ to offset: [0], sizes: [108], strides: [%stride]
+ : memref<1x108xf32> to memref<108xf32, strided<[?]>>
// CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %c1, %c0] : memref<1x100x1xf32,
- strided<[100, 1, 100]>>
+ %0 = memref.load %reinterpret_cast[%idx_2]
+ : memref<108xf32, strided<[?]>>
return
}
-// CHECK-LABEL: func.func private @negative_
diff _non_unit_boundary(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1x100xf32>) {
-func.func private @negative_
diff _non_unit_boundary(
- %src : memref<1x1x1x100xf32>) {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+// CHECK-LABEL: func.func private @negative_
diff _non_unit_dims_order(
+// CHECK-SAME: %[[SRC:.*]]: memref<17x1x1x100xf32>) {
+func.func private @negative_
diff _non_unit_dims_order(
+ %src : memref<17x1x1x100xf32>) {
+ %idx = arith.constant 13 : index
// CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[SRC]]
%reinterpret_cast = memref.reinterpret_cast %src
- to offset: [0], sizes: [100, 1, 1], strides: [1, 100, 100]
- : memref<1x1x1x100xf32> to memref<100x1x1xf32, strided<[1, 100, 100]>>
+ to offset: [0], sizes: [100, 17], strides: [1, 100]
+ : memref<17x1x1x100xf32> to memref<100x17xf32, strided<[1, 100]>>
// CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c1, %c0, %c0] : memref<100x1x1xf32,
- strided<[1, 100, 100]>>
+ %0 = memref.load %reinterpret_cast[%idx, %idx] : memref<100x17xf32,
+ strided<[1, 100]>>
return
}
@@ -524,13 +641,13 @@ func.func private @negative_
diff _non_unit_boundary(
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1x100xf32>) {
func.func private @negative_
diff _non_unit_size(
%src : memref<1x1x1x100xf32>) {
- %c0 = arith.constant 0 : index
- %c98 = arith.constant 98 : index
+ %idx_1 = arith.constant 0 : index
+ %idx_2 = arith.constant 13 : index
// CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[SRC]]
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [1, 99], strides: [99, 1]
: memref<1x1x1x100xf32> to memref<1x99xf32>
// CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %c98] : memref<1x99xf32>
+ %0 = memref.load %reinterpret_cast[%idx_1, %idx_2] : memref<1x99xf32>
return
}
More information about the Mlir-commits
mailing list