[fix][2.8.0]add getMemoryFraction attribute for torch.

Mtest/test_npu.py+18-0

	@@ -192,6 +192,24 @@ class TestNpu(TestCase):
192	torch_npu.npu.empty_cache()	192	torch_npu.npu.empty_cache()
193	torch_npu.npu.reset_peak_memory_stats()	193	torch_npu.npu.reset_peak_memory_stats()
194		194
		195	+ def test_get_per_process_memory_fraction(self):
		196	+ # get the initial memory fraction
		197	+ init_fraction = torch_npu.npu.get_per_process_memory_fraction()
		198	+
		199	+ # set and get the limiting cases
		200	+ torch_npu.npu.set_per_process_memory_fraction(1.0)
		201	+ self.assertEqual(torch_npu.npu.get_per_process_memory_fraction(), 1.0)
		202	+ torch_npu.npu.set_per_process_memory_fraction(0.0)
		203	+ self.assertEqual(torch_npu.npu.get_per_process_memory_fraction(), 0.0)
		204	+
		205	+ # test a few random cases
		206	+ for val in torch.rand(3):
		207	+ torch_npu.npu.set_per_process_memory_fraction(float(val))
		208	+ self.assertEqual(torch_npu.npu.get_per_process_memory_fraction(), float(val))
		209	+
		210	+ # restore the initial memory fraction
		211	+ torch_npu.npu.set_per_process_memory_fraction(init_fraction)
		212	+
195	def test_set_per_process_memory_fraction(self):	213	def test_set_per_process_memory_fraction(self):
196	# test invalid fraction value.	214	# test invalid fraction value.
197	with self.assertRaisesRegex(TypeError, "Invalid type"):	215	with self.assertRaisesRegex(TypeError, "Invalid type"):

Mtest/torch_npu_schema.json+6-0

	@@ -1169,6 +1169,9 @@
1169	"torch_npu.npu.set_option": {	1169	"torch_npu.npu.set_option": {
1170	"signature": "(option)"	1170	"signature": "(option)"
1171	},	1171	},
		1172	+ "torch_npu.npu.get_per_process_memory_fraction": {
		1173	+ "signature": "(device=None) -> float"
		1174	+ },
1172	"torch_npu.npu.set_per_process_memory_fraction": {	1175	"torch_npu.npu.set_per_process_memory_fraction": {
1173	"signature": "(fraction, device=None) -> None"	1176	"signature": "(fraction, device=None) -> None"
1174	},	1177	},
	@@ -1382,6 +1385,9 @@
1382	"torch_npu.npu.memory.reset_peak_memory_stats": {	1385	"torch_npu.npu.memory.reset_peak_memory_stats": {
1383	"signature": "(device=None)"	1386	"signature": "(device=None)"
1384	},	1387	},
		1388	+ "torch_npu.npu.memory.get_per_process_memory_fraction": {
		1389	+ "signature": "(device=None) -> float"
		1390	+ },
1385	"torch_npu.npu.memory.set_per_process_memory_fraction": {	1391	"torch_npu.npu.memory.set_per_process_memory_fraction": {
1386	"signature": "(fraction, device=None) -> None"	1392	"signature": "(fraction, device=None) -> None"
1387	},	1393	},

Mtest/unsupported_test_cases/.pytorch-disabled-tests.json+2-1

	@@ -31598,5 +31598,6 @@
31598	"test_fake_autocast_fft_ifftshift_npu_float32 (__main__.TestFakeTensorPRIVATEUSE1)": ["", [""]],	31598	"test_fake_autocast_fft_ifftshift_npu_float32 (__main__.TestFakeTensorPRIVATEUSE1)": ["", [""]],
31599	"test_replicate (__main__.TestDataParallel)": ["", [""]],	31599	"test_replicate (__main__.TestDataParallel)": ["", [""]],
31600	"test_replicate_buffers (__main__.TestDataParallel)": ["", [""]],	31600	"test_replicate_buffers (__main__.TestDataParallel)": ["", [""]],
31601	- "test_index_put_accumulate_large_tensor_npu (__main__.TestIndexingPRIVATEUSE1)": ["", ["910A"]]	31601	+ "test_index_put_accumulate_large_tensor_npu (__main__.TestIndexingPRIVATEUSE1)": ["", ["910A"]],
		31602	+ "test_serialization_array_with_empty (__main__.TestNpu)": ["", [""]]
