[mlir] Update LLVMIR Fastmath flags use of MLIR BitEnum functionality

[jfurtek]

Summary

This diff updates the LLVMIR dialect Fastmath flags attribute to use recently
added features of BitEnum attributes. Specifically, this diff uses the bit
enum "group" case to represent the fast value as an alias for a combination
of other values (ninf, nnan, ...), instead of using a separate integer
value. (This is in line with LLVM's fastmath flags representation.) This diff
also leverages the printBitEnumPrimaryGroups tblgen field for concise
enum printing.

The BitEnum features were developed for an upcoming diff that adds fastmath
support to the arithmetic dialect. This diff simply applies some of the relevant
new features to the LLVM dialect attribute.

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

39

Nit: can this be ", " (with the trailing space) so we still have it in the generated IR?

[jfurtek]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

39

I think that there is a small "can of worms" here...

1. Adding desired spaces to the separator field works well for printing, but less so for parsing. The generated enum printing function (stringify()) uses llvm::join(), and that would work with ", ". However, unless modified, the generated parse function (symbolize()) would fail if a string without the space is encountered (e.g. enum1,enum2).
2. Another alternative would be to automatically add a space after the separator in the tblgen-erated print() function. However, this is problematic because the other "common" enum separator is a vertical bar. So while enum1, enum2 looks OK, enum1| enum2 is less appealing.

It might work if we allow desired spacing in the tblgen string (separator), use that for printing, and then "strip" the separator value of whitespace for the generated parse() function. (I haven't tried, but it might work.)

I could also just check the separator string in mlir-tblgen and specialize whitespace printing for , and |, with a default for other separators.

Suggestions on the best approach are welcome...

[jfurtek]

• Update handling of separators with spaces for nicer printing.

[jfurtek]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

39

The update addresses this - spaces in the ODS separator field are used for printing, but are stripped away for parsing (so that parsing accepts IR without spaces).

[jfurtek]

• Fix clang formatting

[jfurtek]

• Update handling of separators with spaces for nicer printing.
• Fix clang formatting
• rebase

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 
 def FMFnnan     : I32BitEnumAttrCaseBit<"nnan", 0>;
 def FMFninf     : I32BitEnumAttrCaseBit<"ninf", 1>;
 def FMFnsz      : I32BitEnumAttrCaseBit<"nsz", 2>;
 def FMFarcp     : I32BitEnumAttrCaseBit<"arcp", 3>;
 def FMFcontract : I32BitEnumAttrCaseBit<"contract", 4>;
 def FMFafn      : I32BitEnumAttrCaseBit<"afn", 5>;
 def FMFreassoc  : I32BitEnumAttrCaseBit<"reassoc", 6>;
-def FMFfast     : I32BitEnumAttrCaseBit<"fast", 7>;
+def FMFfast     : I32BitEnumAttrCaseGroup<"fast",
+  [ FMFnnan, FMFninf, FMFnsz, FMFarcp, FMFcontract, FMFafn, FMFreassoc]>;
 
 def FastmathFlags : I32BitEnumAttr<
     "FastmathFlags",
     "LLVM fastmath flags",
     [FMFnnan, FMFninf, FMFnsz, FMFarcp, FMFcontract, FMFafn, FMFreassoc, FMFfast
     ]> {
+  let separator = ", ";
   let cppNamespace = "::mlir::LLVM";
+  let printBitEnumPrimaryGroups = 1;
 }
 
 def LLVM_FMFAttr : DialectAttr<
     LLVM_Dialect,
     CPred<"$_self.isa<::mlir::LLVM::FMFAttr>()">,
     "LLVM fastmath flags"> {
   let storageType = "::mlir::LLVM::FMFAttr";
   let returnType = "::mlir::LLVM::FastmathFlags";

mlir/include/mlir/IR/EnumAttr.td

class BitEnumAttr<I intType, string name, string summary,
     : EnumAttrInfo<name, cases, BitEnumAttrBase<intType, cases, summary>> {
   // Determine "valid" bits from enum cases for error checking
   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.
+  // The delimiter used to separate bit enum cases in strings. Only "|" and
+  // "," (along with optional spaces) are supported due to the use of the
+  // parseSeparatorFn in parameterParser below.
+  // Spaces in the separator string are used for printing, but will be optional
+  // for parsing.
   string separator = "|";
+  assert !or(!ge(!find(separator, "|"), 0), !ge(!find(separator, ","), 0)),
+      "separator must contain '|' or ',' for parameter parsing";
 
   // Parsing function that corresponds to the enum separator. Only
   // "," and "|" are supported by this definition.
-  string parseSeparatorFn = !if(!eq(separator,"|"),"parseOptionalVerticalBar",
-                                                   "parseOptionalComma");
+  string parseSeparatorFn = !if(!ge(!find(separator, "|"), 0),
+                                "parseOptionalVerticalBar",
 
   // 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 {

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

return builder.create<LLVM::GEPOp>(
       ValueRange{cst0, cst0});
 }
 
 bool mlir::LLVM::satisfiesLLVMModule(Operation *op) {
   return op->hasTrait<OpTrait::SymbolTable>() &&
          op->hasTrait<OpTrait::IsIsolatedFromAbove>();
 }
 
-static constexpr const FastmathFlags fastmathFlagsList[] = {
-    // clang-format off
-    FastmathFlags::nnan,
-    FastmathFlags::ninf,
-    FastmathFlags::nsz,
-    FastmathFlags::arcp,
-    FastmathFlags::contract,
-    FastmathFlags::afn,
-    FastmathFlags::reassoc,
-    FastmathFlags::fast,
-    // clang-format on
-};
-
 void FMFAttr::print(AsmPrinter &printer) const {
   printer << "<";
-  auto flags = llvm::make_filter_range(fastmathFlagsList, [&](auto flag) {
-    return bitEnumContains(this->getFlags(), flag);
-  });
-  llvm::interleaveComma(flags, printer,
-                        [&](auto flag) { printer << stringifyEnum(flag); });
+  printer << stringifyFastmathFlags(this->getFlags());
   printer << ">";
 }
 
 Attribute FMFAttr::parse(AsmParser &parser, Type type) {
   if (failed(parser.parseLess()))
     return {};
 
   FastmathFlags flags = {};

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

const std::pair<FastmathFlags, FuncT> handlers[] = {
       // clang-format off
       {FastmathFlags::nnan,     &llvmFMF::setNoNaNs},
       {FastmathFlags::ninf,     &llvmFMF::setNoInfs},
       {FastmathFlags::nsz,      &llvmFMF::setNoSignedZeros},
       {FastmathFlags::arcp,     &llvmFMF::setAllowReciprocal},
       {FastmathFlags::contract, &llvmFMF::setAllowContract},
       {FastmathFlags::afn,      &llvmFMF::setApproxFunc},
       {FastmathFlags::reassoc,  &llvmFMF::setAllowReassoc},
-      {FastmathFlags::fast,     &llvmFMF::setFast},
       // clang-format on
   };
   llvm::FastMathFlags ret;
   auto fmf = op.getFastmathFlags();
   for (auto it : handlers)
     if (bitEnumContains(fmf, it.first))
       (ret.*(it.second))(true);
   return ret;

mlir/test/Dialect/LLVMIR/roundtrip.mlir

// CHECK: {{.*}} = llvm.fneg %arg0 {fastmathFlags = #llvm.fastmath<fast>} : f32
   %6 = llvm.fneg %arg0 {fastmathFlags = #llvm.fastmath<fast>} : f32
 
 // CHECK: {{.*}} = llvm.call @foo(%arg2) {fastmathFlags = #llvm.fastmath<fast>} : (i32) -> !llvm.struct<(i32, f64, i32)>
   %7 = llvm.call @foo(%arg2) {fastmathFlags = #llvm.fastmath<fast>} : (i32) -> !llvm.struct<(i32, f64, i32)>
 
 // CHECK: {{.*}} = llvm.fadd %arg0, %arg1 : f32
   %8 = llvm.fadd %arg0, %arg1 {fastmathFlags = #llvm.fastmath<>} : f32
 // CHECK: {{.*}} = llvm.fadd %arg0, %arg1 {fastmathFlags = #llvm.fastmath<nnan, ninf>} : f32
-  %9 = llvm.fadd %arg0, %arg1 {fastmathFlags = #llvm.fastmath<nnan, ninf>} : f32
+  %9 = llvm.fadd %arg0, %arg1 {fastmathFlags = #llvm.fastmath<nnan,ninf>} : f32
 
 // CHECK: {{.*}} = llvm.fneg %arg0 : f32
   %10 = llvm.fneg %arg0 {fastmathFlags = #llvm.fastmath<>} : f32
   return
 }
 
 module {
   // CHECK-LABEL: @loopOptions

mlir/test/IR/attribute.mlir

 // -----
 
 //===----------------------------------------------------------------------===//
 // Test BitEnumAttr
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: func @allowed_cases_pass
 func.func @allowed_cases_pass() {
-  // CHECK: test.op_with_bit_enum <read,write>
+  // 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>
+  // CHECK: test.op_with_bit_enum <read, execute>
   test.op_with_bit_enum <read,execute>
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @allowed_cases_pass
 func.func @allowed_cases_pass() {
-  // CHECK: test.op_with_bit_enum_vbar <user|group>
+  // CHECK: test.op_with_bit_enum_vbar <user | group>
   "test.op_with_bit_enum_vbar"() {
     value = #test.bit_enum_vbar<user|group>
   } : () -> ()
-  // CHECK: test.op_with_bit_enum_vbar <user|group|other>
+  // CHECK: test.op_with_bit_enum_vbar <user | group | other>
   test.op_with_bit_enum_vbar <user | group | other>
   return
 }
 
 // -----
 
 func.func @disallowed_case_sticky_fail() {
   // expected-error@+2 {{expected test::TestBitEnum to be one of: read, write, execute}}

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

 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 = ",";
+  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";
 }
 
 // Define an enum with a different separator
 def TestBitEnumVerticalBar
     : I32BitEnumAttr<"TestBitEnumVerticalBar", "another test bit enum", [
         I32BitEnumAttrCaseBit<"User", 0, "user">,
         I32BitEnumAttrCaseBit<"Group", 1, "group">,
         I32BitEnumAttrCaseBit<"Other", 2, "other">,
       ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = "test";
-  let separator = "|";
+  let separator = " | ";
 }
 
 def TestBitEnumVerticalBarAttr
     : EnumAttr<Test_Dialect, TestBitEnumVerticalBar, "bit_enum_vbar"> {
   let assemblyFormat = "`<` $value `>`";
 }
 
 // Define an op that contains the enum attribute.

mlir/tools/mlir-tblgen/EnumsGen.cpp

 }
 
 static void emitStrToSymFnForBitEnum(const Record &enumDef, raw_ostream &os) {
   EnumAttr enumAttr(enumDef);
   StringRef enumName = enumAttr.getEnumClassName();
   std::string underlyingType = std::string(enumAttr.getUnderlyingType());
   StringRef strToSymFnName = enumAttr.getStringToSymbolFnName();
   StringRef separator = enumDef.getValueAsString("separator");
+  StringRef separatorTrimmed = separator.trim();
   auto enumerants = enumAttr.getAllCases();
   auto allBitsUnsetCase = getAllBitsUnsetCase(enumerants);
 
   os << formatv("::llvm::Optional<{0}> {1}(::llvm::StringRef str) {{\n",
                 enumName, strToSymFnName);
 
   if (allBitsUnsetCase) {
     os << "  // Special case for all bits unset.\n";
     StringRef caseSymbol = allBitsUnsetCase->getSymbol();
     os << formatv("  if (str == \"{1}\") return {0}::{2};\n\n", enumName,
                   caseSymbol, makeIdentifier(caseSymbol));
   }
 
   // Split the string to get symbols for all the bits.
   os << "  ::llvm::SmallVector<::llvm::StringRef, 2> symbols;\n";
-  os << formatv("  str.split(symbols, \"{0}\");\n\n", separator);
+  // Remove whitespace from the separator string when parsing.
+  os << formatv("  str.split(symbols, \"{0}\");\n\n", separatorTrimmed);
 
   os << formatv("  {0} val = 0;\n", underlyingType);
   os << "  for (auto symbol : symbols) {\n";
 
   // Convert each symbol to the bit ordinal and set the corresponding bit.
-  os << formatv(
-      "    auto bit = llvm::StringSwitch<::llvm::Optional<{0}>>(symbol)\n",
+  os << formatv("    auto bit = "
+                "llvm::StringSwitch<::llvm::Optional<{0}>>(symbol.trim())\n",
                 underlyingType);
   for (const auto &enumerant : enumerants) {
     // Skip the special enumerant for None.
     if (auto val = enumerant.getValue())
       os.indent(6) << formatv(".Case(\"{0}\", {1})\n", enumerant.getStr(), val);
   }
   os.indent(6) << ".Default(::llvm::None);\n";
 
   os << "    if (bit) { val |= *bit; } else { return ::llvm::None; }\n";

mlir/unittests/TableGen/EnumsGenTest.cpp

 }
 
 TEST(EnumsGenTest, GeneratedSymbolToStringFnForBitEnum) {
   EXPECT_EQ(stringifyBitEnumWithNone(BitEnumWithNone::None), "None");
   EXPECT_EQ(stringifyBitEnumWithNone(BitEnumWithNone::Bit0), "Bit0");
   EXPECT_EQ(stringifyBitEnumWithNone(BitEnumWithNone::Bit3), "Bit3");
   EXPECT_EQ(
       stringifyBitEnumWithNone(BitEnumWithNone::Bit0 | BitEnumWithNone::Bit3),
-      "Bit0|Bit3");
+      "Bit0 | Bit3");
 
   EXPECT_EQ(stringifyBitEnum64_Test(BitEnum64_Test::Bit1), "Bit1");
   EXPECT_EQ(
       stringifyBitEnum64_Test(BitEnum64_Test::Bit1 | BitEnum64_Test::Bit57),
       "Bit1|Bit57");
 }
 
 TEST(EnumsGenTest, GeneratedStringToSymbolForBitEnum) {
   EXPECT_EQ(symbolizeBitEnumWithNone("None"), BitEnumWithNone::None);
   EXPECT_EQ(symbolizeBitEnumWithNone("Bit0"), BitEnumWithNone::Bit0);
   EXPECT_EQ(symbolizeBitEnumWithNone("Bit3"), BitEnumWithNone::Bit3);
   EXPECT_EQ(symbolizeBitEnumWithNone("Bit3|Bit0"),
             BitEnumWithNone::Bit3 | BitEnumWithNone::Bit0);
 
   EXPECT_EQ(symbolizeBitEnumWithNone("Bit2"), llvm::None);
-  EXPECT_EQ(symbolizeBitEnumWithNone("Bit3|Bit4"), llvm::None);
+  EXPECT_EQ(symbolizeBitEnumWithNone("Bit3 | Bit4"), llvm::None);
 
   EXPECT_EQ(symbolizeBitEnumWithoutNone("None"), llvm::None);
 }
 
 TEST(EnumsGenTest, GeneratedSymbolToStringFnForGroupedBitEnum) {
   EXPECT_EQ(stringifyBitEnumWithGroup(BitEnumWithGroup::Bit0), "Bit0");
   EXPECT_EQ(stringifyBitEnumWithGroup(BitEnumWithGroup::Bit3), "Bit3");
   EXPECT_EQ(stringifyBitEnumWithGroup(BitEnumWithGroup::Bits0To3),

 TEST(EnumsGenTest, GeneratedSymbolToStringFnForPrimaryGroupBitEnum) {
   EXPECT_EQ(stringifyBitEnumPrimaryGroup(
                 BitEnumPrimaryGroup::Bit0 | BitEnumPrimaryGroup::Bit1 |
                 BitEnumPrimaryGroup::Bit2 | BitEnumPrimaryGroup::Bit3),
             "Bits0To3");
   EXPECT_EQ(stringifyBitEnumPrimaryGroup(BitEnumPrimaryGroup::Bit0 |
                                          BitEnumPrimaryGroup::Bit2 |
                                          BitEnumPrimaryGroup::Bit3),
-            "Bit0,Bit2,Bit3");
+            "Bit0, Bit2, Bit3");
   EXPECT_EQ(stringifyBitEnumPrimaryGroup(BitEnumPrimaryGroup::Bit0 |
                                          BitEnumPrimaryGroup::Bit4 |
                                          BitEnumPrimaryGroup::Bit5),
-            "Bits4And5,Bit0");
+            "Bits4And5, Bit0");
   EXPECT_EQ(stringifyBitEnumPrimaryGroup(
                 BitEnumPrimaryGroup::Bit0 | BitEnumPrimaryGroup::Bit1 |
                 BitEnumPrimaryGroup::Bit2 | BitEnumPrimaryGroup::Bit3 |
                 BitEnumPrimaryGroup::Bit4 | BitEnumPrimaryGroup::Bit5),
             "Bits0To5");
 }
 
 TEST(EnumsGenTest, GeneratedOperator) {

mlir/unittests/TableGen/enums.td

 def Bit0 : I32BitEnumAttrCaseBit<"Bit0", 0>;
 def Bit1 : I32BitEnumAttrCaseBit<"Bit1", 1>;
 def Bit2 : I32BitEnumAttrCaseBit<"Bit2", 2>;
 def Bit3 : I32BitEnumAttrCaseBit<"Bit3", 3>;
 def Bit4 : I32BitEnumAttrCaseBit<"Bit4", 4>;
 def Bit5 : I32BitEnumAttrCaseBit<"Bit5", 5>;
 
 def BitEnumWithNone : I32BitEnumAttr<"BitEnumWithNone", "A test enum",
-                                     [NoBits, Bit0, Bit3]>;
+                                     [NoBits, Bit0, Bit3]> {
+  let separator = " | ";
+}
 
 def BitEnumWithoutNone : I32BitEnumAttr<"BitEnumWithoutNone", "A test enum",
                                         [Bit0, Bit3]>;
 
 def Bits0To3 : I32BitEnumAttrCaseGroup<"Bits0To3",
                                        [Bit0, Bit1, Bit2, Bit3]>;
 def Bits4And5 : I32BitEnumAttrCaseGroup<"Bits4And5",
                                        [Bit4, Bit5]>;
 def Bits0To5 : I32BitEnumAttrCaseGroup<"Bits0To5",
                                        [Bits0To3, Bits4And5]>;
 
 def BitEnumWithGroup : I32BitEnumAttr<"BitEnumWithGroup", "A test enum",
-                                      [Bit0, Bit1, Bit2, Bit3, Bit4, Bits0To3]>;
+                                      [Bit0, Bit1, Bit2, Bit3, Bit4, Bits0To3]> {
+  let separator = "|";
+}
 
 def BitEnumPrimaryGroup : I32BitEnumAttr<"BitEnumPrimaryGroup", "test enum",
                                         [Bit0, Bit1, Bit2, Bit3, Bit4, Bit5,
                                          Bits0To3, Bits4And5, Bits0To5]> {
-  let separator = ",";
+  let separator = ", ";
   let printBitEnumPrimaryGroups = 1;
 }
 
 def BitEnum64_None : I64BitEnumAttrCaseNone<"None">;
 def BitEnum64_57   : I64BitEnumAttrCaseBit<"Bit57", 57>;
 def BitEnum64_1    : I64BitEnumAttrCaseBit<"Bit1", 1>;
 def BitEnum64_Test : I64BitEnumAttr<"BitEnum64_Test", "A 64-bit test enum",
                                     [BitEnum64_None, BitEnum64_1, BitEnum64_57]>;
 
 def PrettyIntEnumCase1: I32EnumAttrCase<"Case1", 1, "case_one">;
 def PrettyIntEnumCase2: I32EnumAttrCase<"Case2", 2, "case_two">;
 
 def PrettyIntEnum: I32EnumAttr<"PrettyIntEnum", "A test enum",
                                [PrettyIntEnumCase1, PrettyIntEnumCase2]>;