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

[yijia1212]

Summary

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
}

[yijia1212]

Revert changes of async.call

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

void mlir::populateAsyncFuncToAsyncRuntimeConversionPatterns(
   patterns.add<AsyncFuncOpLowering, AsyncReturnOpLowering>(ctx, coros);
 
   patterns.add<AwaitTokenOpLowering, AwaitValueOpLowering, AwaitAllOpLowering>(
       ctx, coros, /*should_lower_blocking_wait=*/false);
   patterns.add<YieldOpLowering, AssertOpLowering>(ctx, coros);
 
   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();
       });
 }
 
 void AsyncFuncToAsyncRuntimePass::runOnOperation() {
   ModuleOp module = getOperation();
 
   // Lower async operations to async.runtime operations.
   MLIRContext *ctx = module->getContext();

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

 // CHECK:   cf.br ^[[ORIGIN_ENTRY:.*]]
 
 // CHECK: ^[[ORIGIN_ENTRY]]:
 // CHECK: %[[SAVED:.*]] = async.coro.save %[[HDL]]
 // CHECK: async.runtime.await_and_resume %[[arg1:.*]], %[[HDL]] :
 // 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

 async.func @async_func_passed_memref(%arg0 : !async.value<memref<f32>>) -> !async.token {
   %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>
   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
 
   // ------------------------------------------------------------------------ //
   // Check that simple async.func completes without errors.
   // ------------------------------------------------------------------------ //
   %token0 = async.call @async_func_empty() : () -> !async.token
func.func @main() {
   %token3 = async.call @async_func_passed_memref(%result1) : (!async.value<memref<f32>>) -> !async.token
   async.await %token3 : !async.token
 
   // CHECK: Unranked Memref
   // CHECK-SAME: rank = 0 offset = 0 sizes = [] strides = []
   // 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
 }
 
 func.func private @printMemrefF32(memref<*xf32>)
   attributes { llvm.emit_c_interface }