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
16 #include "tensorflow/core/profiler/utils/xplane_utils.h"
17
18 #include <string>
19
20 #include "absl/container/flat_hash_map.h"
21 #include "absl/strings/string_view.h"
22 #include "absl/types/optional.h"
23 #include "tensorflow/core/platform/test.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 #include "tensorflow/core/profiler/utils/xplane_builder.h"
28 #include "tensorflow/core/profiler/utils/xplane_visitor.h"
29
30 namespace tensorflow {
31 namespace profiler {
32 namespace {
33
CreateEvent(int64_t offset_ps,int64_t duration_ps)34 XEvent CreateEvent(int64_t offset_ps, int64_t duration_ps) {
35 XEvent event;
36 event.set_offset_ps(offset_ps);
37 event.set_duration_ps(duration_ps);
38 return event;
39 }
40
TEST(XPlaneUtilsTest,AddAndRemovePlanes)41 TEST(XPlaneUtilsTest, AddAndRemovePlanes) {
42 XSpace space;
43
44 auto* p1 = FindOrAddMutablePlaneWithName(&space, "p1");
45 EXPECT_EQ(p1, FindPlaneWithName(space, "p1"));
46 auto* p2 = FindOrAddMutablePlaneWithName(&space, "p2");
47 EXPECT_EQ(p2, FindPlaneWithName(space, "p2"));
48 auto* p3 = FindOrAddMutablePlaneWithName(&space, "p3");
49 EXPECT_EQ(p3, FindPlaneWithName(space, "p3"));
50
51 // Removing a plane does not invalidate pointers to other planes.
52
53 RemovePlane(&space, p2);
54 EXPECT_EQ(space.planes_size(), 2);
55 EXPECT_EQ(p1, FindPlaneWithName(space, "p1"));
56 EXPECT_EQ(p3, FindPlaneWithName(space, "p3"));
57
58 RemovePlane(&space, p1);
59 EXPECT_EQ(space.planes_size(), 1);
60 EXPECT_EQ(p3, FindPlaneWithName(space, "p3"));
61
62 RemovePlane(&space, p3);
63 EXPECT_EQ(space.planes_size(), 0);
64 }
65
TEST(XPlaneUtilsTest,RemoveEmptyPlanes)66 TEST(XPlaneUtilsTest, RemoveEmptyPlanes) {
67 XSpace space;
68 RemoveEmptyPlanes(&space);
69 EXPECT_EQ(space.planes_size(), 0);
70
71 auto* plane1 = space.add_planes();
72 plane1->set_name("p1");
73 plane1->add_lines()->set_name("p1l1");
74 plane1->add_lines()->set_name("p1l2");
75
76 auto* plane2 = space.add_planes();
77 plane2->set_name("p2");
78
79 auto* plane3 = space.add_planes();
80 plane3->set_name("p3");
81 plane3->add_lines()->set_name("p3l1");
82
83 auto* plane4 = space.add_planes();
84 plane4->set_name("p4");
85
86 RemoveEmptyPlanes(&space);
87 ASSERT_EQ(space.planes_size(), 2);
88 EXPECT_EQ(space.planes(0).name(), "p1");
89 EXPECT_EQ(space.planes(1).name(), "p3");
90 }
91
TEST(XPlaneUtilsTest,RemoveEmptyLines)92 TEST(XPlaneUtilsTest, RemoveEmptyLines) {
93 XPlane plane;
94 RemoveEmptyLines(&plane);
95 EXPECT_EQ(plane.lines_size(), 0);
96
97 auto* line1 = plane.add_lines();
98 line1->set_name("l1");
99 line1->add_events();
100 line1->add_events();
101
102 auto* line2 = plane.add_lines();
103 line2->set_name("l2");
104
105 auto* line3 = plane.add_lines();
106 line3->set_name("l3");
107 line3->add_events();
108
109 auto* line4 = plane.add_lines();
110 line4->set_name("l4");
111
112 RemoveEmptyLines(&plane);
113 ASSERT_EQ(plane.lines_size(), 2);
114 EXPECT_EQ(plane.lines(0).name(), "l1");
115 EXPECT_EQ(plane.lines(1).name(), "l3");
116 }
117
TEST(XPlaneUtilsTest,RemoveLine)118 TEST(XPlaneUtilsTest, RemoveLine) {
119 XPlane plane;
120 const XLine* line1 = plane.add_lines();
121 const XLine* line2 = plane.add_lines();
122 const XLine* line3 = plane.add_lines();
123 RemoveLine(&plane, line2);
124 ASSERT_EQ(plane.lines_size(), 2);
125 EXPECT_EQ(&plane.lines(0), line1);
126 EXPECT_EQ(&plane.lines(1), line3);
127 }
128
TEST(XPlaneUtilsTest,RemoveEvents)129 TEST(XPlaneUtilsTest, RemoveEvents) {
130 XLine line;
131 const XEvent* event1 = line.add_events();
132 const XEvent* event2 = line.add_events();
133 const XEvent* event3 = line.add_events();
134 const XEvent* event4 = line.add_events();
135 RemoveEvents(&line, {event1, event3});
136 ASSERT_EQ(line.events_size(), 2);
137 EXPECT_EQ(&line.events(0), event2);
138 EXPECT_EQ(&line.events(1), event4);
139 }
140
TEST(XPlaneUtilsTest,SortXPlaneTest)141 TEST(XPlaneUtilsTest, SortXPlaneTest) {
142 XPlane plane;
143 XLine* line = plane.add_lines();
144 *line->add_events() = CreateEvent(200, 100);
145 *line->add_events() = CreateEvent(100, 100);
146 *line->add_events() = CreateEvent(120, 50);
147 *line->add_events() = CreateEvent(120, 30);
148 SortXPlane(&plane);
149 ASSERT_EQ(plane.lines_size(), 1);
150 ASSERT_EQ(plane.lines(0).events_size(), 4);
151 EXPECT_EQ(plane.lines(0).events(0).offset_ps(), 100);
152 EXPECT_EQ(plane.lines(0).events(0).duration_ps(), 100);
153 EXPECT_EQ(plane.lines(0).events(1).offset_ps(), 120);
154 EXPECT_EQ(plane.lines(0).events(1).duration_ps(), 50);
155 EXPECT_EQ(plane.lines(0).events(2).offset_ps(), 120);
156 EXPECT_EQ(plane.lines(0).events(2).duration_ps(), 30);
157 EXPECT_EQ(plane.lines(0).events(3).offset_ps(), 200);
158 EXPECT_EQ(plane.lines(0).events(3).duration_ps(), 100);
159 }
160
161 namespace {
162
CreateXLine(XPlaneBuilder * plane,absl::string_view name,absl::string_view display,int64_t id,int64_t timestamp_ns)163 XLineBuilder CreateXLine(XPlaneBuilder* plane, absl::string_view name,
164 absl::string_view display, int64_t id,
165 int64_t timestamp_ns) {
166 XLineBuilder line = plane->GetOrCreateLine(id);
167 line.SetName(name);
168 line.SetTimestampNs(timestamp_ns);
169 line.SetDisplayNameIfEmpty(display);
170 return line;
171 }
172
CreateXEvent(XPlaneBuilder * plane,XLineBuilder line,absl::string_view event_name,absl::optional<absl::string_view> display,int64_t offset_ns,int64_t duration_ns)173 XEventBuilder CreateXEvent(XPlaneBuilder* plane, XLineBuilder line,
174 absl::string_view event_name,
175 absl::optional<absl::string_view> display,
176 int64_t offset_ns, int64_t duration_ns) {
177 XEventMetadata* event_metadata = plane->GetOrCreateEventMetadata(event_name);
178 if (display) event_metadata->set_display_name(std::string(*display));
179 XEventBuilder event = line.AddEvent(*event_metadata);
180 event.SetOffsetNs(offset_ns);
181 event.SetDurationNs(duration_ns);
182 return event;
183 }
184
185 template <typename T, typename V>
CreateXStats(XPlaneBuilder * plane,T * stats_owner,absl::string_view stats_name,V stats_value)186 void CreateXStats(XPlaneBuilder* plane, T* stats_owner,
187 absl::string_view stats_name, V stats_value) {
188 stats_owner->AddStatValue(*plane->GetOrCreateStatMetadata(stats_name),
189 stats_value);
190 }
191
CheckXLine(const XLine & line,absl::string_view name,absl::string_view display,int64_t start_time_ns,int64_t events_size)192 void CheckXLine(const XLine& line, absl::string_view name,
193 absl::string_view display, int64_t start_time_ns,
194 int64_t events_size) {
195 EXPECT_EQ(line.name(), name);
196 EXPECT_EQ(line.display_name(), display);
197 EXPECT_EQ(line.timestamp_ns(), start_time_ns);
198 EXPECT_EQ(line.events_size(), events_size);
199 }
200
CheckXEvent(const XEvent & event,const XPlane & plane,absl::string_view name,absl::string_view display,int64_t offset_ns,int64_t duration_ns,int64_t stats_size)201 void CheckXEvent(const XEvent& event, const XPlane& plane,
202 absl::string_view name, absl::string_view display,
203 int64_t offset_ns, int64_t duration_ns, int64_t stats_size) {
204 const XEventMetadata& event_metadata =
205 plane.event_metadata().at(event.metadata_id());
206 EXPECT_EQ(event_metadata.name(), name);
207 EXPECT_EQ(event_metadata.display_name(), display);
208 EXPECT_EQ(event.offset_ps(), NanosToPicos(offset_ns));
209 EXPECT_EQ(event.duration_ps(), NanosToPicos(duration_ns));
210 EXPECT_EQ(event.stats_size(), stats_size);
211 }
212 } // namespace
213
TEST(XPlaneUtilsTest,MergeXPlaneTest)214 TEST(XPlaneUtilsTest, MergeXPlaneTest) {
215 XPlane src_plane, dst_plane;
216 constexpr int64_t kLineIdOnlyInSrcPlane = 1LL;
217 constexpr int64_t kLineIdOnlyInDstPlane = 2LL;
218 constexpr int64_t kLineIdInBothPlanes = 3LL; // src start ts < dst start ts
219 constexpr int64_t kLineIdInBothPlanes2 = 4LL; // src start ts > dst start ts
220
221 { // Populate the source plane.
222 XPlaneBuilder src(&src_plane);
223 CreateXStats(&src, &src, "plane_stat1", 1); // only in source.
224 CreateXStats(&src, &src, "plane_stat3", 3.0); // shared by source/dest.
225
226 auto l1 = CreateXLine(&src, "l1", "d1", kLineIdOnlyInSrcPlane, 100);
227 auto e1 = CreateXEvent(&src, l1, "event1", "display1", 1, 2);
228 CreateXStats(&src, &e1, "event_stat1", 2.0);
229 auto e2 = CreateXEvent(&src, l1, "event2", absl::nullopt, 3, 4);
230 CreateXStats(&src, &e2, "event_stat2", 3);
231
232 auto l2 = CreateXLine(&src, "l2", "d2", kLineIdInBothPlanes, 200);
233 auto e3 = CreateXEvent(&src, l2, "event3", absl::nullopt, 5, 7);
234 CreateXStats(&src, &e3, "event_stat3", 2.0);
235 auto e4 = CreateXEvent(&src, l2, "event4", absl::nullopt, 6, 8);
236 CreateXStats(&src, &e4, "event_stat4", 3);
237 CreateXStats(&src, &e4, "event_stat5", 3);
238
239 auto l5 = CreateXLine(&src, "l5", "d5", kLineIdInBothPlanes2, 700);
240 CreateXEvent(&src, l5, "event51", absl::nullopt, 9, 10);
241 CreateXEvent(&src, l5, "event52", absl::nullopt, 11, 12);
242 }
243
244 { // Populate the destination plane.
245 XPlaneBuilder dst(&dst_plane);
246 CreateXStats(&dst, &dst, "plane_stat2", 2); // only in dest
247 CreateXStats(&dst, &dst, "plane_stat3", 4); // shared but different.
248
249 auto l3 = CreateXLine(&dst, "l3", "d3", kLineIdOnlyInDstPlane, 300);
250 auto e5 = CreateXEvent(&dst, l3, "event5", absl::nullopt, 11, 2);
251 CreateXStats(&dst, &e5, "event_stat6", 2.0);
252 auto e6 = CreateXEvent(&dst, l3, "event6", absl::nullopt, 13, 4);
253 CreateXStats(&dst, &e6, "event_stat7", 3);
254
255 auto l2 = CreateXLine(&dst, "l4", "d4", kLineIdInBothPlanes, 400);
256 auto e7 = CreateXEvent(&dst, l2, "event7", absl::nullopt, 15, 7);
257 CreateXStats(&dst, &e7, "event_stat8", 2.0);
258 auto e8 = CreateXEvent(&dst, l2, "event8", "display8", 16, 8);
259 CreateXStats(&dst, &e8, "event_stat9", 3);
260
261 auto l6 = CreateXLine(&dst, "l6", "d6", kLineIdInBothPlanes2, 300);
262 CreateXEvent(&dst, l6, "event61", absl::nullopt, 21, 10);
263 CreateXEvent(&dst, l6, "event62", absl::nullopt, 22, 12);
264 }
265
266 MergePlanes(src_plane, &dst_plane);
267
268 XPlaneVisitor plane(&dst_plane);
269 EXPECT_EQ(dst_plane.lines_size(), 4);
270
271 // Check plane level stats,
272 EXPECT_EQ(dst_plane.stats_size(), 3);
273 absl::flat_hash_map<absl::string_view, absl::string_view> plane_stats;
274 plane.ForEachStat([&](const tensorflow::profiler::XStatVisitor& stat) {
275 if (stat.Name() == "plane_stat1") {
276 EXPECT_EQ(stat.IntValue(), 1);
277 } else if (stat.Name() == "plane_stat2") {
278 EXPECT_EQ(stat.IntValue(), 2);
279 } else if (stat.Name() == "plane_stat3") {
280 // XStat in src_plane overrides the counter-part in dst_plane.
281 EXPECT_EQ(stat.DoubleValue(), 3.0);
282 } else {
283 EXPECT_TRUE(false);
284 }
285 });
286
287 // 3 plane level stats, 9 event level stats.
288 EXPECT_EQ(dst_plane.stat_metadata_size(), 12);
289
290 { // Old lines are untouched.
291 const XLine& line = dst_plane.lines(0);
292 CheckXLine(line, "l3", "d3", 300, 2);
293 CheckXEvent(line.events(0), dst_plane, "event5", "", 11, 2, 1);
294 CheckXEvent(line.events(1), dst_plane, "event6", "", 13, 4, 1);
295 }
296 { // Lines with the same id are merged.
297 // src plane start timestamp > dst plane start timestamp
298 const XLine& line = dst_plane.lines(1);
299 // NOTE: use minimum start time of src/dst.
300 CheckXLine(line, "l4", "d4", 200, 4);
301 CheckXEvent(line.events(0), dst_plane, "event7", "", 215, 7, 1);
302 CheckXEvent(line.events(1), dst_plane, "event8", "display8", 216, 8, 1);
303 CheckXEvent(line.events(2), dst_plane, "event3", "", 5, 7, 1);
304 CheckXEvent(line.events(3), dst_plane, "event4", "", 6, 8, 2);
305 }
306 { // Lines with the same id are merged.
307 // src plane start timestamp < dst plane start timestamp
308 const XLine& line = dst_plane.lines(2);
309 CheckXLine(line, "l6", "d6", 300, 4);
310 CheckXEvent(line.events(0), dst_plane, "event61", "", 21, 10, 0);
311 CheckXEvent(line.events(1), dst_plane, "event62", "", 22, 12, 0);
312 CheckXEvent(line.events(2), dst_plane, "event51", "", 409, 10, 0);
313 CheckXEvent(line.events(3), dst_plane, "event52", "", 411, 12, 0);
314 }
315 { // Lines only in source are "copied".
316 const XLine& line = dst_plane.lines(3);
317 CheckXLine(line, "l1", "d1", 100, 2);
318 CheckXEvent(line.events(0), dst_plane, "event1", "display1", 1, 2, 1);
319 CheckXEvent(line.events(1), dst_plane, "event2", "", 3, 4, 1);
320 }
321 }
322
323 } // namespace
324 } // namespace profiler
325 } // namespace tensorflow
326