add llvm.coro.sizeof intrinsic

[andrewrk]

Summary

This patch adds a new coroutine intrinsic:

llvm.coro.sizeof.i32(i8* %func)

It has the same semantics as llvm.coro.size except that
it can be called from outside the coroutine itself.

This allows the caller to be in control of memory allocation
and deallocation, which is required for the Zig frontend.

With this, the caller would ask for the frame size, allocate
the memory, potentially handle allocation failure, and then
decide whether to actually call the coroutine.

It then does not need to call the llvm.coro.destroy intrinsic;
instead it can simply free the memory. This is useful,
for example, when the caller decides to use an arena allocator
to create N coroutines, and then free them all at once by
destroying the arena.

Here is ex0.ll updated with this different way of calling:

define i8* @f(i32 %n, i8* %alloc) {
entry:
  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
  br label %loop

loop:
  %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ]
  call void @print(i32 %n.val)
  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
  switch i8 %0, label %suspend [i8 0, label %resume]
resume:
  %inc = add i32 %n.val, 1
  br label %loop

suspend:
  call i1 @llvm.coro.end(i8* %hdl, i1 0)
  ret i8* %hdl
}

define i32 @main() {
entry:
  %size = call i32 @llvm.coro.sizeof.i32(i8* bitcast (i8* (i32, i8*)* @f to i8*))
  %alloc = call i8* @malloc(i32 %size)
  %hdl = call i8* @f(i32 4, i8* %alloc)
  call void @llvm.coro.resume(i8* %hdl)
  call void @llvm.coro.resume(i8* %hdl)
  call void @free(i8* %alloc)
  ret i32 0
}

declare token @llvm.coro.id(i32, i8*, i8*, i8*)
declare i32 @llvm.coro.sizeof.i32(i8 *)
declare i8  @llvm.coro.suspend(token, i1)
declare void @llvm.coro.resume(i8*)

declare i8* @llvm.coro.begin(token, i8*)
declare i1 @llvm.coro.end(i8*, i1)

declare noalias i8* @malloc(i32)
declare void @print(i32)
declare void @free(i8*)

Here it is passed through opt ex0.ll -enable-coroutines -O2 -S, with this patch:

%f.Frame = type { void (%f.Frame*)*, void (%f.Frame*)*, i1, i1, i32 }

@f.resumers = private constant [3 x void (%f.Frame*)*] [void (%f.Frame*)* @f.resume, void (%f.Frame*)* @f.destroy, void (%f.Frame*)* @f.cleanup]

define i8* @f(i32 %n, i8* %alloc) {
entry:
  %resume.addr = bitcast i8* %alloc to void (%f.Frame*)**
  store void (%f.Frame*)* @f.resume, void (%f.Frame*)** %resume.addr, align 8
  %destroy.addr = getelementptr inbounds i8, i8* %alloc, i64 8
  %0 = bitcast i8* %destroy.addr to void (%f.Frame*)**
  store void (%f.Frame*)* @f.destroy, void (%f.Frame*)** %0, align 8
  %n.val.spill.addr = getelementptr inbounds i8, i8* %alloc, i64 20
  %1 = bitcast i8* %n.val.spill.addr to i32*
  store i32 %n, i32* %1, align 4
  tail call void @print(i32 %n)
  %index.addr1 = getelementptr inbounds i8, i8* %alloc, i64 17
  %2 = bitcast i8* %index.addr1 to i1*
  store i1 false, i1* %2, align 1
  ret i8* %alloc
}

define i32 @main() local_unnamed_addr {
entry:
  %alloc = tail call i8* @malloc(i32 0)
  %resume.addr.i = bitcast i8* %alloc to void (%f.Frame*)**
  store void (%f.Frame*)* @f.resume, void (%f.Frame*)** %resume.addr.i, align 8
  %destroy.addr.i = getelementptr inbounds i8, i8* %alloc, i64 8
  %0 = bitcast i8* %destroy.addr.i to void (%f.Frame*)**
  store void (%f.Frame*)* @f.destroy, void (%f.Frame*)** %0, align 8
  %n.val.spill.addr.i = getelementptr inbounds i8, i8* %alloc, i64 20
  %1 = bitcast i8* %n.val.spill.addr.i to i32*
  store i32 4, i32* %1, align 4
  tail call void @print(i32 4)
  %index.addr1.i = getelementptr inbounds i8, i8* %alloc, i64 17
  %2 = bitcast i8* %index.addr1.i to i1*
  store i1 false, i1* %2, align 1
  store i32 5, i32* %1, align 4
  tail call void @print(i32 5)
  store i1 false, i1* %2, align 1
  store i32 6, i32* %1, align 4
  tail call void @print(i32 6)
  store i1 false, i1* %2, align 1
  tail call void @free(i8* %alloc)
  ret i32 0
}

