source: codes/icosagcm/trunk/src/dynamics/caldyn_kernels_base.F90 @ 954

Last change on this file since 954 was 954, checked in by adurocher, 5 years ago

trunk : Added metric terms to kernels parameters to avoid Host/GPU transferts

Metric terms are now subroutine parameters instead of module variables in kernel subroutines. Dummy arguments for metric terms are now defined as fixed-size arrays, and arrays dimensions are well known when entering an 'acc data' region. Array descriptors are no longer transferred form host to device each time the data region is executed.
File size: 15.3 KB
Line 
1MODULE caldyn_kernels_base_mod
2  USE icosa
3  USE transfert_mod
4  USE disvert_mod
5  USE omp_para
6  USE trace
7  USE abort_mod
8  IMPLICIT NONE
9  PRIVATE
10
11  INTEGER, PARAMETER,PUBLIC :: energy=1, enstrophy=2
12  TYPE(t_field),POINTER,PUBLIC :: f_out_u(:), f_qu(:), f_qv(:)
13
14  ! temporary shared variables for caldyn
15  TYPE(t_field),POINTER,PUBLIC :: f_pk(:),f_wwuu(:),f_planetvel(:), &
16                                  f_Fel(:), f_gradPhi2(:), f_wil(:), f_Wetadot(:)
17
18  INTEGER, PUBLIC :: caldyn_conserv
19  !$OMP THREADPRIVATE(caldyn_conserv)
20
21  TYPE(t_message),PUBLIC :: req_ps, req_mass, req_theta_rhodz, req_u, req_qu, req_geopot, req_w
22
23  PUBLIC :: compute_geopot, compute_caldyn_vert, compute_caldyn_vert_nh
24
25CONTAINS
26
27  !**************************** Geopotential *****************************
28  SUBROUTINE compute_geopot(rhodz,theta, ps,pk,geopot) 
29    REAL(rstd),INTENT(IN)    :: rhodz(iim*jjm,llm)
30    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm,nqdyn) ! active scalars : theta/entropy, moisture, ...
31    REAL(rstd),INTENT(INOUT) :: ps(iim*jjm)
32    REAL(rstd),INTENT(OUT)   :: pk(iim*jjm,llm)       ! Exner function (compressible) /Lagrange multiplier (Boussinesq)
33    REAL(rstd),INTENT(INOUT) :: geopot(iim*jjm,llm+1) ! geopotential
34
35    INTEGER :: ij,l
36    REAL(rstd) :: Rd, p_ik, exner_ik, temp_ik, qv, chi, Rmix, gv
37    INTEGER    :: ij_omp_begin_ext, ij_omp_end_ext
38
39    CALL trace_start("compute_geopot")
40
41!$OMP BARRIER
42
43    CALL distrib_level(ij_begin_ext,ij_end_ext, ij_omp_begin_ext,ij_omp_end_ext)
44
45    Rd = kappa*cpp
46
47    IF(dysl_geopot) THEN
48      CALL abort_acc("HEVI_scheme/!dysl_geopot")
49#include "../kernels/compute_geopot.k90"
50    ELSE
51    ! Pressure is computed first top-down (temporarily stored in pk)
52    ! Then Exner pressure and geopotential are computed bottom-up
53    ! Works also when caldyn_eta=eta_mass         
54
55    !$acc data present(rhodz(:,:), ps(:), pk(:,:), geopot(:,:), theta(:,:,:)) async
56
57    IF(boussinesq) THEN ! compute geopotential and pk=Lagrange multiplier
58       CALL abort_acc("HEVI_scheme/boussinesq")
59       ! specific volume 1 = dphi/g/rhodz
60       !         IF (is_omp_level_master) THEN ! no openMP on vertical due to dependency
61       DO l = 1,llm
62          !$acc parallel loop
63          !DIR$ SIMD
64          DO ij=ij_omp_begin_ext,ij_omp_end_ext         
65             geopot(ij,l+1) = geopot(ij,l) + g*rhodz(ij,l)
66          ENDDO
67          !$acc end parallel loop
68       ENDDO
69       ! use hydrostatic balance with theta*rhodz to find pk (Lagrange multiplier=pressure)
70       ! uppermost layer
71       !$acc parallel loop
72       !DIR$ SIMD
73       DO ij=ij_begin_ext,ij_end_ext         
74          pk(ij,llm) = ptop + (.5*g)*theta(ij,llm,1)*rhodz(ij,llm)
75       END DO
76       !$acc end parallel loop
77       ! other layers
78       DO l = llm-1, 1, -1
79          !          !$OMP DO SCHEDULE(STATIC)
80          !$acc parallel loop
81          !DIR$ SIMD
82          DO ij=ij_begin_ext,ij_end_ext         
83             pk(ij,l) = pk(ij,l+1) + (.5*g)*(theta(ij,l,1)*rhodz(ij,l)+theta(ij,l+1,1)*rhodz(ij,l+1))
84          END DO
85          !$acc end parallel loop
86       END DO
87       ! now pk contains the Lagrange multiplier (pressure)
88    ELSE ! non-Boussinesq, compute pressure, Exner pressure or temperature, then geopotential
89       ! uppermost layer
90       
91       SELECT CASE(caldyn_thermo)
92          CASE(thermo_theta, thermo_entropy)
93             !$acc parallel loop async
94             !DIR$ SIMD
95             DO ij=ij_omp_begin_ext,ij_omp_end_ext
96                pk(ij,llm) = ptop + (.5*g)*rhodz(ij,llm)
97             END DO
98
99             ! other layers
100             ! We use kernels here instead of "loop seq" + "loop gang vector"
101             ! to be sure the code really abides the standard. In practice,
102             ! it seems like the compiler interchanges the loops.
103             !$acc kernels async
104             DO l = llm-1, 1, -1
105                !DIR$ SIMD
106                DO ij=ij_omp_begin_ext,ij_omp_end_ext         
107                   pk(ij,l) = pk(ij,l+1) + (.5*g)*(rhodz(ij,l)+rhodz(ij,l+1))
108                END DO
109             END DO
110             !$acc end kernels
111             ! surface pressure (for diagnostics)
112             IF(caldyn_eta==eta_lag) THEN
113                !$acc parallel loop async
114                DO ij=ij_omp_begin_ext,ij_omp_end_ext         
115                   ps(ij) = pk(ij,1) + (.5*g)*rhodz(ij,1)
116                END DO
117             END IF
118          CASE(thermo_moist) ! theta(ij,l,2) = qv = mv/md
119             CALL abort_acc("HEVI_scheme/thermo_moist")
120             !$acc parallel loop
121             !DIR$ SIMD
122             DO ij=ij_omp_begin_ext,ij_omp_end_ext
123                pk(ij,llm) = ptop + (.5*g)*rhodz(ij,llm)*(1.+theta(ij,l,2))
124             END DO
125             !$acc end parallel loop
126
127             ! other layers
128             DO l = llm-1, 1, -1
129                !$acc parallel loop
130                !DIR$ SIMD
131                DO ij=ij_omp_begin_ext,ij_omp_end_ext         
132                   pk(ij,l) = pk(ij,l+1) + (.5*g)*(          &
133                        rhodz(ij,l)  *(1.+theta(ij,l,2)) +   &
134                        rhodz(ij,l+1)*(1.+theta(ij,l+1,2)) )
135                END DO
136                !$acc end parallel loop
137             END DO
138             ! surface pressure (for diagnostics)
139             IF(caldyn_eta==eta_lag) THEN
140                !$acc parallel loop
141                DO ij=ij_omp_begin_ext,ij_omp_end_ext         
142                   ps(ij) = pk(ij,1) + (.5*g)*rhodz(ij,1)*(1.+theta(ij,l,2))
143                END DO
144                !$acc end parallel loop
145             END IF
146          END SELECT
147
148        SELECT CASE(caldyn_thermo)
149          CASE(thermo_theta)
150            ! We use kernels here instead of "loop seq" + "loop gang vector"
151            ! to be sure the code really abides the standard. In practice,
152            ! it seems like the compiler interchanges the loops.
153            !$acc kernels async
154            DO l = 1,llm
155               !DIR$ SIMD
156               DO ij=ij_omp_begin_ext,ij_omp_end_ext
157                  p_ik = pk(ij,l)
158                  exner_ik = cpp * (p_ik/preff) ** kappa
159                  pk(ij,l) = exner_ik
160                  ! specific volume v = kappa*theta*pi/p = dphi/g/rhodz
161                  geopot(ij,l+1) = geopot(ij,l) + (g*kappa)*rhodz(ij,l)*theta(ij,l,1)*exner_ik/p_ik
162               ENDDO
163            ENDDO
164            !$acc end kernels
165
166          CASE(thermo_entropy) ! theta is in fact entropy = cpp*log(theta/Treff) = cpp*log(T/Treff) - Rd*log(p/preff)
167             CALL abort_acc("HEVI_scheme/thermo_entropy")
168             DO l = 1,llm
169               !$acc parallel loop
170               !DIR$ SIMD
171               DO ij=ij_omp_begin_ext,ij_omp_end_ext
172                  p_ik = pk(ij,l)
173                  temp_ik = Treff*exp((theta(ij,l,1) + Rd*log(p_ik/preff))/cpp)
174                  pk(ij,l) = temp_ik
175                  ! specific volume v = Rd*T/p = dphi/g/rhodz
176                  geopot(ij,l+1) = geopot(ij,l) + (g*Rd)*rhodz(ij,l)*temp_ik/p_ik
177               ENDDO
178               !$acc end parallel loop
179             ENDDO
180          CASE(thermo_moist) ! theta is moist pseudo-entropy per dry air mass
181             CALL abort_acc("HEVI_scheme/thermo_moist")
182             DO l = 1,llm
183               !$acc parallel loop
184               DO ij=ij_omp_begin_ext,ij_omp_end_ext
185                  p_ik = pk(ij,l)
186                  qv = theta(ij,l,2) ! water vaper mixing ratio = mv/md
187                  Rmix = Rd+qv*Rv
188                  chi = ( theta(ij,l,1) + Rmix*log(p_ik/preff) ) / (cpp + qv*cppv) ! log(T/Treff)
189                  temp_ik = Treff*exp(chi)
190                  pk(ij,l) = temp_ik
191                  ! specific volume v = R*T/p = dphi/g/rhodz
192                  ! R = (Rd + qv.Rv)/(1+qv)
193                  geopot(ij,l+1) = geopot(ij,l) + g*Rmix*rhodz(ij,l)*temp_ik/(p_ik*(1+qv))
194               ENDDO
195               !$acc end parallel loop
196             ENDDO
197          CASE DEFAULT
198             STOP
199        END SELECT
200    END IF
201    !$acc end data
202    END IF ! dysl
203
204    !ym flush geopot
205    !$OMP BARRIER
206
207    CALL trace_end("compute_geopot")
208
209  END SUBROUTINE compute_geopot
210
211  SUBROUTINE compute_caldyn_vert(u, theta, rhodz, convm, wflux, wwuu, dps, dtheta_rhodz, du, bp)
212    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
213    REAL(rstd),INTENT(IN)  :: theta(iim*jjm,llm,nqdyn)
214    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)
215    REAL(rstd),INTENT(INOUT)  :: convm(iim*jjm,llm)  ! mass flux convergence
216
217    REAL(rstd),INTENT(INOUT) :: wflux(iim*jjm,llm+1) ! vertical mass flux (kg/m2/s)
218    REAL(rstd),INTENT(INOUT) :: wwuu(iim*3*jjm,llm+1)
219    REAL(rstd),INTENT(INOUT) :: du(iim*3*jjm,llm)
220    REAL(rstd),INTENT(INOUT) :: dtheta_rhodz(iim*jjm,llm,nqdyn)
221    REAL(rstd),INTENT(OUT) :: dps(iim*jjm)
222    REAL(rstd),INTENT(IN)  :: bp(llm)
223
224    ! temporary variable   
225    INTEGER ::ij,l,iq
226    INTEGER    :: ij_omp_begin, ij_omp_end
227
228    CALL trace_start("compute_caldyn_vert")
229
230    CALL distrib_level(ij_begin,ij_end, ij_omp_begin,ij_omp_end)
231
232    !    REAL(rstd) :: wwuu(iim*3*jjm,llm+1) ! tmp var, don't know why but gain 30% on the whole code in opemp
233    ! need to be understood
234
235    !    wwuu=wwuu_out
236    CALL trace_start("compute_caldyn_vert")
237
238    !$acc data async &
239    !$acc present(rhodz(:,:), u(:,:), wwuu(:,:), wflux(:,:), dps(:), convm(:,:), du(:,:), dtheta_rhodz(:,:,:), theta(:,:,:), bp(:))
240
241
242    !$OMP BARRIER   
243!!! cumulate mass flux convergence from top to bottom
244    !  IF (is_omp_level_master) THEN
245    ! We use kernels here instead of "loop seq" + "loop gang vector"
246    ! to be sure the code really abides the standard. In practice,
247    ! it seems like the compiler interchanges the loops.
248    !$acc kernels async
249    DO  l = llm-1, 1, -1
250       !    IF (caldyn_conserv==energy) CALL test_message(req_qu)
251
252!!$OMP DO SCHEDULE(STATIC)
253       !DIR$ SIMD
254       DO ij=ij_omp_begin,ij_omp_end         
255          convm(ij,l) = convm(ij,l) + convm(ij,l+1)
256       ENDDO
257    ENDDO
258    !$acc end kernels
259    !  ENDIF
260
261    !$OMP BARRIER
262    ! FLUSH on convm
263!!!!!!!!!!!!!!!!!!!!!!!!!
264
265    ! compute dps
266    IF (is_omp_first_level) THEN
267       !$acc parallel loop  async
268       !DIR$ SIMD
269       DO ij=ij_begin,ij_end         
270          ! dps/dt = -int(div flux)dz
271          dps(ij) = convm(ij,1)
272       ENDDO
273    ENDIF
274
275!!! Compute vertical mass flux (l=1,llm+1 done by caldyn_BC)
276    !$acc parallel loop collapse(2) async
277    DO l=ll_beginp1,ll_end
278       !    IF (caldyn_conserv==energy) CALL test_message(req_qu)
279       !DIR$ SIMD
280       DO ij=ij_begin,ij_end         
281          ! w = int(z,ztop,div(flux)dz) + B(eta)dps/dt
282          ! => w>0 for upward transport
283          wflux( ij, l ) = bp(l) * convm( ij, 1 ) - convm( ij, l ) 
284       ENDDO
285    ENDDO
286
287    !--> flush wflux 
288    !$OMP BARRIER
289    !$acc parallel loop collapse(3) async
290    DO iq=1,nqdyn
291       DO l=ll_begin,ll_endm1
292       !DIR$ SIMD
293          DO ij=ij_begin,ij_end         
294             dtheta_rhodz(ij, l, iq) = dtheta_rhodz(ij, l, iq)  -   0.5 * &
295                  ( wflux(ij,l+1) * (theta(ij,l,iq) + theta(ij,l+1,iq)))
296          END DO
297       END DO
298    END DO
299   
300    !$acc parallel loop collapse(3) async
301    DO iq=1,nqdyn
302       DO l=ll_beginp1,ll_end
303          !DIR$ SIMD
304          DO ij=ij_begin,ij_end         
305             dtheta_rhodz(ij, l, iq) = dtheta_rhodz(ij, l, iq)  +   0.5 * &
306                  ( wflux(ij,l) * (theta(ij,l-1,iq) + theta(ij,l,iq) ) )
307          END DO
308       END DO
309    END DO
310
311    ! Compute vertical transport
312    !$acc parallel loop collapse(2) async
313    DO l=ll_beginp1,ll_end
314       !DIR$ SIMD
315       DO ij=ij_begin,ij_end         
316          wwuu(ij+u_right,l) = 0.5*( wflux(ij,l) + wflux(ij+t_right,l)) * (u(ij+u_right,l) - u(ij+u_right,l-1))
317          wwuu(ij+u_lup,l) = 0.5* ( wflux(ij,l) + wflux(ij+t_lup,l)) * (u(ij+u_lup,l) - u(ij+u_lup,l-1))
318          wwuu(ij+u_ldown,l) = 0.5*( wflux(ij,l) + wflux(ij+t_ldown,l)) * (u(ij+u_ldown,l) - u(ij+u_ldown,l-1))
319       ENDDO
320    ENDDO
321
322    !--> flush wwuu 
323    !$OMP BARRIER
324
325    ! Add vertical transport to du
326    !$acc parallel loop collapse(2) async
327    DO l=ll_begin,ll_end
328       !DIR$ SIMD
329       DO ij=ij_begin,ij_end         
330          du(ij+u_right, l )   = du(ij+u_right,l)  - (wwuu(ij+u_right,l+1)+ wwuu(ij+u_right,l)) / (rhodz(ij,l)+rhodz(ij+t_right,l))
331          du(ij+u_lup, l )     = du(ij+u_lup,l)    - (wwuu(ij+u_lup,l+1)  + wwuu(ij+u_lup,l))   / (rhodz(ij,l)+rhodz(ij+t_lup,l))
332          du(ij+u_ldown, l )   = du(ij+u_ldown,l)  - (wwuu(ij+u_ldown,l+1)+ wwuu(ij+u_ldown,l)) / (rhodz(ij,l)+rhodz(ij+t_ldown,l))
333       ENDDO
334    ENDDO
335
336    !  DO l=ll_beginp1,ll_end
337    !!DIR$ SIMD
338    !      DO ij=ij_begin,ij_end         
339    !        wwuu_out(ij+u_right,l) = wwuu(ij+u_right,l)
340    !        wwuu_out(ij+u_lup,l)   = wwuu(ij+u_lup,l)
341    !        wwuu_out(ij+u_ldown,l) = wwuu(ij+u_ldown,l)
342    !     ENDDO
343    !   ENDDO
344
345    !$acc end data
346    CALL trace_end("compute_caldyn_vert")
347
348  END SUBROUTINE compute_caldyn_vert
349
350  SUBROUTINE compute_caldyn_vert_NH(mass,geopot,W,wflux, W_etadot, du,dPhi,dW)
351    REAL(rstd),INTENT(IN) :: mass(iim*jjm,llm)
352    REAL(rstd),INTENT(IN) :: geopot(iim*jjm,llm+1)
353    REAL(rstd),INTENT(IN) :: W(iim*jjm,llm+1)
354    REAL(rstd),INTENT(IN) :: wflux(iim*jjm,llm+1)
355    REAL(rstd),INTENT(INOUT) :: du(iim*3*jjm,llm)
356    REAL(rstd),INTENT(INOUT) :: dPhi(iim*jjm,llm+1)
357    REAL(rstd),INTENT(INOUT) :: dW(iim*jjm,llm+1)
358    REAL(rstd) :: W_etadot(iim*jjm,llm) ! vertical flux of vertical momentum
359    ! local arrays
360    REAL(rstd) :: eta_dot(iim*jjm, llm) ! eta_dot in full layers
361    REAL(rstd) :: wcov(iim*jjm,llm) ! covariant vertical momentum in full layers
362    ! indices and temporary values
363    INTEGER    :: ij, l
364    REAL(rstd) :: wflux_ij, w_ij
365
366    CALL trace_start("compute_caldyn_vert_nh")
367
368    IF(dysl) THEN
369!$OMP BARRIER
370#include "../kernels/caldyn_vert_NH.k90"
371!$OMP BARRIER
372    ELSE
373#define ETA_DOT(ij) eta_dot(ij,1)
374#define WCOV(ij) wcov(ij,1)
375
376    DO l=ll_begin,ll_end
377       ! compute the local arrays
378       !DIR$ SIMD
379       DO ij=ij_begin_ext,ij_end_ext
380          wflux_ij = .5*(wflux(ij,l)+wflux(ij,l+1))
381          w_ij = .5*(W(ij,l)+W(ij,l+1))/mass(ij,l)
382          W_etadot(ij,l) = wflux_ij*w_ij
383          ETA_DOT(ij) = wflux_ij / mass(ij,l)
384          WCOV(ij) = w_ij*(geopot(ij,l+1)-geopot(ij,l))
385       ENDDO
386       ! add NH term to du
387      !DIR$ SIMD
388      DO ij=ij_begin,ij_end
389          du(ij+u_right,l) = du(ij+u_right,l) &
390               - .5*(WCOV(ij+t_right)+WCOV(ij)) &
391               *ne_right*(ETA_DOT(ij+t_right)-ETA_DOT(ij))
392          du(ij+u_lup,l) = du(ij+u_lup,l) &
393               - .5*(WCOV(ij+t_lup)+WCOV(ij)) &
394               *ne_lup*(ETA_DOT(ij+t_lup)-ETA_DOT(ij))
395          du(ij+u_ldown,l) = du(ij+u_ldown,l) &
396               - .5*(WCOV(ij+t_ldown)+WCOV(ij)) &
397               *ne_ldown*(ETA_DOT(ij+t_ldown)-ETA_DOT(ij))
398       END DO
399    ENDDO
400    ! add NH terms to dW, dPhi
401    ! FIXME : TODO top and bottom
402    DO l=ll_beginp1,ll_end ! inner interfaces only
403       !DIR$ SIMD
404       DO ij=ij_begin,ij_end
405          dPhi(ij,l) = dPhi(ij,l) - wflux(ij,l) &
406               * (geopot(ij,l+1)-geopot(ij,l-1))/(mass(ij,l-1)+mass(ij,l))
407       END DO
408    END DO
409    DO l=ll_begin,ll_end
410       !DIR$ SIMD
411       DO ij=ij_begin,ij_end
412          dW(ij,l+1) = dW(ij,l+1) + W_etadot(ij,l) ! update inner+top interfaces
413          dW(ij,l)   = dW(ij,l)   - W_etadot(ij,l) ! update bottom+inner interfaces
414       END DO
415    END DO
416
417#undef ETA_DOT
418#undef WCOV
419
420    END IF ! dysl
421    CALL trace_end("compute_caldyn_vert_nh")
422
423  END SUBROUTINE compute_caldyn_vert_NH
424END MODULE caldyn_kernels_base_mod
Note: See TracBrowser for help on using the repository browser.