HOST内存统计API对齐社区-pytorch-AtomGit

Rtest/allocator/host/test_host_allocator.py→test/allocator/host/test_expandable_host_allocator.py+38-32

	@@ -8,6 +8,9 @@ from torch_npu.testing.testcase import TestCase, run_tests
8	from torch_npu.testing.common_utils import create_common_tensor	8	from torch_npu.testing.common_utils import create_common_tensor
9	from torch_npu.npu.utils import get_cann_version, _is_gte_cann_version	9	from torch_npu.npu.utils import get_cann_version, _is_gte_cann_version
10		10
		11	+RESERVED_BYTES_CURRENT = "reserved_bytes.current"
		12	+ALLOCATED_BYTES_CURRENT = "allocated_bytes.current"
		13	+
11		14
12	def get_skip():	15	def get_skip():
13	try:	16	try:
	@@ -47,26 +50,26 @@ class TestHostCachingAllocator(TestCase):
47	# 申请一个64M的tensor, 会预选绑定80M物理内存	50	# 申请一个64M的tensor, 会预选绑定80M物理内存
48	memory_64m = torch.ones([1024, 1024, 16]).pin_memory()	51	memory_64m = torch.ones([1024, 1024, 16]).pin_memory()
49	# 实际分配的内存是申请大小 = 512 * (size + 32 + 512 - 1) / 512	52	# 实际分配的内存是申请大小 = 512 * (size + 32 + 512 - 1) / 512
50	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), 64 * 1024 * 1024 + 512)	53	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 64 * 1024 * 1024 + 512)
51	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 80 * 1024 * 1024)	54	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 80 * 1024 * 1024)
52		55
53	# 释放64M内存，内存归还内存池，物理内存不变，使用量为0	56	# 释放64M内存，内存归还内存池，物理内存不变，使用量为0
54	memory_64m = None	57	memory_64m = None
55	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 80 * 1024 * 1024)	58	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 80 * 1024 * 1024)
56	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), 0)	59	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 0)
57	# 再次申请4M/8M/16M大小的内存后，总的预留内存应该不变	60	# 再次申请4M/8M/16M大小的内存后，总的预留内存应该不变
58	memory_4m = torch.ones([1024, 1024, 1]).pin_memory()	61	memory_4m = torch.ones([1024, 1024, 1]).pin_memory()
59	memory_8m = torch.ones([1024, 1024, 2]).pin_memory()	62	memory_8m = torch.ones([1024, 1024, 2]).pin_memory()
60	memory_16m = torch.ones([1024, 1024, 4]).pin_memory()	63	memory_16m = torch.ones([1024, 1024, 4]).pin_memory()
61	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), (4 + 8 + 16) * 1024 * 1024 + 512 * 3)	64	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), (4 + 8 + 16) * 1024 * 1024 + 512 * 3)
62	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 80 * 1024 * 1024)	65	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 80 * 1024 * 1024)
63	- # 内存块归还内存池, 执行empty_pin_memory_cache释放所有物理内存防止用例间相互影响	66	+ # 内存块归还内存池, 执行host_empty_cache释放所有物理内存防止用例间相互影响
64	memory_4m = None	67	memory_4m = None
65	memory_8m = None	68	memory_8m = None
66	memory_16m = None	69	memory_16m = None
67	- torch_npu.npu.empty_pin_memory_cache()	70	+ torch_npu.npu.host_empty_cache()
68	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), 0)	71	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 0)
69	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 0)	72	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 0)
70		73
71		74
72	def test_allocate_with_block_merge(self):	75	def test_allocate_with_block_merge(self):
	@@ -75,20 +78,20 @@ class TestHostCachingAllocator(TestCase):
75	memory_40m2 = torch.ones([1024, 1024, 10]).pin_memory()	78	memory_40m2 = torch.ones([1024, 1024, 10]).pin_memory()
76	memory_40m1 = None	79	memory_40m1 = None
77	memory_40m2 = None	80	memory_40m2 = None
78	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 100 * 1024 * 1024)	81	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 100 * 1024 * 1024)
79	# 再申请一个64M的tensor, 不会再申请物理内存	82	# 再申请一个64M的tensor, 不会再申请物理内存
80	memory_64m = torch.ones([1024, 1024, 16], pin_memory=True).pin_memory()	83	memory_64m = torch.ones([1024, 1024, 16], pin_memory=True).pin_memory()
81	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 100 * 1024 * 1024)	84	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 100 * 1024 * 1024)
82	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), 64 * 1024 * 1024 + 512)	85	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 64 * 1024 * 1024 + 512)
83	# 释放空闲block的物理内存, 应该只能释放20M，已分配内存不变	86	# 释放空闲block的物理内存, 应该只能释放20M，已分配内存不变
84	- torch_npu.npu.empty_pin_memory_cache()	87	+ torch_npu.npu.host_empty_cache()
85	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), 64 * 1024 * 1024 + 512)	88	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 64 * 1024 * 1024 + 512)
86	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 80 * 1024 * 1024)	89	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 80 * 1024 * 1024)
87	# 释放所有内存防止用例间相互影响	90	# 释放所有内存防止用例间相互影响
88	memory_64m = None	91	memory_64m = None
89	- torch_npu.npu.empty_pin_memory_cache()	92	+ torch_npu.npu.host_empty_cache()
90	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), 0)	93	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 0)
91	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 0)	94	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 0)
92		95
93		96
94	def test_h2d_inplace(self):	97	def test_h2d_inplace(self):
	@@ -134,14 +137,14 @@ class TestHostCachingAllocator(TestCase):
