xref: /aosp_15_r20/external/tensorflow/tensorflow/core/profiler/utils/xplane_visitor.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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