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

[llvmlistbot at llvm.org]

================
@@ -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)) {
----------------
erick-xanadu wrote:

Hi @tzunghanjuang, you should not use `getNumResults` in this manner:

```
    if (!funcOp.getFunctionBody().empty() && funcOp.getNumResults() != 0) {
```

One of the errors appears to be in the tests themselves. I'll explain here for @matthias-springer:

The tests that involve the transform interpreter will attempt to bufferize the root module which contains the `transform::NamedSequenceOp`. The `transform::NamedSequenceOp` implements `FunctionOpInterface`, but the dialect does not have a `ReturnLike` op. The one-shot bufferizer will attempt to bufferize all nested operations inside the root module which includes the `transform::NamedSequenceOp`. However, this will obviously fail.

So I believe the correct way to solve this is to change the test such that the bufferization only happens in the intended sections of the module. This looks like it can be achieved by specifying a payload target with the option `--debug-payload-root-tag=<string>` or (as I just looked into other tests) using the `transform.structured.match` operation.

Looking at the tests, it seems that `transform.structured.match` is preferred, but I'll let others chime in.


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