[mlir][bufferization] Fix bufferization of repetitive regions

[springerm]

Summary

The previous strategy was too complex and faulty. Op dominance cannot be used to rule out RaW conflicts due to op ordering if the reading op and the conflicting writing op are in a sub repetitive region of the closest enclosing repetitive region of the definition of the read value.

Depends On: D143183

[springerm]

update

[chelini]

Thanks, looks good to me.

[springerm]

update

[springerm]

Updated diff to a better strategy for using op dominance. Also a better explanation and more test cases.

[springerm]

Added Ingo just for FYI since we discussed this earlier.

[ingomueller-net]

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

442–443

I don't fully understand what this does but it looks a bit different from what we had discussed offline: I remember we came to the conclusion that we couldn't use op dominance if the "any common parent region of the read and the write up until the region of the def" was repetitive, right? For this, you'd have to test for any enclosing repetitive region of rRead if it is also an ancestor of uWrite, right? (Maybe your while loop closes too early?)

444

What does this loop do? It doesn't "find" anything, does it, since it doesn't modify anything outside it's scope? The only effect I see is that it asserts...

[springerm]

fix missing assignment in loop

[springerm]

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

442–443

The loop is looking for a repetitive region that is a child region (maybe not direct child) of rDef. Starting from rRead. The loop was not doing what the comment is saying. Added another test case that runs multiple iterations of this loop.

444

Yes, this was broken. An assignment was missing.

[ingomueller-net]

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

442–443

Ah, OK, I think I get it. If the "top-most" repetitive region contains both the read and the write (which is what your code is testing), then there is "any common parent region of the read and the right [...] that is repetitive" (which is what I was expecting). Conversely, if the top-most region doesn't contain both, then there is also no region at a deeper nesting level that would contain both.

[springerm]

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

442–443

Yes exactly.

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

Region *bufferization::getEnclosingRepetitiveRegion(
   return getEnclosingRepetitiveRegion(op->getBlock(), options);
 }
 
 Region *bufferization::getEnclosingRepetitiveRegion(
     Value value, const BufferizationOptions &options) {
   Region *region = value.getParentRegion();
   while (region) {
     Operation *op = region->getParentOp();
+    // Stop lookup when reaching a function.
+    if (isa<FunctionOpInterface>(op))
+      return region;
     if (auto bufferizableOp = options.dynCastBufferizableOp(op))
       if (bufferizableOp.isRepetitiveRegion(region->getRegionNumber()))
         return region;
     region = op->getParentRegion();
   }
   return nullptr;
 }
 
 Region *bufferization::getEnclosingRepetitiveRegion(
     Block *block, const BufferizationOptions &options) {
   Region *region = block->getParent();
   Operation *op = nullptr;
   do {
     op = region->getParentOp();
+    // Stop lookup when reaching a function.
+    if (isa<FunctionOpInterface>(op))
+      return region;
     if (auto bufferizableOp = options.dynCastBufferizableOp(op))
       if (bufferizableOp.isRepetitiveRegion(region->getRegionNumber()))
         return region;
   } while ((region = op->getParentRegion()));
   return nullptr;
 }
 
 Operation *bufferization::getOwnerOfValue(Value value) {

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;
 }
 
-/// Return `true` if op dominance can be used to rule out read-after-write
-/// conflicts wrt. the given reads and writes.
+/// Return `true` if op dominance can be used to rule out a read-after-write
+/// conflicts.
 ///
-/// Op dominance can often be used to rule out potential conflicts such as
-/// "read" happens before "write". E.g., the following IR is not a RaW conflict
-/// because the the read happens *before* the write.
+/// Op dominance can often be used to rule out potential conflicts due to
+/// "read happens before write". E.g., the following IR is not a RaW conflict
+/// because the read happens *before* the write.
 ///
 /// %0 = ... : tensor<?xf32>
 /// "reading_op"(%0) : tensor<?xf32>
 /// %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32>
 ///
 /// This is no longer true inside loops (or repetitive regions). In such cases,
 /// there may not be a meaningful `happensBefore` relationship because ops
 /// could be executed multiple times. E.g.:

 ///   %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32>
 ///   ...
 /// }
 ///
 /// In this example, the definition %0 is in the same repetitive region as
 /// "writing_op", so op dominance can be used to compute the `happensBefore`
 /// relationship.
 ///
-/// Whether op dominance can be used or not is decided as follows: Find the
-/// closest enclosing repetitive region of all buffer writes wrt. the given
-/// tensor reads and writes. (The given sets of reads and writes contain the
-/// entire alias set.) In case of a read, we look at the op that defines the
-/// read value. In case of a write, we look at the op that is writing. If all of
-/// those ops are in the same closest enclosing repetitive region (nullptr in
-/// case of "no repetitive region" found at all), then op dominance can be used.
-/// Otherwise, it cannot be used.
+/// Counter example:
+///
+/// %0 = ... : tensor<?xf32>
+/// %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32>
+/// scf.for ... {
+///   "reading_op"(%1) : tensor<?xf32>
+///   ...
+/// }
 ///
-/// Example: The common enclosing repetitive region is the scf.for loop.
-///          Op dominance can be used.
+/// In this example, the definition %1 is in the same repetitive region as
+/// "writing_op", so op dominance can be used to compute the `happensBefore`
+/// relationship.
+///
+/// Op dominance can be used if:
+/// repetitive_reg(write) == repetitive_reg(def(read))
+///
+/// Example: There is no write, so op dominance can be used.
 /// scf.for ... {
 ///   %0 = tensor.generate
 ///   "read"(%0)
 /// }
 ///
-/// Example: The common enclosing repetitive region is nullptr: There is no
-///          repetitive region around the tensor.generate. Op dominance can be
-///          used.
+/// Example: There is no write, so op dominance can be used.
 /// %0 = tensor.generate
 /// scf.for ... { "read"(%0) }
 ///
-/// Example: The common enclosing repetitive regions of tensor.generate and
-///          "write" differ. Op dominance cannot be used.
+/// Example: The write is in a different repetitive region than %0. Op dominance
+///          cannot be used.
 /// %0 = tensor.generate
 /// scf.for ... {
 ///   "read"(%0)
 ///   "write"(%0)
 /// }
 ///
-/// Example: The common enclosing repetitive regions of tensor.generate and
-///          "write" differ, but there is no read of %0, so op dominance can be
-///          used.
+/// Example: There is no read, so op dominance can be used.
 /// %0 = tensor.generate
 /// scf.for ... {
 ///   "write"(%0)
 /// }
 ///
 /// 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(const DenseSet<OpOperand *> &usesRead,
-                       const DenseSet<OpOperand *> &usesWrite,
+bool canUseOpDominance(OpOperand *uRead,
+                       OpOperand *uWrite,
+                       const SetVector<Value> &definitions,
                        const AnalysisState &state) {
   const BufferizationOptions &options = state.getOptions();
-  std::optional<Region *> commonEnclosingRegion;
-
-  // In case of a write, take the region in which the write takes place.
-  for (OpOperand *uWrite : usesWrite) {
-    Region *r = getEnclosingRepetitiveRegion(uWrite->getOwner(), options);
-    if (!commonEnclosingRegion.has_value()) {
-      commonEnclosingRegion = r;
+  for (Value def : definitions) {
+    // Skip tensors that have no defined contents.
+    if (!state.bufferizesToMemoryWrite(def))
       continue;
-    }
-    if (*commonEnclosingRegion != r)
+    if (r != getEnclosingRepetitiveRegion(def, options))
       return false;
   }
-
-  // In case of a read, take the region which the read value is defined.
-  for (OpOperand *uRead : usesRead) {
-    // Optimization: Skip reads of values that have no defined contents.
-    if (!state.bufferizesToMemoryWrite(uRead->get()))
-      continue;
-    Region *r = getEnclosingRepetitiveRegion(uRead->get(), options);
-    if (!commonEnclosingRegion.has_value()) {
-      commonEnclosingRegion = r;
-      continue;
-    }
-    if (*commonEnclosingRegion != r)
-      return false;
-  }
-
-  return commonEnclosingRegion.has_value();
+  return true;
 }
 
 /// Annotate IR with details about the detected RaW conflict.
 static void annotateConflict(OpOperand *uRead, OpOperand *uConflictingWrite,
                              Value definition) {
   static uint64_t counter = 0;
   Operation *readingOp = uRead->getOwner();
   Operation *conflictingWritingOp = uConflictingWrite->getOwner();

 /// the result of a definition W1. But because of bufferization decisions, R
 /// actually reads another definition 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();
 
-  // Check if op dominance can be used to rule out read-after-write conflicts.
-  bool useDominance = canUseOpDominance(usesRead, usesWrite, state);
-  LLVM_DEBUG(llvm::dbgs() << "\n- useDominance = " << useDominance << "\n");
-
   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, the
     // definition is %0. Note that operations that create an alias but do not
     // bufferize to a memory write (such as ExtractSliceOp) are skipped.
     SetVector<Value> definitions = state.findDefinitions(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) {
+      // Check if op dominance can be used to rule out read-after-write
+      // conflicts.
+      bool useDominance = canUseOpDominance(uRead, uConflictingWrite, definitions, state);
+      LLVM_DEBUG(llvm::dbgs() << "\n- useDominance = " << useDominance << "\n");
+
       LLVM_DEBUG(llvm::dbgs() << "\n- check conflict:\n");
       LLVM_DEBUG(llvm::dbgs()
                  << "  uRead = operand " << uRead->getOperandNumber() << " of "
                  << *uRead->getOwner() << "\n");
       LLVM_DEBUG(llvm::dbgs() << "  unConflictingWrite = operand "
                               << uConflictingWrite->getOperandNumber() << " of "
                               << *uConflictingWrite->getOwner() << "\n");
 

mlir/test/Dialect/SCF/one-shot-bufferize-analysis.mlir

func.func @no_raw_conflict_after_repetitive_use(%arg0: tensor<4xf32>,
   // CHECK: tensor.extract {{.*}} {__inplace_operands_attr__ = ["true", "none"]}
   %6 = tensor.extract %1[%arg4] : tensor<4xf32>
   vector.print %6 : f32
   // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
   %7 = linalg.fill ins(%cst3 : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32>
 
   return %2, %7 : tensor<4xf32>, tensor<4xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @read_of_bbarg_in_repetitive_region(
+func.func @read_of_bbarg_in_repetitive_region(
+    %t: tensor<10xf32>, %a: index, %b: index, %c: index, %cst: f32) {
+  // CHECK: scf.for
+  scf.for %iv = %a to %b step %c {
+    // Must bufferize out-of-place because definition of read is in a different
+    // repetitive region.
+    // CHECK: tensor.extract_slice {{.*}} {__inplace_operands_attr__ = ["false"]}
+    %2 = tensor.extract_slice %t[0][4][1] : tensor<10xf32> to tensor<4xf32>
+    %3 = tensor.extract %2[%a] : tensor<4xf32>
+    vector.print %3 : f32
+    // CHECK: tensor.insert {{.*}} {__inplace_operands_attr__ = ["none", "true", "none"]}
+    %4 = tensor.insert %cst into %2[%a] : tensor<4xf32>
+    %5 = tensor.extract %4[%a] : tensor<4xf32>
+    vector.print %5 : f32
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @read_definition_in_same_repetitive_region_as_write(
+func.func @read_definition_in_same_repetitive_region_as_write(
+    %t: tensor<10xf32>, %a: index, %b: index, %c: index, %cst: f32) {
+  // CHECK: tensor.insert {{.*}} {__inplace_operands_attr__ = ["none", "true", "none"]}
+  %1 = tensor.insert %cst into %t[%a] : tensor<10xf32>
+  // CHECK: scf.for
+  scf.for %iv = %a to %b step %c {
+    // Can bufferize in-place.
+    // CHECK: tensor.extract_slice {{.*}} {__inplace_operands_attr__ = ["true"]}
+    %2 = tensor.extract_slice %1[0][4][1] : tensor<10xf32> to tensor<4xf32>
+    %3 = tensor.extract %2[%a] : tensor<4xf32>
+    vector.print %3 : f32
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @read_definition_in_same_repetitive_region_as_conflicting_write(
+func.func @read_definition_in_same_repetitive_region_as_conflicting_write(
+    %t: tensor<10xf32>, %a: index, %b: index, %c: index, %cst: f32) {
+  // Cannot bufferize in-place according to normal op dominance rules.
+  // CHECK: tensor.insert {{.*}} {__inplace_operands_attr__ = ["none", "false", "none"]}
+  %1 = tensor.insert %cst into %t[%a] : tensor<10xf32>
+  // CHECK: scf.for
+  scf.for %iv = %a to %b step %c {
+    // CHECK: tensor.extract_slice {{.*}} {__inplace_operands_attr__ = ["true"]}
+    %2 = tensor.extract_slice %t[0][4][1] : tensor<10xf32> to tensor<4xf32>
+    %3 = tensor.extract %2[%a] : tensor<4xf32>
+    vector.print %3 : f32
+  }
+  return
+}