[Coroutines] Introduce @llvm.coro.tls.wrapper to block optimizations

[ChuanqiXu]

Summary

This is the original idea to solve the TLS problem in coroutines described in https://discourse.llvm.org/t/address-thread-identification-problems-with-coroutine/62015.

The main discussion of the thread is about how should we handle readnone attribute and this is already done in D127383. The here is the other problem about TLS variables.

@jyknight points out that we could solve this problem and the TLS is not a constant problem in one shot and we get in consensus this is the good direction. So I sent D125291 and D129833.

But now clang15 branch is going to be created in the next week and D125291 has no hopes to land before that. It implies that the TLS bug would remain in clang15. I think it is really a pity that the bug has lived for 2 years.

@rjmccall sent concerns about the automatically merge of TLS addresses so that CoroEarly might not be early enough. But both of us failed to give an example. My thought is that the current one might not be perfect and D125291 should be the direction in the future. And the current one is simple enough and workable in most cases. I think we could land this first to fix the bug. I would revert this one if D125291 or any other similar patches get landed. How do you think about the strategy?

[ychen]

John mentioned this test case,

thread_local int x = 0;
void helper(int *x, int y);
my_coro coroutine() {
  helper(&x, co_await something_suspending());
}

Is replacing the TLS variable with llvm.coro.tls.wrapper in CoroEarly too late to fix the issue?

[ChuanqiXu]

In D130142#3666355, @ychen wrote:

John mentioned this test case,

thread_local int x = 0;
void helper(int *x, int y);
my_coro coroutine() {
  helper(&x, co_await something_suspending());
}

Is replacing the TLS variable with llvm.coro.tls.wrapper in CoroEarly too late to fix the issue?

This example may be fine since C++ standard specifies that the order of evaluation of function arguments is unspecified: https://en.cppreference.com/w/cpp/language/eval_order.

[ychen]

In D130142#3667280, @ChuanqiXu wrote:

In D130142#3666355, @ychen wrote:

John mentioned this test case,

thread_local int x = 0;
void helper(int *x, int y);
my_coro coroutine() {
  helper(&x, co_await something_suspending());
}

Is replacing the TLS variable with llvm.coro.tls.wrapper in CoroEarly too late to fix the issue?

This example may be fine since C++ standard specifies that the order of evaluation of function arguments is unspecified: https://en.cppreference.com/w/cpp/language/eval_order.

That's my understanding as well. Sorry I should've changed it a little,

thread_local int x = 0;
void helper(int *x, int y);
my_coro coroutine() {
  helper(&x, (&x, co_await something_suspending()));
}

The comma operator should have the left-to-right evaluation order. Personally, I hope there is a workaround for Clang 15. I'm just trying to convince myself that this workaround solves both issues in principle (even with potential pessimization with the TLS variables). But this test case makes me hesitate.

[ChuanqiXu]

In D130142#3667305, @ychen wrote:

In D130142#3667280, @ChuanqiXu wrote:

In D130142#3666355, @ychen wrote:

John mentioned this test case,

thread_local int x = 0;
void helper(int *x, int y);
my_coro coroutine() {
  helper(&x, co_await something_suspending());
}

Is replacing the TLS variable with llvm.coro.tls.wrapper in CoroEarly too late to fix the issue?

This example may be fine since C++ standard specifies that the order of evaluation of function arguments is unspecified: https://en.cppreference.com/w/cpp/language/eval_order.

That's my understanding as well. Sorry I should've changed it a little,

thread_local int x = 0;
void helper(int *x, int y);
my_coro coroutine() {
  helper(&x, (&x, co_await something_suspending()));
}

The comma operator should have the left-to-right evaluation order. Personally, I hope there is a workaround for Clang 15. I'm just trying to convince myself that this workaround solves both issues in principle (even with potential pessimization with the TLS variables). But this test case makes me hesitate.

For the specific example, it is fine since &x would be deprecated due to it has no side effect as a comma operator. I guess this example fits your intention more:

thread_local int tls_variable = 0;
void helper(int *x);
int *helper1(int *x);
task coro() {
    helper((helper1(&tls_variable), co_await MyAwaiter()));
}

In this example, &tls_variable would be evaluated before the await expression: https://godbolt.org/z/xse1q57ze

[ChuanqiXu]

Given the bug is important and the patch is relatively small and innocent, I want to land this one before 15.x branch is created if no objection comes in.

[ychen]

In D130142#3675062, @ChuanqiXu wrote:

