compiler: Tweak behaviour of deviceid to respect per-rank DeviceID specified by user

Author: EdCaunt

devito/operator/operator.py

@@ -1393,13 +1393,16 @@ def _get_nbytes(self, i):
     def _physical_deviceid(self):
         if isinstance(self.platform, Device):
             # Get the physical device ID (as CUDA_VISIBLE_DEVICES may be set)
-            rank = self.comm.Get_rank() if self.comm != MPI.COMM_NULL else 0
+            logical_deviceid = self.get('deviceid', -1)
+            if logical_deviceid < 0:
+                rank = self.comm.Get_rank() if self.comm != MPI.COMM_NULL else 0
+                logical_deviceid = rank
 
-            logical_deviceid = max(self.get('deviceid', 0), 0) + rank
-            if self._visible_devices is None:
+            visible_devices = get_visible_devices()
+            if visible_devices is None:
                 return logical_deviceid
             else:
-                return get_visible_devices()[logical_deviceid]
+                return visible_devices[logical_deviceid]
         else:
             return None
 

tests/test_gpu_common.py

@@ -75,9 +75,10 @@ def test_autopad_with_platform_switch(self):
             assert f.shape_allocated[1] == 64
 
 
-class TestEnvironmentVariables:
+class TestDeviceID:
     """
-    Test that environment variables are correctly handled.
+    Test that device IDs and associated environment variables such as
+    CUDA_VISIBLE_DEVICES are correctly handled.
     """
 
     @pytest.mark.parametrize('env_variables', [{"cuda_visible_devices": "1"},
@@ -158,6 +159,26 @@ def test_visible_devices_with_devito_deviceid(self):
             # So should be the second of the two visible devices specified (3)
             assert argmap._physical_deviceid == 3
 
+    @pytest.mark.parallel(mode=2)
+    def test_deviceid_per_rank(self, mode):
+        """
+        Test that Device IDs set by the user on a per-rank basis do not
+        get modifed.
+        """
+        # Reversed order to ensure it is different to default
+        user_set_deviceids = (1, 0)
+
+        grid = Grid(shape=(10, 10))
+        u = Function(name='u', grid=grid)
+
+        rank = grid.distributor.myrank
+        deviceid = user_set_deviceids[rank]
+
+        op = Operator(Eq(u, u+1))
+
+        argmap = op.arguments(deviceid=deviceid)
+        assert argmap._physical_deviceid == deviceid
+
 
 class TestCodeGeneration: