diff --git a/mlir/docs/Dialects/Builtin.md b/mlir/docs/Dialects/Builtin.md --- a/mlir/docs/Dialects/Builtin.md +++ b/mlir/docs/Dialects/Builtin.md @@ -23,6 +23,37 @@ [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]> +``` + +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 --- a/mlir/include/mlir/IR/AttributeSupport.h +++ b/mlir/include/mlir/IR/AttributeSupport.h @@ -149,12 +149,14 @@ 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 --- a/mlir/include/mlir/IR/BuiltinAttributes.h +++ b/mlir/include/mlir/IR/BuiltinAttributes.h @@ -1000,6 +1000,43 @@ return iterator(llvm::seq(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]> +class DistinctAttr + : public detail::StorageUserBase { +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 //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td --- a/mlir/include/mlir/IR/BuiltinDialectBytecode.td +++ b/mlir/include/mlir/IR/BuiltinDialectBytecode.td @@ -181,6 +181,10 @@ DenseElementsAttr:$values )>; +def DistinctAttr : DialectAttribute<(attr + Attribute:$referencedAttr +)>; + // Types // ----- @@ -316,7 +320,8 @@ DenseArrayAttr, DenseIntOrFPElementsAttr, DenseStringElementsAttr, - SparseElementsAttr + SparseElementsAttr, + DistinctAttr ]; } diff --git a/mlir/include/mlir/Support/ThreadLocalStore.h b/mlir/include/mlir/Support/ThreadLocalStore.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Support/ThreadLocalStore.h @@ -0,0 +1,52 @@ +//===- ThreadLocalStore.h - ThreadLocalStore 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 file defines the ThreadLocalStore class. This class provides support for +// allocating thread local objects that have a static lifetime. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_SUPPORT_THREADLOCALSTORE_H +#define MLIR_SUPPORT_THREADLOCALSTORE_H + +#include "mlir/Support/LLVM.h" +#include "llvm/Support/Allocator.h" +#include + +namespace mlir { +/// This class allocates and stores thread local values of the given value type +/// with static lifetime. A thread local bump pointer allocator ensures parallel +/// allocations without any need for further synchronization. +template +class ThreadLocalStore { +public: + ThreadLocalStore() = default; + + ThreadLocalStore(ThreadLocalStore &&) = delete; + ThreadLocalStore(const ThreadLocalStore &) = delete; + ThreadLocalStore &operator=(const ThreadLocalStore &) = delete; + + /// Allocates a value type instance for the current thread. + template + ValueT *create(Args &&...args) { + static_assert(std::is_trivially_destructible_v, + "expected value type to be trivially destructible"); + return new (getStaticAllocator().Allocate()) + ValueT(std::forward(args)...); + } + +private: + /// Returns the static thread local bump pointer allocator. + static llvm::BumpPtrAllocator &getStaticAllocator() { + static LLVM_THREAD_LOCAL llvm::BumpPtrAllocator allocator; + return allocator; + } +}; +} // namespace mlir + +#endif // MLIR_SUPPORT_THREADLOCALSTORE_H diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp --- a/mlir/lib/AsmParser/AttributeParser.cpp +++ b/mlir/lib/AsmParser/AttributeParser.cpp @@ -46,6 +46,7 @@ /// `:` (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 @@ 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,53 @@ return StridedLayoutAttr::get(getContext(), *offset, strides); // return getChecked(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 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 &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"); + 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 --- a/mlir/lib/AsmParser/Parser.h +++ b/mlir/lib/AsmParser/Parser.h @@ -250,6 +250,9 @@ /// 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 --- a/mlir/lib/AsmParser/ParserState.h +++ b/mlir/lib/AsmParser/ParserState.h @@ -36,6 +36,9 @@ DenseMap>> dialectResources; + + /// A map from unique integer identifier to DistinctAttr. + DenseMap distinctAttributes; }; //===----------------------------------------------------------------------===// diff --git a/mlir/lib/AsmParser/TokenKinds.def b/mlir/lib/AsmParser/TokenKinds.def --- a/mlir/lib/AsmParser/TokenKinds.def +++ b/mlir/lib/AsmParser/TokenKinds.def @@ -89,6 +89,7 @@ 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 --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -806,6 +806,8 @@ } else if (llvm::isa(attr)) { return; + } else if (auto distinctAttr = dyn_cast(attr)) { + printAttribute(distinctAttr.getReferencedAttr()); } else if (auto dictAttr = dyn_cast(attr)) { for (const NamedAttribute &nestedAttr : dictAttr.getValue()) { printAttribute(nestedAttr.getName()); @@ -1604,6 +1606,31 @@ 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 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 @@ /// 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 &getDialectInterfaces() { @@ -1758,6 +1788,9 @@ /// 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 @@ } else if (llvm::isa(attr)) { os << "unit"; return; + } else if (auto distinctAttr = llvm::dyn_cast(attr)) { + os << "distinct[" << state.getDistinctState().getId(distinctAttr) << "]<"; + printAttribute(distinctAttr.getReferencedAttr()); + os << '>'; + return; } else if (auto dictAttr = llvm::dyn_cast(attr)) { os << '{'; interleaveComma(dictAttr.getValue(), diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h --- a/mlir/lib/IR/AttributeDetail.h +++ b/mlir/lib/IR/AttributeDetail.h @@ -14,6 +14,7 @@ #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" @@ -349,6 +350,48 @@ 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 + static T get(MLIRContext *context, Args &&...args) { + static_assert(std::is_same_v, + "expects a distinct attribute storage"); + DistinctAttrStorage *storage = DistinctAttributeUniquer::allocateStorage( + context, std::forward(args)...); + storage->initializeAbstractAttribute( + AbstractAttribute::lookup(DistinctAttr::getTypeID(), context)); + return storage; + } + +private: + /// Allocates a distinct attribute storage. + static DistinctAttrStorage *allocateStorage(MLIRContext *context, + Attribute referencedAttr); +}; + } // namespace detail } // namespace mlir diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp --- a/mlir/lib/IR/BuiltinAttributes.cpp +++ b/mlir/lib/IR/BuiltinAttributes.cpp @@ -42,6 +42,7 @@ #define GET_ATTRDEF_LIST #include "mlir/IR/BuiltinAttributes.cpp.inc" >(); + addAttributes(); } //===----------------------------------------------------------------------===// @@ -1745,6 +1746,18 @@ 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 --- a/mlir/lib/IR/BuiltinDialect.cpp +++ b/mlir/lib/IR/BuiltinDialect.cpp @@ -60,6 +60,10 @@ os << "loc"; return AliasResult::OverridableAlias; } + if (llvm::isa(attr)) { + os << "distinct"; + return AliasResult::OverridableAlias; + } return AliasResult::NoAlias; } diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp --- 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 --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -25,6 +25,7 @@ #include "mlir/IR/Location.h" #include "mlir/IR/OpImplementation.h" #include "mlir/IR/Types.h" +#include "mlir/Support/ThreadLocalStore.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/SmallString.h" @@ -247,6 +248,10 @@ DenseMap> dialectReferencingStrAttrs; + /// A store for distinct attributes that allocates every time since the + /// address of the storage serves as unique identifier. + ThreadLocalStore distinctAttrStore; + public: MLIRContextImpl(bool threadingIsEnabled) : threadingIsEnabled(threadingIsEnabled) { @@ -1059,6 +1064,12 @@ return context->getImpl().unknownLocAttr; } +DistinctAttrStorage * +detail::DistinctAttributeUniquer::allocateStorage(MLIRContext *context, + Attribute referencedAttr) { + return context->getImpl().distinctAttrStore.create(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 --- a/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir +++ b/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir @@ -128,6 +128,20 @@ 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 --- /dev/null +++ b/mlir/test/IR/distinct-attr.mlir @@ -0,0 +1,22 @@ +// RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s +// RUN: mlir-opt -allow-unregistered-dialect -mlir-print-local-scope %s | FileCheck %s --check-prefix=CHECK-GENERIC + +// CHECK: #[[DISTINCT0:.*]] = distinct[0]<42 : i32> +// CHECK: #[[DISTINCT1:.*]] = distinct[1]> +// CHECK: #[[DISTINCT2:.*]] = distinct[2]<42 : i32> + +// CHECK: distinct.attr = #[[DISTINCT0]] +// CHECK-GENERIC: distinct.attr = distinct[0]<42 : i32> +"foo.op"() {distinct.attr = distinct[0]<42 : i32>} : () -> () + +// CHECK: distinct.attr = #[[DISTINCT1]] +// CHECK-GENERIC: distinct.attr = distinct[1]> +"foo.op"() {distinct.attr = distinct[1]>} : () -> () + +// CHECK: distinct.attr = #[[DISTINCT0]] +// CHECK-GENERIC: distinct.attr = distinct[0]<42 : i32> +"foo.op"() {distinct.attr = distinct[0]<42 : i32>} : () -> () + +// CHECK: distinct.attr = #[[DISTINCT2]] +// CHECK-GENERIC: distinct.attr = distinct[2]<42 : i32> +"foo.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 --- a/mlir/test/IR/invalid-builtin-attributes.mlir +++ b/mlir/test/IR/invalid-builtin-attributes.mlir @@ -546,3 +546,44 @@ // expected-error@below {{expected '>' to close an array attribute}} #attr = array' to close distinct attribute}} +#attr = distinct[8]<@foo] + +// ----- + +#attr = distinct[0]<42 : i32> +// expected-error@below {{referenced attribute does not match previous definition}} +#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 --- /dev/null +++ b/mlir/test/IR/test-builtin-distinct-attrs.mlir @@ -0,0 +1,77 @@ +// RUN: mlir-opt -allow-unregistered-dialect %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: "foo.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST1]]} + "foo.op"() {distinct.input = #distinct} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST2]], distinct.output = #[[DIST3]]} + "foo.op"() {distinct.input = #distinct1} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST4]], distinct.output = #[[DIST5]]} + "foo.op"() {distinct.input = #distinct2} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST6]], distinct.output = #[[DIST7]]} + "foo.op"() {distinct.input = #distinct3} : () -> () +} + +func.func @foo_2() { + // CHECK: "foo.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST8]]} + "foo.op"() {distinct.input = #distinct} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST2]], distinct.output = #[[DIST9]]} + "foo.op"() {distinct.input = #distinct1} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST4]], distinct.output = #[[DIST10]]} + "foo.op"() {distinct.input = #distinct2} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST6]], distinct.output = #[[DIST11]]} + "foo.op"() {distinct.input = #distinct3} : () -> () +} + +func.func @foo_3() { + // CHECK: "foo.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST12]]} + "foo.op"() {distinct.input = #distinct} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST2]], distinct.output = #[[DIST13]]} + "foo.op"() {distinct.input = #distinct1} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST4]], distinct.output = #[[DIST14]]} + "foo.op"() {distinct.input = #distinct2} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST6]], distinct.output = #[[DIST15]]} + "foo.op"() {distinct.input = #distinct3} : () -> () +} + +func.func @foo_4() { + // CHECK: "foo.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST16]]} + "foo.op"() {distinct.input = #distinct} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST2]], distinct.output = #[[DIST17]]} + "foo.op"() {distinct.input = #distinct1} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST4]], distinct.output = #[[DIST18]]} + "foo.op"() {distinct.input = #distinct2} : () -> () + // CHECK: "foo.op"() {distinct.input = #[[DIST6]], distinct.output = #[[DIST19]]} + "foo.op"() {distinct.input = #distinct3} : () -> () +} diff --git a/mlir/test/IR/test-symbol-rauw.mlir b/mlir/test/IR/test-symbol-rauw.mlir --- 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 @@ 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 --- 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 --- 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 --- /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> { + 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("distinct.input"); + if (!distinctAttr) + return; + op->setAttr("distinct.output", + DistinctAttr::create(distinctAttr.getReferencedAttr())); + }); + } +}; +} // namespace + +namespace mlir { +namespace test { +void registerTestBuiltinDistinctAttributes() { + PassRegistration(); +} +} // namespace test +} // namespace mlir diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp --- a/mlir/tools/mlir-opt/mlir-opt.cpp +++ b/mlir/tools/mlir-opt/mlir-opt.cpp @@ -73,6 +73,7 @@ void registerTestArithEmulateWideIntPass(); void registerTestAliasAnalysisPass(); void registerTestBuiltinAttributeInterfaces(); +void registerTestBuiltinDistinctAttributes(); void registerTestCallGraphPass(); void registerTestCfAssertPass(); void registerTestConstantFold(); @@ -187,6 +188,7 @@ mlir::test::registerTestAliasAnalysisPass(); mlir::test::registerTestArithEmulateWideIntPass(); mlir::test::registerTestBuiltinAttributeInterfaces(); + mlir::test::registerTestBuiltinDistinctAttributes(); mlir::test::registerTestCallGraphPass(); mlir::test::registerTestCfAssertPass(); mlir::test::registerTestConstantFold();