[fix] support memory reuse in launch_host_func

已合并

[fix] support memory reuse in launch_host_func_pending #35787

kita-ikuyo创建于 5月15日从kita-ikuyo/pytorch： v2.9.0 合入到 Ascend/pytorch： v2.9.0

[fix] support memory reuse in launch_host_func_pending #35787

已合并

kita-ikuyo创建于 5月15日

从kita-ikuyo/pytorch v2.9.0合入到Ascend/pytorch v2.9.0

共 4 个文件变更+375-72

Mtest/npu/test_aclgraph_dfx.py+73-51

	@@ -1,14 +1,15 @@
1	import os	1	import os
2	import tempfile	2	import tempfile
3	import time	3	import time
		4	+import unittest
4	from unittest.mock import patch	5	from unittest.mock import patch
5		6
		7	+import torch
6	import torch_npu	8	import torch_npu
		9	+import torch_npu.npu.graphs as npu_graphs
7	from torch_npu.testing.common_utils import SkipIfNotGteCANNVersion	10	from torch_npu.testing.common_utils import SkipIfNotGteCANNVersion
8	from torch_npu.testing.testcase import run_tests, TestCase	11	from torch_npu.testing.testcase import run_tests, TestCase
9		12
10	-import torch
11	-
12		13
13	def wait_until(predicate, timeout=5.0, interval=0.01):	14	def wait_until(predicate, timeout=5.0, interval=0.01):
14	deadline = time.time() + timeout	15	deadline = time.time() + timeout
	@@ -27,11 +28,57 @@ def _resolved_npu_device_index(tensor):
27		28
28		29
29	class TestAclgraphDfx(TestCase):	30	class TestAclgraphDfx(TestCase):
		31	+
		32	+ def test_npugraph_tensor_ptr_spec_uses_raw_pointer_metadata(self):
		33	+ x = torch.arange(6, dtype=torch.float32).reshape(2, 3)
		34	+
		35	+ spec = npu_graphs._make_npugraph_tensor_ptr_spec(x)
		36	+
		37	+ self.assertEqual(spec[0], npu_graphs._NPUGRAPH_TENSOR_PTR_SPEC_MARKER)
		38	+ self.assertEqual(spec[1], x.data_ptr())
		39	+ self.assertEqual(spec[2], x.numel() * x.element_size())
		40	+ self.assertEqual(spec[3], tuple(x.shape))
		41	+ self.assertEqual(spec[4], x.dtype)
		42	+
		43	+ def test_materialize_npugraph_tensor_buffer_spec(self):
		44	+ expected = torch.arange(6, dtype=torch.float32).reshape(2, 3)
		45	+ buffer_spec = (
		46	+ npu_graphs._NPUGRAPH_TENSOR_BUFFER_SPEC_MARKER,
		47	+ bytearray(expected.numpy().tobytes()),
		48	+ tuple(expected.shape),
		49	+ expected.dtype,
		50	+ )
		51	+
		52	+ actual = npu_graphs._materialize_npugraph_tensor_arg(buffer_spec)
		53	+
		54	+ self.assertEqual(actual, expected)
		55	+
		56	+ def test_materialize_npugraph_tensor_buffer_spec_list(self):
		57	+ expected = [
		58	+ torch.arange(6, dtype=torch.float32).reshape(2, 3),
		59	+ torch.arange(4, dtype=torch.int32).reshape(2, 2),
		60	+ ]
		61	+ buffer_specs = [
		62	+ (
		63	+ npu_graphs._NPUGRAPH_TENSOR_BUFFER_SPEC_MARKER,
		64	+ bytearray(tensor.numpy().tobytes()),
		65	+ tuple(tensor.shape),
		66	+ tensor.dtype,
		67	+ )
		68	+ for tensor in expected
		69	+ ]
		70	+
		71	+ actual = npu_graphs._materialize_npugraph_tensor_arg(buffer_specs)
		72	+
		73	+ self.assertEqual(len(actual), len(expected))
		74	+ self.assertEqual(actual[0], expected[0])
		75	+ self.assertEqual(actual[1], expected[1])
		76	+
30	@SkipIfNotGteCANNVersion("8.5.0")	77	@SkipIfNotGteCANNVersion("8.5.0")
31	def test_print_npugraph_tensor(self):	78	def test_print_npugraph_tensor(self):
32	torch.npu.set_device(0)	79	torch.npu.set_device(0)
33	g = torch.npu.NPUGraph()	80	g = torch.npu.NPUGraph()
34	- x = torch.arange(6, dtype=torch.float32, device="npu").reshape(2, 3)	81	+ x = torch.arange(6, dtype=torch.float32, device='npu').reshape(2, 3)
35		82
36	with patch("builtins.print") as mock_print:	83	with patch("builtins.print") as mock_print:
37	with torch.npu.graph(g):	84	with torch.npu.graph(g):
	@@ -40,9 +87,7 @@ class TestAclgraphDfx(TestCase):
40	torch.npu.synchronize()	87	torch.npu.synchronize()
41	self.assertTrue(wait_until(lambda: mock_print.call_count > 0))	88	self.assertTrue(wait_until(lambda: mock_print.call_count > 0))
42		89
43	- printed_messages = [	90	+ printed_messages = [call.args[0] for call in mock_print.call_args_list if call.args]
44	- call.args[0] for call in mock_print.call_args_list if call.args
45	- ]
46	self.assertTrue(any("tensor=tensor(" in msg for msg in printed_messages))	91	self.assertTrue(any("tensor=tensor(" in msg for msg in printed_messages))
47	self.assertTrue(any("shape=(2, 3)" in msg for msg in printed_messages))	92	self.assertTrue(any("shape=(2, 3)" in msg for msg in printed_messages))
48	self.assertTrue(any("dtype=torch.float32" in msg for msg in printed_messages))	93	self.assertTrue(any("dtype=torch.float32" in msg for msg in printed_messages))
	@@ -51,7 +96,7 @@ class TestAclgraphDfx(TestCase):
