From 3785cffbdbcf91901996f61a2fddd5107b16da06 Mon Sep 17 00:00:00 2001
From: Russell Power <power@google.com>
Date: Sat, 22 Dec 2018 18:12:58 -0800
Subject: [PATCH] Disable cost tracking for async kernels.

Async kernels can be modified from both the scheduling thread or a thread-pool,
making it difficult to safely update timing information from the scheduling
thread.

PiperOrigin-RevId: 226640910
---
 tensorflow/core/common_runtime/executor.cc | 33 ++++++----------------
 1 file changed, 9 insertions(+), 24 deletions(-)

diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
index 30ee3a367df..df2ee6c6182 100644
--- a/tensorflow/core/common_runtime/executor.cc
+++ b/tensorflow/core/common_runtime/executor.cc
@@ -135,13 +135,9 @@ struct EdgeInfo {
 // Time the execution of kernels (in CPU cycles).  Used to dynamically identify
 // inexpensive kernels which can be dispatched inline.
 struct KernelTimer {
-  uint64 start_cycles = 0;
+  uint64 start_cycles = profile_utils::CpuUtils::GetCurrentClockCycle();
 
-  void Start() {
-    start_cycles = profile_utils::CpuUtils::GetCurrentClockCycle();
-  }
-
-  uint64 Stop() {
+  uint64 ElapsedCycles() {
     return profile_utils::CpuUtils::GetCurrentClockCycle() - start_cycles;
   }
 };
@@ -1768,24 +1764,11 @@ void ExecutorState::Process(TaggedNode tagged_node, int64 scheduled_nsec) {
           if (completed) Finish();
         };
         nodestats::SetOpStart(stats);
-        KernelTimer timer;
-        if (item.kernel->IsExpensive()) {
-          timer.Start();
-        }
         device->ComputeAsync(async, &state->ctx, done);
-
-        // Update kernel cost based on time spent in the current thread.
-        if (item.kernel->IsExpensive()) {
-          async->UpdateCostEstimate(timer.Stop());
-        }
       } else {
         // Synchronous computes.
         OpKernelContext ctx(&params, item.num_outputs);
         nodestats::SetOpStart(stats);
-        KernelTimer timer;
-        if (item.kernel->IsExpensive()) {
-          timer.Start();
-        }
 
         if (TF_PREDICT_FALSE(
                 MightTrace(item, event_collector_, trace_using_annotations_))) {
@@ -1812,11 +1795,13 @@ void ExecutorState::Process(TaggedNode tagged_node, int64 scheduled_nsec) {
           }
         } else {
           // In the common case, avoid creating any tracing objects.
-          device->Compute(op_kernel, &ctx);
-        }
-
-        if (item.kernel->IsExpensive()) {
-          op_kernel->UpdateCostEstimate(timer.Stop());
+          if (op_kernel->IsExpensive()) {
+            KernelTimer timer;
+            device->Compute(op_kernel, &ctx);
+            op_kernel->UpdateCostEstimate(timer.ElapsedCycles());
+          } else {
+            device->Compute(op_kernel, &ctx);
+          }
         }
 
         nodestats::SetOpEnd(stats);