[mlir][MemRef] Compute unused dimensions of a rank-reducing subviews using strides as well.

[mravishankar]

Summary

For memref.subview operations, when there are more than one
unit-dimensions, the strides need to be used to figure out which of
the unit-dims are actually dropped.

[ftynse]

Sorry for delay.

How does this change the situation for subviews with dynamic strides? Dynamic strides will all have the same static value so it looks like the code will just assume all dimensions are dropped because of that.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

693

Nit: please document and drop trivial braces below.

725

Typo: dimension

730

Did you mean size rather than rank?

732

Nit: spurious empty line

744

756–758

If this can never happen, it should be an assertion. I see you are asserting below, but maybe let's better to assert here and put your comment as assertion message + have different messages for different that currently return None.

mlir/test/Dialect/MemRef/invalid.mlir

365

Nit: please add a newline

[nicolasvasilache]

sorry I should have commented here .. I will pick this up and rewrite some of the basics, this should make the whole thing much smaller code-wise.

[mravishankar]

In D109418#3001440, @ftynse wrote:

Sorry for delay.

How does this change the situation for subviews with dynamic strides? Dynamic strides will all have the same static value so it looks like the code will just assume all dimensions are dropped because of that.

Ill explain by example. Lets say you have two unit dimensions with dynamic strides, s0 and s1, if either of one is dropped, then the result will have only one dynamic stride s0 irrespective of which one is dropped. So this is handled (I also have a test for this I think).

In D109418#3001483, @nicolasvasilache wrote:

sorry I should have commented here .. I will pick this up and rewrite some of the basics, this should make the whole thing much smaller code-wise.

Thats fine, as long as the tests here pass. (In any case, I think the logic here is actually quite simple, it is the simplest of all candidates I could come up with (dont think number of lines of codes correspond to simplicity :) ). I am not tied to the solution though, so please let me know when you have this approach working.

[nicolasvasilache]

Let me unblock this for now and I'll come back to it when I have cycles.

[mravishankar]

Rebase and address comments.

[mravishankar]

@ftynse addressed most of your comments. I am running presubmits, and will submit after they pass. Ill address any other comments you have post-submit.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

756–758

Well, this is used in the verifier as well, so it needs to return on error and not assert. I could have different messages, but that would just be little verbose. Finally its about type mismatch. The caller reports it as such.

[mravishankar]

Change from using DenseMap<int64_t> to std::map<int64_t> since the dynamic stride value is same as tombstone value for DenseMap<int64_t>.

[springerm]

There is a bug in the implementation. It crashes with this test case:

// RUN: mlir-opt %s -canonicalize

func.func @func(%arg0: tensor<?x?x?xf32>, %arg1: index) -> index {
  %0 = bufferization.to_memref %arg0 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
  %c1 = arith.constant 1 : index
  %subview = memref.subview %0[0, 0, 0] [1, %arg1, 1] [1, 1, 1] : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to memref<1x?xf32, strided<[?, ?], offset: ?>>
  %dim = memref.dim %subview, %c1 : memref<1x?xf32, strided<[?, ?], offset: ?>>
  return %dim : index
}

Crashes here:

assert(subview.isDynamicSize(sourceIndex) &&
       "expected dynamic subview size");

It is not clear to me from the comments and the description why we consider strides at all. Couldn't we just make an arbitrary decision which dims are dropped (like tensor.extract_slice)? We could just drop the same strides. (This means that the valid result types of a memref.subview would change a bit.)

In D109418#3001836, @mravishankar wrote:

In D109418#3001440, @ftynse wrote:

Sorry for delay.

How does this change the situation for subviews with dynamic strides? Dynamic strides will all have the same static value so it looks like the code will just assume all dimensions are dropped because of that.

Ill explain by example. Lets say you have two unit dimensions with dynamic strides, s0 and s1, if either of one is dropped, then the result will have only one dynamic stride s0 irrespective of which one is dropped. So this is handled (I also have a test for this I think).

This is also not clear to me yet. When we have static strides, there is some logic to select a specific dim to be dropped. But when we have dynamic strides, we cannot look at the dynamic runtime stride value and just assume that they are all the same. So it looks like we don't really care which stride is dropped. Is this just to "make the result type verify"?

[mravishankar]

In D109418#4613354, @springerm wrote:

There is a bug in the implementation. It crashes with this test case:

// RUN: mlir-opt %s -canonicalize

func.func @func(%arg0: tensor<?x?x?xf32>, %arg1: index) -> index {
  %0 = bufferization.to_memref %arg0 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
  %c1 = arith.constant 1 : index
  %subview = memref.subview %0[0, 0, 0] [1, %arg1, 1] [1, 1, 1] : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to memref<1x?xf32, strided<[?, ?], offset: ?>>
  %dim = memref.dim %subview, %c1 : memref<1x?xf32, strided<[?, ?], offset: ?>>
  return %dim : index
}

Crashes here:

assert(subview.isDynamicSize(sourceIndex) &&
       "expected dynamic subview size");

It is not clear to me from the comments and the description why we consider strides at all. Couldn't we just make an arbitrary decision which dims are dropped (like tensor.extract_slice)? We could just drop the same strides. (This means that the valid result types of a memref.subview would change a bit.)

Nope. Unlike tensor.extract_slice, memref.subviews have strides. And the ordering of the strides need not match the dimensions. For example in column-major ordering the outermost strides are unit-strides....
The best fix for this is we explicitly take as arguments which dimensions are to be dropped. I think I mentioned this at the time I was working on it. Trying to "figure it out" is always going to be harder.

In D109418#3001836, @mravishankar wrote:

In D109418#3001440, @ftynse wrote:

