[mlir][scf] Fix for-loop-peeling crash
#77697
Merged
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-scf @llvm/pr-subscribers-mlir-affine Author: Felix Schneider (ubfx) ChangesBefore applying the peeling patterns, it can happen that the This patch adds an additional check for a constant-zero step and a Fix #75758 Depends on #77691 Full diff: https://github.com/llvm/llvm-project/pull/77697.diff 5 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index f070d048861906..b0a14c9cae4cd2 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -16,6 +16,7 @@
#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index c638646b9c3277..225e4d3194e230 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -24,7 +24,7 @@ def Affine_Dialect : Dialect {
let name = "affine";
let cppNamespace = "::mlir::affine";
let hasConstantMaterializer = 1;
- let dependentDialects = ["arith::ArithDialect"];
+ let dependentDialects = ["arith::ArithDialect", "ub::UBDialect"];
}
// Base class for Affine dialect ops.
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 9fda4861d40a3b..7aa1f768b303f7 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -123,8 +123,8 @@ static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp,
auto ubInt = getConstantIntValue(forOp.getUpperBound());
auto stepInt = getConstantIntValue(forOp.getStep());
- // No specialization necessary if step size is 1.
- if (getConstantIntValue(forOp.getStep()) == static_cast<int64_t>(1))
+ // No specialization necessary if step size is 1. Also bail out in case of a zero step which might have happened during folding.
+ if (stepInt == static_cast<int64_t>(1) || stepInt == static_cast<int64_t>(0))
return failure();
// No specialization necessary if step already divides upper bound evenly.
diff --git a/mlir/test/Dialect/Affine/constant-fold.mlir b/mlir/test/Dialect/Affine/constant-fold.mlir
index 5236b44ddfed96..ffc3946db08df5 100644
--- a/mlir/test/Dialect/Affine/constant-fold.mlir
+++ b/mlir/test/Dialect/Affine/constant-fold.mlir
@@ -64,7 +64,7 @@ func.func @affine_min(%variable: index) -> (index, index) {
// -----
func.func @affine_apply_poison_division_zero() {
- // This is just for mlir::context to load ub dailect
+ // This is just for mlir::context to load ub dialect
%ub = ub.poison : index
%c16 = arith.constant 16 : index
%0 = affine.apply affine_map<(d0)[s0] -> (d0 mod (s0 - s0))>(%c16)[%c16]
@@ -81,3 +81,12 @@ func.func @affine_apply_poison_division_zero() {
// CHECK-NEXT: %[[alloc:.*]] = memref.alloc(%[[poison]], %[[poison]], %[[poison]])
// CHECK-NEXT: %[[load:.*]] = affine.load %[[alloc]][%[[poison]], %[[poison]], %[[poison]]] : memref<?x?x?xi1>
// CHECK-NEXT: affine.store %[[load]], %alloc[%[[poison]], %[[poison]], %[[poison]]] : memref<?x?x?xi1>
+
+// -----
+
+// Check that this doesn't crash because the ub dialect is automatically loaded
+func.func @affine_apply_poison_division_zero_2() -> index {
+ %c16 = arith.constant 16 : index
+ %0 = affine.apply affine_map<(d0)[s0] -> (d0 mod (s0 - s0))>(%c16)[%c16]
+ return %0 : index
+}
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling.mlir b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
index 9a6d1c8c0a14cd..aa8fd7ec7ef84d 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt %s -scf-for-loop-peeling -canonicalize -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -scf-for-loop-peeling=skip-partial=false -canonicalize -split-input-file | FileCheck %s -check-prefix=CHECK-NO-SKIP
+// RUN: mlir-opt %s -scf-for-loop-peeling -canonicalize -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt %s -scf-for-loop-peeling=skip-partial=false -canonicalize -split-input-file -verify-diagnostics | FileCheck %s -check-prefix=CHECK-NO-SKIP
// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (-s0 + s1) mod s2)>
// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>
@@ -289,3 +289,18 @@ func.func @regression(%arg0: memref<i64>, %arg1: index) {
}
return
}
+
+// -----
+
+// Check that this doesn't crash but trigger the verifier.
+func.func @zero_step(%arg0: memref<i64>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %foldto0 = arith.subi %c1, %c1 : index
+ // expected-error @+1 {{'scf.for' op constant step operand must be positive}}
+ scf.for %arg2 = %c0 to %c1 step %foldto0 {
+ %2 = arith.index_cast %arg2 : index to i64
+ memref.store %2, %arg0[] : memref<i64>
+ }
+ return
+}
|
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
|
Unrelated to this bugfix, but invalid step probably shouldn't cause verification failure, as it can appears as folding result in unreachable branches. Can you also add test for negative step, which is also invalid? |
Good point, how could this be done with the
Done! |
matthias-springer
approved these changes
Jan 11, 2024
| if (getConstantIntValue(forOp.getStep()) == static_cast<int64_t>(1)) | ||
| // No specialization necessary if step size is 1. Also bail out in case of an | ||
| // invalid zero or negative step which might have happened during folding. | ||
| if (stepInt && *stepInt <= 1) |
There was a problem hiding this comment.
Interesting, I thought that negative step sizes were allowed, but the documentations says otherwise: The step is a value of same type but required to be positive.
The easiest solution is just to remove corresponding check from verifier and check We can potentially drop the loop entirely and replace results with poisons, but it will introduce dependence on UB dialect and need some discussion, I think. I've created an issue to track #77755 |
Hardcode84
approved these changes
Jan 11, 2024
Before applying the peeling patterns, it can happen that the `ForOp` gets a step of zero during folding. This leads to a division-by-zero down the line. This patch adds an additional check for a constant-zero step and a test. Fix llvm#75758 Depends on llvm#77691
ubfx
force-pushed
the
fix-scf-loop-peeling-crash
branch
2 times, most recently
from
6547003 to
7e5c24c
Compare
justinfargnoli
pushed a commit
to justinfargnoli/llvm-project
that referenced
this pull request
Jan 28, 2024
Before applying the peeling patterns, it can happen that the `ForOp` gets a step of zero during folding. This leads to a division-by-zero down the line. This patch adds an additional check for a constant-zero step and a test. Fix llvm#75758
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.
Before applying the peeling patterns, it can happen that the
ForOpgets a step of zero during folding. This leads to a division-by-zero
down the line.
This patch adds an additional check for a constant-zero step and a
test.
Fix #75758