diff --git a/devito/operator/operator.py b/devito/operator/operator.py index b3690caf74..c1a3941acc 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -1439,7 +1439,9 @@ def _physical_deviceid(self): if logical_deviceid < 0: rank = self.comm.Get_rank() if self.comm != MPI.COMM_NULL else 0 if visible_devices is None: - logical_deviceid = rank + # Mirror the generated code, which sets device `rank % ngpus` + ngpus = self.platform.numa_domains + logical_deviceid = rank % ngpus if ngpus else rank else: logical_deviceid = rank % len(visible_devices) if visible_devices is None: diff --git a/tests/test_gpu_common.py b/tests/test_gpu_common.py index 3efda2820f..6a03f4d6f0 100644 --- a/tests/test_gpu_common.py +++ b/tests/test_gpu_common.py @@ -17,7 +17,9 @@ Conditional, Expression, FindNodes, FindSymbols, Section, retrieve_iteration_tree ) from devito.passes.iet.languages.openmp import OmpIteration -from devito.types import DeviceID, DeviceRM, Lock, NPThreads, PThreadArray, Symbol +from devito.types import ( + DeviceID, DeviceRM, Lock, NPThreads, PThreadArray, Symbol, device_layer +) pytestmark = skipif(['nodevice'], whole_module=True) @@ -163,11 +165,14 @@ def test_visible_devices_mpi(self, visible_devices, mode): argmap1 = op1.arguments() assert argmap1._physical_deviceid == expected - # In default case, physical deviceid will equal rank + # In default case, physical deviceid is `rank % ngpus`, as in the + # generated code, so its memory can be queried also with more ranks + # than devices with switchenv({i: None for i in self.visible_device_envs}): op2 = Operator(eq) argmap2 = op2.arguments() - assert argmap2._physical_deviceid == rank + assert argmap2._physical_deviceid == rank % op2._platform.numa_domains + assert isinstance(argmap2.nbytes_avail_mapper[device_layer], int) def test_visible_devices_with_devito_deviceid(self): """Test interaction between CUDA_VISIBLE_DEVICES and DEVITO_DEVICEID"""