[llvm] 2e4c5d1 - CoroSplit: Fix coroutine splitting for retcon and retcon.once

Arnold Schwaighofer via llvm-commits llvm-commits at lists.llvm.org
Wed Jun 3 12:12:39 PDT 2020


Author: Arnold Schwaighofer
Date: 2020-06-03T12:10:58-07:00
New Revision: 2e4c5d1c483a986dbb3fc6486bdb2f0eb2adc8c8

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

LOG: CoroSplit: Fix coroutine splitting for retcon and retcon.once

Summary:
For retcon and retcon.once coroutines we assume that all uses of spills
can be sunk past coro.begin. This simplifies handling of instructions
that escape the address of an alloca.

The current implementation would have issues if the address of the
alloca is escaped before coro.begin. (It also has issues with casts before and
uses of those casts after the coro.begin instruction)

  %alloca_addr = alloca ...
  %escape  = ptrtoint %alloca_addr
  coro.begin
  store %escape to %alloca_addr

rdar://60272809

Subscribers: hiraditya, modocache, mgrang, llvm-commits

Tags: #llvm

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

Added: 
    llvm/test/Transforms/Coroutines/coro-retcon-frame.ll

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

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index a628c48cb5ca..2d398643a9c3 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -899,6 +899,23 @@ static Instruction *insertSpills(const SpillInfo &Spills, coro::Shape &Shape) {
     FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");      
   SpillBlock->splitBasicBlock(&SpillBlock->front(), "PostSpill");
   Shape.AllocaSpillBlock = SpillBlock;
+
+  // retcon and retcon.once lowering assumes all uses have been sunk.
+  if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce) {
+    // If we found any allocas, replace all of their remaining uses with Geps.
+    Builder.SetInsertPoint(&SpillBlock->front());
+    for (auto &P : Allocas) {
+      auto *G = GetFramePointer(P.second, P.first);
+
+      // We are not using ReplaceInstWithInst(P.first, cast<Instruction>(G))
+      // here, as we are changing location of the instruction.
+      G->takeName(P.first);
+      P.first->replaceAllUsesWith(G);
+      P.first->eraseFromParent();
+    }
+    return FramePtr;
+  }
+
   // If we found any alloca, replace all of their remaining uses with GEP
   // instructions. Because new dbg.declare have been created for these alloca,
   // we also delete the original dbg.declare and replace other uses with undef.
@@ -1482,6 +1499,55 @@ static void eliminateSwiftError(Function &F, coro::Shape &Shape) {
   }
 }
 
+/// retcon and retcon.once conventions assume that all spill uses can be sunk
+/// after the coro.begin intrinsic.
+static void sinkSpillUsesAfterCoroBegin(Function &F, const SpillInfo &Spills,
+                                        CoroBeginInst *CoroBegin) {
+  DominatorTree Dom(F);
+
+  SmallSetVector<Instruction *, 32> ToMove;
+  SmallVector<Instruction *, 32> Worklist;
+
+  // Collect all users that precede coro.begin.
+  for (auto const &Entry : Spills) {
+    auto *SpillDef = Entry.def();
+    for (User *U : SpillDef->users()) {
+      auto Inst = cast<Instruction>(U);
+      if (Inst->getParent() != CoroBegin->getParent() ||
+          Dom.dominates(CoroBegin, Inst))
+        continue;
+      if (ToMove.insert(Inst))
+        Worklist.push_back(Inst);
+    }
+  }
+  // Recursively collect users before coro.begin.
+  while (!Worklist.empty()) {
+    auto *Def = Worklist.back();
+    Worklist.pop_back();
+    for (User *U : Def->users()) {
+      auto Inst = cast<Instruction>(U);
+      if (Dom.dominates(CoroBegin, Inst))
+        continue;
+      if (ToMove.insert(Inst))
+        Worklist.push_back(Inst);
+    }
+  }
+
+  // Sort by dominance.
+  SmallVector<Instruction *, 64> InsertionList(ToMove.begin(), ToMove.end());
+  std::sort(InsertionList.begin(), InsertionList.end(),
+            [&Dom](Instruction *A, Instruction *B) -> bool {
+              // If a dominates b it should preceed (<) b.
+              return Dom.dominates(A, B);
+            });
+
+  Instruction *InsertPt = CoroBegin->getNextNode();
+  for (Instruction *Inst : InsertionList)
+    Inst->moveBefore(InsertPt);
+
+  return;
+}
+
 void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
   eliminateSwiftError(F, Shape);
 
@@ -1618,6 +1684,8 @@ void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
     }
   }
   LLVM_DEBUG(dump("Spills", Spills));
