compiler: Initial investigations to deviceid handling when visible devices specified

Author: EdCaunt

devito/operator/operator.py

@@ -1403,6 +1403,9 @@ 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)
 
         # The amount of space available on the host

devito/types/parallel.py

@@ -275,17 +275,37 @@ class DeviceID(DeviceSymbol):
 
     name = 'deviceid'
 
-    @property
+    @cached_property
     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)}
         elif configuration['deviceid'] != self.default_value:
             return {self.name: configuration['deviceid']}
         else:
             return {self.name: self.default_value}
+        # FIXME: This should first check if CUDA_VISIBLE_DEVICES before the final else
+        # What should this use for multi-GPU runs?
+        # Possibly in the default_value
 
 
 class DeviceRM(DeviceSymbol):