[llvm] f0bfad2 - [Coroutines] Refactor sinkLifetimeStartMarkers

Jun Ma via llvm-commits llvm-commits at lists.llvm.org
Thu Jul 9 03:23:48 PDT 2020


Author: Jun Ma
Date: 2020-07-09T18:23:28+08:00
New Revision: f0bfad2ed9b4a0eec68b71c7df2ee588806788c2

URL: https://github.com/llvm/llvm-project/commit/f0bfad2ed9b4a0eec68b71c7df2ee588806788c2
DIFF: https://github.com/llvm/llvm-project/commit/f0bfad2ed9b4a0eec68b71c7df2ee588806788c2.diff

LOG: [Coroutines] Refactor sinkLifetimeStartMarkers

Differential Revision: https://reviews.llvm.org/D83379

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

Modified: 
    llvm/lib/Transforms/Coroutines/CoroFrame.cpp
    llvm/lib/Transforms/Coroutines/CoroSplit.cpp

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


################################################################################
diff  --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 2d398643a9c3..f55501a05d85 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -1548,6 +1548,75 @@ static void sinkSpillUsesAfterCoroBegin(Function &F, const SpillInfo &Spills,
   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);
 
@@ -1598,6 +1667,7 @@ void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
     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;

diff  --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 0841cebab51c..9c4392e7999b 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -1239,103 +1239,6 @@ static void simplifySuspendPoints(coro::Shape &Shape) {
   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);
@@ -1525,7 +1428,6 @@ static coro::Shape splitCoroutine(Function &F,
     return Shape;
 
   simplifySuspendPoints(Shape);
-  sinkLifetimeStartMarkers(F);
   buildCoroutineFrame(F, Shape);
   replaceFrameSize(Shape);
 

diff  --git a/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime.ll b/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-01.ll
similarity index 93%
rename from llvm/test/Transforms/Coroutines/coro-split-sink-lifetime.ll
rename to llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-01.ll
index 2d6b28a2baf8..9f9c1661138c 100644
--- a/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime.ll
+++ b/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-01.ll
@@ -1,5 +1,5 @@
 ; 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
 
@@ -43,14 +43,14 @@ exit:
 
 ; 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
 

diff  --git a/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-02.ll b/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-02.ll
new file mode 100644
index 000000000000..6695c22a1f09
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-02.ll
@@ -0,0 +1,80 @@
+; 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 at 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 at 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
+


        


More information about the llvm-commits mailing list