[mlir][bufferize][linalg] Fix aliasing reads detection in analysis

[springerm]

Summary

Take into account values that are read via an alias such as ExtractSliceOp.

Note: It is currently not possible to trigger this bug because ops are always analyzed bottom-to-top. However, as the traveral is just a heuristic, and other traversal orders could be specified in the future, this should be fixed now to avoid surprising behavior in the future.

Depends On D116687

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

public:
   /// Return true if `opOperand` bufferizes to a memory write. Return true` if
   /// the op is not bufferizable.
   bool bufferizesToMemoryWrite(OpOperand &opOperand) const;
 
   /// Return true if `opOperand` does neither read nor write but bufferizes to
   /// an alias. Return false if the op is not bufferizable.
   bool bufferizesToAliasOnly(OpOperand &opOperand) const;
 
+  /// Compute all OpOperands that reading `value` and add them to `result`.
+  /// Also takes into account ops that create an alias but do not read by
+  /// themselves (e.g., ExtractSliceOp).
+  void getValueReads(DenseSet<OpOperand *> &result, Value value) const;
+
   /// Return true if the given value is read by an op that bufferizes to a
   /// memory read. Also takes into account ops that create an alias but do not
   /// read by themselves (e.g., ExtractSliceOp).
   bool isValueRead(Value value) const;
 
   /// Starting from `value`, follow the use-def chain in reverse, always
   /// selecting the aliasing OpOperands. Find and return Values for which
   /// `condition` evaluates to true. OpOperands of such matching Values are not

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

if (auto bufferizableOp =
           dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
     return bufferizableOp.bufferizesToAliasOnly(opOperand, *this);
 
   // Unknown op that returns a tensor. The inplace analysis does not support it.
   // Conservatively return false.
   return false;
 }
 
-/// Return true if the given value is read by an op that bufferizes to a memory
-/// read. Also takes into account ops that create an alias but do not read by
+/// Compute all OpOperands that reading `value` and add them to `result`.
+/// Also takes into account ops that create an alias but do not read by
 /// themselves (e.g., ExtractSliceOp).
-bool mlir::linalg::comprehensive_bufferize::BufferizationState::isValueRead(
-    Value value) const {
+void mlir::linalg::comprehensive_bufferize::BufferizationState::getValueReads(
+    DenseSet<OpOperand *> &result, Value value) const {
   assert(value.getType().isa<TensorType>() && "expected TensorType");
   SmallVector<OpOperand *> workingSet;
   for (OpOperand &use : value.getUses())
     workingSet.push_back(&use);
 
   while (!workingSet.empty()) {
     OpOperand *uMaybeReading = workingSet.pop_back_val();
     // Skip over all ops that neither read nor write (but create an alias).
     if (bufferizesToAliasOnly(*uMaybeReading))
       for (OpOperand &use : getAliasingOpResult(*uMaybeReading).getUses())
         workingSet.push_back(&use);
     if (bufferizesToMemoryRead(*uMaybeReading))
-      return true;
-  }
+      result.insert(uMaybeReading);
+  }
 
-  return false;
+/// Return true if the given value is read by an op that bufferizes to a memory
+/// read. Also takes into account ops that create an alias but do not read by
+/// themselves (e.g., ExtractSliceOp).
+bool mlir::linalg::comprehensive_bufferize::BufferizationState::isValueRead(
+    Value value) const {
+  assert(value.getType().isa<TensorType>() && "expected TensorType");
+  DenseSet<OpOperand *> readingOperands;
+  getValueReads(readingOperands, value);
+  return !readingOperands.empty();
 }
 
 // Starting from `value`, follow the use-def chain in reverse, always selecting
 // the aliasing OpOperands. Find and return Values for which `condition`
 // evaluates to true. OpOperands of such matching Values are not traversed any
 // further.
 llvm::SetVector<Value> mlir::linalg::comprehensive_bufferize::
     BufferizationState::findValueInReverseUseDefChain(

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

 /// involving aliases of the given OpOperand are checked.
 static bool wouldCreateReadAfterWriteInterference(
     OpOperand &operand, const DominanceInfo &domInfo, BufferizationState &state,
     const BufferizationAliasInfo &aliasInfo,
     bool checkConsistencyOnly = false) {
   // Helper function to iterate on aliases of `root` and capture the reads.
   auto getAliasingReads = [&](DenseSet<OpOperand *> &res, Value root) {
     aliasInfo.applyOnAliases(root, [&](Value alias) {
-      for (auto &use : alias.getUses())
+      for (auto &use : alias.getUses()) {
+        if (state.bufferizesToMemoryRead(use)) {
           // Read to a value that aliases root.
-        if (state.bufferizesToMemoryRead(use))
           res.insert(&use);
+          continue;
+        }
+        if (state.bufferizesToMemoryWrite(use))
+          // Tensor not read but written. The buffer is completely overwritten.
+          continue;
+
+        OpResult opResult = state.getAliasingOpResult(use);
+        if (!opResult)
+          continue;
+        // Result may be read.
+        state.getValueReads(res, opResult);
+      }
     });
   };
 
   // Helper function to iterate on aliases of `root` and capture the writes.
   auto getAliasingInplaceWrites = [&](DenseSet<OpOperand *> &res, Value root) {
     aliasInfo.applyOnAliases(root, [&](Value alias) {
       for (auto &use : alias.getUses())
         // Inplace write to a value that aliases root.