GPU performance of MUSCL tracer-advection scheme in NEMO

[arporter]

The MUSCL scheme is used in the MED24 configuration and is currently a massive performance bottleneck. This discussion is to analyse the code and see what can be done.

[arporter]

The issue is a very large, doubly-nested loop over tracer species (?) and vertical levels. This loop contains many, separate doubly nested i,j loops. However, it also contains calls to diagnostic routines and these are a problem:

    ! 'dia_ptr_hst' is not available on the 'default' accelerator device, and therefore it cannot be called from within an OMP 
!& Target region. Use the 'device_string' option to specify a different device.
    do jn = 1, kjpt, 1
      ! 'dia_ptr_hst' is not available on the 'default' accelerator device, and therefore it cannot be called from within an OMP 
!& Target region. Use the 'device_string' option to specify a different device.
      do jk = 1, jpkm1, 1
        !$omp target
        !$omp teams loop collapse(2) default(shared) private(ji,jj)
        do jj = loop_start_5, loop_stop_5, 1
          do ji = loop_start_6, loop_stop_6, 1
            zdxt(ji,jj) = (pt(ji + 1,jj,jk,jn,kbb) - pt(ji,jj,jk,jn,kbb)) * umask(ji,jj,jk)
            zdyt(ji,jj) = (pt(ji,jj + 1,jk,jn,kbb) - pt(ji,jj,jk,jn,kbb)) * vmask(ji,jj,jk)
          enddo
        enddo
...
          enddo
        enddo
        !$omp end teams loop
        !$omp end target
        if (l_ptr) then
          call dia_ptr_hst(jn, 'adv', zwy(:,:))
        end if
        if (l_hst) then
          call dia_ar5_hst(jn, 'adv', zwx(:,:), zwy(:,:), ldfin=jk == jpkm1)
        end if
      enddo
      !$omp target
      !$omp teams loop collapse(2) default(shared) private(idx_5,idx_6)
      do idx_5 = loop_start_13, loop_stop_13, 1
...

dia_ptr_hst is actually an interface (in diaptr.f90) that then itself calls one or two other routines..

  subroutine dia_ptr_hst_2d(ktra, cptr, pvflx)
    integer, intent(in) :: ktra
    CHARACTER(LEN = 3), INTENT(IN) :: cptr
    real(kind=wp), dimension(:,:), intent(in) :: pvflx

    call dia_ptr_hst_t(ktra, cptr, pvflx2d=pvflx(:,:), ktvflx=lbnd_ij(pvflx))

  end subroutine dia_ptr_hst_2d

and

  subroutine dia_ptr_hst_t(ktra, cptr, pvflx2d, pvflx3d, ktvflx)
    integer, dimension(2), intent(in) :: ktvflx
    integer, intent(in) :: ktra
    CHARACTER(LEN = 3), INTENT(IN) :: cptr
    REAL(KIND = wp), DIMENSION(ktvflx(1) :, ktvflx(2) :), INTENT(IN), OPTIONAL :: pvflx2d
    REAL(KIND = wp), DIMENSION(ktvflx(1) :, ktvflx(2) :, :), INTENT(IN), OPTIONAL :: pvflx3d
    real(kind=wp), dimension(ntsj - 0:ntej + 0,nbasin) :: zsj
    integer :: jn

    if (PRESENT(pvflx2d)) then
      do jn = 1, nbasin, 1
        zsj(:,jn) = ptr_sj(pvflx2d(:,:), btmsk(:,:,jn))
      enddo
    else
      if (PRESENT(pvflx3d)) then
        do jn = 1, nbasin, 1
          zsj(:,jn) = ptr_sj(pvflx3d(:,:,:), btmsk(:,:,jn))
        enddo
      end if
    end if
    if (cptr == 'adv') then
      if (ktra == jp_tem) then
        call ptr_sum(hstr_all(:,:,jp_tem,1), zsj(:,:))
      end if
      if (ktra == jp_sal) then
        call ptr_sum(hstr_all(:,:,jp_sal,1), zsj(:,:))
      end if
    ...

and ptr_sum is a call to another interface, e.g.:

  subroutine ptr_sum_2d(phstr, pva)
    real(kind=wp), dimension(njs0 - 0:nje0 + 0,nbasin), intent(inout) :: phstr
    real(kind=wp), dimension(ntsj - 0:ntej + 0,nbasin), intent(in) :: pva
   ...
    loop_stop = ntej + 0
    loop_start = ntsj - 0
    !$omp target
    !$omp teams loop collapse(2) default(shared) private(jj,jn)
    do jn = 1, nbasin, 1
      do jj = loop_start, loop_stop, 1
        phstr(jj,jn) = phstr(jj,jn) + pva(jj,jn)
      enddo
    enddo
    !$omp end teams loop
    !$omp end target

  end subroutine ptr_sum_2d

