[mlir][linalg] Allow promotion to use callbacks for alloc/dealloc/copies.

[mravishankar]

Summary

Add options to LinalgPromotion to use callbacks for implementating the
allocation, deallocation of buffers used for the promoted subviews,
and to copy data into and from the original subviews to the allocated
buffers.
Also some misc. cleanup of the code.

Depends On D80188

[mravishankar]

Rebase

[nicolasvasilache]

Thanks Mahesh!

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

144

Could we have defaultAllocBufferCallback, ... callbacks exposed, drop the optional and drop the special if/else casing in the impl?

mlir/test/Dialect/Linalg/promotion_options.mlir

34

nit: newline

[mravishankar]

Addressing comments

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

 /// transformation returns early.
 ///
 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed with
 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
 /// integers, in the range 0..`op.rank` without duplications
 /// (i.e. `[1,1,2]` is an invalid permutation).
 LinalgOp interchange(LinalgOp op, ArrayRef<unsigned> interchangeVector);
 
-/// Promotes the `subViews` into a new buffer allocated at the insertion point
-/// `b`. Promotion occurs in 3 steps:
-///   1. Create a new buffer for a full tile (i.e. not clipped at the boundary).
-///   2. Take a full view on the buffer and `linalg.fill` it with zeros (use
-///      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.
+/// Callback function type used to perform the allocation for the promoted
+/// `subView`. In `boundingSubViewsize` a best attempt is made to find the
+/// smallest constant value for the size of the buffer needed for each
+/// dimension. If that is not possible, contains the dynamic size of the
+/// subview. The call back should return the buffer to use.
+using AllocBufferCallbackFn = std::function<Optional<Value>(
+    OpBuilder &b, SubViewOp subView, ArrayRef<Value> boundingSubViewSize,
+    OperationFolder *folder)>;
+
+/// Callback function type used to deallocate the buffers used to hold the
+/// promoted subview.
+using DeallocBufferCallbackFn =
+    std::function<LogicalResult(OpBuilder &b, Value buffer)>;
+
+/// Callback function type used to insert copy from original subview to subview
+/// of the promoted region for the read operands/subview of promoted region to
+/// original subview for the results. The copy has to happen from `src` to
+/// `dst`.
+using CopyCallbackFn =
+    std::function<LogicalResult(OpBuilder &b, Value src, Value dst)>;
+
 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;
   }
LinalgPromotionOptions &setDynamicBuffers(unsigned 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;
   }
+  /// Callback function to do the allocation of the promoted buffer. If None,
+  /// then the default allocation scheme of allocating a memref<?xi8> buffer
+  /// followed by a view operation is used.
+  Optional<AllocBufferCallbackFn> allocationFn = None;
+  Optional<DeallocBufferCallbackFn> deallocationFn = None;
+  LinalgPromotionOptions &
+  setAllocationDeallocationFns(AllocBufferCallbackFn const &allocFn,
+                               DeallocBufferCallbackFn const &deallocFn) {
+    allocationFn = allocFn;
+    deallocationFn = deallocFn;
+    return *this;
+  }
+
+  /// Callback function to do the copy of data to and from the promoted
+  /// subview. If None then a linalg.copy is used.
+  Optional<CopyCallbackFn> copyInFn = None;
+  Optional<CopyCallbackFn> copyOutFn = None;
+  LinalgPromotionOptions &setCopyInOutFns(CopyCallbackFn const &copyIn,
+                                          CopyCallbackFn const &copyOut) {
+    copyInFn = copyIn;
+    copyOutFn = copyOut;
+    return *this;
+  }
 };
-LinalgOp promoteSubViews(OpBuilder &b, LinalgOp op,
+
+/// Promotes the `subViews` into a new buffer allocated at the insertion point
+/// `b`. Promotion occurs in 3 steps:
+///   1. Create a new buffer for a full tile (i.e. not clipped at the boundary).
+///   2. Take a full view on the buffer.
+///   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 the modified linalg op (the modification happens in place) as well
+/// as all the copy ops created.
+Optional<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
                                    LinalgPromotionOptions options,
                                    OperationFolder *folder = nullptr);
 
 /// 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);
 

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 /// Returns the values obtained by applying `map` to the list of values.
 /// When non-null, the optional pointer `folder` is used to call into the
 /// `createAndFold` builder method. If `folder` is null, the regular `create`
 /// method is called.
 SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
                                        AffineMap map, ArrayRef<Value> values,
                                        OperationFolder *folder = nullptr);
 
