[MLIR] Deduplicate dialect registration by ClassID

[GMNGeoffrey]

Summary

With the move towards dialect registration that does not depend only use
static initialization, we are running into more cases where the dialects
are registered by different methods. For example, TensorFlow still uses
static initialization to register all MLIR core dialects, which prevents
explicit registration of any of them when linking it in. We ran into this
issue in https://github.com/google/iree/pull/982.

To address potential issues with conflicts from non-standard
allocators passed to registerDialectAllocator, made this method
private. Now all dialects can only be registered with their
constructor.

Similarly deduplicates DialectHooks for consistency and makes their
registration follow the same pattern.

[mehdi_amini]

LGTM, but I'd like to hear from River as well.

I am not sure what you mean with:

This potentially introduces some issues if callers using non-standard dialect registration functions passed to registerDialectAllocator and they are silently overridden.

[GMNGeoffrey]

In D76329#1927997, @mehdi_amini wrote:

LGTM, but I'd like to hear from River as well.

Yeah I'd also like Rivers' review

I am not sure what you mean with:

This potentially introduces some issues if callers using non-standard dialect registration functions passed to registerDialectAllocator and they are silently overridden.

If people call registerDialect<Foo>() we're all good as if there's a duplicate, it would have been the same thing anyway. If someone is directly calling registerDialectAllocator (which I don't think anyone is right now) and they pass in their custom allocator that they expect to do something and someone else also registers the same dialect (either with registerDialect or registerDialectAllocator) then we might have surprising behavior. We could move registerDialectAllocator out of the public API, although it's used by the template function, so a bit tricky (maybe just an impl/detail namespace?), but I wanted to hear from others before doing that.

[rriddle]

In D76329#1928036, @GMNGeoffrey wrote:

In D76329#1927997, @mehdi_amini wrote:

LGTM, but I'd like to hear from River as well.

Yeah I'd also like Rivers' review

I am not sure what you mean with:

This potentially introduces some issues if callers using non-standard dialect registration functions passed to registerDialectAllocator and they are silently overridden.

If people call registerDialect<Foo>() we're all good as if there's a duplicate, it would have been the same thing anyway. If someone is directly calling registerDialectAllocator (which I don't think anyone is right now) and they pass in their custom allocator that they expect to do something and someone else also registers the same dialect (either with registerDialect or registerDialectAllocator) then we might have surprising behavior. We could move registerDialectAllocator out of the public API, although it's used by the template function, so a bit tricky (maybe just an impl/detail namespace?), but I wanted to hear from others before doing that.

In general I'm fine with this, as it makes some things easier(we can now make getRegisteredDialect O(1)(modulo locking)). The only problem I can foresee is with the allocator, as you describe. IMO we can just hide that method for now.

mlir/lib/IR/Dialect.cpp

29–30

Can you change these to /// while you are here?

30–32

Can you change this to a MapVector or other stable container? DenseMap does not have a stable insertion order.

39–40

Can you make this a private method inside of Dialect that is only accessible via registerDialect?

[GMNGeoffrey]

Stable maps and private registration.

Use stable-ordered maps. Make allocator functions private static methods.
Make dialect and dialect hooks registry have the same format.

[GMNGeoffrey]

mlir/lib/IR/Dialect.cpp

39–40

Made this and the corresponding hooks setter registration private static methods on the class with friend declarations. Hopefully that's what you meant. In the process, made the registration for hooks and dialects the same format (separate the free template function and the global initialization struct).

[GMNGeoffrey]

River if this LGTY, could you please commit it?

[GMNGeoffrey]

Looks like the failure is a clang format diff. I think this is related to the discussion in https://llvm.discourse.group/t/remove-mlir-custom-clang-format-configuration/647/4? I can reformat as needed, but +1 for unifying LLVM and MLIR configs

[GMNGeoffrey]

Break templates declaration to placate clang format pre-merge.

mlir/include/mlir/IR/Dialect.h

 class Type;
 
 using DialectConstantDecodeHook =
     std::function<bool(const OpaqueElementsAttr, ElementsAttr &)>;
 using DialectConstantFoldHook = std::function<LogicalResult(
     Operation *, ArrayRef<Attribute>, SmallVectorImpl<Attribute> &)>;
 using DialectExtractElementHook =
     std::function<Attribute(const OpaqueElementsAttr, ArrayRef<uint64_t>)>;
+using DialectAllocatorFunction = std::function<void(MLIRContext *)>;
 
 /// Dialects are groups of MLIR operations and behavior associated with the
 /// entire group.  For example, hooks into other systems for constant folding,
 /// default named types for asm printing, etc.
 ///
 /// Instances of the dialect object are global across all MLIRContext's that may
 /// be active in the process.
 ///
private:
 
   /// Flag that specifies whether this dialect allows unregistered types, i.e.
   /// types prefixed with the dialect namespace but not registered with addType.
   /// These types are represented with OpaqueType.
   bool unknownTypesAllowed = false;
 
   /// A collection of registered dialect interfaces.
   DenseMap<ClassID *, std::unique_ptr<DialectInterface>> registeredInterfaces;
-};
-
-using DialectAllocatorFunction = std::function<void(MLIRContext *)>;
-
-/// Registers a specific dialect creation function with the system, typically
-/// used through the DialectRegistration template.
-void registerDialectAllocator(const DialectAllocatorFunction &function);
 
-/// Registers all dialects with the specified MLIRContext.
+  /// Registers a specific dialect creation function with the global registry.
+  /// Used through the registerDialect template.
+  /// Registrations are deduplicated by dialect ClassID and only the first
+  /// registration will be used.
+  static void
+  registerDialectAllocator(const ClassID *classId,
+                           const DialectAllocatorFunction &function);
+  template <typename ConcreteDialect>
+  friend void registerDialect();
+};
+/// Registers all dialects and hooks from the global registries with the
+/// specified MLIRContext.
 void registerAllDialects(MLIRContext *context);
 
 /// Utility to register a dialect. Client can register their dialect with the
 /// global registry by calling registerDialect<MyDialect>();
 template <typename ConcreteDialect> void registerDialect() {
-  registerDialectAllocator([](MLIRContext *ctx) {
-    // Just allocate the dialect, the context takes ownership of it.
+  Dialect::registerDialectAllocator(ClassID::getID<ConcreteDialect>(),
+                                    [](MLIRContext *ctx) {
+                                      // Just allocate the dialect, the context
+                                      // takes ownership of it.
                                       new ConcreteDialect(ctx);
                                     });
 }
 
 /// DialectRegistration provides a global initializer that registers a Dialect
 /// allocation routine.
 ///
 /// Usage:
 ///
 ///   // At namespace scope.
 ///   static DialectRegistration<MyDialect> Unused;
 template <typename ConcreteDialect> struct DialectRegistration {
   DialectRegistration() { registerDialect<ConcreteDialect>(); }
 };
 
 } // namespace mlir
 
 #endif

mlir/include/mlir/IR/DialectHooks.h

 class DialectHooks {
 public:
   // Returns hook to constant fold an operation.
   DialectConstantFoldHook getConstantFoldHook() { return nullptr; }
   // Returns hook to decode opaque constant tensor.
   DialectConstantDecodeHook getDecodeHook() { return nullptr; }
   // Returns hook to extract an element of an opaque constant tensor.
   DialectExtractElementHook getExtractElementHook() { return nullptr; }
+
+private:
+  /// Registers a function that will set hooks in the registered dialects.
+  /// Registrations are deduplicated by dialect ClassID and only the first
+  /// registration will be used.
+  static void registerDialectHooksSetter(const ClassID *classId,
+                                         const DialectHooksSetter &function);
+  template <typename ConcreteHooks>
+  friend void registerDialectHooks(StringRef dialectName);
 };
 
-/// Registers a function that will set hooks in the registered dialects
-/// based on information coming from DialectHooksRegistration.
-void registerDialectHooksSetter(const DialectHooksSetter &function);
+void registerDialectHooksSetter(const ClassID *classId,
+                                const DialectHooksSetter &function);
 
-/// DialectHooksRegistration provides a global initializer that registers
-/// a dialect hooks setter routine.
-/// Usage:
-///
-///   // At namespace scope.
-///   static DialectHooksRegistration<MyHooks, MyDialect> unused;
-template <typename ConcreteHooks> struct DialectHooksRegistration {
-  DialectHooksRegistration(StringRef dialectName) {
-    registerDialectHooksSetter([dialectName](MLIRContext *ctx) {
+/// Utility to register dialect hooks. Client can register their dialect hooks
+/// with the global registry by calling
+/// registerDialectHooks<MyHooks>("dialect_namespace");
+template <typename ConcreteHooks>
+void registerDialectHooks(StringRef dialectName) {
+  DialectHooks::registerDialectHooksSetter(
+      ClassID::getID<ConcreteHooks>(), [dialectName](MLIRContext *ctx) {
         Dialect *dialect = ctx->getRegisteredDialect(dialectName);
         if (!dialect) {
           llvm::errs() << "error: cannot register hooks for unknown dialect '"
                        << dialectName << "'\n";
           abort();
         }
         // Set hooks.
         ConcreteHooks hooks;
         if (auto h = hooks.getConstantFoldHook())
           dialect->constantFoldHook = h;
         if (auto h = hooks.getDecodeHook())
           dialect->decodeHook = h;
         if (auto h = hooks.getExtractElementHook())
           dialect->extractElementHook = h;
       });
-  }
+
+/// DialectHooksRegistration provides a global initializer that registers
+/// a dialect hooks setter routine.
+/// Usage:
+///
+///   // At namespace scope.
+///   static DialectHooksRegistration<MyHooks> Unused("dialect_namespace");
+template <typename ConcreteHooks> struct DialectHooksRegistration {
+  DialectHooksRegistration(StringRef dialectName) {
+    registerDialectHooks<ConcreteHooks>(dialectName);
+  }
 };
 
 } // namespace mlir
 
 #endif

mlir/lib/IR/Dialect.cpp

 //===- Dialect.cpp - Dialect implementation -------------------------------===//
 //
 // 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 "mlir/IR/Dialect.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectHooks.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Regex.h"
 
 using namespace mlir;
 using namespace detail;
 
 DialectAsmParser::~DialectAsmParser() {}
 
 //===----------------------------------------------------------------------===//
 // Dialect Registration
 //===----------------------------------------------------------------------===//
 
-// Registry for all dialect allocation functions.
-static llvm::ManagedStatic<SmallVector<DialectAllocatorFunction, 8>>
+/// Registry for all dialect allocation functions.
+static llvm::ManagedStatic<
+    llvm::MapVector<const ClassID *, DialectAllocatorFunction>>
     dialectRegistry;
 
-// Registry for functions that set dialect hooks.
-static llvm::ManagedStatic<SmallVector<DialectHooksSetter, 8>>
+/// Registry for functions that set dialect hooks.
+static llvm::ManagedStatic<llvm::MapVector<const ClassID *, DialectHooksSetter>>
     dialectHooksRegistry;
 
-/// Registers a specific dialect creation function with the system, typically
-/// used through the DialectRegistration template.
-void mlir::registerDialectAllocator(const DialectAllocatorFunction &function) {
+void Dialect::registerDialectAllocator(
+    const ClassID *classId, const DialectAllocatorFunction &function) {
   assert(function &&
          "Attempting to register an empty dialect initialize function");
-  dialectRegistry->push_back(function);
+  dialectRegistry->insert({classId, function});
 }
 
 /// Registers a function to set specific hooks for a specific dialect, typically
 /// used through the DialectHooksRegistration template.
-void mlir::registerDialectHooksSetter(const DialectHooksSetter &function) {
+void DialectHooks::registerDialectHooksSetter(
+    const ClassID *classId, const DialectHooksSetter &function) {
   assert(
       function &&
       "Attempting to register an empty dialect hooks initialization function");
 
-  dialectHooksRegistry->push_back(function);
+  dialectHooksRegistry->insert({classId, function});
 }
 
-/// Registers all dialects and their const folding hooks with the specified
-/// MLIRContext.
+/// Registers all dialects and hooks from the global registries with the
+/// specified MLIRContext.
 void mlir::registerAllDialects(MLIRContext *context) {
-  for (const auto &fn : *dialectRegistry)
-    fn(context);
-  for (const auto &fn : *dialectHooksRegistry) {
-    fn(context);
+  for (const auto &it : *dialectRegistry)
+    it.second(context);
+  for (const auto &it : *dialectHooksRegistry) {
+    it.second(context);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//
 
 Dialect::Dialect(StringRef name, MLIRContext *context)