[mlir][ods] Added EnumAttr, an AttrDef implementation of enum attributes

[Mogball]

Summary

EnumAttr is a pure TableGen implementation of enum attributes using AttrDef. This is meant as a drop-in replacement for StrEnumAttr, which is soon to be deprecated. StrEnumAttr is often used over IntEnumAttr because its more readable in MLIR assembly formats. However, storing and manipulating strings is not efficient. Defining StrEnumAttr can also be awkward and relies on a lot of special logic in EnumsGen, and has some hidden sharp edges.

Also, EnumAttr stores the enum directly, removing the need to convert to/from integers when calling attribute getters on ops.

[mehdi_amini]

Nice!

mlir/include/mlir/IR/EnumAttr.td

39

Not using the surrounding < and > forces the opaque/generic syntax for the attribute, that's a bit unfortunate though.

Is there a way to get a better printing? (in particular in the declarative assembly used by operations)

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

311

Can you use a declarative assembly here?

[Mogball]

mlir/include/mlir/IR/EnumAttr.td

39

I can add the braces for now, but the reason why I wanted to wait on your strip dialect prefix patch is so that this won't be a concern (when attributes are used in assembly formats). I.e. #test<"enum first"> will be first when parsed as part of an assembly format.

[Mogball]

Update test op with assemblyFormat

[Mogball]

mlir/include/mlir/IR/EnumAttr.td

39

Users can also override the assemblyFormat themselves.

[Mogball]

Updated assembly format with stripped dialect prefix

mlir/test/IR/enum-attr-roundtrip.mlir

18

@mehdi_amini check it out!

[rriddle]

LGTM, but will let Mehdi sign off (as he had some comments).

mlir/include/mlir/IR/EnumAttr.td

37–38

[mehdi_amini]

mlir/include/mlir/IR/EnumAttr.td

39

I think we should just document this better.
To begin with, indicate that the default printer/parser is optimized for uses in the op assembly format.
Provide example of how to use this and example of how it looks like in IR.

mlir/test/IR/enum-attr-roundtrip.mlir

18

Nice! ;)

[Mogball]

Add more documentation

mlir/include/mlir/IR/EnumAttr.td

+//===-- EnumAttr.td - Enum attributes ----------------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ENUM_ATTR
+#define ENUM_ATTR
+
+include "mlir/IR/OpBase.td"
+
+// 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 {$_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)";
+}
+
+// An attribute backed by a C++ enum. The attribute contains a single
+// parameter `value` whose type is the C++ enum class.
+//
+// Example:
+//
+// ```
+// def MyEnum : I32EnumAttr<"MyEnum", "a simple enum", [
+//                            I32EnumAttrCase<"First", 0, "first">,
+//                            I32EnumAttrCase<"Second", 1, "second>]> {
+//   let genSpecializedAttr = 0;
+// }
+//
+// def MyEnumAttr : EnumAttr<MyDialect, MyEnum, "enum">;
+// ```
+//
+// By default, the assembly format of the attribute works best with operation
+// assembly formats. For example:
+//
+// ```
+// def MyOp : Op<MyDialect, "my_op"> {
+//   let arguments = (ins MyEnumAttr:$enum);
+//   let assemblyFormat = "$enum attr-dict";
+// }
+// ```
+//
+// The op will appear in the IR as `my_dialect.my_op first`. However, the
+// generic format of the attribute will be `#my_dialect<"enum first">`. Override
+// the attribute's assembly format as required.
+class EnumAttr<Dialect dialect, EnumAttrInfo enumInfo, string name = "",
+               list <Trait> traits = []>
+    : AttrDef<dialect, enumInfo.className, traits> {
+  let summary = enumInfo.summary;
+
+  // Inherit the C++ namespace from the enum.
+  let cppNamespace = enumInfo.cppNamespace;
+
+  // Define a constant builder for the attribute to convert from C++ enums.
+  let constBuilderCall = cppNamespace # "::" # cppClassName #
+                         "::get($_builder.getContext(), $0)";
+
+  // Op attribute getters should return the underlying C++ enum type.
+  let returnType = enumInfo.cppNamespace # "::" # enumInfo.className;
+
+  // Convert from attribute to the underlying C++ type in op getters.
+  let convertFromStorage = "$_self.getValue()";
+
+  // The enum attribute has one parameter: the C++ enum value.
+  let parameters = (ins EnumParameter<enumInfo>:$value);
+
+  // If a mnemonic was provided, use it to generate a custom assembly format.
+  let mnemonic = name;
+
+  // The default assembly format for enum attributes. Selected to best work with
+  // operation assembly formats.
+  let assemblyFormat = "$value";
+}
+
+#endif // ENUM_ATTR

mlir/test/IR/enum-attr-invalid.mlir

