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