Add MetadataMatcher to help processing xplane.

PiperOrigin-RevId: 285905624
Change-Id: I68e174c654ee975067065dbfcb8b59d08d1fdb00
This commit is contained in:
A. Unique TensorFlower 2019-12-16 20:45:51 -08:00 committed by TensorFlower Gardener
parent b7bd9b0950
commit 8e42b57fc4
5 changed files with 374 additions and 0 deletions

View File

@ -158,6 +158,29 @@ typename Collection::value_type::second_type& LookupOrInsert(
typename Collection::value_type(key, value));
}
// Saves the reverse mapping into reverse. Returns true if values could all be
// inserted.
template <typename M, typename ReverseM>
bool ReverseMap(const M& m, ReverseM* reverse) {
bool all_unique = true;
for (const auto& kv : m) {
if (!InsertOrUpdate(reverse, kv.second, kv.first)) {
all_unique = false;
}
}
return all_unique;
}
// Like ReverseMap above, but returns its output m. Return type has to
// be specified explicitly. Example:
// M::M(...) : m_(...), r_(ReverseMap<decltype(r_)>(m_)) {}
template <typename ReverseM, typename M>
ReverseM ReverseMap(const M& m) {
typename std::remove_const<ReverseM>::type reverse;
ReverseMap(m, &reverse);
return reverse;
}
// Erases the m item identified by the given key, and returns the value
// associated with that key. It is assumed that the value (i.e., the
// mapped_type) is a pointer. Returns null if the key was not found in the

View File

@ -142,3 +142,32 @@ cc_library(
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "metadata_matcher",
srcs = ["metadata_matcher.cc"],
hdrs = ["metadata_matcher.h"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core/profiler/protobuf:xplane_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
"@com_google_absl//absl/types:span",
],
)
tf_cc_test(
name = "metadata_matcher_test",
size = "small",
srcs = ["metadata_matcher_test.cc"],
deps = [
":metadata_matcher",
":xplane_schema",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
],
)

View File

