Merge pull request #10041 from swiftlang/cherry-pick-tmo-retcon-lowering

[SUA][Coroutines] Change LLVM's Coroutine lowering to allow for and use extra hash parameter to llvm.coro.id.retcon.once and llvm.coro.id.retcon (#10015)

Author: bnbarham

llvm/include/llvm/IR/Intrinsics.td

@@ -1674,11 +1674,11 @@ def int_coro_id : DefaultAttrsIntrinsic<[llvm_token_ty],
      NoCapture<ArgIndex<2>>]>;
 def int_coro_id_retcon : Intrinsic<[llvm_token_ty],
     [llvm_i32_ty, llvm_i32_ty, llvm_ptr_ty,
-     llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
+     llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty],
     []>;
 def int_coro_id_retcon_once : Intrinsic<[llvm_token_ty],
     [llvm_i32_ty, llvm_i32_ty, llvm_ptr_ty,
-     llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
+     llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty],
     []>;
 def int_coro_alloc : Intrinsic<[llvm_i1_ty], [llvm_token_ty], []>;
 def int_coro_id_async : Intrinsic<[llvm_token_ty],

llvm/lib/Transforms/Coroutines/CoroInstr.h

@@ -233,7 +233,7 @@ class LLVM_LIBRARY_VISIBILITY CoroIdInst : public AnyCoroIdInst {
 /// This represents either the llvm.coro.id.retcon or
 /// llvm.coro.id.retcon.once instruction.
 class LLVM_LIBRARY_VISIBILITY AnyCoroIdRetconInst : public AnyCoroIdInst {
-  enum { SizeArg, AlignArg, StorageArg, PrototypeArg, AllocArg, DeallocArg };
+  enum { SizeArg, AlignArg, StorageArg, PrototypeArg, AllocArg, DeallocArg, TypeIdArg };
 
 public:
   void checkWellFormed() const;
@@ -267,6 +267,19 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroIdRetconInst : public AnyCoroIdInst {
     return cast<Function>(getArgOperand(DeallocArg)->stripPointerCasts());
   }
 
+  /// Return the TypeId to be used for allocating typed memory
+  ConstantInt *getTypeId() const {
+    if (arg_size() <= TypeIdArg)
+      return nullptr;
+    assert(hasTypeId() && "Invalid number of arguments");
+    return cast<ConstantInt>(getArgOperand(TypeIdArg));
+  }
+
+  /// Return true if TypeId is present in the list of arguments, false otherwise
+  bool hasTypeId() const {
+    return arg_size() == TypeIdArg + 1;
+  }
+
   // Methods to support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const IntrinsicInst *I) {
     auto ID = I->getIntrinsicID();

llvm/lib/Transforms/Coroutines/CoroInternal.h

@@ -130,6 +130,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
     Function *Dealloc;
     BasicBlock *ReturnBlock;
     bool IsFrameInlineInStorage;
+    ConstantInt* TypeId;
   };
 
   struct AsyncLoweringStorage {

llvm/lib/Transforms/Coroutines/Coroutines.cpp

@@ -343,6 +343,7 @@ void coro::Shape::buildFrom(Function &F) {
     this->RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
     this->RetconLowering.ReturnBlock = nullptr;
     this->RetconLowering.IsFrameInlineInStorage = false;
+    this->RetconLowering.TypeId = ContinuationId->getTypeId();
 
     // Determine the result value types, and make sure they match up with
     // the values passed to the suspends.
@@ -465,7 +466,12 @@ Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
     Size = Builder.CreateIntCast(Size,
                                  Alloc->getFunctionType()->getParamType(0),
                                  /*is signed*/ false);
-    auto *Call = Builder.CreateCall(Alloc, Size);
+    ConstantInt* TypeId = RetconLowering.TypeId;
+    CallInst *Call;
+    if (TypeId == nullptr)
+      Call = Builder.CreateCall(Alloc, Size);
+    else
+      Call = Builder.CreateCall(Alloc, {Size, TypeId});
     propagateCallAttrsFromCallee(Call, Alloc);
     addCallToCallGraph(CG, Call, Alloc);
     return Call;
@@ -558,9 +564,14 @@ static void checkWFAlloc(const Instruction *I, Value *V) {
   if (!FT->getReturnType()->isPointerTy())
     fail(I, "llvm.coro.* allocator must return a pointer", F);
 
-  if (FT->getNumParams() != 1 ||
-      !FT->getParamType(0)->isIntegerTy())
-    fail(I, "llvm.coro.* allocator must take integer as only param", F);
+  if (FT->getNumParams() > 2 || FT->getNumParams() == 0)
+    fail(I, "llvm.coro.* allocator must take either one or two params", F);
+
+  if (FT->getNumParams() == 1 && !FT->getParamType(0)->isIntegerTy())
+    fail(I, "llvm.coro.* allocator must take integer as its first param", F);
+
+  if (FT->getNumParams() == 2 && !FT->getParamType(1)->isIntegerTy())
+    fail(I, "llvm.coro.* allocator must take uint64_t as its second param", F);
 }
 
 /// Check that the given value is a well-formed deallocator.

llvm/test/Transforms/Coroutines/coro-retcon-alloca-opaque-ptr.ll

@@ -17,7 +17,7 @@ define {ptr, ptr, i32} @f(ptr %buffer, i32 %n, { i32 } %dummy) {
 ; CHECK-NEXT:    ret { ptr, ptr, i32 } [[TMP2]]
 ;
 entry:
-  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_f, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_f, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   br label %loop
 
@@ -51,7 +51,7 @@ define {ptr, i32} @g(ptr %buffer, i32 %n) {
 ; CHECK-NEXT:    ret { ptr, i32 } [[TMP2]]
 ;
 entry:
-  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_g, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_g, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   br label %loop
 
@@ -73,7 +73,7 @@ cleanup:
   unreachable
 }
 
-declare token @llvm.coro.id.retcon(i32, i32, ptr, ptr, ptr, ptr)
+declare token @llvm.coro.id.retcon(i32, i32, ptr, ptr, ptr, ptr, ...)
 declare ptr @llvm.coro.begin(token, ptr)
 declare i1 @llvm.coro.suspend.retcon.i1(...)
 declare void @llvm.coro.suspend.retcon.isVoid(...)

llvm/test/Transforms/Coroutines/coro-retcon-alloca.ll

@@ -16,7 +16,7 @@ define {ptr, ptr, i32} @f(ptr %buffer, i32 %n) {
 ; CHECK-NEXT:    ret { ptr, ptr, i32 } [[TMP2]]
 ;
 entry:
-  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_f, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_f, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   br label %loop
 
@@ -51,7 +51,7 @@ define {ptr, i32} @g(ptr %buffer, i32 %n) {
 ; CHECK-NEXT:    ret { ptr, i32 } [[TMP2]]
 ;
 entry:
-  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_g, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_g, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   br label %loop
 
@@ -84,7 +84,7 @@ define {ptr, i32} @h(ptr %buffer, i32 %n) {
 ; CHECK-NEXT:    ret { ptr, i32 } [[TMP0]]
 ;
 entry:
-  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_h, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_h, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   br label %loop
 
@@ -117,7 +117,7 @@ define {ptr, i32} @i(ptr %buffer, i32 %n) {
 ; CHECK-NEXT:    ret { ptr, i32 } [[TMP0]]
 ;
 entry:
-  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_i, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_i, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   br label %loop
 
@@ -149,7 +149,7 @@ define {ptr, i32} @j(ptr %buffer, i32 %n) {
 ; CHECK-NEXT:    ret { ptr, i32 } [[TMP0]]
 ;
 entry:
-  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_j, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_j, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   br label %forward
 
@@ -190,7 +190,7 @@ define {ptr, i32} @k(ptr %buffer, i32 %n, i1 %cond) {
 ; CHECK-NEXT:    br label [[CORO_RETURN]]
 ;
 entry:
-  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_g, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon(i32 1024, i32 8, ptr %buffer, ptr @prototype_g, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   br i1 %cond, label %alloca_block, label %non_alloca_block
 
@@ -221,7 +221,7 @@ cleanup:
   unreachable
 }
 
-declare token @llvm.coro.id.retcon(i32, i32, ptr, ptr, ptr, ptr)
+declare token @llvm.coro.id.retcon(i32, i32, ptr, ptr, ptr, ptr, ...)
 declare ptr @llvm.coro.begin(token, ptr)
 declare i1 @llvm.coro.suspend.retcon.i1(...)
 declare void @llvm.coro.suspend.retcon.isVoid(...)

llvm/test/Transforms/Coroutines/coro-retcon-frame.ll

@@ -18,7 +18,7 @@ entry:
   store i64 0, ptr %proj.1, align 8
   store i64 0, ptr %proj.2, align 8
   %escape_addr = ptrtoint ptr %tmp to i64
-  %id = call token @llvm.coro.id.retcon.once(i32 32, i32 8, ptr %buffer, ptr @prototype_f, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon.once(i32 32, i32 8, ptr %buffer, ptr @prototype_f, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   %proj.2.2 = getelementptr inbounds { i64, i64 }, ptr %tmp, i64 0, i32 1
   call void @init(ptr %proj.1)
@@ -49,7 +49,7 @@ end:
 ; CHECK: resume:
 ; CHECK:  call void @use(ptr %0)
 
-declare token @llvm.coro.id.retcon.once(i32, i32, ptr, ptr, ptr, ptr)
+declare token @llvm.coro.id.retcon.once(i32, i32, ptr, ptr, ptr, ptr, ...)
 declare ptr @llvm.coro.begin(token, ptr)
 declare i1 @llvm.coro.suspend.retcon.i1(...)
 declare i1 @llvm.coro.end(ptr, i1, token)

llvm/test/Transforms/Coroutines/coro-retcon-once-private.ll

@@ -7,7 +7,7 @@ target triple = "x86_64-apple-macosx10.12.0"
 
 define internal {ptr, i32} @f(ptr %buffer, ptr %array) {
 entry:
-  %id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   %load = load i32, ptr %array
   %load.pos = icmp sgt i32 %load, 0
@@ -34,7 +34,7 @@ cleanup:
   unreachable
 }
 
-declare token @llvm.coro.id.retcon.once(i32, i32, ptr, ptr, ptr, ptr)
+declare token @llvm.coro.id.retcon.once(i32, i32, ptr, ptr, ptr, ptr, ...)
 declare ptr @llvm.coro.begin(token, ptr)
 declare i1 @llvm.coro.suspend.retcon.i1(...)
 declare i1 @llvm.coro.end(ptr, i1, token)

llvm/test/Transforms/Coroutines/coro-retcon-once-value.ll

@@ -6,7 +6,7 @@ target triple = "x86_64-apple-macosx10.12.0"
 
 define {ptr, i32} @f(ptr %buffer, ptr %array) {
 entry:
-  %id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   %load = load i32, ptr %array
   %load.pos = icmp sgt i32 %load, 0
@@ -47,7 +47,7 @@ entry:
 
 define {ptr, i32} @g(ptr %buffer, ptr %array, i32 %val) {
 entry:
-  %id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype2, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype2, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   %load = load i32, ptr %array
   %load.pos = icmp sgt i32 %load, 0
@@ -93,7 +93,7 @@ entry:
 ;   Unfortunately, we don't seem to fully optimize this right now due
 ;   to some sort of phase-ordering thing.
 
-declare token @llvm.coro.id.retcon.once(i32, i32, ptr, ptr, ptr, ptr)
+declare token @llvm.coro.id.retcon.once(i32, i32, ptr, ptr, ptr, ptr, ...)
 declare ptr @llvm.coro.begin(token, ptr)
 declare i1 @llvm.coro.suspend.retcon.i1(...)
 declare i1 @llvm.coro.end(ptr, i1, token)

llvm/test/Transforms/Coroutines/coro-retcon-once-value2.ll

@@ -6,7 +6,7 @@ target triple = "x86_64-apple-macosx10.12.0"
 define {ptr, ptr} @f(ptr %buffer, ptr %ptr) presplitcoroutine {
 entry:
   %temp = alloca i32, align 4
-  %id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   %oldvalue = load i32, ptr %ptr
   store i32 %oldvalue, ptr %temp
@@ -27,7 +27,7 @@ cleanup:
 define {ptr, ptr} @g(ptr %buffer, ptr %ptr, i8 %val) presplitcoroutine {
 entry:
   %temp = alloca i32, align 4
-  %id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype2, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype2, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   %oldvalue = load i32, ptr %ptr
   store i32 %oldvalue, ptr %temp
@@ -49,7 +49,7 @@ cleanup:
 define {ptr, ptr} @h(ptr %buffer, ptr %ptr) presplitcoroutine {
 entry:
   %temp = alloca i32, align 4
-  %id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype3, ptr @allocate, ptr @deallocate)
+  %id = call token (i32, i32, ptr, ptr, ptr, ptr, ...) @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer, ptr @prototype3, ptr @allocate, ptr @deallocate)
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   %oldvalue = load i32, ptr %ptr
   store i32 %oldvalue, ptr %temp
@@ -68,7 +68,7 @@ cleanup:
 }
 
 
-declare token @llvm.coro.id.retcon.once(i32, i32, ptr, ptr, ptr, ptr)
+declare token @llvm.coro.id.retcon.once(i32, i32, ptr, ptr, ptr, ptr, ...)
 declare ptr @llvm.coro.begin(token, ptr)
 declare i1 @llvm.coro.suspend.retcon.i1(...)
 declare i1 @llvm.coro.end(ptr, i1, token)