[mlir][linalg][bufferize] Store analysis results in BufferizationAliasInfo

[springerm]

Summary

• Store inplace bufferization decisions in inplaceBufferized.
• Remove InPlaceSpec. Use a bool instead.
• Use BufferizableOpInterface::bufferizesToWritableMemory and bufferizesToWritableMemory instead of getInPlace(BlockArgument). The analysis does not care about inplacability of block arguments. It only cares whether the buffer can be written to or not.
• The kInPlaceResultsAttrName op attribute is for testing purposes only.

This commit further decouples BufferizationAliasInfo from other dialects such as SCF.

Depends On D113374

[springerm]

update

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

let methods = [
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           llvm_unreachable("bufferize not implemented");
           return failure();
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
-          Return `true` if the given OpOperand can be written to in-place. This
-          is the case for most ops, but some ops such as ConstantOp may
-          bufferize to non-writable (read-only) memory locations. This method
-          will never be called on OpResults that do not have a tensor type.
+          Return `true` if the given Value can be written to in-place. Value is
+          either an OpResult of this operation or a BlockArgument of a block of
+          this operation.
+
+          Most OpResult buffers can be written to, but some ops such as
+          ConstantOp may bufferize to non-writable (read-only) memory locations.
+          Therefore, by default, this method returns `true` for OpResults. This
+          method will never be called on OpResults that do not have a tensor
+          type.
+
+          Whether a BlockArgument can be written to or not depends on the
+          operation. This method conservatively returns `false`. This method
+          will never be called on BlockArguments that do not have a tensor type.
         }],
         /*retType=*/"bool",
         /*methodName=*/"isWritable",
