[mlir][linalg][bufferize][NFC] Utilize isWritable for FuncOps

This is a cleanup of ModuleBufferization. Instead of storing information about writable function arguments in BufferizationAliasInfo, we can use isWritable and make the decision there, based on dialect-specifc bufferization state.

Differential Revision: https://reviews.llvm.org/D114930

Author: matthias-springer

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

@@ -172,13 +172,6 @@ class BufferizationAliasInfo {
   /// Apply `fun` to all aliases of `v`.
   void applyOnAliases(Value v, function_ref<void(Value)> fun) const;
 
-  // TODO: Move these out of BufferizationAliasInfo.
-  /// Return true if the value is known to bufferize to writable memory.
-  bool bufferizesToWritableMemory(Value v) const;
-
-  /// Specify that the value is known to bufferize to writable memory.
-  void setBufferizesToWritableMemory(Value v);
-
   /// Mark a value as in-place bufferized.
   void markInPlace(OpResult v) { inplaceBufferized.insert(v); }
 
@@ -200,9 +193,6 @@ class BufferizationAliasInfo {
   /// Check that aliasInfo for `v` exists and return a reference to it.
   EquivalenceClassRangeType getAliases(Value v) const;
 
-  /// Set of tensors that are known to bufferize to writable memory.
-  llvm::DenseSet<Value> bufferizeToWritableMemory;
-
   /// Set of all OpResults that were decided to bufferize in-place.
   llvm::DenseSet<OpResult> inplaceBufferized;
 
@@ -429,7 +419,9 @@ struct AllocationHoistingBarrierOnly
     return BufferRelation::None;
   }
 
-  bool isWritable(Operation *op, Value value) const { return false; }
+  bool isWritable(Operation *op, Value value, BufferizationState &state) const {
+    return false;
+  }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {

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

@@ -226,7 +226,8 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
         }],
         /*retType=*/"bool",
         /*methodName=*/"isWritable",
-        /*args=*/(ins "Value":$value),
+        /*args=*/(ins "Value":$value,
+                      "BufferizationState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return value.isa<OpResult>();

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

@@ -42,7 +42,7 @@ struct ConstantOpInterface
     return success();
   }
 
