blob: 47ee502b41fbbd347e6487016ba74400bac9a084 [file] [log] [blame]
MLIR Teamf28e4df2018-11-01 14:26:001//===- LoopFusion.cpp - Code to perform loop fusion -----------------------===//
2//
Mehdi Amini30857102020-01-26 03:58:303// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
Mehdi Amini56222a02019-12-23 17:35:364// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
MLIR Teamf28e4df2018-11-01 14:26:006//
Mehdi Amini56222a02019-12-23 17:35:367//===----------------------------------------------------------------------===//
MLIR Teamf28e4df2018-11-01 14:26:008//
9// This file implements loop fusion.
10//
11//===----------------------------------------------------------------------===//
12
River Riddle1834ad4a2020-04-07 20:58:1213#include "PassDetail.h"
MLIR Teamf28e4df2018-11-01 14:26:0014#include "mlir/Analysis/AffineAnalysis.h"
Uday Bondhuguladfe07b72019-02-23 00:51:0815#include "mlir/Analysis/AffineStructures.h"
MLIR Teamf28e4df2018-11-01 14:26:0016#include "mlir/Analysis/LoopAnalysis.h"
MLIR Team3b692302018-12-17 17:57:1417#include "mlir/Analysis/Utils.h"
Rob Sudermane7084712020-03-20 21:18:4718#include "mlir/Dialect/Affine/IR/AffineOps.h"
MLIR Teamf28e4df2018-11-01 14:26:0019#include "mlir/IR/AffineExpr.h"
20#include "mlir/IR/AffineMap.h"
21#include "mlir/IR/Builders.h"
Andy Davisa560f2c2019-05-24 17:54:2222#include "mlir/Transforms/LoopFusionUtils.h"
MLIR Teamf28e4df2018-11-01 14:26:0023#include "mlir/Transforms/LoopUtils.h"
24#include "mlir/Transforms/Passes.h"
MLIR Teamc4237ae2019-01-18 16:56:2725#include "mlir/Transforms/Utils.h"
MLIR Teamf28e4df2018-11-01 14:26:0026#include "llvm/ADT/DenseMap.h"
MLIR Team3b692302018-12-17 17:57:1427#include "llvm/ADT/DenseSet.h"
28#include "llvm/ADT/SetVector.h"
MLIR Team4eef7952018-12-21 19:06:2329#include "llvm/Support/CommandLine.h"
MLIR Team38c2fe32019-01-14 19:26:2530#include "llvm/Support/Debug.h"
MLIR Team3b692302018-12-17 17:57:1431#include "llvm/Support/raw_ostream.h"
Uday Bondhugula864d9e02019-01-23 17:16:2432#include <iomanip>
Stella Laurenzo1a2ad062019-05-14 01:10:4833#include <sstream>
Nicolas Vasilache258e8d92019-05-03 18:07:3734#define DEBUG_TYPE "affine-loop-fusion"
MLIR Team38c2fe32019-01-14 19:26:2535
MLIR Team3b692302018-12-17 17:57:1436using llvm::SetVector;
MLIR Teamf28e4df2018-11-01 14:26:0037
38using namespace mlir;
39
River Riddle75c21e12019-01-26 06:14:0440static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
41
Uday Bondhugulace7e59532019-03-08 17:21:5242/// Disables fusion profitability check and fuses if valid. Ignore any
43/// additional (redundant) computation tolerance threshold
44/// that would have prevented fusion.
MLIR Teamc4237ae2019-01-18 16:56:2745static llvm::cl::opt<bool>
Uday Bondhugulaeee85362019-03-02 01:42:1346 clMaximalLoopFusion("fusion-maximal",
River Riddle75c21e12019-01-26 06:14:0447 llvm::cl::desc("Enables maximal loop fusion"),
48 llvm::cl::cat(clOptionsCategory));
Uday Bondhugula864d9e02019-01-23 17:16:2449
50/// A threshold in percent of additional computation allowed when fusing.
51static llvm::cl::opt<double> clFusionAddlComputeTolerance(
Uday Bondhugulaeee85362019-03-02 01:42:1352 "fusion-compute-tolerance",
Uday Bondhugulaa1dad3a2019-02-20 02:17:1953 llvm::cl::desc("Fractional increase in additional "
54 "computation tolerated while fusing"),
River Riddle75c21e12019-01-26 06:14:0455 llvm::cl::cat(clOptionsCategory));
MLIR Teamc4237ae2019-01-18 16:56:2756
Uday Bondhugula8be26272019-02-02 01:06:2257static llvm::cl::opt<unsigned> clFusionFastMemorySpace(
Uday Bondhugulaeee85362019-03-02 01:42:1358 "fusion-fast-mem-space",
Uday Bondhugula8be26272019-02-02 01:06:2259 llvm::cl::desc("Faster memory space number to promote fusion buffers to"),
60 llvm::cl::cat(clOptionsCategory));
61
Uday Bondhugulace7e59532019-03-08 17:21:5262// A local buffer of size less than or equal to this size is automatically
63// promoted to fast memory after producer-consumer fusion.
Uday Bondhugulad4b3ff12019-02-27 00:10:1964static llvm::cl::opt<unsigned long long> clFusionLocalBufThreshold(
Uday Bondhugulaeee85362019-03-02 01:42:1365 "fusion-local-buf-threshold",
Uday Bondhugulad4b3ff12019-02-27 00:10:1966 llvm::cl::desc("Threshold size (KiB) for promoting local buffers to fast "
Uday Bondhugula8be26272019-02-02 01:06:2267 "memory space"),
68 llvm::cl::cat(clOptionsCategory));
69
MLIR Teamf28e4df2018-11-01 14:26:0070namespace {
MLIR Team3b692302018-12-17 17:57:1471/// Loop fusion pass. This pass currently supports a greedy fusion policy,
72/// which fuses loop nests with single-writer/single-reader memref dependences
73/// with the goal of improving locality.
74
75// TODO(andydavis) Support fusion of source loop nests which write to multiple
76// memrefs, where each memref can have multiple users (if profitable).
MLIR Teamf28e4df2018-11-01 14:26:0077// TODO(andydavis) Extend this pass to check for fusion preventing dependences,
78// and add support for more general loop fusion algorithms.
MLIR Team3b692302018-12-17 17:57:1479
River Riddle1834ad4a2020-04-07 20:58:1280struct LoopFusion : public AffineLoopFusionBase<LoopFusion> {
Uday Bondhugulace7e59532019-03-08 17:21:5281 LoopFusion(unsigned fastMemorySpace = 0, uint64_t localBufSizeThreshold = 0,
82 bool maximalFusion = false)
River Riddlec6c53442019-02-27 18:59:2983 : localBufSizeThreshold(localBufSizeThreshold),
Uday Bondhugulace7e59532019-03-08 17:21:5284 fastMemorySpace(fastMemorySpace), maximalFusion(maximalFusion) {}
MLIR Teamf28e4df2018-11-01 14:26:0085
River Riddleed5fe202019-02-28 22:50:4286 void runOnFunction() override;
Uday Bondhugula864d9e02019-01-23 17:16:2487
Uday Bondhugulad4b3ff12019-02-27 00:10:1988 // Any local buffers smaller than this size (in bytes) will be created in
Uday Bondhugula8be26272019-02-02 01:06:2289 // `fastMemorySpace` if provided.
Uday Bondhugulad4b3ff12019-02-27 00:10:1990 uint64_t localBufSizeThreshold;
Uday Bondhugula8be26272019-02-02 01:06:2291 Optional<unsigned> fastMemorySpace = None;
Uday Bondhugulace7e59532019-03-08 17:21:5292 // If true, ignore any additional (redundant) computation tolerance threshold
93 // that would have prevented fusion.
94 bool maximalFusion;
Uday Bondhugula8be26272019-02-02 01:06:2295
Uday Bondhugula864d9e02019-01-23 17:16:2496 // The amount of additional computation that is tolerated while fusing
97 // pair-wise as a fraction of the total computation.
98 constexpr static double kComputeToleranceThreshold = 0.30f;
MLIR Teamf28e4df2018-11-01 14:26:0099};
100
MLIR Teamf28e4df2018-11-01 14:26:00101} // end anonymous namespace
102
River Riddle80aca1e2020-04-07 20:56:16103std::unique_ptr<OperationPass<FuncOp>>
Mehdi Amini926fb682019-08-13 02:12:42104mlir::createLoopFusionPass(unsigned fastMemorySpace,
105 uint64_t localBufSizeThreshold, bool maximalFusion) {
Jacques Pienaar79f53b02019-08-17 18:05:35106 return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold,
107 maximalFusion);
Uday Bondhugulad4b3ff12019-02-27 00:10:19108}
MLIR Teamf28e4df2018-11-01 14:26:00109
River Riddle2666b972019-12-18 18:46:16110// TODO(b/117228571) Replace when this is modeled through side-effects/op traits
111static bool isMemRefDereferencingOp(Operation &op) {
112 if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op) ||
113 isa<AffineDmaStartOp>(op) || isa<AffineDmaWaitOp>(op))
114 return true;
115 return false;
116}
117
MLIR Team3b692302018-12-17 17:57:14118namespace {
MLIR Teamf28e4df2018-11-01 14:26:00119
MLIR Team3b692302018-12-17 17:57:14120// LoopNestStateCollector walks loop nests and collects load and store
Chris Lattner456ad6a2018-12-29 00:05:35121// operations, and whether or not an IfInst was encountered in the loop nest.
River Riddlebf9c3812019-02-05 00:24:44122struct LoopNestStateCollector {
Chris Lattnerd9b5bc82019-03-25 02:53:05123 SmallVector<AffineForOp, 4> forOps;
River Riddle99b87c92019-03-27 21:02:02124 SmallVector<Operation *, 4> loadOpInsts;
125 SmallVector<Operation *, 4> storeOpInsts;
River Riddle75553832019-01-29 05:23:53126 bool hasNonForRegion = false;
MLIR Team3b692302018-12-17 17:57:14127
River Riddle99b87c92019-03-27 21:02:02128 void collect(Operation *opToWalk) {
129 opToWalk->walk([&](Operation *op) {
River Riddled5b60ee82019-05-12 01:59:54130 if (isa<AffineForOp>(op))
River Riddleadca3c22019-05-12 00:57:32131 forOps.push_back(cast<AffineForOp>(op));
River Riddle99b87c92019-03-27 21:02:02132 else if (op->getNumRegions() != 0)
River Riddlebf9c3812019-02-05 00:24:44133 hasNonForRegion = true;
Andy Davis2e1187d2019-07-03 17:35:03134 else if (isa<AffineLoadOp>(op))
River Riddle99b87c92019-03-27 21:02:02135 loadOpInsts.push_back(op);
Andy Davis2e1187d2019-07-03 17:35:03136 else if (isa<AffineStoreOp>(op))
River Riddle99b87c92019-03-27 21:02:02137 storeOpInsts.push_back(op);
River Riddlebf9c3812019-02-05 00:24:44138 });
MLIR Team3b692302018-12-17 17:57:14139 }
140};
141
MLIR Team6892ffb2018-12-20 04:42:55142// MemRefDependenceGraph is a graph data structure where graph nodes are
River Riddle8c443672019-07-09 23:17:55143// top-level operations in a FuncOp which contain load/store ops, and edges
MLIR Team6892ffb2018-12-20 04:42:55144// are memref dependences between the nodes.
Kazuaki Ishizaki8bfedb32019-10-20 07:11:03145// TODO(andydavis) Add a more flexible dependence graph representation.
MLIR Team6892ffb2018-12-20 04:42:55146// TODO(andydavis) Add a depth parameter to dependence graph construction.
147struct MemRefDependenceGraph {
148public:
149 // Node represents a node in the graph. A Node is either an entire loop nest
150 // rooted at the top level which contains loads/stores, or a top level
151 // load/store.
152 struct Node {
153 // The unique identifier of this node in the graph.
154 unsigned id;
Amit Sabne70a416d2019-04-09 16:17:40155 // The top-level statement which is (or contains) a load/store.
River Riddle99b87c92019-03-27 21:02:02156 Operation *op;
Chris Lattner5187cfc2018-12-28 05:21:41157 // List of load operations.
River Riddle99b87c92019-03-27 21:02:02158 SmallVector<Operation *, 4> loads;
Chris Lattner456ad6a2018-12-29 00:05:35159 // List of store op insts.
River Riddle99b87c92019-03-27 21:02:02160 SmallVector<Operation *, 4> stores;
161 Node(unsigned id, Operation *op) : id(id), op(op) {}
MLIR Team6892ffb2018-12-20 04:42:55162
163 // Returns the load op count for 'memref'.
River Riddlee62a6952019-12-23 22:45:01164 unsigned getLoadOpCount(Value memref) {
MLIR Team6892ffb2018-12-20 04:42:55165 unsigned loadOpCount = 0;
Chris Lattner456ad6a2018-12-29 00:05:35166 for (auto *loadOpInst : loads) {
Andy Davis2e1187d2019-07-03 17:35:03167 if (memref == cast<AffineLoadOp>(loadOpInst).getMemRef())
MLIR Team6892ffb2018-12-20 04:42:55168 ++loadOpCount;
169 }
170 return loadOpCount;
171 }
172
173 // Returns the store op count for 'memref'.
River Riddlee62a6952019-12-23 22:45:01174 unsigned getStoreOpCount(Value memref) {
MLIR Team6892ffb2018-12-20 04:42:55175 unsigned storeOpCount = 0;
Chris Lattner456ad6a2018-12-29 00:05:35176 for (auto *storeOpInst : stores) {
Andy Davis2e1187d2019-07-03 17:35:03177 if (memref == cast<AffineStoreOp>(storeOpInst).getMemRef())
MLIR Team6892ffb2018-12-20 04:42:55178 ++storeOpCount;
179 }
180 return storeOpCount;
181 }
MLIR Team58aa3832019-02-16 01:12:19182
MLIR Teamd038e342019-03-01 19:50:25183 // Returns all store ops in 'storeOps' which access 'memref'.
River Riddlee62a6952019-12-23 22:45:01184 void getStoreOpsForMemref(Value memref,
River Riddle99b87c92019-03-27 21:02:02185 SmallVectorImpl<Operation *> *storeOps) {
MLIR Team58aa3832019-02-16 01:12:19186 for (auto *storeOpInst : stores) {
Andy Davis2e1187d2019-07-03 17:35:03187 if (memref == cast<AffineStoreOp>(storeOpInst).getMemRef())
MLIR Team58aa3832019-02-16 01:12:19188 storeOps->push_back(storeOpInst);
189 }
190 }
MLIR Teamd038e342019-03-01 19:50:25191
192 // Returns all load ops in 'loadOps' which access 'memref'.
River Riddlee62a6952019-12-23 22:45:01193 void getLoadOpsForMemref(Value memref,
River Riddle99b87c92019-03-27 21:02:02194 SmallVectorImpl<Operation *> *loadOps) {
MLIR Teamd038e342019-03-01 19:50:25195 for (auto *loadOpInst : loads) {
Andy Davis2e1187d2019-07-03 17:35:03196 if (memref == cast<AffineLoadOp>(loadOpInst).getMemRef())
MLIR Teamd038e342019-03-01 19:50:25197 loadOps->push_back(loadOpInst);
198 }
199 }
200
201 // Returns all memrefs in 'loadAndStoreMemrefSet' for which this node
202 // has at least one load and store operation.
River Riddlee62a6952019-12-23 22:45:01203 void getLoadAndStoreMemrefSet(DenseSet<Value> *loadAndStoreMemrefSet) {
204 llvm::SmallDenseSet<Value, 2> loadMemrefs;
MLIR Teamd038e342019-03-01 19:50:25205 for (auto *loadOpInst : loads) {
Andy Davis2e1187d2019-07-03 17:35:03206 loadMemrefs.insert(cast<AffineLoadOp>(loadOpInst).getMemRef());
MLIR Teamd038e342019-03-01 19:50:25207 }
208 for (auto *storeOpInst : stores) {
River Riddle35807bc2019-12-23 05:59:55209 auto memref = cast<AffineStoreOp>(storeOpInst).getMemRef();
MLIR Teamd038e342019-03-01 19:50:25210 if (loadMemrefs.count(memref) > 0)
211 loadAndStoreMemrefSet->insert(memref);
212 }
213 }
MLIR Team6892ffb2018-12-20 04:42:55214 };
215
Kazuaki Ishizaki8bfedb32019-10-20 07:11:03216 // Edge represents a data dependence between nodes in the graph.
MLIR Team6892ffb2018-12-20 04:42:55217 struct Edge {
218 // The id of the node at the other end of the edge.
MLIR Team1e851912019-01-31 00:01:46219 // If this edge is stored in Edge = Node.inEdges[i], then
220 // 'Node.inEdges[i].id' is the identifier of the source node of the edge.
221 // If this edge is stored in Edge = Node.outEdges[i], then
222 // 'Node.outEdges[i].id' is the identifier of the dest node of the edge.
MLIR Team6892ffb2018-12-20 04:42:55223 unsigned id;
MLIR Teama0f3db402019-01-29 17:36:41224 // The SSA value on which this edge represents a dependence.
225 // If the value is a memref, then the dependence is between graph nodes
226 // which contain accesses to the same memref 'value'. If the value is a
227 // non-memref value, then the dependence is between a graph node which
228 // defines an SSA value and another graph node which uses the SSA value
River Riddle99b87c92019-03-27 21:02:02229 // (e.g. a constant operation defining a value which is used inside a loop
MLIR Teama0f3db402019-01-29 17:36:41230 // nest).
River Riddlee62a6952019-12-23 22:45:01231 Value value;
MLIR Team6892ffb2018-12-20 04:42:55232 };
233
234 // Map from node id to Node.
235 DenseMap<unsigned, Node> nodes;
236 // Map from node id to list of input edges.
237 DenseMap<unsigned, SmallVector<Edge, 2>> inEdges;
238 // Map from node id to list of output edges.
239 DenseMap<unsigned, SmallVector<Edge, 2>> outEdges;
MLIR Teamc4237ae2019-01-18 16:56:27240 // Map from memref to a count on the dependence edges associated with that
241 // memref.
River Riddlee62a6952019-12-23 22:45:01242 DenseMap<Value, unsigned> memrefEdgeCount;
MLIR Teama0f3db402019-01-29 17:36:41243 // The next unique identifier to use for newly created graph nodes.
244 unsigned nextNodeId = 0;
MLIR Team6892ffb2018-12-20 04:42:55245
246 MemRefDependenceGraph() {}
247
248 // Initializes the dependence graph based on operations in 'f'.
249 // Returns true on success, false otherwise.
River Riddle8c443672019-07-09 23:17:55250 bool init(FuncOp f);
MLIR Team6892ffb2018-12-20 04:42:55251
252 // Returns the graph node for 'id'.
253 Node *getNode(unsigned id) {
254 auto it = nodes.find(id);
255 assert(it != nodes.end());
256 return &it->second;
257 }
258
MLIR Team9d30b362019-03-29 15:06:25259 // Returns the graph node for 'forOp'.
260 Node *getForOpNode(AffineForOp forOp) {
261 for (auto &idAndNode : nodes)
262 if (idAndNode.second.op == forOp.getOperation())
263 return &idAndNode.second;
264 return nullptr;
265 }
266
River Riddle99b87c92019-03-27 21:02:02267 // Adds a node with 'op' to the graph and returns its unique identifier.
268 unsigned addNode(Operation *op) {
269 Node node(nextNodeId++, op);
MLIR Teama0f3db402019-01-29 17:36:41270 nodes.insert({node.id, node});
271 return node.id;
272 }
273
MLIR Teamc4237ae2019-01-18 16:56:27274 // Remove node 'id' (and its associated edges) from graph.
275 void removeNode(unsigned id) {
276 // Remove each edge in 'inEdges[id]'.
277 if (inEdges.count(id) > 0) {
278 SmallVector<Edge, 2> oldInEdges = inEdges[id];
279 for (auto &inEdge : oldInEdges) {
MLIR Teama0f3db402019-01-29 17:36:41280 removeEdge(inEdge.id, id, inEdge.value);
MLIR Teamc4237ae2019-01-18 16:56:27281 }
282 }
283 // Remove each edge in 'outEdges[id]'.
284 if (outEdges.count(id) > 0) {
285 SmallVector<Edge, 2> oldOutEdges = outEdges[id];
286 for (auto &outEdge : oldOutEdges) {
MLIR Teama0f3db402019-01-29 17:36:41287 removeEdge(id, outEdge.id, outEdge.value);
MLIR Teamc4237ae2019-01-18 16:56:27288 }
289 }
290 // Erase remaining node state.
291 inEdges.erase(id);
292 outEdges.erase(id);
293 nodes.erase(id);
294 }
295
MLIR Teamd7c82442019-01-30 23:53:41296 // Returns true if node 'id' writes to any memref which escapes (or is an
297 // argument to) the function/block. Returns false otherwise.
298 bool writesToLiveInOrEscapingMemrefs(unsigned id) {
MLIR Team71495d52019-01-22 21:23:37299 Node *node = getNode(id);
300 for (auto *storeOpInst : node->stores) {
River Riddle35807bc2019-12-23 05:59:55301 auto memref = cast<AffineStoreOp>(storeOpInst).getMemRef();
River Riddle2bdf33c2020-01-11 16:54:04302 auto *op = memref.getDefiningOp();
MLIR Team58aa3832019-02-16 01:12:19303 // Return true if 'memref' is a block argument.
River Riddle99b87c92019-03-27 21:02:02304 if (!op)
MLIR Teamd7c82442019-01-30 23:53:41305 return true;
MLIR Team58aa3832019-02-16 01:12:19306 // Return true if any use of 'memref' escapes the function.
River Riddle2bdf33c2020-01-11 16:54:04307 for (auto *user : memref.getUsers())
River Riddle8780d8d2019-05-18 18:09:07308 if (!isMemRefDereferencingOp(*user))
MLIR Teamd7c82442019-01-30 23:53:41309 return true;
MLIR Teamd7c82442019-01-30 23:53:41310 }
311 return false;
312 }
313
Diego Caballero34510552019-10-09 17:36:54314 // Returns the unique AffineStoreOp in `node` that meets all the following:
315 // *) store is the only one that writes to a function-local memref live out
316 // of `node`,
317 // *) store is not the source of a self-dependence on `node`.
318 // Otherwise, returns a null AffineStoreOp.
319 AffineStoreOp getUniqueOutgoingStore(Node *node) {
320 AffineStoreOp uniqueStore;
321
322 // Return null if `node` doesn't have any outgoing edges.
323 auto outEdgeIt = outEdges.find(node->id);
324 if (outEdgeIt == outEdges.end())
325 return nullptr;
326
327 const auto &nodeOutEdges = outEdgeIt->second;
328 for (auto *op : node->stores) {
329 auto storeOp = cast<AffineStoreOp>(op);
River Riddle35807bc2019-12-23 05:59:55330 auto memref = storeOp.getMemRef();
Diego Caballero34510552019-10-09 17:36:54331 // Skip this store if there are no dependences on its memref. This means
332 // that store either:
333 // *) writes to a memref that is only read within the same loop nest
334 // (self-dependence edges are not represented in graph at the moment),
335 // *) writes to a function live out memref (function parameter), or
336 // *) is dead.
337 if (llvm::all_of(nodeOutEdges, [=](const Edge &edge) {
338 return (edge.value != memref);
339 }))
340 continue;
341
342 if (uniqueStore)
343 // Found multiple stores to function-local live-out memrefs.
344 return nullptr;
345 // Found first store to function-local live-out memref.
346 uniqueStore = storeOp;
347 }
348
349 return uniqueStore;
350 }
351
MLIR Teamd7c82442019-01-30 23:53:41352 // Returns true if node 'id' can be removed from the graph. Returns false
353 // otherwise. A node can be removed from the graph iff the following
354 // conditions are met:
355 // *) The node does not write to any memref which escapes (or is a
356 // function/block argument).
357 // *) The node has no successors in the dependence graph.
358 bool canRemoveNode(unsigned id) {
359 if (writesToLiveInOrEscapingMemrefs(id))
360 return false;
361 Node *node = getNode(id);
362 for (auto *storeOpInst : node->stores) {
MLIR Teama0f3db402019-01-29 17:36:41363 // Return false if there exist out edges from 'id' on 'memref'.
Andy Davis2e1187d2019-07-03 17:35:03364 if (getOutEdgeCount(id, cast<AffineStoreOp>(storeOpInst).getMemRef()) > 0)
MLIR Teama0f3db402019-01-29 17:36:41365 return false;
MLIR Team71495d52019-01-22 21:23:37366 }
MLIR Teama0f3db402019-01-29 17:36:41367 return true;
MLIR Team71495d52019-01-22 21:23:37368 }
369
MLIR Teamd038e342019-03-01 19:50:25370 // Returns true iff there is an edge from node 'srcId' to node 'dstId' which
371 // is for 'value' if non-null, or for any value otherwise. Returns false
372 // otherwise.
River Riddlee62a6952019-12-23 22:45:01373 bool hasEdge(unsigned srcId, unsigned dstId, Value value = nullptr) {
MLIR Team27d067e2019-01-16 17:55:02374 if (outEdges.count(srcId) == 0 || inEdges.count(dstId) == 0) {
375 return false;
376 }
377 bool hasOutEdge = llvm::any_of(outEdges[srcId], [=](Edge &edge) {
MLIR Teamd038e342019-03-01 19:50:25378 return edge.id == dstId && (!value || edge.value == value);
MLIR Team27d067e2019-01-16 17:55:02379 });
380 bool hasInEdge = llvm::any_of(inEdges[dstId], [=](Edge &edge) {
MLIR Teamd038e342019-03-01 19:50:25381 return edge.id == srcId && (!value || edge.value == value);
MLIR Team27d067e2019-01-16 17:55:02382 });
383 return hasOutEdge && hasInEdge;
384 }
385
MLIR Teama0f3db402019-01-29 17:36:41386 // Adds an edge from node 'srcId' to node 'dstId' for 'value'.
River Riddlee62a6952019-12-23 22:45:01387 void addEdge(unsigned srcId, unsigned dstId, Value value) {
MLIR Teama0f3db402019-01-29 17:36:41388 if (!hasEdge(srcId, dstId, value)) {
389 outEdges[srcId].push_back({dstId, value});
390 inEdges[dstId].push_back({srcId, value});
River Riddle2bdf33c2020-01-11 16:54:04391 if (value.getType().isa<MemRefType>())
MLIR Teama0f3db402019-01-29 17:36:41392 memrefEdgeCount[value]++;
MLIR Team27d067e2019-01-16 17:55:02393 }
MLIR Team6892ffb2018-12-20 04:42:55394 }
395
MLIR Teama0f3db402019-01-29 17:36:41396 // Removes an edge from node 'srcId' to node 'dstId' for 'value'.
River Riddlee62a6952019-12-23 22:45:01397 void removeEdge(unsigned srcId, unsigned dstId, Value value) {
MLIR Team6892ffb2018-12-20 04:42:55398 assert(inEdges.count(dstId) > 0);
399 assert(outEdges.count(srcId) > 0);
River Riddle2bdf33c2020-01-11 16:54:04400 if (value.getType().isa<MemRefType>()) {
MLIR Teama0f3db402019-01-29 17:36:41401 assert(memrefEdgeCount.count(value) > 0);
402 memrefEdgeCount[value]--;
403 }
MLIR Team6892ffb2018-12-20 04:42:55404 // Remove 'srcId' from 'inEdges[dstId]'.
405 for (auto it = inEdges[dstId].begin(); it != inEdges[dstId].end(); ++it) {
MLIR Teama0f3db402019-01-29 17:36:41406 if ((*it).id == srcId && (*it).value == value) {
MLIR Team6892ffb2018-12-20 04:42:55407 inEdges[dstId].erase(it);
408 break;
409 }
410 }
411 // Remove 'dstId' from 'outEdges[srcId]'.
412 for (auto it = outEdges[srcId].begin(); it != outEdges[srcId].end(); ++it) {
MLIR Teama0f3db402019-01-29 17:36:41413 if ((*it).id == dstId && (*it).value == value) {
MLIR Team6892ffb2018-12-20 04:42:55414 outEdges[srcId].erase(it);
415 break;
416 }
417 }
418 }
419
MLIR Teamd038e342019-03-01 19:50:25420 // Returns true if there is a path in the dependence graph from node 'srcId'
421 // to node 'dstId'. Returns false otherwise.
422 bool hasDependencePath(unsigned srcId, unsigned dstId) {
423 // Worklist state is: <node-id, next-output-edge-index-to-visit>
424 SmallVector<std::pair<unsigned, unsigned>, 4> worklist;
425 worklist.push_back({srcId, 0});
426 // Run DFS traversal to see if 'dstId' is reachable from 'srcId'.
427 while (!worklist.empty()) {
428 auto &idAndIndex = worklist.back();
429 // Return true if we have reached 'dstId'.
430 if (idAndIndex.first == dstId)
431 return true;
432 // Pop and continue if node has no out edges, or if all out edges have
433 // already been visited.
434 if (outEdges.count(idAndIndex.first) == 0 ||
435 idAndIndex.second == outEdges[idAndIndex.first].size()) {
436 worklist.pop_back();
437 continue;
438 }
439 // Get graph edge to traverse.
440 Edge edge = outEdges[idAndIndex.first][idAndIndex.second];
441 // Increment next output edge index for 'idAndIndex'.
442 ++idAndIndex.second;
443 // Add node at 'edge.id' to worklist.
444 worklist.push_back({edge.id, 0});
445 }
446 return false;
447 }
448
MLIR Teama0f3db402019-01-29 17:36:41449 // Returns the input edge count for node 'id' and 'memref' from src nodes
MLIR Teamd038e342019-03-01 19:50:25450 // which access 'memref' with a store operation.
River Riddlee62a6952019-12-23 22:45:01451 unsigned getIncomingMemRefAccesses(unsigned id, Value memref) {
MLIR Team6892ffb2018-12-20 04:42:55452 unsigned inEdgeCount = 0;
453 if (inEdges.count(id) > 0)
454 for (auto &inEdge : inEdges[id])
MLIR Teama0f3db402019-01-29 17:36:41455 if (inEdge.value == memref) {
456 Node *srcNode = getNode(inEdge.id);
457 // Only count in edges from 'srcNode' if 'srcNode' accesses 'memref'
MLIR Teamd038e342019-03-01 19:50:25458 if (srcNode->getStoreOpCount(memref) > 0)
MLIR Teama0f3db402019-01-29 17:36:41459 ++inEdgeCount;
460 }
MLIR Team6892ffb2018-12-20 04:42:55461 return inEdgeCount;
462 }
463
MLIR Teamd038e342019-03-01 19:50:25464 // Returns the output edge count for node 'id' and 'memref' (if non-null),
465 // otherwise returns the total output edge count from node 'id'.
River Riddlee62a6952019-12-23 22:45:01466 unsigned getOutEdgeCount(unsigned id, Value memref = nullptr) {
MLIR Team6892ffb2018-12-20 04:42:55467 unsigned outEdgeCount = 0;
468 if (outEdges.count(id) > 0)
469 for (auto &outEdge : outEdges[id])
MLIR Teamd038e342019-03-01 19:50:25470 if (!memref || outEdge.value == memref)
MLIR Team6892ffb2018-12-20 04:42:55471 ++outEdgeCount;
472 return outEdgeCount;
473 }
474
River Riddle99b87c92019-03-27 21:02:02475 // Computes and returns an insertion point operation, before which the
MLIR Teama0f3db402019-01-29 17:36:41476 // the fused <srcId, dstId> loop nest can be inserted while preserving
477 // dependences. Returns nullptr if no such insertion point is found.
River Riddle99b87c92019-03-27 21:02:02478 Operation *getFusedLoopNestInsertionPoint(unsigned srcId, unsigned dstId) {
MLIR Team5c5739d2019-01-25 06:27:40479 if (outEdges.count(srcId) == 0)
River Riddle99b87c92019-03-27 21:02:02480 return getNode(dstId)->op;
MLIR Teama0f3db402019-01-29 17:36:41481
482 // Build set of insts in range (srcId, dstId) which depend on 'srcId'.
River Riddle99b87c92019-03-27 21:02:02483 SmallPtrSet<Operation *, 2> srcDepInsts;
MLIR Teama0f3db402019-01-29 17:36:41484 for (auto &outEdge : outEdges[srcId])
MLIR Teama78edcd2019-02-05 14:57:08485 if (outEdge.id != dstId)
River Riddle99b87c92019-03-27 21:02:02486 srcDepInsts.insert(getNode(outEdge.id)->op);
MLIR Teama0f3db402019-01-29 17:36:41487
488 // Build set of insts in range (srcId, dstId) on which 'dstId' depends.
River Riddle99b87c92019-03-27 21:02:02489 SmallPtrSet<Operation *, 2> dstDepInsts;
MLIR Teama0f3db402019-01-29 17:36:41490 for (auto &inEdge : inEdges[dstId])
MLIR Teama78edcd2019-02-05 14:57:08491 if (inEdge.id != srcId)
River Riddle99b87c92019-03-27 21:02:02492 dstDepInsts.insert(getNode(inEdge.id)->op);
MLIR Teama0f3db402019-01-29 17:36:41493
River Riddle99b87c92019-03-27 21:02:02494 Operation *srcNodeInst = getNode(srcId)->op;
495 Operation *dstNodeInst = getNode(dstId)->op;
MLIR Teama0f3db402019-01-29 17:36:41496
497 // Computing insertion point:
River Riddle99b87c92019-03-27 21:02:02498 // *) Walk all operation positions in Block operation list in the
499 // range (src, dst). For each operation 'op' visited in this search:
500 // *) Store in 'firstSrcDepPos' the first position where 'op' has a
MLIR Teama0f3db402019-01-29 17:36:41501 // dependence edge from 'srcNode'.
River Riddle99b87c92019-03-27 21:02:02502 // *) Store in 'lastDstDepPost' the last position where 'op' has a
MLIR Teama0f3db402019-01-29 17:36:41503 // dependence edge to 'dstNode'.
504 // *) Compare 'firstSrcDepPos' and 'lastDstDepPost' to determine the
River Riddle99b87c92019-03-27 21:02:02505 // operation insertion point (or return null pointer if no such
MLIR Teama0f3db402019-01-29 17:36:41506 // insertion point exists: 'firstSrcDepPos' <= 'lastDstDepPos').
River Riddle99b87c92019-03-27 21:02:02507 SmallVector<Operation *, 2> depInsts;
MLIR Teama0f3db402019-01-29 17:36:41508 Optional<unsigned> firstSrcDepPos;
509 Optional<unsigned> lastDstDepPos;
510 unsigned pos = 0;
511 for (Block::iterator it = std::next(Block::iterator(srcNodeInst));
512 it != Block::iterator(dstNodeInst); ++it) {
River Riddle99b87c92019-03-27 21:02:02513 Operation *op = &(*it);
514 if (srcDepInsts.count(op) > 0 && firstSrcDepPos == None)
MLIR Teama0f3db402019-01-29 17:36:41515 firstSrcDepPos = pos;
River Riddle99b87c92019-03-27 21:02:02516 if (dstDepInsts.count(op) > 0)
MLIR Teama0f3db402019-01-29 17:36:41517 lastDstDepPos = pos;
River Riddle99b87c92019-03-27 21:02:02518 depInsts.push_back(op);
MLIR Teama0f3db402019-01-29 17:36:41519 ++pos;
MLIR Team5c5739d2019-01-25 06:27:40520 }
MLIR Teama0f3db402019-01-29 17:36:41521
522 if (firstSrcDepPos.hasValue()) {
523 if (lastDstDepPos.hasValue()) {
524 if (firstSrcDepPos.getValue() <= lastDstDepPos.getValue()) {
525 // No valid insertion point exists which preserves dependences.
526 return nullptr;
527 }
528 }
529 // Return the insertion point at 'firstSrcDepPos'.
530 return depInsts[firstSrcDepPos.getValue()];
531 }
532 // No dependence targets in range (or only dst deps in range), return
533 // 'dstNodInst' insertion point.
534 return dstNodeInst;
MLIR Team6892ffb2018-12-20 04:42:55535 }
536
MLIR Teama0f3db402019-01-29 17:36:41537 // Updates edge mappings from node 'srcId' to node 'dstId' after 'oldMemRef'
Andy Davis68a8da42019-11-18 19:20:03538 // has been replaced in node at 'dstId' by a private memref depending
539 // on the value of 'createPrivateMemRef'.
River Riddlee62a6952019-12-23 22:45:01540 void updateEdges(unsigned srcId, unsigned dstId, Value oldMemRef,
Andy Davis68a8da42019-11-18 19:20:03541 bool createPrivateMemRef) {
Kazuaki Ishizakifc817b02020-01-20 03:14:37542 // For each edge in 'inEdges[srcId]': add new edge remapping to 'dstId'.
MLIR Team6892ffb2018-12-20 04:42:55543 if (inEdges.count(srcId) > 0) {
544 SmallVector<Edge, 2> oldInEdges = inEdges[srcId];
545 for (auto &inEdge : oldInEdges) {
MLIR Teama0f3db402019-01-29 17:36:41546 // Add edge from 'inEdge.id' to 'dstId' if not for 'oldMemRef'.
547 if (inEdge.value != oldMemRef)
548 addEdge(inEdge.id, dstId, inEdge.value);
MLIR Team6892ffb2018-12-20 04:42:55549 }
550 }
MLIR Teamc4237ae2019-01-18 16:56:27551 // For each edge in 'outEdges[srcId]': remove edge from 'srcId' to 'dstId'.
MLIR Team6892ffb2018-12-20 04:42:55552 if (outEdges.count(srcId) > 0) {
553 SmallVector<Edge, 2> oldOutEdges = outEdges[srcId];
554 for (auto &outEdge : oldOutEdges) {
MLIR Teamc4237ae2019-01-18 16:56:27555 // Remove any out edges from 'srcId' to 'dstId' across memrefs.
556 if (outEdge.id == dstId)
MLIR Teama0f3db402019-01-29 17:36:41557 removeEdge(srcId, outEdge.id, outEdge.value);
MLIR Team6892ffb2018-12-20 04:42:55558 }
559 }
MLIR Teama0f3db402019-01-29 17:36:41560 // Remove any edges in 'inEdges[dstId]' on 'oldMemRef' (which is being
561 // replaced by a private memref). These edges could come from nodes
562 // other than 'srcId' which were removed in the previous step.
Andy Davis68a8da42019-11-18 19:20:03563 if (inEdges.count(dstId) > 0 && createPrivateMemRef) {
MLIR Teama0f3db402019-01-29 17:36:41564 SmallVector<Edge, 2> oldInEdges = inEdges[dstId];
565 for (auto &inEdge : oldInEdges)
566 if (inEdge.value == oldMemRef)
567 removeEdge(inEdge.id, dstId, inEdge.value);
568 }
MLIR Team6892ffb2018-12-20 04:42:55569 }
570
MLIR Teamd038e342019-03-01 19:50:25571 // Update edge mappings for nodes 'sibId' and 'dstId' to reflect fusion
572 // of sibling node 'sidId' into node 'dstId'.
573 void updateEdges(unsigned sibId, unsigned dstId) {
574 // For each edge in 'inEdges[sibId]':
575 // *) Add new edge from source node 'inEdge.id' to 'dstNode'.
576 // *) Remove edge from source node 'inEdge.id' to 'sibNode'.
577 if (inEdges.count(sibId) > 0) {
578 SmallVector<Edge, 2> oldInEdges = inEdges[sibId];
579 for (auto &inEdge : oldInEdges) {
580 addEdge(inEdge.id, dstId, inEdge.value);
581 removeEdge(inEdge.id, sibId, inEdge.value);
582 }
583 }
584
585 // For each edge in 'outEdges[sibId]' to node 'id'
586 // *) Add new edge from 'dstId' to 'outEdge.id'.
587 // *) Remove edge from 'sibId' to 'outEdge.id'.
588 if (outEdges.count(sibId) > 0) {
589 SmallVector<Edge, 2> oldOutEdges = outEdges[sibId];
590 for (auto &outEdge : oldOutEdges) {
591 addEdge(dstId, outEdge.id, outEdge.value);
592 removeEdge(sibId, outEdge.id, outEdge.value);
593 }
594 }
595 }
596
MLIR Team6892ffb2018-12-20 04:42:55597 // Adds ops in 'loads' and 'stores' to node at 'id'.
River Riddle99b87c92019-03-27 21:02:02598 void addToNode(unsigned id, const SmallVectorImpl<Operation *> &loads,
599 const SmallVectorImpl<Operation *> &stores) {
MLIR Team6892ffb2018-12-20 04:42:55600 Node *node = getNode(id);
Chris Lattner456ad6a2018-12-29 00:05:35601 for (auto *loadOpInst : loads)
602 node->loads.push_back(loadOpInst);
603 for (auto *storeOpInst : stores)
604 node->stores.push_back(storeOpInst);
MLIR Team6892ffb2018-12-20 04:42:55605 }
606
MLIR Teamc4237ae2019-01-18 16:56:27607 void clearNodeLoadAndStores(unsigned id) {
608 Node *node = getNode(id);
609 node->loads.clear();
610 node->stores.clear();
611 }
612
MLIR Teamd038e342019-03-01 19:50:25613 // Calls 'callback' for each input edge incident to node 'id' which carries a
614 // memref dependence.
615 void forEachMemRefInputEdge(unsigned id,
616 const std::function<void(Edge)> &callback) {
617 if (inEdges.count(id) > 0)
618 forEachMemRefEdge(inEdges[id], callback);
619 }
Amit Sabne70a416d2019-04-09 16:17:40620
MLIR Teamd038e342019-03-01 19:50:25621 // Calls 'callback' for each output edge from node 'id' which carries a
622 // memref dependence.
623 void forEachMemRefOutputEdge(unsigned id,
624 const std::function<void(Edge)> &callback) {
625 if (outEdges.count(id) > 0)
626 forEachMemRefEdge(outEdges[id], callback);
627 }
Amit Sabne70a416d2019-04-09 16:17:40628
MLIR Teamd038e342019-03-01 19:50:25629 // Calls 'callback' for each edge in 'edges' which carries a memref
630 // dependence.
631 void forEachMemRefEdge(ArrayRef<Edge> edges,
632 const std::function<void(Edge)> &callback) {
633 for (auto &edge : edges) {
634 // Skip if 'edge' is not a memref dependence edge.
River Riddle2bdf33c2020-01-11 16:54:04635 if (!edge.value.getType().isa<MemRefType>())
MLIR Teamd038e342019-03-01 19:50:25636 continue;
637 assert(nodes.count(edge.id) > 0);
638 // Skip if 'edge.id' is not a loop nest.
River Riddled5b60ee82019-05-12 01:59:54639 if (!isa<AffineForOp>(getNode(edge.id)->op))
MLIR Teamd038e342019-03-01 19:50:25640 continue;
641 // Visit current input edge 'edge'.
642 callback(edge);
643 }
644 }
645
MLIR Team6892ffb2018-12-20 04:42:55646 void print(raw_ostream &os) const {
647 os << "\nMemRefDependenceGraph\n";
648 os << "\nNodes:\n";
649 for (auto &idAndNode : nodes) {
650 os << "Node: " << idAndNode.first << "\n";
651 auto it = inEdges.find(idAndNode.first);
652 if (it != inEdges.end()) {
653 for (const auto &e : it->second)
MLIR Teama0f3db402019-01-29 17:36:41654 os << " InEdge: " << e.id << " " << e.value << "\n";
MLIR Team6892ffb2018-12-20 04:42:55655 }
656 it = outEdges.find(idAndNode.first);
657 if (it != outEdges.end()) {
658 for (const auto &e : it->second)
MLIR Teama0f3db402019-01-29 17:36:41659 os << " OutEdge: " << e.id << " " << e.value << "\n";
MLIR Team6892ffb2018-12-20 04:42:55660 }
661 }
662 }
663 void dump() const { print(llvm::errs()); }
664};
665
River Riddle2666b972019-12-18 18:46:16666} // end anonymous namespace
667
Kazuaki Ishizaki8bfedb32019-10-20 07:11:03668// Initializes the data dependence graph by walking operations in 'f'.
MLIR Team6892ffb2018-12-20 04:42:55669// Assigns each node in the graph a node id based on program order in 'f'.
Chris Lattner315a4662018-12-28 21:07:39670// TODO(andydavis) Add support for taking a Block arg to construct the
MLIR Team6892ffb2018-12-20 04:42:55671// dependence graph at a different depth.
River Riddle8c443672019-07-09 23:17:55672bool MemRefDependenceGraph::init(FuncOp f) {
River Riddlee62a6952019-12-23 22:45:01673 DenseMap<Value, SetVector<unsigned>> memrefAccesses;
Chris Lattnerdffc5892018-12-29 23:33:43674
675 // TODO: support multi-block functions.
Chris Lattner46ade282019-03-26 01:02:49676 if (f.getBlocks().size() != 1)
Chris Lattnerdffc5892018-12-29 23:33:43677 return false;
678
River Riddle99b87c92019-03-27 21:02:02679 DenseMap<Operation *, unsigned> forToNodeMap;
680 for (auto &op : f.front()) {
River Riddlec5ecf992019-05-11 22:56:50681 if (auto forOp = dyn_cast<AffineForOp>(op)) {
River Riddle5052bd82019-02-02 00:42:18682 // Create graph node 'id' to represent top-level 'forOp' and record
MLIR Team6892ffb2018-12-20 04:42:55683 // all loads and store accesses it contains.
684 LoopNestStateCollector collector;
River Riddle99b87c92019-03-27 21:02:02685 collector.collect(&op);
River Riddle832567b2019-03-25 17:14:34686 // Return false if a non 'affine.for' region was found (not currently
687 // supported).
River Riddle75553832019-01-29 05:23:53688 if (collector.hasNonForRegion)
MLIR Team6892ffb2018-12-20 04:42:55689 return false;
River Riddle99b87c92019-03-27 21:02:02690 Node node(nextNodeId++, &op);
Chris Lattner456ad6a2018-12-29 00:05:35691 for (auto *opInst : collector.loadOpInsts) {
692 node.loads.push_back(opInst);
River Riddle35807bc2019-12-23 05:59:55693 auto memref = cast<AffineLoadOp>(opInst).getMemRef();
MLIR Team6892ffb2018-12-20 04:42:55694 memrefAccesses[memref].insert(node.id);
695 }
Chris Lattner456ad6a2018-12-29 00:05:35696 for (auto *opInst : collector.storeOpInsts) {
697 node.stores.push_back(opInst);
River Riddle35807bc2019-12-23 05:59:55698 auto memref = cast<AffineStoreOp>(opInst).getMemRef();
MLIR Team6892ffb2018-12-20 04:42:55699 memrefAccesses[memref].insert(node.id);
700 }
River Riddle99b87c92019-03-27 21:02:02701 forToNodeMap[&op] = node.id;
MLIR Team6892ffb2018-12-20 04:42:55702 nodes.insert({node.id, node});
Andy Davis2e1187d2019-07-03 17:35:03703 } else if (auto loadOp = dyn_cast<AffineLoadOp>(op)) {
River Riddleb4992772019-02-04 18:38:47704 // Create graph node for top-level load op.
River Riddle99b87c92019-03-27 21:02:02705 Node node(nextNodeId++, &op);
706 node.loads.push_back(&op);
River Riddle35807bc2019-12-23 05:59:55707 auto memref = cast<AffineLoadOp>(op).getMemRef();
River Riddleb4992772019-02-04 18:38:47708 memrefAccesses[memref].insert(node.id);
709 nodes.insert({node.id, node});
Andy Davis2e1187d2019-07-03 17:35:03710 } else if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
River Riddleb4992772019-02-04 18:38:47711 // Create graph node for top-level store op.
River Riddle99b87c92019-03-27 21:02:02712 Node node(nextNodeId++, &op);
713 node.stores.push_back(&op);
River Riddle35807bc2019-12-23 05:59:55714 auto memref = cast<AffineStoreOp>(op).getMemRef();
River Riddleb4992772019-02-04 18:38:47715 memrefAccesses[memref].insert(node.id);
716 nodes.insert({node.id, node});
River Riddle99b87c92019-03-27 21:02:02717 } else if (op.getNumRegions() != 0) {
River Riddleb4992772019-02-04 18:38:47718 // Return false if another region is found (not currently supported).
719 return false;
River Riddle99b87c92019-03-27 21:02:02720 } else if (op.getNumResults() > 0 && !op.use_empty()) {
River Riddleb4992772019-02-04 18:38:47721 // Create graph node for top-level producer of SSA values, which
722 // could be used by loop nest nodes.
River Riddle99b87c92019-03-27 21:02:02723 Node node(nextNodeId++, &op);
River Riddleb4992772019-02-04 18:38:47724 nodes.insert({node.id, node});
MLIR Teama0f3db402019-01-29 17:36:41725 }
726 }
727
728 // Add dependence edges between nodes which produce SSA values and their
729 // users.
730 for (auto &idAndNode : nodes) {
731 const Node &node = idAndNode.second;
732 if (!node.loads.empty() || !node.stores.empty())
733 continue;
River Riddle99b87c92019-03-27 21:02:02734 auto *opInst = node.op;
River Riddle35807bc2019-12-23 05:59:55735 for (auto value : opInst->getResults()) {
River Riddle2bdf33c2020-01-11 16:54:04736 for (auto *user : value.getUsers()) {
Chris Lattnerd9b5bc82019-03-25 02:53:05737 SmallVector<AffineForOp, 4> loops;
River Riddle8780d8d2019-05-18 18:09:07738 getLoopIVs(*user, &loops);
MLIR Teama0f3db402019-01-29 17:36:41739 if (loops.empty())
740 continue;
River Riddlef9d91532019-03-27 00:05:09741 assert(forToNodeMap.count(loops[0].getOperation()) > 0);
742 unsigned userLoopNestId = forToNodeMap[loops[0].getOperation()];
MLIR Teama0f3db402019-01-29 17:36:41743 addEdge(node.id, userLoopNestId, value);
MLIR Team6892ffb2018-12-20 04:42:55744 }
745 }
MLIR Team6892ffb2018-12-20 04:42:55746 }
747
748 // Walk memref access lists and add graph edges between dependent nodes.
749 for (auto &memrefAndList : memrefAccesses) {
750 unsigned n = memrefAndList.second.size();
751 for (unsigned i = 0; i < n; ++i) {
752 unsigned srcId = memrefAndList.second[i];
753 bool srcHasStore =
754 getNode(srcId)->getStoreOpCount(memrefAndList.first) > 0;
755 for (unsigned j = i + 1; j < n; ++j) {
756 unsigned dstId = memrefAndList.second[j];
757 bool dstHasStore =
758 getNode(dstId)->getStoreOpCount(memrefAndList.first) > 0;
759 if (srcHasStore || dstHasStore)
760 addEdge(srcId, dstId, memrefAndList.first);
761 }
762 }
763 }
764 return true;
765}
766
MLIR Team27d067e2019-01-16 17:55:02767// Removes load operations from 'srcLoads' which operate on 'memref', and
768// adds them to 'dstLoads'.
River Riddlee62a6952019-12-23 22:45:01769static void moveLoadsAccessingMemrefTo(Value memref,
River Riddle99b87c92019-03-27 21:02:02770 SmallVectorImpl<Operation *> *srcLoads,
771 SmallVectorImpl<Operation *> *dstLoads) {
MLIR Team27d067e2019-01-16 17:55:02772 dstLoads->clear();
River Riddle99b87c92019-03-27 21:02:02773 SmallVector<Operation *, 4> srcLoadsToKeep;
MLIR Team27d067e2019-01-16 17:55:02774 for (auto *load : *srcLoads) {
Andy Davis2e1187d2019-07-03 17:35:03775 if (cast<AffineLoadOp>(load).getMemRef() == memref)
MLIR Team27d067e2019-01-16 17:55:02776 dstLoads->push_back(load);
777 else
778 srcLoadsToKeep.push_back(load);
MLIR Team38c2fe32019-01-14 19:26:25779 }
MLIR Team27d067e2019-01-16 17:55:02780 srcLoads->swap(srcLoadsToKeep);
MLIR Team38c2fe32019-01-14 19:26:25781}
782
MLIR Team27d067e2019-01-16 17:55:02783// Returns the innermost common loop depth for the set of operations in 'ops'.
River Riddle99b87c92019-03-27 21:02:02784static unsigned getInnermostCommonLoopDepth(ArrayRef<Operation *> ops) {
MLIR Team27d067e2019-01-16 17:55:02785 unsigned numOps = ops.size();
786 assert(numOps > 0);
787
Chris Lattnerd9b5bc82019-03-25 02:53:05788 std::vector<SmallVector<AffineForOp, 4>> loops(numOps);
MLIR Team27d067e2019-01-16 17:55:02789 unsigned loopDepthLimit = std::numeric_limits<unsigned>::max();
790 for (unsigned i = 0; i < numOps; ++i) {
791 getLoopIVs(*ops[i], &loops[i]);
792 loopDepthLimit =
793 std::min(loopDepthLimit, static_cast<unsigned>(loops[i].size()));
MLIR Team38c2fe32019-01-14 19:26:25794 }
MLIR Team27d067e2019-01-16 17:55:02795
796 unsigned loopDepth = 0;
797 for (unsigned d = 0; d < loopDepthLimit; ++d) {
798 unsigned i;
799 for (i = 1; i < numOps; ++i) {
River Riddle5052bd82019-02-02 00:42:18800 if (loops[i - 1][d] != loops[i][d])
MLIR Team27d067e2019-01-16 17:55:02801 break;
MLIR Team27d067e2019-01-16 17:55:02802 }
803 if (i != numOps)
804 break;
805 ++loopDepth;
806 }
807 return loopDepth;
MLIR Team38c2fe32019-01-14 19:26:25808}
809
MLIR Teamd7c82442019-01-30 23:53:41810// Returns the maximum loop depth at which no dependences between 'loadOpInsts'
811// and 'storeOpInsts' are satisfied.
River Riddle99b87c92019-03-27 21:02:02812static unsigned getMaxLoopDepth(ArrayRef<Operation *> loadOpInsts,
813 ArrayRef<Operation *> storeOpInsts) {
MLIR Teamd7c82442019-01-30 23:53:41814 // Merge loads and stores into the same array.
River Riddle99b87c92019-03-27 21:02:02815 SmallVector<Operation *, 2> ops(loadOpInsts.begin(), loadOpInsts.end());
MLIR Teamd7c82442019-01-30 23:53:41816 ops.append(storeOpInsts.begin(), storeOpInsts.end());
817
818 // Compute the innermost common loop depth for loads and stores.
819 unsigned loopDepth = getInnermostCommonLoopDepth(ops);
820
821 // Return common loop depth for loads if there are no store ops.
822 if (storeOpInsts.empty())
823 return loopDepth;
824
825 // Check dependences on all pairs of ops in 'ops' and store the minimum
826 // loop depth at which a dependence is satisfied.
827 for (unsigned i = 0, e = ops.size(); i < e; ++i) {
828 auto *srcOpInst = ops[i];
829 MemRefAccess srcAccess(srcOpInst);
830 for (unsigned j = 0; j < e; ++j) {
831 auto *dstOpInst = ops[j];
832 MemRefAccess dstAccess(dstOpInst);
833
834 unsigned numCommonLoops =
835 getNumCommonSurroundingLoops(*srcOpInst, *dstOpInst);
836 for (unsigned d = 1; d <= numCommonLoops + 1; ++d) {
837 FlatAffineConstraints dependenceConstraints;
838 // TODO(andydavis) Cache dependence analysis results, check cache here.
Andy Davise33e36f2019-06-10 17:50:08839 DependenceResult result = checkMemrefAccessDependence(
840 srcAccess, dstAccess, d, &dependenceConstraints,
841 /*dependenceComponents=*/nullptr);
842 if (hasDependence(result)) {
MLIR Teamd7c82442019-01-30 23:53:41843 // Store minimum loop depth and break because we want the min 'd' at
844 // which there is a dependence.
845 loopDepth = std::min(loopDepth, d - 1);
846 break;
847 }
848 }
849 }
850 }
851 return loopDepth;
852}
853
MLIR Team8f5f2c72019-02-15 17:32:18854// Sinks all sequential loops to the innermost levels (while preserving
855// relative order among them) and moves all parallel loops to the
856// outermost (while again preserving relative order among them).
857// This can increase the loop depth at which we can fuse a slice, since we are
858// pushing loop carried dependence to a greater depth in the loop nest.
859static void sinkSequentialLoops(MemRefDependenceGraph::Node *node) {
River Riddled5b60ee82019-05-12 01:59:54860 assert(isa<AffineForOp>(node->op));
Andy Davis90d40232019-05-13 13:57:56861 AffineForOp newRootForOp = sinkSequentialLoops(cast<AffineForOp>(node->op));
862 node->op = newRootForOp.getOperation();
MLIR Team8f5f2c72019-02-15 17:32:18863}
864
Uday Bondhugula8be26272019-02-02 01:06:22865// TODO(mlir-team): improve/complete this when we have target data.
River Riddle2666b972019-12-18 18:46:16866static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
Uday Bondhugula8be26272019-02-02 01:06:22867 auto elementType = memRefType.getElementType();
868
869 unsigned sizeInBits;
River Riddlede5a81b2020-03-02 17:18:45870 if (elementType.isIntOrFloat()) {
Uday Bondhugula8be26272019-02-02 01:06:22871 sizeInBits = elementType.getIntOrFloatBitWidth();
872 } else {
873 auto vectorType = elementType.cast<VectorType>();
874 sizeInBits =
875 vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
876 }
877 return llvm::divideCeil(sizeInBits, 8);
878}
879
MLIR Teamc4237ae2019-01-18 16:56:27880// Creates and returns a private (single-user) memref for fused loop rooted
River Riddle5052bd82019-02-02 00:42:18881// at 'forOp', with (potentially reduced) memref size based on the
Uday Bondhugula94a03f82019-01-22 21:58:52882// MemRefRegion written to by 'srcStoreOpInst' at depth 'dstLoopDepth'.
883// TODO(bondhugula): consider refactoring the common code from generateDma and
884// this one.
River Riddlee62a6952019-12-23 22:45:01885static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
886 unsigned dstLoopDepth,
887 Optional<unsigned> fastMemorySpace,
888 uint64_t localBufSizeThreshold) {
River Riddlef9d91532019-03-27 00:05:09889 auto *forInst = forOp.getOperation();
River Riddle5052bd82019-02-02 00:42:18890
891 // Create builder to insert alloc op just before 'forOp'.
River Riddlef1b848e2019-06-05 02:18:23892 OpBuilder b(forInst);
MLIR Teamc4237ae2019-01-18 16:56:27893 // Builder to create constants at the top level.
River Riddlece502af2019-07-08 18:20:26894 OpBuilder top(forInst->getParentOfType<FuncOp>().getBody());
MLIR Teamc4237ae2019-01-18 16:56:27895 // Create new memref type based on slice bounds.
River Riddle35807bc2019-12-23 05:59:55896 auto oldMemRef = cast<AffineStoreOp>(srcStoreOpInst).getMemRef();
River Riddle2bdf33c2020-01-11 16:54:04897 auto oldMemRefType = oldMemRef.getType().cast<MemRefType>();
MLIR Teamc4237ae2019-01-18 16:56:27898 unsigned rank = oldMemRefType.getRank();
899
Uday Bondhugula94a03f82019-01-22 21:58:52900 // Compute MemRefRegion for 'srcStoreOpInst' at depth 'dstLoopDepth'.
Uday Bondhugula0f504142019-02-04 21:48:44901 MemRefRegion region(srcStoreOpInst->getLoc());
River Riddle1e55ae12019-03-08 06:14:47902 bool validRegion = succeeded(region.compute(srcStoreOpInst, dstLoopDepth));
MLIR Teamd42ef782019-03-04 19:01:25903 (void)validRegion;
904 assert(validRegion && "unexpected memref region failure");
River Riddle6859f332019-01-23 22:39:45905 SmallVector<int64_t, 4> newShape;
MLIR Teamc4237ae2019-01-18 16:56:27906 std::vector<SmallVector<int64_t, 4>> lbs;
Uday Bondhugula94a03f82019-01-22 21:58:52907 SmallVector<int64_t, 8> lbDivisors;
MLIR Teamc4237ae2019-01-18 16:56:27908 lbs.reserve(rank);
909 // Query 'region' for 'newShape' and lower bounds of MemRefRegion accessed
Uday Bondhugula94a03f82019-01-22 21:58:52910 // by 'srcStoreOpInst' at depth 'dstLoopDepth'.
MLIR Teamc4237ae2019-01-18 16:56:27911 Optional<int64_t> numElements =
Uday Bondhugula0f504142019-02-04 21:48:44912 region.getConstantBoundingSizeAndShape(&newShape, &lbs, &lbDivisors);
Uday Bondhugula8be26272019-02-02 01:06:22913 assert(numElements.hasValue() &&
914 "non-constant number of elts in local buffer");
MLIR Teamc4237ae2019-01-18 16:56:27915
Uday Bondhugula0f504142019-02-04 21:48:44916 const FlatAffineConstraints *cst = region.getConstraints();
Kazuaki Ishizaki8bfedb32019-10-20 07:11:03917 // 'outerIVs' holds the values that this memory region is symbolic/parametric
Uday Bondhugula94a03f82019-01-22 21:58:52918 // on; this would correspond to loop IVs surrounding the level at which the
919 // slice is being materialized.
River Riddlee62a6952019-12-23 22:45:01920 SmallVector<Value, 8> outerIVs;
Uday Bondhugula94a03f82019-01-22 21:58:52921 cst->getIdValues(rank, cst->getNumIds(), &outerIVs);
922
923 // Build 'rank' AffineExprs from MemRefRegion 'lbs'
MLIR Teamc4237ae2019-01-18 16:56:27924 SmallVector<AffineExpr, 4> offsets;
925 offsets.reserve(rank);
926 for (unsigned d = 0; d < rank; ++d) {
Uday Bondhugula94a03f82019-01-22 21:58:52927 assert(lbs[d].size() == cst->getNumCols() - rank && "incorrect bound size");
928
MLIR Teamc4237ae2019-01-18 16:56:27929 AffineExpr offset = top.getAffineConstantExpr(0);
930 for (unsigned j = 0, e = cst->getNumCols() - rank - 1; j < e; j++) {
931 offset = offset + lbs[d][j] * top.getAffineDimExpr(j);
932 }
Uday Bondhugula94a03f82019-01-22 21:58:52933 assert(lbDivisors[d] > 0);
934 offset =
935 (offset + lbs[d][cst->getNumCols() - 1 - rank]).floorDiv(lbDivisors[d]);
MLIR Teamc4237ae2019-01-18 16:56:27936 offsets.push_back(offset);
937 }
938
939 // Create 'newMemRefType' using 'newShape' from MemRefRegion accessed
940 // by 'srcStoreOpInst'.
Uday Bondhugula8be26272019-02-02 01:06:22941 uint64_t bufSize =
942 getMemRefEltSizeInBytes(oldMemRefType) * numElements.getValue();
943 unsigned newMemSpace;
Uday Bondhugulad4b3ff12019-02-27 00:10:19944 if (bufSize <= localBufSizeThreshold && fastMemorySpace.hasValue()) {
Uday Bondhugula8be26272019-02-02 01:06:22945 newMemSpace = fastMemorySpace.getValue();
946 } else {
947 newMemSpace = oldMemRefType.getMemorySpace();
948 }
River Riddle2acc2202019-10-18 03:08:01949 auto newMemRefType = MemRefType::get(newShape, oldMemRefType.getElementType(),
950 {}, newMemSpace);
MLIR Teamc4237ae2019-01-18 16:56:27951 // Gather alloc operands for the dynamic dimensions of the memref.
River Riddlee62a6952019-12-23 22:45:01952 SmallVector<Value, 4> allocOperands;
MLIR Teamc4237ae2019-01-18 16:56:27953 unsigned dynamicDimCount = 0;
954 for (auto dimSize : oldMemRefType.getShape()) {
955 if (dimSize == -1)
956 allocOperands.push_back(
River Riddleaf1abcc2019-03-25 18:13:31957 top.create<DimOp>(forOp.getLoc(), oldMemRef, dynamicDimCount++));
MLIR Teamc4237ae2019-01-18 16:56:27958 }
959
River Riddle5052bd82019-02-02 00:42:18960 // Create new private memref for fused loop 'forOp'.
MLIR Teama0f3db402019-01-29 17:36:41961 // TODO(andydavis) Create/move alloc ops for private memrefs closer to their
962 // consumer loop nests to reduce their live range. Currently they are added
963 // at the beginning of the function, because loop nests can be reordered
964 // during the fusion pass.
River Riddlee62a6952019-12-23 22:45:01965 Value newMemRef =
River Riddleaf1abcc2019-03-25 18:13:31966 top.create<AllocOp>(forOp.getLoc(), newMemRefType, allocOperands);
MLIR Teamc4237ae2019-01-18 16:56:27967
968 // Build an AffineMap to remap access functions based on lower bound offsets.
969 SmallVector<AffineExpr, 4> remapExprs;
970 remapExprs.reserve(rank);
971 unsigned zeroOffsetCount = 0;
972 for (unsigned i = 0; i < rank; i++) {
973 if (auto constExpr = offsets[i].dyn_cast<AffineConstantExpr>())
974 if (constExpr.getValue() == 0)
975 ++zeroOffsetCount;
Uday Bondhugula94a03f82019-01-22 21:58:52976 auto dimExpr = b.getAffineDimExpr(outerIVs.size() + i);
977
978 auto remapExpr =
979 simplifyAffineExpr(dimExpr - offsets[i], outerIVs.size() + rank, 0);
980 remapExprs.push_back(remapExpr);
MLIR Teamc4237ae2019-01-18 16:56:27981 }
MLIR Team5a91b982019-05-29 21:56:41982 auto indexRemap = zeroOffsetCount == rank
983 ? AffineMap()
River Riddle2acc2202019-10-18 03:08:01984 : AffineMap::get(outerIVs.size() + rank, 0, remapExprs);
MLIR Teamc4237ae2019-01-18 16:56:27985 // Replace all users of 'oldMemRef' with 'newMemRef'.
Uday Bondhugulaaa2cee92019-08-28 00:56:25986 LogicalResult res =
Uday Bondhugula94a03f82019-01-22 21:58:52987 replaceAllMemRefUsesWith(oldMemRef, newMemRef, {}, indexRemap,
988 /*extraOperands=*/outerIVs,
Uday Bondhugula727a50a2019-09-18 18:25:33989 /*symbolOperands=*/{},
River Riddleaf1abcc2019-03-25 18:13:31990 /*domInstFilter=*/&*forOp.getBody()->begin());
Uday Bondhugulaaa2cee92019-08-28 00:56:25991 assert(succeeded(res) &&
992 "replaceAllMemrefUsesWith should always succeed here");
993 (void)res;
MLIR Teamc4237ae2019-01-18 16:56:27994 return newMemRef;
995}
996
Diego Caballero34510552019-10-09 17:36:54997// Checks if node 'srcId' can be safely fused into node 'dstId'. Node 'srcId'
998// may write to multiple memrefs but it is required that only one of them,
Diego Caballero330d1ff2019-12-03 14:09:21999// 'srcLiveOutStoreOp', has output edges.
Diego Caballero34510552019-10-09 17:36:541000// Returns true if 'dstNode's read/write region to 'memref' is a super set of
Diego Caballero330d1ff2019-12-03 14:09:211001// 'srcNode's write region to 'memref' and 'srcId' has only one output edge.
MLIR Team58aa3832019-02-16 01:12:191002// TODO(andydavis) Generalize this to handle more live in/out cases.
1003static bool canFuseSrcWhichWritesToLiveOut(unsigned srcId, unsigned dstId,
Diego Caballero34510552019-10-09 17:36:541004 AffineStoreOp srcLiveOutStoreOp,
MLIR Team58aa3832019-02-16 01:12:191005 MemRefDependenceGraph *mdg) {
Diego Caballero34510552019-10-09 17:36:541006 assert(srcLiveOutStoreOp && "Expected a valid store op");
MLIR Team58aa3832019-02-16 01:12:191007 auto *dstNode = mdg->getNode(dstId);
River Riddlee62a6952019-12-23 22:45:011008 Value memref = srcLiveOutStoreOp.getMemRef();
Diego Caballero330d1ff2019-12-03 14:09:211009 // Return false if 'srcNode' has more than one output edge on 'memref'.
1010 if (mdg->getOutEdgeCount(srcId, memref) > 1)
1011 return false;
MLIR Team58aa3832019-02-16 01:12:191012
Diego Caballero34510552019-10-09 17:36:541013 // Compute MemRefRegion 'srcWriteRegion' for 'srcStoreOp' on 'memref'.
1014 MemRefRegion srcWriteRegion(srcLiveOutStoreOp.getLoc());
1015 if (failed(srcWriteRegion.compute(srcLiveOutStoreOp, /*loopDepth=*/0))) {
MLIR Teamd42ef782019-03-04 19:01:251016 LLVM_DEBUG(llvm::dbgs()
1017 << "Unable to compute MemRefRegion for source operation\n.");
1018 return false;
1019 }
MLIR Team58aa3832019-02-16 01:12:191020 SmallVector<int64_t, 4> srcShape;
1021 // Query 'srcWriteRegion' for 'srcShape' and 'srcNumElements'.
Diego Caballero34510552019-10-09 17:36:541022 // by 'srcStoreOp' at depth 'dstLoopDepth'.
MLIR Team58aa3832019-02-16 01:12:191023 Optional<int64_t> srcNumElements =
1024 srcWriteRegion.getConstantBoundingSizeAndShape(&srcShape);
1025 if (!srcNumElements.hasValue())
1026 return false;
1027
Andy Davis7c1fc9e2019-04-02 13:37:401028 // Compute MemRefRegion 'dstRegion' for 'dstStore/LoadOpInst' on 'memref'.
MLIR Team9d9675f2019-03-28 21:54:491029 // TODO(andydavis) Compute 'unionboundingbox' of all write regions (one for
1030 // each store op in 'dstStoreOps').
Andy Davis7c1fc9e2019-04-02 13:37:401031 SmallVector<Operation *, 2> dstStoreOps;
1032 dstNode->getStoreOpsForMemref(memref, &dstStoreOps);
1033 SmallVector<Operation *, 2> dstLoadOps;
1034 dstNode->getLoadOpsForMemref(memref, &dstLoadOps);
1035
1036 auto *dstOpInst = dstStoreOps.empty() ? dstLoadOps[0] : dstStoreOps[0];
1037 MemRefRegion dstRegion(dstOpInst->getLoc());
1038 if (failed(dstRegion.compute(dstOpInst, /*loopDepth=*/0))) {
MLIR Teamd42ef782019-03-04 19:01:251039 LLVM_DEBUG(llvm::dbgs()
1040 << "Unable to compute MemRefRegion for dest operation\n.");
1041 return false;
1042 }
MLIR Team58aa3832019-02-16 01:12:191043 SmallVector<int64_t, 4> dstShape;
Andy Davis7c1fc9e2019-04-02 13:37:401044 // Query 'dstRegion' for 'dstShape' and 'dstNumElements'.
1045 // by 'dstOpInst' at depth 'dstLoopDepth'.
MLIR Team58aa3832019-02-16 01:12:191046 Optional<int64_t> dstNumElements =
Andy Davis7c1fc9e2019-04-02 13:37:401047 dstRegion.getConstantBoundingSizeAndShape(&dstShape);
MLIR Team58aa3832019-02-16 01:12:191048 if (!dstNumElements.hasValue())
1049 return false;
1050
1051 // Return false if write region is not a superset of 'srcNodes' write
1052 // region to 'memref'.
1053 // TODO(andydavis) Check the shape and lower bounds here too.
1054 if (srcNumElements != dstNumElements)
1055 return false;
1056 return true;
1057}
1058
MLIR Team27d067e2019-01-16 17:55:021059// Checks the profitability of fusing a backwards slice of the loop nest
MLIR Teamd7c82442019-01-30 23:53:411060// surrounding 'srcOpInst' into the loop nest surrounding 'dstLoadOpInsts'.
MLIR Teamd038e342019-03-01 19:50:251061// The argument 'srcStoreOpInst' is used to calculate the storage reduction on
1062// the memref being produced and consumed, which is an input to the cost model.
Kazuaki Ishizaki8bfedb32019-10-20 07:11:031063// For producer-consumer fusion, 'srcStoreOpInst' will be the same as
MLIR Teamd038e342019-03-01 19:50:251064// 'srcOpInst', as we are slicing w.r.t to that producer.
1065// For input-reuse fusion, 'srcOpInst' will be the src loop nest LoadOp which
1066// reads from the same memref as dst loop nest load ops, and 'srcStoreOpInst'
1067// will be the unique store op in the src node, which will be used to check
1068// that the write region is the same after input-reuse fusion.
Uday Bondhugulab4a14432019-01-26 00:00:501069// Returns true if it is profitable to fuse the candidate loop nests. Returns
1070// false otherwise. `dstLoopDepth` is set to the most profitable depth at which
1071// to materialize the source loop nest slice.
MLIR Team38c2fe32019-01-14 19:26:251072// The profitability model executes the following steps:
MLIR Team27d067e2019-01-16 17:55:021073// *) Computes the backward computation slice at 'srcOpInst'. This
1074// computation slice of the loop nest surrounding 'srcOpInst' is
MLIR Team38c2fe32019-01-14 19:26:251075// represented by modified src loop bounds in 'sliceState', which are
MLIR Team27d067e2019-01-16 17:55:021076// functions of loop IVs in the loop nest surrounding 'srcOpInst'.
MLIR Team38c2fe32019-01-14 19:26:251077// *) Computes the cost of unfused src/dst loop nests (currently the cost of a
1078// loop nest is the total number of dynamic operation instances in the loop
1079// nest).
1080// *) Computes the cost of fusing a slice of the src loop nest into the dst
MLIR Team27d067e2019-01-16 17:55:021081// loop nest at various values of dst loop depth, attempting to fuse
Kazuaki Ishizaki8bfedb32019-10-20 07:11:031082// the largest computation slice at the maximal dst loop depth (closest to
1083// the load) to minimize reuse distance and potentially enable subsequent
MLIR Team27d067e2019-01-16 17:55:021084// load/store forwarding.
MLIR Teamd7c82442019-01-30 23:53:411085// NOTE: If the dst loop nest includes multiple loads in 'dstLoadOpInsts' for
MLIR Team27d067e2019-01-16 17:55:021086// the same memref as is written by 'srcOpInst', then the union of slice
1087// loop bounds is used to compute the slice and associated slice cost.
Uday Bondhugulab4a14432019-01-26 00:00:501088// NOTE: 'dstLoopDepth' refers to the loop depth within the destination loop
MLIR Team38c2fe32019-01-14 19:26:251089// nest, at which the src computation slice is inserted/fused.
MLIR Team27d067e2019-01-16 17:55:021090// NOTE: We attempt to maximize the dst loop depth, but there are cases
1091// where a particular setting for 'dstLoopNest' might fuse an unsliced
MLIR Team38c2fe32019-01-14 19:26:251092// loop (within the src computation slice) at a depth which results in
Kazuaki Ishizaki8bfedb32019-10-20 07:11:031093// excessive recomputation (see unit tests for examples).
MLIR Team38c2fe32019-01-14 19:26:251094// *) Compares the total cost of the unfused loop nests to the min cost fused
1095// loop nest computed in the previous step, and returns true if the latter
1096// is lower.
River Riddle99b87c92019-03-27 21:02:021097static bool isFusionProfitable(Operation *srcOpInst, Operation *srcStoreOpInst,
1098 ArrayRef<Operation *> dstLoadOpInsts,
1099 ArrayRef<Operation *> dstStoreOpInsts,
MLIR Team38c2fe32019-01-14 19:26:251100 ComputationSliceState *sliceState,
Uday Bondhugulace7e59532019-03-08 17:21:521101 unsigned *dstLoopDepth, bool maximalFusion) {
Uday Bondhugula06d21d92019-01-25 01:01:491102 LLVM_DEBUG({
1103 llvm::dbgs() << "Checking whether fusion is profitable between:\n";
Uday Bondhugulaa1dad3a2019-02-20 02:17:191104 llvm::dbgs() << " " << *srcOpInst << " and \n";
MLIR Teamd7c82442019-01-30 23:53:411105 for (auto dstOpInst : dstLoadOpInsts) {
Uday Bondhugulaa1dad3a2019-02-20 02:17:191106 llvm::dbgs() << " " << *dstOpInst << "\n";
Uday Bondhugula06d21d92019-01-25 01:01:491107 };
1108 });
Uday Bondhugula864d9e02019-01-23 17:16:241109
MLIR Team38c2fe32019-01-14 19:26:251110 // Compute cost of sliced and unsliced src loop nest.
Chris Lattnerd9b5bc82019-03-25 02:53:051111 SmallVector<AffineForOp, 4> srcLoopIVs;
MLIR Team27d067e2019-01-16 17:55:021112 getLoopIVs(*srcOpInst, &srcLoopIVs);
MLIR Team38c2fe32019-01-14 19:26:251113 unsigned numSrcLoopIVs = srcLoopIVs.size();
1114
1115 // Walk src loop nest and collect stats.
1116 LoopNestStats srcLoopNestStats;
Andy Davis59b68142019-06-18 15:52:091117 if (!getLoopNestStats(srcLoopIVs[0], &srcLoopNestStats))
MLIR Team38c2fe32019-01-14 19:26:251118 return false;
Andy Davis59b68142019-06-18 15:52:091119
MLIR Team38c2fe32019-01-14 19:26:251120 // Compute cost of dst loop nest.
Chris Lattnerd9b5bc82019-03-25 02:53:051121 SmallVector<AffineForOp, 4> dstLoopIVs;
MLIR Teamd7c82442019-01-30 23:53:411122 getLoopIVs(*dstLoadOpInsts[0], &dstLoopIVs);
MLIR Team38c2fe32019-01-14 19:26:251123
1124 LoopNestStats dstLoopNestStats;
Andy Davis59b68142019-06-18 15:52:091125 if (!getLoopNestStats(dstLoopIVs[0], &dstLoopNestStats))
MLIR Team38c2fe32019-01-14 19:26:251126 return false;
1127
MLIR Teamd7c82442019-01-30 23:53:411128 // Compute the maximum loop depth at which we can can insert the src slice
MLIR Teamd038e342019-03-01 19:50:251129 // and still satisfy dest loop nest dependences, for producer-consumer fusion.
1130 unsigned maxDstLoopDepth =
1131 (srcOpInst == srcStoreOpInst)
1132 ? getMaxLoopDepth(dstLoadOpInsts, dstStoreOpInsts)
1133 : dstLoopIVs.size();
MLIR Teamc1ff9e82019-03-06 04:33:301134 if (maxDstLoopDepth == 0) {
1135 LLVM_DEBUG(llvm::dbgs() << "Can't fuse: maxDstLoopDepth == 0 .\n");
MLIR Team27d067e2019-01-16 17:55:021136 return false;
MLIR Teamc1ff9e82019-03-06 04:33:301137 }
MLIR Team27d067e2019-01-16 17:55:021138
1139 // Search for min cost value for 'dstLoopDepth'. At each value of
1140 // 'dstLoopDepth' from 'maxDstLoopDepth' to '1', compute computation slice
1141 // bounds between 'srcOpInst' and each op in 'dstOpinsts' (taking the union
1142 // of these bounds). Next the union slice bounds are used to calculate
1143 // the cost of the slice and the cost of the slice inserted into the dst
1144 // loop nest at 'dstLoopDepth'.
Uday Bondhugula864d9e02019-01-23 17:16:241145 uint64_t minFusedLoopNestComputeCost = std::numeric_limits<uint64_t>::max();
MLIR Teamd038e342019-03-01 19:50:251146 double maxStorageReduction = 0.0;
Uday Bondhugula864d9e02019-01-23 17:16:241147 Optional<uint64_t> sliceMemEstimate = None;
1148
MLIR Team27d067e2019-01-16 17:55:021149 SmallVector<ComputationSliceState, 4> sliceStates;
1150 sliceStates.resize(maxDstLoopDepth);
Uday Bondhugula864d9e02019-01-23 17:16:241151 // The best loop depth at which to materialize the slice.
1152 Optional<unsigned> bestDstLoopDepth = None;
1153
1154 // Compute op instance count for the src loop nest without iteration slicing.
Andy Davis59b68142019-06-18 15:52:091155 uint64_t srcLoopNestCost = getComputeCost(srcLoopIVs[0], srcLoopNestStats);
Uday Bondhugula864d9e02019-01-23 17:16:241156
MLIR Teamb9dde912019-02-06 19:01:101157 // Compute src loop nest write region size.
MLIR Teamd038e342019-03-01 19:50:251158 MemRefRegion srcWriteRegion(srcStoreOpInst->getLoc());
River Riddle1e55ae12019-03-08 06:14:471159 if (failed(srcWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0))) {
MLIR Teamd42ef782019-03-04 19:01:251160 LLVM_DEBUG(llvm::dbgs()
River Riddle99b87c92019-03-27 21:02:021161 << "Unable to compute MemRefRegion for source operation\n.");
MLIR Teamd42ef782019-03-04 19:01:251162 return false;
1163 }
1164
MLIR Teamb9dde912019-02-06 19:01:101165 Optional<int64_t> maybeSrcWriteRegionSizeBytes =
1166 srcWriteRegion.getRegionSize();
1167 if (!maybeSrcWriteRegionSizeBytes.hasValue())
1168 return false;
1169 int64_t srcWriteRegionSizeBytes = maybeSrcWriteRegionSizeBytes.getValue();
1170
Uday Bondhugula864d9e02019-01-23 17:16:241171 // Compute op instance count for the src loop nest.
Andy Davis59b68142019-06-18 15:52:091172 uint64_t dstLoopNestCost = getComputeCost(dstLoopIVs[0], dstLoopNestStats);
MLIR Team27d067e2019-01-16 17:55:021173
MLIR Teamb9dde912019-02-06 19:01:101174 // Evaluate all depth choices for materializing the slice in the destination
1175 // loop nest.
MLIR Team27d067e2019-01-16 17:55:021176 for (unsigned i = maxDstLoopDepth; i >= 1; --i) {
MLIR Teamc1ff9e82019-03-06 04:33:301177 // Compute the union of slice bounds of all ops in 'dstLoadOpInsts'.
Andy Davis1de0f972019-05-29 21:02:141178 if (failed(mlir::computeSliceUnion({srcOpInst}, dstLoadOpInsts,
Andy Davis898cf0e2019-06-17 16:59:351179 /*loopDepth=*/i,
1180 /*numCommonLoops=*/0,
1181 /*isBackwardSlice=*/true,
Andy Davis1de0f972019-05-29 21:02:141182 &sliceStates[i - 1]))) {
MLIR Teamc1ff9e82019-03-06 04:33:301183 LLVM_DEBUG(llvm::dbgs()
Andy Davis1de0f972019-05-29 21:02:141184 << "computeSliceUnion failed for loopDepth: " << i << "\n");
MLIR Teamc1ff9e82019-03-06 04:33:301185 continue;
MLIR Team38c2fe32019-01-14 19:26:251186 }
MLIR Teamc1ff9e82019-03-06 04:33:301187
Andy Davis59b68142019-06-18 15:52:091188 int64_t fusedLoopNestComputeCost;
1189 if (!getFusionComputeCost(srcLoopIVs[0], srcLoopNestStats, dstLoopIVs[0],
1190 dstLoopNestStats, &sliceStates[i - 1],
1191 &fusedLoopNestComputeCost)) {
1192 LLVM_DEBUG(llvm::dbgs() << "Unable to compute fusion compute cost.\n.");
Uday Bondhugula864d9e02019-01-23 17:16:241193 continue;
MLIR Teamc1ff9e82019-03-06 04:33:301194 }
Uday Bondhugula864d9e02019-01-23 17:16:241195
Uday Bondhugula864d9e02019-01-23 17:16:241196 double additionalComputeFraction =
1197 fusedLoopNestComputeCost /
1198 (static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
1199 1;
1200
Amit Sabne70a416d2019-04-09 16:17:401201 // Determine what the slice write MemRefRegion would be, if the src loop
MLIR Teamb9dde912019-02-06 19:01:101202 // nest slice 'sliceStates[i - 1]' were to be inserted into the dst loop
1203 // nest at loop depth 'i'
MLIR Teamd038e342019-03-01 19:50:251204 MemRefRegion sliceWriteRegion(srcStoreOpInst->getLoc());
River Riddle1e55ae12019-03-08 06:14:471205 if (failed(sliceWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0,
1206 &sliceStates[i - 1]))) {
MLIR Teamc1ff9e82019-03-06 04:33:301207 LLVM_DEBUG(llvm::dbgs()
1208 << "Failed to compute slice write region at loopDepth: " << i
1209 << "\n");
MLIR Teamd42ef782019-03-04 19:01:251210 continue;
MLIR Teamc1ff9e82019-03-06 04:33:301211 }
MLIR Teamd42ef782019-03-04 19:01:251212
MLIR Teamb9dde912019-02-06 19:01:101213 Optional<int64_t> maybeSliceWriteRegionSizeBytes =
1214 sliceWriteRegion.getRegionSize();
1215 if (!maybeSliceWriteRegionSizeBytes.hasValue() ||
MLIR Teamc1ff9e82019-03-06 04:33:301216 maybeSliceWriteRegionSizeBytes.getValue() == 0) {
1217 LLVM_DEBUG(llvm::dbgs()
1218 << "Failed to get slice write region size at loopDepth: " << i
1219 << "\n");
MLIR Teamb9dde912019-02-06 19:01:101220 continue;
MLIR Teamc1ff9e82019-03-06 04:33:301221 }
MLIR Teamb9dde912019-02-06 19:01:101222 int64_t sliceWriteRegionSizeBytes =
1223 maybeSliceWriteRegionSizeBytes.getValue();
1224
MLIR Teamd038e342019-03-01 19:50:251225 // If we are fusing for reuse, check that write regions remain the same.
1226 // TODO(andydavis) Write region check should check sizes and offsets in
1227 // each dimension, so that we are sure they are covering the same memref
1228 // region. Also, move this out to a isMemRefRegionSuperSet helper function.
1229 if (srcOpInst != srcStoreOpInst &&
1230 sliceWriteRegionSizeBytes != srcWriteRegionSizeBytes)
1231 continue;
1232
MLIR Teamb9dde912019-02-06 19:01:101233 double storageReduction = static_cast<double>(srcWriteRegionSizeBytes) /
1234 static_cast<double>(sliceWriteRegionSizeBytes);
Uday Bondhugula864d9e02019-01-23 17:16:241235
Uday Bondhugula06d21d92019-01-25 01:01:491236 LLVM_DEBUG({
1237 std::stringstream msg;
1238 msg << " evaluating fusion profitability at depth : " << i << "\n"
Uday Bondhugulad4b3ff12019-02-27 00:10:191239 << std::fixed << std::setprecision(2)
1240 << " additional compute fraction: "
Uday Bondhugula06d21d92019-01-25 01:01:491241 << 100.0 * additionalComputeFraction << "%\n"
1242 << " storage reduction factor: " << storageReduction << "x\n"
1243 << " fused nest cost: " << fusedLoopNestComputeCost << "\n"
Uday Bondhugulaa1dad3a2019-02-20 02:17:191244 << " src write region size: " << srcWriteRegionSizeBytes << "\n"
1245 << " slice write region size: " << sliceWriteRegionSizeBytes
1246 << "\n";
Uday Bondhugula06d21d92019-01-25 01:01:491247 llvm::dbgs() << msg.str();
1248 });
Uday Bondhugula864d9e02019-01-23 17:16:241249
1250 double computeToleranceThreshold =
1251 clFusionAddlComputeTolerance.getNumOccurrences() > 0
1252 ? clFusionAddlComputeTolerance
1253 : LoopFusion::kComputeToleranceThreshold;
1254
1255 // TODO(b/123247369): This is a placeholder cost model.
1256 // Among all choices that add an acceptable amount of redundant computation
1257 // (as per computeToleranceThreshold), we will simply pick the one that
1258 // reduces the intermediary size the most.
1259 if ((storageReduction > maxStorageReduction) &&
Uday Bondhugulace7e59532019-03-08 17:21:521260 (maximalFusion ||
Uday Bondhugula864d9e02019-01-23 17:16:241261 (additionalComputeFraction < computeToleranceThreshold))) {
1262 maxStorageReduction = storageReduction;
MLIR Team27d067e2019-01-16 17:55:021263 bestDstLoopDepth = i;
Uday Bondhugula864d9e02019-01-23 17:16:241264 minFusedLoopNestComputeCost = fusedLoopNestComputeCost;
MLIR Teamb9dde912019-02-06 19:01:101265 sliceMemEstimate = sliceWriteRegionSizeBytes;
MLIR Team38c2fe32019-01-14 19:26:251266 }
1267 }
1268
Uday Bondhugula864d9e02019-01-23 17:16:241269 // A simple cost model: fuse if it reduces the memory footprint. If
1270 // -maximal-fusion is set, fuse nevertheless.
MLIR Team38c2fe32019-01-14 19:26:251271
Uday Bondhugulace7e59532019-03-08 17:21:521272 if (!maximalFusion && !bestDstLoopDepth.hasValue()) {
Uday Bondhugulaa1dad3a2019-02-20 02:17:191273 LLVM_DEBUG(
1274 llvm::dbgs()
1275 << "All fusion choices involve more than the threshold amount of "
1276 "redundant computation; NOT fusing.\n");
MLIR Team38c2fe32019-01-14 19:26:251277 return false;
Uday Bondhugula864d9e02019-01-23 17:16:241278 }
1279
MLIR Teamd42ef782019-03-04 19:01:251280 if (!bestDstLoopDepth.hasValue()) {
1281 LLVM_DEBUG(llvm::dbgs() << "no fusion depth could be evaluated.\n");
1282 return false;
1283 }
Uday Bondhugula864d9e02019-01-23 17:16:241284
1285 // Set dstLoopDepth based on best values from search.
1286 *dstLoopDepth = bestDstLoopDepth.getValue();
1287
1288 LLVM_DEBUG(
Uday Bondhugula06d21d92019-01-25 01:01:491289 llvm::dbgs() << " LoopFusion fusion stats:"
1290 << "\n best loop depth: " << bestDstLoopDepth
Uday Bondhugula864d9e02019-01-23 17:16:241291 << "\n src loop nest compute cost: " << srcLoopNestCost
1292 << "\n dst loop nest compute cost: " << dstLoopNestCost
1293 << "\n fused loop nest compute cost: "
1294 << minFusedLoopNestComputeCost << "\n");
1295
River Riddle5052bd82019-02-02 00:42:181296 auto dstMemSize = getMemoryFootprintBytes(dstLoopIVs[0]);
1297 auto srcMemSize = getMemoryFootprintBytes(srcLoopIVs[0]);
Uday Bondhugula864d9e02019-01-23 17:16:241298
1299 Optional<double> storageReduction = None;
1300
Uday Bondhugulace7e59532019-03-08 17:21:521301 if (!maximalFusion) {
Uday Bondhugula864d9e02019-01-23 17:16:241302 if (!dstMemSize.hasValue() || !srcMemSize.hasValue()) {
1303 LLVM_DEBUG(
1304 llvm::dbgs()
1305 << " fusion memory benefit cannot be evaluated; NOT fusing.\n");
1306 return false;
1307 }
1308
1309 auto srcMemSizeVal = srcMemSize.getValue();
1310 auto dstMemSizeVal = dstMemSize.getValue();
1311
1312 assert(sliceMemEstimate.hasValue() && "expected value");
Uday Bondhugula864d9e02019-01-23 17:16:241313 auto fusedMem = dstMemSizeVal + sliceMemEstimate.getValue();
1314
1315 LLVM_DEBUG(llvm::dbgs() << " src mem: " << srcMemSizeVal << "\n"
1316 << " dst mem: " << dstMemSizeVal << "\n"
1317 << " fused mem: " << fusedMem << "\n"
1318 << " slice mem: " << sliceMemEstimate << "\n");
1319
Jacques Pienaar2fe8ae42019-05-04 02:48:571320 if (static_cast<long>(fusedMem) > srcMemSizeVal + dstMemSizeVal) {
Uday Bondhugula864d9e02019-01-23 17:16:241321 LLVM_DEBUG(llvm::dbgs() << "Fusion is not profitable; NOT fusing.\n");
1322 return false;
1323 }
1324 storageReduction =
1325 100.0 *
1326 (1.0 - fusedMem / (static_cast<double>(srcMemSizeVal) + dstMemSizeVal));
1327 }
1328
1329 double additionalComputeFraction =
1330 100.0 * (minFusedLoopNestComputeCost /
1331 (static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
1332 1);
MLIR Team5c5739d2019-01-25 06:27:401333 (void)additionalComputeFraction;
Uday Bondhugula06d21d92019-01-25 01:01:491334 LLVM_DEBUG({
1335 std::stringstream msg;
1336 msg << " fusion is most profitable at depth " << *dstLoopDepth << " with "
MLIR Team8564b272019-02-22 15:48:591337 << std::setprecision(2) << additionalComputeFraction
Uday Bondhugula06d21d92019-01-25 01:01:491338 << "% redundant computation and a ";
1339 msg << (storageReduction.hasValue()
1340 ? std::to_string(storageReduction.getValue())
1341 : "<unknown>");
1342 msg << "% storage reduction.\n";
1343 llvm::dbgs() << msg.str();
1344 });
Uday Bondhugula864d9e02019-01-23 17:16:241345
MLIR Team27d067e2019-01-16 17:55:021346 // Update return parameter 'sliceState' with 'bestSliceState'.
Uday Bondhugula864d9e02019-01-23 17:16:241347 ComputationSliceState *bestSliceState = &sliceStates[*dstLoopDepth - 1];
MLIR Team27d067e2019-01-16 17:55:021348 sliceState->lbs = bestSliceState->lbs;
1349 sliceState->ubs = bestSliceState->ubs;
1350 sliceState->lbOperands = bestSliceState->lbOperands;
1351 sliceState->ubOperands = bestSliceState->ubOperands;
Uday Bondhugula864d9e02019-01-23 17:16:241352
MLIR Team27d067e2019-01-16 17:55:021353 // Canonicalize slice bound affine maps.
MLIR Team38c2fe32019-01-14 19:26:251354 for (unsigned i = 0; i < numSrcLoopIVs; ++i) {
Nicolas Vasilache0e7a8a92019-01-26 18:41:171355 if (sliceState->lbs[i] != AffineMap()) {
MLIR Team27d067e2019-01-16 17:55:021356 canonicalizeMapAndOperands(&sliceState->lbs[i],
1357 &sliceState->lbOperands[i]);
1358 }
Nicolas Vasilache0e7a8a92019-01-26 18:41:171359 if (sliceState->ubs[i] != AffineMap()) {
MLIR Team27d067e2019-01-16 17:55:021360 canonicalizeMapAndOperands(&sliceState->ubs[i],
1361 &sliceState->ubOperands[i]);
MLIR Team38c2fe32019-01-14 19:26:251362 }
1363 }
1364 return true;
1365}
1366
River Riddle2666b972019-12-18 18:46:161367namespace {
1368
MLIR Teamd038e342019-03-01 19:50:251369// GreedyFusion greedily fuses loop nests which have a producer/consumer or
1370// input-reuse relationship on a memref, with the goal of improving locality.
MLIR Teamf28e4df2018-11-01 14:26:001371//
MLIR Teamd038e342019-03-01 19:50:251372// The steps of the producer-consumer fusion algorithm are as follows:
MLIR Team3b692302018-12-17 17:57:141373//
MLIR Team6892ffb2018-12-20 04:42:551374// *) A worklist is initialized with node ids from the dependence graph.
1375// *) For each node id in the worklist:
Amit Sabne70a416d2019-04-09 16:17:401376// *) Pop an AffineForOp of the worklist. This 'dstAffineForOp' will be a
River Riddle5052bd82019-02-02 00:42:181377// candidate destination AffineForOp into which fusion will be attempted.
1378// *) Add each LoadOp currently in 'dstAffineForOp' into list 'dstLoadOps'.
MLIR Team3b692302018-12-17 17:57:141379// *) For each LoadOp in 'dstLoadOps' do:
Amit Sabne70a416d2019-04-09 16:17:401380// *) Look up dependent loop nests which have a single store op to the same
MLIR Teamd038e342019-03-01 19:50:251381// memref.
1382// *) Check if dependences would be violated by the fusion.
MLIR Team6892ffb2018-12-20 04:42:551383// *) Get a computation slice of 'srcLoopNest', which adjusts its loop
MLIR Team3b692302018-12-17 17:57:141384// bounds to be functions of 'dstLoopNest' IVs and symbols.
1385// *) Fuse the 'srcLoopNest' computation slice into the 'dstLoopNest',
MLIR Teamd038e342019-03-01 19:50:251386// at a loop depth determined by the cost model in 'isFusionProfitable'.
River Riddle99b87c92019-03-27 21:02:021387// *) Add the newly fused load/store operations to the state,
Amit Sabne70a416d2019-04-09 16:17:401388// and also add newly fused load ops to 'dstLoopOps' to be considered
MLIR Team3b692302018-12-17 17:57:141389// as fusion dst load ops in another iteration.
1390// *) Remove old src loop nest and its associated state.
1391//
MLIR Teamd038e342019-03-01 19:50:251392// The steps of the input-reuse fusion algorithm are as follows:
1393//
1394// *) Initialize 'worklist' with node ids from the dependence graph.
1395// *) For each 'dstNode' in the worklist:
1396// *) Find a candidate sibling node 'sibNode' to fuse with 'dstNode' which
1397// loads from the same memref, but which has no dependence paths to/from.
1398// *) Get a computation slice of 'sibLoopNest', which adjusts its loop
1399// bounds to be functions of 'dstLoopNest' IVs and symbols.
1400// *) Fuse the 'sibLoopNest' computation slice into the 'dstLoopNest',
1401// at a loop depth determined by the cost model in 'isFusionProfitable'.
1402// This function also checks that the memref write region of 'sibLoopNest',
1403// is preserved in the fused loop nest.
1404// *) Update graph state to reflect the fusion of 'sibNode' into 'dstNode'.
1405//
River Riddle99b87c92019-03-27 21:02:021406// Given a graph where top-level operations are vertices in the set 'V' and
MLIR Team3b692302018-12-17 17:57:141407// edges in the set 'E' are dependences between vertices, this algorithm
MLIR Team6892ffb2018-12-20 04:42:551408// takes O(V) time for initialization, and has runtime O(V + E).
MLIR Team3b692302018-12-17 17:57:141409//
MLIR Team6892ffb2018-12-20 04:42:551410// This greedy algorithm is not 'maximal' due to the current restriction of
1411// fusing along single producer consumer edges, but there is a TODO to fix this.
MLIR Team3b692302018-12-17 17:57:141412//
1413// TODO(andydavis) Experiment with other fusion policies.
MLIR Team6892ffb2018-12-20 04:42:551414struct GreedyFusion {
1415public:
MLIR Teamd038e342019-03-01 19:50:251416 // The data dependence graph to traverse during fusion.
MLIR Team6892ffb2018-12-20 04:42:551417 MemRefDependenceGraph *mdg;
MLIR Teamd038e342019-03-01 19:50:251418 // Worklist of graph nodes visited during the fusion pass.
MLIR Teama78edcd2019-02-05 14:57:081419 SmallVector<unsigned, 8> worklist;
MLIR Teamd038e342019-03-01 19:50:251420 // Set of graph nodes which are present on the worklist.
MLIR Teama78edcd2019-02-05 14:57:081421 llvm::SmallDenseSet<unsigned, 16> worklistSet;
MLIR Teamd038e342019-03-01 19:50:251422 // Parameter for local buffer size threshold.
1423 unsigned localBufSizeThreshold;
1424 // Parameter for fast memory space.
1425 Optional<unsigned> fastMemorySpace;
Uday Bondhugulace7e59532019-03-08 17:21:521426 // If true, ignore any additional (redundant) computation tolerance threshold
1427 // that would have prevented fusion.
1428 bool maximalFusion;
MLIR Teamf28e4df2018-11-01 14:26:001429
MLIR Teamd038e342019-03-01 19:50:251430 using Node = MemRefDependenceGraph::Node;
1431
1432 GreedyFusion(MemRefDependenceGraph *mdg, unsigned localBufSizeThreshold,
Uday Bondhugulace7e59532019-03-08 17:21:521433 Optional<unsigned> fastMemorySpace, bool maximalFusion)
MLIR Teamd038e342019-03-01 19:50:251434 : mdg(mdg), localBufSizeThreshold(localBufSizeThreshold),
Uday Bondhugulace7e59532019-03-08 17:21:521435 fastMemorySpace(fastMemorySpace), maximalFusion(maximalFusion) {}
MLIR Teamd038e342019-03-01 19:50:251436
1437 // Initializes 'worklist' with nodes from 'mdg'
1438 void init() {
MLIR Teama78edcd2019-02-05 14:57:081439 // TODO(andydavis) Add a priority queue for prioritizing nodes by different
1440 // metrics (e.g. arithmetic intensity/flops-to-bytes ratio).
MLIR Teamd038e342019-03-01 19:50:251441 worklist.clear();
1442 worklistSet.clear();
1443 for (auto &idAndNode : mdg->nodes) {
1444 const Node &node = idAndNode.second;
1445 worklist.push_back(node.id);
1446 worklistSet.insert(node.id);
1447 }
MLIR Team6892ffb2018-12-20 04:42:551448 }
MLIR Team3b692302018-12-17 17:57:141449
MLIR Teamd038e342019-03-01 19:50:251450 // Run the GreedyFusion pass.
1451 // *) First pass through the nodes fuses single-use producer nodes into their
1452 // unique consumer.
1453 // *) Second pass fuses sibling nodes which share no dependence edges.
1454 // *) Third pass fuses any remaining producer nodes into their users.
1455 void run() {
MLIR Teamc1ff9e82019-03-06 04:33:301456 // TODO(andydavis) Run this repeatedly until a fixed-point is reached.
MLIR Teamd038e342019-03-01 19:50:251457 fuseProducerConsumerNodes(/*maxSrcUserCount=*/1);
1458 fuseSiblingNodes();
1459 fuseProducerConsumerNodes(
1460 /*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
1461 eraseUnusedMemRefAllocations();
1462 }
1463
1464 void fuseProducerConsumerNodes(unsigned maxSrcUserCount) {
1465 init();
MLIR Team3b692302018-12-17 17:57:141466 while (!worklist.empty()) {
MLIR Team6892ffb2018-12-20 04:42:551467 unsigned dstId = worklist.back();
MLIR Team3b692302018-12-17 17:57:141468 worklist.pop_back();
MLIR Teama78edcd2019-02-05 14:57:081469 worklistSet.erase(dstId);
1470
MLIR Team6892ffb2018-12-20 04:42:551471 // Skip if this node was removed (fused into another node).
1472 if (mdg->nodes.count(dstId) == 0)
MLIR Team3b692302018-12-17 17:57:141473 continue;
MLIR Team6892ffb2018-12-20 04:42:551474 // Get 'dstNode' into which to attempt fusion.
1475 auto *dstNode = mdg->getNode(dstId);
1476 // Skip if 'dstNode' is not a loop nest.
River Riddled5b60ee82019-05-12 01:59:541477 if (!isa<AffineForOp>(dstNode->op))
MLIR Team3b692302018-12-17 17:57:141478 continue;
MLIR Team8f5f2c72019-02-15 17:32:181479 // Sink sequential loops in 'dstNode' (and thus raise parallel loops)
1480 // while preserving relative order. This can increase the maximum loop
1481 // depth at which we can fuse a slice of a producer loop nest into a
1482 // consumer loop nest.
1483 sinkSequentialLoops(dstNode);
MLIR Team3b692302018-12-17 17:57:141484
River Riddle99b87c92019-03-27 21:02:021485 SmallVector<Operation *, 4> loads = dstNode->loads;
1486 SmallVector<Operation *, 4> dstLoadOpInsts;
River Riddlee62a6952019-12-23 22:45:011487 DenseSet<Value> visitedMemrefs;
MLIR Team6892ffb2018-12-20 04:42:551488 while (!loads.empty()) {
MLIR Team27d067e2019-01-16 17:55:021489 // Get memref of load on top of the stack.
River Riddle35807bc2019-12-23 05:59:551490 auto memref = cast<AffineLoadOp>(loads.back()).getMemRef();
MLIR Teamc4237ae2019-01-18 16:56:271491 if (visitedMemrefs.count(memref) > 0)
1492 continue;
1493 visitedMemrefs.insert(memref);
MLIR Team27d067e2019-01-16 17:55:021494 // Move all loads in 'loads' accessing 'memref' to 'dstLoadOpInsts'.
1495 moveLoadsAccessingMemrefTo(memref, &loads, &dstLoadOpInsts);
MLIR Team6892ffb2018-12-20 04:42:551496 // Skip if no input edges along which to fuse.
1497 if (mdg->inEdges.count(dstId) == 0)
MLIR Team3b692302018-12-17 17:57:141498 continue;
Amit Sabne70a416d2019-04-09 16:17:401499 // Iterate through in-edges for 'dstId' and src node id for any
MLIR Team1e851912019-01-31 00:01:461500 // edges on 'memref'.
1501 SmallVector<unsigned, 2> srcNodeIds;
MLIR Team6892ffb2018-12-20 04:42:551502 for (auto &srcEdge : mdg->inEdges[dstId]) {
1503 // Skip 'srcEdge' if not for 'memref'.
MLIR Teama0f3db402019-01-29 17:36:411504 if (srcEdge.value != memref)
MLIR Team6892ffb2018-12-20 04:42:551505 continue;
MLIR Team1e851912019-01-31 00:01:461506 srcNodeIds.push_back(srcEdge.id);
1507 }
1508 for (unsigned srcId : srcNodeIds) {
1509 // Skip if this node was removed (fused into another node).
1510 if (mdg->nodes.count(srcId) == 0)
1511 continue;
1512 // Get 'srcNode' from which to attempt fusion into 'dstNode'.
1513 auto *srcNode = mdg->getNode(srcId);
MLIR Team6892ffb2018-12-20 04:42:551514 // Skip if 'srcNode' is not a loop nest.
River Riddled5b60ee82019-05-12 01:59:541515 if (!isa<AffineForOp>(srcNode->op))
MLIR Team6892ffb2018-12-20 04:42:551516 continue;
Diego Caballero34510552019-10-09 17:36:541517 // Skip if 'srcNode' has more than one live-out store to a
1518 // function-local memref.
1519 // TODO(andydavis) Support more generic multi-output src loop nests
1520 // fusion.
1521 auto srcStoreOp = mdg->getUniqueOutgoingStore(srcNode);
Andy Davis68a8da42019-11-18 19:20:031522 if (!srcStoreOp) {
1523 // Get the src store op at the deepest loop depth.
1524 // We will use 'LoopFusionUtils::canFuseLoops' to check fusion
1525 // feasibility for loops with multiple stores.
1526 unsigned maxLoopDepth = 0;
1527 for (auto *op : srcNode->stores) {
1528 auto storeOp = cast<AffineStoreOp>(op);
1529 if (storeOp.getMemRef() != memref) {
1530 srcStoreOp = nullptr;
1531 break;
1532 }
1533 unsigned loopDepth = getNestingDepth(*storeOp);
1534 if (loopDepth > maxLoopDepth) {
1535 maxLoopDepth = loopDepth;
1536 srcStoreOp = storeOp;
1537 }
1538 }
1539 if (!srcStoreOp)
1540 continue;
1541 }
1542
Diego Caballero34510552019-10-09 17:36:541543 // Unique outgoing store found must write to 'memref' since 'memref'
1544 // is the one that established the producer-consumer relationship
1545 // between 'srcNode' and 'dstNode'.
1546 assert(srcStoreOp.getMemRef() == memref &&
1547 "Found store to unexpected memref");
Uday Bondhugula864d9e02019-01-23 17:16:241548
MLIR Team58aa3832019-02-16 01:12:191549 // Skip if 'srcNode' writes to any live in or escaping memrefs,
1550 // and cannot be fused.
1551 bool writesToLiveInOrOut =
1552 mdg->writesToLiveInOrEscapingMemrefs(srcNode->id);
1553 if (writesToLiveInOrOut &&
Diego Caballero34510552019-10-09 17:36:541554 !canFuseSrcWhichWritesToLiveOut(srcId, dstId, srcStoreOp, mdg))
MLIR Teamd7c82442019-01-30 23:53:411555 continue;
1556
Kazuaki Ishizaki84a61822019-12-06 13:58:591557 // Don't create a private memref if 'writesToLiveInOrOut'.
Andy Davis68a8da42019-11-18 19:20:031558 bool createPrivateMemref = !writesToLiveInOrOut;
Kazuaki Ishizaki84a61822019-12-06 13:58:591559 // Don't create a private memref if 'srcNode' has in edges on
1560 // 'memref', or if 'dstNode' has out edges on 'memref'.
Andy Davis68a8da42019-11-18 19:20:031561 if (mdg->getIncomingMemRefAccesses(srcNode->id, memref) > 0 ||
1562 mdg->getOutEdgeCount(dstNode->id, memref) > 0) {
1563 createPrivateMemref = false;
1564 }
1565
MLIR Teamd038e342019-03-01 19:50:251566 // Skip if 'srcNode' out edge count on 'memref' > 'maxSrcUserCount'.
1567 if (mdg->getOutEdgeCount(srcNode->id, memref) > maxSrcUserCount)
1568 continue;
1569
River Riddle99b87c92019-03-27 21:02:021570 // Compute an operation list insertion point for the fused loop
MLIR Teama0f3db402019-01-29 17:36:411571 // nest which preserves dependences.
River Riddle99b87c92019-03-27 21:02:021572 Operation *insertPointInst =
MLIR Teama78edcd2019-02-05 14:57:081573 mdg->getFusedLoopNestInsertionPoint(srcNode->id, dstNode->id);
MLIR Teama0f3db402019-01-29 17:36:411574 if (insertPointInst == nullptr)
MLIR Team6892ffb2018-12-20 04:42:551575 continue;
Uday Bondhugula864d9e02019-01-23 17:16:241576
Andy Davis68a8da42019-11-18 19:20:031577 // Compute the innermost common loop depth for dstNode loads/stores.
1578 SmallVector<Operation *, 2> dstOps(dstNode->loads.begin(),
1579 dstNode->loads.end());
1580 dstOps.append(dstNode->stores.begin(), dstNode->stores.end());
1581 unsigned dstLoopDepthTest = getInnermostCommonLoopDepth(dstOps);
1582 // Check the feasibility of fusing src loop nest into dst loop nest
1583 // at loop depths in range [1, dstLoopDepthTest].
1584 // TODO(andydavis) Use slice union computation and union of memref
1585 // read/write regions to cost model and fusion.
1586 bool canFuse = false;
1587 for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
1588 ComputationSliceState sliceUnion;
1589 FusionResult result = mlir::canFuseLoops(
1590 cast<AffineForOp>(srcNode->op), cast<AffineForOp>(dstNode->op),
1591 /*dstLoopDepth=*/i, &sliceUnion);
1592 if (result.value == FusionResult::Success)
1593 canFuse = true;
1594 }
1595
1596 // Skip if fusion is not feasible at all loop depths.
1597 if (!canFuse)
1598 continue;
1599
MLIR Teamd7c82442019-01-30 23:53:411600 // Gather 'dstNode' store ops to 'memref'.
River Riddle99b87c92019-03-27 21:02:021601 SmallVector<Operation *, 2> dstStoreOpInsts;
MLIR Teamd7c82442019-01-30 23:53:411602 for (auto *storeOpInst : dstNode->stores)
Andy Davis2e1187d2019-07-03 17:35:031603 if (cast<AffineStoreOp>(storeOpInst).getMemRef() == memref)
MLIR Teamd7c82442019-01-30 23:53:411604 dstStoreOpInsts.push_back(storeOpInst);
1605
Uday Bondhugulab4a14432019-01-26 00:00:501606 unsigned bestDstLoopDepth;
MLIR Team38c2fe32019-01-14 19:26:251607 mlir::ComputationSliceState sliceState;
MLIR Teama0f3db402019-01-29 17:36:411608 // Check if fusion would be profitable.
Diego Caballero34510552019-10-09 17:36:541609 if (!isFusionProfitable(srcStoreOp, srcStoreOp, dstLoadOpInsts,
1610 dstStoreOpInsts, &sliceState,
Uday Bondhugulace7e59532019-03-08 17:21:521611 &bestDstLoopDepth, maximalFusion))
MLIR Team38c2fe32019-01-14 19:26:251612 continue;
Andy Davis68a8da42019-11-18 19:20:031613
MLIR Team6892ffb2018-12-20 04:42:551614 // Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'.
River Riddle5052bd82019-02-02 00:42:181615 auto sliceLoopNest = mlir::insertBackwardComputationSlice(
Diego Caballero34510552019-10-09 17:36:541616 srcStoreOp, dstLoadOpInsts[0], bestDstLoopDepth, &sliceState);
Chris Lattnerd9b5bc82019-03-25 02:53:051617 if (sliceLoopNest) {
River Riddleaf1abcc2019-03-25 18:13:311618 LLVM_DEBUG(llvm::dbgs() << "\tslice loop nest:\n"
River Riddlef9d91532019-03-27 00:05:091619 << *sliceLoopNest.getOperation() << "\n");
River Riddle5052bd82019-02-02 00:42:181620 // Move 'dstAffineForOp' before 'insertPointInst' if needed.
River Riddleadca3c22019-05-12 00:57:321621 auto dstAffineForOp = cast<AffineForOp>(dstNode->op);
River Riddlef9d91532019-03-27 00:05:091622 if (insertPointInst != dstAffineForOp.getOperation()) {
1623 dstAffineForOp.getOperation()->moveBefore(insertPointInst);
MLIR Teama0f3db402019-01-29 17:36:411624 }
MLIR Teamc4237ae2019-01-18 16:56:271625 // Update edges between 'srcNode' and 'dstNode'.
Andy Davis68a8da42019-11-18 19:20:031626 mdg->updateEdges(srcNode->id, dstNode->id, memref,
1627 createPrivateMemref);
MLIR Teamc4237ae2019-01-18 16:56:271628
1629 // Collect slice loop stats.
1630 LoopNestStateCollector sliceCollector;
River Riddlef9d91532019-03-27 00:05:091631 sliceCollector.collect(sliceLoopNest.getOperation());
MLIR Teamc4237ae2019-01-18 16:56:271632 // Promote single iteration slice loops to single IV value.
River Riddle5052bd82019-02-02 00:42:181633 for (auto forOp : sliceCollector.forOps) {
1634 promoteIfSingleIteration(forOp);
MLIR Team6892ffb2018-12-20 04:42:551635 }
Andy Davis68a8da42019-11-18 19:20:031636 if (createPrivateMemref) {
MLIR Team58aa3832019-02-16 01:12:191637 // Create private memref for 'memref' in 'dstAffineForOp'.
River Riddle99b87c92019-03-27 21:02:021638 SmallVector<Operation *, 4> storesForMemref;
MLIR Team58aa3832019-02-16 01:12:191639 for (auto *storeOpInst : sliceCollector.storeOpInsts) {
Andy Davis2e1187d2019-07-03 17:35:031640 if (cast<AffineStoreOp>(storeOpInst).getMemRef() == memref)
MLIR Team58aa3832019-02-16 01:12:191641 storesForMemref.push_back(storeOpInst);
1642 }
Andy Davis68a8da42019-11-18 19:20:031643 // TODO(andydavis) Use union of memref write regions to compute
1644 // private memref footprint.
River Riddle35807bc2019-12-23 05:59:551645 auto newMemRef = createPrivateMemRef(
MLIR Team58aa3832019-02-16 01:12:191646 dstAffineForOp, storesForMemref[0], bestDstLoopDepth,
1647 fastMemorySpace, localBufSizeThreshold);
1648 visitedMemrefs.insert(newMemRef);
1649 // Create new node in dependence graph for 'newMemRef' alloc op.
1650 unsigned newMemRefNodeId =
River Riddle2bdf33c2020-01-11 16:54:041651 mdg->addNode(newMemRef.getDefiningOp());
MLIR Team58aa3832019-02-16 01:12:191652 // Add edge from 'newMemRef' node to dstNode.
1653 mdg->addEdge(newMemRefNodeId, dstId, newMemRef);
MLIR Teamc4237ae2019-01-18 16:56:271654 }
MLIR Teamc4237ae2019-01-18 16:56:271655
Kazuaki Ishizaki8bfedb32019-10-20 07:11:031656 // Collect dst loop stats after memref privatization transformation.
MLIR Teamc4237ae2019-01-18 16:56:271657 LoopNestStateCollector dstLoopCollector;
River Riddlef9d91532019-03-27 00:05:091658 dstLoopCollector.collect(dstAffineForOp.getOperation());
MLIR Teamc4237ae2019-01-18 16:56:271659
1660 // Add new load ops to current Node load op list 'loads' to
1661 // continue fusing based on new operands.
1662 for (auto *loadOpInst : dstLoopCollector.loadOpInsts) {
River Riddle35807bc2019-12-23 05:59:551663 auto loadMemRef = cast<AffineLoadOp>(loadOpInst).getMemRef();
MLIR Teamc4237ae2019-01-18 16:56:271664 if (visitedMemrefs.count(loadMemRef) == 0)
1665 loads.push_back(loadOpInst);
1666 }
1667
Amit Sabne70a416d2019-04-09 16:17:401668 // Clear and add back loads and stores.
MLIR Teamc4237ae2019-01-18 16:56:271669 mdg->clearNodeLoadAndStores(dstNode->id);
1670 mdg->addToNode(dstId, dstLoopCollector.loadOpInsts,
1671 dstLoopCollector.storeOpInsts);
MLIR Team71495d52019-01-22 21:23:371672 // Remove old src loop nest if it no longer has outgoing dependence
Amit Sabne70a416d2019-04-09 16:17:401673 // edges, and if it does not write to a memref which escapes the
MLIR Team58aa3832019-02-16 01:12:191674 // function. If 'writesToLiveInOrOut' is true, then 'srcNode' has
1675 // been fused into 'dstNode' and write region of 'dstNode' covers
1676 // the write region of 'srcNode', and 'srcNode' has no other users
1677 // so it is safe to remove.
1678 if (writesToLiveInOrOut || mdg->canRemoveNode(srcNode->id)) {
MLIR Teamc4237ae2019-01-18 16:56:271679 mdg->removeNode(srcNode->id);
River Riddle99b87c92019-03-27 21:02:021680 srcNode->op->erase();
MLIR Teama78edcd2019-02-05 14:57:081681 } else {
1682 // Add remaining users of 'oldMemRef' back on the worklist (if not
1683 // already there), as its replacement with a local/private memref
1684 // has reduced dependences on 'oldMemRef' which may have created
1685 // new fusion opportunities.
1686 if (mdg->outEdges.count(srcNode->id) > 0) {
1687 SmallVector<MemRefDependenceGraph::Edge, 2> oldOutEdges =
1688 mdg->outEdges[srcNode->id];
1689 for (auto &outEdge : oldOutEdges) {
1690 if (outEdge.value == memref &&
1691 worklistSet.count(outEdge.id) == 0) {
1692 worklist.push_back(outEdge.id);
1693 worklistSet.insert(outEdge.id);
1694 }
1695 }
1696 }
MLIR Teamc4237ae2019-01-18 16:56:271697 }
MLIR Team3b692302018-12-17 17:57:141698 }
MLIR Team3b692302018-12-17 17:57:141699 }
1700 }
1701 }
MLIR Teamd038e342019-03-01 19:50:251702 }
1703
1704 // Visits each node in the graph, and for each node, attempts to fuse it with
1705 // its sibling nodes (nodes which share a parent, but no dependence edges).
1706 void fuseSiblingNodes() {
1707 init();
1708 while (!worklist.empty()) {
1709 unsigned dstId = worklist.back();
1710 worklist.pop_back();
1711 worklistSet.erase(dstId);
1712
1713 // Skip if this node was removed (fused into another node).
1714 if (mdg->nodes.count(dstId) == 0)
1715 continue;
1716 // Get 'dstNode' into which to attempt fusion.
1717 auto *dstNode = mdg->getNode(dstId);
1718 // Skip if 'dstNode' is not a loop nest.
River Riddled5b60ee82019-05-12 01:59:541719 if (!isa<AffineForOp>(dstNode->op))
MLIR Teamd038e342019-03-01 19:50:251720 continue;
1721 // Attempt to fuse 'dstNode' with its sibling nodes in the graph.
1722 fuseWithSiblingNodes(dstNode);
1723 }
1724 }
1725
1726 // Attempt to fuse 'dstNode' with sibling nodes in the graph.
1727 void fuseWithSiblingNodes(Node *dstNode) {
1728 DenseSet<unsigned> visitedSibNodeIds;
River Riddlee62a6952019-12-23 22:45:011729 std::pair<unsigned, Value> idAndMemref;
MLIR Teamd038e342019-03-01 19:50:251730 while (findSiblingNodeToFuse(dstNode, &visitedSibNodeIds, &idAndMemref)) {
1731 unsigned sibId = idAndMemref.first;
River Riddlee62a6952019-12-23 22:45:011732 Value memref = idAndMemref.second;
MLIR Teamd038e342019-03-01 19:50:251733 // TODO(andydavis) Check that 'sibStoreOpInst' post-dominates all other
1734 // stores to the same memref in 'sibNode' loop nest.
1735 auto *sibNode = mdg->getNode(sibId);
River Riddle99b87c92019-03-27 21:02:021736 // Compute an operation list insertion point for the fused loop
MLIR Teamd038e342019-03-01 19:50:251737 // nest which preserves dependences.
River Riddle99b87c92019-03-27 21:02:021738 assert(sibNode->op->getBlock() == dstNode->op->getBlock());
1739 Operation *insertPointInst =
1740 sibNode->op->isBeforeInBlock(dstNode->op)
MLIR Teamd038e342019-03-01 19:50:251741 ? mdg->getFusedLoopNestInsertionPoint(sibNode->id, dstNode->id)
1742 : mdg->getFusedLoopNestInsertionPoint(dstNode->id, sibNode->id);
1743 if (insertPointInst == nullptr)
1744 continue;
1745
1746 // Check if fusion would be profitable and at what depth.
1747
1748 // Get unique 'sibNode' load op to 'memref'.
River Riddle99b87c92019-03-27 21:02:021749 SmallVector<Operation *, 2> sibLoadOpInsts;
MLIR Teamd038e342019-03-01 19:50:251750 sibNode->getLoadOpsForMemref(memref, &sibLoadOpInsts);
1751 // Currently findSiblingNodeToFuse searches for siblings with one load.
1752 assert(sibLoadOpInsts.size() == 1);
River Riddle99b87c92019-03-27 21:02:021753 Operation *sibLoadOpInst = sibLoadOpInsts[0];
MLIR Teamd038e342019-03-01 19:50:251754 assert(!sibNode->stores.empty());
1755 // TODO(andydavis) Choose the store which postdominates all other stores.
1756 auto *sibStoreOpInst = sibNode->stores.back();
1757
1758 // Gather 'dstNode' load ops to 'memref'.
River Riddle99b87c92019-03-27 21:02:021759 SmallVector<Operation *, 2> dstLoadOpInsts;
MLIR Teamd038e342019-03-01 19:50:251760 dstNode->getLoadOpsForMemref(memref, &dstLoadOpInsts);
1761
1762 // Gather 'dstNode' store ops to 'memref'.
River Riddle99b87c92019-03-27 21:02:021763 SmallVector<Operation *, 2> dstStoreOpInsts;
MLIR Teamd038e342019-03-01 19:50:251764 dstNode->getStoreOpsForMemref(memref, &dstStoreOpInsts);
1765
1766 unsigned bestDstLoopDepth;
1767 mlir::ComputationSliceState sliceState;
1768
1769 // Check if fusion would be profitable.
1770 if (!isFusionProfitable(sibLoadOpInst, sibStoreOpInst, dstLoadOpInsts,
Uday Bondhugulace7e59532019-03-08 17:21:521771 dstStoreOpInsts, &sliceState, &bestDstLoopDepth,
1772 maximalFusion))
MLIR Teamd038e342019-03-01 19:50:251773 continue;
1774
1775 // Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'.
1776 auto sliceLoopNest = mlir::insertBackwardComputationSlice(
1777 sibLoadOpInst, dstLoadOpInsts[0], bestDstLoopDepth, &sliceState);
1778 if (sliceLoopNest != nullptr) {
River Riddleadca3c22019-05-12 00:57:321779 auto dstForInst = cast<AffineForOp>(dstNode->op);
River Riddle99b87c92019-03-27 21:02:021780 // Update operation position of fused loop nest (if needed).
River Riddlef9d91532019-03-27 00:05:091781 if (insertPointInst != dstForInst.getOperation()) {
1782 dstForInst.getOperation()->moveBefore(insertPointInst);
MLIR Teamd038e342019-03-01 19:50:251783 }
1784 // Update data dependence graph state post fusion.
1785 updateStateAfterSiblingFusion(sliceLoopNest, sibNode, dstNode);
1786 }
1787 }
1788 }
1789
MLIR Team9d30b362019-03-29 15:06:251790 // Searches function argument uses and the graph from 'dstNode' looking for a
1791 // fusion candidate sibling node which shares no dependences with 'dstNode'
1792 // but which loads from the same memref. Returns true and sets
1793 // 'idAndMemrefToFuse' on success. Returns false otherwise.
MLIR Teamd038e342019-03-01 19:50:251794 bool findSiblingNodeToFuse(Node *dstNode,
1795 DenseSet<unsigned> *visitedSibNodeIds,
River Riddlee62a6952019-12-23 22:45:011796 std::pair<unsigned, Value> *idAndMemrefToFuse) {
MLIR Team9d30b362019-03-29 15:06:251797 // Returns true if 'sibNode' can be fused with 'dstNode' for input reuse
1798 // on 'memref'.
River Riddlee62a6952019-12-23 22:45:011799 auto canFuseWithSibNode = [&](Node *sibNode, Value memref) {
MLIR Team9d30b362019-03-29 15:06:251800 // Skip if 'outEdge' is not a read-after-write dependence.
1801 // TODO(andydavis) Remove restrict to single load op restriction.
1802 if (sibNode->getLoadOpCount(memref) != 1)
1803 return false;
1804 // Skip if there exists a path of dependent edges between
1805 // 'sibNode' and 'dstNode'.
1806 if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
1807 mdg->hasDependencePath(dstNode->id, sibNode->id))
1808 return false;
1809 // Skip sib node if it loads to (and stores from) the same memref on
1810 // which it also has an input dependence edge.
River Riddlee62a6952019-12-23 22:45:011811 DenseSet<Value> loadAndStoreMemrefSet;
MLIR Team9d30b362019-03-29 15:06:251812 sibNode->getLoadAndStoreMemrefSet(&loadAndStoreMemrefSet);
River Riddlee62a6952019-12-23 22:45:011813 if (llvm::any_of(loadAndStoreMemrefSet, [=](Value memref) {
MLIR Team9d30b362019-03-29 15:06:251814 return mdg->getIncomingMemRefAccesses(sibNode->id, memref) > 0;
1815 }))
1816 return false;
1817
1818 // Check that all stores are to the same memref.
River Riddlee62a6952019-12-23 22:45:011819 DenseSet<Value> storeMemrefs;
MLIR Team9d30b362019-03-29 15:06:251820 for (auto *storeOpInst : sibNode->stores) {
Andy Davis2e1187d2019-07-03 17:35:031821 storeMemrefs.insert(cast<AffineStoreOp>(storeOpInst).getMemRef());
MLIR Team9d30b362019-03-29 15:06:251822 }
1823 if (storeMemrefs.size() != 1)
1824 return false;
1825 return true;
1826 };
1827
1828 // Search for siblings which load the same memref function argument.
River Riddlece502af2019-07-08 18:20:261829 auto fn = dstNode->op->getParentOfType<FuncOp>();
River Riddle54cd6a72019-07-01 17:29:091830 for (unsigned i = 0, e = fn.getNumArguments(); i != e; ++i) {
River Riddle2bdf33c2020-01-11 16:54:041831 for (auto *user : fn.getArgument(i).getUsers()) {
Andy Davis2e1187d2019-07-03 17:35:031832 if (auto loadOp = dyn_cast<AffineLoadOp>(user)) {
MLIR Team9d30b362019-03-29 15:06:251833 // Gather loops surrounding 'use'.
1834 SmallVector<AffineForOp, 4> loops;
River Riddle8780d8d2019-05-18 18:09:071835 getLoopIVs(*user, &loops);
MLIR Team9d30b362019-03-29 15:06:251836 // Skip 'use' if it is not within a loop nest.
1837 if (loops.empty())
1838 continue;
1839 Node *sibNode = mdg->getForOpNode(loops[0]);
1840 assert(sibNode != nullptr);
1841 // Skip 'use' if it not a sibling to 'dstNode'.
1842 if (sibNode->id == dstNode->id)
1843 continue;
1844 // Skip 'use' if it has been visited.
1845 if (visitedSibNodeIds->count(sibNode->id) > 0)
1846 continue;
1847 // Skip 'use' if it does not load from the same memref as 'dstNode'.
River Riddle35807bc2019-12-23 05:59:551848 auto memref = loadOp.getMemRef();
MLIR Team9d30b362019-03-29 15:06:251849 if (dstNode->getLoadOpCount(memref) == 0)
1850 continue;
1851 // Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
1852 if (canFuseWithSibNode(sibNode, memref)) {
1853 visitedSibNodeIds->insert(sibNode->id);
1854 idAndMemrefToFuse->first = sibNode->id;
1855 idAndMemrefToFuse->second = memref;
1856 return true;
1857 }
1858 }
1859 }
1860 }
1861
1862 // Search for siblings by following edges through an intermediate src node.
MLIR Teamd038e342019-03-01 19:50:251863 // Collect candidate 'dstNode' input edges in 'inEdges'.
1864 SmallVector<MemRefDependenceGraph::Edge, 2> inEdges;
1865 mdg->forEachMemRefInputEdge(
1866 dstNode->id, [&](MemRefDependenceGraph::Edge inEdge) {
1867 // Add 'inEdge' if it is a read-after-write dependence.
1868 if (dstNode->getLoadOpCount(inEdge.value) > 0 &&
1869 mdg->getNode(inEdge.id)->getStoreOpCount(inEdge.value) > 0)
1870 inEdges.push_back(inEdge);
1871 });
1872
1873 // Search for sibling nodes to fuse by visiting output edges from each input
1874 // edge in 'inEdges'.
1875 for (auto &inEdge : inEdges) {
1876 // Collect candidate output edges from each node 'inEdge.id' in 'inEdges'.
1877 SmallVector<MemRefDependenceGraph::Edge, 2> outEdges;
1878 mdg->forEachMemRefOutputEdge(
1879 inEdge.id, [&](MemRefDependenceGraph::Edge outEdge) {
1880 unsigned sibNodeId = outEdge.id;
1881 if (visitedSibNodeIds->count(sibNodeId) > 0)
1882 return;
1883 // Skip output edge if not a sibling using the same memref.
1884 if (outEdge.id == dstNode->id || outEdge.value != inEdge.value)
1885 return;
1886 auto *sibNode = mdg->getNode(sibNodeId);
River Riddled5b60ee82019-05-12 01:59:541887 if (!isa<AffineForOp>(sibNode->op))
MLIR Teamd038e342019-03-01 19:50:251888 return;
MLIR Team9d30b362019-03-29 15:06:251889 // Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
1890 if (canFuseWithSibNode(sibNode, outEdge.value)) {
1891 // Add candidate 'outEdge' to sibling node.
1892 outEdges.push_back(outEdge);
MLIR Teamd038e342019-03-01 19:50:251893 }
MLIR Teamd038e342019-03-01 19:50:251894 });
1895
1896 // Add first candidate if any were returned.
1897 if (!outEdges.empty()) {
1898 visitedSibNodeIds->insert(outEdges[0].id);
1899 idAndMemrefToFuse->first = outEdges[0].id;
1900 idAndMemrefToFuse->second = outEdges[0].value;
1901 return true;
1902 }
1903 }
1904 return false;
1905 }
1906
Chris Lattnerd9b5bc82019-03-25 02:53:051907 void updateStateAfterSiblingFusion(AffineForOp sliceLoopNest, Node *sibNode,
1908 Node *dstNode) {
MLIR Teamd038e342019-03-01 19:50:251909 // Update 'sibNode' and 'dstNode' input/output edges to reflect fusion.
1910 mdg->updateEdges(sibNode->id, dstNode->id);
1911
1912 // Collect slice loop stats.
1913 LoopNestStateCollector sliceCollector;
River Riddlef9d91532019-03-27 00:05:091914 sliceCollector.collect(sliceLoopNest.getOperation());
MLIR Teamd038e342019-03-01 19:50:251915 // Promote single iteration slice loops to single IV value.
1916 for (auto forOp : sliceCollector.forOps) {
1917 promoteIfSingleIteration(forOp);
1918 }
1919
Kazuaki Ishizaki8bfedb32019-10-20 07:11:031920 // Collect dst loop stats after memref privatization transformation.
River Riddleadca3c22019-05-12 00:57:321921 auto dstForInst = cast<AffineForOp>(dstNode->op);
MLIR Teamd038e342019-03-01 19:50:251922 LoopNestStateCollector dstLoopCollector;
River Riddlef9d91532019-03-27 00:05:091923 dstLoopCollector.collect(dstForInst.getOperation());
MLIR Teamd038e342019-03-01 19:50:251924 // Clear and add back loads and stores
1925 mdg->clearNodeLoadAndStores(dstNode->id);
1926 mdg->addToNode(dstNode->id, dstLoopCollector.loadOpInsts,
1927 dstLoopCollector.storeOpInsts);
1928 // Remove old sibling loop nest if it no longer has outgoing dependence
1929 // edges, and it does not write to a memref which escapes the
1930 // function.
1931 if (mdg->getOutEdgeCount(sibNode->id) == 0) {
1932 mdg->removeNode(sibNode->id);
River Riddleadca3c22019-05-12 00:57:321933 sibNode->op->erase();
MLIR Teamd038e342019-03-01 19:50:251934 }
1935 }
1936
1937 // Clean up any allocs with no users.
1938 void eraseUnusedMemRefAllocations() {
MLIR Teamc4237ae2019-01-18 16:56:271939 for (auto &pair : mdg->memrefEdgeCount) {
1940 if (pair.second > 0)
1941 continue;
River Riddle35807bc2019-12-23 05:59:551942 auto memref = pair.first;
River Riddle99b87c92019-03-27 21:02:021943 // Skip if there exist other uses (return operation or function calls).
River Riddle2bdf33c2020-01-11 16:54:041944 if (!memref.use_empty())
MLIR Team71495d52019-01-22 21:23:371945 continue;
MLIR Teamc4237ae2019-01-18 16:56:271946 // Use list expected to match the dep graph info.
River Riddle2bdf33c2020-01-11 16:54:041947 auto *op = memref.getDefiningOp();
River Riddle1423acc2019-04-23 21:38:261948 if (isa_and_nonnull<AllocOp>(op))
River Riddle99b87c92019-03-27 21:02:021949 op->erase();
MLIR Teamc4237ae2019-01-18 16:56:271950 }
MLIR Teamf28e4df2018-11-01 14:26:001951 }
MLIR Team3b692302018-12-17 17:57:141952};
1953
1954} // end anonymous namespace
MLIR Teamf28e4df2018-11-01 14:26:001955
River Riddleed5fe202019-02-28 22:50:421956void LoopFusion::runOnFunction() {
Uday Bondhugulad4b3ff12019-02-27 00:10:191957 // Override if a command line argument was provided.
Uday Bondhugula8be26272019-02-02 01:06:221958 if (clFusionFastMemorySpace.getNumOccurrences() > 0) {
1959 fastMemorySpace = clFusionFastMemorySpace.getValue();
1960 }
1961
Uday Bondhugulad4b3ff12019-02-27 00:10:191962 // Override if a command line argument was provided.
1963 if (clFusionLocalBufThreshold.getNumOccurrences() > 0) {
1964 localBufSizeThreshold = clFusionLocalBufThreshold * 1024;
1965 }
1966
Uday Bondhugulace7e59532019-03-08 17:21:521967 if (clMaximalLoopFusion.getNumOccurrences() > 0)
1968 maximalFusion = clMaximalLoopFusion;
1969
MLIR Team6892ffb2018-12-20 04:42:551970 MemRefDependenceGraph g;
Uday Bondhugula02af8c22019-03-05 23:05:341971 if (g.init(getFunction()))
Uday Bondhugulace7e59532019-03-08 17:21:521972 GreedyFusion(&g, localBufSizeThreshold, fastMemorySpace, maximalFusion)
1973 .run();
MLIR Teamf28e4df2018-11-01 14:26:001974}