[mlir] Add SubViewOp::getOrCreateRanges and fix folding pattern

[nicolasvasilache]

Summary

The existing implementation of SubViewOp::getRanges relies on all
offsets/sizes/strides to be dynamic values and does not work in
combination with canonicalization. This revision adds a
SubViewOp::getOrCreateRanges to create the missing constants in the
canonicalized case.

This allows reactivating the fused pass with staged pattern
applications.

However another issue surfaces that the SubViewOp verifier is now too
strict to allow folding. The existing folding pattern is turned into a
canonicalization pattern which rewrites memref_cast + subview into
subview + memref_cast.

The transform-patterns-matmul-to-vector can then be reactivated.

[nicolasvasilache]

Drop spurious methods in Ops.td.

[nicolasvasilache]

Rebase.

[mravishankar]

Not familiar with the folding patterns for memref cast as much, but it seems fine to me. Please let someone with more knowledge of this to review it as well.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2574

Nit: this name is misleading. How about SubViewOpConstantArgumentFolder ?

[ftynse]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2493

Please add documentation.

2503

I don't understand why you do std::prev, and there is no comment to help. Normally, all standard algorithms exclude the upper bound, so this would exclude the element _before_ idx-th.

2709

Could you please explain _how_ is it rewritten?

2714

This needs a test.

[nicolasvasilache]

Rebase.

[nicolasvasilache]

Add doc.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let extraClassDeclaration = [{
     LogicalResult getStaticStrides(SmallVectorImpl<int64_t> &staticStrides);
 
     /// Auxiliary range data structure and helper function that unpacks the
     /// offset, size and stride operands of the SubViewOp into a list of triples.
     /// Such a list of triple is sometimes more convenient to manipulate.
     struct Range {
       Value offset, size, stride;
     };
-    SmallVector<Range, 8> getRanges();
+    SmallVector<Range, 8> getOrCreateRanges(OpBuilder &b, Location loc);
+
+    /// A subview result type can be fully inferred from the source type and the
+    /// static representation of offsets, sizes and strides. Special sentinels
+    /// encode the dynamic case.
+    static Type inferSubViewResultType(MemRefType sourceMemRefType,
+                                       ArrayRef<int64_t> staticOffsets,
+                                       ArrayRef<int64_t> staticSizes,
+                                       ArrayRef<int64_t> staticStrides);
 
     /// Return the rank of the result MemRefType.
     unsigned getRank() { return getType().getRank(); }
 
     static StringRef getStaticOffsetsAttrName() {
       return "static_offsets";
     }
     static StringRef getStaticSizesAttrName() {
       return "static_sizes";
     }
     static StringRef getStaticStridesAttrName() {
       return "static_strides";
     }
     static ArrayRef<StringRef> getSpecialAttrNames() {
       static SmallVector<StringRef, 4> names{
         getStaticOffsetsAttrName(),
         getStaticSizesAttrName(),
         getStaticStridesAttrName(),
         getOperandSegmentSizeAttr()};
       return names;
    }
   }];
 
   let hasCanonicalizer = 1;
