[mlir][OpGen] Cache Identifiers for known attribute names in AbstractOperation.

[rriddle]

Summary

Operations currently rely on the string name of attributes during attribute lookup/removal/replacement, in build methods, and more. This unfortunately means that some of the most used APIs in MLIR require string comparisons, additional hashing(+mutex locking) to construct Identifiers, and more. This revision remedies this by caching identifiers for all of the attributes of the operation in its corresponding AbstractOperation. Just updating the autogenerated usages brings up to a 15% reduction in compile time, greatly reducing the cost of interacting with the attributes of an operation. This number can grow even higher as we use these methods in handwritten C++ code.

Methods for accessing these cached identifiers are exposed via <attr-name>AttrName methods on the derived operation class. Moving forward, users should generally use these methods over raw strings when an attribute name is necessary.

[mehdi_amini]

Awesome :)

[mehdi_amini]

mlir/include/mlir/IR/OperationSupport.h

191

I don't know if this is the right place, but I feel we're missing some doc about how this all fit together

[jpienaar]

OOC what is the memory usage impact?

mlir/lib/IR/MLIRContext.cpp

727

So this initialized on construction rather than first use?

[rriddle]

In D104167#2814917, @jpienaar wrote:

OOC what is the memory usage impact?

The TF model I tested on had an increase in memory usage of ~50kb, an IREE test I ran had an increase of ~28kb.

[rriddle]

rebase

[rriddle]

mlir/include/mlir/IR/OperationSupport.h

191

Beefed up the description, PTAL.

mlir/lib/IR/MLIRContext.cpp

727

Yes. Initializing on first-use adds additional checks to the hot part of the code, which is something that we want. The way the revision is currently written, accessing a cached attribute name is effectively just a load from the AbstractOperation.

[rriddle]

update

[clementval]

This change seems to break flang bulldbots.

[Meinersbur]

The breaking buildbots:

flang-x86_64-windows: https://lab.llvm.org/staging/#/builders/179/builds/129

tools\flang\include\flang/Optimizer/Dialect/FIROps.h.inc(2324): error C2373: 'llvm::StringRef': redefinition; different type modifiers

ppc64le-flang-rhel-clang: https://lab.llvm.org/buildbot/#/builders/21/builds/17760

tools/flang/include/flang/Optimizer/Dialect/FIROps.h.inc:1615:32: error: functions that differ only in their return type cannot be overloaded
    static constexpr StringRef calleeAttrName() { return "callee"; }
                     ~~~~~~~~~ ^
tools/flang/include/flang/Optimizer/Dialect/FIROps.h.inc:1587:22: note: previous definition is here
  ::mlir::Identifier calleeAttrName() {
  ~~~~~~~~~~~~~~~~~~ ^

[awarzynski]

Sadly it also breaks my local build. I've tried reverting, but I'm getting merge conflicts, so that's not really an option. @rriddle , would you be able to fix this?

[rriddle]

In D104167#2836490, @awarzynski wrote:

Sadly it also breaks my local build. I've tried reverting, but I'm getting merge conflicts, so that's not really an option. @rriddle , would you be able to fix this?

Yeah, should be fixed in the next 10-20 minutes. Sorry about that.

[mehdi_amini]

I was trying to fix the FIROps.td, but it is quite messy. We'll need a volunteer from the flang crowd for this: a lot of things should be defined in ODS are actually inline C++ in the .td and it conflicts with ODS (like here). It really does not seem maintainable.

[rriddle]

Sent a minimal fix for check-flang in rG00c93d8801f1, but as Mehdi mentions the flang dialects need a good cleanup. There is a lot of C++ code within the .td files, and the code is using the raw API very frequently instead of the auto-generated ODS niceties (e.g. leading to an over-reliance on the string name of the attribute).

[awarzynski]

Many thanks for addressing this. I've not worked on FIROps.td myself, but I can see how the current design/usage might be troublesome. I'll bring this up, though I doubt that anyone will have the spare bandwidth just now. We do get community members volunteering every now and then and I'll keep this on my list of things to improve.

[clementval]

@mehdi_amini @rriddle I'll do a first pass to extract the big code parts out of the .td. If time permits I'll have a look at the raw API usage.

[mehdi_amini]

Thanks @clementval , I'll sent https://bugs.llvm.org/show_bug.cgi?id=50822 your way then :)

[mehdi_amini]

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

315

This function is unused right now?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

 // multiple stride levels (possibly using AffineMaps to specify multiple levels
 // of striding).
 class AffineDmaStartOp
     : public Op<AffineDmaStartOp, OpTrait::MemRefsNormalizable,
                 OpTrait::VariadicOperands, OpTrait::ZeroResult,
                 AffineMapAccessInterface::Trait> {
 public:
   using Op::Op;
+  static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static void build(OpBuilder &builder, OperationState &result, Value srcMemRef,
                     AffineMap srcMap, ValueRange srcIndices, Value destMemRef,
                     AffineMap dstMap, ValueRange destIndices, Value tagMemRef,
                     AffineMap tagMap, ValueRange tagIndices, Value numElements,
                     Value stride = nullptr, Value elementsPerStride = nullptr);
 
   /// Returns the operand index of the source memref.

 //   affine.dma_wait %tag[%index], %num_elements : memref<1xi32, 2>
 //
 class AffineDmaWaitOp
     : public Op<AffineDmaWaitOp, OpTrait::MemRefsNormalizable,
                 OpTrait::VariadicOperands, OpTrait::ZeroResult,
                 AffineMapAccessInterface::Trait> {
 public:
   using Op::Op;
+  static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static void build(OpBuilder &builder, OperationState &result, Value tagMemRef,
                     AffineMap tagMap, ValueRange tagIndices, Value numElements);
 
   static StringRef getOperationName() { return "affine.dma_wait"; }
 
   /// Returns the Tag MemRef associated with the DMA operation being waited on.
   Value getTagMemRef() { return getOperand(0); }

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

 //
 // TODO: add additional operands to allow source and destination striding, and
 // multiple stride levels.
 // TODO: Consider replacing src/dst memref indices with view memrefs.
 class DmaStartOp
     : public Op<DmaStartOp, OpTrait::VariadicOperands, OpTrait::ZeroResult> {
 public:
   using Op::Op;
+  static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static void build(OpBuilder &builder, OperationState &result, Value srcMemRef,
                     ValueRange srcIndices, Value destMemRef,
                     ValueRange destIndices, Value numElements, Value tagMemRef,
                     ValueRange tagIndices, Value stride = nullptr,
                     Value elementsPerStride = nullptr);
 
   // Returns the source MemRefType for this DMA operation.

 //   ...
 //   ...
 //   dma_wait %tag[%index], %num_elements : memref<1 x i32, (d0) -> (d0), 2>
 //
 class DmaWaitOp
     : public Op<DmaWaitOp, OpTrait::VariadicOperands, OpTrait::ZeroResult> {
 public:
   using Op::Op;
+  static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static void build(OpBuilder &builder, OperationState &result, Value tagMemRef,
                     ValueRange tagIndices, Value numElements);
 
   static StringRef getOperationName() { return "memref.dma_wait"; }
 
   // Returns the Tag MemRef associated with the DMA operation being waited on.
   Value getTagMemRef() { return getOperand(0); }

mlir/include/mlir/IR/OperationSupport.h

public:
   /// This constructor is used by Dialect objects when they register the list of
   /// operations they contain.
   template <typename T>
   static void insert(Dialect &dialect) {
     insert(T::getOperationName(), dialect, TypeID::get<T>(),
            T::getParseAssemblyFn(), T::getPrintAssemblyFn(),
            T::getVerifyInvariantsFn(), T::getFoldHookFn(),
            T::getGetCanonicalizationPatternsFn(), T::getInterfaceMap(),
-           T::getHasTraitFn());
+           T::getHasTraitFn(), T::getAttributeNames());
   }
 
   /// Register a new operation in a Dialect object.
   /// The use of this method is in general discouraged in favor of
   /// 'insert<CustomOp>(dialect)'.
   static void
   insert(StringRef name, Dialect &dialect, TypeID typeID,
          ParseAssemblyFn &&parseAssembly, PrintAssemblyFn &&printAssembly,
          VerifyInvariantsFn &&verifyInvariants, FoldHookFn &&foldHook,
          GetCanonicalizationPatternsFn &&getCanonicalizationPatterns,
-         detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait);
+         detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait,
+         ArrayRef<StringRef> attrNames);
+
+  /// Return the list of cached attribute names registered to this operation.
+  /// The order of attributes cached here is unique to each type of operation,
+  /// and the interpretation of this attribute list should generally be driven
+  /// by the respective operation. In many cases, this caching removes the need
+  /// to use the raw string name of a known attribute.
+  ///
+  /// For example the ODS generator, with an op defining the following
+  /// attributes:
+  ///
+  ///   let arguments = (ins I32Attr:$attr1, I32Attr:$attr2);
+  ///
+  /// ... may produce an order here of ["attr1", "attr2"]. This allows for the
+  /// ODS generator to directly access the cached name for a known attribute,
+  /// greatly simplifying the cost and complexity of attribute usage produced by
+  /// the generator.
+  ///
+  ArrayRef<Identifier> getAttributeNames() const { return attributeNames; }
 
 private:
   AbstractOperation(StringRef name, Dialect &dialect, TypeID typeID,
                     ParseAssemblyFn &&parseAssembly,
                     PrintAssemblyFn &&printAssembly,
                     VerifyInvariantsFn &&verifyInvariants,
                     FoldHookFn &&foldHook,
                     GetCanonicalizationPatternsFn &&getCanonicalizationPatterns,
-                    detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait);
+                    detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait,
+                    ArrayRef<Identifier> attrNames);
 
   /// Give Op access to lookupMutable.
   template <typename ConcreteType, template <typename T> class... Traits>
   friend class Op;
 
   /// Look up the specified operation in the specified MLIRContext and return a
   /// pointer to it if present.  Otherwise, return a null pointer.
   static AbstractOperation *lookupMutable(StringRef opName,
                                           MLIRContext *context);
 
   /// A map of interfaces that were registered to this operation.
   detail::InterfaceMap interfaceMap;
 
   /// Internal callback hooks provided by the op implementation.
   FoldHookFn foldHookFn;
   GetCanonicalizationPatternsFn getCanonicalizationPatternsFn;
   HasTraitFn hasTraitFn;
   ParseAssemblyFn parseAssemblyFn;
   PrintAssemblyFn printAssemblyFn;
   VerifyInvariantsFn verifyInvariantsFn;
+
+  /// A list of attribute names registered to this operation in identifier form.
+  /// This allows for operation classes to use identifiers for attribute
+  /// lookup/creation/etc., as opposed to strings.
+  ArrayRef<Identifier> attributeNames;
 };
 
 //===----------------------------------------------------------------------===//
 // NamedAttrList
 //===----------------------------------------------------------------------===//
 
 /// NamedAttrList is array of NamedAttributes that tracks whether it is sorted
 /// and does some basic work to remain sorted.

mlir/lib/IR/MLIRContext.cpp

AbstractOperation *AbstractOperation::lookupMutable(StringRef opName,
   return nullptr;
 }
 
 void AbstractOperation::insert(
     StringRef name, Dialect &dialect, TypeID typeID,
     ParseAssemblyFn &&parseAssembly, PrintAssemblyFn &&printAssembly,
     VerifyInvariantsFn &&verifyInvariants, FoldHookFn &&foldHook,
     GetCanonicalizationPatternsFn &&getCanonicalizationPatterns,
-    detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait) {
-  AbstractOperation opInfo(name, dialect, typeID, std::move(parseAssembly),
-                           std::move(printAssembly),
-                           std::move(verifyInvariants), std::move(foldHook),
-                           std::move(getCanonicalizationPatterns),
-                           std::move(interfaceMap), std::move(hasTrait));
-
-  auto &impl = dialect.getContext()->getImpl();
+    detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait,
+    ArrayRef<StringRef> attrNames) {
+  MLIRContext *ctx = dialect.getContext();
+  auto &impl = ctx->getImpl();
   assert(impl.multiThreadedExecutionContext == 0 &&
          "Registering a new operation kind while in a multi-threaded execution "
          "context");
+
+  // Register the attribute names of this operation.
+  MutableArrayRef<Identifier> cachedAttrNames;
+  if (!attrNames.empty()) {
+    cachedAttrNames = MutableArrayRef<Identifier>(
+        impl.identifierAllocator.Allocate<Identifier>(attrNames.size()),
+        attrNames.size());
+    for (unsigned i : llvm::seq<unsigned>(0, attrNames.size()))
+      new (&cachedAttrNames[i]) Identifier(Identifier::get(attrNames[i], ctx));
+  }
+
+  // Register the information for this operation.
+  AbstractOperation opInfo(
+      name, dialect, typeID, std::move(parseAssembly), std::move(printAssembly),
+      std::move(verifyInvariants), std::move(foldHook),
+      std::move(getCanonicalizationPatterns), std::move(interfaceMap),
+      std::move(hasTrait), cachedAttrNames);
   if (!impl.registeredOperations.insert({name, std::move(opInfo)}).second) {
     llvm::errs() << "error: operation named '" << name
                  << "' is already registered.\n";
     abort();
   }
 }
 
 AbstractOperation::AbstractOperation(
     StringRef name, Dialect &dialect, TypeID typeID,
     ParseAssemblyFn &&parseAssembly, PrintAssemblyFn &&printAssembly,
     VerifyInvariantsFn &&verifyInvariants, FoldHookFn &&foldHook,
     GetCanonicalizationPatternsFn &&getCanonicalizationPatterns,
-    detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait)
+    detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait,
+    ArrayRef<Identifier> attrNames)
     : name(Identifier::get(name, dialect.getContext())), dialect(dialect),
       typeID(typeID), interfaceMap(std::move(interfaceMap)),
       foldHookFn(std::move(foldHook)),
       getCanonicalizationPatternsFn(std::move(getCanonicalizationPatterns)),
       hasTraitFn(std::move(hasTrait)),
       parseAssemblyFn(std::move(parseAssembly)),
       printAssemblyFn(std::move(printAssembly)),
