[LFRic] [PSyAD] Mismatch between GH_BASIS kernel dummy arguments in the adjoint, and actual arguments provided in the PSy layer.

mo-joshuacolclough commented 3 days ago

Description

During the LFRic adjoint work, an issue was noted when GH_BASIS arguments are added for intended use with intent(in) fields.

The kernel metadata for the problematic case in the linear code was:

type(arg_type) :: meta_args(4) = (/
  arg_type(GH_FIELD*3, GH_REAL, GH_INC, ANY_SPACE_1), &     ! intent(inout) A
  arg_type(GH_FIELD, GH_REAL, GH_READ, ANY_SPACE_2), &      ! intent(in) B
  ! ...
type(func_type) :: meta_funcs(2) = (/
  func_type(ANY_SPACE_2, GH_BASIS), &                         ! {B on A} BASIS
/)

Later within the kernel arguments, the basis is accepted with the shape: basis_B_on_A(3, ndf_B, ndf_A) - {B on A}

Which looks as expected.

Now, in the adjoint code, the kernel metadata becomes:

type(arg_type) :: meta_args(4) = (/
  arg_type(GH_FIELD*3, GH_REAL, GH_READ, ANY_SPACE_1), &     ! intent(in) A
  arg_type(GH_FIELD, GH_REAL, GH_INC, ANY_SPACE_2), &        ! intent(inout) B
  ! ... (NOTE: No change elsewhere in metadata)

Field A becomes in, B becomes inout. In the kernel (generated with PSyAD), the shape arguments are the same as before. This is correct because we are expecting the same basis as the forward code.

The problem

Since B has become inout, PSyclone erroneously generates code which passes the basis function: {B on B} rather than the required {B on A}, to the kernel.

! psy.f90
B_ON_B_basis(:,df_B,df_nodal) = B%vspace%call_function(BASIS, df_B, nodes_B(:,df_nodal))
call kernel( B_ON_B_basis )

This should instead be:

! psy.f90
B_ON_A_basis(:,df_B,df_nodal) = B%vspace%call_function(BASIS, df_B, nodes_A(:,df_nodal))
call kernel( B_ON_A_basis )

The problem results in an incorrect adjoint. This requires manual intervention using PSy-lite code.

Summary

PSyclone has implicitly chosen B as the other part of the basis due to B being the inout GH_INC field in adjoint code, when in reality the basis should choose A in the adjoint case. This requires manual intervention using PSy-lite code.

Tagging @DrTVockerodtMO

DrTVockerodtMO commented 3 days ago

This may be related to https://github.com/stfc/PSyclone/issues/2335, since I believe a similar version to that kernel produced both of these errors.

arporter commented 3 days ago

Thanks @mo-joshuacolclough. This means that the adjoint kernel is no longer following the rules of the LFRic API (see https://psyclone.readthedocs.io/en/latest/dynamo0p3.html#lfric-gh-shape). If the function space on which the basis functions are to be provided is not that of the updated kernel argument, there will have to be some other mechanism to tell PSyclone that. This probably means an extension to the metadata. Probably, the best thing to do will be to follow what's supported for GH_SHAPE = GH_EVALUATOR (where we have an optional GH_EVALUATOR_TARGETS. Whatever we decide, ideally it must be backwards compatible with existing metadata. We could introduce an optional GH_BASIS_TARGETS for instance. This would need clarification with @tommbendall and @TeranIvy.

stfc / PSyclone