[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/ComprehensiveBufferize/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)
           return true;
-        if (isa<scf::IfOp>(op))
-          return true;
         return bufferizableOp.isMemoryWrite(value.cast<OpResult>());
       });
+
   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.
 static bool hasMatchingExtractSliceOp(const BufferizationAliasInfo &aliasInfo,
                                       Value value, InsertSliceOp insertOp) {
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)};
   }
 
+  // TODO: For better bufferization results, this could return `true` only if
+  // there is a memory write in one (or both) of the branches. Since this is not
+  // allowed at the moment, we should never encounter scf.ifs that yield
+  // unmodified tensors. Such scf.yield ops could just fold away.
+  bool isMemoryWrite(Operation *op, OpResult opResult) const {
+    // IfOp results are always considered memory writes in the analysis. This
+    // design decision simplifies the analysis considerably. E.g., consider the
+    // following test case:
+    //
+    // %0 = "some_writing_op" : tensor<?xf32>
+    // %r = scf.if %c -> (tensor<?xf32>) {
+    //   scf.yield %0
+    // } else {
+    //   %1 = "another_writing_op"(%0) : tensor<?xf32>
+    // }
+    // "some_reading_op"(%r)
+    //
+    // "another_writing_op" in the above example should be able to bufferize
+    // inplace in the absence of another read of %0. However, if the scf.if op
+    // would not be considered a "write", the analysis would detect the
+    // following conflict:
+    //
+    // * 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;
+  }
+
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BlockAndValueMapping &bvm,
                           BufferizationAliasInfo &aliasInfo,
                           AllocationCallbacks &allocationFn) const {
     auto ifOp = cast<scf::IfOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
struct YieldOpInterface
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
   OpResult getAliasingOpResult(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);