134	tensor_cpu.copy_(tensor, non_blocking=True)	137	tensor_cpu.copy_(tensor, non_blocking=True)
135	tensor_cpu = None	138	tensor_cpu = None
136		139
137	- torch_npu.npu.empty_pin_memory_cache()	140	+ torch_npu.npu.host_empty_cache()
138	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), 64 * 1024 * 1024 + 512)	141	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 64 * 1024 * 1024 + 512)
139	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 80 * 1024 * 1024)	142	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 80 * 1024 * 1024)
140		143
141	torch.npu.synchronize()	144	torch.npu.synchronize()
142	- torch_npu.npu.empty_pin_memory_cache()	145	+ torch_npu.npu.host_empty_cache()
143	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), 0)	146	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 0)
144	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 0)	147	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 0)
145		148
146		149
147	def test_muti_stream(self):	150	def test_muti_stream(self):
	@@ -169,20 +172,23 @@ class TestHostCachingAllocator(TestCase):
169	self.assertRtolEqual(tensor2_cpu.numpy(), (torch.ones([100, 100, i * 10]).cpu() * 22).numpy())	172	self.assertRtolEqual(tensor2_cpu.numpy(), (torch.ones([100, 100, i * 10]).cpu() * 22).numpy())
170	tensor1_cpu = None	173	tensor1_cpu = None
171	tensor2_cpu = None	174	tensor2_cpu = None
		175	+ tensor = None
		176	+ tensor1 = None
		177	+ tensor2 = None
172		178
173	torch.npu.synchronize()	179	torch.npu.synchronize()
174	- torch_npu.npu.empty_pin_memory_cache()	180	+ torch_npu.npu.host_empty_cache()
175	- self.assertEqual(torch_npu.npu.pin_memory_allocated(), 0)	181	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 0)
176	- self.assertEqual(torch_npu.npu.pin_memory_reserved(), 0)	182	+ self.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 0)
177		183
178		184
179	@classmethod	185	@classmethod
180	def tearDownClass(cls):	186	def tearDownClass(cls):
181	instance = cls()	187	instance = cls()
182	- instance.assertEqual(torch_npu.npu.pin_memory_stats().get("allocated_bytes.peak", 0), 80 * 1024 * 1024 + 512 * 2)	188	+ instance.assertEqual(torch_npu.npu.host_memory_stats().get("allocated_bytes.peak", 0), 80 * 1024 * 1024 + 512 * 2)
183	- instance.assertEqual(torch_npu.npu.pin_memory_stats().get("allocated_bytes.current", 0), 0)	189	+ instance.assertEqual(torch_npu.npu.host_memory_stats().get(ALLOCATED_BYTES_CURRENT, 0), 0)
184	- instance.assertEqual(torch_npu.npu.pin_memory_stats().get("reserved_bytes.peak", 0), 100 * 1024 * 1024)	190	+ instance.assertEqual(torch_npu.npu.host_memory_stats().get("reserved_bytes.peak", 0), 100 * 1024 * 1024)
185	- instance.assertEqual(torch_npu.npu.pin_memory_stats().get("reserved_bytes.current", 0), 0)	191	+ instance.assertEqual(torch_npu.npu.host_memory_stats().get(RESERVED_BYTES_CURRENT, 0), 0)
186		192
187		193
188	if __name__ == '__main__':	194	if __name__ == '__main__':

Atest/allocator/host/test_unexpandable_host_allocator.py+151-0

