GPU attempting to create too many threads

[RicksterB57]

I have been attempting to run a 500x500x200 meter model on a 256x256x128 grid on a GPU and have been encountering this error:
ERROR: LoadError: Number of threads per block exceeds kernel limit (896 > 768).
I have been able to run this model on the same GPU in the past. Any tips on limiting the amount of threads that the model tries to open during the simulation?

[glwagner]

Can you show us the code you are trying to run and also let us know what GPU you are using?

[RicksterB57]

I am using a NVIDIA RTX A5000 GPU. As for the code:

Lx = 500meters;
Ly = 500meters;
Lz = 300meters;
Nx = Ny = 256     # number of points in each of horizontal directions
Nz = 128        # number of points in the vertical direction

arch = GPU()
grid = RectilinearGrid(arch, size = (Nx, Nx, Nz),
                          x = (0, Lx),
                          y = (0, Ly),
                          z = (-Lz, 0));

buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(thermal_expansion = 2e-4,
                haline_contraction = 8e-4));
# Constants
ρ₀ = 1026 # kg m⁻³, average density at the surface of the world ocean
ρₐ = 1.225 # average density of air at sea level
cᴰ = 2.5e-3 # drag coefficient
cᴾ = 3991.0 # specific heat capacity of seawater

time_index = n = Ref(1);

# Defining the functions to update the heat flux and the wind drag every hour
@inline function surface_heat_flux(i, j, grid, clock, model_fields, params)
    n_reference, ρ₀, cᴾ, Qs, Ql = params;
    n = n_reference[];
    @inbounds begin
        Q_s = Qs[1,1,n]
        Q_l = Ql[1,1,n]
    end
    Q = Q_s + Q_l;
    Qᵀ = Q / (ρ₀ * cᴾ);
    return Qᵀ
end

@inline function wind_drag(i, j, grid, clock, model_fields, params)
    n_reference, ρₐ, ρ₀, cᴰ, u10 = params
    n = n_reference[]
    @inbounds begin 
        u_10 = u10[1,1,n]
    end
    Qᵘ = -ρₐ / ρ₀ * cᴰ * u_10 * abs(u_10);
    return Qᵘ
end

function update_time_ind(sim)
    n = Int(floor(time(simulation) / (3600))) + 1
    time_index[] = n
    return nothing
end

time_index_callback = Callback(update_time_ind)

# Boundary Conditions
# if running on gpu need to change the arrays for wind speed and fluxes to the gpu arrays
u_top_bc = FluxBoundaryCondition(wind_drag, discrete_form=true, parameters=(time_index, ρₐ, ρ₀, cᴰ, u10_gpu))
T_top_bc = FluxBoundaryCondition(surface_heat_flux, discrete_form=true, parameters=(time_index, ρ₀, cᴾ, latent_gpu, sensible_gpu))
u_bcs = FieldBoundaryConditions(top = u_top_bc );
T_bcs = FieldBoundaryConditions(top = T_top_bc);

# Define your model here
model = NonhydrostaticModel(; grid, buoyancy, 
                            advection = UpwindBiasedFifthOrder(),
                            timestepper = :RungeKutta3,
                            tracers = (:T, :S),
                            coriolis = FPlane(rotation_rate=7.2921e-5, latitude=72),
                            closure = AnisotropicMinimumDissipation(),
                            boundary_conditions = (u=u_bcs, T=T_bcs))

# Generating random noise for the model in the velocities
Ξ(z) = randn() * z / model.grid.Lz * (1 + z / model.grid.Lz) 
uᵢ(x,y,z) = 1e-6 * Ξ(z)

# Setting the initial conditions for the model
set!(model, u=uᵢ, T=(x, y, z) -> T_ic[findfirst(≥(z), grid.zᵃᵃᶜ)], S=(x, y, z) -> S_ic[findfirst(≥(z), grid.zᵃᵃᶜ)])

simulation = Simulation(model, Δt=1.0, stop_time=5day, verbose=true)
wizard = TimeStepWizard(cfl=0.7, max_change=1.1, max_Δt=5.0)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))


using Oceananigans.Diagnostics
# Print a progress message
progress_message(sim) = @printf("Iteration: %04d, time: %s, Δt: %s, max(|u|) = %.1e ms⁻¹, wall time: %s\n",
                                iteration(sim), prettytime(sim), prettytime(sim.Δt),
                                maximum(abs, sim.model.velocities.u), prettytime(sim.run_wall_time))

add_callback!(simulation, progress_message, IterationInterval(30))
add_callback!(simulation, time_index_callback)


filename = norse_model_output_filename *"storm_model";

simulation.output_writers[:slices] =
    JLD2OutputWriter(model, merge(model.velocities, model.tracers),
                     filename = filename * ".jld2",
                     indices = (:, grid.Ny/2, :),
                     schedule = TimeInterval(20minute),
                     overwrite_existing = true)


#Running the simulation
run!(simulation)

The u10_gpu, latent_gpu, and sensible_gpu are 1d CuArrays that I create to update the boundary conditions. As I mentioned I have run this model before without any errors about too many threads.

[glwagner]

As I mentioned I have run this model before without any errors about too many threads.

What has changed since the code worked?

[glwagner]

I can't run the code to confirm the error because u10_gpu is not defined.

When do you encounter the error? Is it when you build the model, or run the simulation? Or some other time? I'm not sure if the error message gives more information about the specific kernel that is the problem, but it would be helpful to know.

It would also be good to confirm that if you reduce the size of the problem, it does run, as a sanity check.

The information that the code worked previously somehow is frustrating but I don't know how to act on that. If you can pinpoint the different versions of Oceananigans that you were using, we can possibly look into what has changed. That said, I don't think very much at all has changed about how we launch kernels recently. But if the last time you tried was with a very very old version of Oceanangians, there may be some information we can fish out of that.

I tried to investigate the maximum threads issue. It seems like the maximum number of threads per block is usually 1024. Possibly, it could help if you can query your device and establish indeed that the maximum number of threads is 768: https://forums.developer.nvidia.com/t/maximum-number-of-threads-on-thread-block/46392. If it is not, there is a clue of something we might look into (eg perhaps there is an issue with a CUDA library or something).

Any tips on limiting the amount of threads that the model tries to open during the simulation?

I am not sure but I think this may be determined automatically by KernelAbstractions. There may be a way to control it --- it would be interesting to know. No method for controlling it is exposed in the Oceananigans user interface though.

[glwagner]

After a little bit of digging, I think that the threads are determined here:

https://github.com/JuliaGPU/CUDA.jl/blob/a085bbb3d7856dfa929e6cdae04a146a259a2044/src/CUDAKernels.jl#L89

Also the workgroup for 3D kernels is set here in Oceananigans:

Oceananigans.jl/src/Utils/kernel_launching.jl

Line 101 in 6c40d7e

(16, 16)

We might be able to dig deeper and figure out where 896 comes from. I'm curious if the error goes away for a smaller problem (indicating that this issue depends on problem size). I'm also curious if the nuber changes for a different problem, ie if you try 257x256x128. A little experimentation might yield some important clues.

[RicksterB57]

The problem occurs when the simulation right after the simulation begins. In terms of the changes that I have made to the code the only real change I have made is that I have changed the sign of the arrays that I input into my boundary condition functions. I can try experimenting with the size of the domain and grid model and closely examine the error messages that pop up.