[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

41

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",
      TransformOpInterface, TransformEachOpTrait]> {
   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 verification 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 PDL_Operation:$target,
-           ConfinedAttr<I64Attr, [IntPositive]>:$factor);
+           ConfinedAttr<I64Attr, [IntPositive]>:$factor,
+           UnitAttr:$skip_verification);
 
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         memref::AllocOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #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 `skipOverrideVerification`, 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 skipOverrideVerification = false);
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
 

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

 
 //===----------------------------------------------------------------------===//
 // MemRefMultiBufferOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MemRefMultiBufferOp::applyToOne(
     memref::AllocOp target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
-  auto newBuffer = memref::multiBuffer(target, getFactor());
+  auto newBuffer =
+      memref::multiBuffer(target, getFactor(), getSkipVerification());
   if (failed(newBuffer)) {
     Diagnostic diag(target->getLoc(), DiagnosticSeverity::Note);
     diag << "op failed to multibuffer";
     return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
   }
 
   results.push_back(*newBuffer);
   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 skipOverrideVerification) {
   DominanceInfo dom(allocOp->getParentOp());
   LoopLikeOpInterface candidateLoop;
   for (Operation *user : allocOp->getUsers()) {
     auto parentLoop = user->getParentOfType<LoopLikeOpInterface>();
     if (!parentLoop)
       return failure();
+    if (!skipOverrideVerification) {
       /// Make sure there is no loop carried dependency on the allocation.
       if (!overrideBuffer(user, allocOp.getResult()))
         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);
+            return !isa<memref::DeallocOp>(otherUser) &&
+                   !dom.dominates(user, otherUser);
           }))
         continue;
+    } else {
+      if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {
+            return !isa<memref::DeallocOp>(otherUser) &&
+                   !parentLoop->isProperAncestor(otherUser);
+          }))
+        continue;
+    }
     candidateLoop = parentLoop;
     break;
   }
   if (!candidateLoop)
     return failure();
   std::optional<Value> inductionVar = candidateLoop.getSingleInductionVar();
   std::optional<OpFoldResult> lowerBound = candidateLoop.getSingleLowerBound();
   std::optional<OpFoldResult> singleStep = candidateLoop.getSingleStep();

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

-// RUN: mlir-opt %s -test-transform-dialect-interpreter -verify-diagnostics -allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt %s -test-transform-dialect-interpreter -verify-diagnostics -allow-unregistered-dialect -split-input-file | FileCheck %s
 
 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> ((d0 floordiv 4) mod 2)>
 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 
 // CHECK-LABEL: func @multi_buffer
 func.func @multi_buffer(%in: memref<16xf32>) {
   // CHECK: %[[A:.*]] = memref.alloc() : memref<2x4xf32>
   // expected-remark @below {{transformed}}

 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["memref.alloc"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1 = transform.memref.multibuffer %0 {factor = 2 : i64}
   // 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_verification(%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) -> !pdl.operation
+  %1 = transform.memref.multibuffer %0 {factor = 2 : i64, skip_verification}
+  // Verify that the returned handle is usable.
+  transform.test_print_remark_at_operand %1, "transformed" : !pdl.operation
+}