Refactor OperationName to use virtual tables for dispatch (NFC)

[mehdi_amini]

Summary

This streamlines the implementation and makes it so that the virtual tables are in the binary instead of dynamically assembled during initialization.
The dynamic allocation size of op registration is also smaller with this
change.

[mehdi_amini]

Make OperationName::Impl itself having a vtable instead of having an extra level of indirection

[rriddle]

mlir/unittests/Interfaces/InferTypeOpInterfaceTest.cpp

39 ↗

(On Diff #488285)

Debugging?

[mehdi_amini]

Some more cleanup, moved many APIs from RegisteredOperationName to OperationName since they are just dynamic dispatch and the UnregisteredModel implementation does the right thing.

[mehdi_amini]

mlir/unittests/Interfaces/InferTypeOpInterfaceTest.cpp

39 ↗

(On Diff #488285)

Yes, but I figured it makes sense to keep it in the unit-test: it is mildly better than a bad segfault on the next line.

[ashay-github]

Hi @mehdi_amini, this patch is causing MSVC-compiled mlir-opt to crash at runtime with the following stack trace. Could you revert the patch for now to get the build working again please?

Exception Code: 0xC0000005

 #0 0x00007ff6a30cf153 llvm::function_ref<(class mlir::OpAsmParser &, struct mlir::OperationState &)>::callback_fn<class mlir::ParseResult (__cdecl *)(class mlir::OpAsmParser &, struct mlir::OperationState &)>(__int64, class mlir::OpAsmParser &, struct mlir::OperationState &) llvm-project\llvm\include\llvm\ADT\STLFunctionalExtras.h:45:0

 #1 0x00007ff6a85d82aa llvm::function_ref<(class mlir::OpAsmParser &, struct mlir::OperationState &)>::operator()(class mlir::OpAsmParser &, struct mlir::OperationState &) const llvm-project\llvm\include\llvm\ADT\STLFunctionalExtras.h:68:0

 #2 0x00007ff6a8a27468 `anonymous namespace'::CustomOpAsmParser::parseOperation llvm-project\mlir\lib\AsmParser\Parser.cpp:1439:0

 #3 0x00007ff6a8a207e4 `anonymous namespace'::OperationParser::parseCustomOperation llvm-project\mlir\lib\AsmParser\Parser.cpp:1934:0

 #4 0x00007ff6a8a1ce68 `anonymous namespace'::OperationParser::parseOperation llvm-project\mlir\lib\AsmParser\Parser.cpp:1143:0

 #5 0x00007ff6a8a2b366 `anonymous namespace'::TopLevelOperationParser::parse llvm-project\mlir\lib\AsmParser\Parser.cpp:2611:0

 ...

I'm not sure why the pre-merge checks didn't catch this. Perhaps they're not running, since I don't see build results here.

[mehdi_amini]

Sorry for the breakage! Is Exception Code: 0xC0000005 a segfault?
Any chance you could get a debug build symbolized backtrace here?

[mehdi_amini]

I see you posted a symbolized trace here, sorry I was looking directly at the logs on the bot :)

I think I should be able to re-land this with a fix soon

[mehdi_amini]

Green build after re-land: https://lab.llvm.org/buildbot/#/builders/13/builds/30874

[ashay-github]

No worries, thanks for reverting! As I understand, Exception Code: 0xC0000005 does mean a segfault.

mlir/include/mlir/IR/ExtensibleDialect.h

 //===----------------------------------------------------------------------===//
 
 /// The definition of a dynamic op. A dynamic op is an op that is defined at
 /// runtime, and that can be registered at runtime by an extensible dialect (a
 /// dialect inheriting ExtensibleDialect). This class stores the functions that
 /// are in the OperationName class, and in addition defines the TypeID of the op
 /// that will be defined.
 /// Each dynamic operation definition refers to one instance of this class.
-class DynamicOpDefinition {
+class DynamicOpDefinition : public OperationName::InterfaceConcept {
 public:
   /// Create a new op at runtime. The op is registered only after passing it to
   /// the dialect using registerDynamicOp.
   static std::unique_ptr<DynamicOpDefinition>
   get(StringRef name, ExtensibleDialect *dialect,
       OperationName::VerifyInvariantsFn &&verifyFn,
       OperationName::VerifyRegionInvariantsFn &&verifyRegionFn);
   static std::unique_ptr<DynamicOpDefinition>
public:
   }
 
   /// Set the hook populating default attributes.
   void setPopulateDefaultAttrsFn(
       OperationName::PopulateDefaultAttrsFn &&populateDefaultAttrs) {
     populateDefaultAttrsFn = std::move(populateDefaultAttrs);
   }
 
+  LogicalResult foldHook(Operation *op, ArrayRef<Attribute> attrs,
+                         SmallVectorImpl<OpFoldResult> &results) final {
+    return foldHookFn(op, attrs, results);
+  }
+  void getCanonicalizationPatterns(RewritePatternSet &set,
+                                   MLIRContext *context) final {
+    getCanonicalizationPatternsFn(set, context);
+  }
+  bool hasTrait(TypeID id) final { return false; }
+  OperationName::ParseAssemblyFn getParseAssemblyFn() final { return parseFn; }
+  void populateDefaultAttrs(const RegisteredOperationName &name,
+                            NamedAttrList &attrs) final {
+    populateDefaultAttrsFn(name, attrs);
+  }
+  void printAssembly(Operation *op, OpAsmPrinter &printer,
+                     StringRef name) final {
+    printFn(op, printer, name);
+  }
+  LogicalResult verifyInvariants(Operation *op) final { return verifyFn(op); }
+  LogicalResult verifyRegionInvariants(Operation *op) final {
+    return verifyRegionFn(op);
+  }
+  detail::InterfaceMap getInterfaceMap() final {
+    return detail::InterfaceMap::get<>();
+  }
+
 private:
   DynamicOpDefinition(
       StringRef name, ExtensibleDialect *dialect,
       OperationName::VerifyInvariantsFn &&verifyFn,
       OperationName::VerifyRegionInvariantsFn &&verifyRegionFn,
       OperationName::ParseAssemblyFn &&parseFn,
       OperationName::PrintAssemblyFn &&printFn,
       OperationName::FoldHookFn &&foldHookFn,
private:
   llvm::StringMap<DynamicTypeDefinition *> nameToDynTypes;
 
   /// The set of all dynamic attributes registered.
   DenseMap<TypeID, std::unique_ptr<DynamicAttrDefinition>> dynAttrs;
 
   /// This structure allows to get in O(1) a dynamic attribute given its name.
   llvm::StringMap<DynamicAttrDefinition *> nameToDynAttrs;
 
+  /// All the dynamic operations registered.
+  std::vector<std::unique_ptr<DynamicOpDefinition>> dynamicOps;
+
   /// Give DynamicOpDefinition access to allocateTypeID.
   friend DynamicOpDefinition;
 
   /// Allocates a type ID to uniquify operations.
   TypeID allocateTypeID() { return typeIDAllocator.allocate(); }
 
   /// Owns the TypeID generated at runtime for operations.
   TypeIDAllocator typeIDAllocator;

mlir/include/mlir/IR/OpDefinition.h

private:
   getGetCanonicalizationPatternsFn() {
     return &ConcreteType::getCanonicalizationPatterns;
   }
   /// Implementation of `GetHasTraitFn`
   static OperationName::HasTraitFn getHasTraitFn() {
     return
         [](TypeID id) { return op_definition_impl::hasTrait<Traits...>(id); };
   }
-  /// Implementation of `ParseAssemblyFn` OperationName hook.
-  static OperationName::ParseAssemblyFn getParseAssemblyFn() {
-    return &ConcreteType::parse;
-  }
   /// Implementation of `PrintAssemblyFn` OperationName hook.
   static OperationName::PrintAssemblyFn getPrintAssemblyFn() {
     if constexpr (detect_has_print<ConcreteType>::value)
       return [](Operation *op, OpAsmPrinter &p, StringRef defaultDialect) {
         OpState::printOpName(op, p, defaultDialect);
         return cast<ConcreteType>(op).print(p);
       };
     return [](Operation *op, OpAsmPrinter &printer, StringRef defaultDialect) {

mlir/include/mlir/IR/OperationSupport.h

 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/TypeRange.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Support/InterfaceSupport.h"
 #include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/PointerUnion.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include "llvm/Support/TrailingObjects.h"
 #include <memory>
 
 namespace llvm {
 class BitVector;
 } // namespace llvm
 

 class OperationName {
 public:
   using GetCanonicalizationPatternsFn =
       llvm::unique_function<void(RewritePatternSet &, MLIRContext *) const>;
   using FoldHookFn = llvm::unique_function<LogicalResult(
       Operation *, ArrayRef<Attribute>, SmallVectorImpl<OpFoldResult> &) const>;
   using HasTraitFn = llvm::unique_function<bool(TypeID) const>;
   using ParseAssemblyFn =
-      llvm::unique_function<ParseResult(OpAsmParser &, OperationState &) const>;
+      llvm::function_ref<ParseResult(OpAsmParser &, OperationState &)>;
   // Note: RegisteredOperationName is passed as reference here as the derived
   // class is defined below.
   using PopulateDefaultAttrsFn = llvm::unique_function<void(
       const RegisteredOperationName &, NamedAttrList &) const>;
   using PrintAssemblyFn =
       llvm::unique_function<void(Operation *, OpAsmPrinter &, StringRef) const>;
   using VerifyInvariantsFn =
       llvm::unique_function<LogicalResult(Operation *) const>;
   using VerifyRegionInvariantsFn =
       llvm::unique_function<LogicalResult(Operation *) const>;
 
-protected:
   /// This class represents a type erased version of an operation. It contains
   /// all of the components necessary for opaquely interacting with an
   /// operation. If the operation is not registered, some of these components
   /// may not be populated.
+  struct InterfaceConcept {
+    virtual ~InterfaceConcept() = default;
+    virtual LogicalResult foldHook(Operation *, ArrayRef<Attribute>,
+                                   SmallVectorImpl<OpFoldResult> &) = 0;
+    virtual void getCanonicalizationPatterns(RewritePatternSet &,
+                                             MLIRContext *) = 0;
+    virtual bool hasTrait(TypeID) = 0;
+    virtual OperationName::ParseAssemblyFn getParseAssemblyFn() = 0;
+    virtual void populateDefaultAttrs(const RegisteredOperationName &,
+                                      NamedAttrList &) = 0;
+    virtual void printAssembly(Operation *, OpAsmPrinter &, StringRef) = 0;
+    virtual LogicalResult verifyInvariants(Operation *) = 0;
+    virtual LogicalResult verifyRegionInvariants(Operation *) = 0;
+    virtual detail::InterfaceMap getInterfaceMap() = 0;
+  };
+
+  /// Default implementation for unregistered operations.
+  struct UnregisteredOpModel : InterfaceConcept {
+    LogicalResult foldHook(Operation *, ArrayRef<Attribute>,
+                           SmallVectorImpl<OpFoldResult> &) final;
+    void getCanonicalizationPatterns(RewritePatternSet &, MLIRContext *) final;
+    bool hasTrait(TypeID) final;
+    virtual OperationName::ParseAssemblyFn getParseAssemblyFn() final;
+    void populateDefaultAttrs(const RegisteredOperationName &,
+                              NamedAttrList &) final;
+    void printAssembly(Operation *, OpAsmPrinter &, StringRef) final;
+    LogicalResult verifyInvariants(Operation *) final;
+    LogicalResult verifyRegionInvariants(Operation *) final;
+    detail::InterfaceMap getInterfaceMap() final;
+  };
+
+protected:
   struct Impl {
-    Impl(StringAttr name)
-        : name(name), dialect(nullptr), interfaceMap(std::nullopt) {}
+    Impl(StringAttr name, InterfaceConcept *model)
+        : name(name), dialect(nullptr), interfaceMap(std::nullopt),
+          model(model) {}
 
     /// The name of the operation.
     StringAttr name;
 
     //===------------------------------------------------------------------===//
     // Registered Operation Info
 
     /// The following fields are only populated when the operation is
     /// registered.
 
     /// Returns true if the operation has been registered, i.e. if the
     /// registration info has been populated.
     bool isRegistered() const { return dialect; }
 
     /// This is the dialect that this operation belongs to.
     Dialect *dialect;
 
     /// The unique identifier of the derived Op class.
     TypeID typeID;
 
     /// A map of interfaces that were registered to this operation.
     detail::InterfaceMap interfaceMap;
 
     /// Internal callback hooks provided by the op implementation.
-    FoldHookFn foldHookFn;
-    GetCanonicalizationPatternsFn getCanonicalizationPatternsFn;
-    HasTraitFn hasTraitFn;
-    ParseAssemblyFn parseAssemblyFn;
-    PopulateDefaultAttrsFn populateDefaultAttrsFn;
-    PrintAssemblyFn printAssemblyFn;
-    VerifyInvariantsFn verifyInvariantsFn;
-    VerifyRegionInvariantsFn verifyRegionInvariantsFn;
+    InterfaceConcept *model;
 
     /// A list of attribute names registered to this operation in StringAttr
     /// form. This allows for operation classes to use StringAttr for attribute
     /// lookup/creation/etc., as opposed to raw strings.
     ArrayRef<StringAttr> attributeNames;
   };
 
 public:
   OperationName(StringRef name, MLIRContext *context);
 
   /// Return if this operation is registered.
-  bool isRegistered() const { return impl->isRegistered(); }
+  bool isRegistered() const { return getImpl()->isRegistered(); }
 
   /// If this operation is registered, returns the registered information,
   /// std::nullopt otherwise.
   Optional<RegisteredOperationName> getRegisteredInfo() const;
 
   /// Returns true if the operation was registered with a particular trait, e.g.
   /// hasTrait<OperandsAreSignlessIntegerLike>(). Returns false if the operation
   /// is unregistered.
   template <template <typename T> class Trait>
   bool hasTrait() const {
     return hasTrait(TypeID::get<Trait>());
   }
-  bool hasTrait(TypeID traitID) const {
-    return isRegistered() && impl->hasTraitFn(traitID);
-  }
+  bool hasTrait(TypeID traitID) const { return getModel()->hasTrait(traitID); }
 
   /// Returns true if the operation *might* have the provided trait. This
   /// means that either the operation is unregistered, or it was registered with
   /// the provide trait.
   template <template <typename T> class Trait>
   bool mightHaveTrait() const {
     return mightHaveTrait(TypeID::get<Trait>());
   }
   bool mightHaveTrait(TypeID traitID) const {
-    return !isRegistered() || impl->hasTraitFn(traitID);
+    return !isRegistered() || getModel()->hasTrait(traitID);
   }
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this operation, null otherwise. This should not be used
   /// directly.
   template <typename T>
   typename T::Concept *getInterface() const {
-    return impl->interfaceMap.lookup<T>();
+    return getImpl()->interfaceMap.lookup<T>();
   }
 
   /// Returns true if this operation has the given interface registered to it.
   template <typename T>
   bool hasInterface() const {
     return hasInterface(TypeID::get<T>());
   }
   bool hasInterface(TypeID interfaceID) const {
-    return impl->interfaceMap.contains(interfaceID);
+    return getImpl()->interfaceMap.contains(interfaceID);
   }
 
   /// Returns true if the operation *might* have the provided interface. This
   /// means that either the operation is unregistered, or it was registered with
   /// the provide interface.
   template <typename T>
   bool mightHaveInterface() const {
     return mightHaveInterface(TypeID::get<T>());
   }
   bool mightHaveInterface(TypeID interfaceID) const {
     return !isRegistered() || hasInterface(interfaceID);
   }
 
   /// Return the dialect this operation is registered to if the dialect is
   /// loaded in the context, or nullptr if the dialect isn't loaded.
   Dialect *getDialect() const {
-    return isRegistered() ? impl->dialect : impl->name.getReferencedDialect();
+    return isRegistered() ? getImpl()->dialect
+                          : getImpl()->name.getReferencedDialect();
   }
 
   /// Return the name of the dialect this operation is registered to.
   StringRef getDialectNamespace() const;
 
   /// Return the operation name with dialect name stripped, if it has one.
   StringRef stripDialect() const { return getStringRef().split('.').second; }
 
   /// Return the name of this operation. This always succeeds.
   StringRef getStringRef() const { return getIdentifier(); }
 
   /// Return the name of this operation as a StringAttr.
-  StringAttr getIdentifier() const { return impl->name; }
+  StringAttr getIdentifier() const { return getImpl()->name; }
 
   void print(raw_ostream &os) const;
   void dump() const;
 
   /// Represent the operation name as an opaque pointer. (Used to support
   /// PointerLikeTypeTraits).
   void *getAsOpaquePointer() const { return const_cast<Impl *>(impl); }
   static OperationName getFromOpaquePointer(const void *pointer) {
     return OperationName(
         const_cast<Impl *>(reinterpret_cast<const Impl *>(pointer)));
   }
 
   bool operator==(const OperationName &rhs) const { return impl == rhs.impl; }
   bool operator!=(const OperationName &rhs) const { return !(*this == rhs); }
 
 protected:
-  OperationName(Impl *impl) : impl(impl) {}
+  OperationName(Impl *impl) : impl(impl), model(impl ? impl->model : nullptr) {}
+  Impl *getImpl() const { return impl; }
+  InterfaceConcept *getModel() const {
+    assert(impl->model == model && "Inconsistency in model caching");
+    return model;
+  }
+  void setImpl(Impl *rhs) { impl = rhs; }
+  void setModel(InterfaceConcept *rhs) { model = rhs; }
 
+private:
   /// The internal implementation of the operation name.
-  Impl *impl;
+  Impl *impl = nullptr;
+
+  /// Internal callback hooks provided by the op implementation.
+  /// This is set on the impl, but we cache it here as it'll be access often.
+  InterfaceConcept *model = nullptr;
 
   /// Allow access to the Impl struct.
   friend MLIRContextImpl;
+  friend DenseMapInfo<mlir::OperationName>;
+  friend DenseMapInfo<mlir::RegisteredOperationName>;
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, OperationName info) {
   info.print(os);
   return os;
 }
 
 // Make operation names hashable.
 inline llvm::hash_code hash_value(OperationName arg) {
   return llvm::hash_value(arg.getAsOpaquePointer());
 }
 
 //===----------------------------------------------------------------------===//
 // RegisteredOperationName
 //===----------------------------------------------------------------------===//
 
 /// This is a "type erased" representation of a registered operation. This
 /// should only be used by things like the AsmPrinter and other things that need
 /// to be parameterized by generic operation hooks. Most user code should use
 /// the concrete operation types.
 class RegisteredOperationName : public OperationName {
 public:
+  /// Implementation of the InterfaceConcept for operation APIs that forwarded
+  /// to a concrete op implementation.
+  template <typename ConcreteOp> struct Model : public InterfaceConcept {
+    LogicalResult foldHook(Operation *op, ArrayRef<Attribute> attrs,
+                           SmallVectorImpl<OpFoldResult> &results) final {
+      return ConcreteOp::getFoldHookFn()(op, attrs, results);
+    }
+    void getCanonicalizationPatterns(RewritePatternSet &set,
+                                     MLIRContext *context) final {
+      ConcreteOp::getGetCanonicalizationPatternsFn()(set, context);
+    }
+    bool hasTrait(TypeID id) final { return ConcreteOp::getHasTraitFn()(id); }
+    OperationName::ParseAssemblyFn getParseAssemblyFn() final {
+      return ConcreteOp::parse;
+    }
+    void populateDefaultAttrs(const RegisteredOperationName &name,
+                              NamedAttrList &attrs) final {
+      ConcreteOp::populateDefaultAttrs(name, attrs);
+    }
+    void printAssembly(Operation *op, OpAsmPrinter &printer,
+                       StringRef name) final {
+      ConcreteOp::getPrintAssemblyFn()(op, printer, name);
+    }
+    LogicalResult verifyInvariants(Operation *op) final {
+      return ConcreteOp::getVerifyInvariantsFn()(op);
+    }
+    LogicalResult verifyRegionInvariants(Operation *op) final {
+      return ConcreteOp::getVerifyRegionInvariantsFn()(op);
+    }
+    detail::InterfaceMap getInterfaceMap() final {
+      return ConcreteOp::getInterfaceMap();
+    }
+  };
+
   /// Lookup the registered operation information for the given operation.
   /// Returns std::nullopt if the operation isn't registered.
   static Optional<RegisteredOperationName> lookup(StringRef name,
                                                   MLIRContext *ctx);
 
   /// Register a new operation in a Dialect object.
-  /// This constructor is used by Dialect objects when they register the list of
-  /// operations they contain.
-  template <typename T>
-  static void insert(Dialect &dialect) {
-    insert(T::getOperationName(), dialect, TypeID::get<T>(),
-           T::getParseAssemblyFn(), T::getPrintAssemblyFn(),
-           T::getVerifyInvariantsFn(), T::getVerifyRegionInvariantsFn(),
-           T::getFoldHookFn(), T::getGetCanonicalizationPatternsFn(),
-           T::getInterfaceMap(), T::getHasTraitFn(), T::getAttributeNames(),
-           T::getPopulateDefaultAttrsFn());
+  /// This constructor is used by Dialect objects when they register the list
+  /// of operations they contain.
+  template <typename T> static void insert(Dialect &dialect) {
+    static Model<T> model;
+    insert(T::getOperationName(), dialect, TypeID::get<T>(), model,
+           T::getAttributeNames());
   }
   /// The use of this method is in general discouraged in favor of
   /// 'insert<CustomOp>(dialect)'.
-  static void
-  insert(StringRef name, Dialect &dialect, TypeID typeID,
-         ParseAssemblyFn &&parseAssembly, PrintAssemblyFn &&printAssembly,
-         VerifyInvariantsFn &&verifyInvariants,
-         VerifyRegionInvariantsFn &&verifyRegionInvariants,
-         FoldHookFn &&foldHook,
-         GetCanonicalizationPatternsFn &&getCanonicalizationPatterns,
-         detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait,
-         ArrayRef<StringRef> attrNames,
-         PopulateDefaultAttrsFn &&populateDefaultAttrs);
+  static void insert(StringRef name, Dialect &dialect, TypeID typeID,
+                     InterfaceConcept &model, ArrayRef<StringRef> attrNames);
 
   /// Return the dialect this operation is registered to.
-  Dialect &getDialect() const { return *impl->dialect; }
+  Dialect &getDialect() const { return *getImpl()->dialect; }
 
   /// Return the unique identifier of the derived Op class.
-  TypeID getTypeID() const { return impl->typeID; }
+  TypeID getTypeID() const { return getImpl()->typeID; }
 
   /// Use the specified object to parse this ops custom assembly format.
   ParseResult parseAssembly(OpAsmParser &parser, OperationState &result) const;
 
   /// Return the static hook for parsing this operation assembly.
-  const ParseAssemblyFn &getParseAssemblyFn() const {
-    return impl->parseAssemblyFn;
+  ParseAssemblyFn getParseAssemblyFn() const {
+    return getModel()->getParseAssemblyFn();
   }
 
   /// This hook implements the AsmPrinter for this operation.
   void printAssembly(Operation *op, OpAsmPrinter &p,
                      StringRef defaultDialect) const {
-    return impl->printAssemblyFn(op, p, defaultDialect);
+    return getModel()->printAssembly(op, p, defaultDialect);
   }
 
   /// These hooks implement the verifiers for this operation.  It should emits
-  /// an error message and returns failure if a problem is detected, or returns
-  /// success if everything is ok.
+  /// an error message and returns failure if a problem is detected, or
+  /// returns success if everything is ok.
   LogicalResult verifyInvariants(Operation *op) const {
-    return impl->verifyInvariantsFn(op);
+    return getModel()->verifyInvariants(op);
   }
   LogicalResult verifyRegionInvariants(Operation *op) const {
-    return impl->verifyRegionInvariantsFn(op);
+    return getModel()->verifyRegionInvariants(op);
   }
 
-  /// This hook implements a generalized folder for this operation.  Operations
+  /// This hook implements a generalized folder for this operation. Operations
   /// can implement this to provide simplifications rules that are applied by
   /// the Builder::createOrFold API and the canonicalization pass.
   ///
-  /// This is an intentionally limited interface - implementations of this hook
-  /// can only perform the following changes to the operation:
+  /// This is an intentionally limited interface - implementations of this
+  /// hook can only perform the following changes to the operation:
   ///
   ///  1. They can leave the operation alone and without changing the IR, and
   ///     return failure.
-  ///  2. They can mutate the operation in place, without changing anything else
+  ///  2. They can mutate the operation in place, without changing anything
+  ///  else
   ///     in the IR.  In this case, return success.
-  ///  3. They can return a list of existing values that can be used instead of
+  ///  3. They can return a list of existing values that can be used instead
+  ///  of
   ///     the operation.  In this case, fill in the results list and return
   ///     success.  The caller will remove the operation and use those results
   ///     instead.
   ///
   /// This allows expression of some simple in-place canonicalizations (e.g.
   /// "x+0 -> x", "min(x,y,x,z) -> min(x,y,z)", "x+y-x -> y", etc), as well as
   /// generalized constant folding.
   LogicalResult foldHook(Operation *op, ArrayRef<Attribute> operands,
                          SmallVectorImpl<OpFoldResult> &results) const {
-    return impl->foldHookFn(op, operands, results);
+    return getModel()->foldHook(op, operands, results);
   }
 
-  /// This hook returns any canonicalization pattern rewrites that the operation
-  /// supports, for use by the canonicalization pass.
+  /// This hook returns any canonicalization pattern rewrites that the
+  /// operation supports, for use by the canonicalization pass.
   void getCanonicalizationPatterns(RewritePatternSet &results,
                                    MLIRContext *context) const {
-    return impl->getCanonicalizationPatternsFn(results, context);
+    return getModel()->getCanonicalizationPatterns(results, context);
   }
 
-  /// Attach the given models as implementations of the corresponding interfaces
-  /// for the concrete operation.
-  template <typename... Models>
-  void attachInterface() {
-    impl->interfaceMap.insert<Models...>();
+  /// Attach the given models as implementations of the corresponding
+  /// interfaces for the concrete operation.
+  template <typename... Models> void attachInterface() {
+    getImpl()->interfaceMap.insert<Models...>();
   }
 
   /// Returns true if the operation has a particular trait.
-  template <template <typename T> class Trait>
-  bool hasTrait() const {
+  template <template <typename T> class Trait> bool hasTrait() const {
     return hasTrait(TypeID::get<Trait>());
   }
 
   /// Returns true if the operation has a particular trait.
-  bool hasTrait(TypeID traitID) const { return impl->hasTraitFn(traitID); }
+  bool hasTrait(TypeID traitID) const { return getModel()->hasTrait(traitID); }
 
   /// Return the list of cached attribute names registered to this operation.
   /// The order of attributes cached here is unique to each type of operation,
   /// and the interpretation of this attribute list should generally be driven
-  /// by the respective operation. In many cases, this caching removes the need
-  /// to use the raw string name of a known attribute.
+  /// by the respective operation. In many cases, this caching removes the
+  /// need to use the raw string name of a known attribute.
   ///
   /// For example the ODS generator, with an op defining the following
   /// attributes:
   ///
   ///   let arguments = (ins I32Attr:$attr1, I32Attr:$attr2);
   ///
   /// ... may produce an order here of ["attr1", "attr2"]. This allows for the
   /// ODS generator to directly access the cached name for a known attribute,
-  /// greatly simplifying the cost and complexity of attribute usage produced by
-  /// the generator.
+  /// greatly simplifying the cost and complexity of attribute usage produced
+  /// by the generator.
   ///
   ArrayRef<StringAttr> getAttributeNames() const {
-    return impl->attributeNames;
+    return getImpl()->attributeNames;
   }
 
   /// This hook implements the method to populate defaults attributes that are
   /// unset.
   void populateDefaultAttrs(NamedAttrList &attrs) const;
 
   /// Represent the operation name as an opaque pointer. (Used to support
   /// PointerLikeTypeTraits).

 };
 
 } // namespace mlir
 
 namespace llvm {
 template <>
 struct DenseMapInfo<mlir::OperationName> {
   static mlir::OperationName getEmptyKey() {
-    void *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
-    return mlir::OperationName::getFromOpaquePointer(pointer);
+    auto name = mlir::OperationName::getFromOpaquePointer(nullptr);
+    name.impl = static_cast<mlir::OperationName::Impl *>(
+        llvm::DenseMapInfo<void *>::getEmptyKey());
+    return name;
   }
   static mlir::OperationName getTombstoneKey() {
-    void *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
-    return mlir::OperationName::getFromOpaquePointer(pointer);
+    auto name = mlir::OperationName::getFromOpaquePointer(nullptr);
+    name.impl = static_cast<mlir::OperationName::Impl *>(
+        llvm::DenseMapInfo<void *>::getTombstoneKey());
+    return name;
   }
   static unsigned getHashValue(mlir::OperationName val) {
     return DenseMapInfo<void *>::getHashValue(val.getAsOpaquePointer());
   }
   static bool isEqual(mlir::OperationName lhs, mlir::OperationName rhs) {
     return lhs == rhs;
   }
 };
 template <>
 struct DenseMapInfo<mlir::RegisteredOperationName>
     : public DenseMapInfo<mlir::OperationName> {
   static mlir::RegisteredOperationName getEmptyKey() {
-    void *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
-    return mlir::RegisteredOperationName::getFromOpaquePointer(pointer);
+    auto name = mlir::RegisteredOperationName::getFromOpaquePointer(nullptr);
+    name.impl = static_cast<mlir::OperationName::Impl *>(
+        llvm::DenseMapInfo<void *>::getEmptyKey());
+    return name;
   }
   static mlir::RegisteredOperationName getTombstoneKey() {
-    void *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
-    return mlir::RegisteredOperationName::getFromOpaquePointer(pointer);
+    auto name = mlir::RegisteredOperationName::getFromOpaquePointer(nullptr);
+    name.impl = static_cast<mlir::OperationName::Impl *>(
+        llvm::DenseMapInfo<void *>::getTombstoneKey());
+    return name;
   }
 };
 
 template <>
 struct PointerLikeTypeTraits<mlir::OperationName> {
   static inline void *getAsVoidPointer(mlir::OperationName I) {
     return const_cast<void *>(I.getAsOpaquePointer());
   }

mlir/lib/AsmParser/Parser.cpp

Operation *OperationParser::parseGenericOperation() {
 
   OperationState result(srcLocation, name);
   CleanupOpStateRegions guard{result};
 
   // Lazy load dialects in the context as needed.
   if (!result.name.isRegistered()) {
     StringRef dialectName = StringRef(name).split('.').first;
     if (!getContext()->getLoadedDialect(dialectName) &&
-        !getContext()->getOrLoadDialect(dialectName) &&
-        !getContext()->allowsUnregisteredDialects()) {
+        !getContext()->getOrLoadDialect(dialectName)) {
+      if (!getContext()->allowsUnregisteredDialects()) {
         // Emit an error if the dialect couldn't be loaded (i.e., it was not
         // registered) and unregistered dialects aren't allowed.
         emitError("operation being parsed with an unregistered dialect. If "
                   "this is intended, please use -allow-unregistered-dialect "
                   "with the MLIR tool used");
         return nullptr;
       }
+    } else {
+      // Reload the OperationName now that the dialect is loaded.
+      result.name = OperationName(name, getContext());
+    }
   }
 
   // If we are populating the parser state, start a new operation definition.
   if (state.asmState)
     state.asmState->startOperationDefinition(result.name);
 
   if (parseGenericOperationAfterOpName(result))
     return nullptr;

mlir/lib/IR/ExtensibleDialect.cpp

void ExtensibleDialect::registerDynamicAttr(
   addAttribute(typeID, std::move(abstractAttr));
   attrPtr->registerInAttrUniquer();
 }
 
 void ExtensibleDialect::registerDynamicOp(
     std::unique_ptr<DynamicOpDefinition> &&op) {
   assert(op->dialect == this &&
          "trying to register a dynamic op in the wrong dialect");
-  auto hasTraitFn = [](TypeID traitId) { return false; };
-
-  RegisteredOperationName::insert(
-      op->name, *op->dialect, op->typeID, std::move(op->parseFn),
-      std::move(op->printFn), std::move(op->verifyFn),
-      std::move(op->verifyRegionFn), std::move(op->foldHookFn),
-      std::move(op->getCanonicalizationPatternsFn),
-      detail::InterfaceMap::get<>(), std::move(hasTraitFn), {},
-      std::move(op->populateDefaultAttrsFn));
+  RegisteredOperationName::insert(op->name, *op->dialect, op->typeID, *op,
+                                  /*attrNames=*/{});
+  dynamicOps.push_back(std::move(op));
 }
 
 bool ExtensibleDialect::classof(const Dialect *dialect) {
   return const_cast<Dialect *>(dialect)
       ->getRegisteredInterface<IsExtensibleDialect>();
 }
 
 OptionalParseResult ExtensibleDialect::parseOptionalDynamicType(

mlir/lib/IR/MLIRContext.cpp

OperationName::OperationName(StringRef name, MLIRContext *context) {
   bool isMultithreadingEnabled = context->isMultithreadingEnabled();
   if (isMultithreadingEnabled) {
     // Check the registered info map first. In the overwhelmingly common case,
     // the entry will be in here and it also removes the need to acquire any
     // locks.
     auto registeredIt = ctxImpl.registeredOperations.find(name);
     if (LLVM_LIKELY(registeredIt != ctxImpl.registeredOperations.end())) {
       impl = registeredIt->second.impl;
+      model = impl->model;
       return;
     }
 
     llvm::sys::SmartScopedReader<true> contextLock(ctxImpl.operationInfoMutex);
     auto it = ctxImpl.operations.find(name);
     if (it != ctxImpl.operations.end()) {
       impl = &it->second;
+      model = impl->model;
       return;
     }
   }
 
   // Acquire a writer-lock so that we can safely create the new instance.
   ScopedWriterLock lock(ctxImpl.operationInfoMutex, isMultithreadingEnabled);
 
-  auto it = ctxImpl.operations.insert({name, OperationName::Impl(nullptr)});
+  static UnregisteredOpModel unregisteredModel;
+  auto it = ctxImpl.operations.insert(
+      {name, OperationName::Impl(nullptr, &unregisteredModel)});
   if (it.second)
     it.first->second.name = StringAttr::get(context, name);
   impl = &it.first->second;
+  model = impl->model;
 }
 
 StringRef OperationName::getDialectNamespace() const {
   if (Dialect *dialect = getDialect())
     return dialect->getNamespace();
   return getStringRef().split('.').first;
 }
 
+LogicalResult
+OperationName::UnregisteredOpModel::foldHook(Operation *, ArrayRef<Attribute>,
+                                             SmallVectorImpl<OpFoldResult> &) {
+  return failure();
+}
+void OperationName::UnregisteredOpModel::getCanonicalizationPatterns(
+    RewritePatternSet &, MLIRContext *) {}
+bool OperationName::UnregisteredOpModel::hasTrait(TypeID) { return false; }
+
+OperationName::ParseAssemblyFn
+OperationName::UnregisteredOpModel::getParseAssemblyFn() {
+  llvm::report_fatal_error("getParseAssemblyFn hook called on unregistered op");
+}
+void OperationName::UnregisteredOpModel::populateDefaultAttrs(
+    const RegisteredOperationName &, NamedAttrList &) {}
+void OperationName::UnregisteredOpModel::printAssembly(Operation *,
+                                                       OpAsmPrinter &,
+                                                       StringRef) {
+  llvm::report_fatal_error("printAssembly hook called on unregistered op");
+}
+LogicalResult
+OperationName::UnregisteredOpModel::verifyInvariants(Operation *) {
+  return success();
+}
+LogicalResult
+OperationName::UnregisteredOpModel::verifyRegionInvariants(Operation *) {
+  return success();
+}
+detail::InterfaceMap OperationName::UnregisteredOpModel::getInterfaceMap() {
+  return detail::InterfaceMap::get<>();
+}
+
 //===----------------------------------------------------------------------===//
 // RegisteredOperationName
 //===----------------------------------------------------------------------===//
 
 Optional<RegisteredOperationName>
 RegisteredOperationName::lookup(StringRef name, MLIRContext *ctx) {
   auto &impl = ctx->getImpl();
   auto it = impl.registeredOperations.find(name);
   if (it != impl.registeredOperations.end())
     return it->getValue();
   return std::nullopt;
 }
 
-ParseResult
-RegisteredOperationName::parseAssembly(OpAsmParser &parser,
-                                       OperationState &result) const {
-  return impl->parseAssemblyFn(parser, result);
-}
-
 void RegisteredOperationName::populateDefaultAttrs(NamedAttrList &attrs) const {
-  impl->populateDefaultAttrsFn(*this, attrs);
+  getModel()->populateDefaultAttrs(*this, attrs);
 }
 
-void RegisteredOperationName::insert(
-    StringRef name, Dialect &dialect, TypeID typeID,
-    ParseAssemblyFn &&parseAssembly, PrintAssemblyFn &&printAssembly,
-    VerifyInvariantsFn &&verifyInvariants,
-    VerifyRegionInvariantsFn &&verifyRegionInvariants, FoldHookFn &&foldHook,
-    GetCanonicalizationPatternsFn &&getCanonicalizationPatterns,
-    detail::InterfaceMap &&interfaceMap, HasTraitFn &&hasTrait,
-    ArrayRef<StringRef> attrNames,
-    PopulateDefaultAttrsFn &&populateDefaultAttrs) {
+void RegisteredOperationName::insert(StringRef name, Dialect &dialect,
+                                     TypeID typeID, InterfaceConcept &model,
+
+                                     ArrayRef<StringRef> attrNames) {
   MLIRContext *ctx = dialect.getContext();
   auto &ctxImpl = ctx->getImpl();
   assert(ctxImpl.multiThreadedExecutionContext == 0 &&
          "registering a new operation kind while in a multi-threaded execution "
          "context");
 
   // Register the attribute names of this operation.
   MutableArrayRef<StringAttr> cachedAttrNames;
   if (!attrNames.empty()) {
     cachedAttrNames = MutableArrayRef<StringAttr>(
         ctxImpl.abstractDialectSymbolAllocator.Allocate<StringAttr>(
             attrNames.size()),
         attrNames.size());
     for (unsigned i : llvm::seq<unsigned>(0, attrNames.size()))
       new (&cachedAttrNames[i]) StringAttr(StringAttr::get(ctx, attrNames[i]));
   }
 
   // Insert the operation info if it doesn't exist yet.
-  auto it = ctxImpl.operations.insert({name, OperationName::Impl(nullptr)});
-  if (it.second)
+  auto it =
+      ctxImpl.operations.insert({name, OperationName::Impl(nullptr, &model)});
+  if (it.second) {
     it.first->second.name = StringAttr::get(ctx, name);
+  }
   OperationName::Impl &impl = it.first->second;
 
   if (impl.isRegistered()) {
     llvm::errs() << "error: operation named '" << name
                  << "' is already registered.\n";
     abort();
   }
+
+  // Update the registered info for this operation.
+  impl.dialect = &dialect;
+  impl.typeID = typeID;
+  impl.interfaceMap = model.getInterfaceMap();
+  impl.model = &model;
+  impl.attributeNames = cachedAttrNames;
+
   auto emplaced = ctxImpl.registeredOperations.try_emplace(
       name, RegisteredOperationName(&impl));
   assert(emplaced.second && "operation name registration must be successful");
 
   // Add emplaced operation name to the sorted operations container.
   RegisteredOperationName &value = emplaced.first->getValue();
   ctxImpl.sortedRegisteredOperations.insert(
       llvm::upper_bound(ctxImpl.sortedRegisteredOperations, value,
                         [](auto &lhs, auto &rhs) {
                           return lhs.getIdentifier().compare(
                               rhs.getIdentifier());
                         }),
       value);
-
-  // Update the registered info for this operation.
-  impl.dialect = &dialect;
-  impl.typeID = typeID;
-  impl.interfaceMap = std::move(interfaceMap);
-  impl.foldHookFn = std::move(foldHook);
-  impl.getCanonicalizationPatternsFn = std::move(getCanonicalizationPatterns);
-  impl.hasTraitFn = std::move(hasTrait);
-  impl.parseAssemblyFn = std::move(parseAssembly);
-  impl.printAssemblyFn = std::move(printAssembly);
-  impl.verifyInvariantsFn = std::move(verifyInvariants);
-  impl.verifyRegionInvariantsFn = std::move(verifyRegionInvariants);
-  impl.attributeNames = cachedAttrNames;
-  impl.populateDefaultAttrsFn = std::move(populateDefaultAttrs);
 }
 
 //===----------------------------------------------------------------------===//
 // AbstractType
 //===----------------------------------------------------------------------===//
 
 const AbstractType &AbstractType::lookup(TypeID typeID, MLIRContext *context) {
   const AbstractType *type = lookupMutable(typeID, context);