1 /* Copyright 2018 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/internal/cpu/host_tracer.h"
16
17 #include <utility>
18
19 #include "absl/container/flat_hash_map.h"
20 #include "absl/strings/str_split.h"
21 #include "tensorflow/core/common_runtime/step_stats_collector.h"
22 #include "tensorflow/core/platform/env_time.h"
23
24 namespace tensorflow {
25 namespace profiler {
26 namespace cpu {
27
Create(int host_trace_level)28 /* static */ std::unique_ptr<HostTracer> HostTracer::Create(
29 int host_trace_level) {
30 return absl::WrapUnique(new HostTracer(host_trace_level));
31 }
HostTracer(int host_trace_level)32 HostTracer::HostTracer(int host_trace_level)
33 : host_trace_level_(host_trace_level) {}
34
~HostTracer()35 HostTracer::~HostTracer() { Stop().IgnoreError(); }
36
Start()37 Status HostTracer::Start() {
38 if (recording_) {
39 return Status(error::INTERNAL, "TraceMeRecorder already started");
40 }
41 recording_ = TraceMeRecorder::Start(host_trace_level_);
42 if (!recording_) {
43 return Status(error::INTERNAL, "Failed to start TraceMeRecorder");
44 }
45 return Status::OK();
46 }
47
Stop()48 Status HostTracer::Stop() {
49 if (!recording_) {
50 return Status(error::INTERNAL, "TraceMeRecorder not started");
51 }
52 events_ = TraceMeRecorder::Stop();
53 recording_ = false;
54 return Status::OK();
55 }
56
57 constexpr char kUserMetadataMarker = '#';
58
CollectData(RunMetadata * run_metadata)59 Status HostTracer::CollectData(RunMetadata* run_metadata) {
60 auto step_stats_collector =
61 absl::make_unique<StepStatsCollector>(run_metadata->mutable_step_stats());
62 return CollectDataToCollector(step_stats_collector.get());
63 }
64
CollectDataToCollector(StepStatsCollector * step_stats_collector)65 Status HostTracer::CollectDataToCollector(
66 StepStatsCollector* step_stats_collector) {
67 if (events_.empty() && recording_) {
68 events_ = TraceMeRecorder::Collect();
69 }
70 // Pair up start and end events, and add complete events to trace_entries.
71 absl::flat_hash_map<uint64, uint64> end_times;
72 for (const auto& thread : events_) {
73 for (const auto& event : thread.events) {
74 if (event.end_time && !event.start_time) {
75 end_times.emplace(event.activity_id, event.end_time);
76 }
77 }
78 }
79
80 const string cpu_name = "/host:CPU";
81 for (auto& thread : events_) {
82 step_stats_collector->SaveThreadName(cpu_name, thread.thread.tid,
83 thread.thread.name);
84 for (auto& event : thread.events) {
85 if (!event.end_time) {
86 auto it = end_times.find(event.activity_id);
87 if (it != end_times.end()) event.end_time = it->second;
88 }
89 if (event.start_time && event.end_time) {
90 NodeExecStats* ns = new NodeExecStats;
91 if (event.name.back() != kUserMetadataMarker) {
92 ns->set_node_name(std::move(event.name));
93 } else {
94 // Expect the format will be "<name>#<metadata>#"
95 std::vector<absl::string_view> parts =
96 absl::StrSplit(event.name, kUserMetadataMarker);
97 if (parts.size() >= 2) {
98 ns->set_node_name(string(parts[0]));
99 ns->set_timeline_label(string(parts[1]));
100 } else {
101 ns->set_node_name(std::move(event.name));
102 }
103 }
104 ns->set_all_start_micros(event.start_time / EnvTime::kMicrosToNanos);
105 ns->set_all_end_rel_micros((event.end_time - event.start_time) /
106 EnvTime::kMicrosToNanos);
107 ns->set_thread_id(thread.thread.tid);
108 // TODO(fishx): Add thread name to RunMetadata
109 step_stats_collector->Save(cpu_name, ns);
110 }
111 }
112 }
113 events_.clear();
114 step_stats_collector->Finalize();
115 return Status::OK();
116 }
117
118 } // namespace cpu
119 } // namespace profiler
120 } // namespace tensorflow
121