From dbc3c1bd42fc32df642c0219bcd0f94f06f1acb9 Mon Sep 17 00:00:00 2001 From: Kai Sasaki 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 + + // expected-error@+1 {{'test.reify_bound' op invalid dim for shaped type}} + %dim = "test.reify_bound"(%tensor_val) {dim = 1 : i64} : (tensor) -> 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 dim = op.getDim(); + auto shapedType = dyn_cast(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 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 + + // expected-error@+1 {{'test.reify_bound' op dim must be non-negative}} + %dim = "test.reify_bound"(%tensor_val) {dim = -1 : i64} : (tensor) -> 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@+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 dim = op.getDim(); auto shapedType = dyn_cast(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 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 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 + + // expected-error@+1 {{'test.reify_bound' op expected 'dim' attribute for shaped type variable}} + %dim = "test.reify_bound"(%tensor_val) : (tensor) -> 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");