[Mlir-commits] [mlir] 2e0d821 - [mlir][linalg][bufferize] Store analysis results in BufferizationAliasInfo

Matthias Springer llvmlistbot at llvm.org
Wed Nov 10 17:38:12 PST 2021


Author: Matthias Springer
Date: 2021-11-11T10:36:49+09:00
New Revision: 2e0d821bd531e0c3cf449c8785d364adca42efba

URL: https://github.com/llvm/llvm-project/commit/2e0d821bd531e0c3cf449c8785d364adca42efba
DIFF: https://github.com/llvm/llvm-project/commit/2e0d821bd531e0c3cf449c8785d364adca42efba.diff

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

* 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.

Differential Revision: https://reviews.llvm.org/D113375

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
    mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
    mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
index 281873e37a6d..fa97e2a06b81 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
@@ -182,17 +182,26 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
       >,
       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>();
         }]
       >
   ];

diff  --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
index b9021fac3606..75406883612d 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
@@ -91,6 +91,12 @@ class BufferizationAliasInfo {
   /// 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;
@@ -117,6 +123,9 @@ class BufferizationAliasInfo {
   /// 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

diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index 89f6782389af..76126693b228 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
@@ -221,48 +221,20 @@ static Value lookup(const BlockAndValueMapping &bvm, Value key) {
 
 //===----------------------------------------------------------------------===//
 // 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;
 
@@ -272,69 +244,20 @@ static void setInPlaceOpResult(OpResult opResult,
   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
@@ -347,12 +270,6 @@ static void removeBufferizationFuncArguments(BlockArgument bbArg) {
                        LinalgDialect::kInplaceableAttrName);
 }
 
-static InPlaceSpec getInPlace(Value v) {
-  if (auto bbArg = v.dyn_cast<BlockArgument>())
-    return getInPlace(bbArg);
-  return getInPlace(v.cast<OpResult>());
-}
-
 //===----------------------------------------------------------------------===//
 // Printing helpers.
 //===----------------------------------------------------------------------===//
@@ -365,9 +282,6 @@ static void printTensorOrBufferInfo(std::string prefix, Value value,
   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.
@@ -428,12 +342,13 @@ areEquivalentExtractSliceOps(const BufferizationAliasInfo &aliasInfo,
 }
 
 /// 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) {
@@ -460,7 +375,7 @@ BufferizationAliasInfo::BufferizationAliasInfo(Operation *rootOp) {
                  "expected that OpResult has aliasing OpOperand");
           for (OpOperand *operand : operands)
             aliasInfo.unionSets(operand->get(), opResult);
-          setInPlaceOpResult(opResult, InPlaceSpec::True);
+          markInPlace(opResult);
         }
     }
   });
@@ -492,39 +407,33 @@ void BufferizationAliasInfo::insertNewBufferEquivalence(Value newValue,
 /// 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;
-      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;
-    }
+    // 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;
+
+    foundNonWritableBuffer = true;
   });
 
-  if (!foundNonWritableBuffer)
-    LDBG("---->value is writable\n");
+  if (foundNonWritableBuffer)
+    LDBG("--> NON WRITABLE\n");
+  else
+    LDBG("--> WRITABLE\n");
 
   return foundNonWritableBuffer;
 }
@@ -547,7 +456,7 @@ static bool aliasesInPlaceWrite(Value value,
   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;
@@ -562,10 +471,30 @@ static bool aliasesInPlaceWrite(Value value,
   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());
@@ -577,7 +506,8 @@ void BufferizationAliasInfo::bufferizeInPlace(OpResult result,
 
 /// 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
@@ -892,7 +822,7 @@ bool wouldCreateReadAfterWriteInterference(
     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);
     });
   };
@@ -1264,8 +1194,7 @@ static Value getResultBuffer(OpBuilder &b, OpResult result,
   }
 
   // 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(
@@ -1464,9 +1393,7 @@ static LogicalResult bufferize(OpBuilder &b, FuncOp funcOp,
 // 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,
@@ -1500,8 +1427,7 @@ bufferizableInPlaceAnalysisImpl(OpOperand &operand, OpResult result,
 }
 
 /// 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.
@@ -2044,12 +1970,12 @@ LogicalResult mlir::linalg::comprehensive_bufferize::initTensorElimination(
         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).
@@ -2133,12 +2059,12 @@ LogicalResult mlir::linalg::comprehensive_bufferize::
                                               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*/;
       },
@@ -2171,6 +2097,23 @@ static void checkAliasInfoConsistency(FuncOp funcOp,
 }
 #endif
 
+/// 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;
@@ -2200,7 +2143,7 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     if (callerMap.find(funcOp) != callerMap.end())
       for (BlockArgument bbArg : funcOp.getArguments())
         if (bbArg.getType().isa<TensorType>())
-          setInPlaceFuncArgument(bbArg);
+          aliasInfo.setBufferizesToWritableMemory(bbArg);
 
 #ifndef NDEBUG
     checkAliasInfoConsistency(funcOp, domInfo, aliasInfo);
@@ -2226,9 +2169,11 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
         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
@@ -2251,8 +2196,6 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
   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);
@@ -2310,8 +2253,9 @@ struct ConstantOpInterface
     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;
   }
 };
@@ -2515,6 +2459,16 @@ struct TiledLoopOpInterface
     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,
@@ -2804,6 +2758,16 @@ struct ForOpInterface
     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,
@@ -3161,8 +3125,7 @@ struct ExtractSliceOpInterface
 
     // 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);
 
@@ -3240,7 +3203,7 @@ static bool isSourceEquivalentToAMatchingInplaceExtractSliceOp(
     if (extractSliceOp &&
         areEquivalentExtractSliceOps(aliasInfo, extractSliceOp,
                                      insertSliceOp) &&
-        getInPlace(extractSliceOp.result()) == InPlaceSpec::True) {
+        aliasInfo.isInPlace(extractSliceOp->getResult(0))) {
       LDBG("\tfound: " << extractSliceOp.getOperation() << '\n');
       foundOp = true;
     }
@@ -3319,11 +3282,10 @@ struct InsertSliceOpInterface
     //     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.

diff  --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
index fc82522c3ef5..68fe343e412e 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
@@ -1066,7 +1066,7 @@ func @reading_scf_for(%t1: tensor<?xf32> {linalg.inplaceable = true},
     %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


        


More information about the Mlir-commits mailing list