[Mlir-commits] [mlir] [mlir][scf] Fix `for-loop-peeling` crash (PR #77697)
Felix Schneider
llvmlistbot at llvm.org
Wed Jan 10 21:56:30 PST 2024
https://github.com/ubfx updated https://github.com/llvm/llvm-project/pull/77697
>From c58a040092ac716da3b4f8e47004534832fa6ce1 Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Wed, 10 Jan 2024 22:33:02 +0100
Subject: [PATCH 1/3] [mlir][affine] Add dependency on `UBDialect` for
`PoisonAttr`
The folder for `AffineApplyOp` [1] will try creating a `PoisonAttr`
under certain circumstances. However, this will result in a crash
if the `UBDialect` isn't loaded:
`LLVM ERROR: can't create Attribute 'mlir::ub::PoisonAttr' because
storage uniquer isn't initialized ...`
This patch adds a dependency of `AffineDialect` on `UBDialect`.
[1] https://github.com/llvm/llvm-project/blob/6d3ebd831c31d473acb18511949d04038115864a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp#L607
---
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h | 1 +
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td | 2 +-
mlir/test/Dialect/Affine/constant-fold.mlir | 11 ++++++++++-
3 files changed, 12 insertions(+), 2 deletions(-)
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/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
+}
>From 37a69d8444cdd3b6e1eb0d5d5133843b420a320f Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Wed, 10 Jan 2024 23:01:15 +0100
Subject: [PATCH 2/3] [mlir][scf] Fix crash in `for-loop-peeling`
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 https://github.com/llvm/llvm-project/issues/75758
Depends on https://github.com/llvm/llvm-project/pull/77691
---
.../SCF/Transforms/LoopSpecialization.cpp | 4 ++--
mlir/test/Dialect/SCF/for-loop-peeling.mlir | 19 +++++++++++++++++--
2 files changed, 19 insertions(+), 4 deletions(-)
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/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
+}
>From 9d7cc040427c5fe245a4f32c01eaf934601994c0 Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Thu, 11 Jan 2024 06:54:00 +0100
Subject: [PATCH 3/3] Bail on negative steps, fix formatting
---
mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 7aa1f768b303f7..342213507486af 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -123,8 +123,9 @@ 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. 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))
+ // 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)
return failure();
// No specialization necessary if step already divides upper bound evenly.
More information about the Mlir-commits
mailing list