@ -0,0 +1,145 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/profiler/utils/metadata_matcher.h"
#include "absl/strings/string_view.h"
namespace tensorflow {
namespace profiler {
namespace {
using ::tensorflow::profiler::XEvent;
using ::tensorflow::profiler::XPlane;
using ::tensorflow::profiler::XStat;
absl::flat_hash_map<int64, int> CreateEventMetadataMap(
const XPlane& xplane,
const std::vector<std::pair<const absl::Span<const absl::string_view>,
/*first_event_type*/ int>>&
event_type_metadata_maps) {
absl::flat_hash_map<int64, int> id_to_event_type_map;
for (const auto& id_and_event_metadata : xplane.event_metadata()) {
int64 id = id_and_event_metadata.first;
absl::string_view event_name = id_and_event_metadata.second.name();
for (const auto& event_type_metadata_map_and_first_event_type :
event_type_metadata_maps) {
auto event_type_metadata_map =
event_type_metadata_map_and_first_event_type.first;
int first_event_type =
event_type_metadata_map_and_first_event_type.second;
for (int i = 0; i < event_type_metadata_map.size(); ++i) {
if (event_type_metadata_map[i] == event_name) {
id_to_event_type_map[id] = first_event_type + i;
break;
}
}
}
}
return id_to_event_type_map;
}
absl::flat_hash_map<int64, int> CreateStatMetadataMap(
const XPlane& xplane,
const absl::Span<const absl::string_view> stat_type_str_map) {
absl::flat_hash_map<int64, int> id_to_stat_type_map;
for (const auto& id_and_stat_metadata : xplane.stat_metadata()) {
int64 id = id_and_stat_metadata.first;
absl::string_view stat_name = id_and_stat_metadata.second.name();
for (int stat_type = 0; stat_type < stat_type_str_map.size(); ++stat_type) {
if (stat_type_str_map[stat_type] == stat_name) {
id_to_stat_type_map[id] = stat_type;
break;
}
}
}
return id_to_stat_type_map;
}
} // namespace
MetadataMatcher::MetadataMatcher(
const XPlane& xplane,
const std::vector<std::pair<const absl::Span<const absl::string_view>,
/*first_event_type*/ int>>&
event_type_metadata_maps,
const absl::Span<const absl::string_view> stat_type_str_map)
: id_to_event_type_map_(
CreateEventMetadataMap(xplane, event_type_metadata_maps)),
id_to_stat_type_map_(CreateStatMetadataMap(xplane, stat_type_str_map)),
event_type_to_id_map_(gtl::ReverseMap<decltype(event_type_to_id_map_)>(
id_to_event_type_map_)),
stat_type_to_id_map_(gtl::ReverseMap<decltype(stat_type_to_id_map_)>(
id_to_stat_type_map_)) {}
const XStat* MetadataMatcher::GetStat(const XEvent& event,
int stat_type) const {
for (const auto& stat : event.stats()) {
if (GetStatType(stat) == stat_type) {
return &stat;
}
}
return nullptr;
}
absl::optional<std::tuple<const XStat*, const XStat*>>
MetadataMatcher::GetStats(const XEvent& event, int first_stat_type,
int second_stat_type) const {
const XStat* first_stat = nullptr;
const XStat* second_stat = nullptr;
for (const auto& stat : event.stats()) {
if (GetStatType(stat) == first_stat_type) {
first_stat = &stat;
} else if (GetStatType(stat) == second_stat_type) {
second_stat = &stat;
}
}
if (first_stat && second_stat) {
return std::make_tuple(first_stat, second_stat);
}
return absl::nullopt;
}
absl::optional<std::tuple<const XStat*, const XStat*, const XStat*>>
MetadataMatcher::GetStats(const XEvent& event, int first_stat_type,
int second_stat_type, int third_stat_type) const {
const XStat* first_stat = nullptr;
const XStat* second_stat = nullptr;
const XStat* third_stat = nullptr;
for (const auto& stat : event.stats()) {
if (GetStatType(stat) == first_stat_type) {
first_stat = &stat;
} else if (GetStatType(stat) == second_stat_type) {
second_stat = &stat;
} else if (GetStatType(stat) == third_stat_type) {
third_stat = &stat;
}
}
if (first_stat && second_stat && third_stat) {
return std::make_tuple(first_stat, second_stat, third_stat);
}
return absl::nullopt;
}
absl::optional<int64> MetadataMatcher::GetIntStatValue(const XEvent& event,
int stat_type) const {
if (const XStat* stat = GetStat(event, stat_type)) {
return stat->int64_value();
}
return absl::nullopt;
}
} // namespace profiler
} // namespace tensorflow

View File

