[Coroutines] Refactor sinkLifetimeStartMarkers

[junparser]

Summary

D82314 has implemented a method for coroutine frame reducing in sinkLifetimeStartMarkers, which sink lifetime.start marker to multiple users. However, this may cause mismatch. Consider this :

coroutine  A(){ 
  a; 
  lifetime.start(a);
  if (b) 
    {
      ... // do something
      co_await
      use(a);
    }
     use(a);
    lifetime.end(a);
}

Since both user of a do not dominate each other, we just sink lifetime.start for both of them. This may cause buildCoroutineFrame to keep a as local variable. More importantly, it introduces the pattern in resume function such like:

lifetime.start(a);
use(a);
lifetime.start(a);
use(a);
lifetime.end(a);

which cause wrong optimization in later passes.

This patch rewrite sinkLifetimeStartMarkers which only sink lifetime.start markers when all of its user are only used inside one of suspended region.

TestPlan: cppcoro, check-llvm, https://godbolt.org/z/zVx_eB

[lxfind]

Thank you for looking into the fix!

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1578

If I is a BitCastInst, wouldn't it be used by both lifetime.start and lifetime.end intrinsics, and hence has more than one user?

[junparser]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1578

Since sinkLifetimeStartMarkers is called after rewriteMaterializableInstructions, so if BitCastInst both used by lifetime.start and lifetime.end, then I should not cross the suspend point. So I believe that we can use hasOneUse here

[lxfind]

Thank you!

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

static void sinkSpillUsesAfterCoroBegin(Function &F, const SpillInfo &Spills,
 
   Instruction *InsertPt = CoroBegin->getNextNode();
   for (Instruction *Inst : InsertionList)
     Inst->moveBefore(InsertPt);
 
   return;
 }
 
+/// For each local variable that all of its user are only used inside one of
+/// suspended region, we sink their lifetime.start markers to the place where
+/// after the suspend block. Doing so minimizes the lifetime of each variable,
+/// hence minimizing the amount of data we end up putting on the frame.
+static void sinkLifetimeStartMarkers(Function &F, coro::Shape &Shape,
+                                     SuspendCrossingInfo &Checker) {
+  DominatorTree DT(F);
+
+  // Collect all possible basic blocks which may dominate all uses of allocas.
+  SmallPtrSet<BasicBlock *, 4> DomSet;
+  DomSet.insert(&F.getEntryBlock());
+  for (auto *CSI : Shape.CoroSuspends) {
+    BasicBlock *SuspendBlock = CSI->getParent();
+    assert(isSuspendBlock(SuspendBlock) && SuspendBlock->getSingleSuccessor() &&
+           "should have split coro.suspend into its own block");
+    DomSet.insert(SuspendBlock->getSingleSuccessor());
+  }
+
+  for (Instruction &I : instructions(F)) {
+    if (!isa<AllocaInst>(&I))
+      continue;
+
+    for (BasicBlock *DomBB : DomSet) {
+      bool Valid = true;
+      SmallVector<Instruction *, 1> BCInsts;
+
+      auto isUsedByLifetimeStart = [&](Instruction *I) {
+        if (isa<BitCastInst>(I) && I->hasOneUse())
+          if (auto *IT = dyn_cast<IntrinsicInst>(I->user_back()))
+            return IT->getIntrinsicID() == Intrinsic::lifetime_start;
+        return false;
+      };
+
+      for (User *U : I.users()) {
+        Instruction *UI = cast<Instruction>(U);
+        // For all users except lifetime.start markers, if they are all
+        // dominated by one of the basic blocks and do not cross
+        // suspend points as well, then there is no need to spill the
+        // instruction.
+        if (!DT.dominates(DomBB, UI->getParent()) ||
+            Checker.isDefinitionAcrossSuspend(DomBB, U)) {
+          // Skip bitcast used by lifetime.start markers.
+          if (isUsedByLifetimeStart(UI)) {
+            BCInsts.push_back(UI);
+            continue;
+          }
+          Valid = false;
+          break;
+        }
+      }
+      // Sink lifetime.start markers to dominate block when they are
+      // only used outside the region.
+      if (Valid && BCInsts.size() != 0) {
+        auto *NewBitcast = BCInsts[0]->clone();
+        auto *NewLifetime = cast<Instruction>(BCInsts[0]->user_back())->clone();
+        NewLifetime->replaceUsesOfWith(BCInsts[0], NewBitcast);
+        NewBitcast->insertBefore(DomBB->getTerminator());
+        NewLifetime->insertBefore(DomBB->getTerminator());
+
+        // All the outsided lifetime.start markers are no longer necessary.
+        for (Instruction *S : BCInsts) {
+          S->user_back()->eraseFromParent();
+        }
+        break;
+      }
+    }
+  }
+}
+
 void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
   eliminateSwiftError(F, Shape);
 
   if (Shape.ABI == coro::ABI::Switch &&
       Shape.SwitchLowering.PromiseAlloca) {
     Shape.getSwitchCoroId()->clearPromise();
   }
 
