refactor(tests): centralise CUDA runtime handling

- share rank-aware device selection across tests and benchmarks
- query bandwidth information from the active CUDA device
- standardise MPI lifecycle in remaining CUDA tests
- propagate benchmark runtime failures across MPI ranks

Author: ia267

src/backend/backend_runtime.f90

@@ -56,6 +56,27 @@ module m_backend_runtime
 
 contains
 
+#ifdef CUDA
+  subroutine select_cuda_device(nrank, devnum)
+    !! Select a CUDA device round-robin by MPI rank.  The optional result
+    !! returns the device that CUDA reports as current after selection.
+    integer, intent(in) :: nrank
+    integer, optional, intent(out) :: devnum
+
+    integer :: ierr, ndevs, selected_device
+
+    ierr = cudaGetDeviceCount(ndevs)
+    if (ndevs < 1) then
+      error stop 'select_cuda_device: no CUDA devices available'
+    end if
+
+    ierr = cudaSetDevice(mod(nrank, ndevs))
+    ierr = cudaGetDevice(selected_device)
+
+    if (present(devnum)) devnum = selected_device
+  end subroutine select_cuda_device
+#endif
+
   subroutine init(self, mesh, separate_host_allocator)
     class(backend_runtime_t), target, intent(inout) :: self
     type(mesh_t), target, intent(inout) :: mesh
@@ -67,7 +88,7 @@ subroutine init(self, mesh, separate_host_allocator)
 #endif
 
 #ifdef CUDA
-    integer :: ierr, nrank, ndevs, devnum
+    integer :: ierr, nrank
 #endif
 
     dims = mesh%grid%vert_dims
@@ -80,12 +101,7 @@ subroutine init(self, mesh, separate_host_allocator)
 
 #ifdef CUDA
     call MPI_Comm_rank(MPI_COMM_WORLD, nrank, ierr)
-    ierr = cudaGetDeviceCount(ndevs)
-    if (ndevs < 1) then
-      error stop 'backend_runtime_t%init: no CUDA devices available'
-    end if
-    ierr = cudaSetDevice(mod(nrank, ndevs))
-    ierr = cudaGetDevice(devnum)
+    call select_cuda_device(nrank)
 
     self%backend_name = 'CUDA'
     self%cuda_allocator = cuda_allocator_t(dims, SZ)

src/backend/cuda/backend.f90

@@ -202,11 +202,13 @@ subroutine get_device_bw_info_cuda(self, mem_clock_rt, mem_bus_width, &
     integer, intent(out) :: mem_clock_rt
     integer, intent(out) :: mem_bus_width
     logical, intent(out) :: available
-    integer :: ierr
+    integer :: ierr, devnum
 
-    ierr = cudaDeviceGetAttribute(mem_clock_rt, cudaDevAttrMemoryClockRate, 0)
+    ierr = cudaGetDevice(devnum)
+    ierr = cudaDeviceGetAttribute(mem_clock_rt, cudaDevAttrMemoryClockRate, &
+                                  devnum)
     ierr = cudaDeviceGetAttribute(mem_bus_width, &
-                                  cudaDevAttrGlobalMemoryBusWidth, 0)
+                                  cudaDevAttrGlobalMemoryBusWidth, devnum)
     available = .true.
 
   end subroutine get_device_bw_info_cuda

tests/common/test_utils.f90

@@ -5,7 +5,8 @@ module m_test_utils
   implicit none
 
   private
-  public :: initialise_mpi, finalise_test, global_all, global_sum, checkerr, &
+  public :: initialise_mpi, finalise_test, global_all, global_sum, &
+            check_status, checkerr, &
             write_perf_metric, write_perf_minmax_metrics, &
             write_perf_summary, write_perf_minmax_summary, &
             write_device_bw_metric
@@ -48,6 +49,19 @@ subroutine global_sum(value)
                        MPI_COMM_WORLD, ierr)
   end subroutine global_sum
 
+  subroutine check_status(status, label, allpass)
+    !! Mark a test or benchmark as failed when an API returns a non-zero
+    !! status code. The caller decides where ranks must synchronise allpass.
+    integer, intent(in) :: status
+    character(len=*), intent(in) :: label
+    logical, intent(inout) :: allpass
+
+    if (status /= 0) then
+      write (stderr, '(a,": status=",i0)') trim(label), status
+      allpass = .false.
+    end if
+  end subroutine check_status
+
   subroutine finalise_test(allpass, nrank, finalize_mpi)
     !! Report the aggregate test result and finalise MPI, aborting via
     !! `error stop` (non-zero exit status) if any check failed. When nrank is