ptr_sj is also a call to a routine in an interface:

  function ptr_sj_2d(pvflx, pmsk) result(p_fval)
    real(kind=wp), dimension(:,:), intent(in) :: pvflx
    real(kind=wp), dimension(nis0 - nn_hls:nie0 + nn_hls,njs0 - nn_hls:nje0 + nn_hls), intent(in) :: pmsk
    real(kind=wp), dimension(ntsj - 0:ntej + 0) :: p_fval

    call ptr_sj_2d_t(pvflx(:,:), lbnd_ij(pvflx), pmsk(:,:), p_fval(:))

  end function ptr_sj_2d
  subroutine ptr_sj_2d_t(pvflx, ktvflx, pmsk, p_fval)
    integer, dimension(2), intent(in) :: ktvflx
    real(kind=wp), dimension(ktvflx(1):,ktvflx(2):), intent(in) :: pvflx
    real(kind=wp), dimension(nis0 - nn_hls:nie0 + nn_hls,njs0 - nn_hls:nje0 + nn_hls), intent(in) :: pmsk
    real(kind=wp), dimension(ntsj - 0:ntej + 0), intent(out) :: p_fval
  ...

    loop_stop = UBOUND(p_fval, dim=1)
    loop_start = LBOUND(p_fval, dim=1)
    !$omp target
    !$omp teams loop collapse(1) default(shared) private(idx)
    do idx = loop_start, loop_stop, 1
      p_fval(idx) = 0._wp
    enddo
    !$omp end teams loop
    !$omp end target
    loop_stop_1 = ntej + 0
    loop_start_1 = ntsj - 0
    loop_stop_2 = ntei + 0
    loop_start_2 = ntsi - 0
    !$omp target
    !$omp teams loop collapse(1) default(shared) private(ji,jj)
    do jj = loop_start_1, loop_stop_1, 1
      ! Error: The write access to 'p_fval(jj)' in 'p_fval(jj) = p_fval(jj) + pvflx(ji,jj) * pmsk(ji,jj) * tmask_i(ji,jj)' and the 
!& read access to 'p_fval(jj)' in 'p_fval(jj) = p_fval(jj) + pvflx(ji,jj) * pmsk(ji,jj) * tmask_i(ji,jj)' are dependent and cannot 
!& be parallelised. Variable: 'p_fval'. Consider using the "ignore_dependencies_for" transformation option if this is a false 
!& dependency.
      do ji = loop_start_2, loop_stop_2, 1
        p_fval(jj) = p_fval(jj) + pvflx(ji,jj) * pmsk(ji,jj) * tmask_i(ji,jj)
      enddo
    enddo
    !$omp end teams loop
    !$omp end target

  end subroutine ptr_sj_2d_t

[arporter]

Given there's a '3D' version of diaptr_hst, possibly the rank-2 work arrays zwx and zwy could be promoted to rank-3 and the calls to the diagnostics moved outside the loop over k? (Currently PSyclone would refuse to do this for zwx because it is re-used after the k loop.)

[arporter]

Inlining does look feasible but the problem is that it needs to be done recursively as the routines in diaptr depend on other routines in diaptr. Currently, we refuse to module-inline a routine that has any dependencies to its parent module. We could extend this so that if those dependencies are routines then we module-inline them too or fully inline them first.

[francescamele]

Regarding dia_ptr and dia_ar5, these routines are not called in MED24 configuration. We are not outputting any diagnostics within NEMO itself; all diagnostics and analyses are performed in post-processing.

[arporter]

Thanks Francesca. In that case, I think the simplest solution is to try to extend the transformation script to remove these calls. I'll have a go. Longer term, we could add support for setting the values of certain variables (e.g. l_ptr = .FALSE.) and then automatically (#2869) remove dead code.

[arporter]

