[mlir][Linalg] NFC - Refactor and simplify Promotion

[nicolasvasilache]

Summary

This revision introduces LinalgPromotionOptions to more easily control the application of promotion patterns. It also simplifies the different entry points into Promotion in preparation for some behavior change in subsequent revisions.

[nicolasvasilache]

Initial update
Update.

[mravishankar]

Some initial comments.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

101

Should this be a a per operand specification? Suggestion is to have a separate struct that captures per operand specification for this, alignment (and memory space as well. See below).

107

Could we add an option to set memory space here as well? I was touching the same file in a WIP change of mine to add memory space. But you are touching these anyway

267

This seems to diverge from the way patterns are exposed in other places. The convention is AFAIK to use "populate.....Pattern" which takes a OwningRewritePatternList as an argument (by reference) and inserts the pattern. With that you dont need to have the pattern itself exposed just an API to insert the pattern. You could then add as many patterns as you want into that.

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

55

Same here. Maybe allow specification of dynamicBuffers, alignment and memory space per subview that is being promoted?

103

Thanks for the comment. Makes it clear what dynamicBuffers are for.

107

Please add extra option to set memory space as well.

288–292

Please add braces around multi-line statements (disclaimer : I dont know if thats the coding standard for LLVM, but dropping trivial braces for single line body seems fine, but this is not a single line body)

318

Why not just use the pattern above here instead of doing a walk explicitly?

[mravishankar]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

267

Actually I see now how this is supposed to work. Disregard.

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

213

I think these fill needs to have the marker added to it too. If not they can be intercepted by the pattern rewriter which would try to apply transformations on them (AFAIK the marker is supposed to prevent that).
Same is true for the linalg_copy below.

[ftynse]

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

49

Could we have some documentation?

262

