[CIR] Emit CIR builtins: coroAlloc, coroBegin, and coroSize

[Andres-Salamanca]

This PR adds support for emitting the builtins coroAlloc, coroBegin, and coroSize.

[llvmbot]

@llvm/pr-subscribers-clangir

@llvm/pr-subscribers-clang

Author: None (Andres-Salamanca)

Changes

This PR adds support for emitting the builtins coroAlloc, coroBegin, and coroSize.

---

Full diff: https://github.com/llvm/llvm-project/pull/164180.diff

6 Files Affected:

• (modified) clang/include/clang/CIR/MissingFeatures.h (+1-3)
• (modified) clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp (+10-4)
• (modified) clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp (+78-3)
• (modified) clang/lib/CIR/CodeGen/CIRGenFunction.h (+3)
• (modified) clang/lib/CIR/CodeGen/CIRGenModule.h (+2)
• (modified) clang/test/CIR/CodeGen/coro-task.cpp (+18-1)

diff --git a/clang/include/clang/CIR/MissingFeatures.h b/clang/include/clang/CIR/MissingFeatures.h
index 090cf35c2d279..7956ba797f837 100644
--- a/clang/include/clang/CIR/MissingFeatures.h
+++ b/clang/include/clang/CIR/MissingFeatures.h
@@ -149,11 +149,9 @@ struct MissingFeatures {
   static bool zeroSizeRecordMembers() { return false; }
 
   // Coroutines
-  static bool coroAllocBuiltinCall() { return false; }
-  static bool coroBeginBuiltinCall() { return false; }
   static bool coroEndBuiltinCall() { return false; }
-  static bool coroSizeBuiltinCall() { return false; }
   static bool coroutineFrame() { return false; }
+  static bool emitBodyAndFallthrough() { return false; }
 
   // Various handling of deferred processing in CIRGenModule.
   static bool cgmRelease() { return false; }
diff --git a/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp b/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
index ea31871806bd7..92ede62cac630 100644
--- a/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
@@ -449,10 +449,16 @@ RValue CIRGenFunction::emitBuiltinExpr(const GlobalDecl &gd, unsigned builtinID,
   }
   case Builtin::BI__builtin_coro_free:
   case Builtin::BI__builtin_coro_size: {
-    cgm.errorNYI(e->getSourceRange(),
-                 "BI__builtin_coro_free, BI__builtin_coro_size NYI");
-    assert(!cir::MissingFeatures::coroSizeBuiltinCall());
-    return getUndefRValue(e->getType());
+    GlobalDecl gd{fd};
+    mlir::Type ty = cgm.getTypes().getFunctionType(
+        cgm.getTypes().arrangeGlobalDeclaration(gd));
+    const auto *nd = cast<NamedDecl>(gd.getDecl());
+    cir::FuncOp fnOp =
+        cgm.getOrCreateCIRFunction(nd->getName(), ty, gd, /*ForVTable=*/false,
+                                   /*DontDefer=*/false);
+    fnOp.setBuiltin(true);
+    return emitCall(e->getCallee()->getType(), CIRGenCallee::forDirect(fnOp), e,
+                    returnValue);
   }
   }
 
diff --git a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp
index c25cce4ab33b3..86b4e43ea2998 100644
--- a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp
@@ -15,6 +15,7 @@
 #include "clang/AST/StmtCXX.h"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
+#include "clang/CIR/MissingFeatures.h"
 
 using namespace clang;
 using namespace clang::CIRGen;
@@ -23,6 +24,9 @@ struct clang::CIRGen::CGCoroData {
   // Stores the __builtin_coro_id emitted in the function so that we can supply
   // it as the first argument to other builtins.
   cir::CallOp coroId = nullptr;
+
+  // Stores the result of __builtin_coro_begin call.
+  mlir::Value coroBegin = nullptr;
 };
 
 // Defining these here allows to keep CGCoroData private to this file.
@@ -63,6 +67,48 @@ cir::CallOp CIRGenFunction::emitCoroIDBuiltinCall(mlir::Location loc,
                                                nullPtr, nullPtr, nullPtr});
 }
 
