[mlir] Add a builtin distinct attribute

A distinct attribute associates a referenced attribute with 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.

Reviewed By: rriddle, Dinistro, zero9178

Differential Revision: https://reviews.llvm.org/D153360

23 files changed, 490 insertions, 7 deletions

diff --git a/mlir/docs/Dialects/Builtin.md b/mlir/docs/Dialects/Builtin.md
index b39506a..0a9b7ae 100644
--- a/mlir/docs/Dialects/Builtin.md
+++ b/mlir/docs/Dialects/Builtin.md
@@ -23,6 +23,37 @@ 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"]
diff --git a/mlir/include/mlir/IR/AttributeSupport.h b/mlir/include/mlir/IR/AttributeSupport.h
index 796adbe..75ea1ce 100644
--- a/mlir/include/mlir/IR/AttributeSupport.h
+++ b/mlir/include/mlir/IR/AttributeSupport.h
@@ -149,12 +149,14 @@ private:
 
 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:
diff --git a/mlir/include/mlir/IR/BuiltinAttributes.h b/mlir/include/mlir/IR/BuiltinAttributes.h
index 7c41360..d8b102f 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.h
+++ b/mlir/include/mlir/IR/BuiltinAttributes.h
@@ -1001,6 +1001,46 @@ auto SparseElementsAttr::try_value_begin_impl(OverloadToken<T>) const
 }
 
 //===----------------------------------------------------------------------===//
+// 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
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
index fcbb5f4..f50b5dd 100644
--- a/mlir/include/mlir/IR/BuiltinDialectBytecode.td
+++ b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
@@ -181,6 +181,10 @@ def SparseElementsAttr : DialectAttribute<(attr
   DenseElementsAttr:$values
 )>;
 
+def DistinctAttr : DialectAttribute<(attr
+  Attribute:$referencedAttr
+)>;
+
 // Types
 // -----
 
@@ -316,7 +320,8 @@ def BuiltinDialectAttributes : DialectAttributes<"Builtin"> {
     DenseArrayAttr,
     DenseIntOrFPElementsAttr,
     DenseStringElementsAttr,
-    SparseElementsAttr
+    SparseElementsAttr,
+    DistinctAttr
   ];
 }
 
diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index 9501794..20c89d3 100644
--- a/mlir/lib/AsmParser/AttributeParser.cpp
+++ b/mlir/lib/AsmParser/AttributeParser.cpp
@@ -46,6 +46,7 @@ using namespace mlir::detail;
 ///                      `:` (tensor-type | vector-type)
 ///                    | `strided` `<` `[` comma-separated-int-or-question `]`
 ///                      (`,` `offset` `:` integer-literal)? `>`
+///                    | distinct-attribute
 ///                    | extended-attribute
 ///
 Attribute Parser::parseAttribute(Type type) {
@@ -155,6 +156,10 @@ Attribute Parser::parseAttribute(Type type) {
   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();
@@ -1214,3 +1219,54 @@ Attribute Parser::parseStridedLayoutAttr() {
   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();
+}
diff --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h
index 522cf00..01c55f9 100644
--- a/mlir/lib/AsmParser/Parser.h
+++ b/mlir/lib/AsmParser/Parser.h
@@ -250,6 +250,9 @@ public:
   /// Parse an attribute dictionary.
   ParseResult parseAttributeDict(NamedAttrList &attributes);
 
+  /// Parse a distinct attribute.
+  Attribute parseDistinctAttr(Type type);
+
   /// Parse an extended attribute.
   Attribute parseExtendedAttr(Type type);
 
diff --git a/mlir/lib/AsmParser/ParserState.h b/mlir/lib/AsmParser/ParserState.h
index 62e932f..0bf9bff 100644
--- a/mlir/lib/AsmParser/ParserState.h
+++ b/mlir/lib/AsmParser/ParserState.h
@@ -36,6 +36,9 @@ struct SymbolState {
   DenseMap<const OpAsmDialectInterface *,
            llvm::StringMap<std::pair<std::string, AsmDialectResourceHandle>>>
       dialectResources;
+
+  /// A map from unique integer identifier to DistinctAttr.
+  DenseMap<uint64_t, DistinctAttr> distinctAttributes;
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/AsmParser/TokenKinds.def b/mlir/lib/AsmParser/TokenKinds.def
index 1b5aa10..297e074 100644
--- a/mlir/lib/AsmParser/TokenKinds.def
+++ b/mlir/lib/AsmParser/TokenKinds.def
@@ -89,6 +89,7 @@ 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)
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 16ca7f8..603e414 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -806,6 +806,8 @@ private:
     } 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());
