PaddlePaddle
diff --git a/‎paddle/fluid/platform/init.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/platform/init.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/pybind/tensor.cc
Copy file name to clipboard
+163Lines changed: 163 additions & 0 deletions b/‎paddle/fluid/pybind/tensor.cc
Copy file name to clipboard
+163Lines changed: 163 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/xpu_streams_py.cc
Copy file name to clipboard
+94-4Lines changed: 94 additions & 4 deletions b/‎paddle/fluid/pybind/xpu_streams_py.cc
Copy file name to clipboard
+94-4Lines changed: 94 additions & 4 deletions
diff --git a/‎paddle/fluid/pybind/xpu_streams_py.h
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions b/‎paddle/fluid/pybind/xpu_streams_py.h
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions
diff --git a/‎paddle/phi/backends/xpu/enforce_xpu.h
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions b/‎paddle/phi/backends/xpu/enforce_xpu.h
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions
diff --git a/‎paddle/phi/backends/xpu/xpu_context.cc
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/phi/backends/xpu/xpu_context.cc
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
@@ -435,7 +435,7 @@ void InitMemoryMethod() {
     memory_method->allocation_deleter =
         paddle::memory::allocation::Allocator::AllocationDeleter;
 #if defined(PADDLE_WITH_CUSTOM_DEVICE) || defined(PADDLE_WITH_CUDA) || \
-    defined(PADDLE_WITH_HIP)
+    defined(PADDLE_WITH_HIP) || defined(PADDLE_WITH_XPU)
     memory_method->copy_with_stream =
         paddle::memory::Copy<phi::Place, phi::Place>;
 #endif
 
@@ -132,8 +132,15 @@ limitations under the License. */
 #endif
 
 #ifdef PADDLE_WITH_XPU
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include "paddle/fluid/pybind/xpu_streams_py.h"
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/memory/allocation/xpu_ipc_allocator.h"
 #include "paddle/phi/core/platform/device/xpu/xpu_info.h"
 #include "paddle/phi/core/platform/device/xpu/xpu_op_list.h"
+#include "xpu/runtime.h"
+#include "xpu/runtime_ex.h"
 #endif
 
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
@@ -805,6 +812,162 @@ void BindTensor(pybind11::module &m) {  // NOLINT
                     >>> tensor_from_shared = paddle.to_tensor(paddle.base.core.DenseTensor._new_shared_cuda(metainfo))
         )DOC")
 #endif
+#ifdef PADDLE_WITH_XPU
+      .def("_share_buffer_with",
+           [](phi::DenseTensor &self, const phi::DenseTensor src,
+              py::tuple t) {
+             if (!src.meta().is_contiguous()) {
+               PADDLE_THROW(common::errors::InvalidArgument(
+                   "Tensor is not contiguous, cannot call "
+                   "share_buffer_with on it."));
+             }
+
+             // Get the current device ID.
+             int dev_id = platform::GetXPUCurrentDeviceId();
+             paddle::platform::SetXPUDeviceId(dev_id);
+             VLOG(6) << "[DEBUG XPU] _share_buffer_with: current XPU device = "
+                     << dev_id;
+
+             auto *xpu_ipc_allocation =
+                 dynamic_cast<memory::allocation::XpuIpcAllocation *>(
+                     src.Holder().get());
+
+             PADDLE_ENFORCE_NOT_NULL(
+                 xpu_ipc_allocation,
+                 common::errors::PreconditionNotMet(
+                     "Tensor is not Xpu IPC shared tensor. Now only Tensor "
+                     "shared by xpu ipc could use this api."));
+
+             size_t size = t[0].cast<size_t>();
+             auto dtype = static_cast<phi::DataType>(t[1].cast<int>());
+             auto dims = common::make_ddim(
+                 t[2].cast<std::vector<int>>());
+             auto device_id = t[4].cast<int>();
+
+             auto shared_reader_holder =
+                 std::make_shared<memory::allocation::Allocation>(
+                     xpu_ipc_allocation->ptr(),
+                     xpu_ipc_allocation->base_ptr(), size,
+                     phi::XPUPlace(device_id));
+
+             self.ResetHolderWithType(shared_reader_holder, dtype);
+             self.Resize(dims);
+
+             VLOG(6) << "[DEBUG XPU] Reconstructed tensor with buffer shared!";
+           },
+           R"DOC(
+           Deserialize XPU Tensor for existed shared Xpu IPC tensor.
+
+           Params:
+               tensor: Shared Xpu IPC tensor.
+               tuple: contains data size, data type, tensor dims, lod
+                      information, device index.
+           )DOC")
+      .def("_share_xpu",
+           [](phi::DenseTensor &self) {
+             if (!self.IsInitialized() || self.numel() == 0)
+               throw std::runtime_error(
+                   "Tensor not initialized or numel is 0. could not pass to "
+                   "shared memory.");
+
+             // Get the current device ID.
+             int dev_id = platform::GetXPUCurrentDeviceId();
+             paddle::platform::SetXPUDeviceId(dev_id);
+             VLOG(6) << "[DEBUG XPU] _share_xpu: current XPU device = "
+                     << dev_id;
+
+             auto *holder = dynamic_cast<memory::allocation::Allocation *>(
+                 self.Holder().get());
+             PADDLE_ENFORCE_EQ(
+                 phi::is_xpu_place(holder->place()), true,
+                 common::errors::InvalidArgument(
+                     "Tensor is not on XPU. share_xpu only supports XPU "
+                     "tensors."));
+             void *base_ptr = holder->base_ptr();
+             ptrdiff_t offset_bytes = reinterpret_cast<char *>(holder->ptr()) -
+                                      reinterpret_cast<char *>(base_ptr);
+             VLOG(6) << "[DEBUG XPU] _share_xpu: base_ptr = " << base_ptr
+                     << ", offset_bytes = " << offset_bytes;
+             cudaIpcMemHandle_t handle;
+             int ret = cudaIpcGetMemHandle(&handle, base_ptr);
+             VLOG(6) << "[DEBUG XPU] _share_xpu: cudaIpcGetMemHandle returned: "
+                     << ret;
+             PADDLE_ENFORCE_XPU_SUCCESS(ret);
+             // Use the correct size for the IPC handle.
+             auto _handle = py::bytes(
+                 reinterpret_cast<char *>(&handle),
+                 (py::ssize_t)sizeof(cudaIpcMemHandle_t));
+             VLOG(6) << "[DEBUG XPU] _share_xpu: IPC handle (bytes) = "
+                     << _handle;
+             const auto &device_id =
+                 paddle::platform::GetXPUCurrentDeviceId();
+             auto stream = paddle::platform::get_current_stream(device_id);
+             xpu_wait(stream);
+             int type_idx = static_cast<int>(self.type());
+             size_t data_size = self.numel() *
+                 framework::SizeOfType(
+                     framework::TransToProtoVarType(self.type()));
+             VLOG(6) << "[DEBUG XPU] _share_xpu: data_size = " << data_size;
+             return py::make_tuple(_handle,
+                                   (py::size_t)offset_bytes,
+                                   data_size,
+                                   type_idx,
+                                   common::vectorize(self.dims()),
+                                   self.lod(),
+                                   device_id);
+           },
+           R"DOC(
+           Serialize XPU Tensor by IPC.
+
+           Returns:
+               tuple: contains handle, offset, data size, data type,
+                      tensor dims, lod information, and device id.
+           )DOC")
+      .def("_new_shared_xpu",
+           [](py::tuple t) {
+             if (t.size() != 7)
+               throw std::runtime_error(
+                   "Invalid Tensor meta info for shared xpu tensor!");
+
+             // Get the current device ID.
+             int dev_id = platform::GetXPUCurrentDeviceId();
+             paddle::platform::SetXPUDeviceId(dev_id);
+             VLOG(6) << "[DEBUG XPU] _new_shared_xpu: current XPU device = "
+                     << dev_id;
+
+             phi::DenseTensor tensor;
+             const std::string &handle = t[0].cast<std::string>();
+             ptrdiff_t offset_bytes = (ptrdiff_t)t[1].cast<int64_t>();
+             auto device_id = t[6].cast<int>();
+             VLOG(6) << "[DEBUG XPU] _new_shared_xpu: handle = " << handle
+                     << ", offset_bytes = " << offset_bytes;
+             auto base_ptr = memory::allocation::GetIpcBasePtr(handle);
+             size_t size = t[2].cast<size_t>();
+             void *dev = base_ptr.get();
+             dev = reinterpret_cast<char *>(dev) + offset_bytes;
+             auto shared_holder =
+                 std::make_shared<memory::allocation::XpuIpcAllocation>(
+                     dev, size, device_id, std::move(base_ptr));
+             tensor.ResetHolderWithType(
+                 shared_holder,
+                 static_cast<phi::DataType>(t[3].cast<int>()));
+             tensor.Resize(common::make_ddim(
+                 t[4].cast<std::vector<int>>()));
+             VLOG(6) << "[DEBUG XPU] _new_shared_xpu: Reshape tensor dims: "
+                     << tensor.dims();
+             return tensor;
+           },
+           R"DOC(
+           Deserialize XPU Tensor from IPC.
+
+           Params:
+               tuple: contains handle, offset, data size, data type,
+                      tensor dims, lod information, and device index.
+
+           Returns:
+               A new DenseTensor that shares memory via IPC.
+           )DOC")
+#endif  // PADDLE_WITH_XPU
       .def("_share_filename",
            [](phi::DenseTensor &self, bool use_file_descriptor) {
              if (!self.IsInitialized() || self.numel() == 0)
 
@@ -19,19 +19,40 @@
 
 #include "paddle/phi/api/profiler/event.h"
 #include "paddle/phi/core/platform/device_event_base.h"
+
 #if defined(PADDLE_WITH_XPU)
+#include <cuda.h>
+#include <cuda_runtime.h>
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "xpu/runtime.h"
+#include "xpu/runtime_ex.h"
 #endif
 
 namespace py = pybind11;
 
-namespace paddle::pybind {
+namespace paddle {
+namespace platform {
+#ifdef PADDLE_WITH_XPU
+XPUStream get_current_stream(int device_id) {
+  if (device_id == -1) {
+    device_id = phi::backends::xpu::GetXPUCurrentDeviceId();
+  }
+  auto place = phi::XPUPlace(device_id);
+  auto *dev_ctx = static_cast<phi::XPUContext *>(
+      phi::DeviceContextPool::Instance().Get(place));
+  dev_ctx->Wait();
+  return dev_ctx->stream();
+}
+
+#endif
+}  // namespace platform
+namespace pybind {
 void BindXpuStream(py::module *m_ptr) {
   auto &m = *m_ptr;
 
   // Bind Methods
   m.def("_xpu_device_synchronize", [](int device_id) {
-#if defined(PADDLE_WITH_XPU)
+#ifdef PADDLE_WITH_XPU
     if (device_id == -1) {
       device_id = paddle::platform::GetXPUCurrentDeviceId();
     }
@@ -47,6 +68,75 @@ void BindXpuStream(py::module *m_ptr) {
         "Paddle is not compiled with XPU. Cannot visit device synchronize."));
 #endif
   });
-}
+  m.def(
+      "_get_current_stream",
+      [](int device_id) {
+#ifdef PADDLE_WITH_XPU
+        if (device_id == -1) {
+          device_id = paddle::platform::GetXPUCurrentDeviceId();
+        }
+        paddle::platform::SetXPUDeviceId(device_id);
+        return platform::get_current_stream(device_id);
+#else
+        PADDLE_THROW(
+            common::errors::Unavailable("Paddle is not compiled with CUDA. "
+                                        "Cannot visit device synchronize."));
+#endif
+      },
+      py::return_value_policy::reference);
+  m.def("_device_synchronize", [](int device_id) {
+#ifdef PADDLE_WITH_XPU
+    if (device_id == -1) {
+      device_id = paddle::platform::GetXPUCurrentDeviceId();
+    }
+
+    int curr_device_id = paddle::platform::GetXPUCurrentDeviceId();
+    paddle::platform::SetXPUDeviceId(device_id);
+    PADDLE_ENFORCE_XPU_SUCCESS(cudaDeviceSynchronize());
+    paddle::platform::SetXPUDeviceId(curr_device_id);
+#else
+    PADDLE_THROW(common::errors::Unavailable(
+        "Paddle is not compiled with CUDA. Cannot visit device synchronize."));
+#endif
+  });
+
+#ifdef PADDLE_WITH_XPU
+  py::class_<XPUStream>(m, "XPUStream", R"DOC(
+      The handle of the CUDA stream.
+
+      Parameters:
+          device(paddle.CUDAPlace()|int|None, optional): The device which wanted to allocate the stream.
+              If device is None or negative integer, device will be the current device.
+              If device is positive integer, it must less than the device count. Default: None.
+          priority(int|None, optional): The priority of stream. The priority can be 1(high) or 2(normal).
+              If priority is None, the priority is 2(normal). Default: None.
 
-}  // namespace paddle::pybind
+      Examples:
+          .. code-block:: python
+
+              >>> # doctest: +REQUIRES(env:GPU)
+              >>> import paddle
+              >>> s1 = paddle.device.cuda.Stream(paddle.CUDAPlace(0), 1)
+              >>> s2 = paddle.device.cuda.Stream(0, 1)
+              >>> s3 = paddle.device.cuda.Stream()
+
+      )DOC")
+      .def(
+          "synchronize",
+          [](XPUStream &self) { xpu_wait(self); },
+          R"DOC(
+          Waits for stream tasks to complete.
+
+          Examples:
+              .. code-block:: python
+
+                  >>> # doctest: +REQUIRES(env:GPU)
+                  >>> import paddle
+                  >>> s = paddle.device.cuda.Stream(paddle.CUDAPlace(0), 1)
+                  >>> s.synchronize()
+
+          )DOC");
+#endif
+}
+}  // namespace pybind
+}  // namespace paddle
@@ -17,9 +17,24 @@
 #include "pybind11/pybind11.h"
 #include "pybind11/stl.h"
 
+#ifdef PADDLE_WITH_XPU
+#include "paddle/phi/core/xpu_cuda_stream.h"
+#include "xpu/runtime.h"
+#include "xpu/runtime_ex.h"
+#else
+namespace phi {
+class XPUCUDAStream {};
+}  // namespace phi
+#endif
+
 namespace py = pybind11;
 
 namespace paddle {
+namespace platform {
+#ifdef PADDLE_WITH_XPU
+XPUStream get_current_stream(int device_id = -1);
+#endif
+}  // namespace platform
 namespace pybind {
 
 void BindXpuStream(py::module* m);
 
@@ -14,6 +14,11 @@ limitations under the License. */
 
 #pragma once
 
+#ifdef PADDLE_WITH_XPU
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+
 #include "paddle/phi/backends/xpu/xpu_header.h"
 #include "paddle/phi/core/enforce.h"
 #include "xre/cuda_runtime_api.h"
@@ -167,17 +172,24 @@ namespace details {
 template <typename T>
 struct ExternalApiType {};
 
+// Macro to define the external API type specialization
 #define DEFINE_EXTERNAL_API_TYPE(type, success_value) \
   template <>                                         \
   struct ExternalApiType<type> {                      \
     using Type = type;                                \
     static constexpr Type kSuccess = success_value;   \
   }
 
+// Existing specializations for XPU
 DEFINE_EXTERNAL_API_TYPE(int, XPU_SUCCESS);
 #ifdef PADDLE_WITH_XPU_BKCL
 DEFINE_EXTERNAL_API_TYPE(BKCLResult_t, BKCL_SUCCESS);
 #endif
+// Added specialization for cudaError_t to support CUDA errors in XPU
+// enforcement.
+#ifdef PADDLE_WITH_XPU
+DEFINE_EXTERNAL_API_TYPE(cudaError_t, cudaSuccess);
+#endif
 
 #undef DEFINE_EXTERNAL_API_TYPE
 
@@ -224,6 +236,7 @@ DEFINE_EXTERNAL_API_TYPE(BKCLResult_t, BKCL_SUCCESS);
       __THROW_ERROR_INTERNAL__(__summary__);                                   \
     }                                                                          \
   } while (0)
+
 #define PADDLE_ENFORCE_XRE_SUCCESS(COND)                            \
   do {                                                              \
     auto __cond__ = (COND);                                         \
 
@@ -14,6 +14,12 @@
 
 #include "paddle/phi/backends/xpu/xpu_context.h"
 
+#ifdef PADDLE_WITH_XPU
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include "paddle/phi/core/xpu_cuda_stream.h"
+#endif
+
 #include <memory>
 
 #include "glog/logging.h"