Provide interface such as getDeviceStatus, resetPeakStats for NPUPlugg…-pytorch-AtomGit

Mtest/allocator/test_pluggable_allocator_extensions.py+29-2

	@@ -2,6 +2,7 @@ import os
2	import sys	2	import sys
3	import shutil	3	import shutil
4	import subprocess	4	import subprocess
		5	+import ctypes
5	import torch	6	import torch
6	import torch.utils.cpp_extension	7	import torch.utils.cpp_extension
7		8
	@@ -27,6 +28,7 @@ def build_stub(base_dir):
27		28
28	class TestPluggableAllocator(TestCase):	29	class TestPluggableAllocator(TestCase):
29	module = None	30	module = None
		31	+ new_alloc = None
30	build_directory = "allocator/build"	32	build_directory = "allocator/build"
31		33
32	@classmethod	34	@classmethod
	@@ -59,9 +61,9 @@ class TestPluggableAllocator(TestCase):
59	def test_pluggable_allocator(self):	61	def test_pluggable_allocator(self):
60	os_path = os.path.join(TestPluggableAllocator.build_directory, 'pluggable_allocator_extensions.so')	62	os_path = os.path.join(TestPluggableAllocator.build_directory, 'pluggable_allocator_extensions.so')
61	# Load the allocator	63	# Load the allocator
62	- new_alloc = torch_npu.npu.memory.NPUPluggableAllocator(os_path, 'my_malloc', 'my_free')	64	+ TestPluggableAllocator.new_alloc = torch_npu.npu.memory.NPUPluggableAllocator(os_path, 'my_malloc', 'my_free')
63	# Swap the current allocator	65	# Swap the current allocator
64	- torch_npu.npu.memory.change_current_allocator(new_alloc)	66	+ torch_npu.npu.memory.change_current_allocator(TestPluggableAllocator.new_alloc)
65	# This will allocate memory in the device using the new allocator	67	# This will allocate memory in the device using the new allocator
66	self.assertFalse(self.module.check_custom_allocator_used())	68	self.assertFalse(self.module.check_custom_allocator_used())
67	npu_tensor = torch.zeros(10, device='npu')	69	npu_tensor = torch.zeros(10, device='npu')
	@@ -69,6 +71,31 @@ class TestPluggableAllocator(TestCase):
69	self.assertRtolEqual(npu_tensor.cpu().numpy(), cpu_tensor.numpy())	71	self.assertRtolEqual(npu_tensor.cpu().numpy(), cpu_tensor.numpy())
70	self.assertTrue(self.module.check_custom_allocator_used())	72	self.assertTrue(self.module.check_custom_allocator_used())
71		73
		74	+ def test_set_get_device_stats_fn(self):
		75	+ os_path = os.path.join(TestPluggableAllocator.build_directory, 'pluggable_allocator_extensions.so')
		76	+ myallocator = ctypes.CDLL(os_path)
		77	+ get_device_stats_fn = ctypes.cast(getattr(myallocator, "my_get_device_stats"), ctypes.c_void_p).value
		78	+
		79	+ msg = "get_device_stats_fn_ is not define, please set by set_get_device_stats_fn"
		80	+ with self.assertRaisesRegex(RuntimeError, msg):
		81	+ torch.npu.memory_stats_as_nested_dict()
		82	+
		83	+ TestPluggableAllocator.new_alloc.allocator().set_get_device_stats_fn(get_device_stats_fn)
		84	+ self.assertEqual(torch.npu.memory_stats_as_nested_dict()["num_alloc_retries"], 0)
		85	+
		86	+ def test_set_reset_peak_status_fn(self):
		87	+ os_path = os.path.join(TestPluggableAllocator.build_directory, 'pluggable_allocator_extensions.so')
		88	+ myallocator = ctypes.CDLL(os_path)
		89	+ reset_peak_status_fn = ctypes.cast(getattr(myallocator, "my_reset_peak_status"), ctypes.c_void_p).value
		90	+
		91	+ msg = "reset_peak_status_fn_ is not define, please set by set_reset_peak_status_fn"
		92	+ with self.assertRaisesRegex(RuntimeError, msg):
		93	+ torch.npu.reset_peak_memory_stats()
		94	+
		95	+ TestPluggableAllocator.new_alloc.allocator().set_reset_peak_status_fn(reset_peak_status_fn)
		96	+ torch.npu.reset_peak_memory_stats()
		97	+ self.assertEqual(torch.npu.max_memory_allocated(), 0)
		98	+
