From 8d7f18b250a6356623509dee7a4d0636b8937784 Mon Sep 17 00:00:00 2001
From: Jiho Choi <jihochoi@google.com>
Date: Thu, 21 May 2020 14:28:20 -0700
Subject: [PATCH] Add memory cost breakdown per operation type and memory space
to OpMetrics.
PiperOrigin-RevId: 312741555
Change-Id: Id6666a8c6b9d67fe443154fb135cadeeeaecebdc
---
.../core/profiler/protobuf/op_metrics.proto | 15 ++++++++++++++-
1 file changed, 14 insertions(+), 1 deletion(-)
diff --git a/tensorflow/core/profiler/protobuf/op_metrics.proto b/tensorflow/core/profiler/protobuf/op_metrics.proto
index c0f34773e02..af38795b7b2 100644
--- a/tensorflow/core/profiler/protobuf/op_metrics.proto
+++ b/tensorflow/core/profiler/protobuf/op_metrics.proto
@@ -26,7 +26,7 @@ message LayoutAnalysis {
}
// Metrics for an operation (accumulated over all occurrences).
-// Next ID: 19
+// Next ID: 20
message OpMetrics {
// HLO module id. 0 for TF ops.
uint64 hlo_module_id = 13;
@@ -50,6 +50,19 @@ message OpMetrics {
uint64 flops = 2;
// Total bytes accessed.
uint64 bytes_accessed = 5;
+ // Breakdown of memory accessed by operation type and memory space.
+ message MemoryAccessed {
+ enum OperationType {
+ UNKNOWN = 0;
+ READ = 1;
+ WRITE = 2;
+ }
+ OperationType operation_type = 1;
+ // Device-specific id of memory space.
+ uint64 memory_space = 2;
+ uint64 bytes_accessed = 3;
+ }
+ repeated MemoryAccessed memory_accessed_breakdown = 19;
// Total dma stall time in picoseconds.
uint64 dma_stall_ps = 10;
// The data layout for this op. Only set for convolution ops for now.