@francescamele I've made a branch with an example of how to remove subroutine calls: 3285_rm_calls (https://github.com/stfc/PSyclone/compare/3285_rm_calls?expand=1).
It's very simple and doesn't do any checks or attempt to tidy-up the remaining code (e.g. remove empty if blocks).
With this script, PSyclone generates:

    !$omp target
    !$omp teams loop collapse(1) default(shared) private(idx_5,idx_6,idx_7,idx_8,ik,ji,jj,jk,jn,loop_start_19,loop_start_20,loop_start_21,loop_start_22,loop_start_23,loop_start_24,loop_start_25,loop_start_26,loop_stop_19,loop_stop_20,loop_stop_21,loop_stop_22,loop_stop_23,loop_stop_24,loop_stop_25,loop_stop_26,z0u,z0v,z0w,zalpha,zdxt,zdyt,zdzt_kp2,zfw_kp1,zslpx,zslpy,zslpz_kp1,zu,zv,zw,zwx,zwy,zzslpx,zzslpy,zzwx,zzwy)
    do jn = 1, kjpt, 1
      ! Loop cannot be collapsed because it has siblings
      do jk = 1, jpkm1, 1
        do jj = loop_start_5, loop_stop_5, 1
          do ji = loop_start_6, loop_stop_6, 1

which is better than before but still not good because it has parallelised only the outer loop over kjpt which is presumably just tracer species?
@sergisiso I used to have a heuristic for this type of situation where, if a loop contained several nested loops, I refused to parallelise the outermost one and just did the loop nests inside. Do we do anything like that currently?

[arporter]

@francescamele I've updated the utils.py script in the branch now so feel free to try. If you're not using omp_gpu_trans.py then you'll need to update whichever script you are using to call the new remove_calls method.

[arporter]

Actually, I've realised that all I've done is replicate what PSyclone was already doing but in a slightly different way :-(

[arporter]

It may be that the version where we offload the whole jn loop might be better as there may be less CPU-GPU traffic. However, it will run very slowly on the GPU itself as there's so little parallelism.

[sergisiso]

Although we could (probably) solve this by transforming them to extend their rank, that would still not be ideal as the jk loop contains multiple ji/jj loop nests

The increase_rank_and_reorder_nemov5_loops in utils.py may be doing what you need. It does not only increases the rank but also creates a copy of the selected loop inside each inner loop. But you do have to check yourself that you are not breaking dependencies.

(And currently it does it for specific loop names, which are not the same as in your example)

[arporter]

Thanks Sergi. Unfortunately, some of the work arrays are re-used after the jk loop. I can see that they are completely overwritten so it would be safe to rank-increase them but PSyclone can't do that yet (we have an issue to look at coverage of writes to arrays).
There is a final k loop that has a loop-carried dependence (in zdxt) too:

         DO jk = 1, jpkm1
            IF( jk < jpkm1 ) THEN
               DO jj = ntsj-( 0), ntej+(  0 ) ; DO ji = ntsi-( 0), ntei+(  0)
...
                  zslpz_kp1 =                           ( zdxt(ji,jj) + zdzt_kp2 )  &
                  &         * (  0.25_wp + SIGN( 0.25_wp, zdxt(ji,jj) * zdzt_kp2 )  )
...
                  zzwy = pt(ji,jj,jk  ,jn,Kbb) + xind(ji,jj,jk) * zw * zslpx(ji,jj)
...
                  zdxt(ji,jj) = zdzt_kp2
                  zslpx   (ji,jj) = zslpz_kp1
               END DO   ;   END DO
            ELSE
               DO jj = ntsj-( 0), ntej+(  0 ) ; DO ji = ntsi-( 0), ntei+(  0)          !-- vertical advective trend at jpkm1
                  pt(ji,jj,jk,jn,Krhs) =  pt(ji,jj,jk,jn,Krhs) - zwx(ji,jj)    &
                     &                                      * r1_e1e2t(ji,jj) / (e3t_1d(jk) *(1._wp+r3t(ji,jj,Kmm)*tmask(ji,jj,jk)))
               END DO   ;   END DO
            ENDIF
	 END DO                  ! end of jk loop
         !
   END DO                     ! end of tracer loop

i.e. the next iteration of the k loop uses the value of zdxt(ji,jj) computed at the previous k value.

[arporter]

@francescamele I think what would be really helpful for us is if you could do some manual experiments with this particular subroutine to see what works best. Some ideas:

• Create some new, rank-3 work arrays and use them in place of zdxt etc. inside the big jk loop;
• Experiment with having the jn loop offloaded but run in serial (just !$omp target before the loop I think) and then parallelise the inner loops;
• Try 'fissioning' the k-loop to get back the k,j,i triply-nested loops;

I'm not sure what we can do about the last k loop with the loop-carried dependence. Can we alter it so that the k loop is innermost and have that run by just a single thread?

[francescamele]

@arporter Thanks for all the suggestions! I’ll try these approaches and run some manual experiments on the subroutine to see which works best. I’ll keep you updated on the results.