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

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
+}

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();