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