-  bool isWritable(Operation *op, Value value) const {
+  bool isWritable(Operation *op, Value value, BufferizationState &state) const {
     // Memory locations returned by memref::GetGlobalOp may not be written to.
     assert(value.isa<OpResult>());
     return false;

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

@@ -130,15 +130,6 @@ void BufferizationAliasInfo::insertNewBufferEquivalence(Value newValue,
   equivalentInfo.unionSets(newValue, alias);
 }
 
-bool BufferizationAliasInfo::bufferizesToWritableMemory(Value v) const {
-  return bufferizeToWritableMemory.count(v) > 0;
-}
-
-/// Specify that the value is known to bufferize to writable memory.
-void BufferizationAliasInfo::setBufferizesToWritableMemory(Value v) {
-  bufferizeToWritableMemory.insert(v);
-}
-
 /// Return `true` if a value was marked as in-place bufferized.
 bool BufferizationAliasInfo::isInPlace(OpResult opResult) const {
   return inplaceBufferized.contains(opResult);

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

@@ -71,7 +71,7 @@ struct ToTensorOpInterface
     return success();
   }
 
-  bool isWritable(Operation *op, Value value) const {
+  bool isWritable(Operation *op, Value value, BufferizationState &state) const {
     // It is unknown whether the MemRef operand is writable or not.
     return false;
   }

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

@@ -171,15 +171,6 @@ static void setInPlaceOpResult(OpResult opResult, bool inPlace) {
               OpBuilder(op).getStrArrayAttr(inPlaceVector));
 }
 
-/// Set the attribute that triggers inplace bufferization on a FuncOp argument
-/// `bbArg`.
-static void setInPlaceFuncArgument(BlockArgument bbArg, bool inPlace) {
-  auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
-  funcOp.setArgAttr(bbArg.getArgNumber(),
-                    BufferizableOpInterface::kInplaceableAttrName,
-                    BoolAttr::get(bbArg.getContext(), inPlace));
-}
-
 //===----------------------------------------------------------------------===//
 // Printing helpers.
 //===----------------------------------------------------------------------===//
@@ -258,25 +249,22 @@ static bool isInplaceMemoryWrite(OpOperand &opOperand,
 /// Return true if, under current bufferization decisions, the buffer of `value`
 /// is not writable.
 static bool aliasesNonWritableBuffer(Value value,
-                                     const BufferizationAliasInfo &aliasInfo) {
+                                     const BufferizationAliasInfo &aliasInfo,
+                                     BufferizationState &state) {
   LDBG("WRITABILITY ANALYSIS FOR " << printValueInfo(value) << "\n");
   bool foundNonWritableBuffer = false;
   aliasInfo.applyOnAliases(value, [&](Value v) {
-    // Some values are known to be writable.
-    if (aliasInfo.bufferizesToWritableMemory(v))
-      return;
-
     // Query BufferizableOpInterface to see if the OpResult is writable.
     // TODO: Out-of-place bufferized OpResult could be considered writable.
     if (auto bufferizableOp = v.getDefiningOp<BufferizableOpInterface>())
-      if (bufferizableOp && bufferizableOp.isWritable(v))
+      if (bufferizableOp && bufferizableOp.isWritable(v, state))
         return;
 
     // Query BufferizableOpInterface to see if the BlockArgument is writable.
     if (auto bbArg = v.dyn_cast<BlockArgument>())
       if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(
               bbArg.getOwner()->getParentOp()))
-        if (bufferizableOp.isWritable(bbArg))
+        if (bufferizableOp.isWritable(bbArg, state))
           return;
 
     foundNonWritableBuffer = true;
@@ -515,7 +503,8 @@ bool wouldCreateReadAfterWriteInterference(
 /// a write to a non-writable buffer.
 static bool
 wouldCreateWriteToNonWritableBuffer(OpOperand &opOperand, OpResult opResult,
-                                    const BufferizationAliasInfo &aliasInfo) {
+                                    const BufferizationAliasInfo &aliasInfo,
+                                    BufferizationState &state) {
 #ifndef NDEBUG
   SmallVector<OpOperand *> opOperands = getAliasingOpOperand(opResult);
   assert(llvm::find(opOperands, &opOperand) != opOperands.end() &&
@@ -525,9 +514,10 @@ wouldCreateWriteToNonWritableBuffer(OpOperand &opOperand, OpResult opResult,
   // Certain buffers are not writeable:
   //   1. A function bbArg that is not inplaceable or
   //   2. A constant op.
-  assert(!aliasesNonWritableBuffer(opResult, aliasInfo) &&
+  assert(!aliasesNonWritableBuffer(opResult, aliasInfo, state) &&
          "expected that opResult does not alias non-writable buffer");
-  bool nonWritable = aliasesNonWritableBuffer(opOperand.get(), aliasInfo);
+  bool nonWritable =
+      aliasesNonWritableBuffer(opOperand.get(), aliasInfo, state);
   if (!nonWritable)
     return false;
 
@@ -547,10 +537,9 @@ wouldCreateWriteToNonWritableBuffer(OpOperand &opOperand, OpResult opResult,
 //===----------------------------------------------------------------------===//
 
 /// Determine if `operand` can be bufferized in-place with `result`.
-static LogicalResult
-bufferizableInPlaceAnalysisImpl(OpOperand &operand, OpResult result,
-                                BufferizationAliasInfo &aliasInfo,
-                                const DominanceInfo &domInfo) {
+static LogicalResult bufferizableInPlaceAnalysisImpl(
+    OpOperand &operand, OpResult result, BufferizationAliasInfo &aliasInfo,
+    BufferizationState &state, const DominanceInfo &domInfo) {
 #ifndef NDEBUG
   SmallVector<OpOperand *> opOperands = getAliasingOpOperand(result);
   assert(llvm::find(opOperands, &operand) != opOperands.end() &&
@@ -565,7 +554,7 @@ bufferizableInPlaceAnalysisImpl(OpOperand &operand, OpResult result,
                                    << printValueInfo(result) << '\n');
 
   bool foundInterference =
-      wouldCreateWriteToNonWritableBuffer(operand, result, aliasInfo) ||
+      wouldCreateWriteToNonWritableBuffer(operand, result, aliasInfo, state) ||
       wouldCreateReadAfterWriteInterference(operand, result, domInfo,
                                             aliasInfo);
 
@@ -599,6 +588,7 @@ bufferizableInPlaceAnalysisImpl(OpOperand &operand, OpResult result,
 /// RaW dependence violations.
 static LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                                      BufferizationAliasInfo &aliasInfo,
+                                     BufferizationState &state,
                                      const DominanceInfo &domInfo,
                                      unsigned analysisFuzzerSeed = 0) {
   if (analysisFuzzerSeed) {
@@ -615,8 +605,8 @@ static LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
       if (opOperand.get().getType().isa<TensorType>())
         if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
           if (OpResult opResult = bufferizableOp.getAliasingOpResult(opOperand))
-            if (failed(bufferizableInPlaceAnalysisImpl(opOperand, opResult,
-                                                       aliasInfo, domInfo)))
+            if (failed(bufferizableInPlaceAnalysisImpl(
+                    opOperand, opResult, aliasInfo, state, domInfo)))
               return failure();
 
   return success();
@@ -625,6 +615,7 @@ static LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
 /// Analyze all ops that are contained in `op`.
 static LogicalResult inPlaceAnalysis(Operation *op,
                                      BufferizationAliasInfo &aliasInfo,
+                                     BufferizationState &state,
                                      const DominanceInfo &domInfo,
                                      unsigned analysisFuzzerSeed = 0) {
   // Collect ops so we can build our own reverse traversal.
@@ -637,7 +628,7 @@ static LogicalResult inPlaceAnalysis(Operation *op,
     ops.push_back(op);
   });
 
-  return inPlaceAnalysis(ops, aliasInfo, domInfo, analysisFuzzerSeed);
+  return inPlaceAnalysis(ops, aliasInfo, state, domInfo, analysisFuzzerSeed);
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.
@@ -712,15 +703,9 @@ static void
 annotateOpsWithBufferizationMarkers(Operation *op,
                                     const BufferizationAliasInfo &aliasInfo) {
   op->walk([&](Operation *op) {
-    for (OpResult opResult : op->getResults()) {
+    for (OpResult opResult : op->getResults())
       if (opResult.getType().isa<TensorType>())
         setInPlaceOpResult(opResult, aliasInfo.isInPlace(opResult));
-      if (auto funcOp = dyn_cast<FuncOp>(op))
-        for (BlockArgument bbArg : funcOp.getArguments())
-          if (bbArg.getType().isa<TensorType>())
-            setInPlaceFuncArgument(bbArg,
-                                   aliasInfo.bufferizesToWritableMemory(bbArg));
-    }
   });
 }
 
@@ -739,8 +724,8 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
 
   // If the analysis fails, just return.
   Operation *op = funcOp.getOperation();
-  if (failed(
-          inPlaceAnalysis(op, aliasInfo, domInfo, options.analysisFuzzerSeed)))
+  if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
+                             options.analysisFuzzerSeed)))
     return failure();
   equivalenceAnalysis(op, aliasInfo);
 
@@ -750,7 +735,7 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
       if (failed(step->run(funcOp, state, newOps)))
         return failure();
       // Analyze ops that were created by the PostAnalysisStep.
-      if (failed(inPlaceAnalysis(newOps, aliasInfo, domInfo)))
+      if (failed(inPlaceAnalysis(newOps, aliasInfo, state, domInfo)))
         return failure();
       equivalenceAnalysis(newOps, aliasInfo);
     }

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

@@ -193,7 +193,7 @@ struct TiledLoopOpInterface
     return BufferRelation::Equivalent;
   }
 
-  bool isWritable(Operation *op, Value value) const {
+  bool isWritable(Operation *op, Value value, BufferizationState &state) const {
     // Interestingly, linalg::TiledLoopOp's bbArg can **always** be viewed
     // inplace from the perspective of ops nested under:
     //   1. Either the matching iter operand is not bufferized inplace and an

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

@@ -36,6 +36,10 @@ struct ModuleBufferizationState : public DialectBufferizationState {
   /// A mapping of ReturnOp OpOperand indices to equivalent FuncOp BBArg
   /// indices.
   DenseMap<FuncOp, DenseMap<int64_t, int64_t>> equivalentFuncArgs;
+
+  SmallVector<FuncOp> orderedFuncOps;
+
+  DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
 };
 } // namespace
 
@@ -689,6 +693,32 @@ struct FuncOpInterface
     return comprehensive_bufferize::bufferize(&funcOp.body(), state);
   }
 
+  /// Return `true` if the given function argument is writable.
+  bool isWritable(Operation *op, Value value, BufferizationState &state) const {
+    auto funcOp = cast<FuncOp>(op);
+    BlockArgument bbArg = value.dyn_cast<BlockArgument>();
+    assert(bbArg && "expected BlockArgument");
+    ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
+
+    // In a first approximation:
+    // =========================
+    // If the function is called, we can allocate on the caller side which lets
+    // us force inplace arguments at function boundaries.
+    // TODO: do not rely on this behavior.
+    if (moduleState.callerMap.find(funcOp) != moduleState.callerMap.end())
+      return true;
+
+    // Set the function arguments marked with inplaceable to be known as
+    // bufferizing to a writeable memory.
+    BoolAttr inplaceAttr = funcOp.getArgAttrOfType<BoolAttr>(
+        bbArg.getArgNumber(), BufferizableOpInterface::kInplaceableAttrName);
+    if (inplaceAttr && inplaceAttr.getValue())
+      return true;
+
+    // All other function arguments are not writable.
+    return false;
+  }
+
   bool isAllocationHoistingBarrier(Operation *op) const { return true; }
 };
 
@@ -704,46 +734,44 @@ void mlir::linalg::comprehensive_bufferize::std_ext::
   registry.addOpInterface<FuncOp, std_ext::FuncOpInterface>();
 }
 
+/// Set the attribute that triggers inplace bufferization on a FuncOp argument
+/// `bbArg`.
+static void setInPlaceFuncArgument(BlockArgument bbArg, bool inPlace) {
+  auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
+  funcOp.setArgAttr(bbArg.getArgNumber(),
+                    BufferizableOpInterface::kInplaceableAttrName,
+                    BoolAttr::get(bbArg.getContext(), inPlace));
+}
+
+/// Annotate the IR with the result of the analysis. For testing/debugging only.
+static void annotateOpsWithBufferizationMarkers(FuncOp funcOp,
+                                                BufferizationState &state) {
+  auto bufferizableOp = cast<BufferizableOpInterface>(funcOp.getOperation());
+  for (BlockArgument bbArg : funcOp.getArguments())
+    if (bbArg.getType().isa<TensorType>())
+      setInPlaceFuncArgument(bbArg, bufferizableOp.isWritable(bbArg, state));
+}
+
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     ModuleOp moduleOp, const BufferizationOptions &options) {
-  SmallVector<FuncOp> orderedFuncOps;
-  DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
-  if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
-    return failure();
-
   BufferizationState state(moduleOp, options);
   ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
+  if (failed(getFuncOpsOrderedByCalls(moduleOp, moduleState.orderedFuncOps,
+                                      moduleState.callerMap)))
+    return failure();
+
   // Interestingly, all function args that are not visible outside of a module
   // can be fully bufferized inplace by guaranteeing the CallOp is bufferized
   // inplace. Therefore, we just bufferize funcOp as if none of its results were
   // inplaceable, detect which operands are cloned internally and decide what to
   // do at call sites.
-  for (FuncOp funcOp : orderedFuncOps) {
+  for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.body().empty())
       continue;
 
-    // In a first approximation:
-    // =========================
-    // If the function is called, we can allocate on the caller side which lets
-    // us force inplace arguments at function boundaries.
-    // TODO: do not rely on this behavior.
-    if (callerMap.find(funcOp) != callerMap.end())
-      for (BlockArgument bbArg : funcOp.getArguments())
-        if (bbArg.getType().isa<TensorType>())
-          aliasInfo.setBufferizesToWritableMemory(bbArg);
-
-    // Set the function arguments marked with inplaceable to be known as
-    // bufferizing to a writeable memory.
-    for (BlockArgument bbArg : funcOp.getArguments()) {
-      BoolAttr inplaceAttr = funcOp.getArgAttrOfType<BoolAttr>(
-          bbArg.getArgNumber(), BufferizableOpInterface::kInplaceableAttrName);
-      if (inplaceAttr && inplaceAttr.getValue())
-        aliasInfo.setBufferizesToWritableMemory(bbArg);
-    }
-
     // Register extra post analysis steps. These cannot be stored in `options`
     // because `options` is immutable.
     PostAnalysisStepList extraSteps;
@@ -755,12 +783,16 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     // Analyze and bufferize funcOp.
     if (failed(runComprehensiveBufferize(funcOp, options, state, extraSteps)))
       return failure();
+
+    // Add annotations to function arguments.
+    if (options.testAnalysisOnly)
+      annotateOpsWithBufferizationMarkers(funcOp, state);
   }
 
   if (options.testAnalysisOnly)
     return success();
 
-  for (FuncOp funcOp : orderedFuncOps) {
+  for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
     if (failed(bufferizeFuncOpBoundary(funcOp, state)))

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

@@ -204,7 +204,7 @@ struct ForOpInterface
     return equivalentYield ? BufferRelation::Equivalent : BufferRelation::None;
   }
 
-  bool isWritable(Operation *op, Value value) const {
+  bool isWritable(Operation *op, Value value, BufferizationState &state) const {
     // Interestingly, scf::ForOp's bbArg can **always** be viewed
     // inplace from the perspective of ops nested under:
     //   1. Either the matching iter operand is not bufferized inplace and an