[mlir][ods] Extend the EnumAttr tablegen class to support BitEnum attributes

[jfurtek]

Summary

This diff allows the EnumAttr class to be used for bit enum attributes (in
addition to previously supported integer enum attributes). While integer
and bit enum attributes share many common implementation aspects, parsing
bit enum values requires a separate implementation. This is accomplished
by creating empty parser and printer strings in the EnumAttrInfo record,
and having derived classes (specific to bit and integer enums) override with
an appropriate parser/printer string.

To support existing bit enums that may use a vertical bar separator, the
parser is modified to support the | token.

Tests were added for bit enums alongside integer enums.

Future diffs for fastmath attributes in the arithmetic dialect will use these
changes.

(resubmission of earlier abaondoned diff, updated to reflect subsequent changes
in the repository)

[Mogball]

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

326

Add a test for "|" as well?

[jfurtek]

Adding vertical bar test...

[jfurtek]

• Update lexer to support vertical bar token, and add test for vertical bar case

[jfurtek]

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

326

Great suggestion - I had assumed that it would "just work", but I had not made the necessary changes to the lexer. (That is fixed in the updated commit.)

[jfurtek]

• Rebase

mlir/include/mlir/IR/EnumAttr.td

let storageType = !if(genSpecializedAttr,
     baseAttrClass.storageType);
   let returnType = !if(genSpecializedAttr,
     cppNamespace # "::" # className,
     baseAttrClass.returnType);
   let constBuilderCall = !if(genSpecializedAttr,
     cppNamespace # "::" # specializedAttrClassName # "::get($_builder.getContext(), $0)",
     baseAttrClass.constBuilderCall);
   let valueType = baseAttrClass.valueType;
+
+  // C++ type wrapped by attribute
+  string cppType = cppNamespace # "::" # className;
+
+  // Parser and printer code used by the EnumParameter class, to be provided by
+  // derived classes
+  string parameterParser = ?;
+  string parameterPrinter = ?;
 }
 
 // An enum attribute backed by IntegerAttr.
 //
 // Op attributes of this kind are stored as IntegerAttr. Extra verification will
 // be generated on the integer though: only the values of the allowed cases are
 // permitted as the integer value.
 class IntEnumAttrBase<I intType, list<IntEnumAttrCaseBase> cases, string summary> :
     SignlessIntegerAttrBase<intType, summary> {
   let predicate = And<[
     SignlessIntegerAttrBase<intType, summary>.predicate,
     Or<!foreach(case, cases, case.predicate)>]>;
 }
 
 class IntEnumAttr<I intType, string name, string summary,
                   list<IntEnumAttrCaseBase> cases> :
   EnumAttrInfo<name, cases,
     IntEnumAttrBase<intType, cases,
       !if(!empty(summary), "allowed " # intType.summary # " cases: " #
           !interleave(!foreach(case, cases, case.value), ", "),
-          summary)>>;
+          summary)>> {
+  // Parse a keyword and pass it to `stringToSymbol`. Emit an error if a the
+  // symbol is not valid.
+  let parameterParser = [{[&]() -> ::mlir::FailureOr<}] # cppType # [{> {
+    auto loc = $_parser.getCurrentLocation();
+    ::llvm::StringRef enumKeyword;
+    if (::mlir::failed($_parser.parseKeyword(&enumKeyword)))
+      return ::mlir::failure();
+    auto maybeEnum = }] # cppNamespace # "::" #
+                          stringToSymbolFnName # [{(enumKeyword);
+    if (maybeEnum)
+      return *maybeEnum;
+    return {(::mlir::LogicalResult)$_parser.emitError(loc, "expected }] #
+    cppType # [{ to be one of: }] #
+    !interleave(!foreach(enum, enumerants, enum.str), ", ") # [{")};
+  }()}];
+  // Print the enum by calling `symbolToString`.
+  let parameterPrinter = "$_printer << " # symbolToStringFnName # "($_self)";
+}
 
 class I32EnumAttr<string name, string summary, list<I32EnumAttrCase> cases> :
     IntEnumAttr<I32, name, summary, cases> {
   let underlyingType = "uint32_t";
 }
 class I64EnumAttr<string name, string summary, list<I64EnumAttrCase> cases> :
     IntEnumAttr<I64, name, summary, cases> {
   let underlyingType = "uint64_t";
class BitEnumAttr<I intType, string name, string summary,
   int validBits = !foldl(0, cases, value, bitcase, !or(value, bitcase.value));
 
   // We need to return a string because we may concatenate symbols for multiple
   // bits together.
   let symbolToStringFnRetType = "std::string";
 
   // The delimiter used to separate bit enum cases in strings.
   string separator = "|";
+
+  // Parsing function that corresponds to the enum separator. Only
+  // "," and "|" are supported by this definition.
+  string parseSeparatorFn = !if(!eq(separator,"|"),"parseOptionalVerticalBar",
+                                                   "parseOptionalComma");
+  
+  // Parse a keyword and pass it to `stringToSymbol`. Emit an error if a the
+  // symbol is not valid.
+  let parameterParser = [{[&]() -> ::mlir::FailureOr<}] # cppType # [{> {
+    }] # cppType # [{ flags = {};
+    auto loc = $_parser.getCurrentLocation();
+    ::llvm::StringRef enumKeyword;
+    do {
+      if (::mlir::failed($_parser.parseKeyword(&enumKeyword)))
+        return ::mlir::failure();
+      auto maybeEnum = }] # cppNamespace # "::" #
+                            stringToSymbolFnName # [{(enumKeyword);
+      if (!maybeEnum) {
+          return {(::mlir::LogicalResult)$_parser.emitError(loc, "expected }] #
+                  cppType # [{ to be one of: }] #
+                 !interleave(!foreach(enum, enumerants, enum.str),
+                             ", ") # [{")};
+      }
+      flags = flags | *maybeEnum;
+    } while(::mlir::succeeded($_parser.}] # parseSeparatorFn # [{()));
+    return flags;
+  }()}];
+  // Print the enum by calling `symbolToString`.
+  let parameterPrinter = "$_printer << " # symbolToStringFnName # "($_self)";
+
 }
 
 class I32BitEnumAttr<string name, string summary,
                      list<BitEnumAttrCaseBase> cases>
     : BitEnumAttr<I32, name, summary, cases> {
   let underlyingType = "uint32_t";
 }
 
 class I64BitEnumAttr<string name, string summary,
                      list<BitEnumAttrCaseBase> cases>
     : BitEnumAttr<I64, name, summary, cases> {
   let underlyingType = "uint64_t";
 }
 
 // A C++ enum as an attribute parameter. The parameter implements a parser and
 // printer for the enum by dispatching calls to `stringToSymbol` and
 // `symbolToString`.
 class EnumParameter<EnumAttrInfo enumInfo>
     : AttrParameter<enumInfo.cppNamespace # "::" # enumInfo.className,
                     "an enum of type " # enumInfo.className> {
