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/convert/xplane_to_trace_events.h"
17
18 #include <stddef.h>
19
20 #include <algorithm>
21 #include <iterator>
22 #include <string>
23 #include <vector>
24
25 #include "absl/strings/string_view.h"
26 #include "absl/types/optional.h"
27 #include "tensorflow/core/platform/types.h"
28 #include "tensorflow/core/profiler/protobuf/trace_events.pb.h"
29 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
30 #include "tensorflow/core/profiler/utils/tf_xplane_visitor.h"
31 #include "tensorflow/core/profiler/utils/trace_utils.h"
32 #include "tensorflow/core/profiler/utils/xplane_schema.h"
33 #include "tensorflow/core/profiler/utils/xplane_utils.h"
34 #include "tensorflow/core/profiler/utils/xplane_visitor.h"
35
36 namespace tensorflow {
37 namespace profiler {
38
39 namespace {
40
BuildDeviceAndResources(uint32 device_id,const XPlaneVisitor & plane,Device * device)41 void BuildDeviceAndResources(uint32 device_id, const XPlaneVisitor& plane,
42 Device* device) {
43 device->set_name(std::string(plane.Name()));
44 device->set_device_id(device_id);
45
46 bool sort_by_ordinal = (device_id == kHostThreadsDeviceId);
47 int ordinal = 0;
48 plane.ForEachLine([&](const XLineVisitor& line) {
49 uint32 resource_id = line.DisplayId();
50 Resource& resource = (*device->mutable_resources())[resource_id];
51 resource.set_resource_id(resource_id);
52 resource.set_name(std::string(line.DisplayName()));
53 if (sort_by_ordinal) {
54 // When sort_index is absent (i.e. 0), resource id will be used.
55 // Therefore sort_index starts with 1.
56 resource.set_sort_index(++ordinal);
57 }
58 });
59 }
60
ConvertXPlaneToTraceEvents(uint32 device_id,const XPlaneVisitor & xplane,Trace * trace)61 void ConvertXPlaneToTraceEvents(uint32 device_id, const XPlaneVisitor& xplane,
62 Trace* trace) {
63 // Convert devices and resources.
64 BuildDeviceAndResources(device_id, xplane,
65 &(*trace->mutable_devices())[device_id]);
66
67 // Convert events.
68 xplane.ForEachLine([device_id, trace](const XLineVisitor& xline) {
69 uint32 resource_id = xline.DisplayId();
70 xline.ForEachEvent(
71 [device_id, resource_id, trace](const XEventVisitor& xevent) {
72 int64 event_type =
73 xevent.Type().value_or(HostEventType::kUnknownHostEventType);
74 if (IsInternalEvent(event_type)) return;
75 auto* event = trace->add_trace_events();
76 auto& args = *event->mutable_args();
77 event->set_device_id(device_id);
78 event->set_resource_id(resource_id);
79 if (xevent.HasDisplayName()) {
80 event->set_name(std::string(xevent.DisplayName()));
81 args["long_name"] = std::string(xevent.Name());
82 } else {
83 event->set_name(std::string(xevent.Name()));
84 }
85 event->set_timestamp_ps(xevent.TimestampPs());
86 event->set_duration_ps(xevent.DurationPs());
87
88 xevent.ForEachStat([&](const XStatVisitor& stat) {
89 if (stat.ValueCase() == XStat::VALUE_NOT_SET) return;
90 if (IsInternalStat(stat.Type())) return;
91 if (stat.Type() == StatType::kStepName) {
92 event->set_name(stat.ToString());
93 }
94 args[std::string(stat.Name())] = stat.ToString();
95 });
96 });
97 });
98 }
99
100 } // namespace
101
MaybeDropEventsForTraceViewer(Trace * trace,uint32 limit)102 void MaybeDropEventsForTraceViewer(Trace* trace, uint32 limit) {
103 auto* trace_events = trace->mutable_trace_events();
104 size_t trace_event_size = trace_events->size();
105 if (trace_event_size <= limit) return; // Nothing to do.
106 // Sort the events according to start time.
107 std::vector<uint64> timestamps;
108 timestamps.reserve(trace_event_size);
109 for (const auto& event : *trace_events) {
110 timestamps.push_back(event.timestamp_ps());
111 }
112 std::partial_sort(timestamps.begin(), timestamps.begin() + limit,
113 timestamps.end(), std::less<uint64>());
114 uint64 cutoff_timestamp = timestamps[limit - 1];
115 trace_events->erase(std::remove_if(trace_events->begin(), trace_events->end(),
116 [&](const TraceEvent& event) {
117 return event.timestamp_ps() >
118 cutoff_timestamp;
119 }),
120 trace_events->end());
121 }
122
ConvertXSpaceToTraceEvents(const XSpace & xspace,Trace * trace)123 void ConvertXSpaceToTraceEvents(const XSpace& xspace, Trace* trace) {
124 const XPlane* host_plane = FindPlaneWithName(xspace, kHostThreadsPlaneName);
125 if (host_plane != nullptr) {
126 XPlaneVisitor xplane = CreateTfXPlaneVisitor(host_plane);
127 ConvertXPlaneToTraceEvents(kHostThreadsDeviceId, xplane, trace);
128 }
129 std::vector<const XPlane*> device_planes =
130 FindPlanesWithPrefix(xspace, kGpuPlanePrefix);
131 // We don't expect GPU and TPU planes to be present in the same XSpace.
132 if (device_planes.empty()) {
133 device_planes = FindPlanesWithPrefix(xspace, kTpuPlanePrefix);
134 }
135 for (const XPlane* device_plane : device_planes) {
136 XPlaneVisitor xplane = CreateTfXPlaneVisitor(device_plane);
137 uint32 device_id = kFirstDeviceId + xplane.Id();
138 ConvertXPlaneToTraceEvents(device_id, xplane, trace);
139 }
140
141 // Trace viewer (non-streaming) has scalability issues, we need to drop
142 // events to avoid loading failure for trace viewer.
143 constexpr uint64 kMaxEvents = 1000000;
144 MaybeDropEventsForTraceViewer(trace, kMaxEvents);
145 }
146
ConvertXSpaceToTraceEventsString(const XSpace & xspace,std::string * content)147 void ConvertXSpaceToTraceEventsString(const XSpace& xspace,
148 std::string* content) {
149 Trace trace;
150 ConvertXSpaceToTraceEvents(xspace, &trace);
151 trace.SerializeToString(content);
152 }
153
154 } // namespace profiler
155 } // namespace tensorflow
156