[mlir][async] Allow to call async.execute inside async.func

This change added support of calling async execute inside async.func.
Ex.

```
async.func @async_func_call_func() -> !async.token {
  %token = async.execute {
    %c0 = arith.constant 0 : index
    memref.store %arg0, %arg1[%c0] : memref<1xf32>
    async.yield
  }
  async.await %token : !async.token
  return
}
```

Reviewed By: ezhulenev

Differential Revision: https://reviews.llvm.org/D141730

Author: yijia1212

mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp

@@ -840,8 +840,9 @@ void mlir::populateAsyncFuncToAsyncRuntimeConversionPatterns(
 
   target.addDynamicallyLegalOp<AwaitOp, AwaitAllOp, YieldOp, cf::AssertOp>(
       [coros](Operation *op) {
+        auto exec = op->getParentOfType<ExecuteOp>();
         auto func = op->getParentOfType<func::FuncOp>();
-        return coros->find(func) == coros->end();
+        return exec || coros->find(func) == coros->end();
       });
 }
 

mlir/test/Dialect/Async/async-to-async-runtime.mlir

@@ -455,3 +455,31 @@ async.func @async_func_await(%arg0: f32, %arg1: !async.value<f32>)
 // CHECK-SAME: !async.value<f32>
 // CHECK: async.coro.suspend %[[SAVED]]
 // CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
+
+// -----
+// Async execute inside async func
+
+// CHECK-LABEL: @execute_in_async_func
+async.func @execute_in_async_func(%arg0: f32, %arg1: memref<1xf32>)
+             -> !async.token {
+  %token = async.execute {
+    %c0 = arith.constant 0 : index
+    memref.store %arg0, %arg1[%c0] : memref<1xf32>
+    async.yield
+  }
+  async.await %token : !async.token
+  return
+}
+// Call outlind async execute Function
+// CHECK: %[[RES:.*]] = call @async_execute_fn(
+// CHECK-SAME:    %[[VALUE:arg[0-9]+]],
+// CHECK-SAME:    %[[MEMREF:arg[0-9]+]]
+// CHECK-SAME:  ) : (f32, memref<1xf32>) -> !async.token
+
+// Function outlined from the async.execute operation.
+// CHECK-LABEL: func private @async_execute_fn(
+// CHECK-SAME:    %[[VALUE:arg[0-9]+]]: f32,
+// CHECK-SAME:    %[[MEMREF:arg[0-9]+]]: memref<1xf32>
+// CHECK-SAME:  ) -> !async.token
+// CHECK:         %[[CST:.*]] = arith.constant 0 : index
+// CHECK:         memref.store %[[VALUE]], %[[MEMREF]][%[[CST]]]

mlir/test/mlir-cpu-runner/async-func.mlir

@@ -64,6 +64,19 @@ async.func @async_func_passed_memref(%arg0 : !async.value<memref<f32>>) -> !asyn
   return
 }
 
+async.func @async_execute_in_async_func(%arg0 : !async.value<memref<f32>>) -> !async.token {
+  %token0 = async.execute {
+    %unwrapped = async.await %arg0 : !async.value<memref<f32>>
+    %0 = memref.load %unwrapped[] : memref<f32>
+    %1 = arith.addf %0, %0 : f32
+    memref.store %1, %unwrapped[] : memref<f32>
+    async.yield
+  }
+
+  async.await %token0 : !async.token
+  return
+}
+
 
 func.func @main() {
   %false = arith.constant 0 : i1
@@ -140,6 +153,17 @@ func.func @main() {
   // CHECK-NEXT: [0.5]
   call @printMemrefF32(%6) : (memref<*xf32>) -> ()
 
+  // ------------------------------------------------------------------------ //
+  // async.execute inside async.func
+  // ------------------------------------------------------------------------ //
+  %token4 = async.call @async_execute_in_async_func(%result1) : (!async.value<memref<f32>>) -> !async.token
+  async.await %token4 : !async.token
+
+  // CHECK: Unranked Memref
+  // CHECK-SAME: rank = 0 offset = 0 sizes = [] strides = []
+  // CHECK-NEXT: [1]
+  call @printMemrefF32(%6) : (memref<*xf32>) -> ()
+
   memref.dealloc %5 : memref<f32>
 
   return