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
12namespace torch {
13namespace jit {
14namespace fuser {
15namespace cuda {
16namespace inst {
17
18Trace::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
40Trace::~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
49void 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