1 //===- LoopFusion.cpp - Code to perform loop fusion -----------------------===//
2 //
3 // Part of the LLVM 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 "PassDetail.h"
14 #include "mlir/Analysis/AffineAnalysis.h"
15 #include "mlir/Analysis/AffineStructures.h"
16 #include "mlir/Analysis/LoopAnalysis.h"
17 #include "mlir/Analysis/Utils.h"
18 #include "mlir/Dialect/Affine/IR/AffineOps.h"
19 #include "mlir/Dialect/MemRef/IR/MemRef.h"
20 #include "mlir/IR/AffineExpr.h"
21 #include "mlir/IR/AffineMap.h"
22 #include "mlir/IR/Builders.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 namespace mlir;
38
39 namespace {
40 /// Loop fusion pass. This pass currently supports a greedy fusion policy,
41 /// which fuses loop nests with single-writer/single-reader memref dependences
42 /// with the goal of improving locality.
43
44 // TODO: Support fusion of source loop nests which write to multiple
45 // memrefs, where each memref can have multiple users (if profitable).
46 // TODO: Extend this pass to check for fusion preventing dependences,
47 // and add support for more general loop fusion algorithms.
48
49 struct LoopFusion : public AffineLoopFusionBase<LoopFusion> {
50 LoopFusion() = default;
LoopFusion__anon0903053b0111::LoopFusion51 LoopFusion(unsigned fastMemorySpace, uint64_t localBufSizeThresholdBytes,
52 bool maximalFusion) {
53 this->fastMemorySpace = fastMemorySpace;
54 this->localBufSizeThreshold = localBufSizeThresholdBytes / 1024;
55 this->maximalFusion = maximalFusion;
56 }
57
58 void runOnFunction() override;
59 };
60
61 } // end anonymous namespace
62
63 std::unique_ptr<OperationPass<FuncOp>>
createLoopFusionPass(unsigned fastMemorySpace,uint64_t localBufSizeThreshold,bool maximalFusion)64 mlir::createLoopFusionPass(unsigned fastMemorySpace,
65 uint64_t localBufSizeThreshold, bool maximalFusion) {
66 return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold,
67 maximalFusion);
68 }
69
70 namespace {
71
72 // LoopNestStateCollector walks loop nests and collects load and store
73 // operations, and whether or not an IfInst was encountered in the loop nest.
74 struct LoopNestStateCollector {
75 SmallVector<AffineForOp, 4> forOps;
76 SmallVector<Operation *, 4> loadOpInsts;
77 SmallVector<Operation *, 4> storeOpInsts;
78 bool hasNonForRegion = false;
79
collect__anon0903053b0211::LoopNestStateCollector80 void collect(Operation *opToWalk) {
81 opToWalk->walk([&](Operation *op) {
82 if (isa<AffineForOp>(op))
83 forOps.push_back(cast<AffineForOp>(op));
84 else if (op->getNumRegions() != 0)
85 hasNonForRegion = true;
86 else if (isa<AffineReadOpInterface>(op))
87 loadOpInsts.push_back(op);
88 else if (isa<AffineWriteOpInterface>(op))
89 storeOpInsts.push_back(op);
90 });
91 }
92 };
93
94 // MemRefDependenceGraph is a graph data structure where graph nodes are
95 // top-level operations in a FuncOp which contain load/store ops, and edges
96 // are memref dependences between the nodes.
97 // TODO: Add a more flexible dependence graph representation.
98 // TODO: Add a depth parameter to dependence graph construction.
99 struct MemRefDependenceGraph {
100 public:
101 // Node represents a node in the graph. A Node is either an entire loop nest
102 // rooted at the top level which contains loads/stores, or a top level
103 // load/store.
104 struct Node {
105 // The unique identifier of this node in the graph.
106 unsigned id;
107 // The top-level statement which is (or contains) a load/store.
108 Operation *op;
109 // List of load operations.
110 SmallVector<Operation *, 4> loads;
111 // List of store op insts.
112 SmallVector<Operation *, 4> stores;
Node__anon0903053b0211::MemRefDependenceGraph::Node113 Node(unsigned id, Operation *op) : id(id), op(op) {}
114
115 // Returns the load op count for 'memref'.
getLoadOpCount__anon0903053b0211::MemRefDependenceGraph::Node116 unsigned getLoadOpCount(Value memref) {
117 unsigned loadOpCount = 0;
118 for (auto *loadOpInst : loads) {
119 if (memref == cast<AffineReadOpInterface>(loadOpInst).getMemRef())
120 ++loadOpCount;
121 }
122 return loadOpCount;
123 }
124
125 // Returns the store op count for 'memref'.
getStoreOpCount__anon0903053b0211::MemRefDependenceGraph::Node126 unsigned getStoreOpCount(Value memref) {
127 unsigned storeOpCount = 0;
128 for (auto *storeOpInst : stores) {
129 if (memref == cast<AffineWriteOpInterface>(storeOpInst).getMemRef())
130 ++storeOpCount;
131 }
132 return storeOpCount;
133 }
134
135 // Returns all store ops in 'storeOps' which access 'memref'.
getStoreOpsForMemref__anon0903053b0211::MemRefDependenceGraph::Node136 void getStoreOpsForMemref(Value memref,
137 SmallVectorImpl<Operation *> *storeOps) {
138 for (auto *storeOpInst : stores) {
139 if (memref == cast<AffineWriteOpInterface>(storeOpInst).getMemRef())
140 storeOps->push_back(storeOpInst);
141 }
142 }
143
144 // Returns all load ops in 'loadOps' which access 'memref'.
getLoadOpsForMemref__anon0903053b0211::MemRefDependenceGraph::Node145 void getLoadOpsForMemref(Value memref,
146 SmallVectorImpl<Operation *> *loadOps) {
147 for (auto *loadOpInst : loads) {
148 if (memref == cast<AffineReadOpInterface>(loadOpInst).getMemRef())
149 loadOps->push_back(loadOpInst);
150 }
151 }
152
153 // Returns all memrefs in 'loadAndStoreMemrefSet' for which this node
154 // has at least one load and store operation.
getLoadAndStoreMemrefSet__anon0903053b0211::MemRefDependenceGraph::Node155 void getLoadAndStoreMemrefSet(DenseSet<Value> *loadAndStoreMemrefSet) {
156 llvm::SmallDenseSet<Value, 2> loadMemrefs;
157 for (auto *loadOpInst : loads) {
158 loadMemrefs.insert(cast<AffineReadOpInterface>(loadOpInst).getMemRef());
159 }
160 for (auto *storeOpInst : stores) {
161 auto memref = cast<AffineWriteOpInterface>(storeOpInst).getMemRef();
162 if (loadMemrefs.count(memref) > 0)
163 loadAndStoreMemrefSet->insert(memref);
164 }
165 }
166 };
167
168 // Edge represents a data dependence between nodes in the graph.
169 struct Edge {
170 // The id of the node at the other end of the edge.
171 // If this edge is stored in Edge = Node.inEdges[i], then
172 // 'Node.inEdges[i].id' is the identifier of the source node of the edge.
173 // If this edge is stored in Edge = Node.outEdges[i], then
174 // 'Node.outEdges[i].id' is the identifier of the dest node of the edge.
175 unsigned id;
176 // The SSA value on which this edge represents a dependence.
177 // If the value is a memref, then the dependence is between graph nodes
178 // which contain accesses to the same memref 'value'. If the value is a
179 // non-memref value, then the dependence is between a graph node which
180 // defines an SSA value and another graph node which uses the SSA value
181 // (e.g. a constant or load operation defining a value which is used inside
182 // a loop nest).
183 Value value;
184 };
185
186 // Map from node id to Node.
187 DenseMap<unsigned, Node> nodes;
188 // Map from node id to list of input edges.
189 DenseMap<unsigned, SmallVector<Edge, 2>> inEdges;
190 // Map from node id to list of output edges.
191 DenseMap<unsigned, SmallVector<Edge, 2>> outEdges;
192 // Map from memref to a count on the dependence edges associated with that
193 // memref.
194 DenseMap<Value, unsigned> memrefEdgeCount;
195 // The next unique identifier to use for newly created graph nodes.
196 unsigned nextNodeId = 0;
197
MemRefDependenceGraph__anon0903053b0211::MemRefDependenceGraph198 MemRefDependenceGraph() {}
199
200 // Initializes the dependence graph based on operations in 'f'.
201 // Returns true on success, false otherwise.
202 bool init(FuncOp f);
203
204 // Returns the graph node for 'id'.
getNode__anon0903053b0211::MemRefDependenceGraph205 Node *getNode(unsigned id) {
206 auto it = nodes.find(id);
207 assert(it != nodes.end());
208 return &it->second;
209 }
210
211 // Returns the graph node for 'forOp'.
getForOpNode__anon0903053b0211::MemRefDependenceGraph212 Node *getForOpNode(AffineForOp forOp) {
213 for (auto &idAndNode : nodes)
214 if (idAndNode.second.op == forOp.getOperation())
215 return &idAndNode.second;
216 return nullptr;
217 }
218
219 // Adds a node with 'op' to the graph and returns its unique identifier.
addNode__anon0903053b0211::MemRefDependenceGraph220 unsigned addNode(Operation *op) {
221 Node node(nextNodeId++, op);
222 nodes.insert({node.id, node});
223 return node.id;
224 }
225
226 // Remove node 'id' (and its associated edges) from graph.
removeNode__anon0903053b0211::MemRefDependenceGraph227 void removeNode(unsigned id) {
228 // Remove each edge in 'inEdges[id]'.
229 if (inEdges.count(id) > 0) {
230 SmallVector<Edge, 2> oldInEdges = inEdges[id];
231 for (auto &inEdge : oldInEdges) {
232 removeEdge(inEdge.id, id, inEdge.value);
233 }
234 }
235 // Remove each edge in 'outEdges[id]'.
236 if (outEdges.count(id) > 0) {
237 SmallVector<Edge, 2> oldOutEdges = outEdges[id];
238 for (auto &outEdge : oldOutEdges) {
239 removeEdge(id, outEdge.id, outEdge.value);
240 }
241 }
242 // Erase remaining node state.
243 inEdges.erase(id);
244 outEdges.erase(id);
245 nodes.erase(id);
246 }
247
248 // Returns true if node 'id' writes to any memref which escapes (or is an
249 // argument to) the function/block. Returns false otherwise.
writesToLiveInOrEscapingMemrefs__anon0903053b0211::MemRefDependenceGraph250 bool writesToLiveInOrEscapingMemrefs(unsigned id) {
251 Node *node = getNode(id);
252 for (auto *storeOpInst : node->stores) {
253 auto memref = cast<AffineWriteOpInterface>(storeOpInst).getMemRef();
254 auto *op = memref.getDefiningOp();
255 // Return true if 'memref' is a block argument.
256 if (!op)
257 return true;
258 // Return true if any use of 'memref' escapes the function.
259 for (auto *user : memref.getUsers())
260 if (!isa<AffineMapAccessInterface>(*user))
261 return true;
262 }
263 return false;
264 }
265
266 // Returns true iff there is an edge from node 'srcId' to node 'dstId' which
267 // is for 'value' if non-null, or for any value otherwise. Returns false
268 // otherwise.
hasEdge__anon0903053b0211::MemRefDependenceGraph269 bool hasEdge(unsigned srcId, unsigned dstId, Value value = nullptr) {
270 if (outEdges.count(srcId) == 0 || inEdges.count(dstId) == 0) {
271 return false;
272 }
273 bool hasOutEdge = llvm::any_of(outEdges[srcId], [=](Edge &edge) {
274 return edge.id == dstId && (!value || edge.value == value);
275 });
276 bool hasInEdge = llvm::any_of(inEdges[dstId], [=](Edge &edge) {
277 return edge.id == srcId && (!value || edge.value == value);
278 });
279 return hasOutEdge && hasInEdge;
280 }
281
282 // Adds an edge from node 'srcId' to node 'dstId' for 'value'.
addEdge__anon0903053b0211::MemRefDependenceGraph283 void addEdge(unsigned srcId, unsigned dstId, Value value) {
284 if (!hasEdge(srcId, dstId, value)) {
285 outEdges[srcId].push_back({dstId, value});
286 inEdges[dstId].push_back({srcId, value});
287 if (value.getType().isa<MemRefType>())
288 memrefEdgeCount[value]++;
289 }
290 }
291
292 // Removes an edge from node 'srcId' to node 'dstId' for 'value'.
removeEdge__anon0903053b0211::MemRefDependenceGraph293 void removeEdge(unsigned srcId, unsigned dstId, Value value) {
294 assert(inEdges.count(dstId) > 0);
295 assert(outEdges.count(srcId) > 0);
296 if (value.getType().isa<MemRefType>()) {
297 assert(memrefEdgeCount.count(value) > 0);
298 memrefEdgeCount[value]--;
299 }
300 // Remove 'srcId' from 'inEdges[dstId]'.
301 for (auto it = inEdges[dstId].begin(); it != inEdges[dstId].end(); ++it) {
302 if ((*it).id == srcId && (*it).value == value) {
303 inEdges[dstId].erase(it);
304 break;
305 }
306 }
307 // Remove 'dstId' from 'outEdges[srcId]'.
308 for (auto it = outEdges[srcId].begin(); it != outEdges[srcId].end(); ++it) {
309 if ((*it).id == dstId && (*it).value == value) {
310 outEdges[srcId].erase(it);
311 break;
312 }
313 }
314 }
315
316 // Returns true if there is a path in the dependence graph from node 'srcId'
317 // to node 'dstId'. Returns false otherwise.
hasDependencePath__anon0903053b0211::MemRefDependenceGraph318 bool hasDependencePath(unsigned srcId, unsigned dstId) {
319 // Worklist state is: <node-id, next-output-edge-index-to-visit>
320 SmallVector<std::pair<unsigned, unsigned>, 4> worklist;
321 worklist.push_back({srcId, 0});
322 // Run DFS traversal to see if 'dstId' is reachable from 'srcId'.
323 while (!worklist.empty()) {
324 auto &idAndIndex = worklist.back();
325 // Return true if we have reached 'dstId'.
326 if (idAndIndex.first == dstId)
327 return true;
328 // Pop and continue if node has no out edges, or if all out edges have
329 // already been visited.
330 if (outEdges.count(idAndIndex.first) == 0 ||
331 idAndIndex.second == outEdges[idAndIndex.first].size()) {
332 worklist.pop_back();
333 continue;
334 }
335 // Get graph edge to traverse.
336 Edge edge = outEdges[idAndIndex.first][idAndIndex.second];
337 // Increment next output edge index for 'idAndIndex'.
338 ++idAndIndex.second;
339 // Add node at 'edge.id' to worklist.
340 worklist.push_back({edge.id, 0});
341 }
342 return false;
343 }
344
345 // Returns the input edge count for node 'id' and 'memref' from src nodes
346 // which access 'memref' with a store operation.
getIncomingMemRefAccesses__anon0903053b0211::MemRefDependenceGraph347 unsigned getIncomingMemRefAccesses(unsigned id, Value memref) {
348 unsigned inEdgeCount = 0;
349 if (inEdges.count(id) > 0)
350 for (auto &inEdge : inEdges[id])
351 if (inEdge.value == memref) {
352 Node *srcNode = getNode(inEdge.id);
353 // Only count in edges from 'srcNode' if 'srcNode' accesses 'memref'
354 if (srcNode->getStoreOpCount(memref) > 0)
355 ++inEdgeCount;
356 }
357 return inEdgeCount;
358 }
359
360 // Returns the output edge count for node 'id' and 'memref' (if non-null),
361 // otherwise returns the total output edge count from node 'id'.
getOutEdgeCount__anon0903053b0211::MemRefDependenceGraph362 unsigned getOutEdgeCount(unsigned id, Value memref = nullptr) {
363 unsigned outEdgeCount = 0;
364 if (outEdges.count(id) > 0)
365 for (auto &outEdge : outEdges[id])
366 if (!memref || outEdge.value == memref)
367 ++outEdgeCount;
368 return outEdgeCount;
369 }
370
371 /// Return all nodes which define SSA values used in node 'id'.
gatherDefiningNodes__anon0903053b0211::MemRefDependenceGraph372 void gatherDefiningNodes(unsigned id, DenseSet<unsigned> &definingNodes) {
373 for (MemRefDependenceGraph::Edge edge : inEdges[id])
374 // By definition of edge, if the edge value is a non-memref value,
375 // then the dependence is between a graph node which defines an SSA value
376 // and another graph node which uses the SSA value.
377 if (!edge.value.getType().isa<MemRefType>())
378 definingNodes.insert(edge.id);
379 }
380
381 // Computes and returns an insertion point operation, before which the
382 // the fused <srcId, dstId> loop nest can be inserted while preserving
383 // dependences. Returns nullptr if no such insertion point is found.
getFusedLoopNestInsertionPoint__anon0903053b0211::MemRefDependenceGraph384 Operation *getFusedLoopNestInsertionPoint(unsigned srcId, unsigned dstId) {
385 if (outEdges.count(srcId) == 0)
386 return getNode(dstId)->op;
387
388 // Skip if there is any defining node of 'dstId' that depends on 'srcId'.
389 DenseSet<unsigned> definingNodes;
390 gatherDefiningNodes(dstId, definingNodes);
391 if (llvm::any_of(definingNodes, [&](unsigned id) {
392 return hasDependencePath(srcId, id);
393 })) {
394 LLVM_DEBUG(llvm::dbgs()
395 << "Can't fuse: a defining op with a user in the dst "
396 "loop has dependence from the src loop\n");
397 return nullptr;
398 }
399
400 // Build set of insts in range (srcId, dstId) which depend on 'srcId'.
401 SmallPtrSet<Operation *, 2> srcDepInsts;
402 for (auto &outEdge : outEdges[srcId])
403 if (outEdge.id != dstId)
404 srcDepInsts.insert(getNode(outEdge.id)->op);
405
406 // Build set of insts in range (srcId, dstId) on which 'dstId' depends.
407 SmallPtrSet<Operation *, 2> dstDepInsts;
408 for (auto &inEdge : inEdges[dstId])
409 if (inEdge.id != srcId)
410 dstDepInsts.insert(getNode(inEdge.id)->op);
411
412 Operation *srcNodeInst = getNode(srcId)->op;
413 Operation *dstNodeInst = getNode(dstId)->op;
414
415 // Computing insertion point:
416 // *) Walk all operation positions in Block operation list in the
417 // range (src, dst). For each operation 'op' visited in this search:
418 // *) Store in 'firstSrcDepPos' the first position where 'op' has a
419 // dependence edge from 'srcNode'.
420 // *) Store in 'lastDstDepPost' the last position where 'op' has a
421 // dependence edge to 'dstNode'.
422 // *) Compare 'firstSrcDepPos' and 'lastDstDepPost' to determine the
423 // operation insertion point (or return null pointer if no such
424 // insertion point exists: 'firstSrcDepPos' <= 'lastDstDepPos').
425 SmallVector<Operation *, 2> depInsts;
426 Optional<unsigned> firstSrcDepPos;
427 Optional<unsigned> lastDstDepPos;
428 unsigned pos = 0;
429 for (Block::iterator it = std::next(Block::iterator(srcNodeInst));
430 it != Block::iterator(dstNodeInst); ++it) {
431 Operation *op = &(*it);
432 if (srcDepInsts.count(op) > 0 && firstSrcDepPos == None)
433 firstSrcDepPos = pos;
434 if (dstDepInsts.count(op) > 0)
435 lastDstDepPos = pos;
436 depInsts.push_back(op);
437 ++pos;
438 }
439
440 if (firstSrcDepPos.hasValue()) {
441 if (lastDstDepPos.hasValue()) {
442 if (firstSrcDepPos.getValue() <= lastDstDepPos.getValue()) {
443 // No valid insertion point exists which preserves dependences.
444 return nullptr;
445 }
446 }
447 // Return the insertion point at 'firstSrcDepPos'.
448 return depInsts[firstSrcDepPos.getValue()];
449 }
450 // No dependence targets in range (or only dst deps in range), return
451 // 'dstNodInst' insertion point.
452 return dstNodeInst;
453 }
454
455 // Updates edge mappings from node 'srcId' to node 'dstId' after fusing them,
456 // taking into account that:
457 // *) if 'removeSrcId' is true, 'srcId' will be removed after fusion,
458 // *) memrefs in 'privateMemRefs' has been replaced in node at 'dstId' by a
459 // private memref.
updateEdges__anon0903053b0211::MemRefDependenceGraph460 void updateEdges(unsigned srcId, unsigned dstId,
461 const DenseSet<Value> &privateMemRefs, bool removeSrcId) {
462 // For each edge in 'inEdges[srcId]': add new edge remapping to 'dstId'.
463 if (inEdges.count(srcId) > 0) {
464 SmallVector<Edge, 2> oldInEdges = inEdges[srcId];
465 for (auto &inEdge : oldInEdges) {
466 // Add edge from 'inEdge.id' to 'dstId' if it's not a private memref.
467 if (privateMemRefs.count(inEdge.value) == 0)
468 addEdge(inEdge.id, dstId, inEdge.value);
469 }
470 }
471 // For each edge in 'outEdges[srcId]': remove edge from 'srcId' to 'dstId'.
472 // If 'srcId' is going to be removed, remap all the out edges to 'dstId'.
473 if (outEdges.count(srcId) > 0) {
474 SmallVector<Edge, 2> oldOutEdges = outEdges[srcId];
475 for (auto &outEdge : oldOutEdges) {
476 // Remove any out edges from 'srcId' to 'dstId' across memrefs.
477 if (outEdge.id == dstId)
478 removeEdge(srcId, outEdge.id, outEdge.value);
479 else if (removeSrcId) {
480 addEdge(dstId, outEdge.id, outEdge.value);
481 removeEdge(srcId, outEdge.id, outEdge.value);
482 }
483 }
484 }
485 // Remove any edges in 'inEdges[dstId]' on 'oldMemRef' (which is being
486 // replaced by a private memref). These edges could come from nodes
487 // other than 'srcId' which were removed in the previous step.
488 if (inEdges.count(dstId) > 0 && !privateMemRefs.empty()) {
489 SmallVector<Edge, 2> oldInEdges = inEdges[dstId];
490 for (auto &inEdge : oldInEdges)
491 if (privateMemRefs.count(inEdge.value) > 0)
492 removeEdge(inEdge.id, dstId, inEdge.value);
493 }
494 }
495
496 // Update edge mappings for nodes 'sibId' and 'dstId' to reflect fusion
497 // of sibling node 'sibId' into node 'dstId'.
updateEdges__anon0903053b0211::MemRefDependenceGraph498 void updateEdges(unsigned sibId, unsigned dstId) {
499 // For each edge in 'inEdges[sibId]':
500 // *) Add new edge from source node 'inEdge.id' to 'dstNode'.
501 // *) Remove edge from source node 'inEdge.id' to 'sibNode'.
502 if (inEdges.count(sibId) > 0) {
503 SmallVector<Edge, 2> oldInEdges = inEdges[sibId];
504 for (auto &inEdge : oldInEdges) {
505 addEdge(inEdge.id, dstId, inEdge.value);
506 removeEdge(inEdge.id, sibId, inEdge.value);
507 }
508 }
509
510 // For each edge in 'outEdges[sibId]' to node 'id'
511 // *) Add new edge from 'dstId' to 'outEdge.id'.
512 // *) Remove edge from 'sibId' to 'outEdge.id'.
513 if (outEdges.count(sibId) > 0) {
514 SmallVector<Edge, 2> oldOutEdges = outEdges[sibId];
515 for (auto &outEdge : oldOutEdges) {
516 addEdge(dstId, outEdge.id, outEdge.value);
517 removeEdge(sibId, outEdge.id, outEdge.value);
518 }
519 }
520 }
521
522 // Adds ops in 'loads' and 'stores' to node at 'id'.
addToNode__anon0903053b0211::MemRefDependenceGraph523 void addToNode(unsigned id, const SmallVectorImpl<Operation *> &loads,
524 const SmallVectorImpl<Operation *> &stores) {
525 Node *node = getNode(id);
526 for (auto *loadOpInst : loads)
527 node->loads.push_back(loadOpInst);
528 for (auto *storeOpInst : stores)
529 node->stores.push_back(storeOpInst);
530 }
531
clearNodeLoadAndStores__anon0903053b0211::MemRefDependenceGraph532 void clearNodeLoadAndStores(unsigned id) {
533 Node *node = getNode(id);
534 node->loads.clear();
535 node->stores.clear();
536 }
537
538 // Calls 'callback' for each input edge incident to node 'id' which carries a
539 // memref dependence.
forEachMemRefInputEdge__anon0903053b0211::MemRefDependenceGraph540 void forEachMemRefInputEdge(unsigned id,
541 const std::function<void(Edge)> &callback) {
542 if (inEdges.count(id) > 0)
543 forEachMemRefEdge(inEdges[id], callback);
544 }
545
546 // Calls 'callback' for each output edge from node 'id' which carries a
547 // memref dependence.
forEachMemRefOutputEdge__anon0903053b0211::MemRefDependenceGraph548 void forEachMemRefOutputEdge(unsigned id,
549 const std::function<void(Edge)> &callback) {
550 if (outEdges.count(id) > 0)
551 forEachMemRefEdge(outEdges[id], callback);
552 }
553
554 // Calls 'callback' for each edge in 'edges' which carries a memref
555 // dependence.
forEachMemRefEdge__anon0903053b0211::MemRefDependenceGraph556 void forEachMemRefEdge(ArrayRef<Edge> edges,
557 const std::function<void(Edge)> &callback) {
558 for (const auto &edge : edges) {
559 // Skip if 'edge' is not a memref dependence edge.
560 if (!edge.value.getType().isa<MemRefType>())
561 continue;
562 assert(nodes.count(edge.id) > 0);
563 // Skip if 'edge.id' is not a loop nest.
564 if (!isa<AffineForOp>(getNode(edge.id)->op))
565 continue;
566 // Visit current input edge 'edge'.
567 callback(edge);
568 }
569 }
570
print__anon0903053b0211::MemRefDependenceGraph571 void print(raw_ostream &os) const {
572 os << "\nMemRefDependenceGraph\n";
573 os << "\nNodes:\n";
574 for (const auto &idAndNode : nodes) {
575 os << "Node: " << idAndNode.first << "\n";
576 auto it = inEdges.find(idAndNode.first);
577 if (it != inEdges.end()) {
578 for (const auto &e : it->second)
579 os << " InEdge: " << e.id << " " << e.value << "\n";
580 }
581 it = outEdges.find(idAndNode.first);
582 if (it != outEdges.end()) {
583 for (const auto &e : it->second)
584 os << " OutEdge: " << e.id << " " << e.value << "\n";
585 }
586 }
587 }
dump__anon0903053b0211::MemRefDependenceGraph588 void dump() const { print(llvm::errs()); }
589 };
590
591 /// Returns true if node 'srcId' can be removed after fusing it with node
592 /// 'dstId'. The node can be removed if any of the following conditions are met:
593 /// 1. 'srcId' has no output dependences after fusion and no escaping memrefs.
594 /// 2. 'srcId' has no output dependences after fusion, has escaping memrefs
595 /// and the fusion slice is maximal.
596 /// 3. 'srcId' has output dependences after fusion, the fusion slice is
597 /// maximal and the fusion insertion point dominates all the dependences.
canRemoveSrcNodeAfterFusion(unsigned srcId,unsigned dstId,const ComputationSliceState & fusionSlice,Operation * fusedLoopInsPoint,const DenseSet<Value> & escapingMemRefs,MemRefDependenceGraph * mdg)598 static bool canRemoveSrcNodeAfterFusion(
599 unsigned srcId, unsigned dstId, const ComputationSliceState &fusionSlice,
600 Operation *fusedLoopInsPoint, const DenseSet<Value> &escapingMemRefs,
601 MemRefDependenceGraph *mdg) {
602
603 Operation *dstNodeOp = mdg->getNode(dstId)->op;
604 bool hasOutDepsAfterFusion = false;
605
606 for (auto &outEdge : mdg->outEdges[srcId]) {
607 Operation *depNodeOp = mdg->getNode(outEdge.id)->op;
608 // Skip dependence with dstOp since it will be removed after fusion.
609 if (depNodeOp == dstNodeOp)
610 continue;
611
612 // Only fusion within the same block is supported. Use domination analysis
613 // when needed.
614 if (depNodeOp->getBlock() != dstNodeOp->getBlock())
615 return false;
616
617 // Check if the insertion point of the fused loop dominates the dependence.
618 // Otherwise, the src loop can't be removed.
619 if (fusedLoopInsPoint != depNodeOp &&
620 !fusedLoopInsPoint->isBeforeInBlock(depNodeOp)) {
621 LLVM_DEBUG(llvm::dbgs() << "Src loop can't be removed: dst loop doesn't "
622 "dominate dependence\n");
623 return false;
624 }
625
626 hasOutDepsAfterFusion = true;
627 }
628
629 // If src loop has dependences after fusion or it writes to an live-out or
630 // escaping memref, we can only remove it if the fusion slice is maximal so
631 // that all the dependences are preserved.
632 if (hasOutDepsAfterFusion || !escapingMemRefs.empty()) {
633 Optional<bool> isMaximal = fusionSlice.isMaximal();
634 if (!isMaximal.hasValue()) {
635 LLVM_DEBUG(llvm::dbgs() << "Src loop can't be removed: can't determine "
636 "if fusion is maximal\n");
637 return false;
638 }
639
640 if (!isMaximal.getValue()) {
641 LLVM_DEBUG(llvm::dbgs()
642 << "Src loop can't be removed: fusion is not maximal\n");
643 return false;
644 }
645 }
646
647 return true;
648 }
649
650 /// Returns in 'srcIdCandidates' the producer fusion candidates for consumer
651 /// 'dstId'. Candidates are sorted by node id order. This order corresponds to
652 /// the program order when the 'mdg' is created. However, program order is not
653 /// guaranteed and must not be required by the client. Program order won't be
654 /// held if the 'mdg' is reused from a previous fusion step or if the node
655 /// creation order changes in the future to support more advance cases.
656 // TODO: Move this to a loop fusion utility once 'mdg' is also moved.
getProducerCandidates(unsigned dstId,MemRefDependenceGraph * mdg,SmallVectorImpl<unsigned> & srcIdCandidates)657 static void getProducerCandidates(unsigned dstId, MemRefDependenceGraph *mdg,
658 SmallVectorImpl<unsigned> &srcIdCandidates) {
659 // Skip if no input edges along which to fuse.
660 if (mdg->inEdges.count(dstId) == 0)
661 return;
662
663 // Gather memrefs from loads in 'dstId'.
664 auto *dstNode = mdg->getNode(dstId);
665 DenseSet<Value> consumedMemrefs;
666 for (Operation *load : dstNode->loads)
667 consumedMemrefs.insert(cast<AffineReadOpInterface>(load).getMemRef());
668
669 // Traverse 'dstId' incoming edges and gather the nodes that contain a store
670 // to one of the consumed memrefs.
671 for (auto &srcEdge : mdg->inEdges[dstId]) {
672 auto *srcNode = mdg->getNode(srcEdge.id);
673 // Skip if 'srcNode' is not a loop nest.
674 if (!isa<AffineForOp>(srcNode->op))
675 continue;
676
677 if (any_of(srcNode->stores, [&](Operation *op) {
678 auto storeOp = cast<AffineWriteOpInterface>(op);
679 return consumedMemrefs.count(storeOp.getMemRef()) > 0;
680 }))
681 srcIdCandidates.push_back(srcNode->id);
682 }
683
684 std::sort(srcIdCandidates.begin(), srcIdCandidates.end());
685 srcIdCandidates.erase(
686 std::unique(srcIdCandidates.begin(), srcIdCandidates.end()),
687 srcIdCandidates.end());
688 }
689
690 /// Returns in 'producerConsumerMemrefs' the memrefs involved in a
691 /// producer-consumer dependence between 'srcId' and 'dstId'.
692 static void
gatherProducerConsumerMemrefs(unsigned srcId,unsigned dstId,MemRefDependenceGraph * mdg,DenseSet<Value> & producerConsumerMemrefs)693 gatherProducerConsumerMemrefs(unsigned srcId, unsigned dstId,
694 MemRefDependenceGraph *mdg,
695 DenseSet<Value> &producerConsumerMemrefs) {
696 auto *dstNode = mdg->getNode(dstId);
697 auto *srcNode = mdg->getNode(srcId);
698 gatherProducerConsumerMemrefs(srcNode->stores, dstNode->loads,
699 producerConsumerMemrefs);
700 }
701
702 /// Returns in 'escapingMemRefs' the memrefs from affine store ops in node 'id'
703 /// that escape the function. A memref escapes the function if either:
704 /// 1. It's a function argument, or
705 /// 2. It's used by a non-affine op (e.g., std load/store, std call, etc.)
gatherEscapingMemrefs(unsigned id,MemRefDependenceGraph * mdg,DenseSet<Value> & escapingMemRefs)706 void gatherEscapingMemrefs(unsigned id, MemRefDependenceGraph *mdg,
707 DenseSet<Value> &escapingMemRefs) {
708 auto *node = mdg->getNode(id);
709 for (auto *storeOpInst : node->stores) {
710 auto memref = cast<AffineWriteOpInterface>(storeOpInst).getMemRef();
711 if (escapingMemRefs.count(memref))
712 continue;
713 // Check if 'memref' escapes because it's a block argument.
714 if (memref.isa<BlockArgument>()) {
715 escapingMemRefs.insert(memref);
716 continue;
717 }
718 // Check if 'memref' escapes through a non-affine op (e.g., std load/store,
719 // call op, etc.).
720 for (Operation *user : memref.getUsers())
721 if (!isa<AffineMapAccessInterface>(*user))
722 escapingMemRefs.insert(memref);
723 }
724 }
725
726 } // end anonymous namespace
727
728 // Initializes the data dependence graph by walking operations in 'f'.
729 // Assigns each node in the graph a node id based on program order in 'f'.
730 // TODO: Add support for taking a Block arg to construct the
731 // dependence graph at a different depth.
init(FuncOp f)732 bool MemRefDependenceGraph::init(FuncOp f) {
733 LLVM_DEBUG(llvm::dbgs() << "--- Initializing MDG ---\n");
734 DenseMap<Value, SetVector<unsigned>> memrefAccesses;
735
736 // TODO: support multi-block functions.
737 if (!llvm::hasSingleElement(f))
738 return false;
739
740 DenseMap<Operation *, unsigned> forToNodeMap;
741 for (auto &op : f.front()) {
742 if (auto forOp = dyn_cast<AffineForOp>(op)) {
743 // Create graph node 'id' to represent top-level 'forOp' and record
744 // all loads and store accesses it contains.
745 LoopNestStateCollector collector;
746 collector.collect(&op);
747 // Return false if a non 'affine.for' region was found (not currently
748 // supported).
749 if (collector.hasNonForRegion)
750 return false;
751 Node node(nextNodeId++, &op);
752 for (auto *opInst : collector.loadOpInsts) {
753 node.loads.push_back(opInst);
754 auto memref = cast<AffineReadOpInterface>(opInst).getMemRef();
755 memrefAccesses[memref].insert(node.id);
756 }
757 for (auto *opInst : collector.storeOpInsts) {
758 node.stores.push_back(opInst);
759 auto memref = cast<AffineWriteOpInterface>(opInst).getMemRef();
760 memrefAccesses[memref].insert(node.id);
761 }
762 forToNodeMap[&op] = node.id;
763 nodes.insert({node.id, node});
764 } else if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
765 // Create graph node for top-level load op.
766 Node node(nextNodeId++, &op);
767 node.loads.push_back(&op);
768 auto memref = cast<AffineReadOpInterface>(op).getMemRef();
769 memrefAccesses[memref].insert(node.id);
770 nodes.insert({node.id, node});
771 } else if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
772 // Create graph node for top-level store op.
773 Node node(nextNodeId++, &op);
774 node.stores.push_back(&op);
775 auto memref = cast<AffineWriteOpInterface>(op).getMemRef();
776 memrefAccesses[memref].insert(node.id);
777 nodes.insert({node.id, node});
778 } else if (op.getNumRegions() != 0) {
779 // Return false if another region is found (not currently supported).
780 return false;
781 } else if (op.getNumResults() > 0 && !op.use_empty()) {
782 // Create graph node for top-level producer of SSA values, which
783 // could be used by loop nest nodes.
784 Node node(nextNodeId++, &op);
785 nodes.insert({node.id, node});
786 } else if (isa<CallOpInterface>(op)) {
787 // Create graph node for top-level Call Op that takes any argument of
788 // memref type. Call Op that returns one or more memref type results
789 // is already taken care of, by the previous conditions.
790 if (llvm::any_of(op.getOperandTypes(),
791 [&](Type t) { return t.isa<MemRefType>(); })) {
792 Node node(nextNodeId++, &op);
793 nodes.insert({node.id, node});
794 }
795 } else if (auto effectInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
796 // Create graph node for top-level op, which could have a memory write
797 // side effect.
798 SmallVector<MemoryEffects::EffectInstance, 1> effects;
799 effectInterface.getEffects(effects);
800 if (llvm::any_of(effects, [](const MemoryEffects::EffectInstance &it) {
801 return isa<MemoryEffects::Write, MemoryEffects::Free>(
802 it.getEffect());
803 })) {
804 Node node(nextNodeId++, &op);
805 nodes.insert({node.id, node});
806 }
807 }
808 }
809
810 for (auto &idAndNode : nodes) {
811 LLVM_DEBUG(llvm::dbgs() << "Create node " << idAndNode.first << " for:\n"
812 << *(idAndNode.second.op) << "\n");
813 (void)idAndNode;
814 }
815
816 // Add dependence edges between nodes which produce SSA values and their
817 // users. Load ops can be considered as the ones producing SSA values.
818 for (auto &idAndNode : nodes) {
819 const Node &node = idAndNode.second;
820 // Stores don't define SSA values, skip them.
821 if (!node.stores.empty())
822 continue;
823 auto *opInst = node.op;
824 for (auto value : opInst->getResults()) {
825 for (auto *user : value.getUsers()) {
826 SmallVector<AffineForOp, 4> loops;
827 getLoopIVs(*user, &loops);
828 if (loops.empty())
829 continue;
830 assert(forToNodeMap.count(loops[0].getOperation()) > 0);
831 unsigned userLoopNestId = forToNodeMap[loops[0].getOperation()];
832 addEdge(node.id, userLoopNestId, value);
833 }
834 }
835 }
836
837 // Walk memref access lists and add graph edges between dependent nodes.
838 for (auto &memrefAndList : memrefAccesses) {
839 unsigned n = memrefAndList.second.size();
840 for (unsigned i = 0; i < n; ++i) {
841 unsigned srcId = memrefAndList.second[i];
842 bool srcHasStore =
843 getNode(srcId)->getStoreOpCount(memrefAndList.first) > 0;
844 for (unsigned j = i + 1; j < n; ++j) {
845 unsigned dstId = memrefAndList.second[j];
846 bool dstHasStore =
847 getNode(dstId)->getStoreOpCount(memrefAndList.first) > 0;
848 if (srcHasStore || dstHasStore)
849 addEdge(srcId, dstId, memrefAndList.first);
850 }
851 }
852 }
853 return true;
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: 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: 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<AffineWriteOpInterface>(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.getMemorySpaceAsInt();
950 }
951 auto newMemRefType = MemRefType::get(newShape, oldMemRefType.getElementType(),
952 {}, newMemSpace);
953
954 // Create new private memref for fused loop 'forOp'. 'newShape' is always
955 // a constant shape.
956 // TODO: Create/move alloc ops for private memrefs closer to their
957 // consumer loop nests to reduce their live range. Currently they are added
958 // at the beginning of the function, because loop nests can be reordered
959 // during the fusion pass.
960 Value newMemRef = top.create<memref::AllocOp>(forOp.getLoc(), newMemRefType);
961
962 // Build an AffineMap to remap access functions based on lower bound offsets.
963 SmallVector<AffineExpr, 4> remapExprs;
964 remapExprs.reserve(rank);
965 for (unsigned i = 0; i < rank; i++) {
966 auto dimExpr = b.getAffineDimExpr(outerIVs.size() + i);
967
968 auto remapExpr =
969 simplifyAffineExpr(dimExpr - offsets[i], outerIVs.size() + rank, 0);
970 remapExprs.push_back(remapExpr);
971 }
972
973 auto indexRemap =
974 AffineMap::get(outerIVs.size() + rank, 0, remapExprs, forOp.getContext());
975
976 // Replace all users of 'oldMemRef' with 'newMemRef'.
977 LogicalResult res =
978 replaceAllMemRefUsesWith(oldMemRef, newMemRef, {}, indexRemap,
979 /*extraOperands=*/outerIVs,
980 /*symbolOperands=*/{},
981 /*domInstFilter=*/&*forOp.getBody()->begin());
982 assert(succeeded(res) &&
983 "replaceAllMemrefUsesWith should always succeed here");
984 (void)res;
985 return newMemRef;
986 }
987
988 /// Walking from node 'srcId' to node 'dstId' (exclusive of 'srcId' and
989 /// 'dstId'), if there is any non-affine operation accessing 'memref', return
990 /// true. Otherwise, return false.
hasNonAffineUsersOnThePath(unsigned srcId,unsigned dstId,Value memref,MemRefDependenceGraph * mdg)991 static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
992 Value memref,
993 MemRefDependenceGraph *mdg) {
994 auto *srcNode = mdg->getNode(srcId);
995 auto *dstNode = mdg->getNode(dstId);
996 Value::user_range users = memref.getUsers();
997 // For each MemRefDependenceGraph's node that is between 'srcNode' and
998 // 'dstNode' (exclusive of 'srcNodes' and 'dstNode'), check whether any
999 // non-affine operation in the node accesses the 'memref'.
1000 for (auto &idAndNode : mdg->nodes) {
1001 Operation *op = idAndNode.second.op;
1002 // Take care of operations between 'srcNode' and 'dstNode'.
1003 if (srcNode->op->isBeforeInBlock(op) && op->isBeforeInBlock(dstNode->op)) {
1004 // Walk inside the operation to find any use of the memref.
1005 // Interrupt the walk if found.
1006 auto walkResult = op->walk([&](Operation *user) {
1007 // Skip affine ops.
1008 if (isa<AffineMapAccessInterface>(*user))
1009 return WalkResult::advance();
1010 // Find a non-affine op that uses the memref.
1011 if (llvm::is_contained(users, user))
1012 return WalkResult::interrupt();
1013 return WalkResult::advance();
1014 });
1015 if (walkResult.wasInterrupted())
1016 return true;
1017 }
1018 }
1019 return false;
1020 }
1021
1022 /// Check whether a memref value in node 'srcId' has a non-affine that
1023 /// is between node 'srcId' and node 'dstId' (exclusive of 'srcNode' and
1024 /// 'dstNode').
hasNonAffineUsersOnThePath(unsigned srcId,unsigned dstId,MemRefDependenceGraph * mdg)1025 static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
1026 MemRefDependenceGraph *mdg) {
1027 // Collect memref values in node 'srcId'.
1028 auto *srcNode = mdg->getNode(srcId);
1029 llvm::SmallDenseSet<Value, 2> memRefValues;
1030 srcNode->op->walk([&](Operation *op) {
1031 // Skip affine ops.
1032 if (isa<AffineForOp>(op))
1033 return WalkResult::advance();
1034 for (Value v : op->getOperands())
1035 // Collect memref values only.
1036 if (v.getType().isa<MemRefType>())
1037 memRefValues.insert(v);
1038 return WalkResult::advance();
1039 });
1040 // Looking for users between node 'srcId' and node 'dstId'.
1041 for (Value memref : memRefValues)
1042 if (hasNonAffineUsersOnThePath(srcId, dstId, memref, mdg))
1043 return true;
1044 return false;
1045 }
1046
1047 // Checks the profitability of fusing a backwards slice of the loop nest
1048 // surrounding 'srcOpInst' into the loop nest surrounding 'dstLoadOpInsts'.
1049 // The argument 'srcStoreOpInst' is used to calculate the storage reduction on
1050 // the memref being produced and consumed, which is an input to the cost model.
1051 // For producer-consumer fusion, 'srcStoreOpInst' will be the same as
1052 // 'srcOpInst', as we are slicing w.r.t to that producer. For input-reuse
1053 // fusion, 'srcOpInst' will be the src loop nest LoadOp which reads from the
1054 // same memref as dst loop nest load ops, and 'srcStoreOpInst' will be the
1055 // unique store op in the src node, which will be used to check that the write
1056 // region is the same after input-reuse fusion. Computation slices are provided
1057 // in 'depthSliceUnions' for each legal fusion depth. The maximal depth at which
1058 // fusion is legal is provided in 'maxLegalFusionDepth'. Returns true if it is
1059 // profitable to fuse the candidate loop nests. Returns false otherwise.
1060 // `dstLoopDepth` is set to the most profitable depth at which to materialize
1061 // the source loop nest slice.
1062 // The profitability model executes the following steps:
1063 // *) Computes the backward computation slice at 'srcOpInst'. This
1064 // computation slice of the loop nest surrounding 'srcOpInst' is
1065 // represented by modified src loop bounds in 'sliceState', which are
1066 // functions of loop IVs in the loop nest surrounding 'srcOpInst'.
1067 // *) Computes the cost of unfused src/dst loop nests (currently the cost of a
1068 // loop nest is the total number of dynamic operation instances in the loop
1069 // nest).
1070 // *) Computes the cost of fusing a slice of the src loop nest into the dst
1071 // loop nest at various values of dst loop depth, attempting to fuse
1072 // the largest computation slice at the maximal dst loop depth (closest to
1073 // the load) to minimize reuse distance and potentially enable subsequent
1074 // load/store forwarding.
1075 // NOTE: 'dstLoopDepth' refers to the loop depth within the destination loop
1076 // nest, at which the src computation slice is inserted/fused.
1077 // NOTE: We attempt to maximize the dst loop depth, but there are cases
1078 // where a particular setting for 'dstLoopNest' might fuse an unsliced
1079 // loop (within the src computation slice) at a depth which results in
1080 // excessive recomputation (see unit tests for examples).
1081 // *) Compares the total cost of the unfused loop nests to the min cost fused
1082 // loop nest computed in the previous step, and returns true if the latter
1083 // is lower.
1084 // TODO: Extend profitability analysis to support scenarios with multiple
1085 // stores.
isFusionProfitable(Operation * srcOpInst,Operation * srcStoreOpInst,AffineForOp dstForOp,ArrayRef<ComputationSliceState> depthSliceUnions,unsigned maxLegalFusionDepth,unsigned * dstLoopDepth,double computeToleranceThreshold)1086 static bool isFusionProfitable(Operation *srcOpInst, Operation *srcStoreOpInst,
1087 AffineForOp dstForOp,
1088 ArrayRef<ComputationSliceState> depthSliceUnions,
1089 unsigned maxLegalFusionDepth,
1090 unsigned *dstLoopDepth,
1091 double computeToleranceThreshold) {
1092 LLVM_DEBUG({
1093 llvm::dbgs() << "Checking whether fusion is profitable between src op:\n";
1094 llvm::dbgs() << ' ' << *srcOpInst << " and destination loop:\n";
1095 llvm::dbgs() << dstForOp << "\n";
1096 });
1097
1098 if (maxLegalFusionDepth == 0) {
1099 LLVM_DEBUG(llvm::dbgs() << "Can't fuse: maxLegalFusionDepth == 0 .\n");
1100 return false;
1101 }
1102
1103 // Compute cost of sliced and unsliced src loop nest.
1104 SmallVector<AffineForOp, 4> srcLoopIVs;
1105 getLoopIVs(*srcOpInst, &srcLoopIVs);
1106
1107 // Walk src loop nest and collect stats.
1108 LoopNestStats srcLoopNestStats;
1109 if (!getLoopNestStats(srcLoopIVs[0], &srcLoopNestStats))
1110 return false;
1111
1112 // Compute cost of dst loop nest.
1113 LoopNestStats dstLoopNestStats;
1114 if (!getLoopNestStats(dstForOp, &dstLoopNestStats))
1115 return false;
1116
1117 // Search for min cost value for 'dstLoopDepth'. At each value of
1118 // 'dstLoopDepth' from 'maxLegalLoopDepth' to '1', compute computation slice
1119 // bounds between 'srcOpInst' and each op in 'dstOpinsts' (taking the union
1120 // of these bounds). Next the union slice bounds are used to calculate
1121 // the cost of the slice and the cost of the slice inserted into the dst
1122 // loop nest at 'dstLoopDepth'.
1123 uint64_t minFusedLoopNestComputeCost = std::numeric_limits<uint64_t>::max();
1124 double maxStorageReduction = 0.0;
1125 Optional<uint64_t> sliceMemEstimate = None;
1126
1127 // The best loop depth at which to materialize the slice.
1128 Optional<unsigned> bestDstLoopDepth = None;
1129
1130 // Compute op instance count for the src loop nest without iteration slicing.
1131 uint64_t srcLoopNestCost = getComputeCost(srcLoopIVs[0], srcLoopNestStats);
1132
1133 // Compute src loop nest write region size.
1134 MemRefRegion srcWriteRegion(srcStoreOpInst->getLoc());
1135 if (failed(srcWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0))) {
1136 LLVM_DEBUG(llvm::dbgs()
1137 << "Unable to compute MemRefRegion for source operation\n.");
1138 return false;
1139 }
1140
1141 Optional<int64_t> maybeSrcWriteRegionSizeBytes =
1142 srcWriteRegion.getRegionSize();
1143 if (!maybeSrcWriteRegionSizeBytes.hasValue())
1144 return false;
1145 int64_t srcWriteRegionSizeBytes = maybeSrcWriteRegionSizeBytes.getValue();
1146
1147 // Compute op instance count for the src loop nest.
1148 uint64_t dstLoopNestCost = getComputeCost(dstForOp, dstLoopNestStats);
1149
1150 // Evaluate all depth choices for materializing the slice in the destination
1151 // loop nest.
1152 for (unsigned i = maxLegalFusionDepth; i >= 1; --i) {
1153 const ComputationSliceState &slice = depthSliceUnions[i - 1];
1154 // Skip slice union if it wasn't computed for this depth.
1155 if (slice.isEmpty())
1156 continue;
1157
1158 int64_t fusedLoopNestComputeCost;
1159 if (!getFusionComputeCost(srcLoopIVs[0], srcLoopNestStats, dstForOp,
1160 dstLoopNestStats, slice,
1161 &fusedLoopNestComputeCost)) {
1162 LLVM_DEBUG(llvm::dbgs() << "Unable to compute fusion compute cost.\n.");
1163 continue;
1164 }
1165
1166 double additionalComputeFraction =
1167 fusedLoopNestComputeCost /
1168 (static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
1169 1;
1170
1171 // Determine what the slice write MemRefRegion would be, if the src loop
1172 // nest slice 'slice' were to be inserted into the dst loop nest at loop
1173 // depth 'i'.
1174 MemRefRegion sliceWriteRegion(srcStoreOpInst->getLoc());
1175 if (failed(sliceWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0,
1176 &slice))) {
1177 LLVM_DEBUG(llvm::dbgs()
1178 << "Failed to compute slice write region at loopDepth: " << i
1179 << "\n");
1180 continue;
1181 }
1182
1183 Optional<int64_t> maybeSliceWriteRegionSizeBytes =
1184 sliceWriteRegion.getRegionSize();
1185 if (!maybeSliceWriteRegionSizeBytes.hasValue() ||
1186 maybeSliceWriteRegionSizeBytes.getValue() == 0) {
1187 LLVM_DEBUG(llvm::dbgs()
1188 << "Failed to get slice write region size at loopDepth: " << i
1189 << "\n");
1190 continue;
1191 }
1192 int64_t sliceWriteRegionSizeBytes =
1193 maybeSliceWriteRegionSizeBytes.getValue();
1194
1195 // If we are fusing for reuse, check that write regions remain the same.
1196 // TODO: Write region check should check sizes and offsets in
1197 // each dimension, so that we are sure they are covering the same memref
1198 // region. Also, move this out to a isMemRefRegionSuperSet helper function.
1199 if (srcOpInst != srcStoreOpInst &&
1200 sliceWriteRegionSizeBytes != srcWriteRegionSizeBytes)
1201 continue;
1202
1203 double storageReduction = static_cast<double>(srcWriteRegionSizeBytes) /
1204 static_cast<double>(sliceWriteRegionSizeBytes);
1205
1206 LLVM_DEBUG({
1207 std::stringstream msg;
1208 msg << " evaluating fusion profitability at depth : " << i << "\n"
1209 << std::fixed << std::setprecision(2)
1210 << " additional compute fraction: "
1211 << 100.0 * additionalComputeFraction << "%\n"
1212 << " storage reduction factor: " << storageReduction << "x\n"
1213 << " fused nest cost: " << fusedLoopNestComputeCost << "\n"
1214 << " src write region size: " << srcWriteRegionSizeBytes << "\n"
1215 << " slice write region size: " << sliceWriteRegionSizeBytes
1216 << "\n";
1217 llvm::dbgs() << msg.str();
1218 });
1219
1220 // TODO: This is a placeholder cost model.
1221 // Among all choices that add an acceptable amount of redundant computation
1222 // (as per computeToleranceThreshold), we will simply pick the one that
1223 // reduces the intermediary size the most.
1224 if ((storageReduction > maxStorageReduction) &&
1225 (additionalComputeFraction < computeToleranceThreshold)) {
1226 maxStorageReduction = storageReduction;
1227 bestDstLoopDepth = i;
1228 minFusedLoopNestComputeCost = fusedLoopNestComputeCost;
1229 sliceMemEstimate = sliceWriteRegionSizeBytes;
1230 }
1231 }
1232
1233 // A simple cost model: fuse if it reduces the memory footprint.
1234
1235 if (!bestDstLoopDepth.hasValue()) {
1236 LLVM_DEBUG(
1237 llvm::dbgs()
1238 << "All fusion choices involve more than the threshold amount of "
1239 "redundant computation; NOT fusing.\n");
1240 return false;
1241 }
1242
1243 if (!bestDstLoopDepth.hasValue()) {
1244 LLVM_DEBUG(llvm::dbgs() << "no fusion depth could be evaluated.\n");
1245 return false;
1246 }
1247
1248 // Set dstLoopDepth based on best values from search.
1249 *dstLoopDepth = bestDstLoopDepth.getValue();
1250
1251 LLVM_DEBUG(
1252 llvm::dbgs() << " LoopFusion fusion stats:"
1253 << "\n best loop depth: " << bestDstLoopDepth
1254 << "\n src loop nest compute cost: " << srcLoopNestCost
1255 << "\n dst loop nest compute cost: " << dstLoopNestCost
1256 << "\n fused loop nest compute cost: "
1257 << minFusedLoopNestComputeCost << "\n");
1258
1259 auto dstMemSize = getMemoryFootprintBytes(dstForOp);
1260 auto srcMemSize = getMemoryFootprintBytes(srcLoopIVs[0]);
1261
1262 Optional<double> storageReduction = None;
1263
1264 if (!dstMemSize.hasValue() || !srcMemSize.hasValue()) {
1265 LLVM_DEBUG(llvm::dbgs()
1266 << " fusion memory benefit cannot be evaluated; NOT fusing.\n");
1267 return false;
1268 }
1269
1270 auto srcMemSizeVal = srcMemSize.getValue();
1271 auto dstMemSizeVal = dstMemSize.getValue();
1272
1273 assert(sliceMemEstimate.hasValue() && "expected value");
1274 auto fusedMem = dstMemSizeVal + sliceMemEstimate.getValue();
1275
1276 LLVM_DEBUG(llvm::dbgs() << " src mem: " << srcMemSizeVal << "\n"
1277 << " dst mem: " << dstMemSizeVal << "\n"
1278 << " fused mem: " << fusedMem << "\n"
1279 << " slice mem: " << sliceMemEstimate << "\n");
1280
1281 if (static_cast<long>(fusedMem) > srcMemSizeVal + dstMemSizeVal) {
1282 LLVM_DEBUG(llvm::dbgs() << "Fusion is not profitable; NOT fusing.\n");
1283 return false;
1284 }
1285 storageReduction =
1286 100.0 *
1287 (1.0 - fusedMem / (static_cast<double>(srcMemSizeVal) + dstMemSizeVal));
1288
1289 double additionalComputeFraction =
1290 100.0 * (minFusedLoopNestComputeCost /
1291 (static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
1292 1);
1293 (void)additionalComputeFraction;
1294 LLVM_DEBUG({
1295 std::stringstream msg;
1296 msg << " fusion is most profitable at depth " << *dstLoopDepth << " with "
1297 << std::setprecision(2) << additionalComputeFraction
1298 << "% redundant computation and a ";
1299 msg << (storageReduction.hasValue()
1300 ? std::to_string(storageReduction.getValue())
1301 : "<unknown>");
1302 msg << "% storage reduction.\n";
1303 llvm::dbgs() << msg.str();
1304 });
1305
1306 return true;
1307 }
1308
1309 namespace {
1310
1311 // GreedyFusion greedily fuses loop nests which have a producer/consumer or
1312 // input-reuse relationship on a memref, with the goal of improving locality.
1313 //
1314 // The steps of the producer-consumer fusion algorithm are as follows:
1315 //
1316 // *) A worklist is initialized with node ids from the dependence graph.
1317 // *) For each node id in the worklist:
1318 // *) Pop an AffineForOp of the worklist. This 'dstAffineForOp' will be a
1319 // candidate destination AffineForOp into which fusion will be attempted.
1320 // *) Add each LoadOp currently in 'dstAffineForOp' into list 'dstLoadOps'.
1321 // *) For each LoadOp in 'dstLoadOps' do:
1322 // *) Look up dependent loop nests which have a single store op to the same
1323 // memref.
1324 // *) Check if dependences would be violated by the fusion.
1325 // *) Get a computation slice of 'srcLoopNest', which adjusts its loop
1326 // bounds to be functions of 'dstLoopNest' IVs and symbols.
1327 // *) Fuse the 'srcLoopNest' computation slice into the 'dstLoopNest',
1328 // at a loop depth determined by the cost model in 'isFusionProfitable'.
1329 // *) Add the newly fused load/store operations to the state,
1330 // and also add newly fused load ops to 'dstLoopOps' to be considered
1331 // as fusion dst load ops in another iteration.
1332 // *) Remove old src loop nest and its associated state.
1333 //
1334 // The steps of the input-reuse fusion algorithm are as follows:
1335 //
1336 // *) Initialize 'worklist' with node ids from the dependence graph.
1337 // *) For each 'dstNode' in the worklist:
1338 // *) Find a candidate sibling node 'sibNode' to fuse with 'dstNode' which
1339 // loads from the same memref, but which has no dependence paths to/from.
1340 // *) Get a computation slice of 'sibLoopNest', which adjusts its loop
1341 // bounds to be functions of 'dstLoopNest' IVs and symbols.
1342 // *) Fuse the 'sibLoopNest' computation slice into the 'dstLoopNest',
1343 // at a loop depth determined by the cost model in 'isFusionProfitable'.
1344 // This function also checks that the memref write region of 'sibLoopNest',
1345 // is preserved in the fused loop nest.
1346 // *) Update graph state to reflect the fusion of 'sibNode' into 'dstNode'.
1347 //
1348 // Given a graph where top-level operations are vertices in the set 'V' and
1349 // edges in the set 'E' are dependences between vertices, this algorithm
1350 // takes O(V) time for initialization, and has runtime O(V + E).
1351 //
1352 // This greedy algorithm is not 'maximal' due to the current restriction of
1353 // fusing along single producer consumer edges, but there is a TODO: to fix
1354 // this.
1355 //
1356 // TODO: Experiment with other fusion policies.
1357 struct GreedyFusion {
1358 public:
1359 // The data dependence graph to traverse during fusion.
1360 MemRefDependenceGraph *mdg;
1361 // Worklist of graph nodes visited during the fusion pass.
1362 SmallVector<unsigned, 8> worklist;
1363 // Parameter for local buffer size threshold.
1364 unsigned localBufSizeThreshold;
1365 // Parameter for fast memory space.
1366 Optional<unsigned> fastMemorySpace;
1367 // If true, ignore any additional (redundant) computation tolerance threshold
1368 // that would have prevented fusion.
1369 bool maximalFusion;
1370 // The amount of additional computation that is tolerated while fusing
1371 // pair-wise as a fraction of the total computation.
1372 double computeToleranceThreshold;
1373
1374 using Node = MemRefDependenceGraph::Node;
1375
GreedyFusion__anon0903053b0c11::GreedyFusion1376 GreedyFusion(MemRefDependenceGraph *mdg, unsigned localBufSizeThreshold,
1377 Optional<unsigned> fastMemorySpace, bool maximalFusion,
1378 double computeToleranceThreshold)
1379 : mdg(mdg), localBufSizeThreshold(localBufSizeThreshold),
1380 fastMemorySpace(fastMemorySpace), maximalFusion(maximalFusion),
1381 computeToleranceThreshold(computeToleranceThreshold) {}
1382
1383 /// Initializes 'worklist' with nodes from 'mdg'.
init__anon0903053b0c11::GreedyFusion1384 void init() {
1385 // TODO: Add a priority queue for prioritizing nodes by different
1386 // metrics (e.g. arithmetic intensity/flops-to-bytes ratio).
1387 worklist.clear();
1388 for (auto &idAndNode : mdg->nodes) {
1389 const Node &node = idAndNode.second;
1390 worklist.push_back(node.id);
1391 }
1392 }
1393
1394 // Run the GreedyFusion pass.
1395 // *) First pass through the nodes fuses single-use producer nodes into their
1396 // unique consumer.
1397 // *) Second pass fuses sibling nodes which share no dependence edges.
1398 // *) Third pass fuses any remaining producer nodes into their users.
run__anon0903053b0c11::GreedyFusion1399 void run() {
1400 // TODO: Run this repeatedly until a fixed-point is reached.
1401 fuseProducerConsumerNodes(/*maxSrcUserCount=*/1);
1402 fuseSiblingNodes();
1403 fuseProducerConsumerNodes(
1404 /*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
1405 eraseUnusedMemRefAllocations();
1406 }
1407
fuseProducerConsumerNodes__anon0903053b0c11::GreedyFusion1408 void fuseProducerConsumerNodes(unsigned maxSrcUserCount) {
1409 LLVM_DEBUG(llvm::dbgs() << "--- Producer/Consumer Fusion ---\n");
1410 init();
1411 while (!worklist.empty()) {
1412 unsigned dstId = worklist.back();
1413 worklist.pop_back();
1414
1415 // Skip if this node was removed (fused into another node).
1416 if (mdg->nodes.count(dstId) == 0)
1417 continue;
1418 // Get 'dstNode' into which to attempt fusion.
1419 auto *dstNode = mdg->getNode(dstId);
1420 // Skip if 'dstNode' is not a loop nest.
1421 if (!isa<AffineForOp>(dstNode->op))
1422 continue;
1423 // Skip if 'dstNode' is a loop nest returning values.
1424 // TODO: support loop nests that return values.
1425 if (dstNode->op->getNumResults() > 0)
1426 continue;
1427
1428 LLVM_DEBUG(llvm::dbgs() << "Evaluating dst loop " << dstId << "\n");
1429
1430 // Sink sequential loops in 'dstNode' (and thus raise parallel loops)
1431 // while preserving relative order. This can increase the maximum loop
1432 // depth at which we can fuse a slice of a producer loop nest into a
1433 // consumer loop nest.
1434 sinkSequentialLoops(dstNode);
1435 auto dstAffineForOp = cast<AffineForOp>(dstNode->op);
1436
1437 // Try to fuse 'dstNode' with candidate producer loops until a fixed point
1438 // is reached. Fusing two loops may expose new fusion opportunities.
1439 bool dstNodeChanged;
1440 do {
1441 // Gather src loop candidates for 'dstNode' and visit them in "quasi"
1442 // reverse program order to minimize the number of iterations needed to
1443 // reach the fixed point. Note that this is a best effort approach since
1444 // 'getProducerCandidates' does not always guarantee that program order
1445 // in 'srcIdCandidates'.
1446 dstNodeChanged = false;
1447 SmallVector<unsigned, 16> srcIdCandidates;
1448 getProducerCandidates(dstId, mdg, srcIdCandidates);
1449
1450 for (unsigned srcId : llvm::reverse(srcIdCandidates)) {
1451 // Get 'srcNode' from which to attempt fusion into 'dstNode'.
1452 auto *srcNode = mdg->getNode(srcId);
1453 auto srcAffineForOp = cast<AffineForOp>(srcNode->op);
1454 LLVM_DEBUG(llvm::dbgs() << "Evaluating src loop " << srcId
1455 << " for dst loop " << dstId << "\n");
1456
1457 // Skip if 'srcNode' is a loop nest returning values.
1458 // TODO: support loop nests that return values.
1459 if (isa<AffineForOp>(srcNode->op) && srcNode->op->getNumResults() > 0)
1460 continue;
1461
1462 DenseSet<Value> producerConsumerMemrefs;
1463 gatherProducerConsumerMemrefs(srcId, dstId, mdg,
1464 producerConsumerMemrefs);
1465
1466 // Skip if 'srcNode' out edge count on any memref is greater than
1467 // 'maxSrcUserCount'.
1468 if (any_of(producerConsumerMemrefs, [&](Value memref) {
1469 return mdg->getOutEdgeCount(srcNode->id, memref) >
1470 maxSrcUserCount;
1471 }))
1472 continue;
1473
1474 // Gather memrefs in 'srcNode' that are written and escape to the
1475 // function (e.g., memref function arguments, returned memrefs,
1476 // memrefs passed to function calls, etc.).
1477 DenseSet<Value> srcEscapingMemRefs;
1478 gatherEscapingMemrefs(srcNode->id, mdg, srcEscapingMemRefs);
1479
1480 // Skip if there are non-affine operations in between the 'srcNode'
1481 // and 'dstNode' using their memrefs. If so, we wouldn't be able to
1482 // compute a legal insertion point for now. 'srcNode' and 'dstNode'
1483 // memrefs with non-affine operation users would be considered
1484 // escaping memrefs so we can limit this check to only scenarios with
1485 // escaping memrefs.
1486 if (!srcEscapingMemRefs.empty() &&
1487 hasNonAffineUsersOnThePath(srcId, dstId, mdg)) {
1488 LLVM_DEBUG(
1489 llvm::dbgs()
1490 << "Can't fuse: non-affine users in between the loops\n.");
1491 continue;
1492 }
1493
1494 // Compute an operation list insertion point for the fused loop
1495 // nest which preserves dependences.
1496 Operation *fusedLoopInsPoint =
1497 mdg->getFusedLoopNestInsertionPoint(srcNode->id, dstNode->id);
1498 if (fusedLoopInsPoint == nullptr)
1499 continue;
1500
1501 // Compute the innermost common loop depth for dstNode
1502 // producer-consumer loads/stores.
1503 SmallVector<Operation *, 2> dstMemrefOps;
1504 for (Operation *op : dstNode->loads)
1505 if (producerConsumerMemrefs.count(
1506 cast<AffineReadOpInterface>(op).getMemRef()) > 0)
1507 dstMemrefOps.push_back(op);
1508 for (Operation *op : dstNode->stores)
1509 if (producerConsumerMemrefs.count(
1510 cast<AffineWriteOpInterface>(op).getMemRef()))
1511 dstMemrefOps.push_back(op);
1512 unsigned dstLoopDepthTest = getInnermostCommonLoopDepth(dstMemrefOps);
1513
1514 // Check the feasibility of fusing src loop nest into dst loop nest
1515 // at loop depths in range [1, dstLoopDepthTest].
1516 unsigned maxLegalFusionDepth = 0;
1517 SmallVector<ComputationSliceState, 8> depthSliceUnions;
1518 depthSliceUnions.resize(dstLoopDepthTest);
1519 FusionStrategy strategy(FusionStrategy::ProducerConsumer);
1520 for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
1521 FusionResult result = mlir::canFuseLoops(
1522 srcAffineForOp, dstAffineForOp,
1523 /*dstLoopDepth=*/i, &depthSliceUnions[i - 1], strategy);
1524
1525 if (result.value == FusionResult::Success)
1526 maxLegalFusionDepth = i;
1527 }
1528
1529 if (maxLegalFusionDepth == 0) {
1530 LLVM_DEBUG(llvm::dbgs()
1531 << "Can't fuse: fusion is not legal at any depth\n");
1532 continue;
1533 }
1534
1535 // Check if fusion would be profitable. We skip profitability analysis
1536 // for maximal fusion since we already know the maximal legal depth to
1537 // fuse.
1538 unsigned bestDstLoopDepth = maxLegalFusionDepth;
1539 if (!maximalFusion) {
1540 // Retrieve producer stores from the src loop.
1541 SmallVector<Operation *, 2> producerStores;
1542 for (Operation *op : srcNode->stores)
1543 if (producerConsumerMemrefs.count(
1544 cast<AffineWriteOpInterface>(op).getMemRef()))
1545 producerStores.push_back(op);
1546
1547 // TODO: Suppport multiple producer stores in profitability
1548 // analysis. We limit profitability analysis to only scenarios with
1549 // a single producer store for now. Note that some multi-store
1550 // producer scenarios will still go through profitability analysis
1551 // if only one of the stores is involved the producer-consumer
1552 // relationship of the candidate loops.
1553 assert(producerStores.size() > 0 && "Expected producer store");
1554 if (producerStores.size() > 1)
1555 LLVM_DEBUG(llvm::dbgs() << "Skipping profitability analysis. Not "
1556 "supported for this case\n");
1557 else if (!isFusionProfitable(producerStores[0], producerStores[0],
1558 dstAffineForOp, depthSliceUnions,
1559 maxLegalFusionDepth, &bestDstLoopDepth,
1560 computeToleranceThreshold))
1561 continue;
1562 }
1563
1564 assert(bestDstLoopDepth > 0 && "Unexpected loop fusion depth");
1565 ComputationSliceState &bestSlice =
1566 depthSliceUnions[bestDstLoopDepth - 1];
1567 assert(!bestSlice.isEmpty() && "Missing slice union for depth");
1568
1569 // Determine if 'srcId' can be removed after fusion, taking into
1570 // account remaining dependences, escaping memrefs and the fusion
1571 // insertion point.
1572 bool removeSrcNode = canRemoveSrcNodeAfterFusion(
1573 srcId, dstId, bestSlice, fusedLoopInsPoint, srcEscapingMemRefs,
1574 mdg);
1575
1576 DenseSet<Value> privateMemrefs;
1577 for (Value memref : producerConsumerMemrefs) {
1578 // If `memref` is an escaping one, do not create a private memref
1579 // for the below scenarios, since doing so will leave the escaping
1580 // memref unmodified as all the writes originally meant for the
1581 // escaping memref would be performed on the private memref:
1582 // 1. The source is to be removed after fusion,
1583 // OR
1584 // 2. The destination writes to `memref`.
1585 if (srcEscapingMemRefs.count(memref) > 0 &&
1586 (removeSrcNode || dstNode->getStoreOpCount(memref) > 0))
1587 continue;
1588
1589 // Don't create a private memref if 'srcNode' has in edges on
1590 // 'memref' or 'dstNode' has out edges on 'memref'.
1591 if (mdg->getIncomingMemRefAccesses(srcId, memref) > 0 ||
1592 mdg->getOutEdgeCount(dstId, memref) > 0)
1593 continue;
1594
1595 // If 'srcNode' will be removed but it has out edges on 'memref' to
1596 // nodes other than 'dstNode', we have to preserve dependences and
1597 // cannot create a private memref.
1598 if (removeSrcNode &&
1599 any_of(mdg->outEdges[srcId], [&](const auto &edge) {
1600 return edge.value == memref && edge.id != dstId;
1601 }))
1602 continue;
1603
1604 // Create a private version of this memref.
1605 privateMemrefs.insert(memref);
1606 }
1607
1608 // Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'.
1609 fuseLoops(srcAffineForOp, dstAffineForOp, bestSlice);
1610 dstNodeChanged = true;
1611
1612 LLVM_DEBUG(llvm::dbgs()
1613 << "Fused src loop " << srcId << " into dst loop " << dstId
1614 << " at depth " << bestDstLoopDepth << ":\n"
1615 << dstAffineForOp << "\n");
1616
1617 // Move 'dstAffineForOp' before 'insertPointInst' if needed.
1618 if (fusedLoopInsPoint != dstAffineForOp.getOperation())
1619 dstAffineForOp.getOperation()->moveBefore(fusedLoopInsPoint);
1620
1621 // Update edges between 'srcNode' and 'dstNode'.
1622 mdg->updateEdges(srcNode->id, dstNode->id, privateMemrefs,
1623 removeSrcNode);
1624
1625 // Create private memrefs.
1626 if (!privateMemrefs.empty()) {
1627 // Gather stores for all the private-to-be memrefs.
1628 DenseMap<Value, SmallVector<Operation *, 4>> privateMemRefToStores;
1629 dstAffineForOp.walk([&](AffineWriteOpInterface storeOp) {
1630 Value storeMemRef = storeOp.getMemRef();
1631 if (privateMemrefs.count(storeMemRef) > 0)
1632 privateMemRefToStores[storeMemRef].push_back(
1633 storeOp.getOperation());
1634 });
1635
1636 // Replace original memrefs with private memrefs. Note that all the
1637 // loads and stores on these memrefs will be replaced with a new
1638 // loads and stores. Any reference to the original ones becomes
1639 // invalid after this point.
1640 for (auto &memrefToStoresPair : privateMemRefToStores) {
1641 // TODO: Use union of memref write regions to compute
1642 // private memref footprint.
1643 SmallVector<Operation *, 4> &storesForMemref =
1644 memrefToStoresPair.second;
1645 Value newMemRef = createPrivateMemRef(
1646 dstAffineForOp, storesForMemref[0], bestDstLoopDepth,
1647 fastMemorySpace, localBufSizeThreshold);
1648 // Create new node in dependence graph for 'newMemRef' alloc op.
1649 unsigned newMemRefNodeId =
1650 mdg->addNode(newMemRef.getDefiningOp());
1651 // Add edge from 'newMemRef' node to dstNode.
1652 mdg->addEdge(newMemRefNodeId, dstId, newMemRef);
1653 }
1654 // One or more entries for 'newMemRef' alloc op are inserted into
1655 // the DenseMap mdg->nodes. Since an insertion may cause DenseMap to
1656 // reallocate, update dstNode.
1657 dstNode = mdg->getNode(dstId);
1658 }
1659
1660 // Collect dst loop stats after memref privatization transformation.
1661 LoopNestStateCollector dstLoopCollector;
1662 dstLoopCollector.collect(dstAffineForOp.getOperation());
1663
1664 // Clear and add back loads and stores.
1665 mdg->clearNodeLoadAndStores(dstNode->id);
1666 mdg->addToNode(dstId, dstLoopCollector.loadOpInsts,
1667 dstLoopCollector.storeOpInsts);
1668
1669 if (removeSrcNode) {
1670 LLVM_DEBUG(llvm::dbgs()
1671 << "Removing src loop " << srcId << " after fusion\n");
1672 // srcNode is no longer valid after it is removed from mdg.
1673 srcAffineForOp.erase();
1674 mdg->removeNode(srcId);
1675 srcNode = nullptr;
1676 }
1677 }
1678 } while (dstNodeChanged);
1679 }
1680 }
1681
1682 // Visits each node in the graph, and for each node, attempts to fuse it with
1683 // its sibling nodes (nodes which share a parent, but no dependence edges).
fuseSiblingNodes__anon0903053b0c11::GreedyFusion1684 void fuseSiblingNodes() {
1685 LLVM_DEBUG(llvm::dbgs() << "--- Sibling Fusion ---\n");
1686 init();
1687 while (!worklist.empty()) {
1688 unsigned dstId = worklist.back();
1689 worklist.pop_back();
1690
1691 // Skip if this node was removed (fused into another node).
1692 if (mdg->nodes.count(dstId) == 0)
1693 continue;
1694 // Get 'dstNode' into which to attempt fusion.
1695 auto *dstNode = mdg->getNode(dstId);
1696 // Skip if 'dstNode' is not a loop nest.
1697 if (!isa<AffineForOp>(dstNode->op))
1698 continue;
1699 // Attempt to fuse 'dstNode' with its sibling nodes in the graph.
1700 fuseWithSiblingNodes(dstNode);
1701 }
1702 }
1703
1704 // Attempt to fuse 'dstNode' with sibling nodes in the graph.
fuseWithSiblingNodes__anon0903053b0c11::GreedyFusion1705 void fuseWithSiblingNodes(Node *dstNode) {
1706 DenseSet<unsigned> visitedSibNodeIds;
1707 std::pair<unsigned, Value> idAndMemref;
1708 auto dstAffineForOp = cast<AffineForOp>(dstNode->op);
1709
1710 while (findSiblingNodeToFuse(dstNode, &visitedSibNodeIds, &idAndMemref)) {
1711 unsigned sibId = idAndMemref.first;
1712 Value memref = idAndMemref.second;
1713 // TODO: Check that 'sibStoreOpInst' post-dominates all other
1714 // stores to the same memref in 'sibNode' loop nest.
1715 auto *sibNode = mdg->getNode(sibId);
1716 // Compute an operation list insertion point for the fused loop
1717 // nest which preserves dependences.
1718 assert(sibNode->op->getBlock() == dstNode->op->getBlock());
1719 Operation *insertPointInst =
1720 sibNode->op->isBeforeInBlock(dstNode->op)
1721 ? mdg->getFusedLoopNestInsertionPoint(sibNode->id, dstNode->id)
1722 : mdg->getFusedLoopNestInsertionPoint(dstNode->id, sibNode->id);
1723 if (insertPointInst == nullptr)
1724 continue;
1725
1726 // Check if fusion would be profitable and at what depth.
1727
1728 // Get unique 'sibNode' load op to 'memref'.
1729 SmallVector<Operation *, 2> sibLoadOpInsts;
1730 sibNode->getLoadOpsForMemref(memref, &sibLoadOpInsts);
1731 // Currently findSiblingNodeToFuse searches for siblings with one load.
1732 assert(sibLoadOpInsts.size() == 1);
1733 Operation *sibLoadOpInst = sibLoadOpInsts[0];
1734 assert(!sibNode->stores.empty());
1735 // TODO: Choose the store which postdominates all other stores.
1736 auto *sibStoreOpInst = sibNode->stores.back();
1737
1738 // Gather 'dstNode' load ops to 'memref'.
1739 SmallVector<Operation *, 2> dstLoadOpInsts;
1740 dstNode->getLoadOpsForMemref(memref, &dstLoadOpInsts);
1741
1742 SmallVector<AffineForOp, 4> dstLoopIVs;
1743 getLoopIVs(*dstLoadOpInsts[0], &dstLoopIVs);
1744 unsigned dstLoopDepthTest = dstLoopIVs.size();
1745 auto sibAffineForOp = cast<AffineForOp>(sibNode->op);
1746
1747 // Compute loop depth and slice union for fusion.
1748 SmallVector<ComputationSliceState, 8> depthSliceUnions;
1749 depthSliceUnions.resize(dstLoopDepthTest);
1750 unsigned maxLegalFusionDepth = 0;
1751 FusionStrategy strategy(memref);
1752 for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
1753 FusionResult result = mlir::canFuseLoops(
1754 sibAffineForOp, dstAffineForOp,
1755 /*dstLoopDepth=*/i, &depthSliceUnions[i - 1], strategy);
1756
1757 if (result.value == FusionResult::Success)
1758 maxLegalFusionDepth = i;
1759 }
1760
1761 // Skip if fusion is not feasible at any loop depths.
1762 if (maxLegalFusionDepth == 0)
1763 continue;
1764
1765 unsigned bestDstLoopDepth = maxLegalFusionDepth;
1766 if (!maximalFusion) {
1767 // Check if fusion would be profitable.
1768 if (!isFusionProfitable(sibLoadOpInst, sibStoreOpInst, dstAffineForOp,
1769 depthSliceUnions, maxLegalFusionDepth,
1770 &bestDstLoopDepth, computeToleranceThreshold))
1771 continue;
1772 }
1773
1774 assert(bestDstLoopDepth > 0 && "Unexpected loop fusion depth");
1775 assert(!depthSliceUnions[bestDstLoopDepth - 1].isEmpty() &&
1776 "Fusion depth has no computed slice union");
1777 // Check if source loop is being inserted in the innermost
1778 // destination loop. Based on this, the fused loop may be optimized
1779 // further inside `fuseLoops`.
1780 bool isInnermostInsertion = (bestDstLoopDepth == dstLoopDepthTest);
1781 // Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'.
1782 mlir::fuseLoops(sibAffineForOp, dstAffineForOp,
1783 depthSliceUnions[bestDstLoopDepth - 1],
1784 isInnermostInsertion);
1785
1786 auto dstForInst = cast<AffineForOp>(dstNode->op);
1787 // Update operation position of fused loop nest (if needed).
1788 if (insertPointInst != dstForInst.getOperation()) {
1789 dstForInst->moveBefore(insertPointInst);
1790 }
1791 // Update data dependence graph state post fusion.
1792 updateStateAfterSiblingFusion(sibNode, dstNode);
1793 }
1794 }
1795
1796 // Searches function argument uses and the graph from 'dstNode' looking for a
1797 // fusion candidate sibling node which shares no dependences with 'dstNode'
1798 // but which loads from the same memref. Returns true and sets
1799 // 'idAndMemrefToFuse' on success. Returns false otherwise.
findSiblingNodeToFuse__anon0903053b0c11::GreedyFusion1800 bool findSiblingNodeToFuse(Node *dstNode,
1801 DenseSet<unsigned> *visitedSibNodeIds,
1802 std::pair<unsigned, Value> *idAndMemrefToFuse) {
1803 // Returns true if 'sibNode' can be fused with 'dstNode' for input reuse
1804 // on 'memref'.
1805 auto canFuseWithSibNode = [&](Node *sibNode, Value memref) {
1806 // Skip if 'outEdge' is not a read-after-write dependence.
1807 // TODO: Remove restrict to single load op restriction.
1808 if (sibNode->getLoadOpCount(memref) != 1)
1809 return false;
1810 // Skip if there exists a path of dependent edges between
1811 // 'sibNode' and 'dstNode'.
1812 if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
1813 mdg->hasDependencePath(dstNode->id, sibNode->id))
1814 return false;
1815 // Skip sib node if it loads to (and stores from) the same memref on
1816 // which it also has an input dependence edge.
1817 DenseSet<Value> loadAndStoreMemrefSet;
1818 sibNode->getLoadAndStoreMemrefSet(&loadAndStoreMemrefSet);
1819 if (llvm::any_of(loadAndStoreMemrefSet, [=](Value memref) {
1820 return mdg->getIncomingMemRefAccesses(sibNode->id, memref) > 0;
1821 }))
1822 return false;
1823
1824 // Check that all stores are to the same memref.
1825 DenseSet<Value> storeMemrefs;
1826 for (auto *storeOpInst : sibNode->stores) {
1827 storeMemrefs.insert(
1828 cast<AffineWriteOpInterface>(storeOpInst).getMemRef());
1829 }
1830 if (storeMemrefs.size() != 1)
1831 return false;
1832
1833 // Skip if a memref value in one node is used by a non-affine memref
1834 // access that lies between 'dstNode' and 'sibNode'.
1835 if (hasNonAffineUsersOnThePath(dstNode->id, sibNode->id, mdg) ||
1836 hasNonAffineUsersOnThePath(sibNode->id, dstNode->id, mdg))
1837 return false;
1838 return true;
1839 };
1840
1841 // Search for siblings which load the same memref function argument.
1842 auto fn = dstNode->op->getParentOfType<FuncOp>();
1843 for (unsigned i = 0, e = fn.getNumArguments(); i != e; ++i) {
1844 for (auto *user : fn.getArgument(i).getUsers()) {
1845 if (auto loadOp = dyn_cast<AffineReadOpInterface>(user)) {
1846 // Gather loops surrounding 'use'.
1847 SmallVector<AffineForOp, 4> loops;
1848 getLoopIVs(*user, &loops);
1849 // Skip 'use' if it is not within a loop nest.
1850 if (loops.empty())
1851 continue;
1852 Node *sibNode = mdg->getForOpNode(loops[0]);
1853 assert(sibNode != nullptr);
1854 // Skip 'use' if it not a sibling to 'dstNode'.
1855 if (sibNode->id == dstNode->id)
1856 continue;
1857 // Skip 'use' if it has been visited.
1858 if (visitedSibNodeIds->count(sibNode->id) > 0)
1859 continue;
1860 // Skip 'use' if it does not load from the same memref as 'dstNode'.
1861 auto memref = loadOp.getMemRef();
1862 if (dstNode->getLoadOpCount(memref) == 0)
1863 continue;
1864 // Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
1865 if (canFuseWithSibNode(sibNode, memref)) {
1866 visitedSibNodeIds->insert(sibNode->id);
1867 idAndMemrefToFuse->first = sibNode->id;
1868 idAndMemrefToFuse->second = memref;
1869 return true;
1870 }
1871 }
1872 }
1873 }
1874
1875 // Search for siblings by following edges through an intermediate src node.
1876 // Collect candidate 'dstNode' input edges in 'inEdges'.
1877 SmallVector<MemRefDependenceGraph::Edge, 2> inEdges;
1878 mdg->forEachMemRefInputEdge(
1879 dstNode->id, [&](MemRefDependenceGraph::Edge inEdge) {
1880 // Add 'inEdge' if it is a read-after-write dependence.
1881 if (dstNode->getLoadOpCount(inEdge.value) > 0 &&
1882 mdg->getNode(inEdge.id)->getStoreOpCount(inEdge.value) > 0)
1883 inEdges.push_back(inEdge);
1884 });
1885
1886 // Search for sibling nodes to fuse by visiting output edges from each input
1887 // edge in 'inEdges'.
1888 for (auto &inEdge : inEdges) {
1889 // Collect candidate output edges from each node 'inEdge.id' in 'inEdges'.
1890 SmallVector<MemRefDependenceGraph::Edge, 2> outEdges;
1891 mdg->forEachMemRefOutputEdge(
1892 inEdge.id, [&](MemRefDependenceGraph::Edge outEdge) {
1893 unsigned sibNodeId = outEdge.id;
1894 if (visitedSibNodeIds->count(sibNodeId) > 0)
1895 return;
1896 // Skip output edge if not a sibling using the same memref.
1897 if (outEdge.id == dstNode->id || outEdge.value != inEdge.value)
1898 return;
1899 auto *sibNode = mdg->getNode(sibNodeId);
1900 if (!isa<AffineForOp>(sibNode->op))
1901 return;
1902 // Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
1903 if (canFuseWithSibNode(sibNode, outEdge.value)) {
1904 // Add candidate 'outEdge' to sibling node.
1905 outEdges.push_back(outEdge);
1906 }
1907 });
1908
1909 // Add first candidate if any were returned.
1910 if (!outEdges.empty()) {
1911 visitedSibNodeIds->insert(outEdges[0].id);
1912 idAndMemrefToFuse->first = outEdges[0].id;
1913 idAndMemrefToFuse->second = outEdges[0].value;
1914 return true;
1915 }
1916 }
1917 return false;
1918 }
1919
1920 /// Update data dependence graph state to reflect sibling fusion of 'sibNode'
1921 /// into 'dstNode'.
updateStateAfterSiblingFusion__anon0903053b0c11::GreedyFusion1922 void updateStateAfterSiblingFusion(Node *sibNode, Node *dstNode) {
1923 // Update 'sibNode' and 'dstNode' input/output edges to reflect fusion.
1924 mdg->updateEdges(sibNode->id, dstNode->id);
1925
1926 // Collect dst loop stats after memref privatization transformation.
1927 auto dstForInst = cast<AffineForOp>(dstNode->op);
1928 LoopNestStateCollector dstLoopCollector;
1929 dstLoopCollector.collect(dstForInst.getOperation());
1930 // Clear and add back loads and stores
1931 mdg->clearNodeLoadAndStores(dstNode->id);
1932 mdg->addToNode(dstNode->id, dstLoopCollector.loadOpInsts,
1933 dstLoopCollector.storeOpInsts);
1934 // Remove old sibling loop nest if it no longer has outgoing dependence
1935 // edges, and it does not write to a memref which escapes the
1936 // function.
1937 if (mdg->getOutEdgeCount(sibNode->id) == 0) {
1938 mdg->removeNode(sibNode->id);
1939 sibNode->op->erase();
1940 }
1941 }
1942
1943 // Clean up any allocs with no users.
eraseUnusedMemRefAllocations__anon0903053b0c11::GreedyFusion1944 void eraseUnusedMemRefAllocations() {
1945 for (auto &pair : mdg->memrefEdgeCount) {
1946 if (pair.second > 0)
1947 continue;
1948 auto memref = pair.first;
1949 // Skip if there exist other uses (return operation or function calls).
1950 if (!memref.use_empty())
1951 continue;
1952 // Use list expected to match the dep graph info.
1953 auto *op = memref.getDefiningOp();
1954 if (isa_and_nonnull<memref::AllocOp>(op))
1955 op->erase();
1956 }
1957 }
1958 };
1959
1960 } // end anonymous namespace
1961
runOnFunction()1962 void LoopFusion::runOnFunction() {
1963 MemRefDependenceGraph g;
1964 if (!g.init(getFunction()))
1965 return;
1966
1967 Optional<unsigned> fastMemorySpaceOpt;
1968 if (fastMemorySpace.hasValue())
1969 fastMemorySpaceOpt = fastMemorySpace;
1970 unsigned localBufSizeThresholdBytes = localBufSizeThreshold * 1024;
1971 GreedyFusion fusion(&g, localBufSizeThresholdBytes, fastMemorySpaceOpt,
1972 maximalFusion, computeToleranceThreshold);
1973 fusion.run();
1974 }
1975