[CIR][IR] Deprecate cir.yield nosuspend

This changes the `cir.await` operation to expect a `cir.condition` as
the terminator for the ready region. This simplifies the `cir.await`
while also simplifying the `cir.yield`. If `cir.condition` holds a
true value, then the `cir.await` will continue the coroutine, otherwise,
it will suspend its execution.

The `cir.condition` op was also updated to allow `cir.await` as its
parent operation.

ghstack-source-id: 1ebeb2c
Pull Request resolved: #396

Author: sitio-couto

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -593,25 +593,50 @@ def ConditionOp : CIR_Op<"condition", [
 ]> {
   let summary = "Loop continuation condition.";
   let description = [{
-    The `cir.condition` termintes loop's conditional regions. It takes a single
-    `cir.bool` operand. if the operand is true, the loop continues, otherwise
-    it terminates.
+    The `cir.condition` terminates conditional regions. It takes a single
+    `cir.bool` operand and, depending on its value, may branch to different
+    regions:
+
+     - When in the `cond` region of a `cir.loop`, it continues the loop
+       if true, or exits it if false.
+     - When in the `ready` region of a `cir.await`, it branches to the `resume`
+       region when true, and to the `suspend` region when false.
+
+    Example:
+
+    ```mlir
+    cir.loop for(cond : {
+      cir.condition(%arg0) // Branches to `step` region or exits.
+    }, step : {
+      [...]
+    }) {
+      [...]
+    }
+
+    cir.await(user, ready : {
+      cir.condition(%arg0) // Branches to `resume` or `suspend` region.
+    }, suspend : {
+      [...]
+    }, resume : {
+      [...]
+    },)
+    ```
   }];
   let arguments = (ins CIR_BoolType:$condition);
   let assemblyFormat = " `(` $condition `)` attr-dict ";
+  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 def YieldOpKind_FT : I32EnumAttrCase<"Fallthrough", 2, "fallthrough">;
-def YieldOpKind_NS : I32EnumAttrCase<"NoSuspend", 4, "nosuspend">;
 
 def YieldOpKind : I32EnumAttr<
     "YieldOpKind",
     "yield kind",
-    [YieldOpKind_FT, YieldOpKind_NS]> {
+    [YieldOpKind_FT]> {
   let cppNamespace = "::mlir::cir";
 }
 
@@ -630,8 +655,6 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
     Optionally, `cir.yield` can be annotated with extra kind specifiers:
     - `fallthrough`: execution falls to the next region in `cir.switch` case list.
     Only available inside `cir.switch` regions.
-    - `nosuspend`: specific to the `ready` region inside `cir.await` op, it makes
-    control-flow to be transfered back to the parent, preventing suspension.
 
     As a general rule, `cir.yield` must be explicitly used whenever a region has
     more than one block and no terminator, or within `cir.switch` regions not
@@ -651,16 +674,6 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
       }, ...
     ]
 
-    cir.await(init, ready : {
-      // Call std::suspend_always::await_ready
-      %18 = cir.call @_ZNSt14suspend_always11await_readyEv(...)
-      cir.if %18 {
-        // yields back to the parent.
-        cir.yield nosuspend
-      }
-      cir.yield // control-flow to the next region for suspension.
-    }, ...)
-
     cir.scope {
       ...
       cir.yield
@@ -704,9 +717,6 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
     bool isFallthrough() {
       return !isPlain() && *getKind() == YieldOpKind::Fallthrough;
     }
-    bool isNoSuspend() {
-      return !isPlain() && *getKind() == YieldOpKind::NoSuspend;
-    }
   }];
 
   let hasVerifier = 1;

clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp

