pytorch

1
#include <c10/util/thread_name.h>
2

3
#include <algorithm>
4
#include <array>
5

6
#ifndef __GLIBC_PREREQ
7
#define __GLIBC_PREREQ(x, y) 0
8
#endif
9

10
#if defined(__GLIBC__) && __GLIBC_PREREQ(2, 12) && !defined(__APPLE__) && \
11
    !defined(__ANDROID__)
12
#define C10_HAS_PTHREAD_SETNAME_NP
13
#endif
14

15
#ifdef C10_HAS_PTHREAD_SETNAME_NP
16
#include <pthread.h>
17
#endif
18

19
namespace c10 {
20

21
#ifdef C10_HAS_PTHREAD_SETNAME_NP
22
namespace {
23
// pthreads has a limit of 16 characters including the null termination byte.
24
constexpr size_t kMaxThreadName = 15;
25
} // namespace
26
#endif
27

28
void setThreadName(std::string name) {
29
#ifdef C10_HAS_PTHREAD_SETNAME_NP
30
  name.resize(std::min(name.size(), kMaxThreadName));
31

32
  pthread_setname_np(pthread_self(), name.c_str());
33
#endif
34
}
35

36
std::string getThreadName() {
37
#ifdef C10_HAS_PTHREAD_SETNAME_NP
38
  std::array<char, kMaxThreadName + 1> name{};
39
  pthread_getname_np(pthread_self(), name.data(), name.size());
40
  return name.data();
41
#else
42
  return "";
43
#endif
44
}
45

46
} // namespace c10
47