Make GemmThunkShouldHaltAllActivityBeforeRunning easier to understand; NFC

PiperOrigin-RevId: 208110715

3 files changed, 13 insertions, 20 deletions

diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h
index e3e531c2b0..12c81f9bfc 100644
--- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h
+++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h
@@ -52,12 +52,12 @@ class GemmThunk : public Thunk {
                          se::Stream* stream,
                          HloExecutionProfiler* profiler) override;
 
-  // Returns true if we'll perform autotuning if run on the given stream.  If
-  // so, we want the GPU to be quiescent during autotuning, so as not to
-  // introduce noise in our results.
-  bool ShouldHaltAllActivityBeforeRunning(se::Stream* stream) override {
-    return autotune_results_.count(
-               stream->parent()->GetDeviceDescription().name()) == 0;
+  bool WillAutotuneKernel(se::Stream* stream) override {
+    // We will autotune this kernel if we don't already have a autotune result
+    // for the stream device.
+    return autotune_results_.find(
+               stream->parent()->GetDeviceDescription().name()) ==
+           autotune_results_.end();
   }
 
  private:
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
index bb7736efa6..7060837904 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
+++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
@@ -131,9 +131,10 @@ Status GpuExecutable::ExecuteThunks(
       stream->ThenWaitFor(FindOrDie(thunk_to_finish_event, dependency).get());
     }
 
-    // If this thunk requests it, wait for all currently-executing thunks to
-    // finish.  This is useful e.g. if the thunk is about to perform autotuning.
-    if (thunk->ShouldHaltAllActivityBeforeRunning(stream)) {
+    // If this thunk is about to autotune then wait for all currently executing
+    // thunks to finish.  This reduces noise and thus the probability of
+    // choosing a suboptimal algorithm.
+    if (thunk->WillAutotuneKernel(stream)) {
       TF_RETURN_IF_ERROR(main_stream->BlockHostUntilDone());
     }
 
diff --git a/tensorflow/compiler/xla/service/gpu/thunk.h b/tensorflow/compiler/xla/service/gpu/thunk.h
index 4df0bb005b..e68bee035a 100644
--- a/tensorflow/compiler/xla/service/gpu/thunk.h
+++ b/tensorflow/compiler/xla/service/gpu/thunk.h
@@ -82,17 +82,9 @@ class Thunk {
     return Status::OK();
   }
 
-  // Users of Thunk should call ShouldHaltAllActivityBeforeRunning(stream)
-  // before calling ExecuteOnStream(stream).  If it returns true, it's the
-  // user's responsibility to wait for all activity on the GPU to finish before
-  // calling ExecuteOnStream.
-  //
-  // This value is not required to be constant for a given Thunk.  For example,
-  // a Thunk that performs autotuning may return true for its first run and
-  // false thereafter.
-  virtual bool ShouldHaltAllActivityBeforeRunning(se::Stream* /*stream*/) {
-    return false;
-  }
+  // Returns true if this kernel will autotune for the stream device the next
+  // time it is run.
+  virtual bool WillAutotuneKernel(se::Stream* /*stream*/) { return false; }
 
   // Execute the kernel for the thunk on the given stream. This method must be
   // called after Initialize and can be called multiple times over Thunk's