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