; Function Attrs: argmemonly nounwind readonly
declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*) #0

; Function Attrs: nounwind
declare i8* @llvm.coro.begin(token, i8* writeonly) #1

; Function Attrs: nounwind
declare noalias i8* @malloc(i32) local_unnamed_addr #1

declare void @print(i32) local_unnamed_addr

; Function Attrs: nounwind
declare void @free(i8* nocapture) local_unnamed_addr #1

define internal fastcc void @f.resume(%f.Frame* %FramePtr) {
entry.resume:
  %n.val.reload.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i64 0, i32 4
  %n.val.reload1 = load i32, i32* %n.val.reload.addr, align 4
  %inc2 = add i32 %n.val.reload1, 1
  %index.addr1 = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i64 0, i32 3
  store i32 %inc2, i32* %n.val.reload.addr, align 4
  tail call void @print(i32 %inc2)
  store i1 false, i1* %index.addr1, align 1
  ret void
}

define internal fastcc void @f.destroy(%f.Frame* %FramePtr) {
entry.destroy:
  ret void
}

define internal fastcc void @f.cleanup(%f.Frame* %FramePtr) {
entry.cleanup:
  ret void
}

attributes #0 = { argmemonly nounwind readonly }
attributes #1 = { nounwind }

It seems the CoroElide pass was not able to complete its mission.
Perhaps because with this strategy the caller does not have a token to pass
to llvm.coro.alloc. Maybe that could be added?

I would like some advice on where to go from here.

[andrewrk]

Never mind these changes. I've redesigned how coroutines will work in Zig and they no longer depend on these changes.