@@ -0,0 +1,151 @@
	1	+import gc
	2	+
	3	+import torch
	4	+import torch_npu
	5	+from torch_npu.testing.testcase import TestCase, run_tests
	6	+
	7	+
	8	+class TestHostCachingAllocator(TestCase):
	9	+ def test_host_memory_stats(self):
	10	+ # Helper functions
	11	+ def empty_stats():
	12	+ return {
	13	+ "allocated_bytes.allocated": 0,
	14	+ "allocated_bytes.current": 0,
	15	+ "allocated_bytes.freed": 0,
	16	+ "allocated_bytes.peak": 0,
	17	+ "allocation.allocated": 0,
	18	+ "allocation.current": 0,
	19	+ "allocation.freed": 0,
	20	+ "allocation.peak": 0,
	21	+ "host_alloc_time.count": 0,
	22	+ "host_free_time.count": 0,
	23	+ "num_host_alloc": 0,
	24	+ "num_host_free": 0,
	25	+ "reserved_bytes.allocated": 0,
	26	+ "reserved_bytes.current": 0,
	27	+ "reserved_bytes.freed": 0,
	28	+ "reserved_bytes.peak": 0,
	29	+ "segment.allocated": 0,
	30	+ "segment.current": 0,
	31	+ "segment.freed": 0,
	32	+ "segment.peak": 0,
	33	+ }
	34	+
	35	+ def check_stats(expected):
	36	+ stats = torch_npu.npu.host_memory_stats()
	37	+ for k, v in expected.items():
	38	+ self.assertEqual(v, stats[k])
	39	+
	40	+ # Setup the test cleanly
	41	+ alloc1 = 10
	42	+ alloc1_aligned = 16
	43	+ alloc2 = 20
	44	+ alloc2_aligned = 32
	45	+ expected = empty_stats()
	46	+
	47	+ # Reset any lingering state
	48	+ gc.collect()
	49	+ torch_npu.npu.host_empty_cache()
	50	+
	51	+ # Check that stats are empty
	52	+ check_stats(expected)
	53	+
	54	+ # Make first allocation and check stats
	55	+ t1 = torch.ones(alloc1 * 1024, pin_memory=True)
	56	+ self.assertTrue(t1.is_pinned())
	57	+ for prefix in ["segment", "allocation"]:
	58	+ for suffix in ["allocated", "current", "peak"]:
	59	+ expected[prefix + "." + suffix] += 1
	60	+
	61	+ allocation_size1 = alloc1_aligned * 1024 * 4
	62	+ for prefix in ["allocated_bytes", "reserved_bytes"]:
	63	+ for suffix in ["allocated", "current", "peak"]:
	64	+ expected[prefix + "." + suffix] += allocation_size1
	65	+
	66	+ expected["num_host_alloc"] += 1
	67	+ expected["host_alloc_time.count"] += 1
	68	+
	69	+ check_stats(expected)
	70	+
	71	+ # Remove first allocation and check stats
	72	+ del t1
	73	+
	74	+ expected["allocation.current"] -= 1
	75	+ expected["allocation.freed"] += 1
	76	+ expected["allocated_bytes.current"] -= allocation_size1
	77	+ expected["allocated_bytes.freed"] += allocation_size1
	78	+
	79	+ check_stats(expected)
	80	+
	81	+ # Make first allocation again and check reuse
	82	+ t1 = torch.ones(alloc1 * 1024, pin_memory=True)
	83	+ self.assertTrue(t1.is_pinned())
	84	+ for suffix in ["allocated", "current"]:
	85	+ expected["allocation" + "." + suffix] += 1
	86	+
	87	+ allocation_size1 = alloc1_aligned * 1024 * 4
	88	+ for suffix in ["allocated", "current"]:
	89	+ expected["allocated_bytes" + "." + suffix] += allocation_size1
	90	+
	91	+ check_stats(expected)
	92	+
	93	+ # Make second allocation and check stats
	94	+ t2 = torch.ones(alloc2 * 1024, pin_memory=True)
	95	+ self.assertTrue(t2.is_pinned())
	96	+ for prefix in ["segment", "allocation"]:
	97	+ for suffix in ["allocated", "current", "peak"]:
	98	+ expected[prefix + "." + suffix] += 1
	99	+
	100	+ allocation_size2 = alloc2_aligned * 1024 * 4
	101	+ for prefix in ["allocated_bytes", "reserved_bytes"]:
	102	+ for suffix in ["allocated", "current", "peak"]:
	103	+ expected[prefix + "." + suffix] += allocation_size2
	104	+
	105	+ expected["num_host_alloc"] += 1
	106	+ expected["host_alloc_time.count"] += 1
	107	+
	108	+ check_stats(expected)
	109	+
	110	+ # Remove first allocation and check stats
	111	+ del t1
	112	+
	113	+ expected["allocation.current"] -= 1
	114	+ expected["allocation.freed"] += 1
	115	+ expected["allocated_bytes.current"] -= allocation_size1
	116	+ expected["allocated_bytes.freed"] += allocation_size1
	117	+
	118	+ check_stats(expected)
	119	+
	120	+ # Remove second allocation and check stats
	121	+ del t2
	122	+
	123	+ expected["allocation.current"] -= 1
	124	+ expected["allocation.freed"] += 1
	125	+ expected["allocated_bytes.current"] -= allocation_size2
	126	+ expected["allocated_bytes.freed"] += allocation_size2
	127	+
	128	+ check_stats(expected)
	129	+
	130	+ # Empty cache and check stats
	131	+ torch_npu.npu.host_empty_cache()
	132	+ expected["segment.freed"] += expected["segment.current"]
	133	+ expected["segment.current"] = 0
	134	+ expected["reserved_bytes.freed"] += expected["reserved_bytes.current"]
	135	+ expected["reserved_bytes.current"] = 0
	136	+ expected["num_host_free"] = expected["num_host_alloc"]
	137	+ expected["host_free_time.count"] += expected["host_alloc_time.count"]
	138	+
	139	+ check_stats(expected)
	140	+
	141	+ # Finally, check the reset of peak and accumulated stats
	142	+ torch_npu.npu.reset_peak_host_memory_stats()
	143	+ torch_npu.npu.reset_accumulated_host_memory_stats()
	144	+
	145	+ expected = empty_stats()
	146	+
	147	+ check_stats(expected)
	148	+
	149	+
	150	+if __name__ == '__main__':
	151	+ run_tests()

Mtest/torch_npu_schema.json+14-8

	@@ -1094,16 +1094,19 @@
