From 08be643e7426d9303a12ffba090d9b0d7ac6e63f Mon Sep 17 00:00:00 2001 From: Russell Power Date: Wed, 14 Oct 2020 11:03:42 -0700 Subject: [PATCH] Remove ThenExecuteFunction references from eager XLA op dispatch. The cleanup behavior is supplied by default by the GPU and XLA backends. PiperOrigin-RevId: 337127562 Change-Id: I539251f3932703af27e17a59bd39c120624a143d --- tensorflow/compiler/jit/kernels/xla_ops.cc | 24 ---------------------- 1 file changed, 24 deletions(-) diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index 12b40b1c83b..0f0f43cbad6 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -274,18 +274,6 @@ void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) { run_options.set_allocator(allocator); run_options.set_intra_op_thread_pool(&ctx->eigen_cpu_device()); run_options.set_rng_seed(GetXLARandomSeed()); - xla::ThenExecuteFunction then_execute; - if (ctx->op_device_context()) { - then_execute = [&](se::Stream* stream, std::function fn) { - Status status = ctx->op_device_context()->ThenExecute( - down_cast(ctx->device()), stream, std::move(fn)); - if (!status.ok()) { - // This should never happen. - LOG(ERROR) << "ThenExecute failed " << status; - } - }; - run_options.set_then_execute_function(&then_execute); - } Env* env = Env::Default(); auto start_time = env->NowMicros(); @@ -522,18 +510,6 @@ void XlaRunOp::Compute(OpKernelContext* ctx) { run_options.set_allocator(allocator); run_options.set_intra_op_thread_pool(&ctx->eigen_cpu_device()); run_options.set_rng_seed(GetXLARandomSeed()); - xla::ThenExecuteFunction then_execute; - if (ctx->op_device_context()) { - then_execute = [&](se::Stream* stream, std::function fn) { - Status status = ctx->op_device_context()->ThenExecute( - down_cast(ctx->device()), stream, std::move(fn)); - if (!status.ok()) { - // This should never happen. - LOG(ERROR) << "ThenExecute failed " << status; - } - }; - run_options.set_then_execute_function(&then_execute); - } Env* env = Env::Default(); auto start_time = env->NowMicros();