[Mlir-commits] [mlir] 728a8d5 - [mlir] Add a builtin distinct attribute

Tobias Gysi llvmlistbot at llvm.org
Tue Jul 11 00:34:47 PDT 2023


Author: Tobias Gysi
Date: 2023-07-11T07:33:16Z
New Revision: 728a8d5a81f35f24492fcd31928ef95e597a66af

URL: https://github.com/llvm/llvm-project/commit/728a8d5a81f35f24492fcd31928ef95e597a66af
DIFF: https://github.com/llvm/llvm-project/commit/728a8d5a81f35f24492fcd31928ef95e597a66af.diff

LOG: [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

Added: 
    mlir/test/IR/distinct-attr.mlir
    mlir/test/IR/test-builtin-distinct-attrs.mlir
    mlir/test/lib/IR/TestBuiltinDistinctAttributes.cpp

Modified: 
    mlir/docs/Dialects/Builtin.md
    mlir/include/mlir/IR/AttributeSupport.h
    mlir/include/mlir/IR/BuiltinAttributes.h
    mlir/include/mlir/IR/BuiltinDialectBytecode.td
    mlir/lib/AsmParser/AttributeParser.cpp
    mlir/lib/AsmParser/Parser.h
    mlir/lib/AsmParser/ParserState.h
    mlir/lib/AsmParser/TokenKinds.def
    mlir/lib/IR/AsmPrinter.cpp
    mlir/lib/IR/AttributeDetail.h
    mlir/lib/IR/BuiltinAttributes.cpp
    mlir/lib/IR/BuiltinDialect.cpp
    mlir/lib/IR/BuiltinDialectBytecode.cpp
    mlir/lib/IR/MLIRContext.cpp
    mlir/test/Dialect/Builtin/Bytecode/attrs.mlir
    mlir/test/IR/invalid-builtin-attributes.mlir
    mlir/test/IR/test-symbol-rauw.mlir
    mlir/test/IR/test-symbol-uses.mlir
    mlir/test/lib/IR/CMakeLists.txt
    mlir/tools/mlir-opt/mlir-opt.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/docs/Dialects/Builtin.md b/mlir/docs/Dialects/Builtin.md
index b39506a39b5ab6..0a9b7ae8919b55 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 796adbe8531a74..75ea1ce24753c9 100644
--- a/mlir/include/mlir/IR/AttributeSupport.h
+++ b/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:

diff  --git a/mlir/include/mlir/IR/BuiltinAttributes.h b/mlir/include/mlir/IR/BuiltinAttributes.h
index 7c4136021cb529..d8b102f6393816 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.h
+++ b/mlir/include/mlir/IR/BuiltinAttributes.h
@@ -1000,6 +1000,46 @@ auto SparseElementsAttr::try_value_begin_impl(OverloadToken<T>) const
   return iterator<T>(llvm::seq<ptr
diff _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
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
index fcbb5f45acbc51..f50b5dd7ad8226 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 95017942165f71..20c89d342e13f2 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 522cf004e92b54..01c55f97a08c2c 100644
--- a/mlir/lib/AsmParser/Parser.h
+++ b/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);
 

diff  --git a/mlir/lib/AsmParser/ParserState.h b/mlir/lib/AsmParser/ParserState.h
index 62e932ff45cf1d..0bf9bffaf1b600 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 1b5aa10e4ac1dc..297e0745945301 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 16ca7f8c6fc13e..603e4143aae21d 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/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(),

diff  --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 5bf7caa4402770..b76da1b91bff9d 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 9a18f07efed011..c670359b02ec9c 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>();
 }
 
 //===----------------------------------------------------------------------===//
@@ -1745,6 +1746,18 @@ SparseElementsAttr::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// 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 a791b97c12a697..a0f6af389f19af 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 b7a27842e26814..9487e6c50ded90 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 a79c7fb5052e08..bb309d509dbd6c 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -248,6 +248,12 @@ class MLIRContextImpl {
   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 208d331ca0ad31..0f5643aa3bb43e 100644
--- a/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir
+++ b/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir
@@ -125,6 +125,20 @@ module @TestType attributes {
   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
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/IR/distinct-attr.mlir b/mlir/test/IR/distinct-attr.mlir
new file mode 100644
index 00000000000000..a9ec9372445971
--- /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 5343f971537b94..1ff44605cb7ecd 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 at below {{expected '>' to close an array attribute}}
 #attr = array<i8: 1)
+
+// -----
+
+// expected-error at below {{expected '[' after 'distinct'}}
+#attr = distinct<
+
+// -----
+
+// expected-error at below {{expected distinct ID}}
+#attr = distinct[i8
+
+// -----
+
+// expected-error at below {{expected an unsigned 64-bit integer}}
+#attr = distinct[0xAAAABBBBEEEEFFFF1]
+
+// -----
+
+// expected-error at below {{expected ']' to close distinct ID}}
+#attr = distinct[8)
+
+// -----
+
+// expected-error at below {{expected '<' after distinct ID}}
+#attr = distinct[8](
+
+// -----
+
+// expected-error at below {{expected attribute}}
+#attr = distinct[8]<attribute
+
+// -----
+
+// expected-error at below {{expected '>' to close distinct attribute}}
+#attr = distinct[8]<@foo]
+
+// -----
+
+#attr = distinct[0]<42 : i32>
+// expected-error at 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 00000000000000..000d3e5b203d83
--- /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 fe1c80e8f400ee..ba17cf9d104261 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], at symbol_bar, at symbol_foo>],
+        use = [#test.sub_elements_access<[@symbol_foo], at symbol_bar, at 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 75be62778051f3..54e3ef1812510a 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 at below {{symbol_removable function successfully erased}}
 module attributes {sym.outside_use = @symbol_foo } {
-  // expected-remark at +1 {{symbol has 2 uses}}
+  // expected-remark at +1 {{symbol has 3 uses}}
   func.func private @symbol_foo()
 
   // expected-remark at below {{symbol has no uses}}
   // expected-remark at below {{found use of symbol : @symbol_foo}}
-  // expected-remark at below {{symbol contains 2 nested references}}
+  // expected-remark at below {{symbol contains 3 nested references}}
   func.func @symbol_bar() attributes {sym.use = @symbol_foo} {
     // expected-remark at +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 at +1 {{found use of symbol : @symbol_foo}}
+    "foo.op"() { use = distinct[0]<@symbol_foo> } : () -> ()
   }
 
   // expected-remark at below {{symbol has no uses}}

diff  --git a/mlir/test/lib/IR/CMakeLists.txt b/mlir/test/lib/IR/CMakeLists.txt
index 004e728b4b8773..447a2481e8dbad 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 00000000000000..4717ce345fc1fc
--- /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 5b956635b960a1..eb2d0b539ab61d 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();


        


More information about the Mlir-commits mailing list