@@ -1605,6 +1607,31 @@ StringRef SSANameState::uniqueValueName(StringRef 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
 //===----------------------------------------------------------------------===//
 
@@ -1715,6 +1742,9 @@ public:
   /// 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() {
@@ -1758,6 +1788,9 @@ private:
   /// 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;
 
@@ -2106,6 +2139,11 @@ void AsmPrinter::Impl::printAttributeImpl(Attribute attr,
   } 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(),
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 5bf7caa..b76da1b 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -14,11 +14,13 @@
 #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"
@@ -349,6 +351,70 @@ struct StringAttrStorage : public AttributeStorage {
   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
 
diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index 9a18f07..c670359 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -42,6 +42,7 @@ void BuiltinDialect::registerAttributes() {
 #define GET_ATTRDEF_LIST
 #include "mlir/IR/BuiltinAttributes.cpp.inc"
       >();
+  addAttributes<DistinctAttr>();
 }
 
 //===----------------------------------------------------------------------===//
@@ -1746,6 +1747,18 @@ SparseElementsAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 }
 
 //===----------------------------------------------------------------------===//
+// DistinctAttr
+//===----------------------------------------------------------------------===//
+
+DistinctAttr DistinctAttr::create(Attribute referencedAttr) {
+  return Base::get(referencedAttr.getContext(), referencedAttr);
+}
+
+Attribute DistinctAttr::getReferencedAttr() const {
+  return getImpl()->referencedAttr;
+}
+
+//===----------------------------------------------------------------------===//
 // Attribute Utilities
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/lib/IR/BuiltinDialect.cpp b/mlir/lib/IR/BuiltinDialect.cpp
index a791b97..a0f6af3 100644
--- a/mlir/lib/IR/BuiltinDialect.cpp
+++ b/mlir/lib/IR/BuiltinDialect.cpp
@@ -60,6 +60,10 @@ struct BuiltinOpAsmDialectInterface : public OpAsmDialectInterface {
       os << "loc";
       return AliasResult::OverridableAlias;
     }
+    if (llvm::isa<DistinctAttr>(attr)) {
+      os << "distinct";
+      return AliasResult::OverridableAlias;
+    }
     return AliasResult::NoAlias;
   }
 
diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp
index b7a2784..9487e6c 100644
--- a/mlir/lib/IR/BuiltinDialectBytecode.cpp
+++ b/mlir/lib/IR/BuiltinDialectBytecode.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "BuiltinDialectBytecode.h"
+#include "AttributeDetail.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinDialect.h"
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index a79c7fb..bb309d5 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -248,6 +248,12 @@ public:
   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) {
@@ -1064,6 +1070,12 @@ 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;
diff --git a/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir b/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir
index 208d331..0f5643a 100644
--- a/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir
+++ b/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir
@@ -126,6 +126,20 @@ module @TestType attributes {
 } {}
 
 //===----------------------------------------------------------------------===//
+// 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
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/test/IR/distinct-attr.mlir b/mlir/test/IR/distinct-attr.mlir
new file mode 100644
index 0000000..a9ec937
--- /dev/null
+++ b/mlir/test/IR/distinct-attr.mlir
@@ -0,0 +1,22 @@
+// 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>} : () -> ()
diff --git a/mlir/test/IR/invalid-builtin-attributes.mlir b/mlir/test/IR/invalid-builtin-attributes.mlir
index 5343f97..1ff4460 100644
--- a/mlir/test/IR/invalid-builtin-attributes.mlir
+++ b/mlir/test/IR/invalid-builtin-attributes.mlir
@@ -546,3 +546,44 @@ func.func @duplicate_dictionary_attr_key() {
 
 // 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>
diff --git a/mlir/test/IR/test-builtin-distinct-attrs.mlir b/mlir/test/IR/test-builtin-distinct-attrs.mlir
new file mode 100644
index 0000000..000d3e5
--- /dev/null
+++ b/mlir/test/IR/test-builtin-distinct-attrs.mlir
@@ -0,0 +1,77 @@
+// 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} : () -> ()
+}
diff --git a/mlir/test/IR/test-symbol-rauw.mlir b/mlir/test/IR/test-symbol-rauw.mlir
index fe1c80e..ba17cf9 100644
--- a/mlir/test/IR/test-symbol-rauw.mlir
+++ b/mlir/test/IR/test-symbol-rauw.mlir
@@ -1,4 +1,4 @@
-// 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.
@@ -85,11 +85,11 @@ module {
     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
       } : () -> ()
     }
diff --git a/mlir/test/IR/test-symbol-uses.mlir b/mlir/test/IR/test-symbol-uses.mlir
index 75be627..54e3ef1 100644
--- a/mlir/test/IR/test-symbol-uses.mlir
+++ b/mlir/test/IR/test-symbol-uses.mlir
@@ -4,19 +4,21 @@
 // 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}}
diff --git a/mlir/test/lib/IR/CMakeLists.txt b/mlir/test/lib/IR/CMakeLists.txt
index 004e728..447a248 100644
--- a/mlir/test/lib/IR/CMakeLists.txt
+++ b/mlir/test/lib/IR/CMakeLists.txt
@@ -1,6 +1,7 @@
 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
   TestBuiltinAttributeInterfaces.cpp
+  TestBuiltinDistinctAttributes.cpp
   TestClone.cpp
   TestDiagnostics.cpp
   TestDominance.cpp
diff --git a/mlir/test/lib/IR/TestBuiltinDistinctAttributes.cpp b/mlir/test/lib/IR/TestBuiltinDistinctAttributes.cpp
new file mode 100644
index 0000000..4717ce3
--- /dev/null
+++ b/mlir/test/lib/IR/TestBuiltinDistinctAttributes.cpp
@@ -0,0 +1,49 @@
+//===- 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
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 5b95663..eb2d0b5 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -73,6 +73,7 @@ void registerTestAffineLoopParametricTilingPass();
 void registerTestArithEmulateWideIntPass();
 void registerTestAliasAnalysisPass();
 void registerTestBuiltinAttributeInterfaces();
+void registerTestBuiltinDistinctAttributes();
 void registerTestCallGraphPass();
 void registerTestCfAssertPass();
 void registerTestConstantFold();
@@ -188,6 +189,7 @@ void registerTestPasses() {
   mlir::test::registerTestAliasAnalysisPass();
   mlir::test::registerTestArithEmulateWideIntPass();
   mlir::test::registerTestBuiltinAttributeInterfaces();
+  mlir::test::registerTestBuiltinDistinctAttributes();
   mlir::test::registerTestCallGraphPass();
   mlir::test::registerTestCfAssertPass();
   mlir::test::registerTestConstantFold();