Skip to content

Commit a0786fc

Browse files
charleskawczynskicharleskawczynski
committed
Drop field dimension to demo thermo benchmark
Rename TupleOfArrays to FieldArrays Apply formatter wip Pass some unit tests, cleanup wip wip wip Apply formatter
1 parent 18475e5 commit a0786fc

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

49 files changed

+1943
-834
lines changed

‎benchmarks/scripts/thermo_bench_bw.jl

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ using BenchmarkTools
150150
import .TestUtilities as TU;
151151

152152
using Test
153-
@testset "Thermo state" begin
153+
# @testset "Thermo state" begin
154154
FT = Float32
155155
bm = TBB.Benchmark(;problem_size=(63,4,4,1,5400), float_type=FT)
156156
device = ClimaComms.device()
@@ -175,7 +175,8 @@ using Test
175175
)
176176
x = fill((; ts = zero(TBB.PhaseEquil{FT}), nt_core...), cspace)
177177
xv = fill((; ts = nt_ts, nt_core...), cspace)
178-
(_, Nij, _, Nv, Nh) = size(Fields.field_values(x.ts))
178+
fv_ts = Fields.field_values(x.ts)
179+
(_, Nij, _, Nv, Nh) = size(fv_ts)
179180
us = TBB.UniversalSizesStatic(Nv, Nij, Nh)
180181
function to_vec(ξ)
181182
pns = propertynames(ξ)
@@ -186,7 +187,7 @@ using Test
186187
end
187188
return (; zip(propertynames(ξ), dl_vals)...)
188189
end
189-
x_vec = to_vec(xv)
190+
# x_vec = to_vec(xv)
190191

