From 1e22a995276499661c50f9f122d3de498429300c Mon Sep 17 00:00:00 2001 From: Jiho Choi Date: Thu, 28 May 2020 15:30:46 -0700 Subject: [PATCH] Refactor the code by (1) using EventVisitor in EventNode and (2) copying the first child EventNode when creating the virtual EventNode instead creating a new XEvent. PiperOrigin-RevId: 313671525 Change-Id: I421f08025c49a9deda5231184d287535486ac13b --- .../core/profiler/utils/group_events.cc | 68 ++++++++----------- tensorflow/core/profiler/utils/group_events.h | 37 +++++----- 2 files changed, 46 insertions(+), 59 deletions(-) diff --git a/tensorflow/core/profiler/utils/group_events.cc b/tensorflow/core/profiler/utils/group_events.cc index 42961492225..c07a397ee54 100644 --- a/tensorflow/core/profiler/utils/group_events.cc +++ b/tensorflow/core/profiler/utils/group_events.cc @@ -103,16 +103,6 @@ int64 GetEventType(const XPlaneVisitor& visitor, const XEvent& event) { } } -const XStat* GetStat(const XPlaneVisitor& visitor, const XEvent& event, - int64 stat_type) { - for (const auto& stat : event.stats()) { - if (visitor.GetStatType(stat) == stat_type) { - return &stat; - } - } - return nullptr; -} - void SetGroupId(const XPlaneVisitor& visitor, int64 group_id, XEvent* event) { AddOrUpdateIntStat(*visitor.GetStatMetadataId(StatType::kGroupId), group_id, event); @@ -146,8 +136,7 @@ bool NeedsVirtualEventsForAsyncExecutor( bool HasFunctionRun(EventNode* event_node) { for (EventNode* child : event_node->GetChildren()) { - if (child->GetPlaneVisitor().GetEventType(child->GetEvent()) == - HostEventType::kFunctionRun) { + if (child->GetEventVisitor().Type() == HostEventType::kFunctionRun) { return true; } } @@ -156,8 +145,21 @@ bool HasFunctionRun(EventNode* event_node) { } // namespace +EventNode::EventNode(const XPlaneVisitor* plane, XLine* raw_line, + XEvent* raw_event) + : plane_(plane), + visitor_(plane, raw_line, raw_event), + raw_line_(raw_line), + raw_event_(raw_event) {} + +EventNode::EventNode(const EventNode& event_node) + : plane_(event_node.plane_), + visitor_(event_node.plane_, event_node.raw_line_, event_node.raw_event_), + raw_line_(event_node.raw_line_), + raw_event_(event_node.raw_event_) {} + const XStat* EventNode::GetContextStat(int64 stat_type) const { - if (const XStat* stat = GetStat(*visitor_, *event_, stat_type)) { + if (const XStat* stat = visitor_.GetStats(stat_type)) { return stat; } else if (parent_) { return parent_->GetContextStat(stat_type); @@ -168,7 +170,7 @@ const XStat* EventNode::GetContextStat(int64 stat_type) const { std::string EventNode::GetGroupName() const { std::vector name_parts; if (const XStat* graph_type_stat = GetContextStat(StatType::kGraphType)) { - XStatVisitor stat(visitor_, graph_type_stat); + XStatVisitor stat(plane_, graph_type_stat); name_parts.push_back(stat.ToString()); } int64 step_num = group_id_.value_or(0); @@ -184,7 +186,7 @@ std::string EventNode::GetGroupName() const { void EventNode::PropagateGroupId(int64 group_id) { group_id_ = group_id; - SetGroupId(*visitor_, group_id, event_); + SetGroupId(*plane_, group_id, raw_event_); for (const auto& child : children_) { // Skip if it already belongs to a group. Some nodes may be added multiple // times as child (e.g., sometimes async ops are executed synchronously and @@ -196,13 +198,13 @@ void EventNode::PropagateGroupId(int64 group_id) { } void EventNode::AddStepName(absl::string_view step_name) { - AddOrUpdateStrStat(*visitor_->GetStatMetadataId(StatType::kStepName), - step_name, event_); + AddOrUpdateStrStat(*plane_->GetStatMetadataId(StatType::kStepName), step_name, + raw_event_); } void EventNode::SetIsEager(bool is_eager) { - AddOrUpdateIntStat(*visitor_->GetStatMetadataId(StatType::kIsEager), - is_eager ? 1 : 0, event_); + AddOrUpdateIntStat(*plane_->GetStatMetadataId(StatType::kIsEager), + is_eager ? 1 : 0, raw_event_); } bool EventNode::IsEager() { @@ -213,14 +215,9 @@ bool EventNode::IsEager() { FindParent(HostEventType::kEagerKernelExecute) != nullptr; } -bool EventNode::IsNestedIn(EventNode* parent) { - return parent && IsNested(GetEvent(), parent->GetEvent()); -} - EventNode* EventNode::FindParent(int64 event_type) { if (parent_) { - if (GetEventType(parent_->GetPlaneVisitor(), parent_->GetEvent()) == - event_type) { + if (parent_->GetEventVisitor().Type() == event_type) { return parent_; } return parent_->FindParent(event_type); @@ -233,10 +230,11 @@ void EventForest::ConnectIntraThread(const XPlaneVisitor& visitor, for (auto& line : *plane->mutable_lines()) { std::vector parent_nodes; for (auto& event : *line.mutable_events()) { - auto cur_node = absl::make_unique(&visitor, &event); + auto cur_node = absl::make_unique(&visitor, &line, &event); while (!parent_nodes.empty()) { EventNode* parent_node = parent_nodes.back(); - if (cur_node->IsNestedIn(parent_node)) { + if (parent_node->GetEventVisitor().GetTimespan().Includes( + cur_node->GetEventVisitor().GetTimespan())) { parent_node->AddChild(cur_node.get()); break; } else { @@ -357,12 +355,8 @@ void EventForest::CreateVirtualEventsForHostTrainingLoop() { if (!iter_num) continue; EventNode*& virtual_event_node = virtual_event_node_map[step_id][iter_num]; if (!virtual_event_node) { - std::unique_ptr new_virtual_event = - CreateVirtualEvent(*step_id_stat, *iter_num_stat); - auto new_virtual_event_node = absl::make_unique( - &executor_event_node->GetPlaneVisitor(), new_virtual_event.get()); - // virtual_event_container_ keeps new_virtual_event alive. - virtual_event_container_.push_back(std::move(new_virtual_event)); + auto new_virtual_event_node = + absl::make_unique(*executor_event_node); virtual_event_node = new_virtual_event_node.get(); // event_node_map_ keeps new_virtual_event_node alive. event_node_map_[HostEventType::kHostTrainingLoopIteration].push_back( @@ -380,12 +374,8 @@ void EventForest::CreateVirtualEventsForAsyncExecutor() { for (auto& eager_kernel_execute_event_node : *eager_kernel_execute_event_node_list) { if (HasFunctionRun(eager_kernel_execute_event_node.get())) { - auto new_virtual_event = absl::make_unique(); - auto new_virtual_event_node = absl::make_unique( - &eager_kernel_execute_event_node->GetPlaneVisitor(), - new_virtual_event.get()); - // virtual_event_container_ keeps new_virtual_event alive. - virtual_event_container_.push_back(std::move(new_virtual_event)); + auto new_virtual_event_node = + absl::make_unique(*eager_kernel_execute_event_node); virtual_event_node = new_virtual_event_node.get(); // event_node_map_ keeps new_virtual_event_node alive. event_node_map_[HostEventType::kAsyncExecutorTraceContext].push_back( diff --git a/tensorflow/core/profiler/utils/group_events.h b/tensorflow/core/profiler/utils/group_events.h index 4b6fc58e3b8..49a7b349589 100644 --- a/tensorflow/core/profiler/utils/group_events.h +++ b/tensorflow/core/profiler/utils/group_events.h @@ -47,12 +47,10 @@ struct InterThreadConnectInfo { // pointers, a tree of EventNode is formed. class EventNode { public: - // REQUIRED: visitor and event should not be nullptr. - explicit EventNode(const XPlaneVisitor* visitor, XEvent* event) - : visitor_(visitor), event_(event) { - DCHECK(visitor); - DCHECK(event); - } + // REQUIRED: all inputs should not be nullptr. + EventNode(const XPlaneVisitor* plane, XLine* raw_line, XEvent* raw_event); + + EventNode(const EventNode& event_node); EventNode* GetParent() const { return parent_; } @@ -70,9 +68,9 @@ class EventNode { // Sets group_id for this node and its descendants. void PropagateGroupId(int64 group_id); - const XPlaneVisitor& GetPlaneVisitor() const { return *visitor_; } + const XPlaneVisitor& GetPlaneVisitor() const { return *plane_; } - const XEvent& GetEvent() const { return *event_; } + const XEventVisitor& GetEventVisitor() const { return visitor_; } const XStat* GetContextStat(int64 stat_type) const; @@ -89,8 +87,10 @@ class EventNode { EventNode* FindParent(int64 event_type); private: - const XPlaneVisitor* visitor_; - XEvent* event_; + const XPlaneVisitor* plane_; + XEventVisitor visitor_; + XLine* raw_line_; + XEvent* raw_event_; EventNode* parent_ = nullptr; std::vector children_; absl::optional group_id_; @@ -100,8 +100,6 @@ using EventNodeMap = absl::flat_hash_map>>; -using VirtualEventContainer = std::vector>; - using EventGroupNameMap = absl::flat_hash_map; // Creates a forest of EventNode by stitching events in space using the nesting @@ -141,20 +139,19 @@ class EventForest { // Sets the is_eager stat to true for the eagerly executed CPU TF op events. void MarkEagerlyExecutedCpuTfOps(); - // Create virtual events of HostEventType::kHostTrainingLoopIteration and - // event nodes for them. A virtual event is created for each iteration of the - // host training loop and connected to the - // HostEventType::kExecutorStateProcess event nodes of the iteration. + // Create virtual events of HostEventType::kHostTrainingLoopIteration. A + // virtual event is created for each iteration of the host training loop and + // connected to the HostEventType::kExecutorStateProcess events of the + // iteration. void CreateVirtualEventsForHostTrainingLoop(); - // Create virutal events of HostEventType::kAsyncExecutorTraceContext and - // event nodes for them. A virtual event is created for every FunctionRun and - // the following eager ops (e.g., for Keras callback). + // Create virutal events of HostEventType::kAsyncExecutorTraceContext. A + // virtual event is created for every FunctionRun and the following eager ops + // (e.g., for Keras callback). void CreateVirtualEventsForAsyncExecutor(); EventNodeMap event_node_map_; std::vector visitors_; - VirtualEventContainer virtual_event_container_; EventGroupNameMap event_group_name_map_; };