| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | module simple_gpu | ||
| 2 | ! Provenance: | ||
| 3 | ! The trace_orbits_gpu_landreman path is an independent SIMPLE | ||
| 4 | ! implementation of the FIRM3D/Landreman tracing protocol. No FIRM3D | ||
| 5 | ! source code was copied. Reference implementation: | ||
| 6 | ! https://github.com/ColumbiaStellaratorTheory/firm3d | ||
| 7 | ! Related alpha-confinement optimization and loss-threshold protocol: | ||
| 8 | ! https://arxiv.org/abs/2606.19523 | ||
| 9 | ! FIRM3D is MIT-licensed, Copyright (c) 2025 ColumbiaStellaratorTheory; | ||
| 10 | ! the complete upstream notice is in LICENSES/FIRM3D-MIT.txt. | ||
| 11 | |||
| 12 | ! OpenACC GPU offload of the default orbit-tracing hot path: | ||
| 13 | ! Boozer field (isw_field_type=2), one particle per GPU thread. The | ||
| 14 | ! fixed-step symplectic Euler path and allocation-free adaptive RK5(4) | ||
| 15 | ! paths share the same resident field representation. | ||
| 16 | ! | ||
| 17 | ! The CPU integrator dispatches the field evaluation through a procedure | ||
| 18 | ! pointer, which cannot be called inside an OpenACC compute region. This | ||
| 19 | ! module keeps the device-side field evaluation local and reuses the shared | ||
| 20 | ! symplectic-Euler algebra from orbit_symplectic_euler1. Equivalence is | ||
| 21 | ! checked against the CPU path by test_gpu_orbit_bench. | ||
| 22 | use, intrinsic :: iso_fortran_env, only: dp => real64, int64 | ||
| 23 | use util, only: pi, sqrt2, twopi | ||
| 24 | use field_can_mod, only: field_can_t, get_val, get_derivatives, get_derivatives2 | ||
| 25 | use field_can_boozer, only: eval_field_booz_device | ||
| 26 | use orbit_symplectic_base, only: symplectic_integrator_t, & | ||
| 27 | SYMPLECTIC_STEP_OK, SYMPLECTIC_STEP_OUTSIDE_DOMAIN, & | ||
| 28 | SYMPLECTIC_STEP_MAXITER, SYMPLECTIC_STEP_LINEAR_SOLVE, & | ||
| 29 | SYMPLECTIC_STEP_BOUNDARY_LIMITED, symplectic_newton_warning_mode | ||
| 30 | use orbit_symplectic_euler1, only: sympl_euler1_newton_iter | ||
| 31 | use orbit_symplectic_euler1, only: sympl_euler1_extrapolate_field, sympl_euler1_advance_angles | ||
| 32 | use orbit_symplectic_base, only: EXPL_IMPL_EULER, MIDPOINT, CASH_KARP, & | ||
| 33 | DORMAND_PRINCE | ||
| 34 | use orbit_rk_core, only: rk_solve | ||
| 35 | use fortnum_ode_rk54_device, only: rk54_controls4_t, rk54_state4_t, & | ||
| 36 | rk54_initialize4, rk54_request4, rk54_supply4, & | ||
| 37 | RK54_CASH_KARP, RK54_DORMAND_PRINCE_TUNED, RK54_NEED_RHS, & | ||
| 38 | RK54_ACCEPTED, RK54_REJECTED, RK54_FAILED | ||
| 39 | use boozer_sub, only: boozer_state, splint_boozer_rk_device, & | ||
| 40 | BOOZER_SECDERS_RADIAL_MIXED | ||
| 41 | use boozer_rk_tables, only: rk_tables_ready, splint_boozer_rk_table_device | ||
| 42 | use params, only: gpu_num_devices, gpu_pilot_fraction, gpu_work_order | ||
| 43 | !$ use omp_lib, only: omp_get_thread_num | ||
| 44 | #ifdef _OPENACC | ||
| 45 | use openacc, only: acc_get_num_devices, acc_set_device_num, acc_device_nvidia | ||
| 46 | #endif | ||
| 47 | |||
| 48 | implicit none | ||
| 49 | private | ||
| 50 | |||
| 51 | integer, parameter :: maxit = 32 | ||
| 52 | |||
| 53 | public :: trace_orbits_gpu, trace_orbits_gpu_range | ||
| 54 | public :: trace_orbits_gpu_method, trace_orbits_gpu_rk54_range | ||
| 55 | public :: trace_orbits_gpu_midpoint_range | ||
| 56 | public :: trace_orbits_gpu_landreman | ||
| 57 | public :: evaluate_rhs_gpu | ||
| 58 | |||
| 59 | contains | ||
| 60 | |||
| 61 | ✗ | logical function gpu_finite_iterate(x) | |
| 62 | !$acc routine seq | ||
| 63 | real(dp), intent(in) :: x(:) | ||
| 64 | |||
| 65 | integer(int64), parameter :: exponent_mask = & | ||
| 66 | int(z'7FF0000000000000', int64) | ||
| 67 | integer(int64) :: bits | ||
| 68 | integer :: i | ||
| 69 | |||
| 70 | ✗ | gpu_finite_iterate = .false. | |
| 71 | ✗ | do i = 1, size(x) | |
| 72 | ✗ | bits = transfer(x(i), bits) | |
| 73 | ✗ | if (iand(bits, exponent_mask) == exponent_mask) return | |
| 74 | end do | ||
| 75 | ✗ | gpu_finite_iterate = .true. | |
| 76 | ✗ | end function gpu_finite_iterate | |
| 77 | |||
| 78 | 83220 | subroutine gpu_get_derivatives2_radial_mixed(f, pphi) | |
| 79 | ! Euler's two-variable Newton system consumes only derivatives with at | ||
| 80 | ! least one radial or canonical-momentum index. Avoid constructing the | ||
| 81 | ! angular-only Hessian entries used by the five-variable midpoint map. | ||
| 82 | !$acc routine seq | ||
| 83 | type(field_can_t), intent(inout) :: f | ||
| 84 | real(dp), intent(in) :: pphi | ||
| 85 | |||
| 86 | 83220 | call get_derivatives(f, pphi) | |
| 87 | |||
| 88 |
2/2✓ Branch 0 taken 249660 times.
✓ Branch 1 taken 83220 times.
|
332880 | f%d2vpar(1:3) = -f%d2Aph(1:3)/f%ro0 - f%d2hph(1:3)*f%vpar |
| 89 | 83220 | f%d2vpar(1) = f%d2vpar(1) - 2.0_dp*f%dhph(1)*f%dvpar(1) | |
| 90 | f%d2vpar(2) = f%d2vpar(2) - & | ||
| 91 | 83220 | (f%dhph(1)*f%dvpar(2) + f%dhph(2)*f%dvpar(1)) | |
| 92 | f%d2vpar(3) = f%d2vpar(3) - & | ||
| 93 | 83220 | (f%dhph(1)*f%dvpar(3) + f%dhph(3)*f%dvpar(1)) | |
| 94 |
2/2✓ Branch 0 taken 249660 times.
✓ Branch 1 taken 83220 times.
|
332880 | f%d2vpar(1:3) = f%d2vpar(1:3)/f%hph |
| 95 | |||
| 96 |
2/2✓ Branch 0 taken 249660 times.
✓ Branch 1 taken 83220 times.
|
332880 | f%d2H(1:3) = f%vpar*f%d2vpar(1:3) + f%mu*f%d2Bmod(1:3) |
| 97 | 83220 | f%d2H(1) = f%d2H(1) + f%dvpar(1)**2 | |
| 98 | 83220 | f%d2H(2) = f%d2H(2) + f%dvpar(1)*f%dvpar(2) | |
| 99 | 83220 | f%d2H(3) = f%d2H(3) + f%dvpar(1)*f%dvpar(3) | |
| 100 | |||
| 101 | f%d2pth(1:3) = f%d2vpar(1:3)*f%hth + & | ||
| 102 |
2/2✓ Branch 0 taken 249660 times.
✓ Branch 1 taken 83220 times.
|
332880 | f%vpar*f%d2hth(1:3) + f%d2Ath(1:3)/f%ro0 |
| 103 | 83220 | f%d2pth(1) = f%d2pth(1) + 2.0_dp*f%dvpar(1)*f%dhth(1) | |
| 104 | f%d2pth(2) = f%d2pth(2) + f%dvpar(1)*f%dhth(2) + & | ||
| 105 | 83220 | f%dvpar(2)*f%dhth(1) | |
| 106 | f%d2pth(3) = f%d2pth(3) + f%dvpar(1)*f%dhth(3) + & | ||
| 107 | 83220 | f%dvpar(3)*f%dhth(1) | |
| 108 | |||
| 109 |
2/2✓ Branch 0 taken 83220 times.
✓ Branch 1 taken 249660 times.
|
332880 | f%d2vpar(7:9) = -f%dhph/f%hph**2 |
| 110 |
2/2✓ Branch 0 taken 83220 times.
✓ Branch 1 taken 249660 times.
|
332880 | f%d2H(7:9) = f%dvpar(1:3)/f%hph + f%vpar*f%d2vpar(7:9) |
| 111 |
2/2✓ Branch 0 taken 249660 times.
✓ Branch 1 taken 83220 times.
|
332880 | f%d2pth(7:9) = f%dhth/f%hph + f%hth*f%d2vpar(7:9) |
| 112 | 83220 | end subroutine gpu_get_derivatives2_radial_mixed | |
| 113 | |||
| 114 | 41600 | subroutine gpu_newton1(si, f, x, xlast, warning_mode, status, nfev) | |
| 115 | !$acc routine seq | ||
| 116 | type(symplectic_integrator_t), intent(inout) :: si | ||
| 117 | type(field_can_t), intent(inout) :: f | ||
| 118 | real(dp), intent(inout) :: x(2) | ||
| 119 | real(dp), intent(out) :: xlast(2) | ||
| 120 | logical, intent(in) :: warning_mode | ||
| 121 | integer, intent(out) :: status | ||
| 122 | integer, intent(out) :: nfev | ||
| 123 | |||
| 124 | real(dp) :: tolref(2) | ||
| 125 | integer :: kit | ||
| 126 | logical :: converged, linear_failed, boundary_limited | ||
| 127 | logical :: step_boundary_limited | ||
| 128 | |||
| 129 | 41600 | status = SYMPLECTIC_STEP_MAXITER | |
| 130 | 41600 | nfev = 0 | |
| 131 | 41600 | boundary_limited = .false. | |
| 132 | |||
| 133 | 41600 | tolref(1) = 1d0 | |
| 134 | 41600 | tolref(2) = dabs(1d1*boozer_state%torflux/f%ro0) | |
| 135 | |||
| 136 |
1/2✓ Branch 0 taken 83220 times.
✗ Branch 1 not taken.
|
83220 | do kit = 1, maxit |
| 137 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 83220 times.
|
83220 | if (x(1) > 1d0) then |
| 138 | ✗ | status = SYMPLECTIC_STEP_OUTSIDE_DOMAIN | |
| 139 | 41600 | return | |
| 140 | end if | ||
| 141 | ! Transient guard; converged-negative handled by the caller | ||
| 142 | ! (mirrors newton1 in orbit_symplectic, #370). | ||
| 143 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 83220 times.
|
83220 | if (x(1) < 0d0) x(1) = 0.01d0 |
| 144 | |||
| 145 | call eval_field_booz_device(f, x(1), si%z(2), si%z(3), & | ||
| 146 | 83220 | BOOZER_SECDERS_RADIAL_MIXED) | |
| 147 | 83220 | nfev = nfev + 1 | |
| 148 | 83220 | call gpu_get_derivatives2_radial_mixed(f, x(2)) | |
| 149 | call sympl_euler1_newton_iter(si, f, x, tolref, xlast, converged, & | ||
| 150 | 83220 | linear_failed, step_boundary_limited) | |
| 151 |
2/4✓ Branch 0 taken 83220 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 83220 times.
|
83220 | boundary_limited = boundary_limited .or. step_boundary_limited |
| 152 | |||
| 153 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 83220 times.
|
83220 | if (linear_failed) then |
| 154 | ✗ | status = SYMPLECTIC_STEP_LINEAR_SOLVE | |
| 155 | ✗ | return | |
| 156 | end if | ||
| 157 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 83220 times.
|
83220 | if (x(1) > 1d0) then |
| 158 | ✗ | status = SYMPLECTIC_STEP_OUTSIDE_DOMAIN | |
| 159 | ✗ | return | |
| 160 | end if | ||
| 161 |
2/2✓ Branch 1 taken 41600 times.
✓ Branch 2 taken 41620 times.
|
166440 | if (converged) then |
| 162 | 41600 | status = SYMPLECTIC_STEP_OK | |
| 163 | 41600 | return | |
| 164 | end if | ||
| 165 | end do | ||
| 166 | ✗ | if (boundary_limited) then | |
| 167 | ✗ | if (step_boundary_limited .and. & | |
| 168 | abs(1d0 - x(1)) <= max(1d-10, 10d0*si%rtol)) then | ||
| 169 | ✗ | status = SYMPLECTIC_STEP_OUTSIDE_DOMAIN | |
| 170 | else | ||
| 171 | ✗ | status = SYMPLECTIC_STEP_BOUNDARY_LIMITED | |
| 172 | end if | ||
| 173 | else if (warning_mode .and. gpu_finite_iterate(x) .and. & | ||
| 174 | ✗ | gpu_finite_iterate(xlast) .and. gpu_finite_iterate(tolref) .and. & | |
| 175 | all(abs(x - xlast) <= & | ||
| 176 | si%warning_factor*si%rtol*abs(tolref))) then | ||
| 177 | ✗ | status = SYMPLECTIC_STEP_OK | |
| 178 | end if | ||
| 179 | ! Non-convergence diagnostics (CPU writes fort.6601) are omitted on device. | ||
| 180 | end subroutine gpu_newton1 | ||
| 181 | |||
| 182 | 41600 | subroutine gpu_timestep_euler(si, f, warning_mode, ierr, nfev, predictor) | |
| 183 | !$acc routine seq | ||
| 184 | type(symplectic_integrator_t), intent(inout) :: si | ||
| 185 | type(field_can_t), intent(inout) :: f | ||
| 186 | logical, intent(in) :: warning_mode | ||
| 187 | integer, intent(out) :: ierr | ||
| 188 | integer, intent(out) :: nfev | ||
| 189 | real(dp), intent(in), optional :: predictor(2) | ||
| 190 | |||
| 191 | real(dp) :: x(2), xlast(2), accepted_pthold | ||
| 192 | integer :: ktau, newton_status, newton_nfev | ||
| 193 | logical :: crossed | ||
| 194 | |||
| 195 | 41600 | ierr = 0 | |
| 196 | 41600 | nfev = 0 | |
| 197 | 41600 | ktau = 0 | |
| 198 |
2/2✓ Branch 0 taken 41600 times.
✓ Branch 1 taken 41600 times.
|
83200 | do while (ktau < si%ntau) |
| 199 | 41600 | accepted_pthold = si%pthold | |
| 200 | 41600 | si%pthold = f%pth | |
| 201 | |||
| 202 |
2/2✓ Branch 0 taken 41592 times.
✓ Branch 1 taken 8 times.
|
41600 | if (ktau == 0 .and. present(predictor)) then |
| 203 | 41592 | x = predictor | |
| 204 | else | ||
| 205 | 8 | x(1) = si%z(1) | |
| 206 | 8 | x(2) = si%z(4) | |
| 207 | end if | ||
| 208 | |||
| 209 | call gpu_newton1(si, f, x, xlast, warning_mode, newton_status, & | ||
| 210 | 41600 | newton_nfev) | |
| 211 | 41600 | nfev = nfev + newton_nfev | |
| 212 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
41600 | if (newton_status /= SYMPLECTIC_STEP_OK .and. ktau == 0 .and. & |
| 213 | present(predictor)) then | ||
| 214 | ✗ | si%pthold = accepted_pthold | |
| 215 | ✗ | x(1) = si%z(1) | |
| 216 | ✗ | x(2) = si%z(4) | |
| 217 | call gpu_newton1(si, f, x, xlast, warning_mode, newton_status, & | ||
| 218 | ✗ | newton_nfev) | |
| 219 | ✗ | nfev = nfev + newton_nfev | |
| 220 | end if | ||
| 221 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
|
41600 | if (newton_status /= SYMPLECTIC_STEP_OK) then |
| 222 | ✗ | si%pthold = accepted_pthold | |
| 223 | ✗ | ierr = newton_status | |
| 224 | ✗ | return | |
| 225 | end if | ||
| 226 | |||
| 227 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
|
41600 | if (x(1) > 1.0d0) then |
| 228 | ✗ | si%pthold = accepted_pthold | |
| 229 | ✗ | ierr = SYMPLECTIC_STEP_OUTSIDE_DOMAIN | |
| 230 | ✗ | return | |
| 231 | end if | ||
| 232 | 41600 | crossed = .false. | |
| 233 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
|
41600 | if (x(1) < 0.0d0) then |
| 234 | ! The converged radius lies beyond the axis: commit the | ||
| 235 | ! chart switch (r, theta) -> (|r|, theta + pi) (#370). | ||
| 236 | ✗ | x(1) = -x(1) | |
| 237 | ✗ | crossed = .true. | |
| 238 | ✗ | if (x(1) > 1.0d0) then | |
| 239 | ✗ | si%pthold = accepted_pthold | |
| 240 | ✗ | ierr = SYMPLECTIC_STEP_OUTSIDE_DOMAIN | |
| 241 | ✗ | return | |
| 242 | end if | ||
| 243 | ✗ | si%z(2) = si%z(2) + pi | |
| 244 | end if | ||
| 245 | |||
| 246 | 41600 | si%z(1) = x(1) | |
| 247 | 41600 | si%z(4) = x(2) | |
| 248 | |||
| 249 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
|
41600 | if (crossed) then |
| 250 | ! xlast lives in the other chart; extrapolation across the | ||
| 251 | ! flip is invalid, evaluate the field fresh instead. | ||
| 252 | ✗ | call eval_field_booz_device(f, si%z(1), si%z(2), si%z(3), 0) | |
| 253 | ✗ | call get_derivatives(f, si%z(4)) | |
| 254 | ✗ | nfev = nfev + 1 | |
| 255 | else | ||
| 256 | 41600 | call sympl_euler1_extrapolate_field(si, f, x, xlast) | |
| 257 | end if | ||
| 258 | 41600 | call sympl_euler1_advance_angles(si, f) | |
| 259 | |||
| 260 | 41600 | ktau = ktau + 1 | |
| 261 | end do | ||
| 262 | end subroutine gpu_timestep_euler | ||
| 263 | |||
| 264 | 94938 | subroutine gpu_midpoint_system(si, f, x, residual, jacobian) | |
| 265 | ! Five-equation degenerate-variational midpoint system and analytic | ||
| 266 | ! Jacobian, specialized to the resident Boozer field evaluator. | ||
| 267 | !$acc routine seq | ||
| 268 | type(symplectic_integrator_t), intent(in) :: si | ||
| 269 | type(field_can_t), intent(inout) :: f | ||
| 270 | real(dp), intent(in) :: x(5) | ||
| 271 | real(dp), intent(out) :: residual(5), jacobian(5, 5) | ||
| 272 | |||
| 273 | type(field_can_t) :: fmid | ||
| 274 | real(dp) :: dpthmid, pthdotbar | ||
| 275 | |||
| 276 | call eval_field_booz_device(f, x(5), 0.5_dp*(x(2) + si%z(2)), & | ||
| 277 | 94938 | 0.5_dp*(x(3) + si%z(3)), 2) | |
| 278 | 94938 | call get_derivatives2(f, 0.5_dp*(x(4) + si%z(4))) | |
| 279 | |||
| 280 | 94938 | residual(2) = f%dpth(1)*(x(2) - si%z(2)) - si%dt*f%dH(1) | |
| 281 | residual(3) = f%dpth(1)*f%hph*(x(3) - si%z(3)) - & | ||
| 282 | 94938 | si%dt*(f%dpth(1)*f%vpar - f%dH(1)*f%hth) | |
| 283 | residual(4) = f%dpth(1)*(x(4) - si%z(4)) + & | ||
| 284 | 94938 | si%dt*(f%dH(3)*f%dpth(1) - f%dH(1)*f%dpth(3)) | |
| 285 | residual(5) = f%dpth(1)*(f%pth - si%pthold) + 0.5_dp*si%dt*( & | ||
| 286 | 94938 | f%dpth(1)*f%dH(2) - f%dpth(2)*f%dH(1)) | |
| 287 | |||
| 288 | 94938 | jacobian(2, 1) = 0.0_dp | |
| 289 | 94938 | jacobian(2, 5) = f%d2pth(1)*(x(2) - si%z(2)) - si%dt*f%d2H(1) | |
| 290 | jacobian(2, 2:3) = 0.5_dp*(f%d2pth(2:3)*(x(2) - si%z(2)) - & | ||
| 291 |
2/2✓ Branch 0 taken 189876 times.
✓ Branch 1 taken 94938 times.
|
284814 | si%dt*f%d2H(2:3)) |
| 292 | 94938 | jacobian(2, 2) = jacobian(2, 2) + f%dpth(1) | |
| 293 | jacobian(2, 4) = 0.5_dp*(f%d2pth(7)*(x(2) - si%z(2)) - & | ||
| 294 | 94938 | si%dt*f%d2H(7)) | |
| 295 | |||
| 296 | 94938 | jacobian(3, 1) = 0.0_dp | |
| 297 | jacobian(3, 5) = (f%d2pth(1)*f%hph + f%dpth(1)*f%dhph(1))* & | ||
| 298 | (x(3) - si%z(3)) - si%dt*(f%d2pth(1)*f%vpar + & | ||
| 299 | 94938 | f%dpth(1)*f%dvpar(1) - f%d2H(1)*f%hth - f%dH(1)*f%dhth(1)) | |
| 300 | jacobian(3, 2:3) = 0.5_dp*((f%d2pth(2:3)*f%hph + & | ||
| 301 | f%dpth(1)*f%dhph(2:3))*(x(3) - si%z(3)) - si%dt*( & | ||
| 302 | f%d2pth(2:3)*f%vpar + f%dpth(1)*f%dvpar(2:3) - & | ||
| 303 |
2/2✓ Branch 0 taken 189876 times.
✓ Branch 1 taken 94938 times.
|
284814 | f%d2H(2:3)*f%hth - f%dH(1)*f%dhth(2:3))) |
| 304 | 94938 | jacobian(3, 3) = jacobian(3, 3) + f%dpth(1)*f%hph | |
| 305 | jacobian(3, 4) = 0.5_dp*(f%d2pth(7)*f%hph*(x(3) - si%z(3)) - & | ||
| 306 | si%dt*(f%d2pth(7)*f%vpar + f%dpth(1)*f%dvpar(4) - & | ||
| 307 | 94938 | f%d2H(7)*f%hth)) | |
| 308 | |||
| 309 | 94938 | jacobian(4, 1) = 0.0_dp | |
| 310 | jacobian(4, 5) = f%d2pth(1)*(x(4) - si%z(4)) + si%dt*( & | ||
| 311 | f%d2H(3)*f%dpth(1) + f%dH(3)*f%d2pth(1) - & | ||
| 312 | 94938 | f%d2H(1)*f%dpth(3) - f%dH(1)*f%d2pth(3)) | |
| 313 | jacobian(4, 2) = 0.5_dp*(f%d2pth(2)*(x(4) - si%z(4)) + si%dt*( & | ||
| 314 | f%d2H(5)*f%dpth(1) + f%dH(3)*f%d2pth(2) - & | ||
| 315 | 94938 | f%d2H(2)*f%dpth(3) - f%dH(1)*f%d2pth(5))) | |
| 316 | jacobian(4, 3) = 0.5_dp*(f%d2pth(3)*(x(4) - si%z(4)) + si%dt*( & | ||
| 317 | f%d2H(6)*f%dpth(1) + f%dH(3)*f%d2pth(3) - & | ||
| 318 | 94938 | f%d2H(3)*f%dpth(3) - f%dH(1)*f%d2pth(6))) | |
| 319 | jacobian(4, 4) = 0.5_dp*(f%d2pth(7)*(x(4) - si%z(4)) + si%dt*( & | ||
| 320 | f%dH(3)*f%d2pth(7) + f%d2H(9)*f%dpth(1) - & | ||
| 321 | 94938 | f%d2H(7)*f%dpth(3) - f%dH(1)*f%d2pth(9))) + f%dpth(1) | |
| 322 | |||
| 323 | 94938 | jacobian(5, 1) = 0.0_dp | |
| 324 | jacobian(5, 5) = f%d2pth(1)*(f%pth - si%pthold) + & | ||
| 325 | f%dpth(1)*f%dpth(1) + 0.5_dp*si%dt*(f%d2pth(1)*f%dH(2) + & | ||
| 326 | 94938 | f%dpth(1)*f%d2H(2) - f%d2pth(2)*f%dH(1) - f%dpth(2)*f%d2H(1)) | |
| 327 | jacobian(5, 2) = 0.5_dp*(f%d2pth(2)*(f%pth - si%pthold) + & | ||
| 328 | f%dpth(1)*f%dpth(2) + 0.5_dp*si%dt*(f%d2pth(2)*f%dH(2) + & | ||
| 329 | 94938 | f%dpth(1)*f%d2H(4) - f%d2pth(4)*f%dH(1) - f%dpth(2)*f%d2H(2))) | |
| 330 | jacobian(5, 3) = 0.5_dp*(f%d2pth(3)*(f%pth - si%pthold) + & | ||
| 331 | f%dpth(1)*f%dpth(3) + 0.5_dp*si%dt*(f%d2pth(3)*f%dH(2) + & | ||
| 332 | 94938 | f%dpth(1)*f%d2H(5) - f%d2pth(5)*f%dH(1) - f%dpth(2)*f%d2H(3))) | |
| 333 | jacobian(5, 4) = 0.5_dp*(f%d2pth(7)*(f%pth - si%pthold) + & | ||
| 334 | f%dpth(1)*f%dpth(4) + 0.5_dp*si%dt*(f%d2pth(7)*f%dH(2) + & | ||
| 335 | 94938 | f%dpth(1)*f%d2H(8) - f%d2pth(8)*f%dH(1) - f%dpth(2)*f%d2H(7))) | |
| 336 | |||
| 337 | 94938 | fmid = f | |
| 338 | 94938 | dpthmid = f%dpth(1) | |
| 339 | 94938 | pthdotbar = f%dpth(1)*f%dH(2) - f%dpth(2)*f%dH(1) | |
| 340 | 94938 | call eval_field_booz_device(f, x(1), x(2), x(3), 0) | |
| 341 | 94938 | call get_derivatives(f, x(4)) | |
| 342 | 94938 | residual(1) = dpthmid*(f%pth - si%pthold) + si%dt*pthdotbar | |
| 343 | |||
| 344 | 94938 | jacobian(1, 1) = fmid%dpth(1)*f%dpth(1) | |
| 345 | jacobian(1, 2) = 0.5_dp*(fmid%d2pth(2)*(f%pth - si%pthold) + & | ||
| 346 | si%dt*(fmid%d2pth(2)*fmid%dH(2) + fmid%dpth(1)*fmid%d2H(4) - & | ||
| 347 | fmid%dpth(2)*fmid%d2H(2) - fmid%d2pth(4)*fmid%dH(1))) + & | ||
| 348 | 94938 | fmid%dpth(1)*f%dpth(2) | |
| 349 | jacobian(1, 3) = 0.5_dp*(fmid%d2pth(3)*(f%pth - si%pthold) + & | ||
| 350 | si%dt*(fmid%d2pth(3)*fmid%dH(2) + fmid%dpth(1)*fmid%d2H(5) - & | ||
| 351 | fmid%dpth(2)*fmid%d2H(3) - fmid%d2pth(5)*fmid%dH(1))) + & | ||
| 352 | 94938 | fmid%dpth(1)*f%dpth(3) | |
| 353 | jacobian(1, 4) = 0.5_dp*(fmid%d2pth(7)*(f%pth - si%pthold) + & | ||
| 354 | si%dt*(fmid%d2pth(7)*fmid%dH(2) + fmid%dpth(1)*fmid%d2H(8) - & | ||
| 355 | fmid%dpth(2)*fmid%d2H(7) - fmid%d2pth(8)*fmid%dH(1))) + & | ||
| 356 | 94938 | fmid%dpth(1)*f%dpth(4) | |
| 357 | jacobian(1, 5) = fmid%d2pth(1)*(f%pth - si%pthold) + si%dt*( & | ||
| 358 | fmid%d2pth(1)*fmid%dH(2) + fmid%dpth(1)*fmid%d2H(2) - & | ||
| 359 | 94938 | fmid%dpth(2)*fmid%d2H(1) - fmid%d2pth(2)*fmid%dH(1)) | |
| 360 | 94938 | end subroutine gpu_midpoint_system | |
| 361 | |||
| 362 | 41600 | subroutine gpu_newton_midpoint(si, f, x, warning_mode, status, nfev) | |
| 363 | !$acc routine seq | ||
| 364 | type(symplectic_integrator_t), intent(in) :: si | ||
| 365 | type(field_can_t), intent(inout) :: f | ||
| 366 | real(dp), intent(inout) :: x(5) | ||
| 367 | logical, intent(in) :: warning_mode | ||
| 368 | integer, intent(out) :: status, nfev | ||
| 369 | |||
| 370 | integer :: kit, info, ir | ||
| 371 | integer, parameter :: radial_indices(2) = [1, 5] | ||
| 372 | real(dp) :: residual(5), jacobian(5, 5), xlast(5), tolref(5) | ||
| 373 | real(dp) :: residual_abs(5), correction_abs(5), scale, outward | ||
| 374 | logical :: limited | ||
| 375 | |||
| 376 | 41600 | status = SYMPLECTIC_STEP_MAXITER | |
| 377 | 41600 | nfev = 0 | |
| 378 | tolref = [1.0_dp, twopi, twopi, & | ||
| 379 |
2/2✓ Branch 0 taken 208000 times.
✓ Branch 1 taken 41600 times.
|
249600 | max(abs(f%Aph), abs(10.0_dp*boozer_state%torflux/f%ro0)), 1.0_dp] |
| 380 | 41600 | xlast = x | |
| 381 | |||
| 382 |
1/2✓ Branch 0 taken 94938 times.
✗ Branch 1 not taken.
|
94938 | do kit = 1, maxit |
| 383 |
2/4✓ Branch 0 taken 94938 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 94938 times.
|
94938 | if (x(1) > 1.0_dp .or. x(5) > 1.0_dp) then |
| 384 | ✗ | status = SYMPLECTIC_STEP_OUTSIDE_DOMAIN | |
| 385 | 41600 | return | |
| 386 | end if | ||
| 387 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 94938 times.
|
94938 | if (x(1) < 0.0_dp) x(1) = 0.01_dp |
| 388 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 94938 times.
|
94938 | if (x(5) < 0.0_dp) x(5) = 0.01_dp |
| 389 | |||
| 390 | 94938 | call gpu_midpoint_system(si, f, x, residual, jacobian) | |
| 391 | 94938 | nfev = nfev + 2 | |
| 392 |
2/2✓ Branch 0 taken 474690 times.
✓ Branch 1 taken 94938 times.
|
569628 | residual_abs = abs(residual) |
| 393 | 94938 | xlast = x | |
| 394 | 94938 | call rk_solve(5, jacobian, residual, info) | |
| 395 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 94938 times.
|
94938 | if (info /= 0) then |
| 396 | ✗ | status = SYMPLECTIC_STEP_LINEAR_SOLVE | |
| 397 | ✗ | return | |
| 398 | end if | ||
| 399 | |||
| 400 | scale = 1.0_dp | ||
| 401 | 284814 | limited = .false. | |
| 402 |
2/2✓ Branch 0 taken 189876 times.
✓ Branch 1 taken 94938 times.
|
284814 | do ir = 1, 2 |
| 403 | 189876 | outward = -residual(radial_indices(ir)) | |
| 404 |
2/2✓ Branch 0 taken 97897 times.
✓ Branch 1 taken 91979 times.
|
284814 | if (outward > 0.0_dp) then |
| 405 | scale = min(scale, 0.8_dp*max(0.0_dp, & | ||
| 406 | 97897 | 1.0_dp - x(radial_indices(ir)))/outward) | |
| 407 | end if | ||
| 408 | end do | ||
| 409 | 94938 | limited = scale < 1.0_dp | |
| 410 |
2/2✓ Branch 0 taken 474690 times.
✓ Branch 1 taken 94938 times.
|
569628 | x = x - scale*residual |
| 411 |
2/4✓ Branch 0 taken 94938 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 94938 times.
|
94938 | if (x(1) > 1.0_dp .or. x(5) > 1.0_dp) then |
| 412 | ✗ | status = SYMPLECTIC_STEP_OUTSIDE_DOMAIN | |
| 413 | ✗ | return | |
| 414 | end if | ||
| 415 |
2/2✓ Branch 0 taken 474690 times.
✓ Branch 1 taken 94938 times.
|
569628 | correction_abs = abs(x - xlast) |
| 416 | 94938 | tolref(4) = max(abs(x(4)), tolref(4)) | |
| 417 |
10/12✓ Branch 0 taken 94951 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 13 times.
✓ Branch 3 taken 94938 times.
✓ Branch 4 taken 284798 times.
✓ Branch 5 taken 41600 times.
✓ Branch 6 taken 231460 times.
✓ Branch 7 taken 53338 times.
✓ Branch 8 taken 94938 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 41600 times.
✓ Branch 11 taken 53338 times.
|
326411 | if (.not. limited .and. (all(residual_abs < si%atol) .or. & |
| 418 | ✗ | all(correction_abs < si%rtol*tolref))) then | |
| 419 | 41600 | status = SYMPLECTIC_STEP_OK | |
| 420 | 41600 | return | |
| 421 | end if | ||
| 422 | end do | ||
| 423 | |||
| 424 | if (warning_mode .and. gpu_finite_iterate(x) .and. & | ||
| 425 | ✗ | gpu_finite_iterate(xlast) .and. gpu_finite_iterate(tolref) .and. & | |
| 426 | all(abs(x - xlast) <= si%warning_factor*si%rtol*abs(tolref))) & | ||
| 427 | ✗ | status = SYMPLECTIC_STEP_OK | |
| 428 | end subroutine gpu_newton_midpoint | ||
| 429 | |||
| 430 | 41600 | subroutine gpu_timestep_midpoint(si, f, warning_mode, ierr, nfev, predictor) | |
| 431 | !$acc routine seq | ||
| 432 | type(symplectic_integrator_t), intent(inout) :: si | ||
| 433 | type(field_can_t), intent(inout) :: f | ||
| 434 | logical, intent(in) :: warning_mode | ||
| 435 | integer, intent(out) :: ierr, nfev | ||
| 436 | real(dp), intent(in), optional :: predictor(5) | ||
| 437 | |||
| 438 | type(symplectic_integrator_t) :: accepted_integrator | ||
| 439 | type(field_can_t) :: accepted_field | ||
| 440 | real(dp) :: x(5) | ||
| 441 | integer :: ktau, step_status, step_nfev | ||
| 442 | |||
| 443 | 41600 | ierr = 0 | |
| 444 | 41600 | nfev = 0 | |
| 445 |
2/2✓ Branch 0 taken 41600 times.
✓ Branch 1 taken 41600 times.
|
83200 | do ktau = 1, si%ntau |
| 446 | 41600 | accepted_integrator = si | |
| 447 | 41600 | accepted_field = f | |
| 448 | 41600 | si%pthold = f%pth | |
| 449 |
2/2✓ Branch 0 taken 41592 times.
✓ Branch 1 taken 8 times.
|
41600 | if (ktau == 1 .and. present(predictor)) then |
| 450 | 41592 | x = predictor | |
| 451 | else | ||
| 452 |
2/2✓ Branch 0 taken 32 times.
✓ Branch 1 taken 8 times.
|
40 | x(1:4) = si%z |
| 453 | 8 | x(5) = si%z(1) | |
| 454 | end if | ||
| 455 | 41600 | call gpu_newton_midpoint(si, f, x, warning_mode, step_status, step_nfev) | |
| 456 | 41600 | nfev = nfev + step_nfev | |
| 457 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
41600 | if (step_status /= SYMPLECTIC_STEP_OK .and. ktau == 1 .and. & |
| 458 | present(predictor)) then | ||
| 459 | ✗ | si = accepted_integrator | |
| 460 | ✗ | f = accepted_field | |
| 461 | ✗ | x(1:4) = si%z | |
| 462 | ✗ | x(5) = si%z(1) | |
| 463 | call gpu_newton_midpoint(si, f, x, warning_mode, step_status, & | ||
| 464 | ✗ | step_nfev) | |
| 465 | ✗ | nfev = nfev + step_nfev | |
| 466 | end if | ||
| 467 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
|
41600 | if (step_status /= SYMPLECTIC_STEP_OK) then |
| 468 | ✗ | si = accepted_integrator | |
| 469 | ✗ | f = accepted_field | |
| 470 | ✗ | ierr = step_status | |
| 471 | ✗ | return | |
| 472 | end if | ||
| 473 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
|
41600 | if (x(1) < 0.0_dp) then |
| 474 | ✗ | x(1) = -x(1) | |
| 475 | ✗ | x(2) = x(2) + pi | |
| 476 | end if | ||
| 477 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
|
41600 | if (x(1) > 1.0_dp) then |
| 478 | ✗ | si = accepted_integrator | |
| 479 | ✗ | f = accepted_field | |
| 480 | ✗ | ierr = SYMPLECTIC_STEP_OUTSIDE_DOMAIN | |
| 481 | ✗ | return | |
| 482 | end if | ||
| 483 |
2/2✓ Branch 0 taken 166400 times.
✓ Branch 1 taken 41600 times.
|
208000 | si%z = x(1:4) |
| 484 | 41600 | call eval_field_booz_device(f, si%z(1), si%z(2), si%z(3), 0) | |
| 485 | 41600 | call get_val(f, si%z(4)) | |
| 486 | 124800 | nfev = nfev + 1 | |
| 487 | end do | ||
| 488 | 41600 | end subroutine gpu_timestep_midpoint | |
| 489 | |||
| 490 | 16 | subroutine gpu_rhs_canonical(y, mu, ro0, dydt) | |
| 491 | ! Canonical guiding-centre equations used by SIMPLE's CPU RK45 path, | ||
| 492 | ! with the Boozer evaluator bound statically for device compilation. | ||
| 493 | !$acc routine seq | ||
| 494 | real(dp), intent(in) :: y(4) | ||
| 495 | real(dp), intent(in) :: mu, ro0 | ||
| 496 | real(dp), intent(out) :: dydt(4) | ||
| 497 | |||
| 498 | real(dp) :: aphi, daphi, btheta, dbtheta, bphi, dbphi, bmod | ||
| 499 | real(dp) :: dbmod(3), hth, hph, vpar, hprime | ||
| 500 | real(dp) :: inv_bmod, inv_bmod2, inv_hph, inv_dpth | ||
| 501 | real(dp) :: dhth_r, dhth_phi | ||
| 502 | real(dp) :: dhph_r, dhph_theta, dhph_phi | ||
| 503 | real(dp) :: dvpar_r, dvpar_theta, dvpar_phi | ||
| 504 | real(dp) :: dh_r, dh_theta, dh_phi, dpth_r, dpth_phi | ||
| 505 | |||
| 506 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
|
16 | if (rk_tables_ready) then |
| 507 | call splint_boozer_rk_table_device(y(1), y(2), y(3), aphi, & | ||
| 508 | ✗ | daphi, btheta, dbtheta, bphi, dbphi, bmod, dbmod) | |
| 509 | else | ||
| 510 | call splint_boozer_rk_device(y(1), y(2), y(3), aphi, daphi, & | ||
| 511 | 16 | btheta, dbtheta, bphi, dbphi, bmod, dbmod) | |
| 512 | end if | ||
| 513 | 16 | inv_bmod = 1.0_dp/bmod | |
| 514 | 16 | inv_bmod2 = inv_bmod*inv_bmod | |
| 515 | 16 | hth = btheta*inv_bmod | |
| 516 | 16 | hph = bphi*inv_bmod | |
| 517 | 16 | dhth_r = dbtheta*inv_bmod - btheta*dbmod(1)*inv_bmod2 | |
| 518 | 16 | dhth_phi = -btheta*dbmod(3)*inv_bmod2 | |
| 519 | 16 | dhph_r = dbphi*inv_bmod - bphi*dbmod(1)*inv_bmod2 | |
| 520 | 16 | dhph_theta = -bphi*dbmod(2)*inv_bmod2 | |
| 521 | 16 | dhph_phi = -bphi*dbmod(3)*inv_bmod2 | |
| 522 | 16 | inv_hph = 1.0_dp/hph | |
| 523 | 16 | vpar = (y(4) - aphi/ro0)*inv_hph | |
| 524 | 16 | dvpar_r = -(daphi/ro0 + dhph_r*vpar)*inv_hph | |
| 525 | 16 | dvpar_theta = -dhph_theta*vpar*inv_hph | |
| 526 | 16 | dvpar_phi = -dhph_phi*vpar*inv_hph | |
| 527 | 16 | dh_r = vpar*dvpar_r + mu*dbmod(1) | |
| 528 | 16 | dh_theta = vpar*dvpar_theta + mu*dbmod(2) | |
| 529 | 16 | dh_phi = vpar*dvpar_phi + mu*dbmod(3) | |
| 530 | 16 | dpth_r = dvpar_r*hth + vpar*dhth_r + boozer_state%torflux/ro0 | |
| 531 | 16 | dpth_phi = dvpar_phi*hth + vpar*dhth_phi | |
| 532 | 16 | inv_dpth = 1.0_dp/dpth_r | |
| 533 | 16 | hprime = dh_r*inv_dpth | |
| 534 | |||
| 535 | 16 | dydt(1) = -(dh_theta - hth*inv_hph*dh_phi)*inv_dpth | |
| 536 | 16 | dydt(2) = hprime | |
| 537 | 16 | dydt(3) = (vpar - hprime*hth)*inv_hph | |
| 538 | 16 | dydt(4) = -(dh_phi - hprime*dpth_phi) | |
| 539 | 16 | end subroutine gpu_rhs_canonical | |
| 540 | |||
| 541 | ✗ | subroutine gpu_rhs_canonical_cartesian(y, mu, ro0, dydt) | |
| 542 | ! Integrate the radial Boozer plane as (s*cos(theta), s*sin(theta)). | ||
| 543 | ! This removes the theta singularity at the magnetic axis while keeping | ||
| 544 | ! SIMPLE's native canonical momentum and canonical Boozer RHS. | ||
| 545 | !$acc routine seq | ||
| 546 | real(dp), intent(in) :: y(4), mu, ro0 | ||
| 547 | real(dp), intent(out) :: dydt(4) | ||
| 548 | |||
| 549 | real(dp) :: polar(4), polar_rhs(4), radial, radial_x, radial_y | ||
| 550 | |||
| 551 | ✗ | radial = sqrt(y(1)*y(1) + y(2)*y(2)) | |
| 552 | ✗ | polar = [radial, atan2(y(2), y(1)), y(3), y(4)] | |
| 553 | ✗ | call gpu_rhs_canonical(polar, mu, ro0, polar_rhs) | |
| 554 | ✗ | if (radial > tiny(1.0_dp)) then | |
| 555 | ✗ | radial_x = y(1)/radial | |
| 556 | ✗ | radial_y = y(2)/radial | |
| 557 | else | ||
| 558 | ✗ | radial_x = cos(polar(2)) | |
| 559 | ✗ | radial_y = sin(polar(2)) | |
| 560 | end if | ||
| 561 | ✗ | dydt(1) = polar_rhs(1)*radial_x - y(2)*polar_rhs(2) | |
| 562 | ✗ | dydt(2) = polar_rhs(1)*radial_y + y(1)*polar_rhs(2) | |
| 563 | ✗ | dydt(3:4) = polar_rhs(3:4) | |
| 564 | ✗ | end subroutine gpu_rhs_canonical_cartesian | |
| 565 | |||
| 566 | 2 | subroutine evaluate_rhs_gpu(y, mu, ro0, npart, dydt) | |
| 567 | ! Batch entry point used to compare the compact device RHS against the | ||
| 568 | ! host canonical-field implementation on independently evaluated data. | ||
| 569 | integer, intent(in) :: npart | ||
| 570 | real(dp), intent(in) :: y(4, npart), mu(npart), ro0(npart) | ||
| 571 | real(dp), intent(out) :: dydt(4, npart) | ||
| 572 | |||
| 573 | integer :: i | ||
| 574 | |||
| 575 | !$acc parallel loop gang vector copyin(y, mu, ro0) copyout(dydt) | ||
| 576 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 2 times.
|
18 | do i = 1, npart |
| 577 | 18 | call gpu_rhs_canonical(y(:, i), mu(i), ro0(i), dydt(:, i)) | |
| 578 | end do | ||
| 579 | 2 | end subroutine evaluate_rhs_gpu | |
| 580 | |||
| 581 | ✗ | subroutine gpu_trace_rk54(mu, ro0, y, duration, controls, loss_time, & | |
| 582 | nfev, ierr) | ||
| 583 | !$acc routine seq | ||
| 584 | real(dp), intent(in) :: mu, ro0 | ||
| 585 | real(dp), intent(inout) :: y(4) | ||
| 586 | real(dp), intent(in) :: duration | ||
| 587 | type(rk54_controls4_t), intent(in) :: controls | ||
| 588 | real(dp), intent(out) :: loss_time | ||
| 589 | integer, intent(out) :: nfev, ierr | ||
| 590 | |||
| 591 | integer, parameter :: max_attempts = 100000000 | ||
| 592 | type(rk54_state4_t) :: state | ||
| 593 | real(dp) :: t_eval, y_eval(4), dydt(4), h0 | ||
| 594 | integer :: request, attempt, rhs_calls | ||
| 595 | |||
| 596 | ✗ | loss_time = duration | |
| 597 | ✗ | nfev = 0 | |
| 598 | ✗ | ierr = 0 | |
| 599 | ✗ | if (duration <= 0.0_dp) return | |
| 600 | |||
| 601 | ✗ | h0 = min(duration, 1.0e-3_dp*controls%hmax) | |
| 602 | ✗ | call rk54_initialize4(state, 0.0_dp, y, h0) | |
| 603 | |||
| 604 | ✗ | do attempt = 1, max_attempts | |
| 605 | ✗ | if (state%t >= duration) exit | |
| 606 | ✗ | if (state%y(1) >= 1.0_dp) then | |
| 607 | ✗ | loss_time = state%t | |
| 608 | ✗ | exit | |
| 609 | end if | ||
| 610 | |||
| 611 | ✗ | state%h = min(state%h, duration - state%t) | |
| 612 | ✗ | call rk54_request4(state, controls, t_eval, y_eval, request) | |
| 613 | ✗ | rhs_calls = 0 | |
| 614 | ✗ | do while (request == RK54_NEED_RHS) | |
| 615 | ✗ | call gpu_rhs_canonical(y_eval, mu, ro0, dydt) | |
| 616 | ✗ | call rk54_supply4(state, controls, dydt, t_eval, y_eval, request) | |
| 617 | ✗ | rhs_calls = rhs_calls + 1 | |
| 618 | ✗ | if (rhs_calls > 7) then | |
| 619 | ✗ | request = RK54_FAILED | |
| 620 | ✗ | exit | |
| 621 | end if | ||
| 622 | end do | ||
| 623 | |||
| 624 | ✗ | if (request == RK54_FAILED) then | |
| 625 | ✗ | ierr = RK54_FAILED | |
| 626 | ✗ | exit | |
| 627 | end if | ||
| 628 | ✗ | if (request == RK54_ACCEPTED) then | |
| 629 | ✗ | if (state%y(1) < 0.0_dp) then | |
| 630 | ✗ | state%y(1) = -state%y(1) | |
| 631 | ✗ | state%y(2) = state%y(2) + pi | |
| 632 | end if | ||
| 633 | ✗ | if (state%y(1) >= 1.0_dp) then | |
| 634 | ✗ | loss_time = state%t | |
| 635 | ✗ | exit | |
| 636 | end if | ||
| 637 | ✗ | else if (request /= RK54_REJECTED) then | |
| 638 | ✗ | ierr = request | |
| 639 | ✗ | exit | |
| 640 | end if | ||
| 641 | end do | ||
| 642 | ✗ | if (attempt > max_attempts) ierr = RK54_FAILED | |
| 643 | |||
| 644 | ✗ | y = state%y | |
| 645 | ✗ | nfev = state%nfev | |
| 646 | end subroutine gpu_trace_rk54 | ||
| 647 | |||
| 648 | ✗ | subroutine gpu_trace_rk54_cartesian(mu, ro0, y, duration, controls, & | |
| 649 | loss_time, nfev, ierr) | ||
| 650 | !$acc routine seq | ||
| 651 | real(dp), intent(in) :: mu, ro0 | ||
| 652 | real(dp), intent(inout) :: y(4) | ||
| 653 | real(dp), intent(in) :: duration | ||
| 654 | type(rk54_controls4_t), intent(in) :: controls | ||
| 655 | real(dp), intent(out) :: loss_time | ||
| 656 | integer, intent(out) :: nfev, ierr | ||
| 657 | |||
| 658 | integer, parameter :: max_attempts = 100000000 | ||
| 659 | type(rk54_state4_t) :: state | ||
| 660 | real(dp) :: t_eval, y_eval(4), dydt(4), h0, radial | ||
| 661 | integer :: request, attempt, rhs_calls | ||
| 662 | |||
| 663 | ✗ | loss_time = duration | |
| 664 | ✗ | nfev = 0 | |
| 665 | ✗ | ierr = 0 | |
| 666 | ✗ | if (duration <= 0.0_dp) return | |
| 667 | |||
| 668 | ✗ | h0 = min(duration, 1.0e-3_dp*controls%hmax) | |
| 669 | ✗ | call rk54_initialize4(state, 0.0_dp, y, h0) | |
| 670 | |||
| 671 | ✗ | do attempt = 1, max_attempts | |
| 672 | ✗ | if (state%t >= duration) exit | |
| 673 | ✗ | radial = sqrt(state%y(1)*state%y(1) + state%y(2)*state%y(2)) | |
| 674 | ✗ | if (radial >= 1.0_dp) then | |
| 675 | ✗ | loss_time = state%t | |
| 676 | ✗ | exit | |
| 677 | end if | ||
| 678 | |||
| 679 | ✗ | state%h = min(state%h, duration - state%t) | |
| 680 | ✗ | call rk54_request4(state, controls, t_eval, y_eval, request) | |
| 681 | ✗ | rhs_calls = 0 | |
| 682 | ✗ | do while (request == RK54_NEED_RHS) | |
| 683 | ✗ | call gpu_rhs_canonical_cartesian(y_eval, mu, ro0, dydt) | |
| 684 | ✗ | call rk54_supply4(state, controls, dydt, t_eval, y_eval, request) | |
| 685 | ✗ | rhs_calls = rhs_calls + 1 | |
| 686 | ✗ | if (rhs_calls > 7) then | |
| 687 | ✗ | request = RK54_FAILED | |
| 688 | ✗ | exit | |
| 689 | end if | ||
| 690 | end do | ||
| 691 | |||
| 692 | ✗ | if (request == RK54_FAILED) then | |
| 693 | ✗ | ierr = RK54_FAILED | |
| 694 | ✗ | exit | |
| 695 | end if | ||
| 696 | ✗ | if (request == RK54_ACCEPTED) then | |
| 697 | ✗ | radial = sqrt(state%y(1)*state%y(1) + state%y(2)*state%y(2)) | |
| 698 | ✗ | if (radial >= 1.0_dp) then | |
| 699 | ✗ | loss_time = state%t | |
| 700 | ✗ | exit | |
| 701 | end if | ||
| 702 | ✗ | else if (request /= RK54_REJECTED) then | |
| 703 | ✗ | ierr = request | |
| 704 | ✗ | exit | |
| 705 | end if | ||
| 706 | end do | ||
| 707 | ✗ | if (attempt > max_attempts) ierr = RK54_FAILED | |
| 708 | |||
| 709 | ✗ | y = state%y | |
| 710 | ✗ | nfev = state%nfev | |
| 711 | end subroutine gpu_trace_rk54_cartesian | ||
| 712 | |||
| 713 | subroutine gpu_rk_segment_hmax(y, duration, hmax) | ||
| 714 | ! Match FIRM3D's per-segment quarter-turn cap, evaluated at the | ||
| 715 | ! segment's initial state: dt <= (G/B)*pi/(2*v_total). SIMPLE's | ||
| 716 | ! canonical independent variable is physical_time*v_total/sqrt(2). | ||
| 717 | !$acc routine seq | ||
| 718 | real(dp), intent(in) :: y(4), duration | ||
| 719 | real(dp), intent(out) :: hmax | ||
| 720 | |||
| 721 | real(dp) :: aphi, daphi, btheta, dbtheta, bphi, dbphi, bmod | ||
| 722 | real(dp) :: dbmod(3) | ||
| 723 | |||
| 724 | if (rk_tables_ready) then | ||
| 725 | call splint_boozer_rk_table_device(y(1), y(2), y(3), aphi, & | ||
| 726 | daphi, btheta, dbtheta, bphi, dbphi, bmod, dbmod) | ||
| 727 | else | ||
| 728 | call splint_boozer_rk_device(y(1), y(2), y(3), aphi, daphi, & | ||
| 729 | btheta, dbtheta, bphi, dbphi, bmod, dbmod) | ||
| 730 | end if | ||
| 731 | hmax = min(duration, 0.5_dp*pi*abs(bphi/bmod)/sqrt2) | ||
| 732 | end subroutine gpu_rk_segment_hmax | ||
| 733 | |||
| 734 | ✗ | subroutine gpu_rk_segment_hmax_cartesian(y, duration, hmax, & | |
| 735 | momentum_atol_scale) | ||
| 736 | !$acc routine seq | ||
| 737 | real(dp), intent(in) :: y(4), duration | ||
| 738 | real(dp), intent(out) :: hmax, momentum_atol_scale | ||
| 739 | |||
| 740 | real(dp) :: polar(4), aphi, daphi, btheta, dbtheta | ||
| 741 | real(dp) :: bphi, dbphi, bmod, dbmod(3) | ||
| 742 | |||
| 743 | polar = [sqrt(y(1)*y(1) + y(2)*y(2)), atan2(y(2), y(1)), & | ||
| 744 | ✗ | y(3), y(4)] | |
| 745 | ✗ | if (rk_tables_ready) then | |
| 746 | call splint_boozer_rk_table_device(polar(1), polar(2), polar(3), & | ||
| 747 | ✗ | aphi, daphi, btheta, dbtheta, bphi, dbphi, bmod, dbmod) | |
| 748 | else | ||
| 749 | call splint_boozer_rk_device(polar(1), polar(2), polar(3), aphi, & | ||
| 750 | ✗ | daphi, btheta, dbtheta, bphi, dbphi, bmod, dbmod) | |
| 751 | end if | ||
| 752 | ✗ | hmax = min(duration, 0.5_dp*pi*abs(bphi/bmod)/sqrt2) | |
| 753 | ! FIRM3D applies atol=rtol*v_total to v_parallel. In SIMPLE's | ||
| 754 | ! canonical momentum p_phi=v_parallel*h_phi+A_phi/ro0, the equivalent | ||
| 755 | ! absolute scale is sqrt(2)*|h_phi|. | ||
| 756 | ✗ | momentum_atol_scale = sqrt2*abs(bphi/bmod) | |
| 757 | ✗ | end subroutine gpu_rk_segment_hmax_cartesian | |
| 758 | |||
| 759 | ✗ | subroutine gpu_trace_euler_segment(si, f, nsteps, warning_mode, elapsed, & | |
| 760 | nfev, ierr) | ||
| 761 | !$acc routine seq | ||
| 762 | type(symplectic_integrator_t), intent(inout) :: si | ||
| 763 | type(field_can_t), intent(inout) :: f | ||
| 764 | integer, intent(in) :: nsteps | ||
| 765 | logical, intent(in) :: warning_mode | ||
| 766 | real(dp), intent(out) :: elapsed | ||
| 767 | integer, intent(out) :: nfev, ierr | ||
| 768 | |||
| 769 | integer :: k, step_nfev | ||
| 770 | real(dp) :: z_previous(4), predictor(2) | ||
| 771 | logical :: use_predictor | ||
| 772 | |||
| 773 | ✗ | elapsed = 0.0_dp | |
| 774 | ✗ | nfev = 0 | |
| 775 | ✗ | ierr = SYMPLECTIC_STEP_OK | |
| 776 | ✗ | use_predictor = .false. | |
| 777 | ✗ | do k = 1, nsteps | |
| 778 | ✗ | if (use_predictor) then | |
| 779 | ✗ | predictor(1) = 2.0_dp*si%z(1) - z_previous(1) | |
| 780 | ✗ | predictor(2) = 2.0_dp*si%z(4) - z_previous(4) | |
| 781 | use_predictor = predictor(1) >= 0.0_dp .and. & | ||
| 782 | predictor(1) <= 1.0_dp .and. & | ||
| 783 | ✗ | abs(si%z(2) - z_previous(2)) < 0.5_dp*pi | |
| 784 | end if | ||
| 785 | ✗ | z_previous = si%z | |
| 786 | ✗ | if (use_predictor) then | |
| 787 | call gpu_timestep_euler(si, f, warning_mode, ierr, step_nfev, & | ||
| 788 | ✗ | predictor) | |
| 789 | else | ||
| 790 | ✗ | call gpu_timestep_euler(si, f, warning_mode, ierr, step_nfev) | |
| 791 | end if | ||
| 792 | ✗ | nfev = nfev + step_nfev | |
| 793 | ✗ | elapsed = elapsed + si%dt | |
| 794 | ✗ | if (ierr /= SYMPLECTIC_STEP_OK) exit | |
| 795 | ✗ | use_predictor = .true. | |
| 796 | end do | ||
| 797 | ✗ | end subroutine gpu_trace_euler_segment | |
| 798 | |||
| 799 | ✗ | subroutine gpu_trace_midpoint_segment(si, f, nsteps, warning_mode, elapsed, & | |
| 800 | nfev, ierr) | ||
| 801 | !$acc routine seq | ||
| 802 | type(symplectic_integrator_t), intent(inout) :: si | ||
| 803 | type(field_can_t), intent(inout) :: f | ||
| 804 | integer, intent(in) :: nsteps | ||
| 805 | logical, intent(in) :: warning_mode | ||
| 806 | real(dp), intent(out) :: elapsed | ||
| 807 | integer, intent(out) :: nfev, ierr | ||
| 808 | |||
| 809 | integer :: k, step_nfev | ||
| 810 | real(dp) :: z_previous(4), predictor(5) | ||
| 811 | logical :: use_predictor | ||
| 812 | |||
| 813 | ✗ | elapsed = 0.0_dp | |
| 814 | ✗ | nfev = 0 | |
| 815 | ✗ | ierr = SYMPLECTIC_STEP_OK | |
| 816 | ✗ | use_predictor = .false. | |
| 817 | ✗ | do k = 1, nsteps | |
| 818 | ✗ | if (use_predictor) then | |
| 819 | ✗ | predictor(1:4) = 2.0_dp*si%z - z_previous | |
| 820 | ✗ | predictor(5) = 1.5_dp*si%z(1) - 0.5_dp*z_previous(1) | |
| 821 | use_predictor = predictor(1) >= 0.0_dp .and. & | ||
| 822 | predictor(1) <= 1.0_dp .and. predictor(5) >= 0.0_dp .and. & | ||
| 823 | predictor(5) <= 1.0_dp .and. & | ||
| 824 | ✗ | abs(si%z(2) - z_previous(2)) < 0.5_dp*pi | |
| 825 | end if | ||
| 826 | ✗ | z_previous = si%z | |
| 827 | ✗ | if (use_predictor) then | |
| 828 | call gpu_timestep_midpoint(si, f, warning_mode, ierr, step_nfev, & | ||
| 829 | ✗ | predictor) | |
| 830 | else | ||
| 831 | ✗ | call gpu_timestep_midpoint(si, f, warning_mode, ierr, step_nfev) | |
| 832 | end if | ||
| 833 | ✗ | nfev = nfev + step_nfev | |
| 834 | ✗ | elapsed = elapsed + si%dt | |
| 835 | ✗ | if (ierr /= SYMPLECTIC_STEP_OK) exit | |
| 836 | ✗ | use_predictor = .true. | |
| 837 | end do | ||
| 838 | ✗ | end subroutine gpu_trace_midpoint_segment | |
| 839 | |||
| 840 | ✗ | subroutine trace_orbits_gpu_rk54_landreman(si, f, npart, method, & | |
| 841 | total_duration, block_duration, time_scale, loss_tau, maxloss, & | ||
| 842 | ✗ | hmin, loss_step, loss_time, zend, nfev, observed_duration, & | |
| 843 | energy_loss_fraction, warp_nfev_slots) | ||
| 844 | type(symplectic_integrator_t), intent(inout) :: si(npart) | ||
| 845 | type(field_can_t), intent(in) :: f(npart) | ||
| 846 | integer, intent(in) :: npart, method | ||
| 847 | real(dp), intent(in) :: total_duration, block_duration, time_scale | ||
| 848 | real(dp), intent(in) :: loss_tau, maxloss, hmin | ||
| 849 | integer, intent(out) :: loss_step(npart), nfev(npart) | ||
| 850 | real(dp), intent(out) :: loss_time(npart), zend(4, npart) | ||
| 851 | real(dp), intent(out) :: observed_duration, energy_loss_fraction | ||
| 852 | integer(int64), intent(out) :: warp_nfev_slots | ||
| 853 | |||
| 854 | type(rk54_controls4_t) :: controls, controls_i | ||
| 855 | ✗ | real(dp), allocatable :: y(:, :), mu(:), ro0(:), loss_time_slot(:) | |
| 856 | ✗ | real(dp), allocatable :: y_next(:, :), mu_next(:), ro0_next(:) | |
| 857 | ✗ | real(dp), allocatable :: loss_time_next(:), work_key(:) | |
| 858 | ✗ | integer, allocatable :: orbit_status(:), orbit_status_next(:) | |
| 859 | ✗ | integer, allocatable :: nfev_slot(:), nfev_next(:), segment_work(:) | |
| 860 | ✗ | integer, allocatable :: original_index(:), original_index_next(:) | |
| 861 | ✗ | integer, allocatable :: order(:), order_scratch(:) | |
| 862 | real(dp) :: current_time, segment_duration, segment_loss_time | ||
| 863 | real(dp) :: y_local(4), hmax_local, momentum_atol_scale | ||
| 864 | real(dp) :: new_weight, weighted_losses | ||
| 865 | real(dp) :: loss_detection_tolerance, pilot_fraction, pilot_duration | ||
| 866 | integer :: i, ierr, segment_nfev, active_count, original | ||
| 867 | logical :: order_work | ||
| 868 | |||
| 869 | ✗ | if (block_duration <= 0.0_dp .or. time_scale <= 0.0_dp .or. & | |
| 870 | loss_tau <= 0.0_dp) & | ||
| 871 | ✗ | error stop 'invalid Landreman RK segment controls' | |
| 872 | |||
| 873 | allocate (y(4, npart), mu(npart), ro0(npart), & | ||
| 874 | loss_time_slot(npart), y_next(4, npart), mu_next(npart), & | ||
| 875 | ro0_next(npart), loss_time_next(npart), work_key(npart), & | ||
| 876 | orbit_status(npart), orbit_status_next(npart), nfev_slot(npart), & | ||
| 877 | nfev_next(npart), segment_work(npart), original_index(npart), & | ||
| 878 | ✗ | original_index_next(npart), order(npart), order_scratch(npart)) | |
| 879 | ✗ | orbit_status = 0 | |
| 880 | ✗ | loss_time_slot = total_duration | |
| 881 | weighted_losses = 0.0_dp | ||
| 882 | ✗ | do i = 1, npart | |
| 883 | ✗ | y(1, i) = si(i)%z(1)*cos(si(i)%z(2)) | |
| 884 | ✗ | y(2, i) = si(i)%z(1)*sin(si(i)%z(2)) | |
| 885 | ✗ | y(3:4, i) = si(i)%z(3:4) | |
| 886 | ✗ | mu(i) = f(i)%mu | |
| 887 | ✗ | ro0(i) = f(i)%ro0 | |
| 888 | ✗ | original_index(i) = i | |
| 889 | ✗ | if (sqrt(y(1, i)*y(1, i) + y(2, i)*y(2, i)) >= & | |
| 890 | ✗ | 1.0_dp - 64.0_dp*epsilon(1.0_dp)) then | |
| 891 | ✗ | orbit_status(i) = 1 | |
| 892 | ✗ | loss_time_slot(i) = 0.0_dp | |
| 893 | ✗ | weighted_losses = weighted_losses + 1.0_dp | |
| 894 | end if | ||
| 895 | end do | ||
| 896 | ✗ | nfev_slot = 0 | |
| 897 | ✗ | segment_work = 0 | |
| 898 | ✗ | warp_nfev_slots = 0_int64 | |
| 899 | ✗ | current_time = 0.0_dp | |
| 900 | ✗ | loss_detection_tolerance = 1.0e-12_dp/time_scale | |
| 901 | |||
| 902 | ✗ | order_work = gpu_work_order | |
| 903 | #ifndef _OPENACC | ||
| 904 | ✗ | order_work = .false. | |
| 905 | #endif | ||
| 906 | ✗ | pilot_fraction = gpu_pilot_fraction | |
| 907 | ✗ | if (pilot_fraction < 0.0_dp .or. pilot_fraction > 0.25_dp) & | |
| 908 | ✗ | error stop 'gpu_pilot_fraction must be in [0, 0.25]' | |
| 909 | if (.not. order_work .or. npart < 1024) pilot_fraction = 0.0_dp | ||
| 910 | if (order_work) then | ||
| 911 | do i = 1, npart | ||
| 912 | if (orbit_status(i) == 0) then | ||
| 913 | work_key(i) = abs(f(i)%vpar) | ||
| 914 | else | ||
| 915 | work_key(i) = -1.0_dp | ||
| 916 | end if | ||
| 917 | end do | ||
| 918 | call stable_order_descending(work_key, order, order_scratch) | ||
| 919 | call reorder_rk_slots(order, y, mu, ro0, orbit_status, & | ||
| 920 | loss_time_slot, nfev_slot, original_index, y_next, mu_next, & | ||
| 921 | ro0_next, orbit_status_next, loss_time_next, nfev_next, & | ||
| 922 | original_index_next) | ||
| 923 | end if | ||
| 924 | ✗ | active_count = count(orbit_status == 0) | |
| 925 | |||
| 926 | ✗ | controls%rtol = si(1)%rtol | |
| 927 | ✗ | controls%atol = si(1)%rtol | |
| 928 | ✗ | controls%hmin = hmin | |
| 929 | ✗ | if (method == CASH_KARP) then | |
| 930 | ✗ | controls%method = RK54_CASH_KARP | |
| 931 | else | ||
| 932 | ✗ | controls%method = RK54_DORMAND_PRINCE_TUNED | |
| 933 | end if | ||
| 934 | |||
| 935 | !$acc data copy(y, orbit_status, loss_time_slot, nfev_slot, segment_work) & | ||
| 936 | !$acc& copyin(mu, ro0) | ||
| 937 | if (pilot_fraction > 0.0_dp .and. active_count > 0) then | ||
| 938 | pilot_duration = pilot_fraction*min(block_duration, total_duration) | ||
| 939 | #ifdef _OPENACC | ||
| 940 | !$acc parallel loop gang vector default(present) & | ||
| 941 | !$acc& private(controls_i, y_local, segment_loss_time, & | ||
| 942 | !$acc& hmax_local, momentum_atol_scale, segment_nfev, ierr) & | ||
| 943 | !$acc& firstprivate(controls, pilot_duration) | ||
| 944 | #else | ||
| 945 | !$omp parallel do default(shared) schedule(static) & | ||
| 946 | !$omp& private(i, controls_i, y_local, segment_loss_time, & | ||
| 947 | !$omp& hmax_local, momentum_atol_scale, segment_nfev, ierr) | ||
| 948 | #endif | ||
| 949 | do i = 1, active_count | ||
| 950 | y_local = y(:, i) | ||
| 951 | controls_i = controls | ||
| 952 | call gpu_rk_segment_hmax_cartesian(y_local, pilot_duration, & | ||
| 953 | hmax_local, momentum_atol_scale) | ||
| 954 | controls_i%hmax = hmax_local | ||
| 955 | controls_i%atol(4) = controls_i%rtol*momentum_atol_scale | ||
| 956 | call gpu_trace_rk54_cartesian(mu(i), ro0(i), y_local, & | ||
| 957 | pilot_duration, controls_i, segment_loss_time, & | ||
| 958 | segment_nfev, ierr) | ||
| 959 | segment_work(i) = segment_nfev + 1 | ||
| 960 | end do | ||
| 961 | #ifndef _OPENACC | ||
| 962 | !$omp end parallel do | ||
| 963 | #endif | ||
| 964 | !$acc update self(segment_work) | ||
| 965 | nfev_slot(1:active_count) = segment_work(1:active_count) | ||
| 966 | do i = 1, active_count, 32 | ||
| 967 | warp_nfev_slots = warp_nfev_slots + 32_int64*int( & | ||
| 968 | maxval(segment_work(i:min(i + 31, active_count))), int64) | ||
| 969 | end do | ||
| 970 | do i = 1, npart | ||
| 971 | if (orbit_status(i) == 0) then | ||
| 972 | work_key(i) = real(segment_work(i), dp) | ||
| 973 | else | ||
| 974 | work_key(i) = -1.0_dp | ||
| 975 | end if | ||
| 976 | end do | ||
| 977 | call stable_order_descending(work_key, order, order_scratch) | ||
| 978 | call reorder_rk_slots(order, y, mu, ro0, orbit_status, & | ||
| 979 | loss_time_slot, nfev_slot, original_index, y_next, mu_next, & | ||
| 980 | ro0_next, orbit_status_next, loss_time_next, nfev_next, & | ||
| 981 | original_index_next) | ||
| 982 | !$acc update device(y, mu, ro0, orbit_status, loss_time_slot, & | ||
| 983 | !$acc& nfev_slot) | ||
| 984 | end if | ||
| 985 | ✗ | do while (current_time < total_duration) | |
| 986 | ✗ | segment_duration = min(block_duration, total_duration - current_time) | |
| 987 | ✗ | if (weighted_losses/real(npart, dp) > maxloss) then | |
| 988 | ✗ | current_time = current_time + segment_duration | |
| 989 | ✗ | exit | |
| 990 | end if | ||
| 991 | ✗ | new_weight = 0.0_dp | |
| 992 | #ifdef _OPENACC | ||
| 993 | !$acc parallel loop gang vector default(present) & | ||
| 994 | !$acc& private(controls_i, y_local, segment_loss_time, & | ||
| 995 | !$acc& hmax_local, momentum_atol_scale, segment_nfev, ierr) & | ||
| 996 | !$acc& firstprivate(controls, current_time, segment_duration, & | ||
| 997 | !$acc& time_scale, loss_tau, loss_detection_tolerance) & | ||
| 998 | !$acc& copy(new_weight) | ||
| 999 | #else | ||
| 1000 | !$omp parallel do default(shared) schedule(static) & | ||
| 1001 | !$omp& private(i, controls_i, y_local, segment_loss_time, & | ||
| 1002 | !$omp& hmax_local, momentum_atol_scale, segment_nfev, ierr) & | ||
| 1003 | ✗ | !$omp& reduction(+:new_weight) | |
| 1004 | #endif | ||
| 1005 | do i = 1, active_count | ||
| 1006 | segment_work(i) = 0 | ||
| 1007 | if (orbit_status(i) /= 0) cycle | ||
| 1008 | y_local = y(:, i) | ||
| 1009 | if (sqrt(y_local(1)*y_local(1) + y_local(2)*y_local(2)) >= & | ||
| 1010 | 1.0_dp - 64.0_dp*epsilon(1.0_dp)) then | ||
| 1011 | orbit_status(i) = 1 | ||
| 1012 | loss_time_slot(i) = current_time | ||
| 1013 | !$acc atomic update | ||
| 1014 | new_weight = new_weight + & | ||
| 1015 | exp(-loss_time_slot(i)*time_scale/loss_tau) | ||
| 1016 | cycle | ||
| 1017 | end if | ||
| 1018 | controls_i = controls | ||
| 1019 | call gpu_rk_segment_hmax_cartesian(y_local, segment_duration, & | ||
| 1020 | hmax_local, momentum_atol_scale) | ||
| 1021 | controls_i%hmax = hmax_local | ||
| 1022 | controls_i%atol(4) = controls_i%rtol*momentum_atol_scale | ||
| 1023 | call gpu_trace_rk54_cartesian(mu(i), ro0(i), y_local, & | ||
| 1024 | segment_duration, controls_i, segment_loss_time, & | ||
| 1025 | segment_nfev, ierr) | ||
| 1026 | y(:, i) = y_local | ||
| 1027 | segment_work(i) = segment_nfev + 1 | ||
| 1028 | nfev_slot(i) = nfev_slot(i) + segment_work(i) | ||
| 1029 | if (ierr /= 0) then | ||
| 1030 | orbit_status(i) = -ierr | ||
| 1031 | else if (segment_loss_time < segment_duration - & | ||
| 1032 | loss_detection_tolerance) then | ||
| 1033 | orbit_status(i) = 1 | ||
| 1034 | loss_time_slot(i) = current_time + segment_loss_time | ||
| 1035 | !$acc atomic update | ||
| 1036 | new_weight = new_weight + & | ||
| 1037 | exp(-loss_time_slot(i)*time_scale/loss_tau) | ||
| 1038 | end if | ||
| 1039 | end do | ||
| 1040 | #ifndef _OPENACC | ||
| 1041 | !$omp end parallel do | ||
| 1042 | #endif | ||
| 1043 | !$acc update self(segment_work) | ||
| 1044 | ✗ | do i = 1, active_count, 32 | |
| 1045 | warp_nfev_slots = warp_nfev_slots + 32_int64*int( & | ||
| 1046 | ✗ | maxval(segment_work(i:min(i + 31, active_count))), int64) | |
| 1047 | end do | ||
| 1048 | ✗ | weighted_losses = weighted_losses + new_weight | |
| 1049 | ✗ | current_time = current_time + segment_duration | |
| 1050 | ✗ | if (current_time < total_duration .and. & | |
| 1051 | weighted_losses/real(npart, dp) > maxloss) exit | ||
| 1052 | ✗ | if (current_time < total_duration .and. order_work) then | |
| 1053 | !$acc update self(y, orbit_status, loss_time_slot, nfev_slot) | ||
| 1054 | do i = 1, npart | ||
| 1055 | if (orbit_status(i) == 0) then | ||
| 1056 | work_key(i) = real(segment_work(i), dp) | ||
| 1057 | else | ||
| 1058 | work_key(i) = -1.0_dp | ||
| 1059 | end if | ||
| 1060 | end do | ||
| 1061 | call stable_order_descending(work_key, order, order_scratch) | ||
| 1062 | call reorder_rk_slots(order, y, mu, ro0, orbit_status, & | ||
| 1063 | loss_time_slot, nfev_slot, original_index, y_next, mu_next, & | ||
| 1064 | ro0_next, orbit_status_next, loss_time_next, nfev_next, & | ||
| 1065 | original_index_next) | ||
| 1066 | active_count = count(orbit_status == 0) | ||
| 1067 | !$acc update device(y, mu, ro0, orbit_status, loss_time_slot, & | ||
| 1068 | !$acc& nfev_slot) | ||
| 1069 | end if | ||
| 1070 | end do | ||
| 1071 | !$acc end data | ||
| 1072 | |||
| 1073 | ✗ | do i = 1, npart | |
| 1074 | ✗ | original = original_index(i) | |
| 1075 | ✗ | si(original)%z(1) = sqrt(y(1, i)*y(1, i) + y(2, i)*y(2, i)) | |
| 1076 | ✗ | si(original)%z(2) = atan2(y(2, i), y(1, i)) | |
| 1077 | ✗ | si(original)%z(3:4) = y(3:4, i) | |
| 1078 | ✗ | zend(:, original) = si(original)%z | |
| 1079 | ✗ | loss_step(original) = orbit_status(i) | |
| 1080 | ✗ | loss_time(original) = loss_time_slot(i) | |
| 1081 | ✗ | nfev(original) = nfev_slot(i) | |
| 1082 | end do | ||
| 1083 | ✗ | observed_duration = current_time | |
| 1084 | ✗ | energy_loss_fraction = weighted_losses/real(npart, dp) | |
| 1085 | ✗ | end subroutine trace_orbits_gpu_rk54_landreman | |
| 1086 | |||
| 1087 | subroutine stable_order_descending(key, order, scratch) | ||
| 1088 | real(dp), intent(in) :: key(:) | ||
| 1089 | integer, intent(out) :: order(size(key)) | ||
| 1090 | integer, intent(inout) :: scratch(size(key)) | ||
| 1091 | |||
| 1092 | integer :: n, width, left, middle, right, i, j, k | ||
| 1093 | |||
| 1094 | n = size(key) | ||
| 1095 | do i = 1, n | ||
| 1096 | order(i) = i | ||
| 1097 | end do | ||
| 1098 | width = 1 | ||
| 1099 | do while (width < n) | ||
| 1100 | do left = 1, n, 2*width | ||
| 1101 | middle = min(left + width, n + 1) | ||
| 1102 | right = min(left + 2*width - 1, n) | ||
| 1103 | i = left | ||
| 1104 | j = middle | ||
| 1105 | do k = left, right | ||
| 1106 | if (i >= middle) then | ||
| 1107 | scratch(k) = order(j) | ||
| 1108 | j = j + 1 | ||
| 1109 | else if (j > right) then | ||
| 1110 | scratch(k) = order(i) | ||
| 1111 | i = i + 1 | ||
| 1112 | else if (key(order(i)) >= key(order(j))) then | ||
| 1113 | scratch(k) = order(i) | ||
| 1114 | i = i + 1 | ||
| 1115 | else | ||
| 1116 | scratch(k) = order(j) | ||
| 1117 | j = j + 1 | ||
| 1118 | end if | ||
| 1119 | end do | ||
| 1120 | end do | ||
| 1121 | order = scratch | ||
| 1122 | width = 2*width | ||
| 1123 | end do | ||
| 1124 | end subroutine stable_order_descending | ||
| 1125 | |||
| 1126 | subroutine reorder_rk_slots(order, y, mu, ro0, orbit_status, loss_time, & | ||
| 1127 | nfev, original_index, y_next, mu_next, ro0_next, & | ||
| 1128 | orbit_status_next, loss_time_next, nfev_next, original_index_next) | ||
| 1129 | integer, intent(in) :: order(:) | ||
| 1130 | real(dp), intent(inout) :: y(:, :), mu(:), ro0(:), loss_time(:) | ||
| 1131 | integer, intent(inout) :: orbit_status(:), nfev(:), original_index(:) | ||
| 1132 | real(dp), intent(inout) :: y_next(:, :), mu_next(:), ro0_next(:) | ||
| 1133 | real(dp), intent(inout) :: loss_time_next(:) | ||
| 1134 | integer, intent(inout) :: orbit_status_next(:), nfev_next(:) | ||
| 1135 | integer, intent(inout) :: original_index_next(:) | ||
| 1136 | |||
| 1137 | integer :: i, source | ||
| 1138 | |||
| 1139 | do i = 1, size(order) | ||
| 1140 | source = order(i) | ||
| 1141 | y_next(:, i) = y(:, source) | ||
| 1142 | mu_next(i) = mu(source) | ||
| 1143 | ro0_next(i) = ro0(source) | ||
| 1144 | orbit_status_next(i) = orbit_status(source) | ||
| 1145 | loss_time_next(i) = loss_time(source) | ||
| 1146 | nfev_next(i) = nfev(source) | ||
| 1147 | original_index_next(i) = original_index(source) | ||
| 1148 | end do | ||
| 1149 | y = y_next | ||
| 1150 | mu = mu_next | ||
| 1151 | ro0 = ro0_next | ||
| 1152 | orbit_status = orbit_status_next | ||
| 1153 | loss_time = loss_time_next | ||
| 1154 | nfev = nfev_next | ||
| 1155 | original_index = original_index_next | ||
| 1156 | end subroutine reorder_rk_slots | ||
| 1157 | |||
| 1158 | ✗ | subroutine trace_orbits_gpu_symplectic_landreman(si, f, npart, ntimstep, & | |
| 1159 | ✗ | ntau_macro, method, block_duration, time_scale, loss_tau, maxloss, & | |
| 1160 | ✗ | loss_step, loss_time, zend, nfev, observed_duration, & | |
| 1161 | energy_loss_fraction, warp_nfev_slots) | ||
| 1162 | type(symplectic_integrator_t), intent(inout) :: si(npart) | ||
| 1163 | type(field_can_t), intent(inout) :: f(npart) | ||
| 1164 | integer, intent(in) :: npart, ntimstep, method | ||
| 1165 | integer, intent(in) :: ntau_macro(ntimstep) | ||
| 1166 | real(dp), intent(in) :: block_duration, time_scale, loss_tau, maxloss | ||
| 1167 | integer, intent(out) :: loss_step(npart), nfev(npart) | ||
| 1168 | real(dp), intent(out) :: loss_time(npart), zend(4, npart) | ||
| 1169 | real(dp), intent(out) :: observed_duration, energy_loss_fraction | ||
| 1170 | integer(int64), intent(out) :: warp_nfev_slots | ||
| 1171 | |||
| 1172 | ✗ | type(symplectic_integrator_t), allocatable :: si_next(:) | |
| 1173 | ✗ | type(field_can_t), allocatable :: f_next(:) | |
| 1174 | ✗ | integer, allocatable :: orbit_status(:), orbit_status_next(:) | |
| 1175 | ✗ | integer, allocatable :: nfev_next(:), original_index(:) | |
| 1176 | ✗ | integer, allocatable :: order(:), order_scratch(:) | |
| 1177 | ✗ | real(dp), allocatable :: loss_time_next(:), work_key(:) | |
| 1178 | real(dp) :: total_duration, current_time, segment_duration | ||
| 1179 | real(dp) :: elapsed, new_weight, weighted_losses, expected_duration | ||
| 1180 | real(dp) :: schedule_tolerance | ||
| 1181 | integer :: i, it, ierr, segment_nfev, original, source | ||
| 1182 | logical :: warning_mode, order_work | ||
| 1183 | |||
| 1184 | ✗ | if (block_duration <= 0.0_dp .or. time_scale <= 0.0_dp .or. & | |
| 1185 | loss_tau <= 0.0_dp) & | ||
| 1186 | ✗ | error stop 'invalid Landreman symplectic segment controls' | |
| 1187 | ✗ | if (method /= EXPL_IMPL_EULER .and. method /= MIDPOINT) & | |
| 1188 | ✗ | error stop 'invalid Landreman symplectic method' | |
| 1189 | |||
| 1190 | ✗ | total_duration = si(1)%dt*real(sum(ntau_macro(2:ntimstep)), dp) | |
| 1191 | ✗ | schedule_tolerance = 64.0_dp*epsilon(1.0_dp)*max(1.0_dp, total_duration) | |
| 1192 | ✗ | current_time = 0.0_dp | |
| 1193 | ✗ | do it = 2, ntimstep | |
| 1194 | ✗ | segment_duration = si(1)%dt*real(ntau_macro(it), dp) | |
| 1195 | ✗ | expected_duration = min(block_duration, total_duration - current_time) | |
| 1196 | ✗ | if (abs(segment_duration - expected_duration) > schedule_tolerance) & | |
| 1197 | ✗ | error stop 'SIMPLE macrostep schedule does not match Landreman block size' | |
| 1198 | ✗ | current_time = current_time + segment_duration | |
| 1199 | end do | ||
| 1200 | |||
| 1201 | allocate (si_next(npart), f_next(npart), orbit_status(npart), & | ||
| 1202 | orbit_status_next(npart), nfev_next(npart), & | ||
| 1203 | original_index(npart), order(npart), order_scratch(npart), & | ||
| 1204 | ✗ | loss_time_next(npart), work_key(npart)) | |
| 1205 | ✗ | orbit_status = 0 | |
| 1206 | ✗ | loss_time = total_duration | |
| 1207 | ✗ | nfev = 0 | |
| 1208 | ✗ | warp_nfev_slots = 0_int64 | |
| 1209 | ✗ | weighted_losses = 0.0_dp | |
| 1210 | ✗ | do i = 1, npart | |
| 1211 | ✗ | original_index(i) = i | |
| 1212 | ✗ | if (si(i)%z(1) >= 1.0_dp - 64.0_dp*epsilon(1.0_dp)) then | |
| 1213 | ✗ | orbit_status(i) = 1 | |
| 1214 | ✗ | loss_time(i) = 0.0_dp | |
| 1215 | ✗ | weighted_losses = weighted_losses + 1.0_dp | |
| 1216 | end if | ||
| 1217 | end do | ||
| 1218 | |||
| 1219 | ✗ | order_work = gpu_work_order | |
| 1220 | #ifndef _OPENACC | ||
| 1221 | ✗ | order_work = .false. | |
| 1222 | #endif | ||
| 1223 | if (order_work) then | ||
| 1224 | do i = 1, npart | ||
| 1225 | if (orbit_status(i) == 0) then | ||
| 1226 | work_key(i) = abs(f(i)%vpar) | ||
| 1227 | else | ||
| 1228 | work_key(i) = -1.0_dp | ||
| 1229 | end if | ||
| 1230 | end do | ||
| 1231 | call stable_order_descending(work_key, order, order_scratch) | ||
| 1232 | do i = 1, npart | ||
| 1233 | source = order(i) | ||
| 1234 | si_next(i) = si(source) | ||
| 1235 | f_next(i) = f(source) | ||
| 1236 | orbit_status_next(i) = orbit_status(source) | ||
| 1237 | loss_time_next(i) = loss_time(source) | ||
| 1238 | original_index(i) = source | ||
| 1239 | end do | ||
| 1240 | si = si_next | ||
| 1241 | f = f_next | ||
| 1242 | orbit_status = orbit_status_next | ||
| 1243 | loss_time = loss_time_next | ||
| 1244 | end if | ||
| 1245 | ✗ | current_time = 0.0_dp | |
| 1246 | ✗ | warning_mode = symplectic_newton_warning_mode | |
| 1247 | |||
| 1248 | !$acc data copy(si, f, orbit_status, loss_time, nfev) | ||
| 1249 | ✗ | do it = 2, ntimstep | |
| 1250 | ✗ | segment_duration = si(1)%dt*real(ntau_macro(it), dp) | |
| 1251 | ✗ | if (weighted_losses/real(npart, dp) > maxloss) then | |
| 1252 | ✗ | current_time = current_time + segment_duration | |
| 1253 | ✗ | exit | |
| 1254 | end if | ||
| 1255 | ✗ | new_weight = 0.0_dp | |
| 1256 | ✗ | if (method == EXPL_IMPL_EULER) then | |
| 1257 | #ifdef _OPENACC | ||
| 1258 | !$acc parallel loop gang vector default(present) & | ||
| 1259 | !$acc& private(elapsed, segment_nfev, ierr) & | ||
| 1260 | !$acc& firstprivate(warning_mode, current_time, time_scale, loss_tau) & | ||
| 1261 | !$acc& copy(new_weight) | ||
| 1262 | #else | ||
| 1263 | !$omp parallel do default(shared) schedule(static) & | ||
| 1264 | !$omp& private(i, elapsed, segment_nfev, ierr) & | ||
| 1265 | ✗ | !$omp& reduction(+:new_weight) | |
| 1266 | #endif | ||
| 1267 | do i = 1, npart | ||
| 1268 | if (orbit_status(i) /= 0) cycle | ||
| 1269 | call gpu_trace_euler_segment(si(i), f(i), ntau_macro(it), & | ||
| 1270 | warning_mode, elapsed, segment_nfev, ierr) | ||
| 1271 | nfev(i) = nfev(i) + segment_nfev | ||
| 1272 | if (ierr == SYMPLECTIC_STEP_OUTSIDE_DOMAIN) then | ||
| 1273 | orbit_status(i) = 1 | ||
| 1274 | loss_time(i) = current_time + elapsed | ||
| 1275 | !$acc atomic update | ||
| 1276 | new_weight = new_weight + & | ||
| 1277 | exp(-loss_time(i)*time_scale/loss_tau) | ||
| 1278 | else if (ierr /= SYMPLECTIC_STEP_OK) then | ||
| 1279 | orbit_status(i) = -ierr | ||
| 1280 | end if | ||
| 1281 | end do | ||
| 1282 | #ifndef _OPENACC | ||
| 1283 | !$omp end parallel do | ||
| 1284 | #endif | ||
| 1285 | else | ||
| 1286 | #ifdef _OPENACC | ||
| 1287 | !$acc parallel loop gang vector default(present) & | ||
| 1288 | !$acc& private(elapsed, segment_nfev, ierr) & | ||
| 1289 | !$acc& firstprivate(warning_mode, current_time, time_scale, loss_tau) & | ||
| 1290 | !$acc& copy(new_weight) | ||
| 1291 | #else | ||
| 1292 | !$omp parallel do default(shared) schedule(static) & | ||
| 1293 | !$omp& private(i, elapsed, segment_nfev, ierr) & | ||
| 1294 | ✗ | !$omp& reduction(+:new_weight) | |
| 1295 | #endif | ||
| 1296 | do i = 1, npart | ||
| 1297 | if (orbit_status(i) /= 0) cycle | ||
| 1298 | call gpu_trace_midpoint_segment(si(i), f(i), ntau_macro(it), & | ||
| 1299 | warning_mode, elapsed, segment_nfev, ierr) | ||
| 1300 | nfev(i) = nfev(i) + segment_nfev | ||
| 1301 | if (ierr == SYMPLECTIC_STEP_OUTSIDE_DOMAIN) then | ||
| 1302 | orbit_status(i) = 1 | ||
| 1303 | loss_time(i) = current_time + elapsed | ||
| 1304 | !$acc atomic update | ||
| 1305 | new_weight = new_weight + & | ||
| 1306 | exp(-loss_time(i)*time_scale/loss_tau) | ||
| 1307 | else if (ierr /= SYMPLECTIC_STEP_OK) then | ||
| 1308 | orbit_status(i) = -ierr | ||
| 1309 | end if | ||
| 1310 | end do | ||
| 1311 | #ifndef _OPENACC | ||
| 1312 | !$omp end parallel do | ||
| 1313 | #endif | ||
| 1314 | end if | ||
| 1315 | ✗ | weighted_losses = weighted_losses + new_weight | |
| 1316 | ✗ | current_time = current_time + segment_duration | |
| 1317 | ✗ | if (current_time < total_duration .and. & | |
| 1318 | ✗ | weighted_losses/real(npart, dp) > maxloss) exit | |
| 1319 | end do | ||
| 1320 | !$acc end data | ||
| 1321 | |||
| 1322 | ✗ | do i = 1, npart, 32 | |
| 1323 | warp_nfev_slots = warp_nfev_slots + 32_int64*int( & | ||
| 1324 | ✗ | maxval(nfev(i:min(i + 31, npart))), int64) | |
| 1325 | end do | ||
| 1326 | |||
| 1327 | ✗ | do i = 1, npart | |
| 1328 | ✗ | original = original_index(i) | |
| 1329 | ✗ | si_next(original) = si(i) | |
| 1330 | ✗ | f_next(original) = f(i) | |
| 1331 | ✗ | zend(:, original) = si(i)%z | |
| 1332 | ✗ | loss_step(original) = orbit_status(i) | |
| 1333 | ✗ | loss_time_next(original) = loss_time(i) | |
| 1334 | ✗ | nfev_next(original) = nfev(i) | |
| 1335 | end do | ||
| 1336 | ✗ | si = si_next | |
| 1337 | ✗ | f = f_next | |
| 1338 | ✗ | loss_time = loss_time_next | |
| 1339 | ✗ | nfev = nfev_next | |
| 1340 | ✗ | observed_duration = current_time | |
| 1341 | ✗ | energy_loss_fraction = weighted_losses/real(npart, dp) | |
| 1342 | ✗ | end subroutine trace_orbits_gpu_symplectic_landreman | |
| 1343 | |||
| 1344 | ✗ | subroutine trace_orbits_gpu_landreman(si, f, npart, ntimstep, ntau_macro, & | |
| 1345 | method, block_duration, time_scale, loss_tau, maxloss, hmin, & | ||
| 1346 | ✗ | loss_step, loss_time, zend, nfev, observed_duration, & | |
| 1347 | energy_loss_fraction, warp_nfev_slots) | ||
| 1348 | type(symplectic_integrator_t), intent(inout) :: si(npart) | ||
| 1349 | type(field_can_t), intent(inout) :: f(npart) | ||
| 1350 | integer, intent(in) :: npart, ntimstep, method | ||
| 1351 | integer, intent(in) :: ntau_macro(ntimstep) | ||
| 1352 | real(dp), intent(in) :: block_duration, time_scale, loss_tau, maxloss, hmin | ||
| 1353 | integer, intent(out) :: loss_step(npart), nfev(npart) | ||
| 1354 | real(dp), intent(out) :: loss_time(npart), zend(4, npart) | ||
| 1355 | real(dp), intent(out) :: observed_duration, energy_loss_fraction | ||
| 1356 | integer(int64), intent(out) :: warp_nfev_slots | ||
| 1357 | |||
| 1358 | real(dp) :: total_duration | ||
| 1359 | |||
| 1360 | ✗ | total_duration = si(1)%dt*real(sum(ntau_macro(2:ntimstep)), dp) | |
| 1361 | ✗ | select case (method) | |
| 1362 | case (CASH_KARP, DORMAND_PRINCE) | ||
| 1363 | call trace_orbits_gpu_rk54_landreman(si, f, npart, method, & | ||
| 1364 | total_duration, block_duration, time_scale, loss_tau, maxloss, & | ||
| 1365 | hmin, loss_step, loss_time, zend, nfev, observed_duration, & | ||
| 1366 | ✗ | energy_loss_fraction, warp_nfev_slots) | |
| 1367 | case (EXPL_IMPL_EULER, MIDPOINT) | ||
| 1368 | call trace_orbits_gpu_symplectic_landreman(si, f, npart, ntimstep, & | ||
| 1369 | ntau_macro, method, block_duration, time_scale, loss_tau, maxloss, & | ||
| 1370 | loss_step, loss_time, zend, nfev, observed_duration, & | ||
| 1371 | ✗ | energy_loss_fraction, warp_nfev_slots) | |
| 1372 | case default | ||
| 1373 | ✗ | error stop 'trace_orbits_gpu_landreman: unsupported integrator' | |
| 1374 | end select | ||
| 1375 | ✗ | end subroutine trace_orbits_gpu_landreman | |
| 1376 | |||
| 1377 | ✗ | subroutine trace_orbits_gpu_rk54_range(si, f, npart, istart, iend, & | |
| 1378 | ✗ | ntimstep, ntau_macro, method, loss_step, loss_time_out, zend, nfev) | |
| 1379 | type(symplectic_integrator_t), intent(in) :: si(npart) | ||
| 1380 | type(field_can_t), intent(in) :: f(npart) | ||
| 1381 | integer, intent(in) :: npart, istart, iend, ntimstep, method | ||
| 1382 | integer, intent(in) :: ntau_macro(ntimstep) | ||
| 1383 | integer, intent(out) :: loss_step(npart), nfev(npart) | ||
| 1384 | real(dp), intent(out) :: loss_time_out(npart) | ||
| 1385 | real(dp), intent(out) :: zend(4, npart) | ||
| 1386 | |||
| 1387 | type(rk54_controls4_t) :: controls | ||
| 1388 | real(dp) :: duration, loss_time, y(4), elapsed | ||
| 1389 | ✗ | real(dp) :: mu(npart), ro0(npart) | |
| 1390 | integer :: i, it, ierr | ||
| 1391 | |||
| 1392 | ✗ | duration = si(istart)%dt*real(sum(ntau_macro(2:ntimstep)), dp) | |
| 1393 | ✗ | controls%rtol = si(istart)%rtol | |
| 1394 | ✗ | controls%atol = si(istart)%rtol | |
| 1395 | ✗ | controls%hmin = 0.0_dp | |
| 1396 | ✗ | controls%hmax = duration | |
| 1397 | ✗ | if (method == CASH_KARP) then | |
| 1398 | ✗ | controls%method = RK54_CASH_KARP | |
| 1399 | else | ||
| 1400 | ✗ | controls%method = RK54_DORMAND_PRINCE_TUNED | |
| 1401 | end if | ||
| 1402 | |||
| 1403 | ✗ | do i = istart, iend | |
| 1404 | ✗ | mu(i) = f(i)%mu | |
| 1405 | ✗ | ro0(i) = f(i)%ro0 | |
| 1406 | end do | ||
| 1407 | |||
| 1408 | !$acc parallel loop gang vector default(present) & | ||
| 1409 | !$acc& copyin(si(istart:iend), mu(istart:iend), ro0(istart:iend), & | ||
| 1410 | !$acc& ntau_macro) & | ||
| 1411 | !$acc& copyout(loss_step(istart:iend), zend(:, istart:iend), & | ||
| 1412 | !$acc& nfev(istart:iend), loss_time_out(istart:iend)) & | ||
| 1413 | !$acc& private(y, loss_time, elapsed, ierr, it) firstprivate(controls, duration) | ||
| 1414 | ✗ | do i = istart, iend | |
| 1415 | ✗ | y = si(i)%z | |
| 1416 | call gpu_trace_rk54(mu(i), ro0(i), y, duration, controls, loss_time, & | ||
| 1417 | ✗ | nfev(i), ierr) | |
| 1418 | ✗ | loss_step(i) = ntimstep | |
| 1419 | ✗ | if (loss_time < duration) then | |
| 1420 | elapsed = 0.0_dp | ||
| 1421 | ✗ | do it = 2, ntimstep | |
| 1422 | ✗ | elapsed = elapsed + si(i)%dt*real(ntau_macro(it), dp) | |
| 1423 | ✗ | if (loss_time <= elapsed) then | |
| 1424 | ✗ | loss_step(i) = it | |
| 1425 | ✗ | exit | |
| 1426 | end if | ||
| 1427 | end do | ||
| 1428 | end if | ||
| 1429 | ✗ | if (ierr /= 0) loss_step(i) = -ierr | |
| 1430 | ✗ | loss_time_out(i) = loss_time | |
| 1431 | ✗ | zend(:, i) = y | |
| 1432 | end do | ||
| 1433 | ✗ | end subroutine trace_orbits_gpu_rk54_range | |
| 1434 | |||
| 1435 | 1 | subroutine trace_orbits_gpu_midpoint_range(si, f, npart, istart, iend, & | |
| 1436 | 1 | ntimstep, ntau_macro, loss_step, loss_time_out, zend, nfev) | |
| 1437 | type(symplectic_integrator_t), intent(inout) :: si(npart) | ||
| 1438 | type(field_can_t), intent(inout) :: f(npart) | ||
| 1439 | integer, intent(in) :: npart, istart, iend, ntimstep | ||
| 1440 | integer, intent(in) :: ntau_macro(ntimstep) | ||
| 1441 | integer, intent(out) :: loss_step(npart), nfev(npart) | ||
| 1442 | real(dp), intent(out) :: loss_time_out(npart) | ||
| 1443 | real(dp), intent(out) :: zend(4, npart) | ||
| 1444 | |||
| 1445 | integer :: i, it, ktau, ierr, lstep, step_nfev | ||
| 1446 | real(dp) :: elapsed, z_previous(4), predictor(5) | ||
| 1447 | logical :: warning_mode, use_predictor | ||
| 1448 | |||
| 1449 | 1 | warning_mode = symplectic_newton_warning_mode | |
| 1450 | !$acc parallel loop gang vector default(present) & | ||
| 1451 | !$acc& copy(si(istart:iend), f(istart:iend)) copyin(ntau_macro) & | ||
| 1452 | !$acc& copyout(loss_step(istart:iend), zend(:, istart:iend), & | ||
| 1453 | !$acc& nfev(istart:iend), loss_time_out(istart:iend)) & | ||
| 1454 | !$acc& private(it, ktau, ierr, lstep, step_nfev, elapsed, & | ||
| 1455 | !$acc& z_previous, predictor, use_predictor) & | ||
| 1456 | !$acc& firstprivate(warning_mode) | ||
| 1457 |
2/2✓ Branch 0 taken 8 times.
✓ Branch 1 taken 1 time.
|
9 | do i = istart, iend |
| 1458 | 8 | ierr = 0 | |
| 1459 | 8 | lstep = ntimstep | |
| 1460 | 8 | nfev(i) = 0 | |
| 1461 | 8 | elapsed = 0.0_dp | |
| 1462 | 8 | use_predictor = .false. | |
| 1463 |
2/2✓ Branch 0 taken 400 times.
✓ Branch 1 taken 8 times.
|
408 | macro: do it = 2, ntimstep |
| 1464 |
2/2✓ Branch 0 taken 41600 times.
✓ Branch 1 taken 400 times.
|
42000 | do ktau = 1, ntau_macro(it) |
| 1465 |
2/2✓ Branch 0 taken 41592 times.
✓ Branch 1 taken 8 times.
|
41600 | if (use_predictor) then |
| 1466 |
2/2✓ Branch 0 taken 166368 times.
✓ Branch 1 taken 41592 times.
|
207960 | predictor(1:4) = 2.0_dp*si(i)%z - z_previous |
| 1467 | predictor(5) = 1.5_dp*si(i)%z(1) - & | ||
| 1468 | 41592 | 0.5_dp*z_previous(1) | |
| 1469 | use_predictor = predictor(1) >= 0.0_dp .and. & | ||
| 1470 | predictor(1) <= 1.0_dp .and. & | ||
| 1471 | predictor(5) >= 0.0_dp .and. & | ||
| 1472 | predictor(5) <= 1.0_dp .and. & | ||
| 1473 |
5/10✗ Branch 0 not taken.
✓ Branch 1 taken 41592 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 41592 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 41592 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 41592 times.
✗ Branch 8 not taken.
✓ Branch 9 taken 41592 times.
|
41592 | abs(si(i)%z(2) - z_previous(2)) < 0.5_dp*pi |
| 1474 | end if | ||
| 1475 |
2/2✓ Branch 0 taken 166400 times.
✓ Branch 1 taken 41600 times.
|
208000 | z_previous = si(i)%z |
| 1476 |
2/2✓ Branch 0 taken 41592 times.
✓ Branch 1 taken 8 times.
|
41600 | if (use_predictor) then |
| 1477 | call gpu_timestep_midpoint(si(i), f(i), warning_mode, & | ||
| 1478 | 41592 | ierr, step_nfev, predictor) | |
| 1479 | else | ||
| 1480 | call gpu_timestep_midpoint(si(i), f(i), warning_mode, & | ||
| 1481 | 8 | ierr, step_nfev) | |
| 1482 | end if | ||
| 1483 | 41600 | nfev(i) = nfev(i) + step_nfev | |
| 1484 | 41600 | elapsed = elapsed + si(i)%dt | |
| 1485 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
|
41600 | if (ierr /= 0) exit |
| 1486 | 42000 | use_predictor = .true. | |
| 1487 | end do | ||
| 1488 |
1/2✓ Branch 0 taken 400 times.
✗ Branch 1 not taken.
|
408 | if (ierr /= 0) then |
| 1489 | lstep = it | ||
| 1490 | exit macro | ||
| 1491 | end if | ||
| 1492 | end do macro | ||
| 1493 | 8 | loss_step(i) = lstep | |
| 1494 |
2/4✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
|
8 | if (ierr == 0 .and. lstep == ntimstep) & |
| 1495 |
2/2✓ Branch 0 taken 400 times.
✓ Branch 1 taken 8 times.
|
408 | elapsed = real(sum(ntau_macro(2:ntimstep)), dp)*si(i)%dt |
| 1496 | 8 | loss_time_out(i) = elapsed | |
| 1497 |
2/2✓ Branch 0 taken 32 times.
✓ Branch 1 taken 8 times.
|
41 | zend(:, i) = si(i)%z |
| 1498 | end do | ||
| 1499 | 1 | end subroutine trace_orbits_gpu_midpoint_range | |
| 1500 | |||
| 1501 | 1 | subroutine trace_orbits_gpu_range(si, f, npart, istart, iend, ntimstep, & | |
| 1502 | 1 | ntau_macro, loss_step, loss_time, zend, nfev) | |
| 1503 | ! Evolve particles istart:iend on the currently selected OpenACC device. | ||
| 1504 | ! loss_step(i) = first timestep at which the orbit left the plasma (r>1), | ||
| 1505 | ! or ntimstep if confined for the whole trace. zend(:,i) is the final | ||
| 1506 | ! integrator state z = (r, th, ph, pphi). | ||
| 1507 | type(symplectic_integrator_t), intent(inout) :: si(npart) | ||
| 1508 | type(field_can_t), intent(inout) :: f(npart) | ||
| 1509 | integer, intent(in) :: npart, istart, iend, ntimstep | ||
| 1510 | integer, intent(in) :: ntau_macro(ntimstep) | ||
| 1511 | integer, intent(out) :: loss_step(npart) | ||
| 1512 | real(dp), intent(out) :: loss_time(npart) | ||
| 1513 | real(dp), intent(out) :: zend(4, npart) | ||
| 1514 | integer, intent(out) :: nfev(npart) | ||
| 1515 | |||
| 1516 | integer :: i, it, ktau, ierr, lstep, step_nfev | ||
| 1517 | real(dp) :: elapsed, z_previous(4), predictor(2) | ||
| 1518 | logical :: warning_mode, use_predictor | ||
| 1519 | |||
| 1520 | 1 | warning_mode = symplectic_newton_warning_mode | |
| 1521 | |||
| 1522 | !$acc parallel loop gang vector default(present) & | ||
| 1523 | !$acc& copy(si(istart:iend), f(istart:iend)) copyin(ntau_macro) & | ||
| 1524 | !$acc& copyout(loss_step(istart:iend), loss_time(istart:iend), & | ||
| 1525 | !$acc& zend(:, istart:iend), nfev(istart:iend)) & | ||
| 1526 | !$acc& private(it, ktau, ierr, lstep, elapsed, step_nfev, & | ||
| 1527 | !$acc& z_previous, predictor, use_predictor) & | ||
| 1528 | !$acc& firstprivate(warning_mode) | ||
| 1529 |
2/2✓ Branch 0 taken 8 times.
✓ Branch 1 taken 1 time.
|
9 | do i = istart, iend |
| 1530 | 8 | ierr = 0 | |
| 1531 | 8 | lstep = ntimstep | |
| 1532 | 8 | elapsed = 0.0_dp | |
| 1533 | 8 | nfev(i) = 0 | |
| 1534 | 8 | use_predictor = .false. | |
| 1535 |
2/2✓ Branch 0 taken 400 times.
✓ Branch 1 taken 8 times.
|
408 | macro: do it = 2, ntimstep |
| 1536 |
2/2✓ Branch 0 taken 41600 times.
✓ Branch 1 taken 400 times.
|
42000 | do ktau = 1, ntau_macro(it) |
| 1537 |
2/2✓ Branch 0 taken 8 times.
✓ Branch 1 taken 41592 times.
|
41600 | if (use_predictor) then |
| 1538 | 41592 | predictor(1) = 2.0_dp*si(i)%z(1) - z_previous(1) | |
| 1539 | 41592 | predictor(2) = 2.0_dp*si(i)%z(4) - z_previous(4) | |
| 1540 | use_predictor = predictor(1) >= 0.0_dp .and. & | ||
| 1541 | predictor(1) <= 1.0_dp .and. & | ||
| 1542 |
3/6✗ Branch 0 not taken.
✓ Branch 1 taken 41592 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 41592 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 41592 times.
|
41592 | abs(si(i)%z(2) - z_previous(2)) < 0.5_dp*pi |
| 1543 | end if | ||
| 1544 |
2/2✓ Branch 0 taken 166400 times.
✓ Branch 1 taken 41600 times.
|
208000 | z_previous = si(i)%z |
| 1545 |
2/2✓ Branch 0 taken 41592 times.
✓ Branch 1 taken 8 times.
|
41600 | if (use_predictor) then |
| 1546 | call gpu_timestep_euler(si(i), f(i), warning_mode, ierr, & | ||
| 1547 | 41592 | step_nfev, predictor) | |
| 1548 | else | ||
| 1549 | call gpu_timestep_euler(si(i), f(i), warning_mode, ierr, & | ||
| 1550 | 8 | step_nfev) | |
| 1551 | end if | ||
| 1552 | 41600 | nfev(i) = nfev(i) + step_nfev | |
| 1553 | 41600 | elapsed = elapsed + si(i)%dt | |
| 1554 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 41600 times.
|
41600 | if (ierr /= 0) exit |
| 1555 | 42000 | use_predictor = .true. | |
| 1556 | end do | ||
| 1557 |
1/2✓ Branch 0 taken 400 times.
✗ Branch 1 not taken.
|
408 | if (ierr /= 0) then |
| 1558 | lstep = it | ||
| 1559 | exit macro | ||
| 1560 | end if | ||
| 1561 | end do macro | ||
| 1562 | 8 | loss_step(i) = lstep | |
| 1563 |
2/4✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
|
8 | if (ierr == 0 .and. lstep == ntimstep) & |
| 1564 |
2/2✓ Branch 0 taken 400 times.
✓ Branch 1 taken 8 times.
|
408 | elapsed = real(sum(ntau_macro(2:ntimstep)), dp)*si(i)%dt |
| 1565 | 8 | loss_time(i) = elapsed | |
| 1566 | 8 | zend(1, i) = si(i)%z(1) | |
| 1567 | 8 | zend(2, i) = si(i)%z(2) | |
| 1568 | 8 | zend(3, i) = si(i)%z(3) | |
| 1569 | 9 | zend(4, i) = si(i)%z(4) | |
| 1570 | end do | ||
| 1571 | 1 | end subroutine trace_orbits_gpu_range | |
| 1572 | |||
| 1573 | ✗ | subroutine trace_orbits_gpu(si, f, npart, ntimstep, ntau_macro, loss_step, zend) | |
| 1574 | ! Evolve npart pre-initialised particles, splitting the work evenly | ||
| 1575 | ! across all available NVIDIA GPUs (one host thread per device). Falls | ||
| 1576 | ! back to a single device / the host when only one is present. | ||
| 1577 | type(symplectic_integrator_t), intent(inout) :: si(npart) | ||
| 1578 | type(field_can_t), intent(inout) :: f(npart) | ||
| 1579 | integer, intent(in) :: npart, ntimstep | ||
| 1580 | integer, intent(in) :: ntau_macro(ntimstep) | ||
| 1581 | integer, intent(out) :: loss_step(npart) | ||
| 1582 | real(dp), intent(out) :: zend(4, npart) | ||
| 1583 | |||
| 1584 | integer :: ngpu, navail, dev, i0, i1 | ||
| 1585 | ✗ | real(dp), allocatable :: loss_time(:) | |
| 1586 | ✗ | integer, allocatable :: nfev(:) | |
| 1587 | integer :: req | ||
| 1588 | |||
| 1589 | ! Number of devices to use. Default 1: with -gpu=mem:unified the spline | ||
| 1590 | ! coefficient array is shared and migrates between devices on access, so | ||
| 1591 | ! splitting across cards thrashes. Real multi-GPU needs a per-device | ||
| 1592 | ! resident copy of the read-only splines. Opt in with gpu_num_devices. | ||
| 1593 | ✗ | req = gpu_num_devices | |
| 1594 | |||
| 1595 | ✗ | ngpu = 1 | |
| 1596 | ✗ | allocate (loss_time(npart), nfev(npart)) | |
| 1597 | #ifdef _OPENACC | ||
| 1598 | navail = acc_get_num_devices(acc_device_nvidia) | ||
| 1599 | if (navail < 1) navail = 1 | ||
| 1600 | ngpu = min(req, navail) | ||
| 1601 | #ifndef _OPENMP | ||
| 1602 | ! Without OpenMP the device split below runs as one thread with dev = 0, | ||
| 1603 | ! which would trace only the first npart/ngpu particles. | ||
| 1604 | ngpu = 1 | ||
| 1605 | #endif | ||
| 1606 | #endif | ||
| 1607 | if (ngpu <= 1) then | ||
| 1608 | call trace_orbits_gpu_range(si, f, npart, 1, npart, ntimstep, & | ||
| 1609 | ✗ | ntau_macro, loss_step, loss_time, zend, nfev) | |
| 1610 | ✗ | return | |
| 1611 | end if | ||
| 1612 | |||
| 1613 | !$omp parallel num_threads(ngpu) private(dev, i0, i1) | ||
| 1614 | dev = 0 | ||
| 1615 | !$ dev = omp_get_thread_num() | ||
| 1616 | #ifdef _OPENACC | ||
| 1617 | call acc_set_device_num(dev, acc_device_nvidia) | ||
| 1618 | #endif | ||
| 1619 | i0 = int(int(dev, 8)*npart/ngpu) + 1 | ||
| 1620 | i1 = int(int(dev + 1, 8)*npart/ngpu) | ||
| 1621 | if (i1 >= i0) then | ||
| 1622 | call trace_orbits_gpu_range(si, f, npart, i0, i1, ntimstep, & | ||
| 1623 | ntau_macro, loss_step, loss_time, zend, nfev) | ||
| 1624 | end if | ||
| 1625 | !$omp end parallel | ||
| 1626 | ✗ | end subroutine trace_orbits_gpu | |
| 1627 | |||
| 1628 | 2 | subroutine trace_orbits_gpu_method(si, f, npart, ntimstep, ntau_macro, & | |
| 1629 | 2 | method, loss_step, loss_time, zend, nfev) | |
| 1630 | ! Compile-time-specialized kernels sit behind this host-only dispatch; | ||
| 1631 | ! no method branch is executed in the per-stage device hot path. | ||
| 1632 | type(symplectic_integrator_t), intent(inout) :: si(npart) | ||
| 1633 | type(field_can_t), intent(inout) :: f(npart) | ||
| 1634 | integer, intent(in) :: npart, ntimstep, method | ||
| 1635 | integer, intent(in) :: ntau_macro(ntimstep) | ||
| 1636 | integer, intent(out) :: loss_step(npart), nfev(npart) | ||
| 1637 | real(dp), intent(out) :: loss_time(npart) | ||
| 1638 | real(dp), intent(out) :: zend(4, npart) | ||
| 1639 | |||
| 1640 | 3 | select case (method) | |
| 1641 | case (EXPL_IMPL_EULER) | ||
| 1642 | call trace_orbits_gpu_range(si, f, npart, 1, npart, ntimstep, & | ||
| 1643 | 1 | ntau_macro, loss_step, loss_time, zend, nfev) | |
| 1644 | case (MIDPOINT) | ||
| 1645 | call trace_orbits_gpu_midpoint_range(si, f, npart, 1, npart, & | ||
| 1646 | 1 | ntimstep, ntau_macro, loss_step, loss_time, zend, nfev) | |
| 1647 | case (CASH_KARP, DORMAND_PRINCE) | ||
| 1648 | call trace_orbits_gpu_rk54_range(si, f, npart, 1, npart, & | ||
| 1649 | ✗ | ntimstep, ntau_macro, method, loss_step, loss_time, zend, nfev) | |
| 1650 | case default | ||
| 1651 |
2/4✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
2 | error stop 'trace_orbits_gpu_method: unsupported GPU integrator' |
| 1652 | end select | ||
| 1653 | 2 | end subroutine trace_orbits_gpu_method | |
| 1654 | |||
| 1655 | end module simple_gpu | ||
| 1656 |