-      verifyInvariantsFn(std::move(verifyInvariants)) {}
+      verifyInvariantsFn(std::move(verifyInvariants)),
+      attributeNames(attrNames) {}
 
 //===----------------------------------------------------------------------===//
 // AbstractType
 //===----------------------------------------------------------------------===//
 
 const AbstractType &AbstractType::lookup(TypeID typeID, MLIRContext *context) {
   const AbstractType *type = lookupMutable(typeID, context);
   if (!type)

mlir/lib/TableGen/OpClass.cpp

 //===----------------------------------------------------------------------===//
 // OpMethod definitions
 //===----------------------------------------------------------------------===//
 
 void OpMethod::writeDeclTo(raw_ostream &os) const {
   os.indent(2);
   if (isStatic())
     os << "static ";
-  if (properties & MP_Constexpr)
+  if ((properties & MP_Constexpr) == MP_Constexpr)
     os << "constexpr ";
   methodSignature.writeDeclTo(os);
-  if (!isInline())
+  if (!isInline()) {
     os << ";";
-  else {
+  } else {
     os << " {\n";
-    methodBody.writeTo(os);
-    os << "}";
+    methodBody.writeTo(os.indent(2));
+    os.indent(2) << "}";
   }
 }
 
 void OpMethod::writeDefTo(raw_ostream &os, StringRef namePrefix) const {
   // Do not write definition if the method is decl only.
   if (properties & MP_Declaration)
     return;
   // Do not generate separate definition for inline method

mlir/test/mlir-tblgen/op-attribute.td

 def AOp : NS_Op<"a_op", []> {
   let arguments = (ins
       SomeAttr:$aAttr,
       DefaultValuedAttr<SomeAttr, "4.2">:$bAttr,
       OptionalAttr<SomeAttr>:$cAttr
   );
 }
 
+// DECL-LABEL: AOp declarations
+
+// Test attribute name methods
+// ---
+
+// DECL:      static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames()
+// DECL-NEXT:   static ::llvm::StringRef attrNames[] =
+// DECL-SAME:     {::llvm::StringRef("aAttr"), ::llvm::StringRef("bAttr"), ::llvm::StringRef("cAttr")};
+// DECL-NEXT:   return attrNames;
+
+// DECL:      ::mlir::Identifier aAttrAttrName()
+// DECL-NEXT:      return getAttributeNameForIndex(0);
+// DECL:      ::mlir::Identifier aAttrAttrName(::mlir::OperationName name)
+// DECL-NEXT:      return getAttributeNameForIndex(name, 0);
+
+// DECL:      ::mlir::Identifier bAttrAttrName()
+// DECL-NEXT:      return getAttributeNameForIndex(1);
+// DECL:      ::mlir::Identifier bAttrAttrName(::mlir::OperationName name)
+// DECL-NEXT:      return getAttributeNameForIndex(name, 1);
+
+// DECL:      ::mlir::Identifier cAttrAttrName()
+// DECL-NEXT:      return getAttributeNameForIndex(2);
+// DECL:      ::mlir::Identifier cAttrAttrName(::mlir::OperationName name)
+// DECL-NEXT:      return getAttributeNameForIndex(name, 2);
+
 // DEF-LABEL: AOp definitions
 
 // Test verify method
 // ---
 
 // DEF:      ::mlir::LogicalResult AOpAdaptor::verify
 // DEF:      auto tblgen_aAttr = odsAttrs.get("aAttr");
 // DEF-NEXT: if (!tblgen_aAttr) return emitError(loc, "'test.a_op' op ""requires attribute 'aAttr'");
 // DEF:        if (!((some-condition))) return emitError(loc, "'test.a_op' op ""attribute 'aAttr' failed to satisfy constraint: some attribute kind");
 // DEF:      auto tblgen_bAttr = odsAttrs.get("bAttr");
 // DEF-NEXT: if (tblgen_bAttr) {
 // DEF-NEXT:   if (!((some-condition))) return emitError(loc, "'test.a_op' op ""attribute 'bAttr' failed to satisfy constraint: some attribute kind");
 // DEF:      auto tblgen_cAttr = odsAttrs.get("cAttr");
 // DEF-NEXT: if (tblgen_cAttr) {
 // DEF-NEXT:   if (!((some-condition))) return emitError(loc, "'test.a_op' op ""attribute 'cAttr' failed to satisfy constraint: some attribute kind");
 
 // Test getter methods
 // ---
 
 // DEF:      some-attr-kind AOp::aAttrAttr()
-// DEF-NEXT:   (*this)->getAttr("aAttr").template cast<some-attr-kind>()
+// DEF-NEXT:   (*this)->getAttr(aAttrAttrName()).template cast<some-attr-kind>()
 // DEF:      some-return-type AOp::aAttr() {
 // DEF-NEXT:   auto attr = aAttrAttr()
 // DEF-NEXT:   return attr.some-convert-from-storage();
 
 // DEF:      some-attr-kind AOp::bAttrAttr()
-// DEF-NEXT:   return (*this)->getAttr("bAttr").template dyn_cast_or_null<some-attr-kind>()
+// DEF-NEXT:   return (*this)->getAttr(bAttrAttrName()).template dyn_cast_or_null<some-attr-kind>()
 // DEF:      some-return-type AOp::bAttr() {
 // DEF-NEXT:   auto attr = bAttrAttr();
 // DEF-NEXT:   if (!attr)
 // DEF-NEXT:       return some-const-builder-call(::mlir::Builder((*this)->getContext()), 4.2).some-convert-from-storage();
 // DEF-NEXT:   return attr.some-convert-from-storage();
 
 // DEF:      some-attr-kind AOp::cAttrAttr()
-// DEF-NEXT:   return (*this)->getAttr("cAttr").template dyn_cast_or_null<some-attr-kind>()
+// DEF-NEXT:   return (*this)->getAttr(cAttrAttrName()).template dyn_cast_or_null<some-attr-kind>()
 // DEF:      ::llvm::Optional<some-return-type> AOp::cAttr() {
 // DEF-NEXT:   auto attr = cAttrAttr()
 // DEF-NEXT:   return attr ? ::llvm::Optional<some-return-type>(attr.some-convert-from-storage()) : (::llvm::None);
 
 // Test setter methods
 // ---
 
 // DEF:      void AOp::aAttrAttr(some-attr-kind attr) {
-// DEF-NEXT:   (*this)->setAttr("aAttr", attr);
+// DEF-NEXT:   (*this)->setAttr(aAttrAttrName(), attr);
 // DEF:      void AOp::bAttrAttr(some-attr-kind attr) {
-// DEF-NEXT:   (*this)->setAttr("bAttr", attr);
+// DEF-NEXT:   (*this)->setAttr(bAttrAttrName(), attr);
 // DEF:      void AOp::cAttrAttr(some-attr-kind attr) {
-// DEF-NEXT:   (*this)->setAttr("cAttr", attr);
+// DEF-NEXT:   (*this)->setAttr(cAttrAttrName(), attr);
 
 // Test remove methods
 // ---
 
 // DEF: ::mlir::Attribute AOp::removeCAttrAttr() {
-// DEF-NEXT: return (*this)->removeAttr("cAttr");
+// DEF-NEXT: return (*this)->removeAttr(cAttrAttrName());
 
 // Test build methods
 // ---
 
 // DEF:      void AOp::build(
-// DEF:        odsState.addAttribute("aAttr", aAttr);
-// DEF:        odsState.addAttribute("bAttr", bAttr);
+// DEF:        odsState.addAttribute(aAttrAttrName(odsState.name), aAttr);
+// DEF:        odsState.addAttribute(bAttrAttrName(odsState.name), bAttr);
 // DEF:        if (cAttr) {
-// DEF-NEXT:     odsState.addAttribute("cAttr", cAttr);
+// DEF-NEXT:     odsState.addAttribute(cAttrAttrName(odsState.name), cAttr);
 
 // DEF:      void AOp::build(
 // DEF:        some-return-type aAttr, some-return-type bAttr, /*optional*/some-attr-kind cAttr
-// DEF:        odsState.addAttribute("aAttr", some-const-builder-call(odsBuilder, aAttr));
+// DEF:        odsState.addAttribute(aAttrAttrName(odsState.name), some-const-builder-call(odsBuilder, aAttr));
 
 // DEF:      void AOp::build(
 // DEF:        ::llvm::ArrayRef<::mlir::NamedAttribute> attributes
 // DEF:      odsState.addAttributes(attributes);
 
 def SomeTypeAttr : TypeAttrBase<"SomeType", "some type attribute">;
 
 def BOp : NS_Op<"b_op", []> {
let arguments = (ins
     DefaultValuedAttr<SomeI32Enum, "::SomeI32Enum::case5">:$dv_enum_attr
   );
 }
 
 // DECL-LABEL: DOp declarations
 // DECL: static void build({{.*}}, uint32_t i32_attr, ::llvm::APFloat f64_attr, ::llvm::StringRef str_attr, bool bool_attr, ::SomeI32Enum enum_attr, uint32_t dv_i32_attr, ::llvm::APFloat dv_f64_attr, ::llvm::StringRef dv_str_attr = "abc", bool dv_bool_attr = true, ::SomeI32Enum dv_enum_attr = ::SomeI32Enum::case5)
 
 // DEF-LABEL: DOp definitions
-// DEF: odsState.addAttribute("str_attr", odsBuilder.getStringAttr(str_attr));
-// DEF: odsState.addAttribute("dv_str_attr", odsBuilder.getStringAttr(dv_str_attr));
+// DEF: odsState.addAttribute(str_attrAttrName(odsState.name), odsBuilder.getStringAttr(str_attr));
+// DEF: odsState.addAttribute(dv_str_attrAttrName(odsState.name), odsBuilder.getStringAttr(dv_str_attr));
 
 // Test derived type attr.
 // ---
 def DerivedTypeAttrOp : NS_Op<"derived_type_attr_op", []> {
   let results = (outs AnyTensor:$output);
   DerivedTypeAttr element_dtype = DerivedTypeAttr<"return output().getType();">;
 }
 
def UnitAttrOp : NS_Op<"unit_attr_op", []> {
   let arguments = (ins UnitAttr:$attr);
 }
 
 // DEF-LABEL: UnitAttrOp definitions
 // DEF: bool UnitAttrOp::attr() {
 // DEF:   return {{.*}} != nullptr
 
 // DEF: ::mlir::Attribute UnitAttrOp::removeAttrAttr() {
-// DEF-NEXT:   (*this)->removeAttr("attr");
+// DEF-NEXT:   (*this)->removeAttr(attrAttrName());
 
 // DEF: build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, /*optional*/::mlir::UnitAttr attr)
 
 
 // Test elementAttr field of TypedArrayAttr.
 // ---
 
 def SomeTypedArrayAttr : TypedArrayAttrBase<SomeAttr, "SomeAttr array">;

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

 #include "mlir/TableGen/CodeGenHelpers.h"
 #include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/Interfaces.h"
 #include "mlir/TableGen/OpClass.h"
 #include "mlir/TableGen/Operator.h"
 #include "mlir/TableGen/SideEffects.h"
 #include "mlir/TableGen/Trait.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 #include "llvm/TableGen/TableGenBackend.h"
 

 // (variadic or not).
 //
 // {0}: The name of the attribute specifying the segment sizes.
 const char *adapterSegmentSizeAttrInitCode = R"(
   assert(odsAttrs && "missing segment size attribute for op");
   auto sizeAttr = odsAttrs.get("{0}").cast<::mlir::DenseIntElementsAttr>();
 )";
 const char *opSegmentSizeAttrInitCode = R"(
-  auto sizeAttr = (*this)->getAttr("{0}").cast<::mlir::DenseIntElementsAttr>();
+  auto sizeAttr = (*this)->getAttr({0}).cast<::mlir::DenseIntElementsAttr>();
 )";
 const char *attrSizedSegmentValueRangeCalcCode = R"(
   auto sizeAttrValues = sizeAttr.getValues<uint32_t>();
   unsigned start = 0;
   for (unsigned i = 0; i < index; ++i)
     start += *(sizeAttrValues.begin() + i);
   unsigned size = *(sizeAttrValues.begin() + index);
   return {start, size};

 
 private:
   OpEmitter(const Operator &op,
             const StaticVerifierFunctionEmitter &staticVerifierEmitter);
 
   void emitDecl(raw_ostream &os);
   void emitDef(raw_ostream &os);
 
+  // Generate methods for accessing the attribute names of this operation.
+  void genAttrNameGetters();
+
+  // Return the index of the given attribute name. This is a relative ordering
+  // for this name, used in attribute getters.
+  unsigned getAttrNameIndex(StringRef attrName) const;
+
   // Generates the OpAsmOpInterface for this operation if possible.
   void genOpAsmInterface();
 
   // Generates the `getOperationName` method for this op.
   void genOpNameGetter();
 
   // Generates getters for the attributes.
   void genAttrGetters();
private:
   // The C++ code builder for this op
   OpClass opClass;
 
   // The format context for verification code generation.
   FmtContext verifyCtx;
 
   // The emitter containing all of the locally emitted verification functions.
   const StaticVerifierFunctionEmitter &staticVerifierEmitter;
+
+  // A map of attribute names (including implicit attributes) registered to the
+  // current operation, to the relative order in which they were registered.
+  llvm::MapVector<StringRef, unsigned> attributeNames;
 };
 } // end anonymous namespace
 
 // Populate the format context `ctx` with substitutions of attributes, operands
 // and results.
 // - attrGet corresponds to the name of the function to call to get value of
 //   attribute (the generated function call returns an Attribute);
 // - operandGet corresponds to the name of the function with which to retrieve
