(Towards/Close #3315, Closes #3590) New implementation of lfric loop fuse trans

[LonelyCat124]

Initial implementation for the new lfric loop fuse trans setup.
There are still for sure some things to sort out but it seems to work for our tests and solves ANY_SPACE things.

I'd like to have an investigate of being able to fuse more than 2 LFRicLoops, provided they all meet some criteria (i.e. some_space is provided, they're all on the same field etc., or that they all have the same space at runtime) but I'd need to know if that is useful @christophermaynard @MetBenjaminWent ? Its probably a more substantial implementation but I think its possible.

[sergisiso]

Here is the example of min and max builtins that could be easy targets to fuse: src/psyclone/tests/test_files/lfric/15.10.9_min_max_X_builtin.f90

[arporter]

Some notes from our meeting.
I think the check list for whether or not two kernels can be fused is:

• Do the two kernels OPERATE_ON the same thing (dofs or cell columns)?
• Is their iteration space determined by the same field/operator?
• Or, do their iteration spaces correspond to the same function space (static check, will need to resolve 'ANY_SPACE' to an actual space)?
• [Optionally] do the fuse but protect it with a runtime check on the function spaces.

We think most of the the benefit will come from fusing built-ins (particularly those that do MIN/MAX on the same field).
For user-provided kernels that operate on cell-columns, I think they can be fused following the rules above, provided that the DA doesn't determine that a halo-exchange is required between them. (This allows for e.g. a kernel that increments a field on W2 that is also written to in the first kernel and also for stencil accesses.) @christophermaynard thinks the opportunities for doing this kind of fusion will be very limited.

We also discussed 'constant propagation' but realised that often, a setval_c is used to initialise a field that is then incremented by a kernel that loops over cell-columns. Since, for continuous function spaces, different cell-columns will increment the same dof it will not be possible to do 'constant propagation' in this case.

[LonelyCat124]

I think the new implementation does most of this, what I couldn't work out is why is this necessary:

Is their iteration space determined by the same field/operator?

This can be either:

1. to save searching for what the space is if they're both built ins?
2. If we have 2 any_space Kerns without a non-any_space kern on the field to be able to work out the iteration space?

Does that make sense @arporter ?

I am also unsure I've done anything re: " I think they can be fused following the rules above, provided that the DA doesn't determine that a halo-exchange is required between them". I'll have to work that out next.

[arporter]

This can be either:

1. to save searching for what the space is if they're both built ins?
2. If we have 2 any_space Kerns without a non-any_space kern on the field to be able to work out the iteration space?

Yes, sorry, I was a bit slack in how I phrased it. If the iteration space is defined by the same field/operator then we don't have to look at function spaces. We probably still need to worry about stencil accesses. In fact, now that I write that, we need to be careful about kernels that update more than one argument (they are permitted to do this). A first step might be to refuse to fuse such cases.

[LonelyCat124]

I think I've added a check in validate for writing to only one field (I think).
I think it would be good to have more things in validate than apply, however there are challenges due to the different behaviour with conditional_fusion, which we need a lot of information to determine whether its required or not.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (e93a73c) to head (a2f485a).
⚠️ Report is 47 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##            master     #3595    +/-   ##
==========================================
  Coverage   100.00%   100.00%            
==========================================
  Files          403       403            
  Lines        56790     56938   +148     
==========================================
+ Hits         56790     56938   +148     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[LonelyCat124]

I think I have most of the basics done now, but there are some remaining hurdles that I'm not sure I understand.

The first is this test (not implemented in the pushed version) that fails, but I think we wanted to be able to do:

def test_loop_fuse_min_max_same_var():
    '''Test that we fuse two builtins on the same field without any
    other invoke elements.'''
    _, invoke = get_invoke(
        "15.10.9_min_max_X_builtin.f90",
        TEST_API, name="invoke_0", dist_mem=False)
    schedule = invoke.schedule
    ftrans = LFRicLoopFuseTrans()

    ftrans.apply((schedule.children[0], schedule.children[1]))
    ftrans.apply((schedule.children[0], schedule.children[1]))

This fails due to an existing test in validate, where we aren't allowed to fuse 2 reduction kernels (we have a setval_c, a max and a min in the input invoke).
Is there an appropriate way to determine whether this should be allowed?

The other thing I'm not sure about was working out how to avoid fusing loops if there is some dependency between them. I assume I can check if loop1.forward_dependency == loop2 (or something similar)? I'm also unclear how to test that case.

@arporter @sergisiso If you have any idea about either of these let me know.

[LonelyCat124]

I've allowed reductions to be fused if the iteration space variable is the same - not sure if this is correct in general?

[LonelyCat124]

I've implemented some logic locally to disallow fusing of loops that would need a halo exchange if distributed memory is enabled, however it causes a current test to fail (with distributed memory enabled).

The test in question fuses the loops in:

  call invoke(                           &
       testkern_type(a, f1, f2, m1, m2), &
       testkern_type(a, f1, f2, m1, m2)  &
       )

where testkern is:

  type, extends(kernel_type) :: testkern_type
     type(arg_type), dimension(5) :: meta_args =                  &
          (/ arg_type(gh_scalar, gh_real, gh_read),               &
             arg_type(gh_field,  gh_real, gh_inc,  w1),           &
             arg_type(gh_field,  gh_real, gh_read, w2),           &
             arg_type(gh_field,  gh_real, gh_read, w2),           &
             arg_type(gh_field,  gh_real, gh_read, w3, ndata="1") &
           /)
     integer :: operates_on = cell_column
   contains
     procedure, nopass :: code => testkern_code
  end type testkern_type

