[mlir] Add MLIR_USE_FALLBACK_TYPE_IDS macro support for TypeID

[bethebunny]

Adds a macro definition MLIR_USE_FALLBACK_TYPE_IDS. When this is defined, the MLIR_{DECLARE,DEFINE}_EXPLICIT_TYPE_ID functions explicitly fall back to string comparison.

This is useful for complex shared library setups
where it may be difficult to agree on a source of truth for specific type ID resolution. As long as there is a single resolution for registerImplicitTypeID, all type IDs can be reference a shared registration. This way types which are logically shared across multiple DSOs can have the same type ID, even if their definitions are duplicated.

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Stef Lindall (bethebunny)

Changes

Adds a macro definition MLIR_USE_FALLBACK_TYPE_IDS. When this is defined, the MLIR_{DECLARE,DEFINE}_EXPLICIT_TYPE_ID functions explicitly fall back to string comparison.

This is useful for complex shared library setups
where it may be difficult to agree on a source of truth for specific type ID resolution. As long as there is a single resolution for registerImplicitTypeID, all type IDs can be reference a shared registration. This way types which are logically shared across multiple DSOs can have the same type ID, even if their definitions are duplicated.

---

Full diff: https://github.com/llvm/llvm-project/pull/126999.diff

2 Files Affected:

• (modified) mlir/include/mlir/Support/TypeID.h (+74-16)
• (modified) mlir/lib/Support/TypeID.cpp (+3-2)

diff --git a/mlir/include/mlir/Support/TypeID.h b/mlir/include/mlir/Support/TypeID.h
index f27f73fb19928..459e9dae12a9f 100644
--- a/mlir/include/mlir/Support/TypeID.h
+++ b/mlir/include/mlir/Support/TypeID.h
@@ -19,6 +19,7 @@
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Allocator.h"
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include "llvm/Support/TypeName.h"
 
