Make RemovePlaneWithName not invalidate pointers

PiperOrigin-RevId: 338175568
Change-Id: I7e305a262fad4b8c76456d9267163a32ca0d5ae0
This commit is contained in:
Jose Baiocchi 2020-10-20 18:12:23 -07:00 committed by TensorFlower Gardener
parent 0e14b0fdc4
commit c3a4d7ee32
3 changed files with 42 additions and 23 deletions

View File

@ -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",

View File

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

View File

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