Given the bug is important and the patch is relatively small and innocent, I want to land this one before 15.x branch is created if no objection comes in.

Sounds good to me. One last question is: why not just use llvm.threadlocal.address? The plan is to replace llvm.coro.tls.wrapper with llvm.threadlocal.address right? I'd prefer not to add an LLVM intrinsic just for one specific release.

[ChuanqiXu]

In D130142#3677277, @ychen wrote:

In D130142#3675062, @ChuanqiXu wrote:

Given the bug is important and the patch is relatively small and innocent, I want to land this one before 15.x branch is created if no objection comes in.

Sounds good to me. One last question is: why not just use llvm.threadlocal.address? The plan is to replace llvm.coro.tls.wrapper with llvm.threadlocal.address right? I'd prefer not to add an LLVM intrinsic just for one specific release.

Yeah, the plan is going to replace llvm.coro.tls.wrapper with llvm.threadlocal.address. The reason why I introduced a new intrinsics is the semantic differences. llvm.threadlocal.address is a fundamental one and llvm.coro.tls.wrapper is specific to coroutines. I think it might not be too bad. Since we've commented that llvm.coro.tls.wrapper is used internally. I feel like it might not be possible that someone else would use it accidently.

[ychen]

LGTM. Please mention in the docs along the lines that we will "replace llvm.coro.tls.wrapper with llvm.threadlocal.address." in the future.

[ChuanqiXu]

In D130142#3678221, @ychen wrote:

LGTM. Please mention in the docs along the lines that we will "replace llvm.coro.tls.wrapper with llvm.threadlocal.address." in the future.

Thanks! Will do.

[ChuanqiXu]

We failed to land the dependent patch in time. So we could keep patient now.

[ChuanqiXu]

Abandon this one since we landed D129833.

llvm/docs/Coroutines.rst

 abnormally (non-zero).
 
 In a normal coroutine, it is undefined behavior if the coroutine executes
 a call to ``llvm.coro.suspend.retcon`` after resuming abnormally.
 
 In a yield-once coroutine, it is undefined behavior if the coroutine
 executes a call to ``llvm.coro.suspend.retcon`` after resuming in any way.
 
+Coroutine Helper Intrinsics
+------------------------------
+Intrinsics described in this section are used as a helper to show
+the changed properties after we introduced coroutines.
+
+.. _coro.tls.wrapper:
+
+'llvm.coro.tls.wrapper' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+    declare ptr @llvm.coro.tls.wrapper(ptr)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.tls.wrapper``' intrinsic refers to the address of the TLS
+variables. The addresses of TLS variables are thought to be constant in one
+function. But it is not true in coroutines due to a coroutine may resume in
+another thread.
+
+Arguments:
+""""""""""
+
+The address of the TLS variable.
+
+Semantics:
+""""""""""
+
+The `llvm.coro.tls.wrapper` intrinsic would be replaced with its argument
+after we lowered all the coroutines.
+
 Coroutine Transformation Passes
 ===============================
 CoroEarly
 ---------
 The pass CoroEarly lowers coroutine intrinsics that hide the details of the
 structure of the coroutine frame, but, otherwise not needed to be preserved to
 help later coroutine passes. This pass lowers `coro.frame`_, `coro.done`_,
 and `coro.promise`_ intrinsics.

llvm/include/llvm/IR/Intrinsics.td

 
 // Coroutine Lowering Intrinsics. Used internally by coroutine passes.
 
 def int_coro_subfn_addr : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty, llvm_i8_ty],
                                     [IntrReadMem, IntrArgMemOnly,
                                      ReadOnly<ArgIndex<0>>,
                                      NoCapture<ArgIndex<0>>]>;
 
+// A wrapper for TLS variables to avoid incorrect merging in coroutines.
+// Since a coroutine may resume in another thread.
+def int_coro_tls_wrapper : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
+                                     [IntrNoMem, IntrSpeculatable,
+                                      IntrWillReturn]>;
+
 ///===-------------------------- Other Intrinsics --------------------------===//
 //
 def int_trap : Intrinsic<[], [], [IntrNoReturn, IntrCold]>,
                ClangBuiltin<"__builtin_trap">;
 def int_debugtrap : Intrinsic<[]>,
                     ClangBuiltin<"__builtin_debugtrap">;
 def int_ubsantrap : Intrinsic<[], [llvm_i8_ty],
                               [IntrNoReturn, IntrCold, ImmArg<ArgIndex<0>>]>;

