已合并
add torch_npu._C._npu_getCurrentRawStreamNoWait() #29737
liujunzhu创建于 1月19日
add torch_npu._C._npu_getCurrentRawStreamNoWait() #29737
已合并
共 2 个文件变更+22-1
| @@ -18,7 +18,7 @@ class TestNpuStream(TestCase): | |||
| 18 | self.assertTrue(len(stream_instance) == device_number) | 18 | self.assertTrue(len(stream_instance) == device_number) |
| 19 | 19 | ||
| 20 | def test_get_current_stream_interface(self): | 20 | def test_get_current_stream_interface(self): |
| 21 | - from torch_npu._C import _npu_getCurrentRawStream | 21 | + from torch_npu._C import _npu_getCurrentRawStream, _npu_getCurrentRawStreamNoWait |
| 22 | from torch._dynamo.device_interface import get_interface_for_device | 22 | from torch._dynamo.device_interface import get_interface_for_device |
| 23 | 23 | ||
| 24 | device_number = torch.npu.device_count() | 24 | device_number = torch.npu.device_count() |
| @@ -28,8 +28,10 @@ class TestNpuStream(TestCase): | |||
| 28 | with torch.npu.stream(stream): | 28 | with torch.npu.stream(stream): |
| 29 | current_stream = torch.npu.current_stream() | 29 | current_stream = torch.npu.current_stream() |
| 30 | current_raw_stream = _npu_getCurrentRawStream(i) | 30 | current_raw_stream = _npu_getCurrentRawStream(i) |
| 31 | + current_raw_stream_no_wait = _npu_getCurrentRawStreamNoWait(i) | ||
| 31 | interface_raw_stream = get_interface_for_device('npu').get_raw_stream(i) | 32 | interface_raw_stream = get_interface_for_device('npu').get_raw_stream(i) |
| 32 | self.assertTrue(current_stream.npu_stream == current_raw_stream) | 33 | self.assertTrue(current_stream.npu_stream == current_raw_stream) |
| 34 | + self.assertTrue(current_stream.npu_stream == current_raw_stream_no_wait) | ||
| 33 | self.assertTrue(current_stream.npu_stream == interface_raw_stream) | 35 | self.assertTrue(current_stream.npu_stream == interface_raw_stream) |
| 34 | 36 | ||
| 35 | def test_priority(self): | 37 | def test_priority(self): |
| @@ -960,6 +960,24 @@ PyObject* THNPModule_getCurrentStream_raw( | |||
| 960 | END_HANDLE_TH_ERRORS | 960 | END_HANDLE_TH_ERRORS |
| 961 | } | 961 | } |
| 962 | 962 | ||
| 963 | +// Note: The torch_npu._C._npu_getCurrentRawStreamNoWait(device) interface does NOT clear the task queue. | ||
| 964 | +// If tasks are dispatched using both the returned aclrtStream and torch_npu's task queue, | ||
| 965 | +// it may cause ordering issues due to lack of synchronization between the two dispatch paths. | ||
| 966 | +// Users must ensure to use only one of these dispatch methods exclusively. | ||
| 967 | +// If mixed usage is unavoidable, ensure there are no data dependencies between tasks | ||
| 968 | +// and that performance is not sensitive to potential execution reordering. | ||
| 969 | +PyObject* THNPModule_getCurrentRawStreamNoWait_wrap( | ||
| 970 | + PyObject* /* unused */, PyObject* device_index) | ||
| 971 | +{ | ||
| 972 | + HANDLE_TH_ERRORS | ||
| 973 | + TORCH_CHECK( | ||
| 974 | + THPUtils_checkLong(device_index), "invalid argument to getCurrentStream", PTA_ERROR(ErrCode::PARAM)); | ||
| 975 | + int64_t device = THPUtils_unpackLong(device_index); | ||
| 976 | + return PyLong_FromVoidPtr( | ||
| 977 | + c10_npu::getCurrentNPUStreamNoWait(device)); | ||
| 978 | + END_HANDLE_TH_ERRORS | ||
| 979 | +} | ||
| 980 | + | ||
| 963 | PyObject* THNPModule_getDefaultStream_wrap(PyObject *self /* unused */, PyObject *device_index) | 981 | PyObject* THNPModule_getDefaultStream_wrap(PyObject *self /* unused */, PyObject *device_index) |
| 964 | { | 982 | { |
| 965 | HANDLE_TH_ERRORS | 983 | HANDLE_TH_ERRORS |
| @@ -2228,6 +2246,7 @@ static struct PyMethodDef THNPModule_methods[] = { | |||
| 2228 | {"_npu_getDeviceUtilizationRate", (PyCFunction)THNPModule_getDeviceUtilizationRate_wrap, METH_O, nullptr}, | 2246 | {"_npu_getDeviceUtilizationRate", (PyCFunction)THNPModule_getDeviceUtilizationRate_wrap, METH_O, nullptr}, |
| 2229 | {"_npu_getCurrentStream", (PyCFunction)THNPModule_getCurrentStream_wrap, METH_O, nullptr}, | 2247 | {"_npu_getCurrentStream", (PyCFunction)THNPModule_getCurrentStream_wrap, METH_O, nullptr}, |
| 2230 | {"_npu_getCurrentRawStream", (PyCFunction)THNPModule_getCurrentStream_raw, METH_O, nullptr}, | 2248 | {"_npu_getCurrentRawStream", (PyCFunction)THNPModule_getCurrentStream_raw, METH_O, nullptr}, |
| 2249 | + {"_npu_getCurrentRawStreamNoWait", (PyCFunction)THNPModule_getCurrentRawStreamNoWait_wrap, METH_O, nullptr}, | ||
| 2231 | {"_npu_getDefaultStream", (PyCFunction)THNPModule_getDefaultStream_wrap, METH_O, nullptr}, | 2250 | {"_npu_getDefaultStream", (PyCFunction)THNPModule_getDefaultStream_wrap, METH_O, nullptr}, |
| 2232 | {"_npu_setStream", (PyCFunction)THNPModule_setStream_wrap, METH_VARARGS | METH_KEYWORDS, nullptr}, | 2251 | {"_npu_setStream", (PyCFunction)THNPModule_setStream_wrap, METH_VARARGS | METH_KEYWORDS, nullptr}, |
| 2233 | {"_npu_eraseStream", (PyCFunction)THNPModule_npu_eraseStream_wrap, METH_VARARGS | METH_KEYWORDS, nullptr}, | 2252 | {"_npu_eraseStream", (PyCFunction)THNPModule_npu_eraseStream_wrap, METH_VARARGS | METH_KEYWORDS, nullptr}, |