Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

‘Implementing new models’ example doesn’t run on GPU #251

Open
emmabeniston opened this issue Mar 31, 2025 · 5 comments
Open

‘Implementing new models’ example doesn’t run on GPU #251

emmabeniston opened this issue Mar 31, 2025 · 5 comments

Comments

@emmabeniston
Copy link

Hello,

I’ve been trying to run the column model example from the ‘Implementing new models’ documentation but haven't been able to get it to run on a GPU.

This is a trimmed down version of the example which produces an error:

using OceanBioME, Oceananigans
using Oceananigans.Biogeochemistry: AbstractContinuousFormBiogeochemistry
using Oceananigans.Units
import Oceananigans.Biogeochemistry: required_biogeochemical_tracers,
                                            required_biogeochemical_auxiliary_fields,
                                            biogeochemical_drift_velocity
using CUDA
using Oceananigans.Fields: ZeroField, ConstantField

@kwdef struct NutrientPhytoplankton{W} <: AbstractContinuousFormBiogeochemistry
    sinking_velocity :: W = 2 / day                   
end

required_biogeochemical_tracers(::NutrientPhytoplankton) = (:P, )

biogeochemical_drift_velocity(bgc::NutrientPhytoplankton, ::Val{:P}) =
    (u = ZeroField(), v = ZeroField(), w = bgc.sinking_velocity)

grid = RectilinearGrid(GPU(), topology = (Flat, Flat, Bounded), size = (32, ), x = 1, y = 1, z = (-100, 0))

sinking_velocity = ZFaceField(grid)
w_sink(z) = 2 / day * tanh(z / 5)
set!(sinking_velocity, w_sink)

biogeochemistry = Biogeochemistry(NutrientPhytoplankton(; sinking_velocity))

κ = 0.0001

model = NonhydrostaticModel(; grid,
                             biogeochemistry,
                             closure = ScalarDiffusivity(ν = κ; κ))

set!(model, P = 0.01)

const years = 365days

simulation = Simulation(model, Δt = 9minutes, stop_time = 1years)

simulation.output_writers[:tracers] = JLD2OutputWriter(model, model.tracers,
                                                      filename = "column_np.jld2",
                                                      schedule = TimeInterval(1day),
                                                      overwrite_existing = true)

run!(simulation)

And this is the error message I get:

ERROR: GPU compilation of MethodInstance for Oceananigans.Models.NonhydrostaticModels.gpu_compute_Gc!(::KernelAbstractions.CompilerMetadata{…}, ::OffsetArrays.OffsetArray{…}, ::RectilinearGrid{…}, ::Nothing, ::Tuple{…}) failed
KernelError: passing non-bitstype argument

Argument 6 to your kernel function is of type Tuple{Val{1}, Val{:P}, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}, ScalarDiffusivity{Oceananigans.TurbulenceClosures.ExplicitTimeDiscretization, Oceananigans.TurbulenceClosures.ThreeDimensionalFormulation, @NamedTuple{P::Float64}, Float64, @NamedTuple{P::Float64}}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, Nothing, NutrientPhytoplankton{Field{Center, Center, Face, Nothing, RectilinearGrid{Float64, Flat, Flat, Bounded, Oceananigans.Grids.StaticVerticalDiscretization{OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Float64, Float64}, Float64, Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, GPU{CUDABackend}}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.DeviceMemory}}, Float64, FieldBoundaryConditions{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}, @NamedTuple{velocities::@NamedTuple{u::ZeroField{Int64, 3}, v::ZeroField{Int64, 3}, w::ZeroField{Int64, 3}}, tracers::@NamedTuple{P::ZeroField{Int64, 3}}}, @NamedTuple{u::OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}, v::OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}, w::OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}}, @NamedTuple{P::OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}}, @NamedTuple{}, Nothing, @NamedTuple{time::Float64, last_Δt::Float64, last_stage_Δt::Float64, iteration::Int64, stage::Int64}, typeof(Oceananigans.Forcings.zeroforcing)}, which is not a bitstype:
  .7 is of type NutrientPhytoplankton{Field{Center, Center, Face, Nothing, RectilinearGrid{Float64, Flat, Flat, Bounded, Oceananigans.Grids.StaticVerticalDiscretization{OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Float64, Float64}, Float64, Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, GPU{CUDABackend}}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.DeviceMemory}}, Float64, FieldBoundaryConditions{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}} which is not isbits.
    .sinking_velocity is of type Field{Center, Center, Face, Nothing, RectilinearGrid{Float64, Flat, Flat, Bounded, Oceananigans.Grids.StaticVerticalDiscretization{OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Float64, Float64}, Float64, Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, GPU{CUDABackend}}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.DeviceMemory}}, Float64, FieldBoundaryConditions{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}} which is not isbits.
      .data is of type OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.DeviceMemory}} which is not isbits.
        .parent is of type CuArray{Float64, 3, CUDA.DeviceMemory} which is not isbits.
          .data is of type GPUArrays.DataRef{CUDA.Managed{CUDA.DeviceMemory}} which is not isbits.
            .rc is of type GPUArrays.RefCounted{CUDA.Managed{CUDA.DeviceMemory}} which is not isbits.
              .obj is of type CUDA.Managed{CUDA.DeviceMemory} which is not isbits.
                .stream is of type CuStream which is not isbits.
                  .ctx is of type Union{Nothing, CuContext} which is not isbits.
              .finalizer is of type Any which is not isbits.
              .count is of type Base.Threads.Atomic{Int64} which is not isbits.
      .boundary_conditions is of type FieldBoundaryConditions{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}} which is not isbits.


