[Mlir-commits] [mlir] [memref] Simplify loads from reinterpret_casts preserving non-unit dims (PR #202683)
ioana ghiban
llvmlistbot at llvm.org
Fri Jun 19 03:58:06 PDT 2026
https://github.com/ioghiban updated https://github.com/llvm/llvm-project/pull/202683
>From 8a6f396ee3142a03b0e79bb05f26b937222703d9 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Tue, 9 Jun 2026 16:42:38 +0200
Subject: [PATCH 1/4] [memref] Simplify loads from reinterpret_casts preserving
non-unit dims
---
.../Transforms/ElideReinterpretCast.cpp | 225 ++++++-----------
.../MemRef/elide-reinterpret-cast.mlir | 231 ++++++------------
2 files changed, 158 insertions(+), 298 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 41dad1384da75..267990ffaee99 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -17,6 +17,7 @@
#include "llvm/ADT/Repeated.h"
#include <cassert>
#include <optional>
+#include <utility>
namespace mlir {
namespace memref {
@@ -198,42 +199,43 @@ 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 the mapping of preserved non-unit dimensions from the source MemRef
+/// to the result MemRef if both MemRefs have the same non-unit dimensions in
+/// the same order. Unit dimensions may be inserted or removed at any in-bounds
+/// position.
+static std::optional<SmallVector<std::pair<int64_t, int64_t>>>
+getNonUnitDimMapping(MemRefType inputTy, MemRefType outputTy) {
+ 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();
+ SmallVector<std::pair<int64_t, int64_t>> mapping;
+
+ while (inputDim < inputRank && outputDim < outputRank) {
+ if (inputShape[inputDim] == 1) {
+ ++inputDim;
+ continue;
+ }
+ if (outputShape[outputDim] == 1) {
+ ++outputDim;
+ continue;
+ }
-/// 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};
+ if (inputDim == inputRank || outputDim == outputRank)
+ return std::nullopt;
+
+ if (ShapedType::isDynamic(inputShape[inputDim]) ||
+ ShapedType::isDynamic(outputShape[outputDim]) ||
+ inputShape[inputDim] != outputShape[outputDim])
+ return std::nullopt;
+
+ mapping.push_back({inputDim, outputDim});
+ ++inputDim;
+ ++outputDim;
+ }
+ return mapping;
}
static bool hasStaticZeroOffset(memref::ReinterpretCastOp rc) {
@@ -262,13 +264,13 @@ static bool isConstantIndexExplicitlyOutOfBounds(Value idx,
}
/// Examples accepted by this shape restriction:
-/// memref<999xf32> <-> memref<1x1x999xf32>
-/// memref<1x108xf32> <-> memref<1x1x1x108xf32>
-/// memref<100x1xf32> <-> memref<100x1x1xf32>
-/// memref<1> <-> memref<1x1x1>
+/// memref<1x1x1x108xf32> <-> memref<1x108xf32>
+/// memref<100x1xf32> <-> memref<100x1x1xf32>
+/// memref<1x33x40xf32> <-> memref<33x1x1x40xf32>
+/// memref<1> <-> memref<1x1x1>
///
/// General reinterpret_casts are intentionally rejected.
-static bool isPureRankExpansionOrCollapsingRC(memref::ReinterpretCastOp rc) {
+static bool isUnitDimInsertionOrRemovalRC(memref::ReinterpretCastOp rc) {
auto inputTy = cast<MemRefType>(rc.getSource().getType());
auto outputTy = cast<MemRefType>(rc.getResult().getType());
@@ -284,45 +286,16 @@ static bool isPureRankExpansionOrCollapsingRC(memref::ReinterpretCastOp rc) {
llvm::any_of(rc.getStaticStrides(), ShapedType::isDynamic))
return false;
- // 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)
+ // Only unit-dim insertion/removal is accepted. The preserved non-unit
+ // dimensions must have the same static sizes and appear in the same order.
+ if (!getNonUnitDimMapping(inputTy, outputTy))
return false;
-
- // 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 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;
-
return true;
}
-/// 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.
+/// 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 isPureRankExpansionOrCollapsingRC().
+/// Shape restriction gated by isUnitDimInsertionOrRemovalRC().
///
/// 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,10 +343,10 @@ struct RewriteLoadFromReinterpretCast
if (!rc)
return rewriter.notifyMatchFailure(
op, "target is not a memref.reinterpret_cast");
- if (!isPureRankExpansionOrCollapsingRC(rc))
+ if (!isUnitDimInsertionOrRemovalRC(rc))
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!");
@@ -382,66 +354,27 @@ struct RewriteLoadFromReinterpretCast
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());
+
+ std::optional<SmallVector<std::pair<int64_t, int64_t>>> dimMapping =
+ getNonUnitDimMapping(rcInputTy, rcOutputTy);
+ assert(dimMapping && "expected matching non-unit dims");
+
+ // Ensures c0 defined only once.
+ auto getZeroIndex = [&]() -> Value {
+ for (auto [dim, size] : llvm::enumerate(rcOutputTy.getShape()))
+ if (size == 1)
+ return oldIdxs[dim];
+ return 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.
+ Value zeroIndex = getZeroIndex();
+ 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 +405,7 @@ struct ElideReinterpretCastPass
auto rc = op.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
if (!rc)
return true;
- return !isPureRankExpansionOrCollapsingRC(rc);
+ return !isUnitDimInsertionOrRemovalRC(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..22a597e60369a 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -229,8 +229,6 @@ 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>) {
@@ -249,35 +247,20 @@ func.func private @expand_scalar(%src : memref<1xi64>) {
// 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-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>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xi64>
%0 = memref.load %reinterpret_cast[%c0, %c0] : 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-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>
- return
-}
-
-// CHECK-LABEL: func.func private @expand_left_vector_dynamic_index(
+// CHECK-LABEL: func.func private @expand_1d_dynamic_index(
// CHECK-SAME: %[[I:.*]]: index
// CHECK-SAME: %[[SRC:.*]]: memref<999xi64>) {
-func.func private @expand_left_vector_dynamic_index(%i : index,
+func.func private @expand_1d_dynamic_index(%i : index,
%src : memref<999xi64>) {
// CHECK: %[[C0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
@@ -290,136 +273,112 @@ func.func private @expand_left_vector_dynamic_index(%i : index,
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-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>
- return
-}
-
-// CHECK-LABEL: func.func private @partial_expand_left_vector(
-// 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
+// CHECK-LABEL: func.func private @collapse_1d(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1x100xf32>) {
+func.func private @collapse_1d(
+ %src : memref<1x1x1x100xf32>) {
+ // CHECK: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: %[[C1:.*]] = arith.constant 1 : index
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : 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]
- : memref<1x1x999xf32>
+ 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]][%[[C0]], %[[C0]], %[[C0]], %[[C1]]] : memref<1x1x1x100xf32>
+ %0 = memref.load %reinterpret_cast[%c1, %c0, %c0] : memref<100x1x1xf32,
+ strided<[1, 100, 100]>>
return
}
-// CHECK-LABEL: func.func private @partial_collapse_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
+// CHECK-LABEL: func.func private @collapse_1d_moved_unit_dims(
+// CHECK-SAME: %[[I:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<33x1x1x1xf32>) {
+func.func private @collapse_1d_moved_unit_dims(%i : index,
+ %src : memref<33x1x1x1xf32>) {
+ // CHECK: %[[C0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : 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>
+ to offset: [0], sizes: [1, 33], strides: [33, 1]
+ : memref<33x1x1x1xf32> to memref<1x33xf32, strided<[33, 1]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[I]], %[[C0]], %[[C0]], %[[C0]]] : memref<33x1x1x1xf32>
+ %0 = memref.load %reinterpret_cast[%c0, %i]
+ : memref<1x33xf32, strided<[33, 1]>>
return
}
-// CHECK-LABEL: func.func private @expand_right_vector(
-// CHECK-SAME: %[[SRC:.*]]: memref<999xi64>) {
-func.func private @expand_right_vector(%src : memref<999xi64>) {
+// 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<1x33x1x1x40x1x33xf32>) {
+func.func private @collapse_3d_moved_unit_dims(%i : index, %j : index,
+ %k : index, %src : memref<1x33x1x1x40x1x33xf32>) {
// CHECK: %[[C0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : 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,
- strided<[1, 999, 999]>>
+ to offset: [0], sizes: [33, 1, 40, 33, 1, 1],
+ strides: [1320, 1320, 33, 1, 1, 1]
+ : memref<1x33x1x1x40x1x33xf32> to memref<33x1x40x33x1x1xf32,
+ strided<[1320, 1320, 33, 1, 1, 1]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]], %[[C0]], %[[C0]], %[[J]], %[[C0]], %[[K]]] : memref<1x33x1x1x40x1x33xf32>
+ %0 = memref.load %reinterpret_cast[%i, %c0, %j, %k, %c0, %c0]
+ : memref<33x1x40x33x1x1xf32, strided<[1320, 1320, 33, 1, 1, 1]>>
return
}
-// 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-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>
- return
-}
-
-// CHECK-LABEL: func.func private @collapse_right_vector_dynamic_index(
+// CHECK-LABEL: func.func private @expand_1d(
// CHECK-SAME: %[[I:.*]]: index
-// 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-SAME: %[[SRC:.*]]: memref<1x33xf32>) {
+func.func private @expand_1d(%i : index,
+ %src : memref<1x33xf32>) {
+ // CHECK: %[[C0:.*]] = arith.constant 0 : index
+ %c0 = 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>
- %0 = memref.load %reinterpret_cast[%i] : memref<999xi64>
+ to offset: [0], sizes: [1, 33, 1, 1], strides: [33, 1, 1, 1]
+ : memref<1x33xf32> to memref<1x33x1x1xf32>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]]] : memref<1x33xf32>
+ %0 = memref.load %reinterpret_cast[%c0, %i, %c0, %c0]
+ : memref<1x33x1x1xf32>
return
}
-// CHECK-LABEL: func.func private @partial_expand_right_vector(
-// CHECK-SAME: %[[SRC:.*]]: memref<999x1xf32>) {
-func.func private @partial_expand_right_vector(
- %src : memref<999x1xf32>) {
+// CHECK-LABEL: func.func private @expand_1d_moved_unit_dims(
+// CHECK-SAME: %[[I:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<1x33xf32>) {
+func.func private @expand_1d_moved_unit_dims(%i : index,
+ %src : memref<1x33xf32>) {
// CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[C1:.*]] = arith.constant 1 : index
%c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : 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]
- : memref<999x1x1xf32, strided<[1, 999, 999]>>
+ 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]][%[[C0]], %[[I]]] : memref<1x33xf32>
+ %0 = memref.load %reinterpret_cast[%i, %c0, %c0]
+ : memref<33x1x1xf32, strided<[1, 33, 33]>>
return
}
-// CHECK-LABEL: func.func private @partial_collapse_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
+// CHECK-LABEL: func.func private @expand_2d_moved_unit_dims(
+// CHECK-SAME: %[[I:[A-Za-z0-9_]+]]: index
+// CHECK-SAME: %[[J:[A-Za-z0-9_]+]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<1x33x40xf32>) {
+func.func private @expand_2d_moved_unit_dims(%i : index, %j : index,
+ %src : memref<1x33x40xf32>) {
+ // CHECK: %[[C0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : 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,
- strided<[1, 999]>>
+ to offset: [0], sizes: [33, 1, 1, 40], strides: [40, 40, 40, 1]
+ : memref<1x33x40xf32> to memref<33x1x1x40xf32,
+ strided<[40, 40, 40, 1]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]], %[[J]]] : memref<1x33x40xf32>
+ %0 = memref.load %reinterpret_cast[%i, %c0, %c0, %j]
+ : memref<33x1x1x40xf32, strided<[40, 40, 40, 1]>>
return
}
@@ -473,50 +432,18 @@ func.func private @negative_dynamic_stride(%stride0: index,
return
}
-// CHECK-LABEL: func.func private @negative_multiple_non_unit_dims(
+// CHECK-LABEL: func.func private @negative_reordered_non_unit_dims(
// CHECK-SAME: %[[SRC:.*]]: memref<2x1x1x100xf32>) {
-func.func private @negative_multiple_non_unit_dims(
+func.func private @negative_reordered_non_unit_dims(
%src : memref<2x1x1x100xf32>) {
- %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : 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: [100, 2], strides: [1, 100]
+ : memref<2x1x1x100xf32> to memref<100x2xf32, strided<[1, 100]>>
// CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %c1] : memref<2x100xf32>
- 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: %[[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]>>
- // CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %c1, %c0] : memref<1x100x1xf32,
- strided<[100, 1, 100]>>
- 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: %[[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]>>
- // CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c1, %c0, %c0] : memref<100x1x1xf32,
- strided<[1, 100, 100]>>
+ %0 = memref.load %reinterpret_cast[%c1, %c1] : memref<100x2xf32,
+ strided<[1, 100]>>
return
}
>From ca46358a50199a59eb3d16d39620add50045b360 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Tue, 16 Jun 2026 11:33:21 +0200
Subject: [PATCH 2/4] Clear tests diff
---
.../MemRef/elide-reinterpret-cast.mlir | 301 ++++++++++++++----
1 file changed, 244 insertions(+), 57 deletions(-)
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 22a597e60369a..9b2f9678b65c3 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -246,7 +246,7 @@ func.func private @expand_scalar(%src : memref<1xi64>) {
// 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: %[[C0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
@@ -257,10 +257,24 @@ func.func private @collapse_scalar(%src : memref<1x1x1xi64>) {
return
}
-// CHECK-LABEL: func.func private @expand_1d_dynamic_index(
+// 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-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>
+ return
+}
+
+// CHECK-LABEL: func.func private @expand_left_vector_dynamic_index(
// CHECK-SAME: %[[I:.*]]: index
// CHECK-SAME: %[[SRC:.*]]: memref<999xi64>) {
-func.func private @expand_1d_dynamic_index(%i : index,
+func.func private @expand_left_vector_dynamic_index(%i : index,
%src : memref<999xi64>) {
// CHECK: %[[C0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
@@ -273,66 +287,178 @@ func.func private @expand_1d_dynamic_index(%i : index,
return
}
-// CHECK-LABEL: func.func private @collapse_1d(
-// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1x100xf32>) {
-func.func private @collapse_1d(
- %src : memref<1x1x1x100xf32>) {
+// 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-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>
+ return
+}
+
+// CHECK-LABEL: func.func private @partial_expand_left_vector(
+// 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-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]
+ : memref<1x1x999xf32>
+ return
+}
+
+// CHECK-LABEL: func.func private @partial_collapse_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: %[[C1:.*]] = arith.constant 1 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : 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]], %[[C0]], %[[C1]]] : memref<1x1x999xf32>
+ %0 = memref.load %reinterpret_cast[%c0, %c1] : 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
- // CHECK: %[[C1:.*]] = arith.constant 1 : index
%c0 = arith.constant 0 : 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,
+ strided<[1, 999, 999]>>
+ return
+}
+
+// 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-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]][%[[C0]], %[[C0]], %[[C0]], %[[C1]]] : memref<1x1x1x100xf32>
- %0 = memref.load %reinterpret_cast[%c1, %c0, %c0] : memref<100x1x1xf32,
- strided<[1, 100, 100]>>
+ 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>
return
}
-// CHECK-LABEL: func.func private @collapse_1d_moved_unit_dims(
+// CHECK-LABEL: func.func private @collapse_right_vector_dynamic_index(
// CHECK-SAME: %[[I:.*]]: index
-// CHECK-SAME: %[[SRC:.*]]: memref<33x1x1x1xf32>) {
-func.func private @collapse_1d_moved_unit_dims(%i : index,
- %src : memref<33x1x1x1xf32>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
+// 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-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>
+ %0 = memref.load %reinterpret_cast[%i] : memref<999xi64>
+ return
+}
+
+// CHECK-LABEL: func.func private @partial_expand_right_vector(
+// CHECK-SAME: %[[SRC:.*]]: memref<999x1xf32>) {
+func.func private @partial_expand_right_vector(
+ %src : memref<999x1xf32>) {
+ // 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-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
- to offset: [0], sizes: [1, 33], strides: [33, 1]
- : memref<33x1x1x1xf32> to memref<1x33xf32, strided<[33, 1]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[I]], %[[C0]], %[[C0]], %[[C0]]] : memref<33x1x1x1xf32>
- %0 = memref.load %reinterpret_cast[%c0, %i]
- : memref<1x33xf32, strided<[33, 1]>>
+ 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]
+ : memref<999x1x1xf32, strided<[1, 999, 999]>>
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<1x33x1x1x40x1x33xf32>) {
-func.func private @collapse_3d_moved_unit_dims(%i : index, %j : index,
- %k : index, %src : memref<1x33x1x1x40x1x33xf32>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-LABEL: func.func private @partial_collapse_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: %[[C1:.*]] = arith.constant 1 : index
%c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
- to offset: [0], sizes: [33, 1, 40, 33, 1, 1],
- strides: [1320, 1320, 33, 1, 1, 1]
- : memref<1x33x1x1x40x1x33xf32> to memref<33x1x40x33x1x1xf32,
- strided<[1320, 1320, 33, 1, 1, 1]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]], %[[C0]], %[[C0]], %[[J]], %[[C0]], %[[K]]] : memref<1x33x1x1x40x1x33xf32>
- %0 = memref.load %reinterpret_cast[%i, %c0, %j, %k, %c0, %c0]
- : memref<33x1x40x33x1x1xf32, strided<[1320, 1320, 33, 1, 1, 1]>>
+ 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]]] : memref<999x1x1xf32>
+ %0 = memref.load %reinterpret_cast[%c1, %c0] : memref<999x1xf32,
+ strided<[1, 999]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @expand_multiple_non_unit_dims(
+// CHECK-SAME: %[[SRC:.*]]: memref<2x100xf32>) {
+func.func private @expand_multiple_non_unit_dims(
+ %src : memref<2x100xf32>) {
+ // 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-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %src
+ to offset: [0], sizes: [2, 1, 1, 100], strides: [100, 100, 100, 1]
+ : memref<2x100xf32> to memref<2x1x1x100xf32,
+ strided<[100, 100, 100, 1]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C1]]] : memref<2x100xf32>
+ %0 = memref.load %reinterpret_cast[%c0, %c0, %c0, %c1]
+ : memref<2x1x1x100xf32, strided<[100, 100, 100, 1]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @collapse_multiple_non_unit_dims(
+// CHECK-SAME: %[[SRC:.*]]: memref<2x1x1x100xf32>) {
+func.func private @collapse_multiple_non_unit_dims(
+ %src : memref<2x1x1x100xf32>) {
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C0_0:.*]] = arith.constant 0 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %src
+ to offset: [0], sizes: [2, 100], strides: [100, 1]
+ : memref<2x1x1x100xf32> to memref<2x100xf32>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0_0]], %[[C0_0]], %[[C1]]] : memref<2x1x1x100xf32>
+ %0 = memref.load %reinterpret_cast[%c0, %c1] : memref<2x100xf32>
return
}
-// CHECK-LABEL: func.func private @expand_1d(
+// CHECK-LABEL: func.func private @expand_inner_non_unit_dims(
// CHECK-SAME: %[[I:.*]]: index
// CHECK-SAME: %[[SRC:.*]]: memref<1x33xf32>) {
-func.func private @expand_1d(%i : index,
+func.func private @expand_inner_non_unit_dims(%i : index,
%src : memref<1x33xf32>) {
// CHECK: %[[C0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
@@ -346,39 +472,100 @@ func.func private @expand_1d(%i : index,
return
}
-// CHECK-LABEL: func.func private @expand_1d_moved_unit_dims(
+// 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: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : 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]][%[[C0]], %[[C0]], %[[C0]], %[[C1]]] : memref<1x1x1x100xf32>
+ %0 = memref.load %reinterpret_cast[%c0, %c1, %c0] : 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_1d_moved_unit_dims(%i : index,
+func.func private @expand_diff_non_unit_boundary(%i : index,
%src : memref<1x33xf32>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
+ // 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-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]][%[[C0]], %[[I]]] : memref<1x33xf32>
- %0 = memref.load %reinterpret_cast[%i, %c0, %c0]
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C1]]] : memref<1x33xf32>
+ %0 = memref.load %reinterpret_cast[%c1, %c0, %c0]
: memref<33x1x1xf32, strided<[1, 33, 33]>>
return
}
-// CHECK-LABEL: func.func private @expand_2d_moved_unit_dims(
+// 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: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : 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]][%[[C0]], %[[C0]], %[[C0]], %[[C1]]] : memref<1x1x1x100xf32>
+ %0 = memref.load %reinterpret_cast[%c1, %c0, %c0] : 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<1x33x40x33xf32>) {
+func.func private @expand_3d_moved_unit_dims(%i : index, %j : index,
+ %k : index, %src : memref<1x33x40x33xf32>) {
+ // CHECK: %[[C0:.*]] = arith.constant 0 : index
+ %c0 = arith.constant 0 : index
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %src
+ to offset: [0], sizes: [33, 1, 1, 40, 1, 33],
+ strides: [1320, 1320, 1320, 33, 33, 1]
+ : memref<1x33x40x33xf32> to memref<33x1x1x40x1x33xf32,
+ strided<[1320, 1320, 1320, 33, 33, 1]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]], %[[J]], %[[K]]] : memref<1x33x40x33xf32>
+ %0 = memref.load %reinterpret_cast[%i, %c0, %c0, %j, %c0, %k]
+ : memref<33x1x1x40x1x33xf32, strided<[1320, 1320, 1320, 33, 33, 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: %[[SRC:.*]]: memref<1x33x40xf32>) {
-func.func private @expand_2d_moved_unit_dims(%i : index, %j : index,
- %src : memref<1x33x40xf32>) {
+// CHECK-SAME: %[[K:[A-Za-z0-9_]+]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<1x33x1x1x40x1x33xf32>) {
+func.func private @collapse_3d_moved_unit_dims(%i : index, %j : index,
+ %k : index, %src : memref<1x33x1x1x40x1x33xf32>) {
// CHECK: %[[C0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = 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]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]], %[[J]]] : memref<1x33x40xf32>
- %0 = memref.load %reinterpret_cast[%i, %c0, %c0, %j]
- : memref<33x1x1x40xf32, strided<[40, 40, 40, 1]>>
+ to offset: [0], sizes: [33, 1, 40, 33, 1, 1],
+ strides: [1320, 1320, 33, 1, 1, 1]
+ : memref<1x33x1x1x40x1x33xf32> to memref<33x1x40x33x1x1xf32,
+ strided<[1320, 1320, 33, 1, 1, 1]>>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]], %[[C0]], %[[C0]], %[[J]], %[[C0]], %[[K]]] : memref<1x33x1x1x40x1x33xf32>
+ %0 = memref.load %reinterpret_cast[%i, %c0, %j, %k, %c0, %c0]
+ : memref<33x1x40x33x1x1xf32, strided<[1320, 1320, 33, 1, 1, 1]>>
return
}
@@ -432,9 +619,9 @@ func.func private @negative_dynamic_stride(%stride0: index,
return
}
-// CHECK-LABEL: func.func private @negative_reordered_non_unit_dims(
+// CHECK-LABEL: func.func private @negative_diff_non_unit_dims_order(
// CHECK-SAME: %[[SRC:.*]]: memref<2x1x1x100xf32>) {
-func.func private @negative_reordered_non_unit_dims(
+func.func private @negative_diff_non_unit_dims_order(
%src : memref<2x1x1x100xf32>) {
%c1 = arith.constant 1 : index
// CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[SRC]]
>From 23ec7cfda26b55947b9a791fa5fb24ce7669b2a7 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Tue, 16 Jun 2026 12:13:12 +0200
Subject: [PATCH 3/4] Fixup
---
.../Transforms/ElideReinterpretCast.cpp | 19 +++++++++++--------
.../MemRef/elide-reinterpret-cast.mlir | 3 +--
2 files changed, 12 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 267990ffaee99..8ab15d9307420 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -362,16 +362,19 @@ struct RewriteLoadFromReinterpretCast
getNonUnitDimMapping(rcInputTy, rcOutputTy);
assert(dimMapping && "expected matching non-unit dims");
- // Ensures c0 defined only once.
- auto getZeroIndex = [&]() -> Value {
- for (auto [dim, size] : llvm::enumerate(rcOutputTy.getShape()))
- if (size == 1)
- return oldIdxs[dim];
- return arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
- };
+ // 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);
// Initialize new load indices to all 0s.
- Value zeroIndex = getZeroIndex();
SmallVector<Value> rcInputIdxs(rcInputRank, zeroIndex);
for (auto [inputDim, outputDim] : *dimMapping)
rcInputIdxs[inputDim] = oldIdxs[outputDim];
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 9b2f9678b65c3..7df20b8f07da8 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -443,14 +443,13 @@ func.func private @collapse_multiple_non_unit_dims(
%src : memref<2x1x1x100xf32>) {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- // CHECK-DAG: %[[C0_0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
// CHECK-NOT: memref.reinterpret_cast
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [2, 100], strides: [100, 1]
: memref<2x1x1x100xf32> to memref<2x100xf32>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0_0]], %[[C0_0]], %[[C1]]] : memref<2x1x1x100xf32>
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]], %[[C1]]] : memref<2x1x1x100xf32>
%0 = memref.load %reinterpret_cast[%c0, %c1] : memref<2x100xf32>
return
}
>From b149cd05651edc2f1d64021ed4cae0fdd804edcd Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Fri, 19 Jun 2026 11:17:18 +0200
Subject: [PATCH 4/4] Address first round of comments
---
.../Transforms/ElideReinterpretCast.cpp | 112 +++----
.../MemRef/elide-reinterpret-cast.mlir | 308 +++++++++---------
2 files changed, 210 insertions(+), 210 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 8ab15d9307420..0fd1ff83f15b5 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -17,7 +17,6 @@
#include "llvm/ADT/Repeated.h"
#include <cassert>
#include <optional>
-#include <utility>
namespace mlir {
namespace memref {
@@ -199,44 +198,9 @@ struct CopyToScalarLoadAndStore : public OpRewritePattern<memref::CopyOp> {
}
};
-/// Returns the mapping of preserved non-unit dimensions from the source MemRef
-/// to the result MemRef if both MemRefs have the same non-unit dimensions in
-/// the same order. Unit dimensions may be inserted or removed at any in-bounds
-/// position.
-static std::optional<SmallVector<std::pair<int64_t, int64_t>>>
-getNonUnitDimMapping(MemRefType inputTy, MemRefType outputTy) {
- 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();
- SmallVector<std::pair<int64_t, int64_t>> mapping;
-
- while (inputDim < inputRank && outputDim < outputRank) {
- if (inputShape[inputDim] == 1) {
- ++inputDim;
- continue;
- }
- if (outputShape[outputDim] == 1) {
- ++outputDim;
- continue;
- }
-
- if (inputDim == inputRank || outputDim == outputRank)
- return std::nullopt;
-
- if (ShapedType::isDynamic(inputShape[inputDim]) ||
- ShapedType::isDynamic(outputShape[outputDim]) ||
- inputShape[inputDim] != outputShape[outputDim])
- return std::nullopt;
-
- mapping.push_back({inputDim, outputDim});
- ++inputDim;
- ++outputDim;
- }
- return mapping;
-}
+//===----------------------------------------------------------------------===//
+// Load Rewrite Helpers
+//===----------------------------------------------------------------------===//
static bool hasStaticZeroOffset(memref::ReinterpretCastOp rc) {
ArrayRef<int64_t> offsets = rc.getStaticOffsets();
@@ -263,14 +227,22 @@ static bool isConstantIndexExplicitlyOutOfBounds(Value idx,
return idxVal && (*idxVal < 0 || *idxVal >= upperBound);
}
-/// Examples accepted by this shape restriction:
+using NonUnitDimMapping = SmallVector<std::pair<int64_t, int64_t>>;
+
+/// Shape restriction accepting only unit-dim insertion/removal
+/// reinterpret_casts.
+///
+/// Examples accepted:
/// memref<1x1x1x108xf32> <-> memref<1x108xf32>
/// memref<100x1xf32> <-> memref<100x1x1xf32>
/// memref<1x33x40xf32> <-> memref<33x1x1x40xf32>
/// memref<1> <-> memref<1x1x1>
///
-/// General reinterpret_casts are intentionally rejected.
-static bool isUnitDimInsertionOrRemovalRC(memref::ReinterpretCastOp rc) {
+/// 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());
@@ -278,19 +250,47 @@ static bool isUnitDimInsertionOrRemovalRC(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 unit-dim insertion/removal is accepted. The preserved non-unit
- // dimensions must have the same static sizes and appear in the same order.
- if (!getNonUnitDimMapping(inputTy, outputTy))
- return false;
- return true;
+ 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;
+ }
+
+ if (inputDim == inputRank || outputDim == outputRank)
+ return std::nullopt;
+
+ if (ShapedType::isDynamic(inputShape[inputDim]) ||
+ ShapedType::isDynamic(outputShape[outputDim]) ||
+ inputShape[inputDim] != outputShape[outputDim])
+ return std::nullopt;
+
+ mapping.push_back({inputDim, outputDim});
+ ++inputDim;
+ ++outputDim;
+ }
+ return mapping;
}
/// Checks statically known and constant indices accessed by a load from a
@@ -314,8 +314,8 @@ static bool isUnitDimInsertionOrRemovalRC(memref::ReinterpretCastOp rc) {
/// 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 isUnitDimInsertionOrRemovalRC().
+///
+/// Shape restriction gated by getNonUnitDimMapping().
///
/// BEFORE (rank expansion)
/// %view = memref.reinterpret_cast %src
@@ -343,7 +343,8 @@ struct RewriteLoadFromReinterpretCast
if (!rc)
return rewriter.notifyMatchFailure(
op, "target is not a memref.reinterpret_cast");
- if (!isUnitDimInsertionOrRemovalRC(rc))
+ std::optional<NonUnitDimMapping> dimMapping = getNonUnitDimMapping(rc);
+ if (!dimMapping)
return rewriter.notifyMatchFailure(
op, "reinterpret_cast is not a unit-dim insertion/removal preserving "
"non-unit dimensions");
@@ -351,17 +352,12 @@ struct RewriteLoadFromReinterpretCast
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 rcInputRank = rcInputTy.getRank();
SmallVector<Value> oldIdxs(op.getIndices().begin(), op.getIndices().end());
- std::optional<SmallVector<std::pair<int64_t, int64_t>>> dimMapping =
- getNonUnitDimMapping(rcInputTy, rcOutputTy);
- assert(dimMapping && "expected matching non-unit dims");
-
// Prefer reusing an explicit constant-zero index from the old load.
Value zeroIndex;
for (Value idx : oldIdxs) {
@@ -408,7 +404,7 @@ struct ElideReinterpretCastPass
auto rc = op.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
if (!rc)
return true;
- return !isUnitDimInsertionOrRemovalRC(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 7df20b8f07da8..0e23372d754a5 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -232,42 +232,44 @@ func.func private @negative_plain_copy(%src : memref<1x1xf32>,
// 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: %[[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], strides: [1, 1]
: memref<1x1x1xi64> to memref<1x1xi64>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[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
}
@@ -276,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
}
@@ -306,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
}
@@ -324,30 +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: %[[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]], %[[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
}
@@ -355,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
}
@@ -372,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
}
@@ -386,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-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: [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
}
@@ -404,53 +408,52 @@ 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: %[[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]]] : 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
}
// CHECK-LABEL: func.func private @expand_multiple_non_unit_dims(
-// CHECK-SAME: %[[SRC:.*]]: memref<2x100xf32>) {
+// CHECK-SAME: %[[SRC:.*]]: memref<17x100xf32>) {
func.func private @expand_multiple_non_unit_dims(
- %src : memref<2x100xf32>) {
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+ %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: [2, 1, 1, 100], strides: [100, 100, 100, 1]
- : memref<2x100xf32> to memref<2x1x1x100xf32,
+ 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]][%[[C0]], %[[C1]]] : memref<2x100xf32>
- %0 = memref.load %reinterpret_cast[%c0, %c0, %c0, %c1]
- : memref<2x1x1x100xf32, 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 @collapse_multiple_non_unit_dims(
-// CHECK-SAME: %[[SRC:.*]]: memref<2x1x1x100xf32>) {
+// CHECK-SAME: %[[SRC:.*]]: memref<17x1x1x100xf32>) {
func.func private @collapse_multiple_non_unit_dims(
- %src : memref<2x1x1x100xf32>) {
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+ %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: [0], sizes: [2, 100], strides: [100, 1]
- : memref<2x1x1x100xf32> to memref<2x100xf32>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]], %[[C1]]] : memref<2x1x1x100xf32>
- %0 = memref.load %reinterpret_cast[%c0, %c1] : memref<2x100xf32>
+ 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
}
@@ -459,14 +462,14 @@ func.func private @collapse_multiple_non_unit_dims(
// CHECK-SAME: %[[SRC:.*]]: memref<1x33xf32>) {
func.func private @expand_inner_non_unit_dims(%i : index,
%src : memref<1x33xf32>) {
- // 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, 33, 1, 1], strides: [33, 1, 1, 1]
: memref<1x33xf32> to memref<1x33x1x1xf32>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]]] : memref<1x33xf32>
- %0 = memref.load %reinterpret_cast[%c0, %i, %c0, %c0]
+ // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX]], %[[I]]] : memref<1x33xf32>
+ %0 = memref.load %reinterpret_cast[%idx, %i, %idx, %idx]
: memref<1x33x1x1xf32>
return
}
@@ -475,16 +478,16 @@ func.func private @expand_inner_non_unit_dims(%i : index,
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1x100xf32>) {
func.func private @collapse_inner_non_unit_dims(
%src : memref<1x1x1x100xf32>) {
- // 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, 100, 1], strides: [100, 1, 100]
: memref<1x1x1x100xf32> to memref<1x100x1xf32, strided<[100, 1, 100]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]], %[[C1]]] : memref<1x1x1x100xf32>
- %0 = memref.load %reinterpret_cast[%c0, %c1, %c0] : memref<1x100x1xf32,
+ // 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
}
@@ -494,16 +497,16 @@ func.func private @collapse_inner_non_unit_dims(
// CHECK-SAME: %[[SRC:.*]]: memref<1x33xf32>) {
func.func private @expand_diff_non_unit_boundary(%i : index,
%src : memref<1x33xf32>) {
- // 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: [33, 1, 1], strides: [1, 33, 33]
: memref<1x33xf32> to memref<33x1x1xf32, strided<[1, 33, 33]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C1]]] : memref<1x33xf32>
- %0 = memref.load %reinterpret_cast[%c1, %c0, %c0]
+ // 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
}
@@ -512,16 +515,16 @@ func.func private @expand_diff_non_unit_boundary(%i : index,
// CHECK-SAME: %[[SRC:.*]]: memref<1x1x1x100xf32>) {
func.func private @collapse_diff_non_unit_boundary(
%src : memref<1x1x1x100xf32>) {
- // 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: [100, 1, 1], strides: [1, 100, 100]
: memref<1x1x1x100xf32> to memref<100x1x1xf32, strided<[1, 100, 100]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]], %[[C1]]] : memref<1x1x1x100xf32>
- %0 = memref.load %reinterpret_cast[%c1, %c0, %c0] : memref<100x1x1xf32,
+ // 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
}
@@ -530,20 +533,20 @@ func.func private @collapse_diff_non_unit_boundary(
// 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<1x33x40x33xf32>) {
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x22x3xf32>) {
func.func private @expand_3d_moved_unit_dims(%i : index, %j : index,
- %k : index, %src : memref<1x33x40x33xf32>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- %c0 = arith.constant 0 : 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: [33, 1, 1, 40, 1, 33],
- strides: [1320, 1320, 1320, 33, 33, 1]
- : memref<1x33x40x33xf32> to memref<33x1x1x40x1x33xf32,
- strided<[1320, 1320, 1320, 33, 33, 1]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]], %[[J]], %[[K]]] : memref<1x33x40x33xf32>
- %0 = memref.load %reinterpret_cast[%i, %c0, %c0, %j, %c0, %k]
- : memref<33x1x1x40x1x33xf32, strided<[1320, 1320, 1320, 33, 33, 1]>>
+ 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
}
@@ -551,20 +554,20 @@ func.func private @expand_3d_moved_unit_dims(%i : index, %j : index,
// 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<1x33x1x1x40x1x33xf32>) {
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1x1x22x1x3xf32>) {
func.func private @collapse_3d_moved_unit_dims(%i : index, %j : index,
- %k : index, %src : memref<1x33x1x1x40x1x33xf32>) {
- // CHECK: %[[C0:.*]] = arith.constant 0 : index
- %c0 = arith.constant 0 : 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: [33, 1, 40, 33, 1, 1],
- strides: [1320, 1320, 33, 1, 1, 1]
- : memref<1x33x1x1x40x1x33xf32> to memref<33x1x40x33x1x1xf32,
- strided<[1320, 1320, 33, 1, 1, 1]>>
- // CHECK: %[[LOAD:.*]] = memref.load %[[SRC]][%[[C0]], %[[I]], %[[C0]], %[[C0]], %[[J]], %[[C0]], %[[K]]] : memref<1x33x1x1x40x1x33xf32>
- %0 = memref.load %reinterpret_cast[%i, %c0, %j, %k, %c0, %c0]
- : memref<33x1x40x33x1x1xf32, strided<[1320, 1320, 33, 1, 1, 1]>>
+ 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
}
@@ -573,62 +576,63 @@ func.func private @collapse_3d_moved_unit_dims(%i : index, %j : index,
//===----------------------------------------------------------------------===//
// CHECK-LABEL: func.func private @negative_nonzero_offset(
-// CHECK-SAME: %[[SRC:.*]]: memref<1xi64>) {
+// CHECK-SAME: %[[SRC:.*]]: memref<1x100xf32>) {
func.func private @negative_nonzero_offset(
- %src : memref<1xi64>) {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+ %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: [1], sizes: [1, 1, 1], strides: [1, 1, 1]
- : memref<1xi64> to memref<1x1x1xi64, strided<[1, 1, 1], offset: 1>>
+ 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, %c0, %c1]
- : memref<1x1x1xi64, strided<[1, 1, 1], offset: 1>>
+ %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_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-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: [1, %dim], strides: [1, 1]
- : memref<?xi64> to memref<1x?xi64>
+ : memref<?xf32> to memref<1x?xf32>
// CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %i] : memref<1x?xi64>
+ %0 = memref.load %reinterpret_cast[%idx_1, %idx_2] : memref<1x?xf32>
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-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, 1], strides: [%stride0, %stride1]
- : memref<1x108xi64> to memref<1x1xi64, strided<[?, ?]>>
+ to offset: [0], sizes: [108], strides: [%stride]
+ : memref<1x108xf32> to memref<108xf32, strided<[?]>>
// CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%c0, %c1]
- : memref<1x1xi64, strided<[?, ?]>>
+ %0 = memref.load %reinterpret_cast[%idx_2]
+ : memref<108xf32, strided<[?]>>
return
}
// CHECK-LABEL: func.func private @negative_diff_non_unit_dims_order(
-// CHECK-SAME: %[[SRC:.*]]: memref<2x1x1x100xf32>) {
+// CHECK-SAME: %[[SRC:.*]]: memref<17x1x1x100xf32>) {
func.func private @negative_diff_non_unit_dims_order(
- %src : memref<2x1x1x100xf32>) {
- %c1 = arith.constant 1 : index
+ %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, 2], strides: [1, 100]
- : memref<2x1x1x100xf32> to memref<100x2xf32, strided<[1, 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, %c1] : memref<100x2xf32,
+ %0 = memref.load %reinterpret_cast[%idx, %idx] : memref<100x17xf32,
strided<[1, 100]>>
return
}
@@ -637,13 +641,13 @@ func.func private @negative_diff_non_unit_dims_order(
// 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