[mlir] Async: add support for lowering async value operands to LLVM

[ezhulenev]

Summary

Depends On D93592

Add support for async.execute async value unwrapping operands:

%token = async.execute(%async_value as %unwrapped : !async.value<!my.type>) {
  ...
  async.yield
}

[ezhulenev]

Fix style warnings

[ezhulenev]

Rebase

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
   ModuleOp module = execute->getParentOfType<ModuleOp>();
 
   MLIRContext *ctx = module.getContext();
   Location loc = execute.getLoc();
 
   // Collect all outlined function inputs.
   llvm::SetVector<mlir::Value> functionInputs(execute.dependencies().begin(),
                                               execute.dependencies().end());
-  assert(execute.operands().empty() && "operands are not supported");
+  functionInputs.insert(execute.operands().begin(), execute.operands().end());
   getUsedValuesDefinedAbove(execute.body(), functionInputs);
 
   // Collect types for the outlined function inputs and outputs.
   auto typesRange = llvm::map_range(
       functionInputs, [](Value value) { return value.getType(); });
   SmallVector<Type, 4> inputTypes(typesRange.begin(), typesRange.end());
   auto outputTypes = execute.getResultTypes();
 
outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
 
   // Split the entry block before the terminator.
   auto *terminatorOp = entryBlock->getTerminator();
   Block *suspended = terminatorOp->getBlock();
   Block *resume = suspended->splitBlock(terminatorOp);
   addSuspensionPoint(coro, coroSave.getResult(0), terminatorOp, suspended,
                      resume, builder);
 
+  size_t numDependencies = execute.dependencies().size();
+  size_t numOperands = execute.operands().size();
+
   // Await on all dependencies before starting to execute the body region.
   builder.setInsertionPointToStart(resume);
-  for (size_t i = 0; i < execute.dependencies().size(); ++i)
+  for (size_t i = 0; i < numDependencies; ++i)
     builder.create<AwaitOp>(func.getArgument(i));
 
+  // Await on all async value operands and unwrap the payload.
+  SmallVector<Value, 4> unwrappedOperands(numOperands);
+  for (size_t i = 0; i < numOperands; ++i) {
+    Value operand = func.getArgument(numDependencies + i);
+    unwrappedOperands[i] = builder.create<AwaitOp>(loc, operand).result();
+  }
+
   // Map from function inputs defined above the execute op to the function
   // arguments.
   BlockAndValueMapping valueMapping;
   valueMapping.map(functionInputs, func.getArguments());
+  valueMapping.map(execute.body().getArguments(), unwrappedOperands);
 
   // Clone all operations from the execute operation body into the outlined
   // function body.
   for (Operation &op : execute.body().getOps())
     builder.clone(op, valueMapping);
 
   // Replace the original `async.execute` with a call to outlined function.
   ImplicitLocOpBuilder callBuilder(loc, execute);
if (!llvm::all_of(execute.operands().getTypes(), isConvertibleToLlvm)) {
       execute.emitOpError("operands payload must be convertible to LLVM type");
       return WalkResult::interrupt();
     }
     if (!llvm::all_of(execute.results().getTypes(), isConvertibleToLlvm)) {
       execute.emitOpError("results payload must be convertible to LLVM type");
       return WalkResult::interrupt();
     }
 
-    // We currently do not support execute operations that have async value
-    // operands or produce async results.
-    if (!execute.operands().empty()) {
-      execute.emitOpError(
-          "can't outline async.execute op with async value operands");
-      return WalkResult::interrupt();
-    }
-
     outlinedFunctions.insert(outlineExecuteOp(symbolTable, execute));
 
     return WalkResult::advance();
   });
 
   // Failed to outline all async execute operations.
   if (outlineResult.wasInterrupted()) {
     signalPassFailure();

mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir

 // CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]])
 // CHECK: %[[ST_F32:.*]] = llvm.bitcast %[[STORAGE]]
 // CHECK: llvm.store %[[LLVM_CST]], %[[ST_F32]] : !llvm.ptr<float>
 // CHECK: call @mlirAsyncRuntimeEmplaceValue(%[[VALUE]])
 
 // Emplace result token.
 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])
 
