[Mlir-commits] [mlir] [mlir][memref] Fix an invalid dim loop motion crash (PR #74204)
Rik Huijzer
llvmlistbot at llvm.org
Sun Dec 3 00:58:19 PST 2023
https://github.com/rikhuijzer updated https://github.com/llvm/llvm-project/pull/74204
>From c11845000eb41f12167f2a966e9bb7c1bb14c997 Mon Sep 17 00:00:00 2001
From: Rik Huijzer <github at huijzer.xyz>
Date: Sat, 2 Dec 2023 19:29:19 +0100
Subject: [PATCH 1/5] [mlir][memref] Fix an invalid dim loop motion crash
---
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp | 2 --
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 2 --
.../Transforms/loop-invariant-code-motion.mlir | 17 +++++++++++++++++
3 files changed, 17 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index dce96cca016ff..b9c2e61180652 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -893,8 +893,6 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
- // The verifier rejects operations that violate this assertion.
- assert(constantIndex < rankedSourceType.getRank());
return Speculation::Speculatable;
}
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 8970ea1c73b40..c4bbf64381eab 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -686,8 +686,6 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
- // The verifier rejects operations that violate this assertion.
- assert(constantIndex < rankedSourceType.getRank());
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..1b2983916dd2a 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -699,6 +699,23 @@ func.func @speculate_memref_dim_known_rank_known_dim_inbounds(
// -----
+// CHECK-LABEL: @speculate_tensor_dim_known_rank_known_dim_out_of_bounds
+func.func @speculate_tensor_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) {alignment = 64 : i64} : memref<?xi1>
+ %4 = tensor.empty(%c22, %c22, %c22) : tensor<?x?x?xi1>
+ scf.for %arg4 = %c0 to %c22 step %c1 {
+ %dim = memref.dim %alloc, %c22 : memref<?xi1>
+ }
+ spirv.Return
+}
+// CHECK: memref.dim
+// CHECK-NEXT: scf.for
+
+// -----
+
func.func @no_speculate_divui(
// CHECK-LABEL: @no_speculate_divui(
%num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
>From cd154ac875578a16c18e35ff1e4e840fdfe10544 Mon Sep 17 00:00:00 2001
From: Rik Huijzer <github at huijzer.xyz>
Date: Sat, 2 Dec 2023 19:46:20 +0100
Subject: [PATCH 2/5] Set to `NotSpeculatable`
---
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp | 3 +++
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 3 +++
mlir/test/Transforms/loop-invariant-code-motion.mlir | 6 +++---
3 files changed, 9 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index b9c2e61180652..e621ea9d5443e 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -893,6 +893,9 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
+ 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 c4bbf64381eab..54fecc0639028 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -686,6 +686,9 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
+ 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 1b2983916dd2a..8013384b31add 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -709,10 +709,10 @@ func.func @speculate_tensor_dim_known_rank_known_dim_out_of_bounds() {
scf.for %arg4 = %c0 to %c22 step %c1 {
%dim = memref.dim %alloc, %c22 : memref<?xi1>
}
- spirv.Return
+ return
}
-// CHECK: memref.dim
-// CHECK-NEXT: scf.for
+// CHECK: scf.for
+// CHECK-NEXT: memref.dim
// -----
>From 5962f83c6f0a10814aac385785fd09c0e0af34ef Mon Sep 17 00:00:00 2001
From: Rik Huijzer <github at huijzer.xyz>
Date: Sat, 2 Dec 2023 19:50:51 +0100
Subject: [PATCH 3/5] Change name
---
mlir/test/Transforms/loop-invariant-code-motion.mlir | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index 8013384b31add..d5a0c4e459e40 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -699,8 +699,8 @@ func.func @speculate_memref_dim_known_rank_known_dim_inbounds(
// -----
-// CHECK-LABEL: @speculate_tensor_dim_known_rank_known_dim_out_of_bounds
-func.func @speculate_tensor_dim_known_rank_known_dim_out_of_bounds() {
+// CHECK-LABEL: @no_speculate_tensor_dim_known_rank_known_dim_out_of_bounds
+func.func @no_speculate_tensor_dim_known_rank_known_dim_out_of_bounds() {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c22 = arith.constant 22 : index
>From 53e2bcd0cc42b518d45f7d8c71373b3da633ec4a Mon Sep 17 00:00:00 2001
From: Rik Huijzer <github at huijzer.xyz>
Date: Sun, 3 Dec 2023 09:45:51 +0100
Subject: [PATCH 4/5] Fix off-by-one error
---
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp | 2 +-
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 2 +-
.../loop-invariant-code-motion.mlir | 41 ++++++++++++++++---
3 files changed, 38 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index e621ea9d5443e..a397506629cfe 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -893,7 +893,7 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
- if (rankedSourceType.getRank() < constantIndex)
+ 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 54fecc0639028..f15695383d34a 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -686,7 +686,7 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
- if (rankedSourceType.getRank() < constantIndex)
+ 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 d5a0c4e459e40..1415583dde9da 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -699,15 +699,46 @@ func.func @speculate_memref_dim_known_rank_known_dim_inbounds(
// -----
-// CHECK-LABEL: @no_speculate_tensor_dim_known_rank_known_dim_out_of_bounds
-func.func @no_speculate_tensor_dim_known_rank_known_dim_out_of_bounds() {
+// 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) {alignment = 64 : i64} : memref<?xi1>
- %4 = tensor.empty(%c22, %c22, %c22) : tensor<?x?x?xi1>
+ %alloc = memref.alloc(%c22) : memref<?xi1>
scf.for %arg4 = %c0 to %c22 step %c1 {
- %dim = memref.dim %alloc, %c22 : memref<?xi1>
+ %dim = memref.dim %alloc, %c1 : memref<?xi1>
}
return
}
More information about the Mlir-commits
mailing list