1094	"torch_npu.npu.empty_virt_addr_cache": {	1094	"torch_npu.npu.empty_virt_addr_cache": {
1095	"signature": "()"	1095	"signature": "()"
1096	},	1096	},
1097	- "torch_npu.npu.empty_pin_memory_cache": {	1097	+ "torch_npu.npu.host_empty_cache": {
1098	"signature": "()"	1098	"signature": "()"
1099	},	1099	},
1100	- "torch_npu.npu.pin_memory_stats": {	1100	+ "torch_npu.npu.host_memory_stats": {
1101	"signature": "()"	1101	"signature": "()"
1102	},	1102	},
1103	- "torch_npu.npu.pin_memory_allocated": {	1103	+ "torch_npu.npu.host_memory_stats_as_nested_dict": {
1104	"signature": "()"	1104	"signature": "()"
1105	},	1105	},
1106	- "torch_npu.npu.pin_memory_reserved": {	1106	+ "torch_npu.npu.reset_accumulated_host_memory_stats": {
		1107	+ "signature": "()"
		1108	+ },
		1109	+ "torch_npu.npu.reset_peak_host_memory_stats": {
1107	"signature": "()"	1110	"signature": "()"
1108	},	1111	},
1109	"torch_npu.npu.enable_deterministic_with_backward": {	1112	"torch_npu.npu.enable_deterministic_with_backward": {
	@@ -1427,16 +1430,19 @@
1427	"torch_npu.npu.memory.empty_virt_addr_cache": {	1430	"torch_npu.npu.memory.empty_virt_addr_cache": {
1428	"signature": "()"	1431	"signature": "()"
1429	},	1432	},
1430	- "torch_npu.npu.memory.empty_pin_memory_cache": {	1433	+ "torch_npu.npu.memory.host_empty_cache": {
1431	"signature": "()"	1434	"signature": "()"
1432	},	1435	},
1433	- "torch_npu.npu.memory.pin_memory_stats": {	1436	+ "torch_npu.npu.memory.host_memory_stats": {
1434	"signature": "()"	1437	"signature": "()"
1435	},	1438	},
1436	- "torch_npu.npu.memory.pin_memory_allocated": {	1439	+ "torch_npu.npu.memory.host_memory_stats_as_nested_dict": {
1437	"signature": "()"	1440	"signature": "()"
1438	},	1441	},
1439	- "torch_npu.npu.memory.pin_memory_reserved": {	1442	+ "torch_npu.npu.memory.reset_accumulated_host_memory_stats": {
		1443	+ "signature": "()"
		1444	+ },
		1445	+ "torch_npu.npu.memory.reset_peak_host_memory_stats": {
1440	"signature": "()"	1446	"signature": "()"
1441	},	1447	},
1442	"torch_npu.npu.memory.get_allocator_backend": {	1448	"torch_npu.npu.memory.get_allocator_backend": {

Mtorch_npu/csrc/core/npu/CachingHostAllocator.cpp+80-17

	@@ -41,12 +41,6 @@ constexpr size_t kMinBlockSize = 512; // all sizes are rounded t
41	constexpr size_t segmentSize = 20971520; // segment size 20M	41	constexpr size_t segmentSize = 20971520; // segment size 20M
42	constexpr size_t deviceTotalSize = 68719476736; // 64 GB	42	constexpr size_t deviceTotalSize = 68719476736; // 64 GB
43		43
44	-void update_stat(c10::CachingAllocator::Stat &stat, int64_t amount)
45	-{
46	- stat.current += amount;
47	- stat.peak = std::max(stat.current, stat.peak);
48	-}
49	-
50	struct ExpandableBlock;	44	struct ExpandableBlock;
51	using Comparison = bool ()(const ExpandableBlock , const ExpandableBlock *);	45	using Comparison = bool ()(const ExpandableBlock , const ExpandableBlock *);
52	static bool BlockComparatorSize(const ExpandableBlock a, const ExpandableBlock b);	46	static bool BlockComparatorSize(const ExpandableBlock a, const ExpandableBlock b);
	@@ -189,6 +183,16 @@ public:
189	return c10_npu::NPUCachingAllocator::CachingAllocatorConfig::pinned_use_background_threads();	183	return c10_npu::NPUCachingAllocator::CachingAllocatorConfig::pinned_use_background_threads();
190	}	184	}
191		185
		186	+ virtual void resetHostAccumulatedStats()
		187	+ {
		188	+ resetAccumulatedStats();
		189	+ }
		190	+
		191	+ virtual void resetHostPeakStats()
		192	+ {
		193	+ resetPeakStats();
		194	+ }
		195	+
192	private:	196	private:
193	void allocate_host_memory(size_t size, void** ptr) override	197	void allocate_host_memory(size_t size, void** ptr) override
194	{	198	{
	@@ -485,7 +489,7 @@ public:
485	if (params.err == ACL_ERROR_RT_MEMORY_ALLOCATION) {	489	if (params.err == ACL_ERROR_RT_MEMORY_ALLOCATION) {
486	AT_ERROR("Malloc pin memory failed, host out of memory, Rried to allocate size: ", orig_size,	490	AT_ERROR("Malloc pin memory failed, host out of memory, Rried to allocate size: ", orig_size,
487	", total_allocated_memory: ", stats.allocated_bytes.current, ", total_reserved_memory: ", stats.reserved_bytes.current,	491	", total_allocated_memory: ", stats.allocated_bytes.current, ", total_reserved_memory: ", stats.reserved_bytes.current,
488	- ". You might be able to retry after freeing up memory using torch_npu.npu.empty_pin_memory_cache().");	492	+ ". You might be able to retry after freeing up memory using torch_npu.npu.host_empty_cache().");
489	}	493	}
490	return {nullptr, nullptr};	494	return {nullptr, nullptr};
491	}	495	}
	@@ -526,6 +530,26 @@ public:
526	return stats;	530	return stats;
527	}	531	}
528		532
		533	+ void resetHostAccumulatedStats() override
		534	+ {
		535	+ std::lock_guard<std::mutex> lock(stats_mutex_);
		536	+ stats.allocated_bytes.reset_accumulated();
		537	+ stats.reserved_bytes.reset_accumulated();
		538	+
		539	+ stats.host_alloc_time.reset_accumulated();
		540	+ stats.host_free_time.reset_accumulated();
		541	+ }
		542	+
		543	+ void resetHostPeakStats() override
		544	+ {
		545	+ std::lock_guard<std::mutex> lock(stats_mutex_);
		546	+ stats.allocated_bytes.reset_peak();
		547	+ stats.reserved_bytes.reset_peak();
		548	+
		549	+ stats.host_alloc_time.reset_peak();
		550	+ stats.host_free_time.reset_peak();
		551	+ }
		552	+
