Skip to content

[mlir][affine] Guard invalid dim attribute in the test-reify-bound pass #129013

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 14 commits into from
Mar 13, 2025

Conversation

Lewuathe
Copy link
Member

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 #128807

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 llvm#128807
@llvmbot
Copy link
Member

llvmbot commented Feb 27, 2025

@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: Kai Sasaki (Lewuathe)

Changes

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 #128807


Full diff: https://github.com/llvm/llvm-project/pull/129013.diff

2 Files Affected:

  • (added) mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir (+13)
  • (modified) mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp (+7)
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@+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();
 

Copy link

github-actions bot commented Mar 11, 2025

✅ With the latest revision this PR passed the C/C++ code formatter.

@Lewuathe Lewuathe merged commit befa037 into llvm:main Mar 13, 2025
11 checks passed
@Lewuathe Lewuathe deleted the check-dim-attribute-in-test-reify-bound branch March 13, 2025 23:09
@Lewuathe
Copy link
Member Author

@matthias-springer Thank you for the review!

frederik-h pushed a commit to frederik-h/llvm-project that referenced this pull request Mar 18, 2025
…ss (llvm#129013)

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 llvm#128807
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[MLIR] Assertion `*dim < shapedType.getRank() && "invalid dim value"' failed.
3 participants