[mlir][MemRef] Rewrite multi-buffering with proper composable abstractions

[nicolasvasilache]

Summary

Rewrite and document multi-buffering properly:

1. Use IndexingUtils / StaticValueUtils instead of duplicating functionality
2. Properly plumb RewriterBase through.
3. Add support
4. Better debug messages.

This revision is otherwise almost NFC, if it weren't for the extra DeallocOp
support that would previoulsy make multi-buffering fail.

Depends on: D145036

[nicolasvasilache]

Add test and fix Dealloc precondition.

[nicolasvasilache]

Nit

[qcolombet]

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

89

Nit: Dangling For each

93

What's different about this function and getValueOrCreateConstantIndexOp?
Put differently, do we need it?

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

72

Agree, we should fix this at some point. In particular I believe it goes against what we say in the developer guide for MLIR.

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

153

Assuming we need this function, we can drop the comments here.
They are already in the header.

[nicolasvasilache]

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

93

yeah it's annoying... functionality has been duplicated.
Killing redundant uses, the price is an extra comment and a TODO for the future..

mlir/include/mlir/Dialect/Arith/Utils/Utils.h

 };
 
 /// Converts an OpFoldResult to a Value. Returns the fold result if it casts to
 /// a Value or creates a ConstantIndexOp if it casts to an IntegerAttribute.
 /// Other attribute types are not supported.
 Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
                                       OpFoldResult ofr);
 
-/// Create a cast from an index-like value (index or integer) to another
-/// index-like value. If the value type and the target type are the same, it
-/// returns the original value.
-Value getValueOrCreateCastToIndexLike(OpBuilder &b, Location loc,
-                                      Type targetType, Value value);
-
 /// Similar to the other overload, but converts multiple OpFoldResults into
 /// Values.
 SmallVector<Value>
 getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
                                 ArrayRef<OpFoldResult> valueOrAttrVec);
 
+/// Create a cast from an index-like value (index or integer) to another
+/// index-like value. If the value type and the target type are the same, it
+/// returns the original value.
+Value getValueOrCreateCastToIndexLike(OpBuilder &b, Location loc,
+                                      Type targetType, Value value);
+
 /// Converts a scalar value `operand` to type `toType`. If the value doesn't
 /// convert, a warning will be issued and the operand is returned as is (which
 /// will presumably yield a verification issue downstream).
 Value convertScalarToDtype(OpBuilder &b, Location loc, Value operand,
                            Type toType, bool isUnsignedCast);
 
 /// Helper struct to build simple arithmetic quantities with minimal type
 /// inference support.

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

 #define MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES_H
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 class AffineDialect;
 class ModuleOp;
+class RewriterBase;
 
 namespace arith {
 class WideIntEmulationConverter;
 } // namespace arith
 
 namespace func {
 class FuncDialect;
 } // namespace func

 ///   %d = arith.divsi %s, %c3 : index
 ///   %i = arith.remsi %d, %c5 : index
 ///   %sv = memref.subview %0[%i, 0, 0] [1, 4, 128] [1, 1, 1] :
 ///     memref<5x4x128xf32> to memref<4x128xf32, strided<[128, 1], offset: ?>>
 ///   memref.copy %1, %sv : memref<4x128xf32> to memref<4x128xf32, strided<...>>
 ///   "some_use"(%sv) : (memref<4x128xf32, strided<...>) -> ()
 /// }
 /// ```
+FailureOr<memref::AllocOp> multiBuffer(RewriterBase &rewriter,
+                                       memref::AllocOp allocOp,
+                                       unsigned multiplier,
+                                       bool skipOverrideAnalysis = false);
+/// Call into `multiBuffer` with  locally constructed IRRewriter.
 FailureOr<memref::AllocOp> multiBuffer(memref::AllocOp allocOp,
                                        unsigned multiplier,
                                        bool skipOverrideAnalysis = false);
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/Dialect/SCF/IR/SCF.h

 //===- SCFOps.h - Structured Control Flow -----------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines structured control flow operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SCF_SCF_H
 #define MLIR_DIALECT_SCF_SCF_H
 
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/RegionKindInterface.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/ParallelCombiningOpInterface.h"

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

