[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