device caps utils separated from post processing to be used in worker side
PiperOrigin-RevId: 338351871
Change-Id: I1bc1e9e3355999ac2a0ce8f15060e22a35b05813
diff --git a/tensorflow/core/profiler/utils/BUILD b/tensorflow/core/profiler/utils/BUILD
index 1dc8836..7476c5a 100644
--- a/tensorflow/core/profiler/utils/BUILD
+++ b/tensorflow/core/profiler/utils/BUILD
@@ -526,3 +526,21 @@
"@com_google_absl//absl/strings",
],
)
+
+cc_library(
+ name = "device_caps_utils",
+ srcs = ["device_caps_utils.cc"],
+ hdrs = ["device_caps_utils.h"],
+ copts = tf_profiler_copts(),
+ visibility = [":friends"],
+ deps = [
+ ":xplane_builder",
+ ":xplane_schema",
+ ":xplane_visitor",
+ "//tensorflow/core:platform_base",
+ "//tensorflow/core/platform:types",
+ "//tensorflow/core/profiler/protobuf:hardware_types_proto_cc",
+ "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
+ "@com_google_absl//absl/strings",
+ ],
+)
diff --git a/tensorflow/core/profiler/utils/device_caps_utils.cc b/tensorflow/core/profiler/utils/device_caps_utils.cc
new file mode 100644
index 0000000..44ab169
--- /dev/null
+++ b/tensorflow/core/profiler/utils/device_caps_utils.cc
@@ -0,0 +1,82 @@
+/* Copyright 2020 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/device_caps_utils.h"
+
+#include "tensorflow/core/profiler/utils/xplane_builder.h"
+#include "tensorflow/core/profiler/utils/xplane_schema.h"
+#include "tensorflow/core/profiler/utils/xplane_visitor.h"
+
+namespace tensorflow {
+namespace profiler {
+
+void SetDeviceCaps(const DeviceCapabilities& caps, XPlane* plane) {
+ XPlaneBuilder xplane(plane);
+ int clock_rate_in_khz =
+ static_cast<int>(caps.clock_rate_in_ghz() * 1000000.0);
+ xplane.AddStatValue(*xplane.GetOrCreateStatMetadata(
+ GetStatTypeStr(StatType::kDevCapClockRateKHz)),
+ clock_rate_in_khz);
+ xplane.AddStatValue(*xplane.GetOrCreateStatMetadata(
+ GetStatTypeStr(StatType::kDevCapCoreCount)),
+ caps.num_cores());
+ xplane.AddStatValue(*xplane.GetOrCreateStatMetadata(
+ GetStatTypeStr(StatType::kDevCapMemoryBandwidth)),
+ caps.memory_bandwidth());
+ xplane.AddStatValue(*xplane.GetOrCreateStatMetadata(
+ GetStatTypeStr(StatType::kDevCapMemorySize)),
+ caps.memory_size_in_bytes());
+ if (caps.has_compute_capability()) {
+ xplane.AddStatValue(*xplane.GetOrCreateStatMetadata(
+ GetStatTypeStr(StatType::kDevCapComputeCapMajor)),
+ caps.compute_capability().major());
+ xplane.AddStatValue(*xplane.GetOrCreateStatMetadata(
+ GetStatTypeStr(StatType::kDevCapComputeCapMinor)),
+ caps.compute_capability().minor());
+ }
+}
+
+DeviceCapabilities GetDeviceCaps(const XPlane& plane) {
+ DeviceCapabilities caps;
+ XPlaneVisitor xplane(&plane);
+ xplane.ForEachStat([&](const tensorflow::profiler::XStatVisitor& stat) {
+ if (!stat.Type().has_value()) return;
+ switch (stat.Type().value()) {
+ case StatType::kDevCapClockRateKHz:
+ caps.set_clock_rate_in_ghz(stat.IntOrUintValue() * 1000000.0);
+ break;
+ case StatType::kDevCapCoreCount:
+ caps.set_num_cores(stat.IntOrUintValue());
+ break;
+ case StatType::kDevCapMemoryBandwidth:
+ caps.set_memory_bandwidth(stat.IntOrUintValue());
+ break;
+ case StatType::kDevCapMemorySize:
+ caps.set_memory_size_in_bytes(stat.IntOrUintValue());
+ break;
+ case StatType::kDevCapComputeCapMajor:
+ caps.mutable_compute_capability()->set_major(stat.IntOrUintValue());
+ break;
+ case StatType::kDevCapComputeCapMinor:
+ caps.mutable_compute_capability()->set_minor(stat.IntOrUintValue());
+ break;
+ }
+ });
+
+ return caps;
+}
+
+} // namespace profiler
+} // namespace tensorflow
diff --git a/tensorflow/core/profiler/utils/device_caps_utils.h b/tensorflow/core/profiler/utils/device_caps_utils.h
new file mode 100644
index 0000000..5ab116f
--- /dev/null
+++ b/tensorflow/core/profiler/utils/device_caps_utils.h
@@ -0,0 +1,31 @@
+/* Copyright 2020 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_DEVICE_CAP_UTILS_H_
+#define TENSORFLOW_CORE_PROFILER_UTILS_DEVICE_CAP_UTILS_H_
+
+#include "tensorflow/core/profiler/protobuf/hardware_types.pb.h"
+#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
+
+namespace tensorflow {
+namespace profiler {
+
+void SetDeviceCaps(const DeviceCapabilities& caps, XPlane* plane);
+DeviceCapabilities GetDeviceCaps(const XPlane& plane);
+
+} // namespace profiler
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CORE_PROFILER_UTILS_DEVICE_CAP_UTILS_H_
