compiler: Make nbytes_avail_mapper aware of visible devices

Author: EdCaunt

devito/operator/operator.py

@@ -1,3 +1,4 @@
+import os
 from collections import OrderedDict, namedtuple
 from functools import cached_property
 import ctypes
@@ -1388,6 +1389,23 @@ def _get_nbytes(self, i):
 
         return nbytes
 
+    @cached_property
+    def visible_devices(self):
+        device_vars = (
+            'CUDA_VISIBLE_DEVICES',
+            'ROCR_VISIBLE_DEVICES',
+            'HIP_VISIBLE_DEVICES'
+        )
+        for v in device_vars:
+            if v in os.environ:
+                try:
+                    return tuple(int(i) for i in os.environ[v].split(','))
+                except ValueError:
+                    # Visible devices set via UUIDs or other non-integer identifiers
+                    continue
+
+        return None
+
     @cached_property
     def nbytes_avail_mapper(self):
         """
@@ -1402,11 +1420,16 @@ def nbytes_avail_mapper(self):
 
         # The amount of space available on the device
         if isinstance(self.platform, Device):
-            deviceid = max(self.get('deviceid', 0), 0)
-            # FIXME: I think this perhaps picks the wrong device when CUDA_VISIBLE_DEVICES set?
-            # Looks like it uses the physical device ID, not the logical one due to dependence
-            # on Nvidia SMI -> remote into Timewarp and check this
-            mapper[device_layer] = self.platform.memavail(deviceid=deviceid)
+            # 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 = max(self.get('deviceid', 0), 0) + rank
+            if self.visible_devices is not None:
+                physical_deviceid = self.visible_devices[logical_deviceid]
+            else:
+                physical_deviceid = logical_deviceid
+
+            mapper[device_layer] = self.platform.memavail(deviceid=physical_deviceid)
 
         # The amount of space available on the host
         try:

devito/types/parallel.py

@@ -279,23 +279,6 @@ class DeviceID(DeviceSymbol):
     def default_value(self):
         return -1
 
-    @cached_property
-    def visible_devices(self):
-        device_vars = (
-            'CUDA_VISIBLE_DEVICES',
-            'ROCR_VISIBLE_DEVICES',
-            'HIP_VISIBLE_DEVICES'
-        )
-        for v in device_vars:
-            if v in os.environ:
-                try:
-                    return tuple(int(i) for i in os.environ[v].split(','))
-                except ValueError:
-                    # Visible devices set via UUIDs or other non-integer identifiers
-                    continue
-
-        return None
-
     def _arg_values(self, **kwargs):
         if self.name in kwargs:
             return {self.name: kwargs.pop(self.name)}