[mlir][linalg][bufferize] Add `isAllocationHoistingBarrier` to op interface

[springerm]

Summary

This make getResultBuffer in ComprehensiveBufferize independent of the SCF, Affine and Linalg dialects. This commit is in preparating of decoupling op interface implementations from ComprehensiveBufferize.

Depends On D113378

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

 BufferRelation bufferRelation(OpOperand &opOperand);
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc"
 
+namespace mlir {
+namespace linalg {
+namespace comprehensive_bufferize {
+
+/// AllocationHoistingBarrierOnly is an external implementation of
+/// BufferizableOpInterface for ops that are (not yet) bufferizable, but are
+/// known to be allocation hoisting barriers. All interface methods (except for
+/// `isAllocationHoistingBarrier`) are implemented conservatively.
+template <typename OpTy>
+struct AllocationHoistingBarrierOnly
+    : public BufferizableOpInterface::ExternalModel<
+          AllocationHoistingBarrierOnly<OpTy>, OpTy> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
+    return false;
+  }
+
+  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
+                                                OpResult opResult) const {
+    return {};
+  }
+
+  OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
+    return OpResult();
+  }
+
+  BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
+    return BufferRelation::None;
+  }
+
+  bool isWritable(Operation *op, Value value) const { return false; }
+
+  LogicalResult bufferize(Operation *op, OpBuilder &b,
+                          BlockAndValueMapping &bvm,
+                          BufferizationAliasInfo &aliasInfo,
+                          AllocationCallbacks &allocationFn) const {
+    auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
+    if (any_of(op->getOperandTypes(), isaTensor) ||
+        any_of(op->getResultTypes(), isaTensor))
+      return op->emitError() << "unsupported op with tensors";
+    return success();
+  }
+
+  bool isAllocationHoistingBarrier(Operation *op) const { return true; }
+};
+
+} // namespace comprehensive_bufferize
+} // namespace linalg
+} // namespace mlir
+
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td

let methods = [
         }],
         /*retType=*/"bool",
         /*methodName=*/"isWritable",
         /*args=*/(ins "Value":$value),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return value.isa<OpResult>();
         }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return `true` if the op is an allocation hoisting barrier. Buffer
