Revert "[mlir][Affine] Add support for multi-store producer fusion"
[lldb.git] / mlir / lib / Transforms / LoopFusion.cpp
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;
52   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>>
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
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
87   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;
120     Node(unsigned id, Operation *op) : id(id), op(op) {}
121
122     // Returns the load op count for 'memref'.
123     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'.
133     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'.
143     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'.
152     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.
162     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
205   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'.
212   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'.
219   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.
227   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.
234   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.
257   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 the unique AffineWriteOpInterface in `node` that meets all the
274   // following:
275   //   *) store is the only one that writes to a function-local memref live out
276   //      of `node`,
277   //   *) store is not the source of a self-dependence on `node`.
278   // Otherwise, returns a null AffineWriteOpInterface.
279   AffineWriteOpInterface getUniqueOutgoingStore(Node *node) {
280     AffineWriteOpInterface uniqueStore;
281
282     // Return null if `node` doesn't have any outgoing edges.
283     auto outEdgeIt = outEdges.find(node->id);
284     if (outEdgeIt == outEdges.end())
285       return nullptr;
286
287     const auto &nodeOutEdges = outEdgeIt->second;
288     for (auto *op : node->stores) {
289       auto storeOp = cast<AffineWriteOpInterface>(op);
290       auto memref = storeOp.getMemRef();
291       // Skip this store if there are no dependences on its memref. This means
292       // that store either:
293       // *) writes to a memref that is only read within the same loop nest
294       //    (self-dependence edges are not represented in graph at the moment),
295       // *) writes to a function live out memref (function parameter), or
296       // *) is dead.
297       if (llvm::all_of(nodeOutEdges, [=](const Edge &edge) {
298             return (edge.value != memref);
299           }))
300         continue;
301
302       if (uniqueStore)
303         // Found multiple stores to function-local live-out memrefs.
304         return nullptr;
305       // Found first store to function-local live-out memref.
306       uniqueStore = storeOp;
307     }
308
309     return uniqueStore;
310   }
311
312   // Returns true if node 'id' can be removed from the graph. Returns false
313   // otherwise. A node can be removed from the graph iff the following
314   // conditions are met:
315   // *) The node does not write to any memref which escapes (or is a
316   //    function/block argument).
317   // *) The node has no successors in the dependence graph.
318   bool canRemoveNode(unsigned id) {
319     if (writesToLiveInOrEscapingMemrefs(id))
320       return false;
321     Node *node = getNode(id);
322     for (auto *storeOpInst : node->stores) {
323       // Return false if there exist out edges from 'id' on 'memref'.
324       auto storeMemref = cast<AffineWriteOpInterface>(storeOpInst).getMemRef();
325       if (getOutEdgeCount(id, storeMemref) > 0)
326         return false;
327     }
328     return true;
329   }
330
331   // Returns true iff there is an edge from node 'srcId' to node 'dstId' which
332   // is for 'value' if non-null, or for any value otherwise. Returns false
333   // otherwise.
334   bool hasEdge(unsigned srcId, unsigned dstId, Value value = nullptr) {
335     if (outEdges.count(srcId) == 0 || inEdges.count(dstId) == 0) {
336       return false;
337     }
338     bool hasOutEdge = llvm::any_of(outEdges[srcId], [=](Edge &edge) {
339       return edge.id == dstId && (!value || edge.value == value);
340     });
341     bool hasInEdge = llvm::any_of(inEdges[dstId], [=](Edge &edge) {
342       return edge.id == srcId && (!value || edge.value == value);
343     });
344     return hasOutEdge && hasInEdge;
345   }
346
347   // Adds an edge from node 'srcId' to node 'dstId' for 'value'.
348   void addEdge(unsigned srcId, unsigned dstId, Value value) {
349     if (!hasEdge(srcId, dstId, value)) {
350       outEdges[srcId].push_back({dstId, value});
351       inEdges[dstId].push_back({srcId, value});
352       if (value.getType().isa<MemRefType>())
353         memrefEdgeCount[value]++;
354     }
355   }
356
357   // Removes an edge from node 'srcId' to node 'dstId' for 'value'.
358   void removeEdge(unsigned srcId, unsigned dstId, Value value) {
359     assert(inEdges.count(dstId) > 0);
360     assert(outEdges.count(srcId) > 0);
361     if (value.getType().isa<MemRefType>()) {
362       assert(memrefEdgeCount.count(value) > 0);
363       memrefEdgeCount[value]--;
364     }
365     // Remove 'srcId' from 'inEdges[dstId]'.
366     for (auto it = inEdges[dstId].begin(); it != inEdges[dstId].end(); ++it) {
367       if ((*it).id == srcId && (*it).value == value) {
368         inEdges[dstId].erase(it);
369         break;
370       }
371     }
372     // Remove 'dstId' from 'outEdges[srcId]'.
373     for (auto it = outEdges[srcId].begin(); it != outEdges[srcId].end(); ++it) {
374       if ((*it).id == dstId && (*it).value == value) {
375         outEdges[srcId].erase(it);
376         break;
377       }
378     }
379   }
380
381   // Returns true if there is a path in the dependence graph from node 'srcId'
382   // to node 'dstId'. Returns false otherwise.
383   bool hasDependencePath(unsigned srcId, unsigned dstId) {
384     // Worklist state is: <node-id, next-output-edge-index-to-visit>
385     SmallVector<std::pair<unsigned, unsigned>, 4> worklist;
386     worklist.push_back({srcId, 0});
387     // Run DFS traversal to see if 'dstId' is reachable from 'srcId'.
388     while (!worklist.empty()) {
389       auto &idAndIndex = worklist.back();
390       // Return true if we have reached 'dstId'.
391       if (idAndIndex.first == dstId)
392         return true;
393       // Pop and continue if node has no out edges, or if all out edges have
394       // already been visited.
395       if (outEdges.count(idAndIndex.first) == 0 ||
396           idAndIndex.second == outEdges[idAndIndex.first].size()) {
397         worklist.pop_back();
398         continue;
399       }
400       // Get graph edge to traverse.
401       Edge edge = outEdges[idAndIndex.first][idAndIndex.second];
402       // Increment next output edge index for 'idAndIndex'.
403       ++idAndIndex.second;
404       // Add node at 'edge.id' to worklist.
405       worklist.push_back({edge.id, 0});
406     }
407     return false;
408   }
409
410   // Returns the input edge count for node 'id' and 'memref' from src nodes
411   // which access 'memref' with a store operation.
412   unsigned getIncomingMemRefAccesses(unsigned id, Value memref) {
413     unsigned inEdgeCount = 0;
414     if (inEdges.count(id) > 0)
415       for (auto &inEdge : inEdges[id])
416         if (inEdge.value == memref) {
417           Node *srcNode = getNode(inEdge.id);
418           // Only count in edges from 'srcNode' if 'srcNode' accesses 'memref'
419           if (srcNode->getStoreOpCount(memref) > 0)
420             ++inEdgeCount;
421         }
422     return inEdgeCount;
423   }
424
425   // Returns the output edge count for node 'id' and 'memref' (if non-null),
426   // otherwise returns the total output edge count from node 'id'.
427   unsigned getOutEdgeCount(unsigned id, Value memref = nullptr) {
428     unsigned outEdgeCount = 0;
429     if (outEdges.count(id) > 0)
430       for (auto &outEdge : outEdges[id])
431         if (!memref || outEdge.value == memref)
432           ++outEdgeCount;
433     return outEdgeCount;
434   }
435
436   // Computes and returns an insertion point operation, before which the
437   // the fused <srcId, dstId> loop nest can be inserted while preserving
438   // dependences. Returns nullptr if no such insertion point is found.
439   Operation *getFusedLoopNestInsertionPoint(unsigned srcId, unsigned dstId) {
440     if (outEdges.count(srcId) == 0)
441       return getNode(dstId)->op;
442
443     // Build set of insts in range (srcId, dstId) which depend on 'srcId'.
444     SmallPtrSet<Operation *, 2> srcDepInsts;
445     for (auto &outEdge : outEdges[srcId])
446       if (outEdge.id != dstId)
447         srcDepInsts.insert(getNode(outEdge.id)->op);
448
449     // Build set of insts in range (srcId, dstId) on which 'dstId' depends.
450     SmallPtrSet<Operation *, 2> dstDepInsts;
451     for (auto &inEdge : inEdges[dstId])
452       if (inEdge.id != srcId)
453         dstDepInsts.insert(getNode(inEdge.id)->op);
454
455     Operation *srcNodeInst = getNode(srcId)->op;
456     Operation *dstNodeInst = getNode(dstId)->op;
457
458     // Computing insertion point:
459     // *) Walk all operation positions in Block operation list in the
460     //    range (src, dst). For each operation 'op' visited in this search:
461     //   *) Store in 'firstSrcDepPos' the first position where 'op' has a
462     //      dependence edge from 'srcNode'.
463     //   *) Store in 'lastDstDepPost' the last position where 'op' has a
464     //      dependence edge to 'dstNode'.
465     // *) Compare 'firstSrcDepPos' and 'lastDstDepPost' to determine the
466     //    operation insertion point (or return null pointer if no such
467     //    insertion point exists: 'firstSrcDepPos' <= 'lastDstDepPos').
468     SmallVector<Operation *, 2> depInsts;
469     Optional<unsigned> firstSrcDepPos;
470     Optional<unsigned> lastDstDepPos;
471     unsigned pos = 0;
472     for (Block::iterator it = std::next(Block::iterator(srcNodeInst));
473          it != Block::iterator(dstNodeInst); ++it) {
474       Operation *op = &(*it);
475       if (srcDepInsts.count(op) > 0 && firstSrcDepPos == None)
476         firstSrcDepPos = pos;
477       if (dstDepInsts.count(op) > 0)
478         lastDstDepPos = pos;
479       depInsts.push_back(op);
480       ++pos;
481     }
482
483     if (firstSrcDepPos.hasValue()) {
484       if (lastDstDepPos.hasValue()) {
485         if (firstSrcDepPos.getValue() <= lastDstDepPos.getValue()) {
486           // No valid insertion point exists which preserves dependences.
487           return nullptr;
488         }
489       }
490       // Return the insertion point at 'firstSrcDepPos'.
491       return depInsts[firstSrcDepPos.getValue()];
492     }
493     // No dependence targets in range (or only dst deps in range), return
494     // 'dstNodInst' insertion point.
495     return dstNodeInst;
496   }
497
498   // Updates edge mappings from node 'srcId' to node 'dstId' after 'oldMemRef'
499   // has been replaced in node at 'dstId' by a private memref depending
500   // on the value of 'createPrivateMemRef'.
501   void updateEdges(unsigned srcId, unsigned dstId, Value oldMemRef,
502                    bool createPrivateMemRef) {
503     // For each edge in 'inEdges[srcId]': add new edge remapping to 'dstId'.
504     if (inEdges.count(srcId) > 0) {
505       SmallVector<Edge, 2> oldInEdges = inEdges[srcId];
506       for (auto &inEdge : oldInEdges) {
507         // Add edge from 'inEdge.id' to 'dstId' if not for 'oldMemRef'.
508         if (inEdge.value != oldMemRef)
509           addEdge(inEdge.id, dstId, inEdge.value);
510       }
511     }
512     // For each edge in 'outEdges[srcId]': remove edge from 'srcId' to 'dstId'.
513     if (outEdges.count(srcId) > 0) {
514       SmallVector<Edge, 2> oldOutEdges = outEdges[srcId];
515       for (auto &outEdge : oldOutEdges) {
516         // Remove any out edges from 'srcId' to 'dstId' across memrefs.
517         if (outEdge.id == dstId)
518           removeEdge(srcId, outEdge.id, outEdge.value);
519       }
520     }
521     // Remove any edges in 'inEdges[dstId]' on 'oldMemRef' (which is being
522     // replaced by a private memref). These edges could come from nodes
523     // other than 'srcId' which were removed in the previous step.
524     if (inEdges.count(dstId) > 0 && createPrivateMemRef) {
525       SmallVector<Edge, 2> oldInEdges = inEdges[dstId];
526       for (auto &inEdge : oldInEdges)
527         if (inEdge.value == oldMemRef)
528           removeEdge(inEdge.id, dstId, inEdge.value);
529     }
530   }
531
532   // Update edge mappings for nodes 'sibId' and 'dstId' to reflect fusion
533   // of sibling node 'sidId' into node 'dstId'.
534   void updateEdges(unsigned sibId, unsigned dstId) {
535     // For each edge in 'inEdges[sibId]':
536     // *) Add new edge from source node 'inEdge.id' to 'dstNode'.
537     // *) Remove edge from source node 'inEdge.id' to 'sibNode'.
538     if (inEdges.count(sibId) > 0) {
539       SmallVector<Edge, 2> oldInEdges = inEdges[sibId];
540       for (auto &inEdge : oldInEdges) {
541         addEdge(inEdge.id, dstId, inEdge.value);
542         removeEdge(inEdge.id, sibId, inEdge.value);
543       }
544     }
545
546     // For each edge in 'outEdges[sibId]' to node 'id'
547     // *) Add new edge from 'dstId' to 'outEdge.id'.
548     // *) Remove edge from 'sibId' to 'outEdge.id'.
549     if (outEdges.count(sibId) > 0) {
550       SmallVector<Edge, 2> oldOutEdges = outEdges[sibId];
551       for (auto &outEdge : oldOutEdges) {
552         addEdge(dstId, outEdge.id, outEdge.value);
553         removeEdge(sibId, outEdge.id, outEdge.value);
554       }
555     }
556   }
557
558   // Adds ops in 'loads' and 'stores' to node at 'id'.
559   void addToNode(unsigned id, const SmallVectorImpl<Operation *> &loads,
560                  const SmallVectorImpl<Operation *> &stores) {
561     Node *node = getNode(id);
562     for (auto *loadOpInst : loads)
563       node->loads.push_back(loadOpInst);
564     for (auto *storeOpInst : stores)
565       node->stores.push_back(storeOpInst);
566   }
567
568   void clearNodeLoadAndStores(unsigned id) {
569     Node *node = getNode(id);
570     node->loads.clear();
571     node->stores.clear();
572   }
573
574   // Calls 'callback' for each input edge incident to node 'id' which carries a
575   // memref dependence.
576   void forEachMemRefInputEdge(unsigned id,
577                               const std::function<void(Edge)> &callback) {
578     if (inEdges.count(id) > 0)
579       forEachMemRefEdge(inEdges[id], callback);
580   }
581
582   // Calls 'callback' for each output edge from node 'id' which carries a
583   // memref dependence.
584   void forEachMemRefOutputEdge(unsigned id,
585                                const std::function<void(Edge)> &callback) {
586     if (outEdges.count(id) > 0)
587       forEachMemRefEdge(outEdges[id], callback);
588   }
589
590   // Calls 'callback' for each edge in 'edges' which carries a memref
591   // dependence.
592   void forEachMemRefEdge(ArrayRef<Edge> edges,
593                          const std::function<void(Edge)> &callback) {
594     for (const auto &edge : edges) {
595       // Skip if 'edge' is not a memref dependence edge.
596       if (!edge.value.getType().isa<MemRefType>())
597         continue;
598       assert(nodes.count(edge.id) > 0);
599       // Skip if 'edge.id' is not a loop nest.
600       if (!isa<AffineForOp>(getNode(edge.id)->op))
601         continue;
602       // Visit current input edge 'edge'.
603       callback(edge);
604     }
605   }
606
607   void print(raw_ostream &os) const {
608     os << "\nMemRefDependenceGraph\n";
609     os << "\nNodes:\n";
610     for (const auto &idAndNode : nodes) {
611       os << "Node: " << idAndNode.first << "\n";
612       auto it = inEdges.find(idAndNode.first);
613       if (it != inEdges.end()) {
614         for (const auto &e : it->second)
615           os << "  InEdge: " << e.id << " " << e.value << "\n";
616       }
617       it = outEdges.find(idAndNode.first);
618       if (it != outEdges.end()) {
619         for (const auto &e : it->second)
620           os << "  OutEdge: " << e.id << " " << e.value << "\n";
621       }
622     }
623   }
624   void dump() const { print(llvm::errs()); }
625 };
626
627 } // end anonymous namespace
628
629 // Initializes the data dependence graph by walking operations in 'f'.
630 // Assigns each node in the graph a node id based on program order in 'f'.
631 // TODO: Add support for taking a Block arg to construct the
632 // dependence graph at a different depth.
633 bool MemRefDependenceGraph::init(FuncOp f) {
634   DenseMap<Value, SetVector<unsigned>> memrefAccesses;
635
636   // TODO: support multi-block functions.
637   if (!llvm::hasSingleElement(f))
638     return false;
639
640   DenseMap<Operation *, unsigned> forToNodeMap;
641   for (auto &op : f.front()) {
642     if (auto forOp = dyn_cast<AffineForOp>(op)) {
643       // Create graph node 'id' to represent top-level 'forOp' and record
644       // all loads and store accesses it contains.
645       LoopNestStateCollector collector;
646       collector.collect(&op);
647       // Return false if a non 'affine.for' region was found (not currently
648       // supported).
649       if (collector.hasNonForRegion)
650         return false;
651       Node node(nextNodeId++, &op);
652       for (auto *opInst : collector.loadOpInsts) {
653         node.loads.push_back(opInst);
654         auto memref = cast<AffineReadOpInterface>(opInst).getMemRef();
655         memrefAccesses[memref].insert(node.id);
656       }
657       for (auto *opInst : collector.storeOpInsts) {
658         node.stores.push_back(opInst);
659         auto memref = cast<AffineWriteOpInterface>(opInst).getMemRef();
660         memrefAccesses[memref].insert(node.id);
661       }
662       forToNodeMap[&op] = node.id;
663       nodes.insert({node.id, node});
664     } else if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
665       // Create graph node for top-level load op.
666       Node node(nextNodeId++, &op);
667       node.loads.push_back(&op);
668       auto memref = cast<AffineReadOpInterface>(op).getMemRef();
669       memrefAccesses[memref].insert(node.id);
670       nodes.insert({node.id, node});
671     } else if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
672       // Create graph node for top-level store op.
673       Node node(nextNodeId++, &op);
674       node.stores.push_back(&op);
675       auto memref = cast<AffineWriteOpInterface>(op).getMemRef();
676       memrefAccesses[memref].insert(node.id);
677       nodes.insert({node.id, node});
678     } else if (op.getNumRegions() != 0) {
679       // Return false if another region is found (not currently supported).
680       return false;
681     } else if (op.getNumResults() > 0 && !op.use_empty()) {
682       // Create graph node for top-level producer of SSA values, which
683       // could be used by loop nest nodes.
684       Node node(nextNodeId++, &op);
685       nodes.insert({node.id, node});
686     }
687   }
688
689   // Add dependence edges between nodes which produce SSA values and their
690   // users.
691   for (auto &idAndNode : nodes) {
692     const Node &node = idAndNode.second;
693     if (!node.loads.empty() || !node.stores.empty())
694       continue;
695     auto *opInst = node.op;
696     for (auto value : opInst->getResults()) {
697       for (auto *user : value.getUsers()) {
698         SmallVector<AffineForOp, 4> loops;
699         getLoopIVs(*user, &loops);
700         if (loops.empty())
701           continue;
702         assert(forToNodeMap.count(loops[0].getOperation()) > 0);
703         unsigned userLoopNestId = forToNodeMap[loops[0].getOperation()];
704         addEdge(node.id, userLoopNestId, value);
705       }
706     }
707   }
708
709   // Walk memref access lists and add graph edges between dependent nodes.
710   for (auto &memrefAndList : memrefAccesses) {
711     unsigned n = memrefAndList.second.size();
712     for (unsigned i = 0; i < n; ++i) {
713       unsigned srcId = memrefAndList.second[i];
714       bool srcHasStore =
715           getNode(srcId)->getStoreOpCount(memrefAndList.first) > 0;
716       for (unsigned j = i + 1; j < n; ++j) {
717         unsigned dstId = memrefAndList.second[j];
718         bool dstHasStore =
719             getNode(dstId)->getStoreOpCount(memrefAndList.first) > 0;
720         if (srcHasStore || dstHasStore)
721           addEdge(srcId, dstId, memrefAndList.first);
722       }
723     }
724   }
725   return true;
726 }
727
728 // Removes load operations from 'srcLoads' which operate on 'memref', and
729 // adds them to 'dstLoads'.
730 static void moveLoadsAccessingMemrefTo(Value memref,
731                                        SmallVectorImpl<Operation *> *srcLoads,
732                                        SmallVectorImpl<Operation *> *dstLoads) {
733   dstLoads->clear();
734   SmallVector<Operation *, 4> srcLoadsToKeep;
735   for (auto *load : *srcLoads) {
736     if (cast<AffineReadOpInterface>(load).getMemRef() == memref)
737       dstLoads->push_back(load);
738     else
739       srcLoadsToKeep.push_back(load);
740   }
741   srcLoads->swap(srcLoadsToKeep);
742 }
743
744 // Sinks all sequential loops to the innermost levels (while preserving
745 // relative order among them) and moves all parallel loops to the
746 // outermost (while again preserving relative order among them).
747 // This can increase the loop depth at which we can fuse a slice, since we are
748 // pushing loop carried dependence to a greater depth in the loop nest.
749 static void sinkSequentialLoops(MemRefDependenceGraph::Node *node) {
750   assert(isa<AffineForOp>(node->op));
751   AffineForOp newRootForOp = sinkSequentialLoops(cast<AffineForOp>(node->op));
752   node->op = newRootForOp.getOperation();
753 }
754
755 //  TODO: improve/complete this when we have target data.
756 static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
757   auto elementType = memRefType.getElementType();
758
759   unsigned sizeInBits;
760   if (elementType.isIntOrFloat()) {
761     sizeInBits = elementType.getIntOrFloatBitWidth();
762   } else {
763     auto vectorType = elementType.cast<VectorType>();
764     sizeInBits =
765         vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
766   }
767   return llvm::divideCeil(sizeInBits, 8);
768 }
769
770 // Creates and returns a private (single-user) memref for fused loop rooted
771 // at 'forOp', with (potentially reduced) memref size based on the
772 // MemRefRegion written to by 'srcStoreOpInst' at depth 'dstLoopDepth'.
773 // TODO: consider refactoring the common code from generateDma and
774 // this one.
775 static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
776                                  unsigned dstLoopDepth,
777                                  Optional<unsigned> fastMemorySpace,
778                                  uint64_t localBufSizeThreshold) {
779   auto *forInst = forOp.getOperation();
780
781   // Create builder to insert alloc op just before 'forOp'.
782   OpBuilder b(forInst);
783   // Builder to create constants at the top level.
784   OpBuilder top(forInst->getParentOfType<FuncOp>().getBody());
785   // Create new memref type based on slice bounds.
786   auto oldMemRef = cast<AffineWriteOpInterface>(srcStoreOpInst).getMemRef();
787   auto oldMemRefType = oldMemRef.getType().cast<MemRefType>();
788   unsigned rank = oldMemRefType.getRank();
789
790   // Compute MemRefRegion for 'srcStoreOpInst' at depth 'dstLoopDepth'.
791   MemRefRegion region(srcStoreOpInst->getLoc());
792   bool validRegion = succeeded(region.compute(srcStoreOpInst, dstLoopDepth));
793   (void)validRegion;
794   assert(validRegion && "unexpected memref region failure");
795   SmallVector<int64_t, 4> newShape;
796   std::vector<SmallVector<int64_t, 4>> lbs;
797   SmallVector<int64_t, 8> lbDivisors;
798   lbs.reserve(rank);
799   // Query 'region' for 'newShape' and lower bounds of MemRefRegion accessed
800   // by 'srcStoreOpInst' at depth 'dstLoopDepth'.
801   Optional<int64_t> numElements =
802       region.getConstantBoundingSizeAndShape(&newShape, &lbs, &lbDivisors);
803   assert(numElements.hasValue() &&
804          "non-constant number of elts in local buffer");
805
806   const FlatAffineConstraints *cst = region.getConstraints();
807   // 'outerIVs' holds the values that this memory region is symbolic/parametric
808   // on; this would correspond to loop IVs surrounding the level at which the
809   // slice is being materialized.
810   SmallVector<Value, 8> outerIVs;
811   cst->getIdValues(rank, cst->getNumIds(), &outerIVs);
812
813   // Build 'rank' AffineExprs from MemRefRegion 'lbs'
814   SmallVector<AffineExpr, 4> offsets;
815   offsets.reserve(rank);
816   for (unsigned d = 0; d < rank; ++d) {
817     assert(lbs[d].size() == cst->getNumCols() - rank && "incorrect bound size");
818
819     AffineExpr offset = top.getAffineConstantExpr(0);
820     for (unsigned j = 0, e = cst->getNumCols() - rank - 1; j < e; j++) {
821       offset = offset + lbs[d][j] * top.getAffineDimExpr(j);
822     }
823     assert(lbDivisors[d] > 0);
824     offset =
825         (offset + lbs[d][cst->getNumCols() - 1 - rank]).floorDiv(lbDivisors[d]);
826     offsets.push_back(offset);
827   }
828
829   // Create 'newMemRefType' using 'newShape' from MemRefRegion accessed
830   // by 'srcStoreOpInst'.
831   uint64_t bufSize =
832       getMemRefEltSizeInBytes(oldMemRefType) * numElements.getValue();
833   unsigned newMemSpace;
834   if (bufSize <= localBufSizeThreshold && fastMemorySpace.hasValue()) {
835     newMemSpace = fastMemorySpace.getValue();
836   } else {
837     newMemSpace = oldMemRefType.getMemorySpace();
838   }
839   auto newMemRefType = MemRefType::get(newShape, oldMemRefType.getElementType(),
840                                        {}, newMemSpace);
841
842   // Create new private memref for fused loop 'forOp'. 'newShape' is always
843   // a constant shape.
844   // TODO: Create/move alloc ops for private memrefs closer to their
845   // consumer loop nests to reduce their live range. Currently they are added
846   // at the beginning of the function, because loop nests can be reordered
847   // during the fusion pass.
848   Value newMemRef = top.create<AllocOp>(forOp.getLoc(), newMemRefType);
849
850   // Build an AffineMap to remap access functions based on lower bound offsets.
851   SmallVector<AffineExpr, 4> remapExprs;
852   remapExprs.reserve(rank);
853   for (unsigned i = 0; i < rank; i++) {
854     auto dimExpr = b.getAffineDimExpr(outerIVs.size() + i);
855
856     auto remapExpr =
857         simplifyAffineExpr(dimExpr - offsets[i], outerIVs.size() + rank, 0);
858     remapExprs.push_back(remapExpr);
859   }
860
861   auto indexRemap =
862       AffineMap::get(outerIVs.size() + rank, 0, remapExprs, forOp.getContext());
863
864   // Replace all users of 'oldMemRef' with 'newMemRef'.
865   LogicalResult res =
866       replaceAllMemRefUsesWith(oldMemRef, newMemRef, {}, indexRemap,
867                                /*extraOperands=*/outerIVs,
868                                /*symbolOperands=*/{},
869                                /*domInstFilter=*/&*forOp.getBody()->begin());
870   assert(succeeded(res) &&
871          "replaceAllMemrefUsesWith should always succeed here");
872   (void)res;
873   return newMemRef;
874 }
875
876 /// Walking from node 'srcId' to node 'dstId' (exclusive of 'srcId' and
877 /// 'dstId'), if there is any non-affine operation accessing 'memref', return
878 /// false. Otherwise, return true.
879 static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
880                                        Value memref,
881                                        MemRefDependenceGraph *mdg) {
882   auto *srcNode = mdg->getNode(srcId);
883   auto *dstNode = mdg->getNode(dstId);
884   Value::user_range users = memref.getUsers();
885   // For each MemRefDependenceGraph's node that is between 'srcNode' and
886   // 'dstNode' (exclusive of 'srcNodes' and 'dstNode'), check whether any
887   // non-affine operation in the node accesses the 'memref'.
888   for (auto &idAndNode : mdg->nodes) {
889     Operation *op = idAndNode.second.op;
890     // Take care of operations between 'srcNode' and 'dstNode'.
891     if (srcNode->op->isBeforeInBlock(op) && op->isBeforeInBlock(dstNode->op)) {
892       // Walk inside the operation to find any use of the memref.
893       // Interrupt the walk if found.
894       auto walkResult = op->walk([&](Operation *user) {
895         // Skip affine ops.
896         if (isMemRefDereferencingOp(*user))
897           return WalkResult::advance();
898         // Find a non-affine op that uses the memref.
899         if (llvm::is_contained(users, user))
900           return WalkResult::interrupt();
901         return WalkResult::advance();
902       });
903       if (walkResult.wasInterrupted())
904         return true;
905     }
906   }
907   return false;
908 }
909
910 /// Check whether a memref value in node 'srcId' has a non-affine that
911 /// is between node 'srcId' and node 'dstId' (exclusive of 'srcNode' and
912 /// 'dstNode').
913 static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
914                                        MemRefDependenceGraph *mdg) {
915   // Collect memref values in node 'srcId'.
916   auto *srcNode = mdg->getNode(srcId);
917   llvm::SmallDenseSet<Value, 2> memRefValues;
918   srcNode->op->walk([&](Operation *op) {
919     // Skip affine ops.
920     if (isa<AffineForOp>(op))
921       return WalkResult::advance();
922     for (Value v : op->getOperands())
923       // Collect memref values only.
924       if (v.getType().isa<MemRefType>())
925         memRefValues.insert(v);
926     return WalkResult::advance();
927   });
928   // Looking for users between node 'srcId' and node 'dstId'.
929   for (Value memref : memRefValues)
930     if (hasNonAffineUsersOnThePath(srcId, dstId, memref, mdg))
931       return true;
932   return false;
933 }
934
935 // Checks if node 'srcId' can be safely fused into node 'dstId'. Node 'srcId'
936 // may write to multiple memrefs but it is required that only one of them,
937 // 'srcLiveOutStoreOp', has output edges.
938 // Returns true if 'dstNode's read/write region to 'memref' is a super set of
939 // 'srcNode's write region to 'memref' and 'srcId' has only one output edge.
940 // TODO: Generalize this to handle more live in/out cases.
941 static bool
942 canFuseSrcWhichWritesToLiveOut(unsigned srcId, unsigned dstId,
943                                AffineWriteOpInterface srcLiveOutStoreOp,
944                                MemRefDependenceGraph *mdg) {
945   assert(srcLiveOutStoreOp && "Expected a valid store op");
946   auto *dstNode = mdg->getNode(dstId);
947   Value memref = srcLiveOutStoreOp.getMemRef();
948   // Return false if 'srcNode' has more than one output edge on 'memref'.
949   if (mdg->getOutEdgeCount(srcId, memref) > 1)
950     return false;
951
952   // Compute MemRefRegion 'srcWriteRegion' for 'srcStoreOp' on 'memref'.
953   MemRefRegion srcWriteRegion(srcLiveOutStoreOp.getLoc());
954   if (failed(srcWriteRegion.compute(srcLiveOutStoreOp, /*loopDepth=*/0))) {
955     LLVM_DEBUG(llvm::dbgs()
956                << "Unable to compute MemRefRegion for source operation\n.");
957     return false;
958   }
959   SmallVector<int64_t, 4> srcShape;
960   // Query 'srcWriteRegion' for 'srcShape' and 'srcNumElements'.
961   // by 'srcStoreOp' at depth 'dstLoopDepth'.
962   Optional<int64_t> srcNumElements =
963       srcWriteRegion.getConstantBoundingSizeAndShape(&srcShape);
964   if (!srcNumElements.hasValue())
965     return false;
966
967   // Compute MemRefRegion 'dstRegion' for 'dstStore/LoadOpInst' on 'memref'.
968   // TODO: Compute 'unionboundingbox' of all write regions (one for
969   // each store op in 'dstStoreOps').
970   SmallVector<Operation *, 2> dstStoreOps;
971   dstNode->getStoreOpsForMemref(memref, &dstStoreOps);
972   SmallVector<Operation *, 2> dstLoadOps;
973   dstNode->getLoadOpsForMemref(memref, &dstLoadOps);
974
975   auto *dstOpInst = dstStoreOps.empty() ? dstLoadOps[0] : dstStoreOps[0];
976   MemRefRegion dstRegion(dstOpInst->getLoc());
977   if (failed(dstRegion.compute(dstOpInst, /*loopDepth=*/0))) {
978     LLVM_DEBUG(llvm::dbgs()
979                << "Unable to compute MemRefRegion for dest operation\n.");
980     return false;
981   }
982   SmallVector<int64_t, 4> dstShape;
983   // Query 'dstRegion' for 'dstShape' and 'dstNumElements'.
984   // by 'dstOpInst' at depth 'dstLoopDepth'.
985   Optional<int64_t> dstNumElements =
986       dstRegion.getConstantBoundingSizeAndShape(&dstShape);
987   if (!dstNumElements.hasValue())
988     return false;
989
990   // Return false if write region is not a superset of 'srcNodes' write
991   // region to 'memref'.
992   // TODO: Check the shape and lower bounds here too.
993   if (srcNumElements != dstNumElements)
994     return false;
995
996   // Return false if 'memref' is used by a non-affine operation that is
997   // between node 'srcId' and node 'dstId'.
998   if (hasNonAffineUsersOnThePath(srcId, dstId, mdg))
999     return false;
1000
1001   return true;
1002 }
1003
1004 // Checks the profitability of fusing a backwards slice of the loop nest
1005 // surrounding 'srcOpInst' into the loop nest surrounding 'dstLoadOpInsts'.
1006 // The argument 'srcStoreOpInst' is used to calculate the storage reduction on
1007 // the memref being produced and consumed, which is an input to the cost model.
1008 // For producer-consumer fusion, 'srcStoreOpInst' will be the same as
1009 // 'srcOpInst', as we are slicing w.r.t to that producer. For input-reuse
1010 // fusion, 'srcOpInst' will be the src loop nest LoadOp which reads from the
1011 // same memref as dst loop nest load ops, and 'srcStoreOpInst' will be the
1012 // unique store op in the src node, which will be used to check that the write
1013 // region is the same after input-reuse fusion. Computation slices are provided
1014 // in 'depthSliceUnions' for each legal fusion depth. The maximal depth at which
1015 // fusion is legal is provided in 'maxLegalFusionDepth'. Returns true if it is
1016 // profitable to fuse the candidate loop nests. Returns false otherwise.
1017 // `dstLoopDepth` is set to the most profitable depth at which to materialize
1018 // the source loop nest slice.
1019 // The profitability model executes the following steps:
1020 // *) Computes the backward computation slice at 'srcOpInst'. This
1021 //    computation slice of the loop nest surrounding 'srcOpInst' is
1022 //    represented by modified src loop bounds in 'sliceState', which are
1023 //    functions of loop IVs in the loop nest surrounding 'srcOpInst'.
1024 // *) Computes the cost of unfused src/dst loop nests (currently the cost of a
1025 //    loop nest is the total number of dynamic operation instances in the loop
1026 //    nest).
1027 // *) Computes the cost of fusing a slice of the src loop nest into the dst
1028 //    loop nest at various values of dst loop depth, attempting to fuse
1029 //    the largest computation slice at the maximal dst loop depth (closest to
1030 //    the load) to minimize reuse distance and potentially enable subsequent
1031 //    load/store forwarding.
1032 //    NOTE: If the dst loop nest includes multiple loads in 'dstLoadOpInsts' for
1033 //    the same memref as is written by 'srcOpInst', then the union of slice
1034 //    loop bounds is used to compute the slice and associated slice cost.
1035 //    NOTE: 'dstLoopDepth' refers to the loop depth within the destination loop
1036 //    nest, at which the src computation slice is inserted/fused.
1037 //    NOTE: We attempt to maximize the dst loop depth, but there are cases
1038 //    where a particular setting for 'dstLoopNest' might fuse an unsliced
1039 //    loop (within the src computation slice) at a depth which results in
1040 //    excessive recomputation (see unit tests for examples).
1041 // *) Compares the total cost of the unfused loop nests to the min cost fused
1042 //    loop nest computed in the previous step, and returns true if the latter
1043 //    is lower.
1044 static bool isFusionProfitable(Operation *srcOpInst, Operation *srcStoreOpInst,
1045                                ArrayRef<Operation *> dstLoadOpInsts,
1046                                ArrayRef<ComputationSliceState> depthSliceUnions,
1047                                unsigned maxLegalFusionDepth,
1048                                unsigned *dstLoopDepth,
1049                                double computeToleranceThreshold) {
1050   LLVM_DEBUG({
1051     llvm::dbgs() << "Checking whether fusion is profitable between src op:\n";
1052     llvm::dbgs() << ' ' << *srcOpInst << " and destination op(s)\n";
1053     for (auto dstOpInst : dstLoadOpInsts) {
1054       llvm::dbgs() << " " << *dstOpInst << "\n";
1055     };
1056   });
1057
1058   if (maxLegalFusionDepth == 0) {
1059     LLVM_DEBUG(llvm::dbgs() << "Can't fuse: maxLegalFusionDepth == 0 .\n");
1060     return false;
1061   }
1062
1063   // Compute cost of sliced and unsliced src loop nest.
1064   SmallVector<AffineForOp, 4> srcLoopIVs;
1065   getLoopIVs(*srcOpInst, &srcLoopIVs);
1066
1067   // Walk src loop nest and collect stats.
1068   LoopNestStats srcLoopNestStats;
1069   if (!getLoopNestStats(srcLoopIVs[0], &srcLoopNestStats))
1070     return false;
1071
1072   // Compute cost of dst loop nest.
1073   SmallVector<AffineForOp, 4> dstLoopIVs;
1074   getLoopIVs(*dstLoadOpInsts[0], &dstLoopIVs);
1075
1076   LoopNestStats dstLoopNestStats;
1077   if (!getLoopNestStats(dstLoopIVs[0], &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(dstLoopIVs[0], 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     // Skip slice union if it wasn't computed for this depth.
1117     if (depthSliceUnions[i - 1].isEmpty())
1118       continue;
1119
1120     int64_t fusedLoopNestComputeCost;
1121     if (!getFusionComputeCost(srcLoopIVs[0], srcLoopNestStats, dstLoopIVs[0],
1122                               dstLoopNestStats, depthSliceUnions[i - 1],
1123                               &fusedLoopNestComputeCost)) {
1124       LLVM_DEBUG(llvm::dbgs() << "Unable to compute fusion compute cost.\n.");
1125       continue;
1126     }
1127
1128     double additionalComputeFraction =
1129         fusedLoopNestComputeCost /
1130             (static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
1131         1;
1132
1133     // Determine what the slice write MemRefRegion would be, if the src loop
1134     // nest slice 'depthSliceUnions[i - 1]' were to be inserted into the dst
1135     // loop nest at loop depth 'i'.
1136     MemRefRegion sliceWriteRegion(srcStoreOpInst->getLoc());
1137     if (failed(sliceWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0,
1138                                         &depthSliceUnions[i - 1]))) {
1139       LLVM_DEBUG(llvm::dbgs()
1140                  << "Failed to compute slice write region at loopDepth: " << i
1141                  << "\n");
1142       continue;
1143     }
1144
1145     Optional<int64_t> maybeSliceWriteRegionSizeBytes =
1146         sliceWriteRegion.getRegionSize();
1147     if (!maybeSliceWriteRegionSizeBytes.hasValue() ||
1148         maybeSliceWriteRegionSizeBytes.getValue() == 0) {
1149       LLVM_DEBUG(llvm::dbgs()
1150                  << "Failed to get slice write region size at loopDepth: " << i
1151                  << "\n");
1152       continue;
1153     }
1154     int64_t sliceWriteRegionSizeBytes =
1155         maybeSliceWriteRegionSizeBytes.getValue();
1156
1157     // If we are fusing for reuse, check that write regions remain the same.
1158     // TODO: Write region check should check sizes and offsets in
1159     // each dimension, so that we are sure they are covering the same memref
1160     // region. Also, move this out to a isMemRefRegionSuperSet helper function.
1161     if (srcOpInst != srcStoreOpInst &&
1162         sliceWriteRegionSizeBytes != srcWriteRegionSizeBytes)
1163       continue;
1164
1165     double storageReduction = static_cast<double>(srcWriteRegionSizeBytes) /
1166                               static_cast<double>(sliceWriteRegionSizeBytes);
1167
1168     LLVM_DEBUG({
1169       std::stringstream msg;
1170       msg << "  evaluating fusion profitability at depth : " << i << "\n"
1171           << std::fixed << std::setprecision(2)
1172           << "   additional compute fraction: "
1173           << 100.0 * additionalComputeFraction << "%\n"
1174           << "   storage reduction factor: " << storageReduction << "x\n"
1175           << "   fused nest cost: " << fusedLoopNestComputeCost << "\n"
1176           << "   src write region size: " << srcWriteRegionSizeBytes << "\n"
1177           << "   slice write region size: " << sliceWriteRegionSizeBytes
1178           << "\n";
1179       llvm::dbgs() << msg.str();
1180     });
1181
1182     // TODO: This is a placeholder cost model.
1183     // Among all choices that add an acceptable amount of redundant computation
1184     // (as per computeToleranceThreshold), we will simply pick the one that
1185     // reduces the intermediary size the most.
1186     if ((storageReduction > maxStorageReduction) &&
1187         (additionalComputeFraction < computeToleranceThreshold)) {
1188       maxStorageReduction = storageReduction;
1189       bestDstLoopDepth = i;
1190       minFusedLoopNestComputeCost = fusedLoopNestComputeCost;
1191       sliceMemEstimate = sliceWriteRegionSizeBytes;
1192     }
1193   }
1194
1195   // A simple cost model: fuse if it reduces the memory footprint.
1196
1197   if (!bestDstLoopDepth.hasValue()) {
1198     LLVM_DEBUG(
1199         llvm::dbgs()
1200         << "All fusion choices involve more than the threshold amount of "
1201            "redundant computation; NOT fusing.\n");
1202     return false;
1203   }
1204
1205   if (!bestDstLoopDepth.hasValue()) {
1206     LLVM_DEBUG(llvm::dbgs() << "no fusion depth could be evaluated.\n");
1207     return false;
1208   }
1209
1210   // Set dstLoopDepth based on best values from search.
1211   *dstLoopDepth = bestDstLoopDepth.getValue();
1212
1213   LLVM_DEBUG(
1214       llvm::dbgs() << " LoopFusion fusion stats:"
1215                    << "\n  best loop depth: " << bestDstLoopDepth
1216                    << "\n  src loop nest compute cost: " << srcLoopNestCost
1217                    << "\n  dst loop nest compute cost: " << dstLoopNestCost
1218                    << "\n  fused loop nest compute cost: "
1219                    << minFusedLoopNestComputeCost << "\n");
1220
1221   auto dstMemSize = getMemoryFootprintBytes(dstLoopIVs[0]);
1222   auto srcMemSize = getMemoryFootprintBytes(srcLoopIVs[0]);
1223
1224   Optional<double> storageReduction = None;
1225
1226   if (!dstMemSize.hasValue() || !srcMemSize.hasValue()) {
1227     LLVM_DEBUG(llvm::dbgs()
1228                << "  fusion memory benefit cannot be evaluated; NOT fusing.\n");
1229     return false;
1230   }
1231
1232   auto srcMemSizeVal = srcMemSize.getValue();
1233   auto dstMemSizeVal = dstMemSize.getValue();
1234
1235   assert(sliceMemEstimate.hasValue() && "expected value");
1236   auto fusedMem = dstMemSizeVal + sliceMemEstimate.getValue();
1237
1238   LLVM_DEBUG(llvm::dbgs() << "   src mem: " << srcMemSizeVal << "\n"
1239                           << "   dst mem: " << dstMemSizeVal << "\n"
1240                           << "   fused mem: " << fusedMem << "\n"
1241                           << "   slice mem: " << sliceMemEstimate << "\n");
1242
1243   if (static_cast<long>(fusedMem) > srcMemSizeVal + dstMemSizeVal) {
1244     LLVM_DEBUG(llvm::dbgs() << "Fusion is not profitable; NOT fusing.\n");
1245     return false;
1246   }
1247   storageReduction =
1248       100.0 *
1249       (1.0 - fusedMem / (static_cast<double>(srcMemSizeVal) + dstMemSizeVal));
1250
1251   double additionalComputeFraction =
1252       100.0 * (minFusedLoopNestComputeCost /
1253                    (static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
1254                1);
1255   (void)additionalComputeFraction;
1256   LLVM_DEBUG({
1257     std::stringstream msg;
1258     msg << " fusion is most profitable at depth " << *dstLoopDepth << " with "
1259         << std::setprecision(2) << additionalComputeFraction
1260         << "% redundant computation and a ";
1261     msg << (storageReduction.hasValue()
1262                 ? std::to_string(storageReduction.getValue())
1263                 : "<unknown>");
1264     msg << "% storage reduction.\n";
1265     llvm::dbgs() << msg.str();
1266   });
1267
1268   return true;
1269 }
1270
1271 namespace {
1272
1273 // GreedyFusion greedily fuses loop nests which have a producer/consumer or
1274 // input-reuse relationship on a memref, with the goal of improving locality.
1275 //
1276 // The steps of the producer-consumer fusion algorithm are as follows:
1277 //
1278 // *) A worklist is initialized with node ids from the dependence graph.
1279 // *) For each node id in the worklist:
1280 //   *) Pop an AffineForOp of the worklist. This 'dstAffineForOp' will be a
1281 //      candidate destination AffineForOp into which fusion will be attempted.
1282 //   *) Add each LoadOp currently in 'dstAffineForOp' into list 'dstLoadOps'.
1283 //   *) For each LoadOp in 'dstLoadOps' do:
1284 //      *) Look up dependent loop nests which have a single store op to the same
1285 //         memref.
1286 //      *) Check if dependences would be violated by the fusion.
1287 //      *) Get a computation slice of 'srcLoopNest', which adjusts its loop
1288 //         bounds to be functions of 'dstLoopNest' IVs and symbols.
1289 //      *) Fuse the 'srcLoopNest' computation slice into the 'dstLoopNest',
1290 //         at a loop depth determined by the cost model in 'isFusionProfitable'.
1291 //      *) Add the newly fused load/store operations to the state,
1292 //         and also add newly fused load ops to 'dstLoopOps' to be considered
1293 //         as fusion dst load ops in another iteration.
1294 //      *) Remove old src loop nest and its associated state.
1295 //
1296 // The steps of the input-reuse fusion algorithm are as follows:
1297 //
1298 // *) Initialize 'worklist' with node ids from the dependence graph.
1299 // *) For each 'dstNode' in the worklist:
1300 //   *) Find a candidate sibling node 'sibNode' to fuse with 'dstNode' which
1301 //      loads from the same memref, but which has no dependence paths to/from.
1302 //   *) Get a computation slice of 'sibLoopNest', which adjusts its loop
1303 //      bounds to be functions of 'dstLoopNest' IVs and symbols.
1304 //   *) Fuse the 'sibLoopNest' computation slice into the 'dstLoopNest',
1305 //      at a loop depth determined by the cost model in 'isFusionProfitable'.
1306 //      This function also checks that the memref write region of 'sibLoopNest',
1307 //      is preserved in the fused loop nest.
1308 //   *) Update graph state to reflect the fusion of 'sibNode' into 'dstNode'.
1309 //
1310 // Given a graph where top-level operations are vertices in the set 'V' and
1311 // edges in the set 'E' are dependences between vertices, this algorithm
1312 // takes O(V) time for initialization, and has runtime O(V + E).
1313 //
1314 // This greedy algorithm is not 'maximal' due to the current restriction of
1315 // fusing along single producer consumer edges, but there is a TODO: to fix
1316 // this.
1317 //
1318 // TODO: Experiment with other fusion policies.
1319 struct GreedyFusion {
1320 public:
1321   // The data dependence graph to traverse during fusion.
1322   MemRefDependenceGraph *mdg;
1323   // Worklist of graph nodes visited during the fusion pass.
1324   SmallVector<unsigned, 8> worklist;
1325   // Set of graph nodes which are present on the worklist.
1326   llvm::SmallDenseSet<unsigned, 16> worklistSet;
1327   // Parameter for local buffer size threshold.
1328   unsigned localBufSizeThreshold;
1329   // Parameter for fast memory space.
1330   Optional<unsigned> fastMemorySpace;
1331   // If true, ignore any additional (redundant) computation tolerance threshold
1332   // that would have prevented fusion.
1333   bool maximalFusion;
1334   // The amount of additional computation that is tolerated while fusing
1335   // pair-wise as a fraction of the total computation.
1336   double computeToleranceThreshold;
1337
1338   using Node = MemRefDependenceGraph::Node;
1339
1340   GreedyFusion(MemRefDependenceGraph *mdg, unsigned localBufSizeThreshold,
1341                Optional<unsigned> fastMemorySpace, bool maximalFusion,
1342                double computeToleranceThreshold)
1343       : mdg(mdg), localBufSizeThreshold(localBufSizeThreshold),
1344         fastMemorySpace(fastMemorySpace), maximalFusion(maximalFusion),
1345         computeToleranceThreshold(computeToleranceThreshold) {}
1346
1347   // Initializes 'worklist' with nodes from 'mdg'
1348   void init() {
1349     // TODO: Add a priority queue for prioritizing nodes by different
1350     // metrics (e.g. arithmetic intensity/flops-to-bytes ratio).
1351     worklist.clear();
1352     worklistSet.clear();
1353     for (auto &idAndNode : mdg->nodes) {
1354       const Node &node = idAndNode.second;
1355       worklist.push_back(node.id);
1356       worklistSet.insert(node.id);
1357     }
1358   }
1359
1360   // Run the GreedyFusion pass.
1361   // *) First pass through the nodes fuses single-use producer nodes into their
1362   //    unique consumer.
1363   // *) Second pass fuses sibling nodes which share no dependence edges.
1364   // *) Third pass fuses any remaining producer nodes into their users.
1365   void run() {
1366     // TODO: Run this repeatedly until a fixed-point is reached.
1367     fuseProducerConsumerNodes(/*maxSrcUserCount=*/1);
1368     fuseSiblingNodes();
1369     fuseProducerConsumerNodes(
1370         /*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
1371     eraseUnusedMemRefAllocations();
1372   }
1373
1374   void fuseProducerConsumerNodes(unsigned maxSrcUserCount) {
1375     init();
1376     while (!worklist.empty()) {
1377       unsigned dstId = worklist.back();
1378       worklist.pop_back();
1379       worklistSet.erase(dstId);
1380
1381       // Skip if this node was removed (fused into another node).
1382       if (mdg->nodes.count(dstId) == 0)
1383         continue;
1384       // Get 'dstNode' into which to attempt fusion.
1385       auto *dstNode = mdg->getNode(dstId);
1386       // Skip if 'dstNode' is not a loop nest.
1387       if (!isa<AffineForOp>(dstNode->op))
1388         continue;
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
1395       SmallVector<Operation *, 4> loads = dstNode->loads;
1396       SmallVector<Operation *, 4> dstLoadOpInsts;
1397       DenseSet<Value> visitedMemrefs;
1398       while (!loads.empty()) {
1399         // Get memref of load on top of the stack.
1400         auto memref = cast<AffineReadOpInterface>(loads.back()).getMemRef();
1401         if (visitedMemrefs.count(memref) > 0)
1402           continue;
1403         visitedMemrefs.insert(memref);
1404         // Move all loads in 'loads' accessing 'memref' to 'dstLoadOpInsts'.
1405         moveLoadsAccessingMemrefTo(memref, &loads, &dstLoadOpInsts);
1406         // Skip if no input edges along which to fuse.
1407         if (mdg->inEdges.count(dstId) == 0)
1408           continue;
1409         // Iterate through in-edges for 'dstId' and src node id for any
1410         // edges on 'memref'.
1411         SmallVector<unsigned, 2> srcNodeIds;
1412         for (auto &srcEdge : mdg->inEdges[dstId]) {
1413           // Skip 'srcEdge' if not for 'memref'.
1414           if (srcEdge.value != memref)
1415             continue;
1416           srcNodeIds.push_back(srcEdge.id);
1417         }
1418         for (unsigned srcId : srcNodeIds) {
1419           // Skip if this node was removed (fused into another node).
1420           if (mdg->nodes.count(srcId) == 0)
1421             continue;
1422           // Get 'srcNode' from which to attempt fusion into 'dstNode'.
1423           auto *srcNode = mdg->getNode(srcId);
1424           // Skip if 'srcNode' is not a loop nest.
1425           if (!isa<AffineForOp>(srcNode->op))
1426             continue;
1427           // Skip if 'srcNode' has more than one live-out store to a
1428           // function-local memref.
1429           // TODO: Support more generic multi-output src loop nests
1430           // fusion.
1431           auto srcStoreOp = mdg->getUniqueOutgoingStore(srcNode);
1432           if (!srcStoreOp) {
1433             // Get the src store op at the deepest loop depth.
1434             // We will use 'LoopFusionUtils::canFuseLoops' to check fusion
1435             // feasibility for loops with multiple stores.
1436             unsigned maxLoopDepth = 0;
1437             for (auto *op : srcNode->stores) {
1438               auto storeOp = cast<AffineWriteOpInterface>(op);
1439               if (storeOp.getMemRef() != memref) {
1440                 srcStoreOp = nullptr;
1441                 break;
1442               }
1443               unsigned loopDepth = getNestingDepth(storeOp);
1444               if (loopDepth > maxLoopDepth) {
1445                 maxLoopDepth = loopDepth;
1446                 srcStoreOp = storeOp;
1447               }
1448             }
1449             if (!srcStoreOp)
1450               continue;
1451           }
1452
1453           // Unique outgoing store found must write to 'memref' since 'memref'
1454           // is the one that established the producer-consumer relationship
1455           // between 'srcNode' and 'dstNode'.
1456           assert(srcStoreOp.getMemRef() == memref &&
1457                  "Found store to unexpected memref");
1458
1459           // Skip if 'srcNode' writes to any live in or escaping memrefs,
1460           // and cannot be fused.
1461           bool writesToLiveInOrOut =
1462               mdg->writesToLiveInOrEscapingMemrefs(srcNode->id);
1463           if (writesToLiveInOrOut &&
1464               !canFuseSrcWhichWritesToLiveOut(srcId, dstId, srcStoreOp, mdg))
1465             continue;
1466
1467           // Don't create a private memref if 'writesToLiveInOrOut'.
1468           bool createPrivateMemref = !writesToLiveInOrOut;
1469           // Don't create a private memref if 'srcNode' has in edges on
1470           // 'memref', or if 'dstNode' has out edges on 'memref'.
1471           if (mdg->getIncomingMemRefAccesses(srcNode->id, memref) > 0 ||
1472               mdg->getOutEdgeCount(dstNode->id, memref) > 0) {
1473             createPrivateMemref = false;
1474           }
1475
1476           // Skip if 'srcNode' out edge count on 'memref' > 'maxSrcUserCount'.
1477           if (mdg->getOutEdgeCount(srcNode->id, memref) > maxSrcUserCount)
1478             continue;
1479
1480           // Compute an operation list insertion point for the fused loop
1481           // nest which preserves dependences.
1482           Operation *insertPointInst =
1483               mdg->getFusedLoopNestInsertionPoint(srcNode->id, dstNode->id);
1484           if (insertPointInst == nullptr)
1485             continue;
1486
1487           auto srcAffineForOp = cast<AffineForOp>(srcNode->op);
1488           auto dstAffineForOp = cast<AffineForOp>(dstNode->op);
1489
1490           // Compute the innermost common loop depth for dstNode loads/stores.
1491           SmallVector<Operation *, 2> dstMemrefOps;
1492           for (Operation *op : dstNode->loads)
1493             if (cast<AffineReadOpInterface>(op).getMemRef() == memref)
1494               dstMemrefOps.push_back(op);
1495           for (Operation *op : dstNode->stores)
1496             if (cast<AffineWriteOpInterface>(op).getMemRef() == memref)
1497               dstMemrefOps.push_back(op);
1498           unsigned dstLoopDepthTest = getInnermostCommonLoopDepth(dstMemrefOps);
1499
1500           // Check the feasibility of fusing src loop nest into dst loop nest
1501           // at loop depths in range [1, dstLoopDepthTest].
1502           unsigned maxLegalFusionDepth = 0;
1503           SmallVector<ComputationSliceState, 8> depthSliceUnions;
1504           depthSliceUnions.resize(dstLoopDepthTest);
1505           FusionStrategy strategy(FusionStrategy::ProducerConsumer, memref);
1506           for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
1507             FusionResult result = mlir::canFuseLoops(
1508                 srcAffineForOp, dstAffineForOp,
1509                 /*dstLoopDepth=*/i, &depthSliceUnions[i - 1], strategy);
1510
1511             if (result.value == FusionResult::Success)
1512               maxLegalFusionDepth = i;
1513           }
1514
1515           // Skip if fusion is not feasible at any loop depths.
1516           if (maxLegalFusionDepth == 0)
1517             continue;
1518
1519           // Check if fusion would be profitable. We skip profitability analysis
1520           // for maximal fusion since we already know the maximal legal depth to
1521           // fuse.
1522           unsigned bestDstLoopDepth = maxLegalFusionDepth;
1523           if (!maximalFusion &&
1524               !isFusionProfitable(srcStoreOp, srcStoreOp, dstLoadOpInsts,
1525                                   depthSliceUnions, maxLegalFusionDepth,
1526                                   &bestDstLoopDepth, computeToleranceThreshold))
1527             continue;
1528
1529           assert(bestDstLoopDepth > 0 && "Unexpected loop fusion depth");
1530           assert(!depthSliceUnions[bestDstLoopDepth - 1].isEmpty() &&
1531                  "Missing slice union for depth");
1532
1533           // Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'.
1534           fuseLoops(srcAffineForOp, dstAffineForOp,
1535                     depthSliceUnions[bestDstLoopDepth - 1]);
1536
1537           LLVM_DEBUG(llvm::dbgs()
1538                      << "Fused src loop " << srcId << " into dst loop " << dstId
1539                      << " at depth " << bestDstLoopDepth << ":\n"
1540                      << dstAffineForOp << "\n");
1541
1542           // Move 'dstAffineForOp' before 'insertPointInst' if needed.
1543           if (insertPointInst != dstAffineForOp.getOperation())
1544             dstAffineForOp->moveBefore(insertPointInst);
1545
1546           // Update edges between 'srcNode' and 'dstNode'.
1547           mdg->updateEdges(srcNode->id, dstNode->id, memref,
1548                            createPrivateMemref);
1549
1550           // Collect slice loop stats.
1551           LoopNestStateCollector dstForCollector;
1552           dstForCollector.collect(dstAffineForOp);
1553           if (createPrivateMemref) {
1554             // Create private memref for 'memref' in 'dstAffineForOp'.
1555             SmallVector<Operation *, 4> storesForMemref;
1556             for (auto *storeOpInst : dstForCollector.storeOpInsts) {
1557               if (cast<AffineWriteOpInterface>(storeOpInst).getMemRef() ==
1558                   memref)
1559                 storesForMemref.push_back(storeOpInst);
1560             }
1561             // TODO: Use union of memref write regions to compute
1562             // private memref footprint.
1563             auto newMemRef = createPrivateMemRef(
1564                 dstAffineForOp, storesForMemref[0], bestDstLoopDepth,
1565                 fastMemorySpace, localBufSizeThreshold);
1566             visitedMemrefs.insert(newMemRef);
1567             // Create new node in dependence graph for 'newMemRef' alloc op.
1568             unsigned newMemRefNodeId = mdg->addNode(newMemRef.getDefiningOp());
1569             // Add edge from 'newMemRef' node to dstNode.
1570             mdg->addEdge(newMemRefNodeId, dstId, newMemRef);
1571           }
1572
1573           // Collect dst loop stats after memref privatization transformation.
1574           LoopNestStateCollector dstLoopCollector;
1575           dstLoopCollector.collect(dstAffineForOp.getOperation());
1576
1577           // Add new load ops to current Node load op list 'loads' to continue
1578           // fusing based on new operands.
1579           for (auto *loadOpInst : dstLoopCollector.loadOpInsts) {
1580             // NOTE: Change 'loads' to a hash set in case efficiency is an
1581             // issue. We still use a vector since it's expected to be small.
1582             if (!llvm::is_contained(loads, loadOpInst))
1583               loads.push_back(loadOpInst);
1584           }
1585           // Clear visited memrefs after fusion so that previously visited src
1586           // nodes are considered for fusion again in the context of the new
1587           // fused node.
1588           // TODO: This shouldn't be necessary if we visited candidates in the
1589           // dependence graph in post-order or once we fully support multi-store
1590           // producers. Currently, in a multi-store producer scenario such as
1591           // A->B, A->C, B->C, we fail to fuse A+B due to the multiple outgoing
1592           // edges. However, after fusing B+C, A has a single outgoing edge and
1593           // can be fused if we revisit it in the context of the new fused B+C
1594           // node.
1595           visitedMemrefs.clear();
1596
1597           // Clear and add back loads and stores.
1598           mdg->clearNodeLoadAndStores(dstNode->id);
1599           mdg->addToNode(dstId, dstLoopCollector.loadOpInsts,
1600                          dstLoopCollector.storeOpInsts);
1601           // Remove old src loop nest if it no longer has outgoing dependence
1602           // edges, and if it does not write to a memref which escapes the
1603           // function. If 'writesToLiveInOrOut' is true, then 'srcNode' has been
1604           // fused into 'dstNode' and write region of 'dstNode' covers the write
1605           // region of 'srcNode', and 'srcNode' has no other users so it is safe
1606           // to remove.
1607           if (writesToLiveInOrOut || mdg->canRemoveNode(srcNode->id)) {
1608             mdg->removeNode(srcNode->id);
1609             srcNode->op->erase();
1610           } else {
1611             // Add remaining users of 'oldMemRef' back on the worklist (if not
1612             // already there), as its replacement with a local/private memref
1613             // has reduced dependences on 'oldMemRef' which may have created new
1614             // fusion opportunities.
1615             if (mdg->outEdges.count(srcNode->id) > 0) {
1616               SmallVector<MemRefDependenceGraph::Edge, 2> oldOutEdges =
1617                   mdg->outEdges[srcNode->id];
1618               for (auto &outEdge : oldOutEdges) {
1619                 if (outEdge.value == memref &&
1620                     worklistSet.count(outEdge.id) == 0) {
1621                   worklist.push_back(outEdge.id);
1622                   worklistSet.insert(outEdge.id);
1623                 }
1624               }
1625             }
1626           }
1627         }
1628       }
1629     }
1630   }
1631
1632   // Visits each node in the graph, and for each node, attempts to fuse it with
1633   // its sibling nodes (nodes which share a parent, but no dependence edges).
1634   void fuseSiblingNodes() {
1635     init();
1636     while (!worklist.empty()) {
1637       unsigned dstId = worklist.back();
1638       worklist.pop_back();
1639       worklistSet.erase(dstId);
1640
1641       // Skip if this node was removed (fused into another node).
1642       if (mdg->nodes.count(dstId) == 0)
1643         continue;
1644       // Get 'dstNode' into which to attempt fusion.
1645       auto *dstNode = mdg->getNode(dstId);
1646       // Skip if 'dstNode' is not a loop nest.
1647       if (!isa<AffineForOp>(dstNode->op))
1648         continue;
1649       // Attempt to fuse 'dstNode' with its sibling nodes in the graph.
1650       fuseWithSiblingNodes(dstNode);
1651     }
1652   }
1653
1654   // Attempt to fuse 'dstNode' with sibling nodes in the graph.
1655   void fuseWithSiblingNodes(Node *dstNode) {
1656     DenseSet<unsigned> visitedSibNodeIds;
1657     std::pair<unsigned, Value> idAndMemref;
1658     auto dstAffineForOp = cast<AffineForOp>(dstNode->op);
1659
1660     while (findSiblingNodeToFuse(dstNode, &visitedSibNodeIds, &idAndMemref)) {
1661       unsigned sibId = idAndMemref.first;
1662       Value memref = idAndMemref.second;
1663       // TODO: Check that 'sibStoreOpInst' post-dominates all other
1664       // stores to the same memref in 'sibNode' loop nest.
1665       auto *sibNode = mdg->getNode(sibId);
1666       // Compute an operation list insertion point for the fused loop
1667       // nest which preserves dependences.
1668       assert(sibNode->op->getBlock() == dstNode->op->getBlock());
1669       Operation *insertPointInst =
1670           sibNode->op->isBeforeInBlock(dstNode->op)
1671               ? mdg->getFusedLoopNestInsertionPoint(sibNode->id, dstNode->id)
1672               : mdg->getFusedLoopNestInsertionPoint(dstNode->id, sibNode->id);
1673       if (insertPointInst == nullptr)
1674         continue;
1675
1676       // Check if fusion would be profitable and at what depth.
1677
1678       // Get unique 'sibNode' load op to 'memref'.
1679       SmallVector<Operation *, 2> sibLoadOpInsts;
1680       sibNode->getLoadOpsForMemref(memref, &sibLoadOpInsts);
1681       // Currently findSiblingNodeToFuse searches for siblings with one load.
1682       assert(sibLoadOpInsts.size() == 1);
1683       Operation *sibLoadOpInst = sibLoadOpInsts[0];
1684       assert(!sibNode->stores.empty());
1685       // TODO: Choose the store which postdominates all other stores.
1686       auto *sibStoreOpInst = sibNode->stores.back();
1687
1688       // Gather 'dstNode' load ops to 'memref'.
1689       SmallVector<Operation *, 2> dstLoadOpInsts;
1690       dstNode->getLoadOpsForMemref(memref, &dstLoadOpInsts);
1691
1692       SmallVector<AffineForOp, 4> dstLoopIVs;
1693       getLoopIVs(*dstLoadOpInsts[0], &dstLoopIVs);
1694       unsigned dstLoopDepthTest = dstLoopIVs.size();
1695       auto sibAffineForOp = cast<AffineForOp>(sibNode->op);
1696
1697       // Compute loop depth and slice union for fusion.
1698       SmallVector<ComputationSliceState, 8> depthSliceUnions;
1699       depthSliceUnions.resize(dstLoopDepthTest);
1700       unsigned maxLegalFusionDepth = 0;
1701       FusionStrategy strategy(FusionStrategy::Sibling, memref);
1702       for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
1703         FusionResult result = mlir::canFuseLoops(
1704             sibAffineForOp, dstAffineForOp,
1705             /*dstLoopDepth=*/i, &depthSliceUnions[i - 1], strategy);
1706
1707         if (result.value == FusionResult::Success)
1708           maxLegalFusionDepth = i;
1709       }
1710
1711       // Skip if fusion is not feasible at any loop depths.
1712       if (maxLegalFusionDepth == 0)
1713         continue;
1714
1715       unsigned bestDstLoopDepth = dstLoopDepthTest;
1716       if (!maximalFusion) {
1717         // Check if fusion would be profitable.
1718         if (!isFusionProfitable(sibLoadOpInst, sibStoreOpInst, dstLoadOpInsts,
1719                                 depthSliceUnions, maxLegalFusionDepth,
1720                                 &bestDstLoopDepth, computeToleranceThreshold))
1721           continue;
1722       }
1723
1724       assert(bestDstLoopDepth > 0 && "Unexpected loop fusion depth");
1725       assert(!depthSliceUnions[bestDstLoopDepth - 1].isEmpty() &&
1726              "Fusion depth has no computed slice union");
1727
1728       // Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'.
1729       mlir::fuseLoops(sibAffineForOp, dstAffineForOp,
1730                       depthSliceUnions[bestDstLoopDepth - 1]);
1731
1732       auto dstForInst = cast<AffineForOp>(dstNode->op);
1733       // Update operation position of fused loop nest (if needed).
1734       if (insertPointInst != dstForInst.getOperation()) {
1735         dstForInst->moveBefore(insertPointInst);
1736       }
1737       // Update data dependence graph state post fusion.
1738       updateStateAfterSiblingFusion(sibNode, dstNode);
1739     }
1740   }
1741
1742   // Searches function argument uses and the graph from 'dstNode' looking for a
1743   // fusion candidate sibling node which shares no dependences with 'dstNode'
1744   // but which loads from the same memref. Returns true and sets
1745   // 'idAndMemrefToFuse' on success. Returns false otherwise.
1746   bool findSiblingNodeToFuse(Node *dstNode,
1747                              DenseSet<unsigned> *visitedSibNodeIds,
1748                              std::pair<unsigned, Value> *idAndMemrefToFuse) {
1749     // Returns true if 'sibNode' can be fused with 'dstNode' for input reuse
1750     // on 'memref'.
1751     auto canFuseWithSibNode = [&](Node *sibNode, Value memref) {
1752       // Skip if 'outEdge' is not a read-after-write dependence.
1753       // TODO: Remove restrict to single load op restriction.
1754       if (sibNode->getLoadOpCount(memref) != 1)
1755         return false;
1756       // Skip if there exists a path of dependent edges between
1757       // 'sibNode' and 'dstNode'.
1758       if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
1759           mdg->hasDependencePath(dstNode->id, sibNode->id))
1760         return false;
1761       // Skip sib node if it loads to (and stores from) the same memref on
1762       // which it also has an input dependence edge.
1763       DenseSet<Value> loadAndStoreMemrefSet;
1764       sibNode->getLoadAndStoreMemrefSet(&loadAndStoreMemrefSet);
1765       if (llvm::any_of(loadAndStoreMemrefSet, [=](Value memref) {
1766             return mdg->getIncomingMemRefAccesses(sibNode->id, memref) > 0;
1767           }))
1768         return false;
1769
1770       // Check that all stores are to the same memref.
1771       DenseSet<Value> storeMemrefs;
1772       for (auto *storeOpInst : sibNode->stores) {
1773         storeMemrefs.insert(
1774             cast<AffineWriteOpInterface>(storeOpInst).getMemRef());
1775       }
1776       if (storeMemrefs.size() != 1)
1777         return false;
1778
1779       // Skip if a memref value in one node is used by a non-affine memref
1780       // access that lies between 'dstNode' and 'sibNode'.
1781       if (hasNonAffineUsersOnThePath(dstNode->id, sibNode->id, mdg) ||
1782           hasNonAffineUsersOnThePath(sibNode->id, dstNode->id, mdg))
1783         return false;
1784       return true;
1785     };
1786
1787     // Search for siblings which load the same memref function argument.
1788     auto fn = dstNode->op->getParentOfType<FuncOp>();
1789     for (unsigned i = 0, e = fn.getNumArguments(); i != e; ++i) {
1790       for (auto *user : fn.getArgument(i).getUsers()) {
1791         if (auto loadOp = dyn_cast<AffineReadOpInterface>(user)) {
1792           // Gather loops surrounding 'use'.
1793           SmallVector<AffineForOp, 4> loops;
1794           getLoopIVs(*user, &loops);
1795           // Skip 'use' if it is not within a loop nest.
1796           if (loops.empty())
1797             continue;
1798           Node *sibNode = mdg->getForOpNode(loops[0]);
1799           assert(sibNode != nullptr);
1800           // Skip 'use' if it not a sibling to 'dstNode'.
1801           if (sibNode->id == dstNode->id)
1802             continue;
1803           // Skip 'use' if it has been visited.
1804           if (visitedSibNodeIds->count(sibNode->id) > 0)
1805             continue;
1806           // Skip 'use' if it does not load from the same memref as 'dstNode'.
1807           auto memref = loadOp.getMemRef();
1808           if (dstNode->getLoadOpCount(memref) == 0)
1809             continue;
1810           // Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
1811           if (canFuseWithSibNode(sibNode, memref)) {
1812             visitedSibNodeIds->insert(sibNode->id);
1813             idAndMemrefToFuse->first = sibNode->id;
1814             idAndMemrefToFuse->second = memref;
1815             return true;
1816           }
1817         }
1818       }
1819     }
1820
1821     // Search for siblings by following edges through an intermediate src node.
1822     // Collect candidate 'dstNode' input edges in 'inEdges'.
1823     SmallVector<MemRefDependenceGraph::Edge, 2> inEdges;
1824     mdg->forEachMemRefInputEdge(
1825         dstNode->id, [&](MemRefDependenceGraph::Edge inEdge) {
1826           // Add 'inEdge' if it is a read-after-write dependence.
1827           if (dstNode->getLoadOpCount(inEdge.value) > 0 &&
1828               mdg->getNode(inEdge.id)->getStoreOpCount(inEdge.value) > 0)
1829             inEdges.push_back(inEdge);
1830         });
1831
1832     // Search for sibling nodes to fuse by visiting output edges from each input
1833     // edge in 'inEdges'.
1834     for (auto &inEdge : inEdges) {
1835       // Collect candidate output edges from each node 'inEdge.id' in 'inEdges'.
1836       SmallVector<MemRefDependenceGraph::Edge, 2> outEdges;
1837       mdg->forEachMemRefOutputEdge(
1838           inEdge.id, [&](MemRefDependenceGraph::Edge outEdge) {
1839             unsigned sibNodeId = outEdge.id;
1840             if (visitedSibNodeIds->count(sibNodeId) > 0)
1841               return;
1842             // Skip output edge if not a sibling using the same memref.
1843             if (outEdge.id == dstNode->id || outEdge.value != inEdge.value)
1844               return;
1845             auto *sibNode = mdg->getNode(sibNodeId);
1846             if (!isa<AffineForOp>(sibNode->op))
1847               return;
1848             // Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
1849             if (canFuseWithSibNode(sibNode, outEdge.value)) {
1850               // Add candidate 'outEdge' to sibling node.
1851               outEdges.push_back(outEdge);
1852             }
1853           });
1854
1855       // Add first candidate if any were returned.
1856       if (!outEdges.empty()) {
1857         visitedSibNodeIds->insert(outEdges[0].id);
1858         idAndMemrefToFuse->first = outEdges[0].id;
1859         idAndMemrefToFuse->second = outEdges[0].value;
1860         return true;
1861       }
1862     }
1863     return false;
1864   }
1865
1866   /// Update data dependence graph state to reflect sibling fusion of 'sibNode'
1867   /// into 'dstNode'.
1868   void updateStateAfterSiblingFusion(Node *sibNode, Node *dstNode) {
1869     // Update 'sibNode' and 'dstNode' input/output edges to reflect fusion.
1870     mdg->updateEdges(sibNode->id, dstNode->id);
1871
1872     // Collect dst loop stats after memref privatization transformation.
1873     auto dstForInst = cast<AffineForOp>(dstNode->op);
1874     LoopNestStateCollector dstLoopCollector;
1875     dstLoopCollector.collect(dstForInst.getOperation());
1876     // Clear and add back loads and stores
1877     mdg->clearNodeLoadAndStores(dstNode->id);
1878     mdg->addToNode(dstNode->id, dstLoopCollector.loadOpInsts,
1879                    dstLoopCollector.storeOpInsts);
1880     // Remove old sibling loop nest if it no longer has outgoing dependence
1881     // edges, and it does not write to a memref which escapes the
1882     // function.
1883     if (mdg->getOutEdgeCount(sibNode->id) == 0) {
1884       mdg->removeNode(sibNode->id);
1885       sibNode->op->erase();
1886     }
1887   }
1888
1889   // Clean up any allocs with no users.
1890   void eraseUnusedMemRefAllocations() {
1891     for (auto &pair : mdg->memrefEdgeCount) {
1892       if (pair.second > 0)
1893         continue;
1894       auto memref = pair.first;
1895       // Skip if there exist other uses (return operation or function calls).
1896       if (!memref.use_empty())
1897         continue;
1898       // Use list expected to match the dep graph info.
1899       auto *op = memref.getDefiningOp();
1900       if (isa_and_nonnull<AllocOp>(op))
1901         op->erase();
1902     }
1903   }
1904 };
1905
1906 } // end anonymous namespace
1907
1908 void LoopFusion::runOnFunction() {
1909   MemRefDependenceGraph g;
1910   if (!g.init(getFunction()))
1911     return;
1912
1913   Optional<unsigned> fastMemorySpaceOpt;
1914   if (fastMemorySpace.hasValue())
1915     fastMemorySpaceOpt = fastMemorySpace;
1916   unsigned localBufSizeThresholdBytes = localBufSizeThreshold * 1024;
1917   GreedyFusion fusion(&g, localBufSizeThresholdBytes, fastMemorySpaceOpt,
1918                       maximalFusion, computeToleranceThreshold);
1919   fusion.run();
1920 }