[mlir][affine] Fix the crash due to the simultaneous replacement store #90829
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir Author: Kai Sasaki (Lewuathe) Changes
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. Full diff: https://github.com/llvm/llvm-project/pull/90829.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 8b8ed2578ca5cc..eba6446c5f8896 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -866,8 +866,7 @@ static void forwardStoreToLoad(AffineReadOpInterface loadOp,
// 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
diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index 22d394bfcf0979..6432efb6a435bb 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/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() {
@@ -913,3 +914,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
+}
|
|
@llvm/pr-subscribers-mlir-affine Author: Kai Sasaki (Lewuathe) Changes
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. Full diff: https://github.com/llvm/llvm-project/pull/90829.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 8b8ed2578ca5cc..eba6446c5f8896 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -866,8 +866,7 @@ static void forwardStoreToLoad(AffineReadOpInterface loadOp,
// 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
diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index 22d394bfcf0979..6432efb6a435bb 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/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() {
@@ -913,3 +914,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
+}
|
|
Ping~ |
Lewuathe
commented
Sep 5, 2024
There was a problem hiding this comment.
@joker-eph @matthias-springer @River707 @bondhugula Sorry for bothering you. But could review this PR when you get a chance?
matthias-springer
approved these changes
Sep 5, 2024
|
@matthias-springer Thank you! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
AffineScalarReplacementshould 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.