1 | #include <c10/macros/Export.h> |
---|---|
2 | #include <ittnotify.h> |
3 | #include <torch/csrc/profiler/stubs/base.h> |
4 | |
5 | namespace torch { |
6 | namespace profiler { |
7 | __itt_domain* _itt_domain = __itt_domain_create("PyTorch"); |
8 | |
9 | TORCH_API bool itt_is_available() { |
10 | return torch::profiler::impl::ittStubs()->enabled(); |
11 | } |
12 | |
13 | TORCH_API void itt_range_push(const char* msg) { |
14 | __itt_string_handle* hsMsg = __itt_string_handle_create(msg); |
15 | __itt_task_begin(_itt_domain, __itt_null, __itt_null, hsMsg); |
16 | } |
17 | |
18 | TORCH_API void itt_range_pop() { |
19 | __itt_task_end(_itt_domain); |
20 | } |
21 | |
22 | TORCH_API void itt_mark(const char* msg) { |
23 | __itt_string_handle* hsMsg = __itt_string_handle_create(msg); |
24 | __itt_task_begin(_itt_domain, __itt_null, __itt_null, hsMsg); |
25 | __itt_task_end(_itt_domain); |
26 | } |
27 | } // namespace profiler |
28 | } // namespace torch |
29 |