-  // Parse a keyword and pass it to `stringToSymbol`. Emit an error if a the
-  // symbol is not valid.
-  let parser = [{[&]() -> ::mlir::FailureOr<}] # cppType # [{> {
-    auto loc = $_parser.getCurrentLocation();
-    ::llvm::StringRef enumKeyword;
-    if (::mlir::failed($_parser.parseKeyword(&enumKeyword)))
-      return ::mlir::failure();
-    auto maybeEnum = }] # enumInfo.cppNamespace # "::" #
-                          enumInfo.stringToSymbolFnName # [{(enumKeyword);
-    if (maybeEnum)
-      return *maybeEnum;
-    return {(::mlir::LogicalResult)$_parser.emitError(loc, "expected }] #
-    cppType # [{ to be one of: }] #
-    !interleave(!foreach(enum, enumInfo.enumerants, enum.str), ", ") # [{")};
-  }()}];
-  // Print the enum by calling `symbolToString`.
-  let printer = "$_printer << " # enumInfo.symbolToStringFnName # "($_self)";
+  let parser = enumInfo.parameterParser;
+  let printer = enumInfo.parameterPrinter;
 }
 
 // An attribute backed by a C++ enum. The attribute contains a single
 // parameter `value` whose type is the C++ enum class.
 //
 // Example:
 //
 // ```

mlir/include/mlir/IR/OpImplementation.h

public:
   virtual ParseResult parseOptionalPlus() = 0;
 
   /// Parse a '*' token.
   virtual ParseResult parseStar() = 0;
 
   /// Parse a '*' token if present.
   virtual ParseResult parseOptionalStar() = 0;
 
+  /// Parse a '|' token.
+  virtual ParseResult parseVerticalBar() = 0;
+
+  /// Parse a '|' token if present.
+  virtual ParseResult parseOptionalVerticalBar() = 0;
+
   /// Parse a quoted string token.
   ParseResult parseString(std::string *string) {
     auto loc = getCurrentLocation();
     if (parseOptionalString(string))
       return emitError(loc, "expected string");
     return success();
   }
 

mlir/lib/Parser/AsmParserImpl.h

ParseResult parsePlus() override {
     return parser.parseToken(Token::plus, "expected '+'");
   }
 
   /// Parses a '+' token if present.
   ParseResult parseOptionalPlus() override {
     return success(parser.consumeIf(Token::plus));
   }
 
