[mlir] Add a builtin distinct attribute

[gysit]

Summary

A distinct attribute associates a referenced attribute to a unique
identifier. Every call to its create function allocates a new
distinct attribute instance. The address of the attribute instance
temporarily serves as its unique identifier. Similar to the names
of SSA values, the final unique identifiers are generated during
pretty printing.

Examples:
#distinct = distinct[0]<42.0 : f32>
#distinct1 = distinct[1]<42.0 : f32>
#distinct2 = distinct[2]<array<i32: 10, 42>>

This mechanism is meant to generate attributes with a unique
identifier, which can be used to mark groups of operations
that share a common properties such as if they are aliasing.

The design of the distinct attribute ensures minimal memory
footprint per distinct attribute since it only contains a reference
to another attribute. All distinct attributes are stored outside of
the storage uniquer in a thread local store that is part of the
context. It uses one bump pointer allocator per thread to ensure
distinct attributes can be created in-parallel.

[gysit]

This revision implements a custom numbering of distinct attributes during printing an parsing. This has the advantage that distinct attributes can be printing both as aliases or inline on the operation itself.

An alternative may be to use the alias printing mechanism to get a numbering for free. I have a WIP revision that implements this:

https://reviews.llvm.org/D153361

It is a bit more intrusive in terms of the necessary printer and parser changes though and I am unsure if we want to give up on the option of printing every attribute inline (not as an alias). If desired I can merge the WIP commit into this revision.

[gysit]

This are the relevant links for context:

https://discourse.llvm.org/t/rfc-use-attributes-to-model-distinct-llvm-metadata-nodes/69924/
https://discourse.llvm.org/t/open-mlir-meeting-5-24-2023-rfc-discussion-about-distinct-attribute-in-mlir/70856

[Dinistro]

Added some small comment.

mlir/include/mlir/Support/ThreadLocalStore.h

45 ↗

