Add FindOrAddMutableStat to xplane_utils

Remove AddOrUpdate{Int,Str}Stat

PiperOrigin-RevId: 339324309
Change-Id: I2abb355d9311ab9a88dba5bf07f0c15e9a62e01f
This commit is contained in:
Jose Baiocchi 2020-10-27 14:01:06 -07:00 committed by TensorFlower Gardener
parent 6811ac3c58
commit 71aaef09d1
3 changed files with 24 additions and 40 deletions

View File

@ -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() {

View File

@ -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) {

View File

@ -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);