[mlir][linalg][bufferize] Remove special scf::IfOp rules

[springerm]

Summary

Remove some of the special rules for scf::IfOp (not all of them) and encode them in the op interface. This is in preparation of decoupling analysis, bufferization and dialects.

Depends On D113337

[springerm]

no change

[springerm]

no change

[nicolasvasilache]

It is a bit hard to take this as an NFC .. can you maybe rebase it as the last commit so we can see what it would look like there?

[springerm]

rebase + update

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

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));
         equivalentInfo.unionSets(std::get<0>(it), std::get<1>(it));
         equivalentInfo.unionSets(std::get<0>(it), std::get<2>(it));
       }
+
+      // scf::IfOp always bufferizes in-place.
+      for (OpResult opResult : ifOp->getResults())
+        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);

 static Value findLastPrecedingWrite(Value value) {
   SetVector<Value> result =
       findValueInReverseUseDefChain(value, [](Value value) {
         Operation *op = value.getDefiningOp();
         if (!op)
           return true;
         auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
         if (!bufferizableOp)
+          // Unknown op: Assume that it is writing.
           return true;
-        if (isa<scf::IfOp>(op))
-          return true;
-
-        SmallVector<OpOperand *> opOperands =
-            bufferizableOp.getAliasingOpOperand(value.cast<OpResult>());
-        assert(opOperands.size() <= 1 &&
-               "op with multiple aliasing OpOperands not expected");
-
-        if (opOperands.empty())
+        // Check if any of the aliasing OpOperands is writing.
+        if (llvm::any_of(
+                bufferizableOp.getAliasingOpOperand(value.cast<OpResult>()),
+                [](OpOperand *operand) {
+                  return bufferizesToMemoryWrite(*operand);
+                }))
           return true;
-
-        return bufferizesToMemoryWrite(*opOperands.front());
+        // Not a write, continue iterating...
+        return false;
       });
+
   assert(result.size() == 1 && "expected exactly one result");
   return result.front();
 }
 
 /// Return true if `value` is originating from an ExtractSliceOp that matches
 /// the given InsertSliceOp.
 bool BufferizationAliasInfo::hasMatchingExtractSliceOp(
     Value value, InsertSliceOp insertOp) const {
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 && !isa<scf::IfOp>(op);
+  bool needCopy = getInPlace(result) != InPlaceSpec::True;
   if (needCopy) {
-    // Ops such as scf::IfOp can currently not bufferize out-of-place.
+    // 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(
         b, loc, operand, aliasInfo, allocationFns);
     // Do not copy the result of an InitTensorOp.

 } // namespace linalg_ext
 
 namespace scf_ext {
 
 struct IfOpInterface
     : public BufferizableOpInterface::ExternalModel<IfOpInterface, scf::IfOp> {
   SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
                                                 OpResult opResult) const {
+    // IfOps do not have tensor OpOperands. The yielded value can be any SSA
+    // value that is in scope. To allow for use-def chain traversal through
+    // IfOps in the analysis, both corresponding yield values from the then/else
+    // branches are considered to be aliasing with the result.
     auto ifOp = cast<scf::IfOp>(op);
-    // Either one of the corresponding yield values from the then/else branches
-    // may alias with the result.
     size_t resultNum = std::distance(op->getOpResults().begin(),
                                      llvm::find(op->getOpResults(), opResult));
     return {&ifOp.thenYield()->getOpOperand(resultNum),
             &ifOp.elseYield()->getOpOperand(resultNum)};
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BlockAndValueMapping &bvm,

 struct YieldOpInterface
     : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
                                                     scf::YieldOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
-    return false;
+    // scf::YieldOp does usually not bufferize to a memory write. However, there
+    // is a special rule for IfOps. The aliasing OpOperand of an IfOp is the
+    // YieldOp's OpOperand. For analysis purposes, IfOps are considered to be
+    // writing ops, so this method should return `true`.
+    return isa<scf::IfOp>(op->getParentOp());
   }
 
   OpResult getInplaceableOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
+  BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
+    return BufferRelation::Equivalent;
+  }
+
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BlockAndValueMapping &bvm,
                           BufferizationAliasInfo &aliasInfo,
                           AllocationCallbacks &allocationFn) const {
     auto yieldOp = cast<scf::YieldOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);