xref: /aosp_15_r20/external/pytorch/c10/util/AbortHandler.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1*da0073e9SAndroid Build Coastguard Worker #include <c10/macros/Macros.h>
2*da0073e9SAndroid Build Coastguard Worker #include <c10/util/Backtrace.h>
3*da0073e9SAndroid Build Coastguard Worker #include <c10/util/env.h>
4*da0073e9SAndroid Build Coastguard Worker #include <cstdlib>
5*da0073e9SAndroid Build Coastguard Worker #include <exception>
6*da0073e9SAndroid Build Coastguard Worker #include <iostream>
7*da0073e9SAndroid Build Coastguard Worker #include <mutex>
8*da0073e9SAndroid Build Coastguard Worker #include <optional>
9*da0073e9SAndroid Build Coastguard Worker 
10*da0073e9SAndroid Build Coastguard Worker namespace c10 {
11*da0073e9SAndroid Build Coastguard Worker class AbortHandlerHelper {
12*da0073e9SAndroid Build Coastguard Worker  public:
getInstance()13*da0073e9SAndroid Build Coastguard Worker   static AbortHandlerHelper& getInstance() {
14*da0073e9SAndroid Build Coastguard Worker #ifdef _WIN32
15*da0073e9SAndroid Build Coastguard Worker     thread_local
16*da0073e9SAndroid Build Coastguard Worker #endif // _WIN32
17*da0073e9SAndroid Build Coastguard Worker         static AbortHandlerHelper instance;
18*da0073e9SAndroid Build Coastguard Worker     return instance;
19*da0073e9SAndroid Build Coastguard Worker   }
20*da0073e9SAndroid Build Coastguard Worker 
set(std::terminate_handler handler)21*da0073e9SAndroid Build Coastguard Worker   void set(std::terminate_handler handler) {
22*da0073e9SAndroid Build Coastguard Worker     std::lock_guard<std::mutex> lk(mutex);
23*da0073e9SAndroid Build Coastguard Worker     if (!inited) {
24*da0073e9SAndroid Build Coastguard Worker       prev = std::set_terminate(handler);
25*da0073e9SAndroid Build Coastguard Worker       curr = std::get_terminate();
26*da0073e9SAndroid Build Coastguard Worker       inited = true;
27*da0073e9SAndroid Build Coastguard Worker     }
28*da0073e9SAndroid Build Coastguard Worker   }
29*da0073e9SAndroid Build Coastguard Worker 
getPrev()30*da0073e9SAndroid Build Coastguard Worker   std::terminate_handler getPrev() const {
31*da0073e9SAndroid Build Coastguard Worker     return prev;
32*da0073e9SAndroid Build Coastguard Worker   }
33*da0073e9SAndroid Build Coastguard Worker 
34*da0073e9SAndroid Build Coastguard Worker  private:
35*da0073e9SAndroid Build Coastguard Worker   std::terminate_handler prev = nullptr;
36*da0073e9SAndroid Build Coastguard Worker   std::terminate_handler curr = nullptr;
37*da0073e9SAndroid Build Coastguard Worker   bool inited = false;
38*da0073e9SAndroid Build Coastguard Worker   std::mutex mutex;
39*da0073e9SAndroid Build Coastguard Worker   AbortHandlerHelper() = default;
~AbortHandlerHelper()40*da0073e9SAndroid Build Coastguard Worker   ~AbortHandlerHelper() {
41*da0073e9SAndroid Build Coastguard Worker     // Only restore the handler if we are the current one
42*da0073e9SAndroid Build Coastguard Worker     if (inited && curr == std::get_terminate()) {
43*da0073e9SAndroid Build Coastguard Worker       std::set_terminate(prev);
44*da0073e9SAndroid Build Coastguard Worker     }
45*da0073e9SAndroid Build Coastguard Worker   }
46*da0073e9SAndroid Build Coastguard Worker 
47*da0073e9SAndroid Build Coastguard Worker  public:
48*da0073e9SAndroid Build Coastguard Worker   AbortHandlerHelper(AbortHandlerHelper const&) = delete;
49*da0073e9SAndroid Build Coastguard Worker   void operator=(AbortHandlerHelper const&) = delete;
50*da0073e9SAndroid Build Coastguard Worker };
51*da0073e9SAndroid Build Coastguard Worker 
52*da0073e9SAndroid Build Coastguard Worker namespace detail {
terminate_handler()53*da0073e9SAndroid Build Coastguard Worker C10_ALWAYS_INLINE void terminate_handler() {
54*da0073e9SAndroid Build Coastguard Worker   std::cout << "Unhandled exception caught in c10/util/AbortHandler.h" << '\n';
55*da0073e9SAndroid Build Coastguard Worker   auto backtrace = get_backtrace();
56*da0073e9SAndroid Build Coastguard Worker   std::cout << backtrace << '\n' << std::flush;
57*da0073e9SAndroid Build Coastguard Worker   auto prev_handler = AbortHandlerHelper::getInstance().getPrev();
58*da0073e9SAndroid Build Coastguard Worker   if (prev_handler) {
59*da0073e9SAndroid Build Coastguard Worker     prev_handler();
60*da0073e9SAndroid Build Coastguard Worker   } else {
61*da0073e9SAndroid Build Coastguard Worker     std::abort();
62*da0073e9SAndroid Build Coastguard Worker   }
63*da0073e9SAndroid Build Coastguard Worker }
64*da0073e9SAndroid Build Coastguard Worker } // namespace detail
65*da0073e9SAndroid Build Coastguard Worker 
set_terminate_handler()66*da0073e9SAndroid Build Coastguard Worker C10_ALWAYS_INLINE void set_terminate_handler() {
67*da0073e9SAndroid Build Coastguard Worker   bool use_custom_terminate = false;
68*da0073e9SAndroid Build Coastguard Worker   // On Windows it is enabled by default based on
69*da0073e9SAndroid Build Coastguard Worker   // https://github.com/pytorch/pytorch/pull/50320#issuecomment-763147062
70*da0073e9SAndroid Build Coastguard Worker #ifdef _WIN32
71*da0073e9SAndroid Build Coastguard Worker   use_custom_terminate = true;
72*da0073e9SAndroid Build Coastguard Worker #endif // _WIN32
73*da0073e9SAndroid Build Coastguard Worker   auto result = c10::utils::check_env("TORCH_CUSTOM_TERMINATE");
74*da0073e9SAndroid Build Coastguard Worker   if (result != std::nullopt) {
75*da0073e9SAndroid Build Coastguard Worker     use_custom_terminate = result.value();
76*da0073e9SAndroid Build Coastguard Worker   }
77*da0073e9SAndroid Build Coastguard Worker   if (use_custom_terminate) {
78*da0073e9SAndroid Build Coastguard Worker     AbortHandlerHelper::getInstance().set(detail::terminate_handler);
79*da0073e9SAndroid Build Coastguard Worker   }
80*da0073e9SAndroid Build Coastguard Worker }
81*da0073e9SAndroid Build Coastguard Worker } // namespace c10
82