| 85 | } |
| 86 | |
| 87 | unique_handle<cublasHandle_t> *cublasManager(const int deviceId) { |
| 88 | thread_local unique_handle<cublasHandle_t> |
| 89 | handles[DeviceManager::MAX_DEVICES]; |
| 90 | thread_local once_flag initFlags[DeviceManager::MAX_DEVICES]; |
| 91 | |
| 92 | call_once(initFlags[deviceId], [&] { |
| 93 | CUBLAS_CHECK((cublasStatus_t)handles[deviceId].create()); |
| 94 | // TODO(pradeep) When multiple streams per device |
| 95 | // is added to CUDA backend, move the cublasSetStream |
| 96 | // call outside of call_once scope. |
| 97 | CUBLAS_CHECK(cublasSetStream(handles[deviceId], getStream(deviceId))); |
| 98 | #ifdef AF_WITH_FAST_MATH |
| 99 | CUBLAS_CHECK( |
| 100 | cublasSetMathMode(handles[deviceId], CUBLAS_TF32_TENSOR_OP_MATH)); |
| 101 | CUBLAS_CHECK( |
| 102 | cublasSetAtomicsMode(handles[deviceId], CUBLAS_ATOMICS_ALLOWED)); |
| 103 | #endif |
| 104 | }); |
| 105 | |
| 106 | return &handles[deviceId]; |
| 107 | } |
| 108 | |
| 109 | #ifdef WITH_CUDNN |
| 110 | unique_handle<cudnnHandle_t> *nnManager(const int deviceId) { |
no test coverage detected