[Coroutines] Move materialization code into its own utils (#108240)

* Move materialization out of CoroFrame to MaterializationUtils.h
* Move spill related utilities that were used by materialization to
SpillUtils
* Move isSuspendBlock (needed by materialization) to CoroInternal

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

Author: Tyler Nowicki

llvm/lib/Transforms/Coroutines/CMakeLists.txt

@@ -9,6 +9,7 @@ add_llvm_component_library(LLVMCoroutines
   CoroSplit.cpp
   SuspendCrossingInfo.cpp
   SpillUtils.cpp
+  MaterializationUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms/Coroutines

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

@@ -16,10 +16,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
+#include "MaterializationUtils.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/StackLifetime.h"
@@ -36,135 +36,12 @@
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include <algorithm>
-#include <deque>
 #include <optional>
 
 using namespace llvm;
 
 extern cl::opt<bool> UseNewDbgInfoFormat;
 
-// The "coro-suspend-crossing" flag is very noisy. There is another debug type,
-// "coro-frame", which results in leaner debug spew.
-#define DEBUG_TYPE "coro-suspend-crossing"
-
-namespace {
-
-// RematGraph is used to construct a DAG for rematerializable instructions
-// When the constructor is invoked with a candidate instruction (which is
-// materializable) it builds a DAG of materializable instructions from that
-// point.
-// Typically, for each instruction identified as re-materializable across a
-// suspend point, a RematGraph will be created.
-struct RematGraph {
-  // Each RematNode in the graph contains the edges to instructions providing
-  // operands in the current node.
-  struct RematNode {
-    Instruction *Node;
-    SmallVector<RematNode *> Operands;
-    RematNode() = default;
-    RematNode(Instruction *V) : Node(V) {}
-  };
-
-  RematNode *EntryNode;
-  using RematNodeMap =
-      SmallMapVector<Instruction *, std::unique_ptr<RematNode>, 8>;
-  RematNodeMap Remats;
-  const std::function<bool(Instruction &)> &MaterializableCallback;
-  SuspendCrossingInfo &Checker;
-
-  RematGraph(const std::function<bool(Instruction &)> &MaterializableCallback,
-             Instruction *I, SuspendCrossingInfo &Checker)
-      : MaterializableCallback(MaterializableCallback), Checker(Checker) {
-    std::unique_ptr<RematNode> FirstNode = std::make_unique<RematNode>(I);
-    EntryNode = FirstNode.get();
-    std::deque<std::unique_ptr<RematNode>> WorkList;
-    addNode(std::move(FirstNode), WorkList, cast<User>(I));
-    while (WorkList.size()) {
-      std::unique_ptr<RematNode> N = std::move(WorkList.front());
-      WorkList.pop_front();
-      addNode(std::move(N), WorkList, cast<User>(I));
-    }
-  }
-
-  void addNode(std::unique_ptr<RematNode> NUPtr,
-               std::deque<std::unique_ptr<RematNode>> &WorkList,
-               User *FirstUse) {
-    RematNode *N = NUPtr.get();
-    if (Remats.count(N->Node))
-      return;
-
-    // We haven't see this node yet - add to the list
-    Remats[N->Node] = std::move(NUPtr);
-    for (auto &Def : N->Node->operands()) {
-      Instruction *D = dyn_cast<Instruction>(Def.get());
-      if (!D || !MaterializableCallback(*D) ||
-          !Checker.isDefinitionAcrossSuspend(*D, FirstUse))
-        continue;
-
-      if (Remats.count(D)) {
-        // Already have this in the graph
-        N->Operands.push_back(Remats[D].get());
-        continue;
-      }
-
-      bool NoMatch = true;
-      for (auto &I : WorkList) {
-        if (I->Node == D) {
-          NoMatch = false;
-          N->Operands.push_back(I.get());
-          break;
-        }
-      }
-      if (NoMatch) {
-        // Create a new node
-        std::unique_ptr<RematNode> ChildNode = std::make_unique<RematNode>(D);
-        N->Operands.push_back(ChildNode.get());
-        WorkList.push_back(std::move(ChildNode));
-      }
-    }
-  }
-
-#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
-  static std::string getBasicBlockLabel(const BasicBlock *BB) {
-    if (BB->hasName())
-      return BB->getName().str();
-
-    std::string S;
-    raw_string_ostream OS(S);
-    BB->printAsOperand(OS, false);
-    return OS.str().substr(1);
-  }
-
-  void dump() const {
-    dbgs() << "Entry (";
-    dbgs() << getBasicBlockLabel(EntryNode->Node->getParent());
-    dbgs() << ") : " << *EntryNode->Node << "\n";
-    for (auto &E : Remats) {
-      dbgs() << *(E.first) << "\n";
-      for (RematNode *U : E.second->Operands)
-        dbgs() << "  " << *U->Node << "\n";
-    }
-  }
-#endif
-};
-} // end anonymous namespace
-
-namespace llvm {
-
-template <> struct GraphTraits<RematGraph *> {
-  using NodeRef = RematGraph::RematNode *;
-  using ChildIteratorType = RematGraph::RematNode **;
-
-  static NodeRef getEntryNode(RematGraph *G) { return G->EntryNode; }
-  static ChildIteratorType child_begin(NodeRef N) {
-    return N->Operands.begin();
-  }
-  static ChildIteratorType child_end(NodeRef N) { return N->Operands.end(); }
-};
-
-} // end namespace llvm
-
-#undef DEBUG_TYPE // "coro-suspend-crossing"
 #define DEBUG_TYPE "coro-frame"
 
 namespace {
@@ -268,15 +145,6 @@ static void dumpSpills(StringRef Title, const coro::SpillInfo &Spills) {
       I->dump();
   }
 }
-static void dumpRemats(
-    StringRef Title,
-    const SmallMapVector<Instruction *, std::unique_ptr<RematGraph>, 8> &RM) {
-  dbgs() << "------------- " << Title << "--------------\n";
-  for (const auto &E : RM) {
-    E.second->dump();
-    dbgs() << "--\n";
-  }
-}
 
 static void dumpAllocas(const SmallVectorImpl<coro::AllocaInfo> &Allocas) {
   dbgs() << "------------- Allocas --------------\n";
@@ -1634,93 +1502,6 @@ static void rewritePHIs(Function &F) {
     rewritePHIs(*BB);
 }
 
-/// Default materializable callback
-// Check for instructions that we can recreate on resume as opposed to spill
-// the result into a coroutine frame.
-bool coro::defaultMaterializable(Instruction &V) {
-  return (isa<CastInst>(&V) || isa<GetElementPtrInst>(&V) ||
-          isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V));
-}
-
-// For each instruction identified as materializable across the suspend point,
-// and its associated DAG of other rematerializable instructions,
-// recreate the DAG of instructions after the suspend point.
-static void rewriteMaterializableInstructions(
-    const SmallMapVector<Instruction *, std::unique_ptr<RematGraph>, 8>
-        &AllRemats) {
-  // This has to be done in 2 phases
-  // Do the remats and record the required defs to be replaced in the
-  // original use instructions
-  // Once all the remats are complete, replace the uses in the final
-  // instructions with the new defs
-  typedef struct {
-    Instruction *Use;
-    Instruction *Def;
-    Instruction *Remat;
-  } ProcessNode;
-
-  SmallVector<ProcessNode> FinalInstructionsToProcess;
-
-  for (const auto &E : AllRemats) {
-    Instruction *Use = E.first;
-    Instruction *CurrentMaterialization = nullptr;
-    RematGraph *RG = E.second.get();
-    ReversePostOrderTraversal<RematGraph *> RPOT(RG);
-    SmallVector<Instruction *> InstructionsToProcess;
-
-    // If the target use is actually a suspend instruction then we have to
-    // insert the remats into the end of the predecessor (there should only be
-    // one). This is so that suspend blocks always have the suspend instruction
-    // as the first instruction.
-    auto InsertPoint = &*Use->getParent()->getFirstInsertionPt();
-    if (isa<AnyCoroSuspendInst>(Use)) {
-      BasicBlock *SuspendPredecessorBlock =
-          Use->getParent()->getSinglePredecessor();
-      assert(SuspendPredecessorBlock && "malformed coro suspend instruction");
-      InsertPoint = SuspendPredecessorBlock->getTerminator();
-    }
-
-    // Note: skip the first instruction as this is the actual use that we're
-    // rematerializing everything for.
-    auto I = RPOT.begin();
-    ++I;
-    for (; I != RPOT.end(); ++I) {
-      Instruction *D = (*I)->Node;
-      CurrentMaterialization = D->clone();
-      CurrentMaterialization->setName(D->getName());
-      CurrentMaterialization->insertBefore(InsertPoint);
-      InsertPoint = CurrentMaterialization;
-
-      // Replace all uses of Def in the instructions being added as part of this
-      // rematerialization group
-      for (auto &I : InstructionsToProcess)
-        I->replaceUsesOfWith(D, CurrentMaterialization);
-
-      // Don't replace the final use at this point as this can cause problems
-      // for other materializations. Instead, for any final use that uses a
-      // define that's being rematerialized, record the replace values
-      for (unsigned i = 0, E = Use->getNumOperands(); i != E; ++i)
-        if (Use->getOperand(i) == D) // Is this operand pointing to oldval?
-          FinalInstructionsToProcess.push_back(
-              {Use, D, CurrentMaterialization});
-
-      InstructionsToProcess.push_back(CurrentMaterialization);
-    }
-  }
-
-  // Finally, replace the uses with the defines that we've just rematerialized
-  for (auto &R : FinalInstructionsToProcess) {
-    if (auto *PN = dyn_cast<PHINode>(R.Use)) {
-      assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
-                                                "values in the PHINode");
-      PN->replaceAllUsesWith(R.Remat);
-      PN->eraseFromParent();
-      continue;
-    }
-    R.Use->replaceUsesOfWith(R.Def, R.Remat);
-  }
-}
-
 // Splits the block at a particular instruction unless it is the first
 // instruction in the block with a single predecessor.
 static BasicBlock *splitBlockIfNotFirst(Instruction *I, const Twine &Name) {
@@ -1741,10 +1522,6 @@ static void splitAround(Instruction *I, const Twine &Name) {
   splitBlockIfNotFirst(I->getNextNode(), "After" + Name);
 }
 
-static bool isSuspendBlock(BasicBlock *BB) {
-  return isa<AnyCoroSuspendInst>(BB->front());
-}
-
 /// After we split the coroutine, will the given basic block be along
 /// an obvious exit path for the resumption function?
 static bool willLeaveFunctionImmediatelyAfter(BasicBlock *BB,
@@ -1754,7 +1531,7 @@ static bool willLeaveFunctionImmediatelyAfter(BasicBlock *BB,
   if (depth == 0) return false;
 
   // If this is a suspend block, we're about to exit the resumption function.
-  if (isSuspendBlock(BB))
+  if (coro::isSuspendBlock(BB))
     return true;
 
   // Recurse into the successors.
@@ -1995,7 +1772,8 @@ static void sinkLifetimeStartMarkers(Function &F, coro::Shape &Shape,
   DomSet.insert(&F.getEntryBlock());
   for (auto *CSI : Shape.CoroSuspends) {
     BasicBlock *SuspendBlock = CSI->getParent();
-    assert(isSuspendBlock(SuspendBlock) && SuspendBlock->getSingleSuccessor() &&
+    assert(coro::isSuspendBlock(SuspendBlock) &&
+           SuspendBlock->getSingleSuccessor() &&
            "should have split coro.suspend into its own block");
     DomSet.insert(SuspendBlock->getSingleSuccessor());
   }
@@ -2227,68 +2005,6 @@ void coro::salvageDebugInfo(
   }
 }
 
-static void doRematerializations(
-    Function &F, SuspendCrossingInfo &Checker,
-    const std::function<bool(Instruction &)> &MaterializableCallback) {
-  if (F.hasOptNone())
-    return;
-
-  coro::SpillInfo Spills;
-
-  // See if there are materializable instructions across suspend points
-  // We record these as the starting point to also identify materializable
-  // defs of uses in these operations
-  for (Instruction &I : instructions(F)) {
-    if (!MaterializableCallback(I))
-      continue;
-    for (User *U : I.users())
-      if (Checker.isDefinitionAcrossSuspend(I, U))
-        Spills[&I].push_back(cast<Instruction>(U));
-  }
-
-  // Process each of the identified rematerializable instructions
-  // and add predecessor instructions that can also be rematerialized.
-  // This is actually a graph of instructions since we could potentially
-  // have multiple uses of a def in the set of predecessor instructions.
-  // The approach here is to maintain a graph of instructions for each bottom
-  // level instruction - where we have a unique set of instructions (nodes)
-  // and edges between them. We then walk the graph in reverse post-dominator
-  // order to insert them past the suspend point, but ensure that ordering is
-  // correct. We also rely on CSE removing duplicate defs for remats of
-  // different instructions with a def in common (rather than maintaining more
-  // complex graphs for each suspend point)
-
-  // We can do this by adding new nodes to the list for each suspend
-  // point. Then using standard GraphTraits to give a reverse post-order
-  // traversal when we insert the nodes after the suspend
-  SmallMapVector<Instruction *, std::unique_ptr<RematGraph>, 8> AllRemats;
-  for (auto &E : Spills) {
-    for (Instruction *U : E.second) {
-      // Don't process a user twice (this can happen if the instruction uses
-      // more than one rematerializable def)
-      if (AllRemats.count(U))
-        continue;
-
-      // Constructor creates the whole RematGraph for the given Use
-      auto RematUPtr =
-          std::make_unique<RematGraph>(MaterializableCallback, U, Checker);
-
-      LLVM_DEBUG(dbgs() << "***** Next remat group *****\n";
-                 ReversePostOrderTraversal<RematGraph *> RPOT(RematUPtr.get());
-                 for (auto I = RPOT.begin(); I != RPOT.end();
-                      ++I) { (*I)->Node->dump(); } dbgs()
-                 << "\n";);
-
-      AllRemats[U] = std::move(RematUPtr);
-    }
-  }
-
-  // Rewrite materializable instructions to be materialized at the use
-  // point.
-  LLVM_DEBUG(dumpRemats("Materializations", AllRemats));
-  rewriteMaterializableInstructions(AllRemats);
-}
-
 void coro::normalizeCoroutine(Function &F, coro::Shape &Shape,
                               TargetTransformInfo &TTI) {
   // Don't eliminate swifterror in async functions that won't be split.
@@ -2324,8 +2040,8 @@ void coro::normalizeCoroutine(Function &F, coro::Shape &Shape,
       IRBuilder<> Builder(AsyncEnd);
       SmallVector<Value *, 8> Args(AsyncEnd->args());
       auto Arguments = ArrayRef<Value *>(Args).drop_front(3);
-      auto *Call = createMustTailCall(AsyncEnd->getDebugLoc(), MustTailCallFn,
-                                      TTI, Arguments, Builder);
+      auto *Call = coro::createMustTailCall(
+          AsyncEnd->getDebugLoc(), MustTailCallFn, TTI, Arguments, Builder);
       splitAround(Call, "MustTailCall.Before.CoroEnd");
     }
   }

llvm/lib/Transforms/Coroutines/CoroInternal.h

@@ -21,6 +21,7 @@ class CallGraph;
 
 namespace coro {
 
+bool isSuspendBlock(BasicBlock *BB);
 bool declaresAnyIntrinsic(const Module &M);
 bool declaresIntrinsics(const Module &M,
                         const std::initializer_list<StringRef>);

llvm/lib/Transforms/Coroutines/Coroutines.cpp

@@ -100,6 +100,10 @@ static bool isCoroutineIntrinsicName(StringRef Name) {
 }
 #endif
 
+bool coro::isSuspendBlock(BasicBlock *BB) {
+  return isa<AnyCoroSuspendInst>(BB->front());
+}
+
 bool coro::declaresAnyIntrinsic(const Module &M) {
   for (StringRef Name : CoroIntrinsics) {
     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");