MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / load_cuda

Method load_cuda

src/core/impl/comp_node/cuda/comp_node.cpp:952–1007  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

950bool CudaCompNode::is_into_atexit = false;
951#endif
952CompNode::Impl* CudaCompNode::load_cuda(
953 const Locator& locator, const Locator& locator_logical) {
954 int nr_gpu = get_device_count();
955#if MGB_CUDA && defined(WIN32)
956 //! FIXME: windows cuda driver shutdown before call atexit function even
957 //! register atexit function after init cuda driver! as a workround
958 //! recovery resource by OS temporarily, may need remove this after
959 //! upgrade cuda runtime
960 if (!is_into_atexit) {
961 auto err = atexit([] { is_into_atexit = true; });
962 mgb_assert(!err, "failed to register atexit function");
963 }
964#endif
965 mgb_assert(
966 locator.device >= 0 && locator.device < nr_gpu,
967 "request gpu%d out of valid range [0, %d)", locator.device, nr_gpu);
968
969 auto&& sdptr = CudaCompNodeImpl::sd;
970 {
971 MGB_LOCK_GUARD(CudaCompNodeImpl::sd_mtx);
972 if (!sdptr) {
973 // use static storage so object can be safely accessed even after
974 // global finalize
975 using T = CudaCompNodeImpl::StaticData;
976 static std::aligned_storage_t<sizeof(T), alignof(T)> storage;
977 sdptr = new (&storage) T;
978 }
979 }
980 auto&& sd = *sdptr;
981 MGB_LOCK_GUARD(sd.mtx);
982
983 CompNodeImpl* available_node = nullptr;
984 for (int i = 0; i < sd.nr_node; ++i) {
985 auto&& cur = sd.node[i];
986 if (cur.m_initialized) {
987 if (cur.m_locator == locator && cur.m_locator_logical == locator_logical) {
988 return &cur;
989 }
990 } else {
991 available_node = &cur;
992 }
993 }
994
995 if (!available_node) {
996 mgb_assert(
997 sd.nr_node < CompNodeImpl::MAX_NR_COMP_NODE,
998 "too many CompNode allocated");
999 available_node = &sd.node[sd.nr_node++];
1000 }
1001 mgb_assert(locator.device < CompNodeImpl::MAX_NR_DEVICE, "device number too large");
1002
1003 mgb_assert(!available_node->m_initialized);
1004 available_node->init(locator, locator_logical);
1005
1006 return available_node;
1007}
1008
1009void CudaCompNode::try_coalesce_all_free_memory() {

Callers

nothing calls this directly

Calls 2

get_device_countFunction · 0.85
initMethod · 0.45

Tested by

no test coverage detected