Add space for reporting memory bandwidth statistics

PiperOrigin-RevId: 356421433
Change-Id: I52f9d8db3a2fc8209a247f300e79287fd9fa2555
This commit is contained in:
Frank Chen 2021-02-08 21:23:54 -08:00 committed by TensorFlower Gardener
parent 6fd9628434
commit 81e97f2832
5 changed files with 18 additions and 5 deletions

View File

@ -463,6 +463,7 @@ tf_kernel_library(
"//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",
],
)

View File

@ -47,10 +47,12 @@ limitations under the License.
#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 @@ void RecordElementSize(const std::vector<Tensor> element,
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_megabytes_per_sec", mem_bw}});
}
return profiler::TraceMeEncode(
"IteratorGetNextOp::DoCompute",
{{"id", ctx->step_id()}, {"iter_num", ctx->frame_iter().iter_id}});
},
profiler::kInfo);
tensorflow::ResourceTagger tag(kTFDataResourceTag,

View File

@ -357,7 +357,7 @@ double NominalCPUFrequency() {
}
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);

View File

@ -62,6 +62,7 @@ std::size_t MallocExtension_GetAllocatedSize(const void* p);
struct MemoryInfo {
int64 total = 0;
int64 free = 0;
int64 bw_used = 0; // memory bandwidth used across all CPU (in MBs/second)
};
// Retrieves the host memory information. If any of the fields in the returned

View File

@ -192,7 +192,7 @@ double NominalCPUFrequency() {
}
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)) {