1//===- MaterializationUtils.cpp - Builds and manipulates coroutine frame
2//-------------===//
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-exception
7//
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 instructions
33// When the constructor is invoked with a candidate instruction (which is
34// materializable) it builds a DAG of materializable instructions from that
35// point.
36// Typically, for each instruction identified as re-materializable across a
37// suspend point, a RematGraph will be created.
38struct RematGraph {
39 // Each RematNode in the graph contains the edges to instructions providing
40 // 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>(args&: I);
59 EntryNode = FirstNode.get();
60 std::deque<std::unique_ptr<RematNode>> WorkList;
61 addNode(NUPtr: std::move(FirstNode), WorkList, FirstUse: cast<User>(Val: I));
62 while (WorkList.size()) {
63 std::unique_ptr<RematNode> N = std::move(WorkList.front());
64 WorkList.pop_front();
65 addNode(NUPtr: std::move(N), WorkList, FirstUse: cast<User>(Val: 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(Key: N->Node);
74 if (!Inserted)
75 return;
76
77 // We haven't see this node yet - add to the list
78 It->second = std::move(NUPtr);
79 for (auto &Def : N->Node->operands()) {
80 Instruction *D = dyn_cast<Instruction>(Val: Def.get());
81 if (!D || !MaterializableCallback(*D) ||
82 !Checker.isDefinitionAcrossSuspend(I&: *D, U: FirstUse))
83 continue;
84
85 if (auto It = Remats.find(Key: D); It != Remats.end()) {
86 // Already have this in the graph
87 N->Operands.push_back(Elt: 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(Elt: I.get());
96 break;
97 }
98 }
99 if (NoMatch) {
100 // Create a new node
101 std::unique_ptr<RematNode> ChildNode = std::make_unique<RematNode>(args&: D);
102 N->Operands.push_back(Elt: ChildNode.get());
103 WorkList.push_back(x: 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#endif
136};
137
138} // namespace
139
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 phases
158 // Do the remats and record the required defs to be replaced in the
159 // original use instructions
160 // Once all the remats are complete, replace the uses in the final
161 // instructions with the new defs
162 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 to
178 // insert the remats into the end of the predecessor (there should only be
179 // one). This is so that suspend blocks always have the suspend instruction
180 // as the first instruction.
181 BasicBlock::iterator InsertPoint = Use->getParent()->getFirstInsertionPt();
182 if (isa<AnyCoroSuspendInst>(Val: 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're
190 // 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(InsertPos: InsertPoint);
198 InsertPoint = CurrentMaterialization->getIterator();
199
200 // Replace all uses of Def in the instructions being added as part of this
201 // rematerialization group
202 for (auto &I : InstructionsToProcess)
203 I->replaceUsesOfWith(From: D, To: CurrentMaterialization);
204
205 // Don't replace the final use at this point as this can cause problems
206 // for other materializations. Instead, for any final use that uses a
207 // define that's being rematerialized, record the replace values
208 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 Elt: {.Use: Use, .Def: D, .Remat: CurrentMaterialization});
212
213 InstructionsToProcess.push_back(Elt: CurrentMaterialization);
214 }
215 }
216
217 // Finally, replace the uses with the defines that we've just rematerialized
218 for (auto &R : FinalInstructionsToProcess) {
219 if (auto *PN = dyn_cast<PHINode>(Val: R.Use)) {
220 assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
221 "values in the PHINode");
222 PN->replaceAllUsesWith(V: R.Remat);
223 PN->eraseFromParent();
224 continue;
225 }
226 R.Use->replaceUsesOfWith(From: R.Def, To: R.Remat);
227 }
228}
229
230/// Default materializable callback
231// Check for instructions that we can recreate on resume as opposed to spill
232// the result into a coroutine frame.
233bool llvm::coro::defaultMaterializable(Instruction &V) {
234 return (isa<CastInst>(Val: &V) || isa<GetElementPtrInst>(Val: &V) ||
235 isa<BinaryOperator>(Val: &V) || isa<CmpInst>(Val: &V) || isa<SelectInst>(Val: &V));
236}
237
238bool llvm::coro::isTriviallyMaterializable(Instruction &V) {
239 return defaultMaterializable(V);
240}
241
242#ifndef NDEBUG
243static 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#endif
253
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 points
263 // We record these as the starting point to also identify materializable
264 // defs of uses in these operations
265 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(Elt: cast<Instruction>(Val: U));
271 }
272
273 // Process each of the identified rematerializable instructions
274 // and add predecessor instructions that can also be rematerialized.
275 // This is actually a graph of instructions since we could potentially
276 // 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 bottom
278 // level instruction - where we have a unique set of instructions (nodes)
279 // and edges between them. We then walk the graph in reverse post-dominator
280 // order to insert them past the suspend point, but ensure that ordering is
281 // correct. We also rely on CSE removing duplicate defs for remats of
282 // different instructions with a def in common (rather than maintaining more
283 // complex graphs for each suspend point)
284
285 // We can do this by adding new nodes to the list for each suspend
286 // point. Then using standard GraphTraits to give a reverse post-order
287 // traversal when we insert the nodes after the suspend
288 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 uses
292 // more than one rematerializable def)
293 auto [It, Inserted] = AllRemats.try_emplace(Key: U);
294 if (!Inserted)
295 continue;
296
297 // Constructor creates the whole RematGraph for the given Use
298 auto RematUPtr =
299 std::make_unique<RematGraph>(args&: IsMaterializable, args&: U, args&: 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 use
312 // point.
313 LLVM_DEBUG(dumpRemats("Materializations", AllRemats));
314 rewriteMaterializableInstructions(AllRemats);
315}
316