51	def test_print_npugraph_tensor_with_default_message(self):	96	def test_print_npugraph_tensor_with_default_message(self):
52	torch.npu.set_device(0)	97	torch.npu.set_device(0)
53	g = torch.npu.NPUGraph()	98	g = torch.npu.NPUGraph()
54	- x = torch.arange(6, dtype=torch.float32, device="npu").reshape(2, 3)	99	+ x = torch.arange(6, dtype=torch.float32, device='npu').reshape(2, 3)
55		100
56	with patch("builtins.print") as mock_print:	101	with patch("builtins.print") as mock_print:
57	with torch.npu.graph(g):	102	with torch.npu.graph(g):
	@@ -60,16 +105,14 @@ class TestAclgraphDfx(TestCase):
60	torch.npu.synchronize()	105	torch.npu.synchronize()
61	self.assertTrue(wait_until(lambda: mock_print.call_count > 0))	106	self.assertTrue(wait_until(lambda: mock_print.call_count > 0))
62		107
63	- printed_messages = [	108	+ printed_messages = [call.args[0] for call in mock_print.call_args_list if call.args]
64	- call.args[0] for call in mock_print.call_args_list if call.args
65	- ]
66	self.assertTrue(any(msg.startswith("tensor(") for msg in printed_messages))	109	self.assertTrue(any(msg.startswith("tensor(") for msg in printed_messages))
67		110
68	@SkipIfNotGteCANNVersion("8.5.0")	111	@SkipIfNotGteCANNVersion("8.5.0")
69	def test_print_npugraph_tensor_with_args(self):	112	def test_print_npugraph_tensor_with_args(self):
70	torch.npu.set_device(0)	113	torch.npu.set_device(0)
71	g = torch.npu.NPUGraph()	114	g = torch.npu.NPUGraph()
72	- x = torch.arange(6, dtype=torch.float32, device="npu").reshape(2, 3)	115	+ x = torch.arange(6, dtype=torch.float32, device='npu').reshape(2, 3)
73		116
74	with patch("builtins.print") as mock_print:	117	with patch("builtins.print") as mock_print:
75	with torch.npu.graph(g):	118	with torch.npu.graph(g):
	@@ -78,9 +121,7 @@ class TestAclgraphDfx(TestCase):
78	torch.npu.synchronize()	121	torch.npu.synchronize()
79	self.assertTrue(wait_until(lambda: mock_print.call_count > 0))	122	self.assertTrue(wait_until(lambda: mock_print.call_count > 0))
80		123
81	- printed_messages = [	124	+ printed_messages = [call.args[0] for call in mock_print.call_args_list if call.args]
82	- call.args[0] for call in mock_print.call_args_list if call.args
83	- ]
84	self.assertTrue(any("x=tensor(" in msg for msg in printed_messages))	125	self.assertTrue(any("x=tensor(" in msg for msg in printed_messages))
85	self.assertTrue(any("shape=(2, 3)" in msg for msg in printed_messages))	126	self.assertTrue(any("shape=(2, 3)" in msg for msg in printed_messages))
86	self.assertTrue(any("dtype=torch.float32" in msg for msg in printed_messages))	127	self.assertTrue(any("dtype=torch.float32" in msg for msg in printed_messages))
	@@ -90,17 +131,13 @@ class TestAclgraphDfx(TestCase):
90	torch.npu.set_device(0)	131	torch.npu.set_device(0)
91	first_graph = torch.npu.NPUGraph()	132	first_graph = torch.npu.NPUGraph()
92	second_graph = torch.npu.NPUGraph()	133	second_graph = torch.npu.NPUGraph()
93	- x = torch.arange(6, dtype=torch.float32, device="npu").reshape(2, 3)	134	+ x = torch.arange(6, dtype=torch.float32, device='npu').reshape(2, 3)
94	device_index = _resolved_npu_device_index(x)	135	device_index = _resolved_npu_device_index(x)
95		136
96	with tempfile.TemporaryDirectory() as tmpdir:	137	with tempfile.TemporaryDirectory() as tmpdir:
97	save_path = os.path.join(tmpdir, "tensor.pt")	138	save_path = os.path.join(tmpdir, "tensor.pt")
98	- expected_counter_path = os.path.join(	139	+ expected_counter_path = os.path.join(tmpdir, f"tensor_device_{device_index}_0.pt")
99	- tmpdir, f"tensor_device_{device_index}_0.pt"	140	+ expected_second_counter_path = os.path.join(tmpdir, f"tensor_device_{device_index}_1.pt")
100	- )
101	- expected_second_counter_path = os.path.join(
102	- tmpdir, f"tensor_device_{device_index}_1.pt"
103	- )
104		141
105	with torch.npu.graph(first_graph):	142	with torch.npu.graph(first_graph):
106	torch.ops.npu.save_npugraph_tensor(x, save_path=save_path)	143	torch.ops.npu.save_npugraph_tensor(x, save_path=save_path)
	@@ -113,9 +150,7 @@ class TestAclgraphDfx(TestCase):
113	torch.ops.npu.save_npugraph_tensor(x, save_path=save_path)	150	torch.ops.npu.save_npugraph_tensor(x, save_path=save_path)
114	second_graph.replay()	151	second_graph.replay()
115	torch.npu.synchronize()	152	torch.npu.synchronize()
116	- self.assertTrue(	153	+ self.assertTrue(wait_until(lambda: os.path.exists(expected_second_counter_path)))
117	- wait_until(lambda: os.path.exists(expected_second_counter_path))
118	- )
119	self.assertEqual(torch.load(expected_second_counter_path), x.cpu())	154	self.assertEqual(torch.load(expected_second_counter_path), x.cpu())
120		155
121	@SkipIfNotGteCANNVersion("8.5.0")	156	@SkipIfNotGteCANNVersion("8.5.0")
	@@ -123,23 +158,17 @@ class TestAclgraphDfx(TestCase):
