[Mlir-commits] [mlir] 98dceed - [mlir] Make some functions public to use custom TypeIDs

[River Riddle]

Author: Mathieu Fehr
Date: 2021-04-20T10:56:00-07:00
New Revision: 98dceed64bd061ef42272fb84eea8fd2b84083ac

URL: https://github.com/llvm/llvm-project/commit/98dceed64bd061ef42272fb84eea8fd2b84083ac
DIFF: https://github.com/llvm/llvm-project/commit/98dceed64bd061ef42272fb84eea8fd2b84083ac.diff

LOG: [mlir] Make some functions public to use custom TypeIDs

Currently, it is only possible to register an operation or a type
when the TypeID is defined at compile time. Same with InterfaceMaps
which can only be defined with compile-time defined interfaces.

With those changes, it is now possible to register types/operations
with custom TypeIDs. This is necessary to define new operations/types
at runtime.

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/Dialect.h
    mlir/include/mlir/IR/OperationSupport.h
    mlir/include/mlir/IR/TypeSupport.h
    mlir/include/mlir/Support/InterfaceSupport.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Dialect.h b/mlir/include/mlir/IR/Dialect.h
index d3af95521d74..26ae5f9d9033 100644
--- a/mlir/include/mlir/IR/Dialect.h
+++ b/mlir/include/mlir/IR/Dialect.h
@@ -193,6 +193,11 @@ class Dialect {
     (void)std::initializer_list<int>{0, (addType<Args>(), 0)...};
   }
 
+  /// Register a type instance with this dialect.
+  /// The use of this method is in general discouraged in favor of
+  /// 'addTypes<CustomType>()'.
+  void addType(TypeID typeID, AbstractType &&typeInfo);
+
   /// Register a set of attribute classes with this dialect.
   template <typename... Args> void addAttributes() {
     (void)std::initializer_list<int>{0, (addAttribute<Args>(), 0)...};
@@ -231,7 +236,6 @@ class Dialect {
     addType(T::getTypeID(), AbstractType::get<T>(*this));
     detail::TypeUniquer::registerType<T>(context);
   }
-  void addType(TypeID typeID, AbstractType &&typeInfo);
 
   /// The namespace of this dialect.
   StringRef name;

diff  --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index cb82ec9c4714..7cfb979ea2ea 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -172,7 +172,9 @@ class AbstractOperation {
            T::getHasTraitFn());
   }
 
-private:
+  /// Register a new operation in a Dialect object.
+  /// 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,
@@ -180,6 +182,7 @@ class AbstractOperation {
                      GetCanonicalizationPatternsFn getCanonicalizationPatterns,
                      detail::InterfaceMap &&interfaceMap, HasTraitFn hasTrait);
 
+private:
   AbstractOperation(StringRef name, Dialect &dialect, TypeID typeID,
                     ParseAssemblyFn parseAssembly,
                     PrintAssemblyFn printAssembly,

diff  --git a/mlir/include/mlir/IR/TypeSupport.h b/mlir/include/mlir/IR/TypeSupport.h
index c1de58957915..898c26dc42c2 100644
--- a/mlir/include/mlir/IR/TypeSupport.h
+++ b/mlir/include/mlir/IR/TypeSupport.h
@@ -39,6 +39,15 @@ class AbstractType {
     return AbstractType(dialect, T::getInterfaceMap(), T::getTypeID());
   }
 
+  /// This method is used by Dialect objects to register types with
+  /// custom TypeIDs.
+  /// The use of this method is in general discouraged in favor of
+  /// 'get<CustomType>(dialect)';
+  static AbstractType get(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
+                          TypeID typeID) {
+    return AbstractType(dialect, std::move(interfaceMap), typeID);
+  }
+
   /// Return the dialect this type was registered to.
   Dialect &getDialect() const { return const_cast<Dialect &>(dialect); }
 

diff  --git a/mlir/include/mlir/Support/InterfaceSupport.h b/mlir/include/mlir/Support/InterfaceSupport.h
index 6fc61170a592..6af36aa1e622 100644
--- a/mlir/include/mlir/Support/InterfaceSupport.h
+++ b/mlir/include/mlir/Support/InterfaceSupport.h
@@ -188,6 +188,16 @@ class InterfaceMap {
   /// Returns true if the interface map contains an interface for the given id.
   bool contains(TypeID interfaceID) const { return lookup(interfaceID); }
 
+  /// Create an InterfaceMap given with the implementation of the interfaces.
+  /// The use of this constructor is in general discouraged in favor of
+  /// 'InterfaceMap::get<InterfaceA, ...>()'.
+  InterfaceMap(MutableArrayRef<std::pair<TypeID, void *>> elements)
+      : interfaces(elements.begin(), elements.end()) {
+    llvm::sort(interfaces, [](const auto &lhs, const auto &rhs) {
+      return compare(lhs.first, rhs.first);
+    });
+  }
+
 private:
   /// Compare two TypeID instances by comparing the underlying pointer.
   static bool compare(TypeID lhs, TypeID rhs) {
@@ -195,12 +205,6 @@ class InterfaceMap {
   }
 
   InterfaceMap() = default;
-  InterfaceMap(MutableArrayRef<std::pair<TypeID, void *>> elements)
-      : interfaces(elements.begin(), elements.end()) {
-    llvm::sort(interfaces, [](const auto &lhs, const auto &rhs) {
-      return compare(lhs.first, rhs.first);
-    });
-  }
 
   template <typename... Ts>
   static InterfaceMap getImpl(std::tuple<Ts...> *) {