From 4e141b445abe1cbd8aabf63a286f1d7df54a2f00 Mon Sep 17 00:00:00 2001 From: Gerard Gorman Date: Thu, 15 Oct 2026 16:18:10 +0000 Subject: [PATCH 1/2] arch: Raise on a device ID that does not exist An explicit `deviceid` at or beyond the number of devices (visible ones, if e.g. CUDA_VISIBLE_DEVICES is set) was silently ignored by the generated code, which then ran on the current device, while the memory query either raised IndexError (an off-by-one on the visible devices) or wrapped around them. Check it along with the other runtime arguments and raise an InvalidArgument naming the device and the number available --- devito/operator/operator.py | 56 +++++++++++++++++-------------------- devito/types/parallel.py | 4 +++ tests/test_gpu_common.py | 31 ++++++++++++++++---- 3 files changed, 55 insertions(+), 36 deletions(-) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index c1a3941acc..c2ab3e61e9 100644 --- a/devito/operator/operator.py +++ b/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): """ diff --git a/devito/types/parallel.py b/devito/types/parallel.py index b92978de1f..6432e04292 100644 --- a/devito/types/parallel.py +++ b/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): diff --git a/tests/test_gpu_common.py b/tests/test_gpu_common.py index 6a03f4d6f0..fcd125e367 100644 --- a/tests/test_gpu_common.py +++ b/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: From 3c244f802d5bf15018c92103dcd15fccc95ca498 Mon Sep 17 00:00:00 2001 From: Gerard Gorman Date: Thu, 15 Oct 2026 16:19:55 +0000 Subject: [PATCH 2/2] compiler: Select the device before initializing it and only if there is one The OpenACC runtime was initialized before selecting the device, so every rank not using the default device also created a context on it. And with no device found, the MPI device selection computed `rank % 0` --- devito/passes/iet/langbase.py | 19 ++++++++++++------- tests/test_gpu_openacc.py | 14 ++++++++++++++ tests/test_gpu_openmp.py | 3 ++- 3 files changed, 28 insertions(+), 8 deletions(-) diff --git a/devito/passes/iet/langbase.py b/devito/passes/iet/langbase.py index e6eea069f9..14dbe666aa 100644 --- a/devito/passes/iet/langbase.py +++ b/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: diff --git a/tests/test_gpu_openacc.py b/tests/test_gpu_openacc.py index eddae9fb1b..c6d67a5d6b 100644 --- a/tests/test_gpu_openacc.py +++ b/tests/test_gpu_openacc.py @@ -226,6 +226,20 @@ def test_std_max(self): assert '' 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: diff --git a/tests/test_gpu_openmp.py b/tests/test_gpu_openmp.py index 9c2ab3f991..c7989922f4 100644 --- a/tests/test_gpu_openmp.py +++ b/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))