Explore doing the LFRic_apps segmentation pattern automatically win transformation scripts

[sergisiso]

Segmentation (iterating through the cell-columns in 1D segments) is beneficial in multiple places of LFRic_atm. At the moment this optimisation is encoded in the science code (by passing the whole domain in the kernel and choosing the segmented iteration patten there, or by creating a psykal-lite that does the same). However, since it is an optimisation, it would make sense to do in psyclone - in the psykal layer.

We already have a couple of ways to change the loop iteration order, e.g. ChunkLoopTrans or the ColourTiling. We need to explore how these differ from the manually implemented Segmentation and also understand if Segment buffers are created.

[TeranIvy]

The relevant LFRic Apps ticket where this requirement stems from is https://code.metoffice.gov.uk/trac/lfric_apps/ticket/962. I asked the developer to provide more details in this issue.

[dcaseGH]

LFRic apps 962 (above) introduces segmentation to orographic gravity wave drag, and includes OMP parallelization by hand.
I don't know what the ambition of the scope of PSyclone is, but if it is acceptable to allow the manual refactoring of the code to stay as it is, but to control the OMP parallelization through a script, then the functionality is fine as things are.
For example, you can use this script, which is derived from something commonly used I presume:

from psyclone.psyir.nodes import (Call, OMPDirective, Loop)                                                                                                     
from psyclone.transformations import (OMPParallelLoopTrans, TransformationError)                                                                                
                                                                                                                                                                
# Simple script to add OMP with chosen schedule to a loop which we know is safe                                                                                 
# To be applied to psykal_lite_phys_mod.F90                                                                                                                     
                                                                                                                                                                
SAFE_IMPURE_CALLS = ['orographic_drag_kernel_code']                                                                                                             
def trans(psyir):                                                                                                                                               
    loop_idx = 0                                                                                                                                                
    omp_loop_trans = OMPParallelLoopTrans(omp_schedule="dynamic", omp_directive="paralleldo")                                                                   
    for loop in psyir.walk(Loop):                                                                                                                               
        loop_name = f"LOOP{loop_idx} " + loop.debug_string().split('\n')[0]                                                                                     
        loop_idx += 1                                                                                                                                           
        try:                                                                                                                                                    
            # set the routine to be pure as we know it is safe to parallelize                                                                                   
            impure_calls = [c for c in loop.walk(Call) if not c.is_pure]                                                                                        
            for call in impure_calls:                                                                                                                           
                if call.routine.symbol.name in SAFE_IMPURE_CALLS:                                                                                               
                    call.routine.symbol.is_pure = True                                                                                                          
            omp_loop_trans.apply(loop, options={"node-type-check": False})                                                                                      
                                                                                                                                                                
        except TransformationError as err:                                                                                                                      
            print(f"{psyir.name}: Unable to apply OMPParallelLoopTrans to {loop_name}: {err}")                                                                  
                                                                                                                                                                
    directives = psyir.walk(OMPDirective)                                                                                                                       
                                                                                                                                                                
    if len(directives) == 0:                                                                                                                                    
        print(f"{psyir.name}: Unable to add any OpenMP directives")                                                                                             
    else:                                                                                                                                                       
        print(f"{psyir.name}: Successfully added {len(directives)} OpenMP directives")

and of course tested with psyclone -s psykal_lite_phys_mod_transform.py psykal_lite_phys_mod.F90
this will correctly wrap up (only) the loop which is wanted, and set the same variables to be private as was written in the hand-coded version.

If the reviewers agree that this approach is a good one, then PSyclone needs nothing more, and I will continue the review process on LFRic apps tickets.
If you were thinking of something else, then please let me know. Thanks, Dave

[sergisiso]

For reference, the example of segmented loop is in interfaces/physics_schemes_interface/source/psy/psykal_lite_phys_mod.F90:

n_segments = (applicable_points + seg_target - 1) / seg_target

!$omp parallel do default(shared), private(segment, seg_len, l_bound, u_bound), schedule(dynamic)
do segment=1, n_segments
	seg_len = min(seg_target, applicable_points - (segment - 1) * seg_target)
	l_bound = (segment - 1) * seg_target + 1
	u_bound = l_bound + seg_len - 1
	call orographic_drag_kernel_code(                                                &
	                loop_upper_bound, nlayers, du_blk_proxy%data, dv_blk_proxy%data, &
	                ...
	                map_adspc1_sd_orog, seg_len, cell_index(l_bound:u_bound))
end do
!$omp end parallel do

The kernel then allocates several temporary arrays with:

real(r_um), dimension(seg_len) :: ...
real(r_um), dimension(seg_len, nlayers) :: ...

And iterate the cell-columsn with:

    do i= 1, seg_len
      do k = 1, nlayers
        u_on_p(i,k) = real(u_in_w3(map_w3(1,cell_index(i)) + k-1), r_um)
        v_on_p(i,k) = real(v_in_w3(map_w3(1,cell_index(i)) + k-1), r_um)