Revisit Dialect registration: require and store a TypeID on dialects

[mehdi_amini]

Summary

This patch moves the registration to a method in the MLIRContext: getOrCreateDialect<ConcreteDialect>()

This method requires dialect to provide a static getDialectNamespace()
and store a TypeID on the Dialect itself, which allows to lazyily
create a dialect when not yet loaded in the context.
As a side effect, it means that duplicated registration of the same
dialect is not an issue anymore.

To limit the boilerplate, TableGen dialect generation is modified to
emit the constructor entirely and invoke separately a "init()" method
that the user implements.

[mehdi_amini]

Rename init to initialize

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

30–32

Why this change?

mlir/include/mlir/IR/MLIRContext.h

59

nit: new T(this)?

mlir/lib/IR/MLIRContext.cpp

476–477

nit:

auto it  = impl.dialects.insert(insertPt, ctor());
return &**it;?

mlir/lib/Transforms/LoopFusion.cpp

1456 ↗

(On Diff #283809)

Seems unrelated.

[mehdi_amini]

Address comments

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

30–32

Constructor is not longer out-of-line and the forward declaration + the unique_ptr is the usual issue...

mlir/lib/Transforms/LoopFusion.cpp

1456 ↗

(On Diff #283809)

Ahah my sed caught a bit wide apparently, reverted.

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

30–32

Right, forgot that we emit the constructor in the .h now. Can you add a comment here describing why we use a raw pointer? This could no longer be an issue if we generated dialect defs too at some point.

mlir/tools/mlir-tblgen/DialectGen.cpp

67

nit: ::mlir::TypeID::

[tskeith]

This broke the flang build:

../../flang/lib/Optimizer/Dialect/FIRDialect.cpp:17:7: error: no matching constructor for initialization of 'mlir::Dialect'
    : mlir::Dialect("fir", ctx) {
      ^             ~~~~~~~~~~

[rriddle]

In D85495#2203335, @tskeith wrote:

This broke the flang build:

../../flang/lib/Optimizer/Dialect/FIRDialect.cpp:17:7: error: no matching constructor for initialization of 'mlir::Dialect'
    : mlir::Dialect("fir", ctx) {
      ^             ~~~~~~~~~~

Should be fixed in 82fd1392016984c81c6037e147ee2dd36cf91f4c

[mehdi_amini]

In D85495#2203335, @tskeith wrote:

This broke the flang build:

../../flang/lib/Optimizer/Dialect/FIRDialect.cpp:17:7: error: no matching constructor for initialization of 'mlir::Dialect'
    : mlir::Dialect("fir", ctx) {
      ^             ~~~~~~~~~~

Sorry, I should build flang when refactoring mlir. Thanks @rriddle for the quick fix!

mlir/examples/standalone/lib/Standalone/StandaloneDialect.cpp

 
 using namespace mlir;
 using namespace mlir::standalone;
 
 //===----------------------------------------------------------------------===//
 // Standalone dialect.
 //===----------------------------------------------------------------------===//
 
-StandaloneDialect::StandaloneDialect(mlir::MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void StandaloneDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "Standalone/StandaloneOps.cpp.inc"
       >();
 }

mlir/examples/toy/Ch2/mlir/Dialect.cpp

 using namespace mlir::toy;
 
 //===----------------------------------------------------------------------===//
 // ToyDialect
 //===----------------------------------------------------------------------===//
 
 /// Dialect creation, the instance will be owned by the context. This is the
 /// point of registration of custom types and operations for the dialect.
-ToyDialect::ToyDialect(mlir::MLIRContext *ctx) : mlir::Dialect("toy", ctx) {
+ToyDialect::ToyDialect(mlir::MLIRContext *ctx)
+    : mlir::Dialect(getDialectNamespace(), ctx, TypeID::get<ToyDialect>()) {
   addOperations<
 #define GET_OP_LIST
 #include "toy/Ops.cpp.inc"
       >();
 }
 
 //===----------------------------------------------------------------------===//
 // Toy Operations

mlir/examples/toy/Ch3/mlir/Dialect.cpp

 using namespace mlir::toy;
 
 //===----------------------------------------------------------------------===//
 // ToyDialect
 //===----------------------------------------------------------------------===//
 
 /// Dialect creation, the instance will be owned by the context. This is the
 /// point of registration of custom types and operations for the dialect.
-ToyDialect::ToyDialect(mlir::MLIRContext *ctx) : mlir::Dialect("toy", ctx) {
+ToyDialect::ToyDialect(mlir::MLIRContext *ctx)
+    : mlir::Dialect(getDialectNamespace(), ctx, TypeID::get<ToyDialect>()) {
   addOperations<
 #define GET_OP_LIST
 #include "toy/Ops.cpp.inc"
       >();
 }
 
 //===----------------------------------------------------------------------===//
 // Toy Operations

mlir/examples/toy/Ch4/mlir/Dialect.cpp

 };
 
 //===----------------------------------------------------------------------===//
 // ToyDialect
 //===----------------------------------------------------------------------===//
 
 /// Dialect creation, the instance will be owned by the context. This is the
 /// point of registration of custom types and operations for the dialect.
-ToyDialect::ToyDialect(mlir::MLIRContext *ctx) : mlir::Dialect("toy", ctx) {
+ToyDialect::ToyDialect(mlir::MLIRContext *ctx)
+    : mlir::Dialect(getDialectNamespace(), ctx, TypeID::get<ToyDialect>()) {
   addOperations<
 #define GET_OP_LIST
 #include "toy/Ops.cpp.inc"
       >();
   addInterfaces<ToyInlinerInterface>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/examples/toy/Ch5/mlir/Dialect.cpp

 };
 
 //===----------------------------------------------------------------------===//
 // ToyDialect
 //===----------------------------------------------------------------------===//
 
 /// Dialect creation, the instance will be owned by the context. This is the
 /// point of registration of custom types and operations for the dialect.
-ToyDialect::ToyDialect(mlir::MLIRContext *ctx) : mlir::Dialect("toy", ctx) {
+ToyDialect::ToyDialect(mlir::MLIRContext *ctx)
+    : mlir::Dialect(getDialectNamespace(), ctx, TypeID::get<ToyDialect>()) {
   addOperations<
 #define GET_OP_LIST
 #include "toy/Ops.cpp.inc"
       >();
   addInterfaces<ToyInlinerInterface>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/examples/toy/Ch6/mlir/Dialect.cpp

 };
 
 //===----------------------------------------------------------------------===//
 // ToyDialect
 //===----------------------------------------------------------------------===//
 
 /// Dialect creation, the instance will be owned by the context. This is the
 /// point of registration of custom types and operations for the dialect.
-ToyDialect::ToyDialect(mlir::MLIRContext *ctx) : mlir::Dialect("toy", ctx) {
+ToyDialect::ToyDialect(mlir::MLIRContext *ctx)
+    : mlir::Dialect(getDialectNamespace(), ctx, TypeID::get<ToyDialect>()) {
   addOperations<
 #define GET_OP_LIST
 #include "toy/Ops.cpp.inc"
       >();
   addInterfaces<ToyInlinerInterface>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/examples/toy/Ch7/mlir/Dialect.cpp

 };
 
 //===----------------------------------------------------------------------===//
 // ToyDialect
 //===----------------------------------------------------------------------===//
 
 /// Dialect creation, the instance will be owned by the context. This is the
 /// point of registration of custom types and operations for the dialect.
-ToyDialect::ToyDialect(mlir::MLIRContext *ctx) : mlir::Dialect("toy", ctx) {
+ToyDialect::ToyDialect(mlir::MLIRContext *ctx)
+    : mlir::Dialect(getDialectNamespace(), ctx, TypeID::get<ToyDialect>()) {
   addOperations<
 #define GET_OP_LIST
 #include "toy/Ops.cpp.inc"
       >();
   addInterfaces<ToyInlinerInterface>();
   addTypes<StructType>();
 }
 

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

def LLVM_Dialect : Dialect {
   let hasRegionArgAttrVerify = 1;
   let extraClassDeclaration = [{
     ~LLVMDialect();
     const llvm::DataLayout &getDataLayout();
 
   private:
     friend LLVMType;
 
-    std::unique_ptr<detail::LLVMDialectImpl> impl;
+    // This can't be a unique_ptr because the ctor is generated inline
+    // in the class definition at the moment.
+    detail::LLVMDialectImpl *impl;
   }];
 }
 
 // LLVM IR type wrapped in MLIR.
 def LLVM_Type : DialectType<LLVM_Dialect,
                             CPred<"$_self.isa<::mlir::LLVM::LLVMType>()">,
                             "LLVM dialect type">;
 

mlir/include/mlir/Dialect/SDBM/SDBMDialect.h

 #include "mlir/IR/Dialect.h"
 #include "mlir/Support/StorageUniquer.h"
 
 namespace mlir {
 class MLIRContext;
 
 class SDBMDialect : public Dialect {
 public:
-  SDBMDialect(MLIRContext *context) : Dialect(getDialectNamespace(), context) {}
+  SDBMDialect(MLIRContext *context)
+      : Dialect(getDialectNamespace(), context, TypeID::get<SDBMDialect>()) {}
 
   /// Since there are no other virtual methods in this derived class, override
   /// the destructor so that key methods get defined in the corresponding
   /// module.
   ~SDBMDialect() override;
 
   static StringRef getDialectNamespace() { return "sdbm"; }
 

mlir/include/mlir/IR/Dialect.h

 //===- Dialect.h - IR Dialect Description -----------------------*- 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 'dialect' abstraction.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_DIALECT_H
 #define MLIR_IR_DIALECT_H
 
 #include "mlir/IR/OperationSupport.h"
+#include "mlir/Support/TypeID.h"
 
 namespace mlir {
 class DialectAsmParser;
 class DialectAsmPrinter;
 class DialectInterface;
 class OpBuilder;
 class Type;
 
public:
   /// Utility function that returns if the given string is a valid dialect
   /// namespace.
   static bool isValidNamespace(StringRef str);
 
   MLIRContext *getContext() const { return context; }
 
   StringRef getNamespace() const { return name; }
 
+  /// Returns the unique identifier that corresponds to this dialect.
+  TypeID getTypeID() const { return dialectID; }
+
   /// Returns true if this dialect allows for unregistered operations, i.e.
   /// operations prefixed with the dialect namespace but not registered with
   /// addOperation.
   bool allowsUnknownOperations() const { return unknownOpsAllowed; }
 
   /// Return true if this dialect allows for unregistered types, i.e., types
   /// prefixed with the dialect namespace but not registered with addType.
   /// These are represented with OpaqueType.

 protected:
   /// The constructor takes a unique namespace for this dialect as well as the
   /// context to bind to.
   /// Note: The namespace must not contain '.' characters.
   /// Note: All operations belonging to this dialect must have names starting
   ///       with the namespace followed by '.'.
   /// Example:
   ///       - "tf" for the TensorFlow ops like "tf.add".
-  Dialect(StringRef name, MLIRContext *context);
+  Dialect(StringRef name, MLIRContext *context, TypeID id);
 
   /// This method is used by derived classes to add their operations to the set.
   ///
   template <typename... Args> void addOperations() {
     (void)std::initializer_list<int>{
         0, (addOperation(AbstractOperation::get<Args>(*this)), 0)...};
   }
 
template <typename... Args> void addInterfaces() {
     (void)std::initializer_list<int>{
         0, (addInterface(std::make_unique<Args>(this)), 0)...};
   }
 
 private:
   Dialect(const Dialect &) = delete;
   void operator=(Dialect &) = delete;
 
-  /// Register this dialect object with the specified context.  The context
-  /// takes ownership of the heap allocated dialect.
-  void registerDialect(MLIRContext *context);
-
   /// The namespace of this dialect.
   StringRef name;
 
+  /// The unique identifier of the derived Op class, this is used in the context
+  /// to allow registering multiple times the same dialect.
+  TypeID dialectID;
+
   /// This is the context that owns this Dialect object.
   MLIRContext *context;
 
   /// Flag that specifies whether this dialect supports unregistered operations,
   /// i.e. operations prefixed with the dialect namespace but not registered
   /// with addOperation.
   bool unknownOpsAllowed = false;
 
private:
   /// Used through the registerDialect template.
   /// Registrations are deduplicated by dialect TypeID and only the first
   /// registration will be used.
   static void
   registerDialectAllocator(TypeID typeID,
                            const DialectAllocatorFunction &function);
   template <typename ConcreteDialect>
   friend void registerDialect();
+  friend class MLIRContext;
 };
+
 /// Registers all dialects and hooks from the global registries with the
 /// specified MLIRContext.
 /// Note: This method is not thread-safe.
 void registerAllDialects(MLIRContext *context);
 
 /// Utility to register a dialect. Client can register their dialect with the
 /// global registry by calling registerDialect<MyDialect>();
 /// Note: This method is not thread-safe.
 template <typename ConcreteDialect> void registerDialect() {
-  Dialect::registerDialectAllocator(TypeID::get<ConcreteDialect>(),
-                                    [](MLIRContext *ctx) {
-                                      // Just allocate the dialect, the context
-                                      // takes ownership of it.
-                                      new ConcreteDialect(ctx);
-                                    });
+  Dialect::registerDialectAllocator(
+      TypeID::get<ConcreteDialect>(),
+      [](MLIRContext *ctx) { ctx->getOrCreateDialect<ConcreteDialect>(); });
 }
 
 /// 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
 
 namespace llvm {
 /// Provide isa functionality for Dialects.
 template <typename T>
 struct isa_impl<T, ::mlir::Dialect> {
   static inline bool doit(const ::mlir::Dialect &dialect) {
-    return T::getDialectNamespace() == dialect.getNamespace();
+    return mlir::TypeID::get<T>() == dialect.getTypeID();
   }
 };
 } // namespace llvm
 
 #endif

mlir/include/mlir/IR/MLIRContext.h

 //===- MLIRContext.h - MLIR Global Context 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_MLIRCONTEXT_H
 #define MLIR_IR_MLIRCONTEXT_H
 
 #include "mlir/Support/LLVM.h"
+#include "mlir/Support/TypeID.h"
 #include <functional>
 #include <memory>
 #include <vector>
 
 namespace mlir {
 class AbstractOperation;
 class DiagnosticEngine;
 class Dialect;
public:
   Dialect *getRegisteredDialect(StringRef name);
 
   /// Get a registered IR dialect for the given derived dialect type. The
   /// derived type must provide a static 'getDialectNamespace' method.
   template <typename T> T *getRegisteredDialect() {
     return static_cast<T *>(getRegisteredDialect(T::getDialectNamespace()));
   }
 
+  /// Get (or create) a dialect for the given derived dialect type. The derived
+  /// type must provide a static 'getDialectNamespace' method.
+  template <typename T>
+  T *getOrCreateDialect() {
+    return static_cast<T *>(getOrCreateDialect(
+        T::getDialectNamespace(), TypeID::get<T>(), [this]() {
+          std::unique_ptr<T> dialect(new T(this));
+          dialect->dialectID = TypeID::get<T>();
+          return dialect;
+        }));
+  }
+
   /// Return true if we allow to create operation for unregistered dialects.
   bool allowsUnregisteredDialects();
 
   /// Enables creating operations in unregistered dialects.
   void allowUnregisteredDialects(bool allow = true);
 
   /// Return true if multi-threading is enabled by the context.
   bool isMultithreadingEnabled();
public:
 
   /// Returns the storage uniquer used for constructing attribute storage
   /// instances. This should not be used directly.
   StorageUniquer &getAttributeUniquer();
 
 private:
   const std::unique_ptr<MLIRContextImpl> impl;
 
+  /// Get a dialect for the provided namespace and TypeID: abort the program if
+  /// a dialect exist for this namespace with different TypeID. Returns a
+  /// pointer to the dialect owned by the context.
+  Dialect *getOrCreateDialect(StringRef dialectNamespace, TypeID dialectID,
+                              function_ref<std::unique_ptr<Dialect>()> ctor);
+
   MLIRContext(const MLIRContext &) = delete;
   void operator=(const MLIRContext &) = delete;
 };
 
 //===----------------------------------------------------------------------===//
 // MLIRContext CommandLine Options
 //===----------------------------------------------------------------------===//
 
 /// Register a set of useful command-line options that can be used to configure
 /// various flags within the MLIRContext. These flags are used when constructing
 /// an MLIR context for initialization.
 void registerMLIRContextCLOptions();
 
 } // end namespace mlir
 
 #endif // MLIR_IR_MLIRCONTEXT_H

mlir/lib/Dialect/AVX512/IR/AVX512Dialect.cpp

 #include "mlir/Dialect/AVX512/AVX512Dialect.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeUtilities.h"
 
 using namespace mlir;
 
-avx512::AVX512Dialect::AVX512Dialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void avx512::AVX512Dialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/AVX512/AVX512.cpp.inc"
       >();
 }
 
 namespace mlir {
 namespace avx512 {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/AVX512/AVX512.cpp.inc"
 }  // namespace avx512
 } // namespace mlir
 

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

struct AffineInlinerInterface : public DialectInlinerInterface {
   bool shouldAnalyzeRecursively(Operation *op) const final { return true; }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // AffineDialect
 //===----------------------------------------------------------------------===//
 
-AffineDialect::AffineDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void AffineDialect::initialize() {
   addOperations<AffineDmaStartOp, AffineDmaWaitOp,
 #define GET_OP_LIST
 #include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"
                 >();
   addInterfaces<AffineInlinerInterface>();
 }
 
 /// Materialize a single constant operation from a given attribute value with

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 // GPUDialect
 //===----------------------------------------------------------------------===//
 
 bool GPUDialect::isKernel(Operation *op) {
   UnitAttr isKernelAttr = op->getAttrOfType<UnitAttr>(getKernelFuncAttrName());
   return static_cast<bool>(isKernelAttr);
 }
 
-GPUDialect::GPUDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void GPUDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
       >();
 }
 
 LogicalResult GPUDialect::verifyOperationAttribute(Operation *op,
                                                    NamedAttribute attr) {

mlir/lib/Dialect/LLVMIR/IR/LLVMAVX512Dialect.cpp

 #include "mlir/Dialect/LLVMIR/LLVMAVX512Dialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeUtilities.h"
 
 using namespace mlir;
 
-LLVM::LLVMAVX512Dialect::LLVMAVX512Dialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void LLVM::LLVMAVX512Dialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/LLVMIR/LLVMAVX512.cpp.inc"
       >();
 }
 
 namespace mlir {
 namespace LLVM {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMAVX512.cpp.inc"
 } // namespace LLVM
 } // namespace mlir

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

struct LLVMDialectImpl {
   // TODO: this should be moved to some Op equivalent to LLVM module and
   // eventually replaced with a proper MLIR data layout.
   llvm::DataLayout layout;
 };
 } // end namespace detail
 } // end namespace LLVM
 } // end namespace mlir
 
-LLVMDialect::LLVMDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context),
-      impl(new detail::LLVMDialectImpl()) {
+void LLVMDialect::initialize() {
+  impl = new detail::LLVMDialectImpl();
   // clang-format off
   addTypes<LLVMVoidType,
            LLVMHalfType,
            LLVMBFloatType,
            LLVMFloatType,
            LLVMDoubleType,
            LLVMFP128Type,
            LLVMX86FP80Type,

 #define GET_OP_LIST
 #include "mlir/Dialect/LLVMIR/LLVMOps.cpp.inc"
       >();
 
   // Support unknown operations because not all LLVM operations are registered.
   allowUnknownOperations();
 }
 
-LLVMDialect::~LLVMDialect() {}
+LLVMDialect::~LLVMDialect() { delete impl; }
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMOps.cpp.inc"
 
 const llvm::DataLayout &LLVMDialect::getDataLayout() { return impl->layout; }
 
 /// Parse a type registered to this dialect.
 Type LLVMDialect::parseType(DialectAsmParser &parser) const {

mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

static LogicalResult verify(MmaOp op) {
   return op.emitOpError("unimplemented mma.sync variant");
 }
 
 //===----------------------------------------------------------------------===//
 // NVVMDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 
 // TODO: This should be the llvm.nvvm dialect once this is supported.
-NVVMDialect::NVVMDialect(MLIRContext *context) : Dialect("nvvm", context) {
+void NVVMDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
       >();
 
   // Support unknown operations because not all NVVM operations are registered.
   allowUnknownOperations();
 }
 
 namespace mlir {
 namespace NVVM {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
 } // namespace NVVM
 } // namespace mlir
 

mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp

static ParseResult parseROCDLMubufStoreOp(OpAsmParser &parser,
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ROCDLDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 
 // TODO: This should be the llvm.rocdl dialect once this is supported.
-ROCDLDialect::ROCDLDialect(MLIRContext *context) : Dialect("rocdl", context) {
+void ROCDLDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
       >();
 
   // Support unknown operations because not all ROCDL operations are registered.
   allowUnknownOperations();
 }
 
 namespace mlir {
 namespace ROCDL {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
 } // namespace ROCDL
 } // namespace mlir
 

mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp

 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
-mlir::linalg::LinalgDialect::LinalgDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void mlir::linalg::LinalgDialect::initialize() {
   addTypes<RangeType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
       >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

 #include "llvm/ADT/StringSwitch.h"
 #include <cstddef>
 
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.cpp.inc"
 
 using namespace mlir;
 using namespace mlir::omp;
 
-OpenMPDialect::OpenMPDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void OpenMPDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
       >();
 }
 
 //===----------------------------------------------------------------------===//
 // ParallelOp

mlir/lib/Dialect/Quant/IR/QuantOps.cpp

 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/MathExtras.h"
 #include <numeric>
 
 using namespace mlir;
 using namespace mlir::quant;
 using namespace mlir::quant::detail;
 
-QuantizationDialect::QuantizationDialect(MLIRContext *context)
-    : Dialect(/*name=*/"quant", context) {
+void QuantizationDialect::initialize() {
   addTypes<AnyQuantizedType, UniformQuantizedType,
            UniformQuantizedPerAxisType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Quant/QuantOps.cpp.inc"
       >();
 }
 

mlir/lib/Dialect/SCF/SCF.cpp

struct SCFInlinerInterface : public DialectInlinerInterface {
   }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // SCFDialect
 //===----------------------------------------------------------------------===//
 
-SCFDialect::SCFDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void SCFDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/SCF/SCFOps.cpp.inc"
       >();
   addInterfaces<SCFInlinerInterface>();
 }
 
 /// Default callback for IfOp builders. Inserts a yield without arguments.

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

struct SPIRVInlinerInterface : public DialectInlinerInterface {
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // SPIR-V Dialect
 //===----------------------------------------------------------------------===//
 
-SPIRVDialect::SPIRVDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void SPIRVDialect::initialize() {
   addTypes<ArrayType, CooperativeMatrixNVType, ImageType, MatrixType,
            PointerType, RuntimeArrayType, StructType>();
 
   addAttributes<InterfaceVarABIAttr, TargetEnvAttr, VerCapExtAttr>();
 
   // Add SPIR-V ops.
   addOperations<
 #define GET_OP_LIST

mlir/lib/Dialect/Shape/IR/Shape.cpp

if (isErrorPropagationPossible(op->getOperandTypes())) {
     if (!resultTy.isa<ShapeType>())
       return op->emitOpError()
              << "if at least one of the operands can hold error values then "
                 "the result must be of type `shape` to propagate them";
   }
   return success();
 }
 
-ShapeDialect::ShapeDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void ShapeDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
       >();
   addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType,
            WitnessType>();
   // Allow unknown operations during prototyping and testing. As the dialect is
   // still evolving it makes it simple to start with an unregistered ops and

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static LogicalResult verifyCastOp(T op) {
   auto resType = op.getType();
   if (!T::areCastCompatible(opType, resType))
     return op.emitError("operand type ") << opType << " and result type "
                                          << resType << " are cast incompatible";
 
   return success();
 }
 
-StandardOpsDialect::StandardOpsDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void StandardOpsDialect::initialize() {
   addOperations<DmaStartOp, DmaWaitOp,
 #define GET_OP_LIST
 #include "mlir/Dialect/StandardOps/IR/Ops.cpp.inc"
                 >();
   addInterfaces<StdInlinerInterface>();
 }
 
 /// Materialize a single constant operation from a given attribute value with

mlir/lib/Dialect/Vector/VectorOps.cpp

 
 using namespace mlir;
 using namespace mlir::vector;
 
 //===----------------------------------------------------------------------===//
 // VectorDialect
 //===----------------------------------------------------------------------===//
 
-VectorDialect::VectorDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void VectorDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Vector/VectorOps.cpp.inc"
       >();
 }
 
 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.

mlir/lib/IR/Dialect.cpp

void mlir::registerAllDialects(MLIRContext *context) {
   for (const auto &it : *dialectHooksRegistry)
     it.second(context);
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//
 
-Dialect::Dialect(StringRef name, MLIRContext *context)
-    : name(name), context(context) {
+Dialect::Dialect(StringRef name, MLIRContext *context, TypeID id)
+    : name(name), dialectID(id), context(context) {
   assert(isValidNamespace(name) && "invalid dialect namespace");
-  registerDialect(context);
 }
 
 Dialect::~Dialect() {}
 
 /// Verify an attribute from this dialect on the argument at 'argIndex' for
 /// the region at 'regionIndex' on the given operation. Returns failure if
 /// the verification failed, success otherwise. This hook may optionally be
 /// invoked from any operation containing a region.

mlir/lib/IR/MLIRContext.cpp

 //===----------------------------------------------------------------------===//
 // Builtin Dialect
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// A builtin dialect to define types/etc that are necessary for the validity of
 /// the IR.
 struct BuiltinDialect : public Dialect {
-  BuiltinDialect(MLIRContext *context) : Dialect(/*name=*/"", context) {
+  BuiltinDialect(MLIRContext *context)
+      : Dialect(/*name=*/"", context, TypeID::get<BuiltinDialect>()) {
     addAttributes<AffineMapAttr, ArrayAttr, DenseIntOrFPElementsAttr,
                   DenseStringElementsAttr, DictionaryAttr, FloatAttr,
                   SymbolRefAttr, IntegerAttr, IntegerSetAttr, OpaqueAttr,
                   OpaqueElementsAttr, SparseElementsAttr, StringAttr, TypeAttr,
                   UnitAttr>();
     addAttributes<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc, OpaqueLoc,
                   UnknownLoc>();
 
     addTypes<ComplexType, FloatType, FunctionType, IndexType, IntegerType,
              MemRefType, UnrankedMemRefType, NoneType, OpaqueType,
              RankedTensorType, TupleType, UnrankedTensorType, VectorType>();
 
     // TODO: These operations should be moved to a different dialect when they
     // have been fully decoupled from the core.
     addOperations<FuncOp, ModuleOp, ModuleTerminatorOp>();
   }
+  static StringRef getDialectNamespace() { return ""; }
 };
 } // end anonymous namespace.
 
 //===----------------------------------------------------------------------===//
 // Locking Utilities
 //===----------------------------------------------------------------------===//
 
 namespace {
MLIRContext::MLIRContext() : impl(new MLIRContextImpl()) {
   // Initialize values based on the command line flags if they were provided.
   if (clOptions.isConstructed()) {
     disableMultithreading(clOptions->disableThreading);
     printOpOnDiagnostic(clOptions->printOpOnDiagnostic);
     printStackTraceOnDiagnostic(clOptions->printStackTraceOnDiagnostic);
   }
 
   // Register dialects with this context.
-  new BuiltinDialect(this);
+  getOrCreateDialect<BuiltinDialect>();
   registerAllDialects(this);
 
   // Initialize several common attributes and types to avoid the need to lock
   // the context when accessing them.
 
   //// Types.
   /// Floating-point Types.
   impl->bf16Ty = TypeUniquer::get<FloatType>(this, StandardTypes::BF16);
Dialect *MLIRContext::getRegisteredDialect(StringRef name) {
   auto it = llvm::lower_bound(
       impl->dialects, name,
       [](const auto &lhs, StringRef rhs) { return lhs->getNamespace() < rhs; });
   return (it != impl->dialects.end() && (*it)->getNamespace() == name)
              ? (*it).get()
              : nullptr;
 }
 
-/// Register this dialect object with the specified context.  The context
-/// takes ownership of the heap allocated dialect.
-void Dialect::registerDialect(MLIRContext *context) {
-  auto &impl = context->getImpl();
-  std::unique_ptr<Dialect> dialect(this);
-
+/// Get a dialect for the provided namespace and TypeID: abort the program if a
+/// dialect exist for this namespace with different TypeID. Returns a pointer to
+/// the dialect owned by the context.
+Dialect *
+MLIRContext::getOrCreateDialect(StringRef dialectNamespace, TypeID dialectID,
+                                function_ref<std::unique_ptr<Dialect>()> ctor) {
+  auto &impl = getImpl();
   // Get the correct insertion position sorted by namespace.
-  auto insertPt = llvm::lower_bound(
-      impl.dialects, dialect, [](const auto &lhs, const auto &rhs) {
-        return lhs->getNamespace() < rhs->getNamespace();
+  auto insertPt =
+      llvm::lower_bound(impl.dialects, nullptr,
+                        [&](const std::unique_ptr<Dialect> &lhs,
+                            const std::unique_ptr<Dialect> &rhs) {
+                          if (!lhs)
+                            return dialectNamespace < rhs->getNamespace();
+                          return lhs->getNamespace() < dialectNamespace;
                         });
 
   // Abort if dialect with namespace has already been registered.
   if (insertPt != impl.dialects.end() &&
-      (*insertPt)->getNamespace() == getNamespace()) {
-    llvm::report_fatal_error("a dialect with namespace '" + getNamespace() +
+      (*insertPt)->getNamespace() == dialectNamespace) {
+    if ((*insertPt)->getTypeID() == dialectID)
+      return insertPt->get();
+    llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
                              "' has already been registered");
   }
-  impl.dialects.insert(insertPt, std::move(dialect));
+  auto it = impl.dialects.insert(insertPt, ctor());
+  return &**it;
 }
 
 bool MLIRContext::allowsUnregisteredDialects() {
   return impl->allowUnregisteredDialects;
 }
 
 void MLIRContext::allowUnregisteredDialects(bool allowing) {
   impl->allowUnregisteredDialects = allowing;

mlir/test/lib/Dialect/Test/TestDialect.cpp

struct TestInlinerInterface : public DialectInlinerInterface {
   }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
-TestDialect::TestDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+void TestDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
   addInterfaces<TestOpAsmInterface, TestOpFolderDialectInterface,
                 TestInlinerInterface>();
   addTypes<TestType, TestRecursiveType>();
   allowUnknownOperations();

mlir/tools/mlir-tblgen/DialectGen.cpp

 //===----------------------------------------------------------------------===//
 
 /// The code block for the start of a dialect class declaration.
 ///
 /// {0}: The name of the dialect class.
 /// {1}: The dialect namespace.
 static const char *const dialectDeclBeginStr = R"(
 class {0} : public ::mlir::Dialect {
+  explicit {0}(::mlir::MLIRContext *context)
+    : ::mlir::Dialect(getDialectNamespace(), context,
+      ::mlir::TypeID::get<{0}>()) {{
+    initialize();
+  }
+  void initialize();
+  friend class ::mlir::MLIRContext;
 public:
-  explicit {0}(::mlir::MLIRContext *context);
   static ::llvm::StringRef getDialectNamespace() { return "{1}"; }
 )";
 
 /// The code block for the attribute parser/printer hooks.
 static const char *const attrParserDecl = R"(
   /// Parse an attribute registered to this dialect.
   ::mlir::Attribute parseAttribute(::mlir::DialectAsmParser &parser,
                                    ::mlir::Type type) const override;

mlir/unittests/IR/DialectTest.cpp

 //===- DialectTest.cpp - Dialect unit tests -------------------------------===//
 //
 // 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 "gtest/gtest.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 namespace {
 struct TestDialect : public Dialect {
-  TestDialect(MLIRContext *context) : Dialect(/*name=*/"test", context) {}
+  static StringRef getDialectNamespace() { return "test"; };
+  TestDialect(MLIRContext *context)
+      : Dialect(getDialectNamespace(), context, TypeID::get<TestDialect>()) {}
+};
+struct AnotherTestDialect : public Dialect {
+  static StringRef getDialectNamespace() { return "test"; };
+  AnotherTestDialect(MLIRContext *context)
+      : Dialect(getDialectNamespace(), context,
+                TypeID::get<AnotherTestDialect>()) {}
 };
 
 TEST(DialectDeathTest, MultipleDialectsWithSameNamespace) {
   MLIRContext context;
 
   // Registering a dialect with the same namespace twice should result in a
   // failure.
-  new TestDialect(&context);
-  ASSERT_DEATH(new TestDialect(&context), "");
+  context.getOrCreateDialect<TestDialect>();
+  ASSERT_DEATH(context.getOrCreateDialect<AnotherTestDialect>(), "");
 }
 
 } // end namespace