Skip to content

[mlir] Add MLIR_USE_FALLBACK_TYPE_IDS macro support for TypeID #126999

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 15, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 74 additions & 16 deletions mlir/include/mlir/Support/TypeID.h
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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;
Expand Down Expand Up @@ -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 <> \
Expand All @@ -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
Expand Down Expand Up @@ -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 <>
Expand Down
5 changes: 3 additions & 2 deletions mlir/lib/Support/TypeID.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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)