Evaluating a field at certain indices in a GPU simulation

[SamDeAbreu]

HI there!

I am trying to compute a term that has variables in it that must be evaluated at a depth z=h(x,y). That is:

term(x,y) = f(x, y, h(x, y)).

In this case, h is just the mixed-layer depth, which I define to be the depth of maximum salinity gradient. So, I managed to write some code to do this on GPU, but it involves a lot of indexing and use of CUDA.@allowscalar, which obviously slows the code down by a lot. Is there a faster way to do this?

More details on my implementation: Say term(x, y) = var(x, y, h). I define a field:
var_h = Field{Center, Center, Center}(grid)
and then add it to the auxiliary fields of the model (I'm not sure if this is the best way to store this field). I then have a callback that updates it:

S_grad = @at (Center, Center, Center) Field(∂z(model.tracers.S))
compute!(S_grad)
    @inbounds CUDA.@allowscalar begin
        ind_max_grad = argmin(S_grad[:, :, 1:end-1], dims=3)
        # Compute var
        var = .....
        # Compute var at MLD
        sim.model.auxiliary_fields.var_D.data .= var.data[ind_max_grad]

and then the term is computed elsewhere and saved to a netCDF. The indexing in the .= assignment is the slow-down, but I'm not sure how to speed this up. Also, I should note that var_D is a 3D field, but is repeated across the 3rd dimension since it is really a 2D field. I'm not sure if I can add a 2D field to auxiliary fields anyways Any help would be greatly appreciated!

[glwagner]

First you can try

max_dSdz = maximum(∂z(model.tracers.S), dims=3)

or

max_dSdz = maximum(Field(∂z(model.tracers.S)), dims=3)

then max_dSdz should be a 2D field.

If you want to form an object that can be repeatedly compute!'d then I think you want to use Reduction with maximum! (see the docstring for Reduction):

max_dSdz_reduction = Reduction(maximum!, ∂z(model.tracers.S), dims=3)
max_dSdz = Field(max_dSdz_reduction)
compute!(max_dSdz)

To go further with custom kernels I suggest checking out the implementation of MixedLayerDepthOperand and MixedLayerDepthField here:

https://github.com/CliMA/ClimaOcean.jl/blob/main/src/Diagnostics/mixed_layer_depth.jl