Add the acquisition method of cudastream in CustomDevice (#73272)

datutu-L · web-flow · commit ef256258ca70 · 2025-06-15T21:46:38.000+08:00
* Add the acquisition method of cudastream in CustomDevice

* small fix

* small fix

* small fix

* small fix

* small fix
diff --git a/paddle/fluid/pybind/cuda_streams_py.cc b/paddle/fluid/pybind/cuda_streams_py.cc
@@ -24,8 +24,28 @@ namespace py = pybind11;
 
 namespace paddle {
 namespace platform {
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-phi::CUDAStream *get_current_stream(int device_id) {
+#if defined(PADDLE_WITH_CUSTOM_DEVICE)
+PY_STREAM_TYPE get_current_stream(int device_id) {
+  auto dev_types = phi::DeviceManager::GetAllCustomDeviceTypes();
+  if (device_id == -1) {
+    device_id = phi::DeviceManager::GetDevice(dev_types[0]);
+  }
+  auto *custom_context =
+      static_cast<const phi::CustomContext *>(DeviceContextPool::Instance().Get(
+          phi::CustomPlace(dev_types[0], device_id)));
+  return custom_context->GetStream().get();
+}
+
+PY_STREAM_TYPE set_current_stream(PY_STREAM_TYPE stream) {
+  auto *original_stream = get_current_stream(stream->GetPlace().GetDeviceId());
+  auto *custom_context = static_cast<phi::CustomContext *>(
+      DeviceContextPool::Instance().Get(stream->GetPlace()));
+  custom_context->SetStream(std::shared_ptr<phi::stream::Stream>(stream));
+  return original_stream;
+}
+
+#elif defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PY_STREAM_TYPE get_current_stream(int device_id) {
   if (device_id == -1) {
     device_id = phi::backends::gpu::GetCurrentDeviceId();
   }
@@ -34,7 +54,7 @@ phi::CUDAStream *get_current_stream(int device_id) {
   return gpu_context->cuda_stream();
 }
 
-phi::CUDAStream *set_current_stream(phi::CUDAStream *stream) {
+PY_STREAM_TYPE set_current_stream(PY_STREAM_TYPE stream) {
   auto *original_stream = get_current_stream(stream->place().GetDeviceId());
   auto *gpu_context = static_cast<phi::GPUContext *>(
       DeviceContextPool::Instance().Get(stream->place()));
@@ -51,25 +71,27 @@ void BindCudaStream(py::module *m_ptr) {
   m.def(
       "_get_current_stream",
       [](int deviceId) {
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE)
         return platform::get_current_stream(deviceId);
 #else
-        PADDLE_THROW(
-            common::errors::Unavailable("Paddle is not compiled with CUDA. "
-                                        "Cannot visit device synchronize."));
+        PADDLE_THROW(common::errors::Unavailable(
+            "Paddle do not support _get_current_stream "
+            "Cannot visit device synchronize."));
 #endif
       },
       py::return_value_policy::reference);
 
   m.def(
       "_set_current_stream",
-      [](phi::CUDAStream *stream) {
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+      [](PY_STREAM_TYPE stream) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE)
         return platform::set_current_stream(stream);
 #else
-        PADDLE_THROW(
-            common::errors::Unavailable("Paddle is not compiled with CUDA. "
-                                        "Cannot visit device synchronize."));
+        PADDLE_THROW(common::errors::Unavailable(
+            "Paddle do not support _set_current_stream "
+            "Cannot visit device synchronize."));
 #endif
       },
       py::return_value_policy::reference);
@@ -240,6 +262,28 @@ void BindCudaStream(py::module *m_ptr) {
                   >>> ptr = ctypes.c_void_p(cuda_stream)  # convert back to void*
                   >>> print(ptr)
 
+          )DOC")
+      .def_property_readonly(
+          "raw_stream",
+          [](phi::CUDAStream &self) {
+            VLOG(10) << self.raw_stream();
+            return reinterpret_cast<std::uintptr_t>(self.raw_stream());
+          },
+          R"DOC(
+          return the raw cuda stream of type cudaStream_t as type int.
+
+          Examples:
+              .. code-block:: python
+
+                  >>> # doctest: +REQUIRES(env:GPU)
+                  >>> import paddle
+                  >>> import ctypes
+                  >>> raw_stream = paddle.device.cuda.current_stream().raw_stream
+                  >>> print(raw_stream)
+
+                  >>> ptr = ctypes.c_void_p(raw_stream)  # convert back to void*
+                  >>> print(ptr)
+
           )DOC")
       .def_property_readonly(
           "place",
diff --git a/paddle/fluid/pybind/cuda_streams_py.h b/paddle/fluid/pybind/cuda_streams_py.h
@@ -17,6 +17,9 @@
 #include "pybind11/pybind11.h"
 #include "pybind11/stl.h"
 
+#if defined(PADDLE_WITH_CUSTOM_DEVICE)
+#include "paddle/phi/backends/stream.h"
+#endif
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
 #include "paddle/phi/core/cuda_stream.h"
 #else
@@ -25,13 +28,20 @@ class CUDAStream {};
 }  // namespace phi
 #endif
 
+#if defined(PADDLE_WITH_CUSTOM_DEVICE)
+#define PY_STREAM_TYPE phi::stream::Stream*
+#else
+#define PY_STREAM_TYPE phi::CUDAStream*
+#endif
+
 namespace py = pybind11;
 
 namespace paddle {
 namespace platform {
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-phi::CUDAStream* get_current_stream(int device_id = -1);
-phi::CUDAStream* set_current_stream(phi::CUDAStream* stream);
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE)
+PY_STREAM_TYPE get_current_stream(int device_id = -1);
+PY_STREAM_TYPE set_current_stream(PY_STREAM_TYPE stream);
 #endif
 }  // namespace platform
 namespace pybind {
