[MLIR] Remove unnecessary Block argument on MemRefDependenceGraph::init

[bondhugula]

Summary

Remove unnecessary Block argument on MemRefDependenceGraph::init.
block is already a field on MDG.

[bondhugula]

Remove stale comment.

[bondhugula]

Adjust/improve debug messsages.

mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp

public:
   // Map from memref to a count on the dependence edges associated with that
   // memref.
   DenseMap<Value, unsigned> memrefEdgeCount;
   // The next unique identifier to use for newly created graph nodes.
   unsigned nextNodeId = 0;
 
   MemRefDependenceGraph(Block &block) : block(block) {}
 
-  // Initializes the dependence graph based on operations in 'f'.
+  // Initializes the dependence graph based on operations in `block'.
   // Returns true on success, false otherwise.
-  bool init(Block *block);
+  bool init();
 
   // Returns the graph node for 'id'.
   Node *getNode(unsigned id) {
     auto it = nodes.find(id);
     assert(it != nodes.end());
     return &it->second;
   }
 
if (isEscapingMemref(memref, &mdg->block))
       escapingMemRefs.insert(memref);
   }
 }
 
 } // namespace
 
 // Initializes the data dependence graph by walking operations in `block`.
 // Assigns each node in the graph a node id based on program order in 'f'.