let extraClassDeclaration = [{
     // Get steps as OpFoldResult.
     SmallVector<OpFoldResult> getMixedStep() {
       Builder b(getOperation()->getContext());
       return getMixedValues(getStaticStep(), getDynamicStep(), b);
     }
 
     /// Get lower bounds as values.
     SmallVector<Value> getLowerBound(OpBuilder &b) {
-      return getAsValues(b, getLoc(), getMixedLowerBound());
+      return getValueOrCreateConstantIndexOp(b, getLoc(), getMixedLowerBound());
     }
 
     /// Get upper bounds as values.
     SmallVector<Value> getUpperBound(OpBuilder &b) {
-      return getAsValues(b, getLoc(), getMixedUpperBound());
+      return getValueOrCreateConstantIndexOp(b, getLoc(), getMixedUpperBound());
     }
 
     /// Get steps as values.
     SmallVector<Value> getStep(OpBuilder &b) {
-      return getAsValues(b, getLoc(), getMixedStep());
+      return getValueOrCreateConstantIndexOp(b, getLoc(), getMixedStep());
     }
 
     int64_t getRank() { return getStaticLowerBound().size(); }
 
     /// Number of operands controlling the loop: lbs, ubs, steps
     unsigned getNumControlOperands() { return 3 * getRank(); }
 
     /// Number of dynamic operands controlling the loop: lbs, ubs, steps

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

 /// Return true if ofr1 and ofr2 are the same integer constant attribute
 /// values or the same SSA value. Ignore integer bitwitdh and type mismatch
 /// that come from the fact there is no IndexAttr and that IndexType have no
 /// bitwidth.
 bool isEqualConstantIntOrValue(OpFoldResult ofr1, OpFoldResult ofr2);
 bool isEqualConstantIntOrValueArray(ArrayRef<OpFoldResult> ofrs1,
                                     ArrayRef<OpFoldResult> ofrs2);
 
-/// Helper function to convert a vector of `OpFoldResult`s into a vector of
-/// `Value`s. For each `OpFoldResult` in `valueOrAttrVec` return the fold
-/// result if it casts to  a `Value` or create an index-type constant if it
-/// casts to `IntegerAttr`. No other attribute types are supported.
-SmallVector<Value> getAsValues(OpBuilder &b, Location loc,
-                               ArrayRef<OpFoldResult> valueOrAttrVec);
+// To convert an OpFoldResult to a Value of index type, see:
+//   mlir/include/mlir/Dialect/Arith/Utils/Utils.h
+// TODO: find a better common landing place.
+//
+// Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
+//                                       OpFoldResult ofr);
+
+// To convert an OpFoldResult to a Value of index type, see:
+//   mlir/include/mlir/Dialect/Arith/Utils/Utils.h
+// TODO: find a better common landing place.
+//
+// SmallVector<Value>
+// getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
+//                                 ArrayRef<OpFoldResult> valueOrAttrVec);
 
 /// Return a vector of OpFoldResults with the same size a staticValues, but
 /// all elements for which ShapedType::isDynamic is true, will be replaced by
 /// dynamicValues.
 SmallVector<OpFoldResult> getMixedValues(ArrayRef<int64_t> staticValues,
                                          ValueRange dynamicValues, Builder &b);
 
 /// Decompose a vector of mixed static or dynamic values into the

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

FailureOr<linalg::ForallReductionTilingResult> linalg::tileReductionUsingForall(
 
   Operation *tiledOp = nullptr;
 
   SmallVector<OpFoldResult> nonZeroNumThreads =
       llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
         return !isConstantIntValue(ofr, 0);
       }));
   SmallVector<Value> materializedNonZeroNumThreads =
