[Mlir-commits] [mlir] 1fb096e - [MLIR][Affine] Fix affine scalrep - add missing check

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2023-01-12T11:48:25+05:30
New Revision: 1fb096e0e3f78b44991d90850d3189a812341267

URL: https://github.com/llvm/llvm-project/commit/1fb096e0e3f78b44991d90850d3189a812341267
DIFF: https://github.com/llvm/llvm-project/commit/1fb096e0e3f78b44991d90850d3189a812341267.diff

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

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.

Differential Revision: https://reviews.llvm.org/D141255

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Utils/Utils.cpp
    mlir/test/Dialect/Affine/scalrep.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 74a6d961c695f..ece17badbba93 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -638,6 +638,25 @@ LogicalResult mlir::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) {
   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.
@@ -858,7 +877,14 @@ static LogicalResult forwardStoreToLoad(
     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;

diff  --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index 972133a4b4a01..c5862be6a6e97 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -846,3 +846,25 @@ func.func @parallel_surrounding_for() {
 // 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
+}