arch: Match default device memory queries to MPI placement

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:

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"""