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_TRACEME_ENCODE_H_
16 #define TENSORFLOW_CORE_PROFILER_LIB_TRACEME_ENCODE_H_
17
18 #include <string.h>
19
20 #include <initializer_list>
21 #include <string>
22
23 #include "absl/strings/match.h"
24 #include "absl/strings/str_cat.h"
25 #include "absl/strings/string_view.h"
26 #include "tensorflow/core/platform/logging.h"
27 #include "tensorflow/core/platform/macros.h"
28
29 namespace tensorflow {
30 namespace profiler {
31
32 #pragma clang diagnostic push
33 #pragma clang diagnostic ignored "-Wdangling-field"
34
35 // An argument passed to TraceMeEncode.
36 struct TraceMeArg {
37 // This constructor is required because absl::AlphaNum is non-copyable.
38 template <typename Value>
TraceMeArgTraceMeArg39 TraceMeArg(absl::string_view k, Value v) : key(k), value(v) {}
40
41 TF_DISALLOW_COPY_AND_ASSIGN(TraceMeArg);
42
43 absl::string_view key;
44 absl::AlphaNum value;
45 };
46
47 #pragma clang diagnostic pop
48
49 namespace traceme_internal {
50
51 // Copies the contents of str to the address pointed by out.
52 // Returns the address after the copy.
53 // REQUIRED: The address range [out, out + str.size()] must have been allocated.
Append(char * out,absl::string_view str)54 TF_ATTRIBUTE_ALWAYS_INLINE inline char* Append(char* out,
55 absl::string_view str) {
56 DCHECK(!absl::StrContains(str, "#"))
57 << "'#' is not a valid character in TraceMeEncode";
58 const size_t str_size = str.size();
59 if (TF_PREDICT_TRUE(str_size > 0)) {
60 memcpy(out, str.data(), str_size);
61 out += str_size;
62 }
63 return out;
64 }
65
66 // Appends args encoded as TraceMe metadata to name.
AppendArgs(std::string name,std::initializer_list<TraceMeArg> args)67 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string AppendArgs(
68 std::string name, std::initializer_list<TraceMeArg> args) {
69 if (TF_PREDICT_TRUE(args.size() > 0)) {
70 const auto old_size = name.size();
71 auto new_size = old_size + args.size() * 2 + 1;
72 for (const auto& arg : args) {
73 new_size += arg.key.size() + arg.value.size();
74 }
75 name.resize(new_size);
76 char* const begin = &name[0];
77 char* out = begin + old_size;
78 *out++ = '#';
79 for (const auto& arg : args) {
80 out = Append(out, arg.key);
81 *out++ = '=';
82 out = Append(out, arg.value.Piece());
83 *out++ = ',';
84 }
85 *(out - 1) = '#';
86 DCHECK_EQ(out, begin + new_size);
87 }
88 return name;
89 }
90
91 // Appends new_metadata to the metadata part of name.
AppendMetadata(std::string * name,absl::string_view new_metadata)92 TF_ATTRIBUTE_ALWAYS_INLINE inline void AppendMetadata(
93 std::string* name, absl::string_view new_metadata) {
94 if (!TF_PREDICT_FALSE(new_metadata.empty())) {
95 if (!name->empty() && name->back() == '#') { // name already has metadata
96 name->back() = ',';
97 if (TF_PREDICT_TRUE(new_metadata.front() == '#')) {
98 new_metadata.remove_prefix(1);
99 }
100 }
101 name->append(new_metadata.data(), new_metadata.size());
102 }
103 }
104
105 } // namespace traceme_internal
106
107 // Encodes an event name and arguments into TraceMe metadata.
108 // Use within a lambda to avoid expensive operations when tracing is disabled.
109 // Example Usage:
110 // TraceMe trace_me([value1]() {
111 // return TraceMeEncode("my_trace", {{"key1", value1}, {"key2", 42}});
112 // });
TraceMeEncode(std::string name,std::initializer_list<TraceMeArg> args)113 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode(
114 std::string name, std::initializer_list<TraceMeArg> args) {
115 return traceme_internal::AppendArgs(std::move(name), args);
116 }
TraceMeEncode(absl::string_view name,std::initializer_list<TraceMeArg> args)117 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode(
118 absl::string_view name, std::initializer_list<TraceMeArg> args) {
119 return traceme_internal::AppendArgs(std::string(name), args);
120 }
TraceMeEncode(const char * name,std::initializer_list<TraceMeArg> args)121 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode(
122 const char* name, std::initializer_list<TraceMeArg> args) {
123 return traceme_internal::AppendArgs(std::string(name), args);
124 }
125
126 // Encodes arguments into TraceMe metadata.
127 // Use within a lambda to avoid expensive operations when tracing is disabled.
128 // Example Usage:
129 // TraceMe trace_me("my_trace");
130 // ...
131 // trace_me.AppendMetadata([value1]() {
132 // return TraceMeEncode({{"key1", value1}, {"key2", 42}});
133 // });
TraceMeEncode(std::initializer_list<TraceMeArg> args)134 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode(
135 std::initializer_list<TraceMeArg> args) {
136 return traceme_internal::AppendArgs(std::string(), args);
137 }
138
139 // Concatenates op_name and op_type.
TraceMeOp(absl::string_view op_name,absl::string_view op_type)140 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeOp(
141 absl::string_view op_name, absl::string_view op_type) {
142 return absl::StrCat(op_name, ":", op_type);
143 }
144
TraceMeOp(const char * op_name,const char * op_type)145 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeOp(const char* op_name,
146 const char* op_type) {
147 return absl::StrCat(op_name, ":", op_type);
148 }
149
TraceMeOp(std::string && op_name,absl::string_view op_type)150 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeOp(
151 std::string&& op_name, absl::string_view op_type) {
152 absl::StrAppend(&op_name, ":", op_type);
153 return op_name;
154 }
155
156 // Concatenates op_name and op_type.
TraceMeOpOverride(absl::string_view op_name,absl::string_view op_type)157 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeOpOverride(
158 absl::string_view op_name, absl::string_view op_type) {
159 return absl::StrCat("#tf_op=", op_name, ":", op_type, "#");
160 }
161
TraceMeOpOverride(const char * op_name,const char * op_type)162 TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeOpOverride(
163 const char* op_name, const char* op_type) {
164 return absl::StrCat("#tf_op=", op_name, ":", op_type, "#");
165 }
166
167 } // namespace profiler
168 } // namespace tensorflow
169
170 #endif // TENSORFLOW_CORE_PROFILER_LIB_TRACEME_ENCODE_H_
171