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