+  if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce)
+    sinkSpillUsesAfterCoroBegin(F, Spills, Shape.CoroBegin);
   Shape.FrameTy = buildFrameType(F, Shape, Spills);
   Shape.FramePtr = insertSpills(Spills, Shape);
   lowerLocalAllocas(LocalAllocas, DeadInstructions);

diff  --git a/llvm/test/Transforms/Coroutines/coro-retcon-frame.ll b/llvm/test/Transforms/Coroutines/coro-retcon-frame.ll
new file mode 100644
index 000000000000..c7ca8e3a0137
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-retcon-frame.ll
@@ -0,0 +1,63 @@
+; RUN: opt < %s -coro-split -S | FileCheck %s
+
+target datalayout = "p:64:64:64"
+
+declare void @prototype_f(i8*, i1)
+
+declare noalias i8* @allocate(i32 %size)
+declare void @deallocate(i8* %ptr)
+declare void @init(i64 *%ptr)
+declare void @use(i8* %ptr)
+declare void @use_addr_val(i64 %val, {i64, i64}*%addr)
+
+define { i8*, {i64, i64}* } @f(i8* %buffer) "coroutine.presplit"="1" {
+entry:
+  %tmp = alloca { i64, i64 }, align 8
+  %proj.1 = getelementptr inbounds { i64, i64 }, { i64, i64 }* %tmp, i64 0, i32 0
+  %proj.2 = getelementptr inbounds { i64, i64 }, { i64, i64 }* %tmp, i64 0, i32 1
+  store i64 0, i64* %proj.1, align 8
+  store i64 0, i64* %proj.2, align 8
+  %cast = bitcast { i64, i64 }* %tmp to i8*
+  %escape_addr = ptrtoint {i64, i64}* %tmp to i64
+  %id = call token @llvm.coro.id.retcon.once(i32 32, i32 8, i8* %buffer, i8* bitcast (void (i8*, i1)* @prototype_f to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
+  %proj.2.2 = getelementptr inbounds { i64, i64 }, { i64, i64 }* %tmp, i64 0, i32 1
+  call void @init(i64 * %proj.1)
+  call void @init(i64 * %proj.2.2)
+  call void @use_addr_val(i64 %escape_addr, {i64, i64}* %tmp)
+  %abort = call i1 (...) @llvm.coro.suspend.retcon.i1({i64, i64}* %tmp)
+  br i1 %abort, label %end, label %resume
+
+resume:
+  call void @use(i8* %cast)
+  br label %end
+
+end:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  unreachable
+}
+; Make sure we don't lose writes to the frame.
+; CHECK-LABEL: define { i8*, { i64, i64 }* } @f(i8* %buffer) {
+; CHECK:  [[FRAMEPTR:%.*]] = bitcast i8* %buffer to %f.Frame*
+; CHECK:  [[TMP:%.*]] = getelementptr inbounds %f.Frame, %f.Frame* [[FRAMEPTR]], i32 0, i32 0
+; CHECK:  [[PROJ1:%.*]] = getelementptr inbounds { i64, i64 }, { i64, i64 }* [[TMP]], i64 0, i32 0
+; CHECK:  [[PROJ2:%.*]] = getelementptr inbounds { i64, i64 }, { i64, i64 }* [[TMP]], i64 0, i32 1
+; CHECK:  store i64 0, i64* [[PROJ1]]
+; CHECK:  store i64 0, i64* [[PROJ2]]
+; CHECK:  [[ESCAPED_ADDR:%.*]] = ptrtoint { i64, i64 }* [[TMP]] to i64
+; CHECK:  call void @init(i64* [[PROJ1]])
+; CHECK:  call void @init(i64* [[PROJ2]])
+; CHECK:  call void @use_addr_val(i64 [[ESCAPED_ADDR]], { i64, i64 }* [[TMP]])
+
+; CHECK-LABEL: define internal void @f.resume.0(i8* {{.*}} %0, i1 %1) {
+; CHECK:  [[FRAMEPTR:%.*]] = bitcast i8* %0 to %f.Frame*
+; CHECK: resume:
+; CHECK:  [[TMP:%.*]] = getelementptr inbounds %f.Frame, %f.Frame* [[FRAMEPTR]], i32 0, i32 0
+; CHECK:  [[CAST:%.*]] = bitcast { i64, i64 }* [[TMP]] to i8*
+; CHECK:  call void @use(i8* [[CAST]])
+
+declare token @llvm.coro.id.retcon.once(i32, i32, i8*, i8*, i8*, i8*)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.suspend.retcon.i1(...)
+declare i1 @llvm.coro.end(i8*, i1)
+


        


More information about the llvm-commits mailing list