aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/platform/tracing.h
blob: 2b53a64cf19907b48f7d19a7eede011421063ee6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
#ifndef TENSORFLOW_PLATFORM_TRACING_H_
#define TENSORFLOW_PLATFORM_TRACING_H_

// Tracing interface

#include <map>
#include <memory>

#include "tensorflow/core/platform/port.h"  // Must be first
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/platform/thread_annotations.h"

namespace tensorflow {

class NodeExecStats;
class StepStats;

class StepStatsCollector {
 public:
  explicit StepStatsCollector(StepStats* ss);

  void Save(const string& device, NodeExecStats* nt);

  void Swap(StepStats* ss);

 private:
  friend class StepStatsMgr;
  mutex mu_;
  StepStats* step_stats_ GUARDED_BY(mu_);
};

namespace port {

class Tracing {
 public:
  // This enumeration contains the identifiers of all TensorFlow
  // threadscape events and code regions.  Threadscape assigns its
  // own identiers at runtime when we register our events and we
  // cannot know in advance what IDs it will choose.  The "RecordEvent"
  // method and "ScopedActivity" use these event IDs for consistency
  // and remap them to threadscape IDs at runtime.  This enum is limited
  // to 64 values since we use a bitmask to configure which events are
  // enabled.  It must also be kept in step with the code in
  // "Tracing::EventCategoryString".
  enum EventCategory {
    kScheduleClosure = 0,
    kRunClosure = 1,
    kCompute = 2,
    kEventCategoryMax = 3  // sentinel - keep last
  };
  // Note: We currently only support up to 64 categories.
  static_assert(kEventCategoryMax <= 64, "only support up to 64 events");

  // Called by main programs to initialize tracing facilities
  static void Initialize();

  // Return the pathname of the directory where we are writing log files.
  static const char* LogDir();

  // Returns a non-zero identifier which can be used to correlate
  // related events.
  static inline uint64 UniqueId();

  // Returns true if a trace is in progress.  Can be used to reduce tracing
  // overheads in fast-path code.
  static inline bool IsActive();

  // Associate name with the current thread.
  static void RegisterCurrentThread(const char* name);

  // Posts an event with the supplied category and arg.
  static void RecordEvent(EventCategory category, uint64 arg);

  // Traces a region of code.  Posts a tracing "EnterCodeRegion" event
  // when created and an "ExitCodeRegion" event when destroyed.
  class ScopedActivity {
   public:
    explicit ScopedActivity(EventCategory category, uint64 arg);
    ~ScopedActivity();

   private:
    const bool enabled_;
    const int32 region_id_;

    TF_DISALLOW_COPY_AND_ASSIGN(ScopedActivity);
  };

  // Trace collection engine can be registered with this module.
  // If no engine is registered, ScopedAnnotation and TraceMe are no-ops.
  class Engine;
  static void RegisterEngine(Engine*);

  // Forward declaration of the GPU utility classes.
  class ScopedAnnotation;
  class TraceMe;

 private:
  friend class TracingTest;

  static void RegisterEvent(EventCategory id, const char* name);
  static const char* EventCategoryString(EventCategory category);

  //
  // Parses event mask expressions in 'value' of the form:
  //   expr ::= <term> (,<term>)*
  //   term ::= <event> | "!" <event>
  //   event ::= "ALL" | <wait_event> | <other_event>
  //   wait_event ::= "ENewSession" | "ECloseSession" | ...
  //   other_event ::= "Send" | "Wait" | ...
  // ALL denotes all events, <event> turns on tracing for this event, and
  // !<event> turns off tracing for this event.
  // If the expression can be parsed correctly it returns true and sets
  // the event_mask_. Otherwise it returns false and the event_mask_ is left
  // unchanged.
  static bool ParseEventMask(const char* flagname, const string& value);

  // Bit mask of enabled trace categories.
  static uint64 event_mask_;

  // Records the mappings between Threadscape IDs and the "EventCategory" enum.
  static int32 category_id_[kEventCategoryMax];
  static std::map<string, int32>* name_map_;
};

// Trace collection engine that actually implements collection.
class Tracing::Engine {
 public:
  Engine() {}
  virtual ~Engine();

  // Represents an active annotation.
  class Annotation {
   public:
    Annotation() {}
    virtual ~Annotation();
  };

  // Represents an active trace.
  class Tracer {
   public:
    Tracer() {}
    virtual ~Tracer();
  };

 private:
  friend class ScopedAnnotation;
  friend class TraceMe;

  // Register the specified name as an annotation on the current thread.
  // Caller should delete the result to remove the annotation.
  // Annotations from the same thread are destroyed in a LIFO manner.
  // May return nullptr if annotations are not supported.
  virtual Annotation* PushAnnotation(StringPiece name) = 0;

  // Start tracing under the specified label.  Caller should delete the
  // result to stop tracing.
  // May return nullptr if tracing is not supported.
  virtual Tracer* StartTracing(StringPiece label) = 0;
};

// This class permits a user to apply annotation on kernels and memcpys
// when launching them. While an annotation is in scope, all activities
// within that scope get their names replaced by the annotation. The kernel
// name replacement is done when constructing the protobuf for sending out to
// a client (e.g., the stubby requestor) for both API and Activity records.
//
// Ownership: The creator of ScopedAnnotation assumes ownership of the object.
//
// Usage: {
//          ScopedAnnotation annotation("first set of kernels");
//          Kernel1<<<x,y>>>;
//          LaunchKernel2(); // Which eventually launches a cuda kernel.
//        }
// In the above scenario, the GPUProf UI would show 2 kernels with the name
// "first set of kernels" executing -- they will appear as the same kernel.
class Tracing::ScopedAnnotation {
 public:
  explicit ScopedAnnotation(StringPiece name);

 private:
  std::unique_ptr<Engine::Annotation> annotation_;
};

// TODO(opensource): clean up the scoped classes for GPU tracing.
// This class permits user-specified (CPU) tracing activities. A trace
// activity is started when an object of this class is created and stopped
// when the object is destroyed.
class Tracing::TraceMe {
 public:
  explicit TraceMe(StringPiece name);

 private:
  std::unique_ptr<Engine::Tracer> tracer_;
};

}  // namespace port
}  // namespace tensorflow

#if defined(PLATFORM_GOOGLE) && !defined(ANDROID) && !defined(__ANDROID__)
#include "tensorflow/core/platform/google/tracing_impl.h"
#else
#include "tensorflow/core/platform/default/tracing_impl.h"
#endif

#endif  // TENSORFLOW_PLATFORM_TRACING_H_