[mlir][Linalg] NFC - Add support to specify that a tensor value is known to bufferize to writeable memory

[nicolasvasilache]

Summary

This change allows better interop with external clients of comprehensive bufferization functions
but is otherwise NFC for the MLIR pass itself.

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

public:
   /// `alias`.
   void insertNewBufferAlias(Value newValue, Value alias);
 
   /// Insert an info entry for `newValue` and merge its alias set with that of
   /// `alias`. Additionally, merge their equivalence classes.
   void insertNewBufferEquivalence(Value newValue, Value alias);
 
   /// Return true if the buffer to which `operand` would bufferize aliases a
-  /// buffer that is known to not be writeable. This implies that the matching
+  /// buffer that is known to not be writable. This implies that the matching
   /// OpResult cannot be bufferized inplace.
-  bool aliasesNonWriteableBuffer(OpOperand &operand) const;
+  bool aliasesNonWritableBuffer(OpOperand &operand) const;
 
   /// Return true if the buffer to which `operand` would bufferize is equivalent
   /// to some buffer write.
   bool aliasesInPlaceWrite(Value v) const;
 
   /// Set the inPlace bufferization spec to true.
   /// Merge result's and operand's aliasing sets and iterate to a fixed point.
   void bufferizeInPlace(OpResult result, OpOperand &operand,
public:
   /// Return true if the source of an `insertSliceOp` bufferizes to an
   /// equivalent ExtractSliceOp.
   bool isSourceEquivalentToAMatchingInplaceExtractSliceOp(
       tensor::InsertSliceOp insertSliceOp) const;
 
   /// Apply `fun` to all the members of the equivalence class of `v`.
   void applyOnEquivalenceClass(Value v, function_ref<void(Value)> fun) const;
 
+  /// Return true if the value is known to bufferize to writable memory.
+  bool bufferizesToWritableMemory(Value v) const;
+
+  /// Specify that the value is known to bufferize to writable memory.
+  void setBufferizesToWritableMemory(Value v);
+
   /// Print to `os`.
   void printAliases(raw_ostream &os) const;
   void printEquivalences(raw_ostream &os) const;
 
   /// Print to `errs()`.
   void dumpAliases() const;
   void dumpEquivalences() const;
 
private:
   ///
   /// Capture possible cases where `aliasingWriteOp(alias(%opResult))` has no
   /// visible effect on `aliasingReadOp(alias(%opOperand))`.
   bool isClobberedWriteBeforeRead(Operation *opToBufferize,
                                   OpOperand &aliasingRead,
                                   OpOperand &aliasingWrite,
                                   const DominanceInfo &domInfo) const;
 
+  /// Set of tensors that are known to bufferize to writable memory.
+  llvm::DenseSet<Value> bufferizeToWritableMemory;
+
   /// Auxiliary structure to store all the values a given value aliases with.
   /// 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;
 };

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

 /// `alias`. Additionally, merge their equivalence classes.
 void BufferizationAliasInfo::insertNewBufferEquivalence(Value newValue,
                                                         Value alias) {
   insertNewBufferAlias(newValue, alias);
   equivalentInfo.unionSets(newValue, alias);
 }
 
 /// Return true if the buffer to which `operand` would bufferize aliases a
-/// buffer that is known to not be writeable. This implies that the matching
+/// buffer that is known to not be writable. This implies that the matching
 /// OpResult cannot be bufferized inplace.