+          allocations will never be beyond such ops. E.g., ops with certain
+          parallel semantics may be allocation hoisting barriers. The majority
+          of ops, however, is not a barrier. Therefore, this method returns
+          `false` by default.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"isAllocationHoistingBarrier",
+        /*args=*/(ins),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return false;
+        }]
       >
   ];
 
   let extraClassDeclaration = [{
     /// Return `true` if the given OpOperand creates an alias but does neither
     /// read nor write. This implies that `bufferizesToMemoryRead` and
     /// `bufferizesToMemoryWrite` must return `false`. This method will never
     /// be called on OpOperands that do not have a tensor type.

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

static FunctionType getOrCreateBufferizedFunctionType(
   LDBG("FT: " << funcOp.getType() << " -> " << it2.first->second << "\n");
   return it2.first->second;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific scoped alloc/dealloc insertion support.
 //===----------------------------------------------------------------------===//
 
-template <typename... Args>
-Operation *getFirstParentOfType(Value v) {
-  Operation *parent;
-  if (auto bbArg = v.dyn_cast<BlockArgument>())
-    parent = bbArg.getOwner()->getParentOp();
-  else
-    parent = v.getDefiningOp()->getParentOp();
-  while (parent) {
-    if (isa<Args...>(parent))
-      return parent;
-    parent = parent->getParentOp();
-  }
-  return nullptr;
-}
-
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 static Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
                                int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
     return b.createOrFold<memref::DimOp>(loc, source, dim);
   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
     return b.createOrFold<tensor::DimOp>(loc, source, dim);
static MemRefType getAllocationTypeAndShape(OpBuilder &b, Location loc,
   if (!foundDynamicShapes) {
     for (auto dim : enumerate(allocMemRefType.getShape()))
       if (dim.value() == ShapedType::kDynamicSize)
         dynShape.push_back(createOrFoldDimOp(b, loc, shapedValue, dim.index()));
   }
 
   // If the buffer is statically shaped, try to hoist it to the first enclosing
   // parallel region.
-  // TODO: this concept of parallel region and threadlocal needs interfaces.
   // TODO: also hoist in the dynamic case. For now this relies on subsequent
   // calls to LICM and buffer hoisting which will most likely not succeed.
   // TODO: when packing, allocate a static bounding box which will enable more
   // hoisting.
   if (dynShape.empty()) {
-    Operation *parent =
-        getFirstParentOfType<FuncOp, TiledLoopOp, scf::ParallelOp,
-                             AffineParallelOp>(shapedValue);
-    if (parent)
-      b.setInsertionPointToStart(&(parent->getRegion(0).front()));
+    Operation *parent;
+    if (auto bbArg = shapedValue.dyn_cast<BlockArgument>())
+      parent = bbArg.getOwner()->getParentOp();
+    else
+      parent = shapedValue.getDefiningOp()->getParentOp();
+    while (parent) {
+      if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(parent))
+        if (bufferizableOp.isAllocationHoistingBarrier())
+          break;
+      parent = parent->getParentOp();
+    }
+
+    // FuncOp is an allocation hoisting barrier, so the above loop should never
+    // run out of parents.
+    assert(
+        (parent &&
+         cast<BufferizableOpInterface>(parent).isAllocationHoistingBarrier()) &&
+        "expected traversal to end at allocation hoisting barrier");
+
   }
   return allocMemRefType;
 }
 
 /// Create an Allocop/DeAllocOp pair, where the AllocOp is after
 /// `shapedValue.getDefiningOp` (or at the top of the block in case of a
 /// bbArg) and the DeallocOp is at the end of the block.
 static Value createNewAllocDeallocPairForShapedValue(
bool isWritable(Operation *op, Value value) const {
     // inplace from the perspective of ops nested under:
     //   1. Either the matching iter operand is not bufferized inplace and an
     //      alloc + optional copy makes the bbArg itself inplaceable.
     //   2. Or the matching iter operand is bufferized inplace and bbArg just
     //      bufferizes to that too.
     return true;
   }
 
+  bool isAllocationHoistingBarrier(Operation *op) const { return true; }
+
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BlockAndValueMapping &bvm,
                           BufferizationAliasInfo &aliasInfo,
                           AllocationCallbacks &allocationFn) const {
     auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   registry.addOpInterface<tensor::ExtractOp, tensor_ext::ExtractOpInterface>();
   registry.addOpInterface<tensor::InsertSliceOp,
                           tensor_ext::InsertSliceOpInterface>();
   registry.addOpInterface<vector::TransferReadOp,
                           vector_ext::TransferReadOpInterface>();
   registry.addOpInterface<vector::TransferWriteOp,
                           vector_ext::TransferWriteOpInterface>();
 
+  // Ops that are not bufferizable but are allocation hoisting barriers.
+  registry.addOpInterface<FuncOp, AllocationHoistingBarrierOnly<FuncOp>>();
+  registry.addOpInterface<scf::ParallelOp,
+                          AllocationHoistingBarrierOnly<scf::ParallelOp>>();
+  registry.addOpInterface<AffineParallelOp,
+                          AllocationHoistingBarrierOnly<AffineParallelOp>>();
+
   // Register all Linalg structured ops. `LinalgOp` is an interface and it is
   // not possible to attach an external interface to an existing interface.
   // Therefore, attach the `BufferizableOpInterface` to all ops one-by-one.
   LinalgOpInterfaceHelper<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >::registerOpInterface(registry);
 }
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

LinalgComprehensiveModuleBufferize(
       const LinalgComprehensiveModuleBufferize &p) {}
 
   void runOnOperation() override;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<linalg::LinalgDialect, memref::MemRefDialect,
                 tensor::TensorDialect, vector::VectorDialect, scf::SCFDialect,
-                arith::ArithmeticDialect, StandardOpsDialect>();
+                arith::ArithmeticDialect, StandardOpsDialect, AffineDialect>();
     registerBufferizableOpInterfaceExternalModels(registry);
   }
 };
 } // end namespace
 
 static void applyEnablingTransformations(ModuleOp moduleOp) {
   RewritePatternSet patterns(moduleOp.getContext());
   patterns.add<GeneralizePadTensorOpPattern>(moduleOp.getContext());