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/strings/string_view.h" 23 #include "absl/types/optional.h" 24 #include "tensorflow/core/platform/types.h" 25 #include "tensorflow/core/profiler/protobuf/xplane.pb.h" 26 #include "tensorflow/core/profiler/utils/group_events.h" 27 #include "tensorflow/core/profiler/utils/xplane_builder.h" 28 29 namespace tensorflow { 30 namespace profiler { 31 32 // Helper for deriving an XLine from events in another XLine. 33 class DerivedXLineBuilder { 34 public: 35 DerivedXLineBuilder(XPlaneBuilder* plane, int64 line_id, 36 absl::string_view name, int64 timestamp_ns, 37 std::vector<DerivedXLineBuilder*> dependent_lines); 38 ExpandOrAddEvents(const std::vector<XEvent> & event_per_level)39 void ExpandOrAddEvents(const std::vector<XEvent>& event_per_level) { 40 for (size_t level = 0; level < event_per_level.size(); ++level) { 41 ExpandOrAddLevelEvent(event_per_level[level], level); 42 } 43 } 44 ExpandOrAddEvent(const XEvent & event)45 void ExpandOrAddEvent(const XEvent& event) { 46 ExpandOrAddLevelEvent(event, /*level=*/0); 47 } 48 49 // Reset the last events lower than or equal to the given level. 50 void ResetLastEvents(int level = 0); 51 52 private: 53 // If the last event of the given level has the same metadata, expands it to 54 // include the time until the given event's (offset_ps + duration_ps). 55 // Otherwise, adds a new event and clears last_event_by_level_ for the levels 56 // below the given level and all levels of the dependent lines. Clearing 57 // last_event_by_level_ prevents a nested event from growing larger than the 58 // parent event(s). 59 void ExpandOrAddLevelEvent(const XEvent& event, int level); 60 ResetDependentLines()61 void ResetDependentLines() { 62 for (DerivedXLineBuilder* line : dependent_lines_) { 63 line->ResetLastEvents(); 64 } 65 } 66 67 XLineBuilder line_; 68 absl::flat_hash_map<int, absl::optional<XEventBuilder>> last_event_by_level_; 69 std::vector<DerivedXLineBuilder*> dependent_lines_; 70 }; 71 72 using SymbolResolver = std::function<absl::string_view( 73 absl::string_view hlo_module_name, absl::string_view hlo_op)>; 74 75 // Derives TF name scope and op events from the TF op's fully qualified name. 76 void ProcessTfOpEvent(absl::string_view tf_op_full_name, int64 offset_ps, 77 int64 duration_ps, absl::optional<int64> group_id, 78 XPlaneBuilder* plane_builder, 79 DerivedXLineBuilder* tf_name_scope_line_builder, 80 DerivedXLineBuilder* tf_op_line_builder); 81 82 // Derives "Step Info", "Tensorflow Ops", "XLA Ops" and "XLA Module" lines in 83 // an NVIDIA_GPU device trace from data passed as ScopedAnnotations and stored 84 // as XStats in XEvents corresponding to GPU Kernels. Consecutive annotations 85 // with the same value are merged into a single event except for XLA modules. 86 // The device_trace is both input and output. 87 void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver, 88 const GroupMetadataMap& group_metadata_map, 89 XPlane* device_trace, 90 bool step_info_only = false); 91 92 // Derives "Launch Activities Summary" line from host trace. 93 void DeriveEventsFromHostTrace(const XPlane* host_trace, 94 const GroupMetadataMap& group_metadata_map, 95 std::vector<XPlane*> device_traces); 96 97 // Loops through XPlanes of input XSpace, if it is "device" XPlane, generating 98 // derived timelines for the plane by calling DeriveEventsFromAnnotations. 99 void GenerateDerivedTimeLines(const GroupMetadataMap& group_metadata_map, 100 XSpace* space, bool step_info_only = false); 101 102 } // namespace profiler 103 } // namespace tensorflow 104 105 #endif // TENSORFLOW_CORE_PROFILER_UTILS_DERIVED_TIMELINE_H_ 106