| 964 | } |
| 965 | |
| 966 | CpuCompNode::Impl* CpuCompNode::load_cpu(Locator locator, Locator locator_logical) { |
| 967 | #if !MGB_HAVE_THREAD |
| 968 | // use only cpu:default and cpu0:1023 comp node when threading is disabled |
| 969 | mgb_assert( |
| 970 | locator.device == Locator::DEVICE_CPU_DEFAULT || |
| 971 | (locator.device == 0 && locator.stream == 1023)); |
| 972 | locator_logical = {locator_logical.type, locator.device, locator.stream}; |
| 973 | #endif |
| 974 | { |
| 975 | MGB_LOCK_GUARD(sm_pool_mtx); |
| 976 | if (!sm_pool) { |
| 977 | // use static storage so object can be safely accessed even after |
| 978 | // global finalize |
| 979 | static std::aligned_storage_t<sizeof(Pool), alignof(Pool)> storage; |
| 980 | sm_pool = new (&storage) Pool; |
| 981 | } |
| 982 | } |
| 983 | mgb_assert( |
| 984 | locator.device >= 0 || |
| 985 | (locator.device == Locator::DEVICE_CPU_DEFAULT && |
| 986 | locator.stream == 0) || |
| 987 | locator.device == Locator::DEVICE_MULTITHREAD_DEFAULT, |
| 988 | "failed to load cpu for device:%d stream:%d", locator.device, |
| 989 | locator.stream); |
| 990 | MGB_LOCK_GUARD(sm_pool->mtx); |
| 991 | |
| 992 | // encode both device ID and type into a int |
| 993 | mgb_assert( |
| 994 | locator_logical.device >= -1 || |
| 995 | locator_logical.device <= Locator::DEVICE_CPU_DEFAULT); |
| 996 | if (locator_logical.type != CompNode::DeviceType::UNSPEC) { |
| 997 | mgb_assert( |
| 998 | locator_logical.type == CompNode::DeviceType::CPU || |
| 999 | locator_logical.type == CompNode::DeviceType::MULTITHREAD); |
| 1000 | } |
| 1001 | if (locator.type == DeviceType::CPU) { |
| 1002 | auto&& pqueue_weak = sm_pool->physical2queue[{locator.device, locator.stream}]; |
| 1003 | auto pqueue = pqueue_weak.lock(); |
| 1004 | if (!pqueue) { |
| 1005 | pqueue = std::make_shared<WorkerQueue>(locator); |
| 1006 | pqueue_weak = pqueue; |
| 1007 | } |
| 1008 | auto&& pimpl = sm_pool->locator2impl[{locator, locator_logical}]; |
| 1009 | if (!pimpl) { |
| 1010 | mgb_assert( |
| 1011 | sm_pool->nr_used_impl_storage < Pool::MAX_NR_COMP_NODE, |
| 1012 | "too many cpu comp nodes; max %d allowed", Pool::MAX_NR_COMP_NODE); |
| 1013 | pimpl.reset(new (&sm_pool->impl_storage[sm_pool->nr_used_impl_storage++]) |
| 1014 | CompNodeRecorderImpl{locator, locator_logical, pqueue}); |
| 1015 | } |
| 1016 | log_comp_node_created(locator, locator_logical); |
| 1017 | return pimpl.get(); |
| 1018 | } else { |
| 1019 | mgb_assert(locator.type == DeviceType::MULTITHREAD); |
| 1020 | auto&& pqueue_weak = sm_pool->physical2queue_multithead[{ |
| 1021 | locator.device, locator.nr_threads}]; |
| 1022 | auto pqueue = pqueue_weak.lock(); |
| 1023 | if (!pqueue) { |