misc: Move get_visible_devices to archinfo

Author: EdCaunt

devito/arch/archinfo.py

@@ -16,10 +16,11 @@
 from devito.logger import warning
 from devito.tools import as_tuple, all_equal, memoized_func
 
-__all__ = ['platform_registry', 'get_cpu_info', 'get_gpu_info', 'get_nvidia_cc',
-           'get_cuda_path', 'get_hip_path', 'check_cuda_runtime', 'get_m1_llvm_path',
-           'get_advisor_path', 'Platform', 'Cpu64', 'Intel64', 'IntelSkylake', 'Amd',
-           'Arm', 'Power', 'Device', 'NvidiaDevice', 'AmdDevice', 'IntelDevice',
+__all__ = ['platform_registry', 'get_cpu_info', 'get_gpu_info', 'get_visible_devices',
+           'get_nvidia_cc', 'get_cuda_path', 'get_hip_path', 'check_cuda_runtime',
+           'get_m1_llvm_path', 'get_advisor_path', 'Platform', 'Cpu64', 'Intel64',
+           'IntelSkylake', 'Amd', 'Arm', 'Power', 'Device', 'NvidiaDevice',
+           'AmdDevice', 'IntelDevice',
            # Brand-agnostic
            'ANYCPU', 'ANYGPU',
            # Intel CPUs
@@ -488,6 +489,27 @@ def parse_product_arch():
     return None
 
 
+def get_visible_devices():
+    device_vars = (
+        'CUDA_VISIBLE_DEVICES',
+        'ROCR_VISIBLE_DEVICES',
+        'HIP_VISIBLE_DEVICES'
+    )
+    for v in device_vars:
+        try:
+            return tuple(int(i) for i in os.environ[v].split(','))
+        except ValueError:
+            # Visible devices set via UUIDs or other non-integer identifiers.
+            warning("Setting visible devices via UUIDs or other non-integer"
+                    " identifiers is currently unsupported: environment variable"
+                    f" {v}={os.environ[v]} ignored.")
+        except KeyError:
+            # Environment variable not set
+            continue
+
+    return None
+
+
 @memoized_func
 def get_nvidia_cc():
     libnames = ('libcuda.so', 'libcuda.dylib', 'cuda.dll')

devito/operator/operator.py

@@ -1,4 +1,3 @@
-import os
 from collections import OrderedDict, namedtuple
 from functools import cached_property
 import ctypes
@@ -11,7 +10,8 @@
 import sympy
 import numpy as np
 
-from devito.arch import ANYCPU, Device, compiler_registry, platform_registry
+from devito.arch import (ANYCPU, Device, compiler_registry, platform_registry,
+                         get_visible_devices)
 from devito.data import default_allocator
 from devito.exceptions import (CompilationError, ExecutionError, InvalidArgument,
                                InvalidOperator)
@@ -1389,27 +1389,6 @@ 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:
-            try:
-                return tuple(int(i) for i in os.environ[v].split(','))
-            except ValueError:
-                # Visible devices set via UUIDs or other non-integer identifiers.
-                warning("Setting visible devices via UUIDs or other non-integer"
-                        " identifiers is currently unsupported: environment variable"
-                        f" {v}={os.environ[v]} ignored.")
-            except KeyError:
-                # Environment variable not set
-                continue
-
-        return None
-
     @cached_property
     def _physical_deviceid(self):
         if isinstance(self.platform, Device):
@@ -1420,7 +1399,7 @@ def _physical_deviceid(self):
             if self._visible_devices is None:
                 return logical_deviceid
             else:
-                return self._visible_devices[logical_deviceid]
+                return get_visible_devices()[logical_deviceid]
         else:
             return None