[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)

[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

Change the signatures of llvm.coro.id.retcon.once and llvm.coro.id.retcon to pass a hash parameter as a variable argument.
This change will allow IRGen to pass the new runtime function swift_coroFrameAlloc (which calls the typed allocator if TMO is enabled)
along with a type hash to these two intrinsics.

rdar://141236876
(cherry picked from commit fe25ae6)

Author: t-rasmud

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/include/llvm/Transforms/Coroutines/CoroShape.h

@@ -0,0 +1,282 @@
+//===- CoroShape.h - Coroutine info for lowering --------------*- C++ -*---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// This file declares the shape info struct that is required by many coroutine
+// utility methods.
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_COROUTINES_COROSHAPE_H
+#define LLVM_TRANSFORMS_COROUTINES_COROSHAPE_H
+
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Transforms/Coroutines/CoroInstr.h"
+
+namespace llvm {
+
+class CallGraph;
+
+namespace coro {
+
+enum class ABI {
+  /// The "resume-switch" lowering, where there are separate resume and
+  /// destroy functions that are shared between all suspend points.  The
+  /// coroutine frame implicitly stores the resume and destroy functions,
+  /// the current index, and any promise value.
+  Switch,
+
+  /// The "returned-continuation" lowering, where each suspend point creates a
+  /// single continuation function that is used for both resuming and
+  /// destroying.  Does not support promises.
+  Retcon,
+
+  /// The "unique returned-continuation" lowering, where each suspend point
+  /// creates a single continuation function that is used for both resuming
+  /// and destroying.  Does not support promises.  The function is known to
+  /// suspend at most once during its execution, and the return value of
+  /// the continuation is void.
+  RetconOnce,
+
+  /// The "async continuation" lowering, where each suspend point creates a
+  /// single continuation function. The continuation function is available as an
+  /// intrinsic.
+  Async,
+};
+
+// Holds structural Coroutine Intrinsics for a particular function and other
+// values used during CoroSplit pass.
+struct Shape {
+  CoroBeginInst *CoroBegin = nullptr;
+  SmallVector<AnyCoroEndInst *, 4> CoroEnds;
+  SmallVector<CoroSizeInst *, 2> CoroSizes;
+  SmallVector<CoroAlignInst *, 2> CoroAligns;
+  SmallVector<AnyCoroSuspendInst *, 4> CoroSuspends;
+  SmallVector<CoroAwaitSuspendInst *, 4> CoroAwaitSuspends;
+  SmallVector<CallInst *, 2> SymmetricTransfers;
+
+  // Values invalidated by replaceSwiftErrorOps()
+  SmallVector<CallInst *, 2> SwiftErrorOps;
+
+  void clear() {
+    CoroBegin = nullptr;
+    CoroEnds.clear();
+    CoroSizes.clear();
+    CoroAligns.clear();
+    CoroSuspends.clear();
+    CoroAwaitSuspends.clear();
+    SymmetricTransfers.clear();
+
+    SwiftErrorOps.clear();
+
+    FrameTy = nullptr;
+    FramePtr = nullptr;
+    AllocaSpillBlock = nullptr;
+  }
+
+  // Scan the function and collect the above intrinsics for later processing
+  void analyze(Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames,
+               SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
+  // If for some reason, we were not able to find coro.begin, bailout.
+  void invalidateCoroutine(Function &F,
+                           SmallVectorImpl<CoroFrameInst *> &CoroFrames);
+  // Perform ABI related initial transformation
+  void initABI();
+  // Remove orphaned and unnecessary intrinsics
+  void cleanCoroutine(SmallVectorImpl<CoroFrameInst *> &CoroFrames,
+                      SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
+
+  // Field indexes for special fields in the switch lowering.
+  struct SwitchFieldIndex {
+    enum {
+      Resume,
+      Destroy
+
+      // The promise field is always at a fixed offset from the start of
+      // frame given its type, but the index isn't a constant for all
+      // possible frames.
+
+      // The switch-index field isn't at a fixed offset or index, either;
+      // we just work it in where it fits best.
+    };
+  };
+
+  coro::ABI ABI;
+
+  StructType *FrameTy = nullptr;
+  Align FrameAlign;
+  uint64_t FrameSize = 0;
+  Value *FramePtr = nullptr;
+  BasicBlock *AllocaSpillBlock = nullptr;
+
+  struct SwitchLoweringStorage {
+    SwitchInst *ResumeSwitch;
+    AllocaInst *PromiseAlloca;
+    BasicBlock *ResumeEntryBlock;
+    unsigned IndexField;
+    unsigned IndexAlign;
+    unsigned IndexOffset;
+    bool HasFinalSuspend;
+    bool HasUnwindCoroEnd;
+  };
+
+  struct RetconLoweringStorage {
+    Function *ResumePrototype;
+    Function *Alloc;
+    Function *Dealloc;
+    BasicBlock *ReturnBlock;
+    bool IsFrameInlineInStorage;
+  };
+
+  struct AsyncLoweringStorage {
+    Value *Context;
+    CallingConv::ID AsyncCC;
+    unsigned ContextArgNo;
+    uint64_t ContextHeaderSize;
+    uint64_t ContextAlignment;
+    uint64_t FrameOffset; // Start of the frame.
+    uint64_t ContextSize; // Includes frame size.
+    GlobalVariable *AsyncFuncPointer;
+
+    Align getContextAlignment() const { return Align(ContextAlignment); }
+  };
+
+  union {
+    SwitchLoweringStorage SwitchLowering;
+    RetconLoweringStorage RetconLowering;
+    AsyncLoweringStorage AsyncLowering;
+  };
+
+  CoroIdInst *getSwitchCoroId() const {
+    assert(ABI == coro::ABI::Switch);
+    return cast<CoroIdInst>(CoroBegin->getId());
+  }
+
+  AnyCoroIdRetconInst *getRetconCoroId() const {
+    assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce);
+    return cast<AnyCoroIdRetconInst>(CoroBegin->getId());
+  }
+
+  CoroIdAsyncInst *getAsyncCoroId() const {
+    assert(ABI == coro::ABI::Async);
+    return cast<CoroIdAsyncInst>(CoroBegin->getId());
+  }
+
+  unsigned getSwitchIndexField() const {
+    assert(ABI == coro::ABI::Switch);
+    assert(FrameTy && "frame type not assigned");
+    return SwitchLowering.IndexField;
+  }
+  IntegerType *getIndexType() const {
+    assert(ABI == coro::ABI::Switch);
+    assert(FrameTy && "frame type not assigned");
+    return cast<IntegerType>(FrameTy->getElementType(getSwitchIndexField()));
+  }
+  ConstantInt *getIndex(uint64_t Value) const {
+    return ConstantInt::get(getIndexType(), Value);
+  }
+
+  PointerType *getSwitchResumePointerType() const {
+    assert(ABI == coro::ABI::Switch);
+    assert(FrameTy && "frame type not assigned");
+    return cast<PointerType>(FrameTy->getElementType(SwitchFieldIndex::Resume));
+  }
+
+  FunctionType *getResumeFunctionType() const {
+    switch (ABI) {
+    case coro::ABI::Switch:
+      return FunctionType::get(Type::getVoidTy(FrameTy->getContext()),
+                               PointerType::getUnqual(FrameTy->getContext()),
+                               /*IsVarArg=*/false);
+    case coro::ABI::Retcon:
+    case coro::ABI::RetconOnce:
+      return RetconLowering.ResumePrototype->getFunctionType();
+    case coro::ABI::Async:
+      // Not used. The function type depends on the active suspend.
+      return nullptr;
+    }
+
+    llvm_unreachable("Unknown coro::ABI enum");
+  }
+
+  ArrayRef<Type *> getRetconResultTypes() const {
+    assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce);
+    auto FTy = CoroBegin->getFunction()->getFunctionType();
+
+    // The safety of all this is checked by checkWFRetconPrototype.
+    if (auto STy = dyn_cast<StructType>(FTy->getReturnType())) {
+      return STy->elements().slice(1);
+    } else {
+      return ArrayRef<Type *>();
+    }
+  }
+
+  ArrayRef<Type *> getRetconResumeTypes() const {
+    assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce);
+
+    // The safety of all this is checked by checkWFRetconPrototype.
+    auto FTy = RetconLowering.ResumePrototype->getFunctionType();
+    return FTy->params().slice(1);
+  }
+
+  CallingConv::ID getResumeFunctionCC() const {
+    switch (ABI) {
+    case coro::ABI::Switch:
+      return CallingConv::Fast;
+
+    case coro::ABI::Retcon:
+    case coro::ABI::RetconOnce:
+      return RetconLowering.ResumePrototype->getCallingConv();
+    case coro::ABI::Async:
+      return AsyncLowering.AsyncCC;
+    }
+    llvm_unreachable("Unknown coro::ABI enum");
+  }
+
+  AllocaInst *getPromiseAlloca() const {
+    if (ABI == coro::ABI::Switch)
+      return SwitchLowering.PromiseAlloca;
+    return nullptr;
+  }
+
+  BasicBlock::iterator getInsertPtAfterFramePtr() const {
+    if (auto *I = dyn_cast<Instruction>(FramePtr)) {
+      BasicBlock::iterator It = std::next(I->getIterator());
+      It.setHeadBit(true); // Copy pre-RemoveDIs behaviour.
+      return It;
+    }
+    return cast<Argument>(FramePtr)->getParent()->getEntryBlock().begin();
+  }
+
+  /// Allocate memory according to the rules of the active lowering.
+  ///
+  /// \param CG - if non-null, will be updated for the new call
+  Value *emitAlloc(IRBuilder<> &Builder, Value *Size, CallGraph *CG) const;
+
+  /// Deallocate memory according to the rules of the active lowering.
+  ///
+  /// \param CG - if non-null, will be updated for the new call
+  void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const;
+
+  Shape() = default;
+  explicit Shape(Function &F) {
+    SmallVector<CoroFrameInst *, 8> CoroFrames;
+    SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
+
+    analyze(F, CoroFrames, UnusedCoroSaves);
+    if (!CoroBegin) {
+      invalidateCoroutine(F, CoroFrames);
+      return;
+    }
+    cleanCoroutine(CoroFrames, UnusedCoroSaves);
+  }
+};
+
+} // end namespace coro
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_COROUTINES_COROSHAPE_H

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.