[mlir] Async: Add error propagation support to async groups

[ezhulenev]

Summary

Depends On D103109

If any of the tokens/values added to the !async.group switches to the error state, than the group itself switches to the error state.

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

def Async_RuntimeIsErrorOp : Async_Op<"runtime.is_error"> {
   let summary = "returns true if token, value or group is in error state";
   let description = [{
     The `async.runtime.is_error` operation returns true if the token, value or
     group (any of the async runtime values) is in the error state. It is the
     caller responsibility to check error state after the call to `await` or
     resuming after `await_and_resume`.
   }];
 
-  let arguments = (ins Async_AnyValueOrTokenType:$operand);
+  let arguments = (ins Async_AnyAsyncType:$operand);
   let results = (outs I1:$is_error);
 
   let assemblyFormat = "$operand attr-dict `:` type($operand)";
 }
 
 def Async_RuntimeAwaitOp : Async_Op<"runtime.await"> {
   let summary = "blocks the caller thread until the operand becomes available";
   let description = [{

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

 extern "C" void mlirAsyncRuntimeSetValueError(AsyncValue *);
 
 // Returns true if token is in the error state.
 extern "C" bool mlirAsyncRuntimeIsTokenError(AsyncToken *);
 
 // Returns true if value is in the error state.
 extern "C" bool mlirAsyncRuntimeIsValueError(AsyncValue *);
 
+// Returns true if group is in the error state (any of the tokens or values
+// added to the group are in the error state).
+extern "C" bool mlirAsyncRuntimeIsGroupError(AsyncGroup *);
+
 // Blocks the caller thread until the token becomes ready.
 extern "C" void mlirAsyncRuntimeAwaitToken(AsyncToken *);
 
 // Blocks the caller thread until the value becomes ready.
 extern "C" void mlirAsyncRuntimeAwaitValue(AsyncValue *);
 
 // Blocks the caller thread until the elements in the group become ready.
 extern "C" void mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *);

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/StandardOps/Transforms/FuncConversions.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 #define DEBUG_TYPE "convert-async-to-llvm"
 
 using namespace mlir;
 using namespace mlir::async;
 
 //===----------------------------------------------------------------------===//
 // Async Runtime C API declaration.
 //===----------------------------------------------------------------------===//
 
 static constexpr const char *kAddRef = "mlirAsyncRuntimeAddRef";
 static constexpr const char *kDropRef = "mlirAsyncRuntimeDropRef";
 static constexpr const char *kCreateToken = "mlirAsyncRuntimeCreateToken";
 static constexpr const char *kCreateValue = "mlirAsyncRuntimeCreateValue";
 static constexpr const char *kCreateGroup = "mlirAsyncRuntimeCreateGroup";
 static constexpr const char *kEmplaceToken = "mlirAsyncRuntimeEmplaceToken";
 static constexpr const char *kEmplaceValue = "mlirAsyncRuntimeEmplaceValue";
 static constexpr const char *kSetTokenError = "mlirAsyncRuntimeSetTokenError";
 static constexpr const char *kSetValueError = "mlirAsyncRuntimeSetValueError";
 static constexpr const char *kIsTokenError = "mlirAsyncRuntimeIsTokenError";
 static constexpr const char *kIsValueError = "mlirAsyncRuntimeIsValueError";
+static constexpr const char *kIsGroupError = "mlirAsyncRuntimeIsGroupError";
 static constexpr const char *kAwaitToken = "mlirAsyncRuntimeAwaitToken";
 static constexpr const char *kAwaitValue = "mlirAsyncRuntimeAwaitValue";
 static constexpr const char *kAwaitGroup = "mlirAsyncRuntimeAwaitAllInGroup";
 static constexpr const char *kExecute = "mlirAsyncRuntimeExecute";
 static constexpr const char *kGetValueStorage =
     "mlirAsyncRuntimeGetValueStorage";
 static constexpr const char *kAddTokenToGroup =
     "mlirAsyncRuntimeAddTokenToGroup";
struct AsyncAPI {
   }
 
   static FunctionType isValueErrorFunctionType(MLIRContext *ctx) {
     auto value = opaquePointerType(ctx);
     auto i1 = IntegerType::get(ctx, 1);
     return FunctionType::get(ctx, {value}, {i1});
   }
 
+  static FunctionType isGroupErrorFunctionType(MLIRContext *ctx) {
+    auto i1 = IntegerType::get(ctx, 1);
+    return FunctionType::get(ctx, {GroupType::get(ctx)}, {i1});
+  }
+
   static FunctionType awaitTokenFunctionType(MLIRContext *ctx) {
     return FunctionType::get(ctx, {TokenType::get(ctx)}, {});
   }
 
   static FunctionType awaitValueFunctionType(MLIRContext *ctx) {
     auto value = opaquePointerType(ctx);
     return FunctionType::get(ctx, {value}, {});
   }
static void addAsyncRuntimeApiDeclarations(ModuleOp module) {
   addFuncDecl(kCreateValue, AsyncAPI::createValueFunctionType(ctx));
   addFuncDecl(kCreateGroup, AsyncAPI::createGroupFunctionType(ctx));
   addFuncDecl(kEmplaceToken, AsyncAPI::emplaceTokenFunctionType(ctx));
   addFuncDecl(kEmplaceValue, AsyncAPI::emplaceValueFunctionType(ctx));
   addFuncDecl(kSetTokenError, AsyncAPI::setTokenErrorFunctionType(ctx));
   addFuncDecl(kSetValueError, AsyncAPI::setValueErrorFunctionType(ctx));
   addFuncDecl(kIsTokenError, AsyncAPI::isTokenErrorFunctionType(ctx));
   addFuncDecl(kIsValueError, AsyncAPI::isValueErrorFunctionType(ctx));
+  addFuncDecl(kIsGroupError, AsyncAPI::isGroupErrorFunctionType(ctx));
   addFuncDecl(kAwaitToken, AsyncAPI::awaitTokenFunctionType(ctx));
   addFuncDecl(kAwaitValue, AsyncAPI::awaitValueFunctionType(ctx));
   addFuncDecl(kAwaitGroup, AsyncAPI::awaitGroupFunctionType(ctx));
   addFuncDecl(kExecute, AsyncAPI::executeFunctionType(ctx));
   addFuncDecl(kGetValueStorage, AsyncAPI::getValueStorageFunctionType(ctx));
   addFuncDecl(kAddTokenToGroup, AsyncAPI::addTokenToGroupFunctionType(ctx));
   addFuncDecl(kAwaitTokenAndExecute,
               AsyncAPI::awaitTokenAndExecuteFunctionType(ctx));

 class RuntimeSetAvailableOpLowering
     : public OpConversionPattern<RuntimeSetAvailableOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(RuntimeSetAvailableOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    Type operandType = op.operand().getType();
-    rewriter.replaceOpWithNewOp<CallOp>(
-        op, operandType.isa<TokenType>() ? kEmplaceToken : kEmplaceValue,
-        TypeRange(), operands);
+    StringRef apiFuncName =
+        TypeSwitch<Type, StringRef>(op.operand().getType())
+            .Case<TokenType>([](Type) { return kEmplaceToken; })
+            .Case<ValueType>([](Type) { return kEmplaceValue; });
+
+    rewriter.replaceOpWithNewOp<CallOp>(op, apiFuncName, TypeRange(), operands);
+
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Convert async.runtime.set_error to the corresponding runtime API call.
 //===----------------------------------------------------------------------===//
 
 namespace {
 class RuntimeSetErrorOpLowering
     : public OpConversionPattern<RuntimeSetErrorOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(RuntimeSetErrorOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    Type operandType = op.operand().getType();
-    rewriter.replaceOpWithNewOp<CallOp>(
-        op, operandType.isa<TokenType>() ? kSetTokenError : kSetValueError,
-        TypeRange(), operands);
+    StringRef apiFuncName =
+        TypeSwitch<Type, StringRef>(op.operand().getType())
+            .Case<TokenType>([](Type) { return kSetTokenError; })
+            .Case<ValueType>([](Type) { return kSetValueError; });
+
+    rewriter.replaceOpWithNewOp<CallOp>(op, apiFuncName, TypeRange(), operands);
+
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Convert async.runtime.is_error to the corresponding runtime API call.
 //===----------------------------------------------------------------------===//
 
 namespace {
 class RuntimeIsErrorOpLowering : public OpConversionPattern<RuntimeIsErrorOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(RuntimeIsErrorOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    Type operandType = op.operand().getType();
-    rewriter.replaceOpWithNewOp<CallOp>(
-        op, operandType.isa<TokenType>() ? kIsTokenError : kIsValueError,
-        rewriter.getI1Type(), operands);
+    StringRef apiFuncName =
+        TypeSwitch<Type, StringRef>(op.operand().getType())
+            .Case<TokenType>([](Type) { return kIsTokenError; })
+            .Case<GroupType>([](Type) { return kIsGroupError; })
+            .Case<ValueType>([](Type) { return kIsValueError; });
+
+    rewriter.replaceOpWithNewOp<CallOp>(op, apiFuncName, rewriter.getI1Type(),
+                                        operands);
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Convert async.runtime.await to the corresponding runtime API call.
 //===----------------------------------------------------------------------===//
 
 namespace {
 class RuntimeAwaitOpLowering : public OpConversionPattern<RuntimeAwaitOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(RuntimeAwaitOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    Type operandType = op.operand().getType();
-
-    StringRef apiFuncName;
-    if (operandType.isa<TokenType>())
-      apiFuncName = kAwaitToken;
-    else if (operandType.isa<ValueType>())
-      apiFuncName = kAwaitValue;
-    else if (operandType.isa<GroupType>())
-      apiFuncName = kAwaitGroup;
-    else
-      return rewriter.notifyMatchFailure(op, "unsupported async type");
+    StringRef apiFuncName =
+        TypeSwitch<Type, StringRef>(op.operand().getType())
+            .Case<TokenType>([](Type) { return kAwaitToken; })
+            .Case<ValueType>([](Type) { return kAwaitValue; })
+            .Case<GroupType>([](Type) { return kAwaitGroup; });
 
     rewriter.create<CallOp>(op->getLoc(), apiFuncName, TypeRange(), operands);
     rewriter.eraseOp(op);
 
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Convert async.runtime.await_and_resume to the corresponding runtime API call.
 //===----------------------------------------------------------------------===//
 
 namespace {
 class RuntimeAwaitAndResumeOpLowering
     : public OpConversionPattern<RuntimeAwaitAndResumeOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(RuntimeAwaitAndResumeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    Type operandType = op.operand().getType();
-
-    StringRef apiFuncName;
-    if (operandType.isa<TokenType>())
-      apiFuncName = kAwaitTokenAndExecute;
-    else if (operandType.isa<ValueType>())
-      apiFuncName = kAwaitValueAndExecute;
-    else if (operandType.isa<GroupType>())
-      apiFuncName = kAwaitAllAndExecute;
-    else
-      return rewriter.notifyMatchFailure(op, "unsupported async type");
+    StringRef apiFuncName =
+        TypeSwitch<Type, StringRef>(op.operand().getType())
+            .Case<TokenType>([](Type) { return kAwaitTokenAndExecute; })
+            .Case<ValueType>([](Type) { return kAwaitValueAndExecute; })
+            .Case<GroupType>([](Type) { return kAwaitAllAndExecute; });
 
     Value operand = RuntimeAwaitAndResumeOpAdaptor(operands).operand();
     Value handle = RuntimeAwaitAndResumeOpAdaptor(operands).handle();
 
     // A pointer to coroutine resume intrinsic wrapper.
     addResumeFunction(op->getParentOfType<ModuleOp>());
     auto resumeFnTy = AsyncAPI::resumeFunctionType(op->getContext());
     auto resumePtr = rewriter.create<LLVM::AddressOfOp>(

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

if (isInCoroutine) {
       // Split the entry block before the await operation.
       Block *resume = rewriter.splitBlock(suspended, Block::iterator(op));
 
       // Add async.coro.suspend as a suspended block terminator.
       builder.setInsertionPointToEnd(suspended);
       builder.create<CoroSuspendOp>(coroSaveOp.state(), coro.suspend, resume,
                                     coro.cleanup);
 
-      // TODO: Async groups do not yet support runtime errors.
-      if (!std::is_same<AwaitAllOp, AwaitType>::value) {
-        // Split the resume block into error checking and continuation.
-        Block *continuation = rewriter.splitBlock(resume, Block::iterator(op));
 
-        // Check if the awaited value is in the error state.
-        builder.setInsertionPointToStart(resume);
-        auto isError = builder.create<RuntimeIsErrorOp>(
-            loc, rewriter.getI1Type(), operand);
-        builder.create<CondBranchOp>(isError,
-                                     /*trueDest=*/setupSetErrorBlock(coro),
-                                     /*trueArgs=*/ArrayRef<Value>(),
-                                     /*falseDest=*/continuation,
-                                     /*falseArgs=*/ArrayRef<Value>());
+      auto isError =
+          builder.create<RuntimeIsErrorOp>(loc, rewriter.getI1Type(), operand);
 
-        // Make sure that replacement value will be constructed in the
-        // continuation block.
-        rewriter.setInsertionPointToStart(continuation);
-      }
-    }
 
     // Erase or replace the await operation with the new value.
     if (Value replaceWith = getReplacementValue(op, operand, rewriter))
       rewriter.replaceOp(op, replaceWith);
     else
       rewriter.eraseOp(op);
 
     return success();

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

struct AsyncValue : public RefCounted {
   std::vector<std::function<void()>> awaiters;
 };
 
 // Async group provides a mechanism to group together multiple async tokens or
 // values to await on all of them together (wait for the completion of all
 // tokens or values added to the group).
 struct AsyncGroup : public RefCounted {
   AsyncGroup(AsyncRuntime *runtime)
-      : RefCounted(runtime), pendingTokens(0), rank(0) {}
+      : RefCounted(runtime), pendingTokens(0), numErrors(0), rank(0) {}
 
   std::atomic<int> pendingTokens;
+  std::atomic<int> numErrors;
   std::atomic<int> rank;
 
   // Pending awaiters are guarded by a mutex.
   std::mutex mu;
   std::condition_variable cv;
   std::vector<std::function<void()>> awaiters;
 };
 
extern "C" int64_t mlirAsyncRuntimeAddTokenToGroup(AsyncToken *token,
                                                    AsyncGroup *group) {
   std::unique_lock<std::mutex> lockToken(token->mu);
   std::unique_lock<std::mutex> lockGroup(group->mu);
 
   // Get the rank of the token inside the group before we drop the reference.
   int rank = group->rank.fetch_add(1);
   group->pendingTokens.fetch_add(1);
 
-  auto onTokenReady = [group]() {
+  auto onTokenReady = [group, token]() {
+    // Increment the number of errors in the group.
+    if (State(token->state).isError())
+      group->numErrors.fetch_add(1);
+
     // Run all group awaiters if it was the last token in the group.
     if (group->pendingTokens.fetch_sub(1) == 1) {
       group->cv.notify_all();
       for (auto &awaiter : group->awaiters)
         awaiter();
     }
   };
 

 extern "C" bool mlirAsyncRuntimeIsTokenError(AsyncToken *token) {
   return State(token->state).isError();
 }
 
 extern "C" bool mlirAsyncRuntimeIsValueError(AsyncValue *value) {
   return State(value->state).isError();
 }
 
+extern "C" bool mlirAsyncRuntimeIsGroupError(AsyncGroup *group) {
+  return group->numErrors.load() > 0;
+}
+
 extern "C" void mlirAsyncRuntimeAwaitToken(AsyncToken *token) {
   std::unique_lock<std::mutex> lock(token->mu);
   if (!State(token->state).isAvailableOrError())
     token->cv.wait(
         lock, [token] { return State(token->state).isAvailableOrError(); });
 }
 
 extern "C" void mlirAsyncRuntimeAwaitValue(AsyncValue *value) {
void __mlir_runner_init(llvm::StringMap<void *> &exportSymbols) {
   exportSymbol("mlirAsyncRuntimeSetTokenError",
                &mlir::runtime::mlirAsyncRuntimeSetTokenError);
   exportSymbol("mlirAsyncRuntimeSetValueError",
                &mlir::runtime::mlirAsyncRuntimeSetValueError);
   exportSymbol("mlirAsyncRuntimeIsTokenError",
                &mlir::runtime::mlirAsyncRuntimeIsTokenError);
   exportSymbol("mlirAsyncRuntimeIsValueError",
                &mlir::runtime::mlirAsyncRuntimeIsValueError);
+  exportSymbol("mlirAsyncRuntimeIsGroupError",
+               &mlir::runtime::mlirAsyncRuntimeIsGroupError);
   exportSymbol("mlirAsyncRuntimeAwaitToken",
                &mlir::runtime::mlirAsyncRuntimeAwaitToken);
   exportSymbol("mlirAsyncRuntimeAwaitValue",
                &mlir::runtime::mlirAsyncRuntimeAwaitValue);
   exportSymbol("mlirAsyncRuntimeAwaitTokenAndExecute",
                &mlir::runtime::mlirAsyncRuntimeAwaitTokenAndExecute);
   exportSymbol("mlirAsyncRuntimeAwaitValueAndExecute",
                &mlir::runtime::mlirAsyncRuntimeAwaitValueAndExecute);

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

 // CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_0:.*]], ^[[CLEANUP:.*]]
 
 // Suspend coroutine second time waiting for the group.
 // CHECK: ^[[RESUME_0]]:
 // CHECK:   async.runtime.await_and_resume %[[ARG]], %[[HDL]]
 // CHECK:   async.coro.suspend
 // CHECK-SAME: ^[[SUSPEND]], ^[[RESUME_1:.*]], ^[[CLEANUP]]
 
-// Emplace result token.
+// Check the error of the awaited token after resumption.
 // CHECK: ^[[RESUME_1]]:
+// CHECK:   %[[ERR:.*]] = async.runtime.is_error %[[ARG]]
+// CHECK:   cond_br %[[ERR]], ^[[SET_ERROR:.*]], ^[[CONTINUATION:.*]]
+
+// Emplace result token after error checking.
+// CHECK: ^[[CONTINUATION:.*]]:
 // CHECK:   async.runtime.set_available %[[TOKEN]]
 
 // CHECK: ^[[CLEANUP]]:
 // CHECK: ^[[SUSPEND]]:
 
 // -----
 
 // CHECK-LABEL: @execute_and_return_f32

mlir/test/Dialect/Async/runtime.mlir

 
 // CHECK-LABEL: @is_value_error
 func @is_value_error(%arg0: !async.value<f32>) -> i1 {
   // CHECK: %[[ERR:.*]] = async.runtime.is_error %arg0 : !async.value<f32>
   %0 = async.runtime.is_error %arg0 : !async.value<f32>
   return %0 : i1
 }
 
+// CHECK-LABEL: @is_group_error
+func @is_group_error(%arg0: !async.group) -> i1 {
+  // CHECK: %[[ERR:.*]] = async.runtime.is_error %arg0 : !async.group
+  %0 = async.runtime.is_error %arg0 : !async.group
+  return %0 : i1
+}
+
 // CHECK-LABEL: @await_token
 func @await_token(%arg0: !async.token) {
   // CHECK: async.runtime.await %arg0 : !async.token
   async.runtime.await %arg0 : !async.token
   return
 }
 
 // CHECK-LABEL: @await_value

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

func @main() {
 
   // CHECK: 1
   // CHECK: 1
   %err3_0 = async.runtime.is_error %token3 : !async.token
   %err3_1 = async.runtime.is_error %value3 : !async.value<f32>
   vector.print %err3_0 : i1
   vector.print %err3_1 : i1
 
+  // ------------------------------------------------------------------------ //
+  // Check error propagation from a token to the group.
+  // ------------------------------------------------------------------------ //
+
+  %group0 = async.create_group
+
+  %token4 = async.execute {
+    async.yield
+  }
+
+  %token5 = async.execute {
+    assert %false, "error"
+    async.yield
+  }
+
+  %idx0 = async.add_to_group %token4, %group0 : !async.token
+  %idx1 = async.add_to_group %token5, %group0 : !async.token
+
+  async.runtime.await %group0 : !async.group
+
+  // CHECK: 1
+  %err4 = async.runtime.is_error %group0 : !async.group
+  vector.print %err4 : i1
+
   return
 }