[mlir][Linalg] NFC - Drop AliasInfo::existsNonDominatingRead

[nicolasvasilache]

Summary

The case where a non-dominating read can be found is captured by slightly generalizing AliasInfo::wouldCreaateReadAfterWriteInterference.

This simplification will make it easier to implement bufferization across function call.

[nicolasvasilache]

Drop unused function.

[nicolasvasilache]

More cleanups.

[ftynse]

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

411

[nicolasvasilache]

Address comment.

[nicolasvasilache]

Drop dead function.

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

return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
         // ExtractSliceOp is special, when bufferized inplace it just returns an
         // alias to its operand. Its result is never inplaceable on its operand.
         .Case([&](ExtractSliceOp op) { return OpResult(); })
         // Other ops.
         .Default([&](Operation *op) { return OpResult(); });
   // clang-format on
 }
 
+/// Determine which OpOperand* will alias with `result` if the op is bufferized
+/// in place.
+/// Return None if the owner of `opOperand` does not have known
+/// bufferization aliasing behavior, which indicates that the op must allocate
+/// all of its tensor results.
+/// TODO: in the future this may need to evolve towards a list of OpOperand*.
+static Optional<OpOperand *> getAliasingOpOperand(OpResult result) {
+  if (!hasKnownBufferizationAliasingBehavior(result.getDefiningOp()))
+    return None;
+  return TypeSwitch<Operation *, OpOperand *>(result.getDefiningOp())
+      .Case([&](LinalgOp op) {
+        return op.getOutputTensorOperands()[result.getResultNumber()];
+      })
+      .Case([&](ExtractSliceOp op) { return &op->getOpOperand(0); })
+      .Case([&](InsertSliceOp op) { return &op->getOpOperand(1); })
+      .Case([&](vector::TransferWriteOp op) { return &op->getOpOperand(1); })
+      // In the case of scf::ForOp, this currently assumes the iter_args / yield
+      // are 1-1. This may fail and is verified at the end.
+      // TODO: update this.
+      .Case([&](scf::ForOp op) {
+        return &op.getIterOpOperands()[result.getResultNumber()];
+      })
+      .Default([&](Operation *op) {
+        op->dump();
+        llvm_unreachable("unexpected defining op");
+        return nullptr;
+      });
+}
+
 /// Determine which OpResult will alias with `opOperand` if the op is bufferized
 /// in place. This is a superset of `getInplaceableOpResult`.
 /// Return None if the owner of `opOperand` does not have known
 /// bufferization aliasing behavior, which indicates that the op must allocate
 /// all of its tensor results.
 /// TODO: in the future this may need to evolve towards a list of OpResult.
 static Optional<OpResult> getAliasingOpResult(OpOperand &opOperand) {
   if (!hasKnownBufferizationAliasingBehavior(opOperand.getOwner()))
public:
   /// Merge result's and operand's aliasing sets and iterate to a fixed point.
   void bufferizeInPlace(OpResult result, OpOperand &operand,
                         BufferRelation bufferRelation = BufferRelation::None);
 
   /// Set the inPlace bufferization spec to false.
   void bufferizeOutOfPlace(OpResult result);
 
   /// Return true if it is possible to find an inplace write W among the uses of
-  /// aliasInfo[rootWrite], and a read R among the uses of aliasInfo[rootRead],
+  /// aliasInfo[result], and a read R among the uses of aliasInfo[result],
   /// such that W and R interfere.
   /// Such a (W, R) pair is an interference to the inplace bufferization of
   /// rootWrite when:
   ///   1. R is not known properly dominate W (i.e. the effects of the write may
   ///      be visible from R).
   ///   2. one cannot find an intermediate clobbering write `C` to W, such that
   ///      C interleaved between W and R (i.e. W -> C -> R where -> denotes
   ///      dominance).
   bool
-  wouldCreateReadAfterWriteInterference(Value rootWrite, Value rootRead,
-                                        Operation *opToBufferize,
-                                        const DominanceInfo &domInfo) const;
-
-  /// Return true if we find any read to opOperand.get() or any of its aliases,
-  /// that does not dominate opOperand.getOwner().
-  bool existsNonDominatingRead(OpOperand &opOperand,
+  wouldCreateReadAfterWriteInterference(OpResult result,
                                         const DominanceInfo &domInfo) const;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const {
     return equivalentInfo.getLeaderValue(v1) ==
            equivalentInfo.getLeaderValue(v2);
   }
 
   /// Return true if the source of an `insertSliceOp` bufferizes to an
private:
   ///    // Or case 2:
   ///    %rootWrite = opToBufferize(%rootRead)
   ///    aliasingWriteOp(%aliasingWrite = alias(%rootWrite)) // inplace
   ///     aliasingReadOp( %aliasingRead = alias(%rootRead))
   /// ```
   ///
   /// Capture possible cases where `aliasingWriteOp(alias(%rootWrite))` has no
   /// visible effect on `aliasingReadOp(alias(%rootRead))`.
-  bool isClobberedWriteBeforeRead(Operation *opToBufferize, Value rootRead,
-                                  Value rootWrite, OpOperand &aliasingRead,
+  bool isClobberedWriteBeforeRead(Operation *opToBufferize,
+                                  OpOperand &aliasingRead,
                                   OpOperand &aliasingWrite,
                                   const DominanceInfo &domInfo) const;
 
   /// EquivalenceClasses wants comparable elements because it uses std::set.
   /// ValueWrapper wraps Value and uses pointer comparison on the defining op.
   /// This is a poor man's comparison but it's not like UnionFind needs ordering
   /// anyway ..
   struct ValueWrapper {
void BufferizationAliasInfo::bufferizeInPlace(OpResult result,
   LLVM_DEBUG(dump());
 }
 
 /// Set the inPlace bufferization spec to false.
 void BufferizationAliasInfo::bufferizeOutOfPlace(OpResult result) {
   setInPlaceOpResult(result, InPlaceSpec::False);
 }
 
-/// Return true if merging the alias sets of `rootWrite` and `rootRead` would
-/// result in a semantic change in the program (i.e. RAW violation).
-///
-/// This is the case when one can find an inplace write W among the aliases
-/// `rootWrite`, that may become an interference if W were to be bufferized
-/// inplace. A potential interference would be with respect to a read R among
-/// the aliases of `rootRead`.
-///
+/// Return true if it is possible to find an inplace write W among the uses of
+/// aliasInfo[result], and a read R among the uses of aliasInfo[result],
+/// such that W and R interfere.
 /// Such a (W, R) pair is an interference to the inplace bufferization of
-/// rootWrite when R does not properly dominate W (i.e. W may come before R
-/// along some control-flow path).
+/// rootWrite when:
+///   1. R is not known properly dominate W (i.e. the effects of the write may
+///      be visible from R).
+///   2. one cannot find an intermediate clobbering write `C` to W, such that
+///      C interleaved between W and R (i.e. W -> C -> R where -> denotes
+///      dominance).
 bool BufferizationAliasInfo::wouldCreateReadAfterWriteInterference(
-    Value rootWrite, Value rootRead, Operation *opToBufferize,
-    const DominanceInfo &domInfo) const {
+    OpResult result, const DominanceInfo &domInfo) const {
+  Optional<OpOperand *> maybeAliasingOperand = getAliasingOpOperand(result);
+  if (!maybeAliasingOperand)
+    return false;
+
+  Operation *opToBufferize = result.getDefiningOp();
+  Value root = (*maybeAliasingOperand)->get();
   LDBG("----Start wouldCreateReadAfterWriteInterference\n");
+  LDBG("--------rootValue: " << root << "\n");
 
-  // Collect all the inplace write uses of some alias of `rootWrite`.
+  // Collect:
+  //   1. all the inplace write uses of some alias of `root`.
+  //   2. all the write uses that belong to `opToBufferize`.
+  // opToBufferize is not yet inplace, we want to determine if it can be inplace
+  // so we also consider all its write uses, not just the inplace ones.
   DenseSet<OpOperand *> usesWrite;
-  auto &aliasListWrite = getAliasInfoRef(rootWrite);
-  for (Value vWrite : aliasListWrite) {
+  for (Value vWrite : getAliasInfoRef(root)) {
     for (auto &uWrite : vWrite.getUses()) {
-      if (!bufferizesToMemoryWrite(uWrite, InPlaceSpec::True))
+      if (!bufferizesToMemoryWrite(uWrite))
         continue;
+      if (uWrite.getOwner() == opToBufferize ||
+          bufferizesToMemoryWrite(uWrite, InPlaceSpec::True))
         usesWrite.insert(&uWrite);
     }
   }
+  for (Value vWrite : getAliasInfoRef(result))
+    for (auto &uWrite : vWrite.getUses())
+      if (bufferizesToMemoryWrite(uWrite, InPlaceSpec::True))
+        usesWrite.insert(&uWrite);
 
-  // Collect all the read uses of some alias of `rootRead`.
+  // Collect all the reads of some alias of `root`.
+  // opToBufferize is not yet inplace, we want to determine if it can be inplace
+  // so we also consider all read uses of its result.
   DenseSet<OpOperand *> usesRead;
-  auto &aliasListRead = getAliasInfoRef(rootRead);
-  for (Value vRead : aliasListRead) {
-    for (auto &uRead : vRead.getUses()) {
-      if (!bufferizesToMemoryRead(uRead))
-        continue;
+  auto &aliasListRead = getAliasInfoRef(root);
+  for (Value vRead : aliasListRead)
+    for (auto &uRead : vRead.getUses())
+      if (bufferizesToMemoryRead(uRead))
+        usesRead.insert(&uRead);
+  for (Value vRead : getAliasInfoRef(result))
+    for (auto &uRead : vRead.getUses())
+      if (bufferizesToMemoryRead(uRead))
         usesRead.insert(&uRead);
-    }
-  }
 
   for (OpOperand *uRead : usesRead) {
     Operation *aliasingReadOp = uRead->getOwner();
     LDBG("----++++aliasRead  #" << uRead->getOperandNumber()
                                 << " in: " << *aliasingReadOp << '\n');
     for (OpOperand *uWrite : usesWrite) {
-      // Don't consider self-use of the same operand.
-      // Uses within the same op is fine though.
+      // Don't consider self-use of the same operand for interference.
+      // Multiple different uses within the same op is fair game though.
       if (uWrite == uRead)
         continue;
+
       Operation *aliasingWriteOp = uWrite->getOwner();
       LDBG("----    aliasWrite #" << uWrite->getOperandNumber()
                                   << " in: " << *aliasingWriteOp << '\n');
-      // If read and written value already alias, no interference would be added
-      // by bufferizing inplace.
-      if (getAliasInfoRef(uRead->get()).contains(uWrite->get()))
+      // If the candidate write is the one that produces the read value (in the
+      // SSA def-use sense), this is not considered an interference.
+      if (getInplaceableOpResult(*uWrite) == uRead->get())
         continue;
       // If aliasingReadOp properly dominates aliasingWriteOp, the read cannot
       // be affected by the write: there is no interference.
       if (domInfo.properlyDominates(aliasingReadOp, aliasingWriteOp))
         continue;
       // At this point, aliasingWriteOp properly dominates aliasingReadOp or
       // there is no clear dominance and we need to be conservative.
       LDBG("---->found RaW interference\n");
       LDBG("     Interfering  read (op #" << uRead->getOperandNumber()
                                           << "): " << *aliasingReadOp << '\n');
       LDBG("     Interfering write (op #" << uWrite->getOperandNumber()
                                           << "): " << *aliasingWriteOp << '\n');
-      LDBG("         aliases rootRead: " << rootRead << '\n');
-      LDBG("         aliases rootWrite: " << rootWrite << '\n');
       LDBG("---->opportunity to clobber RaW interference\n");
-      if (isClobberedWriteBeforeRead(opToBufferize, rootRead, rootWrite, *uRead,
-                                     *uWrite, domInfo)) {
-
+      if (isClobberedWriteBeforeRead(opToBufferize, *uRead, *uWrite, domInfo)) {
         LDBG("---->clobbered! -> skip\n");
         continue;
       }
 
       LDBG("---->not clobbered -> found an interference\n");
       return true;
     }
   }
   LDBG("----No interference found\n");
   return false;
 }
 
-/// Return true if we find any read to opOperand.get() or any of its aliases,
-/// that does not dominate opOperand.getOwner().
-bool BufferizationAliasInfo::existsNonDominatingRead(
-    OpOperand &opOperand, const DominanceInfo &domInfo) const {
-  LDBG("----Start existsNonDominatingRead\n");
-  Operation *op = opOperand.getOwner();
-  for (Value alias : getAliasInfoRef(opOperand.get())) {
-    for (OpOperand &wantReadUse : alias.getUses()) {
-      LDBG("--------current operand #" << wantReadUse.getOperandNumber() << ": "
-                                       << *(wantReadUse.getOwner()) << '\n');
-      if (!bufferizesToMemoryRead(wantReadUse)) {
-        LDBG("------------not a read -> skip\n");
-        continue;
-      }
-      if (&wantReadUse == &opOperand) {
-        LDBG("------------self-read is not an interference -> skip\n");
-        continue;
-      }
-      if (domInfo.properlyDominates(wantReadUse.getOwner(), op)) {
-        LDBG("------------read properly dominates -> skip\n");
-        continue;
-      }
-      LDBG("----found interfering read of " << wantReadUse.get() << '\n');
-      return true;
-    }
-  }
-  return false;
-}
-
 /// Return true if the source of a `insertSliceOp` bufferizes to an
 /// equivalent ExtractSliceOp.
 bool BufferizationAliasInfo::isSourceEquivalentToAMatchingExtractSliceOp(
     InsertSliceOp insertSliceOp) const {
   auto leaderIt = equivalentInfo.findLeader(insertSliceOp.source());
   for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
        ++mit) {
     if (areEquivalentExtractSliceOps(

 }
 
 /// Return true if there is a write that:
 ///   1. Properly dominates aliasingReadOp.
 ///   2. Is properly dominated by aliasingWriteOp.
 ///   3. Clobbers the write that would be interfering with the read.
 ///
 bool BufferizationAliasInfo::isClobberedWriteBeforeRead(
-    Operation *opToBufferize, Value rootRead, Value rootWrite,
-    OpOperand &aliasingRead, OpOperand &aliasingWrite,
+    Operation *opToBufferize, OpOperand &aliasingRead, OpOperand &aliasingWrite,
     const DominanceInfo &domInfo) const {
   Operation *aliasingReadOp = aliasingRead.getOwner();
   Operation *aliasingWriteOp = aliasingWrite.getOwner();
   assert(!domInfo.properlyDominates(aliasingReadOp, aliasingWriteOp) &&
          "Unexpected aliasingReadOp properly dominates aliasingWriteOp");
 
-  assert(((rootRead.isa<OpResult>() &&
-           rootRead.getDefiningOp() == opToBufferize) ||
-          (rootWrite.isa<OpResult>() &&
-           rootWrite.getDefiningOp() == opToBufferize)) &&
-         "Expected rootRead or rootWrite to be produced by opToBufferize");
-
   // Bail if the write does not dominate the read: it may clobber but only on
   // a strict subset of paths, which is not enough for safety.
   if (!domInfo.dominates(aliasingWriteOp, aliasingReadOp)) {
     LDBG("---->no clobbering: write does not dominate read\n");
     return false;
   }
 
   // The case `opToBufferize` isa ExtractSliceOp is important enough that we
if (wouldCreateAliasingWriteToNonWriteableBuffer)
     LDBG("->the corresponding buffer is not writeable\n");
   else
     LDBG("->bufferizes to writeable inplace buffer\n");
 
   // In any of extractSliceOp.result's aliases, can we find 2 such that we hit
   // an interfering write?
   OpResult r = extractSliceOp->getResult(0);
   OpOperand &s = extractSliceOp->getOpOperand(0);
-  bool foundInterference = wouldCreateAliasingWriteToNonWriteableBuffer ||
-                           // Do not consider (s, s) and (r, r) as all the
-                           // aliasings already exist by construction; we are
-                           // interested in new interfering aliases only.
-                           aliasInfo.wouldCreateReadAfterWriteInterference(
-                               s.get(), r, extractSliceOp, domInfo) ||
-                           aliasInfo.wouldCreateReadAfterWriteInterference(
-                               r, s.get(), extractSliceOp, domInfo);
+  bool foundInterference =
+      wouldCreateAliasingWriteToNonWriteableBuffer ||
+      aliasInfo.wouldCreateReadAfterWriteInterference(r, domInfo);
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(r);
   else
     aliasInfo.bufferizeInPlace(r, s);
 
   LDBG("Done inplace analysis for extract_slice\n");
 
   return success();
bufferizableInPlaceAnalysis(OpOperand &operand, OpResult result,
   int64_t resultNumber = result.getResultNumber();
   (void)resultNumber;
   LDBG('\n');
   LDBG("Inplace analysis for result #" << resultNumber << " (operand #"
                                        << operand.getOperandNumber() << ") in "
                                        << result << '\n');
 
   // `result` must bufferize to a writeable buffer to be a candidate.
-  // This means the use->def chain not backpropagate to a function that is
-  // not inplaceable or to a constant op to be considered.
+  // This means the operand  must not alias either:
+  //   1. a function bbArg that is not inplaceable or
+  //   2. a constant op.
+  // to be considered for inplace bufferization
   bool wouldCreateAliasingWriteToNonWriteableBuffer =
       aliasInfo.aliasesNonWriteableBuffer(operand);
   if (wouldCreateAliasingWriteToNonWriteableBuffer)
     LDBG("->the corresponding buffer is not writeable\n");
   else
     LDBG("->bufferizes to writeable inplace buffer\n");
 
-  Value s = operand.get(), r = result;
+  assert(result == getInplaceableOpResult(operand));
   bool foundInterference =
       wouldCreateAliasingWriteToNonWriteableBuffer ||
-      aliasInfo.existsNonDominatingRead(operand, domInfo) ||
-      // Do not consider (s, s) and (r, r) as all the aliasings already
-      // exist by construction; we are interested in new interfering aliases
-      // only.
-      aliasInfo.wouldCreateReadAfterWriteInterference(s, r, op, domInfo) ||
-      aliasInfo.wouldCreateReadAfterWriteInterference(r, s, op, domInfo);
+      aliasInfo.wouldCreateReadAfterWriteInterference(result, domInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(result);
   else
     // TODO: Atm, all inplace bufferizations yield equivalent tensors. Support
     // more cases on a per-need basis.
     aliasInfo.bufferizeInPlace(
         result, operand, BufferizationAliasInfo::BufferRelation::Equivalent);