[llvm] 9bca4ea - [Coroutines] Allow FramePtr to be an Argument

Nikita Popov via llvm-commits llvm-commits at lists.llvm.org
Mon Mar 7 01:59:03 PST 2022


Author: Nikita Popov
Date: 2022-03-07T10:58:56+01:00
New Revision: 9bca4ea364dd6fd31ad1ca827631760c21ffd93f

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

LOG: [Coroutines] Allow FramePtr to be an Argument

With opaque pointers, after splitRetconCoroutine() the FramePtr
may be an Argument rather than an Instruction. With typed pointers,
this currently doesn't happen because the FramePtr would be a
bitcast instruction.

Fix this by making FramePtr a Value and adding a helper for the
"after FramePtr" insertion point, which would be the start of the
function in the Argument case.

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

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

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

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 53f2756f8c850..372238dcb9ea0 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -1079,7 +1079,7 @@ static void buildFrameDebugInfo(Function &F, coro::Shape &Shape,
 
   DBuilder.insertDeclare(Shape.FramePtr, FrameDIVar,
                          DBuilder.createExpression(), DILoc,
-                         Shape.FramePtr->getNextNode());
+                         Shape.getInsertPtAfterFramePtr());
 }
 
 // Build a struct that will keep state for an active coroutine.
@@ -1523,7 +1523,7 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
   LLVMContext &C = CB->getContext();
   IRBuilder<> Builder(C);
   StructType *FrameTy = Shape.FrameTy;
-  Instruction *FramePtr = Shape.FramePtr;
+  Value *FramePtr = Shape.FramePtr;
   DominatorTree DT(*CB->getFunction());
   SmallDenseMap<llvm::Value *, llvm::AllocaInst *, 4> DbgPtrAllocaCache;
 
@@ -1576,7 +1576,7 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
       // For arguments, we will place the store instruction right after
       // the coroutine frame pointer instruction, i.e. bitcast of
       // coro.begin from i8* to %f.frame*.
-      InsertPt = FramePtr->getNextNode();
+      InsertPt = Shape.getInsertPtAfterFramePtr();
 
       // If we're spilling an Argument, make sure we clear 'nocapture'
       // from the coroutine function.
@@ -1593,7 +1593,7 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
       if (!DT.dominates(CB, I)) {
         // If it is not dominated by CoroBegin, then spill should be
         // inserted immediately after CoroFrame is computed.
-        InsertPt = FramePtr->getNextNode();
+        InsertPt = Shape.getInsertPtAfterFramePtr();
       } else if (auto *II = dyn_cast<InvokeInst>(I)) {
         // If we are spilling the result of the invoke instruction, split
         // the normal edge and insert the spill in the new block.
@@ -1686,10 +1686,10 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
     }
   }
 
-  BasicBlock *FramePtrBB = FramePtr->getParent();
+  BasicBlock *FramePtrBB = Shape.getInsertPtAfterFramePtr()->getParent();
 
-  auto SpillBlock =
-      FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");
+  auto SpillBlock = FramePtrBB->splitBasicBlock(
+      Shape.getInsertPtAfterFramePtr(), "AllocaSpillBB");
   SpillBlock->splitBasicBlock(&SpillBlock->front(), "PostSpill");
   Shape.AllocaSpillBlock = SpillBlock;
 
@@ -1739,7 +1739,7 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
     for (Instruction *I : UsersToUpdate)
       I->replaceUsesOfWith(Alloca, G);
   }
