[XPU] do not use null stream when CDNN_CLUSTER_PARALLEL is ON (#67276)

lj970926 · web-flow · commit c0fced7295ce · 2024-08-14T10:43:16.000+08:00
diff --git a/paddle/phi/core/memory/allocation/allocator_facade.cc b/paddle/phi/core/memory/allocation/allocator_facade.cc
@@ -1893,6 +1893,12 @@ const std::shared_ptr<Allocator>& AllocatorFacade::GetAllocator(
   }
   return m->GetAllocator(place, /* A non-zero num to choose allocator_ */ 1);
 }
+void AllocatorFacade::SetDefaultStream(const phi::XPUPlace& place,
+                                       XPUStream stream) {
+  if (m_->IsStreamSafeCUDAAllocatorUsed()) {
+    m_->SetDefaultStream(place, stream);
+  }
+}
 #endif
 
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
diff --git a/paddle/phi/core/memory/allocation/allocator_facade.h b/paddle/phi/core/memory/allocation/allocator_facade.h
@@ -97,6 +97,7 @@ class AllocatorFacade {
 #elif defined(PADDLE_WITH_XPU)
   TEST_API const std::shared_ptr<Allocator>& GetAllocator(
       const phi::Place& place, XPUStream stream);
+  void SetDefaultStream(const phi::XPUPlace& place, XPUStream stream);
 #endif
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
diff --git a/paddle/phi/core/platform/device_context.cc b/paddle/phi/core/platform/device_context.cc
@@ -33,6 +33,10 @@ limitations under the License. */
 #include "paddle/phi/core/platform/cuda_device_guard.h"
 #endif
 
+#if defined(PADDLE_WITH_XPU)
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#endif
+
 namespace paddle {
 namespace platform {
 
@@ -105,7 +109,12 @@ inline std::unique_ptr<DeviceContext> CreateDeviceContext(
 #endif
   } else if (p.GetType() == phi::AllocationType::XPU) {
 #if defined(PADDLE_WITH_XPU)
-    dev_ctx->SetAllocator(instance.GetAllocator(p).get());
+    auto* xpu_ctx = dynamic_cast<phi::XPUContext*>(dev_ctx);
+    if (!disable_setting_default_stream_for_allocator) {
+      instance.SetDefaultStream(phi::XPUPlace(p.GetDeviceId()),
+                                xpu_ctx->stream());
+    }
+    dev_ctx->SetAllocator(instance.GetAllocator(p, xpu_ctx->stream()).get());
     dev_ctx->SetGenerator(phi::DefaultXPUGenerator(p.GetDeviceId()).get());
 #endif
 #ifdef PADDLE_WITH_CUSTOM_DEVICE

Original file line number	Diff line number	Diff line change
`@@ -1893,6 +1893,12 @@ const std::shared_ptr<Allocator>& AllocatorFacade::GetAllocator(`
`1893`	`1893`	`}`
`1894`	`1894`	`return m->GetAllocator(place, /* A non-zero num to choose allocator_ */ 1);`
`1895`	`1895`	`}`
	`1896`	`+void AllocatorFacade::SetDefaultStream(const phi::XPUPlace& place,`
	`1897`	`+ XPUStream stream) {`
	`1898`	`+ if (m_->IsStreamSafeCUDAAllocatorUsed()) {`
	`1899`	`+ m_->SetDefaultStream(place, stream);`
	`1900`	`+ }`
	`1901`	`+}`
`1896`	`1902`	`#endif`
`1897`	`1903`
`1898`	`1904`	`#ifdef PADDLE_WITH_CUSTOM_DEVICE`