llvm/lib/Transforms/Coroutines/CoroCleanup.cpp

if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
         break;
       case Intrinsic::coro_end:
       case Intrinsic::coro_suspend_retcon:
         if (IsPrivateAndUnprocessed) {
           II->replaceAllUsesWith(UndefValue::get(II->getType()));
         } else
           continue;
         break;
+      case Intrinsic::coro_tls_wrapper:
+        II->replaceAllUsesWith(II->getOperand(0));
+        break;
       case Intrinsic::coro_async_size_replace:
         auto *Target = cast<ConstantStruct>(
             cast<GlobalVariable>(II->getArgOperand(0)->stripPointerCasts())
                 ->getInitializer());
         auto *Source = cast<ConstantStruct>(
             cast<GlobalVariable>(II->getArgOperand(1)->stripPointerCasts())
                 ->getInitializer());
         auto *TargetSize = Target->getOperand(1);
bool Lowerer::lower(Function &F) {
   return Changed;
 }
 
 static bool declaresCoroCleanupIntrinsics(const Module &M) {
   return coro::declaresIntrinsics(
       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.tls.wrapper"});
 }
 
 PreservedAnalyses CoroCleanupPass::run(Module &M,
                                        ModuleAnalysisManager &MAM) {
   if (!declaresCoroCleanupIntrinsics(M))
     return PreservedAnalyses::all();
 
   FunctionAnalysisManager &FAM =

llvm/lib/Transforms/Coroutines/CoroEarly.cpp

 PreservedAnalyses CoroEarlyPass::run(Module &M, ModuleAnalysisManager &) {
   if (!declaresCoroEarlyIntrinsics(M))
     return PreservedAnalyses::all();
 
   Lowerer L(M);
   for (auto &F : M)
     L.lowerEarlyIntrinsics(F);
 
+  // Add @llvm.coro.tls.wrapper to TLS variables to avoid
+  // misoptimization.
+  //
+  // Previously, the optimizer could assume the address of a TLS variable is
+  // same in the same function. The assumption is broken now after we introduced
+  // coroutines.
+  //
+  // Note that we couldn't mark the use in coroutines only due to potentially
+  // inlining. It is too expensive to do reachability analysis here to mark the
+  // use for reachable function only. And the TLS variables in normal functions
+  // should be optimized correctly since @llvm.coro.tls.wrapper is
+  // marked as readnone.
+  auto *MaychangeFn =
+      Intrinsic::getDeclaration(&M, Intrinsic::coro_tls_wrapper);
+  for (auto &GV : M.getGlobalList())
+    if (GV.isThreadLocal())
+      for (auto &Use : llvm::make_early_inc_range(GV.uses())) {
+        auto *UserInst = dyn_cast<Instruction>(Use.getUser());
+        if (!UserInst)
+            continue;
+
+        if (auto *PN = dyn_cast<PHINode>(UserInst)) {
+            auto *IncomingBB = PN->getIncomingBlock(Use);
+            UserInst = IncomingBB->getTerminator();
+        }
+
+        Use.set(CallInst::Create(MaychangeFn, {&GV}, GV.getName(), UserInst));
+      }
+
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   return PA;
 }

llvm/lib/Transforms/Coroutines/Coroutines.cpp

static const char *const CoroIntrinsics[] = {
     "llvm.coro.promise",
     "llvm.coro.resume",
     "llvm.coro.save",
     "llvm.coro.size",
     "llvm.coro.subfn.addr",
     "llvm.coro.suspend",
     "llvm.coro.suspend.async",
     "llvm.coro.suspend.retcon",
+    "llvm.coro.tls.wrapper",
 };
 
 #ifndef NDEBUG
 static bool isCoroutineIntrinsicName(StringRef Name) {
   return Intrinsic::lookupLLVMIntrinsicByName(CoroIntrinsics, Name) != -1;
 }
 #endif
 
 bool coro::declaresAnyIntrinsic(const Module &M) {
   for (StringRef Name : CoroIntrinsics) {
+    if (!isCoroutineIntrinsicName(Name))
+      llvm::outs() << Name << " is not a coro intrinsic.\n";
     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
     if (M.getNamedValue(Name))
       return true;
   }
 
   return false;
 }
 

llvm/test/Transforms/Coroutines/coro-TLS-01.ll

+; Tests that the TLS variables which cross suspend points wouldn't be misoptimized.
+; RUN: opt < %s -S -passes=coro-early,sroa,early-cse,coro-split,coro-cleanup,simplifycfg | FileCheck %s
+; RUN: opt < %s -S -passes='default<O2>' | FileCheck %s
+
+@tls_variable = thread_local global i32 0
+
+define ptr @f() presplitcoroutine {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %i = alloca ptr
+  %j = alloca ptr
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+  store ptr @tls_variable, ptr %i
+  %sus_result = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sus_result, label %suspend [i8 0, label %resume
+                                         i8 1, label %cleanup]
+resume:
+  store ptr @tls_variable, ptr %j
+  %i_value = load ptr, ptr %i
+  %j_value = load ptr, ptr %j
+  %cmp = icmp eq ptr %i_value, %j_value
+  br i1 %cmp, label %same, label %diff
+
+same:
+  call void @print_same()
+  br label %cleanup
+
+diff:
+  call void @print_diff()
+  br label %cleanup
+
+cleanup:
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0)
+  ret ptr %hdl
+}
+
+; CHECK-LABEL: f.resume(
+; CHECK: br i1 %cmp, label %same, label %diff
+; CHECK-EMPTY:
+; CHECK-NEXT: same:
+; CHECK-NEXT:   call void @print_same()
+; CHECK-NEXT:   br label
+; CHECK-EMPTY:
+; CHECK-NEXT: diff:
+; CHECK-NEXT:   call void @print_diff()
+; CHECK-NEXT:   br label
+
+declare void @print_same()
+declare void @print_diff()
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr)
+declare i1 @llvm.coro.end(ptr, i1)
+
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)

