xref: /aosp_15_r20/external/pytorch/aten/src/ATen/core/custom_class.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <typeindex>
4 #include <memory>
5 
6 #include <c10/macros/Export.h>
7 #include <c10/macros/Macros.h>
8 #include <c10/util/Exception.h>
9 
10 namespace c10 {
11 
12 struct ClassType;
13 using ClassTypePtr = std::shared_ptr<ClassType>;
14 
15 TORCH_API c10::ClassTypePtr getCustomClassTypeImpl(const std::type_index &tindex);
16 
17 template <typename T>
getCustomClassType()18 const c10::ClassTypePtr& getCustomClassType() {
19   // Classes are never unregistered from getCustomClassTypeMap and the
20   // hash lookup can be a hot path, so just cache.
21   // For the same reason, it's fine If this ends up getting duplicated across
22   // DSO boundaries for whatever reason.
23   static c10::ClassTypePtr cache = getCustomClassTypeImpl(
24       std::type_index(typeid(T)));
25   return cache;
26 }
27 
28 }
29