Reduce EventNode size: remove raw_event_
Use const_cast<> on the pointer within XEventVisitor.
PiperOrigin-RevId: 452339522
diff --git a/tensorflow/core/profiler/utils/group_events.cc b/tensorflow/core/profiler/utils/group_events.cc
index 3a14ec2..c9ed9b0 100644
--- a/tensorflow/core/profiler/utils/group_events.cc
+++ b/tensorflow/core/profiler/utils/group_events.cc
@@ -277,9 +277,6 @@
return false;
}
-EventNode::EventNode(XEventVisitor visitor, XEvent* raw_event)
- : visitor_(std::move(visitor)), raw_event_(raw_event) {}
-
absl::optional<XStatVisitor> EventNode::GetContextStat(
int64_t stat_type) const {
std::queue<const EventNode*> nodes;
@@ -323,7 +320,8 @@
const XPlaneVisitor& plane = visitor_.Plane();
const XStatMetadata* stat_metadata = plane.GetStatMetadataByType(stat_type);
DCHECK(stat_metadata != nullptr);
- return FindOrAddMutableStat(*stat_metadata, raw_event_);
+ auto* raw_event = const_cast<XEvent*>(&visitor_.RawEvent()); // NOLINT
+ return FindOrAddMutableStat(*stat_metadata, raw_event);
}
void EventNode::SetGroupId(int64_t group_id) {
@@ -404,8 +402,7 @@
for (auto& event : *line.mutable_events()) {
XEventVisitor event_visitor(visitor, &line, &event);
GroupingEventStats stats(event_visitor);
- auto cur_node =
- std::make_unique<EventNode>(std::move(event_visitor), &event);
+ auto cur_node = std::make_unique<EventNode>(std::move(event_visitor));
if (stats.root_level.has_value()) {
cur_node->SetRootLevel(*stats.root_level);
}
diff --git a/tensorflow/core/profiler/utils/group_events.h b/tensorflow/core/profiler/utils/group_events.h
index 0616ccb..c9cb8e7 100644
--- a/tensorflow/core/profiler/utils/group_events.h
+++ b/tensorflow/core/profiler/utils/group_events.h
@@ -20,6 +20,7 @@
#include <functional>
#include <memory>
#include <string>
+#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
@@ -59,8 +60,7 @@
// pointers, a tree of EventNode is formed.
class EventNode {
public:
- // REQUIRED: raw_event should not be nullptr, visitor must wrap raw_event.
- EventNode(XEventVisitor visitor, XEvent* raw_event);
+ explicit EventNode(XEventVisitor visitor) : visitor_(std::move(visitor)) {}
EventNode(const EventNode& event_node) = delete;
EventNode& operator=(const EventNode&) = delete;
@@ -115,7 +115,6 @@
XStat* FindOrAddStatByType(int64_t stat_type);
XEventVisitor visitor_;
- XEvent* raw_event_;
std::vector<EventNode*> parents_;
std::vector<EventNode*> children_;
absl::optional<int64_t> group_id_;
diff --git a/tensorflow/core/profiler/utils/xplane_visitor.h b/tensorflow/core/profiler/utils/xplane_visitor.h
index e046845a..08f3998 100644
--- a/tensorflow/core/profiler/utils/xplane_visitor.h
+++ b/tensorflow/core/profiler/utils/xplane_visitor.h
@@ -152,6 +152,8 @@
const XPlaneVisitor& Plane() const { return *plane_; }
+ const XEvent& RawEvent() const { return *event_; }
+
int64_t Id() const { return event_->metadata_id(); }
absl::string_view Name() const { return metadata_->name(); }
