| 292 | |
| 293 | template <typename DeviceFunc> |
| 294 | Gpu3DLaunchConfig GetGpu3DLaunchConfig(int xdim, int ydim, int zdim, |
| 295 | const Eigen::GpuDevice& d, |
| 296 | DeviceFunc func, |
| 297 | size_t dynamic_shared_memory_size, |
| 298 | int block_size_limit) { |
| 299 | Gpu3DLaunchConfig config; |
| 300 | |
| 301 | if (xdim <= 0 || ydim <= 0 || zdim <= 0) { |
| 302 | return config; |
| 303 | } |
| 304 | |
| 305 | int dev; |
| 306 | int xthreadlimit = 0, ythreadlimit = 0, zthreadlimit = 0; |
| 307 | int xgridlimit = 0, ygridlimit = 0, zgridlimit = 0; |
| 308 | #if GOOGLE_CUDA |
| 309 | cudaGetDevice(&dev); |
| 310 | cudaDeviceGetAttribute(&xthreadlimit, cudaDevAttrMaxBlockDimX, dev); |
| 311 | cudaDeviceGetAttribute(&ythreadlimit, cudaDevAttrMaxBlockDimY, dev); |
| 312 | cudaDeviceGetAttribute(&zthreadlimit, cudaDevAttrMaxBlockDimZ, dev); |
| 313 | cudaDeviceGetAttribute(&xgridlimit, cudaDevAttrMaxGridDimX, dev); |
| 314 | cudaDeviceGetAttribute(&ygridlimit, cudaDevAttrMaxGridDimY, dev); |
| 315 | cudaDeviceGetAttribute(&zgridlimit, cudaDevAttrMaxGridDimZ, dev); |
| 316 | #elif TENSORFLOW_USE_ROCM |
| 317 | hipGetDevice(&dev); |
| 318 | hipDeviceGetAttribute(&xthreadlimit, hipDeviceAttributeMaxBlockDimX, dev); |
| 319 | hipDeviceGetAttribute(&ythreadlimit, hipDeviceAttributeMaxBlockDimY, dev); |
| 320 | hipDeviceGetAttribute(&zthreadlimit, hipDeviceAttributeMaxBlockDimZ, dev); |
| 321 | hipDeviceGetAttribute(&xgridlimit, hipDeviceAttributeMaxGridDimX, dev); |
| 322 | hipDeviceGetAttribute(&ygridlimit, hipDeviceAttributeMaxGridDimY, dev); |
| 323 | hipDeviceGetAttribute(&zgridlimit, hipDeviceAttributeMaxGridDimZ, dev); |
| 324 | #endif |
| 325 | |
| 326 | int block_count = 0; |
| 327 | int thread_per_block = 0; |
| 328 | |
| 329 | #if GOOGLE_CUDA |
| 330 | cudaError_t err = cudaOccupancyMaxPotentialBlockSize( |
| 331 | &block_count, &thread_per_block, func, dynamic_shared_memory_size, |
| 332 | block_size_limit); |
| 333 | CHECK_EQ(err, cudaSuccess); |
| 334 | #elif TENSORFLOW_USE_ROCM |
| 335 | // ROCM TODO re-enable this after hipOccupancyMaxPotentialBlockSize is |
| 336 | // implemented |
| 337 | // hipError_t err = hipOccupancyMaxPotentialBlockSize( |
| 338 | // &block_count, &thread_per_block, func, dynamic_shared_memory_size, |
| 339 | // block_size_limit); |
| 340 | // CHECK_EQ(err, hipSuccess); |
| 341 | |
| 342 | const int physical_thread_count = |
| 343 | d.getNumGpuMultiProcessors() * d.maxGpuThreadsPerMultiProcessor(); |
| 344 | thread_per_block = std::min(1024, d.maxGpuThreadsPerBlock()); |
| 345 | block_count = std::min(DivUp(physical_thread_count, thread_per_block), |
| 346 | d.getNumGpuMultiProcessors()); |
| 347 | #endif |
| 348 | |
| 349 | int threadsx = std::min({xdim, thread_per_block, xthreadlimit}); |
| 350 | int threadsy = |
| 351 | std::min({ydim, std::max(thread_per_block / threadsx, 1), ythreadlimit}); |
no test coverage detected