Method init_cuda_async

src/core/impl/comp_node_env.cpp:166–223

}

void CompNodeEnv::init_cuda_async(
        int dev, CompNode comp_node, const ContinuationCtx<cudaStream_t>& cont) {
    m_comp_node = comp_node;

    mgb_assert(!m_user_data_container && !m_async_init_need_wait);
    m_cuda_env.device = dev;
    m_property.type = DeviceType::CUDA;
    MGB_CUDA_CHECK(cudaGetDeviceProperties(&m_cuda_env.device_prop, dev));
    {
        auto&& prop = m_cuda_env.device_prop;
        m_property.mem_alignment =
                std::max(prop.textureAlignment, prop.texturePitchAlignment);
    }

    std::atomic_bool tid_set{false};
    auto worker = [this, cont, &tid_set]() {
        sys::set_thread_name("async_cuda_init");
        m_async_init_tid = std::this_thread::get_id();
        tid_set.store(true);
        bool stream_done = false;
        MGB_MARK_USED_VAR(stream_done);
        MGB_TRY {
            m_cuda_env.activate();
            MGB_CUDA_CHECK(cudaStreamCreateWithFlags(
                    &m_cuda_env.stream, cudaStreamNonBlocking));
            stream_done = true;

            m_user_data_container = std::make_unique<UserDataContainer>();

#if MGB_ENABLE_DEBUG_UTIL
            nvtxNameCudaStreamA(m_cuda_env.stream, m_comp_node.to_string().c_str());
#endif
            cont.next(m_cuda_env.stream);

            // megdnn is initialized here; must be placed after cont.next()
            // which handles comp node init
            mgb_assert(
                    m_property.mem_alignment ==
                    MegDNNHandle::get(*this).handle()->alignment_requirement());
            auto err = atexit(&CompNode::finalize);
            mgb_assert(!err, "failed to register CompNode::finalize at exit");
        }
        MGB_CATCH(std::exception & exc, {
            mgb_log_error("async cuda init failed: %s", exc.what());
            if (stream_done) {
                cudaStreamDestroy(m_cuda_env.stream);
            }
            cont.err(exc);
            throw;
        })
    };

    m_async_init_need_wait = true;
    m_async_init_future = std::async(std::launch::async, worker);
    while (!tid_set.load())
        std::this_thread::yield();
    mgb_assert(m_async_init_tid != std::this_thread::get_id());
}