[Clang] Introduce [[clang::structured_concurrency]]

Author: yuxuanchen1997

clang/include/clang/Basic/Attr.td

@@ -1212,6 +1212,14 @@ def CoroDisableLifetimeBound : InheritableAttr {
   let SimpleHandler = 1;
 }
 
+def CoroStructuredConcurrencyType : InheritableAttr {
+  let Spellings = [Clang<"coro_structured_concurrency">];
+  let Subjects = SubjectList<[CXXRecord]>;
+  let LangOpts = [CPlusPlus];
+  let Documentation = [CoroStructuredConcurrencyDoc];
+  let SimpleHandler = 1;
+}
+
 // OSObject-based attributes.
 def OSConsumed : InheritableParamAttr {
   let Spellings = [Clang<"os_consumed">];

clang/include/clang/Basic/AttrDocs.td

@@ -8015,6 +8015,26 @@ but do not pass them to the underlying coroutine or pass them by value.
 }];
 }
 
+def CoroStructuredConcurrencyDoc : Documentation {
+  let Category = DocCatDecl;
+  let Content = [{
+The ``[[clang::coro_structured_concurrency]]`` is a class attribute which can be applied
+to a coroutine return type.
+
+When a coroutine function that returns such a type calls another coroutine function,
+the compiler performs heap allocation elision when the following conditions are all met:
+- callee coroutine function returns a type that is annotated with
+  ``[[clang::coro_structured_concurrency]]``.
+- The callee coroutine function is inlined.
+- In caller coroutine, the return value of the callee is a prvalue or an xvalue, and
+- The temporary expression containing the callee coroutine object is immediately co_awaited.
+
+The behavior is undefined if any of the following condition was met:
+- the caller coroutine is destroyed earlier than the callee coroutine.
+
+  }];
+}
+
 def CountedByDocs : Documentation {
   let Category = DocCatField;
   let Content = [{

clang/lib/CodeGen/CGCoroutine.cpp

@@ -12,9 +12,12 @@
 
 #include "CGCleanup.h"
 #include "CodeGenFunction.h"
-#include "llvm/ADT/ScopeExit.h"
+#include "clang/AST/ExprCXX.h"
 #include "clang/AST/StmtCXX.h"
 #include "clang/AST/StmtVisitor.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/IR/Intrinsics.h"
 
 using namespace clang;
 using namespace CodeGen;
@@ -219,17 +222,81 @@ namespace {
     RValue RV;
   };
 }
+
+static MaterializeTemporaryExpr *
+getStructuredConcurrencyOperand(ASTContext &Ctx,
+                                CoroutineSuspendExpr const &S) {
+  auto *E = S.getCommonExpr();
+  auto *Temporary = dyn_cast_or_null<MaterializeTemporaryExpr>(E);
+  if (!Temporary)
+    return nullptr;
+
+  auto *Operator =
+      dyn_cast_or_null<CXXOperatorCallExpr>(Temporary->getSubExpr());
+
+  if (!Operator ||
+      Operator->getOperator() != OverloadedOperatorKind::OO_Coawait ||
+      Operator->getNumArgs() != 1)
+    return nullptr;
+
+  Expr *Arg = Operator->getArg(0);
+  assert(Arg && "Arg to operator co_await should not be null");
+  auto *CalleeRetClass = Arg->getType()->getAsCXXRecordDecl();
+
+  if (!CalleeRetClass ||
+      !CalleeRetClass->hasAttr<CoroStructuredConcurrencyTypeAttr>())
+    return nullptr;
+
+  if (!Arg->isTemporaryObject(Ctx, CalleeRetClass)) {
+    return nullptr;
+  }
+
+  return dyn_cast<MaterializeTemporaryExpr>(Arg);
+}
+
 static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Coro,
                                     CoroutineSuspendExpr const &S,
                                     AwaitKind Kind, AggValueSlot aggSlot,
                                     bool ignoreResult, bool forLValue) {
   auto *E = S.getCommonExpr();
 
+  auto &Builder = CGF.Builder;
+  bool MarkOperandSafeIntrinsic = false;
+
+  auto *TemporaryOperand = [&]() -> MaterializeTemporaryExpr * {
+    bool CurFnRetTyHasAttr = false;
+    if (auto *RetTyPtr = CGF.FnRetTy.getTypePtrOrNull()) {
+      if (auto *CxxRecord = RetTyPtr->getAsCXXRecordDecl()) {
+        CurFnRetTyHasAttr =
+            CxxRecord->hasAttr<CoroStructuredConcurrencyTypeAttr>();
+      }
+    }
+
+    if (CurFnRetTyHasAttr) {
+      return getStructuredConcurrencyOperand(CGF.getContext(), S);
+    }
+    return nullptr;
+  }();
+
+  if (TemporaryOperand) {
+    CGF.TemporaryValues[TemporaryOperand] = nullptr;
+    MarkOperandSafeIntrinsic = true;
+  }
+
   auto CommonBinder =
       CodeGenFunction::OpaqueValueMappingData::bind(CGF, S.getOpaqueValue(), E);
   auto UnbindCommonOnExit =
       llvm::make_scope_exit([&] { CommonBinder.unbind(CGF); });
 
+  if (MarkOperandSafeIntrinsic) {
+    auto It = CGF.TemporaryValues.find(TemporaryOperand);
+    assert(It != CGF.TemporaryValues.end());
+    if (auto Value = It->second)
+      Builder.CreateCall(CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_safe_elide),
+                         {Value});
+    CGF.TemporaryValues.erase(TemporaryOperand);
+  }
+
   auto Prefix = buildSuspendPrefixStr(Coro, Kind);
   BasicBlock *ReadyBlock = CGF.createBasicBlock(Prefix + Twine(".ready"));
   BasicBlock *SuspendBlock = CGF.createBasicBlock(Prefix + Twine(".suspend"));
@@ -241,7 +308,6 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   // Otherwise, emit suspend logic.
   CGF.EmitBlock(SuspendBlock);
 
-  auto &Builder = CGF.Builder;
   llvm::Function *CoroSave = CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_save);
   auto *NullPtr = llvm::ConstantPointerNull::get(CGF.CGM.Int8PtrTy);
   auto *SaveCall = Builder.CreateCall(CoroSave, {NullPtr});
