stfc / PSyclone

Domain-specific compiler and code transformation system for Finite Difference/Volume/Element Earth-system models in Fortran
BSD 3-Clause "New" or "Revised" License
105 stars 28 forks source link

Clarify rules for determining nlayers in LFRic kernels #2636

Closed TeranIvy closed 3 weeks ago

TeranIvy commented 3 months ago

LFRic API seems to return different nlayers depending on whether invoke calls to setval_c built-in and a kernel are joint or not. This was commented on in the LFRic Apps ticket 21, pres_lev_diags_alg_mod.x90. The example pruned-down code is below.

Mock algorithm code

program test_invokes

  ! Description: test single and joint invokes
  use constants_mod, only: r_def
  use field_mod,     only: field_type
  use testkern_mod,  only: testkern_type

  implicit none

  type(field_type) :: plev_heaviside, exner_wth

  ! Case 1: Separate invokes, setval_c before the kernel call
  call invoke( setval_c(plev_heaviside, 1.0_r_def) )
  call invoke( &
       testkern_type(exner_wth, plev_heaviside) )

  ! Case 2: Separate invokes, setval_c after the kernel call
  call invoke( &
       testkern_type(exner_wth, plev_heaviside) )
  call invoke( setval_c(plev_heaviside, 1.0_r_def) )

  ! Case 3: Joint invoke, setval_c before the kernel call
  call invoke( setval_c(plev_heaviside, 1.0_r_def), &
               testkern_type(exner_wth, plev_heaviside) )

  ! Case 3: Joint invoke, setval_c after the kernel call
  call invoke( testkern_type(exner_wth, plev_heaviside), &
               setval_c(plev_heaviside, 1.0_r_def) )

end program test_invokes

Mock kernel code

module testkern_mod

  use argument_mod
  use fs_continuity_mod
  use kernel_mod
  use constants_mod

  implicit none

  type, extends(kernel_type) :: testkern_type
     type(arg_type), dimension(2) :: meta_args =                                   &
          (/ arg_type(GH_FIELD, GH_REAL, GH_READ,      ANY_DISCONTINUOUS_SPACE_1), &
             arg_type(GH_FIELD, GH_REAL, GH_READWRITE, ANY_DISCONTINUOUS_SPACE_2)  &
           /)
     integer :: operates_on = cell_column
   contains
     procedure, nopass :: code => testkern_code
  end type testkern_type

contains

  subroutine testkern_code(nlayers,                &
                          exner,                   &
                          plev_heaviside,          &
                          ndf_in, undf_in, map_in, &
                          ndf_out, undf_out, map_out)

    implicit none

    integer(kind=i_def), intent(in) :: nlayers
    integer(kind=i_def), intent(in) :: ndf_in
    integer(kind=i_def), intent(in), dimension(ndf_in) :: map_in
    integer(kind=i_def), intent(in) :: ndf_out
    integer(kind=i_def), intent(in), dimension(ndf_out) :: map_out
    integer(kind=i_def), intent(in) :: undf_in, undf_out
    real(kind=r_def), intent(in), dimension(undf_in) :: exner
    real(kind=r_def), intent(inout), dimension(undf_out) :: plev_heaviside

  end subroutine testkern_code

end module testkern_mod

Generated PSy-layer code for nlayers

Note that setval_c built-in defines the field access as GH_WRITE (see https://github.com/stfc/PSyclone/blob/master/src/psyclone/parse/lfric_builtins_mod.f90).

When invokes are separate (Case 1 and Case 2), nlayers is determined from the first field in the kernel argument list, exner, which has GH_READ access: nlayers = exner_wth_proxy%vspace%get_nlayers().

When invokes are joint, nlayers is determined from the first field in the first kernel or built-in call in the joint invoke. In Case 3 this is nlayers = plev_heaviside_proxy%vspace%get_nlayers(), and in Case 4 this is nlayers = exner_wth_proxy%vspace%get_nlayers().

Questions/discussion for @arporter, @sergisiso and @hiker

arporter commented 3 months ago

@TeranIvy and I discussed this today and have agreed:

  1. On no account should PSyclone look-up nlayers from field arguments to kernels that operate_on dofs;
  2. For each kernel that has operates_on=cell_column, PSyclone will look-up nlayers for the first argument that is written to.

I've grepped through the tests and, AFAICS, nlayers is only ever looked-up and then supplied to a kernel call so we are free to have a separate variable for each kernel (or, more specifically, the first field that a kernel writes to).

arporter commented 2 months ago

Currently, it is DynCellIterators that is responsible for looking up nlayers and creating a tagged symbol to store it. This will need to be updated to find the first field/operator that is written to by each kernel in the invoke and create an nlayers symbol for each of them. We must have functionality to do some of this since the cell-column iteration space is determined by looking at the field/op arguments that are written to by a kernel.