llvm/docs/Coroutines.rst

 None
 
 Semantics:
 """"""""""
 
 The `coro.size` intrinsic is lowered to a constant representing the size of
 the coroutine frame. 
 
+.. _coro.sizeof:
+
+'llvm.coro.sizeof' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+    declare i32 @llvm.coro.sizeof.i32(i8* %func)
+    declare i64 @llvm.coro.sizeof.i64(i8* %func)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.sizeof``' is the same as `coro.size` except it returns the
+coroutine frame size of the parameter.
+
+Arguments:
+""""""""""
+
+The func argument to '``llvm.coro.sizeof``' must be a constant bitcasted pointer
+to a function defined in the current module.
+
+Semantics:
+""""""""""
+
+Same as `coro.size` except for the function passed as a parameter instead of
+the current function.
+
 .. _coro.begin:
 
 'llvm.coro.begin' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
   declare i8* @llvm.coro.begin(token <id>, i8* <mem>)
 

llvm/include/llvm/IR/Intrinsics.td

 
 def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
                               [IntrReadMem, IntrArgMemOnly, ReadOnly<1>,
                                NoCapture<1>]>;
 def int_coro_end : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_i1_ty], []>;
 
 def int_coro_frame : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>;
 def int_coro_size : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
+def int_coro_sizeof : Intrinsic<[llvm_anyint_ty], [llvm_ptr_ty], [IntrNoMem]>;
 
 def int_coro_save : Intrinsic<[llvm_token_ty], [llvm_ptr_ty], []>;
 def int_coro_suspend : Intrinsic<[llvm_i8_ty], [llvm_token_ty, llvm_i1_ty], []>;
 
 def int_coro_param : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_ptr_ty],
                                [IntrNoMem, ReadNone<0>, ReadNone<1>]>;
 
 // Coroutine Manipulation Intrinsics.

llvm/lib/Transforms/Coroutines/CoroEarly.cpp

bool Lowerer::lowerEarlyIntrinsics(Function &F) {
   CoroIdInst *CoroId = nullptr;
   SmallVector<CoroFreeInst *, 4> CoroFrees;
   for (auto IB = inst_begin(F), IE = inst_end(F); IB != IE;) {
     Instruction &I = *IB++;
     if (auto CS = CallSite(&I)) {
       switch (CS.getIntrinsicID()) {
       default:
         continue;
+      case Intrinsic::coro_sizeof:
+        F.addFnAttr(CORO_SIZEOF_ATTR, HAS_SIZEOF_CALL);
+        continue;
       case Intrinsic::coro_free:
         CoroFrees.push_back(cast<CoroFreeInst>(&I));
         break;
       case Intrinsic::coro_suspend:
         // Make sure that final suspend point is not duplicated as CoroSplit
         // pass expects that there is at most one final suspend point.
         if (cast<CoroSuspendInst>(&I)->isFinal())
           CS.setCannotDuplicate();
struct CoroEarly : public FunctionPass {
   std::unique_ptr<Lowerer> L;
 
   // This pass has work to do only if we find intrinsics we are going to lower
   // in the module.
   bool doInitialization(Module &M) override {
     if (coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.destroy",
                                      "llvm.coro.done", "llvm.coro.end",
                                      "llvm.coro.free", "llvm.coro.promise",
-                                     "llvm.coro.resume", "llvm.coro.suspend"}))
+                                     "llvm.coro.resume", "llvm.coro.suspend",
+                                     "llvm.coro.sizeof" }))
       L = llvm::make_unique<Lowerer>(M);
     return false;
   }
 
   bool runOnFunction(Function &F) override {
     if (!L)
       return false;
 

llvm/lib/Transforms/Coroutines/CoroInstr.h

public:
   static bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_size;
   }
   static bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
+/// This represents the llvm.coro.sizeof instruction.
+class LLVM_LIBRARY_VISIBILITY CoroSizeofInst : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_sizeof;
+  }
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 /// This represents the llvm.coro.end instruction.
 class LLVM_LIBRARY_VISIBILITY CoroEndInst : public IntrinsicInst {
   enum { FrameArg, UnwindArg };
 
 public:
   bool isFallthrough() const { return !isUnwind(); }
   bool isUnwind() const {
     return cast<Constant>(getArgOperand(UnwindArg))->isOneValue();

llvm/lib/Transforms/Coroutines/CoroInternal.h

 // CoroElide pass that triggers a restart of the pipeline by CGPassManager.
 // When CoroSplit pass sees the same coroutine the second time, it splits it up,
 // adds coroutine subfunctions to the SCC to be processed by IPO pipeline.
 
 #define CORO_PRESPLIT_ATTR "coroutine.presplit"
 #define UNPREPARED_FOR_SPLIT "0"
 #define PREPARED_FOR_SPLIT "1"
 
+// CoroEarly pass marks functions that do coro.sizeof
+// and then CoroSplit goes back and lowers this to the frame size.
+#define CORO_SIZEOF_ATTR "coroutine.sizeof"
+#define HAS_SIZEOF_CALL "1"
+
+// CoroSplit adds this attribute so that it can be accessed later
+#define CORO_FRAMESIZE_ATTR "coroutine.framesize"
+
+
 #define CORO_DEVIRT_TRIGGER_FN "coro.devirt.trigger"
 
 namespace coro {
 
 bool declaresIntrinsics(Module &M, std::initializer_list<StringRef>);
 void replaceAllCoroAllocs(CoroBeginInst *CB, bool Replacement);
 void replaceAllCoroFrees(CoroBeginInst *CB, Value *Replacement);
 void replaceCoroFree(CoroIdInst *CoroId, bool Elide);

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

static void removeCoroEnds(coro::Shape &Shape) {
   auto *False = ConstantInt::getFalse(Context);
 
   for (CoroEndInst *CE : Shape.CoroEnds) {
     CE->replaceAllUsesWith(False);
     CE->eraseFromParent();
   }
 }
 
-static void replaceFrameSize(coro::Shape &Shape) {
+static void replaceFrameSize(Function &F, coro::Shape &Shape) {
   if (Shape.CoroSizes.empty())
     return;
 
   // In the same function all coro.sizes should have the same result type.
   auto *SizeIntrin = Shape.CoroSizes.back();
   Module *M = SizeIntrin->getModule();
   const DataLayout &DL = M->getDataLayout();
   auto Size = DL.getTypeAllocSize(Shape.FrameTy);
   auto *SizeConstant = ConstantInt::get(SizeIntrin->getType(), Size);
 
+  F.addFnAttr(CORO_FRAMESIZE_ATTR, std::to_string(Size));
+
   for (CoroSizeInst *CS : Shape.CoroSizes) {
     CS->replaceAllUsesWith(SizeConstant);
     CS->eraseFromParent();
   }
 }
 
 // Create a global constant array containing pointers to functions provided and
 // set Info parameter of CoroBegin to point at this constant. Example:
for (auto B = BB.begin(), E = BB.end(); B != E;) {
     if (&I == CoroBegin)
       break;
     if (DoNotRelocateSet.count(&I))
       continue;
     I.moveBefore(InsertPt);
   }
 }
 
+static void lowerSizeOf(Function &ParentF, CallGraph &CG, CallGraphSCC &SCC) {
+  SmallVector<CoroSizeofInst *, 4> SizeOfInstrs;
+  for (auto IB = inst_begin(ParentF), IE = inst_end(ParentF); IB != IE;) {
+    Instruction &I = *IB++;
+    if (auto CS = CallSite(&I)) {
+      if (CS.getIntrinsicID() == Intrinsic::coro_sizeof) {
+        SizeOfInstrs.push_back(cast<CoroSizeofInst>(&I));
+      }
+    }
+  }
+  for (CoroSizeofInst *CS : SizeOfInstrs) {
+    auto CoroF = cast<Function>(CS->getArgOperand(0)->stripPointerCasts());
+    Attribute Attr = CoroF->getFnAttribute(CORO_FRAMESIZE_ATTR);
+    auto Size = Attr.getValueAsInt();
+    auto *SizeConstant = ConstantInt::get(CS->getType(), Size);
+    CS->replaceAllUsesWith(SizeConstant);
+    CS->eraseFromParent();
+  }
+  ParentF.removeFnAttr(CORO_SIZEOF_ATTR);
+}
+
 static void splitCoroutine(Function &F, CallGraph &CG, CallGraphSCC &SCC) {
   coro::Shape Shape(F);
   if (!Shape.CoroBegin)
     return;
 
   simplifySuspendPoints(Shape);
   relocateInstructionBefore(Shape.CoroBegin, F);
   buildCoroutineFrame(F, Shape);
-  replaceFrameSize(Shape);
+  replaceFrameSize(F, Shape);
 
   // If there are no suspend points, no split required, just remove
   // the allocation and deallocation blocks, they are not needed.
   if (Shape.CoroSuspends.empty()) {
     handleNoSuspendCoroutine(Shape.CoroBegin, Shape.FrameTy);
     removeCoroEnds(Shape);
     postSplitCleanup(F);
     coro::updateCallGraph(F, {}, CG, SCC);
struct CoroSplit : public CallGraphSCCPass {
   }
 
   bool runOnSCC(CallGraphSCC &SCC) override {
     if (!Run)
       return false;
 
     // Find coroutines for processing.
     SmallVector<Function *, 4> Coroutines;
-    for (CallGraphNode *CGN : SCC)
-      if (auto *F = CGN->getFunction())
+    SmallVector<Function *, 4> FnsWithSizeOf;
+    for (CallGraphNode *CGN : SCC) {
+      if (auto *F = CGN->getFunction()) {
         if (F->hasFnAttribute(CORO_PRESPLIT_ATTR))
           Coroutines.push_back(F);
+        if (F->hasFnAttribute(CORO_SIZEOF_ATTR))
+          FnsWithSizeOf.push_back(F);
+      }
+    }
 
-    if (Coroutines.empty())
+    if (Coroutines.empty() && FnsWithSizeOf.empty()) {
       return false;
+    }
 
     CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
+    if (!Coroutines.empty()) {
       createDevirtTriggerFunc(CG, SCC);
 
       for (Function *F : Coroutines) {
         Attribute Attr = F->getFnAttribute(CORO_PRESPLIT_ATTR);
         StringRef Value = Attr.getValueAsString();
         DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F->getName()
                     << "' state: " << Value << "\n");
         if (Value == UNPREPARED_FOR_SPLIT) {
           prepareForSplit(*F, CG);
           continue;
         }
         F->removeFnAttr(CORO_PRESPLIT_ATTR);
         splitCoroutine(*F, CG, SCC);
       }
+    }
+    for (Function *F : FnsWithSizeOf) {
+      lowerSizeOf(*F, CG, SCC);
+    }
     return true;
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     CallGraphSCCPass::getAnalysisUsage(AU);
   }
 
   StringRef getPassName() const override { return "Coroutine Splitting"; }

llvm/lib/Transforms/Coroutines/Coroutines.cpp

 #ifndef NDEBUG
 static bool isCoroutineIntrinsicName(StringRef Name) {
   // NOTE: Must be sorted!
   static const char *const CoroIntrinsics[] = {
       "llvm.coro.alloc",   "llvm.coro.begin",   "llvm.coro.destroy",
       "llvm.coro.done",    "llvm.coro.end",     "llvm.coro.frame",
       "llvm.coro.free",    "llvm.coro.id",      "llvm.coro.param",
       "llvm.coro.promise", "llvm.coro.resume",  "llvm.coro.save",
-      "llvm.coro.size",    "llvm.coro.subfn.addr", "llvm.coro.suspend",
+      "llvm.coro.size",    "llvm.coro.sizeof",  "llvm.coro.subfn.addr",
+      "llvm.coro.suspend",
   };
   return Intrinsic::lookupLLVMIntrinsicByName(CoroIntrinsics, Name) != -1;
 }
 #endif
 
 // Verifies if a module has named values listed. Also, in debug mode verifies
 // that names are intrinsic names.
 bool coro::declaresIntrinsics(Module &M,