[mlir][bufferization] Generalize and rename isMemoryWritten

[springerm]

Summary

The name of the method was confusing. It is bufferizesToMemoryWrite, but from the perspective of OpResults.

bufferizesToMemoryWrite(OpResult) now supports ops with regions that do not have aliasing OpOperands (such as scf.if). These ops no longer need to implement isMemoryWrite.

Depends On: D141683

[nicolasvasilache]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

382

Can we separate this change from this PR?
It seems it is unrelated.

555

Does this really need to be overridable?
If not consider taking out of the interface.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

122–139

Could we have some IR examples to support the cases ?

128

grammo

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

630

missing to

635

Do these correspond to the 3 cases in the doc?
If so, it would be good to give them a Case 1/2/3 to connect them.

635

This feels related to DpsInterface ?
Should be using that more in the impl of BufferizableOpInterface ?

683

can we drop the alwaysIncludeLeaves change in this PR and better justify it in a separate PR?

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

128

Nice red here, above and below!

[springerm]

address comments

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

555

Yes, e.g., tensor.empty must override this. Added an example (Counter-Example) to the InterfaceMethod description.

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

635

We are using it already to provide a default implementation for DPS ops (DstBufferizableOpInterfaceExternalModel). We may be able to do more interesting things if op interfaces could inherit from another op interface.

[aartbik]

LGTM for the very tiny sparse part ;-)

[nicolasvasilache]

Fine to keep the inheritance part for a followup PR

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

635

We may be able to do more interesting things if op interfaces could inherit from another op interface.

See: https://reviews.llvm.org/D140198

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

public:
   /// Return true if `opOperand` bufferizes to a memory read. Return `true` if
   /// the op is not bufferizable.
   bool bufferizesToMemoryRead(OpOperand &opOperand) const;
 
   /// 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 the given `value` bufferizes to a memory write. Return
+  /// true if the value is a block argument. Return `true` if the defining op is
+  /// not bufferizable. Otherwise, consult the BufferizableOpInterface.
+  bool bufferizesToMemoryWrite(Value value) 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;
 
   /// Return true if a copy can always be avoided when allocating a new tensor
   /// for the given OpOperand.
   bool canOmitTensorCopy(OpOperand &opOperand) const;
 
public:
   /// { 2, 7, 8, 5 }
   ///
   /// If `followEquivalentOnly` is set, only equivalent OpOperands are selected.
   SetVector<Value>
   findValueInReverseUseDefChain(Value value,
                                 llvm::function_ref<bool(Value)> condition,
                                 bool followEquivalentOnly = false) const;
 
   /// Find the Values of the last preceding write of a given Value.
   ///
   /// Note: Unknown ops are handled conservatively and assumed to be writes.
   /// Furthermore, BlockArguments are also assumed to be writes. There is no
   /// analysis across block boundaries.
   ///
   /// Note: When reaching an end of the reverse SSA use-def chain, that value
   /// is returned regardless of whether it is a memory write or not.
   SetVector<Value> findLastPrecedingWrite(Value value) const;

 /// This is the default implementation of
 /// BufferizableOpInterface::getBufferType. Should not be called from other
 /// places.
 FailureOr<BaseMemRefType>
 defaultGetBufferType(Value value, const BufferizationOptions &options,
                      const DenseMap<Value, BaseMemRefType> &fixedTypes);
 
 /// This is the default implementation of
+/// BufferizableOpInterface::resultBufferizesToMemoryWrite. Should not be called
+/// from other places.
+bool defaultResultBufferizesToMemoryWrite(OpResult opResult,
+                                          const AnalysisState &state);
+
+/// This is the default implementation of
 /// BufferizableOpInterface::isRepetitiveRegion. Should not be called from other
 /// places.
 bool defaultIsRepetitiveRegion(BufferizableOpInterface bufferizableOp,
                                unsigned index);
 } // namespace detail
 
 } // namespace bufferization
 } // namespace mlir

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

let methods = [
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
           // Does not have to be implemented for OpOperands that do not have an
           // aliasing OpResult.
           llvm_unreachable("bufferizesToMemoryWrite not implemented");
          }]
       >,
       InterfaceMethod<
