Skip to content

Commit

Permalink
Try #1380:
Browse files Browse the repository at this point in the history
  • Loading branch information
bors[bot] authored Jul 20, 2023
2 parents 0e7a5f9 + dff77b0 commit 3517017
Show file tree
Hide file tree
Showing 6 changed files with 89 additions and 34 deletions.
6 changes: 5 additions & 1 deletion src/DataLayouts/cuda.jl
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,11 @@ end
function Base.fill!(
dest::IJFH{S, Nij, A},
val,
) where {S, Nij, A <: CUDA.CuArray}
) where {
S,
Nij,
A <: Union{CUDA.CuArray, SubArray{<:Any, <:Any, <:CUDA.CuArray}},
}
_, _, _, _, Nh = size(dest)
if Nh > 0
CUDA.@cuda threads = (Nij, Nij) blocks = (Nh, 1) knl_fill!(dest, val)
Expand Down
77 changes: 63 additions & 14 deletions src/Operators/integrals.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,31 +7,80 @@
The definite vertical column integral, `col∫field`, of field `field`.
"""
function column_integral_definite!(col∫field::Fields.Field, field::Fields.Field)
column_integral_definite!(col∫field::Fields.Field, field::Fields.Field) =
column_integral_definite!(ClimaComms.device(axes(field)), col∫field, field)

function column_integral_definite!(
::ClimaComms.CUDADevice,
col∫field::Fields.Field,
field::Fields.Field,
)
Ni, Nj, _, _, Nh = size(Fields.field_values(col∫field))
nthreads, nblocks = Spaces._configure_threadblock(Ni * Nj * Nh)
@cuda threads = nthreads blocks = nblocks column_integral_definite_kernel!(
col∫field,
field,
)
end

function column_integral_definite_kernel!(
col∫field::Fields.SpectralElementField,
field::Fields.ExtrudedFiniteDifferenceField,
)
idx = threadIdx().x + (blockIdx().x - 1) * blockDim().x
Ni, Nj, _, _, Nh = size(Fields.field_values(field))
if idx <= Ni * Nj * Nh
i, j, h = Spaces._get_idx((Ni, Nj, Nh), idx)
colfield = Spaces.column(field, i, j, h)
_column_integral_definite!(Spaces.column(col∫field, i, j, h), colfield)
end
return nothing
end

column_integral_definite_kernel!(
col∫field::Fields.PointField,
field::Fields.FiniteDifferenceField,
) = _column_integral_definite!(col∫field, field)

function column_integral_definite!(
::ClimaComms.AbstractCPUDevice,
col∫field::Fields.SpectralElementField,
field::Fields.ExtrudedFiniteDifferenceField,
)
Fields.bycolumn(axes(field)) do colidx
column_integral_definite!(col∫field[colidx], field[colidx])
_column_integral_definite!(col∫field[colidx], field[colidx])
nothing
end
return nothing
end

function column_integral_definite!(
column_integral_definite!(
::ClimaComms.AbstractCPUDevice,
col∫field::Fields.PointField,
field::Fields.FiniteDifferenceField,
) = _column_integral_definite!(col∫field, field)

function _column_integral_definite!(
col∫field::Fields.PointField,
field::Fields.ColumnField,
)
@inbounds col∫field[] = column_integral_definite(field)
return nothing
end
space = axes(field)
Δz_field = Fields.Δz_field(space)
Nv = Spaces.nlevels(space)

function column_integral_definite(field::Fields.ColumnField)
field_data = Fields.field_values(field)
Δz_data = Spaces.Δz_data(axes(field))
Nv = Spaces.nlevels(axes(field))
∫field = zero(eltype(field))
@inbounds for j in 1:Nv
∫field += field_data[j] * Δz_data[j]
col∫field[] = 0
@inbounds for idx in 1:Nv
col∫field[] +=
reduction_getindex(field, idx) * reduction_getindex(Δz_field, idx)
end
return ∫field
return nothing
end

reduction_getindex(column_field, index) = @inbounds getidx(
axes(column_field),
column_field,
Interior(),
index - 1 + left_idx(axes(column_field)),
)

# TODO: add support for indefinite integrals
6 changes: 2 additions & 4 deletions src/Spaces/finitedifference.jl
Original file line number Diff line number Diff line change
Expand Up @@ -212,14 +212,12 @@ Base.@propagate_inbounds function level(
v::PlusHalf,
)
@inbounds local_geometry = level(local_geometry_data(space), v.i + 1)
context = ClimaComms.context(space)
PointSpace(context, local_geometry)
PointSpace(local_geometry)
end
Base.@propagate_inbounds function level(
space::CenterFiniteDifferenceSpace,
v::Int,
)
local_geometry = level(local_geometry_data(space), v)
context = ClimaComms.context(space)
PointSpace(context, local_geometry)
PointSpace(local_geometry)
end
15 changes: 9 additions & 6 deletions src/Spaces/pointspace.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,19 +7,22 @@ local_geometry_data(space::AbstractPointSpace) = space.local_geometry
A zero-dimensional space.
"""
struct PointSpace{C <: ClimaComms.AbstractCommsContext, LG} <:
AbstractPointSpace
context::C
struct PointSpace{LG <: DataLayouts.Data0D} <: AbstractPointSpace
local_geometry::LG
end

