From d78136121eb1ddbc416648f9f08be2e87709b06a Mon Sep 17 00:00:00 2001
From: Mathieu Fehr <mathieu.fehr@gmail.com>
Date: Mon, 18 Oct 2021 15:37:26 +0000
Subject: [PATCH] [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        | 13 +++++++++++
 mlir/test/IR/attribute.mlir           | 31 +++++++++++++++++++++++++++
 mlir/test/lib/Dialect/Test/TestOps.td |  4 ++++
 3 files changed, 48 insertions(+)

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index f555329d374d9..f30693dc76720 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/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 }];
diff --git a/mlir/test/IR/attribute.mlir b/mlir/test/IR/attribute.mlir
index 06ac2cf9c9b80..f8c07e37b63dc 100644
--- a/mlir/test/IR/attribute.mlir
+++ b/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
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 0433c3075ca60..b64083b1d619b 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/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,