Sorry for delay.

How does this change the situation for subviews with dynamic strides? Dynamic strides will all have the same static value so it looks like the code will just assume all dimensions are dropped because of that.

Ill explain by example. Lets say you have two unit dimensions with dynamic strides, s0 and s1, if either of one is dropped, then the result will have only one dynamic stride s0 irrespective of which one is dropped. So this is handled (I also have a test for this I think).

This is also not clear to me yet. When we have static strides, there is some logic to select a specific dim to be dropped. But when we have dynamic strides, we cannot look at the dynamic runtime stride value and just assume that they are all the same. So it looks like we don't really care which stride is dropped. Is this just to "make the result type verify"?

This part is tricky. I think for dynamic strides, you need to drop one dynamic stride value for one dimension dropped... maybe being fast and loose here. Ill look at the test case here.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let extraClassDeclaration = extraBaseClassDeclaration # [{
     std::array<unsigned, 3> getArrayAttrMaxRanks() {
       unsigned rank = getSourceType().getRank();
       return {rank, rank, rank};
     }
 
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
+
+    /// Return the dimensions of the source type that are dropped when
+    /// the result is rank-reduced.
+    llvm::SmallDenseSet<unsigned> getDroppedDims();
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TensorLoadOp

mlir/include/mlir/Interfaces/ViewLikeInterface.td