-  Builder.SetInsertPoint(FramePtr->getNextNode());
+  Builder.SetInsertPoint(Shape.getInsertPtAfterFramePtr());
   for (const auto &A : FrameData.Allocas) {
     AllocaInst *Alloca = A.Alloca;
     if (A.MayWriteBeforeCoroBegin) {

diff  --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
index 9a17068df3a9b..797c8987e9850 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInternal.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -128,7 +128,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
   StructType *FrameTy;
   Align FrameAlign;
   uint64_t FrameSize;
-  Instruction *FramePtr;
+  Value *FramePtr;
   BasicBlock *AllocaSpillBlock;
 
   /// This would only be true if optimization are enabled.
@@ -267,6 +267,12 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
     return nullptr;
   }
 
+  Instruction *getInsertPtAfterFramePtr() const {
+    if (auto *I = dyn_cast<Instruction>(FramePtr))
+      return I->getNextNode();
+    return &cast<Argument>(FramePtr)->getParent()->getEntryBlock().front();
+  }
+
   /// Allocate memory according to the rules of the active lowering.
   ///
   /// \param CG - if non-null, will be updated for the new call

diff  --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 57547f5a936e7..d47b0619fb312 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -1152,7 +1152,8 @@ static void updateCoroFrame(coro::Shape &Shape, Function *ResumeFn,
                             Function *DestroyFn, Function *CleanupFn) {
   assert(Shape.ABI == coro::ABI::Switch);
 
-  IRBuilder<> Builder(Shape.FramePtr->getNextNode());
+  IRBuilder<> Builder(Shape.getInsertPtAfterFramePtr());
+
   auto *ResumeAddr = Builder.CreateStructGEP(
       Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Resume,
       "resume.addr");
@@ -1663,7 +1664,7 @@ static void splitAsyncCoroutine(Function &F, coro::Shape &Shape,
   // Map all uses of llvm.coro.begin to the allocated frame pointer.
   {
     // Make sure we don't invalidate Shape.FramePtr.
-    TrackingVH<Instruction> Handle(Shape.FramePtr);
+    TrackingVH<Value> Handle(Shape.FramePtr);
     Shape.CoroBegin->replaceAllUsesWith(FramePtr);
     Shape.FramePtr = Handle.getValPtr();
   }
@@ -1775,7 +1776,7 @@ static void splitRetconCoroutine(Function &F, coro::Shape &Shape,
   // Map all uses of llvm.coro.begin to the allocated frame pointer.
   {
     // Make sure we don't invalidate Shape.FramePtr.
-    TrackingVH<Instruction> Handle(Shape.FramePtr);
+    TrackingVH<Value> Handle(Shape.FramePtr);
     Shape.CoroBegin->replaceAllUsesWith(RawFramePtr);
     Shape.FramePtr = Handle.getValPtr();
   }

diff  --git a/llvm/test/Transforms/Coroutines/coro-retcon-opaque-ptr.ll b/llvm/test/Transforms/Coroutines/coro-retcon-opaque-ptr.ll
new file mode 100644
index 0000000000000..bdbfd1f17ef03
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-retcon-opaque-ptr.ll
@@ -0,0 +1,98 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -enable-coroutines -passes='default<O2>' -opaque-pointers -S | FileCheck %s
+
+; Same test as coro-retcon.ll, but with opaque pointers enabled.
+
+define ptr @f(ptr %buffer, i32 %n) {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  coro.return:
+; CHECK-NEXT:    store i32 [[N:%.*]], ptr [[BUFFER:%.*]], align 4
+; CHECK-NEXT:    tail call void @print(i32 [[N]])
+; CHECK-NEXT:    ret ptr @f.resume.0
+;
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 8, i32 4, ptr %buffer, ptr @prototype, ptr @allocate, ptr @deallocate)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+  br label %loop
+
+loop:                                             ; preds = %resume, %entry
+  %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ]
+  call void @print(i32 %n.val)
+  %unwind0 = call i1 (...) @llvm.coro.suspend.retcon.i1()
+  br i1 %unwind0, label %cleanup, label %resume
+
+resume:                                           ; preds = %loop
+  %inc = add i32 %n.val, 1
+  br label %loop
+
+cleanup:                                          ; preds = %loop
+  %0 = call i1 @llvm.coro.end(ptr %hdl, i1 false)
+  unreachable
+}
+
+define i32 @main() {
+; CHECK-LABEL: @main(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = alloca [8 x i8], align 4
+; CHECK-NEXT:    store i32 4, ptr [[TMP0]], align 4
+; CHECK-NEXT:    call void @print(i32 4)
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata [[META0:![0-9]+]])
+; CHECK-NEXT:    [[N_VAL_RELOAD_I:%.*]] = load i32, ptr [[TMP0]], align 4, !alias.scope !0
+; CHECK-NEXT:    [[INC_I:%.*]] = add i32 [[N_VAL_RELOAD_I]], 1
+; CHECK-NEXT:    store i32 [[INC_I]], ptr [[TMP0]], align 4, !alias.scope !0
+; CHECK-NEXT:    call void @print(i32 [[INC_I]]), !noalias !0
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata [[META3:![0-9]+]])
+; CHECK-NEXT:    [[N_VAL_RELOAD_I1:%.*]] = load i32, ptr [[TMP0]], align 4, !alias.scope !3
+; CHECK-NEXT:    [[INC_I2:%.*]] = add i32 [[N_VAL_RELOAD_I1]], 1
+; CHECK-NEXT:    call void @print(i32 [[INC_I2]]), !noalias !3
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %0 = alloca [8 x i8], align 4
+  %prepare = call ptr @llvm.coro.prepare.retcon(ptr @f)
+  %cont0 = call ptr %prepare(ptr %0, i32 4)
+  %cont1 = call ptr %cont0(ptr %0, i1 zeroext false)
+  %cont2 = call ptr %cont1(ptr %0, i1 zeroext false)
+  %1 = call ptr %cont2(ptr %0, i1 zeroext true)
+  ret i32 0
+}
+
+define hidden { ptr, ptr } @g(ptr %buffer, ptr %ptr) {
+; CHECK-LABEL: @g(
+; CHECK-NEXT:  coro.return:
+; CHECK-NEXT:    [[TMP0:%.*]] = tail call ptr @allocate(i32 8)
+; CHECK-NEXT:    store ptr [[TMP0]], ptr [[BUFFER:%.*]], align 8
+; CHECK-NEXT:    store ptr [[PTR:%.*]], ptr [[TMP0]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = insertvalue { ptr, ptr } { ptr @g.resume.0, ptr undef }, ptr [[PTR]], 1
+; CHECK-NEXT:    ret { ptr, ptr } [[TMP1]]
+;
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 8, i32 4, ptr %buffer, ptr @g_prototype, ptr @allocate, ptr @deallocate)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+  br label %loop
+
+loop:                                             ; preds = %resume, %entry
+  %unwind0 = call i1 (...) @llvm.coro.suspend.retcon.i1(ptr %ptr)
+  br i1 %unwind0, label %cleanup, label %resume
+
+resume:                                           ; preds = %loop
+  br label %loop
+
+cleanup:                                          ; preds = %loop
+  %0 = call i1 @llvm.coro.end(ptr %hdl, i1 false)
+  unreachable
+}
+
+declare token @llvm.coro.id.retcon(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)
+declare i8* @llvm.coro.prepare.retcon(i8*)
+
+declare i8* @prototype(i8*, i1 zeroext)
+declare {i8*,i8*} @g_prototype(i8*, i1 zeroext)
+
+declare noalias i8* @allocate(i32 %size)
+declare void @deallocate(i8* %ptr)
+
+declare void @print(i32)


        


More information about the llvm-commits mailing list