lib/Transforms/LoopFusion.cpp - platform/external/tensorflow - Git at Google

 //===- LoopFusion.cpp - Code to perform loop fusion -----------------------===//
 //
 // Copyright 2019 The MLIR Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //   http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // =============================================================================
 //
 // This file implements loop fusion.
 //
 //===----------------------------------------------------------------------===//

 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/Analysis/Utils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/StmtVisitor.h"
 #include "mlir/Pass.h"
 #include "mlir/StandardOps/StandardOps.h"
 #include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/raw_ostream.h"

 using llvm::SetVector;

 using namespace mlir;

 namespace {

 /// Loop fusion pass. This pass currently supports a greedy fusion policy,
 /// which fuses loop nests with single-writer/single-reader memref dependences
 /// with the goal of improving locality.

 // TODO(andydavis) Support fusion of source loop nests which write to multiple
 // memrefs, where each memref can have multiple users (if profitable).
 // TODO(andydavis) Extend this pass to check for fusion preventing dependences,
 // and add support for more general loop fusion algorithms.

 struct LoopFusion : public FunctionPass {
   LoopFusion() : FunctionPass(&LoopFusion::passID) {}

   PassResult runOnMLFunction(MLFunction *f) override;
   static char passID;
 };

 } // end anonymous namespace

 char LoopFusion::passID = 0;

 FunctionPass *mlir::createLoopFusionPass() { return new LoopFusion; }

 static void getSingleMemRefAccess(OperationStmt *loadOrStoreOpStmt,
                                   MemRefAccess *access) {
   if (auto loadOp = loadOrStoreOpStmt->dyn_cast<LoadOp>()) {
     access->memref = cast<MLValue>(loadOp->getMemRef());
     access->opStmt = loadOrStoreOpStmt;
     auto loadMemrefType = loadOp->getMemRefType();
     access->indices.reserve(loadMemrefType.getRank());
     for (auto *index : loadOp->getIndices()) {
       access->indices.push_back(cast<MLValue>(index));
     }
   } else {
     assert(loadOrStoreOpStmt->isa<StoreOp>());
     auto storeOp = loadOrStoreOpStmt->dyn_cast<StoreOp>();
     access->opStmt = loadOrStoreOpStmt;
     access->memref = cast<MLValue>(storeOp->getMemRef());
     auto storeMemrefType = storeOp->getMemRefType();
     access->indices.reserve(storeMemrefType.getRank());
     for (auto *index : storeOp->getIndices()) {
       access->indices.push_back(cast<MLValue>(index));
     }
   }
 }

 // FusionCandidate encapsulates source and destination memref access within
 // loop nests which are candidates for loop fusion.
 struct FusionCandidate {
   // Load or store access within src loop nest to be fused into dst loop nest.
   MemRefAccess srcAccess;
   // Load or store access within dst loop nest.
   MemRefAccess dstAccess;
 };

 static FusionCandidate buildFusionCandidate(OperationStmt *srcStoreOpStmt,
                                             OperationStmt *dstLoadOpStmt) {
   FusionCandidate candidate;
   // Get store access for src loop nest.
   getSingleMemRefAccess(srcStoreOpStmt, &candidate.srcAccess);
   // Get load access for dst loop nest.
   getSingleMemRefAccess(dstLoadOpStmt, &candidate.dstAccess);
   return candidate;
 }

 namespace {

 // LoopNestStateCollector walks loop nests and collects load and store
 // operations, and whether or not an IfStmt was encountered in the loop nest.
 class LoopNestStateCollector : public StmtWalker<LoopNestStateCollector> {
 public:
   SmallVector<ForStmt *, 4> forStmts;
   SmallVector<OperationStmt *, 4> loadOpStmts;
   SmallVector<OperationStmt *, 4> storeOpStmts;
   bool hasIfStmt = false;

   void visitForStmt(ForStmt *forStmt) { forStmts.push_back(forStmt); }

   void visitIfStmt(IfStmt *ifStmt) { hasIfStmt = true; }

   void visitOperationStmt(OperationStmt *opStmt) {
     if (opStmt->isa<LoadOp>())
       loadOpStmts.push_back(opStmt);
     if (opStmt->isa<StoreOp>())
       storeOpStmts.push_back(opStmt);
   }
 };

 // GreedyFusionPolicy 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).
 //   *) 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'.
 //   *) For each LoadOp in 'dstLoadOps' do:
 //      *) Lookup dependent loop nests at earlier positions in the MLFunction
 //         which have a single store op to the same memref.
 //      *) 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
 //         bounds to be functions of 'dstLoopNest' IVs and symbols.
 //      *) Fuse the 'srcLoopNest' computation slice into the 'dstLoopNest',
 //         just before the dst load op user.
 //      *) Add the newly fused load/store operation statements to the state,
 //         and also add newly fuse load ops to 'dstLoopOps' to be considered
 //         as fusion dst load ops in another iteration.
 //      *) Remove old src loop nest and its associated state.
 //
 // 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).
 //
 // This greedy algorithm is not 'maximally' 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.
   SmallVector<unsigned, 4> worklist;
   // Mapping from top-level position to StmtInfo.
   DenseMap<unsigned, StmtInfo> posToStmtInfo;
   // Mapping from memref MLValue to set of top-level positions of loop nests
   // which contain load ops on that memref.
   DenseMap<MLValue *, DenseSet<unsigned>> memrefToLoadPosSet;
   // Mapping from memref MLValue to set of top-level positions of loop nests
   // which contain store ops on that memref.
   DenseMap<MLValue *, DenseSet<unsigned>> memrefToStorePosSet;
   // Mapping from top-level loop nest to the set of load ops it contains.
   DenseMap<ForStmt *, SetVector<OperationStmt *>> forStmtToLoadOps;
   // Mapping from top-level loop nest to the set of store ops it contains.
   DenseMap<ForStmt *, SetVector<OperationStmt *>> forStmtToStoreOps;

   GreedyFusionPolicy(MLFunction *f) { init(f); }

   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();
       worklist.pop_back();
       // Skip if 'dstPos' is not tracked (was fused into another loop nest).
       if (posToStmtInfo.count(dstPos) == 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)
         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<OperationStmt *> dstLoadOps = forStmtToLoadOps[dstForStmt];
       while (!dstLoadOps.empty()) {
         auto *dstLoadOpStmt = dstLoadOps.pop_back_val();

         auto dstLoadOp = dstLoadOpStmt->cast<LoadOp>();
         auto *memref = cast<MLValue>(dstLoadOp->getMemRef());
         // Skip if not single src store / dst load pair on 'memref'.
         if (memrefToLoadPosSet[memref].size() != 1 ||
             memrefToStorePosSet[memref].size() != 1)
           continue;
         unsigned srcPos = *memrefToStorePosSet[memref].begin();
         if (srcPos >= dstPos)
           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);
           }
           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<ForStmt>(&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<OperationStmt>(&stmt)) {
         if (auto loadOp = opStmt->dyn_cast<LoadOp>()) {
           // Create StmtInfo for top-level load op.
           addStmtInfoAt(pos, &stmt, /*hasIfStmt=*/false);
           addLoadOpStmtAt(pos, opStmt, /*containingForStmt=*/nullptr);
         }
         if (auto storeOp = opStmt->dyn_cast<StoreOp>()) {
           // Create StmtInfo for top-level store op.
           addStmtInfoAt(pos, &stmt, /*hasIfStmt=*/false);
           addStoreOpStmtAt(pos, opStmt, /*containingForStmt=*/nullptr);
         }
       }
       if (auto *ifStmt = dyn_cast<IfStmt>(&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<LoadOp>();
         auto *memref = cast<MLValue>(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<StoreOp>();
         auto *memref = cast<MLValue>(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<LoadOp>();
         auto *memref = cast<MLValue>(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<StoreOp>();
         auto *memref = cast<MLValue>(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<ForStmt>(posToStmtInfo[pos].stmt));
     return cast<ForStmt>(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<LoadOp>();
     auto *memref = cast<MLValue>(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<StoreOp>();
     auto *memref = cast<MLValue>(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();
   return success();
 }

 static PassRegistration<LoopFusion> pass("loop-fusion", "Fuse loop nests");
	//===- LoopFusion.cpp - Code to perform loop fusion -----------------------===//
	//
	// Copyright 2019 The MLIR Authors.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.
	// =============================================================================
	//
	// This file implements loop fusion.
	//
	//===----------------------------------------------------------------------===//

	#include "mlir/Analysis/AffineAnalysis.h"
	#include "mlir/Analysis/AffineStructures.h"
	#include "mlir/Analysis/LoopAnalysis.h"
	#include "mlir/Analysis/Utils.h"
	#include "mlir/IR/AffineExpr.h"
	#include "mlir/IR/AffineMap.h"
	#include "mlir/IR/Builders.h"
	#include "mlir/IR/BuiltinOps.h"
	#include "mlir/IR/StmtVisitor.h"
	#include "mlir/Pass.h"
	#include "mlir/StandardOps/StandardOps.h"
	#include "mlir/Transforms/LoopUtils.h"
	#include "mlir/Transforms/Passes.h"
	#include "llvm/ADT/DenseMap.h"
	#include "llvm/ADT/DenseSet.h"
	#include "llvm/ADT/SetVector.h"
	#include "llvm/Support/raw_ostream.h"

	using llvm::SetVector;

	using namespace mlir;

	namespace {

	/// Loop fusion pass. This pass currently supports a greedy fusion policy,
	/// which fuses loop nests with single-writer/single-reader memref dependences
	/// with the goal of improving locality.

	// TODO(andydavis) Support fusion of source loop nests which write to multiple
	// memrefs, where each memref can have multiple users (if profitable).
	// TODO(andydavis) Extend this pass to check for fusion preventing dependences,
	// and add support for more general loop fusion algorithms.

	struct LoopFusion : public FunctionPass {
	LoopFusion() : FunctionPass(&LoopFusion::passID) {}

	PassResult runOnMLFunction(MLFunction *f) override;
	static char passID;
	};

	} // end anonymous namespace

	char LoopFusion::passID = 0;

	FunctionPass *mlir::createLoopFusionPass() { return new LoopFusion; }

	static void getSingleMemRefAccess(OperationStmt *loadOrStoreOpStmt,
	MemRefAccess *access) {
	if (auto loadOp = loadOrStoreOpStmt->dyn_cast<LoadOp>()) {
	access->memref = cast<MLValue>(loadOp->getMemRef());
	access->opStmt = loadOrStoreOpStmt;
	auto loadMemrefType = loadOp->getMemRefType();
	access->indices.reserve(loadMemrefType.getRank());
	for (auto *index : loadOp->getIndices()) {
	access->indices.push_back(cast<MLValue>(index));
	}
	} else {
	assert(loadOrStoreOpStmt->isa<StoreOp>());
	auto storeOp = loadOrStoreOpStmt->dyn_cast<StoreOp>();
	access->opStmt = loadOrStoreOpStmt;
	access->memref = cast<MLValue>(storeOp->getMemRef());
	auto storeMemrefType = storeOp->getMemRefType();
	access->indices.reserve(storeMemrefType.getRank());
	for (auto *index : storeOp->getIndices()) {
	access->indices.push_back(cast<MLValue>(index));
	}
	}
	}

	// FusionCandidate encapsulates source and destination memref access within
	// loop nests which are candidates for loop fusion.
	struct FusionCandidate {
	// Load or store access within src loop nest to be fused into dst loop nest.
	MemRefAccess srcAccess;
	// Load or store access within dst loop nest.
	MemRefAccess dstAccess;
	};

	static FusionCandidate buildFusionCandidate(OperationStmt *srcStoreOpStmt,
	OperationStmt *dstLoadOpStmt) {
	FusionCandidate candidate;
	// Get store access for src loop nest.
	getSingleMemRefAccess(srcStoreOpStmt, &candidate.srcAccess);
	// Get load access for dst loop nest.
	getSingleMemRefAccess(dstLoadOpStmt, &candidate.dstAccess);
	return candidate;
	}

	namespace {

	// LoopNestStateCollector walks loop nests and collects load and store
	// operations, and whether or not an IfStmt was encountered in the loop nest.
	class LoopNestStateCollector : public StmtWalker<LoopNestStateCollector> {
	public:
	SmallVector<ForStmt *, 4> forStmts;
	SmallVector<OperationStmt *, 4> loadOpStmts;
	SmallVector<OperationStmt *, 4> storeOpStmts;
	bool hasIfStmt = false;

	void visitForStmt(ForStmt *forStmt) { forStmts.push_back(forStmt); }

	void visitIfStmt(IfStmt *ifStmt) { hasIfStmt = true; }

	void visitOperationStmt(OperationStmt *opStmt) {
	if (opStmt->isa<LoadOp>())
	loadOpStmts.push_back(opStmt);
	if (opStmt->isa<StoreOp>())
	storeOpStmts.push_back(opStmt);
	}
	};

	// GreedyFusionPolicy 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).
	// *) 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'.
	// *) For each LoadOp in 'dstLoadOps' do:
	// *) Lookup dependent loop nests at earlier positions in the MLFunction
	// which have a single store op to the same memref.
	// *) 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
	// bounds to be functions of 'dstLoopNest' IVs and symbols.
	// *) Fuse the 'srcLoopNest' computation slice into the 'dstLoopNest',
	// just before the dst load op user.
	// *) Add the newly fused load/store operation statements to the state,
	// and also add newly fuse load ops to 'dstLoopOps' to be considered
	// as fusion dst load ops in another iteration.
	// *) Remove old src loop nest and its associated state.
	//
	// 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).
	//
	// This greedy algorithm is not 'maximally' 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.
	SmallVector<unsigned, 4> worklist;
	// Mapping from top-level position to StmtInfo.
	DenseMap<unsigned, StmtInfo> posToStmtInfo;
	// Mapping from memref MLValue to set of top-level positions of loop nests
	// which contain load ops on that memref.
	DenseMap<MLValue *, DenseSet<unsigned>> memrefToLoadPosSet;
	// Mapping from memref MLValue to set of top-level positions of loop nests
	// which contain store ops on that memref.
	DenseMap<MLValue *, DenseSet<unsigned>> memrefToStorePosSet;
	// Mapping from top-level loop nest to the set of load ops it contains.
	DenseMap<ForStmt , SetVector<OperationStmt >> forStmtToLoadOps;
	// Mapping from top-level loop nest to the set of store ops it contains.
	DenseMap<ForStmt , SetVector<OperationStmt >> forStmtToStoreOps;

	GreedyFusionPolicy(MLFunction *f) { init(f); }

	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();
	worklist.pop_back();
	// Skip if 'dstPos' is not tracked (was fused into another loop nest).
	if (posToStmtInfo.count(dstPos) == 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)
	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<OperationStmt *> dstLoadOps = forStmtToLoadOps[dstForStmt];
	while (!dstLoadOps.empty()) {
	auto *dstLoadOpStmt = dstLoadOps.pop_back_val();

	auto dstLoadOp = dstLoadOpStmt->cast<LoadOp>();
	auto *memref = cast<MLValue>(dstLoadOp->getMemRef());
	// Skip if not single src store / dst load pair on 'memref'.
	if (memrefToLoadPosSet[memref].size() != 1 \|\|
	memrefToStorePosSet[memref].size() != 1)
	continue;
	unsigned srcPos = *memrefToStorePosSet[memref].begin();
	if (srcPos >= dstPos)
	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);
	}
	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<ForStmt>(&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<OperationStmt>(&stmt)) {
	if (auto loadOp = opStmt->dyn_cast<LoadOp>()) {
	// Create StmtInfo for top-level load op.
	addStmtInfoAt(pos, &stmt, /hasIfStmt=/false);
	addLoadOpStmtAt(pos, opStmt, /containingForStmt=/nullptr);
	}
	if (auto storeOp = opStmt->dyn_cast<StoreOp>()) {
	// Create StmtInfo for top-level store op.
	addStmtInfoAt(pos, &stmt, /hasIfStmt=/false);
	addStoreOpStmtAt(pos, opStmt, /containingForStmt=/nullptr);
	}
	}
	if (auto *ifStmt = dyn_cast<IfStmt>(&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<LoadOp>();
	auto *memref = cast<MLValue>(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<StoreOp>();
	auto *memref = cast<MLValue>(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<LoadOp>();
	auto *memref = cast<MLValue>(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<StoreOp>();
	auto *memref = cast<MLValue>(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<ForStmt>(posToStmtInfo[pos].stmt));
	return cast<ForStmt>(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<LoadOp>();
	auto *memref = cast<MLValue>(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<StoreOp>();
	auto *memref = cast<MLValue>(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();
	return success();
	}

	static PassRegistration<LoopFusion> pass("loop-fusion", "Fuse loop nests");