From c4bdfed6e8008606d9217759a1feb1e4c6bf390c Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Mon, 5 Aug 2019 11:52:54 -0700 Subject: [PATCH] [TF:XLA] Use RunAsync() in TF/XLA launch ops where applicable. PiperOrigin-RevId: 261734859 --- tensorflow/compiler/jit/kernels/xla_ops.cc | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index 788e90ffe99..7cceeb64136 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -366,7 +366,12 @@ void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) { Env* env = Env::Default(); auto start_time = env->NowMicros(); - auto run_result = executable->Run(launch_context.arguments(), run_options); + xla::StatusOr run_result; + if (!stream || platform_info_.platform_id() == se::host::kHostPlatformId) { + run_result = executable->Run(launch_context.arguments(), run_options); + } else { + run_result = executable->RunAsync(launch_context.arguments(), run_options); + } OP_REQUIRES(ctx, run_result.ok(), run_result.status()); auto elapsed = env->NowMicros() - start_time; @@ -550,8 +555,14 @@ void XlaRunOp::Compute(OpKernelContext* ctx) { Env* env = Env::Default(); auto start_time = env->NowMicros(); - auto run_result = - closure.executable()->Run(launch_context.arguments(), run_options); + xla::StatusOr run_result; + if (!stream || platform_info_.platform_id() == se::host::kHostPlatformId) { + run_result = + closure.executable()->Run(launch_context.arguments(), run_options); + } else { + run_result = + closure.executable()->RunAsync(launch_context.arguments(), run_options); + } OP_REQUIRES(ctx, run_result.ok(), run_result.status()); auto elapsed = env->NowMicros() - start_time;