@@ -255,9 +321,9 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
          "expected to be called in coroutine context");
 
   SmallVector<llvm::Value *, 3> SuspendIntrinsicCallArgs;
-  SuspendIntrinsicCallArgs.push_back(
-      CGF.getOrCreateOpaqueLValueMapping(S.getOpaqueValue()).getPointer(CGF));
-
+  auto *BoundAwaiterValue =
+      CGF.getOrCreateOpaqueLValueMapping(S.getOpaqueValue()).getPointer(CGF);
+  SuspendIntrinsicCallArgs.push_back(BoundAwaiterValue);
   SuspendIntrinsicCallArgs.push_back(CGF.CurCoro.Data->CoroBegin);
   SuspendIntrinsicCallArgs.push_back(SuspendWrapper);
 

clang/lib/CodeGen/CGExpr.cpp

@@ -618,7 +618,11 @@ EmitMaterializeTemporaryExpr(const MaterializeTemporaryExpr *M) {
     }
   }
 
-  return MakeAddrLValue(Object, M->getType(), AlignmentSource::Decl);
+  auto Ret = MakeAddrLValue(Object, M->getType(), AlignmentSource::Decl);
+  if (TemporaryValues.contains(M)) {
+    TemporaryValues[M] = Ret.getPointer(*this);
+  }
+  return Ret;
 }
 
 RValue

clang/lib/CodeGen/CodeGenFunction.h

@@ -369,6 +369,9 @@ class CodeGenFunction : public CodeGenTypeCache {
   };
   CGCoroInfo CurCoro;
 
+  llvm::SmallDenseMap<const MaterializeTemporaryExpr *, llvm::Value *>
+      TemporaryValues;
+
   bool isCoroutine() const {
     return CurCoro.Data != nullptr;
   }

clang/test/CodeGenCoroutines/Inputs/utility.h

@@ -0,0 +1,13 @@
+// This is a mock file for <utility>
+
+namespace std {
+
+template <typename T> struct remove_reference { using type = T; };
+template <typename T> struct remove_reference<T &> { using type = T; };
+template <typename T> struct remove_reference<T &&> { using type = T; };
+
+template <typename T>
+constexpr typename std::remove_reference<T>::type&& move(T &&t) noexcept {
+  return static_cast<typename std::remove_reference<T>::type &&>(t);
+}
+}

clang/test/CodeGenCoroutines/coro-structured-concurrency.cpp

