1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7    http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#ifndef TENSORFLOW_PLATFORM_TRACING_H_
17#define TENSORFLOW_PLATFORM_TRACING_H_
18
19// Tracing interface
20
21#include <atomic>
22#include <map>
23#include <memory>
24
25#include "tensorflow/core/lib/core/stringpiece.h"
26#include "tensorflow/core/lib/strings/strcat.h"
27#include "tensorflow/core/platform/macros.h"
28#include "tensorflow/core/platform/mutex.h"
29#include "tensorflow/core/platform/platform.h"
30#include "tensorflow/core/platform/types.h"
31
32namespace tensorflow {
33
34namespace port {
35
36class Tracing {
37 public:
38  // This enumeration contains the identifiers of all TensorFlow
39  // threadscape events and code regions.  Threadscape assigns its
40  // own identifiers at runtime when we register our events and we
41  // cannot know in advance what IDs it will choose.  The "RecordEvent"
42  // method and "ScopedActivity" use these event IDs for consistency
43  // and remap them to threadscape IDs at runtime.  This enum is limited
44  // to 64 values since we use a bitmask to configure which events are
45  // enabled.  It must also be kept in step with the code in
46  // "Tracing::EventCategoryString".
47  enum EventCategory {
48    kScheduleClosure = 0,
49    kRunClosure = 1,
50    kCompute = 2,
51    kEventCategoryMax = 3  // sentinel - keep last
52  };
53  // Note: We currently only support up to 64 categories.
54  static_assert(kEventCategoryMax <= 64, "only support up to 64 events");
55
56  // Called by main programs to initialize tracing facilities
57  static void Initialize();
58
59  // Return the pathname of the directory where we are writing log files.
60  static const char* LogDir();
61
62  // Returns a non-zero identifier which can be used to correlate
63  // related events.
64  static inline uint64 UniqueId();
65
66  // Returns true if a trace is in progress.  Can be used to reduce tracing
67  // overheads in fast-path code.
68  static inline bool IsActive();
69
70  // Associate name with the current thread.
71  static void RegisterCurrentThread(const char* name);
72
73  // Posts an event with the supplied category and arg.
74  static void RecordEvent(EventCategory category, uint64 arg);
75
76  // Traces a region of code.  Posts a tracing "EnterCodeRegion" event
77  // when created and an "ExitCodeRegion" event when destroyed.
78  class ScopedActivity {
79   public:
80    explicit ScopedActivity(EventCategory category, uint64 arg);
81    ~ScopedActivity();
82
83   private:
84#if defined(PLATFORM_GOOGLE)
85    const bool enabled_;
86    const int32 region_id_;
87#endif
88
89    TF_DISALLOW_COPY_AND_ASSIGN(ScopedActivity);
90  };
91
92  // Trace collection engine can be registered with this module.
93  // If no engine is registered, ScopedAnnotation and TraceMe are no-ops.
94  class Engine;
95  static void RegisterEngine(Engine*);
96
97  // Forward declaration of the GPU utility classes.
98  class ScopedAnnotation;
99  class TraceMe;
100
101 private:
102  friend class TracingTest;
103  friend class ScopedAnnotation;
104  friend class TraceMe;
105
106  static std::atomic<Tracing::Engine*> tracing_engine_;
107  static Tracing::Engine* engine() {
108    return tracing_engine_.load(std::memory_order_acquire);
109  }
110
111  static void RegisterEvent(EventCategory id, const char* name);
112  static const char* EventCategoryString(EventCategory category);
113
114  //
115  // Parses event mask expressions in 'value' of the form:
116  //   expr ::= <term> (,<term>)*
117  //   term ::= <event> | "!" <event>
118  //   event ::= "ALL" | <wait_event> | <other_event>
119  //   wait_event ::= "ENewSession" | "ECloseSession" | ...
120  //   other_event ::= "Send" | "Wait" | ...
121  // ALL denotes all events, <event> turns on tracing for this event, and
122  // !<event> turns off tracing for this event.
123  // If the expression can be parsed correctly it returns true and sets
124  // the event_mask_. Otherwise it returns false and the event_mask_ is left
125  // unchanged.
126  static bool ParseEventMask(const char* flagname, const string& value);
127
128  // Bit mask of enabled trace categories.
129  static uint64 event_mask_;
130
131  // Records the mappings between Threadscape IDs and the "EventCategory" enum.
132  static int32 category_id_[kEventCategoryMax];
133  static std::map<string, int32>* name_map_;
134};
135
136// Trace collection engine that actually implements collection.
137class Tracing::Engine {
138 public:
139  Engine() {}
140  virtual ~Engine();
141
142  // Returns true if Tracing is currently enabled.
143  virtual bool IsEnabled() const = 0;
144
145  // Represents an active annotation.
146  class Annotation {
147   public:
148    Annotation() {}
149    virtual ~Annotation();
150  };
151
152  // Represents an active trace.
153  class Tracer {
154   public:
155    Tracer() {}
156    virtual ~Tracer();
157  };
158
159 private:
160  friend class ScopedAnnotation;
161  friend class TraceMe;
162
163  // Register the specified name as an annotation on the current thread.
164  // Caller should delete the result to remove the annotation.
165  // Annotations from the same thread are destroyed in a LIFO manner.
166  // May return nullptr if annotations are not supported.
167  virtual Annotation* PushAnnotation(StringPiece name) = 0;
168
169  // Start tracing under the specified label. Caller should delete the result
170  // to stop tracing.
171  // May return nullptr if tracing is not supported.
172  virtual Tracer* StartTracing(StringPiece label, bool is_expensive) = 0;
173  // Same as above, but implementations can avoid copying the string.
174  virtual Tracer* StartTracing(string&& label, bool is_expensive) {
175    return StartTracing(StringPiece(label), is_expensive);
176  }
177
178  // Backwards compatibility one arg variants (assume is_expensive=true).
179  Tracer* StartTracing(StringPiece label) {
180    return StartTracing(label, /*is_expensive=*/true);
181  }
182  Tracer* StartTracing(string&& label) {
183    return StartTracing(StringPiece(label), /*is_expensive=*/true);
184  }
185};
186
187// This class permits a user to apply annotation on kernels and memcpys
188// when launching them. While an annotation is in scope, all activities
189// within that scope get their names replaced by the annotation. The kernel
190// name replacement is done when constructing the protobuf for sending out to
191// a client (e.g., the stubby requestor) for both API and Activity records.
192//
193// Ownership: The creator of ScopedAnnotation assumes ownership of the object.
194//
195// Usage: {
196//          ScopedAnnotation annotation("first set of kernels");
197//          Kernel1<<<x,y>>>;
198//          LaunchKernel2(); // Which eventually launches a cuda kernel.
199//        }
200// In the above scenario, the GPUProf UI would show 2 kernels with the name
201// "first set of kernels" executing -- they will appear as the same kernel.
202class Tracing::ScopedAnnotation {
203 public:
204  explicit ScopedAnnotation(StringPiece name);
205
206  // If tracing is enabled, set up an annotation with a label of
207  // "<name_part1>:<name_part2>".  Can be cheaper than the
208  // single-argument constructor because the concatenation of the
209  // label string is only done if tracing is enabled.
210  ScopedAnnotation(StringPiece name_part1, StringPiece name_part2);
211
212  // Returns true iff scoped annotations are active.
213  static bool Enabled() {
214    auto e = Tracing::engine();
215    return e && e->IsEnabled();
216  }
217
218 private:
219  std::unique_ptr<Engine::Annotation> annotation_;
220};
221
222// TODO(opensource): clean up the scoped classes for GPU tracing.
223// This class permits user-specified (CPU) tracing activities. A trace
224// activity is started when an object of this class is created and stopped
225// when the object is destroyed.
226class Tracing::TraceMe {
227 public:
228  explicit TraceMe(StringPiece name);
229  TraceMe(StringPiece name, bool is_expensive);
230
231  // If tracing is enabled, set up a traceMe with a label of
232  // "<name_part1>:<name_part2>".  This can be cheaper than the
233  // single-argument constructor because the concatenation of the
234  // label string is only done if tracing is enabled.
235  TraceMe(StringPiece name_part1, StringPiece name_part2);
236  TraceMe(StringPiece name_part1, StringPiece name_part2, bool is_expensive);
237
238 private:
239  std::unique_ptr<Engine::Tracer> tracer_;
240};
241
242inline Tracing::ScopedAnnotation::ScopedAnnotation(StringPiece name) {
243  auto e = Tracing::engine();
244  if (e && e->IsEnabled()) {
245    annotation_.reset(e->PushAnnotation(name));
246  }
247}
248
249inline Tracing::ScopedAnnotation::ScopedAnnotation(StringPiece name_part1,
250                                                   StringPiece name_part2) {
251  auto e = Tracing::engine();
252  if (e && e->IsEnabled()) {
253    annotation_.reset(
254        e->PushAnnotation(strings::StrCat(name_part1, ":", name_part2)));
255  }
256}
257
258inline Tracing::TraceMe::TraceMe(StringPiece name) : TraceMe(name, true) {}
259
260inline Tracing::TraceMe::TraceMe(StringPiece name, bool is_expensive) {
261  auto e = Tracing::engine();
262  if (e && e->IsEnabled()) {
263    tracer_.reset(e->StartTracing(name, is_expensive));
264  }
265}
266
267inline Tracing::TraceMe::TraceMe(StringPiece name_part1, StringPiece name_part2)
268    : TraceMe(name_part1, name_part2, true) {}
269
270inline Tracing::TraceMe::TraceMe(StringPiece name_part1, StringPiece name_part2,
271                                 bool is_expensive) {
272  auto e = Tracing::engine();
273  if (e && e->IsEnabled()) {
274    tracer_.reset(e->StartTracing(strings::StrCat(name_part1, ":", name_part2),
275                                  is_expensive));
276  }
277}
278
279}  // namespace port
280}  // namespace tensorflow
281
282#if defined(PLATFORM_GOOGLE)
283#include "tensorflow/core/platform/google/tracing_impl.h"
284#else
285#include "tensorflow/core/platform/default/tracing_impl.h"
286#endif
287
288#endif  // TENSORFLOW_PLATFORM_TRACING_H_
289