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