| 208 | } |
| 209 | |
| 210 | std::shared_ptr<ReduceBucket> getBuffer(at::ScalarType scalar_type) |
| 211 | { |
| 212 | if (!hasKey(current_buffer_, scalar_type)) { |
| 213 | current_buffer_[scalar_type] = |
| 214 | std::make_shared<ReduceBucket>(initial_bucket_size_, scalar_type); |
| 215 | current_buffer_events_[scalar_type] = |
| 216 | std::make_shared<at::cuda::CUDAEvent>(cudaEventDisableTiming); |
| 217 | |
| 218 | if (enable_double_buffer_) { |
| 219 | shadow_buffer_[scalar_type] = |
| 220 | std::make_shared<ReduceBucket>(initial_bucket_size_, scalar_type); |
| 221 | shadow_buffer_events_[scalar_type] = |
| 222 | std::make_shared<at::cuda::CUDAEvent>(cudaEventDisableTiming); |
| 223 | } |
| 224 | } |
| 225 | |
| 226 | return current_buffer_.at(scalar_type); |
| 227 | } |
| 228 | |
| 229 | std::shared_ptr<at::cuda::CUDAEvent> getEvent(at::ScalarType scalar_type) |
| 230 | { |