diff options
Diffstat (limited to 'tensorflow/core/util/cuda_kernel_helper.h')
-rw-r--r-- | tensorflow/core/util/cuda_kernel_helper.h | 9 |
1 files changed, 9 insertions, 0 deletions
diff --git a/tensorflow/core/util/cuda_kernel_helper.h b/tensorflow/core/util/cuda_kernel_helper.h index 22a6e5dcef..86f974ee26 100644 --- a/tensorflow/core/util/cuda_kernel_helper.h +++ b/tensorflow/core/util/cuda_kernel_helper.h @@ -60,6 +60,15 @@ inline CudaLaunchConfig GetCudaLaunchConfig(int work_element_count, return config; } +template <typename T> +__device__ __host__ inline T ldg(const T* address) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 350 + return __ldg(address); +#else + return *address; +#endif +} + } // namespace tensorflow #endif // GOOGLE_CUDA |