[mlir][bufferize] Add isRepetitiveRegion to BufferizableOpInterface

[springerm]

Summary

This method allows to declare regions as "repetitive" even if the parent op does not implement the RegionBranchOpInterface.

Depends On D133112

[nicolasvasilache]

Looks good but I would also mention in the doc that this is driven by needs from parallel op semantics.
Such ops may have conflicts across threads of execution but do not necessarily implement RegionBranchOpInterface.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 
 namespace detail {
 /// This is the default implementation of
 /// BufferizableOpInterface::getBufferType. Should not be called from other
 /// places.
 FailureOr<BaseMemRefType>
 defaultGetBufferType(Value value, const BufferizationOptions &options,
                      const DenseMap<Value, BaseMemRefType> &fixedTypes);
+
+/// This is the default implementation of
+/// BufferizableOpInterface::isRepetitiveRegion. Should not be called from other
+/// places.
+bool defaultIsRepetitiveRegion(BufferizableOpInterface bufferizableOp,
+                               unsigned index);
 } // namespace detail
 
 } // namespace bufferization
 } // namespace mlir
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

let methods = [
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           assert(getOwnerOfValue(value) == $_op.getOperation() &&
                  "expected that value belongs to this op");
           return bufferization::detail::defaultGetBufferType(
               value, options, fixedTypes);
         }]
       >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return `true` if the given region of this op is repetitive. By default
