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