Introduce a new DialectIdentifier structure, extending Identifier with a Dialect information

[mehdi_amini]

Summary

This class is looking up a dialect prefix on the identifier on initialization
and keeping a pointer to the Dialect when found.

The NamedAttribute key is now a DialectIdentifier.

[rriddle]

The name DialectIdentifier strikes me as weird on first glance, but I think it's slowly being normalized the more I look at it.

I like the idea of making it more efficient to access a context/dialect from a dialect prefixed identifier. Do you have any performance impact numbers for the change to NamedAttribute?

[jpienaar]

Isn't it still just an Identifier? The dialect was always there, it was just previously encoded in the string and is now separate

[mehdi_amini]

In D95418#2524202, @jpienaar wrote:

Isn't it still just an Identifier? The dialect was always there, it was just previously encoded in the string and is now separate

Are you suggesting to not have two classes and just merge the two classes here?

[rriddle]

In D95418#2524208, @mehdi_amini wrote:

In D95418#2524202, @jpienaar wrote:

Isn't it still just an Identifier? The dialect was always there, it was just previously encoded in the string and is now separate

Are you suggesting to not have two classes and just merge the two classes here?

I think that route is going to end up being better, given that it reduces the number of lookups to 1 in the common case. If not that, then diverging from Identifier and still uniquing. Before suggesting to just merge them, I was trying to see what else Identifier is used for aside from Attribute names, but I couldn't find anything.

[mehdi_amini]

Some Locations may be using Identifier like in FileLineColLocationStorage.
An OperationName would get larger as well I think, which would make Operation larger by one pointer (it'd accelerate Operation::getDialect() for unregistered operations, but I don't think we need to optimize this)

[rriddle]

In D95418#2524214, @mehdi_amini wrote:

Some Locations may be using Identifier like in FileLineColLocationStorage.
An OperationName would get larger as well I think, which would make Operation larger by one pointer (it'd accelerate Operation::getDialect() for unregistered operations, but I don't think we need to optimize this)

If we uniqued DialectIdentifier(the Identifier part + the Dialect */MLIRContext *), then OperationName wouldn't grow AFAICS. It would improve memory(no extra pointer per instance of a specific identifier) and reduce lookup time. I don't think we should replace Identifier directly given the other usages, but it still seems interesting to optimize the important cases that do rely on this type of identifier. Having a proper representation for an identifier that can be dialect prefixed seems useful.

[mehdi_amini]