if (Spills.empty())
       break;
 
     // Rewrite materializable instructions to be materialized at the use point.
     LLVM_DEBUG(dump("Materializations", Spills));
     rewriteMaterializableInstructions(Builder, Spills);
     Spills.clear();
   }
 
+  sinkLifetimeStartMarkers(F, Shape, Checker);
   // Collect lifetime.start info for each alloca.
   using LifetimeStart = SmallPtrSet<Instruction *, 2>;
   llvm::DenseMap<Instruction *, std::unique_ptr<LifetimeStart>> LifetimeMap;
   for (Instruction &I : instructions(F)) {
     auto *II = dyn_cast<IntrinsicInst>(&I);
     if (!II || II->getIntrinsicID() != Intrinsic::lifetime_start)
       continue;
 

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

if (!SI->isFinal() && simplifySuspendPoint(SI, Shape.CoroBegin)) {
       continue;
     }
     if (++I == N)
       break;
   }
   S.resize(N);
 }
 
-/// For every local variable that has lifetime intrinsics markers, we sink
-/// their lifetime.start marker to the places where the variable is being
-/// used for the first time. Doing so minimizes the lifetime of each variable,
-/// hence minimizing the amount of data we end up putting on the frame.
-static void sinkLifetimeStartMarkers(Function &F) {
-  DominatorTree Dom(F);
-  for (Instruction &I : instructions(F)) {
-    // We look for this particular pattern:
-    //   %tmpX = alloca %.., align ...
-    //   %0 = bitcast %...* %tmpX to i8*
-    //   call void @llvm.lifetime.start.p0i8(i64 ..., i8* nonnull %0) #2
-    if (!isa<AllocaInst>(&I))
-      continue;
-    // There can be multiple lifetime start markers for the same variable.
-    SmallPtrSet<IntrinsicInst *, 1> LifetimeStartInsts;
-    // SinkBarriers stores all instructions that use this local variable.
-    // When sinking the lifetime start intrinsics, we can never sink past
-    // these barriers.
-    SmallPtrSet<Instruction *, 4> SinkBarriers;
-    bool Valid = true;
-    auto AddSinkBarrier = [&](Instruction *I) {
-      // When adding a new barrier to SinkBarriers, we maintain the case
-      // that no instruction in SinkBarriers dominates another instruction.
-      SmallPtrSet<Instruction *, 1> ToRemove;
-      bool ShouldAdd = true;
-      for (Instruction *S : SinkBarriers) {
-        if (I == S || Dom.dominates(S, I)) {
-          ShouldAdd = false;
-          break;
-        } else if (Dom.dominates(I, S)) {
-          ToRemove.insert(S);
-        }
-      }
-      if (ShouldAdd) {
-        SinkBarriers.insert(I);
-        for (Instruction *R : ToRemove) {
-          SinkBarriers.erase(R);
-        }
-      }
-    };
-    for (User *U : I.users()) {
-      if (!isa<BitCastInst>(U))
-        continue;
-      for (User *CU : U->users()) {
-        // If we see any user of CastInst that's not lifetime start/end
-        // intrinsics, give up because it's too complex.
-        if (auto *CUI = dyn_cast<IntrinsicInst>(CU)) {
-          if (CUI->getIntrinsicID() == Intrinsic::lifetime_start)
-            LifetimeStartInsts.insert(CUI);
-          else if (CUI->getIntrinsicID() == Intrinsic::lifetime_end)
-            AddSinkBarrier(CUI);
-          else
-            Valid = false;
-        } else {
-          Valid = false;
-        }
-      }
-    }
-    if (!Valid || LifetimeStartInsts.empty())
-      continue;
-
-    for (User *U : I.users()) {
-      if (isa<BitCastInst>(U))
-        continue;
-      // Every user of the variable is also a sink barrier.
-      AddSinkBarrier(cast<Instruction>(U));
-    }
-
-    // For each sink barrier, we insert a lifetime start marker right
-    // before it.
-    for (Instruction *S : SinkBarriers) {
-      if (auto *IS = dyn_cast<IntrinsicInst>(S)) {
-        if (IS->getIntrinsicID() == Intrinsic::lifetime_end) {
-          // If we have a lifetime end marker in SinkBarriers, meaning it's
-          // not dominated by any other users, we can safely delete it.
-          IS->eraseFromParent();
-          continue;
-        }
-      }
-      // We find an existing lifetime.start marker that domintes the barrier,
-      // clone it and insert it right before the barrier. We cannot clone an
-      // arbitrary lifetime.start marker because we want to make sure the
-      // BitCast instruction referred in the marker also dominates the barrier.
-      for (const IntrinsicInst *LifetimeStart : LifetimeStartInsts) {
-        if (Dom.dominates(LifetimeStart, S)) {
-          LifetimeStart->clone()->insertBefore(S);
-          break;
-        }
-      }
-    }
-    // All the old lifetime.start markers are no longer necessary.
-    for (IntrinsicInst *S : LifetimeStartInsts) {
-      S->eraseFromParent();
-    }
-  }
-}
-
 static void splitSwitchCoroutine(Function &F, coro::Shape &Shape,
                                  SmallVectorImpl<Function *> &Clones) {
   assert(Shape.ABI == coro::ABI::Switch);
 
   createResumeEntryBlock(F, Shape);
   auto ResumeClone = createClone(F, ".resume", Shape,
                                  CoroCloner::Kind::SwitchResume);
   auto DestroyClone = createClone(F, ".destroy", Shape,
static coro::Shape splitCoroutine(Function &F,
   // up by uses in unreachable blocks, so remove them as a first pass.
   removeUnreachableBlocks(F);
 
   coro::Shape Shape(F);
   if (!Shape.CoroBegin)
     return Shape;
 
   simplifySuspendPoints(Shape);
-  sinkLifetimeStartMarkers(F);
   buildCoroutineFrame(F, Shape);
   replaceFrameSize(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);
   } else {

llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-01.ll

 ; Tests that coro-split will optimize the lifetime.start maker of each local variable,
-; sink them to the places closest to the actual use.
+; sink them to the places after the suspend block.
 ; RUN: opt < %s -coro-split -S | FileCheck %s
 ; RUN: opt < %s -passes=coro-split -S | FileCheck %s
 
 %"struct.std::coroutine_handle" = type { i8* }
 %"struct.std::coroutine_handle.0" = type { %"struct.std::coroutine_handle" }
 %"struct.lean_future<int>::Awaiter" = type { i32, %"struct.std::coroutine_handle.0" }
 
 declare i8* @malloc(i64)
await.ready:
   br label %exit
 exit:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 
 ; CHECK-LABEL: @a.resume(
 ; CHECK:         %testval = alloca i32, align 4
+; CHECK-NEXT:    %0 = bitcast i32* %testval to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 4, i8* %0)
 ; CHECK-NEXT:    getelementptr inbounds %a.Frame
 ; CHECK-NEXT:    getelementptr inbounds %"struct.lean_future<int>::Awaiter"
-; CHECK-NEXT:    %cast1 = bitcast i32* %testval to i8*
 ; CHECK-NEXT:    %val = load i32, i32* %Result
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 4, i8* %cast1)
 ; CHECK-NEXT:    %test = load i32, i32* %testval
 ; CHECK-NEXT:    call void @print(i32 %test)
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 4, i8* %cast1)
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 4, i8* %0)
 ; CHECK-NEXT:    call void @print(i32 %val)
 ; CHECK-NEXT:    ret void
 
 declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
 declare i1 @llvm.coro.alloc(token) #3
 declare noalias nonnull i8* @"\01??2@YAPEAX_K@Z"(i64) local_unnamed_addr
 declare i64 @llvm.coro.size.i64() #5
 declare i8* @llvm.coro.begin(token, i8* writeonly) #3

llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-02.ll

+; Tests that coro-split will optimize the lifetime.start maker of each local variable,
+; sink them to the places after the suspend block.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+; RUN: opt < %s -passes=coro-split -S | FileCheck %s
+
+%"struct.std::coroutine_handle" = type { i8* }
+%"struct.std::coroutine_handle.0" = type { %"struct.std::coroutine_handle" }
+%"struct.lean_future<int>::Awaiter" = type { i32, %"struct.std::coroutine_handle.0" }
+
+declare i1 @getcond()
+declare i8* @malloc(i64)
+declare void @print(i32)
+
+define void @a() "coroutine.presplit"="1" {
+entry:
+  %ref.tmp7 = alloca %"struct.lean_future<int>::Awaiter", align 8
+  %testval = alloca i32
+  %cast = bitcast i32* %testval to i8*
+  ; lifetime of %testval starts here, but not used until await.ready.
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* %cast)
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %alloc = call i8* @malloc(i64 16) #3
+  %vFrame = call noalias nonnull i8* @llvm.coro.begin(token %id, i8* %alloc)
+  %testcond = call i1 @getcond()
+  br i1 %testcond, label %if.suspend, label %else.direct
+
+if.suspend:
+  %save = call token @llvm.coro.save(i8* null)
+  %Result.i19 = getelementptr inbounds %"struct.lean_future<int>::Awaiter", %"struct.lean_future<int>::Awaiter"* %ref.tmp7, i64 0, i32 0
+  %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+  switch i8 %suspend, label %exit [
+    i8 0, label %await.ready
+    i8 1, label %exit
+  ]
+
+else.direct:
+  br label %after.await
+
+await.ready:
+  %StrayCoroSave = call token @llvm.coro.save(i8* null)
+  %val = load i32, i32* %Result.i19
+  %test = load i32, i32* %testval
+  call void @print(i32 %test)
+  call void @print(i32 %val)
+  br label %after.await
+
+after.await:
+  %test1 = load i32, i32* %testval
+  call void @print(i32 %test1)
+  call void @llvm.lifetime.end.p0i8(i64 4, i8*  %cast)
+  br label %exit
+
+exit:
+  call i1 @llvm.coro.end(i8* null, i1 false)
+  ret void
+}
+
+; CHECK-LABEL: @a.resume(
+; CHECK:    %[[VAL:testval.+]] = getelementptr inbounds %a.Frame
+; CHECK-NOT:     %testval = alloca i32, align 4
+; CHECK-NOT:     %[[CAST:.+]] = bitcast i32* %testval to i8*
+; CHECK-NOT:     call void @llvm.lifetime.start.p0i8(i64 4, i8* %[[CAST]])
+; CHECK:         %test = load i32, i32* %[[VAL]]
+; CHECK-NOT:     %test = load i32, i32* %testval
+
+declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
+declare i1 @llvm.coro.alloc(token) #3
+declare noalias nonnull i8* @"\01??2@YAPEAX_K@Z"(i64) local_unnamed_addr
+declare i64 @llvm.coro.size.i64() #5
+declare i8* @llvm.coro.begin(token, i8* writeonly) #3
+declare void @"\01?puts@@YAXZZ"(...)
+declare token @llvm.coro.save(i8*) #3
+declare i8* @llvm.coro.frame() #5
+declare i8 @llvm.coro.suspend(token, i1) #3
+declare void @"\01??3@YAXPEAX@Z"(i8*) local_unnamed_addr #10
+declare i8* @llvm.coro.free(token, i8* nocapture readonly) #2
+declare i1 @llvm.coro.end(i8*, i1) #3
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #4
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #4
+