Add space for reporting memory bandwidth statistics
PiperOrigin-RevId: 355948962
Change-Id: I606e722f5800774783e7abc46be7946a801ac183
diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index e5adb2d..448a5de 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -463,6 +463,7 @@
"//tensorflow/core:session_options",
"//tensorflow/core/kernels:ops_util",
"//tensorflow/core/profiler/lib:traceme",
+ "//tensorflow/core/profiler/lib:traceme_encode",
"@com_google_absl//absl/memory",
],
)
diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc
index f3d2382..dce6ebd 100644
--- a/tensorflow/core/kernels/data/iterator_ops.cc
+++ b/tensorflow/core/kernels/data/iterator_ops.cc
@@ -47,10 +47,12 @@
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/platform/casts.h"
#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mem.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/core/platform/resource.h"
#include "tensorflow/core/profiler/lib/traceme.h"
+#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
@@ -965,9 +967,18 @@
Status IteratorGetNextOp::DoCompute(OpKernelContext* ctx) {
profiler::TraceMe traceme(
[&] {
- return strings::StrCat(
- "IteratorGetNextOp::DoCompute#id=", ctx->step_id(),
- ",iter_num=", ctx->frame_iter().iter_id, "#");
+ int64 mem_bw = port::GetMemoryInfo().bw_used;
+
+ if (mem_bw != INT64_MAX) {
+ return profiler::TraceMeEncode(
+ "IteratorGetNextOp::DoCompute",
+ {{"id", ctx->step_id()},
+ {"iter_num", ctx->frame_iter().iter_id},
+ {"mem_bw_used", mem_bw}});
+ }
+ return profiler::TraceMeEncode(
+ "IteratorGetNextOp::DoCompute",
+ {{"id", ctx->step_id()}, {"iter_num", ctx->frame_iter().iter_id}});
},
profiler::kInfo);
tensorflow::ResourceTagger tag(kTFDataResourceTag,
diff --git a/tensorflow/core/platform/default/port.cc b/tensorflow/core/platform/default/port.cc
index e25ed07..6e82c67 100644
--- a/tensorflow/core/platform/default/port.cc
+++ b/tensorflow/core/platform/default/port.cc
@@ -357,7 +357,7 @@
}
MemoryInfo GetMemoryInfo() {
- MemoryInfo mem_info = {INT64_MAX, INT64_MAX};
+ MemoryInfo mem_info = {INT64_MAX, INT64_MAX, INT64_MAX};
#if defined(__linux__) && !defined(__ANDROID__)
struct sysinfo info;
int err = sysinfo(&info);
diff --git a/tensorflow/core/platform/mem.h b/tensorflow/core/platform/mem.h
index 27ad357..65411ee 100644
--- a/tensorflow/core/platform/mem.h
+++ b/tensorflow/core/platform/mem.h
@@ -62,6 +62,7 @@
struct MemoryInfo {
int64 total = 0;
int64 free = 0;
+ int64 bw_used = 0; // memory bandwidth used across all CPU
};
// Retrieves the host memory information. If any of the fields in the returned
diff --git a/tensorflow/core/platform/windows/port.cc b/tensorflow/core/platform/windows/port.cc
index 16b5a32..256f525 100644
--- a/tensorflow/core/platform/windows/port.cc
+++ b/tensorflow/core/platform/windows/port.cc
@@ -192,7 +192,7 @@
}
MemoryInfo GetMemoryInfo() {
- MemoryInfo mem_info = {INT64_MAX, INT64_MAX};
+ MemoryInfo mem_info = {INT64_MAX, INT64_MAX, INT64_MAX};
MEMORYSTATUSEX statex;
statex.dwLength = sizeof(statex);
if (GlobalMemoryStatusEx(&statex)) {