72	def test_pluggable_allocator_after_init(self):	99	def test_pluggable_allocator_after_init(self):
73	os_path = os.path.join(TestPluggableAllocator.build_directory, 'pluggable_allocator_extensions.so')	100	os_path = os.path.join(TestPluggableAllocator.build_directory, 'pluggable_allocator_extensions.so')
74	# Do an initial memory allocator	101	# Do an initial memory allocator

Mtest/cpp_extensions/pluggable_allocator_extensions.cpp+13-0

	@@ -4,8 +4,10 @@
4		4
5	#include "third_party/acl/inc/acl/acl_base.h"	5	#include "third_party/acl/inc/acl/acl_base.h"
6	#include "third_party/acl/inc/acl/acl_rt.h"	6	#include "third_party/acl/inc/acl/acl_rt.h"
		7	+#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
7		8
8	extern "C" {	9	extern "C" {
		10	+using c10_npu::NPUCachingAllocator::DeviceStats;
9	static bool useflag = false;	11	static bool useflag = false;
10		12
11	void* my_malloc(ssize_t size, int device, aclrtStream stream)	13	void* my_malloc(ssize_t size, int device, aclrtStream stream)
	@@ -27,6 +29,17 @@ bool check_custom_allocator_used()
27	{	29	{
28	return useflag;	30	return useflag;
29	}	31	}
		32	+
		33	+DeviceStats my_get_device_stats(int device)
		34	+{
		35	+ DeviceStats stats;
		36	+ return stats;
		37	+}
		38	+
		39	+void my_reset_peak_status(int device)
		40	+{
		41	+ std::cout<<"resetPeakStatus success!"<<std::endl;
		42	+}
30	}	43	}
31		44
32	PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {	45	PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {

Mtorch_npu/csrc/npu/Module.cpp+18-0

	@@ -275,6 +275,24 @@ void RegisterNpuPluggableAllocator(PyObject* module)
275	std::function<FuncType> func =	275	std::function<FuncType> func =
276	reinterpret_cast<FuncType*>(func_ptr);	276	reinterpret_cast<FuncType*>(func_ptr);
277	self.set_erase_stream_fn(func);	277	self.set_erase_stream_fn(func);
		278	+ })
		279	+ .def(
		280	+ "set_get_device_stats_fn",
		281	+ [](torch::npu::NPUPluggableAllocator::NPUPluggableAllocator& self,
		282	+ uint64_t func_ptr) {
		283	+ using FuncType=c10_npu::NPUCachingAllocator::DeviceStats(int);
		284	+ std::function<FuncType> func =
		285	+ reinterpret_cast<FuncType*>(func_ptr);
		286	+ self.set_get_device_stats_fn(func);
		287	+ })
		288	+ .def(
		289	+ "set_reset_peak_status_fn",
		290	+ [](torch::npu::NPUPluggableAllocator::NPUPluggableAllocator& self,
		291	+ uint64_t func_ptr) {
		292	+ using FuncType = void(int);
		293	+ std::function<FuncType> func =
		294	+ reinterpret_cast<FuncType*>(func_ptr);
		295	+ self.set_reset_peak_status_fn(func);
278	});	296	});
279		297
280	m.def(	298	m.def(

Mtorch_npu/csrc/npu/NPUPluggableAllocator.cpp+22-4

	@@ -74,6 +74,18 @@ void NPUPluggableAllocator::set_erase_stream_fn(
74	erase_stream_fn_ = std::move(erase_stream_fn);	74	erase_stream_fn_ = std::move(erase_stream_fn);
75	}	75	}
76		76
		77	+void NPUPluggableAllocator::set_get_device_stats_fn(
		78	+ std::function<c10_npu::NPUCachingAllocator::DeviceStats(int)> get_device_stats_fn)
		79	+{
		80	+ get_device_stats_fn_ = std::move(get_device_stats_fn);
		81	+}
		82	+
		83	+void NPUPluggableAllocator::set_reset_peak_status_fn(
		84	+ std::function<void(int)> reset_peak_status_fn)
		85	+{
		86	+ reset_peak_status_fn_ = std::move(reset_peak_status_fn);
		87	+}
		88	+
77	void* NPUPluggableAllocator::malloc(	89	void* NPUPluggableAllocator::malloc(
78	size_t size,	90	size_t size,
79	int device,	91	int device,
	@@ -212,8 +224,11 @@ void NPUPluggableAllocator::eraseStream(
212		224
213	c10_npu::NPUCachingAllocator::DeviceStats NPUPluggableAllocator::getDeviceStats(int device)	225	c10_npu::NPUCachingAllocator::DeviceStats NPUPluggableAllocator::getDeviceStats(int device)
214	{	226	{
215	- TORCH_NPU_WARN("NPUPluggableAllocator does not yet support getDeviceStats. "	227	+ if (get_device_stats_fn_) {
216	- "If you need it, please file an issue describing your use case.");	228	+ return get_device_stats_fn_(device);
		229	+ } else {
		230	+ TORCH_CHECK(false, "get_device_stats_fn_ is not define, please set by set_get_device_stats_fn");
		231	+ }
217	}	232	}
218		233
219	void NPUPluggableAllocator::resetAccumulatedStats(int device)	234	void NPUPluggableAllocator::resetAccumulatedStats(int device)
	@@ -224,8 +239,11 @@ void NPUPluggableAllocator::resetAccumulatedStats(int device)
224		239
225	void NPUPluggableAllocator::resetPeakStats(int device)	240	void NPUPluggableAllocator::resetPeakStats(int device)
226	{	241	{
227	- TORCH_NPU_WARN("NPUPluggableAllocator does not yet support resetPeakStats. "	242	+ if (reset_peak_status_fn_) {
228	- "If you need it, please file an issue describing your use case.");	243	+ reset_peak_status_fn_(device);
		244	+ } else {
		245	+ TORCH_CHECK(false, "reset_peak_status_fn_ is not define, please set by set_reset_peak_status_fn");
		246	+ }
229	}	247	}
230		248
231	c10_npu::NPUCachingAllocator::SnapshotInfo NPUPluggableAllocator::snapshot()	249	c10_npu::NPUCachingAllocator::SnapshotInfo NPUPluggableAllocator::snapshot()

Mtorch_npu/csrc/npu/NPUPluggableAllocator.h+4-0

	@@ -45,6 +45,8 @@ struct NPUPluggableAllocator
45	std::function<void(void* ptr, c10_npu::NPUStream stream)> record_stream_fn);	45	std::function<void(void* ptr, c10_npu::NPUStream stream)> record_stream_fn);
46	void set_erase_stream_fn(	46	void set_erase_stream_fn(
47	std::function<void(void* ptr, c10_npu::NPUStream stream)> erase_stream_fn);	47	std::function<void(void* ptr, c10_npu::NPUStream stream)> erase_stream_fn);
		48	+ void set_get_device_stats_fn(std::function<c10_npu::NPUCachingAllocator::DeviceStats(int)> get_device_stats_fn);
		49	+ void set_reset_peak_status_fn(std::function<void(int)> reset_peak_status_fn);
48	void* malloc(size_t size, int device, aclrtStream stream);	50	void* malloc(size_t size, int device, aclrtStream stream);
49		51
50	c10::DataPtr allocate(size_t size) override;	52	c10::DataPtr allocate(size_t size) override;
	@@ -108,6 +110,8 @@ protected:
108	std::function<void(void, size_t*)> base_alloc_fn_;	110	std::function<void(void, size_t*)> base_alloc_fn_;
109	std::function<void(void* ptr, c10_npu::NPUStream stream)> record_stream_fn_;	111	std::function<void(void* ptr, c10_npu::NPUStream stream)> record_stream_fn_;
110	std::function<void(void* ptr, c10_npu::NPUStream stream)> erase_stream_fn_;	112	std::function<void(void* ptr, c10_npu::NPUStream stream)> erase_stream_fn_;
		113	+ std::function<c10_npu::NPUCachingAllocator::DeviceStats(int)> get_device_stats_fn_;
		114	+ std::function<void(int)> reset_peak_status_fn_;
111	std::mutex allocator_mutex_;	115	std::mutex allocator_mutex_;
112	// We do the bookeeping here in order to simplify custom allocators	116	// We do the bookeeping here in order to simplify custom allocators
113	std::unordered_map<void*, _AllocationMetadata> allocation_metadata_;	117	std::unordered_map<void*, _AllocationMetadata> allocation_metadata_;