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/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/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: 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))