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_TRACEME_RECORDER_H_ 16 #define TENSORFLOW_CORE_PROFILER_INTERNAL_TRACEME_RECORDER_H_ 17 18 #include <stddef.h> 19 20 #include <atomic> 21 #include <vector> 22 23 #include "absl/container/flat_hash_map.h" 24 #include "tensorflow/core/platform/macros.h" 25 #include "tensorflow/core/platform/mutex.h" 26 #include "tensorflow/core/platform/thread_annotations.h" 27 #include "tensorflow/core/platform/types.h" 28 29 namespace tensorflow { 30 namespace profiler { 31 namespace internal { 32 33 // Current trace level. 34 // Static atomic so TraceMeRecorder::Active can be fast and non-blocking. 35 // Modified by TraceMeRecorder singleton when tracing starts/stops. 36 extern std::atomic<int> g_trace_level; 37 38 } // namespace internal 39 40 // TraceMeRecorder is a singleton repository of TraceMe events. 41 // It can be safely and cheaply appended to by multiple threads. 42 // 43 // Start() and Stop() must be called in pairs, Stop() returns the events added 44 // since the previous Start(). 45 // 46 // This is the backend for TraceMe instrumentation. 47 // The profiler starts the recorder, the TraceMe destructor records complete 48 // events. TraceMe::ActivityStart records begin events, and TraceMe::ActivityEnd 49 // records end events. The profiler then stops the recorder and finds start/end 50 // pairs. (Unpaired start/end events are discarded at that point). 51 class TraceMeRecorder { 52 public: 53 // An Event is either the start of a TraceMe, the end of a TraceMe, or both. 54 // Times are in ns since the Unix epoch. 55 struct Event { 56 uint64 activity_id; 57 string name; 58 uint64 start_time; // 0 = missing 59 uint64 end_time; // 0 = missing 60 }; 61 struct ThreadInfo { 62 int32 tid; 63 string name; 64 }; 65 struct ThreadEvents { 66 ThreadInfo thread; 67 std::vector<Event> events; 68 }; 69 using Events = std::vector<ThreadEvents>; 70 71 // Starts recording of TraceMe(). 72 // Only traces <= level will be recorded. 73 // Level must be >= 0. If level is 0, no traces will be recorded. Start(int level)74 static bool Start(int level) { return Get()->StartRecording(level); } 75 76 // Stops recording and returns events recorded since Start(). 77 // Events passed to Record after Stop has started will be dropped. Stop()78 static Events Stop() { return Get()->StopRecording(); } 79 80 // Returns whether we're currently recording. Racy, but cheap! 81 static inline bool Active(int level = 1) { 82 return internal::g_trace_level.load(std::memory_order_acquire) >= level; 83 } 84 85 // Default value for trace_level_ when tracing is disabled 86 static constexpr int kTracingDisabled = -1; 87 88 // Records an event. Non-blocking. 89 static void Record(Event event); 90 91 // Returns an activity_id for TraceMe::ActivityStart. 92 static uint64 NewActivityId(); 93 94 private: 95 class ThreadLocalRecorder; 96 97 // Returns singleton. 98 static TraceMeRecorder* Get(); 99 100 TraceMeRecorder() = default; 101 102 TF_DISALLOW_COPY_AND_ASSIGN(TraceMeRecorder); 103 104 void RegisterThread(int32 tid, ThreadLocalRecorder* thread); 105 void UnregisterThread(int32 tid); 106 107 bool StartRecording(int level); 108 Events StopRecording(); 109 110 // Gathers events from all active threads, and clears their buffers. 111 Events Clear() EXCLUSIVE_LOCKS_REQUIRED(mutex_); 112 113 mutex mutex_; 114 // Map of the static container instances (thread_local storage) for each 115 // thread. While active, a ThreadLocalRecorder stores trace events. 116 absl::flat_hash_map<int32, ThreadLocalRecorder*> threads_ GUARDED_BY(mutex_); 117 // Events from threads that died during recording. 118 TraceMeRecorder::Events orphaned_events_ GUARDED_BY(mutex_); 119 }; 120 121 } // namespace profiler 122 } // namespace tensorflow 123 #endif // TENSORFLOW_CORE_PROFILER_INTERNAL_TRACEME_RECORDER_H_ 124