123	torch.npu.set_device(0)	158	torch.npu.set_device(0)
124	first_graph = torch.npu.NPUGraph()	159	first_graph = torch.npu.NPUGraph()
125	second_graph = torch.npu.NPUGraph()	160	second_graph = torch.npu.NPUGraph()
126	- x = torch.arange(6, dtype=torch.float32, device="npu").reshape(2, 3)	161	+ x = torch.arange(6, dtype=torch.float32, device='npu').reshape(2, 3)
127	- y = torch.arange(6, dtype=torch.float32, device="npu").reshape(2, 3) + 1	162	+ y = torch.arange(6, dtype=torch.float32, device='npu').reshape(2, 3) + 1
128	device_index = _resolved_npu_device_index(x)	163	device_index = _resolved_npu_device_index(x)
129		164
130	with tempfile.TemporaryDirectory() as tmpdir:	165	with tempfile.TemporaryDirectory() as tmpdir:
131	save_path = os.path.join(tmpdir, "tensor.pt")	166	save_path = os.path.join(tmpdir, "tensor.pt")
132	- expected_overwrite_path = os.path.join(	167	+ expected_overwrite_path = os.path.join(tmpdir, f"tensor_device_{device_index}.pt")
133	- tmpdir, f"tensor_device_{device_index}.pt"	168	+ unexpected_counter_path = os.path.join(tmpdir, f"tensor_device_{device_index}_0.pt")
134	- )
135	- unexpected_counter_path = os.path.join(
136	- tmpdir, f"tensor_device_{device_index}_0.pt"
137	- )
138		169
139	with torch.npu.graph(first_graph):	170	with torch.npu.graph(first_graph):
140	- torch.ops.npu.save_npugraph_tensor(	171	+ torch.ops.npu.save_npugraph_tensor(x, save_path=save_path, overwrite=True)
141	- x, save_path=save_path, overwrite=True
142	- )
143	first_graph.replay()	172	first_graph.replay()
144	torch.npu.synchronize()	173	torch.npu.synchronize()
145	self.assertTrue(wait_until(lambda: os.path.exists(expected_overwrite_path)))	174	self.assertTrue(wait_until(lambda: os.path.exists(expected_overwrite_path)))
	@@ -147,9 +176,7 @@ class TestAclgraphDfx(TestCase):
147	self.assertEqual(torch.load(expected_overwrite_path), x.cpu())	176	self.assertEqual(torch.load(expected_overwrite_path), x.cpu())
148		177
149	with torch.npu.graph(second_graph):	178	with torch.npu.graph(second_graph):
150	- torch.ops.npu.save_npugraph_tensor(	179	+ torch.ops.npu.save_npugraph_tensor(y, save_path=save_path, overwrite=True)
151	- y, save_path=save_path, overwrite=True
152	- )
153	second_graph.replay()	180	second_graph.replay()
154	torch.npu.synchronize()	181	torch.npu.synchronize()
155	self.assertEqual(torch.load(expected_overwrite_path), y.cpu())	182	self.assertEqual(torch.load(expected_overwrite_path), y.cpu())
	@@ -158,7 +185,7 @@ class TestAclgraphDfx(TestCase):
158	def test_save_npugraph_tensor_with_default_save_path(self):	185	def test_save_npugraph_tensor_with_default_save_path(self):
159	torch.npu.set_device(0)	186	torch.npu.set_device(0)
160	g = torch.npu.NPUGraph()	187	g = torch.npu.NPUGraph()
161	- x = torch.arange(6, dtype=torch.float32, device="npu").reshape(2, 3)	188	+ x = torch.arange(6, dtype=torch.float32, device='npu').reshape(2, 3)
162	device_index = _resolved_npu_device_index(x)	189	device_index = _resolved_npu_device_index(x)
163		190
164	with tempfile.TemporaryDirectory() as tmpdir:	191	with tempfile.TemporaryDirectory() as tmpdir:
	@@ -172,8 +199,7 @@ class TestAclgraphDfx(TestCase):
172		199
173	def default_saved_files():	200	def default_saved_files():
174	return [	201	return [
175	- file_name	202	+ file_name for file_name in os.listdir(tmpdir)
176	- for file_name in os.listdir(tmpdir)
177	if file_name.startswith("tensor_") and file_name.endswith(".pt")	203	if file_name.startswith("tensor_") and file_name.endswith(".pt")
178	]	204	]
179		205
	@@ -188,20 +214,16 @@ class TestAclgraphDfx(TestCase):
188	def test_save_npugraph_tensor_tensor_list(self):	214	def test_save_npugraph_tensor_tensor_list(self):
189	torch.npu.set_device(0)	215	torch.npu.set_device(0)
190	g = torch.npu.NPUGraph()	216	g = torch.npu.NPUGraph()
191	- x = torch.arange(6, dtype=torch.float32, device="npu").reshape(2, 3)	217	+ x = torch.arange(6, dtype=torch.float32, device='npu').reshape(2, 3)
192	- y = torch.arange(4, dtype=torch.float16, device="npu").reshape(2, 2)	218	+ y = torch.arange(4, dtype=torch.float16, device='npu').reshape(2, 2)
193	device_index = _resolved_npu_device_index(x)	219	device_index = _resolved_npu_device_index(x)
194		220
195	with tempfile.TemporaryDirectory() as tmpdir:	221	with tempfile.TemporaryDirectory() as tmpdir:
196	save_path = os.path.join(tmpdir, "tensor_list.pt")	222	save_path = os.path.join(tmpdir, "tensor_list.pt")
197	- expected_path = os.path.join(	223	+ expected_path = os.path.join(tmpdir, f"tensor_list_device_{device_index}_0.pt")
198	- tmpdir, f"tensor_list_device_{device_index}_0.pt"
199	- )
200		224
201	with torch.npu.graph(g):	225	with torch.npu.graph(g):
202	- torch.ops.npu.save_npugraph_tensor.tensor_list(	226	+ torch.ops.npu.save_npugraph_tensor.tensor_list([x, y], save_path=save_path)
203	- [x, y], save_path=save_path
204	- )
205	g.replay()	227	g.replay()
206	torch.npu.synchronize()	228	torch.npu.synchronize()
207	self.assertTrue(wait_until(lambda: os.path.exists(expected_path)))	229	self.assertTrue(wait_until(lambda: os.path.exists(expected_path)))
	@@ -212,5 +234,5 @@ class TestAclgraphDfx(TestCase):