-          /*desc=*/[{
-            Return `true` if the given OpResult is a memory write. This is the
-            case if in the following cases:
+          Return `true` if the given OpResult bufferizes to a memory write.
+          This is the same property as `bufferizesToMemoryWrite`, but from The
+          perspective of OpResults.
 
-            * The corresponding aliasing OpOperand bufferizes to a memory write.
-            * Or: There is no corresponding aliasing OpOperand.
+          This method will never be called on OpResults that do not have a
+          tensor type.
+
+          This method has a default implementation. By default, it returns
+          `true` if any of the following three cases applies.
+
+          1. There is no corresponding aliasing OpOperand.
 
-            If the OpResult has multiple aliasing OpOperands, this method
-            returns `true` if at least one of them bufferizes to a memory write.
+             Example: `tensor.generate ... : tensor<10xf32>`
+             The op fills a newly allocated buffer and bufferizes to a memory
+             write.
+
+             Counter-example: bufferization.alloc_tensor
+             The op just allocates and does not specifiy the data of the tensor,
+             so resultBufferizesToMemoryWrite is overridden to return false.
+
+          2. At least one aliasing OpOperand bufferizes to a memory write.
+
+             Example: `tensor.insert %f into %t[...] : tensor<?xf32>`
+             The destination OpOperand bufferizes to a memory write, so the
+             result also bufferizes to a memory write.
+
+          3. At least one aliasing OpOperand's value is defined inside the
+             defining op of the given OpResult and it is a memory write or the
+             reverse SSA use-def chain ends in the defining op.
+
+             According to this rule, an aliasing OpOperand value that is defined
+             inside this op and is bufferizing to a memory write makes the given
+             OpResult bufferize to a memory write.
+
+             Example:
+             ```
+             %r = scf.if ... -> tensor<?xf32> {
+               %1 = tensor.insert %f into %t[...] : tensor<?xf32>
+               scf.yield %1 : tensor<?xf32>
+             } else { ... }
+             ```
+             The scf.if result bufferizes to a memory write because %1 (an
+             OpResult defined inside the scf.if op) bufferizes to a memory
+             write.
           }],
-          /*retType=*/"bool",
-          /*methodName=*/"isMemoryWrite",
-          /*args=*/(ins "::mlir::OpResult":$opResult,
-                        "const ::mlir::bufferization::AnalysisState &":$state),
-          /*methodBody=*/"",
-          /*defaultImplementation=*/[{
-            auto bufferizableOp =
-                cast<BufferizableOpInterface>($_op.getOperation());
-            SmallVector<OpOperand*> opOperands =
-              bufferizableOp.getAliasingOpOperand(opResult, state);
-            if (opOperands.empty())
-              return true;
-            return llvm::any_of(
-                opOperands,
-                [&](OpOperand *operand) {
-                  return bufferizableOp.bufferizesToMemoryWrite(*operand,
-                                                                state);
-                });
-          }]
+        /*methodName=*/"resultBufferizesToMemoryWrite",
+          assert(opResult.getDefiningOp() == $_op.getOperation() &&
+                 "invalid OpResult");
+          return bufferization::detail::defaultResultBufferizesToMemoryWrite(
+              opResult, state);
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return `true` if the given OpOperand must bufferize in-place. Alias
           sets and inplace attributes will be set up accordingly before making
           any other bufferization decisions. This method will never be called on
           OpOperands that do not have a tensor type.
         }],

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

let arguments = (ins Variadic<Index>:$dynamic_sizes,
                        OptionalAttr<AnyAttr>:$memory_space);
 
   let results = (outs AnyTensor:$result);
 
   let extraClassDeclaration = [{
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
 
-    bool isMemoryWrite(OpResult opResult, const AnalysisState &state);
+    bool resultBufferizesToMemoryWrite(OpResult opResult,
+                                       const AnalysisState &state);
 
     bool bufferizesToAllocation(OpResult opResult) { return true; }
 
     bool bufferizesToMemoryRead(OpOperand &opOperand,
                                 const AnalysisState &state);
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

if (auto bufferizableOp =
           getOptions().dynCastBufferizableOp(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;
 }
 
+bool AnalysisState::bufferizesToMemoryWrite(Value value) const {
+  auto opResult = value.dyn_cast<OpResult>();
+  if (!opResult)
+    return true;
+  auto bufferizableOp = getOptions().dynCastBufferizableOp(value);
+  if (!bufferizableOp)
+    return true;
+  return bufferizableOp.resultBufferizesToMemoryWrite(opResult, *this);
+}
+
 /// 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 AnalysisState::isValueRead(Value value) const {
   assert(value.getType().isa<TensorType>() && "expected TensorType");
   SmallVector<OpOperand *> workingSet;
   for (OpOperand &use : value.getUses())
     workingSet.push_back(&use);
llvm::SetVector<Value> AnalysisState::findValueInReverseUseDefChain(
   }
 
   return result;
 }
 
 // Find the Values of the last preceding write of a given Value.
 llvm::SetVector<Value>
 AnalysisState::findLastPrecedingWrite(Value value) const {
-  return findValueInReverseUseDefChain(value, [&](Value value) {
-    Operation *op = value.getDefiningOp();
-    if (!op)
-      return true;
-    auto bufferizableOp = options.dynCastBufferizableOp(op);
-    if (!bufferizableOp)
-      return true;
-    return bufferizableOp.isMemoryWrite(value.cast<OpResult>(), *this);
-  });
+  return findValueInReverseUseDefChain(
+      value, [&](Value v) { return this->bufferizesToMemoryWrite(v); });
 }
 
 AnalysisState::AnalysisState(const BufferizationOptions &options)
     : AnalysisState(options, TypeID::get<AnalysisState>()) {}
 
 AnalysisState::AnalysisState(const BufferizationOptions &options, TypeID type)
     : options(options), type(type) {
   for (const BufferizationOptions::AnalysisStateInitFn &fn :
#endif // NDEBUG
   if (failed(memrefType))
     return failure();
   ensureToMemrefOpIsValid(value, *memrefType);
   return rewriter
       .create<bufferization::ToMemrefOp>(value.getLoc(), *memrefType, value)
       .getResult();
 }
 
+bool bufferization::detail::defaultResultBufferizesToMemoryWrite(
+    OpResult opResult, const AnalysisState &state) {
+  auto bufferizableOp = cast<BufferizableOpInterface>(opResult.getDefiningOp());
+  SmallVector<OpOperand *> opOperands =
+      bufferizableOp.getAliasingOpOperand(opResult, state);
+
+  // Case 1: OpResults that have no aliasing OpOperand usually bufferize to
+  // memory writes.
+  if (opOperands.empty())
+    return true;
+
+  // Case 2: If an aliasing OpOperand bufferizes to a memory write, the OpResult
+  // may bufferize to a memory write.
+  if (llvm::any_of(opOperands, [&](OpOperand *operand) {
+        return state.bufferizesToMemoryWrite(*operand);
+      }))
+    return true;
+
+  // Case 3: Check if a nested aliasing OpOperand value bufferizes to a memory
+  // write. (Or: The reverse SSA use-def chain ends inside the reigon.) In that
+  // case, the OpResult bufferizes to a memory write. E.g.:
+  //
+  // %0 = "some_writing_op" : tensor<?xf32>
+  // %r = scf.if ... -> tensor<?xf32> {
+  //   scf.yield %0 : tensor<?xf32>
+  // } else {
+  //   %1 = "another_writing_op"(%0) : tensor<?xf32>
+  //   scf.yield %1 : tensor<?xf32>
+  // }
+  // "some_reading_op"(%r)
+  //
+  // %r bufferizes to a memory write because an aliasing OpOperand value (%1)
+  // bufferizes to a memory write and the defining op is inside the scf.if.
+  //
+  // Note: This treatment of surrouding ops is useful for ops that have a
+  // region but no OpOperand such as scf.if or scf.execute_region. It simplifies
+  // the analysis considerably.
+  //
+  // "another_writing_op" in the above example should be able to bufferize
+  // inplace in the absence of another read of %0. However, if the scf.if op
+  // would not be considered a "write", the analysis would detect the
+  // following conflict:
+  //
+  // * read = some_reading_op
+  // * lastWrite = %0  (Note: The last write of %r would be a set: {%0, %1}.)
+  // * conflictingWrite = %1
+  //
+  auto isMemoryWriteInsideOp = [&](Value v) {
+    Operation *op = getOwnerOfValue(v);
+    if (!opResult.getDefiningOp()->isAncestor(op))
+      return false;
+    return state.bufferizesToMemoryWrite(v);
+  };
+  for (OpOperand *operand : opOperands) {
+    if (!state
+             .findValueInReverseUseDefChain(operand->get(),
+                                            isMemoryWriteInsideOp,
+                                            /*followEquivalentOnly=*/false)
+             .empty())
+      return true;
+  }
+  return false;
+}
+
 FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
     Value value, const BufferizationOptions &options,
     const DenseMap<Value, BaseMemRefType> &fixedTypes) {
   assert(value.getType().isa<TensorType>() && "expected tensor type");
 
   // No further analysis is possible for a block argument.
   if (value.isa<BlockArgument>())
     return bufferization::getMemRefType(value, options);

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

if (!dealloc)
     return success();
 
   rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator());
   if (failed(options.createDealloc(rewriter, loc, *alloc)))
     return failure();
   return success();
 }
 
-bool AllocTensorOp::isMemoryWrite(OpResult opResult,
+bool AllocTensorOp::resultBufferizesToMemoryWrite(OpResult opResult,
                                                   const AnalysisState &state) {
   // AllocTensorOps do not write unless they have a `copy` value.
   return static_cast<bool>(getCopy());
 }
 
 bool AllocTensorOp::bufferizesToMemoryRead(OpOperand &opOperand,
                                            const AnalysisState &state) {
   assert(opOperand.getOperandNumber() == getNumOperands() - 1 &&
          "expected copy operand");

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

for (OpResult opResult : op->getOpResults()) {
 
       // If there is no preceding memory write, the tensor contents are
       // undefined.
       // Note: If `findLastPrecedingWrite` reaches the end of the reverse SSA
       // use-def chain, it returns that value, regardless of whether it is a
       // memory write or not.
       SetVector<Value> lastWrites = findLastPrecedingWrite(opResult);
       bool isUndefined = llvm::none_of(lastWrites, [&](Value lastWrite) {
-        if (auto bufferizableOp = getOptions().dynCastBufferizableOp(lastWrite))
-          return bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>(),
-                                              *this);
-        return true;
+        return this->bufferizesToMemoryWrite(lastWrite);
       });
       if (isUndefined)
         for (OpOperand &use : opResult.getUses())
           undefinedTensorUses.insert(&use);
     }
 
     return WalkResult::advance();
   });