ClimaComms.device(space::PointSpace) = ClimaComms.device(space.context)
ClimaComms.context(space::PointSpace) = space.context
# not strictly correct, but it is needed for mapreduce
ClimaComms.device(space::PointSpace) =
ClimaComms.device(ClimaComms.context(space))
ClimaComms.context(space::PointSpace) =
ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded())

#=
PointSpace(x::Geometry.LocalGeometry) =
PointSpace(ClimaComms.CPUSingleThreaded(), x)
PointSpace(x::Geometry.AbstractPoint) =
PointSpace(ClimaComms.CPUSingleThreaded(), x)
=#

function PointSpace(device::ClimaComms.AbstractDevice, x)
context = ClimaComms.SingletonCommsContext(device)
Expand All @@ -34,7 +37,7 @@ function PointSpace(
ArrayType = ClimaComms.array_type(ClimaComms.device(context))
local_geometry_data = DataLayouts.DataF{LG}(Array{FT})
local_geometry_data[] = local_geometry
return PointSpace(context, Adapt.adapt(ArrayType, local_geometry_data))
return PointSpace(Adapt.adapt(ArrayType, local_geometry_data))
end

function PointSpace(
Expand Down
6 changes: 2 additions & 4 deletions src/Spaces/spectralelement.jl
Original file line number Diff line number Diff line change
Expand Up @@ -644,16 +644,14 @@ Base.@propagate_inbounds slab(space::AbstractSpectralElementSpace, h) =

Base.@propagate_inbounds function column(space::SpectralElementSpace1D, i, h)
local_geometry = column(local_geometry_data(space), i, h)
context = ClimaComms.context(space)
PointSpace(context, local_geometry)
PointSpace(local_geometry)
end
Base.@propagate_inbounds column(space::SpectralElementSpace1D, i, j, h) =
column(space, i, h)

Base.@propagate_inbounds function column(space::SpectralElementSpace2D, i, j, h)
local_geometry = column(local_geometry_data(space), i, j, h)
context = ClimaComms.context(space)
PointSpace(context, local_geometry)
PointSpace(local_geometry)
end

# XXX: this cannot take `space` as it must be constructed beforehand so
Expand Down
13 changes: 8 additions & 5 deletions test/Fields/field.jl
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import ClimaCore:
using LinearAlgebra: norm
using Statistics: mean
using ForwardDiff
using CUDA

include(
joinpath(pkgdir(ClimaCore), "test", "TestUtilities", "TestUtilities.jl"),
Expand All @@ -35,7 +36,7 @@ function spectral_space_2D(; n1 = 1, n2 = 1, Nij = 4)
space = Spaces.SpectralElementSpace2D(grid_topology, quad)
return space
end

#=
@testset "1×1 2D domain space" begin
Nij = 4
n1 = n2 = 1
Expand Down Expand Up @@ -269,7 +270,7 @@ end
@testset "FieldVector array_type" begin
device = ClimaComms.device()
context = ClimaComms.SingletonCommsContext(device)
space = TU.PointSpace(Float32; context)
space = TU.SpectralElementSpace1D(Float32; context)
xcenters = Fields.coordinate_field(space).x
y = Fields.FieldVector(x = xcenters)
@test ClimaComms.array_type(y) == ClimaComms.array_type(device)
Expand Down Expand Up @@ -667,13 +668,13 @@ end
C .= Ref(zero(eltype(C)))
@test all(==(0.0), parent(C))
end
=#
function integrate_bycolumn!(∫y, Y)
Fields.bycolumn(axes(Y.y)) do colidx
Operators.column_integral_definite!(∫y[colidx], Y.y[colidx])
nothing
end
end

"""
convergence_rate(err, Δh)
Expand All @@ -698,7 +699,7 @@ convergence_rate(err, Δh) =
col_copy = similar(y[Fields.ColumnIndex((1, 1), 1)])
return Fields.Field(Fields.field_values(col_copy), axes(col_copy))
end
device = ClimaComms.CPUSingleThreaded()
device = ClimaComms.device()
context = ClimaComms.SingletonCommsContext(device)
for zelem in (2^2, 2^3, 2^4, 2^5)
for space in TU.all_spaces(FT; zelem, context)
Expand All @@ -711,7 +712,9 @@ convergence_rate(err, Δh) =
zcf = Fields.coordinate_field(Y.y).z
Δz = Fields.Δz_field(axes(zcf))
Δz_col = Δz[Fields.ColumnIndex((1, 1), 1)]
Δz_1 = parent(Δz_col)[1]
Δz_1 = CUDA.allowscalar() do
parent(Δz_col)[1]
end
key = (space, zelem)
if !haskey(results, key)
results[key] =
Expand Down

0 comments on commit 3517017

Please sign in to comment.