1 /* Copyright 2020 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_CONNECTED_TRACEME_H_ 16 #define TENSORFLOW_CORE_PROFILER_LIB_CONNECTED_TRACEME_H_ 17 18 #include <string> 19 #include <utility> 20 21 #include "absl/strings/string_view.h" 22 #include "absl/types/optional.h" 23 #include "tensorflow/core/profiler/lib/context_types.h" 24 #include "tensorflow/core/profiler/lib/traceme.h" 25 #include "tensorflow/core/profiler/lib/traceme_encode.h" 26 27 namespace tensorflow { 28 namespace profiler { 29 30 /* 31 * TraceMeProducer and TraceMeConsumer are used to correlate TraceMe events on 32 * different threads. TraceMeProducer generates the context information to be 33 * passed to TraceMeConsumer, which consists of the context id and optionally 34 * the context type. They may be provided by the user. Then, the events of the 35 * same context information can be correlated during the analysis. 36 * 37 * Example Usages: 38 * (1) Using the user-provided context type and id. The user is responsible for 39 * providing the same context type and id to TraceMeProducer and 40 * TraceMeConsumer. 41 * [Producer Thread] 42 * // user_context_id is provided by the user. 43 * TraceMeProducer producer( 44 * [&] { return TraceMeEncode("op_dispatch", {{"op_type", "matmul"}}); }, 45 * ContextType::kTfExecutor, user_context_id); 46 * [Consumer Thread] 47 * // user_context_id is provided by the user. 48 * TraceMeConsumer consumer( 49 * [&] { return "op_execute"; }, ContextType::kTfExecutor, user_context_id); 50 * 51 * (2) Using the user-provided context type and generic id. The user is 52 * responsible for passing the TraceMeProducer's context id to 53 * TraceMeConsumer as well as providing the same context type to 54 * TraceMeProducer and TraceMeConsumer. 55 * [Producer Thread] 56 * TraceMeProducer producer( 57 * [&] { return TraceMeEncode("op_dispatch", {{"op_type", "matmul"}}); }, 58 * ContextType::kTfExecutor); 59 * context_id = producer.GetContextId(); 60 * // Pass context_id to the consumer thread. 61 * [Consumer Thread] 62 * // context_id is passed from the producer thread. 63 * TraceMeConsumer consumer( 64 * [&] { return "op_execute"; }, ContextType::kTfExecutor, context_id); 65 * 66 * (3) Using the generic context information. The user is responsible for 67 * passing the TraceMeProducer's context id to TraceMeConsumer. 68 * [Producer Thread] 69 * TraceMeProducer producer( 70 * [&] { return TraceMeEncode("op_dispatch", {{"op_type", "matmul"}}); }); 71 * context_id = producer.GetContextId(); 72 * // Pass context_id to the consumer thread. 73 * [Consumer Thread] 74 * // context_id is passed from the producer thread. 75 * TraceMeConsumer consumer([&] { return "op_execute"; }, context_id); 76 */ 77 class TraceMeProducer { 78 public: 79 template <typename NameT> 80 explicit TraceMeProducer(NameT&& name, 81 ContextType context_type = ContextType::kGeneric, 82 absl::optional<uint64> context_id = absl::nullopt, 83 int level = 2) 84 : context_id_(context_id.has_value() ? context_id.value() 85 : TraceMe::NewActivityId()), 86 trace_me_(std::forward<NameT>(name), level) { 87 trace_me_.AppendMetadata([&] { 88 return TraceMeEncode({{"_pt", context_type}, {"_p", context_id_}}); 89 }); 90 } 91 GetContextId()92 uint64 GetContextId() const { return context_id_; } 93 94 private: 95 uint64 context_id_; 96 TraceMe trace_me_; 97 }; 98 99 class TraceMeConsumer { 100 public: 101 template <typename NameT> 102 TraceMeConsumer(NameT&& name, ContextType context_type, uint64 context_id, 103 int level = 2) trace_me_(std::forward<NameT> (name),level)104 : trace_me_(std::forward<NameT>(name), level) { 105 trace_me_.AppendMetadata([&] { 106 return TraceMeEncode({{"_ct", context_type}, {"_c", context_id}}); 107 }); 108 } 109 110 template <typename NameT> 111 TraceMeConsumer(NameT&& name, uint64 context_id, int level = 2) TraceMeConsumer(std::forward<NameT> (name),ContextType::kGeneric,context_id,level)112 : TraceMeConsumer(std::forward<NameT>(name), ContextType::kGeneric, 113 context_id, level) {} 114 115 private: 116 TraceMe trace_me_; 117 }; 118 119 } // namespace profiler 120 } // namespace tensorflow 121 122 #endif // TENSORFLOW_CORE_PROFILER_LIB_CONNECTED_TRACEME_H_ 123