1 /* Copyright 2019 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_UTILS_XPLANE_VISITOR_H_
16 #define TENSORFLOW_CORE_PROFILER_UTILS_XPLANE_VISITOR_H_
17
18 #include <stddef.h>
19
20 #include <functional>
21 #include <string>
22 #include <vector>
23
24 #include "absl/container/flat_hash_map.h"
25 #include "absl/strings/string_view.h"
26 #include "absl/types/optional.h"
27 #include "tensorflow/core/platform/types.h"
28 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
29 #include "tensorflow/core/profiler/utils/timespan.h"
30
31 namespace tensorflow {
32 namespace profiler {
33
34 class XPlaneVisitor;
35
36 class XStatVisitor {
37 public:
38 // REQUIRED: plane and stat cannot be nullptr.
39 XStatVisitor(const XPlaneVisitor* plane, const XStat* stat);
40
41 // REQUIRED: plane, stat and metadata cannot be nullptr.
42 XStatVisitor(const XPlaneVisitor* plane, const XStat* stat,
43 const XStatMetadata* metadata, absl::optional<int64_t> type);
44
Id()45 int64_t Id() const { return stat_->metadata_id(); }
46
Name()47 absl::string_view Name() const { return metadata_->name(); }
48
Type()49 absl::optional<int64_t> Type() const { return type_; }
50
Description()51 absl::string_view Description() const { return metadata_->description(); }
52
ValueCase()53 XStat::ValueCase ValueCase() const { return stat_->value_case(); }
54
BoolValue()55 bool BoolValue() const { return static_cast<bool>(IntValue()); }
56
IntValue()57 int64_t IntValue() const { return stat_->int64_value(); }
58
UintValue()59 uint64 UintValue() const { return stat_->uint64_value(); }
60
BytesValue()61 absl::string_view BytesValue() const { return stat_->bytes_value(); }
62
IntOrUintValue()63 uint64 IntOrUintValue() const {
64 return ValueCase() == XStat::kUint64Value ? UintValue()
65 : static_cast<uint64>(IntValue());
66 }
67
DoubleValue()68 double DoubleValue() const { return stat_->double_value(); }
69
70 // Returns a string view.
71 // REQUIRED: the value type should be string type or reference type.
72 absl::string_view StrOrRefValue() const;
73
RawStat()74 const XStat& RawStat() const { return *stat_; }
75
76 // Return a string representation of all value type.
77 std::string ToString() const;
78
79 private:
80 const XStat* stat_;
81 const XStatMetadata* metadata_;
82 const XPlaneVisitor* plane_;
83 absl::optional<int64_t> type_;
84 };
85
86 template <class T>
87 class XStatsOwner {
88 public:
89 // REQUIRED: plane and stats_owner cannot be nullptr.
XStatsOwner(const XPlaneVisitor * plane,const T * stats_owner)90 XStatsOwner(const XPlaneVisitor* plane, const T* stats_owner)
91 : plane_(plane), stats_owner_(stats_owner) {}
92
93 // For each stat, call the specified lambda.
94 template <typename ForEachStatFunc>
ForEachStat(ForEachStatFunc && for_each_stat)95 void ForEachStat(ForEachStatFunc&& for_each_stat) const {
96 for (const XStat& stat : stats_owner_->stats()) {
97 for_each_stat(XStatVisitor(plane_, &stat));
98 }
99 }
100
101 // Shortcut to get a specific stat type, nullopt if absent.
102 // This function performs a linear search for the requested stat value.
103 // Prefer ForEachStat above when multiple stat values are necessary.
104 absl::optional<XStatVisitor> GetStat(int64_t stat_type) const;
105
106 // Same as above that skips searching for the stat.
GetStat(int64_t stat_type,const XStatMetadata & stat_metadata)107 absl::optional<XStatVisitor> GetStat(
108 int64_t stat_type, const XStatMetadata& stat_metadata) const {
109 for (const XStat& stat : stats_owner_->stats()) {
110 if (stat.metadata_id() == stat_metadata.id()) {
111 return XStatVisitor(plane_, &stat, &stat_metadata, stat_type);
112 }
113 }
114 return absl::nullopt; // type does not exist in this owner.
115 }
116
117 protected:
plane()118 const XPlaneVisitor* plane() const { return plane_; }
stats_owner()119 const T* stats_owner() const { return stats_owner_; }
120
121 private:
122 const XPlaneVisitor* plane_;
123 const T* stats_owner_;
124 };
125
126 class XEventMetadataVisitor : public XStatsOwner<XEventMetadata> {
127 public:
128 // REQUIRED: plane and metadata cannot be nullptr.
XEventMetadataVisitor(const XPlaneVisitor * plane,const XEventMetadata * metadata)129 XEventMetadataVisitor(const XPlaneVisitor* plane,
130 const XEventMetadata* metadata)
131 : XStatsOwner(plane, metadata) {}
132
Id()133 int64_t Id() const { return metadata()->id(); }
134
Name()135 absl::string_view Name() const { return metadata()->name(); }
136
HasDisplayName()137 bool HasDisplayName() const { return !metadata()->display_name().empty(); }
138
DisplayName()139 absl::string_view DisplayName() const { return metadata()->display_name(); }
140
141 // For each child event metadata, call the specified lambda.
142 template <typename ForEachChildFunc>
143 void ForEachChild(ForEachChildFunc&& for_each_child) const;
144
145 private:
metadata()146 const XEventMetadata* metadata() const { return stats_owner(); }
147 };
148
149 class XEventVisitor : public XStatsOwner<XEvent> {
150 public:
151 // REQUIRED: plane, line and event cannot be nullptr.
152 XEventVisitor(const XPlaneVisitor* plane, const XLine* line,
153 const XEvent* event);
154
Plane()155 const XPlaneVisitor& Plane() const { return *plane_; }
156
RawEvent()157 const XEvent& RawEvent() const { return *event_; }
158
Id()159 int64_t Id() const { return event_->metadata_id(); }
160
Name()161 absl::string_view Name() const { return metadata_->name(); }
162
Type()163 absl::optional<int64_t> Type() const { return type_; }
164
HasDisplayName()165 bool HasDisplayName() const { return !metadata_->display_name().empty(); }
166
DisplayName()167 absl::string_view DisplayName() const { return metadata_->display_name(); }
168
OffsetNs()169 double OffsetNs() const { return PicoToNano(event_->offset_ps()); }
170
OffsetPs()171 int64_t OffsetPs() const { return event_->offset_ps(); }
172
LineTimestampNs()173 int64_t LineTimestampNs() const { return line_->timestamp_ns(); }
174
TimestampNs()175 int64_t TimestampNs() const { return line_->timestamp_ns() + OffsetNs(); }
176
TimestampPs()177 int64_t TimestampPs() const {
178 return NanoToPico(line_->timestamp_ns()) + event_->offset_ps();
179 }
180
DurationNs()181 double DurationNs() const { return PicoToNano(event_->duration_ps()); }
182
DurationPs()183 int64_t DurationPs() const { return event_->duration_ps(); }
184
EndOffsetPs()185 int64_t EndOffsetPs() const {
186 return event_->offset_ps() + event_->duration_ps();
187 }
188
EndTimestampNs()189 int64_t EndTimestampNs() const { return TimestampNs() + DurationNs(); }
190
EndTimestampPs()191 int64_t EndTimestampPs() const { return TimestampPs() + DurationPs(); }
192
NumOccurrences()193 int64_t NumOccurrences() const { return event_->num_occurrences(); }
194
195 bool operator<(const XEventVisitor& other) const {
196 return GetTimespan() < other.GetTimespan();
197 }
198
metadata()199 const XEventMetadata* metadata() const { return metadata_; }
200
Metadata()201 XEventMetadataVisitor Metadata() const {
202 return XEventMetadataVisitor(plane_, metadata_);
203 }
204
GetTimespan()205 Timespan GetTimespan() const { return Timespan(TimestampPs(), DurationPs()); }
206
207 private:
208 const XPlaneVisitor* plane_;
209 const XLine* line_;
210 const XEvent* event_;
211 const XEventMetadata* metadata_;
212 absl::optional<int64_t> type_;
213 };
214
215 class XLineVisitor {
216 public:
217 // REQUIRED: plane and line cannot be nullptr.
XLineVisitor(const XPlaneVisitor * plane,const XLine * line)218 XLineVisitor(const XPlaneVisitor* plane, const XLine* line)
219 : plane_(plane), line_(line) {}
220
Id()221 int64_t Id() const { return line_->id(); }
222
DisplayId()223 int64_t DisplayId() const {
224 return line_->display_id() ? line_->display_id() : line_->id();
225 }
226
Name()227 absl::string_view Name() const { return line_->name(); }
228
DisplayName()229 absl::string_view DisplayName() const {
230 return !line_->display_name().empty() ? line_->display_name()
231 : line_->name();
232 }
233
TimestampNs()234 int64_t TimestampNs() const { return line_->timestamp_ns(); }
235
DurationPs()236 int64_t DurationPs() const { return line_->duration_ps(); }
237
NumEvents()238 size_t NumEvents() const { return line_->events_size(); }
239
240 template <typename ForEachEventFunc>
ForEachEvent(ForEachEventFunc && for_each_event)241 void ForEachEvent(ForEachEventFunc&& for_each_event) const {
242 for (const XEvent& event : line_->events()) {
243 for_each_event(XEventVisitor(plane_, line_, &event));
244 }
245 }
246
247 private:
248 const XPlaneVisitor* plane_;
249 const XLine* line_;
250 };
251
252 using TypeGetter = std::function<absl::optional<int64_t>(absl::string_view)>;
253 using TypeGetterList = std::vector<TypeGetter>;
254
255 class XPlaneVisitor : public XStatsOwner<XPlane> {
256 public:
257 // REQUIRED: plane cannot be nullptr.
258 explicit XPlaneVisitor(
259 const XPlane* plane,
260 const TypeGetterList& event_type_getter_list = TypeGetterList(),
261 const TypeGetterList& stat_type_getter_list = TypeGetterList());
262
Id()263 int64_t Id() const { return plane_->id(); }
264
Name()265 absl::string_view Name() const { return plane_->name(); }
266
NumLines()267 size_t NumLines() const { return plane_->lines_size(); }
268
269 template <typename ForEachLineFunc>
ForEachLine(ForEachLineFunc && for_each_line)270 void ForEachLine(ForEachLineFunc&& for_each_line) const {
271 for (const XLine& line : plane_->lines()) {
272 for_each_line(XLineVisitor(this, &line));
273 }
274 }
275 template <typename ThreadBundle, typename ForEachLineFunc>
ForEachLineInParallel(ForEachLineFunc && for_each_line)276 void ForEachLineInParallel(ForEachLineFunc&& for_each_line) const {
277 ThreadBundle bundle;
278 for (const XLine& line : plane_->lines()) {
279 bundle.Add([this, line = &line, &for_each_line] {
280 for_each_line(XLineVisitor(this, line));
281 });
282 }
283 bundle.JoinAll();
284 }
285
286 template <typename ForEachEventMetadataFunc>
ForEachEventMetadata(ForEachEventMetadataFunc && for_each_event_metadata)287 void ForEachEventMetadata(
288 ForEachEventMetadataFunc&& for_each_event_metadata) {
289 for (const auto& [id, event_metadata] : plane_->event_metadata()) {
290 for_each_event_metadata(XEventMetadataVisitor(this, &event_metadata));
291 }
292 }
293
294 // Returns event metadata given its id. Returns a default value if not found.
295 const XEventMetadata* GetEventMetadata(int64_t event_metadata_id) const;
296
297 // Returns the type of an event given its id.
298 absl::optional<int64_t> GetEventType(int64_t event_metadata_id) const;
299
300 // Returns stat metadata given its id. Returns a default value if not found.
301 const XStatMetadata* GetStatMetadata(int64_t stat_metadata_id) const;
302
303 // Returns stat metadata given its type. Returns nullptr if not found.
304 // Use as an alternative to GetStatMetadata above.
305 const XStatMetadata* GetStatMetadataByType(int64_t stat_type) const;
306
307 // Returns the type of an stat given its id.
308 absl::optional<int64_t> GetStatType(int64_t stat_metadata_id) const;
309
310 private:
311 void BuildEventTypeMap(const XPlane* plane,
312 const TypeGetterList& event_type_getter_list);
313 void BuildStatTypeMap(const XPlane* plane,
314 const TypeGetterList& stat_type_getter_list);
315
316 const XPlane* plane_;
317
318 absl::flat_hash_map<int64_t /*metadata_id*/, int64_t /*EventType*/>
319 event_type_by_id_;
320 absl::flat_hash_map<int64_t /*metadata_id*/, int64_t /*StatType*/>
321 stat_type_by_id_;
322 absl::flat_hash_map<int64_t /*StatType*/, const XStatMetadata*>
323 stat_metadata_by_type_;
324 };
325
326 template <class T>
GetStat(int64_t stat_type)327 absl::optional<XStatVisitor> XStatsOwner<T>::GetStat(int64_t stat_type) const {
328 const auto* stat_metadata = plane_->GetStatMetadataByType(stat_type);
329 if (stat_metadata != nullptr) {
330 return GetStat(stat_type, *stat_metadata);
331 }
332 return absl::nullopt; // type does not exist in this owner.
333 }
334
335 template <typename ForEachChildFunc>
ForEachChild(ForEachChildFunc && for_each_child)336 void XEventMetadataVisitor::ForEachChild(
337 ForEachChildFunc&& for_each_child) const {
338 for (int64_t child_id : metadata()->child_id()) {
339 const auto* event_metadata = plane()->GetEventMetadata(child_id);
340 if (event_metadata != nullptr) {
341 for_each_child(XEventMetadataVisitor(plane(), event_metadata));
342 }
343 }
344 }
345
346 } // namespace profiler
347 } // namespace tensorflow
348
349 #endif // TENSORFLOW_CORE_PROFILER_UTILS_XPLANE_VISITOR_H_
350