Oh I see what you mean, we'd replace the current llvm::StringSet<llvm::BumpPtrAllocator &> identifiers; in the MLIRContext with something like llvm::StringMap<Dialect*, llvm::BumpPtrAllocator &> identifiers;?
(probably a bit more complex, but that's the rough direction?)

Something I need to figure out is what to do on dialect loading in the Context, I propose on such event to loop over the identifiers and if any is prefixed with the loaded dialect then set the Dialect pointer, WDYT?

[mehdi_amini]

Update to fold the pointer into the Identifier storage.

[jpienaar]

Thanks!

mlir/lib/IR/MLIRContext.cpp

717

Remove trivial {} ?

[mehdi_amini]

Remove trivial braces

[rriddle]

Nice. Really clean now!

LGTM after adding examples.

mlir/include/mlir/IR/Identifier.h

58

Can you beef up the main description of the Identifier class? Right now there is no mention of what it means to be "prefixed with a dialect", and it would be good to document some of these invariants(with an example or two). This would allow for us to point to a given location when someone asks what a "dialect prefixed identifier" means.

mlir/lib/IR/MLIRContext.cpp

30

I don't think this is necessary.

[mehdi_amini]

beef up the main description of the Identifier

mlir/include/mlir/IR/Identifier.h

58

PTAL!

[mehdi_amini]

Add an example in the description for getdialect()

[rriddle]

LGTM

mlir/include/mlir/IR/Identifier.h

30–34

57

You use loaded everywhere else in this description, which makes sense given that the dialect has to be loaded for it to exist.

[mehdi_amini]

address comment

[jpienaar]

Seems you retitled the commit, but that got reverted to original description again on push.

mlir/lib/IR/MLIRContext.cpp

493

OOC: if identifierEntry.second is set, can it change here? (E.g., could we rely on/filter based on it, seems like we'd avoid string comparison except where dialect is not set at cost of one extra check to see if it refers to a context or dialect, probably not an important optimization unless we interleaved dialect loading and identified creation much).

[rriddle]

mlir/lib/IR/MLIRContext.cpp

714

Forgot to mention in review, but please rework this to avoid looking up the dialect unless the identifier is being created. We should only be doing the string splicing/dialect lookup when a new identifier is created.

mlir/include/mlir/IR/Identifier.h

 //===- Identifier.h - MLIR Identifier 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_IDENTIFIER_H
 #define MLIR_IR_IDENTIFIER_H
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/StringMapEntry.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
+class Dialect;
 class MLIRContext;
 
 /// This class represents a uniqued string owned by an MLIRContext.  Strings
 /// represented by this type cannot contain nul characters, and may not have a
 /// zero length.
 ///
 /// This is a POD type with pointer size, so it should be passed around by
 /// value.  The underlying data is owned by MLIRContext and is thus immortal for
 /// almost all clients.
 class Identifier {
-  using EntryType = llvm::StringMapEntry<llvm::NoneType>;
+  using EntryType =
+      llvm::StringMapEntry<PointerUnion<Dialect *, MLIRContext *>>;
 
 public:
   /// Return an identifier for the specified string.
   static Identifier get(StringRef str, MLIRContext *context);
   Identifier(const Identifier &) = default;
   Identifier &operator=(const Identifier &other) = default;
 
   /// Return a StringRef for the string.
   StringRef strref() const { return entry->first(); }
 
   /// Identifiers implicitly convert to StringRefs.
   operator StringRef() const { return strref(); }
 
   /// Return an std::string.
   std::string str() const { return strref().str(); }
 
   /// Return a null terminated C string.
   const char *c_str() const { return entry->getKeyData(); }
 
   /// Return a pointer to the start of the string data.
   const char *data() const { return entry->getKeyData(); }
 
   /// Return the number of bytes in this string.
   unsigned size() const { return entry->getKeyLength(); }
 
+  /// Return the dialect registered/loaded in the context for this
+  /// identifier or nullptr if this identifier isn't prefixed with a loaded
+  /// dialect.
+  Dialect *getDialect();
+
+  /// Return the current MLIRContext associated with this identifier.
+  MLIRContext *getContext();
+
   const char *begin() const { return data(); }
   const char *end() const { return entry->getKeyData() + size(); }
 
   bool operator==(Identifier other) const { return entry == other.entry; }
   bool operator!=(Identifier rhs) const { return !(*this == rhs); }
 
   void print(raw_ostream &os) const;
   void dump() const;

mlir/lib/IR/MLIRContext.cpp

 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Support/ThreadLocalCache.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/RWMutex.h"
 #include "llvm/Support/raw_ostream.h"
#endif
   DialectRegistry dialectsRegistry;
 
   /// This is a mapping from operation name to AbstractOperation for registered
   /// operations.
   llvm::StringMap<AbstractOperation> registeredOperations;
 
   /// Identifiers are uniqued by string value and use the internal string set
   /// for storage.
-  llvm::StringSet<llvm::BumpPtrAllocator &> identifiers;
+  llvm::StringMap<PointerUnion<Dialect *, MLIRContext *>,
+                  llvm::BumpPtrAllocator &>
+      identifiers;
   /// A thread local cache of identifiers to reduce lock contention.
-  ThreadLocalCache<llvm::StringMap<llvm::StringMapEntry<llvm::NoneType> *>>
+  ThreadLocalCache<llvm::StringMap<
+      llvm::StringMapEntry<PointerUnion<Dialect *, MLIRContext *>> *>>
       localIdentifierCache;
 
   /// An allocator used for AbstractAttribute and AbstractType objects.
   llvm::BumpPtrAllocator abstractDialectSymbolAllocator;
 
   //===--------------------------------------------------------------------===//
   // Affine uniquing
   //===--------------------------------------------------------------------===//
if (impl.multiThreadedExecutionContext != 0)
       llvm::report_fatal_error(
           "Loading a dialect (" + dialectNamespace +
           ") while in a multi-threaded execution context (maybe "
           "the PassManager): this can indicate a "
           "missing `dependentDialects` in a pass for example.");
 #endif
     dialect = ctor();
     assert(dialect && "dialect ctor failed");
+
+    // Refresh all the identifiers dialect field, this catches cases where a
+    // dialect may be loaded after identifier prefixed with this dialect name
+    // were already created.
+    for (auto &identifierEntry : impl.identifiers)
+      if (identifierEntry.first().startswith(dialectNamespace))
+        identifierEntry.second = dialect.get();
+
     return dialect.get();
   }
 
   // Abort if dialect with namespace has already been registered.
   if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
                              "' has already been registered");
 

 Identifier Identifier::get(StringRef str, MLIRContext *context) {
   // Check invariants after seeing if we already have something in the
   // identifier table - if we already had it in the table, then it already
   // passed invariant checks.
   assert(!str.empty() && "Cannot create an empty identifier");
   assert(str.find('\0') == StringRef::npos &&
          "Cannot create an identifier with a nul character");
 
+  PointerUnion<Dialect *, MLIRContext *> dialectOrContext = context;
+  auto dialectNamePair = str.split('.');
+  if (!dialectNamePair.first.empty()) {
+    if (Dialect *dialect = context->getLoadedDialect(dialectNamePair.first)) {
+      dialectOrContext = dialect;
+    }
+  }
+
   auto &impl = context->getImpl();
   if (!context->isMultithreadingEnabled())
-    return Identifier(&*impl.identifiers.insert(str).first);
+    return Identifier(&*impl.identifiers.insert({str, dialectOrContext}).first);
 
   // Check for an existing instance in the local cache.
   auto *&localEntry = (*impl.localIdentifierCache)[str];
   if (localEntry)
     return Identifier(localEntry);
 
   // Check for an existing identifier in read-only mode.
   {
     llvm::sys::SmartScopedReader<true> contextLock(impl.identifierMutex);
     auto it = impl.identifiers.find(str);
     if (it != impl.identifiers.end()) {
       localEntry = &*it;
       return Identifier(localEntry);
     }
   }
 
   // Acquire a writer-lock so that we can safely create the new instance.
   llvm::sys::SmartScopedWriter<true> contextLock(impl.identifierMutex);
-  auto it = impl.identifiers.insert(str).first;
+  auto it = impl.identifiers.insert({str, dialectOrContext}).first;
   localEntry = &*it;
   return Identifier(localEntry);
 }
 
+Dialect *Identifier::getDialect() {
+  return entry->second.dyn_cast<Dialect *>();
+}
+
+MLIRContext *Identifier::getContext() {
+  if (Dialect *dialect = getDialect())
+    return dialect->getContext();
+  return entry->second.get<MLIRContext *>();
+}
+
 //===----------------------------------------------------------------------===//
 // Type uniquing
 //===----------------------------------------------------------------------===//
 
 /// Returns the storage uniquer used for constructing type storage instances.
 /// This should not be used directly.
 StorageUniquer &MLIRContext::getTypeUniquer() { return getImpl().typeUniquer; }