[mlir][bufferization][NFC] Remove yielded tensor analysis

[matthias-springer]

Remove the yielded tensor analysis. This analysis was used to detect cases where One-Shot Bufferize cannot deallocate buffers. Deallocation has recently been removed from One-Shot Bufferize. Buffers are now deallocated by the buffer deallocation pass. This analysis is no longer needed.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-bufferization

Changes

Remove the yielded tensor analysis. This analysis was used to detect cases where One-Shot Bufferize cannot deallocate buffers. Deallocation has recently been removed from One-Shot Bufferize. Buffers are now deallocated by the buffer deallocation pass. This analysis is no longer needed.

---

Full diff: https://github.com/llvm/llvm-project/pull/67126.diff

4 Files Affected:

• (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h (-7)
• (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h (-12)
• (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (-47)
• (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp (-39)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 1c715f8b9a53ef3..8859f6e2a5bcaf9 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -528,13 +528,6 @@ class AnalysisState {
   /// Return `true` if the given tensor has undefined contents.
   virtual bool hasUndefinedContents(OpOperand *opOperand) const;
 
-  /// Return true if the given tensor (or an aliasing tensor) is yielded from
-  /// the containing block. Also include all aliasing tensors in the same block.
-  ///
-  /// Note: In the absence of an analysis, an implementation may return true for
-  /// any given tensor.
-  virtual bool isTensorYielded(Value tensor) const;
-
   /// Return a reference to the BufferizationOptions.
   const BufferizationOptions &getOptions() const { return options; }
 
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
index 328aff07280a92b..a29af853eb21ba6 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
@@ -101,10 +101,6 @@ class OneShotAnalysisState : public AnalysisState {
   /// and store them in `undefinedTensorUses`.
   void gatherUndefinedTensorUses(Operation *op);
 
-  /// Find all tensors that are yielded/returned from a block and store them in
-  /// `yieldedTensors`. Also include all aliasing tensors in the same block.
-  void gatherYieldedTensors(Operation *op);
-
   int64_t getStatNumTensorOutOfPlace() const { return statNumTensorOutOfPlace; }
   int64_t getStatNumTensorInPlace() const { return statNumTensorInPlace; }
 
@@ -114,10 +110,6 @@ class OneShotAnalysisState : public AnalysisState {
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const override;
 
-  /// Return true if the given tensor (or an aliasing tensor) is yielded from
-  /// the containing block. Also include all aliasing tensors in the same block.
-  bool isTensorYielded(Value tensor) const override;
-
   /// Return true if the buffer of the given tensor value is written to. Must
   /// not be called for values inside not yet analyzed functions.
   bool isValueWritten(Value value) const;
@@ -261,10 +253,6 @@ class OneShotAnalysisState : public AnalysisState {
   int64_t statNumTensorOutOfPlace = 0;
   int64_t statNumTensorInPlace = 0;
 
-  /// A set of all tensors (and maybe aliasing tensors) that yielded from a
-  /// block.
-  DenseSet<Value> yieldedTensors;
-
   /// A set of uses of tensors that have undefined contents.
   DenseSet<OpOperand *> undefinedTensorUses;
 
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 57cd303d2076e73..837b22b4a82a396 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -628,53 +628,6 @@ bool AnalysisState::hasUndefinedContents(OpOperand *opOperand) const {
   return false;
 }
 
-bool AnalysisState::isTensorYielded(Value tensor) const {
-  // In the absence of analysis information, the conservative answer is "true".
-  if (!tensor.getDefiningOp<AllocTensorOp>())
-    return true;
-
-  // For AllocTensorOp results, we can do better: They do not alias with any
-  // preceding value, so we can follow SSA use-def chains and do a simple
-  // analysis.
-  SmallVector<OpOperand *> worklist;
-  DenseSet<OpOperand *> visited;
-  for (OpOperand &use : tensor.getUses())
-    worklist.push_back(&use);
-
-  while (!worklist.empty()) {
-    OpOperand *operand = worklist.pop_back_val();
-    if (visited.contains(operand))
-      continue;
-    visited.insert(operand);
-    Operation *op = operand->getOwner();
-
-    // If the op is not bufferizable, we can safely assume that the value is not
-    // yielded. (When bufferizing that op, it must handle such cases.)
-    if (!options.dynCastBufferizableOp(op))
-      continue;
-
-    // We cannot analyze through ToMemrefOps, so we have to conservatively
-    // assume that the value is yielded.
-    if (isa<ToMemrefOp>(op))
-      return true;
-
-    // Check if the op is returning/yielding.
-    if (isa<RegionBranchTerminatorOpInterface>(op))
-      return true;
-
-    // Add all aliasing Values to the worklist.
-    // Note: In the absence of detailed analysis information (e.g., there may be
-    // no function call analysis information), this `getAliasingValues` is
-    // conservative and may report additional Values as potentially aliasing.
-    for (AliasingValue alias : getAliasingValues(*operand))
-      for (OpOperand &use : alias.value.getUses())
-        worklist.push_back(&use);
-  }
-
-  // No ReturnLike op found: The value is not yielded.
-  return false;
-}
-
 // bufferization.to_memref is not allowed to change the rank.
 static void ensureToMemrefOpIsValid(Value tensor, Type memrefType) {
 #ifndef NDEBUG
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index 09205388a644720..1c85dbb5688be4b 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -181,40 +181,6 @@ void OneShotAnalysisState::createAliasInfoEntry(Value v) {
   equivalentInfo.insert(v);
 }
 
-// Gather yielded tensors in `yieldedTensors` by querying all aliases. This is
-// to ensure that such information is available during bufferization time.
-// Alias information can no longer be queried once we have started modifying
-// the IR.
-void OneShotAnalysisState::gatherYieldedTensors(Operation *op) {
-  op->walk([&](Operation *returnOp) {
-    if (!isa<RegionBranchTerminatorOpInterface>(returnOp) ||
-        !getOptions().isOpAllowed(returnOp))
-      return WalkResult::advance();
-
-    for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
-      Value returnVal = returnValOperand.get();
-      // Skip non-tensor values.
-      if (!isa<TensorType>(returnVal.getType()))
-        continue;
-
-      // Add all aliases of the returned value. But only the ones that are in
-      // the same block.
-      applyOnAliases(returnVal, [&](Value v) {
-        if (auto bbArg = dyn_cast<BlockArgument>(v)) {
-          if (bbArg.getOwner()->getParentOp() == returnOp->getParentOp())
-            yieldedTensors.insert(bbArg);
-          return;
-        }
-        Operation *definingOp = v.getDefiningOp();
-        if (definingOp->getParentOp() == returnOp->getParentOp())
-          yieldedTensors.insert(v);
-      });
-    }
-
-    return WalkResult::advance();
-  });
-}
-
 void OneShotAnalysisState::gatherUndefinedTensorUses(Operation *op) {
   op->walk([&](Operation *op) {
     // Skip unknown ops.
@@ -246,10 +212,6 @@ bool OneShotAnalysisState::isInPlace(OpOperand &opOperand) const {
   return inplaceBufferized.contains(&opOperand);
 }
 
-bool OneShotAnalysisState::isTensorYielded(Value tensor) const {
-  return yieldedTensors.contains(tensor);
-}
-
 bool OneShotAnalysisState::isValueWritten(Value value) const {
   bool isWritten = false;
   applyOnAliases(value, [&](Value val) {
@@ -1328,7 +1290,6 @@ LogicalResult bufferization::analyzeOp(Operation *op,
   bool failedAnalysis = false;
 
   // Gather some extra analysis data.
-  state.gatherYieldedTensors(op);
   state.gatherUndefinedTensorUses(op);
 
   // Analysis verification: After setting up alias/equivalence sets, each op