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 #include "tensorflow/core/profiler/utils/xplane_builder.h"
16
17 #include <algorithm>
18 #include <string>
19 #include <utility>
20
21 #include "absl/container/flat_hash_map.h"
22 #include "absl/strings/string_view.h"
23 #include "absl/types/optional.h"
24 #include "tensorflow/core/platform/types.h"
25 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
26 #include "tensorflow/core/profiler/utils/time_utils.h"
27
28 namespace tensorflow {
29 namespace profiler {
30
XPlaneBuilder(XPlane * plane)31 XPlaneBuilder::XPlaneBuilder(XPlane* plane)
32 : XStatsBuilder<XPlane>(plane, this), plane_(plane) {
33 for (auto& id_and_metadata : *plane->mutable_event_metadata()) {
34 auto& metadata = id_and_metadata.second;
35 last_event_metadata_id_ =
36 std::max<int64>(last_event_metadata_id_, metadata.id());
37 if (!metadata.name().empty()) {
38 event_metadata_by_name_.try_emplace(metadata.name(), &metadata);
39 }
40 }
41 for (auto& id_and_metadata : *plane->mutable_stat_metadata()) {
42 auto& metadata = id_and_metadata.second;
43 last_stat_metadata_id_ =
44 std::max<int64>(last_stat_metadata_id_, metadata.id());
45 if (!metadata.name().empty()) {
46 stat_metadata_by_name_.try_emplace(metadata.name(), &metadata);
47 }
48 }
49 for (XLine& line : *plane->mutable_lines()) {
50 lines_by_id_.try_emplace(line.id(), &line);
51 }
52 }
53
GetOrCreateEventMetadata(int64_t metadata_id)54 XEventMetadata* XPlaneBuilder::GetOrCreateEventMetadata(int64_t metadata_id) {
55 XEventMetadata& metadata = (*plane_->mutable_event_metadata())[metadata_id];
56 metadata.set_id(metadata_id);
57 return &metadata;
58 }
59
CreateEventMetadata()60 XEventMetadata* XPlaneBuilder::CreateEventMetadata() {
61 return GetOrCreateEventMetadata(++last_event_metadata_id_);
62 }
63
GetOrCreateEventMetadata(absl::string_view name)64 XEventMetadata* XPlaneBuilder::GetOrCreateEventMetadata(
65 absl::string_view name) {
66 XEventMetadata*& metadata = event_metadata_by_name_[name];
67 if (metadata == nullptr) {
68 metadata = CreateEventMetadata();
69 metadata->set_name(std::string(name));
70 }
71 return metadata;
72 }
73
GetOrCreateEventMetadata(std::string && name)74 XEventMetadata* XPlaneBuilder::GetOrCreateEventMetadata(std::string&& name) {
75 XEventMetadata*& metadata = event_metadata_by_name_[name];
76 if (metadata == nullptr) {
77 metadata = CreateEventMetadata();
78 metadata->set_name(std::move(name));
79 }
80 return metadata;
81 }
82
GetEventMetadata(absl::string_view name) const83 XEventMetadata* XPlaneBuilder::GetEventMetadata(absl::string_view name) const {
84 auto result = event_metadata_by_name_.find(name);
85 if (result == event_metadata_by_name_.end()) return nullptr;
86 return result->second;
87 }
88
GetStatMetadata(absl::string_view name) const89 XStatMetadata* XPlaneBuilder::GetStatMetadata(absl::string_view name) const {
90 auto result = stat_metadata_by_name_.find(name);
91 if (result == stat_metadata_by_name_.end()) return nullptr;
92 return result->second;
93 }
94
GetOrCreateStatMetadata(int64_t metadata_id)95 XStatMetadata* XPlaneBuilder::GetOrCreateStatMetadata(int64_t metadata_id) {
96 XStatMetadata& metadata = (*plane_->mutable_stat_metadata())[metadata_id];
97 metadata.set_id(metadata_id);
98 return &metadata;
99 }
100
CreateStatMetadata()101 XStatMetadata* XPlaneBuilder::CreateStatMetadata() {
102 return GetOrCreateStatMetadata(++last_stat_metadata_id_);
103 }
104
GetOrCreateStatMetadata(absl::string_view name)105 XStatMetadata* XPlaneBuilder::GetOrCreateStatMetadata(absl::string_view name) {
106 XStatMetadata*& metadata = stat_metadata_by_name_[name];
107 if (metadata == nullptr) {
108 metadata = CreateStatMetadata();
109 metadata->set_name(std::string(name));
110 }
111 return metadata;
112 }
113
GetOrCreateStatMetadata(std::string && name)114 XStatMetadata* XPlaneBuilder::GetOrCreateStatMetadata(std::string&& name) {
115 XStatMetadata*& metadata = stat_metadata_by_name_[name];
116 if (metadata == nullptr) {
117 metadata = CreateStatMetadata();
118 metadata->set_name(std::move(name));
119 }
120 return metadata;
121 }
122
GetOrCreateLine(int64_t line_id)123 XLineBuilder XPlaneBuilder::GetOrCreateLine(int64_t line_id) {
124 XLine*& line = lines_by_id_[line_id];
125 if (line == nullptr) {
126 line = plane_->add_lines();
127 line->set_id(line_id);
128 }
129 return XLineBuilder(line, this);
130 }
131
AddEvent(const XEventMetadata & metadata)132 XEventBuilder XLineBuilder::AddEvent(const XEventMetadata& metadata) {
133 XEvent* event = line_->add_events();
134 event->set_metadata_id(metadata.id());
135 return XEventBuilder(line_, plane_, event);
136 }
137
AddEvent(const XEvent & event)138 XEventBuilder XLineBuilder::AddEvent(const XEvent& event) {
139 XEvent* new_event = line_->add_events();
140 *new_event = event;
141 return XEventBuilder(line_, plane_, new_event);
142 }
143
SetTimestampNsAndAdjustEventOffsets(int64_t timestamp_ns)144 void XLineBuilder::SetTimestampNsAndAdjustEventOffsets(int64_t timestamp_ns) {
145 int64_t offset_ps = NanosToPicos(line_->timestamp_ns() - timestamp_ns);
146 line_->set_timestamp_ns(timestamp_ns);
147 if (offset_ps) {
148 for (auto& event : *line_->mutable_events()) {
149 event.set_offset_ps(event.offset_ps() + offset_ps);
150 }
151 }
152 }
153
154 } // namespace profiler
155 } // namespace tensorflow
156