[Mlir-commits] [mlir] [MLIR] Make `OneShotModuleBufferize` use `OpInterface` (PR #107295)

[Tzung-Han Juang]

================
@@ -291,26 +300,30 @@ static bool hasTensorSignature(func::FuncOp funcOp) {
 /// retrieve the called FuncOp from any func::CallOp.
 static LogicalResult
 getFuncOpsOrderedByCalls(ModuleOp moduleOp,
-                         SmallVectorImpl<func::FuncOp> &orderedFuncOps,
+                         SmallVectorImpl<FunctionOpInterface> &orderedFuncOps,
                          FuncCallerMap &callerMap) {
   // For each FuncOp, the set of functions called by it (i.e. the union of
   // symbols of all nested func::CallOp).
-  DenseMap<func::FuncOp, DenseSet<func::FuncOp>> calledBy;
+  DenseMap<FunctionOpInterface, DenseSet<FunctionOpInterface>> calledBy;
   // For each FuncOp, the number of func::CallOp it contains.
-  DenseMap<func::FuncOp, unsigned> numberCallOpsContainedInFuncOp;
-  WalkResult res = moduleOp.walk([&](func::FuncOp funcOp) -> WalkResult {
-    if (!funcOp.getBody().empty()) {
-      func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
-      if (!returnOp)
-        return funcOp->emitError()
-               << "cannot bufferize a FuncOp with tensors and "
-                  "without a unique ReturnOp";
+  DenseMap<FunctionOpInterface, unsigned> numberCallOpsContainedInFuncOp;
+  WalkResult res = moduleOp.walk([&](FunctionOpInterface funcOp) -> WalkResult {
+    // Only handle ReturnOp if funcOp is exactly the FuncOp type.
+    if(isa<FuncOp>(funcOp)) {
----------------
tzunghanjuang wrote:

Thank you for the suggestion. We have updated `getAssumedUniqueReturnOp` so that it no longer returns `funcOp`. We did not find a directly determine if a `FunctionOpInterface` is shipped with any `ReturnLike` Op, so we used  `funcOp.getNumResults()` to check if there is no return type. 

With this method, we still have to keep the check near `aliasingFuncOpBBArgsAnalysis` and `foldMemRefCasts` or the following example will fail. The example uses `FuncOP` and returns nothing. If we check `funcOp.getNumResults()` too early, the bufferization of the function's body will be entirely skipped and fail the file check. 

https://github.com/llvm/llvm-project/blob/fc7a8936208847b66df9b5f93ba2222b93b9f8b6/mlir/test/Dialect/ControlFlow/one-shot-bufferize-analysis.mlir#L91-L111

https://github.com/llvm/llvm-project/pull/107295