My logic must be wrong/missing something, because I refuse to accept this whilst we don't generate halo exchanges between these loops:

  if (m2_proxy%is_dirty(depth=1)) then
    call m2_proxy%halo_exchange(depth=1)
  end if
  do cell = uninitialised_loop0_start, uninitialised_loop0_stop, 1
    call testkern_code(nlayers_f1, a, f1_data, f2_data, m1_data, m2_data, ndf_w1, undf_w1, map_w1(:,cell), ndf_w2, undf_w2, map_w2(:,cell), ndf_w3, undf_w3, map_w3(:,cell))
  enddo

  ! Set halos dirty/clean for fields modified in the above loop(s)
  call f1_proxy%set_dirty()
  do cell = uninitialised_loop1_start, uninitialised_loop1_stop, 1
    call testkern_code(nlayers_f1, a, f1_data, f2_data, m1_data, m2_data, ndf_w1, undf_w1, map_w1(:,cell), ndf_w2, undf_w2, map_w2(:,cell), ndf_w3, undf_w3, map_w3(:,cell))
  enddo

  ! Set halos dirty/clean for fields modified in the above loop(s)
  call f1_proxy%set_dirty()

so I need to check the halo exchange logic more to check I've not missed something else.

[LonelyCat124]

I realised that wasn't allowing the test to run to completion without fusing - if I allow that then instead I get (the first loop is coloured but not relevat):

    do colour = loop0_start, loop0_stop, 1
      do cell = loop1_start, last_halo_cell_all_colours(colour,max_halo_depth_mesh), 1
        call testkern_code(nlayers_f1, a, f1_data, f2_data, m1_data, m2_data, ndf_w1, undf_w1, map_w1(:,cmap(colour,cell)), ndf_w2, undf_w2, map_w2(:,cmap(colour,cell)), ndf_w3, undf_w3, map_w3(:,cmap(colour,cell)))
      enddo
    enddo

    ! Set halos dirty/clean for fields modified in the above loop(s)
    call f1_proxy%set_dirty()
    call f1_proxy%set_clean(max_halo_depth_mesh - 1)
    do cell = loop2_start, loop2_stop, 1
      call testkern_code(nlayers_f1, a, f1_data, f2_data, m1_data, m2_data, ndf_w1, undf_w1, map_w1(:,cell), ndf_w2, undf_w2, map_w2(:,cell), ndf_w3, undf_w3, map_w3(:,cell))
    enddo

Something in PSyclone works out that the f1 halo exchange isn't "real", which looks like it comes from _add_field_component_halo_exchange and is something like required, _ = exchange.required(ignore_hex_dep=True) on the created Halo exchange. I'll have to check the logic there. The difficulty is I don't want to add a halo exchange into the tree to remove it to do this step ideally, so this may need a variety of refactoring to get right.

Its not required as "we only read annexed dofs and these have been made clean by the redundant computation", however the test states: "Test that we are able to fuse two loops together, perform redundant computation and then colour." which we need to explicitly disallow the fusion first surely? Since to be able to know there isn't a dependency we need to do the redundant computation first? But we need to do the redundant computation on the fused loop I assume (uhoh).

[LonelyCat124]

@arporter @sergisiso Ready for a first look now I think. I think at this stage I'm hopeful that all the cases we want to fuse are handled - if I've missed testing anything you can think of we can add more tests.

I'm not sure about the split of work between validate and apply - apply has a variety of checks to fall through to the conditional fusing but I couldn't think of another way to implement it at the moment.

[arporter]

Just a comment/the $1M question - is it possible to extend one of the LFRic transformation scripts to exercise this functionality for real in (one of) the ITs?

[LonelyCat124]

Just a comment/the $1M question - is it possible to extend one of the LFRic transformation scripts to exercise this functionality for real in (one of) the ITs?

I added it into the everything script and will run the lfric ITs and see.

[LonelyCat124]

Fusing loops currently causes lfric to crash, not quite sure yet, , i'm guessing its fusing loops with allocations or something maybe

[LonelyCat124]

Fixed that bug again - there's still failing test but I want to run ITs again to see if LFRic is happy yet then i'll fix up the tests.

[LonelyCat124]

I fixed the remaining bugs (one being caused by the "improvement" to data_sharing_attribute_mixin). The current rules now for fusing LFRicLoops (on top of whatever LoopFuseTrans does, which is mostly require index alignment I think):

1. No inter-grid kernels.
2. Function spaces must be valid.
3. The function space must be the same or discontinuous.
4. All kernels must only write to a single field.
5. The upper_bound_name must be the same.
6. The halo depths must be the same.
7. All reductions must be on the same field and the second loop can't read the result of a reduction in the first loop.
8. There's no dependency between the loops that could require a halo exchange (regardless of whether it really does or not, this is slightly stricter than the halo exchange logic).
9. If we have a non-discontinuous ANY_SPACE field, we need to be able to work out what the field space is or must be doing conditional fusion.
10. The computed field space for ANY_SPACE fields must be the same space if not discontinuous.

@arporter Does this seem reasonable? I can add it to the docs somewhere as appropriate (LFRic developer guide? Or user guide?)

[LonelyCat124]

Just a comment/the $1M question - is it possible to extend one of the LFRic transformation scripts to exercise this functionality for real in (one of) the ITs?

This now works and compiles and runs lfric_atm without causing the IT to fail so I can only assume its ok?

[arporter]

Does this seem reasonable? I can add it to the docs somewhere as appropriate (LFRic developer guide? Or user guide?)

Thanks Aidan, it does seem reasonable. Probably it should go in the docstring of the class and then we can pull it into the docs without replicating it.