-struct PromotionInfo {
-  Value buffer;
-  Value fullLocalView;
-  Value partialLocalView;
-};
-
-/// Promotes the `subViews` into a new buffer allocated at the insertion point
-/// `b`. For now, promotion occurs in 3 steps:
-///   1. Create a new buffer for a full tile (i.e. not clipped at the boundary).
-///   2. Take a full view on the buffer and `linalg.fill` it with zeros (use
-///      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.
-SmallVector<PromotionInfo, 8>
-promoteSubViews(OpBuilder &b, Location loc, ArrayRef<Value> subViews,
-                bool dynamicBuffers = false, int64_t alignment = 0,
-                OperationFolder *folder = nullptr);
-
 /// Returns all the operands of `linalgOp` that are not views.
 /// Asserts that these operands are value types to allow transformations like
 /// tiling to just use the values when cloning `linalgOp`.
 SmallVector<Value, 4> getAssumedNonViewOperands(LinalgOp linalgOp);
 
 /// Apply the permutation defined by `permutation` to `inVec`.
 /// Element `i` in `inVec` is mapped to location `j = permutation[i]`.
 /// E.g.: for an input vector `inVec = ['a', 'b', 'c']` and a permutation vector

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

 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/FoldUtils.h"
-
-#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 using namespace mlir::scf;
 
-using llvm::SetVector;
+using llvm::MapVector;
 
 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 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) {
+  Optional<int64_t> boundingConst = {};
+  if (auto affineMinOp = size.getDefiningOp<AffineMinOp>()) {
+    for (auto e : affineMinOp.getAffineMap().getResults())
+      if (auto cst = e.dyn_cast<AffineConstantExpr>())
+        boundingConst = boundingConst
+                            ? std::min(boundingConst.getValue(), cst.getValue())
+                            : cst.getValue();
+  } else if (auto constIndexOp = size.getDefiningOp<ConstantOp>()) {
+    if (constIndexOp.getType().isa<IndexType>())
+      boundingConst = constIndexOp.value().cast<IntegerAttr>().getInt();
+  }
+  return boundingConst && *boundingConst >= 0
+             ? b.create<ConstantIndexOp>(loc, *boundingConst)
+             : size;
+}
+
+/// 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,
+                         Optional<unsigned> alignment = None) {
+  auto *ctx = size.getContext();
+  auto width = llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);
+  IntegerAttr alignment_attr;
+  if (alignment.hasValue())
+    alignment_attr =
+        IntegerAttr::get(IntegerType::get(64, ctx), alignment.getValue());
+  if (!dynamicBuffers)
+    if (auto cst = size.getDefiningOp<ConstantIndexOp>())
+      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,
+                   alignment_attr);
+}
+
+/// Default allocation callback function. This allocates a promoted buffer when
+/// no call back to do so is provided. The default is to allocate a
+/// memref<..xi8> and return a view to get a memref type of shape
+/// boundingSubViewSize.
+static Optional<Value>
+allocBufferCallBack(OpBuilder &builder, SubViewOp subView,
+                    ArrayRef<Value> boundingSubViewSize, bool dynamicBuffers,
+                    Optional<unsigned> alignment, OperationFolder *folder) {
+  ShapedType viewType = subView.getType();
+  int64_t rank = viewType.getRank();
+  (void)rank;
+  assert(rank > 0 && boundingSubViewSize.size() == static_cast<size_t>(rank));
+  auto zero = folded_std_constant_index(folder, 0);
+  auto one = folded_std_constant_index(folder, 1);
+
+  Value allocSize = one;
+  for (auto size : llvm::enumerate(boundingSubViewSize))
+    allocSize = folded_std_muli(folder, allocSize, size.value());
+  Value buffer = allocBuffer(viewType.getElementType(), allocSize,
+                             dynamicBuffers, folder, alignment);
+  SmallVector<int64_t, 4> dynSizes(boundingSubViewSize.size(),
+                                   ShapedType::kDynamicSize);
+  Value view = folded_std_view(
+      folder, MemRefType::get(dynSizes, viewType.getElementType()), buffer,
+      zero, boundingSubViewSize);
+  return view;
+}
+
+/// Default implementation of deallocation of the buffer use for promotion. It
+/// expects to get the same value that the default allocation method returned,
+/// i.e. result of a ViewOp.
+static LogicalResult deallocCallBack(OpBuilder &b, Value fullLocalView) {
+  auto viewOp = fullLocalView.getDefiningOp<ViewOp>();
+  assert(viewOp && "expected full local view to be a ViewOp");
+  std_dealloc(viewOp.source());
+  return success();
+}
+
 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;
