[MLIR][Affine] NFC. Fix/improve debug messages for affine analysis/fusion utils (llvm#127164)

Fix/improve debug messages and API signatures for affine
analysis/fusion utils.

Move some warnings under LLVM_DEBUG. These weren't meant to be exposed
during compilation.

Add dump pretty methods for FlatLinearConstraints.

NFC.

Author: bondhugula

mlir/include/mlir/Analysis/FlatLinearValueConstraints.h

@@ -194,6 +194,13 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
     return appendVar(VarKind::Local, num);
   }
 
+  /// A more human-readable version of dump().
+  void dumpPretty() const;
+  /// An easier to read dump of a `row` of the same width as the number of
+  /// columns. `fixedColWidth` ensure that even with a zero coefficient, we
+  /// print spaces so that variables are aligned.
+  void dumpRow(ArrayRef<int64_t> row, bool fixedColWidth = true) const;
+
 protected:
   using VarKind = presburger::VarKind;
 

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

@@ -416,10 +416,10 @@ struct ComputationSliceState {
 //      %v = affine.load %0[%i1] : memref<100xf32>    // 'depSinkAccess'
 //    }
 //
-void getComputationSliceState(Operation *depSourceOp, Operation *depSinkOp,
-                              FlatAffineValueConstraints *dependenceConstraints,
-                              unsigned loopDepth, bool isBackwardSlice,
-                              ComputationSliceState *sliceState);
+void getComputationSliceState(
+    Operation *depSourceOp, Operation *depSinkOp,
+    const FlatAffineValueConstraints &dependenceConstraints, unsigned loopDepth,
+    bool isBackwardSlice, ComputationSliceState *sliceState);
 
 /// Return the number of iterations for the `slicetripCountMap` provided.
 uint64_t getSliceIterationCount(

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

@@ -445,6 +445,58 @@ static bool detectAsFloorDiv(const FlatLinearConstraints &cst, unsigned pos,
   return true;
 }
 
+void FlatLinearConstraints::dumpRow(ArrayRef<int64_t> row,
+                                    bool fixedColWidth) const {
+  unsigned ncols = getNumCols();
+  bool firstNonZero = true;
+  for (unsigned j = 0; j < ncols; j++) {
+    if (j == ncols - 1) {
+      // Constant.
+      if (row[j] == 0 && !firstNonZero) {
+        if (fixedColWidth)
+          llvm::errs().indent(7);
+      } else {
+        llvm::errs() << ((row[j] >= 0) ? "+ " : "") << row[j] << ' ';
+      }
+    } else {
+      std::string var = std::string("c_") + std::to_string(j);
+      if (row[j] == 1)
+        llvm::errs() << "+ " << var << ' ';
+      else if (row[j] == -1)
+        llvm::errs() << "- " << var << ' ';
+      else if (row[j] >= 2)
+        llvm::errs() << "+ " << row[j] << '*' << var << ' ';
+      else if (row[j] <= -2)
+        llvm::errs() << "- " << -row[j] << '*' << var << ' ';
+      else if (fixedColWidth)
+        // Zero coeff.
+        llvm::errs().indent(7);
+      if (row[j] != 0)
+        firstNonZero = false;
+    }
+  }
+}
+
+void FlatLinearConstraints::dumpPretty() const {
+  assert(hasConsistentState());
+  llvm::errs() << "Constraints (" << getNumDimVars() << " dims, "
+               << getNumSymbolVars() << " symbols, " << getNumLocalVars()
+               << " locals), (" << getNumConstraints() << " constraints)\n";
+  auto dumpConstraint = [&](unsigned rowPos, bool isEq) {
+    // Is it the first non-zero entry?
+    SmallVector<int64_t> row =
+        isEq ? getEquality64(rowPos) : getInequality64(rowPos);
+    dumpRow(row);
+    llvm::errs() << (isEq ? "=" : ">=") << " 0\n";
+  };
+
+  for (unsigned i = 0, e = getNumInequalities(); i < e; i++)
+    dumpConstraint(i, /*isEq=*/false);
+  for (unsigned i = 0, e = getNumEqualities(); i < e; i++)
+    dumpConstraint(i, /*isEq=*/true);
+  llvm::errs() << '\n';
+}
+
 std::pair<AffineMap, AffineMap> FlatLinearConstraints::getLowerAndUpperBound(
     unsigned pos, unsigned offset, unsigned num, unsigned symStartPos,
     ArrayRef<AffineExpr> localExprs, MLIRContext *context,
@@ -544,9 +596,9 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
   // Basic simplification.
   normalizeConstraintsByGCD();
 
-  LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for first " << num
-                          << " variables\n");
-  LLVM_DEBUG(dump());
+  LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for variables at positions ["
+                          << offset << ", " << offset + num << ")\n");
+  LLVM_DEBUG(dumpPretty());
 
   // Record computed/detected variables.
   SmallVector<AffineExpr, 8> memo(getNumVars());
@@ -699,12 +751,12 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
         }
       }
     }
