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

[WIP] Adapt Field to run on GPU #746

Closed
wants to merge 5 commits into from
Closed

[WIP] Adapt Field to run on GPU #746

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5f0fcdc
Adds adapt method for fields on the GPU
glwagner May 1, 2020
9f06329
Makes CoordinateBoundaryConditions immutable and implements looser di…
glwagner May 2, 2020
7472168
Grid isbits, no need to adapt
glwagner Jun 16, 2020
300d36a
Merge branch 'master' into glw/adapt-field
glwagner Aug 28, 2020
440b332
Remove dispatch on AbstractArray in 4th order operators
glwagner Sep 1, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion src/BoundaryConditions/coordinate_boundary_conditions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ A set of two `BoundaryCondition`s to be applied along a coordinate x, y, or z.
The `left` boundary condition is applied on the negative or lower side of the coordinate
while the `right` boundary condition is applied on the positive or higher side.
"""
mutable struct CoordinateBoundaryConditions{L, R}
struct CoordinateBoundaryConditions{L, R}
left :: L
right :: R
end
Expand Down Expand Up @@ -37,3 +37,6 @@ setbc!(cbc::CBC, ::Val{:north}, bc) = setfield!(cbc, :right, bc)
@inline getbc(cbc::CBC, ::Val{:north}) = getfield(cbc, :right)
@inline getbc(cbc::CBC, ::Val{:west}) = getfield(cbc, :left)
@inline getbc(cbc::CBC, ::Val{:east}) = getfield(cbc, :right)

Adapt.adapt_structure(to, cbc::CoordinateBoundaryConditions) =
CoordinateBoundaryConditions(Adapt.adapt(to, cbc.left), Adapt.adapt(to, cbc.right))
8 changes: 4 additions & 4 deletions src/Buoyancy/nonlinear_equation_of_state.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,10 @@ end
end

# Dispatch shenanigans
@inline θ_and_sᴬ(i, j, k, θ::AbstractArray, sᴬ::AbstractArray) = @inbounds θ[i, j, k], sᴬ[i, j, k]
@inline θ_and_sᴬ(i, j, k, θ::Number, sᴬ::AbstractArray) = @inbounds θ, sᴬ[i, j, k]
@inline θ_and_sᴬ(i, j, k, θ::AbstractArray, sᴬ::Number) = @inbounds θ[i, j, k], sᴬ
@inline θ_and_sᴬ(i, j, k, θ::Number, sᴬ::Number) = @inbounds θ, sᴬ
@inline θ_and_sᴬ(i, j, k, θ, sᴬ) = @inbounds θ[i, j, k], sᴬ[i, j, k]
@inline θ_and_sᴬ(i, j, k, θ::Number, sᴬ) = @inbounds θ, sᴬ[i, j, k]
@inline θ_and_sᴬ(i, j, k, θ, sᴬ::Number) = @inbounds θ[i, j, k], sᴬ
@inline θ_and_sᴬ(i, j, k, θ::Number, sᴬ::Number) = @inbounds θ, sᴬ

# Basic functionality
@inline ρ′(i, j, k, grid, eos, θ, sᴬ) = @inbounds ρ′(θ_and_sᴬ(i, j, k, θ, sᴬ)..., Zᵃᵃᶜ(i, j, k, grid), eos)
Expand Down
1 change: 1 addition & 0 deletions src/Fields/Fields.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ using Oceananigans.Grids
using Oceananigans.BoundaryConditions

include("field.jl")
include("adapt_field.jl")
include("set!.jl")
include("field_tuples.jl")
include("show_fields.jl")
Expand Down
6 changes: 6 additions & 0 deletions src/Fields/adapt_field.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
using Adapt

Adapt.adapt_structure(to, field::Field{X, Y, Z}) where {X, Y, Z} =
Field{X, Y, Z}(Adapt.adapt(to, field.data),
grid,
Adapt.adapt(to, field.boundary_conditions))
2 changes: 1 addition & 1 deletion src/Fields/field.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,7 @@ const OffsetCuArray = OffsetArray{T, D, <:CuArray} where {T, D}
OffsetArray(Array(parent(f)), f.grid, location(f))

# Endpoint for recursive `datatuple` function:
@inline datatuple(obj::AbstractField) = data(obj)
@inline datatuple(obj::AbstractField) = obj #data(obj)

"Returns `f.data.parent` for `f::Field`."
@inline Base.parent(f::Field) = f.data.parent
Expand Down
6 changes: 3 additions & 3 deletions src/Operators/laplacian_operators.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,6 @@ end
##### Horizontal biharmonic operators
#####

@inline ∇⁴hᶜᶜᵃ(i, j, k, grid, c::AbstractArray) = ∇²hᶜᶜᵃ(i, j, k, grid, ∇²hᶜᶜᵃ, c)
@inline ∇⁴hᶠᶜᵃ(i, j, k, grid, c::AbstractArray) = ∇²hᶠᶜᵃ(i, j, k, grid, ∇²hᶠᶜᵃ, c)
@inline ∇⁴hᶜᶠᵃ(i, j, k, grid, c::AbstractArray) = ∇²hᶜᶠᵃ(i, j, k, grid, ∇²hᶜᶠᵃ, c)
@inline ∇⁴hᶜᶜᵃ(i, j, k, grid, c) = ∇²hᶜᶜᵃ(i, j, k, grid, ∇²hᶜᶜᵃ, c)
@inline ∇⁴hᶠᶜᵃ(i, j, k, grid, c) = ∇²hᶠᶜᵃ(i, j, k, grid, ∇²hᶠᶜᵃ, c)
@inline ∇⁴hᶜᶠᵃ(i, j, k, grid, c) = ∇²hᶜᶠᵃ(i, j, k, grid, ∇²hᶜᶠᵃ, c)