llvm/test/Transforms/Coroutines/coro-TLS-02.ll

+; Tests that the TLS variables which don't cross suspend points would be optimized correctly.
+; RUN: opt < %s -S -passes=coro-early,coro-split,coro-cleanup,sroa,early-cse,simplifycfg | FileCheck %s
+; RUN: opt < %s -S -passes='default<O2>' | FileCheck %s
+
+@tls_variable = thread_local global i32 0
+
+define ptr @f() presplitcoroutine {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %i = alloca ptr
+  %j = alloca ptr
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+  %sus_result = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sus_result, label %suspend [i8 0, label %resume
+                                         i8 1, label %cleanup]
+resume:
+  store ptr @tls_variable, ptr %i
+  store ptr @tls_variable, ptr %j
+  %i_value = load ptr, ptr %i
+  %j_value = load ptr, ptr %j
+  %cmp = icmp eq ptr %i_value, %j_value
+  br i1 %cmp, label %same, label %diff
+
+same:
+  call void @print_same()
+  br label %cleanup
+
+diff:
+  call void @print_diff()
+  br label %cleanup
+
+cleanup:
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0)
+  ret ptr %hdl
+}
+
+; Tests that normal functions could be optimized expectedly.
+define void @normal_func() {
+entry:
+  %i = alloca ptr
+  %j = alloca ptr
+  store ptr @tls_variable, ptr %i
+  store ptr @tls_variable, ptr %j
+  %i_value = load ptr, ptr %i
+  %j_value = load ptr, ptr %j
+  %cmp = icmp eq ptr %i_value, %j_value
+  br i1 %cmp, label %same, label %diff
+
+same:
+  call void @print_same()
+  br label %ret
+
+diff:
+  call void @print_diff()
+  br label %ret
+
+ret:
+  ret void
+}
+
+; CHECK: void @normal_func
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   call void @print_same()
+; CHECk-NEXT:   ret
+
+; CHECK: void @f.resume
+; CHECK-NEXT: [[LABELNAME:.*]]:
+; CHECK-NEXT: call void @print_same()
+; CHECK-NEXT: call void @free(ptr{{.*}}%hdl)
+; CHECK-NEXT: ret void 
+
+declare void @print_same()
+declare void @print_diff()
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr)
+declare i1 @llvm.coro.end(ptr, i1)
+
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)

llvm/test/Transforms/Coroutines/coro-TLS-cleanup.ll

+; Test that coro-cleanup would convert llvm.coro.tls.wrapper intrinsics
+; correctly.
+; RUN: opt < %s -S -passes=coro-cleanup | FileCheck %s
+
+%f.Frame = type { ptr, ptr, ptr, ptr, i1 }
+
+@tls_variable = thread_local global i32 0
[e-mail] = private constant [3 x ptr] [ptr @f.resume, ptr @f.destroy, ptr @f.cleanup]
+
+define ptr @f() {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr @f, ptr @f.resumers)
+  %alloc = call ptr @malloc(i32 40)
+  %i = alloca ptr, align 8
+  %j = alloca ptr, align 8
+  %hdl = call noalias nonnull ptr @llvm.coro.begin(token %id, ptr %alloc)
+  %resume.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 0
+  store ptr @f.resume, ptr %resume.addr, align 8
+  %destroy.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 1
+  store ptr @f.destroy, ptr %destroy.addr, align 8
+  %i.reload.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 2
+  %j.reload.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 3
+  ; CHECK: store ptr @tls_variable, ptr %i.reload.addr, align 8
+  %tls_variable1 = call ptr @llvm.coro.tls.wrapper(ptr @tls_variable)
+  store ptr %tls_variable1, ptr %i.reload.addr, align 8
+  %index.addr2 = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 4
+  store i1 false, ptr %index.addr2, align 1
+  ret ptr %hdl
+}
+
+define internal fastcc void @f.resume(ptr noalias nonnull align 8 dereferenceable(40) %hdl) {
+entry.resume:
+  %i.reload.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 2
+  %j.reload.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 3
+  %tls_variable = call ptr @llvm.coro.tls.wrapper(ptr @tls_variable)
+  ; CHECK: store ptr @tls_variable, ptr %j.reload.addr, align 8
+  store ptr %tls_variable, ptr %j.reload.addr, align 8
+  call void @consume(ptr %i.reload.addr)
+  call void @consume(ptr %j.reload.addr)
+  call void @free(ptr %hdl)
+  ret void
+}
+
+define internal fastcc void @f.destroy(ptr noalias nonnull align 8 dereferenceable(40) %hdl) {
+entry.destroy:
+  %i.reload.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 2
+  %j.reload.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 3
+  call void @free(ptr %hdl)
+  ret void
+}
+
+define internal fastcc void @f.cleanup(ptr noalias nonnull align 8 dereferenceable(40) %hdl) {
+entry.cleanup:
+  %i.reload.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 2
+  %j.reload.addr = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 3
+  call void @free(ptr null)
+  ret void
+}
+
+declare void @consume(ptr)
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr)
+declare i1 @llvm.coro.end(ptr, i1)
+
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)
+
+declare ptr @llvm.coro.tls.wrapper(ptr)

llvm/test/Transforms/Coroutines/coro-TLS-early.ll

+; Tests that coro-early would generate llvm.coro.tls.wrapper intrinsics for TLS
+; variable correctly.
+; RUN: opt < %s -S -passes=coro-early | FileCheck %s
+
+@tls_variable = thread_local global i32 0
+
+define ptr @f() presplitcoroutine {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %i = alloca ptr
+  %j = alloca ptr
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+  ; CHECK: %[[TLS_VARIABLE:.*]] = call ptr @llvm.coro.tls.wrapper(ptr @tls_variable)
+  ; CHECK-NEXT: store ptr %[[TLS_VARIABLE]], ptr %i
+  store ptr @tls_variable, ptr %i
+  %sus_result = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sus_result, label %suspend [i8 0, label %resume
+                                         i8 1, label %cleanup]
+resume:
+  ; CHECK: %[[TLS_VARIABLE2:.*]] = call ptr @llvm.coro.tls.wrapper(ptr @tls_variable)
+  ; CHECK-NEXT:   store ptr %[[TLS_VARIABLE2]], ptr %j
+  ; CHECK-NEXT:   call void @consume(ptr %i)
+  ; CHECK-NEXT:   call void @consume(ptr %j)
+  ; CHECK-NEXT:   %[[PHI_TLS:.*]] = call ptr @llvm.coro.tls.wrapper(ptr @tls_variable)
+  ; CHECK-NEXT:   br label %cleanup
+  ; CHECK-EMPTY:
+  ; CHECK-NEXT: cleanup:
+  ; CHECK-NEXT:   %phi = phi ptr [ %[[PHI_TLS]], %resume ], [ null, %entry ]
+  store ptr @tls_variable, ptr %j
+  call void @consume(ptr %i)
+  call void @consume(ptr %j)
+  br label %cleanup
+
+cleanup:
+  %phi = phi ptr [@tls_variable, %resume], [null, %entry]
+  call void @consume(ptr %phi)
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0)
+  ret ptr %hdl
+}
+
+declare void @consume(ptr)
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr)
+declare i1 @llvm.coro.end(ptr, i1)
+
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)