1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #ifndef TENSORFLOW_CORE_PLATFORM_TRACING_H_
17 #define TENSORFLOW_CORE_PLATFORM_TRACING_H_
18
19 // Tracing interface
20
21 #include <array>
22
23 #include "tensorflow/core/platform/macros.h"
24 #include "tensorflow/core/platform/platform.h"
25 #include "tensorflow/core/platform/stringpiece.h"
26 #include "tensorflow/core/platform/types.h"
27
28 namespace tensorflow {
29 namespace tracing {
30
31 // This enumeration contains the identifiers of all TensorFlow CPU profiler
32 // events. It must be kept in sync with the code in GetEventCategoryName().
33 enum struct EventCategory : unsigned {
34 kScheduleClosure = 0,
35 kRunClosure = 1,
36 kCompute = 2,
37 kNumCategories = 3 // sentinel - keep last
38 };
GetNumEventCategories()39 constexpr unsigned GetNumEventCategories() {
40 return static_cast<unsigned>(EventCategory::kNumCategories);
41 }
42 const char* GetEventCategoryName(EventCategory);
43
44 // Interface for CPU profiler events.
45 class EventCollector {
46 public:
~EventCollector()47 virtual ~EventCollector() {}
48 virtual void RecordEvent(uint64 arg) const = 0;
49 virtual void StartRegion(uint64 arg) const = 0;
50 virtual void StopRegion() const = 0;
51
52 // Annotates the current thread with a name.
53 static void SetCurrentThreadName(const char* name);
54 // Returns whether event collection is enabled.
55 static bool IsEnabled();
56
57 private:
58 friend void SetEventCollector(EventCategory, const EventCollector*);
59 friend const EventCollector* GetEventCollector(EventCategory);
60
61 static std::array<const EventCollector*, GetNumEventCategories()> instances_;
62 };
63 // Set the callback for RecordEvent and ScopedRegion of category.
64 // Not thread safe. Only call while EventCollector::IsEnabled returns false.
65 void SetEventCollector(EventCategory category, const EventCollector* collector);
66
67 // Returns the callback for RecordEvent and ScopedRegion of category if
68 // EventCollector::IsEnabled(), otherwise returns null.
GetEventCollector(EventCategory category)69 inline const EventCollector* GetEventCollector(EventCategory category) {
70 if (EventCollector::IsEnabled()) {
71 return EventCollector::instances_[static_cast<unsigned>(category)];
72 }
73 return nullptr;
74 }
75
76 // Returns a unique id to pass to RecordEvent/ScopedRegion. Never returns zero.
77 uint64 GetUniqueArg();
78
79 // Returns an id for name to pass to RecordEvent/ScopedRegion.
80 uint64 GetArgForName(StringPiece name);
81
82 // Records an atomic event through the currently registered EventCollector.
RecordEvent(EventCategory category,uint64 arg)83 inline void RecordEvent(EventCategory category, uint64 arg) {
84 if (auto collector = GetEventCollector(category)) {
85 collector->RecordEvent(arg);
86 }
87 }
88
89 // Records an event for the duration of the instance lifetime through the
90 // currently registered EventCollector.
91 class ScopedRegion {
92 public:
ScopedRegion(ScopedRegion && other)93 ScopedRegion(ScopedRegion&& other) noexcept // Move-constructible.
94 : collector_(other.collector_) {
95 other.collector_ = nullptr;
96 }
97
ScopedRegion(EventCategory category,uint64 arg)98 ScopedRegion(EventCategory category, uint64 arg)
99 : collector_(GetEventCollector(category)) {
100 if (collector_) {
101 collector_->StartRegion(arg);
102 }
103 }
104
105 // Same as ScopedRegion(category, GetUniqueArg()), but faster if
106 // EventCollector::IsEnabled() returns false.
ScopedRegion(EventCategory category)107 explicit ScopedRegion(EventCategory category)
108 : collector_(GetEventCollector(category)) {
109 if (collector_) {
110 collector_->StartRegion(GetUniqueArg());
111 }
112 }
113
114 // Same as ScopedRegion(category, GetArgForName(name)), but faster if
115 // EventCollector::IsEnabled() returns false.
ScopedRegion(EventCategory category,StringPiece name)116 ScopedRegion(EventCategory category, StringPiece name)
117 : collector_(GetEventCollector(category)) {
118 if (collector_) {
119 collector_->StartRegion(GetArgForName(name));
120 }
121 }
122
~ScopedRegion()123 ~ScopedRegion() {
124 if (collector_) {
125 collector_->StopRegion();
126 }
127 }
128
IsEnabled()129 bool IsEnabled() const { return collector_ != nullptr; }
130
131 private:
132 TF_DISALLOW_COPY_AND_ASSIGN(ScopedRegion);
133
134 const EventCollector* collector_;
135 };
136
137 // Return the pathname of the directory where we are writing log files.
138 const char* GetLogDir();
139
140 } // namespace tracing
141 } // namespace tensorflow
142
143 #if defined(PLATFORM_GOOGLE)
144 #include "tensorflow/core/platform/google/tracing_impl.h"
145 #else
146 #include "tensorflow/core/platform/default/tracing_impl.h"
147 #endif
148
149 #endif // TENSORFLOW_CORE_PLATFORM_TRACING_H_
150