Only bitstypes, which are "plain data" types that are immutable
and contain no references to other values, can be used in GPU kernels.
For more information, see the `Base.isbitstype` function.

Stacktrace:
  [1] check_invocation(job::GPUCompiler.CompilerJob)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/validation.jl:108
  [2] 
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:90
  [3] codegen
    @ ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:82 [inlined]
  [4] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:79
  [5] compile
    @ ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:74 [inlined]
  [6] #1171
    @ ~/.julia/packages/CUDA/RQqFT/src/compiler/compilation.jl:255 [inlined]
  [7] JuliaContext(f::CUDA.var"#1171#1174"{GPUCompiler.CompilerJob{…}}; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:34
  [8] JuliaContext(f::Function)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:25
  [9] compile(job::GPUCompiler.CompilerJob)
    @ CUDA ~/.julia/packages/CUDA/RQqFT/src/compiler/compilation.jl:254
 [10] actual_compilation(cache::Dict{…}, src::Core.MethodInstance, world::UInt64, cfg::GPUCompiler.CompilerConfig{…}, compiler::typeof(CUDA.compile), linker::typeof(CUDA.link))
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/execution.jl:237
 [11] cached_compilation(cache::Dict{…}, src::Core.MethodInstance, cfg::GPUCompiler.CompilerConfig{…}, compiler::Function, linker::Function)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/execution.jl:151
 [12] macro expansion
    @ ~/.julia/packages/CUDA/RQqFT/src/compiler/execution.jl:373 [inlined]
 [13] macro expansion
    @ ./lock.jl:273 [inlined]
 [14] cufunction(f::typeof(Oceananigans.Models.NonhydrostaticModels.gpu_compute_Gc!), tt::Type{…}; kwargs::@Kwargs{…})
    @ CUDA ~/.julia/packages/CUDA/RQqFT/src/compiler/execution.jl:368
 [15] macro expansion
    @ ~/.julia/packages/CUDA/RQqFT/src/compiler/execution.jl:112 [inlined]
 [16] (::KernelAbstractions.Kernel{…})(::Field{…}, ::Vararg{…}; ndrange::Nothing, workgroupsize::Nothing)
    @ CUDA.CUDAKernels ~/.julia/packages/CUDA/RQqFT/src/CUDAKernels.jl:103
 [17] (::KernelAbstractions.Kernel{…})(::Field{…}, ::Vararg{…})
    @ CUDA.CUDAKernels ~/.julia/packages/CUDA/RQqFT/src/CUDAKernels.jl:89
 [18] _launch!(::GPU{…}, ::RectilinearGrid{…}, ::Symbol, ::Function, ::Field{…}, ::RectilinearGrid{…}, ::Vararg{…}; exclude_periphery::Bool, reduced_dimensions::Tuple{}, active_cells_map::Nothing)
    @ Oceananigans.Utils ~/.julia/packages/Oceananigans/2NE0o/src/Utils/kernel_launching.jl:298
 [19] _launch!
    @ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/kernel_launching.jl:275 [inlined]
 [20] launch!
    @ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/kernel_launching.jl:258 [inlined]
 [21] compute_interior_tendency_contributions!(model::NonhydrostaticModel{…}, kernel_parameters::Symbol; active_cells_map::Nothing)
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:133
 [22] compute_interior_tendency_contributions!
    @ ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:57 [inlined]
 [23] compute_tendencies!(model::NonhydrostaticModel{…}, callbacks::Vector{…})
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:35
 [24] #apply_regionally!#56
    @ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/multi_region_transformation.jl:121 [inlined]
 [25] apply_regionally!
    @ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/multi_region_transformation.jl:118 [inlined]
 [26] macro expansion
    @ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/multi_region_transformation.jl:206 [inlined]
 [27] update_state!(model::NonhydrostaticModel{…}, callbacks::Vector{…}; compute_tendencies::Bool)
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:53
 [28] update_state! (repeats 2 times)
    @ ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:20 [inlined]
 [29] initialize!(sim::Simulation{…})
    @ Oceananigans.Simulations ~/.julia/packages/Oceananigans/2NE0o/src/Simulations/run.jl:208
 [30] time_step!(sim::Simulation{…})
    @ Oceananigans.Simulations ~/.julia/packages/Oceananigans/2NE0o/src/Simulations/run.jl:138
 [31] run!(sim::Simulation{…}; pickup::Bool)
    @ Oceananigans.Simulations ~/.julia/packages/Oceananigans/2NE0o/src/Simulations/run.jl:105
 [32] run!(sim::Simulation{…})
    @ Oceananigans.Simulations ~/.julia/packages/Oceananigans/2NE0o/src/Simulations/run.jl:92
 [33] top-level scope
    @ ~/Documents/LagrangianParticleModelDevelopment/oceanbiome_trimmed_example.jl:46
Some type information was truncated. Use `show(err)` to see complete types.

I was wondering what needs to be changed in the example to get it to run on a GPU? The problem seems to be to do with sinking_velocity being a Field.

Thank you
@jagoosw jagoosw mentioned this issue Mar 31, 2025
Open
@jagoosw
Copy link
Collaborator

jagoosw commented Mar 31, 2025

Hi Emma,

The problem is that since the BGC model is being put into GPU kernels, it needs to be adapted and this example script doesn't have an adapt method. You need to write:

using Adapt

import Adapt: adapt_structure

Adapt.adapt_structure(to, bgc::NutrientPhytoplankton) = NutrientPhytoplankton(adapt(to, bgc.sinking_velocity))

These docs kind of explain why adapts are needed.

@jagoosw
Copy link
Collaborator

jagoosw commented Mar 31, 2025

Perhaps we should add this method to the documentation page as a little bit at the end about running on GPU.

You can do this if you want or I can add it to #252

@emmabeniston
Copy link
Author

Thanks Jago, that's sorted it! I'm happy to try and write a paragraph for the documentation.

While we're here, I can't get the example to run on the GPU if I include ScaleNegativeTracers. If I add in

negative_tracer_scaling = ScaleNegativeTracers((:P, ))

biogeochemistry = Biogeochemistry(NutrientPhytoplankton(; sinking_velocity),
                                    modifiers = negative_tracer_scaling)

then I get a similar error to before:

ERROR: GPU compilation of MethodInstance for OceanBioME.gpu_scale_for_negs!(::KernelAbstractions.CompilerMetadata{…}, ::Tuple{…}, ::Vector{…}, ::Float64) failed
KernelError: passing non-bitstype argument

Argument 4 to your kernel function is of type Vector{Float64}, which is not a bitstype:
  .ref is of type MemoryRef{Float64} which is not isbits.
    .mem is of type Memory{Float64} which is not isbits.


Only bitstypes, which are "plain data" types that are immutable
and contain no references to other values, can be used in GPU kernels.
For more information, see the `Base.isbitstype` function.

Stacktrace:
  [1] check_invocation(job::GPUCompiler.CompilerJob)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/validation.jl:108
  [2] 
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:90
  [3] codegen
    @ ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:82 [inlined]
  [4] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:79
  [5] compile
    @ ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:74 [inlined]
  [6] #1171
    @ ~/.julia/packages/CUDA/RQqFT/src/compiler/compilation.jl:255 [inlined]
  [7] JuliaContext(f::CUDA.var"#1171#1174"{GPUCompiler.CompilerJob{…}}; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:34
  [8] JuliaContext(f::Function)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:25
  [9] compile(job::GPUCompiler.CompilerJob)
    @ CUDA ~/.julia/packages/CUDA/RQqFT/src/compiler/compilation.jl:254
 [10] actual_compilation(cache::Dict{…}, src::Core.MethodInstance, world::UInt64, cfg::GPUCompiler.CompilerConfig{…}, compiler::typeof(CUDA.compile), linker::typeof(CUDA.link))
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/execution.jl:237
 [11] cached_compilation(cache::Dict{…}, src::Core.MethodInstance, cfg::GPUCompiler.CompilerConfig{…}, compiler::Function, linker::Function)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/execution.jl:151
 [12] macro expansion
    @ ~/.julia/packages/CUDA/RQqFT/src/compiler/execution.jl:373 [inlined]
 [13] macro expansion
    @ ./lock.jl:273 [inlined]
 [14] cufunction(f::typeof(OceanBioME.gpu_scale_for_negs!), tt::Type{…}; kwargs::@Kwargs{…})
    @ CUDA ~/.julia/packages/CUDA/RQqFT/src/compiler/execution.jl:368
 [15] macro expansion
    @ ~/.julia/packages/CUDA/RQqFT/src/compiler/execution.jl:112 [inlined]
 [16] (::KernelAbstractions.Kernel{…})(::Tuple{…}, ::Vararg{…}; ndrange::Nothing, workgroupsize::Nothing)
    @ CUDA.CUDAKernels ~/.julia/packages/CUDA/RQqFT/src/CUDAKernels.jl:103
 [17] (::KernelAbstractions.Kernel{…})(::Tuple{…}, ::Vararg{…})
    @ CUDA.CUDAKernels ~/.julia/packages/CUDA/RQqFT/src/CUDAKernels.jl:89
 [18] update_biogeochemical_state!(model::NonhydrostaticModel{…}, scale::ScaleNegativeTracers{…})
    @ OceanBioME ~/.julia/packages/OceanBioME/aQK5Q/src/Utils/negative_tracers.jl:138
 [19] update_biogeochemical_state!
    @ ~/.julia/packages/OceanBioME/aQK5Q/src/OceanBioME.jl:153 [inlined]
 [20] update_state!(model::NonhydrostaticModel{…}, callbacks::Vector{…}; compute_tendencies::Bool)
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:51
 [21] update_state!
    @ ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:20 [inlined]
 [22] NonhydrostaticModel(; grid::RectilinearGrid{…}, clock::Clock{…}, advection::Centered{…}, buoyancy::Nothing, coriolis::Nothing, stokes_drift::Nothing, forcing::@NamedTuple{}, closure::ScalarDiffusivity{…}, boundary_conditions::@NamedTuple{}, tracers::Tuple{}, timestepper::Symbol, background_fields::@NamedTuple{}, particles::Nothing, biogeochemistry::OceanBioME.ContinuousBiogeochemistry{…}, velocities::Nothing, hydrostatic_pressure_anomaly::Oceananigans.Models.NonhydrostaticModels.DefaultHydrostaticPressureAnomaly, nonhydrostatic_pressure::Field{…}, diffusivity_fields::Nothing, pressure_solver::Nothing, auxiliary_fields::@NamedTuple{})
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl:236
 [23] top-level scope
    @ ~/Documents/LagrangianParticleModelDevelopment/oceanbiome_trimmed_example.jl:43
Some type information was truncated. Use `show(err)` to see complete types.

I tried doing something similar with adapt by copying

https://github.com/OceanBioME/OceanBioME.jl/blob/f49ceeb823480c9a14900099a03fbeb0a4f5017b/src/Utils/negative_tracers.jl#L46C1-L50C1

but that didn't seem to work. Is there something else I need to do to get ScaleNegativeTracers to work?

@jagoosw
Copy link
Collaborator

jagoosw commented Apr 1, 2025

Okay thank you, that sounds good!

For the ScaleNegativeTracers I think its because the default values for scalefactor are a vector, if you do ScaleNegativeTracers((:P, ), grid) instead then it will automatically make it a CuArray.

@emmabeniston
Copy link
Author

Ahh okay, yes that all works now, thanks so much! I'll have a go at writing something for the documentation

@emmabeniston emmabeniston mentioned this issue Apr 3, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
2 participants
@jagoosw @emmabeniston