arch: Validate device IDs and correct device initialization

devito/operator/operator.py

@@ -1432,38 +1432,34 @@ def _get_nbytes(self, i):
 
     @cached_property
     def _physical_deviceid(self):
-        if isinstance(self.platform, Device):
-            # Get the physical device ID (as CUDA_VISIBLE_DEVICES may be set)
-            logical_deviceid = self.get('deviceid', -1)
-            visible_device_var, visible_devices = get_visible_devices()
-            if logical_deviceid < 0:
-                rank = self.comm.Get_rank() if self.comm != MPI.COMM_NULL else 0
-                if visible_devices is None:
-                    # 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:
-                return logical_deviceid
-            elif len(visible_devices) == 1:
-                # Only one visible device, it's clearly the one we want
-                return visible_devices[0]
-            elif logical_deviceid <= len(visible_devices):
-                # Map the logical device ID to the physical one
-                return visible_devices[logical_deviceid]
-            else:
-                # Logical device ID is out of bounds, likely from oversubscription
-                # Print a warning and map modulo the number of visible devices
-                deviceid = visible_devices[logical_deviceid % len(visible_devices)]
-                warning(f"Logical device ID {logical_deviceid} is out of bounds "
-                        f"for {len(visible_devices)} visible devices"
-                        f" in {visible_device_var}."
-                        f"Mapping to device ID {deviceid} instead.")
-                return visible_devices[logical_deviceid % len(visible_devices)]
-        else:
+        if not isinstance(self.platform, Device):
             return None
 
+        # The logical device ID is relative to the visible devices, if any
+        # (e.g., CUDA_VISIBLE_DEVICES), otherwise to all the devices on the node
+        deviceid = self.get('deviceid', -1)
+        var, visible_devices = get_visible_devices()
+        if visible_devices is None:
+            ngpus = self.platform.numa_domains
+        else:
+            ngpus = len(visible_devices)
+
+        if deviceid < 0:
+            # Mirror the generated code, which sets device `rank % ngpus`
+            rank = self.comm.Get_rank() if self.comm != MPI.COMM_NULL else 0
+            deviceid = rank % ngpus if ngpus else rank
+        elif ngpus is not None and deviceid >= ngpus:
+            visible = f" ({var}={','.join(map(str, visible_devices))})" if var else ""
+            raise InvalidArgument(
+                f"Device {deviceid} does not exist: {ngpus} device(s) available"
+                f"{visible}"
+            )
+
+        if visible_devices is None:
+            return deviceid
+        else:
+            return visible_devices[deviceid]
+
     @cached_property
     def nbytes_avail_mapper(self):
         """

devito/passes/iet/langbase.py

@@ -409,8 +409,8 @@ def initialize(self, iet):
         The initialization follows a pattern which is applicable to virtually
         any target language:
 
-            1. Calling the init function (e.g., `acc_init(...)` for OpenACC)
-            2. Assignment of the target device to a host thread or an MPI rank
+            1. Assignment of the target device to a host thread or an MPI rank
+            2. Calling the init function (e.g., `acc_init(...)` for OpenACC)
             3. Introduction of user-level symbols (e.g., `deviceid` to allow
                users to select a specific device)
 
@@ -437,10 +437,10 @@ def _make_setdevice_seq(iet, nodes=()):
             devicetype = as_list(self.langbb[self.platform])
             deviceid = self.deviceid
 
-            return list(nodes) + [Conditional(
+            return [Conditional(
                 CondNe(deviceid, -1),
                 self.langbb['set-device']([deviceid] + devicetype)
-            )]
+            )] + list(nodes)
 
         def _make_setdevice_mpi(iet, objcomm, nodes=()):
             devicetype = as_list(self.langbb[self.platform])
@@ -453,13 +453,16 @@ def _make_setdevice_mpi(iet, objcomm, nodes=()):
             ngpus, call_ngpus = self.langbb._get_num_devices(self.platform)
 
             osdd_then = self.langbb['set-device']([deviceid] + devicetype)
-            osdd_else = self.langbb['set-device']([rank % ngpus] + devicetype)
+            osdd_else = Conditional(
+                ngpus > 0,
+                self.langbb['set-device']([rank % ngpus] + devicetype)
+            )
 
-            return list(nodes) + [Conditional(
+            return [Conditional(
                 CondNe(deviceid, -1),
                 osdd_then,
                 List(body=[rank_decl, rank_init, call_ngpus, osdd_else]),
-            )]
+            )] + list(nodes)
 
         @singledispatch
         def _initialize(iet):
@@ -471,6 +474,8 @@ def _(iet):
 
             devicetype = as_list(self.langbb[self.platform])
 
+            # The runtime is initialized after selecting the device, as otherwise
+            # it would also create a context on the default device
             try:
                 lang_init = [self.langbb['init'](devicetype)]
             except TypeError:

devito/types/parallel.py

@@ -312,6 +312,10 @@ def _arg_values(self, **kwargs):
         else:
             return {self.name: self.default_value}
 
+    def _arg_check(self, args, *_args, **kwargs):
+        # Raises if the device does not exist
+        args._physical_deviceid  # noqa: B018
+
 
 class DeviceRM(DeviceSymbol):
 

tests/test_gpu_common.py

@@ -201,21 +201,40 @@ def test_deviceid_per_rank(self, mode):
         Test that Device IDs set by the user on a per-rank basis do not
         get modified.
         """
