commit 801b09624f0488132638166fe782be4163269657
author A. Unique TensorFlower <gardener@tensorflow.org> Fri Dec 20 18:58:10 2019 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> Fri Dec 20 19:01:36 2019 -0800
tree 2735529e7b0e4ecc5a7077e6db36e1fdb832ebb3
parent a25c899dcc938c36b2ca8b77393001cd59fd9b97

More accurate input-pipeline analysis for TPU.

PiperOrigin-RevId: 286672368
Change-Id: I14e6e47e79304154404629295b6728857583d781

tensorflow/core/profiler/utils/event_span.cc

diff --git a/tensorflow/core/profiler/utils/event_span.cc b/tensorflow/core/profiler/utils/event_span.cc
index 8c31c55..e6e8fd2 100644
--- a/tensorflow/core/profiler/utils/event_span.cc
+++ b/tensorflow/core/profiler/utils/event_span.cc
@@ -116,17 +116,17 @@
 }
 
 EventType ClassifyCpuEvent(absl::string_view event_name, int64 correlation_id) {
-  if (absl::StartsWithIgnoreCase(event_name, "MEMCPYHtoD"))
+  if (absl::StartsWithIgnoreCase(event_name, "MEMCPYHtoD") ||
+      absl::StrContains(event_name, "Infeed"))
     return HOST_TO_DEVICE;
   if (absl::StartsWithIgnoreCase(event_name, "MEMCPYHtoH")) return HOST_TO_HOST;
   if (correlation_id >= 0 ||
       absl::StartsWithIgnoreCase(event_name, "ExecutorState::Process")) {
     return HOST_PREPARE;
-  } else {
-    if (absl::StartsWithIgnoreCase(event_name, "IteratorGetNext"))
-      return HOST_WAIT_INPUT;
-    return HOST_COMPUTE;
   }
+  if (absl::StartsWithIgnoreCase(event_name, "IteratorGetNext"))
+    return HOST_WAIT_INPUT;
+  return HOST_COMPUTE;
 }
 
 std::string PrintEventType(EventType event_type) {