Add Tensor::stream() for custom device (#73145)

YqGe585 · web-flow · commit badd66d7dd69 · 2025-06-15T19:08:57.000+08:00
* add Tensor::stream() for custom device

* move device_manager.h

* retrigger

* retrigger

* retrigger
diff --git a/paddle/phi/api/include/context_pool.h b/paddle/phi/api/include/context_pool.h
@@ -24,6 +24,7 @@ namespace phi {
 class DeviceContext;
 class CPUContext;
 class GPUContext;
+class CustomContext;
 class Allocator;
 class CUDAStream;
 }  // namespace phi
@@ -44,6 +45,11 @@ struct DefaultDeviceContextType<AllocationType::GPU> {
   using TYPE = phi::GPUContext;
 };
 
+template <>
+struct DefaultDeviceContextType<AllocationType::CUSTOM> {
+  using TYPE = phi::CustomContext;
+};
+
 /**
  * The DeviceContextPool here is just a mirror of the DeviceContextPool in
  * fluid, and does not manage the life cycle of the DeviceContext.
diff --git a/paddle/phi/api/include/tensor.h b/paddle/phi/api/include/tensor.h
@@ -29,6 +29,10 @@ using gpuStream_t = cudaStream_t;
 using gpuStream_t = hipStream_t;
 #endif
 
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/phi/backends/stream.h"
+#endif
+
 #include "paddle/common/layout.h"
 #include "paddle/common/macros.h"
 #include "paddle/phi/common/data_type.h"
@@ -430,6 +434,14 @@ class PADDLE_API Tensor final {
    * @return gpuStream_t
    */
   gpuStream_t stream() const;
+#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
+  /**
+   * @brief Get the stream where the tensor is currently located
+   * This is a deprecated method and may be removed in the future!
+   *
+   * @return stream_t
+   */
+  phi::stream::stream_t stream() const;
 #endif
 
   /**
diff --git a/paddle/phi/api/lib/tensor.cc b/paddle/phi/api/lib/tensor.cc
@@ -24,6 +24,9 @@ limitations under the License. */
 #include "paddle/phi/api/include/context_pool.h"
 #include "paddle/phi/api/lib/data_transform.h"
 #include "paddle/phi/api/lib/utils/allocator.h"
+#include "paddle/phi/backends/custom/custom_context.h"
+#include "paddle/phi/backends/device_manager.h"
+
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/backends/gpu/gpu_info.h"
 #include "paddle/phi/core/dense_tensor.h"
@@ -406,6 +409,20 @@ gpuStream_t Tensor::stream() const {
       GPUPlace(device_id));
   return gpu_context->stream();
 }
+#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
+phi::stream::stream_t Tensor::stream() const {
+  auto dev_types = phi::DeviceManager::GetAllCustomDeviceTypes();
+  for (const auto &dev_type : dev_types) {
+    int device_id = phi::DeviceManager::GetDevice(dev_type);
+    auto *custom_context =
+        DeviceContextPool::Instance().Get<AllocationType::CUSTOM>(
+            phi::CustomPlace(dev_type, device_id));
+    return custom_context->stream();
+  }
+  PADDLE_THROW(common::errors::Unimplemented(
+      "There is no custom device context when calling Tensor::stream()."));
+  return nullptr;
+}
 #endif
 
 const std::string &Tensor::name() const { return name_; }
