Star854
1342
代码
代码
Issues1628
Pull Requests2194
项目讨论
Wiki
分析
Star854
1342
已合并
[feature] aclgraph support PYTORCH_NO_NPU_MEMORY_CACHING=1 2.9.0 #35178
yurongkun创建于 5月9日
[feature] aclgraph support PYTORCH_NO_NPU_MEMORY_CACHING=1 2.9.0 #35178
已合并
yurongkun创建于 5月9日
从yurongkun/pytorch-yu-allno_cache_2.9合入到Ascend/pytorchv2.9.0
共 4 个文件变更+99-18
M.lintrunner.toml+14-0
@@ -103,6 +103,8 @@ exclude_patterns = [
103 'torch_npu/csrc/inductor/aoti_torch/generated/**',103 'torch_npu/csrc/inductor/aoti_torch/generated/**',
104 'torch_npu/csrc/jit/serialization/mobile_bytecode_generated.h',104 'torch_npu/csrc/jit/serialization/mobile_bytecode_generated.h',
105 'torch_npu/csrc/utils/pythoncapi_compat.h',105 'torch_npu/csrc/utils/pythoncapi_compat.h',
106+ 'torch_npu/csrc/core/npu/NPUCachingAllocator.cpp',
107+ 'torch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp',
106 'aten/src/ATen/dlpack.h',108 'aten/src/ATen/dlpack.h',
107]109]
108init_command = [110init_command = [
@@ -421,6 +423,8 @@ exclude_patterns=[
421 'torch_npu/csrc/core/npu/NPUGraph.cpp',423 'torch_npu/csrc/core/npu/NPUGraph.cpp',
422 'torch_npu/csrc/core/npu/NPUGraph.h',424 'torch_npu/csrc/core/npu/NPUGraph.h',
423 'torch_npu/csrc/npu/Graph.cpp',425 'torch_npu/csrc/npu/Graph.cpp',
426+ 'torch_npu/csrc/core/npu/NPUCachingAllocator.cpp',
427+ 'torch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp',
424 'torch_npu/npu/_graph_tree.py',428 'torch_npu/npu/_graph_tree.py',
425 'torch_npu/npu/graphs.py',429 'torch_npu/npu/graphs.py',
426 'torch_npu/utils/_graph_tree.py',430 'torch_npu/utils/_graph_tree.py',
@@ -452,6 +456,8 @@ exclude_patterns = [
452 'torch_npu/csrc/core/npu/NPUGraph.cpp',456 'torch_npu/csrc/core/npu/NPUGraph.cpp',
453 'torch_npu/csrc/core/npu/NPUGraph.h',457 'torch_npu/csrc/core/npu/NPUGraph.h',
454 'torch_npu/csrc/npu/Graph.cpp',458 'torch_npu/csrc/npu/Graph.cpp',
459+ 'torch_npu/csrc/core/npu/NPUCachingAllocator.cpp',
460+ 'torch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp',
455 'torch_npu/npu/_graph_tree.py',461 'torch_npu/npu/_graph_tree.py',
456 'torch_npu/npu/graphs.py',462 'torch_npu/npu/graphs.py',
457 'torch_npu/utils/_graph_tree.py',463 'torch_npu/utils/_graph_tree.py',
@@ -494,6 +500,8 @@ exclude_patterns = [
494 'torch_npu/csrc/core/npu/NPUGraph.cpp',500 'torch_npu/csrc/core/npu/NPUGraph.cpp',
495 'torch_npu/csrc/core/npu/NPUGraph.h',501 'torch_npu/csrc/core/npu/NPUGraph.h',
496 'torch_npu/csrc/npu/Graph.cpp',502 'torch_npu/csrc/npu/Graph.cpp',
503+ 'torch_npu/csrc/core/npu/NPUCachingAllocator.cpp',
504+ 'torch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp',
497 'torch_npu/npu/_graph_tree.py',505 'torch_npu/npu/_graph_tree.py',
498 'torch_npu/npu/graphs.py',506 'torch_npu/npu/graphs.py',
499 'torch_npu/utils/_graph_tree.py',507 'torch_npu/utils/_graph_tree.py',
@@ -584,6 +592,8 @@ exclude_patterns = [
584 'torch_npu/csrc/core/npu/NPUGraph.cpp',592 'torch_npu/csrc/core/npu/NPUGraph.cpp',
585 'torch_npu/csrc/core/npu/NPUGraph.h',593 'torch_npu/csrc/core/npu/NPUGraph.h',
586 'torch_npu/csrc/npu/Graph.cpp',594 'torch_npu/csrc/npu/Graph.cpp',
595+ 'torch_npu/csrc/core/npu/NPUCachingAllocator.cpp',
596+ 'torch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp',
587]597]
588command = [598command = [
589 'python3',599 'python3',
@@ -617,6 +627,8 @@ exclude_patterns = [
617 '**/fb/**',627 '**/fb/**',
618 'torch_npu/csrc/jit/serialization/mobile_bytecode_generated.h',628 'torch_npu/csrc/jit/serialization/mobile_bytecode_generated.h',
619 'torch_npu/csrc/utils/pythoncapi_compat.h',629 'torch_npu/csrc/utils/pythoncapi_compat.h',
630+ 'torch_npu/csrc/core/npu/NPUCachingAllocator.cpp',
631+ 'torch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp',
620]632]
621command = [633command = [
622 'python3',634 'python3',
@@ -1224,6 +1236,8 @@ exclude_patterns = [
1224 'torch_npu/csrc/core/npu/NPUGraph.cpp',1236 'torch_npu/csrc/core/npu/NPUGraph.cpp',
1225 'torch_npu/csrc/core/npu/NPUGraph.h',1237 'torch_npu/csrc/core/npu/NPUGraph.h',
1226 'torch_npu/csrc/npu/Graph.cpp',1238 'torch_npu/csrc/npu/Graph.cpp',
1239+ 'torch_npu/csrc/core/npu/NPUCachingAllocator.cpp',
1240+ 'torch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp',
1227 'torch_npu/npu/_graph_tree.py',1241 'torch_npu/npu/_graph_tree.py',
1228 'torch_npu/npu/graphs.py',1242 'torch_npu/npu/graphs.py',
1229 'torch_npu/utils/_graph_tree.py',1243 'torch_npu/utils/_graph_tree.py',
Atest/npu/test_no_npu_memory_caching.py+58-0
@@ -0,0 +1,58 @@
1+import gc
2+import os
3+ 
4+ 
5+os.environ["PYTORCH_NO_NPU_MEMORY_CACHING"] = "1"
6+ 
7+import torch
8+ 
9+import torch_npu
10+from torch_npu.testing.testcase import run_tests, TestCase
11+ 
12+ 
13+class TestNoNpuMemoryCaching(TestCase):
14+ def test_aclgraph_capture_reset_releases_reserved(self):
15+ # With PYTORCH_NO_NPU_MEMORY_CACHING=1, capture-time allocations must still go through
16+ # caching semantics for address stability, but when the graph is reset/destroyed,
17+ # the pool should release physical memory promptly.
18+ torch.npu.set_device(0)
19+ torch_npu.npu.empty_cache()
20+ gc.collect()
21+ torch_npu.npu.synchronize()
22+ 
23+ g = torch_npu.npu.NPUGraph()
24+ 
25+ x = torch.empty((512, 512), device="npu", dtype=torch.float32)
26+ y = torch.empty((512, 512), device="npu", dtype=torch.float32)
27+ torch_npu.npu.synchronize()
28+ 
29+ s0 = torch_npu.npu.Stream()
30+ s0.wait_stream(torch_npu.npu.current_stream())
31+ with torch_npu.npu.stream(s0):
32+ g.capture_begin()
33+ z = x @ y
34+ g.capture_end()
35+ 
36+ capture_alloc = torch_npu.npu.memory_allocated()
37+ capture_rsv = torch_npu.npu.memory_reserved()
38+ torch_npu.npu.synchronize()
39+ 
40+ g.replay()
41+ replay_alloc = torch_npu.npu.memory_allocated()
42+ replay_rsv = torch_npu.npu.memory_reserved()
43+ 
44+ del z
45+ gc.collect()
46+ torch_npu.npu.synchronize()
47+ 
48+ g.reset()
49+ del g
50+ gc.collect()
51+ torch_npu.npu.synchronize()
52+ 
53+ self.assertEqual(capture_alloc, replay_alloc)
54+ self.assertEqual(capture_rsv, replay_rsv)
55+ 
56+ 
57+if __name__ == "__main__":
58+ run_tests()
Mtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp+19-12
@@ -23,6 +23,7 @@
23#include "torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h"23#include "torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h"
24#include "torch_npu/csrc/core/npu/NPURecovery.h"24#include "torch_npu/csrc/core/npu/NPURecovery.h"
25#include "torch_npu/csrc/core/npu/NPUGuard.h"25#include "torch_npu/csrc/core/npu/NPUGuard.h"
26+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
26#include "NPUBlockHandle.h"27#include "NPUBlockHandle.h"
27#include "torch_npu/csrc/core/npu/NpuVariables.h"28#include "torch_npu/csrc/core/npu/NpuVariables.h"
28#include "torch_npu/csrc/core/npu/GetCANNInfo.h"29#include "torch_npu/csrc/core/npu/GetCANNInfo.h"
@@ -1468,7 +1469,7 @@ public:
1468 1469 
1469 block->allocated = false;1470 block->allocated = false;
1470 1471 
1471- // following logic might modifying underlaying Block, causing the size1472+ // following logic might modifying underlying Block, causing the size
1472 // changed. We store ahead for reporting1473 // changed. We store ahead for reporting
1473 auto orig_block_ptr = block->ptr;1474 auto orig_block_ptr = block->ptr;
1474 auto orig_block_size = block->size;1475 auto orig_block_size = block->size;
@@ -1587,7 +1588,7 @@ public:
1587 std::lock_guard<std::recursive_mutex> lock(mutex);1588 std::lock_guard<std::recursive_mutex> lock(mutex);
1588 block->stream_uses.erase(stream);1589 block->stream_uses.erase(stream);
1589 1590 
1590- // free block, lazy destory block related events1591+ // free block, lazy destroy block related events
1591 for (auto it = npu_events[stream].begin(); it != npu_events[stream].end();) {1592 for (auto it = npu_events[stream].begin(); it != npu_events[stream].end();) {
1592 if (block != it->second) {1593 if (block != it->second) {
1593 it++;1594 it++;
@@ -2003,7 +2004,7 @@ public:
2003 const auto all_blocks = get_all_blocks();2004 const auto all_blocks = get_all_blocks();
2004 2005 
2005 for (const Block * const head_block : all_blocks) {2006 for (const Block * const head_block : all_blocks) {
2006- // For expandable segments, we report one segment for each continguous2007+ // For expandable segments, we report one segment for each contiguous
2007 // mapped range of memory2008 // mapped range of memory
2008 if (head_block->prev && head_block->prev->mapped) {2009 if (head_block->prev && head_block->prev->mapped) {
2009 continue;2010 continue;
@@ -2170,6 +2171,10 @@ public:
2170 // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)2171 // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)
2171 bool inserted = graph_pools_freeable.insert({ mempool_id, it->second.get() }).second;2172 bool inserted = graph_pools_freeable.insert({ mempool_id, it->second.get() }).second;
2172 TORCH_INTERNAL_ASSERT(inserted);2173 TORCH_INTERNAL_ASSERT(inserted);
2174+ if (c10_npu::option::OptionsManager::CheckForceUncached() && captures_underway.empty()) {
2175+ std::shared_ptr<c10::GatheredContext> context = maybeGatherContext(RecordContext::ALL);
2176+ release_cached_blocks(true, context, true);
2177+ }
2173 }2178 }
2174 }2179 }
2175 2180 
@@ -2391,7 +2396,7 @@ private:
2391 // cannot be freed when requested, but fully free pages2396 // cannot be freed when requested, but fully free pages
2392 // of expandable blocks can always be freed.2397 // of expandable blocks can always be freed.
2393 // The logic to track this as statistic is pretty involved,2398 // The logic to track this as statistic is pretty involved,
2394- // so we simply just exclude expandable segements from2399+ // so we simply just exclude expandable segments from
2395 // inactive_split2400 // inactive_split
2396 if (!block->expandable_segment_) {2401 if (!block->expandable_segment_) {
2397 update_stat(stats.inactive_split[stat_type], net_change_inactive_split_blocks);2402 update_stat(stats.inactive_split[stat_type], net_change_inactive_split_blocks);
@@ -3018,7 +3023,7 @@ private:
3018 npu_events[stream].emplace_back(std::move(event), block);3023 npu_events[stream].emplace_back(std::move(event), block);
3019 }3024 }
3020 if (ret_ctx == ACL_ERROR_NONE) {3025 if (ret_ctx == ACL_ERROR_NONE) {
3021- NPU_CHECK_ERROR(aclrtSetCurrentContext(compiler_ctx)); 3026+ NPU_CHECK_ERROR(aclrtSetCurrentContext(compiler_ctx));
3022 // Setting context will exchange device implicitly, so we need to reset the cached device here to ensure consistency.3027 // Setting context will exchange device implicitly, so we need to reset the cached device here to ensure consistency.
3023 NPU_CHECK_ERROR(c10_npu::SetDevice(pre_device));3028 NPU_CHECK_ERROR(c10_npu::SetDevice(pre_device));
3024 }3029 }
@@ -3407,7 +3412,7 @@ public:
3407 3412 
3408 if (block->stream != c10_npu::getCurrentNPUStream(block->device).stream(false)) {3413 if (block->stream != c10_npu::getCurrentNPUStream(block->device).stream(false)) {
3409 // If the Stream applying for tensor block different from3414 // If the Stream applying for tensor block different from
3410- // the stream of submiting event wait task in HCCL synchronize()3415+ // the stream of submitting event wait task in HCCL synchronize()
3411 // method, the recordSteam can not be erased.3416 // method, the recordSteam can not be erased.
3412 // New tensor creation may use the block before HCCL op is complete.3417 // New tensor creation may use the block before HCCL op is complete.
3413 return;3418 return;
@@ -3429,7 +3434,7 @@ public:
3429 3434 
3430 if (block->stream != c10_npu::getCurrentNPUStream(block->device).stream(false) || block->hccl_work_ptr != work_ptr) {3435 if (block->stream != c10_npu::getCurrentNPUStream(block->device).stream(false) || block->hccl_work_ptr != work_ptr) {
3431 // If the Stream applying for tensor block different from3436 // If the Stream applying for tensor block different from
3432- // the stream of submiting event wait task in HCCL synchronize()3437+ // the stream of submitting event wait task in HCCL synchronize()
3433 // method, the recordSteam can not be erased.3438 // method, the recordSteam can not be erased.
3434 // New tensor creation may use the block before HCCL op is complete.3439 // New tensor creation may use the block before HCCL op is complete.
3435 return;3440 return;
@@ -3554,7 +3559,8 @@ public:
3554 void (*deleteFunc)(void *) = &local_raw_delete;3559 void (*deleteFunc)(void *) = &local_raw_delete;
3555 3560 
3556 if (size != 0) {3561 if (size != 0) {
3557- if (c10_npu::option::OptionsManager::CheckForceUncached()) {3562+ if (c10_npu::option::OptionsManager::CheckForceUncached() &&
3563+ (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None)) {
3558 deleteFunc = &uncached_delete;3564 deleteFunc = &uncached_delete;
3559 size_t alloc_size = size + AddPadSize();3565 size_t alloc_size = size + AddPadSize();
3560 NPU_CHECK_ERROR(c10_npu::acl::AclrtMallocAlign32(&devPtr, alloc_size,3566 NPU_CHECK_ERROR(c10_npu::acl::AclrtMallocAlign32(&devPtr, alloc_size,
@@ -3584,7 +3590,8 @@ public:
3584 3590 
3585 size_t aligned = base_addr_aligned_kb * 1024;3591 size_t aligned = base_addr_aligned_kb * 1024;
3586 if (size != 0) {3592 if (size != 0) {
3587- if (c10_npu::option::OptionsManager::CheckForceUncached()) {3593+ if (c10_npu::option::OptionsManager::CheckForceUncached() &&
3594+ (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None)) {
3588 deleteFunc = &uncached_delete;3595 deleteFunc = &uncached_delete;
3589 size_t alloc_size = size + AddPadSize() + aligned;3596 size_t alloc_size = size + AddPadSize() + aligned;
3590 NPU_CHECK_ERROR(c10_npu::acl::AclrtMallocAlign32(&realPtr, alloc_size,3597 NPU_CHECK_ERROR(c10_npu::acl::AclrtMallocAlign32(&realPtr, alloc_size,
@@ -3601,11 +3608,11 @@ public:
3601 3608 
3602 c10::DeleterFnPtr raw_deleter() const override3609 c10::DeleterFnPtr raw_deleter() const override
3603 {3610 {
3604- if (c10_npu::option::OptionsManager::CheckForceUncached()) {3611+ if (c10_npu::option::OptionsManager::CheckForceUncached() &&
3612+ (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None)) {
3605 return &uncached_delete;3613 return &uncached_delete;
3606- } else {
3607- return &local_raw_delete;
3608 }3614 }
3615+ return &local_raw_delete;
3609 }3616 }
3610 3617 
3611 void cacheInfo(int dev_id, size_t *cachedAndFree, size_t *largestBlock) override3618 void cacheInfo(int dev_id, size_t *cachedAndFree, size_t *largestBlock) override
Mtorch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp+8-6
@@ -11,6 +11,7 @@
11#include "torch_npu/csrc/core/npu/NPUFunctions.h"11#include "torch_npu/csrc/core/npu/NPUFunctions.h"
12#include "torch_npu/csrc/framework/utils/OpPreparation.h"12#include "torch_npu/csrc/framework/utils/OpPreparation.h"
13#include "torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h"13#include "torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h"
14+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
14 15 
15#ifndef BUILD_LIBTORCH16#ifndef BUILD_LIBTORCH
16#include "torch_npu/csrc/profiler/npu_profiler.h"17#include "torch_npu/csrc/profiler/npu_profiler.h"
@@ -70,7 +71,7 @@ public:
70 // Memory snapshots have deadlock issues in some scenarios, no longer capture python stacks.71 // Memory snapshots have deadlock issues in some scenarios, no longer capture python stacks.
71 return nullptr;72 return nullptr;
72 }73 }
73- 74+ 
74 void* malloc(size_t size, aclrtStream stream)75 void* malloc(size_t size, aclrtStream stream)
75 {76 {
76 auto context = maybeGatherContext(RecordContext::STATE);77 auto context = maybeGatherContext(RecordContext::STATE);
@@ -490,7 +491,8 @@ public:
490 void* dev_ptr = nullptr;491 void* dev_ptr = nullptr;
491 void (*delete_func)(void*) = &local_raw_delete;492 void (*delete_func)(void*) = &local_raw_delete;
492 493 
493- if (c10_npu::option::OptionsManager::CheckForceUncached()) {494+ if (c10_npu::option::OptionsManager::CheckForceUncached() &&
495+ (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None)) {
494 delete_func = &uncached_delete;496 delete_func = &uncached_delete;
495 if (size != 0) {497 if (size != 0) {
496 size_t alloc_size = size + 32;498 size_t alloc_size = size + 32;
@@ -507,11 +509,11 @@ public:
507 509 
508 c10::DeleterFnPtr raw_deleter() const override510 c10::DeleterFnPtr raw_deleter() const override
509 {511 {
510- if (c10_npu::option::OptionsManager::CheckForceUncached()) {512+ if (c10_npu::option::OptionsManager::CheckForceUncached() &&
513+ (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None)) {
511 return &uncached_delete;514 return &uncached_delete;
512- } else {515+ }
513- return &local_raw_delete;516+ return &local_raw_delete;
514- }
515 }517 }
516 518 
517 // Note [COW/lazy_clone is not supported yet]519 // Note [COW/lazy_clone is not supported yet]