[Mlir-commits] [mlir] [mlir][affine] Guard invalid dim attribute in the test-reify-bound pass (PR #129013)

[Kai Sasaki]

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/2] [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/2] 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();