Address review comments

Author: skatrak

mlir/include/mlir/IR/CommonAttrConstraints.td

@@ -408,26 +408,17 @@ class ElementsAttrBase<Pred condition, string summary> :
   let storageType = [{ ::mlir::ElementsAttr }];
   let returnType = [{ ::mlir::ElementsAttr }];
   let convertFromStorage = "$_self";
-
-  // The underlying C++ value type of each element.
-  string elementReturnType = ?;
-
-  // The number of dimensions represented by the element collection.
-  int rank = 1;
 }
 
 def ElementsAttr : ElementsAttrBase<CPred<"::llvm::isa<::mlir::ElementsAttr>($_self)">,
-                                    "constant vector/tensor attribute"> {
-  let elementReturnType = [{ ::mlir::Attribute }];
-}
+                                    "constant vector/tensor attribute">;
 
 class IntElementsAttrBase<Pred condition, string summary> :
     ElementsAttrBase<And<[CPred<"::llvm::isa<::mlir::DenseIntElementsAttr>($_self)">,
                           condition]>,
                      summary> {
   let storageType = [{ ::mlir::DenseIntElementsAttr }];
   let returnType = [{ ::mlir::DenseIntElementsAttr }];
-  let elementReturnType = [{ ::llvm::APInt }];
 
   let convertFromStorage = "$_self";
 }
@@ -437,7 +428,6 @@ class DenseArrayAttrBase<string denseAttrName, string cppType, string summaryNam
                      summaryName # " dense array attribute"> {
   let storageType = "::mlir::" # denseAttrName;
   let returnType = "::llvm::ArrayRef<" # cppType # ">";
-  let elementReturnType = cppType;
   let constBuilderCall = "$_builder.get" # denseAttrName # "($0)";
 }
 def DenseBoolArrayAttr : DenseArrayAttrBase<"DenseBoolArrayAttr", "bool", "i1">;
@@ -496,8 +486,6 @@ class RankedSignlessIntElementsAttr<int width, list<int> dims> :
   let constBuilderCall = "::mlir::DenseIntElementsAttr::get("
     "::mlir::RankedTensorType::get({" # !interleave(dims, ", ") #
     "}, $_builder.getIntegerType(" # width # ")), ::llvm::ArrayRef($0))";
-  
-  let rank = !size(dims);
 }
 
 class RankedI32ElementsAttr<list<int> dims> :
@@ -513,7 +501,6 @@ class FloatElementsAttr<int width> : ElementsAttrBase<
 
   let storageType = [{ ::mlir::DenseElementsAttr }];
   let returnType = [{ ::mlir::DenseElementsAttr }];
-  let elementReturnType = [{ ::llvm::APFloat }];
 
   // Note that this is only constructing scalar elements attribute.
   let constBuilderCall = "::mlir::DenseElementsAttr::get("
@@ -539,8 +526,6 @@ class RankedFloatElementsAttr<int width, list<int> dims> : ElementsAttrBase<
 
   let storageType = [{ ::mlir::DenseFPElementsAttr }];
   let returnType = [{ ::mlir::DenseFPElementsAttr }];
-  let elementReturnType = [{ ::llvm::APFloat }];
-  let rank = !size(dims);
 
   let constBuilderCall = "::llvm::cast<::mlir::DenseFPElementsAttr>("
   "::mlir::DenseElementsAttr::get("
@@ -559,7 +544,6 @@ def StringElementsAttr : ElementsAttrBase<
 
   let storageType = [{ ::mlir::DenseElementsAttr }];
   let returnType = [{ ::mlir::DenseElementsAttr }];
-  let elementReturnType = [{ ::llvm::SmallString }];
 
   let convertFromStorage = "$_self";
 }

mlir/test/mlir-tblgen/openmp-clause-ops.td

@@ -7,7 +7,7 @@
 // RUN:   %S/../../../llvm/include/llvm/Frontend/OpenMP/OMP.td \
 // RUN:   -I %S/../../../llvm/include > %t/mlir/Dialect/OpenMP/OmpCommon.td
 
-// RUN: mlir-tblgen -gen-openmp-clause-ops -I %S/../../include -I %t %s | FileCheck %s
+// RUN: mlir-tblgen -gen-openmp-clause-ops -I %S/../../include -I %t %s 2>&1 | FileCheck %s
 
 include "mlir/Dialect/OpenMP/OpenMPOpBase.td"
 
@@ -34,7 +34,13 @@ def OpenMP_MyFirstClause : OpenMP_Clause<
     OptionalAttr<DenseI8ArrayAttr>:$opt_int_elems_attr,
 
     // Multi-level composition
-    ConfinedAttr<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$complex_opt_int_attr
+    ConfinedAttr<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$complex_opt_int_attr,
+
+    // ElementsAttrBase-related edge cases.
+    // CHECK: warning: could not infer array-like attribute element type for argument 'elements_attr', will use bare `storageType`
+    ElementsAttr:$elements_attr,
+    // CHECK: warning: could not infer array-like attribute element type for argument 'string_elements_attr', will use bare `storageType`
+    StringElementsAttr:$string_elements_attr
   );
 }
 // CHECK:      struct MyFirstClauseOps {
