| 154 | |
| 155 | #ifndef CPU_ONLY |
| 156 | void SyncedMemory::async_gpu_push(const cudaStream_t& stream) { |
| 157 | check_device(); |
| 158 | CHECK(head_ == HEAD_AT_CPU); |
| 159 | if (gpu_ptr_ == NULL) { |
| 160 | CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_)); |
| 161 | own_gpu_data_ = true; |
| 162 | } |
| 163 | const cudaMemcpyKind put = cudaMemcpyHostToDevice; |
| 164 | CUDA_CHECK(cudaMemcpyAsync(gpu_ptr_, cpu_ptr_, size_, put, stream)); |
| 165 | // Assume caller will synchronize on the stream before use |
| 166 | head_ = SYNCED; |
| 167 | } |
| 168 | #endif |
| 169 | |
| 170 | void SyncedMemory::check_device() { |