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