xref: /aosp_15_r20/external/pytorch/c10/util/TypeIndex.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <c10/util/ConstexprCrc.h>
4 #include <c10/util/IdWrapper.h>
5 #include <c10/util/string_view.h>
6 #include <cstdint>
7 #include <ostream>
8 #include <stdexcept>
9 #include <string>
10 #include <type_traits>
11 
12 namespace c10::util {
13 
14 // TODO Make it work for more compilers
15 
16 // Intel compiler works
17 #if defined(__INTEL_COMPILER)
18 #define C10_TYPENAME_SUPPORTS_CONSTEXPR 0
19 #define C10_TYPENAME_CONSTEXPR
20 
21 // Clang works
22 #elif defined(__clang__)
23 
24 // except for NVCC
25 #if defined(__CUDACC__)
26 #define C10_TYPENAME_SUPPORTS_CONSTEXPR 0
27 #define C10_TYPENAME_CONSTEXPR
28 #else
29 #define C10_TYPENAME_SUPPORTS_CONSTEXPR 1
30 #define C10_TYPENAME_CONSTEXPR constexpr
31 #endif
32 
33 // Windows works
34 #elif defined(_MSC_VER)
35 
36 // except for NVCC
37 #if defined(__CUDACC__)
38 #define C10_TYPENAME_SUPPORTS_CONSTEXPR 0
39 #define C10_TYPENAME_CONSTEXPR
40 #else
41 #define C10_TYPENAME_SUPPORTS_CONSTEXPR 1
42 #define C10_TYPENAME_CONSTEXPR constexpr
43 #endif
44 
45 // GCC works
46 #elif defined(__GNUC__)
47 
48 // except when gcc < 9
49 #if (__GNUC__ < 9) || defined(__CUDACC__)
50 #define C10_TYPENAME_SUPPORTS_CONSTEXPR 0
51 #define C10_TYPENAME_CONSTEXPR
52 #else
53 #define C10_TYPENAME_SUPPORTS_CONSTEXPR 1
54 #define C10_TYPENAME_CONSTEXPR constexpr
55 #endif
56 
57 // some other compiler we don't know about
58 #else
59 #define C10_TYPENAME_SUPPORTS_CONSTEXPR 1
60 #define C10_TYPENAME_CONSTEXPR constexpr
61 #endif
62 
63 struct type_index final : IdWrapper<type_index, uint64_t> {
type_indexfinal64   constexpr explicit type_index(uint64_t checksum) : IdWrapper(checksum) {}
65 
66   // Allow usage in std::map / std::set
67   // TODO Disallow this and rather use std::unordered_map/set everywhere
68   friend constexpr bool operator<(type_index lhs, type_index rhs) noexcept {
69     return lhs.underlyingId() < rhs.underlyingId();
70   }
71 
72   friend std::ostream& operator<<(std::ostream& stream, type_index typeId) {
73     return stream << typeId.underlyingId();
74   }
75 };
76 
77 namespace detail {
78 
79 #if !defined(__clang__) && !defined(_MSC_VER) && defined(__GNUC__) && \
80     __GNUC__ < 5
81 // Getting __PRETTY_FUNCTION__ at compile time only works with GCC >= 5
82 #error "You're running a too old version of GCC. We need GCC 5 or later."
83 #endif
84 
85 #if defined(__clang__) && __clang_major__ < 4
86 // Getting __PRETTY_FUNCTION__ at compile time only works with Clang >= 4
87 #error "You're running a too old version of Clang. We need Clang 4 or later."
88 #endif
89 
extract(string_view prefix,string_view suffix,string_view str)90 inline constexpr string_view extract(
91     string_view prefix,
92     string_view suffix,
93     string_view str) {
94 #if !defined(__CUDA_ARCH__) // CUDA doesn't like std::logic_error in device code
95   return (!str.starts_with(prefix) || !str.ends_with(suffix))
96       ? (throw std::logic_error("Invalid pattern"), string_view())
97       : str.substr(prefix.size(), str.size() - prefix.size() - suffix.size());
98 #else
99   return str.substr(prefix.size(), str.size() - prefix.size() - suffix.size());
100 #endif
101 }
102 
103 template <typename T>
fully_qualified_type_name_impl()104 inline C10_TYPENAME_CONSTEXPR c10::string_view fully_qualified_type_name_impl() {
105 #if defined(_MSC_VER) && !defined(__clang__)
106 #if defined(__NVCC__)
107   return extract(
108       "c10::basic_string_view<char> c10::util::detail::fully_qualified_type_name_impl<",
109       ">()",
110       __FUNCSIG__);
111 #else
112   return extract(
113       "class c10::basic_string_view<char> __cdecl c10::util::detail::fully_qualified_type_name_impl<",
114       ">(void)",
115       __FUNCSIG__);
116 #endif
117 #elif defined(__clang__)
118   return extract(
119       "c10::string_view c10::util::detail::fully_qualified_type_name_impl() [T = ",
120       "]",
121       __PRETTY_FUNCTION__);
122 #elif defined(__GNUC__)
123   return extract(
124 #if C10_TYPENAME_SUPPORTS_CONSTEXPR
125       "constexpr c10::string_view c10::util::detail::fully_qualified_type_name_impl() [with T = ",
126 #else
127       "c10::string_view c10::util::detail::fully_qualified_type_name_impl() [with T = ",
128 #endif
129       "; c10::string_view = c10::basic_string_view<char>]",
130       __PRETTY_FUNCTION__);
131 #endif
132 }
133 
134 #if !defined(__CUDA_ARCH__)
135 template <typename T>
type_index_impl()136 inline constexpr uint64_t type_index_impl() {
137 // Idea: __PRETTY_FUNCTION__ (or __FUNCSIG__ on msvc) contains a qualified name
138 // of this function, including its template parameter, i.e. including the
139 // type we want an id for. We use this name and run crc64 on it to get a type
140 // id.
141 #if defined(_MSC_VER) && !defined(__clang__)
142   return crc64(__FUNCSIG__, sizeof(__FUNCSIG__)).checksum();
143 #elif defined(__clang__)
144   return crc64(__PRETTY_FUNCTION__, sizeof(__PRETTY_FUNCTION__)).checksum();
145 #elif defined(__GNUC__)
146   return crc64(__PRETTY_FUNCTION__, sizeof(__PRETTY_FUNCTION__)).checksum();
147 #endif
148 }
149 #endif
150 
151 } // namespace detail
152 
153 template <typename T>
get_type_index()154 inline constexpr type_index get_type_index() {
155 #if !defined(__CUDA_ARCH__)
156   // To enforce that this is really computed at compile time, we pass the
157   // type index through std::integral_constant.
158   return type_index{std::integral_constant<
159       uint64_t,
160       detail::type_index_impl<std::decay_t<T>>()>::value};
161 #else
162   // There's nothing in theory preventing us from running this on device code
163   // except for nvcc throwing a compiler error if we enable it.
164   return (abort(), type_index(0));
165 #endif
166 }
167 
168 #if !defined(TORCH_PEDANTIC)
169 // Use precomputed hashsum for std::string
170 // Needed to workaround ambiguity in class name resolution
171 // into __PRETTY_FUNCTION__ when abovementioned class is defined in inlined
172 // namespace. In multi-ABI C++ library, `std::string` is an alias to
173 // `std::__cxx11::basic_string<char>` which depending on compiler flags can be
174 // resolved to `basic_string<char>` either in `std` namespace or in
175 // `std::__cxx11` one (`__cxx11` is an inline namespace)
176 template <>
177 inline constexpr type_index get_type_index<std::string>() {
178   // hashsum for std::basic_string<char>
179   return type_index{4193213214807308375ULL};
180 }
181 #endif
182 
183 template <typename T>
184 inline C10_TYPENAME_CONSTEXPR string_view
get_fully_qualified_type_name()185 get_fully_qualified_type_name() noexcept {
186 #if C10_TYPENAME_SUPPORTS_CONSTEXPR
187   constexpr
188 #else
189   static
190 #endif
191       string_view name = detail::fully_qualified_type_name_impl<T>();
192   return name;
193 }
194 } // namespace c10::util
195 
196 C10_DEFINE_HASH_FOR_IDWRAPPER(c10::util::type_index);
197