: def(op.getDef()), op(op),
       staticVerifierEmitter(staticVerifierEmitter) {
   verifyCtx.withOp("(*this->getOperation())");
   verifyCtx.addSubst("_ctxt", "this->getOperation()->getContext()");
 
   genTraits();
 
   // Generate C++ code for various op methods. The order here determines the
   // methods in the generated file.
+  genAttrNameGetters();
   genOpAsmInterface();
   genOpNameGetter();
   genNamedOperandGetters();
   genNamedOperandSetters();
   genNamedResultGetters();
   genNamedRegionGetters();
   genNamedSuccessorGetters();
   genAttrGetters();
void OpEmitter::emitDef(
     const StaticVerifierFunctionEmitter &staticVerifierEmitter) {
   OpEmitter(op, staticVerifierEmitter).emitDef(os);
 }
 
 void OpEmitter::emitDecl(raw_ostream &os) { opClass.writeDeclTo(os); }
 
 void OpEmitter::emitDef(raw_ostream &os) { opClass.writeDefTo(os); }
 
+void OpEmitter::genAttrNameGetters() {
+  // Enumerate the attribute names of this op, assigning each a relative
+  // ordering.
+  auto addAttrName = [&](StringRef name) {
+    unsigned index = attributeNames.size();
+    attributeNames.insert({name, index});
+  };
+  for (const NamedAttribute &namedAttr : op.getAttributes())
+    addAttrName(namedAttr.name);
+  // Include key attributes from several traits as implicitly registered.
+  if (op.getTrait("::mlir::OpTrait::AttrSizedOperandSegments"))
+    addAttrName("operand_segment_sizes");
+  if (op.getTrait("::mlir::OpTrait::AttrSizedResultSegments"))
+    addAttrName("result_segment_sizes");
+
+  // Emit the getAttributeNames method.
+  {
+    auto *method = opClass.addMethodAndPrune(
+        "::llvm::ArrayRef<::llvm::StringRef>", "getAttributeNames",
+        OpMethod::Property(OpMethod::MP_Static | OpMethod::MP_Inline));
+    auto &body = method->body();
+    if (attributeNames.empty()) {
+      body << "  return {};";
+    } else {
+      body << "  static ::llvm::StringRef attrNames[] = {";
+      llvm::interleaveComma(llvm::make_first_range(attributeNames), body,
+                            [&](StringRef attrName) {
+                              body << "::llvm::StringRef(\"" << attrName
+                                   << "\")";
+                            });
+      body << "};\n  return attrNames;";
+    }
+  }
+
+  // Emit the getAttributeNameForIndex methods.
+  {
+    auto *method = opClass.addMethodAndPrune(
+        "::mlir::Identifier", "getAttributeNameForIndex",
+        OpMethod::Property(OpMethod::MP_Private | OpMethod::MP_Inline),
+        "unsigned", "index");
+    method->body()
+        << "  return getAttributeNameForIndex((*this)->getName(), index);";
+  }
+  {
+    auto *method = opClass.addMethodAndPrune(
+        "::mlir::Identifier", "getAttributeNameForIndex",
+        OpMethod::Property(OpMethod::MP_Private | OpMethod::MP_Inline |
+                           OpMethod::MP_Static),
+        "::mlir::OperationName name, unsigned index");
+    method->body() << "assert(index < " << attributeNames.size()
+                   << " && \"invalid attribute index\");\n"
+                      "  return name.getAbstractOperation()"
+                      "->getAttributeNames()[index];";
+  }
+
+  // Generate the <attr>AttrName methods, that expose the attribute names to
+  // users.
+  const char *attrNameMethodBody = "  return getAttributeNameForIndex({0});";
+  for (const std::pair<StringRef, unsigned> &attrIt : attributeNames) {
+    std::string methodName = (attrIt.first + "AttrName").str();
+
+    // Generate the non-static variant.
+    {
+      auto *method =
+          opClass.addMethodAndPrune("::mlir::Identifier", methodName,
+                                    OpMethod::Property(OpMethod::MP_Inline));
+      method->body() << llvm::formatv(attrNameMethodBody, attrIt.second).str();
+    }
+
+    // Generate the static variant.
+    {
+      auto *method = opClass.addMethodAndPrune(
+          "::mlir::Identifier", methodName,
+          OpMethod::Property(OpMethod::MP_Inline | OpMethod::MP_Static),
+          "::mlir::OperationName", "name");
+      method->body() << llvm::formatv(attrNameMethodBody,
+                                      "name, " + Twine(attrIt.second))
+                            .str();
+    }
+  }
+}
+
+unsigned OpEmitter::getAttrNameIndex(StringRef attrName) const {
+  auto it = attributeNames.find(attrName);
+  assert(it != attributeNames.end() && "expected attribute name to have been "
+                                       "registered in genAttrNameGetters");
+  return it->second;
+}
+
 void OpEmitter::genAttrGetters() {
   FmtContext fctx;
   fctx.withBuilder("::mlir::Builder((*this)->getContext())");
 
-  Dialect opDialect = op.getDialect();
   // Emit the derived attribute body.
   auto emitDerivedAttr = [&](StringRef name, Attribute attr) {
-    auto *method = opClass.addMethodAndPrune(attr.getReturnType(), name);
-    if (!method)
-      return;
-    auto &body = method->body();
-    body << "  " << attr.getDerivedCodeBody() << "\n";
+    if (auto *method = opClass.addMethodAndPrune(attr.getReturnType(), name))
+      method->body() << "  " << attr.getDerivedCodeBody() << "\n";
   };
 
   // Emit with return type specified.
   auto emitAttrWithReturnType = [&](StringRef name, Attribute attr) {
     auto *method = opClass.addMethodAndPrune(attr.getReturnType(), name);
     auto &body = method->body();
     body << "  auto attr = " << name << "Attr();\n";
     if (attr.hasDefaultValue()) {
void OpEmitter::genAttrGetters() {
   // referring to the attributes via accessors instead of having to use
   // the string interface for better compile time verification.
   auto emitAttrWithStorageType = [&](StringRef name, Attribute attr) {
     auto *method =
         opClass.addMethodAndPrune(attr.getStorageType(), (name + "Attr").str());
     if (!method)
       return;
     auto &body = method->body();
-    body << "  return (*this)->getAttr(\"" << name << "\").template ";
+    body << "  return (*this)->getAttr(" << name << "AttrName()).template ";
     if (attr.isOptional() || attr.hasDefaultValue())
       body << "dyn_cast_or_null<";
     else
       body << "cast<";
     body << attr.getStorageType() << ">();";
   };
 
-  for (auto &namedAttr : op.getAttributes()) {
-    const auto &name = namedAttr.name;
-    const auto &attr = namedAttr.attr;
-    if (attr.isDerivedAttr()) {
-      emitDerivedAttr(name, attr);
+  for (const NamedAttribute &namedAttr : op.getAttributes()) {
+    if (namedAttr.attr.isDerivedAttr()) {
+      emitDerivedAttr(namedAttr.name, namedAttr.attr);
     } else {
-      emitAttrWithStorageType(name, attr);
-      emitAttrWithReturnType(name, attr);
+      emitAttrWithStorageType(namedAttr.name, namedAttr.attr);
+      emitAttrWithReturnType(namedAttr.name, namedAttr.attr);
     }
   }
 
   auto derivedAttrs = make_filter_range(op.getAttributes(),
                                         [](const NamedAttribute &namedAttr) {
                                           return namedAttr.attr.isDerivedAttr();
                                         });
   if (!derivedAttrs.empty()) {
// Generate method to materialize derived attributes as a DictionaryAttr.
       body << "  ::mlir::MLIRContext* ctx = getContext();\n";
       body << "  ::mlir::Builder odsBuilder(ctx); (void)odsBuilder;\n";
       body << "  return ::mlir::DictionaryAttr::get(";
       body << "  ctx, {\n";
       interleave(
           derivedAttrs, body,
           [&](const NamedAttribute &namedAttr) {
             auto tmpl = namedAttr.attr.getConvertFromStorageCall();
-            body << "    {::mlir::Identifier::get(\"" << namedAttr.name
-                 << "\", ctx),\n"
+            body << "    {" << namedAttr.name << "AttrName(),\n"
                  << tgfmt(tmpl, &fctx.withSelf(namedAttr.name + "()")
                                      .withBuilder("odsBuilder")
                                      .addSubst("_ctx", "ctx"))
                  << "}";
           },
           ",\n");
       body << "});";
     }
   }
 }
 
 void OpEmitter::genAttrSetters() {
   // Generate raw named setter type. This is a wrapper class that allows setting
   // to the attributes via setters instead of having to use the string interface
   // for better compile time verification.
   auto emitAttrWithStorageType = [&](StringRef name, Attribute attr) {
     auto *method = opClass.addMethodAndPrune("void", (name + "Attr").str(),
                                              attr.getStorageType(), "attr");
-    if (!method)
-      return;
-    auto &body = method->body();
-    body << "  (*this)->setAttr(\"" << name << "\", attr);";
+    if (method)
+      method->body() << "  (*this)->setAttr(" << name << "AttrName(), attr);";
   };
 
-  for (auto &namedAttr : op.getAttributes()) {
-    const auto &name = namedAttr.name;
-    const auto &attr = namedAttr.attr;
-    if (!attr.isDerivedAttr())
-      emitAttrWithStorageType(name, attr);
-  }
+  for (const NamedAttribute &namedAttr : op.getAttributes())
+    if (!namedAttr.attr.isDerivedAttr())
+      emitAttrWithStorageType(namedAttr.name, namedAttr.attr);
 }
 
 void OpEmitter::genOptionalAttrRemovers() {
   // Generate methods for removing optional attributes, instead of having to
   // use the string interface. Enables better compile time verification.
   auto emitRemoveAttr = [&](StringRef name) {
     auto upperInitial = name.take_front().upper();
     auto suffix = name.drop_front();
     auto *method = opClass.addMethodAndPrune(
         "::mlir::Attribute", ("remove" + upperInitial + suffix + "Attr").str());
     if (!method)
       return;
-    auto &body = method->body();
-    body << "  return (*this)->removeAttr(\"" << name << "\");";
+    method->body() << "  return (*this)->removeAttr(" << name << "AttrName());";
   };
 
-  for (const auto &namedAttr : op.getAttributes()) {
-    const auto &name = namedAttr.name;
-    const auto &attr = namedAttr.attr;
-    if (attr.isOptional())
-      emitRemoveAttr(name);
-  }
+  for (const NamedAttribute &namedAttr : op.getAttributes())
+    if (namedAttr.attr.isOptional())
+      emitRemoveAttr(namedAttr.name);
 }
 
 // Generates the code to compute the start and end index of an operand or result
 // range.
 template <typename RangeT>
 static void
 generateValueRangeStartAndEnd(Class &opClass, StringRef methodName,
                               int numVariadic, int numNonVariadic,
for (int i = 0; i != numOperands; ++i) {
     } else {
       m = opClass.addMethodAndPrune("::mlir::Value", operand.name);
       m->body() << "  return *getODSOperands(" << i << ").begin();";
     }
   }
 }
 
 void OpEmitter::genNamedOperandGetters() {
+  // Build the code snippet used for initializing the operand_segment_sizes
+  // array.
+  std::string attrSizeInitCode;
+  if (op.getTrait("::mlir::OpTrait::AttrSizedOperandSegments")) {
+    attrSizeInitCode =
+        formatv(opSegmentSizeAttrInitCode, "operand_segment_sizesAttrName()")
+            .str();
+  }
+
   generateNamedOperandGetters(
       op, opClass,
-      /*sizeAttrInit=*/
-      formatv(opSegmentSizeAttrInitCode, "operand_segment_sizes").str(),
+      /*sizeAttrInit=*/attrSizeInitCode,
       /*rangeType=*/"::mlir::Operation::operand_range",
       /*rangeBeginCall=*/"getOperation()->operand_begin()",
       /*rangeSizeCall=*/"getOperation()->getNumOperands()",
       /*getOperandCallPattern=*/"getOperation()->getOperand({0})");
 }
 
 void OpEmitter::genNamedOperandSetters() {
   auto *attrSizedOperands =
       op.getTrait("::mlir::OpTrait::AttrSizedOperandSegments");
   for (int i = 0, e = op.getNumOperands(); i != e; ++i) {
     const auto &operand = op.getOperand(i);
     if (operand.name.empty())
       continue;
     auto *m = opClass.addMethodAndPrune("::mlir::MutableOperandRange",
                                         (operand.name + "Mutable").str());
     auto &body = m->body();
     body << "  auto range = getODSOperandIndexAndLength(" << i << ");\n"
          << "  return ::mlir::MutableOperandRange(getOperation(), "
             "range.first, range.second";
     if (attrSizedOperands)
       body << ", ::mlir::MutableOperandRange::OperandSegment(" << i
            << "u, *getOperation()->getAttrDictionary().getNamed("
-              "\"operand_segment_sizes\"))";
+              "operand_segment_sizesAttrName()))";
     body << ");\n";
   }
 }
 
 void OpEmitter::genNamedResultGetters() {
   const int numResults = op.getNumResults();
   const int numVariadicResults = op.getNumVariableLengthResults();
   const int numNormalResults = numResults - numVariadicResults;
void OpEmitter::genNamedResultGetters() {
   }
 
   if (attrSizedResults && sameVariadicSize) {
     PrintFatalError(op.getLoc(),
                     "op cannot have both 'AttrSizedResultSegments' and "
                     "'SameVariadicResultSize' traits");
   }
 
+  // Build the initializer string for the result segment size attribute.
+  std::string attrSizeInitCode;
+  if (attrSizedResults) {
+    attrSizeInitCode =
+        formatv(opSegmentSizeAttrInitCode, "result_segment_sizesAttrName()")
+            .str();
+  }
+
   generateValueRangeStartAndEnd(
       opClass, "getODSResultIndexAndLength", numVariadicResults,
       numNormalResults, "getOperation()->getNumResults()", attrSizedResults,
-      formatv(opSegmentSizeAttrInitCode, "result_segment_sizes").str(),
-      op.getResults());
+      attrSizeInitCode, op.getResults());
 
   auto *m = opClass.addMethodAndPrune("::mlir::Operation::result_range",
                                       "getODSResults", "unsigned", "index");
   m->body() << formatv(valueRangeReturnCode, "getOperation()->result_begin()",
                        "getODSResultIndexAndLength(index)");
 
   for (int i = 0; i != numResults; ++i) {
     const auto &result = op.getResult(i);
if (!m)
     return;
 
   auto &body = m->body();
 
   // Push all result types to the operation state
   std::string resultType;
   const auto &namedAttr = op.getAttribute(0);
 
-  body << "  for (auto attr : attributes) {\n";
-  body << "    if (attr.first != \"" << namedAttr.name << "\") continue;\n";
+  body << "  auto attrName = " << namedAttr.name << "AttrName("
+       << builderOpState
+       << ".name);\n"
+          "  for (auto attr : attributes) {\n"
+          "    if (attr.first != attrName) continue;\n";
   if (namedAttr.attr.isTypeAttr()) {
     resultType = "attr.second.cast<::mlir::TypeAttr>().getValue()";
   } else {
     resultType = "attr.second.getType()";
   }
 
   // Operands
   body << "  " << builderOpState << ".addOperands(operands);\n";
for (int i = 0, e = op.getNumOperands(); i < e; ++i) {
     if (op.getOperand(i).isOptional())
       body << "  if (" << argName << ")\n  ";
     body << "  " << builderOpState << ".addOperands(" << argName << ");\n";
   }
 
   // If the operation has the operand segment size attribute, add it here.
   if (op.getTrait("::mlir::OpTrait::AttrSizedOperandSegments")) {
     body << "  " << builderOpState
-         << ".addAttribute(\"operand_segment_sizes\", "
-            "odsBuilder.getI32VectorAttr({";
+         << ".addAttribute(operand_segment_sizesAttrName(" << builderOpState
+         << ".name), "
+         << "odsBuilder.getI32VectorAttr({";
     interleaveComma(llvm::seq<int>(0, op.getNumOperands()), body, [&](int i) {
       if (op.getOperand(i).isOptional())
         body << "(" << getArgumentName(op, i) << " ? 1 : 0)";
       else if (op.getOperand(i).isVariadic())
         body << "static_cast<int32_t>(" << getArgumentName(op, i) << ".size())";
       else
         body << "1";
     });
     body << "}));\n";
   }
 
   // Push all attributes to the result.
   for (const auto &namedAttr : op.getAttributes()) {
     auto &attr = namedAttr.attr;
     if (!attr.isDerivedAttr()) {
       bool emitNotNullCheck = attr.isOptional();
-      if (emitNotNullCheck) {
+      if (emitNotNullCheck)
         body << formatv("  if ({0}) ", namedAttr.name) << "{\n";
-      }
+
       if (isRawValueAttr && canUseUnwrappedRawValue(attr)) {
         // If this is a raw value, then we need to wrap it in an Attribute
         // instance.
         FmtContext fctx;
         fctx.withBuilder("odsBuilder");
 
         std::string builderTemplate =
             std::string(attr.getConstBuilderTemplate());
 
         // For StringAttr, its constant builder call will wrap the input in
         // quotes, which is correct for normal string literals, but incorrect
         // here given we use function arguments. So we need to strip the
         // wrapping quotes.
         if (StringRef(builderTemplate).contains("\"$0\""))
           builderTemplate = replaceAllSubstrs(builderTemplate, "\"$0\"", "$0");
 
         std::string value =
             std::string(tgfmt(builderTemplate, &fctx, namedAttr.name));
-        body << formatv("  {0}.addAttribute(\"{1}\", {2});\n", builderOpState,
-                        namedAttr.name, value);
+        body << formatv("  {0}.addAttribute({1}AttrName({0}.name), {2});\n",
+                        builderOpState, namedAttr.name, value);
       } else {
-        body << formatv("  {0}.addAttribute(\"{1}\", {1});\n", builderOpState,
-                        namedAttr.name);
+        body << formatv("  {0}.addAttribute({1}AttrName({0}.name), {1});\n",
+                        builderOpState, namedAttr.name);
       }
-      if (emitNotNullCheck) {
+      if (emitNotNullCheck)
         body << "  }\n";
-      }
-    }
-  }
 
   // Create the correct number of regions.
   for (const NamedRegion &region : op.getRegions()) {
     if (region.isVariadic())
       body << formatv("  for (unsigned i = 0; i < {0}Count; ++i)\n  ",
                       region.name);
 
     body << "  (void)" << builderOpState << ".addRegion();\n";

mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

struct TypeNoLayout : public Type::TypeBase<TypeNoLayout, Type, TypeStorage> {
   static TypeNoLayout get(MLIRContext *ctx) { return Base::get(ctx); }
 };
 
 /// An op that serves as scope for data layout queries with the relevant
 /// attribute attached. This can handle data layout requests for the built-in
 /// types itself.
 struct OpWithLayout : public Op<OpWithLayout, DataLayoutOpInterface::Trait> {
   using Op::Op;
+  static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static StringRef getOperationName() { return "dltest.op_with_layout"; }
 
   DataLayoutSpecInterface getDataLayoutSpec() {
     return getOperation()->getAttrOfType<DataLayoutSpecInterface>(kAttrName);
   }
 
   static unsigned getTypeSizeInBits(Type type, const DataLayout &dataLayout,
static unsigned getTypePreferredAlignment(Type type,
                                             DataLayoutEntryListRef params) {
     return 2 * getTypeABIAlignment(type, dataLayout, params);
   }
 };
 
 struct OpWith7BitByte
     : public Op<OpWith7BitByte, DataLayoutOpInterface::Trait> {
   using Op::Op;
+  static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static StringRef getOperationName() { return "dltest.op_with_7bit_byte"; }
 
   DataLayoutSpecInterface getDataLayoutSpec() {
     return getOperation()->getAttrOfType<DataLayoutSpecInterface>(kAttrName);
   }
 
   // Bytes are assumed to be 7-bit here.