-      getAsValues(b, loc, nonZeroNumThreads);
+      getValueOrCreateConstantIndexOp(b, loc, nonZeroNumThreads);
 
   // 2. Create the ForallOp with an empty region.
   scf::ForallOp forallOp = b.create<scf::ForallOp>(
       loc, getAsOpFoldResult(materializedNonZeroNumThreads),
       (*identityTensor)->getResults(), mapping);
 
   // 3. Calculate the tile offsets and sizes for the subsequent loop that will
   // be nested under `forallOp`.

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
+#include "llvm/Support/Debug.h"
 
 using namespace mlir;
 
+#define DEBUG_TYPE "memref-transforms"
+#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
+
 //===----------------------------------------------------------------------===//
 // MemRefMultiBufferOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MemRefMultiBufferOp::apply(
     transform::TransformResults &transformResults,
     transform::TransformState &state) {
   SmallVector<Operation *> results;
   ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  IRRewriter rewriter(getContext());
   for (auto *op : payloadOps) {
     bool canApplyMultiBuffer = true;
     auto target = cast<memref::AllocOp>(op);
+    LLVM_DEBUG(DBGS() << "Start multibuffer transform op: " << target << "\n";);
     // Skip allocations not used in a loop.
     for (Operation *user : target->getUsers()) {
+      if (isa<memref::DeallocOp>(user))
+        continue;
       auto loop = user->getParentOfType<LoopLikeOpInterface>();
       if (!loop) {
+        LLVM_DEBUG(DBGS() << "--allocation not used in a loop\n";
+                   DBGS() << "----due to user: " << *user;);
         canApplyMultiBuffer = false;
         break;
       }
     }
-    if (!canApplyMultiBuffer)
+    if (!canApplyMultiBuffer) {
+      LLVM_DEBUG(DBGS() << "--cannot apply multibuffering -> Skip\n";);
       continue;
+    }
 
     auto newBuffer =
-        memref::multiBuffer(target, getFactor(), getSkipAnalysis());
-    if (failed(newBuffer))
+        memref::multiBuffer(rewriter, target, getFactor(), getSkipAnalysis());
+
+    if (failed(newBuffer)) {
+      LLVM_DEBUG(DBGS() << "--op failed to multibuffer\n";);
       return emitSilenceableFailure(target->getLoc())
              << "op failed to multibuffer";
+    }
 
     results.push_back(*newBuffer);
   }
   transformResults.set(getResult().cast<OpResult>(), results);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//

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

 //===----------- MultiBuffering.cpp ---------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements multi buffering transformation.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Dominance.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/ValueRange.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "memref-transforms"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 #define DBGSNL() (llvm::dbgs() << "\n")
 
 /// Return true if the op fully overwrite the given `buffer` value.
 static bool overrideBuffer(Operation *op, Value buffer) {
   auto copyOp = dyn_cast<memref::CopyOp>(op);
   if (!copyOp)
     return false;
   return copyOp.getTarget() == buffer;
 }
 
 /// Replace the uses of `oldOp` with the given `val` and for subview uses
 /// propagate the type change. Changing the memref type may require propagating
 /// it through subview ops so we cannot just do a replaceAllUse but need to
 /// propagate the type change and erase old subview ops.
-static void replaceUsesAndPropagateType(Operation *oldOp, Value val,
-                                        OpBuilder &builder) {
-  SmallVector<Operation *> opToDelete;
+static void replaceUsesAndPropagateType(RewriterBase &rewriter,
+                                        Operation *oldOp, Value val) {
+  SmallVector<Operation *> opsToDelete;
   SmallVector<OpOperand *> operandsToReplace;
+
+  // Save the operand to replace / delete later (avoid iterator invalidation).
+  // TODO: can we use an early_inc iterator?
   for (OpOperand &use : oldOp->getUses()) {
+    // Non-subview ops will be replaced by `val`.
     auto subviewUse = dyn_cast<memref::SubViewOp>(use.getOwner());
     if (!subviewUse) {
-      // Save the operand to and replace outside the loop to not invalidate the
-      // iterator.
       operandsToReplace.push_back(&use);
       continue;
     }
-    builder.setInsertionPoint(subviewUse);
+
+    // `subview(old_op)` is replaced by a new `subview(val)`.
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(subviewUse);
     Type newType = memref::SubViewOp::inferRankReducedResultType(
         subviewUse.getType().getShape(), val.getType().cast<MemRefType>(),
         subviewUse.getStaticOffsets(), subviewUse.getStaticSizes(),
         subviewUse.getStaticStrides());
-    Value newSubview = builder.create<memref::SubViewOp>(
+    Value newSubview = rewriter.create<memref::SubViewOp>(
         subviewUse->getLoc(), newType.cast<MemRefType>(), val,
         subviewUse.getMixedOffsets(), subviewUse.getMixedSizes(),
         subviewUse.getMixedStrides());
-    replaceUsesAndPropagateType(subviewUse, newSubview, builder);
-    opToDelete.push_back(use.getOwner());
+
+    // Ouch recursion ... is this really necessary?
+    replaceUsesAndPropagateType(rewriter, subviewUse, newSubview);
+
+    opsToDelete.push_back(use.getOwner());
   }
-  for (OpOperand *operand : operandsToReplace)
+
+  // Perform late replacement.
+  // TODO: can we use an early_inc iterator?
+  for (OpOperand *operand : operandsToReplace) {
+    Operation *op = operand->getOwner();
+    rewriter.startRootUpdate(op);
     operand->set(val);
-  // Clean up old subview ops.
-  for (Operation *op : opToDelete)
-    op->erase();
+    rewriter.finalizeRootUpdate(op);
   }
 
-/// Helper to convert get a value from an OpFoldResult or create it at the
-/// builder insert point.
-static Value getOrCreateValue(OpFoldResult res, OpBuilder &builder,
-                              Location loc) {
-  Value value = res.dyn_cast<Value>();
-  if (value)
-    return value;
-  return builder.create<arith::ConstantIndexOp>(
-      loc, res.dyn_cast<Attribute>().cast<IntegerAttr>().getInt());
+  // Perform late op erasure.
+  // TODO: can we use an early_inc iterator?
+  for (Operation *op : opsToDelete)
+    rewriter.eraseOp(op);
 }
 
 // Transformation to do multi-buffering/array expansion to remove dependencies
 // on the temporary allocation between consecutive loop iterations.
 // Returns success if the transformation happened and failure otherwise.
 // This is not a pattern as it requires propagating the new memref type to its
 // uses and requires updating subview ops.
 FailureOr<memref::AllocOp>
-mlir::memref::multiBuffer(memref::AllocOp allocOp, unsigned multiplier,
+mlir::memref::multiBuffer(RewriterBase &rewriter, memref::AllocOp allocOp,
+                          unsigned multiBufferingFactor,
                           bool skipOverrideAnalysis) {
-  LLVM_DEBUG(DBGS() << "Try multibuffer: " << allocOp << "\n");
+  LLVM_DEBUG(DBGS() << "Start multibuffering: " << allocOp << "\n");
   DominanceInfo dom(allocOp->getParentOp());
   LoopLikeOpInterface candidateLoop;
   for (Operation *user : allocOp->getUsers()) {
     auto parentLoop = user->getParentOfType<LoopLikeOpInterface>();
     if (!parentLoop) {
-      LLVM_DEBUG(DBGS() << "Skip user: no parent loop\n");
+      if (isa<memref::DeallocOp>(user)) {
+        // Allow dealloc outside of any loop.
+        // TODO: The whole precondition function here is very brittle and will
+        // need to rethought an isolated into a cleaner analysis.
+        continue;
+      }
+      LLVM_DEBUG(DBGS() << "--no parent loop -> fail\n");
+      LLVM_DEBUG(DBGS() << "----due to user: " << *user << "\n");
       return failure();
     }
     if (!skipOverrideAnalysis) {
       /// Make sure there is no loop-carried dependency on the allocation.
       if (!overrideBuffer(user, allocOp.getResult())) {
-        LLVM_DEBUG(DBGS() << "Skip user: found loop-carried dependence\n");
+        LLVM_DEBUG(DBGS() << "--Skip user: found loop-carried dependence\n");
         continue;
       }
       // If this user doesn't dominate all the other users keep looking.
       if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {
             return !dom.dominates(user, otherUser);
           })) {
-        LLVM_DEBUG(DBGS() << "Skip user: does not dominate all other users\n");
+        LLVM_DEBUG(
+            DBGS() << "--Skip user: does not dominate all other users\n");
         continue;
       }
     } else {
       if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {
             return !isa<memref::DeallocOp>(otherUser) &&
                    !parentLoop->isProperAncestor(otherUser);
           })) {
         LLVM_DEBUG(
             DBGS()
-            << "Skip user: not all other users are in the parent loop\n");
+            << "--Skip user: not all other users are in the parent loop\n");
         continue;
       }
     }
     candidateLoop = parentLoop;
     break;
   }
+
   if (!candidateLoop) {
     LLVM_DEBUG(DBGS() << "Skip alloc: no candidate loop\n");
     return failure();
   }
+
   std::optional<Value> inductionVar = candidateLoop.getSingleInductionVar();
   std::optional<OpFoldResult> lowerBound = candidateLoop.getSingleLowerBound();
   std::optional<OpFoldResult> singleStep = candidateLoop.getSingleStep();
   if (!inductionVar || !lowerBound || !singleStep) {
     LLVM_DEBUG(DBGS() << "Skip alloc: no single iv, lb or step\n");
     return failure();
   }
 
   if (!dom.dominates(allocOp.getOperation(), candidateLoop)) {
     LLVM_DEBUG(DBGS() << "Skip alloc: does not dominate candidate loop\n");
     return failure();
   }
 
-  OpBuilder builder(candidateLoop);
-  SmallVector<int64_t, 4> newShape(1, multiplier);
-  ArrayRef<int64_t> oldShape = allocOp.getType().getShape();
-  newShape.append(oldShape.begin(), oldShape.end());
-  auto newMemref = MemRefType::get(newShape, allocOp.getType().getElementType(),
-                                   MemRefLayoutAttrInterface(),
-                                   allocOp.getType().getMemorySpace());
-  builder.setInsertionPoint(allocOp);
+  LLVM_DEBUG(DBGS() << "Start multibuffering loop: " << candidateLoop << "\n");
+
+  // 1. Construct the multi-buffered memref type.
+  ArrayRef<int64_t> originalShape = allocOp.getType().getShape();
+  SmallVector<int64_t, 4> multiBufferedShape{multiBufferingFactor};
+  llvm::append_range(multiBufferedShape, originalShape);
+  LLVM_DEBUG(DBGS() << "--original type: " << allocOp.getType() << "\n");
+  MemRefType mbMemRefType = MemRefType::Builder(allocOp.getType())
+                                .setShape(multiBufferedShape)
+                                .setLayout(MemRefLayoutAttrInterface());
+  LLVM_DEBUG(DBGS() << "--multi-buffered type: " << mbMemRefType << "\n");
+
+  // 2. Create the multi-buffered alloc.
   Location loc = allocOp->getLoc();
-  auto newAlloc = builder.create<memref::AllocOp>(loc, newMemref, ValueRange{},
-                                                  allocOp->getAttrs());
-  builder.setInsertionPoint(&candidateLoop.getLoopBody().front(),
-                            candidateLoop.getLoopBody().front().begin());
-
-  SmallVector<Value> operands = {*inductionVar};
-  AffineExpr induc = getAffineDimExpr(0, allocOp.getContext());
-  unsigned dimCount = 1;
-  auto getAffineExpr = [&](OpFoldResult e) -> AffineExpr {
-    if (std::optional<int64_t> constValue = getConstantIntValue(e)) {
-      return getAffineConstantExpr(*constValue, allocOp.getContext());
-    }
-    auto value = getOrCreateValue(e, builder, candidateLoop->getLoc());
-    operands.push_back(value);
-    return getAffineDimExpr(dimCount++, allocOp.getContext());
-  };
-  auto init = getAffineExpr(*lowerBound);
-  auto step = getAffineExpr(*singleStep);
-
-  AffineExpr expr = ((induc - init).floorDiv(step)) % multiplier;
-  auto map = AffineMap::get(dimCount, 0, expr);
-  Value bufferIndex = builder.create<AffineApplyOp>(loc, map, operands);
-  SmallVector<OpFoldResult> offsets, sizes, strides;
-  offsets.push_back(bufferIndex);
-  offsets.append(oldShape.size(), builder.getIndexAttr(0));
-  strides.assign(oldShape.size() + 1, builder.getIndexAttr(1));
-  sizes.push_back(builder.getIndexAttr(1));
-  for (int64_t size : oldShape)
-    sizes.push_back(builder.getIndexAttr(size));
-  auto dstMemref =
-      memref::SubViewOp::inferRankReducedResultType(
-          allocOp.getType().getShape(), newMemref, offsets, sizes, strides)
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(allocOp);
+  auto mbAlloc = rewriter.create<memref::AllocOp>(
+      loc, mbMemRefType, ValueRange{}, allocOp->getAttrs());
+  LLVM_DEBUG(DBGS() << "--multi-buffered alloc: " << mbAlloc << "\n");
+
+  // 3. Within the loop, build the modular leading index (i.e. each loop
+  // iteration %iv accesses slice ((%iv - %lb) / %step) % %mb_factor).
+  rewriter.setInsertionPointToStart(&candidateLoop.getLoopBody().front());
+  Value ivVal = *inductionVar;
+  Value lbVal = getValueOrCreateConstantIndexOp(rewriter, loc, *lowerBound);
+  Value stepVal = getValueOrCreateConstantIndexOp(rewriter, loc, *singleStep);
+  AffineExpr iv, lb, step;
+  bindDims(rewriter.getContext(), iv, lb, step);
+  Value bufferIndex = makeComposedAffineApply(
+      rewriter, loc, ((iv - lb).floorDiv(step)) % multiBufferingFactor,
+      {ivVal, lbVal, stepVal});
+  LLVM_DEBUG(DBGS() << "--multi-buffered indexing: " << bufferIndex << "\n");
+
+  // 4. Build the subview accessing the particular slice, taking modular
+  // rotation into account.
+  int64_t mbMemRefTypeRank = mbMemRefType.getRank();
+  IntegerAttr zero = rewriter.getIndexAttr(0);
+  IntegerAttr one = rewriter.getIndexAttr(1);
+  SmallVector<OpFoldResult> offsets(mbMemRefTypeRank, zero);
+  SmallVector<OpFoldResult> sizes(mbMemRefTypeRank, one);
+  SmallVector<OpFoldResult> strides(mbMemRefTypeRank, one);
+  // Offset is [bufferIndex, 0 ... 0 ].
+  offsets.front() = bufferIndex;
+  // Sizes is [1, original_size_0 ... original_size_n ].
+  for (int64_t i = 0, e = originalShape.size(); i != e; ++i)
+    sizes[1 + i] = rewriter.getIndexAttr(originalShape[i]);
+  // Strides is [1, 1 ... 1 ].
+  auto dstMemref = memref::SubViewOp::inferRankReducedResultType(
+                       originalShape, mbMemRefType, offsets, sizes, strides)
                        .cast<MemRefType>();
-  Value subview = builder.create<memref::SubViewOp>(loc, dstMemref, newAlloc,
+  Value subview = rewriter.create<memref::SubViewOp>(loc, dstMemref, mbAlloc,
                                                      offsets, sizes, strides);
-  replaceUsesAndPropagateType(allocOp, subview, builder);
-  allocOp.erase();
-  return newAlloc;
+  LLVM_DEBUG(DBGS() << "--multi-buffered slice: " << subview << "\n");
+
+  // 5. Due to the recursive nature of replaceUsesAndPropagateType , we need to
+  // handle dealloc uses separately..
+  for (OpOperand &use : llvm::make_early_inc_range(allocOp->getUses())) {
+    auto deallocOp = dyn_cast<memref::DeallocOp>(use.getOwner());
+    if (!deallocOp)
+      continue;
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(deallocOp);
+    auto newDeallocOp =
+        rewriter.create<memref::DeallocOp>(deallocOp->getLoc(), mbAlloc);
+    (void)newDeallocOp;
+    LLVM_DEBUG(DBGS() << "----Created dealloc: " << newDeallocOp << "\n");
+    rewriter.eraseOp(deallocOp);
+  }
+
+  // 6. RAUW with the particular slice, taking modular rotation into account.
+  replaceUsesAndPropagateType(rewriter, allocOp, subview);
+
+  // 7. Finally, erase the old allocOp.
+  rewriter.eraseOp(allocOp);
+
+  return mbAlloc;
+}
+
+FailureOr<memref::AllocOp>
+mlir::memref::multiBuffer(memref::AllocOp allocOp,
+                          unsigned multiBufferingFactor,
+                          bool skipOverrideAnalysis) {
+  IRRewriter rewriter(allocOp->getContext());
+  return multiBuffer(rewriter, allocOp, multiBufferingFactor,
+                     skipOverrideAnalysis);
 }

mlir/lib/Dialect/SCF/IR/SCF.cpp

LogicalResult matchAndRewrite(ForOp op,
                                 PatternRewriter &rewriter) const override {
     for (auto it : llvm::zip(op.getIterOpOperands(), op.getResults())) {
       OpOperand &iterOpOperand = std::get<0>(it);
       auto incomingCast = iterOpOperand.get().getDefiningOp<tensor::CastOp>();
       if (!incomingCast)
         continue;
       // If the dest type of the cast does not preserve static information in
       // the source type.
-      if (!tensor::preservesStaticInformation(incomingCast.getDest().getType(),
-                                         incomingCast.getSource().getType()))
-          continue;
+      if (!tensor::preservesStaticInformation(
+              incomingCast.getDest().getType(),
       if (!std::get<1>(it).hasOneUse())
         continue;
 
       // Create a new ForOp with that iter operand replaced.
       auto newForOp = replaceTensorCastForOpIterArg(rewriter, iterOpOperand,
                                                     incomingCast.getSource());
 
       // Insert outgoing cast and use it to replace the corresponding result.

mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp

 //===- InferTypeOpImpl.cpp - InferType Interface external models *- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
 /// Compute a map that for a given dimension of the expanded type gives the
struct ReifyExpandOrCollapseShapeOp
   LogicalResult
   reifyResultShapes(Operation *op, OpBuilder &b,
                     ReifiedRankedShapedTypeDims &reifiedReturnShapes) const {
     auto loc = op->getLoc();
     auto reshapeOp = cast<OpTy>(op);
     auto resultShape = getReshapeOutputShapeFromInputShape(
         b, loc, reshapeOp.getSrc(), reshapeOp.getResultType().getShape(),
         reshapeOp.getReassociationMaps());
-    reifiedReturnShapes.push_back(getAsValues(b, loc, resultShape));
+    reifiedReturnShapes.push_back(
+        getValueOrCreateConstantIndexOp(b, loc, resultShape));
     return success();
   }
 };
 
 namespace {
 
 struct ReifyPadOp
     : public ReifyRankedShapedTypeOpInterface::ExternalModel<ReifyPadOp,

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

bool isEqualConstantIntOrValueArray(ArrayRef<OpFoldResult> ofrs1,
   if (ofrs1.size() != ofrs2.size())
     return false;
   for (auto [ofr1, ofr2] : llvm::zip_equal(ofrs1, ofrs2))
     if (!isEqualConstantIntOrValue(ofr1, ofr2))
       return false;
   return true;
 }
 
-/// Helper function to convert a vector of `OpFoldResult`s into a vector of
-/// `Value`s. For each `OpFoldResult` in `valueOrAttrVec` return the fold result
-/// if it casts to  a `Value` or create an index-type constant if it casts to
-/// `IntegerAttr`. No other attribute types are supported.
-SmallVector<Value> getAsValues(OpBuilder &b, Location loc,
-                               ArrayRef<OpFoldResult> valueOrAttrVec) {
-  return llvm::to_vector<4>(
-      llvm::map_range(valueOrAttrVec, [&](OpFoldResult value) -> Value {
-        return getValueOrCreateConstantIndexOp(b, loc, value);
-      }));
-}
-
 /// Return a vector of OpFoldResults with the same size a staticValues, but all
 /// elements for which ShapedType::isDynamic is true, will be replaced by
 /// dynamicValues.
 SmallVector<OpFoldResult> getMixedValues(ArrayRef<int64_t> staticValues,
                                          ValueRange dynamicValues, Builder &b) {
   SmallVector<OpFoldResult> res;
   res.reserve(staticValues.size());
   unsigned numDynamic = 0;
   unsigned count = static_cast<unsigned>(staticValues.size());
   for (unsigned idx = 0; idx < count; ++idx) {
     int64_t value = staticValues[idx];
     res.push_back(ShapedType::isDynamic(value)
                       ? OpFoldResult{dynamicValues[numDynamic++]}

mlir/test/Dialect/MemRef/transform-ops.mlir

 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["memref.alloc"]} in %arg1 : (!pdl.operation) -> !transform.op<"memref.alloc">
   %1 = transform.memref.multibuffer %0 {factor = 2 : i64, skip_analysis} : (!transform.op<"memref.alloc">) -> !pdl.operation
   // Verify that the returned handle is usable.
   transform.test_print_remark_at_operand %1, "transformed" : !pdl.operation
 }
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> ((d0 floordiv 4) mod 2)>
+
+// CHECK-LABEL: func @multi_buffer_dealloc
+func.func @multi_buffer_dealloc(%in: memref<16xf32>) {
+  // CHECK: %[[A:.*]] = memref.alloc() : memref<2x4xf32>
+  // expected-remark @below {{transformed}}
+  %tmp = memref.alloc() : memref<4xf32>
+
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: %[[C4:.*]] = arith.constant 4 : index
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %c16 = arith.constant 16 : index
+
+  // CHECK: scf.for %[[IV:.*]] = %[[C0]]
+  scf.for %i0 = %c0 to %c16 step %c4 {
+  // CHECK: %[[I:.*]] = affine.apply #[[$MAP0]](%[[IV]])
+  // CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0] [1, 4] [1, 1] : memref<2x4xf32> to memref<4xf32, strided<[1], offset: ?>>
+    "some_write_read"(%tmp) : (memref<4xf32>) ->()
+  }
+
+  // CHECK-NOT: memref.dealloc {{.*}} : memref<4xf32>
+  // CHECK: memref.dealloc %[[A]] : memref<2x4xf32>
+  memref.dealloc %tmp : memref<4xf32>
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["memref.alloc"]} in %arg1 : (!pdl.operation) -> !transform.op<"memref.alloc">
+  %1 = transform.memref.multibuffer %0 {factor = 2 : i64, skip_analysis} : (!transform.op<"memref.alloc">) -> !pdl.operation
+  // Verify that the returned handle is usable.
+  transform.test_print_remark_at_operand %1, "transformed" : !pdl.operation
+}