Add attr constraint support to constrain IntegerArray attribute values at a specified index.

--

PiperOrigin-RevId: 242209394
diff --git a/include/mlir/IR/OpBase.td b/include/mlir/IR/OpBase.td
index 602d50c..baaaa27 100644
--- a/include/mlir/IR/OpBase.td
+++ b/include/mlir/IR/OpBase.td
@@ -645,6 +645,22 @@
     CPred<"{0}.cast<ArrayAttr>().size() >= " # n>,
     "with at least " # n # " elements">;
 
+class IntArrayNthElemEq<int index, int value> : AttrConstraint<
+    AllOf<[
+      CPred<"{0}.cast<ArrayAttr>().size() > " # index>,
+      CPred<"{0}.cast<ArrayAttr>().getValue()[" # index # "]"
+        ".cast<IntegerAttr>().getInt() == " # value>
+       ]>,
+    "whose " # index # "-th element must be " # value>;
+
+class IntArrayNthElemMinValue<int index, int min> : AttrConstraint<
+    AllOf<[
+      CPred<"{0}.cast<ArrayAttr>().size() > " # index>,
+      CPred<"{0}.cast<ArrayAttr>().getValue()[" # index # "]"
+        ".cast<IntegerAttr>().getInt() >= " # min>
+        ]>,
+    "whose " # index # "-th element must be at least " # min>;
+
 //===----------------------------------------------------------------------===//
 // OpTrait definitions
 //===----------------------------------------------------------------------===//
diff --git a/test/mlir-tblgen/predicate.td b/test/mlir-tblgen/predicate.td
index 6b912af..6a9442a 100644
--- a/test/mlir-tblgen/predicate.td
+++ b/test/mlir-tblgen/predicate.td
@@ -70,3 +70,19 @@
 // CHECK-LABEL: OpG::verify()
 // CHECK:       (tblgen_attr.cast<ArrayAttr>().size() >= 8)
 // CHECK-SAME:    return emitOpError("attribute 'attr' failed to satisfy constraint: array attribute with at least 8 elements");
+
+def OpH : Op<"op_for_arr_value_at_index", []> {
+  let arguments = (ins Confined<ArrayAttr, [IntArrayNthElemEq<0, 8>]>:$attr);
+}
+
+// CHECK-LABEL: OpH::verify()
+// CHECK: (((tblgen_attr.cast<ArrayAttr>().size() > 0)) && ((tblgen_attr.cast<ArrayAttr>().getValue()[0].cast<IntegerAttr>().getInt() == 8)))))
+// CHECK-SAME:    return emitOpError("attribute 'attr' failed to satisfy constraint: array attribute whose 0-th element must be 8");
+
+def OpI: Op<"op_for_arr_min_value_at_index", []> {
+  let arguments = (ins Confined<ArrayAttr, [IntArrayNthElemMinValue<0, 8>]>:$attr);
+}
+
+// CHECK-LABEL: OpI::verify()
+// CHECK: (((tblgen_attr.cast<ArrayAttr>().size() > 0)) && ((tblgen_attr.cast<ArrayAttr>().getValue()[0].cast<IntegerAttr>().getInt() >= 8)))))
+// CHECK-SAME:    return emitOpError("attribute 'attr' failed to satisfy constraint: array attribute whose 0-th element must be at least 8");