1 //===- LoopFusion.cpp - Code to perform loop fusion -----------------------===//
2 //
3 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements loop fusion.
10 //
11 //===----------------------------------------------------------------------===//
12
13 #include "mlir/Analysis/AffineAnalysis.h"
14 #include "mlir/Analysis/AffineStructures.h"
15 #include "mlir/Analysis/LoopAnalysis.h"
16 #include "mlir/Analysis/Utils.h"
17 #include "mlir/Dialect/AffineOps/AffineOps.h"
18 #include "mlir/Dialect/StandardOps/Ops.h"
19 #include "mlir/IR/AffineExpr.h"
20 #include "mlir/IR/AffineMap.h"
21 #include "mlir/IR/Builders.h"
22 #include "mlir/Pass/Pass.h"
23 #include "mlir/Transforms/LoopFusionUtils.h"
24 #include "mlir/Transforms/LoopUtils.h"
25 #include "mlir/Transforms/Passes.h"
26 #include "mlir/Transforms/Utils.h"
27 #include "llvm/ADT/DenseMap.h"
28 #include "llvm/ADT/DenseSet.h"
29 #include "llvm/ADT/SetVector.h"
30 #include "llvm/Support/CommandLine.h"
31 #include "llvm/Support/Debug.h"
32 #include "llvm/Support/raw_ostream.h"
33 #include <iomanip>
34 #include <sstream>
35 #define DEBUG_TYPE "affine-loop-fusion"
36
37 using llvm::SetVector;
38
39 using namespace mlir;
40
41 static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
42
43 /// Disables fusion profitability check and fuses if valid. Ignore any
44 /// additional (redundant) computation tolerance threshold
45 /// that would have prevented fusion.
46 static llvm::cl::opt<bool>
47 clMaximalLoopFusion("fusion-maximal",
48 llvm::cl::desc("Enables maximal loop fusion"),
49 llvm::cl::cat(clOptionsCategory));
50
51 /// A threshold in percent of additional computation allowed when fusing.
52 static llvm::cl::opt<double> clFusionAddlComputeTolerance(
53 "fusion-compute-tolerance",
54 llvm::cl::desc("Fractional increase in additional "
55 "computation tolerated while fusing"),
56 llvm::cl::cat(clOptionsCategory));
57
58 static llvm::cl::opt<unsigned> clFusionFastMemorySpace(
59 "fusion-fast-mem-space",
60 llvm::cl::desc("Faster memory space number to promote fusion buffers to"),
61 llvm::cl::cat(clOptionsCategory));
62
63 // A local buffer of size less than or equal to this size is automatically
64 // promoted to fast memory after producer-consumer fusion.
65 static llvm::cl::opt<unsigned long long> clFusionLocalBufThreshold(
66 "fusion-local-buf-threshold",
67 llvm::cl::desc("Threshold size (KiB) for promoting local buffers to fast "
68 "memory space"),
69 llvm::cl::cat(clOptionsCategory));
70
71 namespace {
72
73 /// Loop fusion pass. This pass currently supports a greedy fusion policy,
74 /// which fuses loop nests with single-writer/single-reader memref dependences
75 /// with the goal of improving locality.
76
77 // TODO(andydavis) Support fusion of source loop nests which write to multiple
78 // memrefs, where each memref can have multiple users (if profitable).
79 // TODO(andydavis) Extend this pass to check for fusion preventing dependences,
80 // and add support for more general loop fusion algorithms.
81
82 struct LoopFusion : public FunctionPass<LoopFusion> {
LoopFusion__anon5a3accbe0111::LoopFusion83 LoopFusion(unsigned fastMemorySpace = 0, uint64_t localBufSizeThreshold = 0,
84 bool maximalFusion = false)
85 : localBufSizeThreshold(localBufSizeThreshold),
86 fastMemorySpace(fastMemorySpace), maximalFusion(maximalFusion) {}
87
88 void runOnFunction() override;
89
90 // Any local buffers smaller than this size (in bytes) will be created in
91 // `fastMemorySpace` if provided.
92 uint64_t localBufSizeThreshold;
93 Optional<unsigned> fastMemorySpace = None;
94 // If true, ignore any additional (redundant) computation tolerance threshold
95 // that would have prevented fusion.
96 bool maximalFusion;
97
98 // The amount of additional computation that is tolerated while fusing
99 // pair-wise as a fraction of the total computation.
100 constexpr static double kComputeToleranceThreshold = 0.30f;
101 };
102
103 } // end anonymous namespace
104
105 std::unique_ptr<OpPassBase<FuncOp>>
createLoopFusionPass(unsigned fastMemorySpace,uint64_t localBufSizeThreshold,bool maximalFusion)106 mlir::createLoopFusionPass(unsigned fastMemorySpace,
107 uint64_t localBufSizeThreshold, bool maximalFusion) {
108 return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold,
109 maximalFusion);
110 }
111
112 // TODO(b/117228571) Replace when this is modeled through side-effects/op traits
isMemRefDereferencingOp(Operation & op)113 static bool isMemRefDereferencingOp(Operation &op) {
114 if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op) ||
115 isa<AffineDmaStartOp>(op) || isa<AffineDmaWaitOp>(op))
116 return true;
117 return false;
118 }
119
120 namespace {
121
122 // LoopNestStateCollector walks loop nests and collects load and store
123 // operations, and whether or not an IfInst was encountered in the loop nest.
124 struct LoopNestStateCollector {
125 SmallVector<AffineForOp, 4> forOps;
126 SmallVector<Operation *, 4> loadOpInsts;
127 SmallVector<Operation *, 4> storeOpInsts;
128 bool hasNonForRegion = false;
129
collect__anon5a3accbe0211::LoopNestStateCollector130 void collect(Operation *opToWalk) {
131 opToWalk->walk([&](Operation *op) {
132 if (isa<AffineForOp>(op))
133 forOps.push_back(cast<AffineForOp>(op));
134 else if (op->getNumRegions() != 0)
135 hasNonForRegion = true;
136 else if (isa<AffineLoadOp>(op))
137 loadOpInsts.push_back(op);
138 else if (isa<AffineStoreOp>(op))
139 storeOpInsts.push_back(op);
140 });
141 }
142 };
143
144 // MemRefDependenceGraph is a graph data structure where graph nodes are
145 // top-level operations in a FuncOp which contain load/store ops, and edges
146 // are memref dependences between the nodes.
147 // TODO(andydavis) Add a more flexible dependence graph representation.
148 // TODO(andydavis) Add a depth parameter to dependence graph construction.
149 struct MemRefDependenceGraph {
150 public:
151 // Node represents a node in the graph. A Node is either an entire loop nest
152 // rooted at the top level which contains loads/stores, or a top level
153 // load/store.
154 struct Node {
155 // The unique identifier of this node in the graph.
156 unsigned id;
157 // The top-level statement which is (or contains) a load/store.
158 Operation *op;
159 // List of load operations.
160 SmallVector<Operation *, 4> loads;
161 // List of store op insts.
162 SmallVector<Operation *, 4> stores;
Node__anon5a3accbe0211::MemRefDependenceGraph::Node163 Node(unsigned id, Operation *op) : id(id), op(op) {}
164
165 // Returns the load op count for 'memref'.
getLoadOpCount__anon5a3accbe0211::MemRefDependenceGraph::Node166 unsigned getLoadOpCount(Value memref) {
167 unsigned loadOpCount = 0;
168 for (auto *loadOpInst : loads) {
169 if (memref == cast<AffineLoadOp>(loadOpInst).getMemRef())
170 ++loadOpCount;
171 }
172 return loadOpCount;
173 }
174
175 // Returns the store op count for 'memref'.
getStoreOpCount__anon5a3accbe0211::MemRefDependenceGraph::Node176 unsigned getStoreOpCount(Value memref) {
177 unsigned storeOpCount = 0;
178 for (auto *storeOpInst : stores) {
179 if (memref == cast<AffineStoreOp>(storeOpInst).getMemRef())
180 ++storeOpCount;
181 }
182 return storeOpCount;
183 }
184
185 // Returns all store ops in 'storeOps' which access 'memref'.
getStoreOpsForMemref__anon5a3accbe0211::MemRefDependenceGraph::Node186 void getStoreOpsForMemref(Value memref,
187 SmallVectorImpl<Operation *> *storeOps) {
188 for (auto *storeOpInst : stores) {
189 if (memref == cast<AffineStoreOp>(storeOpInst).getMemRef())
190 storeOps->push_back(storeOpInst);
191 }
192 }
193
194 // Returns all load ops in 'loadOps' which access 'memref'.
getLoadOpsForMemref__anon5a3accbe0211::MemRefDependenceGraph::Node195 void getLoadOpsForMemref(Value memref,
196 SmallVectorImpl<Operation *> *loadOps) {
197 for (auto *loadOpInst : loads) {
198 if (memref == cast<AffineLoadOp>(loadOpInst).getMemRef())
199 loadOps->push_back(loadOpInst);
200 }
201 }
202
203 // Returns all memrefs in 'loadAndStoreMemrefSet' for which this node
204 // has at least one load and store operation.
getLoadAndStoreMemrefSet__anon5a3accbe0211::MemRefDependenceGraph::Node205 void getLoadAndStoreMemrefSet(DenseSet<Value> *loadAndStoreMemrefSet) {
206 llvm::SmallDenseSet<Value, 2> loadMemrefs;
207 for (auto *loadOpInst : loads) {
208 loadMemrefs.insert(cast<AffineLoadOp>(loadOpInst).getMemRef());
209 }
210 for (auto *storeOpInst : stores) {
211 auto memref = cast<AffineStoreOp>(storeOpInst).getMemRef();
212 if (loadMemrefs.count(memref) > 0)
213 loadAndStoreMemrefSet->insert(memref);
214 }
215 }
216 };
217
218 // Edge represents a data dependence between nodes in the graph.
219 struct Edge {
220 // The id of the node at the other end of the edge.
221 // If this edge is stored in Edge = Node.inEdges[i], then
222 // 'Node.inEdges[i].id' is the identifier of the source node of the edge.
223 // If this edge is stored in Edge = Node.outEdges[i], then
224 // 'Node.outEdges[i].id' is the identifier of the dest node of the edge.
225 unsigned id;
226 // The SSA value on which this edge represents a dependence.
227 // If the value is a memref, then the dependence is between graph nodes
228 // which contain accesses to the same memref 'value'. If the value is a
229 // non-memref value, then the dependence is between a graph node which
230 // defines an SSA value and another graph node which uses the SSA value
231 // (e.g. a constant operation defining a value which is used inside a loop
232 // nest).
233 Value value;
234 };
235
236 // Map from node id to Node.
237 DenseMap<unsigned, Node> nodes;
238 // Map from node id to list of input edges.
239 DenseMap<unsigned, SmallVector<Edge, 2>> inEdges;
240 // Map from node id to list of output edges.
241 DenseMap<unsigned, SmallVector<Edge, 2>> outEdges;
242 // Map from memref to a count on the dependence edges associated with that
243 // memref.
244 DenseMap<Value, unsigned> memrefEdgeCount;
245 // The next unique identifier to use for newly created graph nodes.
246 unsigned nextNodeId = 0;
247
MemRefDependenceGraph__anon5a3accbe0211::MemRefDependenceGraph248 MemRefDependenceGraph() {}
249
250 // Initializes the dependence graph based on operations in 'f'.
251 // Returns true on success, false otherwise.
252 bool init(FuncOp f);
253
254 // Returns the graph node for 'id'.
getNode__anon5a3accbe0211::MemRefDependenceGraph255 Node *getNode(unsigned id) {
256 auto it = nodes.find(id);
257 assert(it != nodes.end());
258 return &it->second;
259 }
260
261 // Returns the graph node for 'forOp'.
getForOpNode__anon5a3accbe0211::MemRefDependenceGraph262 Node *getForOpNode(AffineForOp forOp) {
263 for (auto &idAndNode : nodes)
264 if (idAndNode.second.op == forOp.getOperation())
265 return &idAndNode.second;
266 return nullptr;
267 }
268
269 // Adds a node with 'op' to the graph and returns its unique identifier.
addNode__anon5a3accbe0211::MemRefDependenceGraph270 unsigned addNode(Operation *op) {
271 Node node(nextNodeId++, op);
272 nodes.insert({node.id, node});
273 return node.id;
274 }
275
276 // Remove node 'id' (and its associated edges) from graph.
removeNode__anon5a3accbe0211::MemRefDependenceGraph277 void removeNode(unsigned id) {
278 // Remove each edge in 'inEdges[id]'.
279 if (inEdges.count(id) > 0) {
280 SmallVector<Edge, 2> oldInEdges = inEdges[id];
281 for (auto &inEdge : oldInEdges) {
282 removeEdge(inEdge.id, id, inEdge.value);
283 }
284 }
285 // Remove each edge in 'outEdges[id]'.
286 if (outEdges.count(id) > 0) {
287 SmallVector<Edge, 2> oldOutEdges = outEdges[id];
288 for (auto &outEdge : oldOutEdges) {
289 removeEdge(id, outEdge.id, outEdge.value);
290 }
291 }
292 // Erase remaining node state.
293 inEdges.erase(id);
294 outEdges.erase(id);
295 nodes.erase(id);
296 }
297
298 // Returns true if node 'id' writes to any memref which escapes (or is an
299 // argument to) the function/block. Returns false otherwise.
writesToLiveInOrEscapingMemrefs__anon5a3accbe0211::MemRefDependenceGraph300 bool writesToLiveInOrEscapingMemrefs(unsigned id) {
301 Node *node = getNode(id);
302 for (auto *storeOpInst : node->stores) {
303 auto memref = cast<AffineStoreOp>(storeOpInst).getMemRef();
304 auto *op = memref.getDefiningOp();
305 // Return true if 'memref' is a block argument.
306 if (!op)
307 return true;
308 // Return true if any use of 'memref' escapes the function.
309 for (auto *user : memref.getUsers())
310 if (!isMemRefDereferencingOp(*user))
311 return true;
312 }
313 return false;
314 }
315
316 // Returns the unique AffineStoreOp in `node` that meets all the following:
317 // *) store is the only one that writes to a function-local memref live out
318 // of `node`,
319 // *) store is not the source of a self-dependence on `node`.
320 // Otherwise, returns a null AffineStoreOp.
getUniqueOutgoingStore__anon5a3accbe0211::MemRefDependenceGraph321 AffineStoreOp getUniqueOutgoingStore(Node *node) {
322 AffineStoreOp uniqueStore;
323
324 // Return null if `node` doesn't have any outgoing edges.
325 auto outEdgeIt = outEdges.find(node->id);
326 if (outEdgeIt == outEdges.end())
327 return nullptr;
328
329 const auto &nodeOutEdges = outEdgeIt->second;
330 for (auto *op : node->stores) {
331 auto storeOp = cast<AffineStoreOp>(op);
332 auto memref = storeOp.getMemRef();
333 // Skip this store if there are no dependences on its memref. This means
334 // that store either:
335 // *) writes to a memref that is only read within the same loop nest
336 // (self-dependence edges are not represented in graph at the moment),
337 // *) writes to a function live out memref (function parameter), or
338 // *) is dead.
339 if (llvm::all_of(nodeOutEdges, [=](const Edge &edge) {
340 return (edge.value != memref);
341 }))
342 continue;
343
344 if (uniqueStore)
345 // Found multiple stores to function-local live-out memrefs.
346 return nullptr;
347 // Found first store to function-local live-out memref.
348 uniqueStore = storeOp;
349 }
350
351 return uniqueStore;
352 }
353
354 // Returns true if node 'id' can be removed from the graph. Returns false
355 // otherwise. A node can be removed from the graph iff the following
356 // conditions are met:
357 // *) The node does not write to any memref which escapes (or is a
358 // function/block argument).
359 // *) The node has no successors in the dependence graph.
canRemoveNode__anon5a3accbe0211::MemRefDependenceGraph360 bool canRemoveNode(unsigned id) {
361 if (writesToLiveInOrEscapingMemrefs(id))
362 return false;
363 Node *node = getNode(id);
364 for (auto *storeOpInst : node->stores) {
365 // Return false if there exist out edges from 'id' on 'memref'.
366 if (getOutEdgeCount(id, cast<AffineStoreOp>(storeOpInst).getMemRef()) > 0)
367 return false;
368 }
369 return true;
370 }
371
372 // Returns true iff there is an edge from node 'srcId' to node 'dstId' which
373 // is for 'value' if non-null, or for any value otherwise. Returns false
374 // otherwise.
hasEdge__anon5a3accbe0211::MemRefDependenceGraph375 bool hasEdge(unsigned srcId, unsigned dstId, Value value = nullptr) {
376 if (outEdges.count(srcId) == 0 || inEdges.count(dstId) == 0) {
377 return false;
378 }
379 bool hasOutEdge = llvm::any_of(outEdges[srcId], [=](Edge &edge) {
380 return edge.id == dstId && (!value || edge.value == value);
381 });
382 bool hasInEdge = llvm::any_of(inEdges[dstId], [=](Edge &edge) {
383 return edge.id == srcId && (!value || edge.value == value);
384 });
385 return hasOutEdge && hasInEdge;
386 }
387
388 // Adds an edge from node 'srcId' to node 'dstId' for 'value'.
addEdge__anon5a3accbe0211::MemRefDependenceGraph389 void addEdge(unsigned srcId, unsigned dstId, Value value) {
390 if (!hasEdge(srcId, dstId, value)) {
391 outEdges[srcId].push_back({dstId, value});
392 inEdges[dstId].push_back({srcId, value});
393 if (value.getType().isa<MemRefType>())
394 memrefEdgeCount[value]++;
395 }
396 }
397
398 // Removes an edge from node 'srcId' to node 'dstId' for 'value'.
removeEdge__anon5a3accbe0211::MemRefDependenceGraph399 void removeEdge(unsigned srcId, unsigned dstId, Value value) {
400 assert(inEdges.count(dstId) > 0);
401 assert(outEdges.count(srcId) > 0);
402 if (value.getType().isa<MemRefType>()) {
403 assert(memrefEdgeCount.count(value) > 0);
404 memrefEdgeCount[value]--;
405 }
406 // Remove 'srcId' from 'inEdges[dstId]'.
407 for (auto it = inEdges[dstId].begin(); it != inEdges[dstId].end(); ++it) {
408 if ((*it).id == srcId && (*it).value == value) {
409 inEdges[dstId].erase(it);
410 break;
411 }
412 }
413 // Remove 'dstId' from 'outEdges[srcId]'.
414 for (auto it = outEdges[srcId].begin(); it != outEdges[srcId].end(); ++it) {
415 if ((*it).id == dstId && (*it).value == value) {
416 outEdges[srcId].erase(it);
417 break;
418 }
419 }
420 }
421
422 // Returns true if there is a path in the dependence graph from node 'srcId'
423 // to node 'dstId'. Returns false otherwise.
hasDependencePath__anon5a3accbe0211::MemRefDependenceGraph424 bool hasDependencePath(unsigned srcId, unsigned dstId) {
425 // Worklist state is: <node-id, next-output-edge-index-to-visit>
426 SmallVector<std::pair<unsigned, unsigned>, 4> worklist;
427 worklist.push_back({srcId, 0});
428 // Run DFS traversal to see if 'dstId' is reachable from 'srcId'.
429 while (!worklist.empty()) {
430 auto &idAndIndex = worklist.back();
431 // Return true if we have reached 'dstId'.
432 if (idAndIndex.first == dstId)
433 return true;
434 // Pop and continue if node has no out edges, or if all out edges have
435 // already been visited.
436 if (outEdges.count(idAndIndex.first) == 0 ||
437 idAndIndex.second == outEdges[idAndIndex.first].size()) {
438 worklist.pop_back();
439 continue;
440 }
441 // Get graph edge to traverse.
442 Edge edge = outEdges[idAndIndex.first][idAndIndex.second];
443 // Increment next output edge index for 'idAndIndex'.
444 ++idAndIndex.second;
445 // Add node at 'edge.id' to worklist.
446 worklist.push_back({edge.id, 0});
447 }
448 return false;
449 }
450
451 // Returns the input edge count for node 'id' and 'memref' from src nodes
452 // which access 'memref' with a store operation.
getIncomingMemRefAccesses__anon5a3accbe0211::MemRefDependenceGraph453 unsigned getIncomingMemRefAccesses(unsigned id, Value memref) {
454 unsigned inEdgeCount = 0;
455 if (inEdges.count(id) > 0)
456 for (auto &inEdge : inEdges[id])
457 if (inEdge.value == memref) {
458 Node *srcNode = getNode(inEdge.id);
459 // Only count in edges from 'srcNode' if 'srcNode' accesses 'memref'
460 if (srcNode->getStoreOpCount(memref) > 0)
461 ++inEdgeCount;
462 }
463 return inEdgeCount;
464 }
465
466 // Returns the output edge count for node 'id' and 'memref' (if non-null),
467 // otherwise returns the total output edge count from node 'id'.
getOutEdgeCount__anon5a3accbe0211::MemRefDependenceGraph468 unsigned getOutEdgeCount(unsigned id, Value memref = nullptr) {
469 unsigned outEdgeCount = 0;
470 if (outEdges.count(id) > 0)
471 for (auto &outEdge : outEdges[id])
472 if (!memref || outEdge.value == memref)
473 ++outEdgeCount;
474 return outEdgeCount;
475 }
476
477 // Computes and returns an insertion point operation, before which the
478 // the fused <srcId, dstId> loop nest can be inserted while preserving
479 // dependences. Returns nullptr if no such insertion point is found.
getFusedLoopNestInsertionPoint__anon5a3accbe0211::MemRefDependenceGraph480 Operation *getFusedLoopNestInsertionPoint(unsigned srcId, unsigned dstId) {
481 if (outEdges.count(srcId) == 0)
482 return getNode(dstId)->op;
483
484 // Build set of insts in range (srcId, dstId) which depend on 'srcId'.
485 SmallPtrSet<Operation *, 2> srcDepInsts;
486 for (auto &outEdge : outEdges[srcId])
487 if (outEdge.id != dstId)
488 srcDepInsts.insert(getNode(outEdge.id)->op);
489
490 // Build set of insts in range (srcId, dstId) on which 'dstId' depends.
491 SmallPtrSet<Operation *, 2> dstDepInsts;
492 for (auto &inEdge : inEdges[dstId])
493 if (inEdge.id != srcId)
494 dstDepInsts.insert(getNode(inEdge.id)->op);
495
496 Operation *srcNodeInst = getNode(srcId)->op;
497 Operation *dstNodeInst = getNode(dstId)->op;
498
499 // Computing insertion point:
500 // *) Walk all operation positions in Block operation list in the
501 // range (src, dst). For each operation 'op' visited in this search:
502 // *) Store in 'firstSrcDepPos' the first position where 'op' has a
503 // dependence edge from 'srcNode'.
504 // *) Store in 'lastDstDepPost' the last position where 'op' has a
505 // dependence edge to 'dstNode'.
506 // *) Compare 'firstSrcDepPos' and 'lastDstDepPost' to determine the
507 // operation insertion point (or return null pointer if no such
508 // insertion point exists: 'firstSrcDepPos' <= 'lastDstDepPos').
509 SmallVector<Operation *, 2> depInsts;
510 Optional<unsigned> firstSrcDepPos;
511 Optional<unsigned> lastDstDepPos;
512 unsigned pos = 0;
513 for (Block::iterator it = std::next(Block::iterator(srcNodeInst));
514 it != Block::iterator(dstNodeInst); ++it) {
515 Operation *op = &(*it);
516 if (srcDepInsts.count(op) > 0 && firstSrcDepPos == None)
517 firstSrcDepPos = pos;
518 if (dstDepInsts.count(op) > 0)
519 lastDstDepPos = pos;
520 depInsts.push_back(op);
521 ++pos;
522 }
523
524 if (firstSrcDepPos.hasValue()) {
525 if (lastDstDepPos.hasValue()) {
526 if (firstSrcDepPos.getValue() <= lastDstDepPos.getValue()) {
527 // No valid insertion point exists which preserves dependences.
528 return nullptr;
529 }
530 }
531 // Return the insertion point at 'firstSrcDepPos'.
532 return depInsts[firstSrcDepPos.getValue()];
533 }
534 // No dependence targets in range (or only dst deps in range), return
535 // 'dstNodInst' insertion point.
536 return dstNodeInst;
537 }
538
539 // Updates edge mappings from node 'srcId' to node 'dstId' after 'oldMemRef'
540 // has been replaced in node at 'dstId' by a private memref depending
541 // on the value of 'createPrivateMemRef'.
updateEdges__anon5a3accbe0211::MemRefDependenceGraph542 void updateEdges(unsigned srcId, unsigned dstId, Value oldMemRef,
543 bool createPrivateMemRef) {
544 // For each edge in 'inEdges[srcId]': add new edge remaping to 'dstId'.
545 if (inEdges.count(srcId) > 0) {
546 SmallVector<Edge, 2> oldInEdges = inEdges[srcId];
547 for (auto &inEdge : oldInEdges) {
548 // Add edge from 'inEdge.id' to 'dstId' if not for 'oldMemRef'.
549 if (inEdge.value != oldMemRef)
550 addEdge(inEdge.id, dstId, inEdge.value);
551 }
552 }
553 // For each edge in 'outEdges[srcId]': remove edge from 'srcId' to 'dstId'.
554 if (outEdges.count(srcId) > 0) {
555 SmallVector<Edge, 2> oldOutEdges = outEdges[srcId];
556 for (auto &outEdge : oldOutEdges) {
557 // Remove any out edges from 'srcId' to 'dstId' across memrefs.
558 if (outEdge.id == dstId)
559 removeEdge(srcId, outEdge.id, outEdge.value);
560 }
561 }
562 // Remove any edges in 'inEdges[dstId]' on 'oldMemRef' (which is being
563 // replaced by a private memref). These edges could come from nodes
564 // other than 'srcId' which were removed in the previous step.
565 if (inEdges.count(dstId) > 0 && createPrivateMemRef) {
566 SmallVector<Edge, 2> oldInEdges = inEdges[dstId];
567 for (auto &inEdge : oldInEdges)
568 if (inEdge.value == oldMemRef)
569 removeEdge(inEdge.id, dstId, inEdge.value);
570 }
571 }
572
573 // Update edge mappings for nodes 'sibId' and 'dstId' to reflect fusion
574 // of sibling node 'sidId' into node 'dstId'.
updateEdges__anon5a3accbe0211::MemRefDependenceGraph575 void updateEdges(unsigned sibId, unsigned dstId) {
576 // For each edge in 'inEdges[sibId]':
577 // *) Add new edge from source node 'inEdge.id' to 'dstNode'.
578 // *) Remove edge from source node 'inEdge.id' to 'sibNode'.
579 if (inEdges.count(sibId) > 0) {
580 SmallVector<Edge, 2> oldInEdges = inEdges[sibId];
581 for (auto &inEdge : oldInEdges) {
582 addEdge(inEdge.id, dstId, inEdge.value);
583 removeEdge(inEdge.id, sibId, inEdge.value);
584 }
585 }
586
587 // For each edge in 'outEdges[sibId]' to node 'id'
588 // *) Add new edge from 'dstId' to 'outEdge.id'.
589 // *) Remove edge from 'sibId' to 'outEdge.id'.
590 if (outEdges.count(sibId) > 0) {
591 SmallVector<Edge, 2> oldOutEdges = outEdges[sibId];
592 for (auto &outEdge : oldOutEdges) {
593 addEdge(dstId, outEdge.id, outEdge.value);
594 removeEdge(sibId, outEdge.id, outEdge.value);
595 }
596 }
597 }
598
599 // Adds ops in 'loads' and 'stores' to node at 'id'.
addToNode__anon5a3accbe0211::MemRefDependenceGraph600 void addToNode(unsigned id, const SmallVectorImpl<Operation *> &loads,
601 const SmallVectorImpl<Operation *> &stores) {
602 Node *node = getNode(id);
603 for (auto *loadOpInst : loads)
604 node->loads.push_back(loadOpInst);
605 for (auto *storeOpInst : stores)
606 node->stores.push_back(storeOpInst);
607 }
608
clearNodeLoadAndStores__anon5a3accbe0211::MemRefDependenceGraph609 void clearNodeLoadAndStores(unsigned id) {
610 Node *node = getNode(id);
611 node->loads.clear();
612 node->stores.clear();
613 }
614
615 // Calls 'callback' for each input edge incident to node 'id' which carries a
616 // memref dependence.
forEachMemRefInputEdge__anon5a3accbe0211::MemRefDependenceGraph617 void forEachMemRefInputEdge(unsigned id,
618 const std::function<void(Edge)> &callback) {
619 if (inEdges.count(id) > 0)
620 forEachMemRefEdge(inEdges[id], callback);
621 }
622
623 // Calls 'callback' for each output edge from node 'id' which carries a
624 // memref dependence.
forEachMemRefOutputEdge__anon5a3accbe0211::MemRefDependenceGraph625 void forEachMemRefOutputEdge(unsigned id,
626 const std::function<void(Edge)> &callback) {
627 if (outEdges.count(id) > 0)
628 forEachMemRefEdge(outEdges[id], callback);
629 }
630
631 // Calls 'callback' for each edge in 'edges' which carries a memref
632 // dependence.
forEachMemRefEdge__anon5a3accbe0211::MemRefDependenceGraph633 void forEachMemRefEdge(ArrayRef<Edge> edges,
634 const std::function<void(Edge)> &callback) {
635 for (auto &edge : edges) {
636 // Skip if 'edge' is not a memref dependence edge.
637 if (!edge.value.getType().isa<MemRefType>())
638 continue;
639 assert(nodes.count(edge.id) > 0);
640 // Skip if 'edge.id' is not a loop nest.
641 if (!isa<AffineForOp>(getNode(edge.id)->op))
642 continue;
643 // Visit current input edge 'edge'.
644 callback(edge);
645 }
646 }
647
print__anon5a3accbe0211::MemRefDependenceGraph648 void print(raw_ostream &os) const {
649 os << "\nMemRefDependenceGraph\n";
650 os << "\nNodes:\n";
651 for (auto &idAndNode : nodes) {
652 os << "Node: " << idAndNode.first << "\n";
653 auto it = inEdges.find(idAndNode.first);
654 if (it != inEdges.end()) {
655 for (const auto &e : it->second)
656 os << " InEdge: " << e.id << " " << e.value << "\n";
657 }
658 it = outEdges.find(idAndNode.first);
659 if (it != outEdges.end()) {
660 for (const auto &e : it->second)
661 os << " OutEdge: " << e.id << " " << e.value << "\n";
662 }
663 }
664 }
dump__anon5a3accbe0211::MemRefDependenceGraph665 void dump() const { print(llvm::errs()); }
666 };
667
668 } // end anonymous namespace
669
670 // Initializes the data dependence graph by walking operations in 'f'.
671 // Assigns each node in the graph a node id based on program order in 'f'.
672 // TODO(andydavis) Add support for taking a Block arg to construct the
673 // dependence graph at a different depth.
init(FuncOp f)674 bool MemRefDependenceGraph::init(FuncOp f) {
675 DenseMap<Value, SetVector<unsigned>> memrefAccesses;
676
677 // TODO: support multi-block functions.
678 if (f.getBlocks().size() != 1)
679 return false;
680
681 DenseMap<Operation *, unsigned> forToNodeMap;
682 for (auto &op : f.front()) {
683 if (auto forOp = dyn_cast<AffineForOp>(op)) {
684 // Create graph node 'id' to represent top-level 'forOp' and record
685 // all loads and store accesses it contains.
686 LoopNestStateCollector collector;
687 collector.collect(&op);
688 // Return false if a non 'affine.for' region was found (not currently
689 // supported).
690 if (collector.hasNonForRegion)
691 return false;
692 Node node(nextNodeId++, &op);
693 for (auto *opInst : collector.loadOpInsts) {
694 node.loads.push_back(opInst);
695 auto memref = cast<AffineLoadOp>(opInst).getMemRef();
696 memrefAccesses[memref].insert(node.id);
697 }
698 for (auto *opInst : collector.storeOpInsts) {
699 node.stores.push_back(opInst);
700 auto memref = cast<AffineStoreOp>(opInst).getMemRef();
701 memrefAccesses[memref].insert(node.id);
702 }
703 forToNodeMap[&op] = node.id;
704 nodes.insert({node.id, node});
705 } else if (auto loadOp = dyn_cast<AffineLoadOp>(op)) {
706 // Create graph node for top-level load op.
707 Node node(nextNodeId++, &op);
708 node.loads.push_back(&op);
709 auto memref = cast<AffineLoadOp>(op).getMemRef();
710 memrefAccesses[memref].insert(node.id);
711 nodes.insert({node.id, node});
712 } else if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
713 // Create graph node for top-level store op.
714 Node node(nextNodeId++, &op);
715 node.stores.push_back(&op);
716 auto memref = cast<AffineStoreOp>(op).getMemRef();
717 memrefAccesses[memref].insert(node.id);
718 nodes.insert({node.id, node});
719 } else if (op.getNumRegions() != 0) {
720 // Return false if another region is found (not currently supported).
721 return false;
722 } else if (op.getNumResults() > 0 && !op.use_empty()) {
723 // Create graph node for top-level producer of SSA values, which
724 // could be used by loop nest nodes.
725 Node node(nextNodeId++, &op);
726 nodes.insert({node.id, node});
727 }
728 }
729
730 // Add dependence edges between nodes which produce SSA values and their
731 // users.
732 for (auto &idAndNode : nodes) {
733 const Node &node = idAndNode.second;
734 if (!node.loads.empty() || !node.stores.empty())
735 continue;
736 auto *opInst = node.op;
737 for (auto value : opInst->getResults()) {
738 for (auto *user : value.getUsers()) {
739 SmallVector<AffineForOp, 4> loops;
740 getLoopIVs(*user, &loops);
741 if (loops.empty())
742 continue;
743 assert(forToNodeMap.count(loops[0].getOperation()) > 0);
744 unsigned userLoopNestId = forToNodeMap[loops[0].getOperation()];
745 addEdge(node.id, userLoopNestId, value);
746 }
747 }
748 }
749
750 // Walk memref access lists and add graph edges between dependent nodes.
751 for (auto &memrefAndList : memrefAccesses) {
752 unsigned n = memrefAndList.second.size();
753 for (unsigned i = 0; i < n; ++i) {
754 unsigned srcId = memrefAndList.second[i];
755 bool srcHasStore =
756 getNode(srcId)->getStoreOpCount(memrefAndList.first) > 0;
757 for (unsigned j = i + 1; j < n; ++j) {
758 unsigned dstId = memrefAndList.second[j];
759 bool dstHasStore =
760 getNode(dstId)->getStoreOpCount(memrefAndList.first) > 0;
761 if (srcHasStore || dstHasStore)
762 addEdge(srcId, dstId, memrefAndList.first);
763 }
764 }
765 }
766 return true;
767 }
768
769 // Removes load operations from 'srcLoads' which operate on 'memref', and
770 // adds them to 'dstLoads'.
moveLoadsAccessingMemrefTo(Value memref,SmallVectorImpl<Operation * > * srcLoads,SmallVectorImpl<Operation * > * dstLoads)771 static void moveLoadsAccessingMemrefTo(Value memref,
772 SmallVectorImpl<Operation *> *srcLoads,
773 SmallVectorImpl<Operation *> *dstLoads) {
774 dstLoads->clear();
775 SmallVector<Operation *, 4> srcLoadsToKeep;
776 for (auto *load : *srcLoads) {
777 if (cast<AffineLoadOp>(load).getMemRef() == memref)
778 dstLoads->push_back(load);
779 else
780 srcLoadsToKeep.push_back(load);
781 }
782 srcLoads->swap(srcLoadsToKeep);
783 }
784
785 // Returns the innermost common loop depth for the set of operations in 'ops'.
getInnermostCommonLoopDepth(ArrayRef<Operation * > ops)786 static unsigned getInnermostCommonLoopDepth(ArrayRef<Operation *> ops) {
787 unsigned numOps = ops.size();
788 assert(numOps > 0);
789
790 std::vector<SmallVector<AffineForOp, 4>> loops(numOps);
791 unsigned loopDepthLimit = std::numeric_limits<unsigned>::max();
792 for (unsigned i = 0; i < numOps; ++i) {
793 getLoopIVs(*ops[i], &loops[i]);
794 loopDepthLimit =
795 std::min(loopDepthLimit, static_cast<unsigned>(loops[i].size()));
796 }
797
798 unsigned loopDepth = 0;
799 for (unsigned d = 0; d < loopDepthLimit; ++d) {
800 unsigned i;
801 for (i = 1; i < numOps; ++i) {
802 if (loops[i - 1][d] != loops[i][d])
803 break;
804 }
805 if (i != numOps)
806 break;
807 ++loopDepth;
808 }
809 return loopDepth;
810 }
811
812 // Returns the maximum loop depth at which no dependences between 'loadOpInsts'
813 // and 'storeOpInsts' are satisfied.
getMaxLoopDepth(ArrayRef<Operation * > loadOpInsts,ArrayRef<Operation * > storeOpInsts)814 static unsigned getMaxLoopDepth(ArrayRef<Operation *> loadOpInsts,
815 ArrayRef<Operation *> storeOpInsts) {
816 // Merge loads and stores into the same array.
817 SmallVector<Operation *, 2> ops(loadOpInsts.begin(), loadOpInsts.end());
818 ops.append(storeOpInsts.begin(), storeOpInsts.end());
819
820 // Compute the innermost common loop depth for loads and stores.
821 unsigned loopDepth = getInnermostCommonLoopDepth(ops);
822
823 // Return common loop depth for loads if there are no store ops.
824 if (storeOpInsts.empty())
825 return loopDepth;
826
827 // Check dependences on all pairs of ops in 'ops' and store the minimum
828 // loop depth at which a dependence is satisfied.
829 for (unsigned i = 0, e = ops.size(); i < e; ++i) {
830 auto *srcOpInst = ops[i];
831 MemRefAccess srcAccess(srcOpInst);
832 for (unsigned j = 0; j < e; ++j) {
833 auto *dstOpInst = ops[j];
834 MemRefAccess dstAccess(dstOpInst);
835
836 unsigned numCommonLoops =
837 getNumCommonSurroundingLoops(*srcOpInst, *dstOpInst);
838 for (unsigned d = 1; d <= numCommonLoops + 1; ++d) {
839 FlatAffineConstraints dependenceConstraints;
840 // TODO(andydavis) Cache dependence analysis results, check cache here.
841 DependenceResult result = checkMemrefAccessDependence(
842 srcAccess, dstAccess, d, &dependenceConstraints,
843 /*dependenceComponents=*/nullptr);
844 if (hasDependence(result)) {
845 // Store minimum loop depth and break because we want the min 'd' at
846 // which there is a dependence.
847 loopDepth = std::min(loopDepth, d - 1);
848 break;
849 }
850 }
851 }
852 }
853 return loopDepth;
854 }
855
856 // Sinks all sequential loops to the innermost levels (while preserving
857 // relative order among them) and moves all parallel loops to the
858 // outermost (while again preserving relative order among them).
859 // This can increase the loop depth at which we can fuse a slice, since we are
860 // pushing loop carried dependence to a greater depth in the loop nest.
sinkSequentialLoops(MemRefDependenceGraph::Node * node)861 static void sinkSequentialLoops(MemRefDependenceGraph::Node *node) {
862 assert(isa<AffineForOp>(node->op));
863 AffineForOp newRootForOp = sinkSequentialLoops(cast<AffineForOp>(node->op));
864 node->op = newRootForOp.getOperation();
865 }
866
867 // TODO(mlir-team): improve/complete this when we have target data.
getMemRefEltSizeInBytes(MemRefType memRefType)868 static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
869 auto elementType = memRefType.getElementType();
870
871 unsigned sizeInBits;
872 if (elementType.isIntOrFloat()) {
873 sizeInBits = elementType.getIntOrFloatBitWidth();
874 } else {
875 auto vectorType = elementType.cast<VectorType>();
876 sizeInBits =
877 vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
878 }
879 return llvm::divideCeil(sizeInBits, 8);
880 }
881
882 // Creates and returns a private (single-user) memref for fused loop rooted
883 // at 'forOp', with (potentially reduced) memref size based on the
884 // MemRefRegion written to by 'srcStoreOpInst' at depth 'dstLoopDepth'.
885 // TODO(bondhugula): consider refactoring the common code from generateDma and
886 // this one.
createPrivateMemRef(AffineForOp forOp,Operation * srcStoreOpInst,unsigned dstLoopDepth,Optional<unsigned> fastMemorySpace,uint64_t localBufSizeThreshold)887 static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
888 unsigned dstLoopDepth,
889 Optional<unsigned> fastMemorySpace,
890 uint64_t localBufSizeThreshold) {
891 auto *forInst = forOp.getOperation();
892
893 // Create builder to insert alloc op just before 'forOp'.
894 OpBuilder b(forInst);
895 // Builder to create constants at the top level.
896 OpBuilder top(forInst->getParentOfType<FuncOp>().getBody());
897 // Create new memref type based on slice bounds.
898 auto oldMemRef = cast<AffineStoreOp>(srcStoreOpInst).getMemRef();
899 auto oldMemRefType = oldMemRef.getType().cast<MemRefType>();
900 unsigned rank = oldMemRefType.getRank();
901
902 // Compute MemRefRegion for 'srcStoreOpInst' at depth 'dstLoopDepth'.
903 MemRefRegion region(srcStoreOpInst->getLoc());
904 bool validRegion = succeeded(region.compute(srcStoreOpInst, dstLoopDepth));
905 (void)validRegion;
906 assert(validRegion && "unexpected memref region failure");
907 SmallVector<int64_t, 4> newShape;
908 std::vector<SmallVector<int64_t, 4>> lbs;
909 SmallVector<int64_t, 8> lbDivisors;
910 lbs.reserve(rank);
911 // Query 'region' for 'newShape' and lower bounds of MemRefRegion accessed
912 // by 'srcStoreOpInst' at depth 'dstLoopDepth'.
913 Optional<int64_t> numElements =
914 region.getConstantBoundingSizeAndShape(&newShape, &lbs, &lbDivisors);
915 assert(numElements.hasValue() &&
916 "non-constant number of elts in local buffer");
917
918 const FlatAffineConstraints *cst = region.getConstraints();
919 // 'outerIVs' holds the values that this memory region is symbolic/parametric
920 // on; this would correspond to loop IVs surrounding the level at which the
921 // slice is being materialized.
922 SmallVector<Value, 8> outerIVs;
923 cst->getIdValues(rank, cst->getNumIds(), &outerIVs);
924
925 // Build 'rank' AffineExprs from MemRefRegion 'lbs'
926 SmallVector<AffineExpr, 4> offsets;
927 offsets.reserve(rank);
928 for (unsigned d = 0; d < rank; ++d) {
929 assert(lbs[d].size() == cst->getNumCols() - rank && "incorrect bound size");
930
931 AffineExpr offset = top.getAffineConstantExpr(0);
932 for (unsigned j = 0, e = cst->getNumCols() - rank - 1; j < e; j++) {
933 offset = offset + lbs[d][j] * top.getAffineDimExpr(j);
934 }
935 assert(lbDivisors[d] > 0);
936 offset =
937 (offset + lbs[d][cst->getNumCols() - 1 - rank]).floorDiv(lbDivisors[d]);
938 offsets.push_back(offset);
939 }
940
941 // Create 'newMemRefType' using 'newShape' from MemRefRegion accessed
942 // by 'srcStoreOpInst'.
943 uint64_t bufSize =
944 getMemRefEltSizeInBytes(oldMemRefType) * numElements.getValue();
945 unsigned newMemSpace;
946 if (bufSize <= localBufSizeThreshold && fastMemorySpace.hasValue()) {
947 newMemSpace = fastMemorySpace.getValue();
948 } else {
949 newMemSpace = oldMemRefType.getMemorySpace();
950 }
951 auto newMemRefType = MemRefType::get(newShape, oldMemRefType.getElementType(),
952 {}, newMemSpace);
953 // Gather alloc operands for the dynamic dimensions of the memref.
954 SmallVector<Value, 4> allocOperands;
955 unsigned dynamicDimCount = 0;
956 for (auto dimSize : oldMemRefType.getShape()) {
957 if (dimSize == -1)
958 allocOperands.push_back(
959 top.create<DimOp>(forOp.getLoc(), oldMemRef, dynamicDimCount++));
960 }
961
962 // Create new private memref for fused loop 'forOp'.
963 // TODO(andydavis) Create/move alloc ops for private memrefs closer to their
964 // consumer loop nests to reduce their live range. Currently they are added
965 // at the beginning of the function, because loop nests can be reordered
966 // during the fusion pass.
967 Value newMemRef =
968 top.create<AllocOp>(forOp.getLoc(), newMemRefType, allocOperands);
969
970 // Build an AffineMap to remap access functions based on lower bound offsets.
971 SmallVector<AffineExpr, 4> remapExprs;
972 remapExprs.reserve(rank);
973 unsigned zeroOffsetCount = 0;
974 for (unsigned i = 0; i < rank; i++) {
975 if (auto constExpr = offsets[i].dyn_cast<AffineConstantExpr>())
976 if (constExpr.getValue() == 0)
977 ++zeroOffsetCount;
978 auto dimExpr = b.getAffineDimExpr(outerIVs.size() + i);
979
980 auto remapExpr =
981 simplifyAffineExpr(dimExpr - offsets[i], outerIVs.size() + rank, 0);
982 remapExprs.push_back(remapExpr);
983 }
984 auto indexRemap = zeroOffsetCount == rank
985 ? AffineMap()
986 : AffineMap::get(outerIVs.size() + rank, 0, remapExprs);
987 // Replace all users of 'oldMemRef' with 'newMemRef'.
988 LogicalResult res =
989 replaceAllMemRefUsesWith(oldMemRef, newMemRef, {}, indexRemap,
990 /*extraOperands=*/outerIVs,
991 /*symbolOperands=*/{},
992 /*domInstFilter=*/&*forOp.getBody()->begin());
993 assert(succeeded(res) &&
994 "replaceAllMemrefUsesWith should always succeed here");
995 (void)res;
996 return newMemRef;
997 }
998
999 // Checks if node 'srcId' can be safely fused into node 'dstId'. Node 'srcId'
1000 // may write to multiple memrefs but it is required that only one of them,
1001 // 'srcLiveOutStoreOp', has output edges.
1002 // Returns true if 'dstNode's read/write region to 'memref' is a super set of
1003 // 'srcNode's write region to 'memref' and 'srcId' has only one output edge.
1004 // TODO(andydavis) Generalize this to handle more live in/out cases.
canFuseSrcWhichWritesToLiveOut(unsigned srcId,unsigned dstId,AffineStoreOp srcLiveOutStoreOp,MemRefDependenceGraph * mdg)1005 static bool canFuseSrcWhichWritesToLiveOut(unsigned srcId, unsigned dstId,
1006 AffineStoreOp srcLiveOutStoreOp,
1007 MemRefDependenceGraph *mdg) {
1008 assert(srcLiveOutStoreOp && "Expected a valid store op");
1009 auto *dstNode = mdg->getNode(dstId);
1010 Value memref = srcLiveOutStoreOp.getMemRef();
1011 // Return false if 'srcNode' has more than one output edge on 'memref'.
1012 if (mdg->getOutEdgeCount(srcId, memref) > 1)
1013 return false;
1014
1015 // Compute MemRefRegion 'srcWriteRegion' for 'srcStoreOp' on 'memref'.
1016 MemRefRegion srcWriteRegion(srcLiveOutStoreOp.getLoc());
1017 if (failed(srcWriteRegion.compute(srcLiveOutStoreOp, /*loopDepth=*/0))) {
1018 LLVM_DEBUG(llvm::dbgs()
1019 << "Unable to compute MemRefRegion for source operation\n.");
1020 return false;
1021 }
1022 SmallVector<int64_t, 4> srcShape;
1023 // Query 'srcWriteRegion' for 'srcShape' and 'srcNumElements'.
1024 // by 'srcStoreOp' at depth 'dstLoopDepth'.
1025 Optional<int64_t> srcNumElements =
1026 srcWriteRegion.getConstantBoundingSizeAndShape(&srcShape);
1027 if (!srcNumElements.hasValue())
1028 return false;
1029
1030 // Compute MemRefRegion 'dstRegion' for 'dstStore/LoadOpInst' on 'memref'.
1031 // TODO(andydavis) Compute 'unionboundingbox' of all write regions (one for
1032 // each store op in 'dstStoreOps').
1033 SmallVector<Operation *, 2> dstStoreOps;
1034 dstNode->getStoreOpsForMemref(memref, &dstStoreOps);
1035 SmallVector<Operation *, 2> dstLoadOps;
1036 dstNode->getLoadOpsForMemref(memref, &dstLoadOps);
1037
1038 auto *dstOpInst = dstStoreOps.empty() ? dstLoadOps[0] : dstStoreOps[0];
1039 MemRefRegion dstRegion(dstOpInst->getLoc());
1040 if (failed(dstRegion.compute(dstOpInst, /*loopDepth=*/0))) {
1041 LLVM_DEBUG(llvm::dbgs()
1042 << "Unable to compute MemRefRegion for dest operation\n.");
1043 return false;
1044 }
1045 SmallVector<int64_t, 4> dstShape;
1046 // Query 'dstRegion' for 'dstShape' and 'dstNumElements'.
1047 // by 'dstOpInst' at depth 'dstLoopDepth'.
1048 Optional<int64_t> dstNumElements =
1049 dstRegion.getConstantBoundingSizeAndShape(&dstShape);
1050 if (!dstNumElements.hasValue())
1051 return false;
1052
1053 // Return false if write region is not a superset of 'srcNodes' write
1054 // region to 'memref'.
1055 // TODO(andydavis) Check the shape and lower bounds here too.
1056 if (srcNumElements != dstNumElements)
1057 return false;
1058 return true;
1059 }
1060
1061 // Checks the profitability of fusing a backwards slice of the loop nest
1062 // surrounding 'srcOpInst' into the loop nest surrounding 'dstLoadOpInsts'.
1063 // The argument 'srcStoreOpInst' is used to calculate the storage reduction on
1064 // the memref being produced and consumed, which is an input to the cost model.
1065 // For producer-consumer fusion, 'srcStoreOpInst' will be the same as
1066 // 'srcOpInst', as we are slicing w.r.t to that producer.
1067 // For input-reuse fusion, 'srcOpInst' will be the src loop nest LoadOp which
1068 // reads from the same memref as dst loop nest load ops, and 'srcStoreOpInst'
1069 // will be the unique store op in the src node, which will be used to check
1070 // that the write region is the same after input-reuse fusion.
1071 // Returns true if it is profitable to fuse the candidate loop nests. Returns
1072 // false otherwise. `dstLoopDepth` is set to the most profitable depth at which
1073 // to materialize the source loop nest slice.
1074 // The profitability model executes the following steps:
1075 // *) Computes the backward computation slice at 'srcOpInst'. This
1076 // computation slice of the loop nest surrounding 'srcOpInst' is
1077 // represented by modified src loop bounds in 'sliceState', which are
1078 // functions of loop IVs in the loop nest surrounding 'srcOpInst'.
1079 // *) Computes the cost of unfused src/dst loop nests (currently the cost of a
1080 // loop nest is the total number of dynamic operation instances in the loop
1081 // nest).
1082 // *) Computes the cost of fusing a slice of the src loop nest into the dst
1083 // loop nest at various values of dst loop depth, attempting to fuse
1084 // the largest computation slice at the maximal dst loop depth (closest to
1085 // the load) to minimize reuse distance and potentially enable subsequent
1086 // load/store forwarding.
1087 // NOTE: If the dst loop nest includes multiple loads in 'dstLoadOpInsts' for
1088 // the same memref as is written by 'srcOpInst', then the union of slice
1089 // loop bounds is used to compute the slice and associated slice cost.
1090 // NOTE: 'dstLoopDepth' refers to the loop depth within the destination loop
1091 // nest, at which the src computation slice is inserted/fused.
1092 // NOTE: We attempt to maximize the dst loop depth, but there are cases
1093 // where a particular setting for 'dstLoopNest' might fuse an unsliced
1094 // loop (within the src computation slice) at a depth which results in
1095 // excessive recomputation (see unit tests for examples).
1096 // *) Compares the total cost of the unfused loop nests to the min cost fused
1097 // loop nest computed in the previous step, and returns true if the latter
1098 // is lower.
isFusionProfitable(Operation * srcOpInst,Operation * srcStoreOpInst,ArrayRef<Operation * > dstLoadOpInsts,ArrayRef<Operation * > dstStoreOpInsts,ComputationSliceState * sliceState,unsigned * dstLoopDepth,bool maximalFusion)1099 static bool isFusionProfitable(Operation *srcOpInst, Operation *srcStoreOpInst,
1100 ArrayRef<Operation *> dstLoadOpInsts,
1101 ArrayRef<Operation *> dstStoreOpInsts,
1102 ComputationSliceState *sliceState,
1103 unsigned *dstLoopDepth, bool maximalFusion) {
1104 LLVM_DEBUG({
1105 llvm::dbgs() << "Checking whether fusion is profitable between:\n";
1106 llvm::dbgs() << " " << *srcOpInst << " and \n";
1107 for (auto dstOpInst : dstLoadOpInsts) {
1108 llvm::dbgs() << " " << *dstOpInst << "\n";
1109 };
1110 });
1111
1112 // Compute cost of sliced and unsliced src loop nest.
1113 SmallVector<AffineForOp, 4> srcLoopIVs;
1114 getLoopIVs(*srcOpInst, &srcLoopIVs);
1115 unsigned numSrcLoopIVs = srcLoopIVs.size();
1116
1117 // Walk src loop nest and collect stats.
1118 LoopNestStats srcLoopNestStats;
1119 if (!getLoopNestStats(srcLoopIVs[0], &srcLoopNestStats))
1120 return false;
1121
1122 // Compute cost of dst loop nest.
1123 SmallVector<AffineForOp, 4> dstLoopIVs;
1124 getLoopIVs(*dstLoadOpInsts[0], &dstLoopIVs);
1125
1126 LoopNestStats dstLoopNestStats;
1127 if (!getLoopNestStats(dstLoopIVs[0], &dstLoopNestStats))
1128 return false;
1129
1130 // Compute the maximum loop depth at which we can can insert the src slice
1131 // and still satisfy dest loop nest dependences, for producer-consumer fusion.
1132 unsigned maxDstLoopDepth =
1133 (srcOpInst == srcStoreOpInst)
1134 ? getMaxLoopDepth(dstLoadOpInsts, dstStoreOpInsts)
1135 : dstLoopIVs.size();
1136 if (maxDstLoopDepth == 0) {
1137 LLVM_DEBUG(llvm::dbgs() << "Can't fuse: maxDstLoopDepth == 0 .\n");
1138 return false;
1139 }
1140
1141 // Search for min cost value for 'dstLoopDepth'. At each value of
1142 // 'dstLoopDepth' from 'maxDstLoopDepth' to '1', compute computation slice
1143 // bounds between 'srcOpInst' and each op in 'dstOpinsts' (taking the union
1144 // of these bounds). Next the union slice bounds are used to calculate
1145 // the cost of the slice and the cost of the slice inserted into the dst
1146 // loop nest at 'dstLoopDepth'.
1147 uint64_t minFusedLoopNestComputeCost = std::numeric_limits<uint64_t>::max();
1148 double maxStorageReduction = 0.0;
1149 Optional<uint64_t> sliceMemEstimate = None;
1150
1151 SmallVector<ComputationSliceState, 4> sliceStates;
1152 sliceStates.resize(maxDstLoopDepth);
1153 // The best loop depth at which to materialize the slice.
1154 Optional<unsigned> bestDstLoopDepth = None;
1155
1156 // Compute op instance count for the src loop nest without iteration slicing.
1157 uint64_t srcLoopNestCost = getComputeCost(srcLoopIVs[0], srcLoopNestStats);
1158
1159 // Compute src loop nest write region size.
1160 MemRefRegion srcWriteRegion(srcStoreOpInst->getLoc());
1161 if (failed(srcWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0))) {
1162 LLVM_DEBUG(llvm::dbgs()
1163 << "Unable to compute MemRefRegion for source operation\n.");
1164 return false;
1165 }
1166
1167 Optional<int64_t> maybeSrcWriteRegionSizeBytes =
1168 srcWriteRegion.getRegionSize();
1169 if (!maybeSrcWriteRegionSizeBytes.hasValue())
1170 return false;
1171 int64_t srcWriteRegionSizeBytes = maybeSrcWriteRegionSizeBytes.getValue();
1172
1173 // Compute op instance count for the src loop nest.
1174 uint64_t dstLoopNestCost = getComputeCost(dstLoopIVs[0], dstLoopNestStats);
1175
1176 // Evaluate all depth choices for materializing the slice in the destination
1177 // loop nest.
1178 for (unsigned i = maxDstLoopDepth; i >= 1; --i) {
1179 // Compute the union of slice bounds of all ops in 'dstLoadOpInsts'.
1180 if (failed(mlir::computeSliceUnion({srcOpInst}, dstLoadOpInsts,
1181 /*loopDepth=*/i,
1182 /*numCommonLoops=*/0,
1183 /*isBackwardSlice=*/true,
1184 &sliceStates[i - 1]))) {
1185 LLVM_DEBUG(llvm::dbgs()
1186 << "computeSliceUnion failed for loopDepth: " << i << "\n");
1187 continue;
1188 }
1189
1190 int64_t fusedLoopNestComputeCost;
1191 if (!getFusionComputeCost(srcLoopIVs[0], srcLoopNestStats, dstLoopIVs[0],
1192 dstLoopNestStats, &sliceStates[i - 1],
1193 &fusedLoopNestComputeCost)) {
1194 LLVM_DEBUG(llvm::dbgs() << "Unable to compute fusion compute cost.\n.");
1195 continue;
1196 }
1197
1198 double additionalComputeFraction =
1199 fusedLoopNestComputeCost /
1200 (static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
1201 1;
1202
1203 // Determine what the slice write MemRefRegion would be, if the src loop
1204 // nest slice 'sliceStates[i - 1]' were to be inserted into the dst loop
1205 // nest at loop depth 'i'
1206 MemRefRegion sliceWriteRegion(srcStoreOpInst->getLoc());
1207 if (failed(sliceWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0,
1208 &sliceStates[i - 1]))) {
1209 LLVM_DEBUG(llvm::dbgs()
1210 << "Failed to compute slice write region at loopDepth: " << i
1211 << "\n");
1212 continue;
1213 }
1214
1215 Optional<int64_t> maybeSliceWriteRegionSizeBytes =
1216 sliceWriteRegion.getRegionSize();
1217 if (!maybeSliceWriteRegionSizeBytes.hasValue() ||
1218 maybeSliceWriteRegionSizeBytes.getValue() == 0) {
1219 LLVM_DEBUG(llvm::dbgs()
1220 << "Failed to get slice write region size at loopDepth: " << i
1221 << "\n");
1222 continue;
1223 }
1224 int64_t sliceWriteRegionSizeBytes =
1225 maybeSliceWriteRegionSizeBytes.getValue();
1226
1227 // If we are fusing for reuse, check that write regions remain the same.
1228 // TODO(andydavis) Write region check should check sizes and offsets in
1229 // each dimension, so that we are sure they are covering the same memref
1230 // region. Also, move this out to a isMemRefRegionSuperSet helper function.
1231 if (srcOpInst != srcStoreOpInst &&
1232 sliceWriteRegionSizeBytes != srcWriteRegionSizeBytes)
1233 continue;
1234
1235 double storageReduction = static_cast<double>(srcWriteRegionSizeBytes) /
1236 static_cast<double>(sliceWriteRegionSizeBytes);
1237
1238 LLVM_DEBUG({
1239 std::stringstream msg;
1240 msg << " evaluating fusion profitability at depth : " << i << "\n"
1241 << std::fixed << std::setprecision(2)
1242 << " additional compute fraction: "
1243 << 100.0 * additionalComputeFraction << "%\n"
1244 << " storage reduction factor: " << storageReduction << "x\n"
1245 << " fused nest cost: " << fusedLoopNestComputeCost << "\n"
1246 << " src write region size: " << srcWriteRegionSizeBytes << "\n"
1247 << " slice write region size: " << sliceWriteRegionSizeBytes
1248 << "\n";
1249 llvm::dbgs() << msg.str();
1250 });
1251
1252 double computeToleranceThreshold =
1253 clFusionAddlComputeTolerance.getNumOccurrences() > 0
1254 ? clFusionAddlComputeTolerance
1255 : LoopFusion::kComputeToleranceThreshold;
1256
1257 // TODO(b/123247369): This is a placeholder cost model.
1258 // Among all choices that add an acceptable amount of redundant computation
1259 // (as per computeToleranceThreshold), we will simply pick the one that
1260 // reduces the intermediary size the most.
1261 if ((storageReduction > maxStorageReduction) &&
1262 (maximalFusion ||
1263 (additionalComputeFraction < computeToleranceThreshold))) {
1264 maxStorageReduction = storageReduction;
1265 bestDstLoopDepth = i;
1266 minFusedLoopNestComputeCost = fusedLoopNestComputeCost;
1267 sliceMemEstimate = sliceWriteRegionSizeBytes;
1268 }
1269 }
1270
1271 // A simple cost model: fuse if it reduces the memory footprint. If
1272 // -maximal-fusion is set, fuse nevertheless.
1273
1274 if (!maximalFusion && !bestDstLoopDepth.hasValue()) {
1275 LLVM_DEBUG(
1276 llvm::dbgs()
1277 << "All fusion choices involve more than the threshold amount of "
1278 "redundant computation; NOT fusing.\n");
1279 return false;
1280 }
1281
1282 if (!bestDstLoopDepth.hasValue()) {
1283 LLVM_DEBUG(llvm::dbgs() << "no fusion depth could be evaluated.\n");
1284 return false;
1285 }
1286
1287 // Set dstLoopDepth based on best values from search.
1288 *dstLoopDepth = bestDstLoopDepth.getValue();
1289
1290 LLVM_DEBUG(
1291 llvm::dbgs() << " LoopFusion fusion stats:"
1292 << "\n best loop depth: " << bestDstLoopDepth
1293 << "\n src loop nest compute cost: " << srcLoopNestCost
1294 << "\n dst loop nest compute cost: " << dstLoopNestCost
1295 << "\n fused loop nest compute cost: "
1296 << minFusedLoopNestComputeCost << "\n");
1297
1298 auto dstMemSize = getMemoryFootprintBytes(dstLoopIVs[0]);
1299 auto srcMemSize = getMemoryFootprintBytes(srcLoopIVs[0]);
1300
1301 Optional<double> storageReduction = None;
1302
1303 if (!maximalFusion) {
1304 if (!dstMemSize.hasValue() || !srcMemSize.hasValue()) {
1305 LLVM_DEBUG(
1306 llvm::dbgs()
1307 << " fusion memory benefit cannot be evaluated; NOT fusing.\n");
1308 return false;
1309 }
1310
1311 auto srcMemSizeVal = srcMemSize.getValue();
1312 auto dstMemSizeVal = dstMemSize.getValue();
1313
1314 assert(sliceMemEstimate.hasValue() && "expected value");
1315 auto fusedMem = dstMemSizeVal + sliceMemEstimate.getValue();
1316
1317 LLVM_DEBUG(llvm::dbgs() << " src mem: " << srcMemSizeVal << "\n"
1318 << " dst mem: " << dstMemSizeVal << "\n"
1319 << " fused mem: " << fusedMem << "\n"
1320 << " slice mem: " << sliceMemEstimate << "\n");
1321
1322 if (static_cast<long>(fusedMem) > srcMemSizeVal + dstMemSizeVal) {
1323 LLVM_DEBUG(llvm::dbgs() << "Fusion is not profitable; NOT fusing.\n");
1324 return false;
1325 }
1326 storageReduction =
1327 100.0 *
1328 (1.0 - fusedMem / (static_cast<double>(srcMemSizeVal) + dstMemSizeVal));
1329 }
1330
1331 double additionalComputeFraction =
1332 100.0 * (minFusedLoopNestComputeCost /
1333 (static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
1334 1);
1335 (void)additionalComputeFraction;
1336 LLVM_DEBUG({
1337 std::stringstream msg;
1338 msg << " fusion is most profitable at depth " << *dstLoopDepth << " with "
1339 << std::setprecision(2) << additionalComputeFraction
1340 << "% redundant computation and a ";
1341 msg << (storageReduction.hasValue()
1342 ? std::to_string(storageReduction.getValue())
1343 : "<unknown>");
1344 msg << "% storage reduction.\n";
1345 llvm::dbgs() << msg.str();
1346 });
1347
1348 // Update return parameter 'sliceState' with 'bestSliceState'.
1349 ComputationSliceState *bestSliceState = &sliceStates[*dstLoopDepth - 1];
1350 sliceState->lbs = bestSliceState->lbs;
1351 sliceState->ubs = bestSliceState->ubs;
1352 sliceState->lbOperands = bestSliceState->lbOperands;
1353 sliceState->ubOperands = bestSliceState->ubOperands;
1354
1355 // Canonicalize slice bound affine maps.
1356 for (unsigned i = 0; i < numSrcLoopIVs; ++i) {
1357 if (sliceState->lbs[i] != AffineMap()) {
1358 canonicalizeMapAndOperands(&sliceState->lbs[i],
1359 &sliceState->lbOperands[i]);
1360 }
1361 if (sliceState->ubs[i] != AffineMap()) {
1362 canonicalizeMapAndOperands(&sliceState->ubs[i],
1363 &sliceState->ubOperands[i]);
1364 }
1365 }
1366 return true;
1367 }
1368
1369 namespace {
1370
1371 // GreedyFusion greedily fuses loop nests which have a producer/consumer or
1372 // input-reuse relationship on a memref, with the goal of improving locality.
1373 //
1374 // The steps of the producer-consumer fusion algorithm are as follows:
1375 //
1376 // *) A worklist is initialized with node ids from the dependence graph.
1377 // *) For each node id in the worklist:
1378 // *) Pop an AffineForOp of the worklist. This 'dstAffineForOp' will be a
1379 // candidate destination AffineForOp into which fusion will be attempted.
1380 // *) Add each LoadOp currently in 'dstAffineForOp' into list 'dstLoadOps'.
1381 // *) For each LoadOp in 'dstLoadOps' do:
1382 // *) Look up dependent loop nests which have a single store op to the same
1383 // memref.
1384 // *) Check if dependences would be violated by the fusion.
1385 // *) Get a computation slice of 'srcLoopNest', which adjusts its loop
1386 // bounds to be functions of 'dstLoopNest' IVs and symbols.
1387 // *) Fuse the 'srcLoopNest' computation slice into the 'dstLoopNest',
1388 // at a loop depth determined by the cost model in 'isFusionProfitable'.
1389 // *) Add the newly fused load/store operations to the state,
1390 // and also add newly fused load ops to 'dstLoopOps' to be considered
1391 // as fusion dst load ops in another iteration.
1392 // *) Remove old src loop nest and its associated state.
1393 //
1394 // The steps of the input-reuse fusion algorithm are as follows:
1395 //
1396 // *) Initialize 'worklist' with node ids from the dependence graph.
1397 // *) For each 'dstNode' in the worklist:
1398 // *) Find a candidate sibling node 'sibNode' to fuse with 'dstNode' which
1399 // loads from the same memref, but which has no dependence paths to/from.
1400 // *) Get a computation slice of 'sibLoopNest', which adjusts its loop
1401 // bounds to be functions of 'dstLoopNest' IVs and symbols.
1402 // *) Fuse the 'sibLoopNest' computation slice into the 'dstLoopNest',
1403 // at a loop depth determined by the cost model in 'isFusionProfitable'.
1404 // This function also checks that the memref write region of 'sibLoopNest',
1405 // is preserved in the fused loop nest.
1406 // *) Update graph state to reflect the fusion of 'sibNode' into 'dstNode'.
1407 //
1408 // Given a graph where top-level operations are vertices in the set 'V' and
1409 // edges in the set 'E' are dependences between vertices, this algorithm
1410 // takes O(V) time for initialization, and has runtime O(V + E).
1411 //
1412 // This greedy algorithm is not 'maximal' due to the current restriction of
1413 // fusing along single producer consumer edges, but there is a TODO to fix this.
1414 //
1415 // TODO(andydavis) Experiment with other fusion policies.
1416 struct GreedyFusion {
1417 public:
1418 // The data dependence graph to traverse during fusion.
1419 MemRefDependenceGraph *mdg;
1420 // Worklist of graph nodes visited during the fusion pass.
1421 SmallVector<unsigned, 8> worklist;
1422 // Set of graph nodes which are present on the worklist.
1423 llvm::SmallDenseSet<unsigned, 16> worklistSet;
1424 // Parameter for local buffer size threshold.
1425 unsigned localBufSizeThreshold;
1426 // Parameter for fast memory space.
1427 Optional<unsigned> fastMemorySpace;
1428 // If true, ignore any additional (redundant) computation tolerance threshold
1429 // that would have prevented fusion.
1430 bool maximalFusion;
1431
1432 using Node = MemRefDependenceGraph::Node;
1433
GreedyFusion__anon5a3accbe0711::GreedyFusion1434 GreedyFusion(MemRefDependenceGraph *mdg, unsigned localBufSizeThreshold,
1435 Optional<unsigned> fastMemorySpace, bool maximalFusion)
1436 : mdg(mdg), localBufSizeThreshold(localBufSizeThreshold),
1437 fastMemorySpace(fastMemorySpace), maximalFusion(maximalFusion) {}
1438
1439 // Initializes 'worklist' with nodes from 'mdg'
init__anon5a3accbe0711::GreedyFusion1440 void init() {
1441 // TODO(andydavis) Add a priority queue for prioritizing nodes by different
1442 // metrics (e.g. arithmetic intensity/flops-to-bytes ratio).
1443 worklist.clear();
1444 worklistSet.clear();
1445 for (auto &idAndNode : mdg->nodes) {
1446 const Node &node = idAndNode.second;
1447 worklist.push_back(node.id);
1448 worklistSet.insert(node.id);
1449 }
1450 }
1451
1452 // Run the GreedyFusion pass.
1453 // *) First pass through the nodes fuses single-use producer nodes into their
1454 // unique consumer.
1455 // *) Second pass fuses sibling nodes which share no dependence edges.
1456 // *) Third pass fuses any remaining producer nodes into their users.
run__anon5a3accbe0711::GreedyFusion1457 void run() {
1458 // TODO(andydavis) Run this repeatedly until a fixed-point is reached.
1459 fuseProducerConsumerNodes(/*maxSrcUserCount=*/1);
1460 fuseSiblingNodes();
1461 fuseProducerConsumerNodes(
1462 /*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
1463 eraseUnusedMemRefAllocations();
1464 }
1465
fuseProducerConsumerNodes__anon5a3accbe0711::GreedyFusion1466 void fuseProducerConsumerNodes(unsigned maxSrcUserCount) {
1467 init();
1468 while (!worklist.empty()) {
1469 unsigned dstId = worklist.back();
1470 worklist.pop_back();
1471 worklistSet.erase(dstId);
1472
1473 // Skip if this node was removed (fused into another node).
1474 if (mdg->nodes.count(dstId) == 0)
1475 continue;
1476 // Get 'dstNode' into which to attempt fusion.
1477 auto *dstNode = mdg->getNode(dstId);
1478 // Skip if 'dstNode' is not a loop nest.
1479 if (!isa<AffineForOp>(dstNode->op))
1480 continue;
1481 // Sink sequential loops in 'dstNode' (and thus raise parallel loops)
1482 // while preserving relative order. This can increase the maximum loop
1483 // depth at which we can fuse a slice of a producer loop nest into a
1484 // consumer loop nest.
1485 sinkSequentialLoops(dstNode);
1486
1487 SmallVector<Operation *, 4> loads = dstNode->loads;
1488 SmallVector<Operation *, 4> dstLoadOpInsts;
1489 DenseSet<Value> visitedMemrefs;
1490 while (!loads.empty()) {
1491 // Get memref of load on top of the stack.
1492 auto memref = cast<AffineLoadOp>(loads.back()).getMemRef();
1493 if (visitedMemrefs.count(memref) > 0)
1494 continue;
1495 visitedMemrefs.insert(memref);
1496 // Move all loads in 'loads' accessing 'memref' to 'dstLoadOpInsts'.
1497 moveLoadsAccessingMemrefTo(memref, &loads, &dstLoadOpInsts);
1498 // Skip if no input edges along which to fuse.
1499 if (mdg->inEdges.count(dstId) == 0)
1500 continue;
1501 // Iterate through in-edges for 'dstId' and src node id for any
1502 // edges on 'memref'.
1503 SmallVector<unsigned, 2> srcNodeIds;
1504 for (auto &srcEdge : mdg->inEdges[dstId]) {
1505 // Skip 'srcEdge' if not for 'memref'.
1506 if (srcEdge.value != memref)
1507 continue;
1508 srcNodeIds.push_back(srcEdge.id);
1509 }
1510 for (unsigned srcId : srcNodeIds) {
1511 // Skip if this node was removed (fused into another node).
1512 if (mdg->nodes.count(srcId) == 0)
1513 continue;
1514 // Get 'srcNode' from which to attempt fusion into 'dstNode'.
1515 auto *srcNode = mdg->getNode(srcId);
1516 // Skip if 'srcNode' is not a loop nest.
1517 if (!isa<AffineForOp>(srcNode->op))
1518 continue;
1519 // Skip if 'srcNode' has more than one live-out store to a
1520 // function-local memref.
1521 // TODO(andydavis) Support more generic multi-output src loop nests
1522 // fusion.
1523 auto srcStoreOp = mdg->getUniqueOutgoingStore(srcNode);
1524 if (!srcStoreOp) {
1525 // Get the src store op at the deepest loop depth.
1526 // We will use 'LoopFusionUtils::canFuseLoops' to check fusion
1527 // feasibility for loops with multiple stores.
1528 unsigned maxLoopDepth = 0;
1529 for (auto *op : srcNode->stores) {
1530 auto storeOp = cast<AffineStoreOp>(op);
1531 if (storeOp.getMemRef() != memref) {
1532 srcStoreOp = nullptr;
1533 break;
1534 }
1535 unsigned loopDepth = getNestingDepth(*storeOp);
1536 if (loopDepth > maxLoopDepth) {
1537 maxLoopDepth = loopDepth;
1538 srcStoreOp = storeOp;
1539 }
1540 }
1541 if (!srcStoreOp)
1542 continue;
1543 }
1544
1545 // Unique outgoing store found must write to 'memref' since 'memref'
1546 // is the one that established the producer-consumer relationship
1547 // between 'srcNode' and 'dstNode'.
1548 assert(srcStoreOp.getMemRef() == memref &&
1549 "Found store to unexpected memref");
1550
1551 // Skip if 'srcNode' writes to any live in or escaping memrefs,
1552 // and cannot be fused.
1553 bool writesToLiveInOrOut =
1554 mdg->writesToLiveInOrEscapingMemrefs(srcNode->id);
1555 if (writesToLiveInOrOut &&
1556 !canFuseSrcWhichWritesToLiveOut(srcId, dstId, srcStoreOp, mdg))
1557 continue;
1558
1559 // Don't create a private memref if 'writesToLiveInOrOut'.
1560 bool createPrivateMemref = !writesToLiveInOrOut;
1561 // Don't create a private memref if 'srcNode' has in edges on
1562 // 'memref', or if 'dstNode' has out edges on 'memref'.
1563 if (mdg->getIncomingMemRefAccesses(srcNode->id, memref) > 0 ||
1564 mdg->getOutEdgeCount(dstNode->id, memref) > 0) {
1565 createPrivateMemref = false;
1566 }
1567
1568 // Skip if 'srcNode' out edge count on 'memref' > 'maxSrcUserCount'.
1569 if (mdg->getOutEdgeCount(srcNode->id, memref) > maxSrcUserCount)
1570 continue;
1571
1572 // Compute an operation list insertion point for the fused loop
1573 // nest which preserves dependences.
1574 Operation *insertPointInst =
1575 mdg->getFusedLoopNestInsertionPoint(srcNode->id, dstNode->id);
1576 if (insertPointInst == nullptr)
1577 continue;
1578
1579 // Compute the innermost common loop depth for dstNode loads/stores.
1580 SmallVector<Operation *, 2> dstOps(dstNode->loads.begin(),
1581 dstNode->loads.end());
1582 dstOps.append(dstNode->stores.begin(), dstNode->stores.end());
1583 unsigned dstLoopDepthTest = getInnermostCommonLoopDepth(dstOps);
1584 // Check the feasibility of fusing src loop nest into dst loop nest
1585 // at loop depths in range [1, dstLoopDepthTest].
1586 // TODO(andydavis) Use slice union computation and union of memref
1587 // read/write regions to cost model and fusion.
1588 bool canFuse = false;
1589 for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
1590 ComputationSliceState sliceUnion;
1591 FusionResult result = mlir::canFuseLoops(
1592 cast<AffineForOp>(srcNode->op), cast<AffineForOp>(dstNode->op),
1593 /*dstLoopDepth=*/i, &sliceUnion);
1594 if (result.value == FusionResult::Success)
1595 canFuse = true;
1596 }
1597
1598 // Skip if fusion is not feasible at all loop depths.
1599 if (!canFuse)
1600 continue;
1601
1602 // Gather 'dstNode' store ops to 'memref'.
1603 SmallVector<Operation *, 2> dstStoreOpInsts;
1604 for (auto *storeOpInst : dstNode->stores)
1605 if (cast<AffineStoreOp>(storeOpInst).getMemRef() == memref)
1606 dstStoreOpInsts.push_back(storeOpInst);
1607
1608 unsigned bestDstLoopDepth;
1609 mlir::ComputationSliceState sliceState;
1610 // Check if fusion would be profitable.
1611 if (!isFusionProfitable(srcStoreOp, srcStoreOp, dstLoadOpInsts,
1612 dstStoreOpInsts, &sliceState,
1613 &bestDstLoopDepth, maximalFusion))
1614 continue;
1615
1616 // Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'.
1617 auto sliceLoopNest = mlir::insertBackwardComputationSlice(
1618 srcStoreOp, dstLoadOpInsts[0], bestDstLoopDepth, &sliceState);
1619 if (sliceLoopNest) {
1620 LLVM_DEBUG(llvm::dbgs() << "\tslice loop nest:\n"
1621 << *sliceLoopNest.getOperation() << "\n");
1622 // Move 'dstAffineForOp' before 'insertPointInst' if needed.
1623 auto dstAffineForOp = cast<AffineForOp>(dstNode->op);
1624 if (insertPointInst != dstAffineForOp.getOperation()) {
1625 dstAffineForOp.getOperation()->moveBefore(insertPointInst);
1626 }
1627 // Update edges between 'srcNode' and 'dstNode'.
1628 mdg->updateEdges(srcNode->id, dstNode->id, memref,
1629 createPrivateMemref);
1630
1631 // Collect slice loop stats.
1632 LoopNestStateCollector sliceCollector;
1633 sliceCollector.collect(sliceLoopNest.getOperation());
1634 // Promote single iteration slice loops to single IV value.
1635 for (auto forOp : sliceCollector.forOps) {
1636 promoteIfSingleIteration(forOp);
1637 }
1638 if (createPrivateMemref) {
1639 // Create private memref for 'memref' in 'dstAffineForOp'.
1640 SmallVector<Operation *, 4> storesForMemref;
1641 for (auto *storeOpInst : sliceCollector.storeOpInsts) {
1642 if (cast<AffineStoreOp>(storeOpInst).getMemRef() == memref)
1643 storesForMemref.push_back(storeOpInst);
1644 }
1645 // TODO(andydavis) Use union of memref write regions to compute
1646 // private memref footprint.
1647 auto newMemRef = createPrivateMemRef(
1648 dstAffineForOp, storesForMemref[0], bestDstLoopDepth,
1649 fastMemorySpace, localBufSizeThreshold);
1650 visitedMemrefs.insert(newMemRef);
1651 // Create new node in dependence graph for 'newMemRef' alloc op.
1652 unsigned newMemRefNodeId =
1653 mdg->addNode(newMemRef.getDefiningOp());
1654 // Add edge from 'newMemRef' node to dstNode.
1655 mdg->addEdge(newMemRefNodeId, dstId, newMemRef);
1656 }
1657
1658 // Collect dst loop stats after memref privatization transformation.
1659 LoopNestStateCollector dstLoopCollector;
1660 dstLoopCollector.collect(dstAffineForOp.getOperation());
1661
1662 // Add new load ops to current Node load op list 'loads' to
1663 // continue fusing based on new operands.
1664 for (auto *loadOpInst : dstLoopCollector.loadOpInsts) {
1665 auto loadMemRef = cast<AffineLoadOp>(loadOpInst).getMemRef();
1666 if (visitedMemrefs.count(loadMemRef) == 0)
1667 loads.push_back(loadOpInst);
1668 }
1669
1670 // Clear and add back loads and stores.
1671 mdg->clearNodeLoadAndStores(dstNode->id);
1672 mdg->addToNode(dstId, dstLoopCollector.loadOpInsts,
1673 dstLoopCollector.storeOpInsts);
1674 // Remove old src loop nest if it no longer has outgoing dependence
1675 // edges, and if it does not write to a memref which escapes the
1676 // function. If 'writesToLiveInOrOut' is true, then 'srcNode' has
1677 // been fused into 'dstNode' and write region of 'dstNode' covers
1678 // the write region of 'srcNode', and 'srcNode' has no other users
1679 // so it is safe to remove.
1680 if (writesToLiveInOrOut || mdg->canRemoveNode(srcNode->id)) {
1681 mdg->removeNode(srcNode->id);
1682 srcNode->op->erase();
1683 } else {
1684 // Add remaining users of 'oldMemRef' back on the worklist (if not
1685 // already there), as its replacement with a local/private memref
1686 // has reduced dependences on 'oldMemRef' which may have created
1687 // new fusion opportunities.
1688 if (mdg->outEdges.count(srcNode->id) > 0) {
1689 SmallVector<MemRefDependenceGraph::Edge, 2> oldOutEdges =
1690 mdg->outEdges[srcNode->id];
1691 for (auto &outEdge : oldOutEdges) {
1692 if (outEdge.value == memref &&
1693 worklistSet.count(outEdge.id) == 0) {
1694 worklist.push_back(outEdge.id);
1695 worklistSet.insert(outEdge.id);
1696 }
1697 }
1698 }
1699 }
1700 }
1701 }
1702 }
1703 }
1704 }
1705
1706 // Visits each node in the graph, and for each node, attempts to fuse it with
1707 // its sibling nodes (nodes which share a parent, but no dependence edges).
fuseSiblingNodes__anon5a3accbe0711::GreedyFusion1708 void fuseSiblingNodes() {
1709 init();
1710 while (!worklist.empty()) {
1711 unsigned dstId = worklist.back();
1712 worklist.pop_back();
1713 worklistSet.erase(dstId);
1714
1715 // Skip if this node was removed (fused into another node).
1716 if (mdg->nodes.count(dstId) == 0)
1717 continue;
1718 // Get 'dstNode' into which to attempt fusion.
1719 auto *dstNode = mdg->getNode(dstId);
1720 // Skip if 'dstNode' is not a loop nest.
1721 if (!isa<AffineForOp>(dstNode->op))
1722 continue;
1723 // Attempt to fuse 'dstNode' with its sibling nodes in the graph.
1724 fuseWithSiblingNodes(dstNode);
1725 }
1726 }
1727
1728 // Attempt to fuse 'dstNode' with sibling nodes in the graph.
fuseWithSiblingNodes__anon5a3accbe0711::GreedyFusion1729 void fuseWithSiblingNodes(Node *dstNode) {
1730 DenseSet<unsigned> visitedSibNodeIds;
1731 std::pair<unsigned, Value> idAndMemref;
1732 while (findSiblingNodeToFuse(dstNode, &visitedSibNodeIds, &idAndMemref)) {
1733 unsigned sibId = idAndMemref.first;
1734 Value memref = idAndMemref.second;
1735 // TODO(andydavis) Check that 'sibStoreOpInst' post-dominates all other
1736 // stores to the same memref in 'sibNode' loop nest.
1737 auto *sibNode = mdg->getNode(sibId);
1738 // Compute an operation list insertion point for the fused loop
1739 // nest which preserves dependences.
1740 assert(sibNode->op->getBlock() == dstNode->op->getBlock());
1741 Operation *insertPointInst =
1742 sibNode->op->isBeforeInBlock(dstNode->op)
1743 ? mdg->getFusedLoopNestInsertionPoint(sibNode->id, dstNode->id)
1744 : mdg->getFusedLoopNestInsertionPoint(dstNode->id, sibNode->id);
1745 if (insertPointInst == nullptr)
1746 continue;
1747
1748 // Check if fusion would be profitable and at what depth.
1749
1750 // Get unique 'sibNode' load op to 'memref'.
1751 SmallVector<Operation *, 2> sibLoadOpInsts;
1752 sibNode->getLoadOpsForMemref(memref, &sibLoadOpInsts);
1753 // Currently findSiblingNodeToFuse searches for siblings with one load.
1754 assert(sibLoadOpInsts.size() == 1);
1755 Operation *sibLoadOpInst = sibLoadOpInsts[0];
1756 assert(!sibNode->stores.empty());
1757 // TODO(andydavis) Choose the store which postdominates all other stores.
1758 auto *sibStoreOpInst = sibNode->stores.back();
1759
1760 // Gather 'dstNode' load ops to 'memref'.
1761 SmallVector<Operation *, 2> dstLoadOpInsts;
1762 dstNode->getLoadOpsForMemref(memref, &dstLoadOpInsts);
1763
1764 // Gather 'dstNode' store ops to 'memref'.
1765 SmallVector<Operation *, 2> dstStoreOpInsts;
1766 dstNode->getStoreOpsForMemref(memref, &dstStoreOpInsts);
1767
1768 unsigned bestDstLoopDepth;
1769 mlir::ComputationSliceState sliceState;
1770
1771 // Check if fusion would be profitable.
1772 if (!isFusionProfitable(sibLoadOpInst, sibStoreOpInst, dstLoadOpInsts,
1773 dstStoreOpInsts, &sliceState, &bestDstLoopDepth,
1774 maximalFusion))
1775 continue;
1776
1777 // Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'.
1778 auto sliceLoopNest = mlir::insertBackwardComputationSlice(
1779 sibLoadOpInst, dstLoadOpInsts[0], bestDstLoopDepth, &sliceState);
1780 if (sliceLoopNest != nullptr) {
1781 auto dstForInst = cast<AffineForOp>(dstNode->op);
1782 // Update operation position of fused loop nest (if needed).
1783 if (insertPointInst != dstForInst.getOperation()) {
1784 dstForInst.getOperation()->moveBefore(insertPointInst);
1785 }
1786 // Update data dependence graph state post fusion.
1787 updateStateAfterSiblingFusion(sliceLoopNest, sibNode, dstNode);
1788 }
1789 }
1790 }
1791
1792 // Searches function argument uses and the graph from 'dstNode' looking for a
1793 // fusion candidate sibling node which shares no dependences with 'dstNode'
1794 // but which loads from the same memref. Returns true and sets
1795 // 'idAndMemrefToFuse' on success. Returns false otherwise.
findSiblingNodeToFuse__anon5a3accbe0711::GreedyFusion1796 bool findSiblingNodeToFuse(Node *dstNode,
1797 DenseSet<unsigned> *visitedSibNodeIds,
1798 std::pair<unsigned, Value> *idAndMemrefToFuse) {
1799 // Returns true if 'sibNode' can be fused with 'dstNode' for input reuse
1800 // on 'memref'.
1801 auto canFuseWithSibNode = [&](Node *sibNode, Value memref) {
1802 // Skip if 'outEdge' is not a read-after-write dependence.
1803 // TODO(andydavis) Remove restrict to single load op restriction.
1804 if (sibNode->getLoadOpCount(memref) != 1)
1805 return false;
1806 // Skip if there exists a path of dependent edges between
1807 // 'sibNode' and 'dstNode'.
1808 if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
1809 mdg->hasDependencePath(dstNode->id, sibNode->id))
1810 return false;
1811 // Skip sib node if it loads to (and stores from) the same memref on
1812 // which it also has an input dependence edge.
1813 DenseSet<Value> loadAndStoreMemrefSet;
1814 sibNode->getLoadAndStoreMemrefSet(&loadAndStoreMemrefSet);
1815 if (llvm::any_of(loadAndStoreMemrefSet, [=](Value memref) {
1816 return mdg->getIncomingMemRefAccesses(sibNode->id, memref) > 0;
1817 }))
1818 return false;
1819
1820 // Check that all stores are to the same memref.
1821 DenseSet<Value> storeMemrefs;
1822 for (auto *storeOpInst : sibNode->stores) {
1823 storeMemrefs.insert(cast<AffineStoreOp>(storeOpInst).getMemRef());
1824 }
1825 if (storeMemrefs.size() != 1)
1826 return false;
1827 return true;
1828 };
1829
1830 // Search for siblings which load the same memref function argument.
1831 auto fn = dstNode->op->getParentOfType<FuncOp>();
1832 for (unsigned i = 0, e = fn.getNumArguments(); i != e; ++i) {
1833 for (auto *user : fn.getArgument(i).getUsers()) {
1834 if (auto loadOp = dyn_cast<AffineLoadOp>(user)) {
1835 // Gather loops surrounding 'use'.
1836 SmallVector<AffineForOp, 4> loops;
1837 getLoopIVs(*user, &loops);
1838 // Skip 'use' if it is not within a loop nest.
1839 if (loops.empty())
1840 continue;
1841 Node *sibNode = mdg->getForOpNode(loops[0]);
1842 assert(sibNode != nullptr);
1843 // Skip 'use' if it not a sibling to 'dstNode'.
1844 if (sibNode->id == dstNode->id)
1845 continue;
1846 // Skip 'use' if it has been visited.
1847 if (visitedSibNodeIds->count(sibNode->id) > 0)
1848 continue;
1849 // Skip 'use' if it does not load from the same memref as 'dstNode'.
1850 auto memref = loadOp.getMemRef();
1851 if (dstNode->getLoadOpCount(memref) == 0)
1852 continue;
1853 // Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
1854 if (canFuseWithSibNode(sibNode, memref)) {
1855 visitedSibNodeIds->insert(sibNode->id);
1856 idAndMemrefToFuse->first = sibNode->id;
1857 idAndMemrefToFuse->second = memref;
1858 return true;
1859 }
1860 }
1861 }
1862 }
1863
1864 // Search for siblings by following edges through an intermediate src node.
1865 // Collect candidate 'dstNode' input edges in 'inEdges'.
1866 SmallVector<MemRefDependenceGraph::Edge, 2> inEdges;
1867 mdg->forEachMemRefInputEdge(
1868 dstNode->id, [&](MemRefDependenceGraph::Edge inEdge) {
1869 // Add 'inEdge' if it is a read-after-write dependence.
1870 if (dstNode->getLoadOpCount(inEdge.value) > 0 &&
1871 mdg->getNode(inEdge.id)->getStoreOpCount(inEdge.value) > 0)
1872 inEdges.push_back(inEdge);
1873 });
1874
1875 // Search for sibling nodes to fuse by visiting output edges from each input
1876 // edge in 'inEdges'.
1877 for (auto &inEdge : inEdges) {
1878 // Collect candidate output edges from each node 'inEdge.id' in 'inEdges'.
1879 SmallVector<MemRefDependenceGraph::Edge, 2> outEdges;
1880 mdg->forEachMemRefOutputEdge(
1881 inEdge.id, [&](MemRefDependenceGraph::Edge outEdge) {
1882 unsigned sibNodeId = outEdge.id;
1883 if (visitedSibNodeIds->count(sibNodeId) > 0)
1884 return;
1885 // Skip output edge if not a sibling using the same memref.
1886 if (outEdge.id == dstNode->id || outEdge.value != inEdge.value)
1887 return;
1888 auto *sibNode = mdg->getNode(sibNodeId);
1889 if (!isa<AffineForOp>(sibNode->op))
1890 return;
1891 // Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
1892 if (canFuseWithSibNode(sibNode, outEdge.value)) {
1893 // Add candidate 'outEdge' to sibling node.
1894 outEdges.push_back(outEdge);
1895 }
1896 });
1897
1898 // Add first candidate if any were returned.
1899 if (!outEdges.empty()) {
1900 visitedSibNodeIds->insert(outEdges[0].id);
1901 idAndMemrefToFuse->first = outEdges[0].id;
1902 idAndMemrefToFuse->second = outEdges[0].value;
1903 return true;
1904 }
1905 }
1906 return false;
1907 }
1908
updateStateAfterSiblingFusion__anon5a3accbe0711::GreedyFusion1909 void updateStateAfterSiblingFusion(AffineForOp sliceLoopNest, Node *sibNode,
1910 Node *dstNode) {
1911 // Update 'sibNode' and 'dstNode' input/output edges to reflect fusion.
1912 mdg->updateEdges(sibNode->id, dstNode->id);
1913
1914 // Collect slice loop stats.
1915 LoopNestStateCollector sliceCollector;
1916 sliceCollector.collect(sliceLoopNest.getOperation());
1917 // Promote single iteration slice loops to single IV value.
1918 for (auto forOp : sliceCollector.forOps) {
1919 promoteIfSingleIteration(forOp);
1920 }
1921
1922 // Collect dst loop stats after memref privatization transformation.
1923 auto dstForInst = cast<AffineForOp>(dstNode->op);
1924 LoopNestStateCollector dstLoopCollector;
1925 dstLoopCollector.collect(dstForInst.getOperation());
1926 // Clear and add back loads and stores
1927 mdg->clearNodeLoadAndStores(dstNode->id);
1928 mdg->addToNode(dstNode->id, dstLoopCollector.loadOpInsts,
1929 dstLoopCollector.storeOpInsts);
1930 // Remove old sibling loop nest if it no longer has outgoing dependence
1931 // edges, and it does not write to a memref which escapes the
1932 // function.
1933 if (mdg->getOutEdgeCount(sibNode->id) == 0) {
1934 mdg->removeNode(sibNode->id);
1935 sibNode->op->erase();
1936 }
1937 }
1938
1939 // Clean up any allocs with no users.
eraseUnusedMemRefAllocations__anon5a3accbe0711::GreedyFusion1940 void eraseUnusedMemRefAllocations() {
1941 for (auto &pair : mdg->memrefEdgeCount) {
1942 if (pair.second > 0)
1943 continue;
1944 auto memref = pair.first;
1945 // Skip if there exist other uses (return operation or function calls).
1946 if (!memref.use_empty())
1947 continue;
1948 // Use list expected to match the dep graph info.
1949 auto *op = memref.getDefiningOp();
1950 if (isa_and_nonnull<AllocOp>(op))
1951 op->erase();
1952 }
1953 }
1954 };
1955
1956 } // end anonymous namespace
1957
runOnFunction()1958 void LoopFusion::runOnFunction() {
1959 // Override if a command line argument was provided.
1960 if (clFusionFastMemorySpace.getNumOccurrences() > 0) {
1961 fastMemorySpace = clFusionFastMemorySpace.getValue();
1962 }
1963
1964 // Override if a command line argument was provided.
1965 if (clFusionLocalBufThreshold.getNumOccurrences() > 0) {
1966 localBufSizeThreshold = clFusionLocalBufThreshold * 1024;
1967 }
1968
1969 if (clMaximalLoopFusion.getNumOccurrences() > 0)
1970 maximalFusion = clMaximalLoopFusion;
1971
1972 MemRefDependenceGraph g;
1973 if (g.init(getFunction()))
1974 GreedyFusion(&g, localBufSizeThreshold, fastMemorySpace, maximalFusion)
1975 .run();
1976 }
1977
1978 static PassRegistration<LoopFusion> pass("affine-loop-fusion",
1979 "Fuse loop nests");
1980