Star844
1306
代码
代码
Issues1405
Pull Requests2246
流水线
项目讨论
Wiki
分析
Star844
1306
已合并
支持c++外部流接口getStreamFromExternal #34383
pwc10490创建于 4月25日
支持c++外部流接口getStreamFromExternal #34383
已合并
pwc10490创建于 4月25日
从pwc10490/pytorch_stream_patchv2.9.0合入到Ascend/pytorchv2.9.0
共 14 个文件变更+712-10
Atest/cpp_extensions/external_stream_test.cpp+349-0
@@ -0,0 +1,349 @@
1+#include <torch/extension.h>
2+#include <c10/util/irange.h>
3+ 
4+#include "torch_npu/csrc/core/npu/NPUStream.h"
5+#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
6+#include "torch_npu/csrc/core/npu/NPUEvent.h"
7+#include "torch_npu/csrc/core/npu/NPUGraph.h"
8+#include "third_party/acl/inc/acl/acl.h"
9+ 
10+using namespace at;
11+ 
12+// Get NPU device - assumes torch_npu is already initialized via import
13+static c10::Device get_npu_device() {
14+ return c10::Device(c10::DeviceType::PrivateUse1, 0);
15+}
16+ 
17+// Test 1: External stream creation
18+bool test_external_stream_creation()
19+{
20+ // Create ACL stream directly
21+ aclrtStream acl_stream = nullptr;
22+ auto ret = aclrtCreateStream(&acl_stream);
23+ if (ret != ACL_ERROR_NONE || acl_stream == nullptr) {
24+ return false;
25+ }
26+ 
27+ // Wrap as NPUStream
28+ c10_npu::NPUStream npu_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
29+ 
30+ // Verify stream_id (bit 30 marker for external)
31+ c10::StreamId stream_id = npu_stream.id();
32+ if (!((static_cast<uint64_t>(stream_id) & (1ULL << 30)) != 0)) {
33+ aclrtDestroyStream(acl_stream);
34+ return false;
35+ }
36+ 
37+ // Verify aclrtStream is preserved
38+ if (npu_stream.stream() != acl_stream) {
39+ aclrtDestroyStream(acl_stream);
40+ return false;
41+ }
42+ 
43+ // Default stream should have different encoding (bit 30 = 0)
44+ c10_npu::NPUStream default_stream = c10_npu::getDefaultNPUStream(0);
45+ c10::StreamId default_id = default_stream.id();
46+ if ((static_cast<uint64_t>(default_id) & (1ULL << 30)) != 0) {
47+ aclrtDestroyStream(acl_stream);
48+ return false;
49+ }
50+ 
51+ aclrtDestroyStream(acl_stream);
52+ return true;
53+}
54+ 
55+// Test 2: External stream as current stream
56+bool test_external_stream_as_current()
57+{
58+ aclrtStream acl_stream = nullptr;
59+ auto ret = aclrtCreateStream(&acl_stream);
60+ if (ret != ACL_ERROR_NONE) {
61+ return false;
62+ }
63+ 
64+ c10_npu::NPUStream ext_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
65+ 
66+ // Set current stream to external
67+ c10_npu::setCurrentNPUStream(ext_stream);
68+ 
69+ // Get current stream - should be the external stream
70+ c10_npu::NPUStream current = c10_npu::getCurrentNPUStream(0);
71+ if (current.id() != ext_stream.id()) {
72+ c10_npu::setCurrentNPUStream(c10_npu::getDefaultNPUStream(0));
73+ aclrtDestroyStream(acl_stream);
74+ return false;
75+ }
76+ 
77+ // Reset to default stream
78+ c10_npu::setCurrentNPUStream(c10_npu::getDefaultNPUStream(0));
79+ 
80+ aclrtDestroyStream(acl_stream);
81+ return true;
82+}
83+ 
84+// Test 3: Tensor operations on external stream
85+bool test_operations_on_external_stream()
86+{
87+ auto npu_device = get_npu_device();
88+ 
89+ aclrtStream acl_stream = nullptr;
90+ auto ret = aclrtCreateStream(&acl_stream);
91+ if (ret != ACL_ERROR_NONE) {
92+ return false;
93+ }
94+ 
95+ c10_npu::NPUStream ext_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
96+ 
97+ // Set current stream to external
98+ c10_npu::setCurrentNPUStream(ext_stream);
99+ 
100+ // Execute operations - should work synchronously on external stream
101+ auto a = torch::randn({2, 3}).to(npu_device);
102+ auto b = torch::randn({2, 3}).to(npu_device);
103+ auto c = a + b;
104+ 
105+ bool success = (c.size(0) == 2 && c.size(1) == 3);
106+ 
107+ // Reset to default stream BEFORE tensor destruction
108+ c10_npu::setCurrentNPUStream(c10_npu::getDefaultNPUStream(0));
109+ 
110+ aclrtDestroyStream(acl_stream);
111+ return success;
112+}
113+ 
114+// Test 4: Event record restriction
115+bool test_event_record_restriction()
116+{
117+ aclrtStream acl_stream = nullptr;
118+ auto ret = aclrtCreateStream(&acl_stream);
119+ if (ret != ACL_ERROR_NONE) {
120+ return false;
121+ }
122+ 
123+ c10_npu::NPUStream ext_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
124+ 
125+ c10_npu::NPUEvent event;
126+ 
127+ // record() should throw for external stream
128+ bool threw = false;
129+ try {
130+ event.record(ext_stream);
131+ } catch (const c10::Error&) {
132+ threw = true;
133+ }
134+ 
135+ aclrtDestroyStream(acl_stream);
136+ return threw;
137+}
138+ 
139+// Test 5: Event block restriction
140+bool test_event_block_restriction()
141+{
142+ aclrtStream acl_stream = nullptr;
143+ auto ret = aclrtCreateStream(&acl_stream);
144+ if (ret != ACL_ERROR_NONE) {
145+ return false;
146+ }
147+ 
148+ c10_npu::NPUStream ext_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
149+ 
150+ c10_npu::NPUEvent event;
151+ 
152+ // block() should throw for external stream
153+ bool threw = false;
154+ try {
155+ event.block(ext_stream);
156+ } catch (const c10::Error&) {
157+ threw = true;
158+ }
159+ 
160+ aclrtDestroyStream(acl_stream);
161+ return threw;
162+}
163+ 
164+// Test 6: Graph capture restriction
165+bool test_graph_capture_restriction()
166+{
167+ aclrtStream acl_stream = nullptr;
168+ auto ret = aclrtCreateStream(&acl_stream);
169+ if (ret != ACL_ERROR_NONE) {
170+ return false;
171+ }
172+ 
173+ c10_npu::NPUStream ext_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
174+ 
175+ // Set current stream to external
176+ c10_npu::setCurrentNPUStream(ext_stream);
177+ 
178+ c10_npu::NPUGraph graph;
179+ 
180+ // capture_begin should throw
181+ bool threw = false;
182+ try {
183+ graph.capture_begin();
184+ } catch (const c10::Error&) {
185+ threw = true;
186+ }
187+ 
188+ // Reset to default stream
189+ c10_npu::setCurrentNPUStream(c10_npu::getDefaultNPUStream(0));
190+ 
191+ aclrtDestroyStream(acl_stream);
192+ return threw;
193+}
194+ 
195+// Test 7: Query restriction
196+bool test_query_restriction()
197+{
198+ aclrtStream acl_stream = nullptr;
199+ auto ret = aclrtCreateStream(&acl_stream);
200+ if (ret != ACL_ERROR_NONE) {
201+ return false;
202+ }
203+ 
204+ c10_npu::NPUStream ext_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
205+ 
206+ // query() should throw
207+ bool threw = false;
208+ try {
209+ ext_stream.query();
210+ } catch (const c10::Error&) {
211+ threw = true;
212+ }
213+ 
214+ aclrtDestroyStream(acl_stream);
215+ return threw;
216+}
217+ 
218+// Test 8: Synchronize restriction
219+bool test_synchronize_restriction()
220+{
221+ aclrtStream acl_stream = nullptr;
222+ auto ret = aclrtCreateStream(&acl_stream);
223+ if (ret != ACL_ERROR_NONE) {
224+ return false;
225+ }
226+ 
227+ c10_npu::NPUStream ext_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
228+ 
229+ // synchronize() should throw
230+ bool threw = false;
231+ try {
232+ ext_stream.synchronize();
233+ } catch (const c10::Error&) {
234+ threw = true;
235+ }
236+ 
237+ aclrtDestroyStream(acl_stream);
238+ return threw;
239+}
240+ 
241+// Test 9: isSyncLaunchStream
242+bool test_is_sync_launch_stream()
243+{
244+ aclrtStream acl_stream = nullptr;
245+ auto ret = aclrtCreateStream(&acl_stream);
246+ if (ret != ACL_ERROR_NONE) {
247+ return false;
248+ }
249+ 
250+ c10_npu::NPUStream ext_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
251+ 
252+ // External stream should NOT be sync launch stream
253+ bool is_sync = ext_stream.isSyncLaunchStream();
254+ 
255+ aclrtDestroyStream(acl_stream);
256+ return !is_sync;
257+}
258+ 
259+// Test 10: Multiple external streams
260+bool test_multiple_external_streams()
261+{
262+ aclrtStream acl_stream1 = nullptr;
263+ aclrtStream acl_stream2 = nullptr;
264+ auto ret1 = aclrtCreateStream(&acl_stream1);
265+ auto ret2 = aclrtCreateStream(&acl_stream2);
266+ 
267+ if (ret1 != ACL_ERROR_NONE || ret2 != ACL_ERROR_NONE) {
268+ if (acl_stream1) aclrtDestroyStream(acl_stream1);
269+ if (acl_stream2) aclrtDestroyStream(acl_stream2);
270+ return false;
271+ }
272+ 
273+ c10_npu::NPUStream npu_stream1 = c10_npu::getStreamFromExternal(acl_stream1, 0);
274+ c10_npu::NPUStream npu_stream2 = c10_npu::getStreamFromExternal(acl_stream2, 0);
275+ 
276+ // Stream IDs should be different
277+ bool ids_different = (npu_stream1.id() != npu_stream2.id());
278+ 
279+ // Both should have bit 30 marker
280+ bool both_external = ((static_cast<uint64_t>(npu_stream1.id()) & (1ULL << 30)) != 0) &&
281+ ((static_cast<uint64_t>(npu_stream2.id()) & (1ULL << 30)) != 0);
282+ 
283+ // aclrtStream values should be preserved
284+ bool streams_preserved = (npu_stream1.stream() == acl_stream1) &&
285+ (npu_stream2.stream() == acl_stream2);
286+ 
287+ aclrtDestroyStream(acl_stream1);
288+ aclrtDestroyStream(acl_stream2);
289+ 
290+ return ids_different && both_external && streams_preserved;
291+}
292+ 
293+// Test 11: Same acl_stream returns same NPUStream
294+bool test_same_acl_stream_same_npu_stream()
295+{
296+ aclrtStream acl_stream = nullptr;
297+ auto ret = aclrtCreateStream(&acl_stream);
298+ if (ret != ACL_ERROR_NONE) {
299+ return false;
300+ }
301+ 
302+ // Call getStreamFromExternal twice with same acl_stream
303+ c10_npu::NPUStream npu_stream1 = c10_npu::getStreamFromExternal(acl_stream, 0);
304+ c10_npu::NPUStream npu_stream2 = c10_npu::getStreamFromExternal(acl_stream, 0);
305+ 
306+ // Should return same NPUStream (same stream_id)
307+ bool same = (npu_stream1.id() == npu_stream2.id());
308+ 
309+ aclrtDestroyStream(acl_stream);
310+ return same;
311+}
312+ 
313+// Test 12: Pool stream vs external stream
314+bool test_pool_vs_external_stream()
315+{
316+ aclrtStream acl_stream = nullptr;
317+ auto ret = aclrtCreateStream(&acl_stream);
318+ if (ret != ACL_ERROR_NONE) {
319+ return false;
320+ }
321+ 
322+ c10_npu::NPUStream ext_stream = c10_npu::getStreamFromExternal(acl_stream, 0);
323+ c10_npu::NPUStream pool_stream = c10_npu::getNPUStreamFromPool(0);
324+ 
325+ // External stream has bit 30 = 1
326+ bool ext_has_marker = (static_cast<uint64_t>(ext_stream.id()) & (1ULL << 30)) != 0;
327+ 
328+ // Pool stream has bit 30 = 0
329+ bool pool_no_marker = (static_cast<uint64_t>(pool_stream.id()) & (1ULL << 30)) == 0;
330+ 
331+ aclrtDestroyStream(acl_stream);
332+ return ext_has_marker && pool_no_marker;
333+}
334+ 
335+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
336+{
337+ m.def("test_external_stream_creation", &test_external_stream_creation);
338+ m.def("test_external_stream_as_current", &test_external_stream_as_current);
339+ m.def("test_operations_on_external_stream", &test_operations_on_external_stream);
340+ m.def("test_event_record_restriction", &test_event_record_restriction);
341+ m.def("test_event_block_restriction", &test_event_block_restriction);
342+ m.def("test_graph_capture_restriction", &test_graph_capture_restriction);
343+ m.def("test_query_restriction", &test_query_restriction);
344+ m.def("test_synchronize_restriction", &test_synchronize_restriction);
345+ m.def("test_is_sync_launch_stream", &test_is_sync_launch_stream);
346+ m.def("test_multiple_external_streams", &test_multiple_external_streams);
347+ m.def("test_same_acl_stream_same_npu_stream", &test_same_acl_stream_same_npu_stream);
348+ m.def("test_pool_vs_external_stream", &test_pool_vs_external_stream);
349+}
Mtest/cpp_extensions/setup.py+3-0
@@ -28,6 +28,9 @@ ext_modules = [
28 NpuExtension(28 NpuExtension(
29 'torch_test_cpp_extension.npu_from_blob', ['test_from_blob.cpp'],29 'torch_test_cpp_extension.npu_from_blob', ['test_from_blob.cpp'],
30 extra_compile_args=CXX_FLAGS),30 extra_compile_args=CXX_FLAGS),
31+ NpuExtension(
32+ 'torch_test_cpp_extension.npu_external_stream', ['external_stream_test.cpp'],
33+ extra_compile_args=CXX_FLAGS),
31 NpuExtension(34 NpuExtension(
32 'torch_test_cpp_extension.stable_libtorch', ['test_stable_libtorch.cpp'],35 'torch_test_cpp_extension.stable_libtorch', ['test_stable_libtorch.cpp'],
33 extra_compile_args=CXX_FLAGS),36 extra_compile_args=CXX_FLAGS),
Atest/cpp_extensions/test/test_external_stream.py+85-0
@@ -0,0 +1,85 @@
1+import os
2+import sys
3+from pathlib import Path
4+import unittest
5+ 
6+REPO_ROOT = Path(__file__).resolve().parents[3]
7+CPP_EXTENSIONS_DIR = REPO_ROOT / "test" / "cpp_extensions"
8+ 
9+os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
10+for path in ("", str(REPO_ROOT)):
11+ while path in sys.path:
12+ sys.path.remove(path)
13+ 
14+sys.path.insert(0, str(CPP_EXTENSIONS_DIR))
15+ 
16+import torch
17+import torch_npu
18+from torch_npu.testing.testcase import TestCase, run_tests
19+ 
20+from torch_test_cpp_extension.load_external_stream import load_external_stream_extension
21+ 
22+ext_stream_ext = load_external_stream_extension()
23+ 
24+ 
25+class TestExternalStream(TestCase):
26+ """Tests for external NPU stream functionality
27+ 
28+ These tests verify:
29+ 1. External stream creation via getStreamFromExternal
30+ 2. Stream ID encoding (bit 30 marker for external streams)
31+ 3. Proper restrictions on external streams (event, graph, query, sync)
32+ 4. Basic operations work on external streams
33+ """
34+ 
35+ def test_external_stream_creation(self):
36+ """Test external stream creation and ID encoding"""
37+ self.assertTrue(ext_stream_ext.test_external_stream_creation())
38+ 
39+ def test_external_stream_as_current(self):
40+ """Test setCurrentNPUStream/getCurrentNPUStream with external stream"""
41+ self.assertTrue(ext_stream_ext.test_external_stream_as_current())
42+ 
43+ def test_operations_on_external_stream(self):
44+ """Test tensor operations on external stream"""
45+ self.assertTrue(ext_stream_ext.test_operations_on_external_stream())
46+ 
47+ def test_event_record_restriction(self):
48+ """Test NPUEvent.record() throws for external stream"""
49+ self.assertTrue(ext_stream_ext.test_event_record_restriction())
50+ 
51+ def test_event_block_restriction(self):
52+ """Test NPUEvent.block() throws for external stream"""
53+ self.assertTrue(ext_stream_ext.test_event_block_restriction())
54+ 
55+ def test_graph_capture_restriction(self):
56+ """Test NPUGraph.capture_begin() throws when current stream is external"""
57+ self.assertTrue(ext_stream_ext.test_graph_capture_restriction())
58+ 
59+ def test_query_restriction(self):
60+ """Test NPUStream.query() throws for external stream"""
61+ self.assertTrue(ext_stream_ext.test_query_restriction())
62+ 
63+ def test_synchronize_restriction(self):
64+ """Test NPUStream.synchronize() throws for external stream"""
65+ self.assertTrue(ext_stream_ext.test_synchronize_restriction())
66+ 
67+ def test_is_sync_launch_stream(self):
68+ """Test external stream is not a sync launch stream"""
69+ self.assertTrue(ext_stream_ext.test_is_sync_launch_stream())
70+ 
71+ def test_multiple_external_streams(self):
72+ """Test multiple external streams have unique IDs"""
73+ self.assertTrue(ext_stream_ext.test_multiple_external_streams())
74+ 
75+ def test_same_acl_stream_same_npu_stream(self):
76+ """Test getStreamFromExternal is idempotent"""
77+ self.assertTrue(ext_stream_ext.test_same_acl_stream_same_npu_stream())
78+ 
79+ def test_pool_vs_external_stream(self):
80+ """Test pool stream vs external stream ID encoding"""
81+ self.assertTrue(ext_stream_ext.test_pool_vs_external_stream())
82+ 
83+ 
84+if __name__ == "__main__":
85+ run_tests()
Atest/cpp_extensions/torch_test_cpp_extension/load_external_stream.py+56-0
@@ -0,0 +1,56 @@
1+import importlib
2+import os
3+import subprocess
4+import sys
5+from pathlib import Path
6+ 
7+ 
8+REPO_ROOT = Path(__file__).resolve().parents[3]
9+BUILD_PACKAGES_DIR = REPO_ROOT / "build" / "packages"
10+CPP_EXTENSIONS_DIR = REPO_ROOT / "test" / "cpp_extensions"
11+MODULE_NAME = "torch_test_cpp_extension.npu_external_stream"
12+ 
13+ 
14+def _build_pythonpath_for_subprocess():
15+ parts = []
16+ if BUILD_PACKAGES_DIR.exists():
17+ parts.append(str(BUILD_PACKAGES_DIR))
18+ existing = os.environ.get("PYTHONPATH")
19+ if existing:
20+ filtered = [p for p in existing.split(os.pathsep) if p != str(REPO_ROOT)]
21+ parts.extend(filtered)
22+ return os.pathsep.join(parts)
23+ 
24+ 
25+def _build_extension_inplace():
26+ env = os.environ.copy()
27+ pythonpath = _build_pythonpath_for_subprocess()
28+ if pythonpath:
29+ env["PYTHONPATH"] = pythonpath
30+ 
31+ subprocess.run(
32+ [sys.executable, "setup.py", "build_ext", "--inplace"],
33+ cwd=str(CPP_EXTENSIONS_DIR),
34+ check=True,
35+ env=env,
36+ )
37+ 
38+ 
39+def load_external_stream_extension():
40+ for path in ("", str(REPO_ROOT)):
41+ while path in sys.path:
42+ sys.path.remove(path)
43+ 
44+ os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
45+ 
46+ import torch # noqa: F401
47+ 
48+ import torch_npu # noqa: F401
O
OopenLiBingCI5月23日

此条代码评论区间+44至+48

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。

likedislike
49+ 
50+ try:
O
OopenLiBingCI5月23日

此条代码评论区间+46至+50

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。

likedislike
51+ return importlib.import_module(MODULE_NAME)
52+ except ImportError:
53+ _build_extension_inplace()
54+ sys.modules.pop(MODULE_NAME, None)
55+ importlib.invalidate_caches()
56+ return importlib.import_module(MODULE_NAME)
Mtorch_npu/csrc/aten/ops/StreamAndEventKernelNpu.cpp+5-2
@@ -1,4 +1,5 @@
1#include "torch_npu/csrc/aten/NPUNativeFunctions.h"1#include "torch_npu/csrc/aten/NPUNativeFunctions.h"
2+#include "torch_npu/csrc/core/npu/NPUStreamUtils.h"
2#include "torch_npu/csrc/framework/utils/OpAdapter.h"3#include "torch_npu/csrc/framework/utils/OpAdapter.h"
3 4 
4namespace at_npu {5namespace at_npu {
@@ -7,10 +8,12 @@ namespace native {
7void NPUNativeFunctions::record_stream(at::Tensor& self, c10::Stream stream)8void NPUNativeFunctions::record_stream(at::Tensor& self, c10::Stream stream)
8{9{
9 struct c10::StreamData3 data = stream.pack3();10 struct c10::StreamData3 data = stream.pack3();
11+ auto npu_stream = c10_npu::NPUStream::unpack3(
12+ data.stream_id, data.device_index, data.device_type);
13+ c10_npu::detail::checkNotExternalStream(npu_stream, "record_stream");
10 c10_npu::NPUCachingAllocator::recordStream(14 c10_npu::NPUCachingAllocator::recordStream(
11 self.storage().data_ptr(),15 self.storage().data_ptr(),
12- c10_npu::NPUStream::unpack3(16+ npu_stream);
13- data.stream_id, data.device_index, data.device_type));
14}17}
15 18 
16} // namespace native19} // namespace native
Mtorch_npu/csrc/core/npu/NPUEvent.cpp+3-0
@@ -4,6 +4,7 @@
4#include "torch_npu/csrc/core/npu/NPUGuard.h"4#include "torch_npu/csrc/core/npu/NPUGuard.h"
5#include "torch_npu/csrc/core/npu/NPUException.h"5#include "torch_npu/csrc/core/npu/NPUException.h"
6#include "torch_npu/csrc/core/npu/NPUEventManager.h"6#include "torch_npu/csrc/core/npu/NPUEventManager.h"
7+#include "torch_npu/csrc/core/npu/NPUStreamUtils.h"
7#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"8#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
8#include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h"9#include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h"
9#include "torch_npu/csrc/core/npu/register/OptionsManager.h"10#include "torch_npu/csrc/core/npu/register/OptionsManager.h"
@@ -113,6 +114,7 @@ void NPUEvent::recordOnce(const NPUStream& stream)
113 114 
114void NPUEvent::record(const NPUStream& stream)115void NPUEvent::record(const NPUStream& stream)
115{116{
117+ c10_npu::detail::checkNotExternalStream(stream, "NPUEvent::record");
116 if (!is_created_) {118 if (!is_created_) {
117 createEvent(stream.device_index());119 createEvent(stream.device_index());
118 }120 }
@@ -131,6 +133,7 @@ void NPUEvent::record(const NPUStream& stream)
131 133 
132void NPUEvent::block(const NPUStream& stream)134void NPUEvent::block(const NPUStream& stream)
133{135{
136+ c10_npu::detail::checkNotExternalStream(stream, "NPUEvent::block");
134 if (is_created_) {137 if (is_created_) {
135 // If using multiple task queues or using IPC events across devices in a single process,138 // If using multiple task queues or using IPC events across devices in a single process,
136 // it is necessary to ensure that the enqueued record is dequeued before wait.139 // it is necessary to ensure that the enqueued record is dequeued before wait.
Mtorch_npu/csrc/core/npu/NPUGraph.cpp+15-3
@@ -1,6 +1,7 @@
1#include "torch_npu/csrc/core/npu/NPUGraph.h"1#include "torch_npu/csrc/core/npu/NPUGraph.h"
2#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"2#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
3#include "torch_npu/csrc/core/npu/NPUFunctions.h"3#include "torch_npu/csrc/core/npu/NPUFunctions.h"
4+#include "torch_npu/csrc/core/npu/NPUStreamUtils.h"
4#include "torch_npu/csrc/aten/NPUGeneratorImpl.h"5#include "torch_npu/csrc/aten/NPUGeneratorImpl.h"
5#include "torch_npu/csrc/core/npu/register/OptionRegister.h"6#include "torch_npu/csrc/core/npu/register/OptionRegister.h"
6#include "third_party/acl/inc/acl/error_codes/rt_error_codes.h"7#include "third_party/acl/inc/acl/error_codes/rt_error_codes.h"
@@ -48,11 +49,13 @@ MempoolId_t graph_pool_handle()
48 49 
49void graph_task_group_begin(c10_npu::NPUStream stream)50void graph_task_group_begin(c10_npu::NPUStream stream)
50{51{
52+ c10_npu::detail::checkNotExternalStream(stream, "graph_task_group_begin");
51 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRICaptureTaskGrpBegin(stream));53 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRICaptureTaskGrpBegin(stream));
52}54}
53 55 
54NPUTaskGroupHandle graph_task_group_end(c10_npu::NPUStream stream)56NPUTaskGroupHandle graph_task_group_end(c10_npu::NPUStream stream)
55{57{
58+ c10_npu::detail::checkNotExternalStream(stream, "graph_task_group_end");
56 aclrtTaskGrp group;59 aclrtTaskGrp group;
57 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRICaptureTaskGrpEnd(stream, &group));60 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRICaptureTaskGrpEnd(stream, &group));
58 NPUTaskGroupHandle handle;61 NPUTaskGroupHandle handle;
@@ -62,28 +65,33 @@ NPUTaskGroupHandle graph_task_group_end(c10_npu::NPUStream stream)
62 65 
63void graph_task_update_begin(c10_npu::NPUStream stream, NPUTaskGroupHandle handle)66void graph_task_update_begin(c10_npu::NPUStream stream, NPUTaskGroupHandle handle)
64{67{
68+ c10_npu::detail::checkNotExternalStream(stream, "graph_task_update_begin");
65 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRICaptureTaskUpdateBegin(stream, handle.task_group));69 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRICaptureTaskUpdateBegin(stream, handle.task_group));
66}70}
67 71 
68void graph_task_update_end(c10_npu::NPUStream stream)72void graph_task_update_end(c10_npu::NPUStream stream)
69{73{
74+ c10_npu::detail::checkNotExternalStream(stream, "graph_task_update_end");
70 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRICaptureTaskUpdateEnd(stream));75 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRICaptureTaskUpdateEnd(stream));
71}76}
72 77 
73void super_kernel_scope_begin(const char* scope_name)78void super_kernel_scope_begin(const char* scope_name)
74{79{
75 auto stream = c10_npu::getCurrentNPUStream();80 auto stream = c10_npu::getCurrentNPUStream();
81+ c10_npu::detail::checkNotExternalStream(stream, "super_kernel_scope_begin");
76 NPU_CHECK_ERROR(c10_npu::skapi::AclskScopeBegin(scope_name, stream));82 NPU_CHECK_ERROR(c10_npu::skapi::AclskScopeBegin(scope_name, stream));
77}83}
78 84 
79void super_kernel_scope_end(const char* scope_name)85void super_kernel_scope_end(const char* scope_name)
80{86{
81 auto stream = c10_npu::getCurrentNPUStream();87 auto stream = c10_npu::getCurrentNPUStream();
88+ c10_npu::detail::checkNotExternalStream(stream, "super_kernel_scope_end");
82 NPU_CHECK_ERROR(c10_npu::skapi::AclskScopeEnd(scope_name, stream));89 NPU_CHECK_ERROR(c10_npu::skapi::AclskScopeEnd(scope_name, stream));
83}90}
84 91 
85void launch_callback(c10_npu::NPUStream stream, NPUCallbackFunc func, void *fnData)92void launch_callback(c10_npu::NPUStream stream, NPUCallbackFunc func, void *fnData)
86{93{
94+ c10_npu::detail::checkNotExternalStream(stream, "launch_callback");
87 aclrtCallbackBlockType type = aclrtCallbackBlockType::ACL_CALLBACK_BLOCK;95 aclrtCallbackBlockType type = aclrtCallbackBlockType::ACL_CALLBACK_BLOCK;
88 NPU_CHECK_ERROR(c10_npu::acl::AclrtLaunchCallback(func, fnData, type, stream));96 NPU_CHECK_ERROR(c10_npu::acl::AclrtLaunchCallback(func, fnData, type, stream));
89}97}
@@ -95,11 +103,13 @@ void launch_host_func(c10_npu::NPUStream stream, NPUCallbackFunc func, void *fnD
95 103 
96void subscribe_report(uint64_t threadId, c10_npu::NPUStream stream)104void subscribe_report(uint64_t threadId, c10_npu::NPUStream stream)
97{105{
106+ c10_npu::detail::checkNotExternalStream(stream, "subscribe_report");
98 NPU_CHECK_ERROR(c10_npu::acl::AclrtSubscribeReport(threadId, stream));107 NPU_CHECK_ERROR(c10_npu::acl::AclrtSubscribeReport(threadId, stream));
99}108}
100 109 
101void unsubscribe_report(uint64_t threadId, c10_npu::NPUStream stream)110void unsubscribe_report(uint64_t threadId, c10_npu::NPUStream stream)
102{111{
112+ c10_npu::detail::checkNotExternalStream(stream, "unsubscribe_report");
103 NPU_CHECK_ERROR(c10_npu::acl::AclrtUnSubscribeReport(threadId, stream));113 NPU_CHECK_ERROR(c10_npu::acl::AclrtUnSubscribeReport(threadId, stream));
104}114}
105 115 
@@ -177,6 +187,7 @@ void NPUGraph::capture_begin(MempoolId_t pool, aclmdlRICaptureMode capture_mode,
177 "To capture a new graph, create a new instance.");187 "To capture a new graph, create a new instance.");
178 188 
179 auto stream = c10_npu::getCurrentNPUStream();189 auto stream = c10_npu::getCurrentNPUStream();
190+ c10_npu::detail::checkNotExternalStream(stream, "NPUGraph::capture_begin");
180 191 
181 TORCH_CHECK(stream != c10_npu::getDefaultNPUStream(),192 TORCH_CHECK(stream != c10_npu::getDefaultNPUStream(),
182 "NPU graphs must be captured on a non-default stream. "193 "NPU graphs must be captured on a non-default stream. "
@@ -246,6 +257,8 @@ void NPUGraph::capture_end()
246 NPUGRAPH_LOGD("[NPUGRAPH][Capture] end");257 NPUGRAPH_LOGD("[NPUGRAPH][Capture] end");
247 auto stream = c10_npu::getCurrentNPUStream();258 auto stream = c10_npu::getCurrentNPUStream();
248 259 
260+ c10_npu::detail::checkNotExternalStream(stream, "NPUGraph::capture_end");
261+ 
249 TORCH_CHECK(stream == capture_stream_,262 TORCH_CHECK(stream == capture_stream_,
250 "Capture must end on the same stream it began on.");263 "Capture must end on the same stream it began on.");
251 264 
@@ -291,9 +304,8 @@ void NPUGraph::replay()
291 // model_ri_ may be replayed in any stream.304 // model_ri_ may be replayed in any stream.
292 auto stream = c10_npu::getCurrentNPUStream();305 auto stream = c10_npu::getCurrentNPUStream();
293 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRIExecuteAsync(model_ri_, stream));306 NPU_CHECK_ERROR(c10_npu::acl::AclmdlRIExecuteAsync(model_ri_, stream));
294- // With ASCEND_LAUNCH_BLOCKING enabled, after an aclgraph replay completes,307+ if (c10_npu::option::OptionsManager::CheckBlockingEnable() &&
295- // we need to add an explicit synchronization step.308+ !c10_npu::detail::isExternalStream(stream)) {
296- if (c10_npu::option::OptionsManager::CheckBlockingEnable()) {
297 NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream));309 NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream));
298 }310 }
299}311}
Mtorch_npu/csrc/core/npu/NPUGraphsUtils.cpp+5-0
@@ -1,4 +1,5 @@
1#include "NPUGraphsUtils.h"1#include "NPUGraphsUtils.h"
2+#include "torch_npu/csrc/core/npu/NPUStreamUtils.h"
2#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"3#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
3 4 
4namespace c10_npu {5namespace c10_npu {
@@ -8,6 +9,10 @@ CaptureStatus currentStreamCaptureStatusMayInitCtx()
8 return CaptureStatus::None;9 return CaptureStatus::None;
9 }10 }
10 11 
12+ if (c10_npu::detail::isCurrentStreamExternal()) {
13+ return CaptureStatus::None;
14+ }
15+ 
11 aclmdlRICaptureStatus is_capturing{ACL_MODEL_RI_CAPTURE_STATUS_NONE};16 aclmdlRICaptureStatus is_capturing{ACL_MODEL_RI_CAPTURE_STATUS_NONE};
12 aclmdlRI model_ri;17 aclmdlRI model_ri;
13 auto s = c10_npu::getCurrentNPUStream();18 auto s = c10_npu::getCurrentNPUStream();
Mtorch_npu/csrc/core/npu/NPUStream.cpp+159-5
@@ -3,6 +3,8 @@
3#include <atomic>3#include <atomic>
4#include <cstdint>4#include <cstdint>
5#include <cstring>5#include <cstring>
6+#include <memory>
7+#include <mutex>
6#include <vector>8#include <vector>
7#include <sys/time.h>9#include <sys/time.h>
8#include <unistd.h>10#include <unistd.h>
@@ -15,6 +17,7 @@
15#include "torch_npu/csrc/core/npu/NPUQueue.h"17#include "torch_npu/csrc/core/npu/NPUQueue.h"
16#include "torch_npu/csrc/core/npu/NPUException.h"18#include "torch_npu/csrc/core/npu/NPUException.h"
17#include "torch_npu/csrc/core/npu/register/OptionsManager.h"19#include "torch_npu/csrc/core/npu/register/OptionsManager.h"
20+#include "torch_npu/csrc/core/npu/NPUStreamUtils.h"
18#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"21#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
19#include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h"22#include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h"
20#include "third_party/acl/inc/acl/acl_rt.h"23#include "third_party/acl/inc/acl/acl_rt.h"
@@ -44,7 +47,7 @@ struct LeakyStreamInternals {
44 }47 }
45 48 
46 c10::DeviceIndex device_index = -1;49 c10::DeviceIndex device_index = -1;
47- int32_t stream_id = -1;50+ int64_t stream_id = 0x3FFFFFFF; // Invalid default: bit30=0 (non-external), >191 (non-internal)
48 aclrtStream stream = nullptr;51 aclrtStream stream = nullptr;
49 ::std::unique_ptr<NPUQueueBase> repo = nullptr;52 ::std::unique_ptr<NPUQueueBase> repo = nullptr;
50 bool is_data_preprocess_stream = false;53 bool is_data_preprocess_stream = false;
@@ -89,6 +92,11 @@ static std::array<
89static thread_local std::unique_ptr<LeakyStreamInternals* []> current_streams = nullptr;92static thread_local std::unique_ptr<LeakyStreamInternals* []> current_streams = nullptr;
90 93 
91static std::array<LeakyStreamInternals, kSyncLaunchStreamsPerPool> sync_launch_streams[C10_COMPILE_TIME_MAX_NPUS];94static std::array<LeakyStreamInternals, kSyncLaunchStreamsPerPool> sync_launch_streams[C10_COMPILE_TIME_MAX_NPUS];
95+static std::array<std::vector<std::unique_ptr<LeakyStreamInternals>>, C10_COMPILE_TIME_MAX_NPUS> external_streams;
96+static std::mutex external_streams_mutex;
97+ 
98+static constexpr uint64_t kExternalStreamIdMarker = 1ULL << 30;
99+static constexpr uint64_t kExternalStreamIdIndexMask = kExternalStreamIdMarker - 1;
92 100 
93thread_local aclrtStream tls_prev_stream = nullptr;101thread_local aclrtStream tls_prev_stream = nullptr;
94 102 
@@ -98,6 +106,7 @@ enum class StreamIdType : uint8_t {
98 SYNCLAUNCH = 0x2,106 SYNCLAUNCH = 0x2,
99 NORMAL = 0x3,107 NORMAL = 0x3,
100 HIGH = 0x4,108 HIGH = 0x4,
109+ EXT = 0x5, // External stream type
101};110};
102 111 
103std::ostream& operator<<(std::ostream& stream, StreamIdType s)112std::ostream& operator<<(std::ostream& stream, StreamIdType s)
@@ -118,6 +127,9 @@ std::ostream& operator<<(std::ostream& stream, StreamIdType s)
118 case StreamIdType::SYNCLAUNCH:127 case StreamIdType::SYNCLAUNCH:
119 stream << "SYNCLAUNCH";128 stream << "SYNCLAUNCH";
120 break;129 break;
130+ case StreamIdType::EXT:
131+ stream << "EXT";
132+ break;
121 default:133 default:
122 stream << static_cast<uint8_t>(s);134 stream << static_cast<uint8_t>(s);
123 break;135 break;
@@ -147,20 +159,36 @@ int GetStreamsPerPool()
147 return StreamsPerPool;159 return StreamsPerPool;
148}160}
149 161 
162+static inline bool isExternalStreamId(c10::StreamId s)
163+{
164+ return (static_cast<uint64_t>(s) & kExternalStreamIdMarker) != 0;
165+}
166+ 
150static inline StreamIdType streamIdType(c10::StreamId s)167static inline StreamIdType streamIdType(c10::StreamId s)
151{168{
169+ if (isExternalStreamId(s)) {
170+ return StreamIdType::EXT;
171+ }
152 static int StreamsPerPoolBits = GetStreamsPerPoolBits();172 static int StreamsPerPoolBits = GetStreamsPerPoolBits();
153 return static_cast<StreamIdType>((uint32_t)s >> StreamsPerPoolBits);173 return static_cast<StreamIdType>((uint32_t)s >> StreamsPerPoolBits);
154}174}
155 175 
156static inline size_t streamIdIndex(c10::StreamId s)176static inline size_t streamIdIndex(c10::StreamId s)
157{177{
178+ if (isExternalStreamId(s)) {
179+ return static_cast<size_t>(static_cast<uint64_t>(s) & kExternalStreamIdIndexMask);
180+ }
158 static int StreamsPerPoolBits = GetStreamsPerPoolBits();181 static int StreamsPerPoolBits = GetStreamsPerPoolBits();
159 return static_cast<size_t>((uint32_t)s & ((1 << StreamsPerPoolBits) - 1));182 return static_cast<size_t>((uint32_t)s & ((1 << StreamsPerPoolBits) - 1));
160}183}
161 184 
162c10::StreamId makeStreamId(StreamIdType st, size_t si)185c10::StreamId makeStreamId(StreamIdType st, size_t si)
163{186{
187+ if (st == StreamIdType::EXT) {
188+ TORCH_CHECK(si <= kExternalStreamIdIndexMask,
189+ "Too many external NPU streams are registered.", PTA_ERROR(ErrCode::VALUE));
190+ return static_cast<c10::StreamId>(kExternalStreamIdMarker | si);
191+ }
164 static int StreamsPerPoolBits = GetStreamsPerPoolBits();192 static int StreamsPerPoolBits = GetStreamsPerPoolBits();
165 return static_cast<c10::StreamId>((static_cast<size_t>(st) << StreamsPerPoolBits) | si);193 return static_cast<c10::StreamId>((static_cast<size_t>(st) << StreamsPerPoolBits) | si);
166}194}
@@ -175,6 +203,9 @@ static bool pointer_within(const T* ptr, const A& arr)
175static c10::StreamId NPUStream_getStreamId(const LeakyStreamInternals* ptr)203static c10::StreamId NPUStream_getStreamId(const LeakyStreamInternals* ptr)
176{204{
177 c10::DeviceIndex device_index = ptr->device_index;205 c10::DeviceIndex device_index = ptr->device_index;
206+ if (isExternalStreamId(ptr->stream_id)) {
207+ return static_cast<c10::StreamId>(ptr->stream_id);
208+ }
178 if (ptr == &default_streams[device_index]) {209 if (ptr == &default_streams[device_index]) {
179 return makeStreamId(StreamIdType::DEFAULT, 0);210 return makeStreamId(StreamIdType::DEFAULT, 0);
180 }211 }
@@ -219,6 +250,7 @@ static void initGlobalStreamState()
219 }250 }
220 // Initializes default streams251 // Initializes default streams
221 default_streams[device_id].device_index = device_id;252 default_streams[device_id].device_index = device_id;
253+ default_streams[device_id].stream_id = static_cast<int64_t>(makeStreamId(StreamIdType::DEFAULT, 0));
222 for (const auto p : c10::irange(kMaxStreamPriorities)) {254 for (const auto p : c10::irange(kMaxStreamPriorities)) {
223 npu_counters[p][device_id] = 0;255 npu_counters[p][device_id] = 0;
224 }256 }
@@ -230,6 +262,7 @@ static void initGlobalStreamState()
230 }262 }
231 // Initializes secondary streams263 // Initializes secondary streams
232 secondary_streams[device_id].device_index = device_id;264 secondary_streams[device_id].device_index = device_id;
265+ secondary_streams[device_id].stream_id = static_cast<int64_t>(makeStreamId(StreamIdType::SECONDARY, 0));
233 auto &secondary_streamsi = secondary_streams[device_id];266 auto &secondary_streamsi = secondary_streams[device_id];
234 NPU_CHECK_ERROR(267 NPU_CHECK_ERROR(
235 acl::AclrtCreateStreamWithConfig(&secondary_streamsi.stream, 0, (ACL_STREAM_FAST_LAUNCH | ACL_STREAM_FAST_SYNC)));268 acl::AclrtCreateStreamWithConfig(&secondary_streamsi.stream, 0, (ACL_STREAM_FAST_LAUNCH | ACL_STREAM_FAST_SYNC)));
@@ -245,6 +278,8 @@ static void initDeviceStreamState(c10::DeviceIndex device_index, int p)
245 auto& npu_streami = npu_streams[p][device_index][i];278 auto& npu_streami = npu_streams[p][device_index][i];
246 279 
247 npu_streami.device_index = device_index;280 npu_streami.device_index = device_index;
281+ npu_streami.stream_id = static_cast<int64_t>(makeStreamId(
282+ StreamIdType(static_cast<uint8_t>(StreamIdType::NORMAL) + p), i));
248 283 
249 NPU_CHECK_ERROR(acl::AclrtCreateStreamWithConfig(284 NPU_CHECK_ERROR(acl::AclrtCreateStreamWithConfig(
250 &npu_streami.stream, 0, (ACL_STREAM_FAST_LAUNCH | ACL_STREAM_FAST_SYNC)));285 &npu_streami.stream, 0, (ACL_STREAM_FAST_LAUNCH | ACL_STREAM_FAST_SYNC)));
@@ -304,6 +339,16 @@ LeakyStreamInternals* NPUStream_internals(NPUStream s)
304 StreamIdType st = streamIdType(s.unwrap().id());339 StreamIdType st = streamIdType(s.unwrap().id());
305 size_t si = streamIdIndex(s.unwrap().id());340 size_t si = streamIdIndex(s.unwrap().id());
306 switch (st) {341 switch (st) {
342+ case StreamIdType::EXT: {
343+ size_t ext_si = streamIdIndex(s.unwrap().id());
344+ std::lock_guard<std::mutex> lock(external_streams_mutex);
345+ TORCH_CHECK(ext_si < external_streams[device_index].size(),
346+ "Unrecognized external NPU stream ", s.unwrap(), PTA_ERROR(ErrCode::PARAM));
347+ auto* ptr = external_streams[device_index][ext_si].get();
348+ TORCH_CHECK(ptr != nullptr,
349+ "Invalid external NPU stream ", s.unwrap(), PTA_ERROR(ErrCode::PARAM));
350+ return ptr;
351+ }
307 case StreamIdType::DEFAULT:352 case StreamIdType::DEFAULT:
308 AT_ASSERTM(353 AT_ASSERTM(
309 si == 0,354 si == 0,
@@ -346,6 +391,10 @@ NPUStream NPUStream_fromInternals(const LeakyStreamInternals* ptr)
346 391 
347bool NPUStream::query() const392bool NPUStream::query() const
348{393{
394+ TORCH_CHECK(!isExternalStreamId(this->id()),
395+ "Cannot query external NPU stream status. "
396+ "The caller must track stream completion externally.",
397+ PTA_ERROR(ErrCode::NOT_SUPPORT));
349 c10::DeviceGuard guard{stream_.device()};398 c10::DeviceGuard guard{stream_.device()};
350 acl::aclrtStreamStatus status = acl::ACL_STREAM_STATUS_RESERVED;399 acl::aclrtStreamStatus status = acl::ACL_STREAM_STATUS_RESERVED;
351 NPU_CHECK_ERROR(acl::AclrtStreamQuery(stream(), &status));400 NPU_CHECK_ERROR(acl::AclrtStreamQuery(stream(), &status));
@@ -357,15 +406,22 @@ bool NPUStream::query() const
357 406 
358void NPUStream::synchronize() const407void NPUStream::synchronize() const
359{408{
409+ TORCH_CHECK(!isExternalStreamId(this->id()),
410+ "External NPU stream is not supported in NPUStream::synchronize. "
411+ "The caller is responsible for synchronization of external streams.",
412+ PTA_ERROR(ErrCode::NOT_SUPPORT));
360 c10::DeviceGuard guard{stream_.device()};413 c10::DeviceGuard guard{stream_.device()};
361 NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream()));414 NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream()));
362}415}
363 416 
364aclrtStream NPUStream::stream() const417aclrtStream NPUStream::stream() const
365{418{
419+ auto cur_ptr = NPUStream_internals(*this);
420+ AT_ASSERT(cur_ptr, PTA_ERROR(ErrCode::PTR));
421+ if (isExternalStreamId(this->id())) {
422+ return cur_ptr->stream;
423+ }
366 if (c10_npu::option::OptionsManager::GetPerStreamQueue()) {424 if (c10_npu::option::OptionsManager::GetPerStreamQueue()) {
367- auto cur_ptr = NPUStream_internals(*this);
368- AT_ASSERT(cur_ptr, PTA_ERROR(ErrCode::PTR));
369 if (!this->isSyncLaunchStream() && cur_ptr->repo->CheckInit()) {425 if (!this->isSyncLaunchStream() && cur_ptr->repo->CheckInit()) {
370 NPUStatus ret = cur_ptr->repo->MakeSureQueueEmpty();426 NPUStatus ret = cur_ptr->repo->MakeSureQueueEmpty();
371 if (ret != NPU_STATUS_SUCCESS) {427 if (ret != NPU_STATUS_SUCCESS) {
@@ -385,8 +441,6 @@ aclrtStream NPUStream::stream() const
385 }441 }
386 }442 }
387 }443 }
388- auto cur_ptr = NPUStream_internals(*this);
389- AT_ASSERT(cur_ptr, PTA_ERROR(ErrCode::PTR));
390 return cur_ptr->stream;444 return cur_ptr->stream;
391}445}
392 446 
@@ -427,6 +481,39 @@ NPUStream getDefaultNPUStream(c10::DeviceIndex device_index)
427 return NPUStream_fromInternals(&default_streams[device_index]);481 return NPUStream_fromInternals(&default_streams[device_index]);
428}482}
429 483 
484+NPUStream getStreamFromExternal(aclrtStream stream, c10::DeviceIndex device_index)
485+{
486+ TORCH_CHECK(stream != nullptr, "External NPU stream does not support nullptr.", PTA_ERROR(ErrCode::PARAM));
487+ initNPUStreamsOnce();
488+ if (device_index == -1) {
489+ device_index = current_device();
490+ }
491+ check_npu(device_index);
492+ 
493+ LeakyStreamInternals* ptr = nullptr;
494+ {
495+ std::lock_guard<std::mutex> lock(external_streams_mutex);
496+ auto& streams = external_streams[device_index];
497+ for (const auto i : c10::irange(streams.size())) {
498+ if (streams[i]->stream == stream) {
499+ ptr = streams[i].get();
500+ break;
501+ }
502+ }
503+ 
504+ if (ptr == nullptr) {
505+ auto internals = std::make_unique<LeakyStreamInternals>();
506+ internals->device_index = device_index;
507+ internals->stream_id = static_cast<int64_t>(makeStreamId(StreamIdType::EXT, streams.size()));
508+ internals->stream = stream;
509+ internals->repo.reset();
510+ ptr = internals.get();
511+ streams.emplace_back(std::move(internals));
512+ }
513+ }
514+ return NPUStream_fromInternals(ptr);
515+}
516+ 
430NPUStream getCurrentNPUStream(c10::DeviceIndex device_index)517NPUStream getCurrentNPUStream(c10::DeviceIndex device_index)
431{518{
432 initNPUStreamsOnce();519 initNPUStreamsOnce();
@@ -662,8 +749,15 @@ void enCurrentNPUStream(void* cur_paras, c10::DeviceIndex device_index, NPUStrea
662 if (c10_npu::option::OptionsManager::GetPerStreamQueue()) {749 if (c10_npu::option::OptionsManager::GetPerStreamQueue()) {
663 LeakyStreamInternals* ptr = current_streams[device_index];750 LeakyStreamInternals* ptr = current_streams[device_index];
664 if (task_stream != nullptr) {751 if (task_stream != nullptr) {
752+ // Check before traversing the external_streams
753+ TORCH_CHECK(!isExternalStreamId(task_stream->id()),
754+ "External NPU stream is not supported by task queue enqueue.",
755+ PTA_ERROR(ErrCode::NOT_SUPPORT));
665 ptr = NPUStream_internals(*task_stream);756 ptr = NPUStream_internals(*task_stream);
666 }757 }
758+ TORCH_CHECK(!isExternalStreamId(static_cast<c10::StreamId>(ptr->stream_id)),
759+ "External NPU stream is not supported by task queue enqueue.",
760+ PTA_ERROR(ErrCode::NOT_SUPPORT));
667 // To prevent all taskqueue threads from being created during stream initialization,761 // To prevent all taskqueue threads from being created during stream initialization,
668 // we initialize the taskqueue when enqueueing762 // we initialize the taskqueue when enqueueing
669 // each stream init repo once763 // each stream init repo once
@@ -733,6 +827,9 @@ bool NPUStream::getRepoStopFlag()
733 827 
734bool NPUStream::isSyncLaunchStream() const828bool NPUStream::isSyncLaunchStream() const
735{829{
830+ if (isExternalStreamId(this->id())) {
831+ return false;
832+ }
736 auto ptr = NPUStream_internals(*this);833 auto ptr = NPUStream_internals(*this);
737 AT_ASSERT(ptr, PTA_ERROR(ErrCode::PTR));834 AT_ASSERT(ptr, PTA_ERROR(ErrCode::PTR));
738 return ptr->is_sync_launch;835 return ptr->is_sync_launch;
@@ -749,6 +846,62 @@ aclrtStream NPUStream::stream(const bool need_empty) const
749 return stream();846 return stream();
750}847}
751 848 
849+namespace detail {
850+ 
851+bool isExternalStream(const NPUStream& stream)
852+{
853+ return isExternalStreamId(stream.id());
854+}
855+ 
856+bool isExternalStream(aclrtStream stream, c10::DeviceIndex device_index)
857+{
858+ initNPUStreamsOnce();
859+ std::lock_guard<std::mutex> lock(external_streams_mutex);
860+ if (device_index != -1) {
861+ check_npu(device_index);
862+ const auto& streams = external_streams[device_index];
863+ for (const auto& item : streams) {
864+ if (item->stream == stream) {
865+ return true;
866+ }
867+ }
868+ return false;
869+ }
870+ 
871+ for (const auto device : c10::irange(num_npus)) {
872+ const auto& streams = external_streams[device];
873+ for (const auto& item : streams) {
874+ if (item->stream == stream) {
875+ return true;
876+ }
877+ }
878+ }
879+ return false;
880+}
881+ 
882+bool isCurrentStreamExternal(c10::DeviceIndex device_index)
883+{
884+ return isExternalStreamId(getCurrentNPUStream(device_index).id());
885+}
886+ 
887+void checkNotExternalStream(const NPUStream& stream, const char* api_name)
888+{
889+ TORCH_CHECK(!isExternalStream(stream),
890+ "External NPU stream is not supported in ", api_name,
891+ ". This path requires a torch_npu-managed stream.",
892+ PTA_ERROR(ErrCode::NOT_SUPPORT));
893+}
894+ 
895+void checkCurrentStreamNotExternal(c10::DeviceIndex device_index, const char* api_name)
896+{
897+ TORCH_CHECK(!isCurrentStreamExternal(device_index),
898+ "External NPU stream is not supported in ", api_name,
899+ ". This path requires a torch_npu-managed current stream.",
900+ PTA_ERROR(ErrCode::NOT_SUPPORT));
901+}
902+ 
903+} // namespace detail
904+ 
752void recovery_all_npu_streams(c10::DeviceIndex device_index)905void recovery_all_npu_streams(c10::DeviceIndex device_index)
753{906{
754 if (!initialize_flag[device_index]) {907 if (!initialize_flag[device_index]) {
@@ -784,6 +937,7 @@ static void initDeviceSyncLaunchStream(c10::DeviceIndex device_index)
784 937 
785 sync_streami.device_index = device_index;938 sync_streami.device_index = device_index;
786 sync_streami.is_sync_launch = true;939 sync_streami.is_sync_launch = true;
940+ sync_streami.stream_id = static_cast<int64_t>(makeStreamId(StreamIdType::SYNCLAUNCH, i));
787 941 
788 NPU_CHECK_ERROR(942 NPU_CHECK_ERROR(
789 acl::AclrtCreateStreamWithConfig(&sync_streami.stream, 0, ACL_STREAM_FAST_SYNC));943 acl::AclrtCreateStreamWithConfig(&sync_streami.stream, 0, ACL_STREAM_FAST_SYNC));
Mtorch_npu/csrc/core/npu/NPUStream.h+2-0
@@ -120,6 +120,8 @@ C10_NPU_API NPUStream getNPUStreamFromPool(c10::DeviceIndex device = -1);
120 120 
121C10_NPU_API NPUStream getDefaultNPUStream(c10::DeviceIndex device_index = -1);121C10_NPU_API NPUStream getDefaultNPUStream(c10::DeviceIndex device_index = -1);
122 122 
123+C10_NPU_API NPUStream getStreamFromExternal(aclrtStream stream, c10::DeviceIndex device_index);
124+ 
123C10_NPU_API NPUStream getStreamFromPool(const bool isHighPriority, c10::DeviceIndex device_index);125C10_NPU_API NPUStream getStreamFromPool(const bool isHighPriority, c10::DeviceIndex device_index);
124 126 
125C10_NPU_API NPUStream getCurrentNPUStream(c10::DeviceIndex device_index = -1);127C10_NPU_API NPUStream getCurrentNPUStream(c10::DeviceIndex device_index = -1);
Atorch_npu/csrc/core/npu/NPUStreamUtils.h+19-0
@@ -0,0 +1,19 @@
1+#pragma once
2+ 
3+#include "torch_npu/csrc/core/npu/NPUStream.h"
4+ 
5+namespace c10_npu {
6+namespace detail {
7+ 
8+bool isExternalStream(const NPUStream& stream);
9+ 
10+bool isExternalStream(aclrtStream stream, c10::DeviceIndex device_index = -1);
11+ 
12+bool isCurrentStreamExternal(c10::DeviceIndex device_index = -1);
13+ 
14+void checkNotExternalStream(const NPUStream& stream, const char* api_name);
15+ 
16+void checkCurrentStreamNotExternal(c10::DeviceIndex device_index, const char* api_name);
17+ 
18+} // namespace detail
19+} // namespace c10_npu
Mtorch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.cpp+3-0
@@ -1,6 +1,7 @@
1#include "AsyncTaskQueueInterface.h"1#include "AsyncTaskQueueInterface.h"
2#include "torch_npu/csrc/core/npu/NPUEventManager.h"2#include "torch_npu/csrc/core/npu/NPUEventManager.h"
3#include "torch_npu/csrc/core/npu/NPUGuard.h"3#include "torch_npu/csrc/core/npu/NPUGuard.h"
4+#include "torch_npu/csrc/core/npu/NPUStreamUtils.h"
4#include "torch_npu/csrc/core/npu/register/OptionsManager.h"5#include "torch_npu/csrc/core/npu/register/OptionsManager.h"
5#include <ATen/record_function.h>6#include <ATen/record_function.h>
6#include "torch_npu/csrc/framework/utils/NpuUtils.h"7#include "torch_npu/csrc/framework/utils/NpuUtils.h"
@@ -221,6 +222,7 @@ void EventTask::LaunchRecordTask(c10_npu::NPUStream npuStream, unsigned int flag
221 222 
222aclError LaunchRecordEventTask(aclrtEvent event, c10_npu::NPUStream npuStream, unsigned int flags)223aclError LaunchRecordEventTask(aclrtEvent event, c10_npu::NPUStream npuStream, unsigned int flags)
223{224{
225+ c10_npu::detail::checkNotExternalStream(npuStream, "LaunchRecordEventTask");
224 EventTask recordTask(event);226 EventTask recordTask(event);
225 recordTask.LaunchRecordTask(npuStream, flags);227 recordTask.LaunchRecordTask(npuStream, flags);
226#ifndef BUILD_LIBTORCH228#ifndef BUILD_LIBTORCH
@@ -273,6 +275,7 @@ void EventTask::LaunchWaitTask(c10_npu::NPUStream npuStream, unsigned int flags)
273 275 
274aclError LaunchWaitEventTask(aclrtEvent event, c10_npu::NPUStream npuStream, unsigned int flags)276aclError LaunchWaitEventTask(aclrtEvent event, c10_npu::NPUStream npuStream, unsigned int flags)
275{277{
278+ c10_npu::detail::checkNotExternalStream(npuStream, "LaunchWaitEventTask");
276 EventTask waitTask(event);279 EventTask waitTask(event);
277 waitTask.LaunchWaitTask(npuStream, flags);280 waitTask.LaunchWaitTask(npuStream, flags);
278#ifndef BUILD_LIBTORCH281#ifndef BUILD_LIBTORCH
Mtorch_npu/csrc/distributed/ProcessGroupHCCL.cpp+4-0
@@ -42,6 +42,7 @@
42#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"42#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
43#include "torch_npu/csrc/core/npu/NPUAffinityController.h"43#include "torch_npu/csrc/core/npu/NPUAffinityController.h"
44#include "torch_npu/csrc/core/npu/NPUStream.h"44#include "torch_npu/csrc/core/npu/NPUStream.h"
45+#include "torch_npu/csrc/core/npu/NPUStreamUtils.h"
45#include "torch_npu/csrc/core/npu/register/OptionsManager.h"46#include "torch_npu/csrc/core/npu/register/OptionsManager.h"
46#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"47#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
47#include "torch_npu/csrc/core/npu/interface/OpInterface.h"48#include "torch_npu/csrc/core/npu/interface/OpInterface.h"
@@ -3869,6 +3870,9 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collective(
3869 op_id_++;3870 op_id_++;
3870 3871 
3871 const auto devices = getDeviceList(inputs);3872 const auto devices = getDeviceList(inputs);
3873+ for (const auto& device : devices) {
3874+ c10_npu::detail::checkCurrentStreamNotExternal(device.index(), "ProcessGroupHCCL::collective");
3875+ }
3872 auto key = getKeyFromDevices(devices);3876 auto key = getKeyFromDevices(devices);
3873 HcclCommConfig config = createHcclCommConfigWithOptions();3877 HcclCommConfig config = createHcclCommConfigWithOptions();
3874 std::vector<std::shared_ptr<HCCLComm>> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config);3878 std::vector<std::shared_ptr<HCCLComm>> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config);
Mtorch_npu/csrc/distributed/ProcessGroupLCCL.cpp+4-0
@@ -3,6 +3,7 @@
3#include "torch_npu/csrc/core/NPUBridge.h"3#include "torch_npu/csrc/core/NPUBridge.h"
4#include "torch_npu/csrc/core/npu/DeviceUtils.h"4#include "torch_npu/csrc/core/npu/DeviceUtils.h"
5#include "torch_npu/csrc/core/npu/NPUGuard.h"5#include "torch_npu/csrc/core/npu/NPUGuard.h"
6+#include "torch_npu/csrc/core/npu/NPUStreamUtils.h"
6#include "torch_npu/csrc/framework/FormatHelper.h"7#include "torch_npu/csrc/framework/FormatHelper.h"
7#include "torch_npu/csrc/framework/OpCommand.h"8#include "torch_npu/csrc/framework/OpCommand.h"
8 9 
@@ -214,6 +215,9 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupLCCL::collective(std::vector<at::Tens
214 PostProcess post, c10d::OpType opType)215 PostProcess post, c10d::OpType opType)
215{216{
216 const auto devices = getDeviceList(inputs);217 const auto devices = getDeviceList(inputs);
218+ for (const auto& device : devices) {
219+ c10_npu::detail::checkCurrentStreamNotExternal(device.index(), "ProcessGroupLCCL::collective");
220+ }
217 auto key = getKeyFromDevices(devices);221 auto key = getKeyFromDevices(devices);
218 std::vector<at_npu::lccl::LcclComm> lcclComms;222 std::vector<at_npu::lccl::LcclComm> lcclComms;
219 lcclComms = getLCCLComm(key, devices);223 lcclComms = getLCCLComm(key, devices);