+cir::CallOp CIRGenFunction::emitCoroAllocBuiltinCall(mlir::Location loc) {
+  cir::BoolType boolTy = builder.getBoolTy();
+  cir::IntType int32Ty = builder.getUInt32Ty();
+
+  mlir::Operation *builtin = cgm.getGlobalValue(cgm.builtinCoroAlloc);
+
+  cir::FuncOp fnOp;
+  if (!builtin) {
+    fnOp = cgm.createCIRBuiltinFunction(loc, cgm.builtinCoroAlloc,
+                                        cir::FuncType::get({int32Ty}, boolTy),
+                                        /*FD=*/nullptr);
+    assert(fnOp && "should always succeed");
+  } else {
+    fnOp = cast<cir::FuncOp>(builtin);
+  }
+
+  return builder.createCallOp(
+      loc, fnOp, mlir::ValueRange{curCoro.data->coroId.getResult()});
+}
+
+cir::CallOp
+CIRGenFunction::emitCoroBeginBuiltinCall(mlir::Location loc,
+                                         mlir::Value coroframeAddr) {
+  cir::IntType int32Ty = builder.getUInt32Ty();
+  mlir::Operation *builtin = cgm.getGlobalValue(cgm.builtinCoroBegin);
+
+  cir::FuncOp fnOp;
+  if (!builtin) {
+    fnOp = cgm.createCIRBuiltinFunction(
+        loc, cgm.builtinCoroBegin,
+        cir::FuncType::get({int32Ty, VoidPtrTy}, VoidPtrTy),
+        /*FD=*/nullptr);
+    assert(fnOp && "should always succeed");
+  } else {
+    fnOp = cast<cir::FuncOp>(builtin);
+  }
+
+  return builder.createCallOp(
+      loc, fnOp,
+      mlir::ValueRange{curCoro.data->coroId.getResult(), coroframeAddr});
+}
+
 mlir::LogicalResult
 CIRGenFunction::emitCoroutineBody(const CoroutineBodyStmt &s) {
   mlir::Location openCurlyLoc = getLoc(s.getBeginLoc());
@@ -73,10 +119,39 @@ CIRGenFunction::emitCoroutineBody(const CoroutineBodyStmt &s) {
   cir::CallOp coroId = emitCoroIDBuiltinCall(openCurlyLoc, nullPtrCst);
   createCoroData(*this, curCoro, coroId);
 
-  assert(!cir::MissingFeatures::coroAllocBuiltinCall());
-
-  assert(!cir::MissingFeatures::coroBeginBuiltinCall());
+  // Backend is allowed to elide memory allocations, to help it, emit
+  // auto mem = coro.alloc() ? 0 : ... allocation code ...;
+  cir::CallOp coroAlloc = emitCoroAllocBuiltinCall(openCurlyLoc);
+
+  // Initialize address of coroutine frame to null
+  CanQualType astVoidPtrTy = cgm.getASTContext().VoidPtrTy;
+  mlir::Type allocaTy = convertTypeForMem(astVoidPtrTy);
+  Address coroFrame =
+      createTempAlloca(allocaTy, getContext().getTypeAlignInChars(astVoidPtrTy),
+                       openCurlyLoc, "__coro_frame_addr",
+                       /*ArraySize=*/nullptr);
+
+  mlir::Value storeAddr = coroFrame.getPointer();
+  builder.CIRBaseBuilderTy::createStore(openCurlyLoc, nullPtrCst, storeAddr);
+  cir::IfOp::create(
+      builder, openCurlyLoc, coroAlloc.getResult(),
+      /*withElseRegion=*/false,
+      /*thenBuilder=*/[&](mlir::OpBuilder &b, mlir::Location loc) {
+        builder.CIRBaseBuilderTy::createStore(
+            loc, emitScalarExpr(s.getAllocate()), storeAddr);
+        builder.create<cir::YieldOp>(loc);
+      });
+  curCoro.data->coroBegin =
+      emitCoroBeginBuiltinCall(
+          openCurlyLoc,
+          builder.create<cir::LoadOp>(openCurlyLoc, allocaTy, storeAddr))
+          .getResult();
+
+  // Handle allocation failure if 'ReturnStmtOnAllocFailure' was provided.
+  if (s.getReturnStmtOnAllocFailure())
+    cgm.errorNYI("NYI");
 
   assert(!cir::MissingFeatures::generateDebugInfo());
+  assert(!cir::MissingFeatures::emitBodyAndFallthrough());
   return mlir::success();
 }
diff --git a/clang/lib/CIR/CodeGen/CIRGenFunction.h b/clang/lib/CIR/CodeGen/CIRGenFunction.h
index 3c36f5c697118..f3ef54892bcae 100644
--- a/clang/lib/CIR/CodeGen/CIRGenFunction.h
+++ b/clang/lib/CIR/CodeGen/CIRGenFunction.h
@@ -1252,6 +1252,9 @@ class CIRGenFunction : public CIRGenTypeCache {
   mlir::LogicalResult emitCoroutineBody(const CoroutineBodyStmt &s);
   cir::CallOp emitCoroEndBuiltinCall(mlir::Location loc, mlir::Value nullPtr);
   cir::CallOp emitCoroIDBuiltinCall(mlir::Location loc, mlir::Value nullPtr);
+  cir::CallOp emitCoroAllocBuiltinCall(mlir::Location loc);
+  cir::CallOp emitCoroBeginBuiltinCall(mlir::Location loc,
+                                       mlir::Value coroframeAddr);
 
   void emitDestroy(Address addr, QualType type, Destroyer *destroyer);
 
diff --git a/clang/lib/CIR/CodeGen/CIRGenModule.h b/clang/lib/CIR/CodeGen/CIRGenModule.h
index 1fc116d98a858..186913d1bac9d 100644
--- a/clang/lib/CIR/CodeGen/CIRGenModule.h
+++ b/clang/lib/CIR/CodeGen/CIRGenModule.h
@@ -496,6 +496,8 @@ class CIRGenModule : public CIRGenTypeCache {
                                     bool assumeConvergent = false);
 
   static constexpr const char *builtinCoroId = "__builtin_coro_id";
+  static constexpr const char *builtinCoroAlloc = "__builtin_coro_alloc";
+  static constexpr const char *builtinCoroBegin = "__builtin_coro_begin";
 
   /// Given a builtin id for a function like "__builtin_fabsf", return a
   /// Function* for "fabsf".
diff --git a/clang/test/CIR/CodeGen/coro-task.cpp b/clang/test/CIR/CodeGen/coro-task.cpp
index 1fc7d77be2bce..265325f82d7f7 100644
--- a/clang/test/CIR/CodeGen/coro-task.cpp
+++ b/clang/test/CIR/CodeGen/coro-task.cpp
@@ -106,6 +106,9 @@ co_invoke_fn co_invoke;
 // CIR-NEXT: cir.global external @_ZN5folly4coro9co_invokeE = #cir.zero : !rec_folly3A3Acoro3A3Aco_invoke_fn
 
 // CIR: cir.func builtin private @__builtin_coro_id(!u32i, !cir.ptr<!void>, !cir.ptr<!void>, !cir.ptr<!void>) -> !u32i
+// CIR:  cir.func builtin private @__builtin_coro_alloc(!u32i) -> !cir.bool
+// CIR:  cir.func builtin private @__builtin_coro_size() -> !u64i
+// CIR:  cir.func builtin private @__builtin_coro_begin(!u32i, !cir.ptr<!void>) -> !cir.ptr<!void>
 
 using VoidTask = folly::coro::Task<void>;
 
@@ -114,10 +117,24 @@ VoidTask silly_task() {
 }
 
 // CIR: cir.func coroutine dso_local @_Z10silly_taskv() -> ![[VoidTask]]
-// CHECK: %[[#VoidTaskAddr:]] = cir.alloca ![[VoidTask]], {{.*}}, ["__retval"]
+// CIR: %[[VoidTaskAddr:.*]] = cir.alloca ![[VoidTask]], {{.*}}, ["__retval"]
+// CIR: %[[SavedFrameAddr:.*]] = cir.alloca !cir.ptr<!void>, !cir.ptr<!cir.ptr<!void>>, ["__coro_frame_addr"]
 
 // Get coroutine id with __builtin_coro_id.
 
 // CIR: %[[NullPtr:.*]] = cir.const #cir.ptr<null> : !cir.ptr<!void>
 // CIR: %[[Align:.*]] = cir.const #cir.int<16> : !u32i
 // CIR: %[[CoroId:.*]] = cir.call @__builtin_coro_id(%[[Align]], %[[NullPtr]], %[[NullPtr]], %[[NullPtr]])
+
+// Perform allocation calling operator 'new' depending on __builtin_coro_alloc and
+// call __builtin_coro_begin for the final coroutine frame address.
+
+// CIR: %[[ShouldAlloc:.*]] = cir.call @__builtin_coro_alloc(%[[CoroId]]) : (!u32i) -> !cir.bool
+// CIR: cir.store{{.*}} %[[NullPtr]], %[[SavedFrameAddr]] : !cir.ptr<!void>, !cir.ptr<!cir.ptr<!void>>
+// CIR: cir.if %[[ShouldAlloc]] {
+// CIR:   %[[CoroSize:.*]] = cir.call @__builtin_coro_size() : () -> !u64i
+// CIR:   %[[AllocAddr:.*]] = cir.call @_Znwm(%[[CoroSize]]) : (!u64i) -> !cir.ptr<!void>
+// CIR:   cir.store{{.*}} %[[AllocAddr]], %[[SavedFrameAddr]] : !cir.ptr<!void>, !cir.ptr<!cir.ptr<!void>>
+// CIR: }
+// CIR: %[[Load0:.*]] = cir.load{{.*}} %[[SavedFrameAddr]] : !cir.ptr<!cir.ptr<!void>>, !cir.ptr<!void>
+// CIR: %[[CoroFrameAddr:.*]] = cir.call @__builtin_coro_begin(%[[CoroId]], %[[Load0]])