Commit 3850ab8a authored by Haoyu Zhang's avatar Haoyu Zhang Committed by TensorFlower Gardener
Browse files

Use device private thread pool to launch ops when `gpu_private` thread mode is...

Use device private thread pool to launch ops when `gpu_private` thread mode is turned on in eager runtime.

When setting the environment variables `TF_GPU_THREAD_MODE=gpu_private TF_GPU_THREAD_COUNT=2`, TensorFlow allocates separate thread pools for GPU devices, to reduce contention on GPU kernel launching. Though it's not TF official API, it's been important for TF to achieve good performance.

This CL exposes the FunctionLibraryRuntime's internal runner (which is initialized using the correct device private thread pool if available) to kernel and device ops. When using the instance to run, it will launch ops using the specified runner instead of the default one in EagerContext.

PiperOrigin-RevId: 237132982
parent 11599b4a
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment