| 259 | } |
| 260 | |
| 261 | void CompNodeEnv::init_rocm_async( |
| 262 | int dev, CompNode comp_node, const ContinuationCtx<hipStream_t>& cont) { |
| 263 | m_comp_node = comp_node; |
| 264 | |
| 265 | mgb_assert(!m_user_data_container && !m_async_init_need_wait); |
| 266 | m_rocm_env.device = dev; |
| 267 | m_property.type = DeviceType::ROCM; |
| 268 | MGB_ROCM_CHECK(hipGetDeviceProperties(&m_rocm_env.device_prop, dev)); |
| 269 | { |
| 270 | auto&& prop = m_rocm_env.device_prop; |
| 271 | MGB_MARK_USED_VAR(prop); |
| 272 | //! FIXME: no texure alignment in device property |
| 273 | m_property.mem_alignment = 1u; |
| 274 | } |
| 275 | |
| 276 | std::atomic_bool tid_set{false}; |
| 277 | auto worker = [this, cont, &tid_set]() { |
| 278 | sys::set_thread_name("async_rocm_init"); |
| 279 | m_async_init_tid = std::this_thread::get_id(); |
| 280 | tid_set.store(true); |
| 281 | bool stream_done = false; |
| 282 | MGB_MARK_USED_VAR(stream_done); |
| 283 | MGB_TRY { |
| 284 | m_rocm_env.activate(); |
| 285 | MGB_ROCM_CHECK( |
| 286 | hipStreamCreateWithFlags(&m_rocm_env.stream, hipStreamNonBlocking)); |
| 287 | stream_done = true; |
| 288 | |
| 289 | m_user_data_container = std::make_unique<UserDataContainer>(); |
| 290 | |
| 291 | cont.next(m_rocm_env.stream); |
| 292 | |
| 293 | // megdnn is initialized here; must be placed after cont.next() |
| 294 | // which handles comp node init |
| 295 | mgb_assert( |
| 296 | m_property.mem_alignment == |
| 297 | MegDNNHandle::get(*this).handle()->alignment_requirement()); |
| 298 | auto err = atexit(&CompNode::finalize); |
| 299 | mgb_assert(!err, "failed to register CompNode::finalize at exit"); |
| 300 | } |
| 301 | MGB_CATCH(std::exception & exc, { |
| 302 | mgb_log_error("async rocm init failed: %s", exc.what()); |
| 303 | if (stream_done) { |
| 304 | hipStreamDestroy(m_rocm_env.stream); |
| 305 | } |
| 306 | cont.err(exc); |
| 307 | throw; |
| 308 | }) |
| 309 | }; |
| 310 | |
| 311 | m_async_init_need_wait = true; |
| 312 | m_async_init_future = std::async(std::launch::async, worker); |
| 313 | while (!tid_set.load()) |
| 314 | std::this_thread::yield(); |
| 315 | mgb_assert(m_async_init_tid != std::this_thread::get_id()); |
| 316 | } |
| 317 | #endif |
| 318 | |