[mlir][linalg][bufferize] Add mustBufferizeInPlace to op interface

[springerm]

Summary

This is useful for ops such as scf::IfOp, which always bufferize in-place.

This commit is in preparation of decoupling BufferizationAliasInfo from the SCF dialect.

Depends On D113338

[springerm]

update

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

let methods = [
                 opOperands,
                 [&](OpOperand *operand) {
                   return bufferizableOp.bufferizesToMemoryWrite(*operand);
                 });
           }]
       >,
       InterfaceMethod<
         /*desc=*/[{
+          Return `true` if the given OpResult must bufferize in-place with its
+          corresponding aliasing OpOperand. Alias sets and inplace attributes
+          will be set up accordingly before making any other bufferization
+          decisions. This method will never be called on OpResults that do not
+          have a tensor type.
+
+          Note: This method may not return `true` if the given OpResult does not
+          have an aliasing OpOperand.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"mustBufferizeInPlace",
+        /*args=*/(ins "OpResult":$opResult),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return false;
+        }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
           Return the OpResult that aliases with a given OpOperand when
           bufferized in-place. This method will never be called on OpOperands
           that do not have a tensor type.
         }],
         /*retType=*/"OpResult",
         /*methodName=*/"getAliasingOpResult",
         /*args=*/(ins "OpOperand &":$opOperand),
         /*methodBody=*/"",

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

for (Value v : op->getResults())
         createAliasInfoEntry(v);
     for (Region &r : op->getRegions())
       for (Block &b : r.getBlocks())
         for (auto bbArg : b.getArguments())
           if (bbArg.getType().isa<TensorType>())
             createAliasInfoEntry(bbArg);
   });
 
-  // The return value of an scf::IfOp aliases with both yield values.
-  rootOp->walk([&](scf::IfOp ifOp) {
-    if (ifOp->getNumResults() > 0) {
-      for (auto it : llvm::zip(ifOp.thenYield().results(),
-                               ifOp.elseYield().results(), ifOp.results())) {
-        aliasInfo.unionSets(std::get<0>(it), std::get<1>(it));
-        aliasInfo.unionSets(std::get<0>(it), std::get<2>(it));
-      }
-
-      // scf::IfOp always bufferizes in-place.
-      for (OpResult opResult : ifOp->getResults())
+  // Set up alias sets for OpResults that must bufferize in-place. This should
+  // be done before making any other bufferization decisions.
+  rootOp->walk([&](BufferizableOpInterface bufferizableOp) {
+    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);
         }
+    }
   });
 }
 
 /// 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) {
   aliasInfo.insert(v);
   equivalentInfo.insert(v);

 ///
 /// In the above example, the two TransferWriteOps have already been decided to
 /// bufferize inplace. Bufferizing the ExtractSliceOp inplace would create a
 /// conflict because:
 /// * According to SSA use-def chains, we expect to read the result of %1.
 /// * However, adding an alias {%0, %t} would mean that the second
 ///   TransferWriteOp overwrites the first one. Therefore, the TransferReadOp
 ///   would no longer be reading the result of %1.
+///
+/// If `checkConsistencyOnly` is true, this function checks if there is a
+/// read-after-write conflict without bufferizing `operand` inplace. This would
+/// indicate a problem with the current inplace bufferization decisions.
 bool wouldCreateReadAfterWriteInterference(
     OpOperand &operand, OpResult result, const DominanceInfo &domInfo,
-    const BufferizationAliasInfo &aliasInfo) {
+    const BufferizationAliasInfo &aliasInfo,
+    bool checkConsistencyOnly = false) {
 #ifndef NDEBUG
   SmallVector<OpOperand *> opOperands = getAliasingOpOperand(result);
   assert(llvm::find(opOperands, &operand) != opOperands.end() &&
          "operand and result do not match");
 #endif // NDEBUG
 
   // Helper function to iterate on aliases of `root` and capture the reads.
   auto getAliasingReads = [&](DenseSet<OpOperand *> &res, Value root) {
#endif // NDEBUG
   };
 
   // Collect reads and writes of all aliases of OpOperand and OpResult.
   DenseSet<OpOperand *> usesRead, usesWrite;
   getAliasingReads(usesRead, operand.get());
   getAliasingReads(usesRead, result);
   getAliasingInplaceWrites(usesWrite, operand.get());
   getAliasingInplaceWrites(usesWrite, result);
-  if (bufferizesToMemoryWrite(operand))
+  if (!checkConsistencyOnly && bufferizesToMemoryWrite(operand))
     usesWrite.insert(&operand);
 
   return hasReadAfterWriteInterference(usesRead, usesWrite, domInfo, aliasInfo);
 }
 
 /// Return true if bufferizing `opOperand` inplace with `opResult` would create
 /// a write to a non-writable buffer.
 static bool
return initTensorElimination(
         auto insertSliceOp = cast<InsertSliceOp>(operand.getOwner());
         auto extractOp = b.create<tensor::ExtractSliceOp>(
             loc, insertSliceOp.dest(), insertSliceOp.getMixedOffsets(),
             insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
         return extractOp.result();
       });
 }
 
+/// Assert that the current bufferization decisions are consistent.
+static void checkAliasInfoConsistency(FuncOp funcOp,
+                                      const DominanceInfo &domInfo,
+                                      const BufferizationAliasInfo &aliasInfo) {
+  funcOp.walk([&](Operation *op) {
+    if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+      for (OpOperand &opOperand : op->getOpOperands())
+        if (opOperand.get().getType().isa<TensorType>())
+          if (OpResult opResult = bufferizableOp.getAliasingOpResult(opOperand))
+            // If this assertion fails, there is probably an inconsistent
+            // combination of "mustBufferizeInPlace" decisions.
+            assert(!wouldCreateReadAfterWriteInterference(
+                       opOperand, opResult, domInfo, aliasInfo,
+                       /*checkConsistencyOnly=*/true) &&
+                   "found read after write conflict before running analysis");
+  });
+}
+
 LogicalResult
 mlir::linalg::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();
 
   DominanceInfo domInfo(moduleOp);
   BufferizationAliasInfo aliasInfo(moduleOp);
+
   // Interestingly, all function args that are not visible outside of a module
   // can be fully bufferized inplace by guaranteeing the CallOp is bufferized
   // inplace. Therefore, we just bufferize funcOp as if none of its results were
   // inplaceable, detect which operands are cloned internally and decide what to
   // do at call sites.
   for (FuncOp funcOp : orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.body().empty())
       continue;
 
     // 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);
 
+#ifndef NDEBUG
+    checkAliasInfoConsistency(funcOp, domInfo, aliasInfo);
+#endif // NDEBUG
+
     // 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,
bool isMemoryWrite(Operation *op, OpResult opResult) const {
     // * 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;
   }
 
+  bool mustBufferizeInPlace(Operation *op, OpResult opResult) const {
+    // IfOp results always bufferize in-place. Since they have no OpOperands,
+    // they are mostly ignored by the analysis once alias sets are set up.
+    return true;
+  }
+
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BlockAndValueMapping &bvm,
                           BufferizationAliasInfo &aliasInfo,
                           AllocationCallbacks &allocationFn) const {
     // scf::IfOp is bufferized after scf::YieldOp in the else branch.
     return success();
   }
 };