+  MapVector<unsigned, Value> subViews;
   /// True if the full view should be used for the promoted buffer.
   DenseMap<Value, bool> useFullTileBuffers;
+
+  /// Callback functions for allocation and deallocation of promoted buffers, as
+  /// well as to copy the data into and out of these buffers.
+  AllocBufferCallbackFn allocationFn;
+  DeallocBufferCallbackFn deallocationFn;
+  CopyCallbackFn copyInFn;
+  CopyCallbackFn copyOutFn;
+
   /// Allow the use of dynamicaly-sized buffers.
   bool dynamicBuffers;
   /// Alignment of promoted buffer.
   Optional<unsigned> alignment;
 };
+
+struct PromotionInfo {
+  Value fullLocalView;
+  Value partialLocalView;
+};
 } // namespace
 
 LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
     LinalgOp linalgOp, const LinalgPromotionOptions &options)
-    : subViews(), useFullTileBuffers(), dynamicBuffers(options.dynamicBuffers),
+    : subViews(), dynamicBuffers(options.dynamicBuffers),
       alignment(options.alignment) {
   unsigned nBuffers = linalgOp.getNumInputsAndOutputBuffers();
   auto vUseFullTileBuffers =
       options.useFullTileBuffers.getValueOr(llvm::SmallBitVector());
   vUseFullTileBuffers.resize(nBuffers, options.useFullTileBuffersDefault);
 
-  if (options.operandsToPromote.hasValue()) {
-    for (auto it : llvm::enumerate(options.operandsToPromote.getValue())) {
-      auto *op = linalgOp.getBuffer(it.value()).getDefiningOp();
-      if (auto sv = dyn_cast_or_null<SubViewOp>(op)) {
-        subViews.insert(sv);
-        useFullTileBuffers[sv] = vUseFullTileBuffers[it.index()];
-      }
-    }
-  } else {
-    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);
-        useFullTileBuffers[sv] = vUseFullTileBuffers[idx];
-      }
-    }
-  }
-}
+  for (unsigned idx = 0; idx != nBuffers; ++idx) {
+    if (options.operandsToPromote && !options.operandsToPromote->count(idx))
+      continue;
+      subViews[idx] = 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 = size.getDefiningOp<AffineMinOp>();
-  if (!affineMinOp)
-    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());
-  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,
-                         Optional<unsigned> alignment = None) {
-  auto *ctx = size.getContext();
-  auto width = llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);
-  IntegerAttr alignment_attr;
-  if (alignment.hasValue())
-    alignment_attr =
-        IntegerAttr::get(IntegerType::get(64, ctx), alignment.getValue());
-  if (!dynamicBuffers)
-    if (auto cst = size.getDefiningOp<ConstantIndexOp>())
-      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,
-                   alignment_attr);
+  allocationFn =
+      (options.allocationFn ? *(options.allocationFn)
+                            : [&](OpBuilder &builder, SubViewOp subViewOp,
+                                  ArrayRef<Value> boundingSubViewSize,
+                                  OperationFolder *folder) -> Optional<Value> {
+        return allocBufferCallBack(builder, subViewOp, boundingSubViewSize,
+                                   dynamicBuffers, alignment, folder);
+      });
+  deallocationFn =
+      (options.deallocationFn ? *(options.deallocationFn) : deallocCallBack);
+  auto defaultCopyCallBack = [&](OpBuilder &builder, Value src,
+                                 Value dst) -> LogicalResult {
+    linalg_copy(src, dst);
+    return success();
+  };
+  copyInFn = (options.copyInFn ? *(options.copyInFn) : defaultCopyCallBack);
+  copyOutFn = (options.copyOutFn ? *(options.copyOutFn) : defaultCopyCallBack);
 }
 
 // Performs promotion of a `subView` into a local buffer of the size of the
 // *ranges* of the `subView`. This produces a buffer whose size may be bigger
 // than the actual size of the `subView` at the boundaries.
 // This is related to the full/partial tile problem.
 // Returns a PromotionInfo containing a `buffer`, `fullLocalView` and
 // `partialLocalView` such that:
 //   * `buffer` is always the size of the full tile.
 //   * `fullLocalView` is a dense contiguous view into that buffer.
 //   * `partialLocalView` is a dense non-contiguous slice of `fullLocalView`
 //     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 promoteSubviewAsNewBuffer(OpBuilder &b, Location loc,
