| 47 | #endif |
| 48 | |
| 49 | MemAllocImplHelper::MemAddr MemAllocImplHelper::do_alloc( |
| 50 | size_t size, bool allow_from_parent, bool log_stat_on_error) { |
| 51 | mgb_assert(size); |
| 52 | #if !__DEPLOY_ON_XP_SP2__ |
| 53 | m_mutex.lock(); |
| 54 | #endif |
| 55 | |
| 56 | auto iter = m_free_blk_size.lower_bound(FreeBlock{MemAddr{0, 0}, size}); |
| 57 | if (iter == m_free_blk_size.end()) { |
| 58 | #if !__DEPLOY_ON_XP_SP2__ |
| 59 | m_mutex.unlock(); |
| 60 | #endif |
| 61 | if (!allow_from_parent) { |
| 62 | if (log_stat_on_error) { |
| 63 | print_memory_state(); |
| 64 | } |
| 65 | mgb_throw( |
| 66 | MemAllocError, |
| 67 | "out of memory while requesting %zu bytes; you can try " |
| 68 | "setting MGB_CUDA_RESERVE_MEMORY to reserve all memory. " |
| 69 | "If there are dynamic variables, you can also try enabling " |
| 70 | "graph option `enable_grad_var_static_reshape` so " |
| 71 | "some gradient variables can be statically allocated", |
| 72 | size); |
| 73 | } |
| 74 | return alloc_from_parent(size); |
| 75 | } |
| 76 | size_t remain = iter->first.size - size; |
| 77 | auto alloc_addr = iter->first.addr; |
| 78 | m_free_blk_addr.erase(iter->second.aiter); |
| 79 | m_free_blk_size.erase(iter); |
| 80 | |
| 81 | if (remain) |
| 82 | insert_free_unsafe({alloc_addr + size, remain}); |
| 83 | |
| 84 | #if !__DEPLOY_ON_XP_SP2__ |
| 85 | m_mutex.unlock(); |
| 86 | #endif |
| 87 | return alloc_addr; |
| 88 | } |
| 89 | |
| 90 | void MemAllocImplHelper::merge_free_unsafe(FreeBlock block) { |
| 91 | auto iter = m_free_blk_addr.lower_bound(block.addr.addr); |