| 66 | } |
| 67 | |
| 68 | void CudaGPU::Setup() { |
| 69 | lang_ = kCuda; |
| 70 | ctx_.stream = NULL; // use the default sync stream |
| 71 | |
| 72 | // TODO(wangwei) create one handle for each steam? |
| 73 | // Preserse for future use instead of default sync stream, for concurrency |
| 74 | // cudaStreamCreate(&ctx_.stream); |
| 75 | |
| 76 | #ifdef USE_DIST |
| 77 | CUDA_CHECK(cudaStreamCreateWithFlags(&ctx_.s, cudaStreamNonBlocking)); |
| 78 | CUDA_CHECK(cudaStreamCreateWithFlags(&ctx_.c1, cudaStreamNonBlocking)); |
| 79 | CUDA_CHECK(cudaStreamCreateWithFlags(&ctx_.c2, cudaStreamNonBlocking)); |
| 80 | #endif // USE_DIST |
| 81 | |
| 82 | CUDA_CHECK(cudaSetDevice(id_)); |
| 83 | // use curandCreateGeneratorHost for CudaHost device |
| 84 | CURAND_CHECK( |
| 85 | curandCreateGenerator(&ctx_.curand_generator, CURAND_RNG_PSEUDO_DEFAULT)); |
| 86 | CURAND_CHECK(curandSetStream(ctx_.curand_generator, ctx_.stream)); |
| 87 | auto seed = std::chrono::system_clock::now().time_since_epoch().count(); |
| 88 | SetRandSeed(seed); |
| 89 | // TODO(wangwei) if one generator per stream, then need diff offset per gen? |
| 90 | CURAND_CHECK(curandSetGeneratorOffset(ctx_.curand_generator, 0)); |
| 91 | CUBLAS_CHECK(cublasCreate(&(ctx_.cublas_handle))); |
| 92 | CUBLAS_CHECK(cublasSetStream(ctx_.cublas_handle, ctx_.stream)); |
| 93 | |
| 94 | #ifdef USE_CUDNN |
| 95 | // TODO(wangwei) create one handle for each stream? |
| 96 | auto status = cudnnCreate(&ctx_.cudnn_handle); |
| 97 | CHECK_EQ(status, CUDNN_STATUS_SUCCESS) << cudnnGetErrorString(status); |
| 98 | cudnnSetStream(ctx_.cudnn_handle, ctx_.stream); |
| 99 | #endif // USE_CUDNN |
| 100 | } |
| 101 | |
| 102 | void CudaGPU::SetRandSeed(unsigned seed) { |
| 103 | CHECK(ctx_.curand_generator); |
nothing calls this directly
no outgoing calls
no test coverage detected