31602	}	31603	}

Mtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp+23-0

	@@ -1579,6 +1579,20 @@ public:
1579	}	1579	}
1580	}	1580	}
1581		1581
		1582	+ / get memory fraction limiting maximum allocated memory /
		1583	+ double getMemoryFraction()
		1584	+ {
		1585	+ if (!set_fraction) {
		1586	+ return 1.0;
		1587	+ }
		1588	+
		1589	+ size_t device_free = 0;
		1590	+ size_t device_total = 0;
		1591	+ NPU_CHECK_ERROR(aclrtGetMemInfo(ACL_HBM_MEM, &device_free, &device_total));
		1592	+ return static_cast<double>(allowed_memory_maximum) /
		1593	+ static_cast<double>(device_total);
		1594	+ }
		1595	+
1582	/* * set memory fraction to limit maximum allocated memory * */	1596	/* * set memory fraction to limit maximum allocated memory * */
1583	void setMemoryFraction(double fraction)	1597	void setMemoryFraction(double fraction)
1584	{	1598	{
	@@ -3159,6 +3173,15 @@ public:
3159	device_allocator[block->device]->free(block);	3173	device_allocator[block->device]->free(block);
3160	}	3174	}
3161		3175
		3176	+ double getMemoryFraction(int device) override
		3177	+ {
		3178	+ TORCH_INTERNAL_ASSERT(
		3179	+ 0 <= device && device < device_allocator.size(), "Allocator not initialized for device ",
		3180	+ device, ": did you call init?", PTA_ERROR(ErrCode::PARAM));
		3181	+ NPU_CHECK_ERROR(c10_npu::SetDevice(device));
		3182	+ return device_allocator[device]->getMemoryFraction();
		3183	+ }
		3184	+
