Add GPU section to model_implementation.md #253
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
johnryantaylor
approved these changes
Apr 3, 2025
jagoosw
reviewed
Apr 3, 2025
docs/src/model_implementation.md
Outdated
| @@ -312,5 +314,33 @@ fig | |||
|
|
|||
| We can see in this that some phytoplankton sink to the bottom, and are both remineralized back into nutrients and stored in the sediment. | |||
|
|
|||
| ### Running on a GPU | |||
|
|
|||
| In order to run a BGC model on a GPU, the BGC model must first be `adapted`. After the definition of the BGC `struct`, we need to write: | |||
There was a problem hiding this comment.
It might be good to add a link to this explanation of what is going on here
jagoosw
reviewed
Apr 3, 2025
| adapt(to, bgc.optimal_temperature), | ||
| adapt(to, bgc.mortality_rate), | ||
| adapt(to, bgc.crowding_mortality_rate), | ||
| adapt(to, bgc.sinking_velocity)) |
There was a problem hiding this comment.
Thinking again I think that the sinking velocity can just be nothing in the adapted version since it will never be used inside a kernel (except when passed explitily anyway)
There was a problem hiding this comment.
Do you mean that I should write nothing in place of adapt(to, bgc.sinking_velocity)?
There was a problem hiding this comment.
I tried this but it didn't seem to work for me so have just left as is for now
There was a problem hiding this comment.
ERROR: InvalidIRError: compiling MethodInstance for Oceananigans.Models.NonhydrostaticModels.gpu_compute_Gc!(::KernelAbstractions.CompilerMetadata{…}, ::OffsetArrays.OffsetArray{…}, ::RectilinearGrid{…}, ::Nothing, ::Tuple{…}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to getindex)
Stacktrace:
[1] getindex
@ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/sum_of_arrays.jl:25
[2] Az_qᶜᶜᶠ
@ ~/.julia/packages/Oceananigans/2NE0o/src/Operators/products_between_fields_and_grid_metrics.jl:11
[3] advective_tracer_flux_z
@ ~/.julia/packages/Oceananigans/2NE0o/src/Advection/centered_advective_fluxes.jl:33
[4] _advective_tracer_flux_z
@ ~/.julia/packages/Oceananigans/2NE0o/src/Advection/tracer_advection_operators.jl:4
[5] δzᵃᵃᶜ
@ ~/.julia/packages/Oceananigans/2NE0o/src/Operators/difference_operators.jl:26
[6] div_Uc
@ ~/.julia/packages/Oceananigans/2NE0o/src/Advection/tracer_advection_operators.jl:31
[7] tracer_tendency
@ ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl:253
[8] gpu_compute_Gc!
@ ~/.julia/packages/KernelAbstractions/sWSE0/src/macros.jl:322
[9] gpu_compute_Gc!
@ ./none:0
Hint: catch this exception as `err` and call `code_typed(err; interactive = true)` to introspect the erronous code with Cthulhu.jl
Stacktrace:
[1] check_ir(job::GPUCompiler.CompilerJob{GPUCompiler.PTXCompilerTarget, CUDA.CUDACompilerParams}, args::LLVM.Module)
@ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/validation.jl:167
[2] macro expansion
@ ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:381 [inlined]
[3]
@ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/utils.jl:110
[4] emit_llvm
@ ~/.julia/packages/GPUCompiler/OGnEB/src/utils.jl:108 [inlined]
[5]
@ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:100
[6] codegen
@ ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:82 [inlined]
[7] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
@ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:79
[8] compile
@ ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:74 [inlined]
[9] #1171
@ ~/.julia/packages/CUDA/RQqFT/src/compiler/compilation.jl:255 [inlined]
[10] JuliaContext(f::CUDA.var"#1171#1174"{GPUCompiler.CompilerJob{…}}; kwargs::@Kwargs{})
@ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:34
[11] JuliaContext(f::Function)
@ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/driver.jl:25
[12] compile(job::GPUCompiler.CompilerJob)
@ CUDA ~/.julia/packages/CUDA/RQqFT/src/compiler/compilation.jl:254
[13] 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
[14] cached_compilation(cache::Dict{…}, src::Core.MethodInstance, cfg::GPUCompiler.CompilerConfig{…}, compiler::Function, linker::Function)
@ GPUCompiler ~/.julia/packages/GPUCompiler/OGnEB/src/execution.jl:151
[15] macro expansion
@ ~/.julia/packages/CUDA/RQqFT/src/compiler/execution.jl:373 [inlined]
[16] macro expansion
@ ./lock.jl:273 [inlined]
[17] cufunction(f::typeof(Oceananigans.Models.NonhydrostaticModels.gpu_compute_Gc!), tt::Type{…}; kwargs::@Kwargs{…})
@ CUDA ~/.julia/packages/CUDA/RQqFT/src/compiler/execution.jl:368
[18] macro expansion
@ ~/.julia/packages/CUDA/RQqFT/src/compiler/execution.jl:112 [inlined]
[19] (::KernelAbstractions.Kernel{…})(::Field{…}, ::Vararg{…}; ndrange::Nothing, workgroupsize::Nothing)
@ CUDA.CUDAKernels ~/.julia/packages/CUDA/RQqFT/src/CUDAKernels.jl:103
[20] (::KernelAbstractions.Kernel{…})(::Field{…}, ::Vararg{…})
@ CUDA.CUDAKernels ~/.julia/packages/CUDA/RQqFT/src/CUDAKernels.jl:89
[21] _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
[22] _launch!
@ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/kernel_launching.jl:275 [inlined]
[23] launch!
@ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/kernel_launching.jl:258 [inlined]
[24] 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
[25] compute_interior_tendency_contributions!
@ ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:57 [inlined]
[26] compute_tendencies!(model::NonhydrostaticModel{…}, callbacks::Vector{…})
@ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:35
[27] #apply_regionally!#56
@ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/multi_region_transformation.jl:121 [inlined]
[28] apply_regionally!
@ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/multi_region_transformation.jl:118 [inlined]
[29] macro expansion
@ ~/.julia/packages/Oceananigans/2NE0o/src/Utils/multi_region_transformation.jl:206 [inlined]
[30] 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
[31] update_state! (repeats 2 times)
@ ~/.julia/packages/Oceananigans/2NE0o/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:20 [inlined]
[32] initialize!(sim::Simulation{…})
@ Oceananigans.Simulations ~/.julia/packages/Oceananigans/2NE0o/src/Simulations/run.jl:208
[33] time_step!(sim::Simulation{…})
@ Oceananigans.Simulations ~/.julia/packages/Oceananigans/2NE0o/src/Simulations/run.jl:138
[34] run!(sim::Simulation{…}; pickup::Bool)
@ Oceananigans.Simulations ~/.julia/packages/Oceananigans/2NE0o/src/Simulations/run.jl:105
[35] run!(sim::Simulation{…})
@ Oceananigans.Simulations ~/.julia/packages/Oceananigans/2NE0o/src/Simulations/run.jl:92
[36] top-level scope
@ ~/Documents/Development/LagrangianParticleModelDevelopment/oceanbiome_trimmed_example.jl:54
Some type information was truncated. Use `show(err)` to see complete types.
For clarity, to get this I ran the code from #251 but changed the adapt line to Adapt.adapt_structure(to, bgc::NutrientPhytoplankton) = NutrientPhytoplankton(nothing)
There was a problem hiding this comment.
Oh sorry I just reaslised it gets the bgc velocities here:
https://github.com/CliMA/Oceananigans.jl/blob/224ba94e86f29af457bbf86430c5498f2433ec06/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl#L248C33-L248C62
I thought they were got from somewhere else
jagoosw
reviewed
Apr 3, 2025
docs/src/model_implementation.md
Outdated
| adapt(to, bgc.sinking_velocity)) | ||
| ``` | ||
|
|
||
| Also, in order for `ScaleNegativeTracers` to work on a GPU, we must add `grid` as one of the input arguments. We replace the definition of `negative_tracer_scaling` with: |
There was a problem hiding this comment.
Suggested change
| Also, in order for `ScaleNegativeTracers` to work on a GPU, we must add `grid` as one of the input arguments. We replace the definition of `negative_tracer_scaling` with: | |
| We can add `grid` as the second argument for `ScaleNegativeTracers` so that it automatically works on a GPU. We replace the definition of `negative_tracer_scaling` with: |
jagoosw
reviewed
Apr 3, 2025
docs/src/model_implementation.md
Outdated
| negative_tracer_scaling = ScaleNegativeTracers((:N, :P), grid) | ||
| ``` | ||
|
|
||
| We can then add `GPU()` to the definition of `grid` in the usual way, and the column model above will be able to run on a GPU. |
There was a problem hiding this comment.
Suggested change
| We can then add `GPU()` to the definition of `grid` in the usual way, and the column model above will be able to run on a GPU. |
|
The CI isn't going to run until the maths computers and kelp start working again |
|
Is it okay if I merge main here? |
|
Yeah that's fine! |
jagoosw
reviewed
Apr 8, 2025
jagoosw
reviewed
Apr 8, 2025
|
Sorry, I meant to say, to get the docs to run we need to include the |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
As discussed in #251, I've written a paragraph about running BGC models on a GPU in model_implementation.md. Any feedback would be appreciated!