-    LLVM_DEBUG(llvm::dbgs()
-               << "lb map for pos = " << Twine(pos + offset) << ", expr: ");
-    LLVM_DEBUG(lbMap.dump(););
-    LLVM_DEBUG(llvm::dbgs()
-               << "ub map for pos = " << Twine(pos + offset) << ", expr: ");
-    LLVM_DEBUG(ubMap.dump(););
+
+    LLVM_DEBUG(llvm::dbgs() << "Slice bounds:\n");
+    LLVM_DEBUG(llvm::dbgs() << "lb map for pos = " << Twine(pos + offset)
+                            << ", expr: " << lbMap << '\n');
+    LLVM_DEBUG(llvm::dbgs() << "ub map for pos = " << Twine(pos + offset)
+                            << ", expr: " << ubMap << '\n');
   }
 }
 

mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp

@@ -182,10 +182,9 @@ template bool mlir::affine::isInvariantAccess(AffineStoreOp, AffineForOp);
 DenseSet<Value> mlir::affine::getInvariantAccesses(Value iv,
                                                    ArrayRef<Value> indices) {
   DenseSet<Value> res;
-  for (auto val : indices) {
-    if (isAccessIndexInvariant(iv, val)) {
-      res.insert(val);
-    }
+  for (Value index : indices) {
+    if (isAccessIndexInvariant(iv, index))
+      res.insert(index);
   }
   return res;
 }

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

@@ -249,6 +249,7 @@ bool MemRefDependenceGraph::init() {
   // the memref.
   DenseMap<Value, SetVector<unsigned>> memrefAccesses;
 
+  // Create graph nodes.
   DenseMap<Operation *, unsigned> forToNodeMap;
   for (Operation &op : block) {
     if (auto forOp = dyn_cast<AffineForOp>(op)) {
@@ -300,11 +301,7 @@ bool MemRefDependenceGraph::init() {
     // interface.
   }
 
-  for (auto &idAndNode : nodes) {
-    LLVM_DEBUG(llvm::dbgs() << "Create node " << idAndNode.first << " for:\n"
-                            << *(idAndNode.second.op) << "\n");
-    (void)idAndNode;
-  }
+  LLVM_DEBUG(llvm::dbgs() << "Created " << nodes.size() << " nodes\n");
 
   // Add dependence edges between nodes which produce SSA values and their
   // users. Load ops can be considered as the ones producing SSA values.
@@ -1556,9 +1553,9 @@ mlir::affine::computeSliceUnion(ArrayRef<Operation *> opsA,
   FlatAffineValueConstraints sliceUnionCst;
   assert(sliceUnionCst.getNumDimAndSymbolVars() == 0);
   std::vector<std::pair<Operation *, Operation *>> dependentOpPairs;
-  for (auto *i : opsA) {
+  for (Operation *i : opsA) {
     MemRefAccess srcAccess(i);
-    for (auto *j : opsB) {
+    for (Operation *j : opsB) {
       MemRefAccess dstAccess(j);
       if (srcAccess.memref != dstAccess.memref)
         continue;
@@ -1587,7 +1584,7 @@ mlir::affine::computeSliceUnion(ArrayRef<Operation *> opsA,
 
       // Compute slice bounds for 'srcAccess' and 'dstAccess'.
       ComputationSliceState tmpSliceState;
-      mlir::affine::getComputationSliceState(i, j, &dependenceConstraints,
+      mlir::affine::getComputationSliceState(i, j, dependenceConstraints,
                                              loopDepth, isBackwardSlice,
                                              &tmpSliceState);
 
@@ -1646,8 +1643,10 @@ mlir::affine::computeSliceUnion(ArrayRef<Operation *> opsA,
   }
 
   // Empty union.
-  if (sliceUnionCst.getNumDimAndSymbolVars() == 0)
+  if (sliceUnionCst.getNumDimAndSymbolVars() == 0) {
+    LLVM_DEBUG(llvm::dbgs() << "empty slice union - unexpected\n");
     return SliceComputationResult::GenericFailure;
+  }
 
   // Gather loops surrounding ops from loop nest where slice will be inserted.
   SmallVector<Operation *, 4> ops;
@@ -1686,7 +1685,8 @@ mlir::affine::computeSliceUnion(ArrayRef<Operation *> opsA,
   sliceUnion->ivs.clear();
   sliceUnionCst.getValues(0, numSliceLoopIVs, &sliceUnion->ivs);
 
-  // Set loop nest insertion point to block start at 'loopDepth'.
+  // Set loop nest insertion point to block start at 'loopDepth' for forward
+  // slices, while at the end for backward slices.
   sliceUnion->insertPoint =
       isBackwardSlice
           ? surroundingLoops[loopDepth - 1].getBody()->begin()
@@ -1785,7 +1785,7 @@ const char *const kSliceFusionBarrierAttrName = "slice_fusion_barrier";
 // the other loop nest's IVs, symbols and constants (using 'isBackwardsSlice').
 void mlir::affine::getComputationSliceState(
     Operation *depSourceOp, Operation *depSinkOp,
-    FlatAffineValueConstraints *dependenceConstraints, unsigned loopDepth,
+    const FlatAffineValueConstraints &dependenceConstraints, unsigned loopDepth,
     bool isBackwardSlice, ComputationSliceState *sliceState) {
   // Get loop nest surrounding src operation.
   SmallVector<AffineForOp, 4> srcLoopIVs;
@@ -1804,30 +1804,28 @@ void mlir::affine::getComputationSliceState(
   unsigned pos = isBackwardSlice ? numSrcLoopIVs + loopDepth : loopDepth;
   unsigned num =
       isBackwardSlice ? numDstLoopIVs - loopDepth : numSrcLoopIVs - loopDepth;
-  dependenceConstraints->projectOut(pos, num);
+  FlatAffineValueConstraints sliceCst(dependenceConstraints);
+  sliceCst.projectOut(pos, num);
 
   // Add slice loop IV values to 'sliceState'.
   unsigned offset = isBackwardSlice ? 0 : loopDepth;
   unsigned numSliceLoopIVs = isBackwardSlice ? numSrcLoopIVs : numDstLoopIVs;
-  dependenceConstraints->getValues(offset, offset + numSliceLoopIVs,
-                                   &sliceState->ivs);
+  sliceCst.getValues(offset, offset + numSliceLoopIVs, &sliceState->ivs);
 
   // Set up lower/upper bound affine maps for the slice.
   sliceState->lbs.resize(numSliceLoopIVs, AffineMap());
   sliceState->ubs.resize(numSliceLoopIVs, AffineMap());
 
   // Get bounds for slice IVs in terms of other IVs, symbols, and constants.
-  dependenceConstraints->getSliceBounds(offset, numSliceLoopIVs,
-                                        depSourceOp->getContext(),
-                                        &sliceState->lbs, &sliceState->ubs);
+  sliceCst.getSliceBounds(offset, numSliceLoopIVs, depSourceOp->getContext(),
+                          &sliceState->lbs, &sliceState->ubs);
 
   // Set up bound operands for the slice's lower and upper bounds.
   SmallVector<Value, 4> sliceBoundOperands;
-  unsigned numDimsAndSymbols = dependenceConstraints->getNumDimAndSymbolVars();
+  unsigned numDimsAndSymbols = sliceCst.getNumDimAndSymbolVars();
   for (unsigned i = 0; i < numDimsAndSymbols; ++i) {
-    if (i < offset || i >= offset + numSliceLoopIVs) {
-      sliceBoundOperands.push_back(dependenceConstraints->getValue(i));
-    }
+    if (i < offset || i >= offset + numSliceLoopIVs)
+      sliceBoundOperands.push_back(sliceCst.getValue(i));
   }
 
   // Give each bound its own copy of 'sliceBoundOperands' for subsequent
@@ -2055,16 +2053,18 @@ static std::optional<int64_t> getMemoryFootprintBytes(Block &block,
     if (failed(
             region->compute(opInst,
                             /*loopDepth=*/getNestingDepth(&*block.begin())))) {
-      return opInst->emitError("error obtaining memory region\n");
+      LLVM_DEBUG(opInst->emitError("error obtaining memory region"));
+      return failure();
     }
 
     auto [it, inserted] = regions.try_emplace(region->memref);
     if (inserted) {
       it->second = std::move(region);
     } else if (failed(it->second->unionBoundingBox(*region))) {
-      return opInst->emitWarning(
+      LLVM_DEBUG(opInst->emitWarning(
           "getMemoryFootprintBytes: unable to perform a union on a memory "
-          "region");
+          "region"));
+      return failure();
     }
     return WalkResult::advance();
   });

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

@@ -836,6 +836,15 @@ struct GreedyFusion {
         // Get 'srcNode' from which to attempt fusion into 'dstNode'.
         auto *srcNode = mdg->getNode(srcId);
         auto srcAffineForOp = cast<AffineForOp>(srcNode->op);
+
+        LLVM_DEBUG(llvm::dbgs()
+                   << "Trying to fuse producer loop nest " << srcId
+                   << " with consumer loop nest " << dstId << "\n");
+        LLVM_DEBUG(llvm::dbgs()
+                   << "Producer loop nest:\n"
+                   << *srcNode->op << "\n and consumer loop nest:\n"
+                   << *dstNode->op << '\n');
+
         LLVM_DEBUG(llvm::dbgs() << "Evaluating src loop " << srcId
                                 << " for dst loop " << dstId << "\n");
 
@@ -904,9 +913,11 @@ struct GreedyFusion {
               affine::canFuseLoops(srcAffineForOp, dstAffineForOp,
                                    /*dstLoopDepth=*/i + numSurroundingLoops,
                                    &depthSliceUnions[i - 1], strategy);
-
-          if (result.value == FusionResult::Success)
+          if (result.value == FusionResult::Success) {
             maxLegalFusionDepth = i;
+            LLVM_DEBUG(llvm::dbgs()
+                       << "Found valid slice for depth: " << i << '\n');
+          }
         }
 
         if (maxLegalFusionDepth == 0) {

mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp

@@ -25,7 +25,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include <optional>
 
-#define DEBUG_TYPE "loop-fusion-utils"
+#define DEBUG_TYPE "affine-fusion-utils"
 
 using namespace mlir;
 using namespace mlir::affine;
@@ -49,12 +49,10 @@ static void getLoadAndStoreMemRefAccesses(Operation *opA,
 /// Returns false otherwise.
 static bool isDependentLoadOrStoreOp(Operation *op,
                                      DenseMap<Value, bool> &values) {
-  if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
+  if (auto loadOp = dyn_cast<AffineReadOpInterface>(op))
     return values.count(loadOp.getMemRef()) > 0 && values[loadOp.getMemRef()];
-  }
-  if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
+  if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op))
     return values.count(storeOp.getMemRef()) > 0;
-  }
   return false;
 }