Improvements to the way we use YAXT

[mike-hobson]

Contact Details

@mike-hobson

Are there any linked Issues or Pull Requests?

ticket:4386

Brief description

Our usage of YAXT has led to some very large setup times when running at scale with deep halos. We contacted DKRZ, the authors of YAXT, and they sent a reply with some very useful suggestions that should be investigated.

Further details of the issue.

DKRZ, the authors of YAXT, have made the following very useful suggestions:

On 10 May 2024, Moritz Hanke hanke@dkrz.de wrote:

• overall it looks like a good usage for yaxt; even though the code is quite simple, the functionality it provides is not :-)
• since you generate redists for all halo depth from 1 to size(halo_start), the parameter halo_depth (halo_routing_constructor) does not make much sense anymore
• the halo routing constructor could be made data type independent.
  Instead of generating actual redist, it could generate xmap's for all halo depth. This way, the xmap's could be reused for all datatypes. The function get_redist could generate the actually required redists on-the-fly or reuse a previously generated one (would require a cache of already generated redists and the data type as an additional argument).
• currently the subroutine perform_halo_exchange does not cross-check whether its datatype matches the one used to generate the halo routing object
• if redists are generated on-the-fly, it would also be possible to have optional arguments in perform_halo exchange that contain additional data arrays (even with different datatypes). Then get_redist could generate a single redist for these arrays, which could be exchanged in a single communication call, instead of multiple calls to perform_halo_exchange
• you could use xt_redist_p2p_ext_new instead of xt_redist_p2p_off_new.
  This would save you the need to generate the src/tgt_offsets arrays.

I think this boils down to three main tasks:

1. In infrastructure/source/utilities/halo_comms_mod.F90 we call YAXT to set up both the xt_xmap and the xt_redist for all depths of halo, for all function spaces etc. The xt_xmap is type independent, so could be called once and reused for all field types/kinds. We could split these calls up and only create the xt_redist for each individual field type/kind. (We should look at the relative cost of the calls that generate those quantities to see how useful splitting them would be.)
2. We generate routing tables for every depth of halo up to the maximum halo depth - which is time consuming - and we very rarely use most of those routing tables (maybe just halo_depth=1 and halo_depth=max_halo_depth). Without losing any functionality, we could move to a system where we generate the routing tables on first use, then cache them for later, like we currently do with function space. Any routing tables that are never used, would never be generated.
3. We could try switching our current call to xt_redist_p2p_off_new for a call to xt_redist_p2p_ext_new.

If they are small enough, these pieces of work could be done here, but I suspect (at least) some of the work is big enough to spin into its own ticket(s).