+// RUN: mlir-opt -verify-diagnostics -split-input-file %s
+
+func @test_invalid_enum_case() -> () {
+  // expected-error@+2 {{expected test::TestEnum to be one of: first, second, third}}
+  // expected-error@+1 {{failed to parse TestEnumAttr}}
+  test.op_with_enum #test<"enum fourth">
+}
+
+// -----
+
+func @test_invalid_enum_case() -> () {
+  // expected-error@+1 {{expected test::TestEnum to be one of: first, second, third}}
+  test.op_with_enum fourth
+  // expected-error@+1 {{failed to parse TestEnumAttr}}
+}
+
+// -----
+
+func @test_invalid_attr() -> () {
+  // expected-error@+1 {{op attribute 'value' failed to satisfy constraint: a test enum}}
+  "test.op_with_enum"() {value = 1 : index} : () -> ()
+}
+
+// -----
+
+func @test_parse_invalid_attr() -> () {
+  // expected-error@+2 {{expected valid keyword}}
+  // expected-error@+1 {{failed to parse TestEnumAttr parameter 'value'}}
+  test.op_with_enum 1 : index
+}

mlir/test/IR/enum-attr-roundtrip.mlir

+// RUN: mlir-opt %s | mlir-opt -test-patterns | FileCheck %s
+
+// CHECK-LABEL: @test_enum_attr_roundtrip
+func @test_enum_attr_roundtrip() -> () {
+  // CHECK: value = #test<"enum first">
+  "test.op"() {value = #test<"enum first">} : () -> ()
+  // CHECK: value = #test<"enum second">
+  "test.op"() {value = #test<"enum second">} : () -> ()
+  // CHECK: value = #test<"enum third">
+  "test.op"() {value = #test<"enum third">} : () -> ()
+  return
+}
+
+// CHECK-LABEL: @test_op_with_enum
+func @test_op_with_enum() -> () {
+  // CHECK: test.op_with_enum third
+  test.op_with_enum third
+  return
+}
+
+// CHECK-LABEL: @test_match_op_with_enum
+func @test_match_op_with_enum() -> () {
+  // CHECK: test.op_with_enum third tag 0 : i32
+  test.op_with_enum third tag 0 : i32
+  // CHECK: test.op_with_enum second tag 1 : i32
+  test.op_with_enum first tag 0 : i32
+  return
+}

mlir/test/lib/Dialect/Test/TestAttributes.h

 
 #include "TestTraits.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 
 #include "TestAttrInterfaces.h.inc"
+#include "TestOpEnums.h.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.h.inc"
 
 #endif // MLIR_TESTATTRIBUTES_H

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

 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 class DLTIDialect;
 class RewritePatternSet;
 } // namespace mlir
 
-#include "TestOpEnums.h.inc"
 #include "TestOpInterfaces.h.inc"
 #include "TestOpStructs.h.inc"
 #include "TestOpsDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.h.inc"
 
 namespace test {
 void registerTestDialect(::mlir::DialectRegistry &registry);
 void populateTestReductionPatterns(::mlir::RewritePatternSet &patterns);
 } // namespace test
 
 #endif // MLIR_TESTDIALECT_H

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

 //===-- TestOps.td - Test dialect operation definitions ----*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TEST_OPS
 #define TEST_OPS
 
 include "TestDialect.td"
 include "mlir/Dialect/DLTI/DLTIBase.td"
+include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"

 
 def StringElementsAttrOp : TEST_Op<"string_elements_attr"> {
   let arguments = (ins
       StringElementsAttr:$scalar_string_attr
   );
 }
 
 //===----------------------------------------------------------------------===//
+// Test Enum Attributes
+//===----------------------------------------------------------------------===//
+
+// Define the C++ enum.
+def TestEnum
+    : I32EnumAttr<"TestEnum", "a test enum", [
+        I32EnumAttrCase<"First", 0, "first">,
+        I32EnumAttrCase<"Second", 1, "second">,
+        I32EnumAttrCase<"Third", 2, "third">,
+      ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "test";
+}
+
+// Define the enum attribute.
+def TestEnumAttr : EnumAttr<Test_Dialect, TestEnum, "enum">;
+
+// Define an op that contains the enum attribute.
+def OpWithEnum : TEST_Op<"op_with_enum"> {
+  let arguments = (ins TestEnumAttr:$value, OptionalAttr<AnyAttr>:$tag);
+  let assemblyFormat = "$value (`tag` $tag^)? attr-dict";
+}
+
+// Define a pattern that matches and creates an enum attribute.
+def : Pat<(OpWithEnum ConstantAttr<TestEnumAttr,
+                                   "::test::TestEnum::First">:$value,
+                      ConstantAttr<I32Attr, "0">:$tag),
+          (OpWithEnum ConstantAttr<TestEnumAttr,
+                                   "::test::TestEnum::Second">,
+                      ConstantAttr<I32Attr, "1">)>;
+
+//===----------------------------------------------------------------------===//
 // Test Attribute Constraints
 //===----------------------------------------------------------------------===//
 
 def SymbolRefOp : TEST_Op<"symbol_ref_attr"> {
   let arguments = (ins
     Confined<FlatSymbolRefAttr, [ReferToOp<"FuncOp">]>:$symbol
   );
 }

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_binary(
     ],
 )
 
 ##---------------------------------------------------------------------------##
 
 td_library(
     name = "OpBaseTdFiles",
     srcs = [
+        "include/mlir/IR/EnumAttr.td",
         "include/mlir/IR/OpAsmInterface.td",
         "include/mlir/IR/OpBase.td",
         "include/mlir/IR/RegionKindInterface.td",
         "include/mlir/IR/SymbolInterfaces.td",
         "include/mlir/IR/TensorEncoding.td",
     ],
     includes = ["include"],
 )