[mlir][bufferize] Fix missing copies when writing to a buffer in a loop

[springerm]

Summary

Writes into tensors that are definied outside of a repetitive region, but with the write happening inside of the repetitive region were previously not considered conflicts. This was incorrect.

E.g.:

%0 = ... : tensor<?xf32>
scf.for ... {
  "reading_op"(%0) : tensor<?xf32>
  %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32>
  ...
}

In the above example, "writing_op" should be out-of-place.

This commit fixes the bufferization for any op that declares its repetitive semantics via RegionBranchOpInterface.

Depends On D123790

[nicolasvasilache]

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

451

typo: dominance

[springerm]

This has landed as 9235e597a40b423a298ce415eb922462e7f0b765.

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

do {
     if (a->isProperAncestor(b))
       return false;
     if (domInfo.properlyDominates(a, b))
       return true;
   } while ((a = a->getParentOp()));
   return false;
 }
 
+/// For each given value, find the closest enclosing repetitive region. If this
+/// is the same region for each value, return it. Otherwise return None.
+/// Note: If there is no enclosing repetitive region, return nullptr.
+static Optional<Region *>
+getCommonEnclosingRepetitiveRegion(ArrayRef<Value> values) {
+  if (values.empty())
+    return None;
+  Region *r = getEnclosingRepetitiveRegion(values.front());
+  for (Value value : values.drop_front())
+    if (getEnclosingRepetitiveRegion(value) != r)
+      return None;
+  return r;
+}
+
 /// Annotate IR with details about the detected RaW conflict.
 static void annotateConflict(OpOperand *uRead, OpOperand *uConflictingWrite,
                              Value lastWrite) {
   static uint64_t counter = 0;
   Operation *readingOp = uRead->getOwner();
   Operation *conflictingWritingOp = uConflictingWrite->getOwner();
 
   OpBuilder b(conflictingWritingOp->getContext());

 /// the result of a write W1. But because of bufferization decisions, R actually
 /// reads another write W2.
 static bool hasReadAfterWriteInterference(
     const DenseSet<OpOperand *> &usesRead,
     const DenseSet<OpOperand *> &usesWrite, const DominanceInfo &domInfo,
     AnalysisState &state, const BufferizationAliasInfo &aliasInfo) {
   const BufferizationOptions &options = state.getOptions();
 
+  // Gather all written aliases.
+  SmallVector<Value> writtenAliases;
+  for (OpOperand *uWrite : usesWrite)
+    writtenAliases.push_back(uWrite->get());
+  // Find the inner-most enclosing repetitive region of each alias. If this is
+  // the same region for every alias, save it in `repetitiveRegionOfWrites`.
+  Optional<Region *> repetitiveRegionOfWrites =
+      getCommonEnclosingRepetitiveRegion(writtenAliases);
+
   for (OpOperand *uRead : usesRead) {
     Operation *readingOp = uRead->getOwner();
 
     // Find most recent writes of uRead by following the SSA use-def chain.
     // E.g.:
     //
     // %0 = "writing_op"(%t) : tensor<?x32> -> tensor<?xf32>
     // %1 = "aliasing_op"(%0) : tensor<?x32> -> tensor<?xf32>
     // %2 = "reading_op"(%1) : : tensor<?x32> -> not_a_tensor_type
     //
     // In the above example, if uRead is the OpOperand of reading_op, lastWrite
     // is %0. Note that operations that create an alias but do not write (such
     // as ExtractSliceOp) are skipped.
     SetVector<Value> lastWrites = state.findLastPrecedingWrite(uRead->get());
 
     // Look for conflicting memory writes. Potential conflicts are writes to an
     // alias that have been decided to bufferize inplace.
     for (OpOperand *uConflictingWrite : usesWrite) {
       // Throughout this loop, check for multiple requirements that have to be
       // met for uConflictingWrite to be an actual conflict.
       Operation *conflictingWritingOp = uConflictingWrite->getOwner();
 
+      // Check if conflictingWritingOp is in the same repetitive region as all
+      // written aliases. If this is not the case, there is no meaningful
+      // `happensBefore` relationship because conflictingWritingOp may be
+      // executed multiple times. E.g.:
+      //
+      // %0 = ... : tensor<?xf32>
+      // scf.for ... {
+      //   "reading_op"(%0) : tensor<?xf32>
+      //   %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32>
+      //   ...
+      // }
+      //
+      // In the above example, reading_op happens before writing_op according to
+      // op dominance. However, both ops may happen multiple times; in
+      // particular, the second execution of reading_op happens after the first
+      // execution of writing_op. This is problematic if the tensor they operate
+      // on (%0) is defined outside of the loop.
+      //
+      // Counter example:
+      //
+      // scf.for ... {
+      //   %0 = ... : tensor<?xf32>
+      //   "reading_op"(%0) : tensor<?xf32>
+      //   %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32>
+      //   ...
+      // }
+      //
+      // In this example, %0 is in the same repetitive region as
+      // conflictingWritingOp, so op dominance can be used to compute the
+      // `happensBefore` relationship.
+      //
+      // Note: iter_args of loops are not aliases of their respective block
+      // arguments, so op domanice can be used when analyzing ops that operate
+      // on them.
+      bool canUseOpDominance =
+          repetitiveRegionOfWrites ==
+          getEnclosingRepetitiveRegion(conflictingWritingOp);
+
       // No conflict if the readingOp dominates conflictingWritingOp, i.e., the
       // write is not visible when reading.
-      if (happensBefore(readingOp, conflictingWritingOp, domInfo))
+      //
+      // Note: If ops are executed multiple times (e.g., because they are inside
+      //       a loop), there may be no meaningful `happensBefore` relationship.
+      if (canUseOpDominance &&
+          happensBefore(readingOp, conflictingWritingOp, domInfo))
         continue;
 
       // No conflict if the reading use equals the use of the conflicting write.
-      // A use cannot conflict with itself. Note: Just being the same op is not
-      // enough. It has to be the same use.
-      if (uConflictingWrite == uRead)
+      // A use cannot conflict with itself.
+      //
+      // Note: Just being the same op is not enough. It has to be the same use.
+      // Note: If the op is executed multiple times (e.g., because it is inside
+      //       a loop), it may be conflicting with itself.
+      if (canUseOpDominance && uConflictingWrite == uRead)
         continue;
 
       // No conflict if the op interface says so.
       if (auto bufferizableOp = options.dynCastBufferizableOp(readingOp))
         if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state))
           continue;
 
       if (conflictingWritingOp != readingOp)
         if (auto bufferizableOp =
                 options.dynCastBufferizableOp(conflictingWritingOp))
           if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state))
             continue;
 
       // Ops are not conflicting if they are in mutually exclusive regions.