-        /*args=*/(ins "OpResult":$opResult),
+        /*args=*/(ins "Value":$value),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
-          return true;
+          return value.isa<OpResult>();
         }]
       >
   ];
 
   let extraClassDeclaration = [{
     /// Return `true` if the given OpOperand creates an alias but does neither
     /// read nor write. This implies that `bufferizesToMemoryRead` and
     /// `bufferizesToMemoryWrite` must return `false`. This method will never

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

public:
 
   // TODO: Move these out of BufferizationAliasInfo.
   /// 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);
 
+  /// Mark a value as in-place bufferized.
+  void markInPlace(OpResult v) { inplaceBufferized.insert(v); }
+
+  /// Return `true` if a value was marked as in-place bufferized.
+  bool isInPlace(OpResult opResult) const;
+
   /// 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:
   using EquivalenceClassRangeType = llvm::iterator_range<
       llvm::EquivalenceClasses<Value, ValueComparator>::member_iterator>;
   /// Check that aliasInfo for `v` exists and return a reference to it.
   EquivalenceClassRangeType getAliases(Value v) const;
 
   /// Set of tensors that are known to bufferize to writable memory.
   llvm::DenseSet<Value> bufferizeToWritableMemory;
 
+  /// Set of all OpResults that were decided to bufferize in-place.
+  llvm::DenseSet<OpResult> inplaceBufferized;
+
   /// Auxiliary structure to store all the values a given value may alias with.
   /// Alias information is "may be" conservative: In the presence of branches, a
   /// value may alias with one of multiple other values. The concrete aliasing
   /// value may not even be known at compile time. All such values are
   /// considered to be aliases.
   llvm::EquivalenceClasses<Value, ValueComparator> aliasInfo;
 
   /// Auxiliary structure to store all the equivalent buffer classes. Equivalent

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

static Value lookup(const BlockAndValueMapping &bvm, Value key) {
   }
   LDBG("In func:\n" << *parentOp << "\nNO VALUE FOR KEY: " << key << '\n');
   (void)parentOp;
   return Value();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific attribute manipulation.
-// These could be simplified with helper structs on the side, for now attributes
-// allow simple embedding in the IR which simplifies testing.
-// This could also be folded in BufferizationAliasInfo or a Bufferizer class
-// that uses BufferizationAliasInfo.
+// These are for testing and debugging only. Bufferization information is
+// stored in BufferizationAliasInfo. When run with `testAnalysisOnly`, the IR
+// is annotated with the results of the analysis (copied from
+// BufferizationAliasInfo), so that they can be checked in tests.
 //===----------------------------------------------------------------------===//
 
 /// Attribute marker to specify op results that can be bufferized inPlace.
 constexpr StringLiteral kInPlaceResultsAttrName = "__inplace_results_attr__";
 
-// TODO: proper enum.
-enum class InPlaceSpec {
-  False,
-  True,
-  None,
-};
-
-static StringRef stringify(InPlaceSpec val) {
-  switch (val) {
-  case InPlaceSpec::False:
-    return "false";
-  case InPlaceSpec::True:
-    return "true";
-  case InPlaceSpec::None:
-    return "none";
-  }
-  return "";
-}
-
-static Optional<InPlaceSpec> symbolize(StringRef str) {
-  return StringSwitch<Optional<InPlaceSpec>>(str)
-      .Case("false", InPlaceSpec::False)
-      .Case("true", InPlaceSpec::True)
-      .Case("none", InPlaceSpec::None)
-      .Default(None);
-}
-
 /// Mark whether OpResult can actually be bufferized inplace.
-/// If `inPlace` is `InPlaceSpec::True`, the use-def chain analysis has
-/// guaranteed that no subsequent write would occur to the bufferized
-/// tensor value (i.e. the result can be bufferized inPlace).
-static void setInPlaceOpResult(OpResult opResult,
-                               InPlaceSpec inPlace = InPlaceSpec::True) {
+/// If `inPlace` is `true`, the use-def chain analysis has guaranteed that no
+/// subsequent write would occur to the bufferized tensor value (i.e. the result
+/// can be bufferized inplace).
+static void setInPlaceOpResult(OpResult opResult, bool inPlace) {
   if (!opResult)
     return;
 
   Operation *op = opResult.getOwner();
   auto attr =
       op->getAttr(kInPlaceResultsAttrName).dyn_cast_or_null<ArrayAttr>();
   SmallVector<StringRef> inPlaceVector =
       attr ? SmallVector<StringRef>(
                  llvm::to_vector<4>(attr.getAsValueRange<StringAttr>()))
-           : SmallVector<StringRef>(op->getNumResults(),
-                                    stringify(InPlaceSpec::None));
-  LDBG("->set inPlace=" << stringify(inPlace) << " <- #"
-                        << opResult.getResultNumber() << ": "
-                        << printOperationInfo(op) << "\n");
-  inPlaceVector[opResult.getResultNumber()] = stringify(inPlace);
+           : SmallVector<StringRef>(op->getNumResults(), "false");
+  LDBG("->set inPlace=" << inPlace << " <- #" << opResult.getResultNumber()
+                        << ": " << printOperationInfo(op) << "\n");
+  inPlaceVector[opResult.getResultNumber()] = inPlace ? "true" : "false";
   op->setAttr(kInPlaceResultsAttrName,
               OpBuilder(op).getStrArrayAttr(inPlaceVector));
 }
 
-/// Get the InPlaceSpec attribute entry `kInPlaceResultsAttrName` for
-/// `opResult`. If the result is `InPlaceSpec::True`, the use-def chain analysis
-/// has guaranteed that no subsequent read of the tensor value occurs and the
-/// result can be buferized inPlace.
-/// If no InPlaceSpec attribute has been set for `opResult`, return
-/// InPlaceSpec::None.
-LLVM_ATTRIBUTE_UNUSED static InPlaceSpec getInPlace(OpResult opResult) {
-  if (!opResult)
-    return InPlaceSpec::None;
-
-  Operation *op = opResult.getOwner();
-  auto attr =
-      op->getAttr(kInPlaceResultsAttrName).dyn_cast_or_null<ArrayAttr>();
-  if (!attr)
-    return InPlaceSpec::None;
-
-  // Must return a proper value.
-  return *symbolize(*(attr.getAsValueRange<StringAttr>().begin() +
-                      opResult.getResultNumber()));
-}
-
-/// Get inPlace information for `bbArg`.
-/// FuncOp allow argument attributes, we use those to encode the information.
-/// BlockArgument of other ops delegate to their owner's parent op.
-static InPlaceSpec getInPlace(BlockArgument bbArg) {
-  if (auto funcOp = dyn_cast<FuncOp>(bbArg.getOwner()->getParentOp())) {
-    BoolAttr inplaceAttr = funcOp.getArgAttrOfType<BoolAttr>(
-        bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName);
-    if (!inplaceAttr)
-      return InPlaceSpec::None;
-    return inplaceAttr.getValue() ? InPlaceSpec::True : InPlaceSpec::False;
-  }
-  // Interestingly, scf::ForOp's and TiledLoop's bbArg can **always** be viewed
-  // inplace from the perspective of ops nested under:
-  //   1. Either the matching iter operand is not bufferized inplace and an
-  //      alloc + optional copy makes the bbArg itself inplaceable.
-  //   2. Or the matching iter operand is bufferized inplace and bbArg just
-  //      bufferizes to that too.
-  if (isa<scf::ForOp, TiledLoopOp>(bbArg.getOwner()->getParentOp()))
-    return InPlaceSpec::True;
-  // Unknown cases.
-  return InPlaceSpec::None;
-}
-
 /// Set the attribute that triggers inplace bufferization on a FuncOp argument
 /// `bbArg`.
-static void
-setInPlaceFuncArgument(BlockArgument bbArg,
-                       InPlaceSpec inPlaceSpec = InPlaceSpec::True) {
+static void setInPlaceFuncArgument(BlockArgument bbArg, bool inPlace) {
   auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
-  funcOp.setArgAttr(
-      bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName,
-      BoolAttr::get(bbArg.getContext(), inPlaceSpec == InPlaceSpec::True));
+  funcOp.setArgAttr(bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName,
+                    BoolAttr::get(bbArg.getContext(), inPlace));
 }
 
 /// Remove the attribute that triggers inplace bufferization on a FuncOp
 /// argument `bbArg`.
 static void removeBufferizationFuncArguments(BlockArgument bbArg) {
   auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
   funcOp.removeArgAttr(bbArg.getArgNumber(),
                        LinalgDialect::kBufferLayoutAttrName);
   funcOp.removeArgAttr(bbArg.getArgNumber(),
                        LinalgDialect::kInplaceableAttrName);
 }
 
-static InPlaceSpec getInPlace(Value v) {
-  if (auto bbArg = v.dyn_cast<BlockArgument>())
-    return getInPlace(bbArg);
-  return getInPlace(v.cast<OpResult>());
-}
-
 //===----------------------------------------------------------------------===//
 // Printing helpers.
 //===----------------------------------------------------------------------===//
 
 /// Helper method printing the bufferization information of a buffer / tensor.
 static void printTensorOrBufferInfo(std::string prefix, Value value,
                                     AsmState &state, llvm::raw_ostream &os) {
   if (!value.getType().isa<ShapedType>())
     return;
   os << prefix;
   value.printAsOperand(os, state);
   os << " : " << value.getType();
-  if (getInPlace(value) == InPlaceSpec::None)
-    return;
-  os << " [InPlace=" << stringify(getInPlace(value)) << "]";
 }
 
 /// Print the operation name and bufferization information.
 static std::string printOperationInfo(Operation *op, bool prefix) {
   std::string result;
   llvm::raw_string_ostream os(result);
   AsmState state(op->getParentOfType<mlir::FuncOp>());
   StringRef tab = prefix ? "\n[" DEBUG_TYPE "]\t" : "";
areEquivalentExtractSliceOps(const BufferizationAliasInfo &aliasInfo,
   if (!aliasInfo.areEquivalentBufferizedValues(st.source(), sti.dest()))
     return false;
   if (!sameOffsetsSizesAndStrides(st, sti, isEqualConstantIntOrValue))
     return false;
   return true;
 }
 
 /// Return true if opOperand has been decided to bufferize in-place.
-static bool isInplaceMemoryWrite(OpOperand &opOperand) {
+static bool isInplaceMemoryWrite(OpOperand &opOperand,
+                                 const BufferizationAliasInfo &aliasInfo) {
   // Ops that do not bufferize to a memory write, cannot be write in-place.
   if (!bufferizesToMemoryWrite(opOperand))
     return false;
   OpResult opResult = getAliasingOpResult(opOperand);
-  return opResult && getInPlace(opResult) == InPlaceSpec::True;
+  return opResult && aliasInfo.isInPlace(opResult);
 }
 
 BufferizationAliasInfo::BufferizationAliasInfo(Operation *rootOp) {
   rootOp->walk([&](Operation *op) {
     for (Value v : op->getResults())
       if (v.getType().isa<TensorType>())
         createAliasInfoEntry(v);
     for (Region &r : op->getRegions())
for (OpResult opResult : bufferizableOp->getOpResults()) {
       if (opResult.getType().isa<TensorType>())
         if (bufferizableOp.mustBufferizeInPlace(opResult)) {
           SmallVector<OpOperand *> operands =
               bufferizableOp.getAliasingOpOperand(opResult);
           assert(!operands.empty() &&
                  "expected that OpResult has aliasing OpOperand");
           for (OpOperand *operand : operands)
             aliasInfo.unionSets(operand->get(), opResult);
-          setInPlaceOpResult(opResult, InPlaceSpec::True);
+          markInPlace(opResult);
         }
     }
   });
 }
 
 /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
 /// beginning the alias and equivalence sets only contain `v` itself.
 void BufferizationAliasInfo::createAliasInfoEntry(Value v) {
void BufferizationAliasInfo::insertNewBufferEquivalence(Value newValue,
   insertNewBufferAlias(newValue, alias);
   equivalentInfo.unionSets(newValue, alias);
 }
 
 /// Return true if, under current bufferization decisions, the buffer of `value`
 /// is not writable.
 static bool aliasesNonWritableBuffer(Value value,
                                      const BufferizationAliasInfo &aliasInfo) {
-  LDBG("----Start aliasesNonWritableBuffer\n");
+  LDBG("WRITABILITY ANALYSIS FOR " << printValueInfo(value) << "\n");
   bool foundNonWritableBuffer = false;
   aliasInfo.applyOnAliases(value, [&](Value v) {
-    LDBG("-----------examine: " << printValueInfo(v) << '\n');
-    if (aliasInfo.bufferizesToWritableMemory(v)) {
-      LDBG("-----------Value is known to be writable -> skip: "
-           << printValueInfo(v) << '\n');
+    // Some values are known to be writable.
+    if (aliasInfo.bufferizesToWritableMemory(v))
       return;
-    }
 
-    if (auto bbArg = v.dyn_cast<BlockArgument>()) {
-      if (getInPlace(bbArg) == InPlaceSpec::True) {
-        LDBG("-----------bbArg is writable -> skip: " << printValueInfo(bbArg)
-                                                      << '\n');
+    // Query BufferizableOpInterface to see if the OpResult is writable.
+    // TODO: Out-of-place bufferized OpResult could be considered writable.
+    if (auto bufferizableOp = v.getDefiningOp<BufferizableOpInterface>())
+      if (bufferizableOp && bufferizableOp.isWritable(v))
         return;
-      }
-      LDBG("-----------notWritable bbArg\n");
-      foundNonWritableBuffer = true;
+
+    // Query BufferizableOpInterface to see if the BlockArgument is writable.
+    if (auto bbArg = v.dyn_cast<BlockArgument>())
+      if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(
+              bbArg.getOwner()->getParentOp()))
+        if (bufferizableOp.isWritable(bbArg))
           return;
-    }
 
-    auto bufferizableOp = dyn_cast<BufferizableOpInterface>(v.getDefiningOp());
-    if (!bufferizableOp || !bufferizableOp.isWritable(v.cast<OpResult>())) {
-      // Unknown ops are treated conservatively: Assume that it is illegal to
-      // write to their OpResults in-place.
-      LDBG("-----------notWritable op\n");
-      foundNonWritableBuffer = true;
-      return;
-    }
   });
 
-  if (!foundNonWritableBuffer)
-    LDBG("---->value is writable\n");
+  if (foundNonWritableBuffer)
+    LDBG("--> NON WRITABLE\n");
+  else
+    LDBG("--> WRITABLE\n");
 
   return foundNonWritableBuffer;
 }
 
 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.
 static bool aliasesInPlaceWrite(Value value,
                                 const BufferizationAliasInfo &aliasInfo) {
   LDBG("----Start aliasesInPlaceWrite\n");
   LDBG("-------for : " << printValueInfo(value) << '\n');
   bool foundInplaceWrite = false;
   aliasInfo.applyOnAliases(value, [&](Value v) {
     for (auto &use : v.getUses()) {
-      if (isInplaceMemoryWrite(use)) {
+      if (isInplaceMemoryWrite(use, aliasInfo)) {
         LDBG("-----------wants to bufferize to inPlace write: "
              << printOperationInfo(use.getOwner()) << '\n');
         foundInplaceWrite = true;
         return;
       }
     }
   });
 
   if (!foundInplaceWrite)
     LDBG("----------->does not alias an inplace write\n");
 
   return foundInplaceWrite;
 }
 
+/// Return `true` if a value was marked as in-place bufferized.
+bool BufferizationAliasInfo::isInPlace(OpResult opResult) const {
+  bool inplace = inplaceBufferized.contains(opResult);
+#ifndef NDEBUG
+  if (inplace) {
+    auto bufferizableOp =
+        dyn_cast<BufferizableOpInterface>(opResult.getDefiningOp());
+    assert(bufferizableOp &&
+           "expected that in-place bufferized op is bufferizable");
+    SmallVector<OpOperand *> operands =
+        bufferizableOp.getAliasingOpOperand(opResult);
+    for (OpOperand *operand : operands)
+      assert(areAliasingBufferizedValues(operand->get(), opResult) &&
+             "expected that in-place bufferized OpResult aliases with "
+             "aliasing OpOperand");
+  }
+#endif // NDEBUG
+  return inplace;
+}
+
 /// Set the inPlace bufferization spec to true.
 void BufferizationAliasInfo::bufferizeInPlace(OpResult result,
                                               OpOperand &operand) {
-  setInPlaceOpResult(result, InPlaceSpec::True);
+  markInPlace(result);
   aliasInfo.unionSets(result, operand.get());
   // Dump the updated alias analysis.
   LLVM_DEBUG(dumpAliases());
   if (bufferRelation(operand) == BufferRelation::Equivalent)
     equivalentInfo.unionSets(result, operand.get());
   // Dump the updated equivalence analysis.
   LLVM_DEBUG(dumpEquivalences());
 }
 
 /// Set the inPlace bufferization spec to false.
 void BufferizationAliasInfo::bufferizeOutOfPlace(OpResult result) {
-  setInPlaceOpResult(result, InPlaceSpec::False);
+  if (inplaceBufferized.contains(result))
+    inplaceBufferized.erase(result);
 }
 
 /// 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.
 ///
 /// When reaching the end of a chain (BlockArgument or Value without aliasing
auto getAliasingReads = [&](DenseSet<OpOperand *> &res, Value root) {
     });
   };
 
   // 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.
-        if (isInplaceMemoryWrite(use))
+        if (isInplaceMemoryWrite(use, aliasInfo))
           res.insert(&use);
     });
   };
 
   // Collect reads and writes of all aliases of OpOperand and OpResult.
   DenseSet<OpOperand *> usesRead, usesWrite;
   getAliasingReads(usesRead, operand.get());
   getAliasingReads(usesRead, result);
static Value getResultBuffer(OpBuilder &b, OpResult result,
   if (!llvm::all_of(aliasingOperands, [&](OpOperand *o) {
         return lookup(bvm, o->get()) == operandBuffer;
       })) {
     op->emitError("result buffer is ambiguous");
     return Value();
   }
 
   // If bufferizing out-of-place, allocate a new buffer.
-  bool needCopy = getInPlace(result) != InPlaceSpec::True;
-  if (needCopy) {
+  if (!aliasInfo.isInPlace(result)) {
     // Ops with multiple aliasing operands can currently not bufferize
     // out-of-place.
     assert(
         aliasingOperands.size() == 1 &&
         "ops with multiple aliasing OpOperands cannot bufferize out-of-place");
     Location loc = op->getLoc();
     // Allocate the result buffer.
     Value resultBuffer = createNewAllocDeallocPairForShapedValue(
static LogicalResult bufferize(OpBuilder &b, FuncOp funcOp,
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
-/// Determine if `operand` can be bufferized in-place with `result`. If so, set
-/// InPlaceSpec::True on the result. Otherwise, set InPlaceSpec::False on the
-/// result.
+/// Determine if `operand` can be bufferized in-place with `result`.
 static LogicalResult
 bufferizableInPlaceAnalysisImpl(OpOperand &operand, OpResult result,
                                 BufferizationAliasInfo &aliasInfo,
                                 const DominanceInfo &domInfo) {
 #ifndef NDEBUG
   SmallVector<OpOperand *> opOperands = getAliasingOpOperand(result);
   assert(llvm::find(opOperands, &operand) != opOperands.end() &&
          "operand and result do not match");
else
     aliasInfo.bufferizeInPlace(result, operand);
 
   LDBG("Done inplace analysis for result #" << resultNumber << '\n');
 
   return success();
 }
 
 /// 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.
+/// results.
 ///
 /// Even if an op does not read or write, it may still create an alias when
 /// bufferized in-place. An example of such ops is tensor.extract_slice.
 ///
 /// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace:
 ///
 /// When bufferized out of place, an ExtractSliceOp lowers to alloc + copy. This
 /// cannot change the flow of information for either the source or the
LogicalResult mlir::linalg::comprehensive_bufferize::initTensorElimination(
 
   WalkResult status = funcOp->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Is this a matching OpOperand?
       if (!anchorMatchFunc(operand))
         continue;
 
       SetVector<Value> maybeInitTensor =
-          findValueInReverseUseDefChain(operand.get(), [](Value val) {
+          findValueInReverseUseDefChain(operand.get(), [&](Value val) {
             // Continue traversal until this function returns true.
             OpResult opResult = val.dyn_cast<OpResult>();
             if (!opResult)
               return true;
-            if (getInPlace(opResult) != InPlaceSpec::True)
+            if (!aliasInfo.isInPlace(opResult))
               return true;
             // Only equivalent tensors are supported at the moment.
             // TODO: Support cases such as extract_slice(init_tensor).
             SmallVector<OpOperand *> opOperands =
                 getAliasingOpOperand(opResult);
             if (!llvm::all_of(opOperands, [](OpOperand *operand) {
                   return bufferRelation(*operand) == BufferRelation::Equivalent;
                 }))

 /// Note that the newly inserted ExtractSliceOp may have to bufferize
 /// out-of-place due to RaW conflicts.
 LogicalResult mlir::linalg::comprehensive_bufferize::
     eliminateInsertSliceAnchoredInitTensorOps(FuncOp funcOp,
                                               BufferizationAliasInfo &aliasInfo,
                                               DominanceInfo &domInfo) {
   return initTensorElimination(
       funcOp, aliasInfo, domInfo,
-      [](OpOperand &operand) {
+      [&](OpOperand &operand) {
         auto insertSliceOp = dyn_cast<InsertSliceOp>(operand.getOwner());
         if (!insertSliceOp)
           return false;
         // Only inplace bufferized InsertSliceOps are eligible.
-        if (getInPlace(insertSliceOp->getOpResult(0)) != InPlaceSpec::True)
+        if (!aliasInfo.isInPlace(insertSliceOp->getOpResult(0)))
           return false;
         return &operand == &insertSliceOp->getOpOperand(0) /*source*/;
       },
       [](OpBuilder &b, Location loc, OpOperand &operand) {
         auto insertSliceOp = cast<InsertSliceOp>(operand.getOwner());
         auto extractOp = b.create<tensor::ExtractSliceOp>(
             loc, insertSliceOp.dest(), insertSliceOp.getMixedOffsets(),
             insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
         return extractOp.result();
       });
 }
 
+/// Annotate the IR with the result of the analysis. For testing/debugging only.
+static void
+annotateOpsWithBufferizationMarkers(Operation *op,
+                                    const BufferizationAliasInfo &aliasInfo) {
+  op->walk([&](Operation *op) {
+    for (OpResult opResult : op->getResults()) {
+      if (opResult.getType().isa<TensorType>())
+        setInPlaceOpResult(opResult, aliasInfo.isInPlace(opResult));
+      if (auto funcOp = dyn_cast<FuncOp>(op))
+        for (BlockArgument bbArg : funcOp.getArguments())
+          if (bbArg.getType().isa<TensorType>())
+            setInPlaceFuncArgument(bbArg,
+                                   aliasInfo.bufferizesToWritableMemory(bbArg));
+    }
+  });
+}
+
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     ModuleOp moduleOp, const BufferizationOptions &options) {
   SmallVector<FuncOp> orderedFuncOps;
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
   DenseMap<FuncOp, FunctionType> bufferizedFunctionTypes;
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
for (FuncOp funcOp : orderedFuncOps) {
     // In a first approximation:
     // =========================
     // If the function is called, we can allocate on the caller side which lets
     // us force inplace arguments at function boundaries.
     // TODO: do not rely on this behavior.
     if (callerMap.find(funcOp) != callerMap.end())
       for (BlockArgument bbArg : funcOp.getArguments())
         if (bbArg.getType().isa<TensorType>())
-          setInPlaceFuncArgument(bbArg);
+          aliasInfo.setBufferizesToWritableMemory(bbArg);
 
     // If the analysis fails, just return.
     if (failed(inPlaceAnalysisFuncOpBody(funcOp, aliasInfo, domInfo,
                                          options.analysisFuzzerSeed)))
       return failure();
 
     // Try to eliminate InitTensorOps to avoid new allocations during the
     // bufferization phase.
     if (failed(eliminateInsertSliceAnchoredInitTensorOps(funcOp, aliasInfo,
                                                          domInfo)))
       return failure();
 
     // Bufferization phase.
     if (!options.testAnalysisOnly) {
       BlockAndValueMapping tensorToBufferMap;
       if (failed(bufferizeFuncOpInternals(funcOp, tensorToBufferMap, aliasInfo,
                                           *options.allocationFns,
                                           bufferizedFunctionTypes)))
         return failure();
     }
   }
-  // Don't drop the attributes if we only want to report the analysis.
-  if (options.testAnalysisOnly)
+  // Annotate operations if we only want to report the analysis.
+  if (options.testAnalysisOnly) {
+    annotateOpsWithBufferizationMarkers(moduleOp, aliasInfo);
     return success();
+  }
 
   for (FuncOp funcOp : orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
     if (failed(bufferizeFuncOpBoundary(funcOp, aliasInfo,
                                        bufferizedFunctionTypes)))
       return failure();
 
     if (!options.allowReturnMemref &&
         llvm::any_of(funcOp.getType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {
       funcOp->emitError("memref return type is unsupported");
       return failure();
     }
   }
 
   // Perform a post-processing pass of layout modification at function boundary
   // according to the kBufferLayoutAttrName.
   layoutPostProcessing(moduleOp);
 
   // Post-pass cleanup of inplaceable and buffer_layout attributes.
-  moduleOp.walk(
-      [&](Operation *op) { op->removeAttr(kInPlaceResultsAttrName); });
   moduleOp.walk([&](FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationFuncArguments(bbArg);
   });
 
   return success();
 }
 
LogicalResult bufferize(Operation *op, OpBuilder &b,
     Value memref = b.create<memref::GetGlobalOp>(
         constantOp.getLoc(), globalMemref.type(), globalMemref.getName());
     aliasInfo.insertNewBufferEquivalence(memref, constantOp.getResult());
     map(bvm, constantOp, memref);
 
     return success();
   }
 
-  bool isWritable(Operation *op, OpResult opResult) const {
+  bool isWritable(Operation *op, Value value) const {
     // Memory locations returned by memref::GetGlobalOp may not be written to.
+    assert(value.isa<OpResult>());
     return false;
   }
 };
 
 } // namespace arith_ext
 
 // TODO: Ops in the linalg dialect can directly implement this interface.
 namespace linalg_ext {
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
     return tiledLoopOp.getTiedOpResult(opOperand);
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
+  bool isWritable(Operation *op, Value value) const {
+    // Interestingly, linalg::TiledLoopOp's bbArg can **always** be viewed
+    // inplace from the perspective of ops nested under:
+    //   1. Either the matching iter operand is not bufferized inplace and an
+    //      alloc + optional copy makes the bbArg itself inplaceable.
+    //   2. Or the matching iter operand is bufferized inplace and bbArg just
+    //      bufferizes to that too.
+    return true;
+  }
+
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BlockAndValueMapping &bvm,
                           BufferizationAliasInfo &aliasInfo,
                           AllocationCallbacks &allocationFn) const {
     auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
if (!opOperand.get().getType().isa<RankedTensorType>())
       return OpResult();
     return forOp.getResultForOpOperand(opOperand);
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
+  bool isWritable(Operation *op, Value value) const {
+    // Interestingly, scf::ForOp's bbArg can **always** be viewed
+    // inplace from the perspective of ops nested under:
+    //   1. Either the matching iter operand is not bufferized inplace and an
+    //      alloc + optional copy makes the bbArg itself inplaceable.
+    //   2. Or the matching iter operand is bufferized inplace and bbArg just
+    //      bufferizes to that too.
+    return true;
+  }
+
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BlockAndValueMapping &bvm,
                           BufferizationAliasInfo &aliasInfo,
                           AllocationCallbacks &allocationFn) const {
     // Note: This method is just setting up the mappings for the block arguments
     // and the result buffer. The op is bufferized after the scf::YieldOp.
 
     auto forOp = cast<scf::ForOp>(op);
LogicalResult bufferize(Operation *op, OpBuilder &b,
     if (!srcMemref)
       return failure();
     auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
     auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // If not inplaceable, alloc.
     Value alloc;
-    auto inPlace = getInPlace(extractSliceOp->getResult(0));
-    if (inPlace != InPlaceSpec::True)
+    if (!aliasInfo.isInPlace(extractSliceOp->getResult(0)))
       alloc = createNewAllocDeallocPairForShapedValue(
           b, loc, extractSliceOp.result(), aliasInfo, allocationFn);
 
     // Set insertion point now that potential alloc/dealloc are introduced.
     b.setInsertionPoint(extractSliceOp);
 
     // Bufferize to subview.
     auto subviewMemRefType =
static bool isSourceEquivalentToAMatchingInplaceExtractSliceOp(
   LDBG("isSourceEquivalentToAMatchingInplaceExtractSliceOp: " << *insertSliceOp
                                                               << '\n');
   bool foundOp = false;
   aliasInfo.applyOnEquivalenceClass(insertSliceOp.source(), [&](Value value) {
     auto extractSliceOp = value.getDefiningOp<ExtractSliceOp>();
     if (extractSliceOp &&
         areEquivalentExtractSliceOps(aliasInfo, extractSliceOp,
                                      insertSliceOp) &&
-        getInPlace(extractSliceOp.result()) == InPlaceSpec::True) {
+        aliasInfo.isInPlace(extractSliceOp->getResult(0))) {
       LDBG("\tfound: " << extractSliceOp.getOperation() << '\n');
       foundOp = true;
     }
   });
 
   if (!foundOp)
     LDBG("\tnot equivalent\n");
 
auto subviewMemRefType =
             insertSliceOp.getMixedStrides())
             .cast<MemRefType>();
 
     // A copy of the source buffer is needed if either:
     //   - The producer of `source` is not inplace. This is the case where a
     //     slice is computed out of place into the inplace full tensor.
     //   - The result is not inplace. This is the case where the whole tensor is
     //     cloned and the clone needs to be updated.
-    auto inPlace = getInPlace(insertSliceOp->getResult(0));
     // TODO: Is this necessary?
     if (!isSourceEquivalentToAMatchingInplaceExtractSliceOp(aliasInfo,
                                                             insertSliceOp) ||
-        inPlace != InPlaceSpec::True) {
+        !aliasInfo.isInPlace(insertSliceOp->getResult(0))) {
       LDBG("insert_slice needs extra source copy: " << insertSliceOp.source()
                                                     << " -> copy\n");
       // Take a subview of the dst.
       Value subView = b.create<memref::SubViewOp>(
           loc, subviewMemRefType, dstMemref, insertSliceOp.getMixedOffsets(),
           insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
       // Insert new alias.
       aliasInfo.insertNewBufferAlias(subView, dstMemref);

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

%r, %v3 = scf.for %i = %c0 to %s step %c1 iter_args(%t2 = %t1, %v0 = %v) -> (tensor<?xf32>, vector<5xf32>) {
     // CHECK:      tensor.extract_slice
     // CHECK-SAME: __inplace_results_attr__ = ["true"]
     %e = tensor.extract_slice %t2[%s][%s][1] : tensor<?xf32> to tensor<?xf32>
 
     // Read from %t1 via alias %e.
     %v2 = vector.transfer_read %e[%s], %cst : tensor<?xf32>, vector<5xf32>
     scf.yield %e, %v2 : tensor<?xf32>, vector<5xf32>
   }
-  // CHECK: __inplace_results_attr__ = ["true", "none"]
+  // CHECK: __inplace_results_attr__ = ["true", "false"]
 
   // Use %t3 in some way without reading it, so that it does not get DCE'd.
   // CHECK:      linalg.generic
   // CHECK-SAME: __inplace_results_attr__ = ["true"]
   %o = linalg.generic #trait outs (%t3 : tensor<?xf32>) {
       ^bb(%0: f32) :
         linalg.yield %cst : f32
     } -> (tensor<?xf32>)