[mlir] add option to multi-buffering

[ThomasRaoux]

Summary

Allow user to apply multi-buffering transformation for cases where proving
that there is no loop carried dependency is not trivial. In this case user needs
to ensure that the data are written and read in the same iteration otherwise the
result is incorrect.

[nicolasvasilache]

mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

44

Could we rename this to "skip_analysis" before landing?
This is not an op semantics (verification) concern.

[ThomasRaoux]

rebase and address review comment

[nicolasvasilache]

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

82

typo: that

82

grammo: overridden

mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

def MemRefMultiBufferOp : Op<Transform_Dialect, "memref.multibuffer",
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let summary = "Multibuffers an allocation";
   let description = [{
      Transformation to do multi-buffering/array expansion to remove
      dependencies on the temporary allocation between consecutive loop
      iterations. This transform expands the size of an allocation by
      a given multiplicative factor and fixes up any users of the
      multibuffered allocation.
+     If skip analysis is not set the transformation will only apply
+     if it can prove that there is no data being carried across loop
+     iterations.
 
      #### Return modes
 
      This operation returns the new allocation if multi-buffering
      succeeds, and failure otherwise.
   }];
 
   let arguments =
       (ins Transform_MemRefAllocOp:$target,
-           ConfinedAttr<I64Attr, [IntPositive]>:$factor);
+           ConfinedAttr<I64Attr, [IntPositive]>:$factor,
+           UnitAttr:$skip_analysis);
 
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
 }
 
 #endif // MEMREF_TRANSFORM_OPS

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

 /// ops over narrowe integer types.
 void populateMemRefWideIntEmulationConversions(
     arith::WideIntEmulationConverter &typeConverter);
 
 /// Transformation to do multi-buffering/array expansion to remove dependencies
 /// on the temporary allocation between consecutive loop iterations.
 /// It returns the new allocation if the original allocation was multi-buffered
 /// and returns failure() otherwise.
-/// Example:
+/// When `skipOverrideAnalysis`, the pass will apply the transformation
+/// without checking thwt the buffer is overrided at the beginning of each
+/// iteration. This implies that user knows that there is no data carried across
+/// loop iterations. 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, strided<[128, 1], offset: ?>>
 ///   memref.copy %1, %sv : memref<4x128xf32> to memref<4x128xf32, strided<...>>
 ///   "some_use"(%sv) : (memref<4x128xf32, strided<...>) -> ()
 /// }
 /// ```
 FailureOr<memref::AllocOp> multiBuffer(memref::AllocOp allocOp,
-                                       unsigned multiplier);
+                                       unsigned multiplier,
+                                       bool skipOverrideAnalysis = false);
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
 

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

for (Operation *user : target->getUsers()) {
       if (!loop) {
         canApplyMultiBuffer = false;
         break;
       }
     }
     if (!canApplyMultiBuffer)
       continue;
 
-    auto newBuffer = memref::multiBuffer(target, getFactor());
+    auto newBuffer =
+        memref::multiBuffer(target, getFactor(), getSkipAnalysis());
     if (failed(newBuffer))
       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

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.
-FailureOr<memref::AllocOp> mlir::memref::multiBuffer(memref::AllocOp allocOp,
-                                                     unsigned multiplier) {
+FailureOr<memref::AllocOp>
+mlir::memref::multiBuffer(memref::AllocOp allocOp, unsigned multiplier,
+                          bool skipOverrideAnalysis) {
   LLVM_DEBUG(DBGS() << "Try multibuffer: " << 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");
       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");
         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");
         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");
+        continue;
+      }
+    }
     candidateLoop = parentLoop;
     break;
   }
   if (!candidateLoop) {
     LLVM_DEBUG(DBGS() << "Skip alloc: no candidate loop\n");
     return failure();
   }
   std::optional<Value> inductionVar = candidateLoop.getSingleInductionVar();

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} : (!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
+func.func @multi_buffer_no_analysis(%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>) ->()
+  }
+  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
+}