[Mlir-commits] [mlir] dcfd35d - [mlir]: Added IntArrayNthElemMaxValue & IntArrayNthElemInRange predicates (NFC)

[Aviad Cohen]

Author: Aviad Cohen
Date: 2023-02-07T08:03:20+02:00
New Revision: dcfd35d7042e7979d7da4858542acbaa8600acc1

URL: https://github.com/llvm/llvm-project/commit/dcfd35d7042e7979d7da4858542acbaa8600acc1
DIFF: https://github.com/llvm/llvm-project/commit/dcfd35d7042e7979d7da4858542acbaa8600acc1.diff

LOG: [mlir]: Added IntArrayNthElemMaxValue & IntArrayNthElemInRange predicates (NFC)

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D143331

Added: 
    

Modified: 
    mlir/docs/DefiningDialects/Operations.md
    mlir/include/mlir/IR/OpBase.td
    mlir/test/mlir-tblgen/predicate.td

Removed: 
    


################################################################################
diff  --git a/mlir/docs/DefiningDialects/Operations.md b/mlir/docs/DefiningDialects/Operations.md
index 4b22084822495..03a2adbfcfd54 100644
--- a/mlir/docs/DefiningDialects/Operations.md
+++ b/mlir/docs/DefiningDialects/Operations.md
@@ -296,6 +296,10 @@ Right now, the following primitive constraints are supported:
     element to be equal to `N`
 *   `IntArrayNthElemMinValue<I, N>`: Specifying an integer array attribute's
     `I`-th element to be greater than or equal to `N`
+*   `IntArrayNthElemMaxValue<I, N>`: Specifying an integer array attribute's
+    `I`-th element to be less than or equal to `N`
+*   `IntArrayNthElemInRange<I, M, N>`: Specifying an integer array attribute's
+    `I`-th element to be greater than or equal to `M` and less than or equal to `N`
 
 TODO: Design and implement more primitive constraints
 

diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 7ef6268423e83..e833b0e3aec42 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1739,6 +1739,24 @@ class IntArrayNthElemMinValue<int index, int min> : AttrConstraint<
         ]>,
     "whose " # index # "-th element must be at least " # min>;
 
+class IntArrayNthElemMaxValue<int index, int max> : AttrConstraint<
+    And<[
+      CPred<"$_self.cast<::mlir::ArrayAttr>().size() > " # index>,
+      CPred<"$_self.cast<::mlir::ArrayAttr>()[" # index # "]"
+        ".cast<::mlir::IntegerAttr>().getInt() <= " # max>
+        ]>,
+    "whose " # index # "-th element must be at most " # max>;
+
+class IntArrayNthElemInRange<int index, int min, int max> : AttrConstraint<
+    And<[
+      CPred<"$_self.cast<::mlir::ArrayAttr>().size() > " # index>,
+      CPred<"$_self.cast<::mlir::ArrayAttr>()[" # index # "]"
+        ".cast<::mlir::IntegerAttr>().getInt() >= " # min>,
+      CPred<"$_self.cast<::mlir::ArrayAttr>()[" # index # "]"
+        ".cast<::mlir::IntegerAttr>().getInt() <= " # max>
+        ]>,
+    "whose " # index # "-th element must be at least " # min # " and at most " # max>;
+
 def IsNullAttr : AttrConstraint<
     CPred<"!$_self">, "empty attribute (for optional attributes)">;
 

