[mlir][affine] Guard invalid dim attribute in the test-reify-bound pass (#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 #128807

Author: Lewuathe

Diff for: mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir

@@ -0,0 +1,47 @@
+// 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
+    %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
+}
+
+// -----
+
+func.func @test_invalid_reify_negative_dim(%size: index) -> (index) {
+    %zero = arith.constant 0 : index
+    %tensor_val = tensor.empty(%size) : tensor<?xf32>
+
+    // expected-error@+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@+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
+}
+
+// -----
+
+func.func @test_invalid_reify_without_dim(%size: index) -> (index) {
+    %zero = arith.constant 0 : index
+    %tensor_val = tensor.empty(%size) : tensor<?xf32>
+
+    // expected-error@+1 {{'test.reify_bound' op expected 'dim' attribute for shaped type variable}}
+    %dim = "test.reify_bound"(%tensor_val) : (tensor<?xf32>) -> index
+
+    return %dim: index
+}

Diff for: mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp

@@ -84,6 +84,27 @@ 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 && dim.has_value()) {
+      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");
+        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();