Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Fix instabilities in the launch! function#5745

Merged
simone-silvestri merged 21 commits into
mainfrom
ss/fix-static-size-instability
Jul 1, 2026
Merged

Fix instabilities in the launch! function#5745
simone-silvestri merged 21 commits into
mainfrom
ss/fix-static-size-instability

Conversation

@simone-silvestri

@simone-silvestri simone-silvestri commented Jun 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

this PR is a clean extract of the workaround in PR #5729 to make sure the launch! function is completely type stable.
We always derive our StaticSize from grid size and halo size information, which then need to be compile-time constants.
By adding a GridSize type parameter to the grid we can always infer the StaticSize of kernel and then squash the type instability.

This is not possible for kernels that launch on the active cells map, but in that case, since the kernel is always 1D and the index is computed already with length(map) inside the @index macro, we can safely use a DynamicSize

MWE:

using Oceananigans
using Oceananigans.Utils: work_layout, interior_work_layout, configure_kernel
using KernelAbstractions.NDIteration: StaticSize
using Test

grid = RectilinearGrid(size=(8, 6, 4), halo=(2, 2, 2), extent=(1, 1, 1))
wg,  ws  = @inferred work_layout(grid,:xyz, ())
wg2, ws2 = @inferred interior_work_layout(grid, :xyz, (Center(), Center(), Center()))

on main:

julia> wg,  ws  = @inferred work_layout(grid,:xyz, ())
ERROR: return type Tuple{StaticSize{(16, 16)}, StaticSize{(8, 6, 4)}} does not match inferred return type Tuple{StaticSize, StaticSize}
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:44
 [2] top-level scope
   @ REPL[15]:1

julia> wg2, ws2 = @inferred interior_work_layout(grid, :xyz, (Center(), Center(), Center()))
ERROR: return type Tuple{StaticSize{(16, 16)}, Oceananigans.Utils.OffsetStaticSize{(1:8, 1:6, 1:4)}} does not match inferred return type Tuple{StaticSize, Oceananigans.Utils.OffsetStaticSize}
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:44
 [2] top-level scope
   @ REPL[16]:1

on this PR

julia> wg,  ws  = @inferred work_layout(grid, Val(:xyz), ())
(StaticSize{(16, 16)}(), StaticSize{(8, 6, 4)}())

julia> wg2, ws2 = @inferred interior_work_layout(grid, Val(:xyz), (Center(), Center(), Center()))
(StaticSize{(16, 16)}(), Oceananigans.Utils.OffsetStaticSize{(1:8, 1:6, 1:4)}())

simone-silvestri
simone-silvestri commented Jun 30, 2026
View reviewed changes
Comment thread src/Grids/latitude_longitude_grid.jl Outdated
simone-silvestri
simone-silvestri commented Jun 30, 2026
View reviewed changes
Comment thread src/Grids/orthogonal_spherical_shell_grid.jl Outdated
simone-silvestri
simone-silvestri commented Jun 30, 2026
View reviewed changes
Comment thread src/Grids/rectilinear_grid.jl Outdated
@simone-silvestri

simone-silvestri commented Jun 30, 2026

Copy link
Copy Markdown
Collaborator Author

We are starting to see the light on CPUs:

[2026/06/30 17:09:00.053] INFO    hydrostatic \| CPU \| flat : 479.0 KiB (≤ 2.7 MiB)
[2026/06/30 17:09:51.263] INFO    hydrostatic \| CPU \| immersed : 497.2 KiB (≤ 2.9 MiB)
[2026/06/30 17:11:04.862] INFO    hydrostatic \| CPU \| active_immersed : 329.9 KiB (≤ 3.0 MiB)
[2026/06/30 17:11:21.403] INFO    nonhydrostatic \| CPU \| flat : 826.2 KiB (≤ 2.0 MiB)

GPUs are still quite high, I think it has to do with our custom convert. However, when CPU is up to speed completely I will move to squashing also the GPU allocations

@giordano

giordano commented Jun 30, 2026

Copy link
Copy Markdown
Collaborator

on this PR

julia> wg,  ws  = @inferred work_layout(grid, Val(:xyz), ())
(StaticSize{(16, 16)}(), StaticSize{(8, 6, 4)}())

julia> wg2, ws2 = @inferred interior_work_layout(grid, Val(:xyz), (Center(), Center(), Center()))
(StaticSize{(16, 16)}(), Oceananigans.Utils.OffsetStaticSize{(1:8, 1:6, 1:4)}())

Add that as a test? Ensuring full inference on these things is quite important for high performance workloads.

Also, update the allocations upper bounds? This is a significant drop!

giordano
giordano reviewed Jul 1, 2026
View reviewed changes
Comment thread src/Grids/rectilinear_grid.jl Outdated
Δyᵃᶠᵃ, Δyᵃᶜᵃ, yᵃᶠᵃ, yᵃᶜᵃ, z)
end

@generated function Base.size(grid::RectilinearGrid)

@giordano giordano Jul 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@vchuravy are these generated's a good idea?

@codecov

codecov Bot commented Jul 1, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 81.57895% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.92%. Comparing base (1fa302b) to head (23f7bb6).
⚠️ Report is 18 commits behind head on main.

Files with missing lines Patch % Lines
src/Utils/kernel_launching.jl 76.59% 11 Missing ⚠️
src/Grids/abstract_grid.jl 66.66% 2 Missing ⚠️
...lShellGrids/right_face_folded_kernel_parameters.jl 66.66% 1 Missing ⚠️
Additional details and impacted files 
@@            Coverage Diff             @@
##             main    #5745      +/-   ##
==========================================
- Coverage   73.96%   73.92%   -0.05%     
==========================================
  Files         418      419       +1     
  Lines       25279    25347      +68     
==========================================
+ Hits        18698    18737      +39     
- Misses       6581     6610      +29
Flag Coverage Δ
buildkite 68.58% <77.33%> (-0.03%) ⬇️
julia 68.58% <77.33%> (-0.03%) ⬇️
reactant_1 6.05% <41.33%> (+0.03%) ⬆️
reactant_2 10.75% <52.00%> (-0.03%) ⬇️
reactant_3 9.16% <45.33%> (-0.01%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@simone-silvestri simone-silvestri merged commit 9ab5b28 into main Jul 1, 2026
90 of 91 checks passed
@simone-silvestri simone-silvestri deleted the ss/fix-static-size-instability branch July 1, 2026 19:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@giordano giordano giordano approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@simone-silvestri @giordano