Add RemoveEvents, XPlaneVisitor::EndTimestampsPs and XPlaneBuilder::SetEndTimestampPs

PiperOrigin-RevId: 345347851
Change-Id: I377447dd57959adddcd8229f9e9079346c7090ab

tensorflow/core/profiler/utils/BUILD

@@ -262,6 +262,7 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/strings",
     ],
 )

tensorflow/core/profiler/utils/xplane_builder.h

@@ -190,11 +190,14 @@ class XEventBuilder : public XStatsBuilder<XEvent> {
   void SetDurationPs(int64 duration_ps) {
     event_->set_duration_ps(duration_ps);
   }
-
   void SetDurationNs(int64 duration_ns) {
     SetDurationPs(NanosToPicos(duration_ns));
   }
 
+  void SetEndTimestampPs(int64 end_timestamp_ps) {
+    SetDurationPs(end_timestamp_ps - PicosToNanos(line_->timestamp_ns()) -
+                  event_->offset_ps());
+  }
   void SetEndTimestampNs(int64 end_timestamp_ns) {
     SetDurationPs(NanosToPicos(end_timestamp_ns - line_->timestamp_ns()) -
                   event_->offset_ps());

tensorflow/core/profiler/utils/xplane_utils.cc

@@ -140,6 +140,12 @@ void RemoveLine(XPlane* plane, const XLine* line) {
   Remove(plane->mutable_lines(), line);
 }
 
+void RemoveEvents(XLine* line,
+                  const absl::flat_hash_set<const XEvent*>& events) {
+  RemoveIf(line->mutable_events(),
+           [&](const XEvent* event) { return events.contains(event); });
+}
+
 void RemoveEmptyPlanes(XSpace* space) {
   RemoveIf(space->mutable_planes(),
            [&](const XPlane* plane) { return plane->lines().empty(); });

tensorflow/core/profiler/utils/xplane_utils.h

@@ -18,6 +18,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/strings/string_view.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
@@ -47,6 +48,8 @@ XStat* FindOrAddMutableStat(const XStatMetadata& stat_metadata, XEvent* event);
 
 void RemovePlane(XSpace* space, const XPlane* plane);
 void RemoveLine(XPlane* plane, const XLine* line);
+void RemoveEvents(XLine* line,
+                  const absl::flat_hash_set<const XEvent*>& events);
 
 void RemoveEmptyPlanes(XSpace* space);
 void RemoveEmptyLines(XPlane* plane);

tensorflow/core/profiler/utils/xplane_utils_test.cc

@@ -115,6 +115,29 @@ TEST(XPlaneUtilsTest, RemoveEmptyLines) {
   EXPECT_EQ(plane.lines(1).name(), "l3");
 }
 
+TEST(XPlaneUtilsTest, RemoveLine) {
+  XPlane plane;
+  const XLine* line1 = plane.add_lines();
+  const XLine* line2 = plane.add_lines();
+  const XLine* line3 = plane.add_lines();
+  RemoveLine(&plane, line2);
+  ASSERT_EQ(plane.lines_size(), 2);
+  EXPECT_EQ(&plane.lines(0), line1);
+  EXPECT_EQ(&plane.lines(1), line3);
+}
+
+TEST(XPlaneUtilsTest, RemoveEvents) {
+  XLine line;
+  const XEvent* event1 = line.add_events();
+  const XEvent* event2 = line.add_events();
+  const XEvent* event3 = line.add_events();
+  const XEvent* event4 = line.add_events();
+  RemoveEvents(&line, {event1, event3});
+  ASSERT_EQ(line.events_size(), 2);
+  EXPECT_EQ(&line.events(0), event2);
+  EXPECT_EQ(&line.events(1), event4);
+}
+
 TEST(XPlaneUtilsTest, SortXPlaneTest) {
   XPlane plane;
   XLine* line = plane.add_lines();

tensorflow/core/profiler/utils/xplane_visitor.h

@@ -100,6 +100,17 @@ class XStatsOwner {
   // Prefer ForEachStat above when multiple stat values are necessary.
   absl::optional<XStatVisitor> GetStat(int64 stat_type) const;
 
+  // Same as above that skips searching for the stat.
+  absl::optional<XStatVisitor> GetStat(
+      int64 stat_type, const XStatMetadata& stat_metadata) const {
+    for (const XStat& stat : stats_owner_->stats()) {
+      if (stat.metadata_id() == stat_metadata.id()) {
+        return XStatVisitor(plane_, &stat, &stat_metadata, stat_type);
+      }
+    }
+    return absl::nullopt;  // type does not exist in this owner.
+  }
+
  protected:
   const XPlaneVisitor* plane() const { return plane_; }
   const T* stats_owner() const { return stats_owner_; }
@@ -165,6 +176,7 @@ class XEventVisitor : public XStatsOwner<XEvent> {
   int64 EndOffsetPs() const {
     return event_->offset_ps() + event_->duration_ps();
   }
+  int64 EndTimestampPs() const { return TimestampPs() + DurationPs(); }
 
   int64 NumOccurrences() const { return event_->num_occurrences(); }
 
@@ -285,11 +297,7 @@ template <class T>
 absl::optional<XStatVisitor> XStatsOwner<T>::GetStat(int64 stat_type) const {
   const auto* stat_metadata = plane_->GetStatMetadataByType(stat_type);
   if (stat_metadata != nullptr) {
-    for (const XStat& stat : stats_owner_->stats()) {
-      if (stat.metadata_id() == stat_metadata->id()) {
-        return XStatVisitor(plane_, &stat, stat_metadata, stat_type);
-      }
-    }
+    return GetStat(stat_type, *stat_metadata);
   }
   return absl::nullopt;  // type does not exist in this owner.
 }