[mlir][affine] Fix the crash due to the simultaneous replacement store (#90829)

`AffineScalarReplacement` should forward the memref store op to load op
only if the store op reaches the load. But it now checks the
reachability only if these ops are in the same block, which causes the
crash reported in #76309.

We need to check the reachability even if they are both in the same
block, which rescues the case where consecutive store operations are
written before the load op.

Author: Lewuathe

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

@@ -860,8 +860,7 @@ static void forwardStoreToLoad(
     // 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))
+    if (!mustReachAtInnermost(srcAccess, destAccess))
       continue;
 
     // 4. Ensure there is no intermediate operation which could replace the

mlir/test/Dialect/Affine/scalrep.mlir

@@ -5,6 +5,7 @@
 // CHECK-DAG: [[$MAP2:#map[0-9]*]] = affine_map<(d0, d1) -> (d1)>
 // CHECK-DAG: [[$MAP3:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 - 1)>
 // CHECK-DAG: [[$MAP4:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
+// CHECK-DAG: [[$IDENT:#map[0-9]*]] = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: func @simple_store_load() {
 func.func @simple_store_load() {
@@ -931,3 +932,23 @@ func.func @cross_block() {
   %69 = affine.load %alloc_99[%c10] : memref<13xi1>
   return
 }
+
+#map1 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @consecutive_store
+func.func @consecutive_store() {
+  // CHECK: %[[CST:.*]] = arith.constant
+  %tmp = arith.constant 1.1 : f16
+  // CHECK: %[[ALLOC:.*]] = memref.alloc
+  %alloc_66 = memref.alloc() : memref<f16, 1>
+  affine.for %arg2 = 4 to 6 {
+    affine.for %arg3 = #map1(%arg2) to #map1(%arg2) step 4 {
+      // CHECK: affine.store %[[CST]], %[[ALLOC]][]
+      affine.store %tmp, %alloc_66[] : memref<f16, 1>
+      // CHECK-NOT: affine.store %[[CST]], %[[ALLOC]][]
+      affine.store %tmp, %alloc_66[] : memref<f16, 1>
+      %270 = affine.load %alloc_66[] : memref<f16, 1>
+    }
+  }
+  return
+}