File tree Expand file tree Collapse file tree 1 file changed +9
-2
lines changed Expand file tree Collapse file tree 1 file changed +9
-2
lines changed Original file line number Diff line number Diff line change @@ -21,18 +21,25 @@ ::sycl::queue* DeviceManager::GetQueue(const DeviceOrd& device_spec) const {
2121 (collective::IsDistributed ());
2222 DeviceRegister& device_register = GetDevicesRegister ();
2323 if (not_use_default_selector) {
24- const int device_idx =
25- collective::IsDistributed () ? collective::GetRank () : device_spec.ordinal ;
2624 if (device_spec.IsSyclDefault ()) {
2725 auto & devices = device_register.devices ;
26+ const int device_idx = collective::IsDistributed ()
27+ ? collective::GetRank () % devices.size ()
28+ : device_spec.ordinal ;
2829 CHECK_LT (device_idx, devices.size ());
2930 queue_idx = device_idx;
3031 } else if (device_spec.IsSyclCPU ()) {
3132 auto & cpu_devices_idxes = device_register.cpu_devices_idxes ;
33+ const int device_idx = collective::IsDistributed ()
34+ ? collective::GetRank () % cpu_devices_idxes.size ()
35+ : device_spec.ordinal ;
3236 CHECK_LT (device_idx, cpu_devices_idxes.size ());
3337 queue_idx = cpu_devices_idxes[device_idx];
3438 } else if (device_spec.IsSyclGPU ()) {
3539 auto & gpu_devices_idxes = device_register.gpu_devices_idxes ;
40+ const int device_idx = collective::IsDistributed ()
41+ ? collective::GetRank () % gpu_devices_idxes.size ()
42+ : device_spec.ordinal ;
3643 CHECK_LT (device_idx, gpu_devices_idxes.size ());
3744 queue_idx = gpu_devices_idxes[device_idx];
3845 } else {
You can’t perform that action at this time.
0 commit comments