[Mlir-commits] [mlir] [mlir][affine] Guard invalid dim attribute in the test-reify-bound pass (PR #129013)
Kai Sasaki
llvmlistbot at llvm.org
Wed Mar 12 19:54:35 PDT 2025
https://github.com/Lewuathe updated https://github.com/llvm/llvm-project/pull/129013
>From dbc3c1bd42fc32df642c0219bcd0f94f06f1acb9 Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Thu, 27 Feb 2025 15:51:01 +0900
Subject: [PATCH 1/4] [mlir][affine] Guard invalid dim attribute in the
test-reify-bound pass
Computing the bound of affine op (ValueBoundsConstraintSet::computeBound)
crashes due to the invalid dim value given to the op.
It is necessary for the pass to check the dim attribute not to be
greater than the rank of the input type.
Fixes https://github.com/llvm/llvm-project/issues/128807
---
.../Dialect/Affine/invalid-reify-bound-dim.mlir | 13 +++++++++++++
.../lib/Dialect/Affine/TestReifyValueBounds.cpp | 7 +++++++
2 files changed, 20 insertions(+)
create mode 100644 mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
diff --git a/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir b/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
new file mode 100644
index 0000000000000..8c878b2664042
--- /dev/null
+++ b/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds))' -verify-diagnostics
+
+// -----
+
+func.func @test_invalid_reify_dim(%size: index) -> (index) {
+ %zero = arith.constant 0 : index
+ %tensor_val = tensor.empty(%size) : tensor<?xf32>
+
+ // expected-error at +1 {{'test.reify_bound' op invalid dim for shaped type}}
+ %dim = "test.reify_bound"(%tensor_val) {dim = 1 : i64} : (tensor<?xf32>) -> index
+
+ return %dim: index
+}
diff --git a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
index 891b3bab8629d..fc44da4b53865 100644
--- a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
@@ -84,6 +84,13 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
auto boundType = op.getBoundType();
Value value = op.getVar();
std::optional<int64_t> dim = op.getDim();
+ auto shapedType = dyn_cast<ShapedType>(value.getType());
+ if (shapedType && shapedType.hasRank() && dim.has_value() &&
+ dim.value() >= shapedType.getRank()) {
+ op->emitOpError("invalid dim for shaped type");
+ return WalkResult::interrupt();
+ }
+
bool constant = op.getConstant();
bool scalable = op.getScalable();
>From 8ca082930f6fcea704e36757dbd2831da4acf477 Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Tue, 11 Mar 2025 14:22:40 +0900
Subject: [PATCH 2/4] Post review follow-up
---
.../Affine/invalid-reify-bound-dim.mlir | 28 +++++++++++++++++--
.../Dialect/Affine/TestReifyValueBounds.cpp | 16 +++++++++--
2 files changed, 38 insertions(+), 6 deletions(-)
diff --git a/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir b/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
index 8c878b2664042..ff16f1d76a9ca 100644
--- a/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
+++ b/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
@@ -1,6 +1,4 @@
-// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds))' -verify-diagnostics
-
-// -----
+// RUN: mlir-opt -split-input-file %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds))' -verify-diagnostics
func.func @test_invalid_reify_dim(%size: index) -> (index) {
%zero = arith.constant 0 : index
@@ -11,3 +9,27 @@ func.func @test_invalid_reify_dim(%size: index) -> (index) {
return %dim: index
}
+
+// -----
+
+func.func @test_invalid_reify_negative_dim(%size: index) -> (index) {
+ %zero = arith.constant 0 : index
+ %tensor_val = tensor.empty(%size) : tensor<?xf32>
+
+ // expected-error at +1 {{'test.reify_bound' op dim must be non-negative}}
+ %dim = "test.reify_bound"(%tensor_val) {dim = -1 : i64} : (tensor<?xf32>) -> index
+
+ return %dim: index
+}
+
+// -----
+
+func.func @test_invalid_reify_int_value(%size: index) -> (index) {
+ %zero = arith.constant 0 : index
+ %int_val = arith.constant 1 : index
+
+ // expected-error at +1 {{'test.reify_bound' op unexpected 'dim' attribute for index variable}}
+ %dim = "test.reify_bound"(%int_val) {dim = 1 : i64} : (index) -> index
+
+ return %dim: index
+}
diff --git a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
index fc44da4b53865..975171f405e46 100644
--- a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
@@ -85,11 +85,21 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
Value value = op.getVar();
std::optional<int64_t> dim = op.getDim();
auto shapedType = dyn_cast<ShapedType>(value.getType());
- if (shapedType && shapedType.hasRank() && dim.has_value() &&
- dim.value() >= shapedType.getRank()) {
- op->emitOpError("invalid dim for shaped type");
+ if (!shapedType && dim.has_value()) {
+ op->emitOpError("dim specified for non-shaped type");
return WalkResult::interrupt();
}
+ if (shapedType && shapedType.hasRank() && dim.has_value()) {
+ if (dim.value() < 0) {
+ op->emitOpError("dim must be non-negative");
+ return WalkResult::interrupt();
+ }
+
+ if(dim.value() >= shapedType.getRank()) {
+ op->emitOpError("invalid dim for shaped type rank");
+ return WalkResult::interrupt();
+ }
+ }
bool constant = op.getConstant();
bool scalable = op.getScalable();
>From c7c85d000f63f628755e8f5db291b5cf4479e57d Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Tue, 11 Mar 2025 14:47:00 +0900
Subject: [PATCH 3/4] Fix format
---
mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
index 975171f405e46..4231e3917972c 100644
--- a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
@@ -95,7 +95,7 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
return WalkResult::interrupt();
}
- if(dim.value() >= shapedType.getRank()) {
+ if (dim.value() >= shapedType.getRank()) {
op->emitOpError("invalid dim for shaped type rank");
return WalkResult::interrupt();
}
>From 730e0f8a691f3f4077073d312136c05ae1f2136b Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Wed, 12 Mar 2025 14:34:40 +0900
Subject: [PATCH 4/4] Add the case without dim attributes
---
.../test/Dialect/Affine/invalid-reify-bound-dim.mlir | 12 ++++++++++++
.../test/lib/Dialect/Affine/TestReifyValueBounds.cpp | 4 ++++
2 files changed, 16 insertions(+)
diff --git a/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir b/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
index ff16f1d76a9ca..e923ac19c5177 100644
--- a/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
+++ b/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
@@ -33,3 +33,15 @@ func.func @test_invalid_reify_int_value(%size: index) -> (index) {
return %dim: index
}
+
+// -----
+
+func.func @test_invalid_reify_without_dim(%size: index) -> (index) {
+ %zero = arith.constant 0 : index
+ %tensor_val = tensor.empty(%size) : tensor<?xf32>
+
+ // expected-error at +1 {{'test.reify_bound' op expected 'dim' attribute for shaped type variable}}
+ %dim = "test.reify_bound"(%tensor_val) : (tensor<?xf32>) -> index
+
+ return %dim: index
+}
\ No newline at end of file
diff --git a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
index 4231e3917972c..e1c0e878523a9 100644
--- a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
@@ -89,6 +89,10 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
op->emitOpError("dim specified for non-shaped type");
return WalkResult::interrupt();
}
+ if (shapedType && !dim.has_value()) {
+ op->emitOpError("dim not specified for shaped type");
+ return WalkResult::interrupt();
+ }
if (shapedType && shapedType.hasRank() && dim.has_value()) {
if (dim.value() < 0) {
op->emitOpError("dim must be non-negative");
More information about the Mlir-commits
mailing list