-// TODO: Add support for taking a Block arg to construct the
-// dependence graph at a different depth.
-bool MemRefDependenceGraph::init(Block *block) {
+bool MemRefDependenceGraph::init() {
   LLVM_DEBUG(llvm::dbgs() << "--- Initializing MDG ---\n");
   // Map from a memref to the set of ids of the nodes that have ops accessing
   // the memref.
   DenseMap<Value, SetVector<unsigned>> memrefAccesses;
 
   DenseMap<Operation *, unsigned> forToNodeMap;
-  for (Operation &op : *block) {
+  for (Operation &op : block) {
     if (auto forOp = dyn_cast<AffineForOp>(op)) {
       // Create graph node 'id' to represent top-level 'forOp' and record
       // all loads and store accesses it contains.
       LoopNestStateCollector collector;
       collector.collect(&op);
       // Return false if a region holding op other than 'affine.for' and
       // 'affine.if' was found (not currently supported).
       if (collector.hasNonAffineRegionOp)
for (auto &idAndNode : nodes) {
     const Node &node = idAndNode.second;
     // Stores don't define SSA values, skip them.
     if (!node.stores.empty())
       continue;
     Operation *opInst = node.op;
     for (Value value : opInst->getResults()) {
       for (Operation *user : value.getUsers()) {
         // Ignore users outside of the block.
-        if (block->getParent()->findAncestorOpInRegion(*user)->getBlock() !=
-            block)
+        if (block.getParent()->findAncestorOpInRegion(*user)->getBlock() !=
+            &block)
           continue;
         SmallVector<AffineForOp, 4> loops;
         getAffineForIVs(*user, &loops);
         if (loops.empty())
           continue;
         assert(forToNodeMap.count(loops[0]) > 0 && "missing mapping");
         unsigned userLoopNestId = forToNodeMap[loops[0]];
         addEdge(node.id, userLoopNestId, value);
static bool isFusionProfitable(Operation *srcOpInst, Operation *srcStoreOpInst,
                                double computeToleranceThreshold) {
   LLVM_DEBUG({
     llvm::dbgs() << "Checking whether fusion is profitable between src op:\n";
     llvm::dbgs() << ' ' << *srcOpInst << " and destination loop:\n";
     llvm::dbgs() << dstForOp << "\n";
   });
 
   if (maxLegalFusionDepth == 0) {
-    LLVM_DEBUG(llvm::dbgs() << "Can't fuse: maxLegalFusionDepth == 0 .\n");
+    LLVM_DEBUG(llvm::dbgs() << "Can't fuse: maxLegalFusionDepth is 0\n");
     return false;
   }
 
   // Compute cost of sliced and unsliced src loop nest.
   SmallVector<AffineForOp, 4> srcLoopIVs;
   getAffineForIVs(*srcOpInst, &srcLoopIVs);
 
   // Walk src loop nest and collect stats.
static bool isFusionProfitable(Operation *srcOpInst, Operation *srcStoreOpInst,
 
   // Compute op instance count for the src loop nest without iteration slicing.
   uint64_t srcLoopNestCost = getComputeCost(srcLoopIVs[0], srcLoopNestStats);
 
   // Compute src loop nest write region size.
   MemRefRegion srcWriteRegion(srcStoreOpInst->getLoc());
   if (failed(srcWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0))) {
     LLVM_DEBUG(llvm::dbgs()
-               << "Unable to compute MemRefRegion for source operation\n.");
+               << "Unable to compute MemRefRegion for source operation\n");
     return false;
   }
 
   std::optional<int64_t> maybeSrcWriteRegionSizeBytes =
       srcWriteRegion.getRegionSize();
   if (!maybeSrcWriteRegionSizeBytes.has_value())
     return false;
   int64_t srcWriteRegionSizeBytes = *maybeSrcWriteRegionSizeBytes;
 
   // Compute op instance count for the src loop nest.
   uint64_t dstLoopNestCost = getComputeCost(dstForOp, dstLoopNestStats);
 
   // Evaluate all depth choices for materializing the slice in the destination
   // loop nest.
   for (unsigned i = maxLegalFusionDepth; i >= 1; --i) {
     const ComputationSliceState &slice = depthSliceUnions[i - 1];
     // Skip slice union if it wasn't computed for this depth.
     if (slice.isEmpty())
       continue;
 
     int64_t fusedLoopNestComputeCost;
     if (!getFusionComputeCost(srcLoopIVs[0], srcLoopNestStats, dstForOp,
                               dstLoopNestStats, slice,
                               &fusedLoopNestComputeCost)) {
-      LLVM_DEBUG(llvm::dbgs() << "Unable to compute fusion compute cost.\n.");
+      LLVM_DEBUG(llvm::dbgs() << "Unable to compute fusion compute cost\n");
       continue;
     }
 
     double additionalComputeFraction =
         fusedLoopNestComputeCost /
             (static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
         1;
 
do {
         // and 'dstNode' using their memrefs. If so, we wouldn't be able to
         // compute a legal insertion point for now. 'srcNode' and 'dstNode'
         // memrefs with non-affine operation users would be considered
         // escaping memrefs so we can limit this check to only scenarios with
         // escaping memrefs.
         if (!srcEscapingMemRefs.empty() &&
             hasNonAffineUsersOnThePath(srcId, dstId, mdg)) {
           LLVM_DEBUG(llvm::dbgs()
-                     << "Can't fuse: non-affine users in between the loops\n.");
+                     << "Can't fuse: non-affine users in between the loops\n");
           continue;
         }
 
         // Compute an operation list insertion point for the fused loop
         // nest which preserves dependences.
         Operation *fusedLoopInsPoint =
             mdg->getFusedLoopNestInsertionPoint(srcNode->id, dstNode->id);
         if (fusedLoopInsPoint == nullptr)
public:
   }
 };
 
 } // namespace
 
 /// Run fusion on `block`.
 void LoopFusion::runOnBlock(Block *block) {
   MemRefDependenceGraph g(*block);
-  if (!g.init(block))
+  if (!g.init()) {
+    LLVM_DEBUG(llvm::dbgs() << "MDG init failed\n");
     return;
+  }
 
   std::optional<unsigned> fastMemorySpaceOpt;
   if (fastMemorySpace.hasValue())
     fastMemorySpaceOpt = fastMemorySpace;
   unsigned localBufSizeThresholdBytes = localBufSizeThreshold * 1024;
   GreedyFusion fusion(&g, localBufSizeThresholdBytes, fastMemorySpaceOpt,
                       maximalFusion, computeToleranceThreshold);