已合并
[feat] add ~mempool() with emptyCache #39221
[feat] add ~mempool() with emptyCache #39221
已合并
xuyun15创建于 6月24日
8 个文件变更+221-33
@@ -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 false1664 // 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 Event1668 // 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) const2182 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 called2857+ /**
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 to2874+ 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 blocks2877+ // 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+ }
atomgit-bot
atomgit-botatomgit-bot6月24日

🟡 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。

likedislike
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 well3073 // 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#ifndef BUILD_LIBTORCH3100#ifndef BUILD_LIBTORCH
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#endif3105#endif
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) override3463+ // 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+ }
atomgit-bot
atomgit-botatomgit-bot6月24日

🔵 Low Priority

MemPool::~MemPool() 调用 emptyCache(id_)NPUAllocator::emptyCache(MempoolId_t)emptyCache(true, true, mempool_id)。后者遍历 c10_npu::GetUsedDevices() 返回的所有设备(第 3442-3449 行),对每个设备执行 SetDevice + npuSynchronizeDevice + release_cached_blocks

但 MemPool 的 device_ 在构造时已固定为单设备(第 4056 行),且 releasePool(device_, id_) 仅操作该设备。对于其他设备,release_cached_blocks 在遍历 graph_pools_freeable 时找不到匹配的 mempool_id,实际是空操作,但 npuSynchronizeDeviceSetDevice 的开销已经发生。

在多设备环境中,析构一个单设备 MemPool 会触发所有设备的同步,造成不必要的延迟。建议将 emptyCache(MempoolId_t) 的实现限定为仅操作 pool 所在的设备,或新增仅针对单设备的 emptyCache(device, mempool_id) 重载。

建议:考虑新增仅针对指定设备的 emptyCache 重载,或在 emptyCache(MempoolId_t) 中利用 MemPool.device() 方法限定操作设备范围,避免不必要的全设备遍历和同步。

likedislike
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) override3500 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 
3925NpuCachingAllocator caching_allocator;4002NpuCachingAllocator 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}
atomgit-bot
atomgit-botatomgit-bot6月24日

🟡 Medium Priority

MemPool::~MemPool() 无条件调用 NPUCachingAllocator::releasePool(device_, id_),这会进入 DeviceCachingAllocator::releasePoolget_private_pool(mempool_id),而 get_private_pool 中有 TORCH_INTERNAL_ASSERT(it != graph_pools.end())(第 2177 行)。

如果 MemPool 被创建后从未通过 torch.npu.use_mem_pool() 激活(即从未调用 beginAllocateToPool),则池在 graph_pools 中不存在,导致 ASSERT 崩溃。

在 Python 中,以下场景可能触发:

  • 用户创建了 MemPool 但因逻辑分支未使用就删除
  • Python GC 回收未使用的 MemPool 对象
  • 异常路径中创建了 MemPool 但未进入 use_mem_pool 上下文

建议:在析构函数中捕获可能的异常或在调用前检查池状态。例如可以在 releasePool 中将对不存在的池改为 TORCH_CHECK 或静默返回,或者在析构函数中使用 try-catch 保护。

likedislike
4019 4104 
4020MempoolId_t MemPool::id()4105MempoolId_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 here4126// Note that active_mempool_ is a global variable here
4031// and not inside MemPoolContext class, because in windows we4127// 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 this306 // 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 initialized322// 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+ 
367C10_NPU_API inline void emptyCache(bool check_error = true)384C10_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+ 
550bool checkConfigExpandableSegments();574bool checkConfigExpandableSegments();
551 575 
552bool isConfig1GPageSizeEnable();576bool 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();
atomgit-bot
atomgit-botatomgit-bot6月24日

🟡 Medium Priority

变更行为:MemPool 新增了析构函数 ~MemPool()(第 4059–4063 行),其中调用 releasePool(device_, id_)emptyCache(id_) 来释放池资源。同时显式声明了移动构造/移动赋值为 = default(头文件第 582–584 行)。

受影响的行为/契约:C++ 默认移动构造只是逐成员拷贝,移动后源对象的 device_id_ 保持不变。当源对象和目标对象先后析构时,会对同一个 {device_, id_} 调用两次 releasePool

失败模式:

  1. 第一次 ~MemPool() 调用 releasePool,将 use_count 减至 0,将池插入 graph_pools_freeable,然后 emptyCache 释放块并从 graph_pools 中擦除池。
  2. 第二次 ~MemPool()(被移动的源对象)再次调用 releasePool(device_, id_),其中 get_private_pool(mempool_id)graph_pools 中找不到该池,触发 TORCH_INTERNAL_ASSERT(it != graph_pools.end()),导致进程崩溃。

触发条件:C++ 侧对 MemPool 对象执行移动操作(例如 MemPool b = std::move(a);)。Python 绑定使用 shared_ptr 管理 MemPool,因此 Python 侧一般不会触发移动;但 C++ 内部或未来扩展中可能触发。

建议:自定义移动构造/赋值将源对象的 id_ 重置为 {0,0},并在析构函数开头处理 {0,0} 池的情况(early return)。或者直接删除移动语义(= delete),若确认 MemPool 不需要移动。

likedislike
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 
577private:609private:
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 previous618// 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 from41 // 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 
47void graph_task_group_begin(c10_npu::NPUStream stream)46void 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#include <torch/csrc/utils/pybind.h>4#include <torch/csrc/utils/pybind.h>
5 5 
6#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"6#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
7+#include "torch_npu/csrc/utils/LazyInit.h"
7 8 
8template <typename T>9template <typename T>
9using shared_ptr_class_ = py::class_<T, std::shared_ptr<T>>;10using 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>>;
11void TORCH_NPU_API THNPMemPool_init(PyObject* module) {12void 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+ 
239void NPUPluggableAllocator::emptyCache(bool check_error)245void 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
6import sys6import sys
7import os7import os
8import stat8import stat
9-import platform9+from typing import Optional, Tuple
10-from typing import Any, Dict, Optional, Tuple, Union
11 10 
12import torch_npu11import torch_npu
13from torch_npu.utils._error_code import ErrCode, pta_error12from 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 yield831 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 ctx836 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):