1 | #include <instrumentation.h> |
2 | |
3 | #include <c10/macros/Export.h> |
4 | |
5 | #ifdef _WIN32 |
6 | #include <c10/util/win32-headers.h> |
7 | #else |
8 | #include <pthread.h> |
9 | #include <unistd.h> |
10 | #endif |
11 | |
12 | namespace torch { |
13 | namespace jit { |
14 | namespace fuser { |
15 | namespace cuda { |
16 | namespace inst { |
17 | |
18 | Trace::Trace() { |
19 | const char* trace_filename = getenv("PYTORCH_NVFUSER_TRACE" ); |
20 | if (trace_filename != nullptr) { |
21 | log_file_ = fopen(trace_filename, "w" ); |
22 | TORCH_CHECK(log_file_ != nullptr, "Can't open trace file" ); |
23 | |
24 | // Disable the file stream buffering, since it may result |
25 | // in torn writes in multi-threaded tracing |
26 | setbuf(log_file_, nullptr); |
27 | |
28 | // Print the trace prologue |
29 | // (including a dummy TRACE_START event) |
30 | fprintf(log_file_, "{\n\"traceEvents\": [\n" ); |
31 | start_timestamp_ = Clock::now(); |
32 | logEvent('I', "TRACE_START" ); |
33 | } |
34 | |
35 | if (isOptionDisabled(DisableOption::Nvtx)) { |
36 | record_nvtx_range_ = false; |
37 | } |
38 | } |
39 | |
40 | Trace::~Trace() { |
41 | if (log_file_ != nullptr) { |
42 | // Print trace epilogue |
43 | logEvent('I', "TRACE_END" , ' '); |
44 | fprintf(log_file_, "],\n\"displayTimeUnit\": \"ms\"\n}\n" ); |
45 | fclose(log_file_); |
46 | } |
47 | } |
48 | |
49 | void Trace::logEvent(char ph, const char* name, char sep) { |
50 | const std::chrono::duration<double> d = Clock::now() - start_timestamp_; |
51 | const double elapsed = d.count() * 1e6; |
52 | |
53 | #ifdef _WIN32 |
54 | const unsigned int pid = GetCurrentProcessId(); |
55 | const unsigned int tid = GetCurrentThreadId(); |
56 | #else |
57 | const unsigned int pid = getpid(); |
58 | const unsigned int tid = std::hash<pthread_t>{}(pthread_self()); |
59 | #endif // _WIN32 |
60 | |
61 | fprintf( |
62 | log_file_, |
63 | "{ \"name\": \"%s\", \"ph\": \"%c\", \"pid\": %u, \"tid\": %u, \"ts\": %.0f }%c\n" , |
64 | name, |
65 | ph, |
66 | pid, |
67 | tid, |
68 | elapsed, |
69 | sep); |
70 | } |
71 | |
72 | } // namespace inst |
73 | } // namespace cuda |
74 | } // namespace fuser |
75 | } // namespace jit |
76 | } // namespace torch |
77 | |