-                                               SubViewOp subView,
-                                               bool dynamicBuffers,
-                                               Optional<unsigned> alignment,
-                                               OperationFolder *folder) {
-  auto zero = folded_std_constant_index(folder, 0);
-  auto one = folded_std_constant_index(folder, 1);
-
+static Optional<PromotionInfo>
+promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, SubViewOp subView,
+                          LinalgOpInstancePromotionOptions const &options,
   auto viewType = subView.getType();
   auto rank = viewType.getRank();
-  Value allocSize = one;
-  SmallVector<Value, 8> fullSizes, partialSizes;
+  SmallVector<Value, 4> fullSizes, partialSizes;
   fullSizes.reserve(rank);
   partialSizes.reserve(rank);
   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));
+    partialSizes.push_back(folded_std_dim(folder, subView, en.index()));
   }
   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,
-      zero, fullSizes);
+  // If a callback is not specified, then use the default implementation for
+  // allocating the promoted buffer.
+  Optional<Value> fullLocalView =
+      options.allocationFn(b, subView, fullSizes, folder);
+  if (!fullLocalView)
+    return {};
+  auto zero = folded_std_constant_index(folder, 0);
+  auto one = folded_std_constant_index(folder, 1);
   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};
+      folded_std_subview(folder, *fullLocalView, zeros, partialSizes, ones);
+  return PromotionInfo{*fullLocalView, partialLocalView};
 }
 
-static SmallVector<PromotionInfo, 8>
+static Optional<MapVector<unsigned, PromotionInfo>>
 promoteSubViews(OpBuilder &b, Location loc,
                 LinalgOpInstancePromotionOptions options,
                 OperationFolder *folder) {
   if (options.subViews.empty())
     return {};
 
   ScopedContext scope(b, loc);
-  SmallVector<PromotionInfo, 8> res;
-  res.reserve(options.subViews.size());
-  DenseMap<Value, PromotionInfo> promotionInfoMap;
-  for (auto v : options.subViews) {
-    SubViewOp subView = cast<SubViewOp>(v.getDefiningOp());
-    auto promotionInfo = promoteSubviewAsNewBuffer(
-        b, loc, subView, options.dynamicBuffers, options.alignment, folder);
-    promotionInfoMap.insert(std::make_pair(subView.getResult(), promotionInfo));
-    res.push_back(promotionInfo);
-  }
+  MapVector<unsigned, PromotionInfo> promotionInfoMap;
 
   for (auto v : options.subViews) {
-    SubViewOp subView = cast<SubViewOp>(v.getDefiningOp());
-    auto info = promotionInfoMap.find(v);
-    if (info == promotionInfoMap.end())
-      continue;
+    SubViewOp subView = cast<SubViewOp>(v.second.getDefiningOp());
+    Optional<PromotionInfo> promotionInfo =
+        promoteSubviewAsNewBuffer(b, loc, subView, options, folder);
+    if (!promotionInfo)
+      return {};
+    promotionInfoMap[v.first] = *promotionInfo;
+
     // Only fill the buffer if the full local view is used
-    if (!options.useFullTileBuffers[v])
+    if (!options.useFullTileBuffers[v.second])
       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);
+    linalg_fill(promotionInfo->fullLocalView, fillVal);
   }
 