529	bool record_event(void* ptr, void* ctx, c10::Stream s) override	553	bool record_event(void* ptr, void* ctx, c10::Stream s) override
530	{	554	{
531	c10_npu::NPUStream stream = c10_npu::NPUStream(s);	555	c10_npu::NPUStream stream = c10_npu::NPUStream(s);
	@@ -563,6 +587,8 @@ private:
563		587
564	std::mutex mutex;	588	std::mutex mutex;
565		589
		590	+ std::mutex stats_mutex_;
		591	+
566	// allocated or in use by a stream	592	// allocated or in use by a stream
567	ska::flat_hash_map<void , ExpandableBlock > ptr_to_block_;	593	ska::flat_hash_map<void , ExpandableBlock > ptr_to_block_;
568		594
	@@ -649,7 +675,15 @@ private:
649		675
650	void unmap_block(ExpandableBlock *block)	676	void unmap_block(ExpandableBlock *block)
651	{	677	{
		678	+ auto start = std::chrono::steady_clock::now();
652	auto unmapped = block->expandable_segment_->unmap(SegmentRange{ block->ptr_, block->size_ });	679	auto unmapped = block->expandable_segment_->unmap(SegmentRange{ block->ptr_, block->size_ });
		680	+ auto end = std::chrono::steady_clock::now();
		681	+ auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
		682	+ // Update the statistics on the time spent on AclrtUnmapMem/AclrtFreePhysical
		683	+ {
		684	+ std::lock_guard<std::mutex> lock(stats_mutex_);
		685	+ stats.host_free_time.increase(duration.count());
		686	+ }
653	if (unmapped.size == 0) {	687	if (unmapped.size == 0) {
654	return;	688	return;
655	}	689	}
	@@ -676,8 +710,10 @@ private:
676	block->ptr_ = unmapped.ptr;	710	block->ptr_ = unmapped.ptr;
677	block->size_ = unmapped.size;	711	block->size_ = unmapped.size;
678	block->mapped = false;	712	block->mapped = false;
679	- update_stat(stats.reserved_bytes, -static_cast<std::int64_t>(unmapped.size));	713	+ {
680	-	714	+ std::lock_guard<std::mutex> lock(stats_mutex_);
		715	+ stats.reserved_bytes.decrease(unmapped.size);
		716	+ }
681	try_merge_blocks(block, block->prev, *block->pool);	717	try_merge_blocks(block, block->prev, *block->pool);
682	try_merge_blocks(block, block->next, *block->pool);	718	try_merge_blocks(block, block->next, *block->pool);
683	block->pool->unmapped.insert(block);	719	block->pool->unmapped.insert(block);
	@@ -734,7 +770,10 @@ private:
734	{	770	{
735	AT_ASSERT(!block->allocated_ && block->event_count_ == 0, PTA_ERROR(ErrCode::VALUE));	771	AT_ASSERT(!block->allocated_ && block->event_count_ == 0, PTA_ERROR(ErrCode::VALUE));
736		772
737	- update_stat(stats.allocated_bytes, -static_cast<std::int64_t>(block->size_));	773	+ {
		774	+ std::lock_guard<std::mutex> lock(stats_mutex_);
		775	+ stats.allocated_bytes.decrease(block->size_);
		776	+ }
738	auto &pool = *block->pool;	777	auto &pool = *block->pool;
739	const std::array<ExpandableBlock *, 2> merge_candidates = { block->prev, block->next };	778	const std::array<ExpandableBlock *, 2> merge_candidates = { block->prev, block->next };
740	for (ExpandableBlock *merge_candidate : merge_candidates) {	779	for (ExpandableBlock *merge_candidate : merge_candidates) {
	@@ -928,7 +967,10 @@ private:
928		967
929	try_merge_blocks(to_map, to_map->prev, pool);	968	try_merge_blocks(to_map, to_map->prev, pool);
930	try_merge_blocks(to_map, to_map->next, pool);	969	try_merge_blocks(to_map, to_map->next, pool);
931	- update_stat(stats.reserved_bytes, static_cast<std::int64_t>(mapped_range.size));	970	+ {
		971	+ std::lock_guard<std::mutex> lock(stats_mutex_);
		972	+ stats.reserved_bytes.increase(mapped_range.size);
		973	+ }
932	pool.blocks.insert(to_map);	974	pool.blocks.insert(to_map);
933	return true;	975	return true;
934	}	976	}
	@@ -941,6 +983,7 @@ private:
941	// unmapped -> free -> *	983	// unmapped -> free -> *
942	// free -> unmapped -> *	984	// free -> unmapped -> *
943		985
		986	+ auto start = std::chrono::steady_clock::now();
944	if (!candidate->mapped && !map_block(candidate, std::min(candidate->size_, size), pool)) {	987	if (!candidate->mapped && !map_block(candidate, std::min(candidate->size_, size), pool)) {
945	return nullptr;	988	return nullptr;
946	}	989	}
	@@ -960,6 +1003,13 @@ private:
960	candidate = new_candidate;	1003	candidate = new_candidate;
961	}	1004	}
962	pool->blocks.erase(candidate);	1005	pool->blocks.erase(candidate);
		1006	+ auto end = std::chrono::steady_clock::now();
		1007	+ auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
		1008	+ // Update the statistics on the time spent on AclrtMallocPhysical/AclrtMapMem
		1009	+ {
		1010	+ std::lock_guard<std::mutex> lock(stats_mutex_);
		1011	+ stats.host_alloc_time.increase(duration.count());
		1012	+ }
963	return candidate;	1013	return candidate;
964	}	1014	}
965		1015
	@@ -993,7 +1043,10 @@ private:
993	block->allocated_ = true;	1043	block->allocated_ = true;
994	block->requested_size = orig_size;	1044	block->requested_size = orig_size;
995	ptr_to_block_[block->ptr_] = block;	1045	ptr_to_block_[block->ptr_] = block;
996	- update_stat(stats.allocated_bytes, static_cast<std::int64_t>(block->size_));	1046	+ {
		1047	+ std::lock_guard<std::mutex> lock(stats_mutex_);
		1048	+ stats.allocated_bytes.increase(block->size_);
		1049	+ }
997	return block;	1050	return block;
998	}	1051	}
999		1052
	@@ -1049,6 +1102,21 @@ struct NPUCachingHostAllocator final
1049	impl_->empty_cache();	1102	impl_->empty_cache();
1050	}	1103	}
1051		1104
		1105	+ at::HostStats get_stats() override
		1106	+ {
		1107	+ return impl_->getHostStats();
		1108	+ }
		1109	+
		1110	+ void reset_accumulated_stats() override
		1111	+ {
		1112	+ impl_->resetHostAccumulatedStats();
		1113	+ }
		1114	+
		1115	+ void reset_peak_stats() override
		1116	+ {
		1117	+ impl_->resetHostPeakStats();
		1118	+ }
		1119	+
