Simplify device validation implementation following Devito patterns

Co-authored-by: ggorman <5394691+ggorman@users.noreply.github.com>

Author: Copilot

devito/passes/iet/langbase.py

@@ -3,7 +3,6 @@
 from abc import ABC
 
 import cgen as c
-import numpy as np
 
 from devito.data import FULL
 from devito.ir import (DummyExpr, Call, Conditional, Expression, List, Prodder,
@@ -13,7 +12,7 @@
 from devito.passes import is_on_device
 from devito.passes.iet.engine import iet_pass
 from devito.symbolics import Byref, CondNe, SizeOf
-from devito.types.relational import Ge
+from sympy import Ge
 from devito.tools import as_list, is_integer, prod
 from devito.types import Symbol, QueueID, Wildcard
 
@@ -63,7 +62,8 @@ def _get_num_devices(cls, platform):
         Get the number of accessible devices.
         Returns a tuple of (ngpus_symbol, call_to_get_num_devices).
         """
-        ngpus = Symbol(name='ngpus', dtype=np.int32)
+        from devito.types import Symbol
+        ngpus = Symbol(name='ngpus', dtype='int32')
         devicetype = as_list(cls[platform])
         call_ngpus = cls['num-devices'](devicetype, retobj=ngpus)
         return ngpus, call_ngpus
@@ -434,22 +434,19 @@ def _make_setdevice_seq(iet, nodes=()):
 
             # Add device validation check
             ngpus, call_ngpus = self.langbb._get_num_devices(self.platform)
-            
-            # Create validation: if deviceid >= num_devices, print error and exit
-            validation_check = Conditional(
+
+            validation = Conditional(
                 Ge(deviceid, ngpus),
                 List(body=[
-                    Call('printf', ['"%s: Error - Requested device ID %d does not exist. '
-                                  'Only %d device(s) available. Check CUDA_VISIBLE_DEVICES '
-                                  'and container GPU configuration.\\n"', 
-                                  self.langbb['name'], deviceid, ngpus]),
+                    Call('printf', ['"%s: Error - device %d >= %d devices\\n"',
+                                    self.langbb['name'], deviceid, ngpus]),
                     Call('exit', [1])
                 ])
             )
 
             device_setup = List(body=[
                 call_ngpus,
-                validation_check,
+                validation,
                 self.langbb['set-device']([deviceid] + devicetype)
             ])
 
@@ -468,21 +465,19 @@ def _make_setdevice_mpi(iet, objcomm, nodes=()):
 
             ngpus, call_ngpus = self.langbb._get_num_devices(self.platform)
 
-            # Add device validation check for explicit device ID
-            validation_check = Conditional(
+            # Add device validation for explicit device ID
+            validation = Conditional(
                 Ge(deviceid, ngpus),
                 List(body=[
-                    Call('printf', ['"%s: Error - Requested device ID %d does not exist. '
-                                  'Only %d device(s) available. Check CUDA_VISIBLE_DEVICES '
-                                  'and container GPU configuration.\\n"', 
-                                  self.langbb['name'], deviceid, ngpus]),
+                    Call('printf', ['"%s: Error - device %d >= %d devices\\n"',
+                                    self.langbb['name'], deviceid, ngpus]),
                     Call('exit', [1])
                 ])
             )
 
             osdd_then = List(body=[
                 call_ngpus,
-                validation_check,
+                validation,
                 self.langbb['set-device']([deviceid] + devicetype)
             ])
             osdd_else = self.langbb['set-device']([rank % ngpus] + devicetype)

tests/test_gpu_openacc.py

@@ -208,20 +208,19 @@ def test_device_validation_error_message(self):
 
         op = Operator(Eq(u.forward, u + 1), platform='nvidiaX', language='openacc')
 
-        # Check that the generated code contains device validation with informative error
+        # Check that the generated code contains device validation
         code = str(op)
-        
+
         # Should contain device count check
         assert 'acc_get_num_devices' in code, "Missing OpenACC device count check"
-        
+
         # Should contain validation condition
-        assert 'deviceid >= ngpus' in code, "Missing OpenACC device ID validation condition"
-        
-        # Should contain helpful error message components
-        assert 'does not exist' in code, "Missing 'does not exist' error message"
-        assert 'CUDA_VISIBLE_DEVICES' in code, "Missing CUDA_VISIBLE_DEVICES guidance"
-        assert 'container GPU configuration' in code, "Missing container guidance"
-        
+        assert 'deviceid >= ngpus' in code, "Missing OpenACC device ID " + \
+                                            "validation condition"
+
+        # Should contain error message
+        assert 'Error - device' in code, "Missing error message"
+
         # Should contain exit call to prevent undefined behavior
         assert 'exit(1)' in code, "Missing exit call on validation failure"
 

tests/test_gpu_openmp.py

@@ -25,7 +25,7 @@ def test_init_omp_env(self):
         assert 'if (deviceid != -1)' in init_code
         assert 'int ngpus = omp_get_num_devices()' in init_code
         assert 'if (deviceid >= ngpus)' in init_code
-        assert 'does not exist' in init_code
+        assert 'Error - device' in init_code
         assert 'omp_set_default_device(deviceid)' in init_code
 
     @pytest.mark.parallel(mode=1)
@@ -36,13 +36,14 @@ def test_init_omp_env_w_mpi(self, mode):
 
         op = Operator(Eq(u.forward, u.dx+1), language='openmp')
 
-        # With device validation, the MPI case also includes validation for explicit deviceid
+        # With device validation, the MPI case also includes validation for explicit
+        # deviceid
         init_code = str(op.body.init[0].body[0])
         assert 'if (deviceid != -1)' in init_code
         assert 'int ngpus = omp_get_num_devices()' in init_code
         # For MPI case with explicit deviceid, should have validation
         assert 'if (deviceid >= ngpus)' in init_code
-        assert 'does not exist' in init_code
+        assert 'Error - device' in init_code
         # Should still have MPI rank-based assignment in else clause
         assert 'int rank = 0' in init_code
         assert 'MPI_Comm_rank(comm,&rank)' in init_code
@@ -56,20 +57,18 @@ def test_device_validation_error_message(self):
 
         op = Operator(Eq(u.forward, u.dx+1), language='openmp')
 
-        # Check that the generated code contains device validation with informative error
+        # Check that the generated code contains device validation
         code = str(op)
-        
+
         # Should contain device count check
         assert 'omp_get_num_devices()' in code, "Missing device count check"
-        
+
         # Should contain validation condition
         assert 'deviceid >= ngpus' in code, "Missing device ID validation condition"
-        
-        # Should contain helpful error message components
-        assert 'does not exist' in code, "Missing 'does not exist' error message"
-        assert 'CUDA_VISIBLE_DEVICES' in code, "Missing CUDA_VISIBLE_DEVICES guidance"
-        assert 'container GPU configuration' in code, "Missing container guidance"
-        
+
+        # Should contain error message
+        assert 'Error - device' in code, "Missing error message"
+
         # Should contain exit call to prevent undefined behavior
         assert 'exit(1)' in code, "Missing exit call on validation failure"