[llvm-branch-commits] [mlir] 61422c8 - [mlir] Async: add support for lowering async value operands to LLVM
Eugene Zhulenev via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Fri Dec 25 02:29:24 PST 2020
Author: Eugene Zhulenev
Date: 2020-12-25T02:25:20-08:00
New Revision: 61422c8b661c506b25b5ac705cdbb2265af120d1
URL: https://github.com/llvm/llvm-project/commit/61422c8b661c506b25b5ac705cdbb2265af120d1
DIFF: https://github.com/llvm/llvm-project/commit/61422c8b661c506b25b5ac705cdbb2265af120d1.diff
LOG: [mlir] Async: add support for lowering async value operands to LLVM
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
}
```
Reviewed By: csigg
Differential Revision: https://reviews.llvm.org/D93598
Added:
Modified:
mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir
mlir/test/mlir-cpu-runner/async-value.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
index f1d6264606be..517c8d2c6f56 100644
--- a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
+++ b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
@@ -586,7 +586,7 @@ outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
// 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.
@@ -636,15 +636,26 @@ outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
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.
@@ -1069,14 +1080,6 @@ void ConvertAsyncToLLVMPass::runOnOperation() {
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();
diff --git a/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir b/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir
index dce0cf89628f..75b323490cbb 100644
--- a/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir
+++ b/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir
@@ -252,3 +252,54 @@ func @execute_and_return_f32() -> f32 {
// 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]])
+
+
diff --git a/mlir/test/mlir-cpu-runner/async-value.mlir b/mlir/test/mlir-cpu-runner/async-value.mlir
index 44b3b29e8491..31ed8860bf08 100644
--- a/mlir/test/mlir-cpu-runner/async-value.mlir
+++ b/mlir/test/mlir-cpu-runner/async-value.mlir
@@ -44,7 +44,7 @@ func @main() {
// ------------------------------------------------------------------------ //
%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>
}
@@ -53,8 +53,25 @@ func @main() {
// 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
More information about the llvm-branch-commits
mailing list