Add MetadataMatcher to help processing xplane.

PiperOrigin-RevId: 285905624
Change-Id: I68e174c654ee975067065dbfcb8b59d08d1fdb00

tensorflow/core/lib/gtl/map_util.h

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

tensorflow/core/profiler/utils/BUILD

@@ -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",
+    ],
+)

tensorflow/core/profiler/utils/metadata_matcher.cc

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

tensorflow/core/profiler/utils/metadata_matcher.h

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

tensorflow/core/profiler/utils/metadata_matcher_test.cc

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