tests/performance/perf_cuda_penta.f90

@@ -13,7 +13,10 @@ program perf_cuda_penta
   use m_cuda_common, only: SZ
   use m_cuda_exec_dist, only: exec_dist_penta_compact
   use m_cuda_tdsops, only: cuda_tdsops_t, cuda_tdsops_init
-  use m_test_utils, only: initialise_mpi, write_perf_minmax_metrics, &
+  use m_backend_runtime, only: select_cuda_device
+  use m_test_utils, only: initialise_mpi, finalise_test, global_all, &
+                          check_status, &
+                          write_perf_minmax_metrics, &
                           write_perf_minmax_summary, write_device_bw_metric
 
   implicit none
@@ -31,23 +34,26 @@ program perf_cuda_penta
   integer :: n, n_block, n_halo, n_iters, n_warmup, n_glob, ndof, nrank, nproc
   integer :: pprev, pnext
   integer :: ierr
-  integer :: ndevs, devnum
+  integer :: devnum
   integer :: memClockRt, memBusWidth
   type(dim3) :: blocks, threads
   real(dp) :: dx
+  logical :: allpass = .true.
 
   call initialise_mpi(nrank, nproc, pprev, pnext)
   if (nrank == 0) then
     print *, 'Performance benchmark for compact10_penta (Lele 10 penta 1st-deriv)'
     print *, 'Scheme: alpha=0.5 beta=0.05, non-periodic, single-GPU Thomas'
     print *, 'Ranks:', nproc
   end if
-  call select_device()
+  call select_cuda_device(nrank, devnum)
   call configure_benchmark()
   call allocate_fields()
   call setup_backend()
-  call run_case('penta_lele10_nonper', dx, penta_bw)
-  call finalise()
+  call global_all(allpass)
+  if (allpass) call run_case('penta_lele10_nonper', dx, penta_bw)
+  call global_all(allpass)
+  call finalise_test(allpass, nrank)
 
 contains
 
@@ -72,12 +78,6 @@ subroutine allocate_fields()
     u_recv_e_dev = 0._dp
   end subroutine allocate_fields
 
-  subroutine select_device()
-    ierr = cudaGetDeviceCount(ndevs)
-    ierr = cudaSetDevice(mod(nrank, ndevs))
-    ierr = cudaGetDevice(devnum)
-  end subroutine select_device
-
   subroutine setup_backend()
     integer :: i, j, k
 
@@ -101,7 +101,9 @@ subroutine setup_backend()
     threads = dim3(SZ, 1, 1)
 
     ierr = cudaDeviceGetAttribute(memClockRt, cudaDevAttrMemoryClockRate, devnum)
+    call check_status(ierr, 'query CUDA memory clock rate', allpass)
     ierr = cudaDeviceGetAttribute(memBusWidth, cudaDevAttrGlobalMemoryBusWidth, devnum)
+    call check_status(ierr, 'query CUDA memory bus width', allpass)
   end subroutine setup_backend
 
   subroutine run_case(case_name, delta, consumed_bw)
@@ -119,12 +121,14 @@ subroutine run_case(case_name, delta, consumed_bw)
       call run_kernel()
     end do
     call sync_backend()
+    if (.not. allpass) return
 
     call start_timer(tstart)
     do iter = 1, n_iters
       call run_kernel()
     end do
     call sync_backend()
