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

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

trunk : GPU implementation with OpenACC ( merge from glcp.idris.fr )
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)
212    USE disvert_mod, ONLY : bp
213    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
214    REAL(rstd),INTENT(IN)  :: theta(iim*jjm,llm,nqdyn)
215    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)
216    REAL(rstd),INTENT(INOUT)  :: convm(iim*jjm,llm)  ! mass flux convergence
217
218    REAL(rstd),INTENT(INOUT) :: wflux(iim*jjm,llm+1) ! vertical mass flux (kg/m2/s)
219    REAL(rstd),INTENT(INOUT) :: wwuu(iim*3*jjm,llm+1)
220    REAL(rstd),INTENT(INOUT) :: du(iim*3*jjm,llm)
221    REAL(rstd),INTENT(INOUT) :: dtheta_rhodz(iim*jjm,llm,nqdyn)
222    REAL(rstd),INTENT(OUT) :: dps(iim*jjm)
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.