+// -----
+
+// RUN: mlir-opt %s -split-input-file -convert-async-to-llvm | FileCheck %s
+
+func @async_value_operands() {
+  // CHECK: %[[RET:.*]]:2 = call @async_execute_fn
+  %token, %result = async.execute -> !async.value<f32> {
+    %c0 = constant 123.0 : f32
+    async.yield %c0 : f32
+  }
+
+  // CHECK: %[[TOKEN:.*]] = call @async_execute_fn_0(%[[RET]]#1)
+  %token0 = async.execute(%result as %value: !async.value<f32>) {
+    %0 = addf %value, %value : f32
+    async.yield
+  }
+
+  // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])
+  async.await %token0 : !async.token
+
+  return
+}
+
+// Function outlined from the first async.execute operation.
+// CHECK-LABEL: func private @async_execute_fn()
+
+// Function outlined from the second async.execute operation.
+// CHECK-LABEL: func private @async_execute_fn_0(%arg0: !llvm.ptr<i8>)
+// CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken()
+// CHECK: %[[HDL:.*]] = llvm.call @llvm.coro.begin
+
+// Suspend coroutine in the beginning.
+// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]],
+// CHECK: llvm.call @llvm.coro.suspend
+
+// Suspend coroutine second time waiting for the async operand.
+// CHECK: llvm.call @llvm.coro.save
+// CHECK: call @mlirAsyncRuntimeAwaitValueAndExecute(%arg0, %[[HDL]],
+// CHECK: llvm.call @llvm.coro.suspend
+
+// Get the operand value storage, cast to f32 and add the value.
+// CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%arg0)
+// CHECK: %[[ST_F32:.*]] = llvm.bitcast %[[STORAGE]]
+// CHECK: %[[LOADED:.*]] = llvm.load %[[ST_F32]] :  !llvm.ptr<float>
+// CHECK: %[[CASTED:.*]] = llvm.mlir.cast %[[LOADED]] : !llvm.float to f32
+// CHECK: addf %[[CASTED]], %[[CASTED]] : f32
+
+// Emplace result token.
+// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])
+
+

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

func @main() {
   // CHECK: 456.789
   vector.print %4 : f32
 
   // ------------------------------------------------------------------------ //
   // Memref allocated inside async.execute region.
   // ------------------------------------------------------------------------ //
   %token2, %result2 = async.execute[%token0] -> !async.value<memref<f32>> {
     %5 = alloc() : memref<f32>
-    %c0 = constant 987.654 : f32
+    %c0 = constant 0.25 : f32
     store %c0, %5[]: memref<f32>
     async.yield %5 : memref<f32>
   }
   %6 = async.await %result2 : !async.value<memref<f32>>
   %7 = memref_cast %6 :  memref<f32> to memref<*xf32>
 
   // CHECK: Unranked Memref
   // CHECK-SAME: rank = 0 offset = 0 sizes = [] strides = []
-  // CHECK-NEXT: [987.654]
+  // CHECK-NEXT: [0.25]
   call @print_memref_f32(%7): (memref<*xf32>) -> ()
+
+  // ------------------------------------------------------------------------ //
+  // Memref passed as async.execute operand.
+  // ------------------------------------------------------------------------ //
+  %token3 = async.execute(%result2 as %unwrapped : !async.value<memref<f32>>) {
+    %8 = load %unwrapped[]: memref<f32>
+    %9 = addf %8, %8 : f32
+    store %9, %unwrapped[]: memref<f32>
+    async.yield
+  }
+  async.await %token3 : !async.token
+
+  // CHECK: Unranked Memref
+  // CHECK-SAME: rank = 0 offset = 0 sizes = [] strides = []
+  // CHECK-NEXT: [0.5]
+  call @print_memref_f32(%7): (memref<*xf32>) -> ()
+
   dealloc %6 : memref<f32>
 
   return
 }
 
 func private @print_memref_f32(memref<*xf32>)
   attributes { llvm.emit_c_interface }