+    if (.not. allpass) return
     call stop_timer(tend)
 
     call collect_perf_stats(tend - tstart, consumed_bw, achieved_bw, &
@@ -180,7 +184,9 @@ end subroutine collect_perf_stats
 
   subroutine sync_backend()
     ierr = cudaDeviceSynchronize()
+    call check_status(ierr, 'synchronise CUDA device', allpass)
     call MPI_Barrier(MPI_COMM_WORLD, ierr)
+    call global_all(allpass)
   end subroutine sync_backend
 
   subroutine start_timer(t)
@@ -195,8 +201,4 @@ subroutine stop_timer(t)
     call cpu_time(t)
   end subroutine stop_timer
 
-  subroutine finalise()
-    call MPI_Finalize(ierr)
-  end subroutine finalise
-
 end program perf_cuda_penta

tests/performance/perf_cuda_reorder.f90

@@ -6,7 +6,10 @@ program perf_cuda_reorder
   use m_cuda_kernels_reorder, only: reorder_x2y, reorder_x2z, reorder_y2x, &
                                     reorder_y2z, reorder_z2x, reorder_z2y, &
                                     reorder_c2x, reorder_x2c
-  use m_test_utils, only: write_perf_metric, write_perf_summary, &
+  use m_backend_runtime, only: select_cuda_device
+  use m_test_utils, only: initialise_mpi, finalise_test, global_all, &
+                          check_status, &
+                          write_perf_metric, write_perf_summary, &
                           write_device_bw_metric
 
   implicit none
@@ -29,12 +32,17 @@ program perf_cuda_reorder
 
   integer :: ierr
   integer :: n_block, ndof
+  integer :: nrank, nproc, devnum
   integer :: memClockRt, memBusWidth
   real(dp), allocatable :: u_i(:, :, :)
   real(dp), device, allocatable :: u_i_d(:, :, :), u_o_d(:, :, :), &
                                    u_temp_d(:, :, :)
   real(dp), device, allocatable :: u_c_d(:, :, :)
   type(dim3) :: blocks, threads
+  logical :: allpass = .true.
+
+  call initialise_mpi(nrank, nproc)
+  call select_cuda_device(nrank, devnum)
 
   n_block = ny*nz/SZ
   ndof = nx*ny*nz
@@ -47,20 +55,26 @@ program perf_cuda_reorder
   call random_number(u_i)
   u_i_d = u_i
 
-  ierr = cudaDeviceGetAttribute(memClockRt, cudaDevAttrMemoryClockRate, 0)
+  ierr = cudaDeviceGetAttribute(memClockRt, cudaDevAttrMemoryClockRate, &
+                                devnum)
+  call check_status(ierr, 'query CUDA memory clock rate', allpass)
   ierr = cudaDeviceGetAttribute(memBusWidth, &
-                                cudaDevAttrGlobalMemoryBusWidth, 0)
-
-  call run_case('cuda_reorder_x2y', CASE_X2Y)
-  call run_case('cuda_reorder_x2z', CASE_X2Z)
-  call run_case('cuda_reorder_y2x', CASE_Y2X)
-  call run_case('cuda_reorder_y2z', CASE_Y2Z)
-  call run_case('cuda_reorder_z2x', CASE_Z2X)
-  call run_case('cuda_reorder_z2y', CASE_Z2Y)
-  call run_case('cuda_reorder_x2c', CASE_X2C)
-  call run_case('cuda_reorder_c2x', CASE_C2X)
-
-  call write_device_bw_metric(memClockRt, memBusWidth)
+                                cudaDevAttrGlobalMemoryBusWidth, devnum)
+  call check_status(ierr, 'query CUDA memory bus width', allpass)
+  call global_all(allpass)
+
+  if (allpass) call run_case('cuda_reorder_x2y', CASE_X2Y)
+  if (allpass) call run_case('cuda_reorder_x2z', CASE_X2Z)
+  if (allpass) call run_case('cuda_reorder_y2x', CASE_Y2X)
+  if (allpass) call run_case('cuda_reorder_y2z', CASE_Y2Z)
+  if (allpass) call run_case('cuda_reorder_z2x', CASE_Z2X)
+  if (allpass) call run_case('cuda_reorder_z2y', CASE_Z2Y)
+  if (allpass) call run_case('cuda_reorder_x2c', CASE_X2C)
+  if (allpass) call run_case('cuda_reorder_c2x', CASE_C2X)
+
+  if (allpass) call write_device_bw_metric(memClockRt, memBusWidth)
+  call global_all(allpass)
+  call finalise_test(allpass, nrank)
 
 contains
 
@@ -74,17 +88,20 @@ subroutine run_case(label, case_id)
     print *, 'Performance test:', trim(label)
 
     call prepare_input(case_id)
+    if (.not. allpass) return
 
     do iter = 1, n_warmup
       call launch_kernel(case_id)
     end do
     call sync_device()
+    if (.not. allpass) return
 
     call cpu_time(tstart)
     do iter = 1, n_iters
       call launch_kernel(case_id)
     end do
     call sync_device()
+    if (.not. allpass) return
     call cpu_time(tend)
 
     call write_perf_metric(label, tend - tstart, n_iters, ndof, consumed_bw)