diff  --git a/mlir/test/mlir-tblgen/predicate.td b/mlir/test/mlir-tblgen/predicate.td
index 07d9b4ffb7ba6..e9ab67b035f45 100644
--- a/mlir/test/mlir-tblgen/predicate.td
+++ b/mlir/test/mlir-tblgen/predicate.td
@@ -89,7 +89,23 @@ def OpI: NS_Op<"op_for_arr_min_value_at_index", []> {
 // CHECK: (((tblgen_attr.cast<::mlir::ArrayAttr>().size() > 0)) && ((tblgen_attr.cast<::mlir::ArrayAttr>()[0].cast<::mlir::IntegerAttr>().getInt() >= 8)))))
 // CHECK-NEXT: "attribute 'attr' failed to satisfy constraint: array attribute whose 0-th element must be at least 8"
 
-def OpJ: NS_Op<"op_for_TCopVTEtAreSameAt", [
+def OpJ: NS_Op<"op_for_arr_max_value_at_index", []> {
+  let arguments = (ins ConfinedAttr<ArrayAttr, [IntArrayNthElemMaxValue<0, 8>]>:$attr);
+}
+
+// CHECK-LABEL: OpJAdaptor::verify
+// CHECK: (((tblgen_attr.cast<::mlir::ArrayAttr>().size() > 0)) && ((tblgen_attr.cast<::mlir::ArrayAttr>()[0].cast<::mlir::IntegerAttr>().getInt() <= 8)))))
+// CHECK-NEXT: "attribute 'attr' failed to satisfy constraint: array attribute whose 0-th element must be at most 8"
+
+def OpK: NS_Op<"op_for_arr_in_range_at_index", []> {
+  let arguments = (ins ConfinedAttr<ArrayAttr, [IntArrayNthElemInRange<0, 4, 8>]>:$attr);
+}
+
+// CHECK-LABEL: OpKAdaptor::verify
+// CHECK: (((tblgen_attr.cast<::mlir::ArrayAttr>().size() > 0)) && ((tblgen_attr.cast<::mlir::ArrayAttr>()[0].cast<::mlir::IntegerAttr>().getInt() >= 4)) && ((tblgen_attr.cast<::mlir::ArrayAttr>()[0].cast<::mlir::IntegerAttr>().getInt() <= 8)))))
+// CHECK-NEXT: "attribute 'attr' failed to satisfy constraint: array attribute whose 0-th element must be at least 4 and at most 8"
+
+def OpL: NS_Op<"op_for_TCopVTEtAreSameAt", [
                 PredOpTrait<"operands indexed at 0, 2, 3 should all have "
                  "the same type", TCopVTEtAreSameAt<[0, 2, 3]>>]> {
   let arguments = (ins
@@ -101,28 +117,28 @@ def OpJ: NS_Op<"op_for_TCopVTEtAreSameAt", [
   );
 }
 
-// CHECK-LABEL: OpJAdaptor::verify
+// CHECK-LABEL: OpLAdaptor::verify
 // CHECK:      ::llvm::all_equal(::llvm::map_range(
 // CHECK-SAME:   ::mlir::ArrayRef<unsigned>({0, 2, 3}),
 // CHECK-SAME:   [this](unsigned i) { return getElementTypeOrSelf(this->getOperand(i)); }))
 // CHECK: "failed to verify that operands indexed at 0, 2, 3 should all have the same type"
 
-def OpK : NS_Op<"op_for_AnyTensorOf", []> {
+def OpM : NS_Op<"op_for_AnyTensorOf", []> {
   let arguments = (ins TensorOf<[F32, I32]>:$x);
 }
 
-// CHECK-LABEL: OpK::verify
+// CHECK-LABEL: OpM::verify
 // CHECK: auto valueGroup0 = getODSOperands(0);
 // CHECK: for (auto v : valueGroup0) {
 // CHECK: if (::mlir::failed([[$TENSOR_INTEGER_FLOAT_CONSTRAINT]]
 
-def OpL : NS_Op<"op_for_StringEscaping", []> {
+def OpN : NS_Op<"op_for_StringEscaping", []> {
   let arguments = (ins
     StringBasedAttr<CPred<"$_self.cast<StringAttr>().getValue() == \"foo\"">,
                     "only value \"foo\" is allowed">:$s
   );
 }
 
-// CHECK-LABEL: OpLAdaptor::verify
+// CHECK-LABEL: OpNAdaptor::verify
 // CHECK: getValue() == "foo"
 // CHECK-NEXT: only value \"foo\" is allowed