+          this information is queried from the `RegionBranchOpInterface`. Ops
+          that do not implement this inferface can override this method to
+          declare regions as repetitive.
+
+          The RaW conflict detection of One-Shot Analysis is more strict inside
+          repetitive regions: Op dominance cannot always be used to rule out
+          certain potential conflicts (e.g., a conflicting write happening after
+          a read), because there may not be a meaningful ordering of certain ops
+          that are executed multiple times. This is described in more detail in
+          documentation of One-Shot Analysis.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"isRepetitiveRegion",
+        /*args=*/(ins "unsigned":$index),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return mlir::bufferization::detail::defaultIsRepetitiveRegion(
+              cast<BufferizableOpInterface>($_op.getOperation()), index);
+        }]
+      >
   ];
 
   let extraClassDeclaration = [{
     /// Resolve out-of-place tensor OpOperands with explicit allocations in the
     /// form of `bufferization.alloc_tensor` ops.
     LogicalResult resolveTensorOpOperandConflicts(
         RewriterBase &rewriter, const AnalysisState &state);
 

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

 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "llvm/Support/Debug.h"
 
 //===----------------------------------------------------------------------===//
 // BufferizableOpInterface
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace bufferization {
bufferization::getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
   auto rankedTensorType = tensorType.cast<RankedTensorType>();
   auto memorySpaceAttr = IntegerAttr::get(
       IntegerType::get(tensorType.getContext(), 64), memorySpace);
   MemRefLayoutAttrInterface layout = {};
   return MemRefType::get(rankedTensorType.getShape(),
                          rankedTensorType.getElementType(), layout,
                          memorySpaceAttr);
 }
+
+bool bufferization::detail::defaultIsRepetitiveRegion(
+    BufferizableOpInterface bufferizableOp, unsigned index) {
+  assert(index < bufferizableOp->getNumRegions() && "invalid region index");
+  auto regionInterface =
+      dyn_cast<RegionBranchOpInterface>(bufferizableOp.getOperation());
+  if (!regionInterface)
+    return false;
+  return regionInterface.isRepetitiveRegion(index);
+}

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;
 }
 
+static Region *
+getEnclosingRepetitiveRegion(Operation *op,
+                             const BufferizationOptions &options) {
+  while (Region *region = op->getParentRegion()) {
+    op = region->getParentOp();
+    if (auto bufferizableOp = options.dynCastBufferizableOp(op))
+      if (bufferizableOp.isRepetitiveRegion(region->getRegionNumber()))
+        return region;
+  }
+  return nullptr;
+}
+
+static Region *
+getEnclosingRepetitiveRegion(Value value, const BufferizationOptions &options) {
+  Region *region = value.getParentRegion();
+  while (region) {
+    Operation *op = region->getParentOp();
+    if (auto bufferizableOp = options.dynCastBufferizableOp(op))
+      if (bufferizableOp.isRepetitiveRegion(region->getRegionNumber()))
+        return region;
+    region = op->getParentRegion();
+  }
+  return nullptr;
+}
+
 /// 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) {
+getCommonEnclosingRepetitiveRegion(ArrayRef<Value> values,
+                                   const BufferizationOptions &options) {
   if (values.empty())
     return None;
-  Region *r = getEnclosingRepetitiveRegion(values.front());
+  Region *r = getEnclosingRepetitiveRegion(values.front(), options);
   for (Value value : values.drop_front())
     if (getEnclosingRepetitiveRegion(value) != r)
       return None;
   return r;
 }
 
 /// Return `true` if the given tensor value is a memory write. Most values are
 /// tensor writes, but ops that define a tensor SSA value without specifying its
static bool hasReadAfterWriteInterference(
   // Gather all written aliases. Skip over aliases that are not actual writes.
   SmallVector<Value> writtenAliases;
   for (OpOperand *uWrite : usesWrite)
     if (isMemoryWrite(uWrite->get(), state))
       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);
+      getCommonEnclosingRepetitiveRegion(writtenAliases, options);
 
   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>
for (OpOperand *uConflictingWrite : usesWrite) {
       // on them.
       //
       // Note: If `writtenAliases` is empty, there are no memory writes outside
       // of the repetitive region of conflictingWritingOp, which means that all
       // relevant aliases are inside the same repetitive region.
       bool canUseOpDominance =
           writtenAliases.empty() ||
           repetitiveRegionOfWrites ==
-              getEnclosingRepetitiveRegion(conflictingWritingOp);
+              getEnclosingRepetitiveRegion(conflictingWritingOp, options);
 
       // No conflict if the readingOp dominates conflictingWritingOp, i.e., the
       // write is not visible when reading.
       //
       // 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))

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

 /// }
 static void
 resolveUsesInRepetitiveRegions(Operation *op,
                                const BufferizationOptions &options) {
   IRRewriter rewriter(op->getContext());
   AnalysisState state(options);
 
   // Look for repetitive ops (loops).
-  op->walk([&](RegionBranchOpInterface regionBranchOp) {
-    // Skip non-bufferizable ops.
-    auto bufferizableOp = options.dynCastBufferizableOp(regionBranchOp);
-    if (!bufferizableOp)
+  op->walk([&](BufferizableOpInterface bufferizableOp) {
+    // Skip filtered ops.
+    if (!options.isOpAllowed(bufferizableOp.getOperation()))
       return WalkResult::advance();
 
-    // Find all operands that are also used inside of a repetitve region of this
-    // op.
-    for (OpOperand &opOperand : regionBranchOp->getOpOperands()) {
+    // Find all operands that are also used inside of a repetitive region of
+    // this op.
+    for (OpOperand &opOperand : bufferizableOp->getOpOperands()) {
       Value operand = opOperand.get();
       // Skip non-tensor operands.
       if (!operand.getType().isa<TensorType>())
         continue;
       // Skip operands that do not bufferize to memory writes.
       if (!bufferizableOp.bufferizesToMemoryWrite(opOperand, state))
         continue;
 
       // Gather all uses inside repetitive regions.
       SmallVector<OpOperand *> usesInsideRegion;
       for (OpOperand &use : operand.getUses()) {
         Operation *owner = use.getOwner();
-        if (!regionBranchOp->isProperAncestor(owner))
+        if (!bufferizableOp->isProperAncestor(owner))
           continue;
-        for (Region &r : regionBranchOp->getRegions()) {
+        for (Region &r : bufferizableOp->getRegions()) {
           if (r.findAncestorOpInRegion(*owner) &&
-              regionBranchOp.isRepetitiveRegion(r.getRegionNumber())) {
+              bufferizableOp.isRepetitiveRegion(r.getRegionNumber())) {
             usesInsideRegion.push_back(&use);
             break;
           }
         }
       }
       // Nothing to do if the operand is not used inside a repetitive region.
       if (usesInsideRegion.empty())
         continue;
 
       // Insert a tensor copy and replace all uses inside of repetitive regions.
-      rewriter.setInsertionPoint(regionBranchOp);
+      rewriter.setInsertionPoint(bufferizableOp);
       auto tensorCopy = rewriter.create<AllocTensorOp>(
-          regionBranchOp->getLoc(), operand.getType().cast<TensorType>(),
+          bufferizableOp->getLoc(), operand.getType().cast<TensorType>(),
           /*dynamicSizes=*/ValueRange(),
           /*copy=*/operand, /*memory_space=*/IntegerAttr());
       for (OpOperand *use : usesInsideRegion)
         use->set(tensorCopy);
     }
 
     return WalkResult::advance();
   });

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

cc_library(
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":AllocationOpInterface",
         ":ArithmeticDialect",
         ":BufferizableOpInterfaceIncGen",
         ":BufferizationBaseIncGen",
         ":BufferizationOpsIncGen",
+        ":ControlFlowInterfaces",
         ":CopyOpInterface",
         ":FuncDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefDialect",
         ":SparseTensorDialect",
         ":Support",
         ":TensorDialect",