Saving a custom output field

[lgloege]

Hi! So I am trying to calculate the gas transfer velocity at each time step and save the average.

I keep getting this error message:

Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.

Which usually means I need to wrap my code in

CUDA.@allowscalar begin
    <your_code_here>
end

I did this in my compute_k_local(mod) function but am still getting this error. I've tried a few other things, but am not at a loss. My understanding is when I create a custom function it needs to return a Field, it can not be an CuArray or OffsetArray.

Any help would be greatly appreciated. Thanks!

Here is the relevant bits of my code:

coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation);

simulation = Simulation(coupled_model; Δt=8minutes, stop_time=10days);

function compute_k_local(mod)
    CUDA.@allowscalar begin
        T = mod.tracers.T
        u10 = coupled_model.interfaces.exchanger.exchange_atmosphere_state.u # u wind component
        v10 = coupled_model.interfaces.exchanger.exchange_atmosphere_state.v # v wind component
    
        SST = view(T, :, :, mod.grid.Nz) 
            
        A = 2116.8
        B = 136.25
        C = 4.7353
        D = 0.092307
        E = 0.0007555
        SCALE_FACTOR = 0.26
        
        k_local_field = Field{Center, Center, Nothing}(ocean.model.grid);
        
        k_local = SCALE_FACTOR .* (interior(u10).^2 .+ interior(v10).^2) .* (((A .- B.*interior(SST) .+ C.*interior(SST).^2 - D.*interior(SST).^3 + E.*interior(SST).^4)./660).^(-0.5))
    
        copyto!(interior(k_local_field), k_local);
    end
    
    return k_local_field 
end

base = merge(ocean.model.tracers, ocean.model.velocities);

funcs = (kw=compute_k_local);

ocean_outputs = merge(base, funcs);

ocean.output_writers[:surface] = JLD2Writer(ocean.model, ocean_outputs;
                                            schedule = TimeInterval(1days),
                                            filename = "global_1deg_30days_sim",
                                            indices = (:, :, grid.Nz),
                                            overwrite_existing = true)

run!(simulation)

[glwagner]

Hm, as a side note, you actually should never use @allowscalar unless you are literally extracting a scalar (eg a[1, 1, 1]).

Another note is that you should not create a field inside a function which computes a diagnostic. This will allocate memory for a field every time you need to compute the diagnostic. Instead you should allocate only the memory needed for the diagnostic, and then recompute the field as the simulation goes along.

You can implement this kind of workflow yourself using a function for sure. However, for your case, it seems you can use abstract operations directly and avoid the function approach altogether. To present a simplified example:

u10 = coupled_model.interfaces.exchanger.exchange_atmosphere_state.u # u wind component
v10 = coupled_model.interfaces.exchanger.exchange_atmosphere_state.v # v wind component

SCALE_FACTOR = 0.26       
k = SCALE_FACTOR * (u10^2 + v10^2)

ocean.output_writers[:surface] = JLD2Writer(ocean.model, (; k);
                                            schedule = TimeInterval(1days),
                                            filename = "global_1deg_30days_sim",
                                            indices = (:, :, grid.Nz),
                                            overwrite_existing = true)

If you would like to go about a more manual approach, you should look into using KernelFunctionOperation. This will allow you to implement a function-based approach to your computation which is compatible with Field.

As sort of a general rule, you should not find yourself using broadcasting for computations very often, except for the trivial ones like parent(a) .= parent(b) that are sometimes required for annoying forcing or boundary condition functions that have to be updated with callbacks.

I believe if you take this new approach you will not have issues with scalar operations.

[lgloege]

Thank @glwagner , I appreciate the quick reply and explanation. Took your advice and got rid of the functions, now my code looks like this

A,B,C,D,E = 2116.8, 136.25, 4.7353, 0.092307, 0.0007555
SCALE_FACTOR = 0.26

taux = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.x_momentum;
tauy = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.y_momentum;
u10 = coupled_model.interfaces.exchanger.exchange_atmosphere_state.u;
v10 = coupled_model.interfaces.exchanger.exchange_atmosphere_state.v;
SST = view(ocean.model.tracers.T, :, :, ocean.model.grid.Nz);
Sc = A - B*SST + C*SST^2 - D*SST^3 + E*SST^4;
ws = sqrt(u10^2 + v10^2);
kw = SCALE_FACTOR * ws^2 * (Sc/660)^(-0.5);


custom_outputs = (;kw, ws, Sc, taux, tauy);

ocean_outputs = merge(ocean.model.tracers, coupled_model.ocean.model.velocities, custom_outputs);

