[mlir][linalg][bufferize][NFC] Conflict detection funcs take OpOperand and OpResult

[springerm]

Summary

By doing so, it is not necessary to get the OpOperand a second time via getAliasingOpOperand. Also, code slightly more readable because we do not have to deal with Optional<> return value.

Depends On D111380

[nicolasvasilache]

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

2272

Please do not drop these comments as they are not trivial to infer.
A lot of load-bearing semantic is now carried by BufferRelation::None which hides this discussion.
We need to find a way to document this to avoid surprises in the future.

[springerm]

Complete rewrite. Commit message to be updated.

[nicolasvasilache]

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

2299

This is fishy: there may be ops in the future that have an input which aliases only a part of the result.
When bufferized, these don't have enough room to fit the result.

Putting a blocker until this is discussed.

[springerm]

rebase

[springerm]

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

2299

Currently, the entire machinery assumes (conservatively) that two tensors alias either entirely or not at all. It's not just getAliasingOpResult, but also wouldCreateReadAfterWriteInterference.

This is fishy: there may be ops in the future that have an input which aliases only a part of the result. When bufferized, these don't have enough room to fit the result.

I'm trying to understand when there would not be enough room to fit the result. E.g., this would be an op where the result tensor is larger than the argument tensor? This would require some redesign. There has to be an allocation somewhere. We currently avoid these cases by having explicit "output" tensors (like in linalg.generic). Then we would always have a suitable tensor (with the right size) to alias with.

If we have to support such ops that we cannot change, we have to find a different solution.

However, my thinking is: If getAliasingOpResult returns an OpResult and getInplaceableOpResult returns null (that's presumably the case you are concerned about; and that's also the case for ExtractSliceOp), the OpOperand is guaranteed to *not* bufferize to a memory write.

"Proof" by contradiction: Let's assume that:

• getAliasingOpResult returns an OpResult
• getInplaceableOpResult returns null
• The OpOperand bufferizes to a memory write

Since getInplaceableOpResult is null and the operation is writing, it must create a new allocation and copy, i.e., bufferize out-of-place. But in that case, it cannot have an aliasing OpResult. Does this make sense?

[springerm]

address comments

[springerm]

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

2299

As discussed, special handling of ExtractSliceOp for the moment.

[springerm]

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

2272

Moved into the body of mlir::linalg::inPlaceAnalysis.

[nicolasvasilache]

seeing some feature creep here, please split and land the trivial pieces out first and then let's iterate on a more robust API / preconditions.

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

564

This fix is independent, please split it in a separate revision.

692

this NFC API change can be in its own CL, please do that and land it without review to reduce the amount of changes in this revision.

741

this NFC API change can be in its own CL, please do that and land it without review to reduce the amount of changes in this revision.

880–883

I don't understand why this API needs to change in this revision.
Losing the connection between opOperand and opResult is problematic to me here and below.

At the very least we should have strong documentation and loud assertion about the cases we allow.

977

losing the connection between opOperand and opResult is problematic to me here.

982

Since we lost the connection between opOperand and opResult, opOperand may be readonly and it does not make sense to call aliasesNonWritableBuffer on it.

2284

Let's call this one bufferizableInPlaceAnalysisImpl and add some assert that operand and result are related and cannot be random things. Atm we lose this information and it is unclear who should call what.

Let's keep the ExtractSliceOp version and make it call the Impl.

2339

Let's not jump the step here, call a bufferizableInPlaceAnalysis(ExtractSliceOp which can in turn call the lower level bufferizableInPlaceAnalysisImpl.

[springerm]

rebase

[springerm]

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

880–883

This was just for convenience. Since we already have the OpOperand in the caller. We could add another assert here, expand the comment to say that they have to match, or just abandon this revision. Either one is fine with me.

I just think that the Optional<OpOperand *> leads to unnecessarily complex (hard to parse understand) code.

982

aliasesNonWritableBuffer can be called on any Value. Should not be a problem.

The main change around this location has moved to D111379. While the new version may cover a few cases that may never appear in reality, I find it easier to understand.

In particular, I was thinking for a while that the existing conflict detection code non-writable buffer is missing a case.

%r = tensor.extract_slice %t

* %t aliases in-place write
* %r aliases non-writable buffer

This case is not covered by the existing code. The new one (D111379) does. I now believe that it is not possible to construct such an example. However, this was not immediately obvious to me. So for the sake of simplicity/understandability, my suggestion would be to go with a version similar to D111379.

[springerm]

Added asserts to functions.

[springerm]

address comments

[nicolasvasilache]

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

880–883

with the assert, this WFM, thanks!

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

