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 #include "tensorflow/core/profiler/utils/xplane_visitor.h"
16
17 #include <string>
18 #include <utility>
19
20 #include "absl/container/flat_hash_map.h"
21 #include "absl/strings/str_cat.h"
22 #include "absl/strings/string_view.h"
23 #include "absl/types/optional.h"
24 #include "tensorflow/core/platform/logging.h"
25 #include "tensorflow/core/platform/types.h"
26 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
27
28 namespace tensorflow {
29 namespace profiler {
30
XStatVisitor(const XPlaneVisitor * plane,const XStat * stat)31 XStatVisitor::XStatVisitor(const XPlaneVisitor* plane, const XStat* stat)
32 : XStatVisitor(plane, stat, plane->GetStatMetadata(stat->metadata_id()),
33 plane->GetStatType(stat->metadata_id())) {}
34
XStatVisitor(const XPlaneVisitor * plane,const XStat * stat,const XStatMetadata * metadata,absl::optional<int64_t> type)35 XStatVisitor::XStatVisitor(const XPlaneVisitor* plane, const XStat* stat,
36 const XStatMetadata* metadata,
37 absl::optional<int64_t> type)
38 : stat_(stat), metadata_(metadata), plane_(plane), type_(type) {}
39
ToString() const40 std::string XStatVisitor::ToString() const {
41 switch (stat_->value_case()) {
42 case XStat::kInt64Value:
43 return absl::StrCat(stat_->int64_value());
44 case XStat::kUint64Value:
45 return absl::StrCat(stat_->uint64_value());
46 case XStat::kDoubleValue:
47 return absl::StrCat(stat_->double_value());
48 case XStat::kStrValue:
49 return stat_->str_value();
50 case XStat::kBytesValue:
51 return "<opaque bytes>";
52 case XStat::kRefValue:
53 return plane_->GetStatMetadata(stat_->ref_value())->name();
54 case XStat::VALUE_NOT_SET:
55 return "";
56 }
57 }
58
StrOrRefValue() const59 absl::string_view XStatVisitor::StrOrRefValue() const {
60 switch (stat_->value_case()) {
61 case XStat::kStrValue:
62 return stat_->str_value();
63 case XStat::kRefValue:
64 return plane_->GetStatMetadata(stat_->ref_value())->name();
65 case XStat::kInt64Value:
66 case XStat::kUint64Value:
67 case XStat::kDoubleValue:
68 case XStat::kBytesValue:
69 case XStat::VALUE_NOT_SET:
70 return absl::string_view();
71 }
72 }
73
XEventVisitor(const XPlaneVisitor * plane,const XLine * line,const XEvent * event)74 XEventVisitor::XEventVisitor(const XPlaneVisitor* plane, const XLine* line,
75 const XEvent* event)
76 : XStatsOwner<XEvent>(plane, event),
77 plane_(plane),
78 line_(line),
79 event_(event),
80 metadata_(plane->GetEventMetadata(event_->metadata_id())),
81 type_(plane->GetEventType(event_->metadata_id())) {}
82
XPlaneVisitor(const XPlane * plane,const TypeGetterList & event_type_getter_list,const TypeGetterList & stat_type_getter_list)83 XPlaneVisitor::XPlaneVisitor(const XPlane* plane,
84 const TypeGetterList& event_type_getter_list,
85 const TypeGetterList& stat_type_getter_list)
86 : XStatsOwner<XPlane>(this, plane), plane_(plane) {
87 BuildEventTypeMap(plane, event_type_getter_list);
88 BuildStatTypeMap(plane, stat_type_getter_list);
89 }
90
BuildEventTypeMap(const XPlane * plane,const TypeGetterList & event_type_getter_list)91 void XPlaneVisitor::BuildEventTypeMap(
92 const XPlane* plane, const TypeGetterList& event_type_getter_list) {
93 for (const auto& event_metadata : plane->event_metadata()) {
94 uint64 metadata_id = event_metadata.first;
95 const auto& metadata = event_metadata.second;
96 for (const auto& event_type_getter : event_type_getter_list) {
97 absl::optional<int64_t> event_type = event_type_getter(metadata.name());
98 if (event_type.has_value()) {
99 auto result = event_type_by_id_.emplace(metadata_id, *event_type);
100 DCHECK(result.second); // inserted
101 break;
102 }
103 }
104 }
105 }
106
GetEventMetadata(int64_t event_metadata_id) const107 const XEventMetadata* XPlaneVisitor::GetEventMetadata(
108 int64_t event_metadata_id) const {
109 const auto& event_metadata_by_id = plane_->event_metadata();
110 const auto it = event_metadata_by_id.find(event_metadata_id);
111 if (it != event_metadata_by_id.end()) return &it->second;
112 return &XEventMetadata::default_instance();
113 }
114
GetEventType(int64_t event_metadata_id) const115 absl::optional<int64_t> XPlaneVisitor::GetEventType(
116 int64_t event_metadata_id) const {
117 const auto it = event_type_by_id_.find(event_metadata_id);
118 if (it != event_type_by_id_.end()) return it->second;
119 return absl::nullopt;
120 }
121
BuildStatTypeMap(const XPlane * plane,const TypeGetterList & stat_type_getter_list)122 void XPlaneVisitor::BuildStatTypeMap(
123 const XPlane* plane, const TypeGetterList& stat_type_getter_list) {
124 for (const auto& stat_metadata : plane->stat_metadata()) {
125 uint64 metadata_id = stat_metadata.first;
126 const auto& metadata = stat_metadata.second;
127 for (const auto& stat_type_getter : stat_type_getter_list) {
128 absl::optional<int64_t> stat_type = stat_type_getter(metadata.name());
129 if (stat_type.has_value()) {
130 auto result = stat_type_by_id_.emplace(metadata_id, *stat_type);
131 DCHECK(result.second); // inserted
132 stat_metadata_by_type_.emplace(*stat_type, &metadata);
133 break;
134 }
135 }
136 }
137 }
138
GetStatMetadata(int64_t stat_metadata_id) const139 const XStatMetadata* XPlaneVisitor::GetStatMetadata(
140 int64_t stat_metadata_id) const {
141 const auto& stat_metadata_by_id = plane_->stat_metadata();
142 const auto it = stat_metadata_by_id.find(stat_metadata_id);
143 if (it != stat_metadata_by_id.end()) return &it->second;
144 return &XStatMetadata::default_instance();
145 }
146
GetStatType(int64_t stat_metadata_id) const147 absl::optional<int64_t> XPlaneVisitor::GetStatType(
148 int64_t stat_metadata_id) const {
149 const auto it = stat_type_by_id_.find(stat_metadata_id);
150 if (it != stat_type_by_id_.end()) return it->second;
151 return absl::nullopt;
152 }
153
GetStatMetadataByType(int64_t stat_type) const154 const XStatMetadata* XPlaneVisitor::GetStatMetadataByType(
155 int64_t stat_type) const {
156 const auto it = stat_metadata_by_type_.find(stat_type);
157 if (it != stat_metadata_by_type_.end()) return it->second;
158 return nullptr;
159 }
160
161 } // namespace profiler
162 } // namespace tensorflow
163