-        # 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))
 
+        # Reversed order to ensure it is different to default
+        rank = grid.distributor.myrank
+        ngpus = op._platform.numa_domains
+        deviceid = (ngpus - 1 - rank) % ngpus
+
         with switchenv({i: None for i in self.visible_device_envs}):
             argmap = op.arguments(deviceid=deviceid)
             assert argmap._physical_deviceid == deviceid
 
+    @pytest.mark.parametrize('visible_devices', [None, '0,1'])
+    def test_nonexistent_deviceid(self, visible_devices):
+        """
+        Test that requesting a device that does not exist raises an error before
+        anything is launched.
+        """
+        grid = Grid(shape=(10, 10))
+        u = Function(name='u', grid=grid)
+
+        op = Operator(Eq(u, u+1), name='NonexistentDevice')
+
+        temp_env = {i: None for i in self.visible_device_envs}
+        temp_env['CUDA_VISIBLE_DEVICES'] = visible_devices
+        with switchenv(temp_env):
+            ngpus = op._platform.numa_domains if visible_devices is None else 2
+            with pytest.raises(InvalidArgument, match=f'Device {ngpus} does not exist'):
+                op.apply(deviceid=ngpus)
+
+            assert np.all(u.data == 0)
+
 
 class TestCodeGeneration:
 

tests/test_gpu_openacc.py

@@ -226,6 +226,20 @@ def test_std_max(self):
 
         assert '<algorithm>' in str(op)
 
+    def test_init_acc_env(self):
+        grid = Grid(shape=(3, 3, 3))
+
+        u = TimeFunction(name='u', grid=grid)
+
+        op = Operator(Eq(u.forward, u.dx+1), platform='nvidiaX', language='openacc',
+                      name='InitAcc')
+
+        # The device is selected before initializing the runtime
+        init = op.body.init[0].body
+        assert str(init[0]) ==\
+            'if (deviceid != -1)\n{\n  acc_set_device_num(deviceid,acc_device_nvidia);\n}'
+        assert str(init[1]) == 'acc_init(acc_device_nvidia);'
+
 
 class TestOperator:
 

tests/test_gpu_openmp.py

@@ -39,7 +39,8 @@ def test_init_omp_env_w_mpi(self, mode):
              '{\n  int rank = 0;\n'
              '  MPI_Comm_rank(comm,&rank);\n'
              '  int ngpus = omp_get_num_devices();\n'
-             '  omp_set_default_device((rank)%(ngpus));\n}')
+             '  if (ngpus > 0)\n'
+             '  {\n    omp_set_default_device((rank)%(ngpus));\n  }\n}')
 
     def test_basic(self):
         grid = Grid(shape=(3, 3, 3))