-      if (insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp))
+      //
+      // Note: If ops are executed multiple times (e.g., because they are inside
+      //       a loop), mutually exclusive regions may be executed multiple
+      //       times.
+      if (canUseOpDominance &&
+          insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp))
         continue;
 
       // Check all possible last writes.
       for (Value lastWrite : lastWrites) {
         // No conflict if the conflicting write happens before the last
         // write.
         if (Operation *writingOp = lastWrite.getDefiningOp()) {
           if (happensBefore(conflictingWritingOp, writingOp, domInfo))

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

func @write_after_select_no_conflict(
   // CHECK-SAME:   {__inplace_operands_attr__ = ["none", "true", "none"]}
   %w = tensor.insert %cst into %s[%idx] : tensor<?xf32>
   //      CHECK: tensor.extract
   // CHECK-SAME:   {__inplace_operands_attr__ = ["true", "none"]}
   %f = tensor.extract %w[%idx] : tensor<?xf32>
 
   return %f, %w : f32, tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @write_to_same_tensor_in_loop_out_of_place(
+func @write_to_same_tensor_in_loop_out_of_place(
+    %A : tensor<?xf32> {linalg.inplaceable = true},
+    %B : tensor<?xf32> {linalg.inplaceable = true},
+    %lb : index, %ub : index, %step : index, %sz: index)
+  -> (tensor<?xf32>)
+{
+  // CHECK: scf.for {{.*}} {
+  %r0 = scf.for %i = %lb to %ub step %step iter_args(%t = %A) -> (tensor<?xf32>) {
+    %i2 = arith.index_cast %i : index to i32
+    %i3 = arith.sitofp %i2 : i32 to f32
+    // The tensor.insert is out-of-place because the %B is written multiple
+    // times inside a loop.
+    //      CHECK: tensor.insert
+    // CHECK-SAME:   {__inplace_operands_attr__ = ["none", "false", "none"]}
+    %B2 = tensor.insert %i3 into %B[%i] : tensor<?xf32>
+    //      CHECK: tensor.insert_slice
+    // CHECK-SAME:   {__inplace_operands_attr__ = ["true", "true", "none", "none"]}
+    %A2 = tensor.insert_slice %B2 into %t[%i][%sz][1] : tensor<?xf32> into tensor<?xf32>
+    scf.yield %A2 : tensor<?xf32>
+  }
+  // CHECK: } {__inplace_operands_attr__ = ["none", "none", "none", "true"]}
+
+  return %r0 : tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @write_to_same_tensor_in_loop_in_place(
+func @write_to_same_tensor_in_loop_in_place(
+    %A : tensor<?xf32> {linalg.inplaceable = true},
+    %lb : index, %ub : index, %step : index, %sz: index)
+  -> (tensor<?xf32>)
+{
+  // CHECK: scf.for {{.*}} {
+  %r0 = scf.for %i = %lb to %ub step %step iter_args(%t = %A) -> (tensor<?xf32>) {
+    %B = linalg.init_tensor [%sz] : tensor<?xf32>
+    %i2 = arith.index_cast %i : index to i32
+    %i3 = arith.sitofp %i2 : i32 to f32
+    // The tensor.insert is in-place because the %B is defined inside the loop.
+    //      CHECK: tensor.insert
+    // CHECK-SAME:   {__inplace_operands_attr__ = ["none", "true", "none"]}
+    %B2 = tensor.insert %i3 into %B[%i] : tensor<?xf32>
+    //      CHECK: tensor.insert_slice
+    // CHECK-SAME:   {__inplace_operands_attr__ = ["true", "true", "none", "none"]}
+    %A2 = tensor.insert_slice %B2 into %t[%i][%sz][1] : tensor<?xf32> into tensor<?xf32>
+    scf.yield %A2 : tensor<?xf32>
+  }
+  // CHECK: } {__inplace_operands_attr__ = ["none", "none", "none", "true"]}
+
+  return %r0 : tensor<?xf32>
+}