[mlir] Add a builtin distinct attribute

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.

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

Author: gysit

mlir/docs/Dialects/Builtin.md

@@ -23,6 +23,37 @@ Operations.
 
 [include "Dialects/BuiltinLocationAttributes.md"]
 
+## DistinctAttribute
+
+A DistinctAttribute associates an attribute to 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"]

mlir/include/mlir/IR/AttributeSupport.h

@@ -149,12 +149,14 @@ class AbstractAttribute {
 
 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:

mlir/include/mlir/IR/BuiltinAttributes.h

@@ -1000,6 +1000,43 @@ auto SparseElementsAttr::try_value_begin_impl(OverloadToken<T>) const
   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 to 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>>
+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 to a
+  /// unique identifier.
+  static DistinctAttr create(Attribute referencedAttr);
+};
+
 //===----------------------------------------------------------------------===//
 // StringAttr
 //===----------------------------------------------------------------------===//

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
   ];
 }
 

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();
+}

mlir/lib/AsmParser/Parser.h

@@ -250,6 +250,9 @@ class Parser {
   /// Parse an attribute dictionary.
   ParseResult parseAttributeDict(NamedAttrList &attributes);
 
+  /// Parse a distinct attribute.
+  Attribute parseDistinctAttr(Type type);
+
   /// Parse an extended attribute.
   Attribute parseExtendedAttr(Type type);
 

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;
 };
 
 //===----------------------------------------------------------------------===//

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)

mlir/lib/IR/AsmPrinter.cpp

@@ -806,6 +806,8 @@ class DummyAliasDialectAsmPrinter : public DialectAsmPrinter {
     } 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());
@@ -1604,6 +1606,31 @@ StringRef SSANameState::uniqueValueName(StringRef name) {
   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
 //===----------------------------------------------------------------------===//
@@ -1715,6 +1742,9 @@ class AsmStateImpl {
   /// 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 @@ class AsmStateImpl {
   /// 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(),

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