do {
     if (a->isProperAncestor(b))
       return false;
     if (domInfo.properlyDominates(a, b))
       return true;
   } while ((a = a->getParentOp()));
   return false;
 }
 
-/// Return `true` if the given tensor value is a memory write. Most values are
-/// tensor writes, but ops that define a tensor SSA value without specifying its
-/// contents (e.g., alloc_tensor) are not.
-static bool isMemoryWrite(Value value, const AnalysisState &state) {
-  auto opResult = value.dyn_cast<OpResult>();
-  if (!opResult)
-    return true;
-  auto bufferizableOp = state.getOptions().dynCastBufferizableOp(value);
-  if (!bufferizableOp)
-    return true;
-  return bufferizableOp.isMemoryWrite(opResult, state);
-}
-
 /// Return `true` if op dominance can be used to rule out read-after-write
 /// conflicts wrt. the given reads and writes.
 ///
 /// Op dominance can often be used to rule out potential conflicts such as
 /// "read" happens before "write". E.g., the following IR is not a RaW conflict
 /// because the the read happens *before* the write.
 ///
 /// %0 = ... : tensor<?xf32>
for (OpOperand *uWrite : usesWrite) {
     }
     if (*commonEnclosingRegion != r)
       return false;
   }
 
   // In case of a read, take the region which the read value is defined.
   for (OpOperand *uRead : usesRead) {
     // Optimization: Skip reads of values that have no defined contents.
-    if (!isMemoryWrite(uRead->get(), state))
+    if (!state.bufferizesToMemoryWrite(uRead->get()))
       continue;
     Region *r = getEnclosingRepetitiveRegion(uRead->get(), options);
     if (!commonEnclosingRegion.has_value()) {
       commonEnclosingRegion = r;
       continue;
     }
     if (*commonEnclosingRegion != r)
       return false;

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

getAliasingOpOperand(Operation *op, OpResult opResult,
     assert(executeRegionOp.getRegion().getBlocks().size() == 1 &&
            "expected exactly 1 block");
     auto yieldOp = dyn_cast<scf::YieldOp>(
         executeRegionOp.getRegion().front().getTerminator());
     assert(yieldOp && "expected scf.yield terminator in scf.execute_region");
     return {&yieldOp->getOpOperand(resultNum)};
   }
 
-  // TODO: For better bufferization results, this could return `true` only if
-  // there is a memory write in the region.
-  bool isMemoryWrite(Operation *op, OpResult opResult,
-                     const AnalysisState &state) const {
-    // Similar to scf.if, results of this op are always considered memory writes
-    // in the analysis. This is a useful pattern for all ops that have tensor
-    // OpResults but no tensor OpOperands. By default, `isMemoryWrite` is
-    // implemented in terms of `bufferizesToMemoryWrite`, which does not work on
-    // ops without OpOperands.
-    return true;
-  }
-
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
     assert(executeRegionOp.getRegion().getBlocks().size() == 1 &&
            "only 1 block supported");
     auto yieldOp =
         cast<scf::YieldOp>(executeRegionOp.getRegion().front().getTerminator());
     TypeRange newResultTypes(yieldOp.getResults());
getAliasingOpOperand(Operation *op, OpResult opResult,
     // branches are considered to be aliasing with the result.
     auto ifOp = cast<scf::IfOp>(op);
     size_t resultNum = std::distance(op->getOpResults().begin(),
                                      llvm::find(op->getOpResults(), opResult));
     return {&ifOp.thenYield()->getOpOperand(resultNum),
             &ifOp.elseYield()->getOpOperand(resultNum)};
   }
 
-  // TODO: For better bufferization results, this could return `true` only if
-  // there is a memory write in one (or both) of the branches. Since this is not
-  // allowed at the moment, we should never encounter scf.ifs that yield
-  // unmodified tensors. Such scf.yield ops could just fold away.
-  bool isMemoryWrite(Operation *op, OpResult opResult,
-                     const AnalysisState &state) const {
-    // IfOp results are always considered memory writes in the analysis. This
-    // design decision simplifies the analysis considerably. E.g., consider the
-    // following test case:
-    //
-    // %0 = "some_writing_op" : tensor<?xf32>
-    // %r = scf.if %c -> (tensor<?xf32>) {
-    //   scf.yield %0
-    // } else {
-    //   %1 = "another_writing_op"(%0) : tensor<?xf32>
-    // }
-    // "some_reading_op"(%r)
-    //
-    // "another_writing_op" in the above example should be able to bufferize
-    // inplace in the absence of another read of %0. However, if the scf.if op
-    // would not be considered a "write", the analysis would detect the
-    // following conflict:
-    //
-    // * read = some_reading_op
-    // * lastWrite = %0  (Note: The last write of %r would be a set: {%0, %1}.)
-    // * conflictingWrite = %1
-    //
-    // For more details, check the "scf.IfOp" section of the design document.
-    return true;
-  }
-
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     OpBuilder::InsertionGuard g(rewriter);
     auto ifOp = cast<scf::IfOp>(op);
 
     // Compute bufferized result types.
     SmallVector<Type> newTypes;
     for (Value result : ifOp.getResults()) {

mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp

getAliasingOpOperand(Operation *op, OpResult opResult,
     assert(assumingOp.getDoRegion().getBlocks().size() == 1 &&
            "expected exactly 1 block");
     auto yieldOp = dyn_cast<shape::AssumingYieldOp>(
         assumingOp.getDoRegion().front().getTerminator());
     assert(yieldOp && "expected shape.assuming_yield terminator");
     return {&yieldOp->getOpOperand(resultNum)};
   }
 
-  // TODO: For better bufferization results, this could return `true` only if
-  // there is a memory write in the region.
-  bool isMemoryWrite(Operation *op, OpResult opResult,
-                     const AnalysisState &state) const {
-    // Similar to scf.if, results of this op are always considered memory writes
-    // in the analysis. This is a useful pattern for all ops that have tensor
-    // OpResults but no tensor OpOperands. By default, `isMemoryWrite` is
-    // implemented in terms of `bufferizesToMemoryWrite`, which does not work on
-    // ops without OpOperands.
-    return true;
-  }
-
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto assumingOp = cast<shape::AssumingOp>(op);
     assert(assumingOp.getDoRegion().getBlocks().size() == 1 &&
            "only 1 block supported");
     auto yieldOp = cast<shape::AssumingYieldOp>(
         assumingOp.getDoRegion().front().getTerminator());
 

mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp

BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 };
 
 struct NewOpInterface
     : public BufferizableOpInterface::ExternalModel<NewOpInterface,
                                                     sparse_tensor::NewOp> {
-  bool isMemoryWrite(Operation *op, OpResult opResult,
+  bool resultBufferizesToMemoryWrite(Operation *op, OpResult opResult,
                                      const AnalysisState &state) const {
     // NewOps allocate but do not write.
     return false;
   }
 
   bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
     return true;
   }
 };