Add support for async f! and j!

[charleskawczynski]

@tapios suggested adding this if it's not too difficult, and here it is:

This PR adds support for asynchronous f! and j! on both CPU and GPU.

We can try it out by passing the comms context, but we don't need to use it if we don't want.

[tapios]

Nice! When you know how much it affects performance, please let us know.

[simonbyrne]

Trying it out here: https://buildkite.com/clima/climaatmos-target-gpu-simulations/builds/122

[charleskawczynski]

New build at:

• https://buildkite.com/clima/climaatmos-ci/builds/14157
• https://buildkite.com/clima/climaatmos-target-gpu-simulations/builds/123

[charleskawczynski]

Interesting, we're hitting maxiter reached in saturation_adjustment:. Ah, right, it was reduced to 3. 🤷🏻

[charleskawczynski]

I guess, maybe, we need a smaller timestep? (we could try increasing maxiter, but I assume it was reduced for good reason).

[charleskawczynski]

Update on this: the failures in ClimaAtmos were because the device was not synchronized, adding proper syncs did fix this.

I've updated the PR in a couple ways:

• The multiple streams version is now in a CUDA extension, so that we don't need to directly depend on CUDA
• I've kept both the multiple streams and multiple tasks for CUDADevices (we should compare them)
• I've added a multithreaded version for the CPU

We should compare the streams vs tasks on the CUDA side. The multiple streams version with events do seem to be executing T_exp! and T_lim! asynchronously on separate streams, however, are not executing concurrently. The reason being that we're probably utilizing too many resources, and the scheduler is deciding to execute them serially.

One way we can try getting concurrent execution is by reducing the number of blocks in our kernel launches, but we might only want to (somehow) do that in T_exp! and T_lim!, rather than globally.

So, all in all, this isn't immediately offering a performance improvement on GPUs yet. The CPU multithreaded implementation might be an improvement, but it's not yet been tested/benchmarked. Also, there could be interplay with threads being used in the broadcast kernels (which, perhaps, we could/should just remove).

[charleskawczynski]

Further update: we will not add async support with T_exp! and T_lim!, as fusing these allowed us to eliminate dss calls. Instead, we'll add async support with f! and j!.

I've updated the PR to reflect this.

[charleskawczynski]

It could be that CUDA doesn't like making these streams at every RHS evaluation. Maybe we can try caching this in the timestepper?

[charleskawczynski]

This could be really good for strong scaling