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