1 /* Copyright 2018 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 #ifndef TENSORFLOW_CORE_PROFILER_INTERNAL_CPU_TRACEME_RECORDER_H_ 16 #define TENSORFLOW_CORE_PROFILER_INTERNAL_CPU_TRACEME_RECORDER_H_ 17 18 #include <atomic> 19 #include <deque> 20 #include <memory> 21 #include <string> 22 #include <vector> 23 24 #include "absl/container/flat_hash_map.h" 25 #include "tensorflow/core/platform/macros.h" 26 #include "tensorflow/core/platform/mutex.h" 27 #include "tensorflow/core/platform/thread_annotations.h" 28 #include "tensorflow/core/platform/types.h" 29 30 namespace tensorflow { 31 namespace profiler { 32 namespace internal { 33 34 // Current trace level. 35 // Static atomic so TraceMeRecorder::Active can be fast and non-blocking. 36 // Modified by TraceMeRecorder singleton when tracing starts/stops. 37 TF_EXPORT extern std::atomic<int> g_trace_level; 38 39 } // namespace internal 40 41 // TraceMeRecorder is a singleton repository of TraceMe events. 42 // It can be safely and cheaply appended to by multiple threads. 43 // 44 // Start() and Stop() must be called in pairs, Stop() returns the events added 45 // since the previous Start(). 46 // 47 // This is the backend for TraceMe instrumentation. 48 // The profiler starts the recorder, the TraceMe destructor records complete 49 // events. TraceMe::ActivityStart records start events, and TraceMe::ActivityEnd 50 // records end events. The profiler then stops the recorder and finds start/end 51 // pairs. (Unpaired start/end events are discarded at that point). 52 class TraceMeRecorder { 53 public: 54 // An Event is either the start of a TraceMe, the end of a TraceMe, or both. 55 // Times are in ns since the Unix epoch. 56 // A negative time encodes the activity_id used to pair up the start of an 57 // event with its end. 58 struct Event { IsCompleteEvent59 bool IsComplete() const { return start_time > 0 && end_time > 0; } IsStartEvent60 bool IsStart() const { return end_time < 0; } IsEndEvent61 bool IsEnd() const { return start_time < 0; } 62 ActivityIdEvent63 int64 ActivityId() const { 64 if (IsStart()) return -end_time; 65 if (IsEnd()) return -start_time; 66 return 1; // complete 67 } 68 69 std::string name; 70 int64 start_time; 71 int64 end_time; 72 }; 73 struct ThreadInfo { 74 uint32 tid; 75 std::string name; 76 }; 77 struct ThreadEvents { 78 ThreadInfo thread; 79 std::deque<Event> events; 80 }; 81 using Events = std::vector<ThreadEvents>; 82 83 // Starts recording of TraceMe(). 84 // Only traces <= level will be recorded. 85 // Level must be >= 0. If level is 0, no traces will be recorded. Start(int level)86 static bool Start(int level) { return Get()->StartRecording(level); } 87 88 // Stops recording and returns events recorded since Start(). 89 // Events passed to Record after Stop has started will be dropped. Stop()90 static Events Stop() { return Get()->StopRecording(); } 91 92 // Returns whether we're currently recording. Racy, but cheap! 93 static inline bool Active(int level = 1) { 94 return internal::g_trace_level.load(std::memory_order_acquire) >= level; 95 } 96 97 // Default value for trace_level_ when tracing is disabled 98 static constexpr int kTracingDisabled = -1; 99 100 // Records an event. Non-blocking. 101 static void Record(Event&& event); 102 103 // Returns an activity_id for TraceMe::ActivityStart. 104 static int64 NewActivityId(); 105 106 private: 107 class ThreadLocalRecorder; 108 class ThreadLocalRecorderWrapper; 109 110 // Returns singleton. 111 static TraceMeRecorder* Get(); 112 113 TraceMeRecorder() = default; 114 115 TF_DISALLOW_COPY_AND_ASSIGN(TraceMeRecorder); 116 117 void RegisterThread(uint32 tid, std::shared_ptr<ThreadLocalRecorder> thread); 118 void UnregisterThread(uint32 tid); 119 120 bool StartRecording(int level); 121 Events StopRecording(); 122 123 // Clears events from all active threads that were added due to Record 124 // racing with StopRecording. 125 void Clear() TF_EXCLUSIVE_LOCKS_REQUIRED(mutex_); 126 127 // Gathers events from all active threads, and clears their buffers. 128 TF_MUST_USE_RESULT Events Consume() TF_EXCLUSIVE_LOCKS_REQUIRED(mutex_); 129 130 mutex mutex_; 131 // A ThreadLocalRecorder stores trace events. Ownership is shared with 132 // ThreadLocalRecorderWrapper, which is allocated in thread_local storage. 133 // ThreadLocalRecorderWrapper creates the ThreadLocalRecorder and registers it 134 // with TraceMeRecorder on the first TraceMe executed on a thread while 135 // tracing is active. If the thread is destroyed during tracing, the 136 // ThreadLocalRecorder is marked inactive but remains alive until tracing 137 // stops so the events can be retrieved. 138 absl::flat_hash_map<uint32, std::shared_ptr<ThreadLocalRecorder>> threads_ 139 TF_GUARDED_BY(mutex_); 140 }; 141 142 } // namespace profiler 143 } // namespace tensorflow 144 145 #endif // TENSORFLOW_CORE_PROFILER_INTERNAL_CPU_TRACEME_RECORDER_H_ 146