+  /// Parse a '|' token.
+  virtual ParseResult parseVerticalBar() override {
+    return parser.parseToken(Token::vertical_bar, "expected '|'");
+  }
+
+  /// Parse a '|' token if present.
+  virtual ParseResult parseOptionalVerticalBar() override {
+    return success(parser.consumeIf(Token::vertical_bar));
+  }
+
   /// Parses a quoted string token if present.
   ParseResult parseOptionalString(std::string *string) override {
     if (!parser.getToken().is(Token::string))
       return failure();
 
     if (string)
       *string = parser.getToken().getStringValue();
     parser.consumeToken();

mlir/lib/Parser/TokenKinds.def

 TOK_PUNCTUATION(less, "<")
 TOK_PUNCTUATION(minus, "-")
 TOK_PUNCTUATION(plus, "+")
 TOK_PUNCTUATION(question, "?")
 TOK_PUNCTUATION(r_brace, "}")
 TOK_PUNCTUATION(r_paren, ")")
 TOK_PUNCTUATION(r_square, "]")
 TOK_PUNCTUATION(star, "*")
+TOK_PUNCTUATION(vertical_bar, "|")
 
 // Keywords.  These turn "foo" into Token::kw_foo enums.
 
 // NOTE: Please key these alphabetized to make it easier to find something in
 // this list and to cater to OCD.
 TOK_KEYWORD(affine_map)
 TOK_KEYWORD(affine_set)
 TOK_KEYWORD(attributes)

mlir/test/IR/attribute.mlir

func @disallowed_case7_fail() {
   // expected-error @+1 {{allowed 64-bit signless integer cases: 5, 10}}
   %0 = "test.i64_enum_attr"() {attr = 5: i32} : () -> i32
   return
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
+// Test BitEnumAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @allowed_cases_pass
+func @allowed_cases_pass() {
+  // CHECK: test.op_with_bit_enum <read,write>
+  "test.op_with_bit_enum"() {value = #test.bit_enum<read, write>} : () -> ()
+  // CHECK: test.op_with_bit_enum <read,execute>
+  test.op_with_bit_enum <read,execute>
+  return
+}
+
+// -----
+
+func @disallowed_case_sticky_fail() {
+  // expected-error@+2 {{expected test::TestBitEnum to be one of: read, write, execute}}
+  // expected-error@+1 {{failed to parse TestBitEnumAttr}}
+  "test.op_with_bit_enum"() {value = #test.bit_enum<sticky>} : () -> ()
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // Test FloatElementsAttr
 //===----------------------------------------------------------------------===//
 
 func @correct_type_pass() {
   "test.float_elements_attr"() {
     // CHECK: scalar_f32_attr = dense<5.000000e+00> : tensor<2xf32>
     // CHECK: tensor_f64_attr = dense<6.000000e+00> : tensor<4x8xf64>
     scalar_f32_attr = dense<5.0> : tensor<2xf32>,

mlir/test/lib/Dialect/Test/TestDialect.cpp

 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 
 // Include this before the using namespace lines below to
 // test that we don't have namespace dependencies.
 #include "TestOpsDialect.cpp.inc"
 
 using namespace mlir;
 using namespace test;

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

 def : Pat<(OpWithEnum ConstantAttr<TestEnumAttr,
                                    "::test::TestEnum::First">:$value,
                       ConstantAttr<I32Attr, "0">:$tag),
           (OpWithEnum ConstantAttr<TestEnumAttr,
                                    "::test::TestEnum::Second">,
                       ConstantAttr<I32Attr, "1">)>;
 
 //===----------------------------------------------------------------------===//
+// Test Bit Enum Attributes
+//===----------------------------------------------------------------------===//
+
+// Define the C++ enum.
+def TestBitEnum
+    : I32BitEnumAttr<"TestBitEnum", "a test bit enum", [
+        I32BitEnumAttrCaseBit<"Read", 0, "read">,
+        I32BitEnumAttrCaseBit<"Write", 1, "write">,
+        I32BitEnumAttrCaseBit<"Execute", 2, "execute">,
+      ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "test";
+  let separator = ",";
+}
+
+// Define the enum attribute.
+def TestBitEnumAttr : EnumAttr<Test_Dialect, TestBitEnum, "bit_enum"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
+// Define an op that contains the enum attribute.
+def OpWithBitEnum : TEST_Op<"op_with_bit_enum"> {
+  let arguments = (ins TestBitEnumAttr:$value, OptionalAttr<AnyAttr>:$tag);
+  let assemblyFormat = "$value (`tag` $tag^)? attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
 // Test Attribute Constraints
 //===----------------------------------------------------------------------===//
 
 def SymbolRefOp : TEST_Op<"symbol_ref_attr"> {
   let arguments = (ins
     Confined<FlatSymbolRefAttr, [ReferToOp<"FuncOp">]>:$symbol
   );
 }