From 6892ffb8965c2849565fedfbbf60f05e475c9858 Mon Sep 17 00:00:00 2001 From: MLIR Team Date: Wed, 19 Dec 2018 20:42:55 -0800 Subject: [PATCH] Improve loop fusion algorithm by using a memref dependence graph. Fixed TODO for reduction fusion unit test. PiperOrigin-RevId: 226277226 --- mlir/lib/Transforms/LoopFusion.cpp | 592 +++++++++++++++----------- mlir/test/Transforms/loop-fusion.mlir | 21 +- 2 files changed, 350 insertions(+), 263 deletions(-) diff --git a/mlir/lib/Transforms/LoopFusion.cpp b/mlir/lib/Transforms/LoopFusion.cpp index 521fca8979fa..6393fa6069d8 100644 --- a/mlir/lib/Transforms/LoopFusion.cpp +++ b/mlir/lib/Transforms/LoopFusion.cpp @@ -130,24 +130,270 @@ public: } }; -// GreedyFusionPolicy greedily fuses loop nests which have a producer/consumer +// MemRefDependenceGraph is a graph data structure where graph nodes are +// top-level statements in an MLFunction which contain load/store ops, and edges +// are memref dependences between the nodes. +// TODO(andydavis) Add a depth parameter to dependence graph construction. +struct MemRefDependenceGraph { +public: + // Node represents a node in the graph. A Node is either an entire loop nest + // rooted at the top level which contains loads/stores, or a top level + // load/store. + struct Node { + // The unique identifier of this node in the graph. + unsigned id; + // The top-level statment which is (or contains) loads/stores. + Statement *stmt; + // List of load op stmts. + SmallVector loads; + // List of store op stmts. + SmallVector stores; + Node(unsigned id, Statement *stmt) : id(id), stmt(stmt) {} + + // Returns the load op count for 'memref'. + unsigned getLoadOpCount(MLValue *memref) { + unsigned loadOpCount = 0; + for (auto *loadOpStmt : loads) { + if (memref == cast(loadOpStmt->cast()->getMemRef())) + ++loadOpCount; + } + return loadOpCount; + } + + // Returns the store op count for 'memref'. + unsigned getStoreOpCount(MLValue *memref) { + unsigned storeOpCount = 0; + for (auto *storeOpStmt : stores) { + if (memref == cast(storeOpStmt->cast()->getMemRef())) + ++storeOpCount; + } + return storeOpCount; + } + }; + + // Edge represents a memref data dependece between nodes in the graph. + struct Edge { + // The id of the node at the other end of the edge. + unsigned id; + // The memref on which this edge represents a dependence. + MLValue *memref; + }; + + // Map from node id to Node. + DenseMap nodes; + // Map from node id to list of input edges. + DenseMap> inEdges; + // Map from node id to list of output edges. + DenseMap> outEdges; + + MemRefDependenceGraph() {} + + // Initializes the dependence graph based on operations in 'f'. + // Returns true on success, false otherwise. + bool init(MLFunction *f); + + // Returns the graph node for 'id'. + Node *getNode(unsigned id) { + auto it = nodes.find(id); + assert(it != nodes.end()); + return &it->second; + } + + // Adds an edge from node 'srcId' to node 'dstId' for 'memref'. + void addEdge(unsigned srcId, unsigned dstId, MLValue *memref) { + outEdges[srcId].push_back({dstId, memref}); + inEdges[dstId].push_back({srcId, memref}); + } + + // Removes an edge from node 'srcId' to node 'dstId' for 'memref'. + void removeEdge(unsigned srcId, unsigned dstId, MLValue *memref) { + assert(inEdges.count(dstId) > 0); + assert(outEdges.count(srcId) > 0); + // Remove 'srcId' from 'inEdges[dstId]'. + for (auto it = inEdges[dstId].begin(); it != inEdges[dstId].end(); ++it) { + if ((*it).id == srcId && (*it).memref == memref) { + inEdges[dstId].erase(it); + break; + } + } + // Remove 'dstId' from 'outEdges[srcId]'. + for (auto it = outEdges[srcId].begin(); it != outEdges[srcId].end(); ++it) { + if ((*it).id == dstId && (*it).memref == memref) { + outEdges[srcId].erase(it); + break; + } + } + } + + // Returns the input edge count for node 'id' and 'memref'. + unsigned getInEdgeCount(unsigned id, MLValue *memref) { + unsigned inEdgeCount = 0; + if (inEdges.count(id) > 0) + for (auto &inEdge : inEdges[id]) + if (inEdge.memref == memref) + ++inEdgeCount; + return inEdgeCount; + } + + // Returns the output edge count for node 'id' and 'memref'. + unsigned getOutEdgeCount(unsigned id, MLValue *memref) { + unsigned outEdgeCount = 0; + if (outEdges.count(id) > 0) + for (auto &outEdge : outEdges[id]) + if (outEdge.memref == memref) + ++outEdgeCount; + return outEdgeCount; + } + + // Returns the min node id of all output edges from node 'id'. + unsigned getMinOutEdgeNodeId(unsigned id) { + unsigned minId = std::numeric_limits::max(); + if (outEdges.count(id) > 0) + for (auto &outEdge : outEdges[id]) + minId = std::min(minId, outEdge.id); + return minId; + } + + // Updates edge mappings from node 'srcId' to node 'dstId' and removes + // state associated with node 'srcId'. + void updateEdgesAndRemoveSrcNode(unsigned srcId, unsigned dstId) { + // For each edge in 'inEdges[srcId]': add new edge remaping to 'dstId'. + if (inEdges.count(srcId) > 0) { + SmallVector oldInEdges = inEdges[srcId]; + for (auto &inEdge : oldInEdges) { + // Remove edge from 'inEdge.id' to 'srcId'. + removeEdge(inEdge.id, srcId, inEdge.memref); + // Add edge from 'inEdge.id' to 'dstId'. + addEdge(inEdge.id, dstId, inEdge.memref); + } + } + // For each edge in 'outEdges[srcId]': add new edge remaping to 'dstId'. + if (outEdges.count(srcId) > 0) { + SmallVector oldOutEdges = outEdges[srcId]; + for (auto &outEdge : oldOutEdges) { + // Remove edge from 'srcId' to 'outEdge.id'. + removeEdge(srcId, outEdge.id, outEdge.memref); + // Add edge from 'dstId' to 'outEdge.id' (if 'outEdge.id' != 'dstId'). + if (outEdge.id != dstId) + addEdge(dstId, outEdge.id, outEdge.memref); + } + } + // Remove 'srcId' from graph state. + inEdges.erase(srcId); + outEdges.erase(srcId); + nodes.erase(srcId); + } + + // Adds ops in 'loads' and 'stores' to node at 'id'. + void addToNode(unsigned id, const SmallVectorImpl &loads, + const SmallVectorImpl &stores) { + Node *node = getNode(id); + for (auto *loadOpStmt : loads) + node->loads.push_back(loadOpStmt); + for (auto *storeOpStmt : stores) + node->stores.push_back(storeOpStmt); + } + + void print(raw_ostream &os) const { + os << "\nMemRefDependenceGraph\n"; + os << "\nNodes:\n"; + for (auto &idAndNode : nodes) { + os << "Node: " << idAndNode.first << "\n"; + auto it = inEdges.find(idAndNode.first); + if (it != inEdges.end()) { + for (const auto &e : it->second) + os << " InEdge: " << e.id << " " << e.memref << "\n"; + } + it = outEdges.find(idAndNode.first); + if (it != outEdges.end()) { + for (const auto &e : it->second) + os << " OutEdge: " << e.id << " " << e.memref << "\n"; + } + } + } + void dump() const { print(llvm::errs()); } +}; + +// Intializes the data dependence graph by walking statements in 'f'. +// Assigns each node in the graph a node id based on program order in 'f'. +// TODO(andydavis) Add support for taking a StmtBlock arg to construct the +// dependence graph at a different depth. +bool MemRefDependenceGraph::init(MLFunction *f) { + unsigned id = 0; + DenseMap> memrefAccesses; + for (auto &stmt : *f) { + if (auto *forStmt = dyn_cast(&stmt)) { + // Create graph node 'id' to represent top-level 'forStmt' and record + // all loads and store accesses it contains. + LoopNestStateCollector collector; + collector.walkForStmt(forStmt); + // Return false if IfStmts are found (not currently supported). + if (collector.hasIfStmt) + return false; + Node node(id++, &stmt); + for (auto *opStmt : collector.loadOpStmts) { + node.loads.push_back(opStmt); + auto *memref = cast(opStmt->cast()->getMemRef()); + memrefAccesses[memref].insert(node.id); + } + for (auto *opStmt : collector.storeOpStmts) { + node.stores.push_back(opStmt); + auto *memref = cast(opStmt->cast()->getMemRef()); + memrefAccesses[memref].insert(node.id); + } + nodes.insert({node.id, node}); + } + if (auto *opStmt = dyn_cast(&stmt)) { + if (auto loadOp = opStmt->dyn_cast()) { + // Create graph node for top-level load op. + Node node(id++, &stmt); + node.loads.push_back(opStmt); + auto *memref = cast(opStmt->cast()->getMemRef()); + memrefAccesses[memref].insert(node.id); + nodes.insert({node.id, node}); + } + if (auto storeOp = opStmt->dyn_cast()) { + // Create graph node for top-level store op. + Node node(id++, &stmt); + node.stores.push_back(opStmt); + auto *memref = cast(opStmt->cast()->getMemRef()); + memrefAccesses[memref].insert(node.id); + nodes.insert({node.id, node}); + } + } + // Return false if IfStmts are found (not currently supported). + if (isa(&stmt)) + return false; + } + + // Walk memref access lists and add graph edges between dependent nodes. + for (auto &memrefAndList : memrefAccesses) { + unsigned n = memrefAndList.second.size(); + for (unsigned i = 0; i < n; ++i) { + unsigned srcId = memrefAndList.second[i]; + bool srcHasStore = + getNode(srcId)->getStoreOpCount(memrefAndList.first) > 0; + for (unsigned j = i + 1; j < n; ++j) { + unsigned dstId = memrefAndList.second[j]; + bool dstHasStore = + getNode(dstId)->getStoreOpCount(memrefAndList.first) > 0; + if (srcHasStore || dstHasStore) + addEdge(srcId, dstId, memrefAndList.first); + } + } + } + return true; +} + +// GreedyFusion greedily fuses loop nests which have a producer/consumer // relationship on a memref, with the goal of improving locality. Currently, // this the producer/consumer relationship is required to be unique in the // MLFunction (there are TODOs to relax this constraint in the future). // // The steps of the algorithm are as follows: // -// *) Initialize. While visiting each statement in the MLFunction do: -// *) Assign each top-level ForStmt a 'position' which is its initial -// position in the MLFunction's StmtBlock at the start of the pass. -// *) Gather memref load/store state aggregated by top-level statement. For -// example, all loads and stores contained in a loop nest are aggregated -// under the loop nest's top-level ForStmt. -// *) Add each top-level ForStmt to a worklist. -// -// *) Run. The algorithm processes the worklist with the following steps: -// *) The worklist is processed in reverse order (starting from the last -// top-level ForStmt in the MLFunction). +// *) A worklist is initialized with node ids from the dependence graph. +// *) For each node id in the worklist: // *) Pop a ForStmt of the worklist. This 'dstForStmt' will be a candidate // destination ForStmt into which fusion will be attempted. // *) Add each LoadOp currently in 'dstForStmt' into list 'dstLoadOps'. @@ -157,7 +403,7 @@ public: // *) Check if dependences would be violated by the fusion. For example, // the src loop nest may load from memrefs which are different than // the producer-consumer memref between src and dest loop nests. -// *) Get a computation slice of 'srcLoopNest', which adjust its loop +// *) Get a computation slice of 'srcLoopNest', which adjusts its loop // bounds to be functions of 'dstLoopNest' IVs and symbols. // *) Fuse the 'srcLoopNest' computation slice into the 'dstLoopNest', // just before the dst load op user. @@ -168,268 +414,112 @@ public: // // Given a graph where top-level statements are vertices in the set 'V' and // edges in the set 'E' are dependences between vertices, this algorithm -// takes O(V) time for initialization, and has runtime O(V * E). -// TODO(andydavis) Reduce this time complexity to O(V + E). +// takes O(V) time for initialization, and has runtime O(V + E). // -// This greedy algorithm is not 'maximally' but there is a TODO to fix this. +// This greedy algorithm is not 'maximal' due to the current restriction of +// fusing along single producer consumer edges, but there is a TODO to fix this. // // TODO(andydavis) Experiment with other fusion policies. -struct GreedyFusionPolicy { - // Convenience wrapper with information about 'stmt' ready to access. - struct StmtInfo { - Statement *stmt; - bool isOrContainsIfStmt = false; - }; - // The worklist of top-level loop nest positions. +// TODO(andydavis) Add support for fusing for input reuse (perhaps by +// constructing a graph with edges which represent loads from the same memref +// in two different loop nestst. +struct GreedyFusion { +public: + MemRefDependenceGraph *mdg; SmallVector worklist; - // Mapping from top-level position to StmtInfo. - DenseMap posToStmtInfo; - // Mapping from memref MLValue to set of top-level positions of loop nests - // which contain load ops on that memref. - DenseMap> memrefToLoadPosSet; - // Mapping from memref MLValue to set of top-level positions of loop nests - // which contain store ops on that memref. - DenseMap> memrefToStorePosSet; - // Mapping from top-level loop nest to the set of load ops it contains. - DenseMap> forStmtToLoadOps; - // Mapping from top-level loop nest to the set of store ops it contains. - DenseMap> forStmtToStoreOps; - GreedyFusionPolicy(MLFunction *f) { init(f); } + GreedyFusion(MemRefDependenceGraph *mdg) : mdg(mdg) { + // Initialize worklist with nodes from 'mdg'. + worklist.resize(mdg->nodes.size()); + std::iota(worklist.begin(), worklist.end(), 0); + } void run() { - if (hasIfStmts()) - return; - while (!worklist.empty()) { - // Pop the position of a loop nest into which fusion will be attempted. - unsigned dstPos = worklist.back(); + unsigned dstId = worklist.back(); worklist.pop_back(); - // Skip if 'dstPos' is not tracked (was fused into another loop nest). - if (posToStmtInfo.count(dstPos) == 0) + // Skip if this node was removed (fused into another node). + if (mdg->nodes.count(dstId) == 0) continue; - // Get the top-level ForStmt at 'dstPos'. - auto *dstForStmt = getForStmtAtPos(dstPos); - // Skip if this ForStmt contains no load ops. - if (forStmtToLoadOps.count(dstForStmt) == 0) + // Get 'dstNode' into which to attempt fusion. + auto *dstNode = mdg->getNode(dstId); + // Skip if 'dstNode' is not a loop nest. + if (!isa(dstNode->stmt)) continue; - // Greedy Policy: iterate through load ops in 'dstForStmt', greedily - // fusing in src loop nests which have a single store op on the same - // memref, until a fixed point is reached where there is nothing left to - // fuse. - SetVector dstLoadOps = forStmtToLoadOps[dstForStmt]; - while (!dstLoadOps.empty()) { - auto *dstLoadOpStmt = dstLoadOps.pop_back_val(); - - auto dstLoadOp = dstLoadOpStmt->cast(); - auto *memref = cast(dstLoadOp->getMemRef()); - // Skip if not single src store / dst load pair on 'memref'. - if (memrefToLoadPosSet[memref].size() != 1 || - memrefToStorePosSet[memref].size() != 1) + SmallVector loads = dstNode->loads; + while (!loads.empty()) { + auto *dstLoadOpStmt = loads.pop_back_val(); + auto *memref = + cast(dstLoadOpStmt->cast()->getMemRef()); + // Skip 'dstLoadOpStmt' if multiple loads to 'memref' in 'dstNode'. + if (dstNode->getLoadOpCount(memref) != 1) continue; - unsigned srcPos = *memrefToStorePosSet[memref].begin(); - if (srcPos >= dstPos) + // Skip if no input edges along which to fuse. + if (mdg->inEdges.count(dstId) == 0) continue; - auto *srcForStmt = getForStmtAtPos(srcPos); - // Skip if 'srcForStmt' has more than one store op. - if (forStmtToStoreOps[srcForStmt].size() > 1) - continue; - // Skip if fusion would violated dependences between 'memref' access - // for loop nests between 'srcPos' and 'dstPos': - // For each src load op: check for store ops in range (srcPos, dstPos). - // For each src store op: check for load ops in range (srcPos, dstPos). - if (moveWouldViolateDependences(srcPos, dstPos)) - continue; - auto *srcStoreOpStmt = forStmtToStoreOps[srcForStmt].front(); - // Build fusion candidate out of 'srcStoreOpStmt' and 'dstLoadOpStmt'. - FusionCandidate candidate = - buildFusionCandidate(srcStoreOpStmt, dstLoadOpStmt); - // Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'. - auto *sliceLoopNest = mlir::insertBackwardComputationSlice( - &candidate.srcAccess, &candidate.dstAccess); - if (sliceLoopNest != nullptr) { - // Remove 'srcPos' mappings from 'state'. - moveAccessesAndRemovePos(srcPos, dstPos); - // Record all load/store accesses in 'sliceLoopNest' at 'dstPos'. - LoopNestStateCollector collector; - collector.walkForStmt(sliceLoopNest); - // Record mappings for loads and stores from 'collector'. - for (auto *opStmt : collector.loadOpStmts) { - addLoadOpStmtAt(dstPos, opStmt, dstForStmt); - // Add newly fused load ops to 'dstLoadOps' to be considered for - // fusion on subsequent iterations. - dstLoadOps.insert(opStmt); + // Iterate through in edges for 'dstId'. + for (auto &srcEdge : mdg->inEdges[dstId]) { + // Skip 'srcEdge' if not for 'memref'. + if (srcEdge.memref != memref) + continue; + auto *srcNode = mdg->getNode(srcEdge.id); + // Skip if 'srcNode' is not a loop nest. + if (!isa(srcNode->stmt)) + continue; + // Skip if 'srcNode' has more than one store to 'memref'. + if (srcNode->getStoreOpCount(memref) != 1) + continue; + // Skip 'srcNode' if it has out edges on 'memref' other than 'dstId'. + if (mdg->getOutEdgeCount(srcNode->id, memref) != 1) + continue; + // Skip 'srcNode' if it has in dependence edges. NOTE: This is overly + // TODO(andydavis) Track dependence type with edges, and just check + // for WAW dependence edge here. + if (mdg->getInEdgeCount(srcNode->id, memref) != 0) + continue; + // Skip if 'srcNode' has out edges to other memrefs after 'dstId'. + if (mdg->getMinOutEdgeNodeId(srcNode->id) != dstId) + continue; + // Get unique 'srcNode' store op. + auto *srcStoreOpStmt = srcNode->stores.front(); + // Build fusion candidate out of 'srcStoreOpStmt' and 'dstLoadOpStmt'. + FusionCandidate candidate = + buildFusionCandidate(srcStoreOpStmt, dstLoadOpStmt); + // Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'. + auto *sliceLoopNest = mlir::insertBackwardComputationSlice( + &candidate.srcAccess, &candidate.dstAccess); + if (sliceLoopNest != nullptr) { + // Remove edges between 'srcNode' and 'dstNode' and remove 'srcNode' + mdg->updateEdgesAndRemoveSrcNode(srcNode->id, dstNode->id); + // Record all load/store accesses in 'sliceLoopNest' at 'dstPos'. + LoopNestStateCollector collector; + collector.walkForStmt(sliceLoopNest); + mdg->addToNode(dstId, collector.loadOpStmts, + collector.storeOpStmts); + // Add new load ops to current Node load op list 'loads' to + // continue fusing based on new operands. + for (auto *loadOpStmt : collector.loadOpStmts) + loads.push_back(loadOpStmt); + // Promote single iteration loops to single IV value. + for (auto *forStmt : collector.forStmts) { + promoteIfSingleIteration(forStmt); + } + // Remove old src loop nest. + cast(srcNode->stmt)->erase(); } - for (auto *opStmt : collector.storeOpStmts) { - addStoreOpStmtAt(dstPos, opStmt, dstForStmt); - } - for (auto *forStmt : collector.forStmts) { - promoteIfSingleIteration(forStmt); - } - // Remove old src loop nest. - srcForStmt->erase(); } } } } - - // Walk MLFunction 'f' assigning each top-level statement a position, and - // gathering state on load and store ops. - void init(MLFunction *f) { - unsigned pos = 0; - for (auto &stmt : *f) { - if (auto *forStmt = dyn_cast(&stmt)) { - // Record all loads and store accesses in 'forStmt' at 'pos'. - LoopNestStateCollector collector; - collector.walkForStmt(forStmt); - // Create StmtInfo for 'forStmt' for top-level loop nests. - addStmtInfoAt(pos, forStmt, collector.hasIfStmt); - // Record mappings for loads and stores from 'collector'. - for (auto *opStmt : collector.loadOpStmts) { - addLoadOpStmtAt(pos, opStmt, forStmt); - } - for (auto *opStmt : collector.storeOpStmts) { - addStoreOpStmtAt(pos, opStmt, forStmt); - } - // Add 'pos' associated with 'forStmt' to worklist. - worklist.push_back(pos); - } - if (auto *opStmt = dyn_cast(&stmt)) { - if (auto loadOp = opStmt->dyn_cast()) { - // Create StmtInfo for top-level load op. - addStmtInfoAt(pos, &stmt, /*hasIfStmt=*/false); - addLoadOpStmtAt(pos, opStmt, /*containingForStmt=*/nullptr); - } - if (auto storeOp = opStmt->dyn_cast()) { - // Create StmtInfo for top-level store op. - addStmtInfoAt(pos, &stmt, /*hasIfStmt=*/false); - addStoreOpStmtAt(pos, opStmt, /*containingForStmt=*/nullptr); - } - } - if (auto *ifStmt = dyn_cast(&stmt)) { - addStmtInfoAt(pos, &stmt, /*hasIfStmt=*/true); - } - ++pos; - } - } - - // Check if fusing loop nest at 'srcPos' into the loop nest at 'dstPos' - // would violated any dependences w.r.t other loop nests in that range. - bool moveWouldViolateDependences(unsigned srcPos, unsigned dstPos) { - // Lookup src ForStmt at 'srcPos'. - auto *srcForStmt = getForStmtAtPos(srcPos); - // For each src load op: check for store ops in range (srcPos, dstPos). - if (forStmtToLoadOps.count(srcForStmt) > 0) { - for (auto *opStmt : forStmtToLoadOps[srcForStmt]) { - auto loadOp = opStmt->cast(); - auto *memref = cast(loadOp->getMemRef()); - for (unsigned pos = srcPos + 1; pos < dstPos; ++pos) { - if (memrefToStorePosSet.count(memref) > 0 && - memrefToStorePosSet[memref].count(pos) > 0) - return true; - } - } - } - // For each src store op: check for load ops in range (srcPos, dstPos). - if (forStmtToStoreOps.count(srcForStmt) > 0) { - for (auto *opStmt : forStmtToStoreOps[srcForStmt]) { - auto storeOp = opStmt->cast(); - auto *memref = cast(storeOp->getMemRef()); - for (unsigned pos = srcPos + 1; pos < dstPos; ++pos) { - if (memrefToLoadPosSet.count(memref) > 0 && - memrefToLoadPosSet[memref].count(pos) > 0) - return true; - } - } - } - return false; - } - - // Update mappings of memref loads and stores at 'srcPos' to 'dstPos'. - void moveAccessesAndRemovePos(unsigned srcPos, unsigned dstPos) { - // Lookup ForStmt at 'srcPos'. - auto *srcForStmt = getForStmtAtPos(srcPos); - // Move load op accesses from src to dst. - if (forStmtToLoadOps.count(srcForStmt) > 0) { - for (auto *opStmt : forStmtToLoadOps[srcForStmt]) { - auto loadOp = opStmt->cast(); - auto *memref = cast(loadOp->getMemRef()); - // Remove 'memref' to 'srcPos' mapping. - memrefToLoadPosSet[memref].erase(srcPos); - } - } - // Move store op accesses from src to dst. - if (forStmtToStoreOps.count(srcForStmt) > 0) { - for (auto *opStmt : forStmtToStoreOps[srcForStmt]) { - auto storeOp = opStmt->cast(); - auto *memref = cast(storeOp->getMemRef()); - // Remove 'memref' to 'srcPos' mapping. - memrefToStorePosSet[memref].erase(srcPos); - } - } - // Remove old state. - forStmtToLoadOps.erase(srcForStmt); - forStmtToStoreOps.erase(srcForStmt); - posToStmtInfo.erase(srcPos); - } - - ForStmt *getForStmtAtPos(unsigned pos) { - assert(posToStmtInfo.count(pos) > 0); - assert(isa(posToStmtInfo[pos].stmt)); - return cast(posToStmtInfo[pos].stmt); - } - - void addStmtInfoAt(unsigned pos, Statement *stmt, bool hasIfStmt) { - StmtInfo stmtInfo; - stmtInfo.stmt = stmt; - stmtInfo.isOrContainsIfStmt = hasIfStmt; - // Add mapping from 'pos' to StmtInfo for 'forStmt'. - posToStmtInfo[pos] = stmtInfo; - } - - // Adds the following mappings: - // *) 'containingForStmt' to load 'opStmt' - // *) 'memref' of load 'opStmt' to 'topLevelPos'. - void addLoadOpStmtAt(unsigned topLevelPos, OperationStmt *opStmt, - ForStmt *containingForStmt) { - if (containingForStmt != nullptr) { - // Add mapping from 'containingForStmt' to 'opStmt' for load op. - forStmtToLoadOps[containingForStmt].insert(opStmt); - } - auto loadOp = opStmt->cast(); - auto *memref = cast(loadOp->getMemRef()); - // Add mapping from 'memref' to 'topLevelPos' for load. - memrefToLoadPosSet[memref].insert(topLevelPos); - } - - // Adds the following mappings: - // *) 'containingForStmt' to store 'opStmt' - // *) 'memref' of store 'opStmt' to 'topLevelPos'. - void addStoreOpStmtAt(unsigned topLevelPos, OperationStmt *opStmt, - ForStmt *containingForStmt) { - if (containingForStmt != nullptr) { - // Add mapping from 'forStmt' to 'opStmt' for store op. - forStmtToStoreOps[containingForStmt].insert(opStmt); - } - auto storeOp = opStmt->cast(); - auto *memref = cast(storeOp->getMemRef()); - // Add mapping from 'memref' to 'topLevelPos' for store. - memrefToStorePosSet[memref].insert(topLevelPos); - } - - bool hasIfStmts() { - for (auto &pair : posToStmtInfo) - if (pair.second.isOrContainsIfStmt) - return true; - return false; - } }; } // end anonymous namespace PassResult LoopFusion::runOnMLFunction(MLFunction *f) { - GreedyFusionPolicy(f).run(); + MemRefDependenceGraph g; + if (g.init(f)) + GreedyFusion(&g).run(); return success(); } diff --git a/mlir/test/Transforms/loop-fusion.mlir b/mlir/test/Transforms/loop-fusion.mlir index d0de62e8a067..a668e181cc19 100644 --- a/mlir/test/Transforms/loop-fusion.mlir +++ b/mlir/test/Transforms/loop-fusion.mlir @@ -35,20 +35,17 @@ mlfunc @should_fuse_raw_dep_for_locality() { // CHECK: [[MAP0:#map[0-9]+]] = (d0) -> (d0) -// TODO(andydavis) Turn this into a proper reduction when constraints on -// the current greedy fusion policy are relaxed. // CHECK-LABEL: mlfunc @should_fuse_reduction_to_pointwise() { mlfunc @should_fuse_reduction_to_pointwise() { %a = alloc() : memref<10x10xf32> %b = alloc() : memref<10xf32> %c = alloc() : memref<10xf32> - %d = alloc() : memref<10xf32> %cf7 = constant 7.0 : f32 for %i0 = 0 to 10 { for %i1 = 0 to 10 { - %v0 = load %d[%i0] : memref<10xf32> + %v0 = load %b[%i0] : memref<10xf32> %v1 = load %a[%i0, %i1] : memref<10x10xf32> %v3 = addf %v0, %v1 : f32 store %v3, %b[%i0] : memref<10xf32> @@ -62,15 +59,15 @@ mlfunc @should_fuse_reduction_to_pointwise() { // Should fuse in entire inner loop on %i1 from source loop nest, as %i1 // is not used in the access function of the store/load on %b. // CHECK: for %i0 = 0 to 10 { - // CHECK-NEXT: %4 = affine_apply [[MAP0]](%i0) + // CHECK-NEXT: %3 = affine_apply [[MAP0]](%i0) // CHECK-NEXT: for %i1 = 0 to 10 { - // CHECK-NEXT: %5 = load %3[%4] : memref<10xf32> - // CHECK-NEXT: %6 = load %0[%4, %i1] : memref<10x10xf32> - // CHECK-NEXT: %7 = addf %5, %6 : f32 - // CHECK-NEXT: store %7, %1[%4] : memref<10xf32> - // CHECK-NEXT: } - // CHECK-NEXT: %8 = load %1[%i0] : memref<10xf32> - // CHECK-NEXT: store %8, %2[%i0] : memref<10xf32> + // CHECK-NEXT: %4 = load %1[%3] : memref<10xf32> + // CHECK-NEXT: %5 = load %0[%3, %i1] : memref<10x10xf32> + // CHECK-NEXT: %6 = addf %4, %5 : f32 + // CHECK-NEXT: store %6, %1[%3] : memref<10xf32> + // CHECK-NEXT: } + // CHECK-NEXT: %7 = load %1[%i0] : memref<10xf32> + // CHECK-NEXT: store %7, %2[%i0] : memref<10xf32> // CHECK-NEXT: } // CHECK-NEXT: return return