+  // Copy data into the promoted buffers. Use callback if provided.
   for (auto v : options.subViews) {
-    auto info = promotionInfoMap.find(v);
+    auto info = promotionInfoMap.find(v.first);
     if (info == promotionInfoMap.end())
       continue;
-    linalg_copy(cast<SubViewOp>(v.getDefiningOp()),
-                info->second.partialLocalView);
+    if (failed(options.copyInFn(b, cast<SubViewOp>(v.second.getDefiningOp()),
+                                info->second.partialLocalView)))
+      return {};
   }
-  return res;
+  return promotionInfoMap;
 }
 
-static void promoteSubViews(OpBuilder &b, LinalgOp op,
-                            LinalgOpInstancePromotionOptions options,
-                            OperationFolder *folder) {
+static Optional<LinalgOp>
+promoteSubViews(OpBuilder &b, LinalgOp op,
   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");
+      return {};
   }
 
   // 1. Promote the specified views and use them in the new op.
   auto loc = op.getLoc();
-  auto promotedBufferAndViews = promoteSubViews(b, loc, options, folder);
+  auto promotedBuffersAndViews = promoteSubViews(b, loc, options, folder);
+  if (!promotedBuffersAndViews ||
+      promotedBuffersAndViews->size() != options.subViews.size())
+    return {};
+
+  // 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.  Keep a reference to output buffers;
   SmallVector<Value, 8> opViews;
   opViews.reserve(op.getNumInputsAndOutputs());
   SmallVector<std::pair<Value, Value>, 8> writebackViews;
-  writebackViews.reserve(promotedBufferAndViews.size());
-  unsigned promotedIdx = 0;
-  for (auto view : op.getInputsAndOutputBuffers()) {
-    if (options.subViews.count(view) != 0) {
-      if (options.useFullTileBuffers[view])
-        opViews.push_back(promotedBufferAndViews[promotedIdx].fullLocalView);
+  writebackViews.reserve(promotedBuffersAndViews->size());
+  for (auto view : llvm::enumerate(op.getInputsAndOutputBuffers())) {
+    if (options.subViews.count(view.index()) != 0) {
+      if (options.useFullTileBuffers[view.value()])
+        opViews.push_back(
+            (*promotedBuffersAndViews)[view.index()].fullLocalView);
       else
-        opViews.push_back(promotedBufferAndViews[promotedIdx].partialLocalView);
+        opViews.push_back(
+            (*promotedBuffersAndViews)[view.index()].partialLocalView);
+      if (view.index() >= op.getNumInputs())
         writebackViews.emplace_back(std::make_pair(
-          view, promotedBufferAndViews[promotedIdx].partialLocalView));
-      promotedIdx++;
+            view.value(),
+            (*promotedBuffersAndViews)[view.index()].partialLocalView));
     } else {
-      opViews.push_back(view);
+      opViews.push_back(view.value());
     }
   }
-
-  // 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.
-  // 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)
-    if (originalOutputs.count(viewAndPartialLocalView.first))
-      linalg_copy(viewAndPartialLocalView.second,
-                  viewAndPartialLocalView.first);
+  for (auto viewAndPartialLocalView : writebackViews) {
+    if (failed(options.copyOutFn(b, viewAndPartialLocalView.second,
+                                 viewAndPartialLocalView.first)))
+      return {};
+  }
 
   // 4. Dealloc all local buffers.
-  for (const auto &pi : promotedBufferAndViews)
-    std_dealloc(pi.buffer);
+  for (const auto &pi : *promotedBuffersAndViews) {
+    options.deallocationFn(b, pi.second.fullLocalView);
+  }
+  return op;
 }
 
 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) {
       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,
+Optional<LinalgOp> mlir::linalg::promoteSubViews(OpBuilder &b,
+                                                 LinalgOp linalgOp,
                                                  LinalgPromotionOptions options,
                                                  OperationFolder *folder) {
   LinalgOpInstancePromotionOptions linalgOptions(linalgOp, options);
-  ::promoteSubViews(
+  return ::promoteSubViews(
       b, linalgOp, LinalgOpInstancePromotionOptions(linalgOp, options), folder);
-  return linalgOp;
 }
 
 namespace {
 struct LinalgPromotionPass : public LinalgPromotionBase<LinalgPromotionPass> {
   LinalgPromotionPass() = default;
   LinalgPromotionPass(bool dynamicBuffers) {
     this->dynamicBuffers = dynamicBuffers;
   }

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

: RewritePattern(opName, {}, benefit, context), marker(marker),
       options(options) {}
 
 LogicalResult mlir::linalg::LinalgBasePromotionPattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   if (failed(marker.checkAndNotify(rewriter, op)))
     return failure();
   if (failed(promoteSubviewsPrecondition(op, options)))
     return failure();
-  rewriter.updateRootInPlace(op, [&]() {
-    auto promotedOp = promoteSubViews(rewriter, op, options);
-    (void)promotedOp;
-    assert(promotedOp && "Unexpected pattern failure");
-    marker.replaceLinalgMarker(rewriter, op);
-  });
+
+  // TODO: We cannot use root update here. This pattern is creating other ops,
+  // so if the promotion fails, those need to be cleaned up, which doesnt seem
+  // to be happening here. So to fail properly, we should be cloning the op and
+  // deleting the previous op. This needs more investigation.
+  rewriter.startRootUpdate(op);
+  Optional<LinalgOp> promotedOp = promoteSubViews(rewriter, op, options);
+  if (!promotedOp) {
+    rewriter.cancelRootUpdate(op);
+    return op->emitError("subview promotion failed");
+  }
+  rewriter.finalizeRootUpdate(op);
   return success();
 }
 
 mlir::linalg::LinalgBaseVectorizationPattern::LinalgBaseVectorizationPattern(
     StringRef opName, MLIRContext *context, LinalgMarker marker,
     PatternBenefit benefit)
     : RewritePattern(opName, {}, benefit, context), marker(marker) {}
 

mlir/test/Dialect/Linalg/promotion_options.mlir

+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-linalg-promotion-options -split-input-file | FileCheck %s
+
+func @gemm(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)
+{
+   linalg.matmul(%a, %b, %c) {__internal_linalg_transform__ = "START"}
+     : memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
+   return
+}
+
+//      CHECK: func @gemm
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+//  CHECK-DAG: %[[C42:.+]] = constant 4.200000e+01 : f32
+//      CHECK: scf.for
+//      CHECK:   scf.for
+//      CHECK:     scf.for
+//      CHECK:       %[[T7:.+]] = subview %[[ARG0]]
+//      CHECK:       %[[T12:.+]] = subview %[[ARG1]]
+//      CHECK:       %[[T17:.+]] = subview %[[ARG2]]
+//      CHECK:       %[[T18:.+]] = alloc(%{{.*}}, %{{.*}}) : memref<?x?xf32, 3>
+//      CHECK:       %[[T19:.+]] = subview %[[T18]]
+//      CHECK:       %[[T20:.+]] = alloc(%{{.*}}, %{{.*}}) : memref<?x?xf32, 3>
+//      CHECK:       %[[T21:.+]] = subview %[[T20]]
+//      CHECK:       linalg.fill(%[[T19]], %[[C42]])
+//      CHECK:       linalg.copy(%[[T7]], %[[T19]])
+//      CHECK:       linalg.fill(%[[T21]], %[[C42]])
+//      CHECK:       linalg.copy(%[[T17]], %[[T21]])
+//      CHECK:       linalg.matmul(%[[T19]], %[[T12]], %[[T21]])
+//  CHECK-NOT:       linalg.fill
+//      CHECK:       linalg.copy(%[[T21]], %[[T17]])
+//      CHECK:       dealloc %[[T18]]
+//      CHECK:       dealloc %[[T20]]

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

Option<bool> testMatmulToVectorPatterns1dTiling{
           "1-d tiling"),
       llvm::cl::init(false)};
   Option<bool> testMatmulToVectorPatterns2dTiling{
       *this, "test-matmul-to-vector-patterns-tile-2d",
       llvm::cl::desc(
           "Test a fused pass that applies patterns from matmul to vectors via "
           "2-d tiling"),
       llvm::cl::init(false)};
+  Option<bool> testPromotionOptions{*this, "test-linalg-promotion-options",
+                                    llvm::cl::desc("Test promotion options"),
+                                    llvm::cl::init(false)};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   OwningRewritePatternList patterns;
 
   //===--------------------------------------------------------------------===//
static void fillL1TilingAndMatmulToVectorPatterns(
 
   patternsVector.emplace_back(
       LinalgVectorizationPattern<MatmulOp>(context, LinalgMarker({"VEC"})));
   patternsVector.back()
       .insert<LinalgVectorizationPattern<FillOp>,
               LinalgVectorizationPattern<CopyOp>>(context);
 }
 
+//===----------------------------------------------------------------------===//
+// Test promotion callbacks
+//===----------------------------------------------------------------------===//
+
+// Allocation call back
+static Optional<Value> allocCallBackFn(OpBuilder &b, SubViewOp subView,
+                                       ArrayRef<Value> boundingSubViewSize,
+                                       OperationFolder *folder) {
+  SmallVector<int64_t, 4> shape(boundingSubViewSize.size(), -1);
+  return b
+      .create<AllocOp>(subView.getLoc(),
+                       MemRefType::get(shape,
+                                       subView.getType().getElementType(),
+                                       /*affineMapComposition =*/{}, 3),
+                       boundingSubViewSize)
+      .getResult();
+}
+
+// Deallocation callback
+static LogicalResult deallocCallBackFn(OpBuilder &b, Value buffer) {
+  b.create<DeallocOp>(buffer.getLoc(), buffer);
+  return success();
+}
+
+// Copy in call back
+static LogicalResult copyCallBackFn(OpBuilder &b, Value src, Value dst,
+                                    bool isOutput) {
+  auto floatType = src.getType().cast<MemRefType>().getElementType();
+  if (!floatType.isa<FloatType>())
+    return failure();
+  if (!isOutput)
+    b.create<FillOp>(
+        src.getLoc(), dst,
+        b.create<ConstantOp>(src.getLoc(), FloatAttr::get(floatType, 42.0)));
+  b.create<CopyOp>(src.getLoc(), src, dst);
+  return success();
+}
+
+void fillPromotionCallBackPatterns(MLIRContext *context,
+                                   OwningRewritePatternList &patterns) {
+  patterns.insert<LinalgTilingPattern<MatmulOp>>(
+      context, LinalgTilingOptions().setTileSizes({16, 16, 16}),
+      LinalgMarker({"START"}, "PROMOTE"));
+  patterns.insert<LinalgPromotionPattern<MatmulOp>>(
+      context,
+      LinalgPromotionOptions()
+          .setOperandsToPromote({0, 2})
+          .setUseFullTileBuffers({false, false})
+          .setAllocationDeallocationFns(allocCallBackFn, deallocCallBackFn)
+          .setCopyInOutFns(
+              [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
+                copyCallBackFn(b, src, dst, false);
+                return success();
+              },
+              [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
+                copyCallBackFn(b, src, dst, true);
+                return success();
+              }),
+      LinalgMarker({"PROMOTE"}));
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnFunction() {
   if (testPatterns) {
     applyPatterns(getFunction());
+    return;
+  }
+  if (testPromotionOptions) {
+    OwningRewritePatternList patterns;
+    fillPromotionCallBackPatterns(&getContext(), patterns);
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
   } else {
     SmallVector<OwningRewritePatternList, 4> stage1Patterns;
     if (testMatmulToVectorPatterns1dTiling) {
       fillL1TilingAndMatmulToVectorPatterns(getFunction(), "START",
                                             stage1Patterns);
     } else if (testMatmulToVectorPatterns2dTiling) {
       stage1Patterns.emplace_back(
           LinalgTilingPattern<MatmulOp>(&getContext(),