-  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TanhOp
 //===----------------------------------------------------------------------===//
 
 def TanhOp : FloatUnaryOp<"tanh"> {
   let summary = "hyperbolic tangent of the specified value";

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

static LinalgOp fuse(Value producedView, LinalgOp producer, LinalgOp consumer,
   LLVM_DEBUG(dbgs() << "Producer Idx: " << producerIdx
                     << ", producer map: " << producerMap << "\n");
 
   unsigned nPar = producer.getNumParallelLoops();
   unsigned nRed = producer.getNumReductionLoops();
   unsigned nWin = producer.getNumWindowLoops();
   SmallVector<SubViewOp::Range, 8> loopRanges(nPar + nRed + nWin);
 
+  OpBuilder b(consumer.getOperation());
+  auto loc = consumer.getLoc();
   // Iterate over dimensions identified by the producer map for `producerIdx`.
   // This defines a subset of the loop ranges that we need to complete later.
   for (auto en : llvm::enumerate(producerMap.getResults())) {
     unsigned posInProducerLoop = en.value().cast<AffineDimExpr>().getPosition();
-    loopRanges[posInProducerLoop] = subView.getRanges()[en.index()];
+    loopRanges[posInProducerLoop] =
+        subView.getOrCreateRanges(b, loc)[en.index()];
   }
 
-  OpBuilder b(consumer.getOperation());
-  auto loc = consumer.getLoc();
   // Iterate over all dimensions. For the dimensions not identified by the
   // producer map for `producerIdx`, we need to explicitly compute the view that
   // defines the loop ranges using the `producer`.
   for (unsigned i = 0, nLoops = loopRanges.size(); i < nLoops; ++i) {
     if (loopRanges[i].offset)
       LLVM_DEBUG(llvm::dbgs()
                  << "existing LoopRange: " << loopRanges[i] << "\n");
     else {

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

static PromotionInfo promoteSubviewAsNewBuffer(OpBuilder &b, Location loc,
   auto one = folded_std_constant_index(folder, 1);
 
   auto viewType = subView.getType();
   auto rank = viewType.getRank();
   Value allocSize = one;
   SmallVector<Value, 8> fullSizes, partialSizes;
   fullSizes.reserve(rank);
   partialSizes.reserve(rank);
-  for (auto en : llvm::enumerate(subView.getRanges())) {
+  for (auto en : llvm::enumerate(subView.getOrCreateRanges(b, loc))) {
     auto rank = en.index();
     auto rangeValue = en.value();
     // Try to extract a tight constant.
     LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n");
     Value size = extractSmallestConstantBoundingSize(b, loc, rangeValue.size);
     LLVM_DEBUG(llvm::dbgs() << "Extracted tightest: " << size << "\n");
     allocSize = folded_std_muli(folder, allocSize, size);
     fullSizes.push_back(size);
     partialSizes.push_back(folded_std_dim(folder, subView, rank));
   }
   SmallVector<int64_t, 4> dynSizes(fullSizes.size(), -1);
   auto buffer = allocBuffer(viewType.getElementType(), allocSize,
                             dynamicBuffers, folder, alignment);
   auto fullLocalView = folded_std_view(
       folder, MemRefType::get(dynSizes, viewType.getElementType()), buffer,
-      folded_std_constant_index(folder, 0), fullSizes);
+      zero, fullSizes);
   SmallVector<Value, 4> zeros(fullSizes.size(), zero);
   SmallVector<Value, 4> ones(fullSizes.size(), one);
   auto partialLocalView =
       folded_std_subview(folder, fullLocalView, zeros, partialSizes, ones);
   return PromotionInfo{buffer, fullLocalView, partialLocalView};
 }
 
 static SmallVector<PromotionInfo, 8>

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

Wrapper operator*(Wrapper a, int64_t b) {
   return Wrapper(a.v * b);
 }
 } // end namespace saturated_arith
 } // end namespace
 
 /// A subview result type can be fully inferred from the source type and the
 /// static representation of offsets, sizes and strides. Special sentinels
 /// encode the dynamic case.
-static Type inferSubViewResultType(MemRefType sourceMemRefType,
+Type SubViewOp::inferSubViewResultType(MemRefType sourceMemRefType,
                                        ArrayRef<int64_t> staticOffsets,
                                        ArrayRef<int64_t> staticSizes,
                                        ArrayRef<int64_t> staticStrides) {
   unsigned rank = sourceMemRefType.getRank();
   (void)rank;
   assert(staticOffsets.size() == rank &&
          "unexpected staticOffsets size mismatch");
   assert(staticSizes.size() == rank && "unexpected staticSizes size mismatch");
   assert(staticStrides.size() == rank &&
          "unexpected staticStrides size mismatch");
 

 
 void mlir::SubViewOp::build(OpBuilder &b, OperationState &result, Value source,
                             ArrayRef<int64_t> staticOffsets,
                             ArrayRef<int64_t> staticSizes,
                             ArrayRef<int64_t> staticStrides, ValueRange offsets,
                             ValueRange sizes, ValueRange strides,
                             ArrayRef<NamedAttribute> attrs) {
   auto sourceMemRefType = source.getType().cast<MemRefType>();
-  auto resultType = inferSubViewResultType(sourceMemRefType, staticOffsets,
-                                           staticSizes, staticStrides);
+  auto resultType = SubViewOp::inferSubViewResultType(
+      sourceMemRefType, staticOffsets, staticSizes, staticStrides);
   build(b, result, resultType, source, offsets, sizes, strides,
         b.getI64ArrayAttr(staticOffsets), b.getI64ArrayAttr(staticSizes),
         b.getI64ArrayAttr(staticStrides));
   result.addAttributes(attrs);
 }
 
 /// Build a SubViewOp with all dynamic entries: `staticOffsets`, `staticSizes`
 /// and `staticStrides` are  automatically filled with source-memref-rank
if (failed(verifySubViewOpPart(op, "size", op.getStaticSizesAttrName(),
                                  op.sizes())))
     return failure();
   if (failed(verifySubViewOpPart(
           op, "stride", op.getStaticStridesAttrName(), op.static_strides(),
           ShapedType::isDynamicStrideOrOffset, op.strides())))
     return failure();
 
   // Verify result type against inferred type.
-  auto expectedType = inferSubViewResultType(
+  auto expectedType = SubViewOp::inferSubViewResultType(
       op.getBaseMemRefType(), extractFromI64ArrayAttr(op.static_offsets()),
       extractFromI64ArrayAttr(op.static_sizes()),
       extractFromI64ArrayAttr(op.static_strides()));
   if (op.getType() != expectedType)
     return op.emitError("expected result type to be ") << expectedType;
 
   return success();
 }
 
 raw_ostream &mlir::operator<<(raw_ostream &os, SubViewOp::Range &range) {
   return os << "range " << range.offset << ":" << range.size << ":"
             << range.stride;
 }
 
-SmallVector<SubViewOp::Range, 8> SubViewOp::getRanges() {
+static Value getOrCreateSubViewPart(OpBuilder &b, Location loc,
+                                    ValueRange values, ArrayAttr attr,
+                                    llvm::function_ref<bool(int64_t)> isDynamic,
+                                    unsigned idx) {
+  auto cst = attr.getValue()[idx].cast<IntegerAttr>().getInt();
+  if (!isDynamic(cst))
+    return b.create<ConstantIndexOp>(loc, cst);
+  if (idx == 0)
+    return values[0];
+  unsigned numDynamicEntriesUpToIdx = std::count_if(
+      attr.getValue().begin(), std::prev(attr.getValue().begin() + idx),
+      [&](Attribute attr) {
+        return isDynamic(attr.cast<IntegerAttr>().getInt());
+      });
+  return values[idx - numDynamicEntriesUpToIdx];
+}
+
+SmallVector<SubViewOp::Range, 8> SubViewOp::getOrCreateRanges(OpBuilder &b,
+                                                              Location loc) {
   SmallVector<Range, 8> res;
   unsigned rank = getType().getRank();
   res.reserve(rank);
-  for (unsigned i = 0; i < rank; ++i)
-    res.emplace_back(Range{*(offsets().begin() + i), *(sizes().begin() + i),
-                           *(strides().begin() + i)});
+  for (unsigned idx = 0; idx < rank; ++idx) {
+    res.emplace_back(Range{
+        getOrCreateSubViewPart(b, loc, offsets(),
+                               static_offsets().cast<ArrayAttr>(),
+                               ShapedType::isDynamicStrideOrOffset, idx),
+        getOrCreateSubViewPart(b, loc, sizes(),
+                               static_sizes().cast<ArrayAttr>(),
+                               ShapedType::isDynamic, idx),
+        getOrCreateSubViewPart(b, loc, strides(),
+                               static_strides().cast<ArrayAttr>(),
+                               ShapedType::isDynamicStrideOrOffset, idx),
+    });
+  }
   return res;
 }
 
 LogicalResult
 SubViewOp::getStaticStrides(SmallVectorImpl<int64_t> &staticStrides) {
   if (!strides().empty())
     return failure();
   staticStrides = extractFromI64ArrayAttr(static_strides());
for (unsigned cstIdx = 0, valIdx = 0, e = constantValues.size();
       }
       // Remains dynamic move to next value.
       ++valIdx;
     }
   }
 }
 
 /// Pattern to rewrite a subview op with constant arguments.
-class SubViewOpFolder final : public OpRewritePattern<SubViewOp> {
+class SubViewOpConstantFolder final : public OpRewritePattern<SubViewOp> {
 public:
   using OpRewritePattern<SubViewOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(SubViewOp subViewOp,
                                 PatternRewriter &rewriter) const override {
     // No constant operand, just return;
     if (llvm::none_of(subViewOp.getOperands(), [](Value operand) {
           return matchPattern(operand, m_ConstantIndex());
if (ss != st)
       if (MemRefType::isDynamicStrideOrOffset(ss) &&
           !MemRefType::isDynamicStrideOrOffset(st))
         return false;
   }
 
   return true;
 }
 
-OpFoldResult SubViewOp::fold(ArrayRef<Attribute>) {
-  auto folds = [](Operation *op) {
-    bool folded = false;
-    for (OpOperand &operand : op->getOpOperands()) {
-      auto castOp = operand.get().getDefiningOp<MemRefCastOp>();
-      if (castOp && canFoldIntoConsumerOp(castOp)) {
-        operand.set(castOp.getOperand());
-        folded = true;
-      }
+/// Pattern to rewrite a subview op with MemRefCast arguments.
+class SubViewOpMemRefCastFolder final : public OpRewritePattern<SubViewOp> {
+public:
+  using OpRewritePattern<SubViewOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SubViewOp subViewOp,
+                                PatternRewriter &rewriter) const override {
+    // Any constant operand, just return to let SubViewOpConstantFolder kick in.
+    if (llvm::any_of(subViewOp.getOperands(), [](Value operand) {
+          return matchPattern(operand, m_ConstantIndex());
+        }))
+      return failure();
+
+    auto castOp =
+        dyn_cast_or_null<MemRefCastOp>(subViewOp.source().getDefiningOp());
+    if (!castOp)
+      return failure();
+
+    Type resultType = SubViewOp::inferSubViewResultType(
+        castOp.source().getType().cast<MemRefType>(),
+        extractFromI64ArrayAttr(subViewOp.static_offsets()),
+        extractFromI64ArrayAttr(subViewOp.static_sizes()),
+        extractFromI64ArrayAttr(subViewOp.static_strides()));
+
+    if (resultType == subViewOp.getType()) {
+      rewriter.replaceOpWithNewOp<SubViewOp>(
+          subViewOp, resultType, castOp.source(), subViewOp.offsets(),
+          subViewOp.sizes(), subViewOp.strides(), subViewOp.static_offsets(),
+          subViewOp.static_sizes(), subViewOp.static_strides());
+      return success();
     }
-    return folded ? success() : failure();
-  };
 
-  if (succeeded(folds(*this)))
-    return getResult();
-  return {};
+    if (!canFoldIntoConsumerOp(castOp))
+      return failure();
+
+    Value newSubView = rewriter.create<SubViewOp>(
+        subViewOp.getLoc(), resultType, castOp.source(), subViewOp.offsets(),
+        subViewOp.sizes(), subViewOp.strides(), subViewOp.static_offsets(),
+        subViewOp.static_sizes(), subViewOp.static_strides());
+    rewriter.replaceOpWithNewOp<MemRefCastOp>(subViewOp, subViewOp.getType(),
+                                              newSubView);
+    return success();
   }
+};
 
 void SubViewOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                             MLIRContext *context) {
-  results.insert<SubViewOpFolder>(context);
+  results.insert<SubViewOpConstantFolder, SubViewOpMemRefCastFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // TensorCastOp
 //===----------------------------------------------------------------------===//
 
 bool TensorCastOp::areCastCompatible(Type a, Type b) {
   auto aT = a.dyn_cast<TensorType>();

mlir/test/Dialect/Linalg/transform-patterns-matmul-to-vector.mlir

-// TODO: this needs a fix to land before being reactivated.
-// RUN: ls
-// R_UN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-1d | FileCheck %s
-// R_UN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-2d | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-1d | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-2d | FileCheck %s
 
 func @matmul(%A: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                   %B: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                   %C: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>) {
   linalg.matmul(%A, %B, %C) {__internal_linalg_transform__ = "START"} :
     memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
     memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
     memref<1584x1584xf32, offset: 0, strides: [1584, 1]>