212	self.assertEqual(saved[1], y.cpu())	234	self.assertEqual(saved[1], y.cpu())
213		235
214		236
215	-if __name__ == "__main__":	237	+if __name__ == '__main__':
216	- run_tests()	238	+ run_tests()

Mtorch_npu/csrc/npu/Graph.cpp+217-12

	@@ -1,7 +1,12 @@
1	#include <deque>	1	#include <deque>
		2	+#include <cstdint>
		3	+#include <cstring>
2	#include <thread>	4	#include <thread>
		5	+#include <utility>
3	#include <vector>	6	#include <vector>
4		7
		8	+#include <ATen/ATen.h>
		9	+
5	#include "torch_npu/csrc/npu/Graph.h"	10	#include "torch_npu/csrc/npu/Graph.h"
6		11
7	#include "op_plugin/OpApiInterface.h"	12	#include "op_plugin/OpApiInterface.h"
	@@ -19,7 +24,8 @@ constexpr auto pendingCallRetryInterval = std::chrono::milliseconds(1);
19	static ThreadArgs* threadArgs = nullptr;	24	static ThreadArgs* threadArgs = nullptr;
20	static uint64_t threadId = -1;	25	static uint64_t threadId = -1;
21	static std::mutex pendingCallbacksMutex;	26	static std::mutex pendingCallbacksMutex;
22	-static std::deque<PyFuncStruct*> pendingCallbacksQueue;	27	+using PendingCallbackEntry = std::pair<PendingCallPayload*, std::vector<at::Tensor>>;
		28	+static std::deque<PendingCallbackEntry> pendingCallbacksQueue;
23	static std::atomic<bool> pendingCallScheduled{false};	29	static std::atomic<bool> pendingCallScheduled{false};
24		30
25	void process_callback(void arg)	31	void process_callback(void arg)
	@@ -60,6 +66,9 @@ void LaunchCallFunc(void *userData)
