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

Method load_rocm

src/core/impl/comp_node/rocm/comp_node.cpp:572–615  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

570}
571
572CompNode::Impl* ROCmCompNode::load_rocm(
573 const Locator& locator, const Locator& locator_logical) {
574 int nr_gpu = get_device_count();
575 mgb_assert(
576 locator.device >= 0 && locator.device < nr_gpu,
577 "request gpu%d out of valid range [0, %d)", locator.device, nr_gpu);
578
579 auto&& sdptr = ROCmCompNodeImpl::sd;
580 {
581 MGB_LOCK_GUARD(ROCmCompNodeImpl::sd_mtx);
582 if (!sdptr) {
583 // use static storage so object can be safely accessed even after
584 // global finalize
585 using T = ROCmCompNodeImpl::StaticData;
586 static std::aligned_storage_t<sizeof(T), alignof(T)> storage;
587 sdptr = new (&storage) T;
588 }
589 }
590 auto&& sd = *sdptr;
591 MGB_LOCK_GUARD(sd.mtx);
592
593 CompNodeImpl* available_node = nullptr;
594 for (int i = 0; i < sd.nr_node; ++i) {
595 auto&& cur = sd.node[i];
596 if (cur.m_initialized) {
597 if (cur.m_locator == locator && cur.m_locator_logical == locator_logical) {
598 return &cur;
599 }
600 } else {
601 available_node = &cur;
602 }
603 }
604
605 if (!available_node) {
606 mgb_assert(sd.nr_node < sd.MAX_NR_COMP_NODE, "too many CompNode allocated");
607 mgb_assert(locator.device < sd.MAX_NR_COMP_NODE, "device number too large");
608 available_node = &sd.node[sd.nr_node++];
609 }
610
611 mgb_assert(!available_node->m_initialized);
612 available_node->init(locator, locator_logical);
613
614 return available_node;
615}
616
617void ROCmCompNode::try_coalesce_all_free_memory() {
618 // TODO: optimized implementation

Callers

nothing calls this directly

Calls 2

get_device_countFunction · 0.85
initMethod · 0.45

Tested by

no test coverage detected