3162	void setMemoryFraction(double fraction, int device) override	3185	void setMemoryFraction(double fraction, int device) override
3163	{	3186	{
3164	TORCH_INTERNAL_ASSERT(0 <= device && device < device_allocator.size(), "Allocator not initialized for device ",	3187	TORCH_INTERNAL_ASSERT(0 <= device && device < device_allocator.size(), "Allocator not initialized for device ",

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

	@@ -201,6 +201,7 @@ public:
201	virtual void raw_delete(void* ptr) = 0;	201	virtual void raw_delete(void* ptr) = 0;
202	virtual void init(int device_count) = 0;	202	virtual void init(int device_count) = 0;
203	virtual bool initialized() = 0;	203	virtual bool initialized() = 0;
		204	+ virtual double getMemoryFraction(int device) = 0;
204	virtual void setMemoryFraction(double fraction, int device) = 0;	205	virtual void setMemoryFraction(double fraction, int device) = 0;
205	virtual void emptyCacheImpl(bool check_error, bool free_physical) = 0;	206	virtual void emptyCacheImpl(bool check_error, bool free_physical) = 0;
206	virtual void emptyCache(bool check_error) = 0;	207	virtual void emptyCache(bool check_error) = 0;
	@@ -305,6 +306,11 @@ inline void init()
305	return get()->init(device_count);	306	return get()->init(device_count);
306	}	307	}
307		308
		309	+inline double getMemoryFraction(int device)
		310	+{
		311	+ return get()->getMemoryFraction(device);
		312	+}
		313	+
308	inline void setMemoryFraction(double fraction, int device)	314	inline void setMemoryFraction(double fraction, int device)
309	{	315	{
310	return get()->setMemoryFraction(fraction, device);	316	return get()->setMemoryFraction(fraction, device);

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

	@@ -1104,6 +1104,24 @@ PyObject THNPModule_is_jit_compile_false_wrap(PyObject self, PyObject *noargs)
1104	END_HANDLE_TH_ERRORS	1104	END_HANDLE_TH_ERRORS
1105	}	1105	}
1106		1106
		1107	+PyObject* THNPModule_getMemoryFraction(PyObject* _unused, PyObject* args)
		1108	+{
		1109	+ HANDLE_TH_ERRORS
		1110	+ PyObject* device_o = nullptr;
		1111	+ if (!PyArg_ParseTuple(args, "O", &device_o)) {
		1112	+ THPUtils_invalidArguments(
		1113	+ args,
		1114	+ nullptr,
		1115	+ "get_memory_fraction",
		1116	+ 1,
		1117	+ "(int device);");
		1118	+ return nullptr;
		1119	+ }
		1120	+ int64_t device_index = PyLong_AsLongLong(device_o);
		1121	+ return PyFloat_FromDouble(c10_npu::NPUCachingAllocator::getMemoryFraction(device_index));
		1122	+ END_HANDLE_TH_ERRORS
		1123	+}
		1124	+
1107	PyObject* THNPModule_setMemoryFraction(PyObject _unused, PyObject args)	1125	PyObject* THNPModule_setMemoryFraction(PyObject _unused, PyObject args)
1108	{	1126	{
1109	HANDLE_TH_ERRORS	1127	HANDLE_TH_ERRORS
	@@ -2269,6 +2287,7 @@ static struct PyMethodDef THNPModule_methods[] = {
2269	{"_npu_eraseStream", (PyCFunction)THNPModule_npu_eraseStream_wrap, METH_VARARGS \| METH_KEYWORDS, nullptr},	2287	{"_npu_eraseStream", (PyCFunction)THNPModule_npu_eraseStream_wrap, METH_VARARGS \| METH_KEYWORDS, nullptr},
2270	{"_npu_isCurrentStreamCapturing", (PyCFunction)THNPModule_isCurrentStreamCapturing_wrap, METH_NOARGS, nullptr},	2288	{"_npu_isCurrentStreamCapturing", (PyCFunction)THNPModule_isCurrentStreamCapturing_wrap, METH_NOARGS, nullptr},
2271	{"_npu_is_jit_compile_false", (PyCFunction)THNPModule_is_jit_compile_false_wrap, METH_NOARGS, nullptr},	2289	{"_npu_is_jit_compile_false", (PyCFunction)THNPModule_is_jit_compile_false_wrap, METH_NOARGS, nullptr},
		2290	+ {"_npu_getMemoryFraction", (PyCFunction) THNPModule_getMemoryFraction, METH_VARARGS, nullptr},
2272	{"_npu_setMemoryFraction", (PyCFunction) THNPModule_setMemoryFraction, METH_VARARGS, nullptr},	2291	{"_npu_setMemoryFraction", (PyCFunction) THNPModule_setMemoryFraction, METH_VARARGS, nullptr},
2273	{"_npu_emptyCache", (PyCFunction) THNPModule_emptyCache, METH_NOARGS, nullptr},	2292	{"_npu_emptyCache", (PyCFunction) THNPModule_emptyCache, METH_NOARGS, nullptr},
2274	{"_npu_hostEmptyCache", (PyCFunction) THNPModule_npu_hostEmptyCache, METH_NOARGS, nullptr},	2293	{"_npu_hostEmptyCache", (PyCFunction) THNPModule_npu_hostEmptyCache, METH_NOARGS, nullptr},

Mtorch_npu/csrc/npu/NPUPluggableAllocator.cpp+8-0

	@@ -175,6 +175,14 @@ bool NPUPluggableAllocator::initialized()
175	return initialized_;	175	return initialized_;
176	}	176	}
177		177
		178	+double NPUPluggableAllocator::getMemoryFraction(int device)
		179	+{
		180	+ TORCH_CHECK(
		181	+ false,
		182	+ "NPUPluggableAllocator does not yet support getMemoryFraction. "
		183	+ "If you need it, please file an issue describing your use case.");
		184	+}
		185	+
178	void NPUPluggableAllocator::setMemoryFraction(double fraction, int device)	186	void NPUPluggableAllocator::setMemoryFraction(double fraction, int device)
179	{	187	{
180	if (memory_fraction_fn_) {	188	if (memory_fraction_fn_) {

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

	@@ -58,6 +58,7 @@ struct NPUPluggableAllocator
58	void raw_delete(void* ptr) override;	58	void raw_delete(void* ptr) override;
59	void init(int device_count) override;	59	void init(int device_count) override;
60	bool initialized() override;	60	bool initialized() override;
		61	+ double getMemoryFraction(int device) override;
61	void setMemoryFraction(double fraction, int device) override;	62	void setMemoryFraction(double fraction, int device) override;
62	void emptyCacheImpl(bool check_error, bool free_physical) override;	63	void emptyCacheImpl(bool check_error, bool free_physical) override;
63	void emptyCache(bool check_error) override;	64	void emptyCache(bool check_error) override;

Mtorch_npu/dynamo/trace_rule.py+2-0

Mtorch_npu/npu/__init__.py+1-0

	@@ -31,6 +31,7 @@ __all__ = [
31	"initial_seed",	31	"initial_seed",
32	"caching_allocator_alloc",	32	"caching_allocator_alloc",
33	"caching_allocator_delete",	33	"caching_allocator_delete",
		34	+ "get_per_process_memory_fraction",
34	"set_per_process_memory_fraction",	35	"set_per_process_memory_fraction",
35	"empty_cache",	36	"empty_cache",
36	"empty_virt_addr_cache",	37	"empty_virt_addr_cache",

Mtorch_npu/npu/memory.py+16-0

	@@ -18,6 +18,7 @@ from ._memory_viz import memory as _memory, segments as _segments
18	__all__ = [	18	__all__ = [
19	"caching_allocator_alloc",	19	"caching_allocator_alloc",
20	"caching_allocator_delete",	20	"caching_allocator_delete",
		21	+ "get_per_process_memory_fraction",
21	"set_per_process_memory_fraction",	22	"set_per_process_memory_fraction",
22	"empty_cache",	23	"empty_cache",
23	"empty_virt_addr_cache",	24	"empty_virt_addr_cache",
	@@ -123,6 +124,21 @@ def caching_allocator_delete(mem_ptr):
123	torch_npu._C._npu_npuCachingAllocator_raw_delete(mem_ptr)	124	torch_npu._C._npu_npuCachingAllocator_raw_delete(mem_ptr)
124		125
125		126
		127	+def get_per_process_memory_fraction(device=None) -> float:
		128	+ r"""Get memory fraction for a process.
		129	+ Args:
		130	+ device (torch.device or int, optional): selected device. If it is
		131	+ ``None`` the default NPU device is used.
		132	+ Returns:
		133	+ memory fraction, in range 0~1. Allowed memory equals total_memory * fraction.
		134	+ """
		135	+ _lazy_init()
		136	+ if device is None:
		137	+ device = torch_npu.npu.current_device()
		138	+ device = _get_device_index(device)
		139	+ return torch_npu._C._npu_getMemoryFraction(device)
		140	+
		141	+
126	def set_per_process_memory_fraction(fraction, device=None) -> None:	142	def set_per_process_memory_fraction(fraction, device=None) -> None:
127	r"""Set memory fraction for a process.	143	r"""Set memory fraction for a process.
128	The fraction is used to limit an caching allocator to allocated memory on a NPU device.	144	The fraction is used to limit an caching allocator to allocated memory on a NPU device.

	@@ -32,6 +32,7 @@ torch_non_c_binding_in_graph_functions_npu = dict.fromkeys(
32	"torch.npu.memory.reset_max_memory_cached",	32	"torch.npu.memory.reset_max_memory_cached",
33	"torch.npu.memory.reset_peak_host_memory_stats",	33	"torch.npu.memory.reset_peak_host_memory_stats",
34	"torch.npu.memory.reset_peak_memory_stats",	34	"torch.npu.memory.reset_peak_memory_stats",
		35	+ "torch.npu.memory.get_per_process_memory_fraction",
35	"torch.npu.memory.set_per_process_memory_fraction",	36	"torch.npu.memory.set_per_process_memory_fraction",
36	"torch.npu.random.manual_seed_all",	37	"torch.npu.random.manual_seed_all",
37	"torch.npu.random.manual_seed",	38	"torch.npu.random.manual_seed",
	@@ -67,6 +68,7 @@ torch_c_binding_in_graph_functions_npu = dict.fromkeys(
67	"torch_npu._C._npu_resetPeakMemoryStats",	68	"torch_npu._C._npu_resetPeakMemoryStats",
68	"torch_npu._C._npu_set_sync_debug_mode",	69	"torch_npu._C._npu_set_sync_debug_mode",
69	"torch_npu._C._npu_setDevice",	70	"torch_npu._C._npu_setDevice",
		71	+ "torch_npu._C._npu_getMemoryFraction",
70	"torch_npu._C._npu_setMemoryFraction",	72	"torch_npu._C._npu_setMemoryFraction",
71	"torch_npu._C._npu_synchronize",	73	"torch_npu._C._npu_synchronize",
72	"torch_npu._C._npu_resetAccumulatedMemoryStats",	74	"torch_npu._C._npu_resetAccumulatedMemoryStats",