add torch.npu.memory._set_allocator_settings(only support expandable

Atest/test_npu_expandable_segments.py+222-0

@@ -0,0 +1,222 @@
	1	+import os
	2	+import gc
	3	+import shutil
	4	+import threading
	5	+import subprocess
	6	+
	7	+import torch
	8	+import torch.utils.cpp_extension
	9	+from torch.utils.data import Dataset, DataLoader
	10	+import torch.nn as nn
	11	+from torch.testing._internal.common_utils import TestCase, run_tests, instantiate_parametrized_tests, parametrize
	12	+import torch_npu
	13	+
	14	+PYTORCH_INSTALL_PATH = os.path.dirname(os.path.realpath(torch.__file__))
	15	+PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.realpath(torch_npu.__file__))
	16	+os.environ["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False"
	17	+
	18	+
	19	+def create_build_path(build_directory):
	20	+ if os.path.exists(build_directory):
	21	+ shutil.rmtree(build_directory, ignore_errors=True)
	22	+ os.makedirs(build_directory, exist_ok=True)
	23	+
	24	+
	25	+def build_stub(base_dir):
	26	+ build_stub_cmd = ["sh", os.path.join(base_dir, 'third_party/acl/libs/build_stub.sh')]
	27	+ if subprocess.call(build_stub_cmd) != 0:
	28	+ raise RuntimeError('Failed to build stub: {}'.format(build_stub_cmd))
	29	+
	30	+
	31	+class TestPluggableAllocator(TestCase):
	32	+ torch.npu.memory._set_allocator_settings("expandable_segments:True")
	33	+ module = None
	34	+ new_alloc = None
	35	+ build_directory = "allocator/build"
	36	+ conv = nn.Conv1d(1024, 256, 4, stride=4).to("npu")
	37	+ deconv = nn.ConvTranspose1d(256, 1024, 4, stride=4).to("npu")
	38	+
	39	+ @classmethod
	40	+ def setUpClass(cls):
	41	+ os_path = os.path.join(cls.build_directory, 'pluggable_allocator_extensions.so')
	42	+ if os.path.exists(os_path):
	43	+ cls.new_alloc = torch_npu.npu.memory.NPUPluggableAllocator(os_path, 'my_malloc', 'my_free')
	44	+ return
	45	+
	46	+ # Build Extension
	47	+ BASE_DIR = os.path.abspath("./../")
	48	+ build_stub(BASE_DIR)
	49	+ create_build_path(cls.build_directory)
	50	+ CANN_LIB_PATH = os.path.join(BASE_DIR, 'third_party/acl/libs')
	51	+ extra_ldflags = []
	52	+ extra_ldflags.append("-lascendcl")
	53	+ extra_ldflags.append(f"-L{CANN_LIB_PATH}")
	54	+ extra_ldflags.append("-lc10")
	55	+ extra_ldflags.append(f"-L{PYTORCH_INSTALL_PATH}")
	56	+ extra_include_paths = ["cpp_extensions"]
	57	+ extra_include_paths.append(os.path.join(PYTORCH_NPU_INSTALL_PATH, 'include'))
	58	+
	59	+ cls.module = torch.utils.cpp_extension.load(
	60	+ name="pluggable_allocator_extensions",
	61	+ sources=[
	62	+ "cpp_extensions/pluggable_allocator_extensions.cpp"
	63	+ ],
	64	+ extra_include_paths=extra_include_paths,
	65	+ extra_cflags=["-g"],
	66	+ extra_ldflags=extra_ldflags,
	67	+ build_directory=cls.build_directory,
	68	+ verbose=True,
	69	+ )
	70	+ # Load the allocator
	71	+ cls.new_alloc = torch_npu.npu.memory.NPUPluggableAllocator(os_path, 'my_malloc', 'my_free')
	72	+
	73	+ def test_pluggable_allocator(self):
	74	+ torch.npu.memory._set_allocator_settings("expandable_segments:False")
	75	+ with torch.npu.use_mem_pool(torch.npu.MemPool(TestPluggableAllocator.new_alloc._allocator)):
	76	+ x = torch.empty((7500, 1024, 1024), device="npu")
	77	+ del x
	78	+ torch.npu.memory._set_allocator_settings("expandable_segments:True")
	79	+
	80	+ @staticmethod
	81	+ def conv_operation(x):
	82	+ return TestPluggableAllocator.deconv(TestPluggableAllocator.conv(x) + 0.005)
	83	+
	84	+ @staticmethod
	85	+ def conv_with_allocator(x):
	86	+ torch.npu.memory._set_allocator_settings("expandable_segments:False")
	87	+ with torch.npu.use_mem_pool(torch.npu.MemPool(TestPluggableAllocator.new_alloc._allocator)):
	88	+ x = TestPluggableAllocator.conv_operation(x)
	89	+ torch.npu.memory._set_allocator_settings("expandable_segments:True")
	90	+ return x
	91	+
	92	+ @parametrize("task_queue_enable", [0, 1, 2])
	93	+ def test_task_queue(self, task_queue_enable):
	94	+ os.environ["TASK_QUEUE_ENABLE"] = str(task_queue_enable)
	95	+ input_data = torch.randn(1, 1024, 96, dtype=torch.float32, device="npu")
	96	+ x1 = input_data
	97	+ for _ in range(5):
	98	+ x1 = self.conv_operation(x1)
	99	+ x1 = self.conv_with_allocator(x1)
	100	+ x2 = input_data
	101	+ for _ in range(10):
	102	+ x2 = self.conv_operation(x2)
	103	+ self.assertEqual(x1, x2)
	104	+ os.environ["TASK_QUEUE_ENABLE"] = "1"
	105	+
	106	+ def test_thread_share(self):
	107	+ lock = threading.Lock()
	108	+
	109	+ def worker(name, shared_tensor):
	110	+ torch.npu.synchronize()
	111	+ with lock:
	112	+ shared_tensor.sub_(1)
	113	+ torch.npu.synchronize()
	114	+ torch.npu.memory._set_allocator_settings("expandable_segments:False")
	115	+ with torch.npu.use_mem_pool(torch.npu.MemPool(TestPluggableAllocator.new_alloc._allocator)):
	116	+ input_data = torch.zeros((4, 4), dtype=torch.float32, device="npu")
	117	+ with lock:
	118	+ input_data.add_(1)
	119	+ th = threading.Thread(target=worker, args=("thread1", input_data))
	120	+ th.start()
	121	+ th.join()
	122	+ self.assertEqual(input_data, torch.zeros((4, 4), dtype=torch.float32, device="npu"))
	123	+
	124	+ def test_mul_stream(self):
	125	+ input_data = torch.randn(1, 1024, 96, dtype=torch.float32, device="npu")
	126	+ x1 = input_data
	127	+ x2 = input_data
	128	+ stream1, stream2 = torch.npu.Stream(), torch.npu.Stream()
	129	+ events = [torch.npu.Event(False, False) for _ in range(3)]
	130	+
	131	+ with torch.npu.stream(stream1):
	132	+ x1 = self.conv_with_allocator(x1)
	133	+ events[0].record()
	134	+
	135	+ with torch.npu.stream(stream2):
	136	+ events[0].wait(stream2)
	137	+ x2 = self.conv_operation(x2)
	138	+ events[1].record()
	139	+
	140	+ with torch.npu.stream(stream1):
	141	+ events[1].wait(stream1)
	142	+ x1 = self.conv_with_allocator(x1)
	143	+ events[2].record()
	144	+
	145	+ with torch.npu.stream(stream2):
	146	+ events[2].wait(stream2)
	147	+ x2 = self.conv_operation(x2)
	148	+
	149	+ torch.npu.synchronize()
	150	+ self.assertEqual(x1, x2)
	151	+
	152	+ def test_mul_stream_with_threads(self):
	153	+ input_data = torch.randn(1, 1024, 96, dtype=torch.float32, device="npu")
	154	+ events = [torch.npu.Event(False, False) for _ in range(3)]
	155	+
	156	+ def stream_worker(data, stream, event_sequence):
	157	+ """Generic stream worker function"""
	158	+ with torch.npu.stream(stream):
	159	+ for event, operation in event_sequence:
	160	+ event.wait(stream)
	161	+ data = operation(data)
	162	+ events[event_sequence.index((event, operation)) + 1].record()
	163	+ return data
	164	+
	165	+ # Define operation sequences for two streams
	166	+ stream1_ops = [
	167	+ (events[0], self.conv_with_allocator),
	168	+ (events[1], self.conv_operation)
	169	+ ]
	170	+ stream2_ops = [
	171	+ (events[0], self.conv_operation),
	172	+ (events[2], self.conv_with_allocator)
	173	+ ]
	174	+
	175	+ result_container = {}
	176	+ stream2 = torch.npu.Stream()
	177	+
	178	+ def thread_func():
	179	+ result_container["x2"] = stream_worker(input_data, stream2, stream2_ops)
	180	+
	181	+ thread = threading.Thread(target=thread_func)
	182	+ thread.start()
	183	+
	184	+ stream1 = torch.npu.Stream()
	185	+ x1 = stream_worker(input_data, stream1, stream1_ops)
	186	+
	187	+ thread.join()
	188	+ torch.npu.synchronize()
	189	+
	190	+ self.assertEqual(x1, result_container["x2"])
	191	+
	192	+ def test_dict_data_loader(self):
	193	+ class DictDataset(Dataset):
	194	+ def __len__(self):
	195	+ return 4
	196	+
	197	+ def __getitem__(self, idx):
	198	+ torch.npu.memory._set_allocator_settings("expandable_segments:False")
	199	+ with torch.npu.use_mem_pool(torch.npu.MemPool(TestPluggableAllocator.new_alloc._allocator)):
	200	+ ret_dict = {
	201	+ "a_tensor": torch.randn(4, 2, dtype=torch.float32, device="npu"),
	202	+ "another_dict": {"a_number": idx}
	203	+ }
	204	+ torch.npu.memory._set_allocator_settings("expandable_segments:True")
	205	+ return ret_dict
	206	+
	207	+ class TestDictDataLoader():
	208	+ def __init__(self):
	209	+ self.dataset = DictDataset()
	210	+
	211	+ def test_memory(self):
	212	+ loader = DataLoader(self.dataset, batch_size=2)
	213	+ for sample in loader:
	214	+ print(f'sample: {sample}')
	215	+
	216	+ loader = TestDictDataLoader()
	217	+ loader.test_memory()
	218	+
	219	+instantiate_parametrized_tests(TestPluggableAllocator)
	220	+
	221	+if __name__ == '__main__':
	222	+ run_tests()

Mtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp+18-2

	@@ -890,7 +890,7 @@ public:
890	return *s_instance;	890	return *s_instance;
891	}	891	}
892		892
893	- void parseArgs(const char *env);	893	+ void parseArgs(const char *env, std::set<std::string> supported_settings = {});
894		894
895	private:	895	private:
896	size_t m_max_split_size;	896	size_t m_max_split_size;
	@@ -1031,7 +1031,7 @@ size_t CachingAllocatorConfig::parsePageSize(const std::vector<std::string> &con
1031	return i + 2; // 返回最后处理的索引位置	1031	return i + 2; // 返回最后处理的索引位置
1032	}	1032	}
1033		1033
1034	-void CachingAllocatorConfig::parseArgs(const char *env)	1034	+void CachingAllocatorConfig::parseArgs(const char *env, std::set<std::string> supported_settings)
1035	{	1035	{
1036	// If empty, set the default values	1036	// If empty, set the default values
1037	m_max_split_size = std::numeric_limits<size_t>::max();	1037	m_max_split_size = std::numeric_limits<size_t>::max();
	@@ -1045,6 +1045,12 @@ void CachingAllocatorConfig::parseArgs(const char *env)
1045	lexArgs(env, config);	1045	lexArgs(env, config);
1046		1046
1047	for (size_t i = 0; i < config.size(); i++) {	1047	for (size_t i = 0; i < config.size(); i++) {
		1048	+ // If supported_settings is not empty,
		1049	+ // check if the setting is supported by torch_npu.npu.memory._set_allocator_settings().
		1050	+ if (!supported_settings.empty() && supported_settings.count(config[i]) == 0) {
		1051	+ TORCH_CHECK(false, "torch_npu.npu.memory._set_allocator_settings() unsupported setting: ", config[i],
		1052	+ OPS_ERROR(ErrCode::VALUE));
		1053	+ }
1048	if (config[i].compare("max_split_size_mb") == 0) {	1054	if (config[i].compare("max_split_size_mb") == 0) {
1049	i = parseMaxSplitSize(config, i);	1055	i = parseMaxSplitSize(config, i);
1050	} else if (config[i].compare("garbage_collection_threshold") == 0) {	1056	} else if (config[i].compare("garbage_collection_threshold") == 0) {
	@@ -1089,6 +1095,16 @@ bool isConfig1GPageSizeEnable()
1089	return CachingAllocatorConfig::page_size_1g_enable();	1095	return CachingAllocatorConfig::page_size_1g_enable();
1090	}	1096	}
1091		1097
		1098	+void setAllocatorSettings(const std::string& settings)
		1099	+{
		1100	+ ASCEND_LOGI("setAllocatorSettings: %s.", settings.c_str());
		1101	+ // Empty NPU task queue before changing the allocator settings.
		1102	+ NPUStatus ret = c10_npu::emptyAllNPUStream();
		1103	+ TORCH_CHECK(ret == NPU_STATUS_SUCCESS, "Failed to empty NPU task queue, ret:", ret, PTA_ERROR(ErrCode::INTERNAL));
		1104	+ // Only support expandable_segments setting.
		1105	+ CachingAllocatorConfig::instance().parseArgs(settings.c_str(), {"expandable_segments"});
		1106	+}
		1107	+
1092	// To prevent the deadlock situation, temporarily release the lock.	1108	// To prevent the deadlock situation, temporarily release the lock.
1093	//	1109	//
1094	// Deadlock Scenario Description:	1110	// Deadlock Scenario Description:

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

	@@ -474,6 +474,8 @@ bool checkConfigExpandableSegments();
474		474
475	bool isConfig1GPageSizeEnable();	475	bool isConfig1GPageSizeEnable();
476		476
		477	+C10_NPU_API void setAllocatorSettings(const std::string& settings);
		478	+
477	} // namespace NPUCachingAllocator	479	} // namespace NPUCachingAllocator
478	} // namespace c10_npu	480	} // namespace c10_npu
479		481

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

	@@ -1359,6 +1359,15 @@ PyObject* THNPModule_npuCachingAllocator_raw_delete(PyObject *_unused, PyObject
1359	END_HANDLE_TH_ERRORS	1359	END_HANDLE_TH_ERRORS
1360	}	1360	}
1361		1361
		1362	+PyObject* THNPModule_npuCachingAllocator_set_allocator_settings(PyObject _unused, PyObject arg)
		1363	+{
		1364	+ HANDLE_TH_ERRORS
		1365	+ std::string settings = THPUtils_unpackString(arg);
		1366	+ c10_npu::NPUCachingAllocator::setAllocatorSettings(settings);
		1367	+ END_HANDLE_TH_ERRORS
		1368	+ Py_RETURN_NONE;
		1369	+}
		1370	+
1362	PyObject* THNPModule_getAllocatorBackend(PyObject _unused, PyObject noargs)	1371	PyObject* THNPModule_getAllocatorBackend(PyObject _unused, PyObject noargs)
1363	{	1372	{
1364	HANDLE_TH_ERRORS	1373	HANDLE_TH_ERRORS
	@@ -1970,6 +1979,7 @@ static struct PyMethodDef THNPModule_methods[] = {
1970	{"_npu_attach_out_of_memory_observer", THNPModule_attachOutOfMemoryObserver, METH_O, nullptr},	1979	{"_npu_attach_out_of_memory_observer", THNPModule_attachOutOfMemoryObserver, METH_O, nullptr},
1971	{"_npu_npuCachingAllocator_raw_alloc", (PyCFunction)THNPModule_npuCachingAllocator_raw_alloc, METH_VARARGS, nullptr},	1980	{"_npu_npuCachingAllocator_raw_alloc", (PyCFunction)THNPModule_npuCachingAllocator_raw_alloc, METH_VARARGS, nullptr},
1972	{"_npu_npuCachingAllocator_raw_delete", (PyCFunction)THNPModule_npuCachingAllocator_raw_delete, METH_O, nullptr},	1981	{"_npu_npuCachingAllocator_raw_delete", (PyCFunction)THNPModule_npuCachingAllocator_raw_delete, METH_O, nullptr},
		1982	+ {"_npu_npuCachingAllocator_set_allocator_settings", (PyCFunction)THNPModule_npuCachingAllocator_set_allocator_settings, METH_O, nullptr},
1973	{"_npu_getAllocatorBackend", (PyCFunction)THNPModule_getAllocatorBackend, METH_NOARGS, nullptr},	1983	{"_npu_getAllocatorBackend", (PyCFunction)THNPModule_getAllocatorBackend, METH_NOARGS, nullptr},
1974	{"_npu_lock_mutex", (PyCFunction)THNPModule_npuLockMutex, METH_NOARGS, nullptr},	1984	{"_npu_lock_mutex", (PyCFunction)THNPModule_npuLockMutex, METH_NOARGS, nullptr},
1975	{"_npu_unlock_mutex", (PyCFunction)THNPModule_npuUnlockMutex, METH_NOARGS, nullptr},	1985	{"_npu_unlock_mutex", (PyCFunction)THNPModule_npuUnlockMutex, METH_NOARGS, nullptr},

Mtorch_npu/npu/memory.py+6-0

	@@ -586,6 +586,12 @@ def get_allocator_backend() -> str:
586	return torch_npu._C._npu_getAllocatorBackend()	586	return torch_npu._C._npu_getAllocatorBackend()
587		587
588		588
		589	+def _set_allocator_settings(settings) -> None:
		590	+ r"""Sets the allocator settings. Only support expandable_segments:True or False.
		591	+ """
		592	+ return torch_npu._C._npu_npuCachingAllocator_set_allocator_settings(settings)
		593	+
		594	+
589	class _NPUAllocator:	595	class _NPUAllocator:
590	r"""Wrapper over internal NPU memory allocators."""	596	r"""Wrapper over internal NPU memory allocators."""
591		597