60		66
61	namespace {	67	namespace {
62		68
		69	+constexpr const char* kNpuGraphTensorPtrSpecMarker = "host_tensor_ptr";
		70	+constexpr const char* kNpuGraphTensorBufferSpecMarker = "host_tensor_buffer";
		71	+
63	int PendingCallHandler(void *arg);	72	int PendingCallHandler(void *arg);
64		73
65	bool TrySchedulePendingCall()	74	bool TrySchedulePendingCall()
	@@ -82,25 +91,215 @@ void StartPendingCallRetryThread()
82	}).detach();	91	}).detach();
83	}	92	}
84		93
		94	+bool IsNpuGraphTensorPtrSpec(PyObject* obj)
		95	+{
		96	+ if (!PyTuple_Check(obj) \|\| PyTuple_GET_SIZE(obj) != 5) {
		97	+ return false;
		98	+ }
		99	+ PyObject* marker = PyTuple_GET_ITEM(obj, 0);
		100	+ return PyUnicode_Check(marker) &&
		101	+ PyUnicode_CompareWithASCIIString(marker, kNpuGraphTensorPtrSpecMarker) == 0;
		102	+}
		103	+
		104	+bool CollectNpuGraphTensorPtrSpec(PyObject* obj, PendingCallPayload* payload)
		105	+{
		106	+ PyObject* ptrObj = PyTuple_GET_ITEM(obj, 1);
		107	+ PyObject* nbytesObj = PyTuple_GET_ITEM(obj, 2);
		108	+ PyObject* shapeObj = PyTuple_GET_ITEM(obj, 3);
		109	+ PyObject* dtypeObj = PyTuple_GET_ITEM(obj, 4);
		110	+
		111	+ auto dataPtr = static_cast<uintptr_t>(PyLong_AsUnsignedLongLong(ptrObj));
		112	+ if (PyErr_Occurred()) {
		113	+ return false;
		114	+ }
		115	+ auto nbytes = static_cast<Py_ssize_t>(PyLong_AsSsize_t(nbytesObj));
		116	+ if (PyErr_Occurred()) {
		117	+ return false;
		118	+ }
		119	+ if (nbytes < 0) {
		120	+ PyErr_SetString(PyExc_ValueError, "npugraph tensor buffer size must be non-negative");
		121	+ return false;
		122	+ }
		123	+ payload->pendingTensorData.emplace_back(dataPtr, nbytes, shapeObj, dtypeObj);
		124	+ return true;
		125	+}
		126	+
		127	+bool CollectPendingTensorData(PyObject* obj, PendingCallPayload* payload)
		128	+{
		129	+ if (IsNpuGraphTensorPtrSpec(obj)) {
		130	+ return CollectNpuGraphTensorPtrSpec(obj, payload);
		131	+ }
		132	+
		133	+ if (PyTuple_Check(obj)) {
		134	+ Py_ssize_t size = PyTuple_GET_SIZE(obj);
		135	+ for (Py_ssize_t i = 0; i < size; ++i) {
		136	+ if (!CollectPendingTensorData(PyTuple_GET_ITEM(obj, i), payload)) {
		137	+ return false;
		138	+ }
		139	+ }
		140	+ return true;
		141	+ }
		142	+
		143	+ if (PyList_Check(obj)) {
		144	+ Py_ssize_t size = PyList_GET_SIZE(obj);
		145	+ for (Py_ssize_t i = 0; i < size; ++i) {
		146	+ if (!CollectPendingTensorData(PyList_GET_ITEM(obj, i), payload)) {
		147	+ return false;
		148	+ }
		149	+ }
		150	+ return true;
		151	+ }
		152	+
		153	+ return true;
		154	+}
		155	+
		156	+void CopyPendingTensorData(PendingCallPayload* payload, std::vector<at::Tensor>& copiedTensors)
		157	+{
		158	+ try {
		159	+ copiedTensors.clear();
		160	+ copiedTensors.reserve(payload->pendingTensorData.size());
		161	+ for (const auto& tensorData : payload->pendingTensorData) {
		162	+ at::Tensor copiedTensor = at::empty(
		163	+ {static_cast<int64_t>(tensorData.nbytes)},
		164	+ at::TensorOptions().dtype(at::kByte).device(at::kCPU));
		165	+ if (tensorData.nbytes > 0) {
		166	+ std::memcpy(
		167	+ copiedTensor.data_ptr(),
		168	+ reinterpret_cast<void*>(tensorData.dataPtr),
		169	+ static_cast<size_t>(tensorData.nbytes));
		170	+ }
		171	+ copiedTensors.emplace_back(std::move(copiedTensor));
		172	+ }
		173	+ } catch (...) {
		174	+ copiedTensors.clear();
		175	+ }
		176	+}
		177	+
		178	+PyObject* MaterializeNpuGraphTensorBufferSpec(
		179	+ PendingCallPayload* payload,
		180	+ const std::vector<at::Tensor>& copiedTensors,
		181	+ size_t& tensorDataIndex)
		182	+{
		183	+ if (tensorDataIndex >= payload->pendingTensorData.size() \|\|
		184	+ tensorDataIndex >= copiedTensors.size()) {
		185	+ PyErr_SetString(PyExc_RuntimeError, "npugraph tensor buffer index is out of range");
		186	+ return nullptr;
		187	+ }
		188	+ auto& tensorData = payload->pendingTensorData[tensorDataIndex];
		189	+ const auto& copiedTensor = copiedTensors[tensorDataIndex++];
		190	+ if (!copiedTensor.defined()) {
		191	+ PyErr_NoMemory();
		192	+ return nullptr;
		193	+ }
		194	+
		195	+ PyObject* buffer = PyMemoryView_FromMemory(
		196	+ static_cast<char*>(copiedTensor.data_ptr()),
		197	+ tensorData.nbytes,
		198	+ PyBUF_WRITE);
		199	+ if (buffer == nullptr) {
		200	+ return nullptr;
		201	+ }
		202	+
		203	+ PyObject* marker = PyUnicode_FromString(kNpuGraphTensorBufferSpecMarker);
		204	+ if (marker == nullptr) {
		205	+ Py_DECREF(buffer);
		206	+ return nullptr;
		207	+ }
		208	+ PyObject* bufferSpec = PyTuple_Pack(4, marker, buffer, tensorData.shape, tensorData.dtype);
		209	+ Py_DECREF(marker);
		210	+ Py_DECREF(buffer);
		211	+ return bufferSpec;
		212	+}
		213	+
		214	+PyObject* MaterializePendingArg(
		215	+ PyObject* obj,
		216	+ PendingCallPayload* payload,
		217	+ const std::vector<at::Tensor>& copiedTensors,
		218	+ size_t& tensorDataIndex)
		219	+{
		220	+ if (IsNpuGraphTensorPtrSpec(obj)) {
		221	+ return MaterializeNpuGraphTensorBufferSpec(payload, copiedTensors, tensorDataIndex);
		222	+ }
		223	+
		224	+ if (PyTuple_Check(obj)) {
		225	+ Py_ssize_t size = PyTuple_GET_SIZE(obj);
		226	+ PyObject* tuple = PyTuple_New(size);
		227	+ if (tuple == nullptr) {
		228	+ return nullptr;
		229	+ }
		230	+ for (Py_ssize_t i = 0; i < size; ++i) {
		231	+ PyObject* item = MaterializePendingArg(PyTuple_GET_ITEM(obj, i), payload, copiedTensors, tensorDataIndex);
		232	+ if (item == nullptr) {
		233	+ Py_DECREF(tuple);
		234	+ return nullptr;
		235	+ }
		236	+ PyTuple_SET_ITEM(tuple, i, item);
		237	+ }
		238	+ return tuple;
		239	+ }
		240	+
		241	+ if (PyList_Check(obj)) {
		242	+ Py_ssize_t size = PyList_GET_SIZE(obj);
		243	+ PyObject* list = PyList_New(size);
		244	+ if (list == nullptr) {
		245	+ return nullptr;
		246	+ }
		247	+ for (Py_ssize_t i = 0; i < size; ++i) {
		248	+ PyObject* item = MaterializePendingArg(PyList_GET_ITEM(obj, i), payload, copiedTensors, tensorDataIndex);
		249	+ if (item == nullptr) {
		250	+ Py_DECREF(list);
		251	+ return nullptr;
		252	+ }
		253	+ PyList_SET_ITEM(list, i, item);
		254	+ }
		255	+ return list;
		256	+ }
		257	+
		258	+ Py_INCREF(obj);
		259	+ return obj;
		260	+}
		261	+
		262	+PyObject* MaterializePendingArgs(PendingCallPayload* payload, const std::vector<at::Tensor>& copiedTensors)
		263	+{
		264	+ size_t tensorDataIndex = 0;
		265	+ PyObject* materializedArgs = MaterializePendingArg(
		266	+ payload->pyFuncData.pyFuncArgs, payload, copiedTensors, tensorDataIndex);
		267	+ if (materializedArgs == nullptr) {
		268	+ PyErr_WriteUnraisable(payload->pyFuncData.pyFunc);
		269	+ PyErr_Clear();
		270	+ return nullptr;
		271	+ }
		272	+ return materializedArgs;
		273	+}
		274	+
