-
Notifications
You must be signed in to change notification settings - Fork 11.6k
/
Inliner.cpp
799 lines (690 loc) · 30.2 KB
/
Inliner.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
//===- Inliner.cpp - Pass to inline function calls ------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements a basic inlining algorithm that operates bottom up over
// the Strongly Connect Components(SCCs) of the CallGraph. This enables a more
// incremental propagation of inlining decisions from the leafs to the roots of
// the callgraph.
//
//===----------------------------------------------------------------------===//
#include "PassDetail.h"
#include "mlir/Analysis/CallGraph.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/InliningUtils.h"
#include "mlir/Transforms/Passes.h"
#include "llvm/ADT/SCCIterator.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/Parallel.h"
#define DEBUG_TYPE "inlining"
using namespace mlir;
/// This function implements the default inliner optimization pipeline.
static void defaultInlinerOptPipeline(OpPassManager &pm) {
pm.addPass(createCanonicalizerPass());
}
//===----------------------------------------------------------------------===//
// Symbol Use Tracking
//===----------------------------------------------------------------------===//
/// Walk all of the used symbol callgraph nodes referenced with the given op.
static void walkReferencedSymbolNodes(
Operation *op, CallGraph &cg, SymbolTableCollection &symbolTable,
DenseMap<Attribute, CallGraphNode *> &resolvedRefs,
function_ref<void(CallGraphNode *, Operation *)> callback) {
auto symbolUses = SymbolTable::getSymbolUses(op);
assert(symbolUses && "expected uses to be valid");
Operation *symbolTableOp = op->getParentOp();
for (const SymbolTable::SymbolUse &use : *symbolUses) {
auto refIt = resolvedRefs.insert({use.getSymbolRef(), nullptr});
CallGraphNode *&node = refIt.first->second;
// If this is the first instance of this reference, try to resolve a
// callgraph node for it.
if (refIt.second) {
auto *symbolOp = symbolTable.lookupNearestSymbolFrom(symbolTableOp,
use.getSymbolRef());
auto callableOp = dyn_cast_or_null<CallableOpInterface>(symbolOp);
if (!callableOp)
continue;
node = cg.lookupNode(callableOp.getCallableRegion());
}
if (node)
callback(node, use.getUser());
}
}
//===----------------------------------------------------------------------===//
// CGUseList
namespace {
/// This struct tracks the uses of callgraph nodes that can be dropped when
/// use_empty. It directly tracks and manages a use-list for all of the
/// call-graph nodes. This is necessary because many callgraph nodes are
/// referenced by SymbolRefAttr, which has no mechanism akin to the SSA `Use`
/// class.
struct CGUseList {
/// This struct tracks the uses of callgraph nodes within a specific
/// operation.
struct CGUser {
/// Any nodes referenced in the top-level attribute list of this user. We
/// use a set here because the number of references does not matter.
DenseSet<CallGraphNode *> topLevelUses;
/// Uses of nodes referenced by nested operations.
DenseMap<CallGraphNode *, int> innerUses;
};
CGUseList(Operation *op, CallGraph &cg, SymbolTableCollection &symbolTable);
/// Drop uses of nodes referred to by the given call operation that resides
/// within 'userNode'.
void dropCallUses(CallGraphNode *userNode, Operation *callOp, CallGraph &cg);
/// Remove the given node from the use list.
void eraseNode(CallGraphNode *node);
/// Returns true if the given callgraph node has no uses and can be pruned.
bool isDead(CallGraphNode *node) const;
/// Returns true if the given callgraph node has a single use and can be
/// discarded.
bool hasOneUseAndDiscardable(CallGraphNode *node) const;
/// Recompute the uses held by the given callgraph node.
void recomputeUses(CallGraphNode *node, CallGraph &cg);
/// Merge the uses of 'lhs' with the uses of the 'rhs' after inlining a copy
/// of 'lhs' into 'rhs'.
void mergeUsesAfterInlining(CallGraphNode *lhs, CallGraphNode *rhs);
private:
/// Decrement the uses of discardable nodes referenced by the given user.
void decrementDiscardableUses(CGUser &uses);
/// A mapping between a discardable callgraph node (that is a symbol) and the
/// number of uses for this node.
DenseMap<CallGraphNode *, int> discardableSymNodeUses;
/// A mapping between a callgraph node and the symbol callgraph nodes that it
/// uses.
DenseMap<CallGraphNode *, CGUser> nodeUses;
/// A symbol table to use when resolving call lookups.
SymbolTableCollection &symbolTable;
};
} // end anonymous namespace
CGUseList::CGUseList(Operation *op, CallGraph &cg,
SymbolTableCollection &symbolTable)
: symbolTable(symbolTable) {
/// A set of callgraph nodes that are always known to be live during inlining.
DenseMap<Attribute, CallGraphNode *> alwaysLiveNodes;
// Walk each of the symbol tables looking for discardable callgraph nodes.
auto walkFn = [&](Operation *symbolTableOp, bool allUsesVisible) {
for (Operation &op : symbolTableOp->getRegion(0).getOps()) {
// If this is a callgraph operation, check to see if it is discardable.
if (auto callable = dyn_cast<CallableOpInterface>(&op)) {
if (auto *node = cg.lookupNode(callable.getCallableRegion())) {
SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(&op);
if (symbol && (allUsesVisible || symbol.isPrivate()) &&
symbol.canDiscardOnUseEmpty()) {
discardableSymNodeUses.try_emplace(node, 0);
}
continue;
}
}
// Otherwise, check for any referenced nodes. These will be always-live.
walkReferencedSymbolNodes(&op, cg, symbolTable, alwaysLiveNodes,
[](CallGraphNode *, Operation *) {});
}
};
SymbolTable::walkSymbolTables(op, /*allSymUsesVisible=*/!op->getBlock(),
walkFn);
// Drop the use information for any discardable nodes that are always live.
for (auto &it : alwaysLiveNodes)
discardableSymNodeUses.erase(it.second);
// Compute the uses for each of the callable nodes in the graph.
for (CallGraphNode *node : cg)
recomputeUses(node, cg);
}
void CGUseList::dropCallUses(CallGraphNode *userNode, Operation *callOp,
CallGraph &cg) {
auto &userRefs = nodeUses[userNode].innerUses;
auto walkFn = [&](CallGraphNode *node, Operation *user) {
auto parentIt = userRefs.find(node);
if (parentIt == userRefs.end())
return;
--parentIt->second;
--discardableSymNodeUses[node];
};
DenseMap<Attribute, CallGraphNode *> resolvedRefs;
walkReferencedSymbolNodes(callOp, cg, symbolTable, resolvedRefs, walkFn);
}
void CGUseList::eraseNode(CallGraphNode *node) {
// Drop all child nodes.
for (auto &edge : *node)
if (edge.isChild())
eraseNode(edge.getTarget());
// Drop the uses held by this node and erase it.
auto useIt = nodeUses.find(node);
assert(useIt != nodeUses.end() && "expected node to be valid");
decrementDiscardableUses(useIt->getSecond());
nodeUses.erase(useIt);
discardableSymNodeUses.erase(node);
}
bool CGUseList::isDead(CallGraphNode *node) const {
// If the parent operation isn't a symbol, simply check normal SSA deadness.
Operation *nodeOp = node->getCallableRegion()->getParentOp();
if (!isa<SymbolOpInterface>(nodeOp))
return MemoryEffectOpInterface::hasNoEffect(nodeOp) && nodeOp->use_empty();
// Otherwise, check the number of symbol uses.
auto symbolIt = discardableSymNodeUses.find(node);
return symbolIt != discardableSymNodeUses.end() && symbolIt->second == 0;
}
bool CGUseList::hasOneUseAndDiscardable(CallGraphNode *node) const {
// If this isn't a symbol node, check for side-effects and SSA use count.
Operation *nodeOp = node->getCallableRegion()->getParentOp();
if (!isa<SymbolOpInterface>(nodeOp))
return MemoryEffectOpInterface::hasNoEffect(nodeOp) && nodeOp->hasOneUse();
// Otherwise, check the number of symbol uses.
auto symbolIt = discardableSymNodeUses.find(node);
return symbolIt != discardableSymNodeUses.end() && symbolIt->second == 1;
}
void CGUseList::recomputeUses(CallGraphNode *node, CallGraph &cg) {
Operation *parentOp = node->getCallableRegion()->getParentOp();
CGUser &uses = nodeUses[node];
decrementDiscardableUses(uses);
// Collect the new discardable uses within this node.
uses = CGUser();
DenseMap<Attribute, CallGraphNode *> resolvedRefs;
auto walkFn = [&](CallGraphNode *refNode, Operation *user) {
auto discardSymIt = discardableSymNodeUses.find(refNode);
if (discardSymIt == discardableSymNodeUses.end())
return;
if (user != parentOp)
++uses.innerUses[refNode];
else if (!uses.topLevelUses.insert(refNode).second)
return;
++discardSymIt->second;
};
walkReferencedSymbolNodes(parentOp, cg, symbolTable, resolvedRefs, walkFn);
}
void CGUseList::mergeUsesAfterInlining(CallGraphNode *lhs, CallGraphNode *rhs) {
auto &lhsUses = nodeUses[lhs], &rhsUses = nodeUses[rhs];
for (auto &useIt : lhsUses.innerUses) {
rhsUses.innerUses[useIt.first] += useIt.second;
discardableSymNodeUses[useIt.first] += useIt.second;
}
}
void CGUseList::decrementDiscardableUses(CGUser &uses) {
for (CallGraphNode *node : uses.topLevelUses)
--discardableSymNodeUses[node];
for (auto &it : uses.innerUses)
discardableSymNodeUses[it.first] -= it.second;
}
//===----------------------------------------------------------------------===//
// CallGraph traversal
//===----------------------------------------------------------------------===//
namespace {
/// This class represents a specific callgraph SCC.
class CallGraphSCC {
public:
CallGraphSCC(llvm::scc_iterator<const CallGraph *> &parentIterator)
: parentIterator(parentIterator) {}
/// Return a range over the nodes within this SCC.
std::vector<CallGraphNode *>::iterator begin() { return nodes.begin(); }
std::vector<CallGraphNode *>::iterator end() { return nodes.end(); }
/// Reset the nodes of this SCC with those provided.
void reset(const std::vector<CallGraphNode *> &newNodes) { nodes = newNodes; }
/// Remove the given node from this SCC.
void remove(CallGraphNode *node) {
auto it = llvm::find(nodes, node);
if (it != nodes.end()) {
nodes.erase(it);
parentIterator.ReplaceNode(node, nullptr);
}
}
private:
std::vector<CallGraphNode *> nodes;
llvm::scc_iterator<const CallGraph *> &parentIterator;
};
} // end anonymous namespace
/// Run a given transformation over the SCCs of the callgraph in a bottom up
/// traversal.
static LogicalResult runTransformOnCGSCCs(
const CallGraph &cg,
function_ref<LogicalResult(CallGraphSCC &)> sccTransformer) {
llvm::scc_iterator<const CallGraph *> cgi = llvm::scc_begin(&cg);
CallGraphSCC currentSCC(cgi);
while (!cgi.isAtEnd()) {
// Copy the current SCC and increment so that the transformer can modify the
// SCC without invalidating our iterator.
currentSCC.reset(*cgi);
++cgi;
if (failed(sccTransformer(currentSCC)))
return failure();
}
return success();
}
namespace {
/// This struct represents a resolved call to a given callgraph node. Given that
/// the call does not actually contain a direct reference to the
/// Region(CallGraphNode) that it is dispatching to, we need to resolve them
/// explicitly.
struct ResolvedCall {
ResolvedCall(CallOpInterface call, CallGraphNode *sourceNode,
CallGraphNode *targetNode)
: call(call), sourceNode(sourceNode), targetNode(targetNode) {}
CallOpInterface call;
CallGraphNode *sourceNode, *targetNode;
};
} // end anonymous namespace
/// Collect all of the callable operations within the given range of blocks. If
/// `traverseNestedCGNodes` is true, this will also collect call operations
/// inside of nested callgraph nodes.
static void collectCallOps(iterator_range<Region::iterator> blocks,
CallGraphNode *sourceNode, CallGraph &cg,
SymbolTableCollection &symbolTable,
SmallVectorImpl<ResolvedCall> &calls,
bool traverseNestedCGNodes) {
SmallVector<std::pair<Block *, CallGraphNode *>, 8> worklist;
auto addToWorklist = [&](CallGraphNode *node,
iterator_range<Region::iterator> blocks) {
for (Block &block : blocks)
worklist.emplace_back(&block, node);
};
addToWorklist(sourceNode, blocks);
while (!worklist.empty()) {
Block *block;
std::tie(block, sourceNode) = worklist.pop_back_val();
for (Operation &op : *block) {
if (auto call = dyn_cast<CallOpInterface>(op)) {
// TODO: Support inlining nested call references.
CallInterfaceCallable callable = call.getCallableForCallee();
if (SymbolRefAttr symRef = callable.dyn_cast<SymbolRefAttr>()) {
if (!symRef.isa<FlatSymbolRefAttr>())
continue;
}
CallGraphNode *targetNode = cg.resolveCallable(call, symbolTable);
if (!targetNode->isExternal())
calls.emplace_back(call, sourceNode, targetNode);
continue;
}
// If this is not a call, traverse the nested regions. If
// `traverseNestedCGNodes` is false, then don't traverse nested call graph
// regions.
for (auto &nestedRegion : op.getRegions()) {
CallGraphNode *nestedNode = cg.lookupNode(&nestedRegion);
if (traverseNestedCGNodes || !nestedNode)
addToWorklist(nestedNode ? nestedNode : sourceNode, nestedRegion);
}
}
}
}
//===----------------------------------------------------------------------===//
// Inliner
//===----------------------------------------------------------------------===//
namespace {
/// This class provides a specialization of the main inlining interface.
struct Inliner : public InlinerInterface {
Inliner(MLIRContext *context, CallGraph &cg,
SymbolTableCollection &symbolTable)
: InlinerInterface(context), cg(cg), symbolTable(symbolTable) {}
/// Process a set of blocks that have been inlined. This callback is invoked
/// *before* inlined terminator operations have been processed.
void
processInlinedBlocks(iterator_range<Region::iterator> inlinedBlocks) final {
// Find the closest callgraph node from the first block.
CallGraphNode *node;
Region *region = inlinedBlocks.begin()->getParent();
while (!(node = cg.lookupNode(region))) {
region = region->getParentRegion();
assert(region && "expected valid parent node");
}
collectCallOps(inlinedBlocks, node, cg, symbolTable, calls,
/*traverseNestedCGNodes=*/true);
}
/// Mark the given callgraph node for deletion.
void markForDeletion(CallGraphNode *node) { deadNodes.insert(node); }
/// This method properly disposes of callables that became dead during
/// inlining. This should not be called while iterating over the SCCs.
void eraseDeadCallables() {
for (CallGraphNode *node : deadNodes)
node->getCallableRegion()->getParentOp()->erase();
}
/// The set of callables known to be dead.
SmallPtrSet<CallGraphNode *, 8> deadNodes;
/// The current set of call instructions to consider for inlining.
SmallVector<ResolvedCall, 8> calls;
/// The callgraph being operated on.
CallGraph &cg;
/// A symbol table to use when resolving call lookups.
SymbolTableCollection &symbolTable;
};
} // namespace
/// Returns true if the given call should be inlined.
static bool shouldInline(ResolvedCall &resolvedCall) {
// Don't allow inlining terminator calls. We currently don't support this
// case.
if (resolvedCall.call->hasTrait<OpTrait::IsTerminator>())
return false;
// Don't allow inlining if the target is an ancestor of the call. This
// prevents inlining recursively.
if (resolvedCall.targetNode->getCallableRegion()->isAncestor(
resolvedCall.call->getParentRegion()))
return false;
// Otherwise, inline.
return true;
}
/// Attempt to inline calls within the given scc. This function returns
/// success if any calls were inlined, failure otherwise.
static LogicalResult inlineCallsInSCC(Inliner &inliner, CGUseList &useList,
CallGraphSCC ¤tSCC) {
CallGraph &cg = inliner.cg;
auto &calls = inliner.calls;
// A set of dead nodes to remove after inlining.
SmallVector<CallGraphNode *, 1> deadNodes;
// Collect all of the direct calls within the nodes of the current SCC. We
// don't traverse nested callgraph nodes, because they are handled separately
// likely within a different SCC.
for (CallGraphNode *node : currentSCC) {
if (node->isExternal())
continue;
// Don't collect calls if the node is already dead.
if (useList.isDead(node)) {
deadNodes.push_back(node);
} else {
collectCallOps(*node->getCallableRegion(), node, cg, inliner.symbolTable,
calls, /*traverseNestedCGNodes=*/false);
}
}
// Try to inline each of the call operations. Don't cache the end iterator
// here as more calls may be added during inlining.
bool inlinedAnyCalls = false;
for (unsigned i = 0; i != calls.size(); ++i) {
ResolvedCall it = calls[i];
bool doInline = shouldInline(it);
CallOpInterface call = it.call;
LLVM_DEBUG({
if (doInline)
llvm::dbgs() << "* Inlining call: " << call << "\n";
else
llvm::dbgs() << "* Not inlining call: " << call << "\n";
});
if (!doInline)
continue;
Region *targetRegion = it.targetNode->getCallableRegion();
// If this is the last call to the target node and the node is discardable,
// then inline it in-place and delete the node if successful.
bool inlineInPlace = useList.hasOneUseAndDiscardable(it.targetNode);
LogicalResult inlineResult = inlineCall(
inliner, call, cast<CallableOpInterface>(targetRegion->getParentOp()),
targetRegion, /*shouldCloneInlinedRegion=*/!inlineInPlace);
if (failed(inlineResult)) {
LLVM_DEBUG(llvm::dbgs() << "** Failed to inline\n");
continue;
}
inlinedAnyCalls = true;
// If the inlining was successful, Merge the new uses into the source node.
useList.dropCallUses(it.sourceNode, call.getOperation(), cg);
useList.mergeUsesAfterInlining(it.targetNode, it.sourceNode);
// then erase the call.
call.erase();
// If we inlined in place, mark the node for deletion.
if (inlineInPlace) {
useList.eraseNode(it.targetNode);
deadNodes.push_back(it.targetNode);
}
}
for (CallGraphNode *node : deadNodes) {
currentSCC.remove(node);
inliner.markForDeletion(node);
}
calls.clear();
return success(inlinedAnyCalls);
}
//===----------------------------------------------------------------------===//
// InlinerPass
//===----------------------------------------------------------------------===//
namespace {
class InlinerPass : public InlinerBase<InlinerPass> {
public:
InlinerPass();
InlinerPass(const InlinerPass &) = default;
InlinerPass(std::function<void(OpPassManager &)> defaultPipeline);
InlinerPass(std::function<void(OpPassManager &)> defaultPipeline,
llvm::StringMap<OpPassManager> opPipelines);
void runOnOperation() override;
private:
/// Attempt to inline calls within the given scc, and run simplifications,
/// until a fixed point is reached. This allows for the inlining of newly
/// devirtualized calls. Returns failure if there was a fatal error during
/// inlining.
LogicalResult inlineSCC(Inliner &inliner, CGUseList &useList,
CallGraphSCC ¤tSCC, MLIRContext *context);
/// Optimize the nodes within the given SCC with one of the held optimization
/// pass pipelines. Returns failure if an error occurred during the
/// optimization of the SCC, success otherwise.
LogicalResult optimizeSCC(CallGraph &cg, CGUseList &useList,
CallGraphSCC ¤tSCC, MLIRContext *context);
/// Optimize the nodes within the given SCC in parallel. Returns failure if an
/// error occurred during the optimization of the SCC, success otherwise.
LogicalResult optimizeSCCAsync(MutableArrayRef<CallGraphNode *> nodesToVisit,
MLIRContext *context);
/// Optimize the given callable node with one of the pass managers provided
/// with `pipelines`, or the default pipeline. Returns failure if an error
/// occurred during the optimization of the callable, success otherwise.
LogicalResult optimizeCallable(CallGraphNode *node,
llvm::StringMap<OpPassManager> &pipelines);
/// Attempt to initialize the options of this pass from the given string.
/// Derived classes may override this method to hook into the point at which
/// options are initialized, but should generally always invoke this base
/// class variant.
LogicalResult initializeOptions(StringRef options) override;
/// An optional function that constructs a default optimization pipeline for
/// a given operation.
std::function<void(OpPassManager &)> defaultPipeline;
/// A map of operation names to pass pipelines to use when optimizing
/// callable operations of these types. This provides a specialized pipeline
/// instead of the default. The vector size is the number of threads used
/// during optimization.
SmallVector<llvm::StringMap<OpPassManager>, 8> opPipelines;
};
} // end anonymous namespace
InlinerPass::InlinerPass() : InlinerPass(defaultInlinerOptPipeline) {}
InlinerPass::InlinerPass(std::function<void(OpPassManager &)> defaultPipeline)
: defaultPipeline(defaultPipeline) {
opPipelines.push_back({});
// Initialize the pass options with the provided arguments.
if (defaultPipeline) {
OpPassManager fakePM("__mlir_fake_pm_op");
defaultPipeline(fakePM);
llvm::raw_string_ostream strStream(defaultPipelineStr);
fakePM.printAsTextualPipeline(strStream);
}
}
InlinerPass::InlinerPass(std::function<void(OpPassManager &)> defaultPipeline,
llvm::StringMap<OpPassManager> opPipelines)
: InlinerPass(std::move(defaultPipeline)) {
if (opPipelines.empty())
return;
// Update the option for the op specific optimization pipelines.
for (auto &it : opPipelines) {
std::string pipeline;
llvm::raw_string_ostream pipelineOS(pipeline);
pipelineOS << it.getKey() << "(";
it.second.printAsTextualPipeline(pipelineOS);
pipelineOS << ")";
opPipelineStrs.addValue(pipeline);
}
this->opPipelines.emplace_back(std::move(opPipelines));
}
void InlinerPass::runOnOperation() {
CallGraph &cg = getAnalysis<CallGraph>();
auto *context = &getContext();
// The inliner should only be run on operations that define a symbol table,
// as the callgraph will need to resolve references.
Operation *op = getOperation();
if (!op->hasTrait<OpTrait::SymbolTable>()) {
op->emitOpError() << " was scheduled to run under the inliner, but does "
"not define a symbol table";
return signalPassFailure();
}
// Run the inline transform in post-order over the SCCs in the callgraph.
SymbolTableCollection symbolTable;
Inliner inliner(context, cg, symbolTable);
CGUseList useList(getOperation(), cg, symbolTable);
LogicalResult result = runTransformOnCGSCCs(cg, [&](CallGraphSCC &scc) {
return inlineSCC(inliner, useList, scc, context);
});
if (failed(result))
return signalPassFailure();
// After inlining, make sure to erase any callables proven to be dead.
inliner.eraseDeadCallables();
}
LogicalResult InlinerPass::inlineSCC(Inliner &inliner, CGUseList &useList,
CallGraphSCC ¤tSCC,
MLIRContext *context) {
// Continuously simplify and inline until we either reach a fixed point, or
// hit the maximum iteration count. Simplifying early helps to refine the cost
// model, and in future iterations may devirtualize new calls.
unsigned iterationCount = 0;
do {
if (failed(optimizeSCC(inliner.cg, useList, currentSCC, context)))
return failure();
if (failed(inlineCallsInSCC(inliner, useList, currentSCC)))
break;
} while (++iterationCount < maxInliningIterations);
return success();
}
LogicalResult InlinerPass::optimizeSCC(CallGraph &cg, CGUseList &useList,
CallGraphSCC ¤tSCC,
MLIRContext *context) {
// Collect the sets of nodes to simplify.
SmallVector<CallGraphNode *, 4> nodesToVisit;
for (auto *node : currentSCC) {
if (node->isExternal())
continue;
// Don't simplify nodes with children. Nodes with children require special
// handling as we may remove the node during simplification. In the future,
// we should be able to handle this case with proper node deletion tracking.
if (node->hasChildren())
continue;
// We also won't apply simplifications to nodes that can't have passes
// scheduled on them.
auto *region = node->getCallableRegion();
if (!region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
continue;
nodesToVisit.push_back(node);
}
if (nodesToVisit.empty())
return success();
// Optimize each of the nodes within the SCC in parallel.
// NOTE: This is simple now, because we don't enable optimizing nodes within
// children. When we remove this restriction, this logic will need to be
// reworked.
if (context->isMultithreadingEnabled()) {
if (failed(optimizeSCCAsync(nodesToVisit, context)))
return failure();
// Otherwise, we are optimizing within a single thread.
} else {
for (CallGraphNode *node : nodesToVisit) {
if (failed(optimizeCallable(node, opPipelines[0])))
return failure();
}
}
// Recompute the uses held by each of the nodes.
for (CallGraphNode *node : nodesToVisit)
useList.recomputeUses(node, cg);
return success();
}
LogicalResult
InlinerPass::optimizeSCCAsync(MutableArrayRef<CallGraphNode *> nodesToVisit,
MLIRContext *context) {
// Ensure that there are enough pipeline maps for the optimizer to run in
// parallel.
size_t numThreads = llvm::hardware_concurrency().compute_thread_count();
if (opPipelines.size() != numThreads) {
// Reserve before resizing so that we can use a reference to the first
// element.
opPipelines.reserve(numThreads);
opPipelines.resize(numThreads, opPipelines.front());
}
// Ensure an analysis manager has been constructed for each of the nodes.
// This prevents thread races when running the nested pipelines.
for (CallGraphNode *node : nodesToVisit)
getAnalysisManager().nest(node->getCallableRegion()->getParentOp());
// An index for the current node to optimize.
std::atomic<unsigned> nodeIt(0);
// Optimize the nodes of the SCC in parallel.
ParallelDiagnosticHandler optimizerHandler(context);
return llvm::parallelTransformReduce(
llvm::seq<size_t>(0, numThreads), success(),
[](LogicalResult lhs, LogicalResult rhs) {
return success(succeeded(lhs) && succeeded(rhs));
},
[&](size_t index) {
LogicalResult result = success();
for (auto e = nodesToVisit.size(); nodeIt < e && succeeded(result);) {
// Get the next available operation index.
unsigned nextID = nodeIt++;
if (nextID >= e)
break;
// Set the order for this thread so that diagnostics will be
// properly ordered, and reset after optimization has finished.
optimizerHandler.setOrderIDForThread(nextID);
result = optimizeCallable(nodesToVisit[nextID], opPipelines[index]);
optimizerHandler.eraseOrderIDForThread();
}
return result;
});
}
LogicalResult
InlinerPass::optimizeCallable(CallGraphNode *node,
llvm::StringMap<OpPassManager> &pipelines) {
Operation *callable = node->getCallableRegion()->getParentOp();
StringRef opName = callable->getName().getStringRef();
auto pipelineIt = pipelines.find(opName);
if (pipelineIt == pipelines.end()) {
// If a pipeline didn't exist, use the default if possible.
if (!defaultPipeline)
return success();
OpPassManager defaultPM(opName);
defaultPipeline(defaultPM);
pipelineIt = pipelines.try_emplace(opName, std::move(defaultPM)).first;
}
return runPipeline(pipelineIt->second, callable);
}
LogicalResult InlinerPass::initializeOptions(StringRef options) {
if (failed(Pass::initializeOptions(options)))
return failure();
// Initialize the default pipeline builder to use the option string.
if (!defaultPipelineStr.empty()) {
std::string defaultPipelineCopy = defaultPipelineStr;
defaultPipeline = [=](OpPassManager &pm) {
(void)parsePassPipeline(defaultPipelineCopy, pm);
};
} else if (defaultPipelineStr.getNumOccurrences()) {
defaultPipeline = nullptr;
}
// Initialize the op specific pass pipelines.
llvm::StringMap<OpPassManager> pipelines;
for (StringRef pipeline : opPipelineStrs) {
// Skip empty pipelines.
if (pipeline.empty())
continue;
// Pipelines are expected to be of the form `<op-name>(<pipeline>)`.
size_t pipelineStart = pipeline.find_first_of('(');
if (pipelineStart == StringRef::npos || !pipeline.consume_back(")"))
return failure();
StringRef opName = pipeline.take_front(pipelineStart);
OpPassManager pm(opName);
if (failed(parsePassPipeline(pipeline.drop_front(1 + pipelineStart), pm)))
return failure();
pipelines.try_emplace(opName, std::move(pm));
}
opPipelines.assign({std::move(pipelines)});
return success();
}
std::unique_ptr<Pass> mlir::createInlinerPass() {
return std::make_unique<InlinerPass>();
}
std::unique_ptr<Pass>
mlir::createInlinerPass(llvm::StringMap<OpPassManager> opPipelines) {
return std::make_unique<InlinerPass>(defaultInlinerOptPipeline,
std::move(opPipelines));
}
std::unique_ptr<Pass>
createInlinerPass(llvm::StringMap<OpPassManager> opPipelines,
std::function<void(OpPassManager &)> defaultPipelineBuilder) {
return std::make_unique<InlinerPass>(std::move(defaultPipelineBuilder),
std::move(opPipelines));
}