@@ -157,6 +174,8 @@ end subroutine launch_kernel
 
   subroutine sync_device()
     ierr = cudaDeviceSynchronize()
+    call check_status(ierr, 'synchronise CUDA device', allpass)
+    call global_all(allpass)
   end subroutine sync_device
 
 end program perf_cuda_reorder

tests/performance/perf_cuda_transeq.f90

@@ -7,7 +7,10 @@ program perf_cuda_transeq
   use m_cuda_exec_dist, only: exec_dist_transeq_3fused
   use m_cuda_sendrecv, only: sendrecv_fields
   use m_cuda_tdsops, only: cuda_tdsops_t
-  use m_test_utils, only: initialise_mpi, write_perf_minmax_metrics, &
+  use m_backend_runtime, only: select_cuda_device
+  use m_test_utils, only: initialise_mpi, finalise_test, global_all, &
+                          check_status, &
+                          write_perf_minmax_metrics, &
                           write_perf_minmax_summary, write_device_bw_metric
 
   implicit none
@@ -34,19 +37,22 @@ program perf_cuda_transeq
   integer :: n, n_block, n_halo, n_iters, n_warmup, n_glob, ndof, nrank, nproc
   integer :: pprev, pnext
   integer :: ierr
-  integer :: ndevs, devnum
+  integer :: devnum
   integer :: memClockRt, memBusWidth
   type(dim3) :: blocks, threads
   real(dp) :: dx_per, nu
+  logical :: allpass = .true.
 
   call initialise_mpi(nrank, nproc, pprev, pnext)
   if (nrank == 0) print *, 'Performance benchmark with', nproc, 'ranks'
-  call select_device()
+  call select_cuda_device(nrank, devnum)
   call configure_benchmark()
   call allocate_fields()
   call setup_backend()
-  call run_case('periodic', dx_per, periodic_bw)
-  call finalise()
+  call global_all(allpass)
+  if (allpass) call run_case('periodic', dx_per, periodic_bw)
+  call global_all(allpass)
+  call finalise_test(allpass, nrank)
 
 contains
 
@@ -80,12 +86,6 @@ subroutine allocate_fields()
     allocate (d2u_recv_s_dev(SZ, 1, n_block), d2u_recv_e_dev(SZ, 1, n_block))
   end subroutine allocate_fields
 
-  subroutine select_device()
-    ierr = cudaGetDeviceCount(ndevs)
-    ierr = cudaSetDevice(mod(nrank, ndevs))
-    ierr = cudaGetDevice(devnum)
-  end subroutine select_device
-
   subroutine setup_backend()
     integer :: i, j, k
 
@@ -112,7 +112,9 @@ subroutine setup_backend()
 
     ierr = cudaDeviceGetAttribute(memClockRt, cudaDevAttrMemoryClockRate, &
                                   devnum)
+    call check_status(ierr, 'query CUDA memory clock rate', allpass)
     ierr = cudaDeviceGetAttribute(memBusWidth, cudaDevAttrGlobalMemoryBusWidth, devnum)
+    call check_status(ierr, 'query CUDA memory bus width', allpass)
   end subroutine setup_backend
 
   subroutine run_case(case_name, delta, consumed_bw)
@@ -130,12 +132,14 @@ subroutine run_case(case_name, delta, consumed_bw)
       call run_kernel()
     end do
     call sync_backend()
+    if (.not. allpass) return
 
     call start_timer(tstart)
     do iter = 1, n_iters
       call run_kernel()
     end do
     call sync_backend()
+    if (.not. allpass) return
     call stop_timer(tend)
 
     call collect_perf_stats(tend - tstart, consumed_bw, achieved_bw, &
@@ -216,7 +220,9 @@ end subroutine collect_perf_stats
 
   subroutine sync_backend()
     ierr = cudaDeviceSynchronize()
+    call check_status(ierr, 'synchronise CUDA device', allpass)
     call MPI_Barrier(MPI_COMM_WORLD, ierr)
+    call global_all(allpass)
   end subroutine sync_backend
 
   subroutine start_timer(t)
@@ -231,8 +237,4 @@ subroutine stop_timer(t)
     call cpu_time(t)
   end subroutine stop_timer
 
-  subroutine finalise()
-    call MPI_Finalize(ierr)
-  end subroutine finalise
-
 end program perf_cuda_transeq