[Mlir-commits] [mlir] c9c1b3c - [mlir][memref] Fix an invalid dim loop motion crash (#74204)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Dec 3 23:58:03 PST 2023
Author: Rik Huijzer
Date: 2023-12-04T08:57:59+01:00
New Revision: c9c1b3c37fa5d5c617068afe67afcafacd58a241
URL: https://github.com/llvm/llvm-project/commit/c9c1b3c37fa5d5c617068afe67afcafacd58a241
DIFF: https://github.com/llvm/llvm-project/commit/c9c1b3c37fa5d5c617068afe67afcafacd58a241.diff
LOG: [mlir][memref] Fix an invalid dim loop motion crash (#74204)
Fixes https://github.com/llvm/llvm-project/issues/73382.
This PR suggests to replace two assertions that were introduced in
https://github.com/llvm/llvm-project/commit/adabce41185910227ca276a1cfd22e76443dd238
(https://reviews.llvm.org/D135748). According to the enum definition of
`NotSpeculatable`, an op that invokes undefined behavior is
`NotSpeculatable`.
https://github.com/llvm/llvm-project/blob/0c06e8745f131d867c566f4d35a7a04e24b4a075/mlir/include/mlir/Interfaces/SideEffectInterfaces.h#L248-L258
and both `tensor.dim` and `memref.dim` state that "If the dimension
index is out of bounds, the behavior is undefined."
So therefore it seems to me that `DimOp::getSpeculatability()` should
return `NotSpeculatable` if the dimension index is out of bounds.
The added test is just a simplified version of
https://github.com/llvm/llvm-project/issues/73382.
Added:
Modified:
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
mlir/test/Transforms/loop-invariant-code-motion.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index dce96cca016ff..a397506629cfe 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -893,8 +893,9 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
- // The verifier rejects operations that violate this assertion.
- assert(constantIndex < rankedSourceType.getRank());
+ if (rankedSourceType.getRank() <= constantIndex)
+ return Speculation::NotSpeculatable;
+
return Speculation::Speculatable;
}
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 8970ea1c73b40..f15695383d34a 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -686,8 +686,9 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
- // The verifier rejects operations that violate this assertion.
- assert(constantIndex < rankedSourceType.getRank());
+ if (rankedSourceType.getRank() <= constantIndex)
+ return Speculation::NotSpeculatable;
+
return Speculation::Speculatable;
}
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index e9b1e22359001..1415583dde9da 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -699,6 +699,54 @@ func.func @speculate_memref_dim_known_rank_known_dim_inbounds(
// -----
+// CHECK-LABEL: @speculate_memref_dim_known_rank_known_dim_inbounds
+func.func @speculate_memref_dim_known_rank_known_dim_inbounds() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c22 = arith.constant 22 : index
+ %alloc = memref.alloc(%c22) : memref<?xi1>
+ scf.for %arg4 = %c0 to %c22 step %c1 {
+ %dim = memref.dim %alloc, %c0 : memref<?xi1>
+ }
+ return
+}
+// CHECK: memref.dim
+// CHECK-NEXT: scf.for
+
+// -----
+
+// CHECK-LABEL: @speculate_tensor_dim_known_rank_known_dim_inbounds
+func.func @speculate_tensor_dim_known_rank_known_dim_inbounds() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c22 = arith.constant 22 : index
+ %t = tensor.empty(%c22, %c22) : tensor<?x?xi1>
+ scf.for %arg4 = %c0 to %c22 step %c1 {
+ %dim = tensor.dim %t, %c1 : tensor<?x?xi1>
+ }
+ return
+}
+// CHECK: tensor.dim
+// CHECK-NEXT: scf.for
+
+// -----
+
+// CHECK-LABEL: @no_speculate_memref_dim_known_rank_known_dim_out_of_bounds
+func.func @no_speculate_memref_dim_known_rank_known_dim_out_of_bounds() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c22 = arith.constant 22 : index
+ %alloc = memref.alloc(%c22) : memref<?xi1>
+ scf.for %arg4 = %c0 to %c22 step %c1 {
+ %dim = memref.dim %alloc, %c1 : memref<?xi1>
+ }
+ return
+}
+// CHECK: scf.for
+// CHECK-NEXT: memref.dim
+
+// -----
+
func.func @no_speculate_divui(
// CHECK-LABEL: @no_speculate_divui(
%num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
More information about the Mlir-commits
mailing list