public:
   ///      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(
       Operation *opToBufferize, DenseSet<OpOperand *> &usesRead,
       DenseSet<OpOperand *> &usesWrite, const DominanceInfo &domInfo) const;
 
-  /// Return true if bufferizing `opResult` inplace would create a write to a
-  /// non-writable buffer.
-  bool wouldCreateWriteToNonWritableBuffer(OpResult opResult) const;
+  /// Return true if bufferizing `opOperand` inplace with `opResult` would
+  /// create a write to a non-writable buffer.
+  bool wouldCreateWriteToNonWritableBuffer(OpOperand &opOperand,
+                                           OpResult opResult) const;
 
   /// Assume that result bufferizes in-place with one of the operation's
   /// operands. Return true if it is possible to find an inplace write W (resp.
   /// a read R) among the uses of `aliasInfo[result]`, and a read R (resp. an
   /// inplace write W) among the uses of
   /// `aliasInfo[getAliasingOpOperand(result)]`, such that W and R interfere.
   /// Interference detection is needed to determine which cases may bufferize
   /// inplace without interferences. Such cases comprise:
   ///
   /// ```
   ///  %0 = op_to_bufferize(%1)
   ///  read(%1)
   ///
   ///  %0 = op_to_bufferize(%1)
   ///  write(%0)
   ///  read(%1)
   ///
   ///  %0 = op_to_bufferize(%1)
   ///  write(%1)
   ///  read(%0)
   /// ```
   bool
-  wouldCreateReadAfterWriteInterference(OpResult result,
+  wouldCreateReadAfterWriteInterference(OpOperand &operand, 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);
   }
 
private:
   /// These are the conservative cases that can further decompose into
   /// "equivalent" buffer relationships.
   llvm::EquivalenceClasses<ValueWrapper> aliasInfo;
 
   /// Auxiliary structure to store all the equivalent buffer classes.
   llvm::EquivalenceClasses<ValueWrapper> equivalentInfo;
 };
 
-/// Analyze the `ops` to determine which OpResults are inplaceable:
+/// Analyze the `ops` to determine which OpResults are inplaceable.
 LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                               BufferizationAliasInfo &aliasInfo,
                               const DominanceInfo &domInfo);
 
 /// Bufferize one particular op.
 /// `bufferizedFunctionTypes` (resp. `globalCreator`) are expected to be
 /// non-null if `op` is a CallOpInterface (resp. GlobalCreator).
 LogicalResult

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

 /// 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([&](tensor::CastOp op) { return &op->getOpOperand(0); })
       .Case([&](ConstantOp op) { return nullptr; })
       .Case([&](ExtractSliceOp op) { return &op->getOpOperand(0); })
       // 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()];
       })
       .Case([&](InitTensorOp op) { return nullptr; })
