[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

339

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.
+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;
+
+  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_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@+1 {{invalid kind of attribute specified}}
+  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 #test<"enum third">
+  test.op_with_enum #test<"enum third">
+  return
+}
+
+// CHECK-LABEL: @test_match_op_with_enum
+func @test_match_op_with_enum() -> () {
+  // CHECK: test.op_with_enum #test<"enum third"> tag 0 : i32
+  test.op_with_enum #test<"enum third"> tag 0 : i32
+  // CHECK: test.op_with_enum #test<"enum second"> tag 1 : i32
+  test.op_with_enum #test<"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);
 } // end 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 "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"],
 )