@@ -0,0 +1,84 @@
+// This file tests the coro_structured_concurrency attribute semantics. 
+// RUN: %clang_cc1 -std=c++20 -disable-llvm-passes -emit-llvm %s -o - | FileCheck %s
+
+#include "Inputs/coroutine.h"
+#include "Inputs/utility.h"
+
+template <typename T>
+struct [[clang::coro_structured_concurrency]] Task {
+  struct promise_type {
+    struct FinalAwaiter {
+      bool await_ready() const noexcept { return false; }
+
+      template <typename P>
+      std::coroutine_handle<> await_suspend(std::coroutine_handle<P> coro) noexcept {
+        if (!coro)
+          return std::noop_coroutine();
+        return coro.promise().continuation;
+      }
+      void await_resume() noexcept {}
+    };
+
+    Task get_return_object() noexcept {
+      return std::coroutine_handle<promise_type>::from_promise(*this);
+    }
+
+    std::suspend_always initial_suspend() noexcept { return {}; }
+    FinalAwaiter final_suspend() noexcept { return {}; }
+    void unhandled_exception() noexcept {}
+    void return_value(T x) noexcept {
+      value = x;
+    }
+
+    std::coroutine_handle<> continuation;
+    T value;
+  };
+
+  Task(std::coroutine_handle<promise_type> handle) : handle(handle) {}
+  ~Task() {
+    if (handle)
+      handle.destroy();
+  }
+
+  struct Awaiter {
+    Awaiter(Task *t) : task(t) {}
+    bool await_ready() const noexcept { return false; }
+    void await_suspend(std::coroutine_handle<void> continuation) noexcept {}
+    T await_resume() noexcept {
+      return task->handle.promise().value;
+    }
+
+    Task *task;
+  };
+
+  auto operator co_await() {
+    return Awaiter{this};
+  }
+
+private:
+  std::coroutine_handle<promise_type> handle;
+};
+
+// CHECK-LABEL: define{{.*}} @_Z6calleev 
+Task<int> callee() {
+  co_return 1;
+}
+
+// CHECK-LABEL: define{{.*}} @_Z8elidablev 
+Task<int> elidable() {
+  // CHECK: %[[TARK_OBJ:.+]] = alloca %struct.Task
+  // CHECK: call void @llvm.coro.safe.elide(ptr %[[TARK_OBJ:.+]])
+  co_return co_await callee();
+}
+
+// CHECK-LABEL: define{{.*}} @_Z11nonelidablev 
+Task<int> nonelidable() {
+  // CHECK: %[[TARK_OBJ:.+]] = alloca %struct.Task
+  auto t = callee();
+  // Because we aren't co_awaiting a prvalue, we cannot elide here.
+  // CHECK-NOT: call void @llvm.coro.safe.elide(ptr %[[TARK_OBJ:.+]])
+  co_await t;
+  co_await std::move(t);
+  
+  co_return 1;
+}

clang/test/Misc/pragma-attribute-supported-attributes-list.test

@@ -62,6 +62,7 @@
 // CHECK-NEXT: CoroLifetimeBound (SubjectMatchRule_record)
 // CHECK-NEXT: CoroOnlyDestroyWhenComplete (SubjectMatchRule_record)
 // CHECK-NEXT: CoroReturnType (SubjectMatchRule_record)
+// CHECK-NEXT: CoroStructuredConcurrencyType (SubjectMatchRule_record)
 // CHECK-NEXT: CoroWrapper (SubjectMatchRule_function)
 // CHECK-NEXT: DLLExport (SubjectMatchRule_function, SubjectMatchRule_variable, SubjectMatchRule_record, SubjectMatchRule_objc_interface)
 // CHECK-NEXT: DLLImport (SubjectMatchRule_function, SubjectMatchRule_variable, SubjectMatchRule_record, SubjectMatchRule_objc_interface)

llvm/include/llvm/IR/Intrinsics.td

@@ -1733,6 +1733,9 @@ def int_coro_subfn_addr : DefaultAttrsIntrinsic<
     [IntrReadMem, IntrArgMemOnly, ReadOnly<ArgIndex<0>>,
      NoCapture<ArgIndex<0>>]>;
 
+def int_coro_safe_elide : DefaultAttrsIntrinsic<
+    [], [llvm_ptr_ty], []>;
+
 ///===-------------------------- Other Intrinsics --------------------------===//
 //
 // TODO: We should introduce a new memory kind fo traps (and other side effects

llvm/lib/Transforms/Coroutines/CoroCleanup.cpp

@@ -8,10 +8,11 @@
 
 #include "llvm/Transforms/Coroutines/CoroCleanup.h"
 #include "CoroInternal.h"
+#include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/PassManager.h"
-#include "llvm/IR/Function.h"
 #include "llvm/Transforms/Scalar/SimplifyCFG.h"
 
 using namespace llvm;
@@ -80,7 +81,7 @@ bool Lowerer::lower(Function &F) {
         } else
           continue;
         break;
-      case Intrinsic::coro_async_size_replace:
+      case Intrinsic::coro_async_size_replace: {
         auto *Target = cast<ConstantStruct>(
             cast<GlobalVariable>(II->getArgOperand(0)->stripPointerCasts())
                 ->getInitializer());
@@ -98,6 +99,9 @@ bool Lowerer::lower(Function &F) {
         Target->replaceAllUsesWith(NewFuncPtrStruct);
         break;
       }
+      case Intrinsic::coro_safe_elide:
+        break;
+      }
       II->eraseFromParent();
       Changed = true;
     }
@@ -111,7 +115,8 @@ static bool declaresCoroCleanupIntrinsics(const Module &M) {
       M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr",
           "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon",
           "llvm.coro.id.async", "llvm.coro.id.retcon.once",
-          "llvm.coro.async.size.replace", "llvm.coro.async.resume"});
+          "llvm.coro.async.size.replace", "llvm.coro.async.resume",
+          "llvm.coro.safe.elide"});
 }
 
 PreservedAnalyses CoroCleanupPass::run(Module &M,