comm_env_setup_roles Subroutine

   subroutine comm_env_setup_roles(use_io_server)
      !! Phase 2: Assign roles (I/O server vs compute), bind GPU
      !! Call this after reading config, before decomp/solver init.
      logical, intent(in) :: use_io_server
         !! Dedicate one rank per node as I/O server

      type(comm_t) :: node_comm
      integer :: node_rank, node_size, gpu_rank, i, n_dev
      integer, allocatable :: world_ranks(:)

      ! Compute node-local rank for GPU binding and I/O server assignment
      node_comm = comm_global%split()
      node_rank = node_comm%rank()
      node_size = node_comm%size()

      if (use_io_server .and. cached_size > 1 .and. node_size > 1) then
         ! Last rank on each node becomes the I/O server
         cached_is_io_server = (node_rank == node_size - 1)

         ! Broadcast I/O server's world rank to all ranks on this node
         cached_io_server_rank = cached_rank
         call bcast(node_comm, cached_io_server_rank, 1, node_size - 1)

         ! Gather world ranks of all node-local processes to build
         ! the compute rank list (exclude the I/O server)
         ! Use bcast from each rank as a simple allgather substitute
         allocate (world_ranks(0:node_size - 1))
         do i = 0, node_size - 1
            world_ranks(i) = cached_rank
            call bcast(node_comm, world_ranks(i), 1, i)
         end do

         cached_node_n_compute = node_size - 1
         allocate (cached_node_compute_ranks(cached_node_n_compute))
         do i = 0, node_size - 2
            cached_node_compute_ranks(i + 1) = world_ranks(i)
         end do
         deallocate (world_ranks)

         ! Create compute communicator (excludes I/O ranks)
         if (cached_is_io_server) then
            comm_compute = comm_global%split_by(1)  ! color 1 = I/O
         else
            comm_compute = comm_global%split_by(0)  ! color 0 = compute
         end if

         cached_compute_rank = comm_compute%rank()
         cached_compute_size = comm_compute%size()

         ! GPU binding: compute ranks bind to their node-local rank
         gpu_rank = node_rank
      else
         ! No I/O server: all ranks are compute
         cached_is_io_server = .false.
         cached_io_server_rank = -1
         comm_compute = comm_global
         cached_compute_rank = cached_rank
         cached_compute_size = cached_size
         cached_node_n_compute = node_size
         allocate (cached_node_compute_ranks(node_size))
         do i = 0, node_size - 1
            cached_node_compute_ranks(i + 1) = cached_rank - node_rank + i
         end do
         gpu_rank = node_rank
      end if

      call node_comm%finalize()

      ! Bind GPU (compute ranks only). Two API calls are needed when both
      ! `do concurrent` (stdpar) and `!$omp target` regions co-exist:
      !
      !   * acc_set_device_num — selects the OpenACC runtime device, which
      !     is what stdpar (NVHPC, Cray) uses to dispatch `do concurrent`
      !     loops. Active whenever the compiler exposes the OpenACC
      !     runtime — see RDB_HAS_OPENACC_RUNTIME in compiler_flags.
      !   * omp_set_default_device — selects the OpenMP target device,
      !     used by `!$omp target` regions on the dc-openmp backend.
      !     Active when -mp= / -fopenmp is on.
      !
      ! On the OpenACC backend (-acc=gpu) only the first fires; on the
      ! OpenMP backend both fire. Either way, every rank lands on the
      ! right GPU under mpirun. acc_device_default lets the runtime
      ! pick the appropriate device kind (NVIDIA, AMD, etc.).
      if (.not. cached_is_io_server) then
#ifdef RDB_HAS_OPENACC_RUNTIME
         ! Clamp the node-local GPU index to the number of devices this
         ! process can actually see. With all GPUs visible this is a no-op
         ! (mod(node_rank, ndev) == node_rank for node_rank < ndev). When the
         ! launcher pins one GPU per rank via CUDA_VISIBLE_DEVICES, only one
         ! device is visible, so this binds device 0 (the rank's own GPU) and
         ! the CUDA-aware-MPI primary context also lands there -- no leftover
         ! context on the global device 0. Also makes ranks > GPUs (over-
         ! subscription) bind round-robin instead of failing.
         !
         ! REVISIT: this clamp may be unnecessary -- NVHPC's acc_set_device_num
         ! might already tolerate an out-of-range device num under one-GPU
         ! CUDA_VISIBLE_DEVICES pinning (clamping internally). It's untested
         ! (the runs that confirmed the fix used this clamp). If a no-clamp
         ! build + pin reports a clean bind, drop this. Kept for now because
         ! it's a no-op when all GPUs are visible and the spec calls
         ! out-of-range device nums undefined, so it's the portable choice.
         n_dev = acc_get_num_devices(acc_device_default)
         if (n_dev > 0) gpu_rank = mod(gpu_rank, n_dev)
         call acc_set_device_num(gpu_rank, acc_device_default)
#endif
#ifdef _OPENMP
         call omp_set_default_device(gpu_rank)
#endif
         call print_gpu_binding(gpu_rank)
      end if

   end subroutine comm_env_setup_roles