85	int PendingCallHandler(void *arg)	275	int PendingCallHandler(void *arg)
86	{	276	{
87	- std::deque<PyFuncStruct*> readyCallbacks;	277	+ std::deque<PendingCallbackEntry> readyCallbacks;
88	{	278	{
89	std::lock_guard<std::mutex> lock(pendingCallbacksMutex);	279	std::lock_guard<std::mutex> lock(pendingCallbacksMutex);
90	pendingCallScheduled.store(false, std::memory_order_release);	280	pendingCallScheduled.store(false, std::memory_order_release);
91	readyCallbacks.swap(pendingCallbacksQueue);	281	readyCallbacks.swap(pendingCallbacksQueue);
92	}	282	}
93		283
94	- for (auto* data : readyCallbacks) {	284	+ for (auto& callback : readyCallbacks) {
95	- if (data == nullptr) {	285	+ auto* payload = callback.first;
		286	+ auto& copiedTensors = callback.second;
		287	+ if (payload == nullptr) {
96	continue;	288	continue;
97	}	289	}
98	- PyObject* result = PyObject_CallObject(data->pyFunc, data->pyFuncArgs);	290	+ PyObject* materializedArgs = MaterializePendingArgs(payload, copiedTensors);
		291	+ if (materializedArgs == nullptr) {
		292	+ copiedTensors.clear();
		293	+ continue;
		294	+ }
		295	+ PyObject* result = PyObject_CallObject(payload->pyFuncData.pyFunc, materializedArgs);
99	if (result != nullptr) {	296	if (result != nullptr) {
100	Py_XDECREF(result);	297	Py_XDECREF(result);
101	} else {	298	} else {
102	- PyErr_WriteUnraisable(data->pyFunc);	299	+ PyErr_WriteUnraisable(payload->pyFuncData.pyFunc);
103	}	300	}
		301	+ Py_DECREF(materializedArgs);
		302	+ copiedTensors.clear();
104	}	303	}
105		304
106	bool shouldScheduleAgain = false;	305	bool shouldScheduleAgain = false;
	@@ -121,15 +320,18 @@ int PendingCallHandler(void *arg)
121		320
122	void LaunchCallbackViaPendingCall(void *userData)	321	void LaunchCallbackViaPendingCall(void *userData)
123	{	322	{
124	- auto* data = static_cast<PyFuncStruct*>(userData);	323	+ auto* payload = static_cast<PendingCallPayload*>(userData);
125	- if (data == nullptr) {	324	+ if (payload == nullptr) {
126	return;	325	return;
127	}	326	}
128		327
		328	+ std::vector<at::Tensor> copiedTensors;
		329	+ CopyPendingTensorData(payload, copiedTensors);
		330	+
129	bool shouldSchedule = false;	331	bool shouldSchedule = false;
130	{	332	{
131	std::lock_guard<std::mutex> lock(pendingCallbacksMutex);	333	std::lock_guard<std::mutex> lock(pendingCallbacksMutex);
132	- pendingCallbacksQueue.emplace_back(data);	334	+ pendingCallbacksQueue.emplace_back(payload, std::move(copiedTensors));
133	if (!pendingCallScheduled.exchange(true, std::memory_order_acq_rel)) {	335	if (!pendingCallScheduled.exchange(true, std::memory_order_acq_rel)) {
134	shouldSchedule = true;	336	shouldSchedule = true;
135	}	337	}
	@@ -340,9 +542,12 @@ void TORCH_NPU_API THNPGraph_init(PyObject* module) {
340	auto func = (*py_func).ptr();	542	auto func = (*py_func).ptr();
341	auto userDataList = (*py_data).ptr();	543	auto userDataList = (*py_data).ptr();
342	auto stream = THNPUtils_PyObject_to_NPUStream((*py_stream).ptr());	544	auto stream = THNPUtils_PyObject_to_NPUStream((*py_stream).ptr());
343	- auto data = std::make_unique<PyFuncStruct>(func, userDataList);	545	+ auto payload = std::make_unique<PendingCallPayload>(func, userDataList);
344	- c10_npu::launch_host_func(stream, LaunchCallbackViaPendingCall, data.get());	546	+ if (!CollectPendingTensorData(userDataList, payload.get())) {
345	- (void)data.release();	547	+ throw py::error_already_set();
		548	+ }
		549	+ c10_npu::launch_host_func(stream, LaunchCallbackViaPendingCall, payload.get());
		550	+ (void)payload.release();
346	})	551	})
347	.def("_subscribe_report", [](py::object py_stream) {	552	.def("_subscribe_report", [](py::object py_stream) {
348	auto stream = (*py_stream).ptr();	553	auto stream = (*py_stream).ptr();

Mtorch_npu/csrc/npu/Graph.h+35-0

	@@ -1,4 +1,6 @@
1	#include <torch/csrc/python_headers.h>	1	#include <torch/csrc/python_headers.h>
		2	+#include <cstdint>
		3	+#include <vector>
2	#include <pybind11/chrono.h>	4	#include <pybind11/chrono.h>
3	#include <torch/csrc/jit/python/pybind_utils.h>	5	#include <torch/csrc/jit/python/pybind_utils.h>
4	#include <torch/csrc/utils/pybind.h>	6	#include <torch/csrc/utils/pybind.h>
	@@ -7,6 +9,20 @@
7	#include "third_party/acl/inc/acl/acl_rt.h"	9	#include "third_party/acl/inc/acl/acl_rt.h"
8	#include "third_party/acl/inc/acl/acl_sk.h"	10	#include "third_party/acl/inc/acl/acl_sk.h"
9		11
		12	+struct PendingTensorData {
		13	+ PendingTensorData(uintptr_t dataPtr, Py_ssize_t nbytes, PyObject* shape, PyObject* dtype)
		14	+ : dataPtr(dataPtr), nbytes(nbytes), shape(shape), dtype(dtype)
		15	+ {
		16	+ Py_XINCREF(shape);
		17	+ Py_XINCREF(dtype);
		18	+ }
		19	+
		20	+ uintptr_t dataPtr = 0;
		21	+ Py_ssize_t nbytes = 0;
		22	+ PyObject* shape = nullptr;
		23	+ PyObject* dtype = nullptr;
		24	+};
		25	+
10	struct PyFuncStruct {	26	struct PyFuncStruct {
11	PyFuncStruct(PyObject pyFunc, PyObject pyFuncArgs)	27	PyFuncStruct(PyObject pyFunc, PyObject pyFuncArgs)
12	: pyFunc(pyFunc), pyFuncArgs(pyFuncArgs)	28	: pyFunc(pyFunc), pyFuncArgs(pyFuncArgs)
	@@ -25,6 +41,25 @@ struct PyFuncStruct {
25	PyObject* pyFuncArgs = nullptr;	41	PyObject* pyFuncArgs = nullptr;
26	};	42	};
27		43
		44	+struct PendingCallPayload {
		45	+ PendingCallPayload(PyObject* pyFunc, PyObject* pyFuncArgs)
		46	+ : pyFuncData(pyFunc, pyFuncArgs)
		47	+ {
		48	+ }
		49	+
		50	+ ~PendingCallPayload()
		51	+ {
		52	+ Py_CLEAR(pyFuncData.pyFuncArgs);
		53	+ for (auto& tensorData : pendingTensorData) {
		54	+ Py_XDECREF(tensorData.shape);
		55	+ Py_XDECREF(tensorData.dtype);
		56	+ }
		57	+ }
		58	+
		59	+ PyFuncStruct pyFuncData;
		60	+ std::vector<PendingTensorData> pendingTensorData;
		61	+};
		62	+
28	struct ThreadArgs {	63	struct ThreadArgs {
29	ThreadArgs(aclrtContext context, bool exitFlag)	64	ThreadArgs(aclrtContext context, bool exitFlag)
30	: context(context), exitFlag(exitFlag) {}	65	: context(context), exitFlag(exitFlag) {}

Mtorch_npu/npu/graphs.py+50-9

	@@ -122,6 +122,8 @@ _save_npugraph_tensor_lock = threading.Lock()
122	_save_npugraph_tensor_counters = {}	122	_save_npugraph_tensor_counters = {}
123	_save_tensor_streams: Dict[int, "torch_npu.npu.Stream"] = {}	123	_save_tensor_streams: Dict[int, "torch_npu.npu.Stream"] = {}
124	_save_tensor_stream_lock = threading.Lock()	124	_save_tensor_stream_lock = threading.Lock()
		125	+_NPUGRAPH_TENSOR_PTR_SPEC_MARKER = "host_tensor_ptr"
		126	+_NPUGRAPH_TENSOR_BUFFER_SPEC_MARKER = "host_tensor_buffer"
125		127
126		128
127	def _get_save_tensor_stream(device_index: int):	129	def _get_save_tensor_stream(device_index: int):
	@@ -167,6 +169,7 @@ def _build_save_npugraph_tensor_path(save_path=None, device_index=None, overwrit
167		169
168		170
169	def _print_callback_pending(tensor_name, tensor_arg):	171	def _print_callback_pending(tensor_name, tensor_arg):
		172	+ tensor_arg = _materialize_npugraph_tensor_arg(tensor_arg)
170	output = str(tensor_arg)	173	output = str(tensor_arg)
171	if tensor_name is not None:	174	if tensor_name is not None:
172	output = f"{tensor_name}={output}"	175	output = f"{tensor_name}={output}"
	@@ -175,9 +178,35 @@ def _print_callback_pending(tensor_name, tensor_arg):
175		178
176		179
177	def _save_callback_pending(tensor_arg, str_arg):	180	def _save_callback_pending(tensor_arg, str_arg):
		181	+ tensor_arg = _materialize_npugraph_tensor_arg(tensor_arg)
178	torch.save(tensor_arg, str_arg)	182	torch.save(tensor_arg, str_arg)
179		183
180		184
		185	+def _make_npugraph_tensor_ptr_spec(tensor):
		186	+ return (
		187	+ _NPUGRAPH_TENSOR_PTR_SPEC_MARKER,
		188	+ tensor.data_ptr(),
		189	+ tensor.numel() * tensor.element_size(),
		190	+ tuple(tensor.shape),
		191	+ tensor.dtype,
		192	+ )
		193	+
		194	+
		195	+def _materialize_npugraph_tensor_arg(tensor_arg):
		196	+ if (
		197	+ isinstance(tensor_arg, tuple)
		198	+ and len(tensor_arg) == 4
		199	+ and tensor_arg[0] == _NPUGRAPH_TENSOR_BUFFER_SPEC_MARKER
		200	+ ):
		201	+ _, buffer, shape, dtype = tensor_arg
		202	+ if len(buffer) == 0:
		203	+ return torch.empty(shape, dtype=dtype)
		204	+ return torch.frombuffer(buffer, dtype=dtype).reshape(shape)
		205	+ if isinstance(tensor_arg, list):
		206	+ return [_materialize_npugraph_tensor_arg(arg) for arg in tensor_arg]
		207	+ return tensor_arg
		208	+
		209	+
181	def _validate_tensor_list(inputs):	210	def _validate_tensor_list(inputs):
182	if not isinstance(inputs, (list, tuple)):	211	if not isinstance(inputs, (list, tuple)):
183	raise TypeError(f"input must be Tensor or TensorList, but got {type(inputs).__name__}")	212	raise TypeError(f"input must be Tensor or TensorList, but got {type(inputs).__name__}")
	@@ -197,9 +226,12 @@ def _print_npugraph_tensor_impl(input, tensor_name=None):
197	return	226	return
198		227
199	device = input.device	228	device = input.device
200	- if device.type != "npu":	229	+ if device.type == "cpu":
201	_print_callback_pending(tensor_name, input)	230	_print_callback_pending(tensor_name, input)
202	return	231	return
		232	+
		233	+ if device.type != "npu":
		234	+ return
203		235
204	device_index = device.index	236	device_index = device.index
205	save_stream = _get_save_tensor_stream(device_index)	237	save_stream = _get_save_tensor_stream(device_index)
	@@ -212,7 +244,8 @@ def _print_npugraph_tensor_impl(input, tensor_name=None):
212	with torch.npu.stream(save_stream):	244	with torch.npu.stream(save_stream):
213	# Wait for the original stream to complete before D2H	245	# Wait for the original stream to complete before D2H
214	event1.wait()	246	event1.wait()
215	- cpu_arg = input.to("cpu", non_blocking=True)	247	+ cpu_tensor = input.to("cpu", non_blocking=True)
		248	+ cpu_arg = _make_npugraph_tensor_ptr_spec(cpu_tensor)
216	# Get current stream inside context (which is save_stream)	249	# Get current stream inside context (which is save_stream)
217	current_stream = torch.npu.current_stream()	250	current_stream = torch.npu.current_stream()
218	torch_npu.npu._launch_host_func_pending(	251	torch_npu.npu._launch_host_func_pending(
	@@ -232,9 +265,12 @@ def _save_npugraph_tensor_impl(input, save_path=None, overwrite=False):
232	return	265	return
233		266
234	device = input.device	267	device = input.device
235	- if device.type != "npu":	268	+ if device.type == "cpu":
236	torch.save(input, _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))	269	torch.save(input, _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))
237	return	270	return
		271	+
		272	+ if device.type != "npu":
		273	+ return
238		274
239	device_index = device.index	275	device_index = device.index
240	save_stream = _get_save_tensor_stream(device_index)	276	save_stream = _get_save_tensor_stream(device_index)
	@@ -248,7 +284,8 @@ def _save_npugraph_tensor_impl(input, save_path=None, overwrite=False):
248	with torch.npu.stream(save_stream):	284	with torch.npu.stream(save_stream):
249	# Wait for the original stream to complete before D2H	285	# Wait for the original stream to complete before D2H
250	event1.wait()	286	event1.wait()
251	- cpu_arg = input.to("cpu", non_blocking=True)	287	+ cpu_tensor = input.to("cpu", non_blocking=True)
		288	+ cpu_arg = _make_npugraph_tensor_ptr_spec(cpu_tensor)
252	# Get current stream inside context (which is save_stream)	289	# Get current stream inside context (which is save_stream)
253	current_stream = torch.npu.current_stream()	290	current_stream = torch.npu.current_stream()
254	torch_npu.npu._launch_host_func_pending(	291	torch_npu.npu._launch_host_func_pending(
	@@ -265,9 +302,12 @@ def _save_npugraph_tensor_impl(input, save_path=None, overwrite=False):
265		302
266	def _save_npugraph_tensor_tensor_list_impl(input, save_path=None, overwrite=False):	303	def _save_npugraph_tensor_tensor_list_impl(input, save_path=None, overwrite=False):
267	device = _validate_tensor_list(input)	304	device = _validate_tensor_list(input)
268	- if device.type != "npu":	305	+ if device.type == "cpu":
269	torch.save(list(input), _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))	306	torch.save(list(input), _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))
270	return	307	return
		308	+
		309	+ if device.type != "npu":
		310	+ return
271		311
272	device_index = device.index	312	device_index = device.index
273	save_stream = _get_save_tensor_stream(device_index)	313	save_stream = _get_save_tensor_stream(device_index)
	@@ -281,7 +321,8 @@ def _save_npugraph_tensor_tensor_list_impl(input, save_path=None, overwrite=Fals
281	with torch.npu.stream(save_stream):	321	with torch.npu.stream(save_stream):
282	# Wait for the original stream to complete before D2H	322	# Wait for the original stream to complete before D2H
283	event1.wait()	323	event1.wait()
284	- cpu_args = [tensor.to("cpu", non_blocking=True) for tensor in input]	324	+ cpu_tensors = [tensor.to("cpu", non_blocking=True) for tensor in input]
		325	+ cpu_args = [_make_npugraph_tensor_ptr_spec(tensor) for tensor in cpu_tensors]
285	# Get current stream inside context (which is save_stream)	326	# Get current stream inside context (which is save_stream)
286	current_stream = torch.npu.current_stream()	327	current_stream = torch.npu.current_stream()
287	torch_npu.npu._launch_host_func_pending(	328	torch_npu.npu._launch_host_func_pending(
	@@ -329,7 +370,7 @@ if not hasattr(torch.ops.npu, "print_npugraph_tensor"):
329	_npu_lib.define("print_npugraph_tensor(Tensor input, *, str? tensor_name=None) -> ()")	370	_npu_lib.define("print_npugraph_tensor(Tensor input, *, str? tensor_name=None) -> ()")
330		371
331	_npu_lib.impl("print_npugraph_tensor", _print_npugraph_tensor_impl, "PrivateUse1")	372	_npu_lib.impl("print_npugraph_tensor", _print_npugraph_tensor_impl, "PrivateUse1")
332	- _npu_lib.impl("print_npugraph_tensor", lambda input, *, tensor_name=None: None, "CPU")	373	+ _npu_lib.impl("print_npugraph_tensor", _print_npugraph_tensor_impl, "CPU")
333		374
334	has_side_effect(torch.ops.npu.print_npugraph_tensor.default)	375	has_side_effect(torch.ops.npu.print_npugraph_tensor.default)
335		376
	@@ -343,9 +384,9 @@ if not hasattr(torch.ops.npu, "save_npugraph_tensor"):
343	_npu_lib.define("save_npugraph_tensor.tensor_list(Tensor[] input, *, str? save_path=None, bool overwrite=False) -> ()")	384	_npu_lib.define("save_npugraph_tensor.tensor_list(Tensor[] input, *, str? save_path=None, bool overwrite=False) -> ()")
344		385
345	_npu_lib.impl("save_npugraph_tensor", _save_npugraph_tensor_impl, "PrivateUse1")	386	_npu_lib.impl("save_npugraph_tensor", _save_npugraph_tensor_impl, "PrivateUse1")
346	- _npu_lib.impl("save_npugraph_tensor", lambda input, *, save_path=None, overwrite=False: None, "CPU")	387	+ _npu_lib.impl("save_npugraph_tensor", _save_npugraph_tensor_impl, "CPU")
347	_npu_lib.impl("save_npugraph_tensor.tensor_list", _save_npugraph_tensor_tensor_list_impl, "PrivateUse1")	388	_npu_lib.impl("save_npugraph_tensor.tensor_list", _save_npugraph_tensor_tensor_list_impl, "PrivateUse1")
348	- _npu_lib.impl("save_npugraph_tensor.tensor_list", lambda input, *, save_path=None, overwrite=False: None, "CPU")	389	+ _npu_lib.impl("save_npugraph_tensor.tensor_list", _save_npugraph_tensor_tensor_list_impl, "CPU")
349		390
350	has_side_effect(torch.ops.npu.save_npugraph_tensor.default)	391	has_side_effect(torch.ops.npu.save_npugraph_tensor.default)
351	has_side_effect(torch.ops.npu.save_npugraph_tensor.tensor_list)	392	has_side_effect(torch.ops.npu.save_npugraph_tensor.tensor_list)