1 /* Copyright 2019 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_LIB_SCOPED_ANNOTATION_H_ 16 #define TENSORFLOW_CORE_PROFILER_LIB_SCOPED_ANNOTATION_H_ 17 18 #include <stddef.h> 19 20 #include <atomic> 21 #include <utility> 22 23 #include "absl/strings/string_view.h" 24 #include "tensorflow/core/platform/macros.h" 25 #include "tensorflow/core/platform/platform.h" 26 #include "tensorflow/core/platform/types.h" 27 #if !defined(IS_MOBILE_PLATFORM) 28 #include "tensorflow/core/profiler/internal/cpu/annotation_stack.h" 29 #endif 30 31 namespace tensorflow { 32 namespace profiler { 33 34 // Adds an annotation to all activities for the duration of the instance 35 // lifetime through the currently registered TraceCollector. 36 // 37 // Usage: { 38 // ScopedAnnotation annotation("my kernels"); 39 // Kernel1<<<x,y>>>; 40 // LaunchKernel2(); // Launches a CUDA kernel. 41 // } 42 // This will add 'my kernels' to both kernels in the profiler UI 43 class ScopedAnnotation { 44 public: ScopedAnnotation(absl::string_view name)45 explicit ScopedAnnotation(absl::string_view name) { 46 #if !defined(IS_MOBILE_PLATFORM) 47 if (TF_PREDICT_FALSE(AnnotationStack::IsEnabled())) { 48 old_length_ = AnnotationStack::PushAnnotation(name); 49 } 50 #endif 51 } 52 ScopedAnnotation(const char * name)53 explicit ScopedAnnotation(const char* name) 54 : ScopedAnnotation(absl::string_view(name)) {} 55 ScopedAnnotation(const string & name)56 explicit ScopedAnnotation(const string& name) { 57 #if !defined(IS_MOBILE_PLATFORM) 58 if (TF_PREDICT_FALSE(AnnotationStack::IsEnabled())) { 59 old_length_ = AnnotationStack::PushAnnotation(name); 60 } 61 #endif 62 } 63 ScopedAnnotation(string && name)64 explicit ScopedAnnotation(string&& name) { 65 #if !defined(IS_MOBILE_PLATFORM) 66 if (TF_PREDICT_FALSE(AnnotationStack::IsEnabled())) { 67 old_length_ = AnnotationStack::PushAnnotation(std::move(name)); 68 } 69 #endif 70 } 71 72 template <typename NameGeneratorT> ScopedAnnotation(NameGeneratorT name_generator)73 explicit ScopedAnnotation(NameGeneratorT name_generator) { 74 #if !defined(IS_MOBILE_PLATFORM) 75 if (TF_PREDICT_FALSE(AnnotationStack::IsEnabled())) { 76 old_length_ = AnnotationStack::PushAnnotation(name_generator()); 77 } 78 #endif 79 } 80 81 // Pops the name passed in the constructor from the current annotation. ~ScopedAnnotation()82 ~ScopedAnnotation() { 83 // TODO(b/137971921): without this memory fence, two presubmit tests will 84 // fail probably due to compiler in that presubmit config. 85 std::atomic_thread_fence(std::memory_order_acquire); 86 #if !defined(IS_MOBILE_PLATFORM) 87 if (TF_PREDICT_FALSE(old_length_ != kInvalidLength)) { 88 AnnotationStack::PopAnnotation(old_length_); 89 } 90 #endif 91 } 92 IsEnabled()93 static bool IsEnabled() { 94 #if !defined(IS_MOBILE_PLATFORM) 95 return AnnotationStack::IsEnabled(); 96 #else 97 return false; 98 #endif 99 } 100 101 private: 102 // signals that annotation is disabled at the constructor. 103 static constexpr size_t kInvalidLength = static_cast<size_t>(-1); 104 105 TF_DISALLOW_COPY_AND_ASSIGN(ScopedAnnotation); 106 107 size_t old_length_ = kInvalidLength; 108 }; 109 110 } // namespace profiler 111 } // namespace tensorflow 112 113 #endif // TENSORFLOW_CORE_PROFILER_LIB_SCOPED_ANNOTATION_H_ 114