bufferizesToMemoryWrite(OpOperand &opOperand,
   // If we have a matching OpResult, this is a write.
   // Additionally allow to restrict to only inPlace write, if so specified.
   return inPlaceSpec == InPlaceSpec::None ||
          getInPlace(opResult) == inPlaceSpec;
 }
 
 /// Returns the relationship between the operand and the its corresponding
 /// OpResult that it may alias with.
 static BufferRelation bufferRelation(OpOperand &operand) {
   return TypeSwitch<Operation *, BufferRelation>(operand.getOwner())
       // ExtractSliceOp returns a subview of the original tensor.
       .Case([&](ExtractSliceOp op) { return BufferRelation::None; })
       // All other ops: Buffers are equivalent.
       .Default([&](Operation *op) { return BufferRelation::Equivalent; });
 }
 
 //===----------------------------------------------------------------------===//

 void BufferizationAliasInfo::insertNewBufferEquivalence(Value newValue,
                                                         Value alias) {
   insertNewBufferAlias(newValue, alias);
   equivalentInfo.unionSets(newValue, alias);
 }
 
 /// Return true if, under current bufferization decisions, the buffer of `value`
 /// is not writable.
 bool BufferizationAliasInfo::aliasesNonWritableBuffer(Value value) const {
   LDBG("----Start aliasesNonWritableBuffer\n");
   for (Value v : getAliases(value)) {
     LDBG("-----------examine: " << printValueInfo(v) << '\n');
     if (bufferizesToWritableMemory(v)) {
       LDBG("-----------Value is known to be writable -> skip: "
            << printValueInfo(v) << '\n');
       continue;
     }

 /// such that W and R interfere.
 /// Such a (W, R) pair is an interference to the inplace bufferization of
 /// opResult when:
 ///   1. R is not known to 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(
-    OpResult result, const DominanceInfo &domInfo) const {
-  Optional<OpOperand *> maybeAliasingOperand = getAliasingOpOperand(result);
-  if (!maybeAliasingOperand)
-    return false;
+    OpOperand &operand, OpResult result, const DominanceInfo &domInfo) const {
+  assert(getAliasingOpOperand(result) == &operand &&
+         "operand and result do not match");
 
   Operation *opToBufferize = result.getDefiningOp();
   Value opResult = result;
-  Value opOperand = (*maybeAliasingOperand)->get();
+  Value opOperand = operand.get();
 
   LDBG("----Start wouldCreateReadAfterWriteInterference\n");
   LDBG("--------consider all aliases to root read: "
        << printValueInfo(opOperand) << "\n");
   LDBG("--------consider all aliases to root write: "
        << printValueInfo(opResult) << "\n");
 
   /// Helper function to iterate on aliases of `root` and capture the reads.
bool BufferizationAliasInfo::wouldCreateReadAfterWriteInterference(
   //  read(%3)
   // ```
   DenseSet<OpOperand *> usesRead, usesWrite;
   LDBG("--------\n");
   LDBG("--------Test reads(opOperand) vs writes(opResult)\n");
   getAliasingReads(usesRead, opOperand);
   getAliasingInplaceWrites(usesWrite, opResult);
   // Additionally, `result` is not yet bufferized and we need to check for
-  // interferences as if it were bufferized inplace: add `maybeAliasingOperand`
-  // if it is a write. This handles the case:
+  // interferences as if it were bufferized inplace: add `operand` if it is a
+  // write. This handles the case:
   //
   // ```
   //  %0 = op_to_bufferize_maybe_inplace(%1)
   //  %2 = some_alias(%1)
   //  read(%2)
   // ```
-  if (bufferizesToMemoryWrite(**maybeAliasingOperand))
-    usesWrite.insert(*maybeAliasingOperand);
+  if (bufferizesToMemoryWrite(operand))
+    usesWrite.insert(&operand);
   if (wouldCreateReadAfterWriteInterference(opToBufferize, usesRead, usesWrite,
                                             domInfo))
     return true;
 
   // Check if we can find any interference between writes to
   // aliases[`opOperand`] and reads to aliases[`opResult`]. This handles the
   // case:
   //
bool BufferizationAliasInfo::wouldCreateReadAfterWriteInterference(
   usesRead.clear();
   usesWrite.clear();
   getAliasingReads(usesRead, opResult);
   getAliasingInplaceWrites(usesWrite, opOperand);
   return wouldCreateReadAfterWriteInterference(opToBufferize, usesRead,
                                                usesWrite, domInfo);
 }
 
-/// Return true if bufferizing `opResult` inplace would create a write to a
-/// non-writable buffer.
+/// Return true if bufferizing `opOperand` inplace with `opResult` would create
+/// a write to a non-writable buffer.
 bool BufferizationAliasInfo::wouldCreateWriteToNonWritableBuffer(
-    OpResult opResult) const {
-  Optional<OpOperand *> maybeAliasingOperand = getAliasingOpOperand(opResult);
-  if (!maybeAliasingOperand || !*maybeAliasingOperand)
-    return false;
+    OpOperand &opOperand, OpResult opResult) const {
+  assert(getAliasingOpOperand(opResult) == &opOperand &&
+         "operand and result do not match");
 
   // Certain buffers are not writeable:
   //   1. A function bbArg that is not inplaceable or
   //   2. A constant op.
   assert(!aliasesNonWritableBuffer(opResult) &&
          "expected that opResult does not alias non-writable buffer");
-  bool nonWritable = aliasesNonWritableBuffer((*maybeAliasingOperand)->get());
+  bool nonWritable = aliasesNonWritableBuffer(opOperand.get());
   if (!nonWritable)
     return false;
 
   // This is a problem only if the buffer is written to via some alias.
   bool hasWrite = aliasesInPlaceWrite(opResult) ||
-                  aliasesInPlaceWrite((*maybeAliasingOperand)->get()) ||
-                  bufferizesToMemoryWrite(**maybeAliasingOperand);
+                  aliasesInPlaceWrite(opOperand.get()) ||
+                  bufferizesToMemoryWrite(opOperand);
   if (!hasWrite)
     return false;
 
   LDBG("->the corresponding buffer is not writeable\n");
   return true;
 }
 
 /// Return true if the source of a `insertSliceOp` bufferizes to an
bufferizableInPlaceAnalysisImpl(OpOperand &operand, OpResult result,
   int64_t resultNumber = result.getResultNumber();
   (void)resultNumber;
   LDBG('\n');
   LDBG("Inplace analysis for <- #" << resultNumber << " -> #"
                                    << operand.getOperandNumber() << " in "
                                    << printValueInfo(result) << '\n');
 
   bool foundInterference =
-      aliasInfo.wouldCreateWriteToNonWritableBuffer(result) ||
-      aliasInfo.wouldCreateReadAfterWriteInterference(result, domInfo);
+      aliasInfo.wouldCreateWriteToNonWritableBuffer(operand, result) ||
+      aliasInfo.wouldCreateReadAfterWriteInterference(operand, result, domInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(result);
   else
     aliasInfo.bufferizeInPlace(result, operand);
 
   LDBG("Done inplace analysis for result #" << resultNumber << '\n');
 
   return success();
 }
 
 ///
 /// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace.
 /// ===========================================================
 ///
 /// When bufferized out of place, a ExtractSlice lowers to alloc + copy. This
 /// cannot change the flow of information for either the source or the
 /// result buffers.
 ///
 /// When bufferized inplace, a ExtractSliceOp does not by itself create any read
 /// or write from memory. Instead, it has the effect of merging the alias sets
 /// of the source and the result buffers.
 ///
 /// An analysis is required to ensure inplace bufferization would not result in
 /// RaW dependence violations.
 static LogicalResult
 bufferizableInPlaceAnalysis(ExtractSliceOp extractSliceOp,
                             BufferizationAliasInfo &aliasInfo,
                             const DominanceInfo &domInfo) {
   return bufferizableInPlaceAnalysisImpl(extractSliceOp->getOpOperand(0),
                                          extractSliceOp->getOpResult(0),
                                          aliasInfo, domInfo);
 }
 
 /// Determine if `operand` can be bufferized in-place with one of the op's
 /// results. If so, set InPlaceSpec::True on the result. Otherwise, set
 /// InPlaceSpec::False on the result.
 static LogicalResult
 bufferizableInPlaceAnalysis(OpOperand &operand,
                             BufferizationAliasInfo &aliasInfo,
                             const DominanceInfo &domInfo) {
   OpResult result = getInplaceableOpResult(operand);
   if (!result)
     return success();
   return bufferizableInPlaceAnalysisImpl(operand, result, aliasInfo, domInfo);
 }
 
-/// Analyze the `ops` to determine which OpResults are inplaceable:
-///   1. First, analyze InsertSliceOp greedily: we almost never want to
-///      bufferize the tensor "inserted into" to become out-of-place.
-///   2. Walk the other ops in reverse. This is a good starter heuristic.
-///      ExtractSliceOps are interleaved with other ops in traversal order.
-///
+/// Analyze the `ops` to determine which OpResults are inplaceable. Walk ops in
+/// reverse and bufferize ops greedily. This is a good starter heuristic.
+/// ExtractSliceOps are interleaved with other ops in traversal order.
 LogicalResult mlir::linalg::inPlaceAnalysis(SmallVector<Operation *> &ops,
                                             BufferizationAliasInfo &aliasInfo,
                                             const DominanceInfo &domInfo) {
   // Walk ops in reverse for better interference analysis.
   for (Operation *op : reverse(ops)) {
     for (OpOperand &opOperand : op->getOpOperands())
       if (failed(bufferizableInPlaceAnalysis(opOperand, aliasInfo, domInfo)))
         return failure();
 
     // Special logic to analyze ExtractSliceOp.
     // Note that ExtractSliceOp analysis needs to be interleaved with other ops
     // to properly capture aliases.
     // Walk ExtractSliceOps in reverse for better clobbering analysis behavior:
     // it is easier to detect clobbers of smaller slices before larger ones.
-    if (auto extractSliceOp = dyn_cast<ExtractSliceOp>(op)) {
+    if (auto extractSliceOp = dyn_cast<ExtractSliceOp>(op))
       if (failed(
               bufferizableInPlaceAnalysis(extractSliceOp, aliasInfo, domInfo)))
         return failure();
-      continue;
-    }
   }
+
   return success();
 }
 
 /// Analyze the `funcOp` body to determine which OpResults are inplaceable.
 static LogicalResult
 inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
                           const DominanceInfo &domInfo) {
   LLVM_DEBUG(llvm::dbgs() << "\n\n");
   LDBG("Begin InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
   assert(funcOp && funcOp->getNumRegions() > 0 && !funcOp.body().empty() &&
          "expected a funcOp definition with a body");
 
   // Collect ops so we can build our own reverse traversal.
   SmallVector<Operation *> ops;
   funcOp.walk([&](Operation *op) {
     // No tensors => no buffers.
     if (none_of(op->getOperandTypes(), isaTensor) &&
         none_of(op->getResultTypes(), isaTensor))
       return;
     ops.push_back(op);