If you want to call this function from a pattern, you are not allowed to do this (unless it's the root operation of the pattern and you surround it with a proper transaction start/stop)

[nicolasvasilache]

Address reviews.

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

101

yes but in a future revision please, this wants to remain NFC.

107

yes but in a future revision please, this wants to remain NFC.
Feel free to rebase a memory space change on top of this.
I expect a few followup revisions will add a few more per operand options.

267

Ack

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

55

yes but in a future revision please, this wants to remain NFC.

107

yes but in a future revision please, this wants to remain NFC.

213

I am envisioning is an analysis would place markers that will trigger patterns.
Anything unmarked would not be rewritten in the first place.

I think we can experiment with different scenarios in the future.
In any case, changing this now would make the revision non-NFC.

262

This is indeed called from under

rewriter.updateRootInPlace(op, [&]() {

and the prereq checks that it will succeed.

I can revert to the old behavior if you prefer but it seems like a legitimate use to me?

318

Because it is unfortunately not strictly equivalent ..
The folder that is passed here does additional cleanups on the fly that cannot be applied as a pattern.
This is one of the reasons that motivates rewriting transforms as staged pattern applications (a few revisions down the line).

Once those are in we can drop the folder and use the patterns.

[mravishankar]

Could you let me know when you are done with most refactoring so that I can pick up some other things that intersected with this refactoring work. I dont want to be touching the same code will its being modified.

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

213

Ok, lets sync about this offline. Maybe the way I was thinking about this is reverse from what you mentioned. By reading the code, I was thinking that you do a transformation and set a marker, preventing it from being transformed again. But what is being done is that it is setting the maker for the next stage of transformation to kick in. I need to account for this in IREE codegen. Will take a look after all the refactoring is done.

[rriddle]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

96

operandsToPromote = DenseSet<unsigned>(operands.begin(), operands.end())?

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

75

nit: Can we early exit here?

[mehdi_amini]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

96

@nicolasvasilache ping?

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 ///      float zero for now).
 ///   3. Take a partial slice of the full view in step 2. and copy into it.
 /// Infers statically sized buffers from subViews unless `dynamicBuffers` is
 /// true.
 ///
 /// Returns a list of PromotionInfo which hold the promoted buffer and the
 /// full and partial views indexing into the buffer.
 // TODO: revisit dynamicBuffers option.
-LinalgOp promoteSubViewOperands(OpBuilder &b, LinalgOp op,
-                                llvm::SetVector<Value> subViews,
-                                bool dynamicBuffers = false,
-                                int64_t alignment = 0,
-                                OperationFolder *folder = nullptr);
+struct LinalgPromotionOptions {
+  /// Indices of subViews to promote. If `None`, try to promote all operands.
+  Optional<DenseSet<unsigned>> operandsToPromote = None;
+  LinalgPromotionOptions &setOperandsToPromote(ArrayRef<int64_t> operands) {
+    operandsToPromote = DenseSet<unsigned>();
+    operandsToPromote->insert(operands.begin(), operands.end());
+    return *this;
+  }
+  /// Allow the use of dynamicaly-sized buffers.
+  bool dynamicBuffers = false;
+  LinalgPromotionOptions &setDynamicBuffers(unsigned dynamic) {
+    dynamicBuffers = dynamic;
+    return *this;
+  }
+  /// Alignment of promoted buffer. If `None` do not specify alignment.
+  Optional<unsigned> alignment = None;
+  LinalgPromotionOptions &setAlignment(unsigned align) {
+    alignment = align;
+    return *this;
+  }
+};
+LinalgOp promoteSubViews(OpBuilder &b, LinalgOp op,
+                         LinalgPromotionOptions options,
 
 /// Emit a suitable vector form for a Linalg op with fully static shape.
 void vectorizeLinalgOp(OpBuilder &builder, Operation *op);
 
 /// Emits a loop nest of `LoopTy` with the proper body for `op`.
 template <typename LoopTy, typename ConcreteOp>
 Optional<LinalgLoops> linalgLowerOpToLoops(OpBuilder &builder, Operation *op);
 

 //===----------------------------------------------------------------------===//
 /// Emits a `generic` or `indexed_generic` operation with the `indexing_maps`
 /// and `iterator_types` permutated according to `permutation`.
 LogicalResult
 interchangeGenericLinalgOpPrecondition(Operation *op,
                                        ArrayRef<unsigned> interchangeVector);
 
 /// Promote std.subviews feeding linalg operations.
-LogicalResult promoteSubviewsLinalgOpPrecondition(
-    Operation *op, Optional<DenseSet<unsigned>> operandIndicesToPromote = None);
+LogicalResult promoteSubviewsPrecondition(Operation *op,
+                                          LinalgPromotionOptions options);
 
 /// Rewrite a linalg.generic into a suitable vector.contraction op.
 LogicalResult vectorizeLinalgOpPrecondition(Operation *op);
 
 //===----------------------------------------------------------------------===//
 // Transformations exposed as rewrite patterns.
 //===----------------------------------------------------------------------===//
 // Marker used as attribute name in generated Linalg rewriting transformations.
LinalgInterchangePattern(MLIRContext *context,
                            PatternBenefit benefit = 1)
       : LinalgBaseInterchangePattern(OpTy::getOperationName(), context,
                                      interchangeVector, marker, benefit) {}
 };
 
 ///
 /// Linalg promotion patterns.
 ///
-/// Apply the `promoteSubViewOperands` transformation as a pattern.
+/// Apply the `promoteSubViews` transformation as a pattern.
 /// `marker` controls LinalgTransformMarker matching and update when specified.
-/// See `promoteSubViewOperands` for more details.
+/// See `promoteSubViews` for more details.
 struct LinalgBasePromotionPattern : public RewritePattern {
   LinalgBasePromotionPattern(StringRef opName, MLIRContext *context,
-                             ArrayRef<unsigned> operandsToPromote = {},
-                             unsigned alignment = 0,
+                             LinalgPromotionOptions options,
                              LinalgMarker marker = LinalgMarker(),
                              PatternBenefit benefit = 1);
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgMarker marker;
-  /// Indices of subViews to promote.
-  SmallVector<unsigned, 4> operandsToPromote;
-  /// Alignment of promoted buffer.
-  unsigned alignment;
+  /// Promotion options.
+  LinalgPromotionOptions options;
 };
 
 template <typename OpTy>
 struct LinalgPromotionPattern : public LinalgBasePromotionPattern {
-  LinalgPromotionPattern(MLIRContext *context,
-                         ArrayRef<unsigned> operandsToPromote = {},
-                         unsigned alignment = 0,
+  LinalgPromotionPattern(MLIRContext *context, LinalgPromotionOptions options,
                          LinalgMarker marker = LinalgMarker(),
                          PatternBenefit benefit = 1)
-      : LinalgBasePromotionPattern(OpTy::getOperationName(), context,
-                                   operandsToPromote, alignment, marker,
-                                   benefit) {}
-  LinalgPromotionPattern(MLIRContext *context,
-                         ArrayRef<unsigned> operandsToPromote,
-                         LinalgMarker marker = LinalgMarker(),
-                         PatternBenefit benefit = 1)
-      : LinalgPromotionPattern(context, operandsToPromote, 0, marker, benefit) {
-  }
-  LinalgPromotionPattern(MLIRContext *context, unsigned alignment,
-                         LinalgMarker marker = LinalgMarker(),
-                         PatternBenefit benefit = 1)
-      : LinalgPromotionPattern(context, {}, alignment, marker, benefit) {}
-  LinalgPromotionPattern(MLIRContext *context, LinalgMarker marker,
-                         PatternBenefit benefit = 1)
-      : LinalgPromotionPattern(context, {}, 0, marker, benefit) {}
+      : LinalgBasePromotionPattern(OpTy::getOperationName(), context, options,
+                                   marker, benefit) {}
 };
 
 ///
 /// Linalg vectorization patterns.
 ///
 /// Apply the `vectorizeLinalgOp` transformation as a pattern.
 /// `marker` controls LinalgTransformMarker matching and update when specified.
 /// See `vectorizeLinalgOp` for more details.
struct LinalgLoweringPattern : public RewritePattern {
   // TODO: Move implementation to .cpp once named ops are auto-generated.
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
     if (!linalgOp)
       return failure();
     if (failed(marker.checkAndNotify(rewriter, linalgOp)))
       return failure();
-    if (failed(promoteSubviewsLinalgOpPrecondition(op)))
-      return failure();
 
     if (loweringType == LinalgLoweringType::LibraryCall) {
       // TODO: Move lowering to library calls here.
       return failure();
     } else if (loweringType == LinalgLoweringType::Loops) {
       if (failed(linalgOpToLoops<OpTy>(rewriter, op)))
         return failure();
     } else if (loweringType == LinalgLoweringType::AffineLoops) {

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

 using folded_affine_min = FoldedValueBuilder<AffineMinOp>;
 using folded_linalg_range = FoldedValueBuilder<linalg::RangeOp>;
 using folded_std_dim = FoldedValueBuilder<DimOp>;
 using folded_std_subview = FoldedValueBuilder<SubViewOp>;
 using folded_std_view = FoldedValueBuilder<ViewOp>;
 
 #define DEBUG_TYPE "linalg-promotion"
 
-/// If `size` comes from an AffineMinOp and one of the dimensions of AffineMin
+namespace {
+
+/// Helper struct that captures the information required to apply the
+/// transformation on each op. This bridges the abstraction gap with the
+/// user-facing API which exposes positional arguments to control which operands
+/// are promoted.
+struct LinalgOpInstancePromotionOptions {
+  LinalgOpInstancePromotionOptions(LinalgOp op,
+                                   const LinalgPromotionOptions &options);
+  /// SubViews to promote.
+  SetVector<Value> subViews;
+  /// Allow the use of dynamicaly-sized buffers.
+  bool dynamicBuffers;
+  /// Alignment of promoted buffer.
+  Optional<unsigned> alignment;
+};
+} // namespace
+
+LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
+    LinalgOp linalgOp, const LinalgPromotionOptions &options)
+    : subViews(), dynamicBuffers(options.dynamicBuffers),
+      alignment(options.alignment) {
+  if (options.operandsToPromote.hasValue()) {
+    for (unsigned idx : options.operandsToPromote.getValue()) {
+      auto *op = linalgOp.getBuffer(idx).getDefiningOp();
+      if (auto sv = dyn_cast_or_null<SubViewOp>(op))
+        subViews.insert(sv);
+    }
+  } else {
+    unsigned nBuffers = linalgOp.getNumInputsAndOutputBuffers();
+    for (unsigned idx = 0; idx < nBuffers; ++idx) {
+      auto *op = linalgOp.getBuffer(idx).getDefiningOp();
+      if (auto sv = dyn_cast_or_null<SubViewOp>(op))
+        subViews.insert(sv);
+    }
+  }
+}
+
+/// If `size` comes from an AffineMinOp and one of the values of AffineMinOp
 /// is a constant then return a new value set to the smallest such constant.
 /// Otherwise return size.
 static Value extractSmallestConstantBoundingSize(OpBuilder &b, Location loc,
                                                  Value size) {
   auto affineMinOp = dyn_cast_or_null<AffineMinOp>(size.getDefiningOp());
   if (!affineMinOp)
     return size;
-  if (!llvm::any_of(affineMinOp.getAffineMap().getResults(), [](AffineExpr e) {
-        return e.dyn_cast<AffineConstantExpr>();
-      }))
-    return size;
   int64_t minConst = std::numeric_limits<int64_t>::max();
   for (auto e : affineMinOp.getAffineMap().getResults())
     if (auto cst = e.dyn_cast<AffineConstantExpr>())
       minConst = std::min(minConst, cst.getValue());
-  assert(minConst != std::numeric_limits<int64_t>::max());
-  return b.create<ConstantIndexOp>(loc, minConst);
+  return (minConst == std::numeric_limits<int64_t>::max())
+             ? size
+             : b.create<ConstantIndexOp>(loc, minConst);
 }
 
+/// Alloc a new buffer of `size`. If `dynamicBuffers` is true allocate exactly
+/// the size needed, otherwise try to allocate a static bounding box.
 static Value allocBuffer(Type elementType, Value size, bool dynamicBuffers,
-                         OperationFolder *folder, int64_t alignment = 0) {
+                         OperationFolder *folder,
+                         Optional<unsigned> alignment = None) {
   auto *ctx = size.getContext();
   auto width = llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);
   IntegerAttr alignment_attr;
-  if (alignment)
-    alignment_attr = IntegerAttr::get(IntegerType::get(64, ctx), alignment);
+  if (alignment.hasValue())
+    alignment_attr =
+        IntegerAttr::get(IntegerType::get(64, ctx), alignment.getValue());
   if (!dynamicBuffers)
     if (auto cst = dyn_cast_or_null<ConstantIndexOp>(size.getDefiningOp()))
       return std_alloc(
           MemRefType::get(width * cst.getValue(), IntegerType::get(8, ctx)),
           ValueRange{}, alignment_attr);
   Value mul =
       folded_std_muli(folder, folded_std_constant_index(folder, width), size);
   return std_alloc(MemRefType::get(-1, IntegerType::get(8, ctx)), mul,

 //     that corresponds to the size of `subView` and accounting for boundary
 //     effects.
 // The point of the full tile buffer is that constant static tile sizes are
 // folded and result in a buffer type with statically known size and alignment
 // properties.
 // To account for general boundary effects, padding must be performed on the
 // boundary tiles. For now this is done with an unconditional `fill` op followed
 // by a partial `copy` op.
-static PromotionInfo promoteFullTileBuffer(OpBuilder &b, Location loc,
+static PromotionInfo promoteSubviewAsNewBuffer(OpBuilder &b, Location loc,
                                                SubViewOp subView,
                                                bool dynamicBuffers,
-                                           int64_t alignment,
+                                               Optional<unsigned> alignment,
                                                OperationFolder *folder) {
   auto zero = folded_std_constant_index(folder, 0);
   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())) {
     auto rank = en.index();
     auto rangeValue = en.value();
-    // Try to extract a tight constant
+    // 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,
       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};
 }
 
-SmallVector<PromotionInfo, 8>
-mlir::linalg::promoteSubViews(OpBuilder &b, Location loc,
-                              ArrayRef<Value> subViews, bool dynamicBuffers,
-                              int64_t alignment, OperationFolder *folder) {
-  if (subViews.empty())
+static SmallVector<PromotionInfo, 8>
+promoteSubViews(OpBuilder &b, Location loc,
+                LinalgOpInstancePromotionOptions options,
+                OperationFolder *folder) {
+  if (options.subViews.empty())
     return {};
 
   ScopedContext scope(b, loc);
   SmallVector<PromotionInfo, 8> res;
-  res.reserve(subViews.size());
+  res.reserve(options.subViews.size());
   DenseMap<Value, PromotionInfo> promotionInfoMap;
-  for (auto v : subViews) {
+  for (auto v : options.subViews) {
     SubViewOp subView = cast<SubViewOp>(v.getDefiningOp());
-    auto promotionInfo = promoteFullTileBuffer(b, loc, subView, dynamicBuffers,
-                                               alignment, folder);
+    auto promotionInfo = promoteSubviewAsNewBuffer(
+        b, loc, subView, options.dynamicBuffers, options.alignment, folder);
     promotionInfoMap.insert(std::make_pair(subView.getResult(), promotionInfo));
     res.push_back(promotionInfo);
   }
 
-  for (auto v : subViews) {
+  for (auto v : options.subViews) {
     SubViewOp subView = cast<SubViewOp>(v.getDefiningOp());
     auto info = promotionInfoMap.find(v);
     if (info == promotionInfoMap.end())
       continue;
     Value fillVal;
     if (auto t = subView.getType().getElementType().dyn_cast<FloatType>())
       fillVal = folded_std_constant(folder, FloatAttr::get(t, 0.0));
     else if (auto t =
                  subView.getType().getElementType().dyn_cast<IntegerType>())
       fillVal = folded_std_constant_int(folder, 0, t);
     // TODO(ntv): fill is only necessary if `promotionInfo` has a full local
     // view that is different from the partial local view and we are on the
     // boundary.
     linalg_fill(info->second.fullLocalView, fillVal);
   }
 
-  for (auto v : subViews) {
+  for (auto v : options.subViews) {
     auto info = promotionInfoMap.find(v);
     if (info == promotionInfoMap.end())
       continue;
     linalg_copy(cast<SubViewOp>(v.getDefiningOp()),
                 info->second.partialLocalView);
   }
   return res;
 }
 
-LinalgOp mlir::linalg::promoteSubViewOperands(OpBuilder &b, LinalgOp op,
-                                              SetVector<Value> subViews,
-                                              bool dynamicBuffers,
-                                              int64_t alignment,
-                                              OperationFolder *folder) {
+static void promoteSubViews(OpBuilder &b, LinalgOp op,
+                            LinalgOpInstancePromotionOptions options,
   assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
 
   if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) {
     // TODO(ntv): add a level of indirection to linalg.generic.
     if (convOp.padding())
       llvm_unreachable("Unexpected conv with padding");
   }
 
   // 1. Promote the specified views and use them in the new op.
-  ScopedContext scope(b, op.getLoc());
-  auto promotedBufferAndViews =
-      promoteSubViews(b, op.getLoc(), subViews.getArrayRef(), dynamicBuffers,
-                      alignment, folder);
+  auto loc = op.getLoc();
+  auto promotedBufferAndViews = promoteSubViews(b, loc, options, folder);
   SmallVector<Value, 8> opViews;
   opViews.reserve(op.getNumInputsAndOutputs());
   SmallVector<std::pair<Value, Value>, 8> writebackViews;
-  writebackViews.reserve(subViews.size());
+  writebackViews.reserve(promotedBufferAndViews.size());
   unsigned promotedIdx = 0;
   for (auto view : op.getInputsAndOutputBuffers()) {
-    if (subViews.count(view) != 0) {
+    if (options.subViews.count(view) != 0) {
       opViews.push_back(promotedBufferAndViews[promotedIdx].fullLocalView);
       writebackViews.emplace_back(std::make_pair(
           view, promotedBufferAndViews[promotedIdx].partialLocalView));
       promotedIdx++;
     } else {
       opViews.push_back(view);
     }
   }
 
   // 2. Append all other operands as they appear, this enforces that such
   // operands are not views. This is to support cases such as FillOp taking
   // extra scalars etc.
-  auto operands = getAssumedNonViewOperands(op);
-  opViews.append(operands.begin(), operands.end());
-  LinalgOp res = op.clone(b, op.getLoc(), opViews);
+  // Keep a reference to output buffers;
+  DenseSet<Value> originalOutputs(op.getOutputBuffers().begin(),
+                                  op.getOutputBuffers().end());
+  op.getOperation()->setOperands(0, opViews.size(), opViews);
 
+  OpBuilder::InsertionGuard guard(b);
+  b.setInsertionPointAfter(op);
+  ScopedContext scope(b, loc);
   // 3. Emit write-back for the promoted output views: copy the partial view.
-  for (auto viewAndPartialLocalView : writebackViews) {
-    // WARNING: MUST use the old op to determine whether the operand view is an
-    // output.
-    bool isOutput =
-        op.getIndexOfOutputBuffer(viewAndPartialLocalView.first).hasValue();
-    if (isOutput)
+  for (auto viewAndPartialLocalView : writebackViews)
+    if (originalOutputs.count(viewAndPartialLocalView.first))
       linalg_copy(viewAndPartialLocalView.second,
                   viewAndPartialLocalView.first);
-  }
 
-  // 4. Dealloc local buffers.
+  // 4. Dealloc all local buffers.
   for (const auto &pi : promotedBufferAndViews)
     std_dealloc(pi.buffer);
-
-  return res;
-}
-
-static void promoteSubViews(FuncOp f, bool dynamicBuffers) {
-  SmallVector<LinalgOp, 8> toErase;
-  OperationFolder folder(f.getContext());
-  f.walk([dynamicBuffers, &folder, &toErase](LinalgOp op) {
-    if (!op.hasBufferSemantics())
-      return;
-
-    // TODO(ntv) some heuristic here to decide what to promote. Atm only float
-    // and integer buffers can be promoted.
-    SetVector<Value> subViews;
-    OpBuilder b(op);
-    for (auto it : op.getInputsAndOutputBuffers())
-      if (auto sv = dyn_cast_or_null<SubViewOp>(it.getDefiningOp()))
-        if (sv.getType().getElementType().isSignlessIntOrFloat())
-          subViews.insert(sv);
-    if (!subViews.empty()) {
-      promoteSubViewOperands(b, op, subViews, dynamicBuffers, 0, &folder);
-      toErase.push_back(op);
-    }
-  });
-  for (auto op : toErase)
-    op.erase();
 }
 
-LogicalResult mlir::linalg::promoteSubviewsLinalgOpPrecondition(
-    Operation *op, llvm::Optional<DenseSet<unsigned>> operandIndicesToPromote) {
+LogicalResult
+mlir::linalg::promoteSubviewsPrecondition(Operation *op,
+                                          LinalgPromotionOptions options) {
   LinalgOp linOp = dyn_cast<LinalgOp>(op);
   // Transformation applies to buffers only.
   if (!linOp || !linOp.hasBufferSemantics())
     return failure();
+  // Check that at least one of the requested operands is indeed a subview.
   for (auto en : llvm::enumerate(linOp.getInputsAndOutputBuffers())) {
     auto sv = isa_and_nonnull<SubViewOp>(en.value().getDefiningOp());
-    if (sv && (!operandIndicesToPromote.hasValue() ||
-               operandIndicesToPromote->count(en.index())))
+    if (sv) {
+      if (!options.operandsToPromote.hasValue() ||
+          options.operandsToPromote->count(en.index()))
         return success();
     }
+  }
+  // TODO: Check all subviews requested are bound by a static constant.
+  // TODO: Check that the total footprint fits within a given size.
   return failure();
 }
 
+LinalgOp mlir::linalg::promoteSubViews(OpBuilder &b, LinalgOp linalgOp,
+                                       LinalgPromotionOptions options,
+                                       OperationFolder *folder) {
+  LinalgOpInstancePromotionOptions linalgOptions(linalgOp, options);
+  ::promoteSubViews(
+      b, linalgOp, LinalgOpInstancePromotionOptions(linalgOp, options), folder);
+  return linalgOp;
+}
+
 namespace {
 struct LinalgPromotionPass : public LinalgPromotionBase<LinalgPromotionPass> {
   LinalgPromotionPass() = default;
   LinalgPromotionPass(bool dynamicBuffers) {
     this->dynamicBuffers = dynamicBuffers;
   }
 
   void runOnFunction() override {
-    promoteSubViews(getFunction(), dynamicBuffers);
+    OperationFolder folder(&getContext());
+    getFunction().walk([this, &folder](LinalgOp op) {
+      auto options = LinalgPromotionOptions().setDynamicBuffers(dynamicBuffers);
+      if (failed(promoteSubviewsPrecondition(op, options)))
+        return;
+      LLVM_DEBUG(llvm::dbgs() << "Promote: " << *(op.getOperation()) << "\n");
+      OpBuilder b(op);
+      promoteSubViews(b, op, options, &folder);
+    });
   }
 };
 } // namespace
 
+// TODO: support more transformation options in the pass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgPromotionPass(bool dynamicBuffers) {
   return std::make_unique<LinalgPromotionPass>(dynamicBuffers);
 }
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgPromotionPass() {
   return std::make_unique<LinalgPromotionPass>();
 }

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

rewriter.updateRootInPlace(op, [&]() {
     interchange(linalgOp, interchangeVector);
     // New marker if specified.
     marker.replaceLinalgMarker(rewriter, op);
   });
   return success();
 }
 
 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
-    StringRef opName, MLIRContext *context,
-    ArrayRef<unsigned> operandsToPromote, unsigned alignment,
+    StringRef opName, MLIRContext *context, LinalgPromotionOptions options,
     LinalgMarker marker, PatternBenefit benefit)
     : RewritePattern(opName, {}, benefit, context), marker(marker),
-      operandsToPromote(operandsToPromote.begin(), operandsToPromote.end()),
-      alignment(alignment) {}
+      options(options) {}
 
 LogicalResult mlir::linalg::LinalgBasePromotionPattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
-  LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-  if (!linalgOp)
+  if (failed(marker.checkAndNotify(rewriter, op)))
     return failure();
-  if (failed(marker.checkAndNotify(rewriter, linalgOp)))
+  if (failed(promoteSubviewsPrecondition(op, options)))
     return failure();
-  if (operandsToPromote.empty()) {
-    if (failed(promoteSubviewsLinalgOpPrecondition(op, llvm::None)))
-      return failure();
-  } else {
-    DenseSet<unsigned> set;
-    set.insert(operandsToPromote.begin(), operandsToPromote.end());
-    if (failed(promoteSubviewsLinalgOpPrecondition(op, set)))
-      return failure();
-  }
-
-  llvm::SetVector<Value> subViews;
-  if (!operandsToPromote.empty()) {
-    for (unsigned idx : operandsToPromote) {
-      auto *op = linalgOp.getBuffer(idx).getDefiningOp();
-      if (auto sv = dyn_cast_or_null<SubViewOp>(op))
-        subViews.insert(sv);
-    }
-  } else {
-    unsigned nBuffers = linalgOp.getNumInputsAndOutputBuffers();
-    for (unsigned idx = 0; idx < nBuffers; ++idx) {
-      auto *op = linalgOp.getBuffer(idx).getDefiningOp();
-      if (auto sv = dyn_cast_or_null<SubViewOp>(op))
-        subViews.insert(sv);
-    }
-  }
-
-  auto promotedOp =
-      promoteSubViewOperands(rewriter, op, subViews, /*dynamicBuffers=*/false,
-                             /*alignment=*/alignment);
-  marker.replaceLinalgMarker(rewriter, promotedOp.getOperation());
-  rewriter.eraseOp(op);
+  rewriter.updateRootInPlace(op, [&]() {
+    auto promotedOp = promoteSubViews(rewriter, op, options);
+    (void)promotedOp;
+    assert(promotedOp && "Unexpected pattern failure");
+    marker.replaceLinalgMarker(rewriter, op);
+  });
   return success();
 }
 
 mlir::linalg::LinalgBaseVectorizationPattern::LinalgBaseVectorizationPattern(
     StringRef opName, MLIRContext *context, LinalgMarker marker,
     PatternBenefit benefit)
     : RewritePattern(opName, {}, benefit, context), marker(marker) {}
 

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

patterns.insert<LinalgInterchangePattern<IndexedGenericOp>>(
       ctx,
       /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
       LinalgMarker({}, "PERMUTED"));
 
   //===--------------------------------------------------------------------===//
   // Linalg subview operands promotion.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgPromotionPattern<MatmulOp>>(
-      ctx, LinalgMarker({"_promote_views_"}, "_views_promoted_"));
+      ctx, LinalgPromotionOptions(),
+      LinalgMarker({"_promote_views_"}, "_views_promoted_"));
   patterns.insert<LinalgPromotionPattern<MatmulOp>>(
-      ctx,
-      /*operandsToPromote=*/ArrayRef<unsigned>{0},
+      ctx, LinalgPromotionOptions().setOperandsToPromote({0}),
       LinalgMarker({"_promote_first_view_"}, "_first_view_promoted_"));
   patterns.insert<LinalgPromotionPattern<FillOp>>(
-      ctx,
-      /*operandsToPromote=*/ArrayRef<unsigned>{0},
-      /*alignment=*/32,
+      ctx, LinalgPromotionOptions().setOperandsToPromote({0}).setAlignment(32),
       LinalgMarker({"_promote_views_aligned_"}, "_views_aligned_promoted_"));
 
   applyPatternsAndFoldGreedily(funcOp, patterns);
 
   // Drop the marker.
   funcOp.walk([](LinalgOp op) {
     op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
   });