[mlir] Async: check awaited operand error state after sync await

[ezhulenev]

Summary

Previously only await inside the async function (coroutine after lowering to async runtime) would check the error state

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

void AsyncRuntimePolicyBasedRefCountingPass::initializeDefaultPolicy() {
   policy.push_back([](OpOperand &operand) -> FailureOr<int> {
     Operation *op = operand.getOwner();
     Type type = operand.get().getType();
 
     bool isToken = type.isa<TokenType>();
     bool isGroup = type.isa<GroupType>();
     bool isValue = type.isa<ValueType>();
 
-    // Drop reference after async token or group await (sync await)
-    if (auto await = dyn_cast<RuntimeAwaitOp>(op))
-      return (isToken || isGroup) ? -1 : 0;
-
     // Drop reference after async token or group error check (coro await).
     if (auto await = dyn_cast<RuntimeIsErrorOp>(op))
       return (isToken || isGroup) ? -1 : 0;
 
     // Drop reference after async value load.
     if (auto load = dyn_cast<RuntimeLoadOp>(op))
       return isValue ? -1 : 0;
 

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

matchAndRewrite(AwaitType op, ArrayRef<Value> operands,
     // Check if await operation is inside the outlined coroutine function.
     auto func = op->template getParentOfType<FuncOp>();
     auto outlined = outlinedFunctions.find(func);
     const bool isInCoroutine = outlined != outlinedFunctions.end();
 
     Location loc = op->getLoc();
     Value operand = AwaitAdaptor(operands).operand();
 
+    Type i1 = rewriter.getI1Type();
+
     // Inside regular functions we use the blocking wait operation to wait for
     // the async object (token, value or group) to become available.
-    if (!isInCoroutine)
-      rewriter.create<RuntimeAwaitOp>(loc, operand);
+    if (!isInCoroutine) {
+      ImplicitLocOpBuilder builder(loc, op, rewriter.getListener());
+      builder.create<RuntimeAwaitOp>(loc, operand);
+
+      // Assert that the awaited operands is not in the error state.
+      Value isError = builder.create<RuntimeIsErrorOp>(i1, operand);
+      Value notError = builder.create<XOrOp>(
+          isError,
+          builder.create<ConstantOp>(loc, i1, builder.getIntegerAttr(i1, 1)));
+
+      builder.create<AssertOp>(notError,
+                               "Awaited async operand is in error state");
+    }
 
     // Inside the coroutine we convert await operation into coroutine suspension
     // point, and resume execution asynchronously.
     if (isInCoroutine) {
       CoroMachinery &coro = outlined->getSecond();
       Block *suspended = op->getBlock();
 
       ImplicitLocOpBuilder builder(loc, op, rewriter.getListener());
if (isInCoroutine) {
       builder.create<CoroSuspendOp>(coroSaveOp.state(), coro.suspend, resume,
                                     coro.cleanup);
 
       // 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);
+      auto isError = builder.create<RuntimeIsErrorOp>(loc, i1, operand);
       builder.create<CondBranchOp>(isError,
                                    /*trueDest=*/setupSetErrorBlock(coro),
                                    /*trueArgs=*/ArrayRef<Value>(),
                                    /*falseDest=*/continuation,
                                    /*falseArgs=*/ArrayRef<Value>());
 
       // Make sure that replacement value will be constructed in the
       // continuation block.
void AsyncToAsyncRuntimePass::runOnOperation() {
   runtimeTarget.addDynamicallyLegalDialect<scf::SCFDialect>([&](Operation *op) {
     auto walkResult = op->walk([&](Operation *nested) {
       bool isAsync = isa<async::AsyncDialect>(nested->getDialect());
       return isAsync && isInCoroutine(nested) ? WalkResult::interrupt()
                                               : WalkResult::advance();
     });
     return !walkResult.wasInterrupted();
   });
-  runtimeTarget.addLegalOp<BranchOp, CondBranchOp>();
+  runtimeTarget
+      .addLegalOp<AssertOp, XOrOp, ConstantOp, BranchOp, CondBranchOp>();
 
   // Assertions must be converted to runtime errors inside async functions.
   runtimeTarget.addDynamicallyLegalOp<AssertOp>([&](AssertOp op) -> bool {
     auto func = op->getParentOfType<FuncOp>();
     return outlinedFunctions.find(func) == outlinedFunctions.end();
   });
 
   if (eliminateBlockingAwaitOps)

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

 func @execute_no_async_args(%arg0: f32, %arg1: memref<1xf32>) {
   // CHECK: %[[TOKEN:.*]] = call @async_execute_fn(%arg0, %arg1)
   %token = async.execute {
     %c0 = constant 0 : index
     memref.store %arg0, %arg1[%c0] : memref<1xf32>
     async.yield
   }
   // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])
+  // CHECK: %[[IS_ERROR:.*]] = call @mlirAsyncRuntimeIsTokenError(%[[TOKEN]])
+  // CHECK: %[[TRUE:.*]] = constant true
+  // CHECK: %[[NOT_ERROR:.*]] = xor %[[IS_ERROR]], %[[TRUE]] : i1
+  // CHECK: assert %[[NOT_ERROR]]
   // CHECK-NEXT: return
   async.await %token : !async.token
   return
 }
 
 // Function outlined from the async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>)
 // CHECK-SAME: -> !llvm.ptr<i8>
%token1 = async.execute {
       async.yield
     }
     async.await %token1 : !async.token
 
     memref.store %arg1, %arg2[%c0] : memref<1xf32>
     async.yield
   }
   // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])
-  // CHECK-NEXT: return
+  // CHECK: %[[IS_ERROR:.*]] = call @mlirAsyncRuntimeIsTokenError(%[[TOKEN]])
+  // CHECK: %[[TRUE:.*]] = constant true
+  // CHECK: %[[NOT_ERROR:.*]] = xor %[[IS_ERROR]], %[[TRUE]] : i1
+  // CHECK: assert %[[NOT_ERROR]]
   async.await %token0 : !async.token
   return
 }
 
 // Function outlined from the inner async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>)
 // CHECK-SAME: -> !llvm.ptr<i8>
 // CHECK: %[[RET_0:.*]] = call @mlirAsyncRuntimeCreateToken()

mlir/test/Dialect/Async/async-runtime-policy-based-ref-counting.mlir

 // RUN: mlir-opt %s -async-runtime-policy-based-ref-counting | FileCheck %s
 
 // CHECK-LABEL: @token_await
 // CHECK:         %[[TOKEN:.*]]: !async.token
 func @token_await(%arg0: !async.token) {
   // CHECK: async.runtime.await %[[TOKEN]]
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK-NOT: async.runtime.drop_ref
   async.runtime.await %arg0 : !async.token
   return
 }
 
 // CHECK-LABEL: @group_await
 // CHECK:         %[[GROUP:.*]]: !async.group
 func @group_await(%arg0: !async.group) {
   // CHECK: async.runtime.await %[[GROUP]]
-  // CHECK: async.runtime.drop_ref %[[GROUP]] {count = 1 : i32}
+  // CHECK-NOT: async.runtime.drop_ref
   async.runtime.await %arg0 : !async.group
   return
 }
 
 // CHECK-LABEL: @add_token_to_group
 // CHECK:         %[[GROUP:.*]]: !async.group
 // CHECK:         %[[TOKEN:.*]]: !async.token
 func @add_token_to_group(%arg0: !async.group, %arg1: !async.token) {

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

%token1 = async.execute {
       async.yield
     }
     async.await %token1 : !async.token
 
     memref.store %arg1, %arg2[%c0] : memref<1xf32>
     async.yield
   }
   // CHECK: async.runtime.await %[[TOKEN]]
+  // CHECK: %[[IS_ERROR:.*]] = async.runtime.is_error %[[TOKEN]]
+  // CHECK: %[[TRUE:.*]] = constant true
+  // CHECK: %[[NOT_ERROR:.*]] = xor %[[IS_ERROR]], %[[TRUE]] : i1
+  // CHECK: assert %[[NOT_ERROR]]
   // CHECK-NEXT: return
   async.await %token0 : !async.token
   return
 }
 
 // Function outlined from the inner async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn
 // CHECK-SAME: -> !async.token