Extend loop-fusion's slicing utility + other fixes / updates

- refactor toAffineFromEq and the code surrounding it; refactor code into
  FlatAffineConstraints::getSliceBounds
- add FlatAffineConstraints methods to detect identifiers as mod's and div's of other
  identifiers
- add FlatAffineConstraints::getConstantLower/UpperBound
- Address b/122118218 (don't assert on invalid fusion depths cmdline flags -
  instead, don't do anything; change cmdline flags
  src-loop-depth -> fusion-src-loop-depth
- AffineExpr/Map print method update: don't fail on null instances (since we have
  a wrapper around a pointer, it's avoidable); rationale: dump/print methods should
  never fail if possible.
- Update memref-dataflow-opt to add an optimization to avoid a unnecessary call to
  IsRangeOneToOne when it's trivially going to be true.
- Add additional test cases to exercise the new support
- update a few existing test cases since the maps are now generated uniformly with
  all destination loop operands appearing for the backward slice
- Fix projectOut - fix wrong range for getBestElimCandidate.
- Fix for getConstantBoundOnDimSize() - didn't show up in any test cases since
  we didn't have any non-hyperrectangular ones.

PiperOrigin-RevId: 228265152

1 files changed, 12 insertions, 5 deletions

diff --git a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
index adf91b76276..1a1502de738 100644
--- a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
+++ b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
@@ -117,6 +117,8 @@ void MemRefDataFlowOpt::visitOperationInst(OperationInst *opInst) {
     storeOps.push_back(storeOpInst);
   }
 
+  unsigned loadOpDepth = getNestingDepth(*loadOpInst);
+
   // 1. Check if there is a dependence satisfied at depth equal to the depth
   // of the loop body of the innermost common surrounding loop of the storeOp
   // and loadOp.
@@ -165,11 +167,16 @@ void MemRefDataFlowOpt::visitOperationInst(OperationInst *opInst) {
       //   a [i] = ...
       //   for (j ...)
       //      ... = a[j]
-      MemRefRegion region;
-      getMemRefRegion(loadOpInst, nsLoops, &region);
-      if (!region.getConstraints()->isRangeOneToOne(
-              /*start=*/0, /*limit=*/loadOp->getMemRefType().getRank()))
-        break;
+      // If storeOpInst and loadOpDepth at the same nesting depth, the load Op
+      // is trivially loading from a single location at that depth; so there
+      // isn't a need to call isRangeOneToOne.
+      if (getNestingDepth(*storeOpInst) < loadOpDepth) {
+        MemRefRegion region;
+        getMemRefRegion(loadOpInst, nsLoops, &region);
+        if (!region.getConstraints()->isRangeOneToOne(
+                /*start=*/0, /*limit=*/loadOp->getMemRefType().getRank()))
+          break;
+      }
 
       // After all these conditions, we have a candidate for forwarding!
       fwdingCandidates.push_back(storeOpInst);