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