Star846
1320
代码
代码
Issues1500
Pull Requests2300
流水线
项目讨论
Wiki
分析
Star846
1320
已合并
[PROF]fix profiler report npu mem err in multi thread #14621
mei-feiyao创建于 2024年9月14日
[PROF]fix profiler report npu mem err in multi thread #14621
已合并
mei-feiyao创建于 2024年9月14日
从refs/pull/14621/head合入到master
共 2 个文件变更+122-98
Atest/npu/test_allocator_multi_thread_prof.py+61-0
@@ -0,0 +1,61 @@
1+import threading
2+import torch
3+from torch.utils.data import Dataset
4+ 
5+from torch_npu.testing.testcase import TestCase, run_tests
6+ 
7+ 
8+class RandomDataset(Dataset):
9+ def __init__(self, size, length):
10+ self.len = length
11+ self.data = torch.randn(length, size)
12+ 
13+ def __len__(self):
14+ return self.len
15+
16+ def __getitem__(self, index):
17+ return self.data[index].clone()
18+ 
19+ 
20+def thread_worker(data, device, results, index):
21+ for _ in range(50):
22+ data = data.to(device, non_blocking=True)
23+ results.append(data)
24+ 
25+ 
26+class TestModel:
27+ 
28+ def run(self):
29+ torch.npu.set_device(0)
30+ device = torch.device(f'npu:{0}')
31+ 
32+ dataset = RandomDataset(1000, 1000)
33+ 
34+ results = [None] * len(dataset)
35+ 
36+ threads = []
37+ batch_size = 64
38+ for i in range(0, len(dataset), batch_size):
39+ batch = dataset[i:i + batch_size]
40+ thread = threading.Thread(target=thread_worker, args=(batch, device, results, i))
41+ threads.append(thread)
42+ thread.start()
43+ 
44+ for thread in threads:
45+ thread.join()
46+ 
47+ 
48+class AllocatorMultiThreadProf(TestCase):
49+ test_model = TestModel()
50+ 
51+ def test_model_run_succ(self):
52+ res = True
53+ try:
54+ self.test_model.run()
55+ except Exception as e:
56+ res = False
57+ self.assertEqual(res, True)
58+ 
59+ 
60+if __name__ == "__main__":
61+ run_tests()
Mtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp+61-98
@@ -22,8 +22,8 @@
22#include "NPUBlockHandle.h"22#include "NPUBlockHandle.h"
23#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"23#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
24#include "torch_npu/csrc/core/npu/NPUEvent.h"24#include "torch_npu/csrc/core/npu/NPUEvent.h"
25-#ifndef BUILD_LIBTORCH
26#include "torch_npu/csrc/profiler/npu_profiler.h"25#include "torch_npu/csrc/profiler/npu_profiler.h"
26+#ifndef BUILD_LIBTORCH
27#include "torch_npu/csrc/sanitizer/NPUTrace.h"27#include "torch_npu/csrc/sanitizer/NPUTrace.h"
28#endif28#endif
29 29 
@@ -909,7 +909,8 @@ class DeviceCachingAllocator {
909 // All public methods (except the above) acquire the allocator mutex.909 // All public methods (except the above) acquire the allocator mutex.
910 // Thus, do not call a public method from another public method.910 // Thus, do not call a public method from another public method.
911 911 
912- Block* malloc(int device, size_t orig_size, aclrtStream stream) {912+ Block* malloc(int device, size_t orig_size, aclrtStream stream, uint8_t allocator_type = 0)
913+ {
913 // done outside the lock because we don't know what locks the recorder needs914 // done outside the lock because we don't know what locks the recorder needs
914 // to have...915 // to have...
915 auto context = maybeGatherContext(RecordContext::STATE);916 auto context = maybeGatherContext(RecordContext::STATE);
@@ -1059,14 +1060,16 @@ class DeviceCachingAllocator {
1059 1060 
1060 bool split_remainder = should_split(params.block, params.size());1061 bool split_remainder = should_split(params.block, params.size());
1061 return alloc_found_block(1062 return alloc_found_block(
1062- std::move(params), orig_size, std::move(context), split_remainder);1063+ std::move(params), orig_size, std::move(context), split_remainder, allocator_type);
1063 }1064 }
1064 1065 
1065 Block* alloc_found_block(1066 Block* alloc_found_block(
1066 AllocParams params,1067 AllocParams params,
1067 size_t orig_size,1068 size_t orig_size,
1068 std::shared_ptr<c10::GatheredContext> context,1069 std::shared_ptr<c10::GatheredContext> context,
1069- bool split_remainder) {1070+ bool split_remainder,
1071+ uint8_t allocator_type)
1072+ {
1070 auto size = params.size();1073 auto size = params.size();
1071 auto device = params.device();1074 auto device = params.device();
1072 auto pool = params.pool;1075 auto pool = params.pool;
@@ -1158,11 +1161,27 @@ class DeviceCachingAllocator {
1158 stats.reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,1161 stats.reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1159 stats.allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current);1162 stats.allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current);
1160 1163 
1164+#ifndef BUILD_LIBTORCH
1165+ torch_npu::profiler::reportMemoryDataToNpuProfiler({
1166+ static_cast<int8_t>(c10::DeviceType::PrivateUse1),
1167+ block->device,
1168+ static_cast<uint8_t>(torch_npu::profiler::MemoryDataType::MEMORY_MALLOC),
1169+ allocator_type,
1170+ reinterpret_cast<int64_t>(block->ptr),
1171+ block->size,
1172+ stats.allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1173+ stats.reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1174+ stats.active_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1175+ reinterpret_cast<int64_t>(block->stream)}
1176+ );
1177+#endif
1178+ 
1161 return block;1179 return block;
1162}1180}
1163 1181 
1164 1182 
1165- void free(Block* block) {1183+ void free(Block* block, uint8_t allocator_type = 0)
1184+ {
1166 std::shared_ptr<c10::GatheredContext> context =1185 std::shared_ptr<c10::GatheredContext> context =
1167 maybeGatherContext(RecordContext::ALL);1186 maybeGatherContext(RecordContext::ALL);
1168 std::lock_guard<std::recursive_mutex> lock(mutex);1187 std::lock_guard<std::recursive_mutex> lock(mutex);
@@ -1194,13 +1213,27 @@ class DeviceCachingAllocator {
1194 if (!block->stream_uses.empty() && c10_npu::NpuSysCtrl::GetInstance().GetInitFlag()) {1213 if (!block->stream_uses.empty() && c10_npu::NpuSysCtrl::GetInstance().GetInitFlag()) {
1195 insert_events(block);1214 insert_events(block);
1196 } else {1215 } else {
1197- free_block(block, context);1216+ free_block(block, context, allocator_type);
1198 }1217 }
1199 1218 
1200 ASCEND_LOGD("PTA CachingAllocator free: free = %zu, cached = %lu, allocated = %lu",1219 ASCEND_LOGD("PTA CachingAllocator free: free = %zu, cached = %lu, allocated = %lu",
1201 orig_block_size,1220 orig_block_size,
1202 stats.reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,1221 stats.reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1203 stats.allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current);1222 stats.allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current);
1223+#ifndef BUILD_LIBTORCH
1224+ torch_npu::profiler::reportMemoryDataToNpuProfiler({
1225+ static_cast<int8_t>(c10::DeviceType::PrivateUse1),
1226+ block->device,
1227+ static_cast<uint8_t>(torch_npu::profiler::MemoryDataType::MEMORY_FREE),
1228+ allocator_type,
1229+ reinterpret_cast<int64_t>(orig_block_ptr),
1230+ -orig_block_size,
1231+ stats.allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1232+ stats.reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1233+ stats.active_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1234+ reinterpret_cast<int64_t>(block->stream)}
1235+ );
1236+#endif
1204 }1237 }
1205 1238 
1206 void* getBaseAllocation(Block* block, size_t* outSize) {1239 void* getBaseAllocation(Block* block, size_t* outSize) {
@@ -1398,11 +1431,6 @@ class DeviceCachingAllocator {
1398 }1431 }
1399 }1432 }
1400 1433 
1401- DeviceStats get_stats()
1402- {
1403- return stats;
1404- }
1405- 
1406 private:1434 private:
1407 1435 
1408 // All private methods do not acquire the allocator mutex.1436 // All private methods do not acquire the allocator mutex.
@@ -1563,7 +1591,8 @@ class DeviceCachingAllocator {
1563 /** moves a block into a pool of cached free blocks **/1591 /** moves a block into a pool of cached free blocks **/
1564 void free_block(1592 void free_block(
1565 Block* block,1593 Block* block,
1566- const std::shared_ptr<c10::GatheredContext>& context)1594+ const std::shared_ptr<c10::GatheredContext>& context,
1595+ uint8_t allocator_type = 0)
1567 {1596 {
1568 AT_ASSERT(!block->allocated && block->event_count == 0, PTA_ERROR(ErrCode::VALUE));1597 AT_ASSERT(!block->allocated && block->event_count == 0, PTA_ERROR(ErrCode::VALUE));
1569 1598 
@@ -1577,6 +1606,7 @@ class DeviceCachingAllocator {
1577 1606 
1578 block->context_when_allocated = nullptr;1607 block->context_when_allocated = nullptr;
1579 size_t original_block_size = block->size;1608 size_t original_block_size = block->size;
1609+ auto orig_block_ptr = block->ptr;
1580 size_t requested_size = block->requested_size;1610 size_t requested_size = block->requested_size;
1581 1611 
1582 auto& pool = *block->pool;1612 auto& pool = *block->pool;
@@ -1621,6 +1651,20 @@ class DeviceCachingAllocator {
1621 stats.requested_bytes[stat_type],1651 stats.requested_bytes[stat_type],
1622 -static_cast<std::int64_t>(requested_size));1652 -static_cast<std::int64_t>(requested_size));
1623 });1653 });
1654+#ifndef BUILD_LIBTORCH
1655+ torch_npu::profiler::reportMemoryDataToNpuProfiler({
1656+ static_cast<int8_t>(c10::DeviceType::PrivateUse1),
1657+ block->device,
1658+ static_cast<uint8_t>(torch_npu::profiler::MemoryDataType::MEMORY_BLOCK_FREE),
1659+ allocator_type,
1660+ reinterpret_cast<int64_t>(orig_block_ptr),
1661+ -original_block_size,
1662+ stats.allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1663+ stats.reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1664+ stats.active_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
1665+ reinterpret_cast<int64_t>(block->stream)}
1666+ );
1667+#endif
1624 }1668 }
1625 1669 
1626 /** combine previously split blocks. returns the size of the subsumed block, or 0 on failure. **/1670 /** combine previously split blocks. returns the size of the subsumed block, or 0 on failure. **/
@@ -2281,20 +2325,7 @@ class NpuCachingAllocator : public NPUAllocator {
2281 "Allocator not initialized for device ", device, ": did you call init?",2325 "Allocator not initialized for device ", device, ": did you call init?",
2282 PTA_ERROR(ErrCode::PARAM));2326 PTA_ERROR(ErrCode::PARAM));
2283 Block* block = device_allocator[device]->malloc(device, size, stream);2327 Block* block = device_allocator[device]->malloc(device, size, stream);
2284-#ifndef BUILD_LIBTORCH2328+ 
2285- torch_npu::profiler::reportMemoryDataToNpuProfiler({
2286- static_cast<int8_t>(c10::DeviceType::PrivateUse1),
2287- block->device,
2288- static_cast<uint8_t>(torch_npu::profiler::MemoryDataType::MEMORY_MALLOC),
2289- static_cast<uint8_t>(torch_npu::profiler::MemoryAllocatorType::ALLOCATOR_INNER),
2290- reinterpret_cast<int64_t>(block->ptr),
2291- block->size,
2292- device_allocator[device]->get_stats().allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2293- device_allocator[device]->get_stats().reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2294- device_allocator[device]->get_stats().active_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2295- reinterpret_cast<int64_t>(block->stream)}
2296- );
2297-#endif
2298 add_allocated_block(block);2329 add_allocated_block(block);
2299 *devPtr = static_cast<void*>(block->ptr);2330 *devPtr = static_cast<void*>(block->ptr);
2300#ifndef BUILD_LIBTORCH2331#ifndef BUILD_LIBTORCH
@@ -2324,34 +2355,6 @@ class NpuCachingAllocator : public NPUAllocator {
2324 auto orig_block_ptr = block->ptr;2355 auto orig_block_ptr = block->ptr;
2325 auto orig_block_size = block->size;2356 auto orig_block_size = block->size;
2326 device_allocator[block->device]->free(block);2357 device_allocator[block->device]->free(block);
2327-#ifndef BUILD_LIBTORCH
2328- if (block->stream_uses.empty() || !c10_npu::NpuSysCtrl::GetInstance().GetInitFlag()) {
2329- torch_npu::profiler::reportMemoryDataToNpuProfiler({
2330- static_cast<int8_t>(c10::DeviceType::PrivateUse1),
2331- block->device,
2332- static_cast<uint8_t>(torch_npu::profiler::MemoryDataType::MEMORY_BLOCK_FREE),
2333- static_cast<uint8_t>(torch_npu::profiler::MemoryAllocatorType::ALLOCATOR_INNER),
2334- reinterpret_cast<int64_t>(orig_block_ptr),
2335- -orig_block_size,
2336- device_allocator[block->device]->get_stats().allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2337- device_allocator[block->device]->get_stats().reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2338- device_allocator[block->device]->get_stats().active_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2339- reinterpret_cast<int64_t>(block->stream)}
2340- );
2341- }
2342- torch_npu::profiler::reportMemoryDataToNpuProfiler({
2343- static_cast<int8_t>(c10::DeviceType::PrivateUse1),
2344- block->device,
2345- static_cast<uint8_t>(torch_npu::profiler::MemoryDataType::MEMORY_FREE),
2346- static_cast<uint8_t>(torch_npu::profiler::MemoryAllocatorType::ALLOCATOR_INNER),
2347- reinterpret_cast<int64_t>(orig_block_ptr),
2348- -orig_block_size,
2349- device_allocator[block->device]->get_stats().allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2350- device_allocator[block->device]->get_stats().reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2351- device_allocator[block->device]->get_stats().active_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2352- reinterpret_cast<int64_t>(block->stream)}
2353- );
2354-#endif
2355 }2358 }
2356 2359 
2357 void setMemoryFraction(double fraction, int device) override2360 void setMemoryFraction(double fraction, int device) override
@@ -2641,22 +2644,9 @@ void* MallocBlock(size_t size, void *stream, int device) {
2641 }2644 }
2642 AT_ASSERT(caching_allocator.device_allocator[device], PTA_ERROR(ErrCode::NOT_FOUND));2645 AT_ASSERT(caching_allocator.device_allocator[device], PTA_ERROR(ErrCode::NOT_FOUND));
2643 AT_ASSERT(stream, PTA_ERROR(ErrCode::NOT_FOUND));2646 AT_ASSERT(stream, PTA_ERROR(ErrCode::NOT_FOUND));
2644- auto block = caching_allocator.device_allocator[device]->malloc(device, size, stream);2647+ auto block = caching_allocator.device_allocator[device]->malloc(device, size, stream,
2648+ static_cast<uint8_t>(torch_npu::profiler::MemoryAllocatorType::ALLOCATOR_EXTERNAL));
2645 AT_ASSERT(block, PTA_ERROR(ErrCode::NOT_FOUND));2649 AT_ASSERT(block, PTA_ERROR(ErrCode::NOT_FOUND));
2646-#ifndef BUILD_LIBTORCH
2647- torch_npu::profiler::reportMemoryDataToNpuProfiler({
2648- static_cast<int8_t>(c10::DeviceType::PrivateUse1),
2649- block->device,
2650- static_cast<uint8_t>(torch_npu::profiler::MemoryDataType::MEMORY_MALLOC),
2651- static_cast<uint8_t>(torch_npu::profiler::MemoryAllocatorType::ALLOCATOR_EXTERNAL),
2652- reinterpret_cast<int64_t>(block->ptr),
2653- block->size,
2654- caching_allocator.device_allocator[device]->get_stats().allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2655- caching_allocator.device_allocator[device]->get_stats().reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2656- caching_allocator.device_allocator[device]->get_stats().active_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2657- reinterpret_cast<int64_t>(block->stream)}
2658- );
2659-#endif
2660 return reinterpret_cast<void*>(block);2650 return reinterpret_cast<void*>(block);
2661}2651}
2662 2652 
@@ -2667,35 +2657,8 @@ void FreeBlock(void *handle) {
2667 AT_ASSERT(caching_allocator.device_allocator[block->device], PTA_ERROR(ErrCode::NOT_FOUND));2657 AT_ASSERT(caching_allocator.device_allocator[block->device], PTA_ERROR(ErrCode::NOT_FOUND));
2668 auto orig_block_ptr = block->ptr;2658 auto orig_block_ptr = block->ptr;
2669 auto orig_block_size = block->size;2659 auto orig_block_size = block->size;
2670- caching_allocator.device_allocator[block->device]->free(block);2660+ caching_allocator.device_allocator[block->device]->free(block,
2671-#ifndef BUILD_LIBTORCH2661+ static_cast<uint8_t>(torch_npu::profiler::MemoryAllocatorType::ALLOCATOR_EXTERNAL));
2672- if (block->stream_uses.empty() || !c10_npu::NpuSysCtrl::GetInstance().GetInitFlag()) {
2673- torch_npu::profiler::reportMemoryDataToNpuProfiler({
2674- static_cast<int8_t>(c10::DeviceType::PrivateUse1),
2675- block->device,
2676- static_cast<uint8_t>(torch_npu::profiler::MemoryDataType::MEMORY_BLOCK_FREE),
2677- static_cast<uint8_t>(torch_npu::profiler::MemoryAllocatorType::ALLOCATOR_EXTERNAL),
2678- reinterpret_cast<int64_t>(orig_block_ptr),
2679- -orig_block_size,
2680- caching_allocator.device_allocator[block->device]->get_stats().allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2681- caching_allocator.device_allocator[block->device]->get_stats().reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2682- caching_allocator.device_allocator[block->device]->get_stats().active_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2683- reinterpret_cast<int64_t>(block->stream)}
2684- );
2685- }
2686- torch_npu::profiler::reportMemoryDataToNpuProfiler({
2687- static_cast<int8_t>(c10::DeviceType::PrivateUse1),
2688- block->device,
2689- static_cast<uint8_t>(torch_npu::profiler::MemoryDataType::MEMORY_FREE),
2690- static_cast<uint8_t>(torch_npu::profiler::MemoryAllocatorType::ALLOCATOR_EXTERNAL),
2691- reinterpret_cast<int64_t>(orig_block_ptr),
2692- -orig_block_size,
2693- caching_allocator.device_allocator[block->device]->get_stats().allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2694- caching_allocator.device_allocator[block->device]->get_stats().reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2695- caching_allocator.device_allocator[block->device]->get_stats().active_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,
2696- reinterpret_cast<int64_t>(block->stream)}
2697- );
2698-#endif
2699}2662}
2700 2663 
2701void* GetBlockPtr(const void *handle) {2664void* GetBlockPtr(const void *handle) {