@ -0,0 +1,108 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_PROFILER_UTILS_METADATA_MATCHER_H_
#define TENSORFLOW_CORE_PROFILER_UTILS_METADATA_MATCHER_H_
#include "absl/container/flat_hash_map.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "absl/types/span.h"
#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
namespace tensorflow {
namespace profiler {
// Builds mapping between metadata ids and interesting event and stat types.
// Event and stat types are represented in integer ids. Multiple spans of event
// types can be passed with offset values (i.e., first_event_type) to be
// used to calculate integer ids for event types. Spans and offset values are
// expected to result in a unique integer id for each event type.
class MetadataMatcher {
public:
explicit MetadataMatcher(
const XPlane& xplane,
const std::vector<std::pair<const absl::Span<const absl::string_view>,
/*first_event_type*/ int>>&
event_type_metadata_maps,
const absl::Span<const absl::string_view> stat_type_str_map);
// Returns EventType if input is one of interesting event types.
// Otherwise, it returns kUnknownEventType.
int GetEventType(const XEvent& xevent) const {
return gtl::FindWithDefault(id_to_event_type_map_, xevent.metadata_id(),
/*kUnknownEventType*/ 0);
}
// Overload of GetEventType function.
// Returns EventType if input is one of interesting event types.
// Otherwise, it returns kUnknownEventType.
int GetEventType(int64 metadata_id) const {
return gtl::FindWithDefault(id_to_event_type_map_, metadata_id,
/*kUnknownEventType*/ 0);
}
// Returns metadata id if xplane has the input event type.
absl::optional<int64> GetEventMetadataId(int event_type) const {
if (const int64* id = gtl::FindOrNull(event_type_to_id_map_, event_type)) {
return *id;
}
return absl::nullopt;
}
// Returns StatType if input is one of interesting stat types.
// Otherwise, it returns kUnknownStatType.
int GetStatType(const XStat& xstat) const {
return gtl::FindWithDefault(id_to_stat_type_map_, xstat.metadata_id(),
/*kUnknownStatType*/ 0);
}
// Returns metadata id if xplane has the input stat type.
absl::optional<int64> GetStatMetadataId(int stat_type) const {
if (const int64* id = gtl::FindOrNull(stat_type_to_id_map_, stat_type)) {
return *id;
}
return absl::nullopt;
}
const XStat* GetStat(const XEvent& event, int stat_type) const;
absl::optional<std::tuple<const XStat*, const XStat*>> GetStats(
const XEvent& event, int first_stat_type, int second_stat_type) const;
absl::optional<std::tuple<const XStat*, const XStat*, const XStat*>> GetStats(
const XEvent& event, int first_stat_type, int second_stat_type,
int third_stat_type) const;
absl::optional<int64> GetIntStatValue(const XEvent& event,
int stat_type) const;
private:
// Maps from metada ids to interesting event and stat types.
// Uninteresting event and stat types are not cached in these maps and
// considered to be kUnknown*.
const absl::flat_hash_map<int64, int> id_to_event_type_map_;
const absl::flat_hash_map<int64, int> id_to_stat_type_map_;
// Reverse of the above.
const absl::flat_hash_map<int, int64> event_type_to_id_map_;
const absl::flat_hash_map<int, int64> stat_type_to_id_map_;
};
} // namespace profiler
} // namespace tensorflow
#endif // TENSORFLOW_CORE_PROFILER_UTILS_METADATA_MATCHER_H_

View File

@ -0,0 +1,69 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/profiler/utils/metadata_matcher.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/profiler/utils/xplane_schema.h"
namespace tensorflow {
namespace profiler {
namespace {
using ::tensorflow::profiler::XEventMetadata;
using ::tensorflow::profiler::XPlane;
using ::tensorflow::profiler::XStatMetadata;
TEST(MetadataMatcherTest, GetHostEventTypeTest) {
for (int event_type = HostEventType::kFirstHostEventType;
event_type <= HostEventType::kLastHostEventType; ++event_type) {
XPlane xplane;
XEventMetadata& metadata = (*xplane.mutable_event_metadata())[0];
metadata.set_id(0);
metadata.set_name(std::string(
GetHostEventTypeStr(static_cast<HostEventType>(event_type))));
MetadataMatcher metadata_matcher(
xplane,
{{GetHostEventTypeStrMap(), HostEventType::kFirstHostEventType}},
GetStatTypeStrMap());
XEvent event;
event.set_metadata_id(0);
EXPECT_EQ(metadata_matcher.GetEventType(event), event_type);
}
}
TEST(MetadataMatcherTest, GetStatTypeTest) {
for (int stat_type = StatType::kFirstStatType;
stat_type <= StatType::kLastStatType; ++stat_type) {
XPlane xplane;
XStatMetadata& metadata = (*xplane.mutable_stat_metadata())[0];
metadata.set_id(0);
metadata.set_name(
std::string(GetStatTypeStr(static_cast<StatType>(stat_type))));
MetadataMatcher metadata_matcher(
xplane,
{{GetHostEventTypeStrMap(), HostEventType::kFirstHostEventType}},
GetStatTypeStrMap());
XStat stat;
stat.set_metadata_id(0);
EXPECT_EQ(metadata_matcher.GetStatType(stat), stat_type);
}
}
} // namespace
} // namespace profiler
} // namespace tensorflow