[llvm] [llvm/llvm-project][Coroutines] Move spill related methods to a Spill utils (PR #107884)

via llvm-commits llvm-commits at lists.llvm.org
Mon Sep 9 09:13:15 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-coroutines

Author: Tyler Nowicki (TylerNowicki)

<details>
<summary>Changes</summary>

* Move code related to spilling into SpillUtils to help cleanup CoroFrame

See RFC for more info: https://discourse.llvm.org/t/rfc-abi-objects-for-coroutines/81057

---

Patch is 58.58 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/107884.diff


4 Files Affected:

- (modified) llvm/lib/Transforms/Coroutines/CMakeLists.txt (+1) 
- (modified) llvm/lib/Transforms/Coroutines/CoroFrame.cpp (+30-611) 
- (added) llvm/lib/Transforms/Coroutines/SpillUtils.cpp (+631) 
- (added) llvm/lib/Transforms/Coroutines/SpillUtils.h (+60) 


``````````diff
diff --git a/llvm/lib/Transforms/Coroutines/CMakeLists.txt b/llvm/lib/Transforms/Coroutines/CMakeLists.txt
index ddc4caa82b0e98..c6508174a7f109 100644
--- a/llvm/lib/Transforms/Coroutines/CMakeLists.txt
+++ b/llvm/lib/Transforms/Coroutines/CMakeLists.txt
@@ -8,6 +8,7 @@ add_llvm_component_library(LLVMCoroutines
   CoroFrame.cpp
   CoroSplit.cpp
   SuspendCrossingInfo.cpp
+  SpillUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms/Coroutines
diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 7a631a646f3f8b..c6027a7de2f806 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -16,27 +16,22 @@
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
+#include "SpillUtils.h"
 #include "SuspendCrossingInfo.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallString.h"
-#include "llvm/Analysis/CFG.h"
-#include "llvm/Analysis/PtrUseVisitor.h"
 #include "llvm/Analysis/StackLifetime.h"
-#include "llvm/Config/llvm-config.h"
-#include "llvm/IR/CFG.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Module.h"
 #include "llvm/Support/Debug.h"
-#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/OptimizedStructLayout.h"
-#include "llvm/Support/circular_raw_ostream.h"
-#include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
@@ -175,24 +170,17 @@ template <> struct GraphTraits<RematGraph *> {
 namespace {
 class FrameTypeBuilder;
 // Mapping from the to-be-spilled value to all the users that need reload.
-using SpillInfo = SmallMapVector<Value *, SmallVector<Instruction *, 2>, 8>;
-struct AllocaInfo {
-  AllocaInst *Alloca;
-  DenseMap<Instruction *, std::optional<APInt>> Aliases;
-  bool MayWriteBeforeCoroBegin;
-  AllocaInfo(AllocaInst *Alloca,
-             DenseMap<Instruction *, std::optional<APInt>> Aliases,
-             bool MayWriteBeforeCoroBegin)
-      : Alloca(Alloca), Aliases(std::move(Aliases)),
-        MayWriteBeforeCoroBegin(MayWriteBeforeCoroBegin) {}
-};
 struct FrameDataInfo {
   // All the values (that are not allocas) that needs to be spilled to the
   // frame.
-  SpillInfo Spills;
+  coro::SpillInfo &Spills;
   // Allocas contains all values defined as allocas that need to live in the
   // frame.
-  SmallVector<AllocaInfo, 8> Allocas;
+  SmallVectorImpl<coro::AllocaInfo> &Allocas;
+
+  FrameDataInfo(coro::SpillInfo &Spills,
+                SmallVectorImpl<coro::AllocaInfo> &Allocas)
+      : Spills(Spills), Allocas(Allocas) {}
 
   SmallVector<Value *, 8> getAllDefs() const {
     SmallVector<Value *, 8> Defs;
@@ -271,7 +259,7 @@ struct FrameDataInfo {
 } // namespace
 
 #ifndef NDEBUG
-static void dumpSpills(StringRef Title, const SpillInfo &Spills) {
+static void dumpSpills(StringRef Title, const coro::SpillInfo &Spills) {
   dbgs() << "------------- " << Title << " --------------\n";
   for (const auto &E : Spills) {
     E.first->dump();
@@ -290,7 +278,7 @@ static void dumpRemats(
   }
 }
 
-static void dumpAllocas(const SmallVectorImpl<AllocaInfo> &Allocas) {
+static void dumpAllocas(const SmallVectorImpl<coro::AllocaInfo> &Allocas) {
   dbgs() << "------------- Allocas --------------\n";
   for (const auto &A : Allocas) {
     A.Alloca->dump();
@@ -538,7 +526,7 @@ void FrameTypeBuilder::addFieldForAllocas(const Function &F,
     return StackLifetimeAnalyzer.getLiveRange(AI1).overlaps(
         StackLifetimeAnalyzer.getLiveRange(AI2));
   };
-  auto GetAllocaSize = [&](const AllocaInfo &A) {
+  auto GetAllocaSize = [&](const coro::AllocaInfo &A) {
     std::optional<TypeSize> RetSize = A.Alloca->getAllocationSize(DL);
     assert(RetSize && "Variable Length Arrays (VLA) are not supported.\n");
     assert(!RetSize->isScalable() && "Scalable vectors are not yet supported");
@@ -1119,367 +1107,6 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
   return FrameTy;
 }
 
-// We use a pointer use visitor to track how an alloca is being used.
-// The goal is to be able to answer the following three questions:
-// 1. Should this alloca be allocated on the frame instead.
-// 2. Could the content of the alloca be modified prior to CoroBegn, which would
-// require copying the data from alloca to the frame after CoroBegin.
-// 3. Is there any alias created for this alloca prior to CoroBegin, but used
-// after CoroBegin. In that case, we will need to recreate the alias after
-// CoroBegin based off the frame. To answer question 1, we track two things:
-//   a. List of all BasicBlocks that use this alloca or any of the aliases of
-//   the alloca. In the end, we check if there exists any two basic blocks that
-//   cross suspension points. If so, this alloca must be put on the frame. b.
-//   Whether the alloca or any alias of the alloca is escaped at some point,
-//   either by storing the address somewhere, or the address is used in a
-//   function call that might capture. If it's ever escaped, this alloca must be
-//   put on the frame conservatively.
-// To answer quetion 2, we track through the variable MayWriteBeforeCoroBegin.
-// Whenever a potential write happens, either through a store instruction, a
-// function call or any of the memory intrinsics, we check whether this
-// instruction is prior to CoroBegin. To answer question 3, we track the offsets
-// of all aliases created for the alloca prior to CoroBegin but used after
-// CoroBegin. std::optional is used to be able to represent the case when the
-// offset is unknown (e.g. when you have a PHINode that takes in different
-// offset values). We cannot handle unknown offsets and will assert. This is the
-// potential issue left out. An ideal solution would likely require a
-// significant redesign.
-namespace {
-struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
-  using Base = PtrUseVisitor<AllocaUseVisitor>;
-  AllocaUseVisitor(const DataLayout &DL, const DominatorTree &DT,
-                   const coro::Shape &CoroShape,
-                   const SuspendCrossingInfo &Checker,
-                   bool ShouldUseLifetimeStartInfo)
-      : PtrUseVisitor(DL), DT(DT), CoroShape(CoroShape), Checker(Checker),
-        ShouldUseLifetimeStartInfo(ShouldUseLifetimeStartInfo) {
-    for (AnyCoroSuspendInst *SuspendInst : CoroShape.CoroSuspends)
-      CoroSuspendBBs.insert(SuspendInst->getParent());
-  }
-
-  void visit(Instruction &I) {
-    Users.insert(&I);
-    Base::visit(I);
-    // If the pointer is escaped prior to CoroBegin, we have to assume it would
-    // be written into before CoroBegin as well.
-    if (PI.isEscaped() &&
-        !DT.dominates(CoroShape.CoroBegin, PI.getEscapingInst())) {
-      MayWriteBeforeCoroBegin = true;
-    }
-  }
-  // We need to provide this overload as PtrUseVisitor uses a pointer based
-  // visiting function.
-  void visit(Instruction *I) { return visit(*I); }
-
-  void visitPHINode(PHINode &I) {
-    enqueueUsers(I);
-    handleAlias(I);
-  }
-
-  void visitSelectInst(SelectInst &I) {
-    enqueueUsers(I);
-    handleAlias(I);
-  }
-
-  void visitStoreInst(StoreInst &SI) {
-    // Regardless whether the alias of the alloca is the value operand or the
-    // pointer operand, we need to assume the alloca is been written.
-    handleMayWrite(SI);
-
-    if (SI.getValueOperand() != U->get())
-      return;
-
-    // We are storing the pointer into a memory location, potentially escaping.
-    // As an optimization, we try to detect simple cases where it doesn't
-    // actually escape, for example:
-    //   %ptr = alloca ..
-    //   %addr = alloca ..
-    //   store %ptr, %addr
-    //   %x = load %addr
-    //   ..
-    // If %addr is only used by loading from it, we could simply treat %x as
-    // another alias of %ptr, and not considering %ptr being escaped.
-    auto IsSimpleStoreThenLoad = [&]() {
-      auto *AI = dyn_cast<AllocaInst>(SI.getPointerOperand());
-      // If the memory location we are storing to is not an alloca, it
-      // could be an alias of some other memory locations, which is difficult
-      // to analyze.
-      if (!AI)
-        return false;
-      // StoreAliases contains aliases of the memory location stored into.
-      SmallVector<Instruction *, 4> StoreAliases = {AI};
-      while (!StoreAliases.empty()) {
-        Instruction *I = StoreAliases.pop_back_val();
-        for (User *U : I->users()) {
-          // If we are loading from the memory location, we are creating an
-          // alias of the original pointer.
-          if (auto *LI = dyn_cast<LoadInst>(U)) {
-            enqueueUsers(*LI);
-            handleAlias(*LI);
-            continue;
-          }
-          // If we are overriding the memory location, the pointer certainly
-          // won't escape.
-          if (auto *S = dyn_cast<StoreInst>(U))
-            if (S->getPointerOperand() == I)
-              continue;
-          if (auto *II = dyn_cast<IntrinsicInst>(U))
-            if (II->isLifetimeStartOrEnd())
-              continue;
-          // BitCastInst creats aliases of the memory location being stored
-          // into.
-          if (auto *BI = dyn_cast<BitCastInst>(U)) {
-            StoreAliases.push_back(BI);
-            continue;
-          }
-          return false;
-        }
-      }
-
-      return true;
-    };
-
-    if (!IsSimpleStoreThenLoad())
-      PI.setEscaped(&SI);
-  }
-
-  // All mem intrinsics modify the data.
-  void visitMemIntrinsic(MemIntrinsic &MI) { handleMayWrite(MI); }
-
-  void visitBitCastInst(BitCastInst &BC) {
-    Base::visitBitCastInst(BC);
-    handleAlias(BC);
-  }
-
-  void visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
-    Base::visitAddrSpaceCastInst(ASC);
-    handleAlias(ASC);
-  }
-
-  void visitGetElementPtrInst(GetElementPtrInst &GEPI) {
-    // The base visitor will adjust Offset accordingly.
-    Base::visitGetElementPtrInst(GEPI);
-    handleAlias(GEPI);
-  }
-
-  void visitIntrinsicInst(IntrinsicInst &II) {
-    // When we found the lifetime markers refers to a
-    // subrange of the original alloca, ignore the lifetime
-    // markers to avoid misleading the analysis.
-    if (!IsOffsetKnown || !Offset.isZero())
-      return Base::visitIntrinsicInst(II);
-    switch (II.getIntrinsicID()) {
-    default:
-      return Base::visitIntrinsicInst(II);
-    case Intrinsic::lifetime_start:
-      LifetimeStarts.insert(&II);
-      LifetimeStartBBs.push_back(II.getParent());
-      break;
-    case Intrinsic::lifetime_end:
-      LifetimeEndBBs.insert(II.getParent());
-      break;
-    }
-  }
-
-  void visitCallBase(CallBase &CB) {
-    for (unsigned Op = 0, OpCount = CB.arg_size(); Op < OpCount; ++Op)
-      if (U->get() == CB.getArgOperand(Op) && !CB.doesNotCapture(Op))
-        PI.setEscaped(&CB);
-    handleMayWrite(CB);
-  }
-
-  bool getShouldLiveOnFrame() const {
-    if (!ShouldLiveOnFrame)
-      ShouldLiveOnFrame = computeShouldLiveOnFrame();
-    return *ShouldLiveOnFrame;
-  }
-
-  bool getMayWriteBeforeCoroBegin() const { return MayWriteBeforeCoroBegin; }
-
-  DenseMap<Instruction *, std::optional<APInt>> getAliasesCopy() const {
-    assert(getShouldLiveOnFrame() && "This method should only be called if the "
-                                     "alloca needs to live on the frame.");
-    for (const auto &P : AliasOffetMap)
-      if (!P.second)
-        report_fatal_error("Unable to handle an alias with unknown offset "
-                           "created before CoroBegin.");
-    return AliasOffetMap;
-  }
-
-private:
-  const DominatorTree &DT;
-  const coro::Shape &CoroShape;
-  const SuspendCrossingInfo &Checker;
-  // All alias to the original AllocaInst, created before CoroBegin and used
-  // after CoroBegin. Each entry contains the instruction and the offset in the
-  // original Alloca. They need to be recreated after CoroBegin off the frame.
-  DenseMap<Instruction *, std::optional<APInt>> AliasOffetMap{};
-  SmallPtrSet<Instruction *, 4> Users{};
-  SmallPtrSet<IntrinsicInst *, 2> LifetimeStarts{};
-  SmallVector<BasicBlock *> LifetimeStartBBs{};
-  SmallPtrSet<BasicBlock *, 2> LifetimeEndBBs{};
-  SmallPtrSet<const BasicBlock *, 2> CoroSuspendBBs{};
-  bool MayWriteBeforeCoroBegin{false};
-  bool ShouldUseLifetimeStartInfo{true};
-
-  mutable std::optional<bool> ShouldLiveOnFrame{};
-
-  bool computeShouldLiveOnFrame() const {
-    // If lifetime information is available, we check it first since it's
-    // more precise. We look at every pair of lifetime.start intrinsic and
-    // every basic block that uses the pointer to see if they cross suspension
-    // points. The uses cover both direct uses as well as indirect uses.
-    if (ShouldUseLifetimeStartInfo && !LifetimeStarts.empty()) {
-      // If there is no explicit lifetime.end, then assume the address can
-      // cross suspension points.
-      if (LifetimeEndBBs.empty())
-        return true;
-
-      // If there is a path from a lifetime.start to a suspend without a
-      // corresponding lifetime.end, then the alloca's lifetime persists
-      // beyond that suspension point and the alloca must go on the frame.
-      llvm::SmallVector<BasicBlock *> Worklist(LifetimeStartBBs);
-      if (isManyPotentiallyReachableFromMany(Worklist, CoroSuspendBBs,
-                                             &LifetimeEndBBs, &DT))
-        return true;
-
-      // Addresses are guaranteed to be identical after every lifetime.start so
-      // we cannot use the local stack if the address escaped and there is a
-      // suspend point between lifetime markers. This should also cover the
-      // case of a single lifetime.start intrinsic in a loop with suspend point.
-      if (PI.isEscaped()) {
-        for (auto *A : LifetimeStarts) {
-          for (auto *B : LifetimeStarts) {
-            if (Checker.hasPathOrLoopCrossingSuspendPoint(A->getParent(),
-                                                          B->getParent()))
-              return true;
-          }
-        }
-      }
-      return false;
-    }
-    // FIXME: Ideally the isEscaped check should come at the beginning.
-    // However there are a few loose ends that need to be fixed first before
-    // we can do that. We need to make sure we are not over-conservative, so
-    // that the data accessed in-between await_suspend and symmetric transfer
-    // is always put on the stack, and also data accessed after coro.end is
-    // always put on the stack (esp the return object). To fix that, we need
-    // to:
-    //  1) Potentially treat sret as nocapture in calls
-    //  2) Special handle the return object and put it on the stack
-    //  3) Utilize lifetime.end intrinsic
-    if (PI.isEscaped())
-      return true;
-
-    for (auto *U1 : Users)
-      for (auto *U2 : Users)
-        if (Checker.isDefinitionAcrossSuspend(*U1, U2))
-          return true;
-
-    return false;
-  }
-
-  void handleMayWrite(const Instruction &I) {
-    if (!DT.dominates(CoroShape.CoroBegin, &I))
-      MayWriteBeforeCoroBegin = true;
-  }
-
-  bool usedAfterCoroBegin(Instruction &I) {
-    for (auto &U : I.uses())
-      if (DT.dominates(CoroShape.CoroBegin, U))
-        return true;
-    return false;
-  }
-
-  void handleAlias(Instruction &I) {
-    // We track all aliases created prior to CoroBegin but used after.
-    // These aliases may need to be recreated after CoroBegin if the alloca
-    // need to live on the frame.
-    if (DT.dominates(CoroShape.CoroBegin, &I) || !usedAfterCoroBegin(I))
-      return;
-
-    if (!IsOffsetKnown) {
-      AliasOffetMap[&I].reset();
-    } else {
-      auto Itr = AliasOffetMap.find(&I);
-      if (Itr == AliasOffetMap.end()) {
-        AliasOffetMap[&I] = Offset;
-      } else if (Itr->second && *Itr->second != Offset) {
-        // If we have seen two different possible values for this alias, we set
-        // it to empty.
-        AliasOffetMap[&I].reset();
-      }
-    }
-  }
-};
-} // namespace
-
-// We need to make room to insert a spill after initial PHIs, but before
-// catchswitch instruction. Placing it before violates the requirement that
-// catchswitch, like all other EHPads must be the first nonPHI in a block.
-//
-// Split away catchswitch into a separate block and insert in its place:
-//
-//   cleanuppad <InsertPt> cleanupret.
-//
-// cleanupret instruction will act as an insert point for the spill.
-static Instruction *splitBeforeCatchSwitch(CatchSwitchInst *CatchSwitch) {
-  BasicBlock *CurrentBlock = CatchSwitch->getParent();
-  BasicBlock *NewBlock = CurrentBlock->splitBasicBlock(CatchSwitch);
-  CurrentBlock->getTerminator()->eraseFromParent();
-
-  auto *CleanupPad =
-      CleanupPadInst::Create(CatchSwitch->getParentPad(), {}, "", CurrentBlock);
-  auto *CleanupRet =
-      CleanupReturnInst::Create(CleanupPad, NewBlock, CurrentBlock);
-  return CleanupRet;
-}
-
-static BasicBlock::iterator getSpillInsertionPt(const coro::Shape &Shape,
-                                                Value *Def,
-                                                const DominatorTree &DT) {
-  BasicBlock::iterator InsertPt;
-  if (auto *Arg = dyn_cast<Argument>(Def)) {
-    // For arguments, we will place the store instruction right after
-    // the coroutine frame pointer instruction, i.e. coro.begin.
-    InsertPt = Shape.getInsertPtAfterFramePtr();
-
-    // If we're spilling an Argument, make sure we clear 'nocapture'
-    // from the coroutine function.
-    Arg->getParent()->removeParamAttr(Arg->getArgNo(), Attribute::NoCapture);
-  } else if (auto *CSI = dyn_cast<AnyCoroSuspendInst>(Def)) {
-    // Don't spill immediately after a suspend; splitting assumes
-    // that the suspend will be followed by a branch.
-    InsertPt = CSI->getParent()->getSingleSuccessor()->getFirstNonPHIIt();
-  } else {
-    auto *I = cast<Instruction>(Def);
-    if (!DT.dominates(Shape.CoroBegin, I)) {
-      // If it is not dominated by CoroBegin, then spill should be
-      // inserted immediately after CoroFrame is computed.
-      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.
-      auto *NewBB = SplitEdge(II->getParent(), II->getNormalDest());
-      InsertPt = NewBB->getTerminator()->getIterator();
-    } else if (isa<PHINode>(I)) {
-      // Skip the PHINodes and EH pads instructions.
-      BasicBlock *DefBlock = I->getParent();
-      if (auto *CSI = dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
-        InsertPt = splitBeforeCatchSwitch(CSI)->getIterator();
-      else
-        InsertPt = DefBlock->getFirstInsertionPt();
-    } else {
-      assert(!I->isTerminator() && "unexpected terminator");
-      // For all other values, the spill is placed immediately after
-      // the definition.
-      InsertPt = I->getNextNode()->getIterator();
-    }
-  }
-
-  return InsertPt;
-}
-
 // Replace all alloca and SSA values that are accessed across suspend points
 // with GetElementPointer from coroutine frame + loads and stores. Create an
 // AllocaSpillBB that will become the new entry block for the resume parts of
@@ -1564,7 +1191,7 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
     auto SpillAlignment = Align(FrameData.getAlign(Def));
     // Create a store instruction storing the value into the
     // coroutine frame.
-    BasicBlock::iterator InsertPt = getSpillInsertionPt(Shape, Def, DT);
+    BasicBlock::iterator InsertPt = coro::getSpillInsertionPt(Shape, Def, DT);
 
     Type *ByValTy = nullptr;
     if (auto *Arg = dyn_cast<Argument>(Def)) {
@@ -2015,13 +1642,6 @@ bool coro::defaultMaterializable(Instruction &V) {
           isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V));
 }
 
-// Check for structural coroutine intrinsics that should not be spilled into
-// the coroutine frame.
-static bool isCoroutineStructureIntrinsic(Instruction &I) {
-  return isa<CoroIdInst>(&I) ||...
[truncated]

``````````

</details>


https://github.com/llvm/llvm-project/pull/107884


More information about the llvm-commits mailing list