From 71aaef09d1a407a362a216a8f77147b905d6ed3a Mon Sep 17 00:00:00 2001
From: Jose Baiocchi <jbaiocchi@google.com>
Date: Tue, 27 Oct 2020 14:01:06 -0700
Subject: [PATCH] Add FindOrAddMutableStat to xplane_utils

Remove AddOrUpdate{Int,Str}Stat

PiperOrigin-RevId: 339324309
Change-Id: I2abb355d9311ab9a88dba5bf07f0c15e9a62e01f
---
 .../core/profiler/utils/group_events.cc       | 38 ++++++++++---------
 .../core/profiler/utils/xplane_utils.cc       | 20 ++--------
 tensorflow/core/profiler/utils/xplane_utils.h |  6 +--
 3 files changed, 24 insertions(+), 40 deletions(-)

diff --git a/tensorflow/core/profiler/utils/group_events.cc b/tensorflow/core/profiler/utils/group_events.cc
index a78fe5a1513..7c32b3d0753 100644
--- a/tensorflow/core/profiler/utils/group_events.cc
+++ b/tensorflow/core/profiler/utils/group_events.cc
@@ -405,8 +405,9 @@ std::string EventNode::GetGroupName() const {
 
 void EventNode::SetGroupId(int64 group_id) {
   group_id_ = group_id;
-  AddOrUpdateIntStat(*plane_->GetStatMetadataId(StatType::kGroupId), group_id,
-                     raw_event_);
+  FindOrAddMutableStat(*plane_->GetStatMetadataId(StatType::kGroupId),
+                       raw_event_)
+      ->set_int64_value(group_id);
 }
 
 void EventNode::PropagateGroupId(int64 group_id,
@@ -435,31 +436,32 @@ void EventNode::PropagateGroupId(int64 group_id,
 }
 
 void EventNode::AddStepName(absl::string_view step_name) {
-  AddOrUpdateStrStat(*plane_->GetStatMetadataId(StatType::kStepName), step_name,
-                     raw_event_);
+  FindOrAddMutableStat(*plane_->GetStatMetadataId(StatType::kStepName),
+                       raw_event_)
+      ->set_str_value(step_name.data(), step_name.size());
 }
 
 void EventNode::AddSelectedGroupIds(
     const GroupMetadataMap& group_metadata_map) {
+  const auto& group_metadata = group_metadata_map.at(*group_id_);
   std::vector<int64> group_ids;
-  group_ids.reserve(1 + group_metadata_map.at(*group_id_).parents.size() +
-                    group_metadata_map.at(*group_id_).children.size());
+  group_ids.reserve(1 + group_metadata.parents.size() +
+                    group_metadata.children.size());
   group_ids.push_back(*group_id_);
-  group_ids.insert(group_ids.end(),
-                   group_metadata_map.at(*group_id_).parents.begin(),
-                   group_metadata_map.at(*group_id_).parents.end());
-  group_ids.insert(group_ids.end(),
-                   group_metadata_map.at(*group_id_).children.begin(),
-                   group_metadata_map.at(*group_id_).children.end());
-  AddOrUpdateStrStat(
-      *plane_->GetStatMetadataId(StatType::kSelectedGroupIds),
-      absl::StrCat("?selected_group_ids=", absl::StrJoin(group_ids, ",")),
-      raw_event_);
+  group_ids.insert(group_ids.end(), group_metadata.parents.begin(),
+                   group_metadata.parents.end());
+  group_ids.insert(group_ids.end(), group_metadata.children.begin(),
+                   group_metadata.children.end());
+  FindOrAddMutableStat(*plane_->GetStatMetadataId(StatType::kSelectedGroupIds),
+                       raw_event_)
+      ->set_str_value(
+          absl::StrCat("?selected_group_ids=", absl::StrJoin(group_ids, ",")));
 }
 
 void EventNode::SetIsEager(bool is_eager) {
-  AddOrUpdateIntStat(*plane_->GetStatMetadataId(StatType::kIsEager),
-                     is_eager ? 1 : 0, raw_event_);
+  FindOrAddMutableStat(*plane_->GetStatMetadataId(StatType::kIsEager),
+                       raw_event_)
+      ->set_int64_value(is_eager ? 1 : 0);
 }
 
 bool EventNode::IsEager() {
diff --git a/tensorflow/core/profiler/utils/xplane_utils.cc b/tensorflow/core/profiler/utils/xplane_utils.cc
index 4af4fb79491..68f7c127e2f 100644
--- a/tensorflow/core/profiler/utils/xplane_utils.cc
+++ b/tensorflow/core/profiler/utils/xplane_utils.cc
@@ -117,29 +117,15 @@ bool IsNested(const XEvent& event, const XEvent& parent) {
   return XEventTimespan(parent).Includes(XEventTimespan(event));
 }
 
-void AddOrUpdateIntStat(int64 metadata_id, int64 value, XEvent* event) {
+XStat* FindOrAddMutableStat(int64 metadata_id, XEvent* event) {
   for (auto& stat : *event->mutable_stats()) {
     if (stat.metadata_id() == metadata_id) {
-      stat.set_int64_value(value);
-      return;
+      return &stat;
     }
   }
   XStat* stat = event->add_stats();
   stat->set_metadata_id(metadata_id);
-  stat->set_int64_value(value);
-}
-
-void AddOrUpdateStrStat(int64 metadata_id, absl::string_view value,
-                        XEvent* event) {
-  for (auto& stat : *event->mutable_stats()) {
-    if (stat.metadata_id() == metadata_id) {
-      stat.set_str_value(std::string(value));
-      return;
-    }
-  }
-  XStat* stat = event->add_stats();
-  stat->set_metadata_id(metadata_id);
-  stat->set_str_value(std::string(value));
+  return stat;
 }
 
 void RemovePlane(XSpace* space, const XPlane* plane) {
diff --git a/tensorflow/core/profiler/utils/xplane_utils.h b/tensorflow/core/profiler/utils/xplane_utils.h
index 77abb2c53d7..2183c1151dc 100644
--- a/tensorflow/core/profiler/utils/xplane_utils.h
+++ b/tensorflow/core/profiler/utils/xplane_utils.h
@@ -44,11 +44,7 @@ std::vector<XPlane*> FindMutablePlanesWithPrefix(XSpace* space,
 bool IsNested(const tensorflow::profiler::XEvent& event,
               const tensorflow::profiler::XEvent& parent);
 
-void AddOrUpdateIntStat(int64 metadata_id, int64 value,
-                        tensorflow::profiler::XEvent* event);
-
-void AddOrUpdateStrStat(int64 metadata_id, absl::string_view value,
-                        tensorflow::profiler::XEvent* event);
+XStat* FindOrAddMutableStat(int64 metadata_id, XEvent* event);
 
 void RemovePlane(XSpace* space, const XPlane* plane);