Make RemovePlaneWithName not invalidate pointers
PiperOrigin-RevId: 338175568 Change-Id: I7e305a262fad4b8c76456d9267163a32ca0d5ae0
This commit is contained in:
parent
0e14b0fdc4
commit
c3a4d7ee32
@ -240,8 +240,7 @@ cc_library(
|
||||
":trace_utils",
|
||||
":xplane_builder",
|
||||
":xplane_visitor",
|
||||
"//tensorflow/core:platform_base",
|
||||
"//tensorflow/core/platform:types",
|
||||
"//tensorflow/core:lib",
|
||||
"//tensorflow/core/profiler/protobuf:xplane_proto_cc",
|
||||
"@com_google_absl//absl/container:flat_hash_map",
|
||||
"@com_google_absl//absl/strings",
|
||||
@ -255,7 +254,7 @@ tf_cc_test(
|
||||
":xplane_builder",
|
||||
":xplane_utils",
|
||||
":xplane_visitor",
|
||||
"//tensorflow/core:platform_base",
|
||||
"//tensorflow/core:lib",
|
||||
"//tensorflow/core:test",
|
||||
"//tensorflow/core:test_main",
|
||||
"//tensorflow/core/profiler/protobuf:xplane_proto_cc",
|
||||
@ -276,7 +275,7 @@ cc_library(
|
||||
":xplane_builder",
|
||||
":xplane_schema",
|
||||
":xplane_utils",
|
||||
"//tensorflow/core/platform:types",
|
||||
"//tensorflow/core:lib",
|
||||
"//tensorflow/core/profiler/protobuf:xplane_proto_cc",
|
||||
"@com_google_absl//absl/container:flat_hash_map",
|
||||
"@com_google_absl//absl/strings",
|
||||
|
@ -24,6 +24,7 @@ limitations under the License.
|
||||
#include "absl/strings/string_view.h"
|
||||
#include "tensorflow/core/platform/env_time.h"
|
||||
#include "tensorflow/core/platform/logging.h"
|
||||
#include "tensorflow/core/platform/protobuf.h"
|
||||
#include "tensorflow/core/platform/types.h"
|
||||
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
|
||||
#include "tensorflow/core/profiler/utils/timespan.h"
|
||||
@ -34,6 +35,24 @@ namespace tensorflow {
|
||||
namespace profiler {
|
||||
namespace {
|
||||
|
||||
// Removes all elements from 'array' for which pred is true.
|
||||
template <typename T, typename Pred>
|
||||
void RemoveIf(protobuf::RepeatedPtrField<T>* array, Pred&& pred) {
|
||||
T** const begin = array->mutable_data();
|
||||
T** const end = begin + array->size();
|
||||
T** write = begin;
|
||||
while (write < end && !pred(*write)) ++write;
|
||||
if (write == end) return;
|
||||
// 'write' points to the first element to be removed.
|
||||
for (T** scan = write + 1; scan < end; ++scan) {
|
||||
if (!pred(*scan)) std::swap(*scan, *write++);
|
||||
}
|
||||
const int new_size = write - begin;
|
||||
const int size = array->size();
|
||||
DCHECK_GE(size, new_size);
|
||||
array->DeleteSubrange(new_size, size - new_size);
|
||||
}
|
||||
|
||||
// Creates a Timespan from an XEvent.
|
||||
// WARNING: This should only be used when comparing events from the same XLine.
|
||||
Timespan XEventTimespan(const XEvent& event) {
|
||||
@ -113,28 +132,18 @@ void AddOrUpdateStrStat(int64 metadata_id, absl::string_view value,
|
||||
}
|
||||
|
||||
void RemovePlaneWithName(XSpace* space, absl::string_view name) {
|
||||
auto* planes = space->mutable_planes();
|
||||
planes->erase(
|
||||
std::remove_if(planes->begin(), planes->end(),
|
||||
[&](const XPlane& plane) { return plane.name() == name; }),
|
||||
planes->end());
|
||||
RemoveIf(space->mutable_planes(),
|
||||
[&](const XPlane* plane) { return plane->name() == name; });
|
||||
}
|
||||
|
||||
void RemoveEmptyPlanes(XSpace* space) {
|
||||
auto* planes = space->mutable_planes();
|
||||
planes->erase(std::remove_if(planes->begin(), planes->end(),
|
||||
[&](const XPlane& plane) {
|
||||
return plane.lines_size() == 0;
|
||||
}),
|
||||
planes->end());
|
||||
RemoveIf(space->mutable_planes(),
|
||||
[&](const XPlane* plane) { return plane->lines().empty(); });
|
||||
}
|
||||
|
||||
void RemoveEmptyLines(XPlane* plane) {
|
||||
auto* lines = plane->mutable_lines();
|
||||
lines->erase(std::remove_if(
|
||||
lines->begin(), lines->end(),
|
||||
[&](const XLine& line) { return line.events_size() == 0; }),
|
||||
lines->end());
|
||||
RemoveIf(plane->mutable_lines(),
|
||||
[&](const XLine* line) { return line->events().empty(); });
|
||||
}
|
||||
|
||||
bool XEventsComparator::operator()(const XEvent* a, const XEvent* b) const {
|
||||
|
@ -56,17 +56,28 @@ TEST(XPlaneUtilsTest, RemovePlaneWithName) {
|
||||
RemovePlaneWithName(&space, "non-exist");
|
||||
EXPECT_EQ(space.planes_size(), 0);
|
||||
|
||||
space.add_planes()->set_name("p1");
|
||||
space.add_planes()->set_name("p2");
|
||||
space.add_planes()->set_name("p3");
|
||||
auto* p1 = space.add_planes();
|
||||
p1->set_name("p1");
|
||||
auto* p2 = space.add_planes();
|
||||
p2->set_name("p2");
|
||||
auto* p3 = space.add_planes();
|
||||
p3->set_name("p3");
|
||||
|
||||
RemovePlaneWithName(&space, "non-exist");
|
||||
EXPECT_EQ(space.planes_size(), 3);
|
||||
|
||||
RemovePlaneWithName(&space, "p2");
|
||||
EXPECT_EQ(space.planes_size(), 2);
|
||||
EXPECT_EQ(p1, FindPlaneWithName(space, "p1"));
|
||||
EXPECT_EQ(p3, FindPlaneWithName(space, "p3"));
|
||||
|
||||
RemovePlaneWithName(&space, "p1");
|
||||
EXPECT_EQ(space.planes_size(), 1);
|
||||
EXPECT_EQ(p3, FindPlaneWithName(space, "p3"));
|
||||
|
||||
RemovePlaneWithName(&space, "p1");
|
||||
EXPECT_EQ(space.planes_size(), 1);
|
||||
|
||||
RemovePlaneWithName(&space, "p3");
|
||||
EXPECT_EQ(space.planes_size(), 0);
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user