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