316 lines · cpp
1//===- MaterializationUtils.cpp - Builds and manipulates coroutine frame2//-------------===//3//4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.5// See https://llvm.org/LICENSE.txt for license information.6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception7//8//===----------------------------------------------------------------------===//9// This file contains classes used to materialize insts after suspends points.10//===----------------------------------------------------------------------===//11 12#include "llvm/Transforms/Coroutines/MaterializationUtils.h"13#include "CoroInternal.h"14#include "llvm/ADT/PostOrderIterator.h"15#include "llvm/IR/Dominators.h"16#include "llvm/IR/InstIterator.h"17#include "llvm/IR/Instruction.h"18#include "llvm/IR/ModuleSlotTracker.h"19#include "llvm/Transforms/Coroutines/SpillUtils.h"20#include <deque>21 22using namespace llvm;23 24using namespace coro;25 26// The "coro-suspend-crossing" flag is very noisy. There is another debug type,27// "coro-frame", which results in leaner debug spew.28#define DEBUG_TYPE "coro-suspend-crossing"29 30namespace {31 32// RematGraph is used to construct a DAG for rematerializable instructions33// When the constructor is invoked with a candidate instruction (which is34// materializable) it builds a DAG of materializable instructions from that35// point.36// Typically, for each instruction identified as re-materializable across a37// suspend point, a RematGraph will be created.38struct RematGraph {39 // Each RematNode in the graph contains the edges to instructions providing40 // operands in the current node.41 struct RematNode {42 Instruction *Node;43 SmallVector<RematNode *> Operands;44 RematNode() = default;45 RematNode(Instruction *V) : Node(V) {}46 };47 48 RematNode *EntryNode;49 using RematNodeMap =50 SmallMapVector<Instruction *, std::unique_ptr<RematNode>, 8>;51 RematNodeMap Remats;52 const std::function<bool(Instruction &)> &MaterializableCallback;53 SuspendCrossingInfo &Checker;54 55 RematGraph(const std::function<bool(Instruction &)> &MaterializableCallback,56 Instruction *I, SuspendCrossingInfo &Checker)57 : MaterializableCallback(MaterializableCallback), Checker(Checker) {58 std::unique_ptr<RematNode> FirstNode = std::make_unique<RematNode>(I);59 EntryNode = FirstNode.get();60 std::deque<std::unique_ptr<RematNode>> WorkList;61 addNode(std::move(FirstNode), WorkList, cast<User>(I));62 while (WorkList.size()) {63 std::unique_ptr<RematNode> N = std::move(WorkList.front());64 WorkList.pop_front();65 addNode(std::move(N), WorkList, cast<User>(I));66 }67 }68 69 void addNode(std::unique_ptr<RematNode> NUPtr,70 std::deque<std::unique_ptr<RematNode>> &WorkList,71 User *FirstUse) {72 RematNode *N = NUPtr.get();73 auto [It, Inserted] = Remats.try_emplace(N->Node);74 if (!Inserted)75 return;76 77 // We haven't see this node yet - add to the list78 It->second = std::move(NUPtr);79 for (auto &Def : N->Node->operands()) {80 Instruction *D = dyn_cast<Instruction>(Def.get());81 if (!D || !MaterializableCallback(*D) ||82 !Checker.isDefinitionAcrossSuspend(*D, FirstUse))83 continue;84 85 if (auto It = Remats.find(D); It != Remats.end()) {86 // Already have this in the graph87 N->Operands.push_back(It->second.get());88 continue;89 }90 91 bool NoMatch = true;92 for (auto &I : WorkList) {93 if (I->Node == D) {94 NoMatch = false;95 N->Operands.push_back(I.get());96 break;97 }98 }99 if (NoMatch) {100 // Create a new node101 std::unique_ptr<RematNode> ChildNode = std::make_unique<RematNode>(D);102 N->Operands.push_back(ChildNode.get());103 WorkList.push_back(std::move(ChildNode));104 }105 }106 }107 108#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)109 static void dumpBasicBlockLabel(const BasicBlock *BB,110 ModuleSlotTracker &MST) {111 if (BB->hasName()) {112 dbgs() << BB->getName();113 return;114 }115 116 dbgs() << MST.getLocalSlot(BB);117 }118 119 void dump() const {120 BasicBlock *BB = EntryNode->Node->getParent();121 Function *F = BB->getParent();122 123 ModuleSlotTracker MST(F->getParent());124 MST.incorporateFunction(*F);125 126 dbgs() << "Entry (";127 dumpBasicBlockLabel(BB, MST);128 dbgs() << ") : " << *EntryNode->Node << "\n";129 for (auto &E : Remats) {130 dbgs() << *(E.first) << "\n";131 for (RematNode *U : E.second->Operands)132 dbgs() << " " << *U->Node << "\n";133 }134 }135#endif136};137 138} // namespace139 140template <> struct llvm::GraphTraits<RematGraph *> {141 using NodeRef = RematGraph::RematNode *;142 using ChildIteratorType = RematGraph::RematNode **;143 144 static NodeRef getEntryNode(RematGraph *G) { return G->EntryNode; }145 static ChildIteratorType child_begin(NodeRef N) {146 return N->Operands.begin();147 }148 static ChildIteratorType child_end(NodeRef N) { return N->Operands.end(); }149};150 151// For each instruction identified as materializable across the suspend point,152// and its associated DAG of other rematerializable instructions,153// recreate the DAG of instructions after the suspend point.154static void rewriteMaterializableInstructions(155 const SmallMapVector<Instruction *, std::unique_ptr<RematGraph>, 8>156 &AllRemats) {157 // This has to be done in 2 phases158 // Do the remats and record the required defs to be replaced in the159 // original use instructions160 // Once all the remats are complete, replace the uses in the final161 // instructions with the new defs162 typedef struct {163 Instruction *Use;164 Instruction *Def;165 Instruction *Remat;166 } ProcessNode;167 168 SmallVector<ProcessNode> FinalInstructionsToProcess;169 170 for (const auto &E : AllRemats) {171 Instruction *Use = E.first;172 Instruction *CurrentMaterialization = nullptr;173 RematGraph *RG = E.second.get();174 ReversePostOrderTraversal<RematGraph *> RPOT(RG);175 SmallVector<Instruction *> InstructionsToProcess;176 177 // If the target use is actually a suspend instruction then we have to178 // insert the remats into the end of the predecessor (there should only be179 // one). This is so that suspend blocks always have the suspend instruction180 // as the first instruction.181 BasicBlock::iterator InsertPoint = Use->getParent()->getFirstInsertionPt();182 if (isa<AnyCoroSuspendInst>(Use)) {183 BasicBlock *SuspendPredecessorBlock =184 Use->getParent()->getSinglePredecessor();185 assert(SuspendPredecessorBlock && "malformed coro suspend instruction");186 InsertPoint = SuspendPredecessorBlock->getTerminator()->getIterator();187 }188 189 // Note: skip the first instruction as this is the actual use that we're190 // rematerializing everything for.191 auto I = RPOT.begin();192 ++I;193 for (; I != RPOT.end(); ++I) {194 Instruction *D = (*I)->Node;195 CurrentMaterialization = D->clone();196 CurrentMaterialization->setName(D->getName());197 CurrentMaterialization->insertBefore(InsertPoint);198 InsertPoint = CurrentMaterialization->getIterator();199 200 // Replace all uses of Def in the instructions being added as part of this201 // rematerialization group202 for (auto &I : InstructionsToProcess)203 I->replaceUsesOfWith(D, CurrentMaterialization);204 205 // Don't replace the final use at this point as this can cause problems206 // for other materializations. Instead, for any final use that uses a207 // define that's being rematerialized, record the replace values208 for (unsigned i = 0, E = Use->getNumOperands(); i != E; ++i)209 if (Use->getOperand(i) == D) // Is this operand pointing to oldval?210 FinalInstructionsToProcess.push_back(211 {Use, D, CurrentMaterialization});212 213 InstructionsToProcess.push_back(CurrentMaterialization);214 }215 }216 217 // Finally, replace the uses with the defines that we've just rematerialized218 for (auto &R : FinalInstructionsToProcess) {219 if (auto *PN = dyn_cast<PHINode>(R.Use)) {220 assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "221 "values in the PHINode");222 PN->replaceAllUsesWith(R.Remat);223 PN->eraseFromParent();224 continue;225 }226 R.Use->replaceUsesOfWith(R.Def, R.Remat);227 }228}229 230/// Default materializable callback231// Check for instructions that we can recreate on resume as opposed to spill232// the result into a coroutine frame.233bool llvm::coro::defaultMaterializable(Instruction &V) {234 return (isa<CastInst>(&V) || isa<GetElementPtrInst>(&V) ||235 isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V));236}237 238bool llvm::coro::isTriviallyMaterializable(Instruction &V) {239 return defaultMaterializable(V);240}241 242#ifndef NDEBUG243static void dumpRemats(244 StringRef Title,245 const SmallMapVector<Instruction *, std::unique_ptr<RematGraph>, 8> &RM) {246 dbgs() << "------------- " << Title << "--------------\n";247 for (const auto &E : RM) {248 E.second->dump();249 dbgs() << "--\n";250 }251}252#endif253 254void coro::doRematerializations(255 Function &F, SuspendCrossingInfo &Checker,256 std::function<bool(Instruction &)> IsMaterializable) {257 if (F.hasOptNone())258 return;259 260 coro::SpillInfo Spills;261 262 // See if there are materializable instructions across suspend points263 // We record these as the starting point to also identify materializable264 // defs of uses in these operations265 for (Instruction &I : instructions(F)) {266 if (!IsMaterializable(I))267 continue;268 for (User *U : I.users())269 if (Checker.isDefinitionAcrossSuspend(I, U))270 Spills[&I].push_back(cast<Instruction>(U));271 }272 273 // Process each of the identified rematerializable instructions274 // and add predecessor instructions that can also be rematerialized.275 // This is actually a graph of instructions since we could potentially276 // have multiple uses of a def in the set of predecessor instructions.277 // The approach here is to maintain a graph of instructions for each bottom278 // level instruction - where we have a unique set of instructions (nodes)279 // and edges between them. We then walk the graph in reverse post-dominator280 // order to insert them past the suspend point, but ensure that ordering is281 // correct. We also rely on CSE removing duplicate defs for remats of282 // different instructions with a def in common (rather than maintaining more283 // complex graphs for each suspend point)284 285 // We can do this by adding new nodes to the list for each suspend286 // point. Then using standard GraphTraits to give a reverse post-order287 // traversal when we insert the nodes after the suspend288 SmallMapVector<Instruction *, std::unique_ptr<RematGraph>, 8> AllRemats;289 for (auto &E : Spills) {290 for (Instruction *U : E.second) {291 // Don't process a user twice (this can happen if the instruction uses292 // more than one rematerializable def)293 auto [It, Inserted] = AllRemats.try_emplace(U);294 if (!Inserted)295 continue;296 297 // Constructor creates the whole RematGraph for the given Use298 auto RematUPtr =299 std::make_unique<RematGraph>(IsMaterializable, U, Checker);300 301 LLVM_DEBUG(dbgs() << "***** Next remat group *****\n";302 ReversePostOrderTraversal<RematGraph *> RPOT(RematUPtr.get());303 for (auto I = RPOT.begin(); I != RPOT.end();304 ++I) { (*I)->Node->dump(); } dbgs()305 << "\n";);306 307 It->second = std::move(RematUPtr);308 }309 }310 311 // Rewrite materializable instructions to be materialized at the use312 // point.313 LLVM_DEBUG(dumpRemats("Materializations", AllRemats));314 rewriteMaterializableInstructions(AllRemats);315}316