[mlir] Refactor AbstractOperation and OperationName

The current implementation is quite clunky; OperationName stores either an Identifier
or an AbstractOperation that corresponds to an operation. This has several problems:

* OperationNames created before and after an operation are registered are different
* Accessing the identifier name/dialect/etc. from an OperationName are overly branchy
  - they need to dyn_cast a PointerUnion to check the state

This commit refactors this such that we create a single information struct for every
operation name, even operations that aren't registered yet. When an OperationName is
created for an unregistered operation, we only populate the name field. When the
operation is registered, we populate the remaining fields. With this we now have two
new classes: OperationName and RegisteredOperationName. These both point to the
same underlying operation information struct, but only RegisteredOperationName can
assume that the operation is actually registered. This leads to a much cleaner API, and
we can also move some AbstractOperation functionality directly to OperationName.

Differential Revision: https://reviews.llvm.org/D114049

Author: River707

mlir/include/mlir/IR/Builders.h

@@ -408,8 +408,8 @@ class OpBuilder : public Builder {
 
 private:
   /// Helper for sanity checking preconditions for create* methods below.
-  void checkHasAbstractOperation(const OperationName &name) {
-    if (LLVM_UNLIKELY(!name.getAbstractOperation()))
+  void checkHasRegisteredInfo(const OperationName &name) {
+    if (LLVM_UNLIKELY(!name.isRegistered()))
       llvm::report_fatal_error(
           "Building op `" + name.getStringRef() +
           "` but it isn't registered in this MLIRContext: the dialect may not "
@@ -423,7 +423,7 @@ class OpBuilder : public Builder {
   template <typename OpTy, typename... Args>
   OpTy create(Location location, Args &&...args) {
     OperationState state(location, OpTy::getOperationName());
-    checkHasAbstractOperation(state.name);
+    checkHasRegisteredInfo(state.name);
     OpTy::build(*this, state, std::forward<Args>(args)...);
     auto *op = createOperation(state);
     auto result = dyn_cast<OpTy>(op);
@@ -440,7 +440,7 @@ class OpBuilder : public Builder {
     // Create the operation without using 'createOperation' as we don't want to
     // insert it yet.
     OperationState state(location, OpTy::getOperationName());
-    checkHasAbstractOperation(state.name);
+    checkHasRegisteredInfo(state.name);
     OpTy::build(*this, state, std::forward<Args>(args)...);
     Operation *op = Operation::create(state);
 

mlir/include/mlir/IR/Dialect.h

@@ -114,7 +114,7 @@ class Dialect {
 
   /// Return the hook to parse an operation registered to this dialect, if any.
   /// By default this will lookup for registered operations and return the
-  /// `parse()` method registered on the AbstractOperation. Dialects can
+  /// `parse()` method registered on the RegisteredOperationName. Dialects can
   /// override this behavior and handle unregistered operations as well.
   virtual Optional<ParseOpHook> getParseOperationHook(StringRef opName) const;
 
@@ -194,7 +194,7 @@ class Dialect {
   ///
   template <typename... Args> void addOperations() {
     (void)std::initializer_list<int>{
-        0, (AbstractOperation::insert<Args>(*this), 0)...};
+        0, (RegisteredOperationName::insert<Args>(*this), 0)...};
   }
 
   /// Register a set of type classes with this dialect.

mlir/include/mlir/IR/MLIRContext.h

@@ -20,14 +20,14 @@ class ThreadPool;
 } // end namespace llvm
 
 namespace mlir {
-class AbstractOperation;
 class DebugActionManager;
 class DiagnosticEngine;
 class Dialect;
 class DialectRegistry;
 class InFlightDiagnostic;
 class Location;
 class MLIRContextImpl;
+class RegisteredOperationName;
 class StorageUniquer;
 
 /// MLIRContext is the top-level object for a collection of MLIR operations. It
@@ -172,7 +172,7 @@ class MLIRContext {
   /// Return information about all registered operations.  This isn't very
   /// efficient: typically you should ask the operations about their properties
   /// directly.
-  std::vector<AbstractOperation *> getRegisteredOperations();
+  std::vector<RegisteredOperationName> getRegisteredOperations();
 
   /// Return true if this operation name is registered in this context.
   bool isOperationRegistered(StringRef name);

mlir/include/mlir/IR/OpDefinition.h

@@ -191,7 +191,7 @@ class OpState {
   Operation *state;
 
   /// Allow access to internal hook implementation methods.
-  friend AbstractOperation;
+  friend RegisteredOperationName;
 };
 
 // Allow comparing operators.
@@ -1585,8 +1585,8 @@ class Op : public OpState, public Traits<ConcreteType>... {
 
   /// Return true if this "op class" can match against the specified operation.
   static bool classof(Operation *op) {
-    if (auto *abstractOp = op->getAbstractOperation())
-      return TypeID::get<ConcreteType>() == abstractOp->typeID;
+    if (auto info = op->getRegisteredInfo())
+      return TypeID::get<ConcreteType>() == info->getTypeID();
 #ifndef NDEBUG
     if (op->getName().getStringRef() == ConcreteType::getOperationName())
       llvm::report_fatal_error(
@@ -1628,13 +1628,13 @@ class Op : public OpState, public Traits<ConcreteType>... {
   /// for the concrete operation.
   template <typename... Models>
   static void attachInterface(MLIRContext &context) {
-    AbstractOperation *abstract = AbstractOperation::lookupMutable(
+    Optional<RegisteredOperationName> info = RegisteredOperationName::lookup(
         ConcreteType::getOperationName(), &context);
-    if (!abstract)
+    if (!info)
       llvm::report_fatal_error(
           "Attempting to attach an interface to an unregistered operation " +
           ConcreteType::getOperationName() + ".");
-    abstract->interfaceMap.insert<Models...>();
+    info->attachInterface<Models...>();
   }
 
 private:
@@ -1673,10 +1673,10 @@ class Op : public OpState, public Traits<ConcreteType>... {
     return detail::InterfaceMap::template get<Traits<ConcreteType>...>();
   }
 
-  /// Return the internal implementations of each of the AbstractOperation
+  /// Return the internal implementations of each of the OperationName
   /// hooks.
-  /// Implementation of `FoldHookFn` AbstractOperation hook.
-  static AbstractOperation::FoldHookFn getFoldHookFn() {
+  /// Implementation of `FoldHookFn` OperationName hook.
+  static OperationName::FoldHookFn getFoldHookFn() {
     return getFoldHookFnImpl<ConcreteType>();
   }
   /// The internal implementation of `getFoldHookFn` above that is invoked if
@@ -1685,7 +1685,7 @@ class Op : public OpState, public Traits<ConcreteType>... {
   static std::enable_if_t<llvm::is_one_of<OpTrait::OneResult<ConcreteOpT>,
                                           Traits<ConcreteOpT>...>::value &&
                               detect_has_single_result_fold<ConcreteOpT>::value,
-                          AbstractOperation::FoldHookFn>
+                          OperationName::FoldHookFn>
   getFoldHookFnImpl() {
     return [](Operation *op, ArrayRef<Attribute> operands,
               SmallVectorImpl<OpFoldResult> &results) {
@@ -1698,7 +1698,7 @@ class Op : public OpState, public Traits<ConcreteType>... {
   static std::enable_if_t<!llvm::is_one_of<OpTrait::OneResult<ConcreteOpT>,
                                            Traits<ConcreteOpT>...>::value &&
                               detect_has_fold<ConcreteOpT>::value,
-                          AbstractOperation::FoldHookFn>
+                          OperationName::FoldHookFn>
   getFoldHookFnImpl() {
     return [](Operation *op, ArrayRef<Attribute> operands,
               SmallVectorImpl<OpFoldResult> &results) {
@@ -1710,7 +1710,7 @@ class Op : public OpState, public Traits<ConcreteType>... {
   template <typename ConcreteOpT>
   static std::enable_if_t<!detect_has_single_result_fold<ConcreteOpT>::value &&
                               !detect_has_fold<ConcreteOpT>::value,
-                          AbstractOperation::FoldHookFn>
+                          OperationName::FoldHookFn>
   getFoldHookFnImpl() {
     return [](Operation *op, ArrayRef<Attribute> operands,
               SmallVectorImpl<OpFoldResult> &results) {
@@ -1754,29 +1754,29 @@ class Op : public OpState, public Traits<ConcreteType>... {
     return result;
   }
 
-  /// Implementation of `GetCanonicalizationPatternsFn` AbstractOperation hook.
-  static AbstractOperation::GetCanonicalizationPatternsFn
+  /// Implementation of `GetCanonicalizationPatternsFn` OperationName hook.
+  static OperationName::GetCanonicalizationPatternsFn
   getGetCanonicalizationPatternsFn() {
     return &ConcreteType::getCanonicalizationPatterns;
   }
   /// Implementation of `GetHasTraitFn`
-  static AbstractOperation::HasTraitFn getHasTraitFn() {
+  static OperationName::HasTraitFn getHasTraitFn() {
     return
         [](TypeID id) { return op_definition_impl::hasTrait<Traits...>(id); };
   }
-  /// Implementation of `ParseAssemblyFn` AbstractOperation hook.
-  static AbstractOperation::ParseAssemblyFn getParseAssemblyFn() {
+  /// Implementation of `ParseAssemblyFn` OperationName hook.
+  static OperationName::ParseAssemblyFn getParseAssemblyFn() {
     return &ConcreteType::parse;
   }
-  /// Implementation of `PrintAssemblyFn` AbstractOperation hook.
-  static AbstractOperation::PrintAssemblyFn getPrintAssemblyFn() {
+  /// Implementation of `PrintAssemblyFn` OperationName hook.
+  static OperationName::PrintAssemblyFn getPrintAssemblyFn() {
     return getPrintAssemblyFnImpl<ConcreteType>();
   }
   /// The internal implementation of `getPrintAssemblyFn` that is invoked when
   /// the concrete operation does not define a `print` method.
   template <typename ConcreteOpT>
   static std::enable_if_t<!detect_has_print<ConcreteOpT>::value,
-                          AbstractOperation::PrintAssemblyFn>
+                          OperationName::PrintAssemblyFn>
   getPrintAssemblyFnImpl() {
     return [](Operation *op, OpAsmPrinter &printer, StringRef defaultDialect) {
       return OpState::print(op, printer);
@@ -1786,7 +1786,7 @@ class Op : public OpState, public Traits<ConcreteType>... {
   /// the concrete operation defines a `print` method.
   template <typename ConcreteOpT>
   static std::enable_if_t<detect_has_print<ConcreteOpT>::value,
-                          AbstractOperation::PrintAssemblyFn>
+                          OperationName::PrintAssemblyFn>
   getPrintAssemblyFnImpl() {
     return &printAssembly;
   }
@@ -1795,8 +1795,8 @@ class Op : public OpState, public Traits<ConcreteType>... {
     OpState::printOpName(op, p, defaultDialect);
     return cast<ConcreteType>(op).print(p);
   }
-  /// Implementation of `VerifyInvariantsFn` AbstractOperation hook.
-  static AbstractOperation::VerifyInvariantsFn getVerifyInvariantsFn() {
+  /// Implementation of `VerifyInvariantsFn` OperationName hook.
+  static OperationName::VerifyInvariantsFn getVerifyInvariantsFn() {
     return &verifyInvariants;
   }
 
@@ -1816,7 +1816,7 @@ class Op : public OpState, public Traits<ConcreteType>... {
   }
 
   /// Allow access to internal implementation methods.
-  friend AbstractOperation;
+  friend RegisteredOperationName;
 };
 
 /// This class represents the base of an operation interface. See the definition
@@ -1836,22 +1836,22 @@ class OpInterface
 protected:
   /// Returns the impl interface instance for the given operation.
   static typename InterfaceBase::Concept *getInterfaceFor(Operation *op) {
-    // Access the raw interface from the abstract operation.
-    auto *abstractOp = op->getAbstractOperation();
-    if (abstractOp) {
-      if (auto *opIface = abstractOp->getInterface<ConcreteType>())
+    OperationName name = op->getName();
+
+    // Access the raw interface from the operation info.
+    if (Optional<RegisteredOperationName> rInfo = name.getRegisteredInfo()) {
+      if (auto *opIface = rInfo->getInterface<ConcreteType>())
         return opIface;
       // Fallback to the dialect to provide it with a chance to implement this
       // interface for this operation.
-      return abstractOp->dialect.getRegisteredInterfaceForOp<ConcreteType>(
+      return rInfo->getDialect().getRegisteredInterfaceForOp<ConcreteType>(
           op->getName());
     }
     // Fallback to the dialect to provide it with a chance to implement this
     // interface for this operation.
-    Dialect *dialect = op->getName().getDialect();
-    return dialect ? dialect->getRegisteredInterfaceForOp<ConcreteType>(
-                         op->getName())
-                   : nullptr;
+    if (Dialect *dialect = name.getDialect())
+      return dialect->getRegisteredInterfaceForOp<ConcreteType>(name);
+    return nullptr;
   }
 
   /// Allow access to `getInterfaceFor`.

mlir/include/mlir/IR/Operation.h

@@ -57,14 +57,14 @@ class alignas(8) Operation final
   OperationName getName() { return name; }
 
   /// If this operation has a registered operation description, return it.
-  /// Otherwise return null.
-  const AbstractOperation *getAbstractOperation() {
-    return getName().getAbstractOperation();
+  /// Otherwise return None.
+  Optional<RegisteredOperationName> getRegisteredInfo() {
+    return getName().getRegisteredInfo();
   }
 
   /// Returns true if this operation has a registered operation description,
   /// otherwise false.
-  bool isRegistered() { return getAbstractOperation(); }
+  bool isRegistered() { return getName().isRegistered(); }
 
   /// Remove this operation from its parent block and delete it.
   void erase();
@@ -468,16 +468,14 @@ class alignas(8) Operation final
   /// Returns true if the operation was registered with a particular trait, e.g.
   /// hasTrait<OperandsAreSignlessIntegerLike>().
   template <template <typename T> class Trait> bool hasTrait() {
-    const AbstractOperation *abstractOp = getAbstractOperation();
-    return abstractOp ? abstractOp->hasTrait<Trait>() : false;
+    return name.hasTrait<Trait>();
   }
 
-  /// Returns true if the operation is *might* have the provided trait. This
+  /// 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 AbstractOperation *abstractOp = getAbstractOperation();
-    return abstractOp ? abstractOp->hasTrait<Trait>() : true;
+    return name.mightHaveTrait<Trait>();
   }
 
   //===--------------------------------------------------------------------===//