ocean.output_writers[:surface] = JLD2Writer(ocean.model, ocean_outputs;
                                            schedule = TimeInterval(1days),
                                            filename = "global_1deg_30days_sim",
                                            indices = (:, :, grid.Nz),
                                            overwrite_existing = true)

Now i'm getting this error

[ Info: Initializing simulation...
Kernel invocation uses too much parameter memory.
38.609 KiB exceeds the 31.996 KiB limit imposed by sm_89 [/](https://ood-bouchet.ycrc.yale.edu/) PTX v8.8.
...

Google leads me to believe I am using too much memory on the GPU? Don't think it's because i'm testing code in a jupyter notebook.

any thoughts on how to resolve this? or is my only solution to use multiple GPUs.

Thanks again for the help :)

[glwagner]

Ah yes the infamous parameter memory! This is not GPU memory, so its not solved by using multiple GPUs. It basically has to do with the complexity of the function you're trying to compile. There are some tips on how to solve this here:

https://clima.github.io/OceananigansDocumentation/stable/simulation_tips/#Complex-diagnostics-using-computed-Fields-may-not-work-on-GPUs

PS always use sqrt rather than ^(0.5). I see you have a ^(-0.5) that could be written 1 / sqrt instead

[glwagner]

Can you incorporate the number "660" into scale_factor? That will also help reduce parameter usage a little.

[glwagner]

Also I think we need a function for sea_surface_temperature(coupled_model) bc this method will not wokrk for all flux configurations. cc @simone-silvestri

[lgloege]

Thank you so much @glwagner !!

Here is what is working (or least not giving parameter memory errors yet)

sea_surface_temperature(coupled_model) = view(coupled_model.ocean.model.tracers.T, :, :, coupled_model.ocean.model.grid.Nz);

A,B,C,D,E = 2116.8, 136.25, 4.7353, 0.092307, 0.0007555
SCALE_FACTOR = 0.26 * sqrt(660) 

taux = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.x_momentum;
tauy = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.y_momentum;
u10 = coupled_model.interfaces.exchanger.exchange_atmosphere_state.u;
v10 = coupled_model.interfaces.exchanger.exchange_atmosphere_state.v;
SST = sea_surface_temperature(coupled_model)
Sc = (((E*SST - D)*SST + C)*SST - B)*SST + A; # via Horner's method
ws = sqrt(u10^2 + v10^2);
kw = SCALE_FACTOR * (u10^2 + v10^2) / sqrt(Sc);


custom_outputs = (;kw, ws, Sc, taux, tauy);

ocean_outputs = merge(ocean.model.tracers, coupled_model.ocean.model.velocities, custom_outputs);

ocean.output_writers[:surface] = JLD2Writer(ocean.model, ocean_outputs;
                                            schedule = TimeInterval(1days),
                                            filename = "global_1deg_30days_sim",
                                            indices = (:, :, grid.Nz),
                                            overwrite_existing = true)

• lumped the 660 in with my scale parameter,
• replaced ^(0.5) with sqrt()
• made a function for SST
• Changed Schmidt number (Sc) calculation to use Horner's method

I appreciate all your help on this!

Couple questions for my own edification

1. Why is sqrt() always preferred over ^(0.5) ?
2. Why do I need to make an SST function? What is the benefit?

[glwagner]

1. sqrt calls an optimized routine for square roots. ^(0.5) calls a more general routine for computing the power between a scalar and any number. The more general routine is slower.
2. Sorry for being unclear -- you don't need to define this function. The reason we (ClimaOcean more broadly) needs it is because the SST is not always the top grid point of the ocean model. It's possible to formulate air-sea fluxes in terms of a diagnostically computed / modeled SST. In such a formulation (which seems to make much more sense to me), the top grid point of the ocean model is taken more literally (the ocean temperature averaged over the top dz meters). The SST, which is the temperature of the top few mm or cm of the ocean, must be computed.

[glwagner]

(The SST is usually colder than the temperature averaged over the top 1-5 meters).

[glwagner]

Also amazing that you were able to do it without any intermediate fields or a manual KFO, pretty good!

[lgloege]

Thanks for all the help @glwagner

Is it possible to see the parameter memory error again if I increase the spatial resolution or does this error just mean the expression I am evaluating is too complex? Just trying to fully understand the meaning of the error so I know how to resolve it in the future

[glwagner]

It's not a function of resolution. The "memory" is being consumed by type information, not data. It has to do with the complexity of the objects in terms of their representation in assmbly code. This is independent of resolution. For example, a pointer to memory consumes X bytes, regardless of how big the block of memory it refers to.

[glwagner]

this may help: https://chatgpt.com/share/69173ad8-cb34-8002-b58f-c15ee7fb84dc