let description = [{
     This interface is useful to factor out common behavior and provide support
     for carrying or injecting static behavior through the use of the static
     attributes.
   }];
 
   let cppNamespace = "::mlir";
 
   let methods = [
-    InterfaceMethod<
+    StaticInterfaceMethod<
       /*desc=*/[{
         Return the number of leading operands before the `offsets`, `sizes` and
         and `strides` operands.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getOffsetSizeAndStrideStartOperandIndex",
       /*args=*/(ins)
     >,

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

matchAndRewrite(memref::SubViewOp subViewOp, ArrayRef<Value> operands,
     extracted = sourceMemRef.alignedPtr(rewriter, loc);
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc,
         LLVM::LLVMPointerType::get(targetElementTy,
                                    viewMemRefType.getMemorySpaceAsInt()),
         extracted);
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
-    auto shape = viewMemRefType.getShape();
-    auto inferredShape = inferredType.getShape();
-    size_t inferredShapeRank = inferredShape.size();
-    size_t resultShapeRank = shape.size();
-    llvm::SmallDenseSet<unsigned> unusedDims =
-        computeRankReductionMask(inferredShape, shape).getValue();
+    size_t inferredShapeRank = inferredType.getRank();
+    size_t resultShapeRank = viewMemRefType.getRank();
 
     // Extract strides needed to compute offset.
     SmallVector<Value, 4> strideValues;
     strideValues.reserve(inferredShapeRank);
     for (unsigned i = 0; i < inferredShapeRank; ++i)
       strideValues.push_back(sourceMemRef.stride(rewriter, loc, i));
 
     // Offset.
if (!ShapedType::isDynamicStrideOrOffset(offset)) {
       targetMemRef.setOffset(rewriter, loc, baseOffset);
     }
 
     // Update sizes and strides.
     SmallVector<OpFoldResult> mixedSizes = subViewOp.getMixedSizes();
     SmallVector<OpFoldResult> mixedStrides = subViewOp.getMixedStrides();
     assert(mixedSizes.size() == mixedStrides.size() &&
            "expected sizes and strides of equal length");
+    llvm::SmallDenseSet<unsigned> unusedDims = subViewOp.getDroppedDims();
     for (int i = inferredShapeRank - 1, j = resultShapeRank - 1;
          i >= 0 && j >= 0; --i) {
       if (unusedDims.contains(i))
         continue;
 
       // `i` may overflow subViewOp.getMixedSizes because of trailing semantics.
       // In this case, the size is guaranteed to be interpreted as Dim and the
       // stride as 1.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

static LogicalResult verify(DimOp op) {
   } else if (type.isa<UnrankedMemRefType>()) {
     // Assume index to be in range.
   } else {
     llvm_unreachable("expected operand with memref type");
   }
   return success();
 }
 
+static llvm::DenseMap<int64_t, unsigned>
+getNumOccurences(ArrayRef<int64_t> vals) {
+  llvm::DenseMap<int64_t, unsigned> numOccurences;
+  for (auto val : vals) {
+    numOccurences[val]++;
+  }
+  return numOccurences;
+}
+
+/// Given the type of the un-rank reduced subview result type and the
+/// rank-reduced result type, computes the dropped dimensions. This accounts for
+/// cases where there are multiple unit-dims, but only a subset of those are
+/// dropped. For MemRefTypes these can be disambiguated using the strides. If a
+/// dimension is dropped the stride must be dropped too.
+static llvm::Optional<llvm::SmallDenseSet<unsigned>>
+computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType,
+                               ArrayAttr staticSizes) {
+  llvm::SmallDenseSet<unsigned> unusedDims;
+  if (originalType.getRank() == reducedType.getRank())
+    return unusedDims;
+
+  for (auto dim : llvm::enumerate(staticSizes))
+    if (dim.value().cast<IntegerAttr>().getInt() == 1)
+      unusedDims.insert(dim.index());
+  SmallVector<int64_t> originalStrides, candidateStrides;
+  int64_t originalOffset, candidateOffset;
+  if (failed(
+          getStridesAndOffset(originalType, originalStrides, originalOffset)) ||
+      failed(
+          getStridesAndOffset(reducedType, candidateStrides, candidateOffset)))
+    return llvm::None;
+
+  // For memrefs, a dimenion is truly dropped if its corresponding stride is
+  // also dropped. This is particularly important when more than one of the dims
+  // is 1. Track the number of occurences of the strides in the original type
+  // and the candidate type. For each unused dim that stride should not be
+  // present in the candidate type. Note that there could be multiple dimensions
+  // that have the same rank. We dont need to exactly figure out which dim
+  // corresponds to which stride, we just need to verify that the
+  //
+  // number of reptitions of a stride in the original + number of unused dims
+  // with that stride == number of repititions of a stride in the candidate.
+  llvm::DenseMap<int64_t, unsigned> currUnaccountedStrides =
+      getNumOccurences(originalStrides);
+  llvm::DenseMap<int64_t, unsigned> candidateStridesNumOccurences =
+      getNumOccurences(candidateStrides);
+  llvm::SmallDenseSet<unsigned> prunedUnusedDims;
+  for (unsigned dim : unusedDims) {
+    int64_t originalStride = originalStrides[dim];
+    if (currUnaccountedStrides[originalStride] >
+        candidateStridesNumOccurences[originalStride]) {
+      // This is dim can be treated as dropped.
+      currUnaccountedStrides[originalStride]--;
+      continue;
+    }
+    if (currUnaccountedStrides[originalStride] ==
+        candidateStridesNumOccurences[originalStride]) {
+      // The stride for this is not dropped. Keep as is.
+      prunedUnusedDims.insert(dim);
+      continue;
+    }
+    if (currUnaccountedStrides[originalStride] <
+        candidateStridesNumOccurences[originalStride]) {
+      // This should never happen. Cant have a stride in the reduced rank type
+      // that wasnt in the original one.
+      return llvm::None;
+    }
+  }
+
+  for (auto prunedDim : prunedUnusedDims)
+    unusedDims.erase(prunedDim);
+  if (unusedDims.size() + reducedType.getRank() != originalType.getRank())
+    return llvm::None;
+  return unusedDims;
+}
+
+llvm::SmallDenseSet<unsigned> SubViewOp::getDroppedDims() {
+  MemRefType sourceType = getSourceType();
+  MemRefType resultType = getType();
+  llvm::Optional<llvm::SmallDenseSet<unsigned>> unusedDims =
+      computeMemRefRankReductionMask(sourceType, resultType, static_sizes());
+  assert(unusedDims && "unable to find unused dims of subview");
+  return *unusedDims;
+}
+
 OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
   // All forms of folding require a known index.
   auto index = operands[1].dyn_cast_or_null<IntegerAttr>();
   if (!index)
     return {};
 
   // Folding for unranked types (UnrankedMemRefType) is not supported.
   auto memrefType = source().getType().dyn_cast<MemRefType>();
OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
   if (auto alloca = dyn_cast_or_null<AllocaOp>(definingOp))
     return *(alloca.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto view = dyn_cast_or_null<ViewOp>(definingOp))
     return *(view.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
+  if (auto subview = dyn_cast_or_null<SubViewOp>(definingOp)) {
+    llvm::SmallDenseSet<unsigned> unusedDims = subview.getDroppedDims();
+    unsigned resultIndex = 0;
+    unsigned sourceRank = subview.getSourceType().getRank();
+    unsigned sourceIndex = 0;
+    for (auto i : llvm::seq<unsigned>(0, sourceRank)) {
+      if (unusedDims.count(i))
+        continue;
+      if (resultIndex == unsignedIndex) {
+        sourceIndex = i;
+        break;
+      }
+      resultIndex++;
+    }
+    assert(subview.isDynamicSize(sourceIndex) &&
+           "expected dynamic subview size");
+    return subview.getDynamicSize(sourceIndex);
+  }
+
   if (auto sizeInterface =
           dyn_cast_or_null<OffsetSizeAndStrideOpInterface>(definingOp)) {
     assert(sizeInterface.isDynamicSize(unsignedIndex) &&
            "Expected dynamic subview size");
     return sizeInterface.getDynamicSize(unsignedIndex);
   }
 
   // dim(memrefcast) -> dim
enum SubViewVerificationResult {
   AffineMapMismatch
 };
 
 /// Checks if `original` Type type can be rank reduced to `reduced` type.
 /// This function is slight variant of `is subsequence` algorithm where
 /// not matching dimension must be 1.
 static SubViewVerificationResult
 isRankReducedType(Type originalType, Type candidateReducedType,
-                  std::string *errMsg = nullptr) {
+                  ArrayAttr staticSizes, std::string *errMsg = nullptr) {
   if (originalType == candidateReducedType)
     return SubViewVerificationResult::Success;
   if (!originalType.isa<MemRefType>())
     return SubViewVerificationResult::Success;
   if (originalType.isa<MemRefType>() && !candidateReducedType.isa<MemRefType>())
     return SubViewVerificationResult::Success;
 
   ShapedType originalShapedType = originalType.cast<ShapedType>();
   ShapedType candidateReducedShapedType =
       candidateReducedType.cast<ShapedType>();
 
   // Rank and size logic is valid for all ShapedTypes.
   ArrayRef<int64_t> originalShape = originalShapedType.getShape();
   ArrayRef<int64_t> candidateReducedShape =
       candidateReducedShapedType.getShape();
   unsigned originalRank = originalShape.size(),
            candidateReducedRank = candidateReducedShape.size();
   if (candidateReducedRank > originalRank)
     return SubViewVerificationResult::RankTooLarge;
 
+  MemRefType original = originalType.cast<MemRefType>();
+  MemRefType candidateReduced = candidateReducedType.cast<MemRefType>();
+
   auto optionalUnusedDimsMask =
-      computeRankReductionMask(originalShape, candidateReducedShape);
+      computeMemRefRankReductionMask(original, candidateReduced, staticSizes);
 
   // Sizes cannot be matched in case empty vector is returned.
   if (!optionalUnusedDimsMask.hasValue())
     return SubViewVerificationResult::SizeMismatch;
 
   if (originalShapedType.getElementType() !=
       candidateReducedShapedType.getElementType())
     return SubViewVerificationResult::ElemTypeMismatch;
 
   // Strided layout logic is relevant for MemRefType only.
-  MemRefType original = originalType.cast<MemRefType>();
-  MemRefType candidateReduced = candidateReducedType.cast<MemRefType>();
   if (original.getMemorySpace() != candidateReduced.getMemorySpace())
     return SubViewVerificationResult::MemSpaceMismatch;
-
-  llvm::SmallDenseSet<unsigned> unusedDims = optionalUnusedDimsMask.getValue();
-  auto inferredType =
-      getProjectedMap(getStridedLinearLayoutMap(original), unusedDims);
-  AffineMap candidateLayout;
-  if (candidateReduced.getAffineMaps().empty())
-    candidateLayout = getStridedLinearLayoutMap(candidateReduced);
-  else
-    candidateLayout = candidateReduced.getAffineMaps().front();
-  assert(inferredType.getNumResults() == 1 &&
-         candidateLayout.getNumResults() == 1);
-  if (inferredType.getNumSymbols() != candidateLayout.getNumSymbols() ||
-      inferredType.getNumDims() != candidateLayout.getNumDims()) {
-    if (errMsg) {
-      llvm::raw_string_ostream os(*errMsg);
-      os << "inferred type: " << inferredType;
-    }
-    return SubViewVerificationResult::AffineMapMismatch;
-  }
-  // Check that the difference of the affine maps simplifies to 0.
-  AffineExpr diffExpr =
-      inferredType.getResult(0) - candidateLayout.getResult(0);
-  diffExpr = simplifyAffineExpr(diffExpr, inferredType.getNumDims(),
-                                inferredType.getNumSymbols());
-  auto cst = diffExpr.dyn_cast<AffineConstantExpr>();
-  if (!(cst && cst.getValue() == 0)) {
-    if (errMsg) {
-      llvm::raw_string_ostream os(*errMsg);
-      os << "inferred type: " << inferredType;
-    }
-    return SubViewVerificationResult::AffineMapMismatch;
-  }
   return SubViewVerificationResult::Success;
 }
 
 template <typename OpTy>
 static LogicalResult produceSubViewErrorMsg(SubViewVerificationResult result,
                                             OpTy op, Type expectedType,
                                             StringRef errMsg = "") {
   auto memrefType = expectedType.cast<ShapedType>();
static LogicalResult verify(SubViewOp op) {
 
   // Verify result type against inferred type.
   auto expectedType = SubViewOp::inferResultType(
       baseType, extractFromI64ArrayAttr(op.static_offsets()),
       extractFromI64ArrayAttr(op.static_sizes()),
       extractFromI64ArrayAttr(op.static_strides()));
 
   std::string errMsg;
-  auto result = isRankReducedType(expectedType, subViewType, &errMsg);
+  auto result =
+      isRankReducedType(expectedType, subViewType, op.static_sizes(), &errMsg);
   return produceSubViewErrorMsg(result, op, expectedType, errMsg);
 }
 
 raw_ostream &mlir::operator<<(raw_ostream &os, Range &range) {
   return os << "range " << range.offset << ":" << range.size << ":"
             << range.stride;
 }
 

mlir/lib/Dialect/MemRef/Transforms/FoldSubViewOps.cpp

resolveSourceIndices(Location loc, PatternRewriter &rewriter,
                      SmallVectorImpl<Value> &sourceIndices) {
   SmallVector<OpFoldResult> mixedOffsets = subViewOp.getMixedOffsets();
   SmallVector<OpFoldResult> mixedSizes = subViewOp.getMixedSizes();
   SmallVector<OpFoldResult> mixedStrides = subViewOp.getMixedStrides();
 
   SmallVector<Value> useIndices;
   // Check if this is rank-reducing case. Then for every unit-dim size add a
   // zero to the indices.
-  ArrayRef<int64_t> resultShape = subViewOp.getType().getShape();
   unsigned resultDim = 0;
-  for (auto size : llvm::enumerate(mixedSizes)) {
-    auto attr = size.value().dyn_cast<Attribute>();
-    // Check if this dimension has been dropped, i.e. the size is 1, but the
-    // associated dimension is not 1.
-    if (attr && attr.cast<IntegerAttr>().getInt() == 1 &&
-        (resultDim >= resultShape.size() || resultShape[resultDim] != 1))
+  llvm::SmallDenseSet<unsigned> unusedDims = subViewOp.getDroppedDims();
+  for (auto dim : llvm::seq<unsigned>(0, subViewOp.getSourceType().getRank())) {
+    if (unusedDims.count(dim))
       useIndices.push_back(rewriter.create<ConstantIndexOp>(loc, 0));
-    else if (resultDim < resultShape.size()) {
+    else
       useIndices.push_back(indices[resultDim++]);
-    }
-  }
   if (useIndices.size() != mixedOffsets.size())
     return failure();
   sourceIndices.resize(useIndices.size());
   for (auto index : llvm::seq<size_t>(0, mixedOffsets.size())) {
     SmallVector<Value> dynamicOperands;
     AffineExpr expr = rewriter.getAffineDimExpr(0);
     unsigned numSymbols = 0;
     dynamicOperands.push_back(useIndices[index]);

 static Value getMemRefOperand(vector::TransferReadOp op) { return op.source(); }
 
 static Value getMemRefOperand(memref::StoreOp op) { return op.memref(); }
 
 static Value getMemRefOperand(vector::TransferWriteOp op) {
   return op.source();
 }
 
+/// Given the permutation map of the original
+/// `vector.transfer_read`/`vector.transfer_write` operations compute the
+/// permutation map to use after the subview is folded with it.
+static AffineMap getPermutationMap(MLIRContext *context,
+                                   memref::SubViewOp subViewOp,
+                                   AffineMap currPermutationMap) {
+  llvm::SmallDenseSet<unsigned> unusedDims = subViewOp.getDroppedDims();
+  SmallVector<AffineExpr> exprs;
+  unsigned resultIdx = 0;
+  int64_t sourceRank = subViewOp.getSourceType().getRank();
+  for (auto dim : llvm::seq<int64_t>(0, sourceRank)) {
+    if (unusedDims.count(dim))
+      continue;
+    exprs.push_back(getAffineDimExpr(resultIdx++, context));
+  }
+  auto resultDimToSourceDimMap = AffineMap::get(sourceRank, 0, exprs, context);
+  return currPermutationMap.compose(resultDimToSourceDimMap);
+}
+
 //===----------------------------------------------------------------------===//
 // Patterns
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Merges subview operation with load/transferRead operation.
 template <typename OpTy>
 class LoadOpOfSubViewFolder final : public OpRewritePattern<OpTy> {

 }
 
 template <>
 void LoadOpOfSubViewFolder<vector::TransferReadOp>::replaceOp(
     vector::TransferReadOp loadOp, memref::SubViewOp subViewOp,
     ArrayRef<Value> sourceIndices, PatternRewriter &rewriter) const {
   rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
       loadOp, loadOp.getVectorType(), subViewOp.source(), sourceIndices,
-      loadOp.permutation_map(), loadOp.padding(), loadOp.in_boundsAttr());
+      getPermutationMap(rewriter.getContext(), subViewOp,
+                        loadOp.permutation_map()),
+      loadOp.padding(), loadOp.in_boundsAttr());
 }
 
 template <>
 void StoreOpOfSubViewFolder<memref::StoreOp>::replaceOp(
     memref::StoreOp storeOp, memref::SubViewOp subViewOp,
     ArrayRef<Value> sourceIndices, PatternRewriter &rewriter) const {
   rewriter.replaceOpWithNewOp<memref::StoreOp>(
       storeOp, storeOp.value(), subViewOp.source(), sourceIndices);
 }
 
 template <>
 void StoreOpOfSubViewFolder<vector::TransferWriteOp>::replaceOp(
     vector::TransferWriteOp transferWriteOp, memref::SubViewOp subViewOp,
     ArrayRef<Value> sourceIndices, PatternRewriter &rewriter) const {
   rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
       transferWriteOp, transferWriteOp.vector(), subViewOp.source(),
-      sourceIndices, transferWriteOp.permutation_map(),
+      sourceIndices,
+      getPermutationMap(rewriter.getContext(), subViewOp,
+                        transferWriteOp.permutation_map()),
       transferWriteOp.in_boundsAttr());
 }
 } // namespace
 
 template <typename OpTy>
 LogicalResult
 LoadOpOfSubViewFolder<OpTy>::matchAndRewrite(OpTy loadOp,
                                              PatternRewriter &rewriter) const {

mlir/test/Dialect/Linalg/loops.mlir

 //       CHECKPARALLEL:         %[[aff2:.*]] = affine.apply #[[$stride1Dilation1]](%[[arg4]], %[[arg7]])
 //       CHECKPARALLEL:         %[[aff3:.*]] = affine.apply #[[$stride1Dilation1]](%[[arg5]], %[[arg8]])
 //       CHECKPARALLEL:         %[[vb:.*]] = memref.load %[[arg0]][%[[aff]], %[[aff2]], %[[aff3]]] : memref<?x?x?xf32>
 //       CHECKPARALLEL:         %[[va:.*]] = memref.load %[[arg1]][%[[arg6]], %[[arg7]], %[[arg8]]] : memref<?x?x?xf32>
 //       CHECKPARALLEL:         %[[vc:.*]] = memref.load %[[arg2]][%[[arg3]], %[[arg4]], %[[arg5]]] : memref<?x?x?xf32>
 //       CHECKPARALLEL:         %[[inc:.*]] = mulf %[[vb]], %[[va]] : f32
 //       CHECKPARALLEL:         %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
 //       CHECKPARALLEL:         store %[[res]], %[[arg2]][%[[arg3]], %[[arg4]], %[[arg5]]] : memref<?x?x?xf32>
+
+// -----
+
+func @lower_to_loops_with_rank_reducing_subviews(
+    %arg0 : memref<?xi32>, %arg1 : memref<?x?xi32>, %arg2 : index,
+    %arg3 : index, %arg4 : index) {
+  %0 = memref.subview %arg0[%arg2] [%arg3] [1]
+      : memref<?xi32> to memref<?xi32, offset: ?, strides: [1]>
+  %1 = memref.subview %arg1[0, %arg4] [1, %arg3] [1, 1]
+      : memref<?x?xi32> to memref<?xi32, offset: ?, strides : [1]>
+  linalg.copy(%0, %1)
+      : memref<?xi32, offset: ?, strides: [1]>, memref<?xi32, offset: ?, strides: [1]>
+  return
+}
+// CHECK-LABEL: func @lower_to_loops_with_rank_reducing_subviews
+//       CHECK:   scf.for %[[IV:.+]] = %{{.+}} to %{{.+}} step %{{.+}} {
+//       CHECK:     %[[VAL:.+]] = memref.load %{{.+}}[%[[IV]]]
+//       CHECK:     memref.store %[[VAL]], %{{.+}}[%[[IV]]]
+//       CHECK:   }
+
+// CHECKPARALLEL-LABEL: func @lower_to_loops_with_rank_reducing_subviews
+//       CHECKPARALLEL:   scf.parallel (%[[IV:.+]]) = (%{{.+}}) to (%{{.+}}) step (%{{.+}}) {
+//       CHECKPARALLEL:     %[[VAL:.+]] = memref.load %{{.+}}[%[[IV]]]
+//       CHECKPARALLEL:     memref.store %[[VAL]], %{{.+}}[%[[IV]]]
+//       CHECKPARALLEL:   }

mlir/test/Dialect/MemRef/canonicalize.mlir

 //       CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1]
 //  CHECK-SAME:      [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1]
 //  CHECK-SAME:      : memref<?x?x?xf32> to memref<4x?xf32
 //       CHECK:   %[[RESULT:.+]] = memref.cast %[[SUBVIEW]]
 //       CHECK:   return %[[RESULT]]
 
 // -----
 
+func @multiple_reducing_dims(%arg0 : memref<1x384x384xf32>,
+    %arg1 : index, %arg2 : index, %arg3 : index) -> memref<?xf32, offset: ?, strides: [1]>
+{
+  %c1 = constant 1 : index
+  %0 = memref.subview %arg0[0, %arg1, %arg2] [1, %c1, %arg3] [1, 1, 1] : memref<1x384x384xf32> to memref<?x?xf32, offset: ?, strides: [384, 1]>
+  %1 = memref.subview %0[0, 0] [1, %arg3] [1, 1] : memref<?x?xf32, offset: ?, strides: [384, 1]> to memref<?xf32, offset: ?, strides: [1]>
+  return %1 : memref<?xf32, offset: ?, strides: [1]>
+}
+//   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
+//   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1)[s0] -> (d0 * 384 + s0 + d1)>
+//       CHECK: func @multiple_reducing_dims
+//       CHECK:   %[[REDUCED1:.+]] = memref.subview %{{.+}}[0, %{{.+}}, %{{.+}}] [1, 1, %{{.+}}] [1, 1, 1]
+//  CHECK-SAME:       : memref<1x384x384xf32> to memref<1x?xf32, #[[MAP1]]>
+//       CHECK:   %[[REDUCED2:.+]] = memref.subview %[[REDUCED1]][0, 0] [1, %{{.+}}] [1, 1]
+//  CHECK-SAME:       : memref<1x?xf32, #[[MAP1]]> to memref<?xf32, #[[MAP0]]>
+
+// -----
+
+func @multiple_reducing_dims_dynamic(%arg0 : memref<?x?x?xf32>,
+    %arg1 : index, %arg2 : index, %arg3 : index) -> memref<?xf32, offset: ?, strides: [1]>
+{
+  %c1 = constant 1 : index
+  %0 = memref.subview %arg0[0, %arg1, %arg2] [1, %c1, %arg3] [1, 1, 1] : memref<?x?x?xf32> to memref<?x?xf32, offset: ?, strides: [?, 1]>
+  %1 = memref.subview %0[0, 0] [1, %arg3] [1, 1] : memref<?x?xf32, offset: ?, strides: [?, 1]> to memref<?xf32, offset: ?, strides: [1]>
+  return %1 : memref<?xf32, offset: ?, strides: [1]>
+}
+//   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
+//   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
+//       CHECK: func @multiple_reducing_dims_dynamic
+//       CHECK:   %[[REDUCED1:.+]] = memref.subview %{{.+}}[0, %{{.+}}, %{{.+}}] [1, 1, %{{.+}}] [1, 1, 1]
+//  CHECK-SAME:       : memref<?x?x?xf32> to memref<1x?xf32, #[[MAP1]]>
+//       CHECK:   %[[REDUCED2:.+]] = memref.subview %[[REDUCED1]][0, 0] [1, %{{.+}}] [1, 1]
+//  CHECK-SAME:       : memref<1x?xf32, #[[MAP1]]> to memref<?xf32, #[[MAP0]]>
+
+// -----
+
+func @multiple_reducing_dims_all_dynamic(%arg0 : memref<?x?x?xf32, offset: ?, strides: [?, ?, ?]>,
+    %arg1 : index, %arg2 : index, %arg3 : index) -> memref<?xf32, offset: ?, strides: [?]>
+{
+  %c1 = constant 1 : index
+  %0 = memref.subview %arg0[0, %arg1, %arg2] [1, %c1, %arg3] [1, 1, 1]
+      : memref<?x?x?xf32, offset: ?, strides: [?, ?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+  %1 = memref.subview %0[0, 0] [1, %arg3] [1, 1] : memref<?x?xf32, offset: ?, strides: [?, ?]> to memref<?xf32, offset: ?, strides: [?]>
+  return %1 : memref<?xf32, offset: ?, strides: [?]>
+}
+//   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+//   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+//   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
+//       CHECK: func @multiple_reducing_dims_all_dynamic
+//       CHECK:   %[[REDUCED1:.+]] = memref.subview %{{.+}}[0, %{{.+}}, %{{.+}}] [1, 1, %{{.+}}] [1, 1, 1]
+//  CHECK-SAME:       : memref<?x?x?xf32, #[[MAP2]]> to memref<1x?xf32, #[[MAP1]]>
+//       CHECK:   %[[REDUCED2:.+]] = memref.subview %[[REDUCED1]][0, 0] [1, %{{.+}}] [1, 1]
+//  CHECK-SAME:       : memref<1x?xf32, #[[MAP1]]> to memref<?xf32, #[[MAP0]]>
+
+
+// -----
+
 // CHECK-LABEL: @clone_before_dealloc
 // CHECK-SAME: %[[ARG:.*]]: memref<?xf32>
 func @clone_before_dealloc(%arg0: memref<?xf32>) -> memref<?xf32> {
   // CHECK-NEXT: return %[[ARG]]
   %0 = memref.clone %arg0 : memref<?xf32> to memref<?xf32>
   memref.dealloc %arg0 : memref<?xf32>
   return %0 : memref<?xf32>
 }

 // CHECK:           %[[COLLAPSED:.*]] = memref.collapse_shape %[[INPUT]]
 // CHECK_SAME:        {{\[\[}}0], [1, 2, 3]] : memref<?x512x1x?xf32> into memref<?x?xf32>
 // CHECK:           return %[[COLLAPSED]] : memref<?x?xf32>
 func @collapse_after_memref_cast(%arg0 : memref<?x512x1x?xf32>) -> memref<?x?xf32> {
   %dynamic = memref.cast %arg0: memref<?x512x1x?xf32> to memref<?x?x?x?xf32>
   %collapsed = memref.collapse_shape %dynamic [[0], [1, 2, 3]] : memref<?x?x?x?xf32> into memref<?x?xf32>
   return %collapsed : memref<?x?xf32>
 }
-

mlir/test/Dialect/MemRef/fold-subview-ops.mlir

 //  CHECK-DAG:   %[[C0:.+]] = constant 0 : index
 //  CHECK-DAG:   %[[I1:.+]] = affine.apply #[[MAP]](%[[ARG13]])[%[[ARG7]], %[[ARG1]]]
 //  CHECK-DAG:   %[[I2:.+]] = affine.apply #[[MAP]](%[[ARG14]])[%[[ARG8]], %[[ARG2]]]
 //  CHECK-DAG:   %[[I3:.+]] = affine.apply #[[MAP]](%[[C0]])[%[[ARG9]], %[[ARG3]]]
 //  CHECK-DAG:   %[[I4:.+]] = affine.apply #[[MAP]](%[[ARG15]])[%[[ARG10]], %[[ARG4]]]
 //  CHECK-DAG:   %[[I5:.+]] = affine.apply #[[MAP]](%[[ARG16]])[%[[ARG11]], %[[ARG5]]]
 //  CHECK-DAG:   %[[I6:.+]] = affine.apply #[[MAP]](%[[C0]])[%[[ARG12]], %[[ARG6]]]
 //      CHECK:   memref.load %[[ARG0]][%[[I1]], %[[I2]], %[[I3]], %[[I4]], %[[I5]], %[[I6]]]
+
+// -----
+
+func @fold_vector_transfer_read_with_rank_reduced_subview(
+    %arg0 : memref<?x?x?xf32, offset: ?, strides: [?, ?, ?]>,
+    %arg1: index, %arg2 : index, %arg3 : index, %arg4: index, %arg5 : index,
+    %arg6 : index) -> vector<4xf32> {
+  %cst = constant 0.0 : f32
+  %0 = memref.subview %arg0[0, %arg1, %arg2] [1, %arg3, %arg4] [1, 1, 1]
+      : memref<?x?x?xf32, offset: ?, strides: [?, ?, ?]> to
+        memref<?x?xf32, offset: ?, strides: [?, ?]>
+  %1 = vector.transfer_read %0[%arg5, %arg6], %cst {in_bounds = [true]}
+      : memref<?x?xf32, offset: ?, strides: [?, ?]>, vector<4xf32>
+  return %1 : vector<4xf32>
+}
+//   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
+//   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
+//   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1)>
+//       CHECK: func @fold_vector_transfer_read_with_rank_reduced_subview
+//  CHECK-SAME:    %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?xf32, #[[MAP0]]>
+//  CHECK-SAME:    %[[ARG1:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG2:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG3:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG4:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG5:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG6:[a-zA-Z0-9]+]]: index
+//   CHECK-DAG:    %[[C0:.+]] = constant 0 : index
+//   CHECK-DAG:    %[[IDX0:.+]] = affine.apply #[[MAP1]](%[[ARG5]])[%[[ARG1]]]
+//   CHECK-DAG:    %[[IDX1:.+]] = affine.apply #[[MAP1]](%[[ARG6]])[%[[ARG2]]]
+//       CHECK:    vector.transfer_read %[[ARG0]][%[[C0]], %[[IDX0]], %[[IDX1]]]
+//  CHECK-SAME:        permutation_map = #[[MAP2]]
+
+// -----
+
+func @fold_vector_transfer_write_with_rank_reduced_subview(
+    %arg0 : memref<?x?x?xf32, offset: ?, strides: [?, ?, ?]>,
+    %arg1 : vector<4xf32>, %arg2: index, %arg3 : index, %arg4 : index,
+    %arg5: index, %arg6 : index, %arg7 : index) {
+  %cst = constant 0.0 : f32
+  %0 = memref.subview %arg0[0, %arg2, %arg3] [1, %arg4, %arg5] [1, 1, 1]
+      : memref<?x?x?xf32, offset: ?, strides: [?, ?, ?]> to
+        memref<?x?xf32, offset: ?, strides: [?, ?]>
+  vector.transfer_write %arg1, %0[%arg6, %arg7] {in_bounds = [true]}
+      : vector<4xf32>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+  return
+}
+//   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
+//   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
+//   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1)>
+//       CHECK: func @fold_vector_transfer_write_with_rank_reduced_subview
+//  CHECK-SAME:    %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?xf32, #[[MAP0]]>
+//  CHECK-SAME:    %[[ARG1:[a-zA-Z0-9]+]]: vector<4xf32>
+//  CHECK-SAME:    %[[ARG2:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG3:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG4:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG5:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG6:[a-zA-Z0-9]+]]: index
+//  CHECK-SAME:    %[[ARG7:[a-zA-Z0-9]+]]: index
+//   CHECK-DAG:    %[[C0:.+]] = constant 0 : index
+//   CHECK-DAG:    %[[IDX0:.+]] = affine.apply #[[MAP1]](%[[ARG6]])[%[[ARG2]]]
+//   CHECK-DAG:    %[[IDX1:.+]] = affine.apply #[[MAP1]](%[[ARG7]])[%[[ARG3]]]
+//   CHECK-DAG:    vector.transfer_write %[[ARG1]], %[[ARG0]][%[[C0]], %[[IDX0]], %[[IDX1]]]
+//  CHECK-SAME:        permutation_map = #[[MAP2]]

mlir/test/Dialect/MemRef/invalid.mlir

 
 func @collapse_shape_illegal_mixed_memref_2(%arg0 : memref<?x4x5xf32>)
     -> memref<?x?xf32> {
   // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
   %0 = memref.collapse_shape %arg0 [[0], [1, 2]]
       : memref<?x4x5xf32> into memref<?x?xf32>
   return %0 : memref<?x?xf32>
 }
+
+// -----
+
+func @static_stride_to_dynamic_stride(%arg0 : memref<?x?x?xf32>, %arg1 : index,
+    %arg2 : index) -> memref<?x?xf32, offset:?, strides: [?, ?]> {
+  // expected-error @+1 {{expected result type to be 'memref<1x?x?xf32, affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>>' or a rank-reduced version. (mismatch of result sizes)}}
+  %0 = memref.subview %arg0[0, 0, 0] [1, %arg1, %arg2] [1, 1, 1] : memref<?x?x?xf32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+  return %0 : memref<?x?xf32, offset: ?, strides: [?, ?]>
+}
 No newline at end of file

mlir/test/IR/invalid-ops.mlir

: memref<8x16x4xf32> to
       memref<8x?x4xf32, offset: 0, strides:[?, ?, 4]>
   return
 }
 
 // -----
 
 func @invalid_subview(%arg0 : index, %arg1 : index, %arg2 : index) {
   %0 = memref.alloc() : memref<8x16x4xf32>
-  // expected-error@+1 {{expected result type to be 'memref<?x?x?xf32, affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>>' or a rank-reduced version. (mismatch of result affine map)}}
-  %1 = memref.subview %0[%arg0, %arg1, %arg2][%arg0, %arg1, %arg2][%arg0, %arg1, %arg2]
-    : memref<8x16x4xf32> to
-      memref<?x?x?xf32, offset: ?, strides: [64, 4, 1]>
-  return
-}
-
-// -----
-
-func @invalid_subview(%arg0 : index, %arg1 : index, %arg2 : index) {
-  %0 = memref.alloc() : memref<8x16x4xf32>
   // expected-error@+1 {{expected result element type to be 'f32'}}
   %1 = memref.subview %0[0, 0, 0][8, 16, 4][1, 1, 1]
     : memref<8x16x4xf32> to
       memref<8x16x4xi32>
   return
 }
 
 // -----
func @invalid_rank_reducing_subview(%arg0 : index, %arg1 : index, %arg2 : index) {
   %1 = memref.subview %0[0, 2, 0][8, 16, 4][1, 1, 1]
     : memref<8x16x4xf32> to memref<16x4xf32>
   return
 }
 
 // -----
 
 func @invalid_rank_reducing_subview(%arg0 : memref<?x?xf32>, %arg1 : index, %arg2 : index) {
-  // expected-error@+1 {{expected result type to be 'memref<?x1xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>>' or a rank-reduced version. (mismatch of result affine map)}}
+  // expected-error@+1 {{expected result type to be 'memref<?x1xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>>' or a rank-reduced version. (mismatch of result sizes)}}
   %0 = memref.subview %arg0[0, %arg1][%arg2, 1][1, 1] : memref<?x?xf32> to memref<?xf32>
   return
 }
 
 // -----
 
-// The affine map affine_map<(d0)[s0, s1, s2] -> (d0 * s1 + s0)> has an extra unused symbol.
-func @invalid_rank_reducing_subview(%arg0 : memref<?x?xf32>, %arg1 : index, %arg2 : index) {
-  // expected-error@+1 {{expected result type to be 'memref<?x1xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>>' or a rank-reduced version. (mismatch of result affine map) inferred type: (d0)[s0, s1] -> (d0 * s1 + s0)}}
-  %0 = memref.subview %arg0[0, %arg1][%arg2, 1][1, 1] : memref<?x?xf32> to memref<?xf32, affine_map<(d0)[s0, s1, s2] -> (d0 * s1 + s0)>>
-  return
-}
-
-// -----
-
 func @invalid_memref_cast(%arg0 : memref<12x4x16xf32, offset:0, strides:[64, 16, 1]>) {
   // expected-error@+1{{operand type 'memref<12x4x16xf32, affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 16 + d2)>>' and result type 'memref<12x4x16xf32, affine_map<(d0, d1, d2) -> (d0 * 128 + d1 * 32 + d2 * 2)>>' are cast incompatible}}
   %0 = memref.cast %arg0 : memref<12x4x16xf32, offset:0, strides:[64, 16, 1]> to memref<12x4x16xf32, offset:0, strides:[128, 32, 2]>
   return
 }
 
 // -----