[coroutines] Experimenting with compile-time size check

[GorNishanov]

Summary

! Not intended to be merged into the trunk at the moment !

Adding a coro_size_check(IntConstant) intrinsic.

Presense of this intrinsic in the body of the coroutine inhibits Halo optimization of both into this coroutine and this coroutine into others.

The intent is that this intrinsics will be emitted into the Coroutine IR when a constract is specified on operator new for the coroutine indicating the maximum allowed size. To be used when coroutine state is stored in a fixed size buffer.

At the moment this uses report_fatal_error, in further iteration diagnistic will be improve. Currently it emits two kind of errors:

error in backend: In function '_Z1fv' with coroutine frame size 24 front end size estimate is not a constant.

or

In function '_Z1fv' coroutine frame size 24 exceeds the frontend estimate of 15.

Intrinsic can be extended to include more precise location information so that it can point precisely back to file name + line number.

Here is a background for the intended use by the frontend:

The coroutine machinery requests memory via overloaded operator new on the coroutine promise.

Let’s say we put a contract there:

template <size_t MaxSize>
void *operator new(size_t sz,

                 Storage<MaxSize> put_it_here)[[expects:sz <= MaxSize]] {
return &put_it_here;

}

Note that in “[[expects: sz <= MaxSize]]” contract sz is a compile time constant, but, it is a constant that is known at translation time, but not necessarily during semantic analysis time.

Normally, it is QoI if contract checking is done at compile time as opposed to a runtime.

I am wondering if having some wording that mandates compile checking of some contracts we will have desired portable guaranteed way of given a hard compile time error if precondition on operator new is not satisfied.

Here is one way to give the guarantee: (Somewhere in the coroutine wording when describing how memory for the coroutine state is acquired).

“If operator new has a precondition contract AND that contract condition would be a constant expression if “sz” argument were a constant time expression, then, the contract will be checked even in the Contract checking “Off” mode and violation will make the program ill-formed.”

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_noop : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>;
 def int_coro_size : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
+def int_coro_size_chk : Intrinsic<[], [llvm_anyint_ty], []>;
 
 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.

lib/Transforms/Coroutines/CoroCleanup.cpp

if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
         II->replaceAllUsesWith(ConstantInt::getTrue(Context));
         break;
       case Intrinsic::coro_id:
         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
         break;
       case Intrinsic::coro_subfn_addr:
         lowerSubFn(Builder, cast<CoroSubFnInst>(II));
         break;
