| @@ -0,0 +1,48 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 3 | + | ||
| 4 | + | ||
| 5 | +class TestMemPoolDestructor(TestCase): | ||
| 6 | + def test_empty_cache_mempool_destruction(self): | ||
| 7 | + # make sure clear other task mem | ||
| 8 | + torch.npu.synchronize() | ||
| 9 | + torch.npu.empty_cache() | ||
| 10 | + | ||
| 11 | + pool_memory_reserved = 0 | ||
| 12 | + pool = torch.npu.MemPool() | ||
| 13 | + with torch.npu.use_mem_pool(pool): | ||
| 14 | + tensor = torch.full((1024 ** 3,), 53, dtype=torch.uint8, device='npu') | ||
| 15 | + pool_memory_reserved = torch.npu.memory_reserved() | ||
| 16 | + del tensor | ||
| 17 | + | ||
| 18 | + torch.npu.synchronize() | ||
| 19 | + torch.npu.empty_cache() | ||
| 20 | + pool_before_del_memory_reserved = torch.npu.memory_reserved() | ||
| 21 | + self.assertGreater(pool_memory_reserved, pool_before_del_memory_reserved / 2) | ||
| 22 | + self.assertGreater(pool_before_del_memory_reserved, pool_memory_reserved / 2) | ||
| 23 | + | ||
| 24 | + # should free mempool cache | ||
| 25 | + del pool | ||
| 26 | + pool_after_del_memory_reserved = torch.npu.memory_reserved() | ||
| 27 | + self.assertGreater(pool_memory_reserved / 2, pool_after_del_memory_reserved) | ||
| 28 | + | ||
| 29 | + # normally empty cache after mempool processing | ||
| 30 | + tensor_2 = torch.full((1024 ** 3,), 53, dtype=torch.uint8, device='npu') | ||
| 31 | + del tensor_2 | ||
| 32 | + torch.npu.synchronize() | ||
| 33 | + torch.npu.empty_cache() | ||
| 34 | + end_memory_reserved = torch.npu.memory_reserved() | ||
| 35 | + self.assertGreater(pool_memory_reserved / 2, end_memory_reserved) | ||
| 36 | + | ||
| 37 | + def test_empty_cache_mempool_destruction_abnormal(self): | ||
| 38 | + test_result = True | ||
| 39 | + try: | ||
| 40 | + pool = torch.npu.MemPool() | ||
| 41 | + del pool | ||
| 42 | + except Exception as e: | ||
| 43 | + test_result = False | ||
| 44 | + self.assertTrue(test_result) | ||
| 45 | + | ||
| 46 | + | ||
| 47 | +if __name__ == '__main__': | ||
| 48 | + run_tests() | ||
| @@ -1657,9 +1657,10 @@ public: | |||
| 1657 | } | 1657 | } |
| 1658 | 1658 | ||
| 1659 | /* * returns cached blocks to the system allocator * */ | 1659 | /* * returns cached blocks to the system allocator * */ |
| 1660 | - void emptyCache(int device, bool check_error, bool free_physical) | 1660 | + void emptyCache(int device, bool check_error, bool free_physical, MempoolId_t mempool_id = {0, 0}) |
| 1661 | { | 1661 | { |
| 1662 | - TORCH_NPU_MEMORY_LOGI("emptyCache: device=%d, check_error=%d, free_physical=%d", device, check_error, free_physical); | 1662 | + TORCH_NPU_MEMORY_LOGI("emptyCache: device=%d, check_error=%d, free_physical=%d, mempool_id={%lu, %lu}", |
| 1663 | + device, check_error, free_physical, mempool_id.first, mempool_id.second); | ||
| 1663 | // when exec emptyCache in torch_npu.npu.check_uce_in_memory(), check_error is false | 1664 | // when exec emptyCache in torch_npu.npu.check_uce_in_memory(), check_error is false |
| 1664 | bool prev_need_check_error = need_check_error; | 1665 | bool prev_need_check_error = need_check_error; |
| 1665 | need_check_error = check_error; | 1666 | need_check_error = check_error; |
| @@ -1667,8 +1668,10 @@ public: | |||
| 1667 | // Make sure event deque from taskqueue, then synchronize Event | 1668 | // Make sure event deque from taskqueue, then synchronize Event |
| 1668 | c10_npu::npuSynchronizeDevice(need_check_error); | 1669 | c10_npu::npuSynchronizeDevice(need_check_error); |
| 1669 | std::lock_guard<std::recursive_mutex> lock(mutex); | 1670 | std::lock_guard<std::recursive_mutex> lock(mutex); |
| 1670 | - c10_npu::NPUWorkspaceAllocator::emptyCache(device, need_check_error); | 1671 | + if (mempool_id.first == 0 && mempool_id.second == 0) { |
| 1671 | - release_cached_blocks(check_error, context, free_physical); | 1672 | + c10_npu::NPUWorkspaceAllocator::emptyCache(device, need_check_error); |
| 1673 | + } | ||
| 1674 | + release_cached_blocks(check_error, context, free_physical, mempool_id); | ||
| 1672 | need_check_error = prev_need_check_error; | 1675 | need_check_error = prev_need_check_error; |
| 1673 | TORCH_NPU_MEMORY_LOGI("emptyCache success, device=%d.", device); | 1676 | TORCH_NPU_MEMORY_LOGI("emptyCache success, device=%d.", device); |
| 1674 | } | 1677 | } |
| @@ -2168,6 +2171,14 @@ public: | |||
| 2168 | } | 2171 | } |
| 2169 | } | 2172 | } |
| 2170 | 2173 | ||
| 2174 | + void createOrIncrefPool(MempoolId_t mempool_id) | ||
| 2175 | + { | ||
| 2176 | + // Create a PrivatePool object if it does not exist yet | ||
| 2177 | + // and increment its use_count | ||
| 2178 | + std::lock_guard<std::recursive_mutex> lock(mutex); | ||
| 2179 | + create_or_incref_pool(mempool_id); | ||
| 2180 | + } | ||
| 2181 | + | ||
| 2171 | PrivatePool* get_private_pool(MempoolId_t mempool_id) const | 2182 | PrivatePool* get_private_pool(MempoolId_t mempool_id) const |
| 2172 | { | 2183 | { |
| 2173 | auto it = graph_pools.find(mempool_id); | 2184 | auto it = graph_pools.find(mempool_id); |
| @@ -2843,22 +2854,53 @@ private: | |||
| 2843 | return true; | 2854 | return true; |
| 2844 | } | 2855 | } |
| 2845 | 2856 | ||
| 2846 | - // npuSynchronizeDevice must be executed before this function can be called | 2857 | + /** |
| 2847 | - bool release_cached_blocks(bool check_error, const std::shared_ptr<c10::GatheredContext> &context, bool free_physical) | 2858 | + * 1) If mempool_id is {0,0} (the default pool) and there are no |
| 2859 | + * currently capturing memory pools, free the default pool's blocks | ||
| 2860 | + * and also free the blocks of the freeable private pools. | ||
| 2861 | + * | ||
| 2862 | + * If mempool_id corresponds to a private pool that is freeable, | ||
| 2863 | + * call synchronize_and_free_events() on that private pool. Free the | ||
| 2864 | + * blocks of all freeable private pools, including this one. | ||
| 2865 | + * | ||
| 2866 | + * 2) npuSynchronizeDevice must be executed before this function can be called | ||
| 2867 | + */ | ||
| 2868 | + bool release_cached_blocks( | ||
| 2869 | + bool check_error, | ||
| 2870 | + const std::shared_ptr<c10::GatheredContext>& context, | ||
| 2871 | + bool free_physical, | ||
| 2872 | + MempoolId_t mempool_id = {0, 0}) | ||
| 2848 | { | 2873 | { |
| 2849 | - // First ensure that all blocks that can't currently be allocated due to | 2874 | + if (mempool_id.first == 0 && mempool_id.second == 0 && captures_underway.empty()) { |
| 2850 | - // outstanding events are returned to the pool. | 2875 | + // If there is no active mempool, we work on releasing *all* blocks. |
| 2851 | - synchronize_and_free_events(check_error, context); | ||
| 2852 | 2876 | ||
| 2853 | - // Free all non-split cached blocks | 2877 | + // First ensure that all blocks that can't currently be allocated due to |
| 2854 | - release_blocks(large_blocks, context, free_physical); | 2878 | + // outstanding events are returned to the pool. |
| 2855 | - release_blocks(small_blocks, context, free_physical); | 2879 | + synchronize_and_free_events(check_error, context); |
| 2880 | + | ||
| 2881 | + // Free all non-split cached blocks to system allocator | ||
| 2882 | + release_blocks(large_blocks, context, free_physical); | ||
| 2883 | + release_blocks(small_blocks, context, free_physical); | ||
| 2884 | + } | ||
| 2856 | 2885 | ||
| 2857 | for (auto it = graph_pools_freeable.begin(); it != graph_pools_freeable.end();) { | 2886 | for (auto it = graph_pools_freeable.begin(); it != graph_pools_freeable.end();) { |
| 2858 | - // See notifyCaptureDestroy for the strategy here. | 2887 | + if (mempool_id.first != 0 || mempool_id.second != 0) { |
| 2888 | + if (it->first == mempool_id) { | ||
| 2889 | + // If there is an active mempool, we sync only the events | ||
| 2890 | + // associated with the pool | ||
| 2891 | + synchronize_and_free_events(check_error, context, it->second); | ||
| 2892 | + } else { | ||
| 2893 | + // otherwise we move on | ||
| 2894 | + ++it; | ||
| 2895 | + continue; | ||
| 2896 | + } | ||
| 2897 | + } | ||
| 2898 | + | ||
| 2859 | TORCH_INTERNAL_ASSERT(it->second->use_count == 0); | 2899 | TORCH_INTERNAL_ASSERT(it->second->use_count == 0); |
| 2900 | + | ||
| 2860 | release_blocks(it->second->small_blocks, context, free_physical); | 2901 | release_blocks(it->second->small_blocks, context, free_physical); |
| 2861 | release_blocks(it->second->large_blocks, context, free_physical); | 2902 | release_blocks(it->second->large_blocks, context, free_physical); |
| 2903 | + | ||
| 2862 | if (it->second->npuMalloc_count == 0) { | 2904 | if (it->second->npuMalloc_count == 0) { |
| 2863 | auto erase_count = graph_pools.erase(it->first); | 2905 | auto erase_count = graph_pools.erase(it->first); |
| 2864 | TORCH_INTERNAL_ASSERT(erase_count == 1); | 2906 | TORCH_INTERNAL_ASSERT(erase_count == 1); |
| @@ -3023,7 +3065,10 @@ private: | |||
| 3023 | return event_pool->get(idx); | 3065 | return event_pool->get(idx); |
| 3024 | } | 3066 | } |
| 3025 | 3067 | ||
| 3026 | - void synchronize_and_free_events(bool check_error, const std::shared_ptr<c10::GatheredContext> &context) | 3068 | + void synchronize_and_free_events( |
| 3069 | + bool check_error, | ||
| 3070 | + const std::shared_ptr<c10::GatheredContext> &context, | ||
| 3071 | + PrivatePool* pool = nullptr) | ||
| 3027 | { | 3072 | { |
| 3028 | // This function syncs, so capture should not be underway. Might as well | 3073 | // This function syncs, so capture should not be underway. Might as well |
| 3029 | // make sure capture-deferred end of life events get processed too. | 3074 | // make sure capture-deferred end of life events get processed too. |
| @@ -3031,10 +3076,17 @@ private: | |||
| 3031 | insert_events_deferred_until_no_capture(context); | 3076 | insert_events_deferred_until_no_capture(context); |
| 3032 | 3077 | ||
| 3033 | // Synchronize on outstanding events and then free associated blocks. | 3078 | // Synchronize on outstanding events and then free associated blocks. |
| 3034 | - for (auto &st : npu_events) { | 3079 | + for (auto it = npu_events.begin(); it != npu_events.end();) { |
| 3035 | - for (auto &e : st.second) { | 3080 | + for (auto e = it->second.begin(); e != it->second.end();) { |
| 3036 | - EventPool::Event event = std::move(e.first); | 3081 | + Block *block = e->second; |
| 3037 | - Block *block = e.second; | 3082 | + |
| 3083 | + // If a pool was passed, only synchronize the events | ||
| 3084 | + // that are associated with the pool, otherwise move on | ||
| 3085 | + if (pool && block->pool->owner_PrivatePool != pool) { | ||
| 3086 | + ++e; | ||
| 3087 | + continue; | ||
| 3088 | + } | ||
| 3089 | + EventPool::Event event = std::move(e->first); | ||
| 3038 | auto err = aclrtSynchronizeEvent(*event); | 3090 | auto err = aclrtSynchronizeEvent(*event); |
| 3039 | if (err != ACL_ERROR_NONE) { | 3091 | if (err != ACL_ERROR_NONE) { |
| 3040 | if (check_error) { | 3092 | if (check_error) { |
| @@ -3046,20 +3098,26 @@ private: | |||
| 3046 | TORCH_NPU_MEMORY_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event.get()); | 3098 | TORCH_NPU_MEMORY_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event.get()); |
| 3047 | } | 3099 | } |
| 3048 | 3100 | ||
| 3049 | - const c10_npu::impl::PyCallbackTrigger *trigger = c10_npu::impl::NPUTrace::getTrace(); | 3101 | + const auto *trigger = c10_npu::impl::NPUTrace::getTrace(); |
| 3050 | if (C10_UNLIKELY(trigger)) { | 3102 | if (C10_UNLIKELY(trigger)) { |
| 3051 | trigger->traceNpuEventSynchronization(reinterpret_cast<uintptr_t>(event.get())); | 3103 | trigger->traceNpuEventSynchronization(reinterpret_cast<uintptr_t>(event.get())); |
| 3052 | } | 3104 | } |
| 3053 | 3105 | ||
| 3054 | - | ||
| 3055 | block->event_count--; | 3106 | block->event_count--; |
| 3056 | if (block->event_count == 0) { | 3107 | if (block->event_count == 0) { |
| 3057 | free_block(block, context); | 3108 | free_block(block, context); |
| 3058 | } | 3109 | } |
| 3110 | + e = it->second.erase(e); | ||
| 3111 | + } | ||
| 3112 | + | ||
| 3113 | + // If the events deque is empty, only then erase the | ||
| 3114 | + // cuda event from the events map | ||
| 3115 | + if (it->second.empty()) { | ||
| 3116 | + it = npu_events.erase(it); | ||
| 3117 | + } else { | ||
| 3118 | + ++it; | ||
| 3059 | } | 3119 | } |
| 3060 | } | 3120 | } |
| 3061 | - | ||
| 3062 | - npu_events.clear(); | ||
| 3063 | } | 3121 | } |
| 3064 | 3122 | ||
| 3065 | void remove_npugraph_stream_uses(Block *block) | 3123 | void remove_npugraph_stream_uses(Block *block) |
| @@ -3402,7 +3460,8 @@ public: | |||
| 3402 | } | 3460 | } |
| 3403 | } | 3461 | } |
| 3404 | 3462 | ||
| 3405 | - void emptyCacheImpl(bool check_error, bool free_physical) override | 3463 | + // uninherited function |
| 3464 | + void emptyCache(bool check_error, bool free_physical, MempoolId_t mempool_id) | ||
| 3406 | { | 3465 | { |
| 3407 | TORCH_NPU_MEMORY_LOGD("Begin empty cache with check_error = %d", check_error); | 3466 | TORCH_NPU_MEMORY_LOGD("Begin empty cache with check_error = %d", check_error); |
| 3408 | int32_t current_device = 0; | 3467 | int32_t current_device = 0; |
| @@ -3418,7 +3477,7 @@ public: | |||
| 3418 | } else { | 3477 | } else { |
| 3419 | NPU_CHECK_WARN(c10_npu::SetDevice(device_idx)); | 3478 | NPU_CHECK_WARN(c10_npu::SetDevice(device_idx)); |
| 3420 | } | 3479 | } |
| 3421 | - device_allocator[device_idx]->emptyCache(device_idx, check_error, free_physical); | 3480 | + device_allocator[device_idx]->emptyCache(device_idx, check_error, free_physical, mempool_id); |
| 3422 | } | 3481 | } |
| 3423 | if (check_error) { | 3482 | if (check_error) { |
| 3424 | NPU_CHECK_ERROR(c10_npu::MaybeSetDevice(current_device)); | 3483 | NPU_CHECK_ERROR(c10_npu::MaybeSetDevice(current_device)); |
| @@ -3428,6 +3487,16 @@ public: | |||
| 3428 | TORCH_NPU_MEMORY_LOGD("End empty cache with check_error = %d", check_error); | 3487 | TORCH_NPU_MEMORY_LOGD("End empty cache with check_error = %d", check_error); |
| 3429 | } | 3488 | } |
| 3430 | 3489 | ||
| 3490 | + void emptyCache(MempoolId_t mempool_id) override | ||
| 3491 | + { | ||
| 3492 | + emptyCache(true, true, mempool_id); | ||
| 3493 | + } | ||
🔵 Low Priority
但 MemPool 的 在多设备环境中,析构一个单设备 MemPool 会触发所有设备的同步,造成不必要的延迟。建议将 建议:考虑新增仅针对指定设备的 emptyCache 重载,或在 emptyCache(MempoolId_t) 中利用 MemPool.device() 方法限定操作设备范围,避免不必要的全设备遍历和同步。 ![]() ![]() | |||
| 3494 | + | ||
| 3495 | + void emptyCacheImpl(bool check_error, bool free_physical) override | ||
| 3496 | + { | ||
| 3497 | + emptyCache(check_error, free_physical, {0, 0}); | ||
| 3498 | + } | ||
| 3499 | + | ||
| 3431 | void emptyCache(bool check_error) override | 3500 | void emptyCache(bool check_error) override |
| 3432 | { | 3501 | { |
| 3433 | emptyCacheImpl(check_error, true); | 3502 | emptyCacheImpl(check_error, true); |
| @@ -3920,6 +3989,14 @@ public: | |||
| 3920 | { | 3989 | { |
| 3921 | device_allocator[device]->buildServerMemMapForHccl(hcclComm); | 3990 | device_allocator[device]->buildServerMemMapForHccl(hcclComm); |
| 3922 | } | 3991 | } |
| 3992 | + | ||
| 3993 | + void createOrIncrefPool( | ||
| 3994 | + c10::DeviceIndex device, | ||
| 3995 | + MempoolId_t mempool_id) | ||
| 3996 | + { | ||
| 3997 | + assertValidDevice(device); | ||
| 3998 | + device_allocator[device]->createOrIncrefPool(std::move(mempool_id)); | ||
| 3999 | + } | ||
| 3923 | }; | 4000 | }; |
| 3924 | 4001 | ||
| 3925 | NpuCachingAllocator caching_allocator; | 4002 | NpuCachingAllocator caching_allocator; |
| @@ -4015,6 +4092,14 @@ MemPool::MemPool(NPUCachingAllocator::NPUAllocator *allocator, bool is_user_crea | |||
| 4015 | } else { | 4092 | } else { |
| 4016 | id_ = { uuid_++, 0 }; | 4093 | id_ = { uuid_++, 0 }; |
| 4017 | } | 4094 | } |
| 4095 | + device_ = c10_npu::current_device(); | ||
| 4096 | + NPUCachingAllocator::createOrIncrefPool(device_, id_); | ||
| 4097 | +} | ||
| 4098 | + | ||
| 4099 | +MemPool::~MemPool() { | ||
| 4100 | + // not support UseOnOom in this version | ||
| 4101 | + NPUCachingAllocator::releasePool(device_, id_); | ||
| 4102 | + NPUCachingAllocator::emptyCache(id_); | ||
| 4018 | } | 4103 | } |
🟡 Medium Priority
如果 在 Python 中,以下场景可能触发:
建议:在析构函数中捕获可能的异常或在调用前检查池状态。例如可以在 releasePool 中将对不存在的池改为 TORCH_CHECK 或静默返回,或者在析构函数中使用 try-catch 保护。 ![]() ![]() | |||
| 4019 | 4104 | ||
| 4020 | MempoolId_t MemPool::id() | 4105 | MempoolId_t MemPool::id() |
| @@ -4027,6 +4112,17 @@ NPUCachingAllocator::NPUAllocator *MemPool::allocator() | |||
| 4027 | return allocator_; | 4112 | return allocator_; |
| 4028 | } | 4113 | } |
| 4029 | 4114 | ||
| 4115 | +c10::DeviceIndex MemPool::device() { | ||
| 4116 | + return device_; | ||
| 4117 | +} | ||
| 4118 | + | ||
| 4119 | +MempoolId_t MemPool::graph_pool_handle(bool is_user_created) { | ||
| 4120 | + if (is_user_created) { | ||
| 4121 | + return {0, uid_++}; | ||
| 4122 | + } | ||
| 4123 | + return {uuid_++, 0}; | ||
| 4124 | +} | ||
| 4125 | + | ||
| 4030 | // Note that active_mempool_ is a global variable here | 4126 | // Note that active_mempool_ is a global variable here |
| 4031 | // and not inside MemPoolContext class, because in windows we | 4127 | // and not inside MemPoolContext class, because in windows we |
| 4032 | // can't use __declspec(dllexport) and __declspec(thread) | 4128 | // can't use __declspec(dllexport) and __declspec(thread) |
| @@ -306,6 +306,17 @@ public: | |||
| 306 | // calling the allocator so it is unsafe to try to acquire the GIL in this | 306 | // calling the allocator so it is unsafe to try to acquire the GIL in this |
| 307 | // callback. | 307 | // callback. |
| 308 | virtual void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) = 0; | 308 | virtual void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) = 0; |
| 309 | + // start from torch 2.7, not support input 'allocator' in this version | ||
| 310 | + // will gradually fill in missing abilities | ||
| 311 | + virtual void createOrIncrefPool(c10::DeviceIndex, MempoolId_t) | ||
| 312 | + { | ||
| 313 | + TORCH_CHECK( | ||
| 314 | + false, | ||
| 315 | + name(), | ||
| 316 | + " does not yet support createOrIncrefPool. " | ||
| 317 | + "If you need it, please file an issue describing your use case."); | ||
| 318 | + } | ||
| 319 | + virtual void emptyCache(MempoolId_t mempool_id) = 0; | ||
| 309 | }; | 320 | }; |
| 310 | 321 | ||
| 311 | // Allocator object, statically initialized | 322 | // Allocator object, statically initialized |
| @@ -364,6 +375,12 @@ inline void emptyCacheImpl(bool check_error = true, bool free_physical = true) | |||
| 364 | return get()->emptyCacheImpl(check_error, free_physical); | 375 | return get()->emptyCacheImpl(check_error, free_physical); |
| 365 | } | 376 | } |
| 366 | 377 | ||
| 378 | +// overload, align to cuda, not outer(C10_NPU_API) in this version | ||
| 379 | +inline void emptyCache(MempoolId_t mempool_id) | ||
| 380 | +{ | ||
| 381 | + return get()->emptyCache(mempool_id); | ||
| 382 | +} | ||
| 383 | + | ||
| 367 | C10_NPU_API inline void emptyCache(bool check_error = true) | 384 | C10_NPU_API inline void emptyCache(bool check_error = true) |
| 368 | { | 385 | { |
| 369 | return get()->emptyCache(check_error); | 386 | return get()->emptyCache(check_error); |
| @@ -547,6 +564,13 @@ inline void buildServerMemMapForHccl(int device, std::shared_ptr<c10d_npu::HCCLC | |||
| 547 | return get()->buildServerMemMapForHccl(device, hcclComm); | 564 | return get()->buildServerMemMapForHccl(device, hcclComm); |
| 548 | } | 565 | } |
| 549 | 566 | ||
| 567 | +inline void createOrIncrefPool( | ||
| 568 | + c10::DeviceIndex device, | ||
| 569 | + MempoolId_t mempool_id) | ||
| 570 | +{ | ||
| 571 | + get()->createOrIncrefPool(device, mempool_id); | ||
| 572 | +} | ||
| 573 | + | ||
| 550 | bool checkConfigExpandableSegments(); | 574 | bool checkConfigExpandableSegments(); |
| 551 | 575 | ||
| 552 | bool isConfig1GPageSizeEnable(); | 576 | bool isConfig1GPageSizeEnable(); |
| @@ -571,8 +595,16 @@ struct C10_NPU_API MemPool { | |||
| 571 | NPUCachingAllocator::NPUAllocator* allocator = nullptr, | 595 | NPUCachingAllocator::NPUAllocator* allocator = nullptr, |
| 572 | bool is_user_created = true); | 596 | bool is_user_created = true); |
| 573 | 597 | ||
| 598 | + MemPool(const MemPool&) = delete; | ||
| 599 | + MemPool(MemPool&&) = default; | ||
| 600 | + MemPool& operator=(const MemPool&) = delete; | ||
| 601 | + MemPool& operator=(MemPool&&) = default; | ||
| 602 | + ~MemPool(); | ||
🟡 Medium Priority 变更行为: 受影响的行为/契约:C++ 默认移动构造只是逐成员拷贝,移动后源对象的 失败模式:
触发条件:C++ 侧对 建议:自定义移动构造/赋值将源对象的 id_ 重置为 {0,0},并在析构函数开头处理 {0,0} 池的情况(early return)。或者直接删除移动语义(= delete),若确认 MemPool 不需要移动。 ![]() ![]() | |||
| 603 | + | ||
| 574 | MempoolId_t id(); | 604 | MempoolId_t id(); |
| 575 | NPUCachingAllocator::NPUAllocator* allocator(); | 605 | NPUCachingAllocator::NPUAllocator* allocator(); |
| 606 | + c10::DeviceIndex device(); | ||
| 607 | + static MempoolId_t graph_pool_handle(bool is_user_created = true); | ||
| 576 | 608 | ||
| 577 | private: | 609 | private: |
| 578 | static std::atomic<CaptureId_t> uid_; | 610 | static std::atomic<CaptureId_t> uid_; |
| @@ -580,6 +612,7 @@ private: | |||
| 580 | NPUCachingAllocator::NPUAllocator* allocator_; | 612 | NPUCachingAllocator::NPUAllocator* allocator_; |
| 581 | bool is_user_created_; | 613 | bool is_user_created_; |
| 582 | MempoolId_t id_; | 614 | MempoolId_t id_; |
| 615 | + c10::DeviceIndex device_; | ||
| 583 | }; | 616 | }; |
| 584 | 617 | ||
| 585 | // MemPoolContext holds the currently active pool and stashes the previous | 618 | // MemPoolContext holds the currently active pool and stashes the previous |
| @@ -40,8 +40,7 @@ MempoolId_t graph_pool_handle() | |||
| 40 | { | 40 | { |
| 41 | // Sets just the second value, to distinguish it from MempoolId_ts created from | 41 | // Sets just the second value, to distinguish it from MempoolId_ts created from |
| 42 | // aclmdlRICaptureGetInfo id_s in capture_begin. | 42 | // aclmdlRICaptureGetInfo id_s in capture_begin. |
| 43 | - auto new_pool = c10_npu::MemPool(); | 43 | + return c10_npu::MemPool::graph_pool_handle(); |
| 44 | - return new_pool.id(); | ||
| 45 | } | 44 | } |
| 46 | 45 | ||
| 47 | void graph_task_group_begin(c10_npu::NPUStream stream) | 46 | void graph_task_group_begin(c10_npu::NPUStream stream) |
| @@ -193,8 +192,7 @@ void NPUGraph::capture_begin(MempoolId_t pool, aclmdlRICaptureMode capture_mode, | |||
| 193 | } else { | 192 | } else { |
| 194 | // User did not ask us to share a mempool. Create graph pool handle using is_user_created=false. | 193 | // User did not ask us to share a mempool. Create graph pool handle using is_user_created=false. |
| 195 | // Sets just the first value, to distinguish it from MempoolId_ts created by graph_pool_handle(). | 194 | // Sets just the first value, to distinguish it from MempoolId_ts created by graph_pool_handle(). |
| 196 | - auto mempool = c10_npu::MemPool({}, false); | 195 | + mempool_id_ = c10_npu::MemPool::graph_pool_handle(false); |
| 197 | - mempool_id_ = mempool.id(); | ||
| 198 | TORCH_INTERNAL_ASSERT(mempool_id_.first > 0); | 196 | TORCH_INTERNAL_ASSERT(mempool_id_.first > 0); |
| 199 | } | 197 | } |
| 200 | 198 | ||
| @@ -4,6 +4,7 @@ | |||
| 4 | 4 | ||
| 5 | 5 | ||
| 6 | 6 | ||
| 7 | + | ||
| 7 | 8 | ||
| 8 | template <typename T> | 9 | template <typename T> |
| 9 | using shared_ptr_class_ = py::class_<T, std::shared_ptr<T>>; | 10 | using shared_ptr_class_ = py::class_<T, std::shared_ptr<T>>; |
| @@ -11,7 +12,11 @@ using shared_ptr_class_ = py::class_<T, std::shared_ptr<T>>; | |||
| 11 | void TORCH_NPU_API THNPMemPool_init(PyObject* module) { | 12 | void TORCH_NPU_API THNPMemPool_init(PyObject* module) { |
| 12 | auto torch_C_m = py::handle(module).cast<py::module>(); | 13 | auto torch_C_m = py::handle(module).cast<py::module>(); |
| 13 | shared_ptr_class_<::c10_npu::MemPool>(torch_C_m, "_MemPool") | 14 | shared_ptr_class_<::c10_npu::MemPool>(torch_C_m, "_MemPool") |
| 14 | - .def(py::init<c10_npu::NPUCachingAllocator::NPUAllocator*, bool>()) | 15 | + .def(py::init( |
| 16 | + [](c10_npu::NPUCachingAllocator::NPUAllocator* allocator, bool is_user_created) { | ||
| 17 | + torch_npu::utils::npu_lazy_init(); // init npu before construct mempool | ||
| 18 | + return std::make_shared<::c10_npu::MemPool>(allocator, is_user_created); | ||
| 19 | + })) | ||
| 15 | .def_property_readonly("id", &::c10_npu::MemPool::id) | 20 | .def_property_readonly("id", &::c10_npu::MemPool::id) |
| 16 | .def_property_readonly("allocator", &::c10_npu::MemPool::allocator); | 21 | .def_property_readonly("allocator", &::c10_npu::MemPool::allocator); |
| 17 | shared_ptr_class_<::c10_npu::MemPoolContext>(torch_C_m, "_MemPoolContext") | 22 | shared_ptr_class_<::c10_npu::MemPoolContext>(torch_C_m, "_MemPoolContext") |
| @@ -236,6 +236,12 @@ void NPUPluggableAllocator::emptyCacheImpl(bool check_error, bool free_physical) | |||
| 236 | "If you need it, please file an issue describing your use case."); | 236 | "If you need it, please file an issue describing your use case."); |
| 237 | } | 237 | } |
| 238 | 238 | ||
| 239 | +void NPUPluggableAllocator::emptyCache(c10_npu::MempoolId_t mempool_id) | ||
| 240 | +{ | ||
| 241 | + TORCH_NPU_WARN("NPUPluggableAllocator does not yet support emptyCache with mempool_id. " | ||
| 242 | + "If you need it, please file an issue describing your use case."); | ||
| 243 | +} | ||
| 244 | + | ||
| 239 | void NPUPluggableAllocator::emptyCache(bool check_error) | 245 | void NPUPluggableAllocator::emptyCache(bool check_error) |
| 240 | { | 246 | { |
| 241 | if (reset_fn_) { | 247 | if (reset_fn_) { |
| @@ -66,6 +66,7 @@ struct NPUPluggableAllocator | |||
| 66 | void setMemoryFraction(double fraction, int device) override; | 66 | void setMemoryFraction(double fraction, int device) override; |
| 67 | void emptyCacheImpl(bool check_error, bool free_physical) override; | 67 | void emptyCacheImpl(bool check_error, bool free_physical) override; |
| 68 | void emptyCache(bool check_error) override; | 68 | void emptyCache(bool check_error) override; |
| 69 | + void emptyCache(c10_npu::MempoolId_t mempool_id = {0, 0}) override; | ||
| 69 | void emptyVirtAddrCache(bool check_error) override; | 70 | void emptyVirtAddrCache(bool check_error) override; |
| 70 | void cacheInfo(int dev_id, size_t* cachedAndFree, size_t* largestBlock) override; | 71 | void cacheInfo(int dev_id, size_t* cachedAndFree, size_t* largestBlock) override; |
| 71 | void* getBaseAllocation(void* ptr, size_t* size) override; | 72 | void* getBaseAllocation(void* ptr, size_t* size) override; |
| @@ -6,8 +6,7 @@ import pickle | |||
| 6 | import sys | 6 | import sys |
| 7 | import os | 7 | import os |
| 8 | import stat | 8 | import stat |
| 9 | -import platform | 9 | +from typing import Optional, Tuple |
| 10 | -from typing import Any, Dict, Optional, Tuple, Union | ||
| 11 | 10 | ||
| 12 | import torch_npu | 11 | import torch_npu |
| 13 | from torch_npu.utils._error_code import ErrCode, pta_error | 12 | from torch_npu.utils._error_code import ErrCode, pta_error |
| @@ -102,7 +101,7 @@ def caching_allocator_alloc(size, device=None, stream=None): | |||
| 102 | if not isinstance(stream, int): | 101 | if not isinstance(stream, int): |
| 103 | raise TypeError('Invalid type for stream argument, must be ' | 102 | raise TypeError('Invalid type for stream argument, must be ' |
| 104 | '`torch_npu.npu.Stream` or `int` representing a pointer ' | 103 | '`torch_npu.npu.Stream` or `int` representing a pointer ' |
| 105 | - 'to a exisiting stream' + pta_error(ErrCode.TYPE)) | 104 | + 'to a existing stream' + pta_error(ErrCode.TYPE)) |
| 106 | with torch_npu.npu.device(device): | 105 | with torch_npu.npu.device(device): |
| 107 | return torch_npu._C._npu_npuCachingAllocator_raw_alloc(size, stream) | 106 | return torch_npu._C._npu_npuCachingAllocator_raw_alloc(size, stream) |
| 108 | 107 | ||
| @@ -832,6 +831,8 @@ def use_mem_pool(pool: MemPool, device=None): | |||
| 832 | yield | 831 | yield |
| 833 | finally: | 832 | finally: |
| 834 | torch_npu._C._npu_endAllocateCurrentStreamToPool(device_index, pool.id) | 833 | torch_npu._C._npu_endAllocateCurrentStreamToPool(device_index, pool.id) |
| 834 | + # after endAllocate, need call releasePool to achieve the right count(use_count--) | ||
| 835 | + torch_npu._C._npu_releasePool(device_index, pool.id) | ||
| 835 | del ctx | 836 | del ctx |
| 836 | 837 | ||
| 837 | 838 | ||
| @@ -856,7 +857,7 @@ def _record_memory_history(enabled="all", *args, **kwargs): | |||
| 856 | Args: | 857 | Args: |
| 857 | enabled (Literal[None, "state", "all"], optional): | 858 | enabled (Literal[None, "state", "all"], optional): |
| 858 | `None`, disable recording memory history. | 859 | `None`, disable recording memory history. |
| 859 | - `"state"`, keep information for currenly allocated memory. | 860 | + `"state"`, keep information for currently allocated memory. |
| 860 | `"all"`, additionally keep a history of all alloc/free calls. | 861 | `"all"`, additionally keep a history of all alloc/free calls. |
| 861 | Defaults to "all". | 862 | Defaults to "all". |
| 862 | context (Literal[None, "state", "alloc", "all"], optional): | 863 | context (Literal[None, "state", "alloc", "all"], optional): |


🟡 Medium Priority
release_cached_blocks函数中,当mempool_id非零(即针对特定池的释放路径,由MemPool::~MemPool()触发)并且匹配到graph_pools_freeable中的条目时,直接调用synchronize_and_free_events(check_error, context, it->second)(第 2866 行)。但synchronize_and_free_events内部第一行就是TORCH_INTERNAL_ASSERT(captures_underway.empty())(第 3049 行)。而在默认
{0, 0}路径中(第 2856 行),调用synchronize_and_free_events前有captures_underway.empty()守卫。但特定池路径(第 2864-2870 行)缺少该守卫。触发条件:当同一设备上存在另一个正在进行中的 graph capture(属于不同 pool)时,
MemPool::~MemPool()析构函数会触发此路径,导致进程 ASSERT 崩溃。虽然emptyCache调用前已执行npuSynchronizeDevice,但这只同步设备执行,不会清空captures_underway列表。建议:在调用 synchronize_and_free_events 前增加 captures_underway.empty() 守卫,与默认路径行为保持一致。如果 captures_underway 非空,应跳过本次同步释放,或仅在不触发 ASSERT 的前提下释放 blocks。