xref: /aosp_15_r20/external/pytorch/c10/util/thread_name.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1*da0073e9SAndroid Build Coastguard Worker #include <c10/util/thread_name.h>
2*da0073e9SAndroid Build Coastguard Worker 
3*da0073e9SAndroid Build Coastguard Worker #include <algorithm>
4*da0073e9SAndroid Build Coastguard Worker #include <array>
5*da0073e9SAndroid Build Coastguard Worker 
6*da0073e9SAndroid Build Coastguard Worker #ifndef __GLIBC_PREREQ
7*da0073e9SAndroid Build Coastguard Worker #define __GLIBC_PREREQ(x, y) 0
8*da0073e9SAndroid Build Coastguard Worker #endif
9*da0073e9SAndroid Build Coastguard Worker 
10*da0073e9SAndroid Build Coastguard Worker #if defined(__GLIBC__) && __GLIBC_PREREQ(2, 12) && !defined(__APPLE__) && \
11*da0073e9SAndroid Build Coastguard Worker     !defined(__ANDROID__)
12*da0073e9SAndroid Build Coastguard Worker #define C10_HAS_PTHREAD_SETNAME_NP
13*da0073e9SAndroid Build Coastguard Worker #endif
14*da0073e9SAndroid Build Coastguard Worker 
15*da0073e9SAndroid Build Coastguard Worker #ifdef C10_HAS_PTHREAD_SETNAME_NP
16*da0073e9SAndroid Build Coastguard Worker #include <pthread.h>
17*da0073e9SAndroid Build Coastguard Worker #endif
18*da0073e9SAndroid Build Coastguard Worker 
19*da0073e9SAndroid Build Coastguard Worker namespace c10 {
20*da0073e9SAndroid Build Coastguard Worker 
21*da0073e9SAndroid Build Coastguard Worker #ifdef C10_HAS_PTHREAD_SETNAME_NP
22*da0073e9SAndroid Build Coastguard Worker namespace {
23*da0073e9SAndroid Build Coastguard Worker // pthreads has a limit of 16 characters including the null termination byte.
24*da0073e9SAndroid Build Coastguard Worker constexpr size_t kMaxThreadName = 15;
25*da0073e9SAndroid Build Coastguard Worker } // namespace
26*da0073e9SAndroid Build Coastguard Worker #endif
27*da0073e9SAndroid Build Coastguard Worker 
setThreadName(std::string name)28*da0073e9SAndroid Build Coastguard Worker void setThreadName(std::string name) {
29*da0073e9SAndroid Build Coastguard Worker #ifdef C10_HAS_PTHREAD_SETNAME_NP
30*da0073e9SAndroid Build Coastguard Worker   name.resize(std::min(name.size(), kMaxThreadName));
31*da0073e9SAndroid Build Coastguard Worker 
32*da0073e9SAndroid Build Coastguard Worker   pthread_setname_np(pthread_self(), name.c_str());
33*da0073e9SAndroid Build Coastguard Worker #endif
34*da0073e9SAndroid Build Coastguard Worker }
35*da0073e9SAndroid Build Coastguard Worker 
getThreadName()36*da0073e9SAndroid Build Coastguard Worker std::string getThreadName() {
37*da0073e9SAndroid Build Coastguard Worker #ifdef C10_HAS_PTHREAD_SETNAME_NP
38*da0073e9SAndroid Build Coastguard Worker   std::array<char, kMaxThreadName + 1> name{};
39*da0073e9SAndroid Build Coastguard Worker   pthread_getname_np(pthread_self(), name.data(), name.size());
40*da0073e9SAndroid Build Coastguard Worker   return name.data();
41*da0073e9SAndroid Build Coastguard Worker #else
42*da0073e9SAndroid Build Coastguard Worker   return "";
43*da0073e9SAndroid Build Coastguard Worker #endif
44*da0073e9SAndroid Build Coastguard Worker }
45*da0073e9SAndroid Build Coastguard Worker 
46*da0073e9SAndroid Build Coastguard Worker } // namespace c10
47