+      case Intrinsic::coro_size_chk:
+        // Just remove it.
+        break;
       }
       II->eraseFromParent();
       Changed = true;
     }
   }
 
   if (Changed) {
     // After replacement were made we can cleanup the function body a little.

lib/Transforms/Coroutines/CoroElide.cpp

 bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
                             DominatorTree &DT) {
   CoroBegins.clear();
   CoroAllocs.clear();
   CoroFrees.clear();
   ResumeAddr.clear();
   DestroyAddr.clear();
 
+  bool HasSizeChecksIntrinsics = false;
+
+  // TODO: For now, just the mere presense of the coro_size_chk disables
+  // the Halo optimization. We need to investigate if less blunt approach
+  // is possible.
+  for (Instruction &I : instructions(CoroId->getFunction()))
+    if (isa<CoroSizeChkInst>(&I)) {
+      LLVM_DEBUG(dbgs() << "disabling Halo optimization "
+                           " due to presence of coro_size_chk\n");
+      HasSizeChecksIntrinsics = false;
+      break;
+    }
+
   // Collect all coro.begin and coro.allocs associated with this coro.id.
   for (User *U : CoroId->users()) {
     if (auto *CB = dyn_cast<CoroBeginInst>(U))
       CoroBegins.push_back(CB);
     else if (auto *CA = dyn_cast<CoroAllocInst>(U))
       CoroAllocs.push_back(CA);
     else if (auto *CF = dyn_cast<CoroFreeInst>(U))
       CoroFrees.push_back(CF);
bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
   bool ShouldElide = shouldElide(CoroId->getFunction(), DT);
 
   auto *DestroyAddrConstant = ConstantExpr::getExtractValue(
       Resumers,
       ShouldElide ? CoroSubFnInst::CleanupIndex : CoroSubFnInst::DestroyIndex);
 
   replaceWithConstant(DestroyAddrConstant, DestroyAddr);
 
-  if (ShouldElide) {
+  if (ShouldElide && !HasSizeChecksIntrinsics) {
     auto *FrameTy = getFrameType(cast<Function>(ResumeAddrConstant));
     elideHeapAllocations(CoroId->getFunction(), FrameTy, AA);
     coro::replaceCoroFree(CoroId, /*Elide=*/true);
   }
 
   return true;
 }
 

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.size instruction.
+class LLVM_LIBRARY_VISIBILITY CoroSizeChkInst : public IntrinsicInst {
+  enum { SizeEst };
+
+  ConstantInt *getRawIndex() const {
+    return cast<ConstantInt>(getArgOperand(SizeEst));
+  }
+public:
+  bool isConstant() const {
+    return isa<ConstantInt>(getArgOperand(SizeEst));
+  }
+
+  int64_t getSizeEstimate() const {
+    return getRawIndex()->getValue().getSExtValue();
+  }
+
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_size_chk;
+  }
+  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();

lib/Transforms/Coroutines/CoroInternal.h

 };
 
 // Holds structural Coroutine Intrinsics for a particular function and other
 // values used during CoroSplit pass.
 struct LLVM_LIBRARY_VISIBILITY Shape {
   CoroBeginInst *CoroBegin;
   SmallVector<CoroEndInst *, 4> CoroEnds;
   SmallVector<CoroSizeInst *, 2> CoroSizes;
+  SmallVector<CoroSizeChkInst *, 1> CoroSizeChks;
   SmallVector<CoroSuspendInst *, 4> CoroSuspends;
 
   // Field Indexes for known coroutine frame fields.
   enum {
     ResumeField,
     DestroyField,
     PromiseField,
     IndexField,

lib/Transforms/Coroutines/CoroSplit.cpp

static void replaceFrameSize(coro::Shape &Shape) {
   const DataLayout &DL = M->getDataLayout();
   auto Size = DL.getTypeAllocSize(Shape.FrameTy);
   auto *SizeConstant = ConstantInt::get(SizeIntrin->getType(), Size);
 
   for (CoroSizeInst *CS : Shape.CoroSizes) {
     CS->replaceAllUsesWith(SizeConstant);
     CS->eraseFromParent();
   }
+
+  for (CoroSizeChkInst *CS : Shape.CoroSizeChks) {
+    if (!CS->isConstant())
+      report_fatal_error("In function '" + CS->getFunction()->getName() +
+        "' with coroutine frame size " + Twine(Size) +
+        " front end size estimate is not a constant.");
+    else if ((int64_t)Size > CS->getSizeEstimate())
+      report_fatal_error("In function '" + CS->getFunction()->getName() +
+        "' coroutine frame size " + Twine(Size) + " exceeds the frontend estimate of "
+        + Twine(CS->getSizeEstimate()) + ".");
+
+    // Note that we do not remove the check intrinsic, so that it can be
+    // used to inhibit Halo optimization.
+  }
 }
 
 // Create a global constant array containing pointers to functions provided and
 // set Info parameter of CoroBegin to point at this constant. Example:
 //
 //   @f.resumers = internal constant [2 x void(%f.frame*)*]
 //                    [void(%f.frame*)* @f.resume, void(%f.frame*)* @f.destroy]
 //   define void @f() {

lib/Transforms/Coroutines/Coroutines.cpp

void coro::Shape::buildFrom(Function &F) {
   for (Instruction &I : instructions(F)) {
     if (auto II = dyn_cast<IntrinsicInst>(&I)) {
       switch (II->getIntrinsicID()) {
       default:
         continue;
       case Intrinsic::coro_size:
         CoroSizes.push_back(cast<CoroSizeInst>(II));
         break;
+      case Intrinsic::coro_size_chk:
+        CoroSizeChks.push_back(cast<CoroSizeChkInst>(II));
+        break;
       case Intrinsic::coro_frame:
         CoroFrames.push_back(cast<CoroFrameInst>(II));
         break;
       case Intrinsic::coro_save:
         // After optimizations, coro_suspends using this coro_save might have
         // been removed, remember orphaned coro_saves to remove them later.
         if (II->use_empty())
           UnusedCoroSaves.push_back(cast<CoroSaveInst>(II));