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 #ifndef TENSORFLOW_CORE_PROFILER_UTILS_DERIVED_TIMELINE_H_ 16 #define TENSORFLOW_CORE_PROFILER_UTILS_DERIVED_TIMELINE_H_ 17 18 #include <cstdint> 19 #include <functional> 20 #include <optional> 21 #include <string> 22 #include <vector> 23 24 #include "absl/container/flat_hash_map.h" 25 #include "absl/container/flat_hash_set.h" 26 #include "absl/strings/string_view.h" 27 #include "tensorflow/core/profiler/protobuf/xplane.pb.h" 28 #include "tensorflow/core/profiler/utils/group_events.h" 29 #include "tensorflow/core/profiler/utils/timespan.h" 30 #include "tensorflow/core/profiler/utils/xplane_builder.h" 31 32 namespace tensorflow { 33 namespace profiler { 34 35 // Helper for deriving XEvents. 36 class DerivedXEventBuilder { 37 public: 38 DerivedXEventBuilder(XEventBuilder event, std::optional<int64_t> group_id); 39 40 bool ShouldExpand(const XEventMetadata& event_metadata, 41 std::optional<int64_t> group_id) const; 42 43 void Expand(Timespan event_span); 44 45 private: 46 XEventBuilder event_; 47 std::optional<int64_t> group_id_; 48 }; 49 50 // Helper for deriving an XLine from events in another XLine. 51 class DerivedXLineBuilder { 52 public: 53 DerivedXLineBuilder(XPlaneBuilder* plane, int64_t line_id, 54 absl::string_view name, int64_t timestamp_ns, 55 std::vector<DerivedXLineBuilder*> dependent_lines); 56 Line()57 XLineBuilder& Line() { return line_; } 58 59 // Either merges event with the last event or creates a new event on this 60 // XLine. group_id and low_level_event_name may be passed to separate 61 // consecutive invocations of the same event, depending on the XEvent type: 62 // TF-op, TF name scope: both group_id and low_level_event_name are used. 63 // HLO-op, step: only group_id is used. 64 // HLO module, source: both group_id and low_level_event_name are NOT used. 65 void ExpandOrAddEvent(const XEventMetadata& event_metadata, 66 Timespan event_span, std::optional<int64_t> group_id); 67 68 // The multi-level version of ExpandOrAddEvent. Here, the XEvents at different 69 // levels all share the same group_id and low_level_event_name. 70 void ExpandOrAddEvents( 71 const std::vector<XEventMetadata*>& events_metadata_per_level, 72 Timespan event_span, std::optional<int64_t> group_id); 73 74 // Reset the last events lower than or equal to the given level. 75 void ResetLastEvents(int level = 0); 76 77 private: 78 // If the last event of the given level has the same metadata, expands it to 79 // include the time until the given event's end time. 80 // Otherwise, adds a new event and clears last_event_by_level_ for the levels 81 // below the given level and all levels of the dependent lines. Clearing 82 // last_event_by_level_ prevents a nested event from growing larger than the 83 // parent event(s). 84 void ExpandOrAddLevelEvent(const XEventMetadata& event_metadata, 85 Timespan event_span, 86 std::optional<int64_t> group_id, int level); 87 88 const XStatMetadata* group_id_stat_metadata_ = nullptr; 89 const XStatMetadata* level_stat_metadata_ = nullptr; 90 XLineBuilder line_; 91 absl::flat_hash_map<int, std::optional<DerivedXEventBuilder>> 92 last_event_by_level_; 93 std::vector<DerivedXLineBuilder*> dependent_lines_; 94 }; 95 96 struct Symbol { 97 absl::string_view tf_op_name; 98 std::string source_info; 99 }; 100 101 using SymbolResolver = std::function<Symbol(std::optional<uint64_t> program_id, 102 absl::string_view hlo_module_name, 103 absl::string_view hlo_op)>; 104 105 // Derives TF name scope and op events from the TF op's fully qualified name 106 // with the name of the originating low-level event. 107 void ProcessTfOpEvent(absl::string_view tf_op_full_name, Timespan event_span, 108 std::optional<int64_t> group_id, 109 XPlaneBuilder& plane_builder, 110 DerivedXLineBuilder& tf_name_scope_line_builder, 111 DerivedXLineBuilder& tf_op_line_builder); 112 113 // Adds step names from GroupMetadataMap to "Steps" line in plane. 114 // The event name is updated when converted to trace events. 115 void AddGroupMetadataToStepEvents(const GroupMetadataMap& group_metadata_map, 116 XLineBuilder& line); 117 118 // Derives "Steps" line from group_id XStat in XEvents. 119 void DeriveStepEventsFromGroups(const GroupMetadataMap& group_metadata_map, 120 XPlane* device_trace); 121 122 // Derives "TensorFlow Ops", "TensorFlow Name Scope", "XLA Ops" and "XLA Module" 123 // lines in an NVIDIA_GPU device trace from data passed as ScopedAnnotations and 124 // stored as XStats in XEvents corresponding to GPU Kernels. Consecutive 125 // annotations with the same value are merged into a single event except for XLA 126 // modules. The device_trace is both input and output. 127 void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver, 128 XPlane* device_trace); 129 130 // Derives "Launch Activities Summary" line from host trace. 131 void DeriveEventsFromHostTrace(const XPlane* host_trace, 132 const GroupMetadataMap& group_metadata_map, 133 std::vector<XPlane*> device_traces); 134 135 // Loops through XPlanes of input XSpace, if it is "device" XPlane, generating 136 // derived timelines for the plane by calling DeriveEventsFromAnnotations. 137 void GenerateDerivedTimeLines(const GroupMetadataMap& group_metadata_map, 138 XSpace* space); 139 140 // Derives `Tensorflow Ops`, `Tensorflow Name Scope` and `Source Code` lines 141 // from device_trace. 142 void DeriveLinesFromStats(tensorflow::profiler::XPlane* device_trace); 143 144 } // namespace profiler 145 } // namespace tensorflow 146 147 #endif // TENSORFLOW_CORE_PROFILER_UTILS_DERIVED_TIMELINE_H_ 148