[MLIR] NFC cleanup/modernize memref-dataflow-opt / getNestingDepth

[bondhugula]

Summary

Bring code to date with with later changes to the core infrastructure /
coding style.

mlir/include/mlir/Analysis/Utils.h

 
 /// Populates 'loops' with IVs of the loops surrounding 'op' ordered from
 /// the outermost 'affine.for' operation to the innermost one.
 //  TODO(bondhugula): handle 'affine.if' ops.
 void getLoopIVs(Operation &op, SmallVectorImpl<AffineForOp> *loops);
 
 /// Returns the nesting depth of this operation, i.e., the number of loops
 /// surrounding this operation.
-unsigned getNestingDepth(Operation &op);
+unsigned getNestingDepth(Operation *op);
 
 /// Returns in 'sequentialLoops' all sequential loops in loop nest rooted
 /// at 'forOp'.
 void getSequentialLoops(AffineForOp forOp,
                         llvm::SmallDenseSet<Value, 8> *sequentialLoops);
 
 /// ComputationSliceState aggregates loop IVs, loop bound AffineMaps and their
 /// associated operands for a set of loops within a loop nest (typically the

mlir/lib/Analysis/Utils.cpp

LogicalResult mlir::computeSliceUnion(ArrayRef<Operation *> opsA,
   std::vector<std::pair<Operation *, Operation *>> dependentOpPairs;
   for (unsigned i = 0, numOpsA = opsA.size(); i < numOpsA; ++i) {
     MemRefAccess srcAccess(opsA[i]);
     for (unsigned j = 0, numOpsB = opsB.size(); j < numOpsB; ++j) {
       MemRefAccess dstAccess(opsB[j]);
       if (srcAccess.memref != dstAccess.memref)
         continue;
       // Check if 'loopDepth' exceeds nesting depth of src/dst ops.
-      if ((!isBackwardSlice && loopDepth > getNestingDepth(*opsA[i])) ||
-          (isBackwardSlice && loopDepth > getNestingDepth(*opsB[j]))) {
+      if ((!isBackwardSlice && loopDepth > getNestingDepth(opsA[i])) ||
+          (isBackwardSlice && loopDepth > getNestingDepth(opsB[j]))) {
         LLVM_DEBUG(llvm::dbgs() << "Invalid loop depth\n.");
         return failure();
       }
 
       bool readReadAccesses = isa<AffineLoadOp>(srcAccess.opInst) &&
                               isa<AffineLoadOp>(dstAccess.opInst);
       FlatAffineConstraints dependenceConstraints;
       // Check dependence between 'srcAccess' and 'dstAccess'.

 unsigned MemRefAccess::getRank() const {
   return memref.getType().cast<MemRefType>().getRank();
 }
 
 bool MemRefAccess::isStore() const { return isa<AffineStoreOp>(opInst); }
 
 /// Returns the nesting depth of this statement, i.e., the number of loops
 /// surrounding this statement.
-unsigned mlir::getNestingDepth(Operation &op) {
-  Operation *currOp = &op;
+unsigned mlir::getNestingDepth(Operation *op) {
+  Operation *currOp = op;
   unsigned depth = 0;
   while ((currOp = currOp->getParentOp())) {
     if (isa<AffineForOp>(currOp))
       depth++;
   }
   return depth;
 }
 
if (!isa<AffineLoadOp>(opInst) && !isa<AffineStoreOp>(opInst)) {
       // Neither load nor a store op.
       return WalkResult::advance();
     }
 
     // Compute the memref region symbolic in any IVs enclosing this block.
     auto region = std::make_unique<MemRefRegion>(opInst->getLoc());
     if (failed(
             region->compute(opInst,
-                            /*loopDepth=*/getNestingDepth(*block.begin())))) {
+                            /*loopDepth=*/getNestingDepth(&*block.begin())))) {
       return opInst->emitError("error obtaining memory region\n");
     }
 
     auto it = regions.find(region->memref);
     if (it == regions.end()) {
       regions[region->memref] = std::move(region);
     } else if (failed(it->second->unionBoundingBox(*region))) {
       return opInst->emitWarning(
auto walkResult = forOp.walk([&](Operation *opInst) -> WalkResult {
     return WalkResult::advance();
   });
 
   // Stop early if the loop has unknown ops with side effects.
   if (walkResult.wasInterrupted())
     return false;
 
   // Dep check depth would be number of enclosing loops + 1.
-  unsigned depth = getNestingDepth(*forOp.getOperation()) + 1;
+  unsigned depth = getNestingDepth(forOp) + 1;
 
   // Check dependences between all pairs of ops in 'loadAndStoreOpInsts'.
   for (auto *srcOpInst : loadAndStoreOpInsts) {
     MemRefAccess srcAccess(srcOpInst);
     for (auto *dstOpInst : loadAndStoreOpInsts) {
       MemRefAccess dstAccess(dstOpInst);
       FlatAffineConstraints dependenceConstraints;
       DependenceResult result = checkMemrefAccessDependence(

mlir/lib/Transforms/LoopFusion.cpp

while (!worklist.empty()) {
             // feasibility for loops with multiple stores.
             unsigned maxLoopDepth = 0;
             for (auto *op : srcNode->stores) {
               auto storeOp = cast<AffineStoreOp>(op);
               if (storeOp.getMemRef() != memref) {
                 srcStoreOp = nullptr;
                 break;
               }
-              unsigned loopDepth = getNestingDepth(*storeOp);
+              unsigned loopDepth = getNestingDepth(storeOp);
               if (loopDepth > maxLoopDepth) {
                 maxLoopDepth = loopDepth;
                 srcStoreOp = storeOp;
               }
             }
             if (!srcStoreOp)
               continue;
           }

mlir/lib/Transforms/MemRefDataFlowOpt.cpp

 /// store to load forwarding, elimination of dead stores, and dead allocs.
 std::unique_ptr<OperationPass<FuncOp>> mlir::createMemRefDataFlowOptPass() {
   return std::make_unique<MemRefDataFlowOpt>();
 }
 
 // This is a straightforward implementation not optimized for speed. Optimize
 // if needed.
 void MemRefDataFlowOpt::forwardStoreToLoad(AffineLoadOp loadOp) {
-  Operation *loadOpInst = loadOp.getOperation();
-
-  // First pass over the use list to get minimum number of surrounding
+  // First pass over the use list to get the minimum number of surrounding
   // loops common between the load op and the store op, with min taken across
   // all store ops.
   SmallVector<Operation *, 8> storeOps;
-  unsigned minSurroundingLoops = getNestingDepth(*loadOpInst);
+  unsigned minSurroundingLoops = getNestingDepth(loadOp);
   for (auto *user : loadOp.getMemRef().getUsers()) {
     auto storeOp = dyn_cast<AffineStoreOp>(user);
     if (!storeOp)
       continue;
-    auto *storeOpInst = storeOp.getOperation();
-    unsigned nsLoops = getNumCommonSurroundingLoops(*loadOpInst, *storeOpInst);
+    unsigned nsLoops = getNumCommonSurroundingLoops(*loadOp, *storeOp);
     minSurroundingLoops = std::min(nsLoops, minSurroundingLoops);
-    storeOps.push_back(storeOpInst);
+    storeOps.push_back(storeOp);
   }
 
   // The list of store op candidates for forwarding that satisfy conditions
   // (1) and (2) above - they will be filtered later when checking (3).
   SmallVector<Operation *, 8> fwdingCandidates;
 
   // Store ops that have a dependence into the load (even if they aren't
   // forwarding candidates). Each forwarding candidate will be checked for a
   // post-dominance on these. 'fwdingCandidates' are a subset of depSrcStores.
   SmallVector<Operation *, 8> depSrcStores;
 
-  for (auto *storeOpInst : storeOps) {
-    MemRefAccess srcAccess(storeOpInst);
-    MemRefAccess destAccess(loadOpInst);
+  for (auto *storeOp : storeOps) {
+    MemRefAccess srcAccess(storeOp);
+    MemRefAccess destAccess(loadOp);
     // Find stores that may be reaching the load.
     FlatAffineConstraints dependenceConstraints;
-    unsigned nsLoops = getNumCommonSurroundingLoops(*loadOpInst, *storeOpInst);
+    unsigned nsLoops = getNumCommonSurroundingLoops(*loadOp, *storeOp);
     unsigned d;
     // Dependences at loop depth <= minSurroundingLoops do NOT matter.
     for (d = nsLoops + 1; d > minSurroundingLoops; d--) {
       DependenceResult result = checkMemrefAccessDependence(
           srcAccess, destAccess, d, &dependenceConstraints,
           /*dependenceComponents=*/nullptr);
       if (hasDependence(result))
         break;
     }
     if (d == minSurroundingLoops)
       continue;
 
     // Stores that *may* be reaching the load.
-    depSrcStores.push_back(storeOpInst);
+    depSrcStores.push_back(storeOp);
 
     // 1. Check if the store and the load have mathematically equivalent
     // affine access functions; this implies that they statically refer to the
     // same single memref element. As an example this filters out cases like:
     //     store %A[%i0 + 1]
     //     load %A[%i0]
     //     store %A[%M]
     //     load %A[%N]
     // Use the AffineValueMap difference based memref access equality checking.
     if (srcAccess != destAccess)
       continue;
 
     // 2. The store has to dominate the load op to be candidate.
-    if (!domInfo->dominates(storeOpInst, loadOpInst))
+    if (!domInfo->dominates(storeOp, loadOp))
       continue;
 
     // We now have a candidate for forwarding.
-    fwdingCandidates.push_back(storeOpInst);
+    fwdingCandidates.push_back(storeOp);
   }
 
   // 3. Of all the store op's that meet the above criteria, the store that
   // postdominates all 'depSrcStores' (if one exists) is the unique store
   // providing the value to the load, i.e., provably the last writer to that
   // memref loc.
   // Note: this can be implemented in a cleaner way with postdominator tree
   // traversals. Consider this for the future if needed.
   Operation *lastWriteStoreOp = nullptr;
-  for (auto *storeOpInst : fwdingCandidates) {
+  for (auto *storeOp : fwdingCandidates) {
     if (llvm::all_of(depSrcStores, [&](Operation *depStore) {
-          return postDomInfo->postDominates(storeOpInst, depStore);
+          return postDomInfo->postDominates(storeOp, depStore);
         })) {
-      lastWriteStoreOp = storeOpInst;
+      lastWriteStoreOp = storeOp;
       break;
     }
   }
   if (!lastWriteStoreOp)
     return;
 
   // Perform the actual store to load forwarding.
   Value storeVal = cast<AffineStoreOp>(lastWriteStoreOp).getValueToStore();
   loadOp.replaceAllUsesWith(storeVal);
   // Record the memref for a later sweep to optimize away.
   memrefsToErase.insert(loadOp.getMemRef());
   // Record this to erase later.
-  loadOpsToErase.push_back(loadOpInst);
+  loadOpsToErase.push_back(loadOp);
 }
 
 void MemRefDataFlowOpt::runOnFunction() {
   // Only supports single block functions at the moment.
   FuncOp f = getFunction();
   if (f.getBlocks().size() != 1) {
     markAllAnalysesPreserved();
     return;
   }
 
   domInfo = &getAnalysis<DominanceInfo>();
   postDomInfo = &getAnalysis<PostDominanceInfo>();
 
   loadOpsToErase.clear();
   memrefsToErase.clear();
 
-  // Walk all load's and perform load/store forwarding.
+  // Walk all load's and perform store to load forwarding.
   f.walk([&](AffineLoadOp loadOp) { forwardStoreToLoad(loadOp); });
 
   // Erase all load op's whose results were replaced with store fwd'ed ones.
-  for (auto *loadOp : loadOpsToErase) {
+  for (auto *loadOp : loadOpsToErase)
     loadOp->erase();
-  }
 
   // Check if the store fwd'ed memrefs are now left with only stores and can
   // thus be completely deleted. Note: the canonicalize pass should be able
   // to do this as well, but we'll do it here since we collected these anyway.
   for (auto memref : memrefsToErase) {
     // If the memref hasn't been alloc'ed in this function, skip.
-    Operation *defInst = memref.getDefiningOp();
-    if (!defInst || !isa<AllocOp>(defInst))
+    Operation *defOp = memref.getDefiningOp();
+    if (!defOp || !isa<AllocOp>(defOp))
       // TODO(mlir-team): if the memref was returned by a 'call' operation, we
       // could still erase it if the call had no side-effects.
       continue;
-    if (llvm::any_of(memref.getUsers(), [&](Operation *ownerInst) {
-          return (!isa<AffineStoreOp>(ownerInst) && !isa<DeallocOp>(ownerInst));
+    if (llvm::any_of(memref.getUsers(), [&](Operation *ownerOp) {
+          return (!isa<AffineStoreOp>(ownerOp) && !isa<DeallocOp>(ownerOp));
         }))
       continue;
 
     // Erase all stores, the dealloc, and the alloc on the memref.
     for (auto *user : llvm::make_early_inc_range(memref.getUsers()))
       user->erase();
-    defInst->erase();
+    defOp->erase();
   }
 }

mlir/lib/Transforms/Utils/LoopUtils.cpp

uint64_t mlir::affineDataCopyGenerate(Block::iterator begin,
   assert(begin->getBlock() == std::prev(end)->getBlock() &&
          "Inconsistent block begin/end args");
   assert(end != end->getBlock()->end() && "end can't be the block terminator");
 
   Block *block = begin->getBlock();
 
   // Copies will be generated for this depth, i.e., symbolic in all loops
   // surrounding the this block range.
-  unsigned copyDepth = getNestingDepth(*begin);
+  unsigned copyDepth = getNestingDepth(&*begin);
 
   LLVM_DEBUG(llvm::dbgs() << "Generating copies at depth " << copyDepth
                           << "\n");
   LLVM_DEBUG(llvm::dbgs() << "from begin: " << *begin << "\n");
   LLVM_DEBUG(llvm::dbgs() << "to inclusive end: " << *std::prev(end) << "\n");
 
   // List of memory regions to copy for. We need a map vector to have a
   // guaranteed iteration order to write test cases. CHECK-DAG doesn't help here

mlir/test/lib/Transforms/TestLoopFusion.cpp

 //===- TestLoopFusion.cpp - Test loop fusion ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass to test various loop fusion utility functions.
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Analysis/AffineAnalysis.h"
-#include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
-#include "mlir/IR/Builders.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/LoopFusionUtils.h"
 #include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Passes.h"
 
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/Support/CommandLine.h"
-#include "llvm/Support/Debug.h"
-
 #define DEBUG_TYPE "test-loop-fusion"
 
 using namespace mlir;
 
 static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
 
 static llvm::cl::opt<bool> clTestDependenceCheck(
     "test-loop-fusion-dependence-check",
static unsigned getBlockIndex(Operation &op) {
   return index;
 }
 
 // Returns a string representation of 'sliceUnion'.
 static std::string getSliceStr(const mlir::ComputationSliceState &sliceUnion) {
   std::string result;
   llvm::raw_string_ostream os(result);
   // Slice insertion point format [loop-depth, operation-block-index]
-  unsigned ipd = getNestingDepth(*sliceUnion.insertPoint);
+  unsigned ipd = getNestingDepth(&*sliceUnion.insertPoint);
   unsigned ipb = getBlockIndex(*sliceUnion.insertPoint);
   os << "insert point: (" << std::to_string(ipd) << ", " << std::to_string(ipb)
      << ")";
   assert(sliceUnion.lbs.size() == sliceUnion.ubs.size());
   os << " loop bounds: ";
   for (unsigned k = 0, e = sliceUnion.lbs.size(); k < e; ++k) {
     os << '[';
     sliceUnion.lbs[k].print(os);