1 files changed, 205 insertions, 0 deletions
diff --git a/tensorflow/core/platform/tracing.h b/tensorflow/core/platform/tracing.h
new file mode 100644
index 0000000000..2b53a64cf1
--- /dev/null
+++ b/tensorflow/core/platform/tracing.h
@@ -0,0 +1,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_