[mlir][memref] Add transformation to do loop multi-buffering

[ThomasRaoux]

Summary

This transformation is useful to break dependency between consecutive loop iterations by increasing the size of a temporary buffer.
This is usually combined with heavy software pipelining.
Started a discourse thread in case we need this needs further discussions: https://discourse.llvm.org/t/loop-double-buffering-multi-buffering/59979

[ThomasRaoux]

clang-format

[bondhugula]

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

68–70 ↗

(On Diff #407388)

This should in general have nothing to do with an scf::for op -- you aren't transforming the loop in any way. Just use LoopLikeOpInterface.

99–104 ↗

(On Diff #407388)

Avoid generating lower-level arithmetic ops. Instead, just generate a modulo using affine.apply.

[ThomasRaoux]

Use affine.apply for arithmetic

[ThomasRaoux]

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

68–70 ↗

(On Diff #407388)

I need to access the single induction variable as well as the lower bound and step. This is not possible with LoopLikeInterface as it cannot be applied to any loops.

99–104 ↗

(On Diff #407388)

Good point, I didn't realize this was a valid affine map. Changed it.

[bondhugula]

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

77

Please document the return value.

77

This is a utility and should be returning LogicalResult instead of bool I think.

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

65 ↗

(On Diff #407796)

Please document the return value.

66 ↗

(On Diff #407796)

You don't need a triple doc comment here. Just the more detailed version on the declaration will suffice.

68–70 ↗

(On Diff #407388)

This should be pretty straightforward: you can just have the op interface method returning the IV (say getInductionVar()) return None for multi-dim loops (like for scf.parallel) and have this method fail in such a case just like you already do at line 87. It's going to be as simple and strictly more general. For the other two things you mention (lower bound and step), that's straightforward as well but the interface method will have to materialize some trivial IR (for eg. when it's a constant as in the case of affine.for):

// For IV.
Optional<Value> getSingleInductionVar();

1. lower bound

Optional<Value> getOrCreateLowerBound(OpBuilder &b);
Just return the lower bound SSA value (or materialize the step in case of affine.for)

2. step

Optional<Value> getOrCreateStep(OpBuilder &b);

With the above, it would be strictly more general than restricting it to scf::ForOp and as simple. And in your case of interest here, you'll get output IR identical to what you have in the test cases below.

mlir/test/Dialect/MemRef/multibuffering.mlir

67–68 ↗

(On Diff #407796)

You don't need the %{{.*}} = part.

[bondhugula]

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

77

The utilities are almost always named imperatively instead of a gerund. multiBuffering -> multiBuffer?

[ThomasRaoux]

Address review feedback

[ThomasRaoux]

I switched to using loop interface and added a test for affine.for op. Please take another look.

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

68–70 ↗

(On Diff #407388)

I added the interface and moved the code to use it. This means it now works on affine.for op. I added a test for it.

[bondhugula]

The comments that I had have been addressed. Dismissing change request.

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

23

Define using dyn_cast to avoid a double check in assert mode.

[bondhugula]

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

29

You are also collecting and erasing ops here.

36–53

Invert the conditional?

if (!....) {
   push_back
   continue;
}

...

70

Unchecked use of dyn_cast result. This should just be a single cast?

130–131

strides: append -> assign (since this is empty at this stage).

[ThomasRaoux]

Address review comments

[ThomasRaoux]

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

70

OpFoldResult/PointerUnion don't have cast they only have dyn_cast, I see it used like that in the rest of the code

[mravishankar]

LGTM!

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

89

Nit: if this user doesnt dominate...

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

def AffineApplyOp : Affine_Op<"apply", [NoSideEffect]> {
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def AffineForOp : Affine_Op<"for",
     [ImplicitAffineTerminator, RecursiveSideEffects,
-     DeclareOpInterfaceMethods<LoopLikeOpInterface>]> {
+     DeclareOpInterfaceMethods<LoopLikeOpInterface,
+     ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep"]>]> {
   let summary = "for operation";
   let description = [{
     Syntax:
 
     ```
     operation   ::= `affine.for` ssa-id `=` lower-bound `to` upper-bound
                     (`step` integer-literal)? `{` op* `}`
 

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

 namespace tensor {
 class TensorDialect;
 } // namespace tensor
 namespace vector {
 class VectorDialect;
 } // namespace vector
 
 namespace memref {
-
+class AllocOp;
 //===----------------------------------------------------------------------===//
 // Patterns
 //===----------------------------------------------------------------------===//
 
 /// Collects a set of patterns to rewrite ops within the memref dialect.
 void populateExpandOpsPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns for folding memref.subview ops into consumer load/store ops
 /// into `patterns`.
 void populateFoldSubViewOpPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns that resolve `memref.dim` operations with values that are
 /// defined by operations that implement the
 /// `ReifyRankedShapeTypeShapeOpInterface`, in terms of shapes of its input
 /// operands.
 void populateResolveRankedShapeTypeResultDimsPatterns(
     RewritePatternSet &patterns);
 
 /// Appends patterns that resolve `memref.dim` operations with values that are
 /// defined by operations that implement the `InferShapedTypeOpInterface`, in
 /// terms of shapes of its input operands.
 void populateResolveShapedTypeResultDimsPatterns(RewritePatternSet &patterns);
 
+/// Transformation to do multi-buffering/array expansion to remove dependencies
+/// on the temporary allocation between consecutive loop iterations.
+/// It return success if the allocation was multi-buffered and returns failure()
+/// otherwise.
+/// Example:
+/// ```
+/// %0 = memref.alloc() : memref<4x128xf32>
+/// scf.for %iv = %c1 to %c1024 step %c3 {
+///   memref.copy %1, %0 : memref<4x128xf32> to memref<4x128xf32>
+///   "some_use"(%0) : (memref<4x128xf32>) -> ()
+/// }
+/// ```
+/// into:
+/// ```
+/// %0 = memref.alloc() : memref<5x4x128xf32>
+/// scf.for %iv = %c1 to %c1024 step %c3 {
+///   %s = arith.subi %iv, %c1 : index
+///   %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, #map0>
+///   memref.copy %1, %sv : memref<4x128xf32> to memref<4x128xf32, #map0>
+///   "some_use"(%sv) : (memref<4x128xf32, $map0>) -> ()
+/// }
+/// ```
+LogicalResult multiBuffer(memref::AllocOp allocOp, unsigned multiplier);
+
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 /// Creates an instance of the ExpandOps pass that legalizes memref dialect ops
 /// to be convertible to LLVM. For example, `memref.reshape` gets converted to
 /// `memref_reinterpret_cast`.
 std::unique_ptr<Pass> createExpandOpsPass();

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

def ExecuteRegionOp : SCF_Op<"execute_region"> {
 
   let hasCanonicalizer = 1;
 
   let hasFolder = 0;
   let hasVerifier = 1;
 }
 
 def ForOp : SCF_Op<"for",
-      [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+      [DeclareOpInterfaceMethods<LoopLikeOpInterface,
+       ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep"]>,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveSideEffects]> {
   let summary = "for operation";
   let description = [{
     The "scf.for" operation represents a loop taking 3 SSA value as operands
     that represent the lower bound, upper bound and step respectively. The
     operation defines an SSA value for its induction variable. It has one

mlir/include/mlir/Interfaces/LoopLikeInterface.td

InterfaceMethod<[{
       "::mlir::Region &", "getLoopBody"
     >,
     InterfaceMethod<[{
         Moves the given vector of operations out of the loop. The vector is
         sorted topologically.
       }],
       "::mlir::LogicalResult", "moveOutOfLoop", (ins "::mlir::ArrayRef<::mlir::Operation *>":$ops)
     >,
+    InterfaceMethod<[{
+        If there is a single induction variable return it, otherwise return
+        llvm::None.
+      }],
+      /*retTy=*/"::mlir::Optional<::mlir::Value>",
+      /*methodName=*/"getSingleInductionVar",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return llvm::None;
+      }]
+    >,
+    InterfaceMethod<[{
+        Return the single lower bound value or attribute if it exist, otherwise
+        return llvm::None.
+      }],
+      /*retTy=*/"::mlir::Optional<::mlir::OpFoldResult>",
+      /*methodName=*/"getSingleLowerBound",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return llvm::None;
+      }]
+    >,
+    InterfaceMethod<[{
+        Return the single step value or attribute if it exist, otherwise
+        return llvm::None.
+      }],
+      /*retTy=*/"::mlir::Optional<::mlir::OpFoldResult>",
+      /*methodName=*/"getSingleStep",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return llvm::None;
+      }]
+    >,
   ];
 }
 
 #endif // MLIR_INTERFACES_LOOPLIKEINTERFACE

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 }
 
 Region &AffineForOp::getLoopBody() { return region(); }
 
 bool AffineForOp::isDefinedOutsideOfLoop(Value value) {
   return !region().isAncestor(value.getParentRegion());
 }
 
+Optional<Value> AffineForOp::getSingleInductionVar() {
+  return getInductionVar();
+}
+
+Optional<OpFoldResult> AffineForOp::getSingleLowerBound() {
+  if (!hasConstantLowerBound())
+    return llvm::None;
+  OpBuilder b(getContext());
+  return OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()));
+}
+
+Optional<OpFoldResult> AffineForOp::getSingleStep() {
+  OpBuilder b(getContext());
+  return OpFoldResult(b.getI64IntegerAttr(getStep()));
+}
+
 LogicalResult AffineForOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
   for (auto *op : ops)
     op->moveBefore(*this);
   return success();
 }
 
 /// Returns true if the provided value is the induction variable of a
 /// AffineForOp.

mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefTransforms
   ComposeSubView.cpp
   ExpandOps.cpp
   FoldSubViewOps.cpp
+  MultiBuffer.cpp
   NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAffineUtils
   MLIRArithmetic
   MLIRInferTypeOpInterface
+  MLIRLoopLikeInterface
   MLIRMemRef
   MLIRPass
   MLIRStandard
   MLIRTensor
   MLIRTransforms
   MLIRVector
 )
 

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/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+
+using namespace mlir;
+
+/// Return true if the op fully overwrite the given `buffer` value.
+static bool overrideBuffer(Operation *op, Value buffer) {
+  return isa<memref::CopyOp>(op) && cast<memref::CopyOp>(op).target() == 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.
+static void replaceUsesAndPropagateType(Operation *oldOp, Value val,
+                                        OpBuilder &builder) {
+  SmallVector<Operation *> opToDelete;
+  SmallVector<OpOperand *> operandsToReplace;
+  for (OpOperand &use : oldOp->getUses()) {
+    auto subviewUse = dyn_cast<memref::SubViewOp>(use.getOwner());
+    if (subviewUse) {
+      builder.setInsertionPoint(subviewUse);
+      Type newType = memref::SubViewOp::inferRankReducedResultType(
+          subviewUse.getType().getRank(), val.getType().cast<MemRefType>(),
+          extractFromI64ArrayAttr(subviewUse.static_offsets()),
+          extractFromI64ArrayAttr(subviewUse.static_sizes()),
+          extractFromI64ArrayAttr(subviewUse.static_strides()));
+      Value newSubview = builder.create<memref::SubViewOp>(
+          subviewUse->getLoc(), newType.cast<MemRefType>(), val,
+          subviewUse.getMixedOffsets(), subviewUse.getMixedSizes(),
+          subviewUse.getMixedStrides());
+      replaceUsesAndPropagateType(subviewUse, newSubview, builder);
+      opToDelete.push_back(use.getOwner());
+      continue;
+    }
+    // Save the operand to and replace outside the loop to not invalidate the
+    // iterator.
+    operandsToReplace.push_back(&use);
+  }
+  for (OpOperand *operand : operandsToReplace)
+    operand->set(val);
+  // Clean up old subview ops.
+  for (Operation *op : opToDelete)
+    op->erase();
+}
+
+/// 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());
+}
+
+// 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.
+LogicalResult mlir::memref::multiBuffer(memref::AllocOp allocOp,
+                                        unsigned multiplier) {
+  DominanceInfo dom(allocOp->getParentOp());
+  LoopLikeOpInterface candidateLoop;
+  for (Operation *user : allocOp->getUsers()) {
+    auto parentLoop = user->getParentOfType<LoopLikeOpInterface>();
+    if (!parentLoop)
+      return failure();
+    /// Make sure there is no loop carried dependency on the allocation.
+    if (!overrideBuffer(user, allocOp.getResult()))
+      continue;
+    // If the user doesn't dominate all the other users keep looking.
+    if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {
+          return otherUser->getNumRegions() != 0 ||
+                 !dom.dominates(user, otherUser);
+        }))
+      continue;
+    candidateLoop = parentLoop;
+    break;
+  }
+  if (!candidateLoop)
+    return failure();
+  llvm::Optional<Value> inductionVar = candidateLoop.getSingleInductionVar();
+  llvm::Optional<OpFoldResult> lowerBound = candidateLoop.getSingleLowerBound();
+  llvm::Optional<OpFoldResult> singleStep = candidateLoop.getSingleStep();
+  if (!inductionVar || !lowerBound || !singleStep)
+    return failure();
+  OpBuilder builder(candidateLoop);
+  Value stepValue =
+      getOrCreateValue(*singleStep, builder, candidateLoop->getLoc());
+  Value lowerBoundValue =
+      getOrCreateValue(*lowerBound, builder, candidateLoop->getLoc());
+  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);
+  Location loc = allocOp->getLoc();
+  auto newAlloc = builder.create<memref::AllocOp>(loc, newMemref);
+  builder.setInsertionPoint(&candidateLoop.getLoopBody().front(),
+                            candidateLoop.getLoopBody().front().begin());
+  AffineExpr induc = getAffineDimExpr(0, allocOp.getContext());
+  AffineExpr init = getAffineDimExpr(1, allocOp.getContext());
+  AffineExpr step = getAffineDimExpr(2, allocOp.getContext());
+  AffineExpr expr = ((induc - init).floorDiv(step)) % multiplier;
+  auto map = AffineMap::get(3, 0, expr);
+  std::array<Value, 3> operands = {*inductionVar, lowerBoundValue, stepValue};
+  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.append(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().getRank(), newMemref, offsets, sizes, strides)
+          .cast<MemRefType>();
+  Value subview = builder.create<memref::SubViewOp>(loc, dstMemref, newAlloc,
+                                                    offsets, sizes, strides);
+  replaceUsesAndPropagateType(allocOp, subview, builder);
+  allocOp.erase();
+  return success();
+}

mlir/lib/Dialect/SCF/SCF.cpp

if (std::get<1>(e).getType() != std::get<2>(e).getType())
                            << "th iter region arg and defined value";
 
     i++;
   }
 
   return RegionBranchOpInterface::verifyTypes(*this);
 }
 
+Optional<Value> ForOp::getSingleInductionVar() { return getInductionVar(); }
+
+Optional<OpFoldResult> ForOp::getSingleLowerBound() {
+  return OpFoldResult(getLowerBound());
+}
+
+Optional<OpFoldResult> ForOp::getSingleStep() {
+  return OpFoldResult(getStep());
+}
+
 /// Prints the initialization list in the form of
 ///   <prefix>(%inner = %outer, %inner2 = %outer2, <...>)
 /// where 'inner' values are assumed to be region arguments and 'outer' values
 /// are regular SSA values.
 static void printInitializationList(OpAsmPrinter &p,
                                     Block::BlockArgListType blocksArgs,
                                     ValueRange initializers,
                                     StringRef prefix = "") {

mlir/test/Dialect/MemRef/multibuffer.mlir

+// RUN: mlir-opt %s -allow-unregistered-dialect -test-multi-buffering=multiplier=5 -cse -split-input-file | FileCheck %s
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 128 + s0 + d1)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (((d0 - d1) floordiv d2) mod 5)>
+
+// CHECK-LABEL: func @multi_buffer
+func @multi_buffer(%a: memref<1024x1024xf32>) {
+// CHECK-DAG: %[[A:.*]] = memref.alloc() : memref<5x4x128xf32>
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+  %0 = memref.alloc() : memref<4x128xf32>
+  %c1024 = arith.constant 1024 : index
+  %c1 = arith.constant 1 : index
+  %c3 = arith.constant 3 : index
+// CHECK: scf.for %[[IV:.*]] = %[[C1]]
+  scf.for %arg2 = %c1 to %c1024 step %c3 {
+// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]], %[[C1]], %[[C3]])
+// CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0, 0] [1, 4, 128] [1, 1, 1] : memref<5x4x128xf32> to memref<4x128xf32, #[[$MAP0]]>
+   %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
+    memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
+// CHECK: memref.copy %{{.*}}, %[[SV]] : memref<4x128xf32, #{{.*}}> to memref<4x128xf32, #[[$MAP0]]>
+   memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
+// CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, #[[$MAP0]]>) -> ()
+    "some_use"(%0) : (memref<4x128xf32>) -> ()
+// CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, #[[$MAP0]]>) -> ()
+   "some_use"(%0) : (memref<4x128xf32>) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @multi_buffer_affine
+func @multi_buffer_affine(%a: memref<1024x1024xf32>) {
+// CHECK-DAG: %[[A:.*]] = memref.alloc() : memref<5x4x128xf32>
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+  %0 = memref.alloc() : memref<4x128xf32>
+  %c1024 = arith.constant 1024 : index
+  %c1 = arith.constant 1 : index
+  %c3 = arith.constant 3 : index
+// CHECK: affine.for %[[IV:.*]] = 1
+  affine.for %arg2 = 1 to 1024 step 3 {
+// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]], %[[C1]], %[[C3]])
+// CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0, 0] [1, 4, 128] [1, 1, 1] : memref<5x4x128xf32> to memref<4x128xf32, #[[$MAP0]]>
+   %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
+    memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
+// CHECK: memref.copy %{{.*}}, %[[SV]] : memref<4x128xf32, #{{.*}}> to memref<4x128xf32, #[[$MAP0]]>
+   memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
+// CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, #[[$MAP0]]>) -> ()
+    "some_use"(%0) : (memref<4x128xf32>) -> ()
+// CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, #[[$MAP0]]>) -> ()
+   "some_use"(%0) : (memref<4x128xf32>) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 128 + s0 + d1)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (((d0 - d1) floordiv d2) mod 5)>
+
+// CHECK-LABEL: func @multi_buffer_subview_use
+func @multi_buffer_subview_use(%a: memref<1024x1024xf32>) {
+// CHECK-DAG: %[[A:.*]] = memref.alloc() : memref<5x4x128xf32>
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+  %0 = memref.alloc() : memref<4x128xf32>
+  %c1024 = arith.constant 1024 : index
+  %c1 = arith.constant 1 : index
+  %c3 = arith.constant 3 : index
+// CHECK: scf.for %[[IV:.*]] = %[[C1]]
+  scf.for %arg2 = %c1 to %c1024 step %c3 {
+// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]], %[[C1]], %[[C3]])
+// CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0, 0] [1, 4, 128] [1, 1, 1] : memref<5x4x128xf32> to memref<4x128xf32, #[[$MAP0]]>
+   %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
+    memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
+// CHECK: memref.copy %{{.*}}, %[[SV]] : memref<4x128xf32, #{{.*}}> to memref<4x128xf32, #[[$MAP0]]>
+   memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
+// CHECK: %[[SV1:.*]] = memref.subview %[[SV]][0, 1] [4, 127] [1, 1] : memref<4x128xf32, #[[$MAP0]]> to memref<4x127xf32, #[[$MAP0]]>
+   %s = memref.subview %0[0, 1] [4, 127] [1, 1] :
+      memref<4x128xf32> to memref<4x127xf32, affine_map<(d0, d1) -> (d0 * 128 + d1 + 1)>>
+// CHECK: "some_use"(%[[SV1]]) : (memref<4x127xf32, #[[$MAP0]]>) -> ()
+   "some_use"(%s) : (memref<4x127xf32, affine_map<(d0, d1) -> (d0 * 128 + d1 + 1)>>) -> ()
+// CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, #[[$MAP0]]>) -> ()
+   "some_use"(%0) : (memref<4x128xf32>) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @multi_buffer_negative
+func @multi_buffer_negative(%a: memref<1024x1024xf32>) {
+// CHECK-NOT: %{{.*}} = memref.alloc() : memref<5x4x128xf32>
+//     CHECK: %{{.*}} = memref.alloc() : memref<4x128xf32>
+  %0 = memref.alloc() : memref<4x128xf32>
+  %c1024 = arith.constant 1024 : index
+  %c0 = arith.constant 0 : index
+  %c3 = arith.constant 3 : index
+  scf.for %arg2 = %c0 to %c1024 step %c3 {
+   "blocking_use"(%0) : (memref<4x128xf32>) -> ()
+   %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
+    memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
+   memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
+   "some_use"(%0) : (memref<4x128xf32>) -> ()
+  }
+  return
+}
+

mlir/test/lib/Dialect/MemRef/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRMemRefTestPasses
   TestComposeSubView.cpp
+  TestMultiBuffer.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
+  MLIRMemRef
   MLIRMemRefTransforms
   MLIRTestDialect
   )
 
 target_include_directories(MLIRMemRefTestPasses
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../Test
   ${CMAKE_CURRENT_BINARY_DIR}/../Test
   )

mlir/test/lib/Dialect/MemRef/TestMultiBuffer.cpp

+//===- TestComposeSubView.cpp - Test composed subviews --------------------===//
+//
+// 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/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+struct TestMultiBufferingPass
+    : public PassWrapper<TestMultiBufferingPass, OperationPass<FuncOp>> {
+  TestMultiBufferingPass() = default;
+  TestMultiBufferingPass(const TestMultiBufferingPass &pass)
+      : PassWrapper(pass) {}
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect>();
+  }
+  StringRef getArgument() const final { return "test-multi-buffering"; }
+  StringRef getDescription() const final {
+    return "Test multi buffering transformation";
+  }
+  void runOnOperation() override;
+  Option<unsigned> multiplier{
+      *this, "multiplier",
+      llvm::cl::desc(
+          "Decide how many versions of the buffer should be created,"),
+      llvm::cl::init(2)};
+};
+
+void TestMultiBufferingPass::runOnOperation() {
+  SmallVector<memref::AllocOp> allocs;
+  getOperation().walk(
+      [&allocs](memref::AllocOp alloc) { allocs.push_back(alloc); });
+  for (memref::AllocOp alloc : allocs)
+    (void)multiBuffer(alloc, multiplier);
+}
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestMultiBuffering() {
+  PassRegistration<TestMultiBufferingPass>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestGpuSerializeToHsacoPass();
 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestComposeSubView();
+void registerTestMultiBuffering();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgControlFuseByExpansion();
 void registerTestLinalgDistribution();

 #endif
   mlir::test::registerTestComprehensiveFunctionBufferize();
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestExpandTanhPass();
   mlir::test::registerTestComposeSubView();
+  mlir::test::registerTestMultiBuffering();
   mlir::test::registerTestGpuParallelLoopMappingPass();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLinalgCodegenStrategy();
   mlir::test::registerTestLinalgControlFuseByExpansion();
   mlir::test::registerTestLinalgDistribution();
   mlir::test::registerTestLinalgElementwiseFusion();

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     includes = ["include"],
     deps = [
         ":Affine",
         ":AffineUtils",
         ":ArithmeticDialect",
         ":ArithmeticTransforms",
         ":IR",
         ":InferTypeOpInterface",
+        ":LoopLikeInterface",
         ":MemRefDialect",
         ":MemRefPassIncGen",
         ":Pass",
         ":StandardOps",
         ":Support",
         ":TensorDialect",
         ":Transforms",
         ":VectorOps",

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

 cc_library(
     name = "TestMemRef",
     srcs = glob(["lib/Dialect/MemRef/*.cpp"]),
     defines = ["MLIR_CUDA_CONVERSIONS_ENABLED"],
     includes = ["lib/Dialect/Test"],
     deps = [
         ":TestDialect",
         "//mlir:Affine",
+        "//mlir:MemRefDialect",
         "//mlir:MemRefTransforms",
         "//mlir:Pass",
         "//mlir:Transforms",
     ],
 )
 
 cc_library(
     name = "TestSCF",