[mlir][memref] Fix hoist-static-allocs option of buffer-results-to-out-params when function parameters are returned (#102093)

buffer-results-to-out-params pass will have a nullptr-referencing error
when hoist-static-allocs option is on, when the return value of a
function is a parameter of the function. This PR fixes this issue.

Author: Menooker

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

@@ -120,7 +120,8 @@ static LogicalResult updateReturnOps(func::FuncOp func,
     }
     OpBuilder builder(op);
     for (auto [orig, arg] : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {
-      if (hoistStaticAllocs && isa<memref::AllocOp>(orig.getDefiningOp()) &&
+      if (hoistStaticAllocs &&
+          isa_and_nonnull<memref::AllocOp>(orig.getDefiningOp()) &&
           mlir::cast<MemRefType>(orig.getType()).hasStaticShape()) {
         orig.replaceAllUsesWith(arg);
         orig.getDefiningOp()->erase();

mlir/test/Transforms/buffer-results-to-out-params-elim.mlir

@@ -34,4 +34,18 @@ func.func @basic_dynamic(%d: index) -> (memref<?xf32>) {
   %b = memref.alloc(%d) : memref<?xf32>
   "test.source"(%b)  : (memref<?xf32>) -> ()
   return %b : memref<?xf32>
-}
+}
+
+// -----
+
+// no change due to writing to func args
+// CHECK-LABEL:   func @return_arg(
+// CHECK-SAME:        %[[ARG0:.*]]: memref<128x256xf32>, %[[ARG1:.*]]: memref<128x256xf32>, %[[ARG2:.*]]: memref<128x256xf32>) {
+// CHECK:           "test.source"(%[[ARG0]], %[[ARG1]])
+// CHECK:           memref.copy
+// CHECK:           return
+// CHECK:         }
+func.func @return_arg(%arg0: memref<128x256xf32>, %arg1: memref<128x256xf32>) -> memref<128x256xf32> {
+  "test.source"(%arg0, %arg1)  : (memref<128x256xf32>, memref<128x256xf32>) -> ()
+  return %arg0 : memref<128x256xf32>
+}