[mlir][func][bufferization][NFC] Simplify implementation

[springerm]

Summary

The bufferization implementation of func.func and func.call can be simplified. It still contained code that was necessary when One-Shot Bufferize removed return values. This functionality has been extracted into a separate pass a while ago.

Depends On: D157863

[springerm]

more simplifications

[mehdi_amini]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

546 ↗

(On Diff #550227)

There is no way for a static method to be overridden, the body is supposed to use information from the concrete op class here otherwise it does not make sense.
Shall we make it a non-static method?

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

LogicalResult bufferization::bufferizeOp(Operation *op,
 
   // Gather all bufferizable ops in top-to-bottom order.
   //
   // We should ideally know the exact memref type of all operands when
   // bufferizing an op. (This is the case when bufferizing top-to-bottom.)
   // Otherwise, we have to use a memref type with a fully dynamic layout map to
   // avoid copies. We are currently missing patterns for layout maps to
   // canonicalize away (or canonicalize to more precise layouts).
-  //
-  // FuncOps must be bufferized before their bodies, so add them to the worklist
-  // first.
   SmallVector<Operation *> worklist;
-  op->walk([&](func::FuncOp funcOp) {
-    if (hasTensorSemantics(funcOp))
-      worklist.push_back(funcOp);
-  });
   op->walk<WalkOrder::PostOrder>([&](Operation *op) {
-    if (hasTensorSemantics(op) && !isa<func::FuncOp>(op))
+    if (hasTensorSemantics(op))
       worklist.push_back(op);
   });
 
   // Keep track of all erased ops.
   DenseSet<Operation *> erasedOps;
 
   // Bufferize all ops.
   BufferizationRewriter rewriter(op->getContext(), erasedOps, toMemrefOps,

mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp

AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
     for (int64_t resultIdx : aliasingReturnVals)
       result.addAlias({callOp->getOpResult(resultIdx),
                        equivalent.has_value() ? BufferRelation::Equivalent
                                               : BufferRelation::Unknown,
                        /*isDefinite=*/equivalent.has_value()});
     return result;
   }
 
+  FailureOr<BaseMemRefType>
+  getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+    auto callOp = cast<func::CallOp>(op);
+    FuncOp funcOp = getCalledFunction(callOp);
+    assert(funcOp && "expected CallOp to a FuncOp");
+
+    // The callee was already bufferized, so we can directly take the type from
+    // its signature.
+    FunctionType funcType = funcOp.getFunctionType();
+    return cast<BaseMemRefType>(
+        funcType.getResult(cast<OpResult>(value).getResultNumber()));
+  }
+
   /// All function arguments are writable. It is the responsibility of the
   /// CallOp to insert buffer copies where necessary.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     func::CallOp callOp = cast<func::CallOp>(op);
-    unsigned numResults = callOp.getNumResults();
-    unsigned numOperands = callOp->getNumOperands();
-    FuncOp funcOp = getCalledFunction(callOp);
-    assert(funcOp && "expected CallOp to a FuncOp");
-    FunctionType funcType = funcOp.getFunctionType();
-
-    // Result types of the bufferized CallOp.
-    SmallVector<Type> resultTypes;
-    // Replacement values for the existing CallOp. These are usually the results
-    // of the bufferized CallOp, unless a tensor result folds onto an operand.
-    SmallVector<Value> replacementValues(numResults, Value());
-    // For non-tensor results: A mapping from return val indices of the old
-    // CallOp to return val indices of the bufferized CallOp.
-    SmallVector<std::optional<unsigned>> retValMapping(numResults,
-                                                       std::nullopt);
-    // Operands of the bufferized CallOp.
-    SmallVector<Value> newOperands(numOperands, Value());
 
     // 1. Compute the result types of the new CallOp.
-    for (const auto &it : llvm::enumerate(callOp.getResultTypes())) {
-      unsigned returnValIdx = it.index();
-      Type returnType = it.value();
+    SmallVector<Type> resultTypes;
+    for (Value result : callOp.getResults()) {
+      Type returnType = result.getType();
       if (!isa<TensorType>(returnType)) {
         // Non-tensor values are returned.
-        retValMapping[returnValIdx] = resultTypes.size();
         resultTypes.push_back(returnType);
         continue;
       }
 
       // Returning a memref.
-      retValMapping[returnValIdx] = resultTypes.size();
-      resultTypes.push_back(funcType.getResult(resultTypes.size()));
+      FailureOr<BaseMemRefType> resultType =
+          bufferization::getBufferType(result, options);
+      if (failed(resultType))
+        return failure();
+      resultTypes.push_back(*resultType);
     }
 
-    // 2. Rewrite tensor operands as memrefs based on `bufferizedFuncType`.
-    for (OpOperand &opOperand : callOp->getOpOperands()) {
-      unsigned idx = opOperand.getOperandNumber();
-      Value tensorOperand = opOperand.get();
+    // 2. Rewrite tensor operands as memrefs based on type of the already
+    //    bufferized callee.
+    SmallVector<Value> newOperands;
+    FuncOp funcOp = getCalledFunction(callOp);
+    assert(funcOp && "expected CallOp to a FuncOp");
+    FunctionType funcType = funcOp.getFunctionType();
 
+    for (OpOperand &opOperand : callOp->getOpOperands()) {
       // Non-tensor operands are just copied.
-      if (!isa<TensorType>(tensorOperand.getType())) {
-        newOperands[idx] = tensorOperand;
+      if (!isa<TensorType>(opOperand.get().getType())) {
+        newOperands.push_back(opOperand.get());
         continue;
       }
 
       // Retrieve buffers for tensor operands.
-      Value buffer = newOperands[idx];
-      if (!buffer) {
-        FailureOr<Value> maybeBuffer =
-            getBuffer(rewriter, opOperand.get(), options);
-        if (failed(maybeBuffer))
-          return failure();
-        buffer = *maybeBuffer;
-      }
+      Value buffer = *maybeBuffer;
 
       // Caller / callee type mismatch is handled with a CastOp.
-      auto memRefType = funcType.getInput(idx);
+      auto memRefType = funcType.getInput(opOperand.getOperandNumber());
       // Since we don't yet have a clear layout story, to_memref may
       // conservatively turn tensors into more dynamic memref than necessary.
       // If the memref type of the callee fails, introduce an extra memref.cast
       // that will either canonicalize away or fail compilation until we can do
       // something better.
       if (buffer.getType() != memRefType) {
         assert(
             memref::CastOp::areCastCompatible(buffer.getType(), memRefType) &&
             "CallOp::bufferize: cast incompatible");
         Value castBuffer = rewriter.create<memref::CastOp>(callOp.getLoc(),
                                                            memRefType, buffer);
         buffer = castBuffer;
       }
-      newOperands[idx] = buffer;
+      newOperands.push_back(buffer);
     }
 
     // 3. Create the new CallOp.
     Operation *newCallOp = rewriter.create<func::CallOp>(
         callOp.getLoc(), funcOp.getSymName(), resultTypes, newOperands);
     newCallOp->setAttrs(callOp->getAttrs());
-    // Get replacement values.
-    for (unsigned i = 0; i < replacementValues.size(); ++i) {
-      if (replacementValues[i])
-        continue;
-      replacementValues[i] = newCallOp->getResult(*retValMapping[i]);
-    }
 
     // 4. Replace the old op with the new op.
-    replaceOpWithBufferizedValues(rewriter, callOp, replacementValues);
+    replaceOpWithBufferizedValues(rewriter, callOp, newCallOp->getResults());
 
     return success();
   }
 };
 
 struct ReturnOpInterface
     : public BufferizableOpInterface::ExternalModel<ReturnOpInterface,
                                                     func::ReturnOp> {
#endif // NDEBUG
 
     // ReturnOps are bufferized as part of FuncOps.
     return success();
   }
 };
 
 struct FuncOpInterface
     : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
+  FailureOr<BaseMemRefType>
+  getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+    auto funcOp = cast<FuncOp>(op);
+    auto bbArg = cast<BlockArgument>(value);
+    // Unstructured control flow is not supported.
+    assert(bbArg.getOwner() == &funcOp.getBody().front() &&
+           "expected that block argument belongs to first block");
+    return getBufferizedFunctionArgType(funcOp, bbArg.getArgNumber(), options);
+  }
+
   /// Rewrite function bbArgs and return values into buffer form. This function
   /// bufferizes the function signature and the ReturnOp. When the entire
   /// function body has been bufferized, function return types can be switched
   /// to more concise memref types as part of `foldMemRefCasts`.
   ///
   /// All function bbArgs are writable unless they are explicitly marked as
   /// read-only. Callers must insert copies when needed.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
for (BlockArgument &bbArg : frontBlock.getArguments()) {
         continue;
 
       // Collect all uses of the bbArg.
       SmallVector<OpOperand *> bbArgUses;
       for (OpOperand &use : bbArg.getUses())
         bbArgUses.push_back(&use);
 
       // Change the bbArg type to memref.
-      Type memrefType =
-          getBufferizedFunctionArgType(funcOp, bbArg.getArgNumber(), options);
-      bbArg.setType(memrefType);
+      FailureOr<BaseMemRefType> memrefType =
+          bufferization::getBufferType(bbArg, options);
+      if (failed(memrefType))
+        return failure();
+      bbArg.setType(*memrefType);
 
       // Replace all uses of the original tensor bbArg.
       rewriter.setInsertionPointToStart(&frontBlock);
       if (!bbArgUses.empty()) {
         // Insert to_tensor because the remaining function body has not been
         // bufferized yet.
         Value toTensorOp =
             rewriter.create<bufferization::ToTensorOp>(funcOp.getLoc(), bbArg);