@@ -428,28 +428,8 @@ buildSuspendExpression(CIRGenFunction &CGF, CGCoroData &Coro,
       CGF.getLoc(S.getSourceRange()), Kind,
       /*readyBuilder=*/
       [&](mlir::OpBuilder &b, mlir::Location loc) {
-        auto *cond = S.getReadyExpr();
-        cond = cond->IgnoreParens();
-        mlir::Value condV = CGF.evaluateExprAsBool(cond);
-
-        builder.create<mlir::cir::IfOp>(
-            loc, condV, /*withElseRegion=*/false,
-            /*thenBuilder=*/
-            [&](mlir::OpBuilder &b, mlir::Location loc) {
-              // If expression is ready, no need to suspend,
-              // `YieldOpKind::NoSuspend` tells control flow to return to
-              // parent, no more regions to be executed.
-              builder.create<mlir::cir::YieldOp>(
-                  loc, mlir::cir::YieldOpKind::NoSuspend);
-            });
-
-        if (!condV) {
-          awaitBuild = mlir::failure();
-          return;
-        }
-
-        // Signals the parent that execution flows to next region.
-        builder.create<mlir::cir::YieldOp>(loc);
+        Expr *condExpr = S.getReadyExpr()->IgnoreParens();
+        builder.createCondition(CGF.evaluateExprAsBool(condExpr));
       },
       /*suspendBuilder=*/
       [&](mlir::OpBuilder &b, mlir::Location loc) {

clang/lib/CIR/Dialect/IR/CIRDialect.cpp

@@ -14,6 +14,7 @@
 #include "clang/CIR/Dialect/IR/CIRAttrs.h"
 #include "clang/CIR/Dialect/IR/CIROpsEnums.h"
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
+#include "llvm/Support/ErrorHandling.h"
 #include <optional>
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -246,13 +247,19 @@ LogicalResult BreakOp::verify() {
 
 void ConditionOp::getSuccessorRegions(
     ArrayRef<Attribute> operands, SmallVectorImpl<RegionSuccessor> &regions) {
-  auto loopOp = cast<LoopOp>(getOperation()->getParentOp());
-
   // TODO(cir): The condition value may be folded to a constant, narrowing
   // down its list of possible successors.
-  // Condition may branch to the body or to the parent op.
-  regions.emplace_back(&loopOp.getBody(), loopOp.getBody().getArguments());
-  regions.emplace_back(loopOp->getResults());
+
+  // Parent is a loop: condition may branch to the body or to the parent op.
+  if (auto loopOp = dyn_cast<LoopOp>(getOperation()->getParentOp())) {
+    regions.emplace_back(&loopOp.getBody(), loopOp.getBody().getArguments());
+    regions.emplace_back(loopOp->getResults());
+  }
+
+  // Parent is an await: condition may branch to resume or suspend regions.
+  auto await = cast<AwaitOp>(getOperation()->getParentOp());
+  regions.emplace_back(&await.getResume(), await.getResume().getArguments());
+  regions.emplace_back(&await.getSuspend(), await.getSuspend().getArguments());
 }
 
 MutableOperandRange
@@ -261,6 +268,12 @@ ConditionOp::getMutableSuccessorOperands(RegionBranchPoint point) {
   return MutableOperandRange(getOperation(), 0, 0);
 }
 
+LogicalResult ConditionOp::verify() {
+  if (!isa<LoopOp, AwaitOp>(getOperation()->getParentOp()))
+    return emitOpError("condition must be within a conditional region");
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
@@ -786,34 +799,6 @@ void TernaryOp::build(OpBuilder &builder, OperationState &result, Value cond,
 //===----------------------------------------------------------------------===//
 
 mlir::LogicalResult YieldOp::verify() {
-  auto isDominatedByProperAwaitRegion = [&](Operation *parentOp,
-                                            mlir::Region *currRegion) {
-    while (!llvm::isa<cir::FuncOp>(parentOp)) {
-      auto awaitOp = dyn_cast<cir::AwaitOp>(parentOp);
-      if (awaitOp) {
-        if (currRegion && currRegion == &awaitOp.getResume()) {
-          emitOpError() << "kind 'nosuspend' can only be used in 'ready' and "
-                           "'suspend' regions";
-          return false;
-        }
-        return true;
-      }
-
-      currRegion = parentOp->getParentRegion();
-      parentOp = parentOp->getParentOp();
-    }
-
-    emitOpError() << "shall be dominated by 'cir.await'";
-    return false;
-  };
-
-  if (isNoSuspend()) {
-    if (!isDominatedByProperAwaitRegion(getOperation()->getParentOp(),
-                                        getOperation()->getParentRegion()))
-      return mlir::failure();
-    return mlir::success();
-  }
-
   if (isFallthrough()) {
     if (!llvm::isa<SwitchOp>(getOperation()->getParentOp()))
       return emitOpError() << "fallthrough only expected within 'cir.switch'";
@@ -2223,7 +2208,11 @@ void AwaitOp::getSuccessorRegions(mlir::RegionBranchPoint point,
   regions.push_back(RegionSuccessor(&this->getResume()));
 }
 
-LogicalResult AwaitOp::verify() { return success(); }
+LogicalResult AwaitOp::verify() {
+  if (!isa<ConditionOp>(this->getReady().back().getTerminator()))
+    return emitOpError("ready region must end with cir.condition");
+  return success();
+}
 
 //===----------------------------------------------------------------------===//
 // CIR defined traits

clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp

@@ -1391,8 +1391,6 @@ class CIRSwitchOpLowering
           case mlir::cir::YieldOpKind::Fallthrough:
             fallthroughYieldOp = yieldOp;
             break;
-          default:
-            return op->emitError("invalid yield kind in case statement");
           }
         }
       }

clang/test/CIR/CodeGen/coro-task.cpp

@@ -206,10 +206,7 @@ VoidTask silly_task() {
 // CHECK:       %[[#TmpCallRes:]] = cir.call @_ZNSt14suspend_always11await_readyEv(%[[#SuspendAlwaysAddr]])
 // CHECK:       cir.yield %[[#TmpCallRes]] : !cir.bool
 // CHECK:     }
-// CHECK:     cir.if %[[#ReadyVeto]] {
-// CHECK:       cir.yield nosuspend
-// CHECK:     }
-// CHECK:     cir.yield
+// CHECK:     cir.condition(%[[#ReadyVeto]])
 
 // Second region `suspend` contains the actual suspend logic.
 //

clang/test/CIR/IR/await.cir

@@ -0,0 +1,22 @@
+// RUN: cir-opt %s -o %t.cir
+// RUN: FileCheck --input-file=%t.cir %s
+
+cir.func coroutine @checkPrintParse(%arg0 : !cir.bool) {
+  cir.await(user, ready : {
+    cir.condition(%arg0)
+  }, suspend : {
+    cir.yield
+  }, resume : {
+    cir.yield
+  },)
+  cir.return
+}
+
+// CHECK:  cir.func coroutine @checkPrintParse
+// CHECK:  cir.await(user, ready : {
+// CHECK:    cir.condition(%arg0)
+// CHECK:  }, suspend : {
+// CHECK:    cir.yield
+// CHECK:  }, resume : {
+// CHECK:    cir.yield
+// CHECK:  },)

clang/test/CIR/IR/invalid.cir

@@ -494,27 +494,12 @@ cir.func coroutine @bad_task() { // expected-error {{coroutine body must use at
 
 // -----
 
-cir.func coroutine @bad_yield() {
+cir.func coroutine @missing_condition() {
   cir.scope {
-    cir.await(user, ready : {
+    cir.await(user, ready : { // expected-error {{ready region must end with cir.condition}}
       cir.yield
     }, suspend : {
       cir.yield
-    }, resume : {
-      cir.yield nosuspend // expected-error {{kind 'nosuspend' can only be used in 'ready' and 'suspend' regions}}
-    },)
-  }
-  cir.return
-}
-
-// -----
-
-cir.func coroutine @good_yield() {
-  cir.scope {
-    cir.await(user, ready : {
-      cir.yield nosuspend
-    }, suspend : {
-      cir.yield nosuspend
     }, resume : {
       cir.yield
     },)