[MLIR][Affine] Fix affine scalrep - add missing check

[bondhugula]

Summary

This fixes https://github.com/llvm/llvm-project/issues/59461

Add missing check in affine-scalrep pass that led to scalrep assert or
wrong scalrep when dead affine region ops existed in the same block.

[bondhugula]

Remove extra new line.

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

LogicalResult mlir::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) {
   AffineMap ivMap = AffineMap::get(origLbMap.getNumDims() + 1,
                                    origLbMap.getNumSymbols(), newIVExpr);
   canonicalizeMapAndOperands(&ivMap, &lbOperands);
   Operation *newIV = opBuilder.create<AffineApplyOp>(loc, ivMap, lbOperands);
   op.getInductionVar().replaceAllUsesExcept(newIV->getResult(0), newIV);
   return success();
 }
 
+/// Returns true if the memory operation of `destAccess` depends on `srcAccess`
+/// inside of the innermost common surrounding affine loop between the two
+/// accesses.
+static bool mustReachAtInnermost(const MemRefAccess &srcAccess,
+                                 const MemRefAccess &destAccess) {
+  // Affine dependence analysis is possible only if both ops in the same
+  // AffineScope.
+  if (getAffineScope(srcAccess.opInst) != getAffineScope(destAccess.opInst))
+    return false;
+
+  unsigned nsLoops =
+      getNumCommonSurroundingLoops(*srcAccess.opInst, *destAccess.opInst);
+  FlatAffineValueConstraints dependenceConstraints;
+  DependenceResult result = checkMemrefAccessDependence(
+      srcAccess, destAccess, nsLoops + 1, &dependenceConstraints,
+      /*dependenceComponents=*/nullptr);
+  return hasDependence(result);
+}
+
 /// Returns true if `srcMemOp` may have an effect on `destMemOp` within the
 /// scope of the outermost `minSurroundingLoops` loops that surround them.
 /// `srcMemOp` and `destMemOp` are expected to be affine read/write ops.
 static bool mayHaveEffect(Operation *srcMemOp, Operation *destMemOp,
                           unsigned minSurroundingLoops) {
   MemRefAccess srcAccess(srcMemOp);
   MemRefAccess destAccess(destMemOp);
 
for (auto *user : loadOp.getMemRef().getUsers()) {
     // 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(storeOp, loadOp))
       continue;
 
-    // 3. Ensure there is no intermediate operation which could replace the
+    // 3. The store must reach the load. Access function equivalence only
+    // guarantees this for accesses in the same block. The load could be in a
+    // nested block that is unreachable.
+    if (storeOp->getBlock() != loadOp->getBlock() &&
+        !mustReachAtInnermost(srcAccess, destAccess))
+      continue;
+
+    // 4. Ensure there is no intermediate operation which could replace the
     // value in memory.
     if (!mlir::hasNoInterveningEffect<MemoryEffects::Write>(storeOp, loadOp))
       continue;
 
     // We now have a candidate for forwarding.
     assert(lastWriteStoreOp == nullptr &&
            "multiple simultaneous replacement stores");
     lastWriteStoreOp = storeOp;

mlir/test/Dialect/Affine/scalrep.mlir

 // CHECK:       %[[C7:.*]] = arith.constant 7.000000e+00 : f32
 // CHECK-NEXT:  affine.parallel (%{{.*}}) = (0) to (10) {
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:      arith.addf %[[C7]], %[[C7]] : f32
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
 }
+
+// CHECK-LABEL: func.func @dead_affine_region_op
+func.func @dead_affine_region_op() {
+  %c1 = arith.constant 1 : index
+  %alloc = memref.alloc() : memref<15xi1>
+  %true = arith.constant true
+  affine.store %true, %alloc[%c1] : memref<15xi1>
+  // Dead store.
+  affine.store %true, %alloc[%c1] : memref<15xi1>
+  // This affine.if is dead.
+  affine.if affine_set<(d0, d1, d2, d3) : ((d0 + 1) mod 8 >= 0, d0 * -8 >= 0)>(%c1, %c1, %c1, %c1){
+    // No forwarding will happen.
+    affine.load %alloc[%c1] : memref<15xi1>
+  }
+  // CHECK-NEXT: arith.constant
+  // CHECK-NEXT: memref.alloc
+  // CHECK-NEXT: arith.constant
+  // CHECK-NEXT: affine.store
+  // CHECK-NEXT: affine.if
+  // CHECK-NEXT:   affine.load
+  return
+}