[mlir][bufferize] Fix bug in module equivalence analysis

[springerm]

Summary

CallOp result are not equivalent to an OpOperand if the OpOperand bufferizes out-of-place.

Depends On D126253

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

 }
 
 /// Gather equivalence info of CallOps.
 /// Note: This only adds new equivalence info if the called function was already
 /// analyzed.
 // TODO: This does not handle cyclic function call graphs etc.
 static void equivalenceAnalysis(func::FuncOp funcOp,
                                 BufferizationAliasInfo &aliasInfo,
-                                FuncAnalysisState &funcState) {
+                                OneShotAnalysisState &state) {
+  FuncAnalysisState &funcState = getFuncAnalysisState(state);
   funcOp->walk([&](func::CallOp callOp) {
     func::FuncOp calledFunction = getCalledFunction(callOp);
     assert(calledFunction && "could not retrieved called func::FuncOp");
 
     // No equivalence info available for the called function.
     if (!funcState.equivalentFuncArgs.count(calledFunction))
       return WalkResult::skip();
 
     for (auto it : funcState.equivalentFuncArgs[calledFunction]) {
       int64_t returnIdx = it.first;
       int64_t bbargIdx = it.second;
+      if (!state.isInPlace(callOp->getOpOperand(bbargIdx)))
+        continue;
       Value returnVal = callOp.getResult(returnIdx);
       Value argVal = callOp->getOperand(bbargIdx);
       aliasInfo.unionEquivalenceClasses(returnVal, argVal);
     }
 
     return WalkResult::advance();
   });
 }
for (func::FuncOp funcOp : orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
-    equivalenceAnalysis(funcOp, aliasInfo, funcState);
+    equivalenceAnalysis(funcOp, aliasInfo, state);
 
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, state)))
       return failure();
 
     // Run some extra function analyses.
     if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, state)) ||
         failed(funcOpBbArgReadWriteAnalysis(funcOp, state)))

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

 func.func @call_func_with_non_tensor_return(
     %t0: tensor<?xf32> {bufferization.writable = false}) -> (f32, tensor<?xf32>) {
   // CHECK: %[[alloc:.*]] = memref.alloc
   // CHECK-DAG: memref.copy %[[arg0]], %[[alloc]]
   // CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]
   // CHECK: %[[call:.*]] = call @inner_func(%[[casted]])
   %0, %1 = call @inner_func(%t0) : (tensor<?xf32>) -> (tensor<?xf32>, f32)
 
-  // Note: The tensor return value has folded away.
-  // CHECK: return %[[call]] : f32
+  // Note: The tensor return value cannot fold away because the CallOp
+  // bufferized out-of-place.
+  // CHECK: return %[[call]], %[[alloc]] : f32, memref<?xf32>
   return %1, %0 : f32, tensor<?xf32>
 }
 
 // -----
 
 // A chain of function calls. The last function f0 is potentially writing to the
 // buffer. This becomes a problem when bufferizing main and a copy must be
 // inserted then. (No copies in the other functions.)