(On Diff #532946)

Just to be sure: If I understand this correctly, each template instantiation will have its own bump allocator. Would it be more efficient to combine them into one bump allocator per thread? If so, this shouldn't be a static member function. Instead, it could just be a static function.

mlir/lib/AsmParser/AttributeParser.cpp

1235

Ultra NIT: I prefer "ID" over "Id", as it is an abbreviation.

[mehdi_amini]

mlir/include/mlir/Support/ThreadLocalStore.h

46 ↗

(On Diff #532946)

That is leaking memory for the entire process duration, it's not clear to me why this is needed?

mlir/lib/AsmParser/AttributeParser.cpp

1263

Can you print the existing definition?

[gysit]

Address comments and rebase.

[gysit]

@mehdi_amini Thanks for the review comments. ThreadLocalCache indeed already implements a thread local storage with a non-static lifetime. I changed my code to use it.

[mehdi_amini]

mlir/lib/IR/AttributeDetail.h

417

Why do we need all of this instead of just defining DistinctAttr in ODS as any other attribute?

mlir/test/IR/distinct-attr.mlir

3

Can you stick to using the test dialect and avoid -allow-unregistered-dialect ? This isn't really good practice to use this option: it disables some safety net.

[gysit]

mlir/lib/IR/AttributeDetail.h

417

The distinct attribute uses a different uniquer which cannot be generated using tablegen:

class DistinctAttr
    : public detail::StorageUserBase<DistinctAttr, Attribute,
                                     detail::DistinctAttrStorage,
                                     detail::DistinctAttributeUniquer> {

There is at least currently no way to set the uniquer to something else than the default storage uniquer. The distinct attribute uses a separate DistinctAttributeUniquer though. In theory ODS could be changed to support this but I don't think we want this since there will be only one builtin DistinctAttr and users are not supposed to create there own distinct attribute types.

I can try to put a comment somewhere that explains this design decision. Probably on the distinct attribute declaration.

[mehdi_amini]

mlir/lib/IR/AttributeDetail.h

417

Oh right I forgot about the uniquer at some point...

[gysit]

Address comments.

[Mogball]

Seems useful!

[zero9178]

LGTM

[gysit]

Are there any more comments? @rriddle @Mogball @mehdi_amini maybe?

[Dinistro]

I did another pass over the revision and it's LGTM!
Not sure if the others have additional comments, though.

[rriddle]

This patch looks like it's missing bytecode support, have you tested that case? Failure there will be a miscompile, not a hard error.

mlir/include/mlir/IR/BuiltinAttributes.h

1012–1041

All attributes should be defined in ODS, this should be moved over. The custom allocator could be solved other ways (e.g. a parameter that is a key that we have a custom allocator/counter/etc for inside of the context).

[gysit]

In D153360#4472092, @rriddle wrote:

This patch looks like it's missing bytecode support, have you tested that case? Failure there will be a miscompile, not a hard error.

The revision has two changes related to bytecode:

• mlir/include/mlir/IR/BuiltinDialectBytecode.td
• mlir/test/Dialect/Builtin/Bytecode/attrs.mlir

I would expect these roundtrip tests fail if the implementation is broken/incomplete?

mlir/include/mlir/IR/BuiltinAttributes.h

1012–1041

Let me rephrase to make sure I understand correctly. You would introduce a new tablegen attribute parameter type derived from AttrOrTypeParameter that uses a custom allocator and a custom printer and parser? A distinct attribute could then use a parameter of this new key type which for example contains the address of an allocation made using a custom allocator. A custom printer/parser then converts the address to a deterministic identifier. Does that make sense?

I think such a solution improves composability since a user attribute could contain such a key parameter and no boxing into a distinct attribute is needed.

There are also downsides:

• The distinctness mechanism is less encapsulated since the get method will take a key parameter and there will be an accessor for the key parameter.
• The storage usage may increase since there is an extra allocation for the key. An attribute that contains such a key will thus use 2 x 8 bytes while the proposed distinct attribute uses only 8 bytes.

[rriddle]

I'm okay with landing this now, given we should only ever have a single distinct attribute type.

[gysit]

Rebase and improve comments.

mlir/docs/Dialects/Builtin.md

 ## Location Attributes
 
 A subset of the builtin attribute values correspond to
 [source locations](../Diagnostics.md/#source-locations), that may be attached to
 Operations.
 
 [include "Dialects/BuiltinLocationAttributes.md"]
 
+## DistinctAttribute
+
+A DistinctAttribute associates an attribute with a unique identifier.
+As a result, multiple DistinctAttribute instances may point to the same
+attribute. Every call to the `create` function allocates a new
+DistinctAttribute instance. The address of the attribute instance serves as a
+temporary unique identifier. Similar to the names of SSA values, the final
+unique identifiers are generated during pretty printing. This delayed
+numbering ensures the printed identifiers are deterministic even if
+multiple DistinctAttribute instances are created in-parallel.
+
+Syntax:
+
+```
+distinct-id ::= integer-literal
+distinct-attribute ::= `distinct` `[` distinct-id `]<` attribute `>`
+```
+
+Examples:
+
+```mlir
+#distinct = distinct[0]<42.0 : f32>
+#distinct1 = distinct[1]<42.0 : f32>
+#distinct2 = distinct[2]<array<i32: 10, 42>>
+```
+
+This mechanism is meant to generate attributes with a unique
+identifier, which can be used to mark groups of operations that share a
+common property. For example, groups of aliasing memory operations may be
+marked using one DistinctAttribute instance per alias group.
+
 ## Operations
 
 [include "Dialects/BuiltinOps.md"]
 
 ## Types
 
 [include "Dialects/BuiltinTypes.md"]
 
 ## Type Interfaces
 
 [include "Dialects/BuiltinTypeInterfaces.md"]

mlir/include/mlir/IR/AttributeSupport.h

 };
 
 //===----------------------------------------------------------------------===//
 // AttributeStorage
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 class AttributeUniquer;
+class DistinctAttributeUniquer;
 } // namespace detail
 
 /// Base storage class appearing in an attribute. Derived storage classes should
 /// only be constructed within the context of the AttributeUniquer.
 class alignas(8) AttributeStorage : public StorageUniquer::BaseStorage {
   friend detail::AttributeUniquer;
+  friend detail::DistinctAttributeUniquer;
   friend StorageUniquer;
 
 public:
   /// Return the abstract descriptor for this attribute.
   const AbstractAttribute &getAbstractAttribute() const {
     assert(abstractAttribute && "Malformed attribute storage object.");
     return *abstractAttribute;
   }

mlir/include/mlir/IR/BuiltinAttributes.h

std::function<T(ptrdiff_t)> mapFn =
             return *std::next(valueIt, i);
         // Otherwise, return the zero value.
         return zeroValue;
       };
   return iterator<T>(llvm::seq<ptrdiff_t>(0, getNumElements()).begin(), mapFn);
 }
 
 //===----------------------------------------------------------------------===//
+// DistinctAttr
+//===----------------------------------------------------------------------===//
+
+namespace detail {
+struct DistinctAttrStorage;
+class DistinctAttributeUniquer;
+} // namespace detail
+
+/// An attribute that associates a referenced attribute with a unique
+/// identifier. Every call to the create function allocates a new distinct
+/// attribute instance. The address of the attribute instance serves as a
+/// temporary identifier. Similar to the names of SSA values, the final
+/// identifiers are generated during pretty printing. This delayed numbering
+/// ensures the printed identifiers are deterministic even if multiple distinct
+/// attribute instances are created in-parallel.
+///
+/// Examples:
+///
+/// #distinct = distinct[0]<42.0 : f32>
+/// #distinct1 = distinct[1]<42.0 : f32>
+/// #distinct2 = distinct[2]<array<i32: 10, 42>>
+///
+/// NOTE: The distinct attribute cannot be defined using ODS since it uses a
+/// custom distinct attribute uniquer that cannot be set from ODS.
+class DistinctAttr
+    : public detail::StorageUserBase<DistinctAttr, Attribute,
+                                     detail::DistinctAttrStorage,
+                                     detail::DistinctAttributeUniquer> {
+public:
+  using Base::Base;
+
+  /// Returns the referenced attribute.
+  Attribute getReferencedAttr() const;
+
+  /// Creates a distinct attribute that associates a referenced attribute with a
+  /// unique identifier.
+  static DistinctAttr create(Attribute referencedAttr);
+};
+
+//===----------------------------------------------------------------------===//
 // StringAttr
 //===----------------------------------------------------------------------===//
 
 /// Define comparisons for StringAttr against nullptr and itself to avoid the
 /// StringRef overloads from being chosen when not desirable.
 inline bool operator==(StringAttr lhs, std::nullptr_t) { return !lhs; }
 inline bool operator!=(StringAttr lhs, std::nullptr_t) {
   return static_cast<bool>(lhs);

mlir/include/mlir/IR/BuiltinDialectBytecode.td

 
 def DenseIntElementsAttr : WithType<"DenseIntElementsAttr", Attribute>;
 def SparseElementsAttr : DialectAttribute<(attr
   ShapedType:$type,
   DenseIntElementsAttr:$indices,
   DenseElementsAttr:$values
 )>;
 
+def DistinctAttr : DialectAttribute<(attr
+  Attribute:$referencedAttr
+)>;
+
 // Types
 // -----
 
 def IntegerType : DialectType<(type
   // Yes not pretty,
   WithParser<"succeeded($_reader.readVarInt($_var))",
     WithBuilder<"$_args",
     WithPrinter<"$_writer.writeVarInt($_name.getWidth() << 2 | $_name.getSignedness())",
let elems = [
     FusedLoc,
     FusedLocWithMetadata,
     NameLoc,
     UnknownLoc,
     DenseResourceElementsAttr,
     DenseArrayAttr,
     DenseIntOrFPElementsAttr,
     DenseStringElementsAttr,
-    SparseElementsAttr
+    SparseElementsAttr,
+    DistinctAttr
   ];
 }
 
 def BuiltinDialectTypes : DialectTypes<"Builtin"> {
   let elems = [
     IntegerType,
     IndexType,
     FunctionType,

mlir/lib/AsmParser/AttributeParser.cpp

 ///                    | `{` (attribute-entry (`,` attribute-entry)*)? `}`
 ///                    | symbol-ref-id (`::` symbol-ref-id)*
 ///                    | `dense` `<` tensor-literal `>` `:`
 ///                      (tensor-type | vector-type)
 ///                    | `sparse` `<` attribute-value `,` attribute-value `>`
 ///                      `:` (tensor-type | vector-type)
 ///                    | `strided` `<` `[` comma-separated-int-or-question `]`
 ///                      (`,` `offset` `:` integer-literal)? `>`
+///                    | distinct-attribute
 ///                    | extended-attribute
 ///
 Attribute Parser::parseAttribute(Type type) {
   switch (getToken().getKind()) {
   // Parse an AffineMap or IntegerSet attribute.
   case Token::kw_affine_map: {
     consumeToken(Token::kw_affine_map);
 
Attribute Parser::parseAttribute(Type type) {
   // Parse a sparse elements attribute.
   case Token::kw_sparse:
     return parseSparseElementsAttr(type);
 
   // Parse a strided layout attribute.
   case Token::kw_strided:
     return parseStridedLayoutAttr();
 
+  // Parse a distinct attribute.
+  case Token::kw_distinct:
+    return parseDistinctAttr(type);
+
   // Parse a string attribute.
   case Token::string: {
     auto val = getToken().getStringValue();
     consumeToken(Token::string);
     // Parse the optional trailing colon type if one wasn't explicitly provided.
     if (!type && consumeIf(Token::colon) && !(type = parseType()))
       return Attribute();
 
Attribute Parser::parseStridedLayoutAttr() {
   if (!offset || failed(parseToken(Token::greater, "expected '>'")))
     return nullptr;
 
   if (failed(StridedLayoutAttr::verify(errorEmitter, *offset, strides)))
     return nullptr;
   return StridedLayoutAttr::get(getContext(), *offset, strides);
   // return getChecked<StridedLayoutAttr>(loc,getContext(), *offset, strides);
 }
+
+/// Parse a distinct attribute.
+///
+///  distinct-attribute ::= `distinct`
+///                         `[` integer-literal `]<` attribute-value `>`
+///
+Attribute Parser::parseDistinctAttr(Type type) {
+  consumeToken(Token::kw_distinct);
+  if (parseToken(Token::l_square, "expected '[' after 'distinct'"))
+    return {};
+
+  // Parse the distinct integer identifier.
+  Token token = getToken();
+  if (parseToken(Token::integer, "expected distinct ID"))
+    return {};
+  std::optional<uint64_t> value = token.getUInt64IntegerValue();
+  if (!value) {
+    emitError("expected an unsigned 64-bit integer");
+    return {};
+  }
+
+  // Parse the referenced attribute.
+  if (parseToken(Token::r_square, "expected ']' to close distinct ID") ||
+      parseToken(Token::less, "expected '<' after distinct ID"))
+    return {};
+  Attribute referencedAttr = parseAttribute(type);
+  if (!referencedAttr) {
+    emitError("expected attribute");
+    return {};
+  }
+
+  // Add the distinct attribute to the parser state, if it has not been parsed
+  // before. Otherwise, check if the parsed reference attribute matches the one
+  // found in the parser state.
+  DenseMap<uint64_t, DistinctAttr> &distinctAttrs =
+      state.symbols.distinctAttributes;
+  auto it = distinctAttrs.find(*value);
+  if (it == distinctAttrs.end()) {
+    DistinctAttr distinctAttr = DistinctAttr::create(referencedAttr);
+    it = distinctAttrs.try_emplace(*value, distinctAttr).first;
+  } else if (it->getSecond().getReferencedAttr() != referencedAttr) {
+    emitError("referenced attribute does not match previous definition: ")
+        << it->getSecond().getReferencedAttr();
+    return {};
+  }
+
+  if (parseToken(Token::greater, "expected '>' to close distinct attribute"))
+    return {};
+
+  return it->getSecond();
+}

mlir/lib/AsmParser/Parser.h

if (Attribute parsedAttr = parseAttribute(type)) {
       return success();
     }
     return failure();
   }
 
   /// Parse an attribute dictionary.
   ParseResult parseAttributeDict(NamedAttrList &attributes);
 
+  /// Parse a distinct attribute.
+  Attribute parseDistinctAttr(Type type);
+
   /// Parse an extended attribute.
   Attribute parseExtendedAttr(Type type);
 
   /// Parse a float attribute.
   Attribute parseFloatAttr(Type type, bool isNegative);
 
   /// Parse a decimal or a hexadecimal literal, which can be either an integer
   /// or a float attribute.

mlir/lib/AsmParser/ParserState.h

struct SymbolState {
   /// A map from type alias identifier to Type.
   llvm::StringMap<Type> typeAliasDefinitions;
 
   /// A map of dialect resource keys to the resolved resource name and handle
   /// to use during parsing.
   DenseMap<const OpAsmDialectInterface *,
            llvm::StringMap<std::pair<std::string, AsmDialectResourceHandle>>>
       dialectResources;
+
+  /// A map from unique integer identifier to DistinctAttr.
+  DenseMap<uint64_t, DistinctAttr> distinctAttributes;
 };
 
 //===----------------------------------------------------------------------===//
 // ParserState
 //===----------------------------------------------------------------------===//
 
 /// This class refers to all of the state maintained globally by the parser,
 /// such as the current lexer position etc.

mlir/lib/AsmParser/TokenKinds.def

 TOK_KEYWORD(affine_set)
 TOK_KEYWORD(array)
 TOK_KEYWORD(attributes)
 TOK_KEYWORD(bf16)
 TOK_KEYWORD(ceildiv)
 TOK_KEYWORD(complex)
 TOK_KEYWORD(dense)
 TOK_KEYWORD(dense_resource)
+TOK_KEYWORD(distinct)
 TOK_KEYWORD(f16)
 TOK_KEYWORD(f32)
 TOK_KEYWORD(f64)
 TOK_KEYWORD(f80)
 TOK_KEYWORD(f8E5M2)
 TOK_KEYWORD(f8E4M3FN)
 TOK_KEYWORD(f8E5M2FNUZ)
 TOK_KEYWORD(f8E4M3FNUZ)

mlir/lib/IR/AsmPrinter.cpp

private:
   void printAndVisitNestedAliasesImpl(Attribute attr, bool elideType) {
     if (!isa<BuiltinDialect>(attr.getDialect())) {
       attr.getDialect().printAttribute(attr, *this);
 
       // Process the builtin attributes.
     } else if (llvm::isa<AffineMapAttr, DenseArrayAttr, FloatAttr, IntegerAttr,
                          IntegerSetAttr, UnitAttr>(attr)) {
       return;
+    } else if (auto distinctAttr = dyn_cast<DistinctAttr>(attr)) {
+      printAttribute(distinctAttr.getReferencedAttr());
     } else if (auto dictAttr = dyn_cast<DictionaryAttr>(attr)) {
       for (const NamedAttribute &nestedAttr : dictAttr.getValue()) {
         printAttribute(nestedAttr.getName());
         printAttribute(nestedAttr.getValue());
       }
     } else if (auto arrayAttr = dyn_cast<ArrayAttr>(attr)) {
       for (Attribute nestedAttr : arrayAttr.getValue())
         printAttribute(nestedAttr);
if (!usedNames.count(name)) {
     }
   }
 
   usedNames.insert(name, char());
   return name;
 }
 
 //===----------------------------------------------------------------------===//
+// DistinctState
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class manages the state for distinct attributes.
+class DistinctState {
+public:
+  /// Returns a unique identifier for the given distinct attribute.
+  uint64_t getId(DistinctAttr distinctAttr);
+
+private:
+  uint64_t distinctCounter = 0;
+  DenseMap<DistinctAttr, uint64_t> distinctAttrMap;
+};
+} // namespace
+
+uint64_t DistinctState::getId(DistinctAttr distinctAttr) {
+  auto [it, inserted] =
+      distinctAttrMap.try_emplace(distinctAttr, distinctCounter);
+  if (inserted)
+    distinctCounter++;
+  return it->getSecond();
+}
+
+//===----------------------------------------------------------------------===//
 // Resources
 //===----------------------------------------------------------------------===//
 
 AsmParsedResourceEntry::~AsmParsedResourceEntry() = default;
 AsmResourceBuilder::~AsmResourceBuilder() = default;
 AsmResourceParser::~AsmResourceParser() = default;
 AsmResourcePrinter::~AsmResourcePrinter() = default;
 
public:
   }
 
   /// Get the state used for aliases.
   AliasState &getAliasState() { return aliasState; }
 
   /// Get the state used for SSA names.
   SSANameState &getSSANameState() { return nameState; }
 
+  /// Get the state used for distinct attribute identifiers.
+  DistinctState &getDistinctState() { return distinctState; }
+
   /// Return the dialects within the context that implement
   /// OpAsmDialectInterface.
   DialectInterfaceCollection<OpAsmDialectInterface> &getDialectInterfaces() {
     return interfaces;
   }
 
   /// Return the non-dialect resource printers.
   auto getResourcePrinters() {
private:
   DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> dialectResources;
 
   /// The state used for attribute and type aliases.
   AliasState aliasState;
 
   /// The state used for SSA value names.
   SSANameState nameState;
 
+  /// The state used for distinct attribute identifiers.
+  DistinctState distinctState;
+
   /// Flags that control op output.
   OpPrintingFlags printerFlags;
 
   /// An optional location map to be populated.
   AsmState::LocationMap *locationMap;
 
   // Allow direct access to the impl fields.
   friend AsmState;
void AsmPrinter::Impl::printAttributeImpl(Attribute attr,
   if (!isa<BuiltinDialect>(attr.getDialect())) {
     printDialectAttribute(attr);
   } else if (auto opaqueAttr = llvm::dyn_cast<OpaqueAttr>(attr)) {
     printDialectSymbol(os, "#", opaqueAttr.getDialectNamespace(),
                        opaqueAttr.getAttrData());
   } else if (llvm::isa<UnitAttr>(attr)) {
     os << "unit";
     return;
+  } else if (auto distinctAttr = llvm::dyn_cast<DistinctAttr>(attr)) {
+    os << "distinct[" << state.getDistinctState().getId(distinctAttr) << "]<";
+    printAttribute(distinctAttr.getReferencedAttr());
+    os << '>';
+    return;
   } else if (auto dictAttr = llvm::dyn_cast<DictionaryAttr>(attr)) {
     os << '{';
     interleaveComma(dictAttr.getValue(),
                     [&](NamedAttribute attr) { printNamedAttribute(attr); });
     os << '}';
 
   } else if (auto intAttr = llvm::dyn_cast<IntegerAttr>(attr)) {
     Type intType = intAttr.getType();

mlir/lib/IR/AttributeDetail.h

 //===- AttributeDetail.h - MLIR Affine Map details Class --------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This holds implementation details of Attribute.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef ATTRIBUTEDETAIL_H_
 #define ATTRIBUTEDETAIL_H_
 
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/AttributeSupport.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Support/StorageUniquer.h"
+#include "mlir/Support/ThreadLocalCache.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/Support/TrailingObjects.h"
 
 namespace mlir {
 namespace detail {
 
 //===----------------------------------------------------------------------===//
struct StringAttrStorage : public AttributeStorage {
   Type type;
   /// The raw string value.
   StringRef value;
   /// If the string value contains a dialect namespace prefix (e.g.
   /// dialect.blah), this is the dialect referenced.
   Dialect *referencedDialect;
 };
 
+//===----------------------------------------------------------------------===//
+// DistinctAttr
+//===----------------------------------------------------------------------===//
+
+/// An attribute to store a distinct reference to another attribute.
+struct DistinctAttrStorage : public AttributeStorage {
+  using KeyTy = Attribute;
+
+  DistinctAttrStorage(Attribute referencedAttr)
+      : referencedAttr(referencedAttr) {}
+
+  /// Returns the referenced attribute as key.
+  KeyTy getAsKey() const { return KeyTy(referencedAttr); }
+
+  /// The referenced attribute.
+  Attribute referencedAttr;
+};
+
+/// A specialized attribute uniquer for distinct attributes that always
+/// allocates since the distinct attribute instances use the address of their
+/// storage as unique identifier.
+class DistinctAttributeUniquer {
+public:
+  /// Creates a distinct attribute storage. Allocates every time since the
+  /// address of the storage serves as unique identifier.
+  template <typename T, typename... Args>
+  static T get(MLIRContext *context, Args &&...args) {
+    static_assert(std::is_same_v<typename T::ImplType, DistinctAttrStorage>,
+                  "expects a distinct attribute storage");
+    DistinctAttrStorage *storage = DistinctAttributeUniquer::allocateStorage(
+        context, std::forward<Args>(args)...);
+    storage->initializeAbstractAttribute(
+        AbstractAttribute::lookup(DistinctAttr::getTypeID(), context));
+    return storage;
+  }
+
+private:
+  /// Allocates a distinct attribute storage.
+  static DistinctAttrStorage *allocateStorage(MLIRContext *context,
+                                              Attribute referencedAttr);
+};
+
+/// An allocator for distinct attribute storage instances. It uses thread local
+/// bump pointer allocators stored in a thread local cache to ensure the storage
+/// is freed after the destruction of the distinct attribute allocator.
+class DistinctAttributeAllocator {
+public:
+  DistinctAttributeAllocator() = default;
+
+  DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete;
+  DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete;
+  DistinctAttributeAllocator &
+  operator=(const DistinctAttributeAllocator &) = delete;
+
+  /// Allocates a distinct attribute storage using a thread local bump pointer
+  /// allocator to enable synchronization free parallel allocations.
+  DistinctAttrStorage *allocate(Attribute referencedAttr) {
+    return new (allocatorCache.get().Allocate<DistinctAttrStorage>())
+        DistinctAttrStorage(referencedAttr);
+  }
+
+private:
+  ThreadLocalCache<llvm::BumpPtrAllocator> allocatorCache;
+};
 } // namespace detail
 } // namespace mlir
 
 #endif // ATTRIBUTEDETAIL_H_

mlir/lib/IR/BuiltinAttributes.cpp

 // BuiltinDialect
 //===----------------------------------------------------------------------===//
 
 void BuiltinDialect::registerAttributes() {
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/IR/BuiltinAttributes.cpp.inc"
       >();
+  addAttributes<DistinctAttr>();
 }
 
 //===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
 /// Helper function that does either an in place sort or sorts from source array
 /// into destination. If inPlace then storage is both the source and the
for (size_t i = 0, e = numSparseIndices; i != e; ++i) {
     if (!ElementsAttr::isValidIndex(type, index))
       return emitIndexError(i, index);
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// DistinctAttr
+//===----------------------------------------------------------------------===//
+
+DistinctAttr DistinctAttr::create(Attribute referencedAttr) {
+  return Base::get(referencedAttr.getContext(), referencedAttr);
+}
+
+Attribute DistinctAttr::getReferencedAttr() const {
+  return getImpl()->referencedAttr;
+}
+
+//===----------------------------------------------------------------------===//
 // Attribute Utilities
 //===----------------------------------------------------------------------===//
 
 AffineMap mlir::makeStridedLinearLayoutMap(ArrayRef<int64_t> strides,
                                            int64_t offset,
                                            MLIRContext *context) {
   AffineExpr expr;
   unsigned nSymbols = 0;

mlir/lib/IR/BuiltinDialect.cpp

AliasResult getAlias(Attribute attr, raw_ostream &os) const override {
     if (llvm::isa<IntegerSetAttr>(attr)) {
       os << "set";
       return AliasResult::OverridableAlias;
     }
     if (llvm::isa<LocationAttr>(attr)) {
       os << "loc";
       return AliasResult::OverridableAlias;
     }
+    if (llvm::isa<DistinctAttr>(attr)) {
+      os << "distinct";
+      return AliasResult::OverridableAlias;
+    }
     return AliasResult::NoAlias;
   }
 
   AliasResult getAlias(Type type, raw_ostream &os) const final {
     if (auto tupleType = llvm::dyn_cast<TupleType>(type)) {
       if (tupleType.size() > 16) {
         os << "tuple";
         return AliasResult::OverridableAlias;

mlir/lib/IR/BuiltinDialectBytecode.cpp

 //===- BuiltinDialectBytecode.cpp - Builtin Bytecode Implementation -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "BuiltinDialectBytecode.h"
+#include "AttributeDetail.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectResourceBlobManager.h"
 #include "llvm/ADT/TypeSwitch.h"
 

mlir/lib/IR/MLIRContext.cpp

#endif
   StringAttr emptyStringAttr;
 
   /// Map of string attributes that may reference a dialect, that are awaiting
   /// that dialect to be loaded.
   llvm::sys::SmartMutex<true> dialectRefStrAttrMutex;
   DenseMap<StringRef, SmallVector<StringAttrStorage *>>
       dialectReferencingStrAttrs;
 
+  /// A distinct attribute allocator that allocates every time since the
+  /// address of the distinct attribute storage serves as unique identifier. The
+  /// allocator is thread safe and frees the allocated storage after its
+  /// destruction.
+  DistinctAttributeAllocator distinctAttributeAllocator;
+
 public:
   MLIRContextImpl(bool threadingIsEnabled)
       : threadingIsEnabled(threadingIsEnabled) {
     if (threadingIsEnabled) {
       ownedThreadPool = std::make_unique<llvm::ThreadPool>();
       threadPool = ownedThreadPool.get();
     }
   }

 UnitAttr UnitAttr::get(MLIRContext *context) {
   return context->getImpl().unitAttr;
 }
 
 UnknownLoc UnknownLoc::get(MLIRContext *context) {
   return context->getImpl().unknownLocAttr;
 }
 
+DistinctAttrStorage *
+detail::DistinctAttributeUniquer::allocateStorage(MLIRContext *context,
+                                                  Attribute referencedAttr) {
+  return context->getImpl().distinctAttributeAllocator.allocate(referencedAttr);
+}
+
 /// Return empty dictionary.
 DictionaryAttr DictionaryAttr::getEmpty(MLIRContext *context) {
   return context->getImpl().emptyDictionaryAttr;
 }
 
 void StringAttrStorage::initialize(MLIRContext *context) {
   // Check for a dialect namespace prefix, if there isn't one we don't need to
   // do any additional initialization.

mlir/test/Dialect/Builtin/Bytecode/attrs.mlir

 
 // CHECK-LABEL: @TestType
 module @TestType attributes {
   // CHECK: bytecode.type = i178
   bytecode.type = i178
 } {}
 
 //===----------------------------------------------------------------------===//
+// DistinctAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestDistinct
+module @TestDistinct attributes {
+  // CHECK: bytecode.distinct = distinct[0]<42 : i32>
+  // CHECK: bytecode.distinct2 = distinct[0]<42 : i32>
+  // CHECK: bytecode.distinct3 = distinct[1]<42 : i32>
+  bytecode.distinct = distinct[0]<42 : i32>,
+  bytecode.distinct2 = distinct[0]<42 : i32>,
+  bytecode.distinct3 = distinct[1]<42 : i32>
+} {}
+
+//===----------------------------------------------------------------------===//
 // CallSiteLoc
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @TestLocCallSite
 module @TestLocCallSite attributes {
   // CHECK: bytecode.loc = loc(callsite("foo" at "mysource.cc":10:8))
   bytecode.loc = loc(callsite("foo" at "mysource.cc":10:8))
 } {}

mlir/test/IR/distinct-attr.mlir

+// RUN: mlir-opt %s | FileCheck %s
+// RUN: mlir-opt -mlir-print-local-scope %s | FileCheck %s --check-prefix=CHECK-GENERIC
+
+// CHECK: #[[DISTINCT0:.*]] = distinct[0]<42 : i32>
+// CHECK: #[[DISTINCT1:.*]] = distinct[1]<array<i32: 10, 42>>
+// CHECK: #[[DISTINCT2:.*]] = distinct[2]<42 : i32>
+
+// CHECK:         distinct.attr = #[[DISTINCT0]]
+// CHECK-GENERIC: distinct.attr = distinct[0]<42 : i32>
+"test.op"() {distinct.attr = distinct[0]<42 : i32>} : () -> ()
+
+// CHECK:         distinct.attr = #[[DISTINCT1]]
+// CHECK-GENERIC: distinct.attr = distinct[1]<array<i32: 10, 42>>
+"test.op"() {distinct.attr = distinct[1]<array<i32: 10, 42>>} : () -> ()
+
+// CHECK:         distinct.attr = #[[DISTINCT0]]
+// CHECK-GENERIC: distinct.attr = distinct[0]<42 : i32>
+"test.op"() {distinct.attr = distinct[0]<42 : i32>} : () -> ()
+
+// CHECK:         distinct.attr = #[[DISTINCT2]]
+// CHECK-GENERIC: distinct.attr = distinct[2]<42 : i32>
+"test.op"() {distinct.attr = distinct[42]<42 : i32>} : () -> ()

mlir/test/IR/invalid-builtin-attributes.mlir

 
 // expected-error@below {{expected ':' after dense array type}}
 #attr = array<i8)
 
 // -----
 
 // expected-error@below {{expected '>' to close an array attribute}}
 #attr = array<i8: 1)
+
+// -----
+
+// expected-error@below {{expected '[' after 'distinct'}}
+#attr = distinct<
+
+// -----
+
+// expected-error@below {{expected distinct ID}}
+#attr = distinct[i8
+
+// -----
+
+// expected-error@below {{expected an unsigned 64-bit integer}}
+#attr = distinct[0xAAAABBBBEEEEFFFF1]
+
+// -----
+
+// expected-error@below {{expected ']' to close distinct ID}}
+#attr = distinct[8)
+
+// -----
+
+// expected-error@below {{expected '<' after distinct ID}}
+#attr = distinct[8](
+
+// -----
+
+// expected-error@below {{expected attribute}}
+#attr = distinct[8]<attribute
+
+// -----
+
+// expected-error@below {{expected '>' to close distinct attribute}}
+#attr = distinct[8]<@foo]
+
+// -----
+
+#attr = distinct[0]<42 : i32>
+// expected-error@below {{referenced attribute does not match previous definition: 42 : i32}}
+#attr1 = distinct[0]<43 : i32>

mlir/test/IR/test-builtin-distinct-attrs.mlir

+// RUN: mlir-opt %s -test-distinct-attrs | FileCheck %s
+
+// CHECK: #[[DIST0:.*]] = distinct[0]<42 : i32>
+// CHECK: #[[DIST1:.*]] = distinct[1]<42 : i32>
+#distinct = distinct[0]<42 : i32>
+// CHECK: #[[DIST2:.*]] = distinct[2]<42 : i32>
+// CHECK: #[[DIST3:.*]] = distinct[3]<42 : i32>
+#distinct1 = distinct[1]<42 : i32>
+// CHECK: #[[DIST4:.*]] = distinct[4]<43 : i32>
+// CHECK: #[[DIST5:.*]] = distinct[5]<43 : i32>
+#distinct2 = distinct[2]<43 : i32>
+// CHECK: #[[DIST6:.*]] = distinct[6]<@foo_1>
+// CHECK: #[[DIST7:.*]] = distinct[7]<@foo_1>
+#distinct3 = distinct[3]<@foo_1>
+
+// Copies made for foo_2
+// CHECK: #[[DIST8:.*]] = distinct[8]<42 : i32>
+// CHECK: #[[DIST9:.*]] = distinct[9]<42 : i32>
+// CHECK: #[[DIST10:.*]] = distinct[10]<43 : i32>
+// CHECK: #[[DIST11:.*]] = distinct[11]<@foo_1>
+
+// Copies made for foo_3
+// CHECK: #[[DIST12:.*]] = distinct[12]<42 : i32>
+// CHECK: #[[DIST13:.*]] = distinct[13]<42 : i32>
+// CHECK: #[[DIST14:.*]] = distinct[14]<43 : i32>
+// CHECK: #[[DIST15:.*]] = distinct[15]<@foo_1>
+
+// Copies made for foo_4
+// CHECK: #[[DIST16:.*]] = distinct[16]<42 : i32>
+// CHECK: #[[DIST17:.*]] = distinct[17]<42 : i32>
+// CHECK: #[[DIST18:.*]] = distinct[18]<43 : i32>
+// CHECK: #[[DIST19:.*]] = distinct[19]<@foo_1>
+
+// CHECK: @foo_1
+func.func @foo_1() {
+  // CHECK: "test.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST1]]}
+  "test.op"() {distinct.input = #distinct} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST2]], distinct.output = #[[DIST3]]}
+  "test.op"() {distinct.input = #distinct1} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST4]], distinct.output = #[[DIST5]]}
+  "test.op"() {distinct.input = #distinct2} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST6]], distinct.output = #[[DIST7]]}
+  "test.op"() {distinct.input = #distinct3} : () -> ()
+}
+
+func.func @foo_2() {
+  // CHECK: "test.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST8]]}
+  "test.op"() {distinct.input = #distinct} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST2]], distinct.output = #[[DIST9]]}
+  "test.op"() {distinct.input = #distinct1} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST4]], distinct.output = #[[DIST10]]}
+  "test.op"() {distinct.input = #distinct2} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST6]], distinct.output = #[[DIST11]]}
+  "test.op"() {distinct.input = #distinct3} : () -> ()
+}
+
+func.func @foo_3() {
+  // CHECK: "test.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST12]]}
+  "test.op"() {distinct.input = #distinct} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST2]], distinct.output = #[[DIST13]]}
+  "test.op"() {distinct.input = #distinct1} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST4]], distinct.output = #[[DIST14]]}
+  "test.op"() {distinct.input = #distinct2} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST6]], distinct.output = #[[DIST15]]}
+  "test.op"() {distinct.input = #distinct3} : () -> ()
+}
+
+func.func @foo_4() {
+  // CHECK: "test.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST16]]}
+  "test.op"() {distinct.input = #distinct} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST2]], distinct.output = #[[DIST17]]}
+  "test.op"() {distinct.input = #distinct1} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST4]], distinct.output = #[[DIST18]]}
+  "test.op"() {distinct.input = #distinct2} : () -> ()
+  // CHECK: "test.op"() {distinct.input = #[[DIST6]], distinct.output = #[[DIST19]]}
+  "test.op"() {distinct.input = #distinct3} : () -> ()
+}

mlir/test/IR/test-symbol-rauw.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -test-symbol-rauw -split-input-file | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -mlir-print-local-scope %s -test-symbol-rauw -split-input-file | FileCheck %s
 
 // Symbol references to the module itself don't affect uses of symbols within
 // its table.
 // CHECK: module
 // CHECK-SAME: @symbol_foo
 module attributes {sym.outside_use = @symbol_foo } {
   // CHECK: func private @replaced_foo
   func.func private @symbol_foo() attributes {sym.new_name = "replaced_foo" }

 module {
     // CHECK: func private @replaced_foo
     func.func private @symbol_foo() attributes {sym.new_name = "replaced_foo" }
 
     // CHECK: func @symbol_bar
     func.func @symbol_bar() {
       // CHECK: foo.op
       // CHECK-SAME: non_symbol_attr,
-      // CHECK-SAME: use = [#test.sub_elements_access<[@replaced_foo], @symbol_bar, @replaced_foo>],
+      // CHECK-SAME: use = [#test.sub_elements_access<[@replaced_foo], @symbol_bar, @replaced_foo>, distinct[0]<@replaced_foo>],
       // CHECK-SAME: z_non_symbol_attr_3
       "foo.op"() {
         non_symbol_attr,
-        use = [#test.sub_elements_access<[@symbol_foo],@symbol_bar,@symbol_foo>],
+        use = [#test.sub_elements_access<[@symbol_foo],@symbol_bar,@symbol_foo>, distinct[0]<@symbol_foo>],
         z_non_symbol_attr_3
       } : () -> ()
     }
 }
 
 // -----
 
 module {

mlir/test/IR/test-symbol-uses.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -test-symbol-uses -split-input-file -verify-diagnostics
 
 // Symbol references to the module itself don't affect uses of symbols within
 // its table.
 // expected-remark@below {{symbol_removable function successfully erased}}
 module attributes {sym.outside_use = @symbol_foo } {
-  // expected-remark@+1 {{symbol has 2 uses}}
+  // expected-remark@+1 {{symbol has 3 uses}}
   func.func private @symbol_foo()
 
   // expected-remark@below {{symbol has no uses}}
   // expected-remark@below {{found use of symbol : @symbol_foo}}
-  // expected-remark@below {{symbol contains 2 nested references}}
+  // expected-remark@below {{symbol contains 3 nested references}}
   func.func @symbol_bar() attributes {sym.use = @symbol_foo} {
     // expected-remark@+1 {{found use of symbol : @symbol_foo}}
     "foo.op"() {
       non_symbol_attr,
-      use = [{ nested_symbol = [@symbol_foo]}],
+      use = [{nested_symbol = [@symbol_foo]}],
       z_other_non_symbol_attr
     } : () -> ()
+    // expected-remark@+1 {{found use of symbol : @symbol_foo}}
+    "foo.op"() { use = distinct[0]<@symbol_foo> } : () -> ()
   }
 
   // expected-remark@below {{symbol has no uses}}
   func.func private @symbol_removable()
 
   // expected-remark@+1 {{symbol has 1 use}}
   func.func private @symbol_baz()
 

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
   TestBuiltinAttributeInterfaces.cpp
+  TestBuiltinDistinctAttributes.cpp
   TestClone.cpp
   TestDiagnostics.cpp
   TestDominance.cpp
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp
   TestLazyLoading.cpp
   TestOpaqueLoc.cpp

mlir/test/lib/IR/TestBuiltinDistinctAttributes.cpp

+//===- TestBuiltinDistinctAttributes.cpp - Test DistinctAttributes --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestDialect.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// This is a distinct attribute test pass that tests if distinct attributes can
+/// be created in parallel in a deterministic way.
+struct DistinctAttributesPass
+    : public PassWrapper<DistinctAttributesPass, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DistinctAttributesPass)
+
+  StringRef getArgument() const final { return "test-distinct-attrs"; }
+  StringRef getDescription() const final {
+    return "Test parallel creation of distinct attributes";
+  }
+
+  void runOnOperation() override {
+    auto funcOp = getOperation();
+
+    /// Walk all operations and create a distinct output attribute given a
+    /// distinct input attribute.
+    funcOp->walk([](Operation *op) {
+      auto distinctAttr = op->getAttrOfType<DistinctAttr>("distinct.input");
+      if (!distinctAttr)
+        return;
+      op->setAttr("distinct.output",
+                  DistinctAttr::create(distinctAttr.getReferencedAttr()));
+    });
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestBuiltinDistinctAttributes() {
+  PassRegistration<DistinctAttributesPass>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerInliner();
 void registerMemRefBoundCheck();
 void registerPatternsTestPass();
 void registerSimpleParametricTilingPass();
 void registerTestAffineLoopParametricTilingPass();
 void registerTestArithEmulateWideIntPass();
 void registerTestAliasAnalysisPass();
 void registerTestBuiltinAttributeInterfaces();
+void registerTestBuiltinDistinctAttributes();
 void registerTestCallGraphPass();
 void registerTestCfAssertPass();
 void registerTestConstantFold();
 void registerTestControlFlowSink();
 void registerTestGpuSerializeToCubinPass();
 void registerTestGpuSerializeToHsacoPass();
 void registerTestDataLayoutPropagation();
 void registerTestDataLayoutQuery();
void registerTestPasses() {
   mlir::test::registerInliner();
   mlir::test::registerMemRefBoundCheck();
   mlir::test::registerPatternsTestPass();
   mlir::test::registerSimpleParametricTilingPass();
   mlir::test::registerTestAffineLoopParametricTilingPass();
   mlir::test::registerTestAliasAnalysisPass();
   mlir::test::registerTestArithEmulateWideIntPass();
   mlir::test::registerTestBuiltinAttributeInterfaces();
+  mlir::test::registerTestBuiltinDistinctAttributes();
   mlir::test::registerTestCallGraphPass();
   mlir::test::registerTestCfAssertPass();
   mlir::test::registerTestConstantFold();
   mlir::test::registerTestControlFlowSink();
   mlir::test::registerTestDiagnosticsPass();
   mlir::test::registerTestDialectConversionPasses();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToCubinPass();