未验证 提交 35af5818 编写于 作者: W wanghuancoder 提交者: GitHub

refine SynchronizeAllDevice (#53370)

上级 1bd468e2
...@@ -95,28 +95,34 @@ std::vector<std::vector<MemEvent>> GetMemEvents() { ...@@ -95,28 +95,34 @@ std::vector<std::vector<MemEvent>> GetMemEvents() {
void SynchronizeAllDevice() { void SynchronizeAllDevice() {
#ifdef PADDLE_WITH_CUDA #ifdef PADDLE_WITH_CUDA
int pre_device_id = GetCurrentDeviceId();
int count = GetGPUDeviceCount(); int count = GetGPUDeviceCount();
for (int i = 0; i < count; i++) { for (int i = 0; i < count; i++) {
SetDeviceId(i); SetDeviceId(i);
PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize()); PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
} }
SetDeviceId(pre_device_id);
#endif #endif
#ifdef PADDLE_WITH_HIP #ifdef PADDLE_WITH_HIP
int pre_device_id = GetCurrentDeviceId();
int count = GetGPUDeviceCount(); int count = GetGPUDeviceCount();
for (int i = 0; i < count; i++) { for (int i = 0; i < count; i++) {
SetDeviceId(i); SetDeviceId(i);
PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceSynchronize()); PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceSynchronize());
} }
SetDeviceId(pre_device_id);
#endif #endif
#ifdef PADDLE_WITH_CUSTOM_DEVICE #ifdef PADDLE_WITH_CUSTOM_DEVICE
auto dev_types = phi::DeviceManager::GetAllCustomDeviceTypes(); auto dev_types = phi::DeviceManager::GetAllCustomDeviceTypes();
for (const auto &dev_type : dev_types) { for (const auto &dev_type : dev_types) {
int pre_device_id = phi::DeviceManager::GetDevice(dev_type);
auto dev_cnt = phi::DeviceManager::GetDeviceCount(dev_type); auto dev_cnt = phi::DeviceManager::GetDeviceCount(dev_type);
for (size_t i = 0; i < dev_cnt; i++) { for (size_t i = 0; i < dev_cnt; i++) {
auto place = paddle::platform::CustomPlace(dev_type, i); auto place = paddle::platform::CustomPlace(dev_type, i);
phi::DeviceManager::SetDevice(place); phi::DeviceManager::SetDevice(place);
phi::DeviceManager::SynchronizeDevice(place); phi::DeviceManager::SynchronizeDevice(place);
} }
phi::DeviceManager::SetDevice(dev_type, pre_device_id);
} }
#endif #endif
} }
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册