@@ -46,13 +52,15 @@ def OpenMP_MyFirstClause : OpenMP_Clause<
 // CHECK-NEXT:   ::llvm::SmallVector<::mlir::Attribute> strArrayAttr;
 // CHECK-NEXT:   ::llvm::SmallVector<::llvm::APInt> anyintElemsAttr;
 // CHECK-NEXT:   ::llvm::SmallVector<::llvm::APFloat> floatNdElemsAttr;
-// CHECK-NEXT:   int floatNdElemsAttrDims[3];
 
 // CHECK-NEXT:   ::mlir::BoolAttr optBoolAttr;
 // CHECK-NEXT:   ::llvm::SmallVector<::mlir::Attribute> optIntArrayAttr;
 // CHECK-NEXT:   ::llvm::SmallVector<int8_t> optIntElemsAttr;
 
 // CHECK-NEXT:   ::mlir::IntegerAttr complexOptIntAttr;
+
+// CHECK-NEXT:   ::mlir::ElementsAttr elementsAttr;
+// CHECK-NEXT:   ::mlir::DenseElementsAttr stringElementsAttr;
 // CHECK-NEXT: }
 
 def OpenMP_MySecondClause : OpenMP_Clause<

mlir/tools/mlir-tblgen/OmpOpGen.cpp

@@ -14,6 +14,7 @@
 
 #include "mlir/TableGen/CodeGenHelpers.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
@@ -184,40 +185,59 @@ static void verifyClause(Record *op, Record *clause) {
 /// All kinds of values are represented as `mlir::Value` fields, whereas
 /// attributes are represented based on their `storageType`.
 ///
+/// \param[in] name The name of the argument.
 /// \param[in] init The `DefInit` object representing the argument.
 /// \param[out] rank Number of levels of array nesting associated with the
 ///                  type.
 ///
 /// \return the name of the base type to represent elements of the argument
 ///         type.
-static StringRef translateArgumentType(Init *init, int &rank) {
+static StringRef translateArgumentType(ArrayRef<SMLoc> loc, StringInit *name,
+                                       Init *init, int &rank) {
   Record *def = cast<DefInit>(init)->getDef();
-  bool isAttr = false, isValue = false;
 
-  for (auto [sc, _] : def->getSuperClasses()) {
-    std::string scName = sc->getNameInitAsString();
-    if (scName == "OptionalAttr")
-      return translateArgumentType(def->getValue("baseAttr")->getValue(), rank);
-
-    if (scName == "TypedArrayAttrBase") {
-      ++rank;
-      return translateArgumentType(def->getValue("elementAttr")->getValue(),
-                                   rank);
-    }
-
-    if (scName == "ElementsAttrBase") {
-      rank += def->getValueAsInt("rank");
-      return def->getValueAsString("elementReturnType").trim();
-    }
-
-    if (scName == "Attr")
-      isAttr = true;
-    else if (scName == "TypeConstraint")
-      isValue = true;
-    else if (scName == "Variadic")
-      ++rank;
+  llvm::StringSet superClasses;
+  for (auto [sc, _] : def->getSuperClasses())
+    superClasses.insert(sc->getNameInitAsString());
+
+  // Handle wrapper-style superclasses.
+  if (superClasses.contains("OptionalAttr"))
+    return translateArgumentType(loc, name,
+                                 def->getValue("baseAttr")->getValue(), rank);
+
+  if (superClasses.contains("TypedArrayAttrBase"))
+    return translateArgumentType(
+        loc, name, def->getValue("elementAttr")->getValue(), ++rank);
+
+  // Handle ElementsAttrBase superclasses.
+  if (superClasses.contains("ElementsAttrBase")) {
+    // TODO: Support properly obtaining rank from ranked types.
+    ++rank;
+
+    if (superClasses.contains("IntElementsAttrBase"))
+      return "::llvm::APInt";
+    if (superClasses.contains("FloatElementsAttr") ||
+        superClasses.contains("RankedFloatElementsAttr"))
+      return "::llvm::APFloat";
+    if (superClasses.contains("DenseArrayAttrBase"))
+      return stripPrefixAndSuffix(def->getValueAsString("returnType"),
+                                  {"::llvm::ArrayRef<"}, {">"});
+
+    // Reset the rank in the case where the base type cannot be inferred, so
+    // that the bare storageType is used instead of a vector.
+    rank = 0;
+    PrintWarning(
+        loc,
+        "could not infer array-like attribute element type for argument '" +
+            name->getAsUnquotedString() + "', will use bare `storageType`");
   }
 
+  // Handle simple attribute and value types.
+  bool isAttr = superClasses.contains("Attr");
+  bool isValue = superClasses.contains("TypeConstraint");
+  if (superClasses.contains("Variadic"))
+    ++rank;
+
   if (isValue) {
     assert(!isAttr &&
            "argument can't be simultaneously a value and an attribute");
@@ -246,7 +266,8 @@ static void genClauseOpsStruct(Record *clause, raw_ostream &os) {
   for (auto [name, arg] :
        zip_equal(arguments->getArgNames(), arguments->getArgs())) {
     int rank = 0;
-    StringRef baseType = translateArgumentType(arg, rank);
+    StringRef baseType =
+        translateArgumentType(clause->getLoc(), name, arg, rank);
 
     if (rank > 0)
       os << "  ::llvm::SmallVector<" << baseType << ">";