[mlir] Add AnyAttrOf tablegen attribute constraint

AnyAttrOf, similar to AnyTypeOf, expects the attribute to be one of the
given attributes.
For instance, `AnyAttrOf<[I32Attr, StrAttr]>` expects either a `I32Attr`,
or a `StrAttr`.

Reviewed By: rriddle

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

mlir/include/mlir/IR/OpBase.td

@@ -955,6 +955,19 @@ def AnyAttr : Attr<CPred<"true">, "any attribute"> {
   let constBuilderCall = "$0";
 }
 
+// Any attribute from the given list
+class AnyAttrOf<list<Attr> allowedAttrs, string summary = "",
+                string cppClassName = "::mlir::Attribute",
+                string fromStorage = "$_self"> : Attr<
+    // Satisfy any of the allowed attribute's condition
+    Or<!foreach(allowedattr, allowedAttrs, allowedattr.predicate)>,
+    !if(!eq(summary, ""),
+        !interleave(!foreach(t, allowedAttrs, t.summary), " or "),
+        summary)> {
+    let returnType = cppClassName;
+    let convertFromStorage = fromStorage;
+}
+
 def BoolAttr : Attr<CPred<"$_self.isa<::mlir::BoolAttr>()">, "bool attribute"> {
   let storageType = [{ ::mlir::BoolAttr }];
   let returnType = [{ bool }];

mlir/test/IR/attribute.mlir

@@ -1,5 +1,36 @@
 // RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -verify-diagnostics | FileCheck %s
 
+//===----------------------------------------------------------------------===//
+// Test AnyAttrOf attributes
+//===----------------------------------------------------------------------===//
+
+func @any_attr_of_pass() {
+  "test.any_attr_of_i32_str"() {
+    // CHECK: attr = 3 : i32
+    attr = 3 : i32
+  } : () -> ()
+
+  "test.any_attr_of_i32_str"() {
+    // CHECK: attr = "string_data"
+    attr = "string_data"
+  } : () -> ()
+
+  return
+}
+
+// -----
+
+func @any_attr_of_fail() {
+  // expected-error @+1 {{'test.any_attr_of_i32_str' op attribute 'attr' failed to satisfy constraint: 32-bit signless integer attribute or string attribute}}
+  "test.any_attr_of_i32_str"() {
+    attr = 3 : i64
+  } : () -> ()
+
+  return
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // Test integer attributes
 //===----------------------------------------------------------------------===//

mlir/test/lib/Dialect/Test/TestOps.td

@@ -187,6 +187,10 @@ def MixedNormalVariadicResults : TEST_Op<
 // Test Attributes
 //===----------------------------------------------------------------------===//
 
+def AnyAttrOfOp : TEST_Op<"any_attr_of_i32_str"> {
+  let arguments = (ins AnyAttrOf<[I32Attr, StrAttr]>:$attr);
+}
+
 def NonNegIntAttrOp : TEST_Op<"non_negative_int_attr"> {
   let arguments = (ins
       Confined<I32Attr, [IntNonNegative]>:$i32attr,