-bool BufferizationAliasInfo::aliasesNonWriteableBuffer(
+bool BufferizationAliasInfo::aliasesNonWritableBuffer(
     OpOperand &operand) const {
-  LDBG("----Start aliasesNonWriteableBuffer\n");
+  LDBG("----Start aliasesNonWritableBuffer\n");
   LDBG("-------for -> #" << operand.getOperandNumber() << ": "
                          << printOperationInfo(operand.getOwner()) << '\n');
   for (Value v : getAliases(operand.get())) {
     LDBG("-----------examine: " << printValueInfo(v) << '\n');
+    if (bufferizesToWritableMemory(v)) {
+      LDBG("-----------Value is known to be writeable -> skip: "
+           << printValueInfo(v) << '\n');
+      continue;
+    }
+
     if (auto bbArg = v.dyn_cast<BlockArgument>()) {
       if (getInPlace(bbArg) == InPlaceSpec::True) {
         LDBG("-----------bbArg is writeable -> skip: " << printValueInfo(bbArg)
                                                        << '\n');
         continue;
       }
       LDBG("-----------notWriteable\n");
       return true;
     }
 
     if (Operation *op = v.getDefiningOp()) {
       if (isa<ConstantOp>(op) || !hasKnownBufferizationAliasingBehavior(op)) {
-        LDBG("-----------notWriteable\n");
+        LDBG("-----------notWritable\n");
         return true;
       }
     }
   }
-  LDBG("---->operand is writeable\n");
+  LDBG("---->operand is writable\n");
   return false;
 }
 
+bool BufferizationAliasInfo::bufferizesToWritableMemory(Value v) const {
+  return bufferizeToWritableMemory.count(v) > 0;
+}
+
+/// Specify that the value is known to bufferize to writable memory.
+void BufferizationAliasInfo::setBufferizesToWritableMemory(Value v) {
+  bufferizeToWritableMemory.insert(v);
+}
+
 /// Return true if the buffer to which `operand` would bufferize is equivalent
 /// to some buffer write.
 bool BufferizationAliasInfo::aliasesInPlaceWrite(Value value) const {
   LDBG("----Start aliasesInPlaceWrite\n");
   LDBG("-------for : " << printValueInfo(value) << '\n');
   for (Value v : getAliases(value)) {
     for (auto &use : v.getUses()) {
       if (bufferizesToMemoryWrite(use, InPlaceSpec::True)) {

 bufferizableInPlaceAnalysis(ExtractSliceOp extractSliceOp,
                             BufferizationAliasInfo &aliasInfo,
                             const DominanceInfo &domInfo) {
   LDBG('\n');
   LDBG("Inplace analysis for extract_slice: "
        << printOperationInfo(extractSliceOp) << '\n');
 
   // If `extractSliceOp` were to be bufferized inplace, it cannot end up
-  // aliasing a write into a non-writeable buffer.
-  bool wouldCreateAliasingWriteToNonWriteableBuffer =
+  // aliasing a write into a non-writable buffer.
+  bool wouldCreateAliasingWriteToNonWritableBuffer =
       aliasInfo.aliasesInPlaceWrite(extractSliceOp.result()) &&
-      aliasInfo.aliasesNonWriteableBuffer(extractSliceOp->getOpOperand(0));
+      aliasInfo.aliasesNonWritableBuffer(extractSliceOp->getOpOperand(0));
 
-  if (wouldCreateAliasingWriteToNonWriteableBuffer)
-    LDBG("->the corresponding buffer is not writeable\n");
+  if (wouldCreateAliasingWriteToNonWritableBuffer)
+    LDBG("->the corresponding buffer is not writable\n");
   else
-    LDBG("->bufferizes to writeable inplace buffer\n");
+    LDBG("->bufferizes to writable 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 ||
+      wouldCreateAliasingWriteToNonWritableBuffer ||
       aliasInfo.wouldCreateReadAfterWriteInterference(r, domInfo);
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(r);
   else
     aliasInfo.bufferizeInPlace(r, s);
 
   LDBG("Done inplace analysis for extract_slice\n");
 
bufferizableInPlaceAnalysis(OpOperand &operand, OpResult result,
 
   int64_t resultNumber = result.getResultNumber();
   (void)resultNumber;
   LDBG('\n');
   LDBG("Inplace analysis for <- #" << resultNumber << " -> #"
                                    << operand.getOperandNumber() << " in "
                                    << printValueInfo(result) << '\n');
 
-  // `result` must bufferize to a writeable buffer to be a candidate.
+  // `result` must bufferize to a writable buffer to be a candidate.
   // 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");
+  bool wouldCreateAliasingWriteToNonWritableBuffer =
+      aliasInfo.aliasesNonWritableBuffer(operand);
+  if (wouldCreateAliasingWriteToNonWritableBuffer)
+    LDBG("->the corresponding buffer is not writable\n");
   else
-    LDBG("->bufferizes to writeable inplace buffer\n");
+    LDBG("->bufferizes to writable inplace buffer\n");
 
   assert(result == getInplaceableOpResult(operand));
   bool foundInterference =
-      wouldCreateAliasingWriteToNonWriteableBuffer ||
+      wouldCreateAliasingWriteToNonWritableBuffer ||
       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(
inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
   funcOp.walk([&](Operation *op) {
     // No tensors => no buffers.
     if (none_of(op->getOperandTypes(), isaTensor) &&
         none_of(op->getResultTypes(), isaTensor))
       return;
     ops.push_back(op);
   });
 
+  // Set the function arguments marked with inplaceable to be known as
+  // bufferizing to a writeable memory.
+  for (BlockArgument bbArg : funcOp.getArguments()) {
+    BoolAttr inplaceAttr = funcOp.getArgAttrOfType<BoolAttr>(
+        bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName);
+    if (inplaceAttr && inplaceAttr.getValue())
+      aliasInfo.setBufferizesToWritableMemory(bbArg);
+  }
+
   LogicalResult res = inPlaceAnalysis(ops, aliasInfo, domInfo);
   LDBG("End InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
 
   return res;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization entry-point for functions.