@@ -100,6 +101,8 @@ namespace mlir {
 ///     uses the name of the type, it may not be used for types defined in
 ///     anonymous namespaces (which is asserted when it can be detected). String
 ///     names do not provide any guarantees on uniqueness in these contexts.
+///   - This behavior may be forced even in the presence of explicit declarations
+///     by specifying `MLIR_USE_FALLBACK_TYPE_IDS`.
 ///
 class TypeID {
   /// This class represents the storage of a type info object.
@@ -161,9 +164,30 @@ namespace detail {
 class FallbackTypeIDResolver {
 protected:
   /// Register an implicit type ID for the given type name.
-  static TypeID registerImplicitTypeID(StringRef name);
+  LLVM_ALWAYS_EXPORT static TypeID registerImplicitTypeID(StringRef name);
 };
 
+template <typename T>
+struct is_fully_resolved_t {
+  /// Trait to check if `U` is fully resolved. We use this to verify that `T` is
+  /// fully resolved when trying to resolve a TypeID. We don't technically need
+  /// to have the full definition of `T` for the fallback, but it does help
+  /// prevent situations where a forward declared type uses this fallback even
+  /// though there is a strong definition for the TypeID in the location where
+  /// `T` is defined.
+  template <typename U>
+  using is_fully_resolved_trait = decltype(sizeof(U));
+  template <typename U>
+  using is_fully_resolved = llvm::is_detected<is_fully_resolved_trait, U>;
+  static constexpr bool value = is_fully_resolved<T>::value;
+};
+
+template <typename T>
+constexpr bool is_fully_resolved() {
+  /// Helper function for is_fully_resolved_t.
+  return is_fully_resolved_t<T>::value;
+}
+
 /// This class provides a resolver for getting the ID for a given class T. This
 /// allows for the derived type to specialize its resolution behavior. The
 /// default implementation uses the string name of the type to resolve the ID.
@@ -178,19 +202,8 @@ class FallbackTypeIDResolver {
 template <typename T, typename Enable = void>
 class TypeIDResolver : public FallbackTypeIDResolver {
 public:
-  /// Trait to check if `U` is fully resolved. We use this to verify that `T` is
-  /// fully resolved when trying to resolve a TypeID. We don't technically need
-  /// to have the full definition of `T` for the fallback, but it does help
-  /// prevent situations where a forward declared type uses this fallback even
-  /// though there is a strong definition for the TypeID in the location where
-  /// `T` is defined.
-  template <typename U>
-  using is_fully_resolved_trait = decltype(sizeof(U));
-  template <typename U>
-  using is_fully_resolved = llvm::is_detected<is_fully_resolved_trait, U>;
-
   static TypeID resolveTypeID() {
-    static_assert(is_fully_resolved<T>::value,
+    static_assert(is_fully_resolved<T>(),
                   "TypeID::get<> requires the complete definition of `T`");
     static TypeID id = registerImplicitTypeID(llvm::getTypeName<T>());
     return id;
@@ -246,7 +259,7 @@ TypeID TypeID::get() {
 // circumstances a hard-to-catch runtime bug when a TypeID is hidden in two
 // different shared libraries and instances of the same class only gets the same
 // TypeID inside a given DSO.
-#define MLIR_DECLARE_EXPLICIT_TYPE_ID(CLASS_NAME)                              \
+#define MLIR_DECLARE_EXPLICIT_SELF_OWNING_TYPE_ID(CLASS_NAME)                  \
   namespace mlir {                                                             \
   namespace detail {                                                           \
   template <>                                                                  \
@@ -260,13 +273,57 @@ TypeID TypeID::get() {
   } /* namespace detail */                                                     \
   } /* namespace mlir */
 
-#define MLIR_DEFINE_EXPLICIT_TYPE_ID(CLASS_NAME)                               \
+#define MLIR_DEFINE_EXPLICIT_SELF_OWNING_TYPE_ID(CLASS_NAME)                   \
   namespace mlir {                                                             \
   namespace detail {                                                           \
   SelfOwningTypeID TypeIDResolver<CLASS_NAME>::id = {};                        \
   } /* namespace detail */                                                     \
   } /* namespace mlir */
 
+
+/// Declare/define an explicit specialization for TypeID using the string
+/// comparison fallback. This is useful for complex shared library setups
+/// where it may be difficult to agree on a source of truth for specific
+/// type ID resolution. As long as there is a single resolution for
+/// registerImplicitTypeID, all type IDs can be reference a shared
+/// registration. This way types which are logically shared across multiple
+/// DSOs can have the same type ID, even if their definitions are duplicated.
+#define MLIR_DECLARE_EXPLICIT_FALLBACK_TYPE_ID(CLASS_NAME)                     \
+  namespace mlir {                                                             \
+  namespace detail {                                                           \
+  template <>                                                                  \
+  class TypeIDResolver<CLASS_NAME> : public FallbackTypeIDResolver {           \
+  public:                                                                      \
+    static TypeID resolveTypeID() {                                            \
+      static_assert(is_fully_resolved<CLASS_NAME>(),                           \
+                    "TypeID::get<> requires the complete definition of `T`");  \
+      static TypeID id =                                                       \
+          registerImplicitTypeID(llvm::getTypeName<CLASS_NAME>());             \
+      return id;                                                               \
+    }                                                                          \
+  };                                                                           \
+  } /* namespace detail */                                                     \
+  } /* namespace mlir */
+
+#define MLIR_DEFINE_EXPLICIT_FALLBACK_TYPE_ID(CLASS_NAME)
+
+
+#ifndef MLIR_USE_FALLBACK_TYPE_IDS
+#define MLIR_USE_FALLBACK_TYPE_IDS false
+#endif
+
+#if MLIR_USE_FALLBACK_TYPE_IDS
+#define MLIR_DECLARE_EXPLICIT_TYPE_ID(CLASS_NAME)                              \
+  MLIR_DECLARE_EXPLICIT_FALLBACK_TYPE_ID(CLASS_NAME)
+#define MLIR_DEFINE_EXPLICIT_TYPE_ID(CLASS_NAME)                               \
+  MLIR_DEFINE_EXPLICIT_FALLBACK_TYPE_ID(CLASS_NAME)
+#else
+#define MLIR_DECLARE_EXPLICIT_TYPE_ID(CLASS_NAME)                              \
+  MLIR_DECLARE_EXPLICIT_SELF_OWNING_TYPE_ID(CLASS_NAME)
+#define MLIR_DEFINE_EXPLICIT_TYPE_ID(CLASS_NAME)                               \
+  MLIR_DEFINE_EXPLICIT_SELF_OWNING_TYPE_ID(CLASS_NAME)
+#endif /* MLIR_USE_FALLBACK_TYPE_IDS */
+
 // Declare/define an explicit, **internal**, specialization of TypeID for the
 // given class. This is useful for providing an explicit specialization of
 // TypeID for a class that is known to be internal to a specific library. It
@@ -331,7 +388,8 @@ class alignas(8) SelfOwningTypeID {
 //===----------------------------------------------------------------------===//
 
 /// Explicitly register a set of "builtin" types.
-MLIR_DECLARE_EXPLICIT_TYPE_ID(void)
+/// `void` must be self-owning, it can't be fully resolved.
+MLIR_DECLARE_EXPLICIT_SELF_OWNING_TYPE_ID(void)
 
 namespace llvm {
 template <>
diff --git a/mlir/lib/Support/TypeID.cpp b/mlir/lib/Support/TypeID.cpp
index e499e2f283633..01ad910113018 100644
--- a/mlir/lib/Support/TypeID.cpp
+++ b/mlir/lib/Support/TypeID.cpp
@@ -80,7 +80,8 @@ struct ImplicitTypeIDRegistry {
 };
 } // end namespace
 
-TypeID detail::FallbackTypeIDResolver::registerImplicitTypeID(StringRef name) {
+LLVM_ALWAYS_EXPORT TypeID
+detail::FallbackTypeIDResolver::registerImplicitTypeID(StringRef name) {
   static ImplicitTypeIDRegistry registry;
   return registry.lookupOrInsert(name);
 }
@@ -89,4 +90,4 @@ TypeID detail::FallbackTypeIDResolver::registerImplicitTypeID(StringRef name) {
 // Builtin TypeIDs
 //===----------------------------------------------------------------------===//
 
-MLIR_DEFINE_EXPLICIT_TYPE_ID(void)
+MLIR_DEFINE_EXPLICIT_SELF_OWNING_TYPE_ID(void)

[lattner]

Cool, thank you for implementing and iterating on this!

[dukebw]

LGTM. Thank you!

[github-actions]

@bethebunny Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!

[joker-eph]

While adding the macro here won't hurt much, the feature seems quite incomplete to me: there is no CMake option to enable it for example. Also it isn't documented right now.
More importantly, with respect to the end goal of making MLIR resilient to shared-libraries builds where TypeID are duplicated, it's not clear that this would cover uses of TypeID::get<T>() ; that is all of these are still likely broken in the scenario you're trying to address: https://github.com/search?q=repo%3Allvm%2Fllvm-project%20TypeID%3A%3Aget&type=code

[bethebunny]

Thank you so much for the feedback!

• My project uses Bazel, I'm much less familiar with CMake, but I'll look into how those are defined and send a followup PR.
• Where do you think is the right place to add more documentation? I learned mainly through the docs in TypeID.h so I added some threads there, but I definitely admit that it would be quite hard for someone to meaningfully follow this path themselves still.
• I can confirm that this completely resolves the issue in my specific case.
  • Consider a tree of SOs using a tree of MLIR dialects
  • The property we want to achieve is that for library A and B, regardless of where they exist in this tree, if they both use a shared MLIR type T, then TypeID::get<T_A>() == TypeID::get<T_B>().
    • Note that this is not true when A and B separately compile definitions of T with explicit type IDs. In this case each type ID is the memory location of separate global statics.
  • In practice this has proven intractable to achieve when there's a large tree of SOs using many different dialects, and in particular with visibility=hidden. For each dialect you need to designate an SO which owns this dialect's types, and ensure that no downstream SO decides that it should re-define them.
  • With fallback type IDs you have the same problem, but only 1 occurrence of it: you need all libraries to agree on a single definition of registerImplicitTypeID.
    • This is relatively practical to achieve since each library needs to depend on mlir:IR anyway.
    • In this case, when T_A and T_B are defined, each are assigned static TypeIDs by calling registerImplicitTypeID, which internally recognizes that T_A and T_B are the same type.
    • Now TypeID::get<T_A>() == TypeID::get<T_B>() because they were both assigned the same type ID by the central registry.
  • Empirically, this patch has been a breakthrough that has allowed me to use dialects in cases I was never able to achieve before!! I'm very excited that this was able to merge, and happy to do followup work to make it more useful to others.

[bethebunny]

Another thought, there's probably some clever way to actually write a test for this 🤔 I'll noodle on how to do that.