1052	void setAllocator()	1120	void setAllocator()
1053	{	1121	{
1054	std::call_once(init_config, [this] {	1122	std::call_once(init_config, [this] {
	@@ -1097,11 +1165,6 @@ bool ptr_exist(void* ptr)
1097	return getNPUCachingHostAllocator().impl_->ptr_check(ptr);	1165	return getNPUCachingHostAllocator().impl_->ptr_check(ptr);
1098	}	1166	}
1099		1167
1100	-at::HostStats CachingHostAllocator_getStats()
1101	-{
1102	- return getNPUCachingHostAllocator().impl_->getHostStats();
1103	-}
1104	-
1105	// the host memory is not allocated by malloc	1168	// the host memory is not allocated by malloc
1106	aclError process_unregistered_mem_location_type(c10_npu::NPUStream stream, aclrtMemcpyKind kind)	1169	aclError process_unregistered_mem_location_type(c10_npu::NPUStream stream, aclrtMemcpyKind kind)
1107	{	1170	{

Mtorch_npu/csrc/core/npu/CachingHostAllocator.h+0-2

	@@ -54,8 +54,6 @@ inline at::DataPtr HostAlloc(size_t size)
54	return at::getHostAllocator(at::kPrivateUse1)->allocate(size);	54	return at::getHostAllocator(at::kPrivateUse1)->allocate(size);
55	}	55	}
56		56
57	-TORCH_NPU_API at::HostStats CachingHostAllocator_getStats();
58	-
59	c10::Allocator* getPinnedMemoryAllocator();	57	c10::Allocator* getPinnedMemoryAllocator();
60		58
61	} // namespace at_npu::native	59	} // namespace at_npu::native

Mtorch_npu/csrc/core/npu/NPUAllocatorConfig.cpp+1-1

	@@ -9,7 +9,7 @@
9	namespace c10_npu {	9	namespace c10_npu {
10	namespace NPUCachingAllocator {	10	namespace NPUCachingAllocator {
11	const std::string pinMemoryExpandableMinCannVersion = "8.5.0";	11	const std::string pinMemoryExpandableMinCannVersion = "8.5.0";
12	-const std::string pinMemoryExpandableMinDriverVersion = "25.4.0";	12	+const std::string pinMemoryExpandableMinDriverVersion = "25.5.0";
13	const std::string cannModule = "CANN";	13	const std::string cannModule = "CANN";
14		14
15	bool isDigit(std::string str)	15	bool isDigit(std::string str)

Mtorch_npu/csrc/npu/Module.cpp+51-7

	@@ -1124,35 +1124,77 @@ PyObject* THNPModule_emptyCache(PyObject _unused, PyObject noargs)
1124	Py_RETURN_NONE;	1124	Py_RETURN_NONE;
1125	}	1125	}
1126		1126
1127	-PyObject* THNPModule_npu_emptyPinMemoryCache(PyObject _unused, PyObject noargs)	1127	+PyObject* THNPModule_npu_hostEmptyCache(PyObject _unused, PyObject noargs)
1128	{	1128	{
1129	HANDLE_TH_ERRORS	1129	HANDLE_TH_ERRORS
1130	- at_npu::native::CachingHostAllocator_emptyCache();	1130	+ at::getHostAllocator(at::kPrivateUse1)->empty_cache();
1131	END_HANDLE_TH_ERRORS	1131	END_HANDLE_TH_ERRORS
1132	Py_RETURN_NONE;	1132	Py_RETURN_NONE;
1133	}	1133	}
1134		1134
1135	-PyObject* THNPModule_npu_pinMemoryStats(PyObject _unused, PyObject noargs)	1135	+PyObject* THNPModule_npu_hostMemoryStats(PyObject _unused, PyObject noargs)
1136	{	1136	{
1137	HANDLE_TH_ERRORS	1137	HANDLE_TH_ERRORS
1138		1138
1139	- const auto statToDict = [](const c10::CachingAllocator::Stat& stat) {	1139	+ using at::HostStats;
		1140	+ using c10::CachingAllocator::DurationStat;
		1141	+ using c10::CachingAllocator::Stat;
		1142	+ using c10::CachingAllocator::StatArray;
		1143	+ using c10::CachingAllocator::StatType;
		1144	+
		1145	+ const auto statToDict = [](const Stat& stat) {
1140	py::dict dict;	1146	py::dict dict;
		1147	+
1141	dict["current"] = stat.current;	1148	dict["current"] = stat.current;
1142	dict["peak"] = stat.peak;	1149	dict["peak"] = stat.peak;
		1150	+ dict["allocated"] = stat.allocated;
		1151	+ dict["freed"] = stat.freed;
1143	return dict;	1152	return dict;
1144	};	1153	};
1145		1154
1146	- const at::HostStats stats = at_npu::native::CachingHostAllocator_getStats();	1155	+ const auto durationStatToDict = [](const DurationStat& stat) {
		1156	+ py::dict dict;
		1157	+
		1158	+ dict["total"] = stat.total;
		1159	+ dict["max"] = stat.max;
		1160	+ dict["min"] = stat.min;
		1161	+ dict["count"] = stat.count;
		1162	+ dict["avg"] = stat.count == 0 ? 0 : stat.total / stat.count;
		1163	+ return dict;
		1164	+ };
		1165	+
		1166	+ const HostStats stats = at::getHostAllocator(at::kPrivateUse1)->get_stats();
1147		1167
1148	py::dict result;	1168	py::dict result;
		1169	+ result["num_host_alloc"] = stats.num_host_alloc;
		1170	+ result["num_host_free"] = stats.num_host_free;
		1171	+ result["allocation"] = statToDict(stats.allocation);
		1172	+ result["segment"] = statToDict(stats.segment);
1149	result["allocated_bytes"] = statToDict(stats.allocated_bytes);	1173	result["allocated_bytes"] = statToDict(stats.allocated_bytes);
1150	result["reserved_bytes"] = statToDict(stats.reserved_bytes);	1174	result["reserved_bytes"] = statToDict(stats.reserved_bytes);
		1175	+ result["host_alloc_time"] = durationStatToDict(stats.host_alloc_time);
		1176	+ result["host_free_time"] = durationStatToDict(stats.host_free_time);
1151		1177
1152	return result.release().ptr();	1178	return result.release().ptr();
1153	END_HANDLE_TH_ERRORS	1179	END_HANDLE_TH_ERRORS
1154	}	1180	}
1155		1181
		1182	+PyObject* THNPModule_npu_resetAccumulatedHostMemoryStats(PyObject* _unused, PyObject* noargs)
		1183	+{
		1184	+ HANDLE_TH_ERRORS
		1185	+ at::getHostAllocator(at::kPrivateUse1)->reset_accumulated_stats();
		1186	+ END_HANDLE_TH_ERRORS
		1187	+ Py_RETURN_NONE;
		1188	+}
		1189	+
		1190	+PyObject* THNPModule_npu_resetPeakHostMemoryStats(PyObject* _unused, PyObject* noargs)
		1191	+{
		1192	+ HANDLE_TH_ERRORS
		1193	+ at::getHostAllocator(at::kPrivateUse1)->reset_peak_stats();
		1194	+ END_HANDLE_TH_ERRORS
		1195	+ Py_RETURN_NONE;
		1196	+}
		1197	+
1156	PyObject* THNPModule_npu_ipc_collect(PyObject _unused, PyObject noargs)	1198	PyObject* THNPModule_npu_ipc_collect(PyObject _unused, PyObject noargs)
1157	{	1199	{
1158	HANDLE_TH_ERRORS	1200	HANDLE_TH_ERRORS
	@@ -2201,8 +2243,10 @@ static struct PyMethodDef THNPModule_methods[] = {
2201	{"_npu_is_jit_compile_false", (PyCFunction)THNPModule_is_jit_compile_false_wrap, METH_NOARGS, nullptr},	2243	{"_npu_is_jit_compile_false", (PyCFunction)THNPModule_is_jit_compile_false_wrap, METH_NOARGS, nullptr},
2202	{"_npu_setMemoryFraction", (PyCFunction) THNPModule_setMemoryFraction, METH_VARARGS, nullptr},	2244	{"_npu_setMemoryFraction", (PyCFunction) THNPModule_setMemoryFraction, METH_VARARGS, nullptr},
2203	{"_npu_emptyCache", (PyCFunction) THNPModule_emptyCache, METH_NOARGS, nullptr},	2245	{"_npu_emptyCache", (PyCFunction) THNPModule_emptyCache, METH_NOARGS, nullptr},
2204	- {"_npu_emptyPinMemoryCache", (PyCFunction) THNPModule_npu_emptyPinMemoryCache, METH_NOARGS, nullptr},	2246	+ {"_npu_hostEmptyCache", (PyCFunction) THNPModule_npu_hostEmptyCache, METH_NOARGS, nullptr},
2205	- {"_npu_pinMemoryStats", (PyCFunction) THNPModule_npu_pinMemoryStats, METH_NOARGS, nullptr},	2247	+ {"_npu_hostMemoryStats", (PyCFunction) THNPModule_npu_hostMemoryStats, METH_NOARGS, nullptr},
		2248	+ {"_npu_resetAccumulatedHostMemoryStats", (PyCFunction) THNPModule_npu_resetAccumulatedHostMemoryStats, METH_NOARGS, nullptr},
		2249	+ {"_npu_resetPeakHostMemoryStats", (PyCFunction) THNPModule_npu_resetPeakHostMemoryStats, METH_NOARGS, nullptr},
2206	{"_npu_ipc_collect", (PyCFunction) THNPModule_npu_ipc_collect, METH_NOARGS, nullptr},	2250	{"_npu_ipc_collect", (PyCFunction) THNPModule_npu_ipc_collect, METH_NOARGS, nullptr},
2207	{"_npu_emptyVirtAddrCache", (PyCFunction) THNPModule_emptyVirtAddrCache, METH_NOARGS, nullptr},	2251	{"_npu_emptyVirtAddrCache", (PyCFunction) THNPModule_emptyVirtAddrCache, METH_NOARGS, nullptr},
2208	{"_npu_memoryStats", (PyCFunction) THNPModule_memoryStats, METH_O, nullptr},	2252	{"_npu_memoryStats", (PyCFunction) THNPModule_memoryStats, METH_O, nullptr},

Mtorch_npu/npu/__init__.py+5-4

	@@ -127,10 +127,11 @@ __all__ = [
127	"obfuscation_finalize",	127	"obfuscation_finalize",
128	"obfuscation_calculate",	128	"obfuscation_calculate",
129	"set_op_timeout_ms",	129	"set_op_timeout_ms",
130	- "empty_pin_memory_cache",	130	+ "host_empty_cache",
131	- "pin_memory_stats",	131	+ "host_memory_stats",
132	- "pin_memory_allocated",	132	+ "host_memory_stats_as_nested_dict",
133	- "pin_memory_reserved"	133	+ "reset_accumulated_host_memory_stats",
		134	+ "reset_peak_host_memory_stats"
134	]	135	]
135		136
136	from typing import Tuple, Union, List, cast, Optional	137	from typing import Tuple, Union, List, cast, Optional

Mtorch_npu/npu/memory.py+63-30

	@@ -41,10 +41,11 @@ __all__ = [
41	"MemPool",	41	"MemPool",
42	"MemPoolContext",	42	"MemPoolContext",
43	"use_mem_pool",	43	"use_mem_pool",
44	- "empty_pin_memory_cache",	44	+ "host_empty_cache",
45	- "pin_memory_stats",	45	+ "host_memory_stats",
46	- "pin_memory_allocated",	46	+ "host_memory_stats_as_nested_dict",
47	- "pin_memory_reserved",	47	+ "reset_accumulated_host_memory_stats",
		48	+ "reset_peak_host_memory_stats"
48	]	49	]
49		50
50	if not hasattr(torch_npu._C, "_npu_NPUAllocator"):	51	if not hasattr(torch_npu._C, "_npu_NPUAllocator"):
	@@ -163,41 +164,50 @@ def empty_cache():
163	torch_npu._C._npu_emptyCache()	164	torch_npu._C._npu_emptyCache()
164		165
165		166
166	-def empty_pin_memory_cache():	167	+def host_empty_cache():
167	r"""	168	r"""
168	Releases all unoccupied cached memory currently held by the caching	169	Releases all unoccupied cached memory currently held by the caching
169	allocator so that those can be used in other NPU application.	170	allocator so that those can be used in other NPU application.
170	"""	171	"""
171	- torch_npu._C._npu_emptyPinMemoryCache()	172	+ torch_npu._C._npu_hostEmptyCache()
172		173
173		174
174	-def pin_memory_reserved():	175	+def host_memory_stats():
175	- r"""	176	+ r"""Return a dictionary of pinned (host) allocator statistics.
176	- Returns the current pin memory managed by the caching allocator in bytes.
177	- """
178	- return pin_memory_stats().get('reserved_bytes.current', 0)
179		177
		178	+ Core statistics (host pinned allocator):
180		179
181	-def pin_memory_allocated():	180	+ - ``"allocation.{current,peak,allocated,freed}"``:
182	- r"""	181	+ pinned blocks owned by the allocator (active + cached). Grows when a new
183	- Returns the current pin memory occupied by tensors in bytes.	182	+ block is created via NPU and shrinks when cached blocks are returned.
184	- """	183	+ - ``"reserved_bytes.{current,peak,allocated,freed}"``:
185	- return pin_memory_stats().get('allocated_bytes.current', 0)	184	+ bytes of pinned blocks owned by the allocator (active + cached), using
		185	+ the rounded block size requested from NPU.
		186	+ - ``"segment.{current,peak,allocated,freed}"``:
		187	+ blocks currently checked out to callers (increments on handout, decrements
		188	+ when the block becomes reusable after stream deps finish).
		189	+ - ``"allocated_bytes.{current,peak,allocated,freed}"``:
		190	+ bytes corresponding to active blocks.
186		191
		192	+ Metric type:
187		193
188	-def pin_memory_stats():	194	+ - ``current``: current value.
189	- """Returns a dictionary of pin memory allocator statistics.	195	+ - ``peak``: maximum value.
190	- The return value of this function is a dictionary of statistics, each of	196	+ - ``allocated``: historical total increase.
191	- which is a non-negative integer.	197	+ - ``freed``: historical total decrease.
192	- Core statistics:	198	+
193	- - ``"allocated_bytes.{current,peak}"``:	199	+ Event/timing counters:
194	- amount of allocated memory.	200	+
195	- - ``"reserved_bytes.{current,peak}"``:	201	+ - ``"num_host_alloc"`` / ``"num_host_free"``: blocks created to grow the
196	- amount of reserved memory.	202	+ pool / cached blocks returned to NPU (matches allocations allocated/freed).
197	- For these core statistics, values are broken down as follows.	203	+ - ``"host_alloc_time.{total,max,min,count,avg}"``: time in host alloc calls
198	- Pool type:	204	+ when growing the pool (microseconds).
199	- - ``current``: current value of this metric.	205	+ - ``"host_free_time.{total,max,min,count,avg}"``: time in host free calls
200	- - ``peak``: maximum value of this metric.	206	+ when cached blocks are returned (microseconds).
		207	+
		208	+ Block sizes are rounded up to the next power of two before calling NPU, so
		209	+ byte stats reflect the rounded size. Peak values are aggregated per bucket
		210	+ and are a best-effort approximation of the true peak.
201	"""	211	"""
202	result = []	212	result = []
203		213
	@@ -210,12 +220,35 @@ def pin_memory_stats():
210	else:	220	else:
211	result.append((prefix, obj))	221	result.append((prefix, obj))
212		222
213	- stats = torch_npu._C._npu_pinMemoryStats()	223	+ stats = host_memory_stats_as_nested_dict()
214	_recurse_add_to_result("", stats)	224	_recurse_add_to_result("", stats)
215	result.sort()	225	result.sort()
216	return collections.OrderedDict(result)	226	return collections.OrderedDict(result)
217		227
218		228
		229	+def host_memory_stats_as_nested_dict():
		230	+ r"""Return the result of :func:`~torch_npu.npu.host_memory_stats` as a nested dictionary."""
		231	+ return torch_npu._C._npu_hostMemoryStats()
		232	+
		233	+
		234	+def reset_accumulated_host_memory_stats():
		235	+ r"""Reset the "accumulated" (historical) stats tracked by the host memory allocator.
		236	+
		237	+ See :func:`~torch_npu.npu.host_memory_stats` for details. Accumulated stats correspond to
		238	+ the `"allocated"` and `"freed"` keys in each individual stat dict.
		239	+ """
		240	+ return torch_npu._C._npu_resetAccumulatedHostMemoryStats()
		241	+
		242	+
		243	+def reset_peak_host_memory_stats():
		244	+ r"""Reset the "peak" stats tracked by the host memory allocator.
		245	+
		246	+ See :func:`~torch_npu.npu.host_memory_stats` for details. Peak stats correspond to the
		247	+ `"peak"` key in each individual stat dict.
		248	+ """
		249	+ return torch_npu._C._npu_resetPeakHostMemoryStats()
		250	+
		251	+
219	def empty_virt_addr_cache():	252	def empty_virt_addr_cache():
220	r"""Light-weight version of empty_cache(). It only unmaps virtual address,	253	r"""Light-weight version of empty_cache(). It only unmaps virtual address,
221	and store the free physical handles for later malloc.	254	and store the free physical handles for later malloc.