diff --git a/paddle/fluid/pybind/cuda_streams_py.cc b/paddle/fluid/pybind/cuda_streams_py.cc index 706012f4a44793d76fa2dd343a209345c395693b..311fb872ac10394fbb29c4ee2c90801c2816269d 100644 --- a/paddle/fluid/pybind/cuda_streams_py.cc +++ b/paddle/fluid/pybind/cuda_streams_py.cc @@ -202,6 +202,28 @@ void BindCudaStream(py::module *m_ptr) { )DOC", py::arg("event") = nullptr) + .def_property_readonly( + "cuda_stream", + [](paddle::platform::stream::CUDAStream &self) { + VLOG(10) << self.raw_stream(); + return reinterpret_cast(self.raw_stream()); + }, + R"DOC( + retrun the raw cuda stream of type cudaStream_t as type int. + + Examples: + .. code-block:: python + + # required: gpu + import paddle + import ctypes + cuda_stream = paddle.device.cuda.current_stream().cuda_stream + print(cuda_stream) + + ptr = ctypes.c_void_p(cuda_stream) # convert back to void* + print(ptr) + + )DOC") #endif .def("__init__", [](paddle::platform::stream::CUDAStream &self, diff --git a/python/paddle/fluid/tests/unittests/test_cuda_stream_event.py b/python/paddle/fluid/tests/unittests/test_cuda_stream_event.py index ec024105f899d9179f5e74a70664872f802a6dcd..30bc00c9d9427ab0d54f4e22dfbedeb1cd541b1c 100644 --- a/python/paddle/fluid/tests/unittests/test_cuda_stream_event.py +++ b/python/paddle/fluid/tests/unittests/test_cuda_stream_event.py @@ -14,6 +14,7 @@ from paddle.device import cuda import paddle +import ctypes import unittest import numpy as np @@ -156,5 +157,14 @@ class TestStreamGuard(unittest.TestCase): None) +class TestRawStream(unittest.TestCase): + def test_cuda_stream(self): + if paddle.is_compiled_with_cuda(): + cuda_stream = paddle.device.cuda.current_stream().cuda_stream + print(cuda_stream) + self.assertTrue(type(cuda_stream) is int) + ptr = ctypes.c_void_p(cuda_stream) + + if __name__ == "__main__": unittest.main()