191192
x_aos = fill((; ρ_read = FT(0), ρ_write = FT(0)), cspace)
192193
x_soa = (;
@@ -199,20 +200,21 @@ using Test
199200
@. x_aos.ρ_write = 7
200201
TBB.singlefield_bc!(x_soa, us; nreps=1, n_trials = 1)
201202
TBB.singlefield_bc!(x_aos, us; nreps=1, n_trials = 1)
202-
203+
203204
TBB.thermo_func_bc!(x, us; nreps=1, n_trials = 1)
204-
TBB.thermo_func_sol!(x_vec, us; nreps=1, n_trials = 1)
205+
# TBB.thermo_func_sol!(x_vec, us; nreps=1, n_trials = 1)
205206

206-
rc = Fields.rcompare(x_vec, to_vec(x))
207-
rc || Fields.@rprint_diff(x_vec, to_vec(x)) # test correctness (should print nothing)
208-
@test rc # test correctness
207+
# rc = Fields.rcompare(x_vec, to_vec(x))
208+
# rc || Fields.@rprint_diff(x_vec, to_vec(x)) # test correctness (should print nothing)
209+
# @test rc # test correctness
209210

210-
TBB.singlefield_bc!(x_soa, us; nreps=100, bm)
211-
TBB.singlefield_bc!(x_aos, us; nreps=100, bm)
211+
# TBB.singlefield_bc!(x_soa, us; nreps=100, bm)
212+
# TBB.singlefield_bc!(x_aos, us; nreps=100, bm)
212213
TBB.thermo_func_bc!(x, us; nreps=100, bm)
213-
TBB.thermo_func_sol!(x_vec, us; nreps=100, bm)
214+
@info "Success!"
215+
# TBB.thermo_func_sol!(x_vec, us; nreps=100, bm)
214216

215217
TBB.tabulate_benchmark(bm)
216218

217-
end
219+
# end
218220
#! format: on

‎ext/ClimaCoreCUDAExt.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ import ClimaCore.Utilities: cart_ind, linear_ind
1717
import ClimaCore.RecursiveApply:
1818
⊠, ⊞, ⊟, radd, rmul, rsub, rdiv, rmap, rzero, rmin, rmax
1919
import ClimaCore.DataLayouts: get_N, get_Nv, get_Nij, get_Nij, get_Nh
20+
import ClimaCore.DataLayouts: universal_size, UniversalSize
21+
import ClimaCore.DataLayouts: ArraySize
2022

2123
include(joinpath("cuda", "cuda_utils.jl"))
2224
include(joinpath("cuda", "data_layouts.jl"))

‎ext/cuda/data_layouts.jl

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,17 @@ import CUDA
1313
parent_array_type(::Type{<:CUDA.CuArray{T, N, B} where {N}}) where {T, B} =
1414
CUDA.CuArray{T, N, B} where {N}
1515

16+
# Can we remove this?
17+
# parent_array_type(
18+
# ::Type{<:CUDA.CuArray{T, N, B} where {N}},
19+
# ::Val{ND},
20+
# ) where {T, B, ND} = CUDA.CuArray{T, ND, B}
21+
22+
parent_array_type(
23+
::Type{<:CUDA.CuArray{T, N, B} where {N}},
24+
as::ArraySize,
25+
) where {T, B} = CUDA.CuArray{T, ndims(as), B}
26+
1627
# Ensure that both parent array types have the same memory buffer type.
1728
promote_parent_array_type(
1829
::Type{CUDA.CuArray{T1, N, B} where {N}},
@@ -53,3 +64,16 @@ function Adapt.adapt_structure(
5364
end,
5465
)
5566
end
67+
68+
import Adapt
69+
import CUDA
70+
function Adapt.adapt_structure(
71+
to::CUDA.KernelAdaptor,
72+
bc::DataLayouts.NonExtrudedBroadcasted{Style},
73+
) where {Style}
74+
DataLayouts.NonExtrudedBroadcasted{Style}(
75+
adapt_f(to, bc.f),
76+
Adapt.adapt(to, bc.args),
77+
Adapt.adapt(to, bc.axes),
78+
)
79+
end

‎ext/cuda/data_layouts_copyto.jl

Lines changed: 56 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -1,87 +1,60 @@
1+
import ClimaCore.DataLayouts:
2+
to_non_extruded_broadcasted, has_uniform_datalayouts
13
DataLayouts._device_dispatch(x::CUDA.CuArray) = ToCUDA()
24

3-
function knl_copyto!(dest, src)
5+
# function Base.copyto!(
6+
# dest::VIJFH{S, Nv, Nij, Nh},
7+
# bc::DataLayouts.BroadcastedUnionVIJFH{S, Nv, Nij, Nh},
8+
# ::ToCUDA,
9+
# ) where {S, Nv, Nij, Nh}
10+
# if Nv > 0 && Nh > 0
11+
# us = DataLayouts.UniversalSize(dest)
12+
# n = prod(DataLayouts.universal_size(us))
13+
# if has_uniform_datalayouts(bc)
14+
# bc′ = to_non_extruded_broadcasted(bc)
15+
# auto_launch!(knl_copyto_linear!, (dest, bc′, us), n; auto = true)
16+
# else
17+
# auto_launch!(knl_copyto_cart!, (dest, bc, us), n; auto = true)
18+
# end
19+
# end
20+
# return dest
21+
# end
422

5-
i = CUDA.threadIdx().x
6-
j = CUDA.threadIdx().y
7-
8-
h = CUDA.blockIdx().x
9-
v = CUDA.blockDim().z * (CUDA.blockIdx().y - 1) + CUDA.threadIdx().z
10-
11-
if v <= size(dest, 4)
12-
I = CartesianIndex((i, j, 1, v, h))
13-
@inbounds dest[I] = src[I]
23+
function knl_copyto_linear!(dest::AbstractData, bc, us)
24+
@inbounds begin
25+
tidx = thread_index()
26+
if tidx ≤ get_N(us)
27+
dest[tidx] = bc[tidx]
28+
end
1429
end
1530
return nothing
1631
end
1732

18-
function Base.copyto!(
19-
dest::IJFH{S, Nij, Nh},
20-
bc::DataLayouts.BroadcastedUnionIJFH{S, Nij, Nh},
21-
::ToCUDA,
22-
) where {S, Nij, Nh}
23-
if Nh > 0
24-
auto_launch!(
25-
knl_copyto!,
26-
(dest, bc);
27-
threads_s = (Nij, Nij),
28-
blocks_s = (Nh, 1),
29-
)
30-
end
31-
return dest
32-
end
33-
34-
function Base.copyto!(
35-
dest::VIJFH{S, Nv, Nij, Nh},
36-
bc::DataLayouts.BroadcastedUnionVIJFH{S, Nv, Nij, Nh},
37-
::ToCUDA,
38-
) where {S, Nv, Nij, Nh}
39-
if Nv > 0 && Nh > 0
40-
Nv_per_block = min(Nv, fld(256, Nij * Nij))
41-
Nv_blocks = cld(Nv, Nv_per_block)
42-
auto_launch!(
43-
knl_copyto!,
44-
(dest, bc);
45-
threads_s = (Nij, Nij, Nv_per_block),
46-
blocks_s = (Nh, Nv_blocks),
47-
)
48-
end
49-
return dest
33+
function knl_copyto_linear!(dest::DataF, bc, us)
34+
tidx = thread_index()
35+
@inbounds dest[] = bc[tidx]
36+
return nothing
5037
end
5138

52-
function Base.copyto!(
53-
dest::VF{S, Nv},
54-
bc::DataLayouts.BroadcastedUnionVF{S, Nv},
55-
::ToCUDA,
56-
) where {S, Nv}
57-
if Nv > 0
58-
auto_launch!(
59-
knl_copyto!,
60-
(dest, bc);
61-
threads_s = (1, 1),
62-
blocks_s = (1, Nv),
63-
)
39+
function knl_copyto_flat!(dest::AbstractData, bc, us)
40+
@inbounds begin
41+
tidx = thread_index()
42+
if tidx ≤ get_N(us)
43+
n = size(dest)
44+
I = kernel_indexes(tidx, n)
45+
dest[I] = bc[I]
46+
end
6447
end
65-
return dest
66-
end
67-
68-
function Base.copyto!(
69-
dest::DataF{S},
70-
bc::DataLayouts.BroadcastedUnionDataF{S},
71-
::ToCUDA,
72-
) where {S}
73-
auto_launch!(knl_copyto!, (dest, bc); threads_s = (1, 1), blocks_s = (1, 1))
74-
return dest
48+
return nothing
7549
end
7650

77-
import ClimaCore.DataLayouts: isascalar
78-
function knl_copyto_flat!(dest::AbstractData, bc, us)
51+
function knl_copyto_flat!(dest::DataF, bc, us)
7952
@inbounds begin
8053
tidx = thread_index()
8154
if tidx ≤ get_N(us)
8255
n = size(dest)
8356
I = kernel_indexes(tidx, n)
84-
dest[I] = bc[I]
57+
dest[] = bc[I]
8558
end
8659
end
8760
return nothing
@@ -90,22 +63,32 @@ end
9063
function cuda_copyto!(dest::AbstractData, bc)
9164
(_, _, Nv, _, Nh) = DataLayouts.universal_size(dest)
9265
us = DataLayouts.UniversalSize(dest)
66+
n = prod(DataLayouts.universal_size(us))
9367
if Nv > 0 && Nh > 0
94-
nitems = prod(DataLayouts.universal_size(dest))
95-
auto_launch!(knl_copyto_flat!, (dest, bc, us), nitems; auto = true)
68+
if has_uniform_datalayouts(bc)
69+
bc′ = to_non_extruded_broadcasted(bc)
70+
auto_launch!(
71+
knl_copyto_linear!,
72+
(dest, bc′, us),
73+
nitems;
74+
auto = true,
75+
)
76+
else
77+
auto_launch!(knl_copyto_flat!, (dest, bc, us), nitems; auto = true)
78+
end
9679
end
9780
return dest
9881
end
9982

10083
# TODO: can we use CUDA's luanch configuration for all data layouts?
10184
# Currently, it seems to have a slight performance degradation.
10285
#! format: off
103-
# Base.copyto!(dest::IJFH{S, Nij}, bc::DataLayouts.BroadcastedUnionIJFH{S, Nij, Nh}, ::ToCUDA) where {S, Nij, Nh} = cuda_copyto!(dest, bc)
86+
Base.copyto!(dest::IJFH{S, Nij}, bc::DataLayouts.BroadcastedUnionIJFH{S, Nij, Nh}, ::ToCUDA) where {S, Nij, Nh} = cuda_copyto!(dest, bc)
10487
Base.copyto!(dest::IFH{S, Ni, Nh}, bc::DataLayouts.BroadcastedUnionIFH{S, Ni, Nh}, ::ToCUDA) where {S, Ni, Nh} = cuda_copyto!(dest, bc)
10588
Base.copyto!(dest::IJF{S, Nij}, bc::DataLayouts.BroadcastedUnionIJF{S, Nij}, ::ToCUDA) where {S, Nij} = cuda_copyto!(dest, bc)
10689
Base.copyto!(dest::IF{S, Ni}, bc::DataLayouts.BroadcastedUnionIF{S, Ni}, ::ToCUDA) where {S, Ni} = cuda_copyto!(dest, bc)
10790
Base.copyto!(dest::VIFH{S, Nv, Ni, Nh}, bc::DataLayouts.BroadcastedUnionVIFH{S, Nv, Ni, Nh}, ::ToCUDA) where {S, Nv, Ni, Nh} = cuda_copyto!(dest, bc)
108-
# Base.copyto!(dest::VIJFH{S, Nv, Nij, Nh}, bc::DataLayouts.BroadcastedUnionVIJFH{S, Nv, Nij, Nh}, ::ToCUDA) where {S, Nv, Nij, Nh} = cuda_copyto!(dest, bc)
109-
# Base.copyto!(dest::VF{S, Nv}, bc::DataLayouts.BroadcastedUnionVF{S, Nv}, ::ToCUDA) where {S, Nv} = cuda_copyto!(dest, bc)
110-
# Base.copyto!(dest::DataF{S}, bc::DataLayouts.BroadcastedUnionDataF{S}, ::ToCUDA) where {S} = cuda_copyto!(dest, bc)
91+
Base.copyto!(dest::VIJFH{S, Nv, Nij, Nh}, bc::DataLayouts.BroadcastedUnionVIJFH{S, Nv, Nij, Nh}, ::ToCUDA) where {S, Nv, Nij, Nh} = cuda_copyto!(dest, bc)
92+
Base.copyto!(dest::VF{S, Nv}, bc::DataLayouts.BroadcastedUnionVF{S, Nv}, ::ToCUDA) where {S, Nv} = cuda_copyto!(dest, bc)
93+
Base.copyto!(dest::DataF{S}, bc::DataLayouts.BroadcastedUnionDataF{S}, ::ToCUDA) where {S} = cuda_copyto!(dest, bc)
11194
#! format: on

‎ext/cuda/data_layouts_fill.jl

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,11 @@ function knl_fill_flat!(dest::AbstractData, val, us)
1010
return nothing
1111
end
1212

13+
function knl_fill_flat!(dest::DataF, val, us)
14+
@inbounds dest[] = val
15+
return nothing
16+
end
17+
1318
function cuda_fill!(dest::AbstractData, val)
1419
(_, _, Nv, _, Nh) = DataLayouts.universal_size(dest)
1520
us = DataLayouts.UniversalSize(dest)

‎ext/cuda/topologies_dss.jl

Lines changed: 28 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,9 @@ function dss_load_perimeter_data_kernel!(
4848
if gidx ≤ prod(sizep)
4949
(level, p, fidx, elem) = cart_ind(sizep, gidx).I
5050
(ip, jp) = perimeter[p]
51-
data_idx = linear_ind(sized, (level, ip, jp, fidx, elem))
52-
pperimeter_data[level, p, fidx, elem] = pdata[data_idx]
51+
data_idx = linear_ind(sized, (level, ip, jp, elem))
52+
pperimeter_data.arrays[fidx][level, p, elem] =
53+
pdata.arrays[fidx][data_idx]
5354
end
5455
return nothing
5556
end
@@ -89,7 +90,8 @@ function dss_unload_perimeter_data_kernel!(
8990
(level, p, fidx, elem) = cart_ind(sizep, gidx).I
9091
(ip, jp) = perimeter[p]
9192
data_idx = linear_ind(sized, (level, ip, jp, fidx, elem))
92-
pdata[data_idx] = pperimeter_data[level, p, fidx, elem]
93+
pdata.arrays[fidx][data_idx] =
94+
pperimeter_data.arrays[fidx][level, p, elem]
9395
end
9496
return nothing
9597
end
@@ -148,12 +150,12 @@ function dss_local_kernel!(
148150
for idx in st:(en - 1)
149151
(lidx, vert) = local_vertices[idx]
150152
ip = perimeter_vertex_node_index(vert)
151-
sum_data += pperimeter_data[level, ip, fidx, lidx]
153+
sum_data += pperimeter_data.arrays[fidx][level, ip, lidx]
152154
end
153155
for idx in st:(en - 1)
154156
(lidx, vert) = local_vertices[idx]
155157
ip = perimeter_vertex_node_index(vert)
156-
pperimeter_data[level, ip, fidx, lidx] = sum_data
158+
pperimeter_data.arrays[fidx][level, ip, lidx] = sum_data
157159
end
158160
elseif gidx ≤ nlevels * nfidx * (nlocalvertices + nlocalfaces) # interior faces
159161
nfacedof = div(nperimeter - 4, 4)
@@ -169,10 +171,10 @@ function dss_local_kernel!(
169171
ip1 = inc1 == 1 ? first1 + i - 1 : first1 - i + 1
170172
ip2 = inc2 == 1 ? first2 + i - 1 : first2 - i + 1
171173
val =
172-
pperimeter_data[level, ip1, fidx, lidx1] +
173-
pperimeter_data[level, ip2, fidx, lidx2]
174-
pperimeter_data[level, ip1, fidx, lidx1] = val
175-
pperimeter_data[level, ip2, fidx, lidx2] = val
174+
pperimeter_data.arrays[fidx][level, ip1, lidx1] +
175+
pperimeter_data.arrays[fidx][level, ip2, lidx2]
176+
pperimeter_data.arrays[fidx][level, ip1, lidx1] = val
177+
pperimeter_data.arrays[fidx][level, ip2, lidx2] = val
176178
end
177179
end
178180

@@ -254,7 +256,7 @@ function dss_transform_kernel!(
254256
if gidx ≤ nlevels * nperimeter * nlocalelems
255257
sizet = (nlevels, nperimeter, nlocalelems)
256258
sizet_data = (nlevels, Nq, Nq, nfid, nelems)
257-
sizet_wt = (Nq, Nq, 1, nelems)
259+
sizet_wt = (Nq, Nq, nelems)
258260
sizet_metric = (nlevels, Nq, Nq, nmetric, nelems)
259261

260262
(level, p, localelemno) = cart_ind(sizet, gidx).I
@@ -267,26 +269,28 @@ function dss_transform_kernel!(
267269
pperimeter_data[level, p, fidx, elem] = pdata[data_idx] * weight
268270
end
269271
for fidx in covariant12fidx
270-
data_idx1 = linear_ind(sizet_data, (level, ip, jp, fidx, elem))
271-
data_idx2 = linear_ind(sizet_data, (level, ip, jp, fidx + 1, elem))
272-
(idx11, idx12, idx21, idx22) =
273-
Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
272+
data_idx = linear_ind(sizet_data, (level, ip, jp, elem))
273+
(idx11, idx12, idx21, idx22) = (1, 2, 3, 4)
274+
# Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
274275
pperimeter_data[level, p, fidx, elem] =
275276
(
276-
p∂ξ∂x[idx11] * pdata[data_idx1] +
277-
p∂ξ∂x[idx12] * pdata[data_idx2]
277+
p∂ξ∂x.arrays[idx11][data_idx] *
278+
pdata.arrays[fidx][data_idx] +
279+
p∂ξ∂x.arrays[idx12][data_idx] *
280+
pdata.arrays[fidx + 1][data_idx]
278281
) * weight
279282
pperimeter_data[level, p, fidx + 1, elem] =
280283
(
281-
p∂ξ∂x[idx21] * pdata[data_idx1] +
282-
p∂ξ∂x[idx22] * pdata[data_idx2]
284+
p∂ξ∂x.arrays[idx21][data_idx] *
285+
pdata.arrays[fidx][data_idx] +
286+
p∂ξ∂x.arrays[idx22][data_idx] *
287+
pdata.arrays[fidx + 1][data_idx]
283288
) * weight
284289
end
285290
for fidx in contravariant12fidx
286-
data_idx1 = linear_ind(sizet_data, (level, ip, jp, fidx, elem))
287-
data_idx2 = linear_ind(sizet_data, (level, ip, jp, fidx + 1, elem))
288-
(idx11, idx12, idx21, idx22) =
289-
Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
291+
data_idx = linear_ind(sizet_data, (level, ip, jp, elem))
292+
(idx11, idx12, idx21, idx22) = (1, 2, 3, 4)
293+
# Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
290294
pperimeter_data[level, p, fidx, elem] =
291295
(
292296
p∂x∂ξ[idx11] * pdata[data_idx1] +
@@ -683,7 +687,8 @@ function load_from_recv_buffer_kernel!(
683687
lidx = recv_buf_idx[irecv, 1]
684688
ip = recv_buf_idx[irecv, 2]
685689
idx = level + ((fidx - 1) + (irecv - 1) * nfid) * nlevels
686-
CUDA.@atomic pperimeter_data[level, ip, fidx, lidx] += recv_data[idx]
690+
CUDA.@atomic pperimeter_data.arrays[fidx][level, ip, lidx] +=
691+
recv_data[idx]
687692
end
688693
return nothing
689694
end

0 commit comments

Comments
 (0)