Skip to content

Commit 4417516

Browse files
charleskawczynskicharleskawczynski
committed
Use prescribed thread-block configurations
1 parent add3b1b commit 4417516

12 files changed

+331
-171
lines changed

‎ext/ClimaCoreCUDAExt.jl

+1
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ 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: UniversalSize
2021

2122
include(joinpath("cuda", "cuda_utils.jl"))
2223
include(joinpath("cuda", "data_layouts.jl"))

‎ext/cuda/cuda_utils.jl

+6
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,12 @@ function auto_launch!(
9090
return nothing
9191
end
9292

93+
function threads_via_occupancy(f!::F!, args) where {F!}
94+
kernel = CUDA.@cuda always_inline = true launch = false f!(args...)
95+
config = CUDA.launch_configuration(kernel.fun)
96+
return config.threads
97+
end
98+
9399
"""
94100
thread_index()
95101

‎ext/cuda/data_layouts.jl

+1
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ include("data_layouts_fill.jl")
2929
include("data_layouts_copyto.jl")
3030
include("data_layouts_fused_copyto.jl")
3131
include("data_layouts_mapreduce.jl")
32+
include("data_layouts_threadblock.jl")
3233

3334
adapt_f(to, f::F) where {F} = Adapt.adapt(to, f)
3435
adapt_f(to, ::Type{F}) where {F} = (x...) -> F(x...)

‎ext/cuda/data_layouts_copyto.jl

+17-90
Original file line numberDiff line numberDiff line change
@@ -1,111 +1,38 @@
11
DataLayouts._device_dispatch(x::CUDA.CuArray) = ToCUDA()
22

3-
function knl_copyto!(dest, src)
4-
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))
3+
function knl_copyto!(dest, src, us)
4+
I = universal_index(dest)
5+
if is_valid_index(dest, I, us)
136
@inbounds dest[I] = src[I]
147
end
158
return nothing
169
end
1710

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
50-
end
51-
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-
)
64-
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
75-
end
76-
77-
import ClimaCore.DataLayouts: isascalar
78-
function knl_copyto_flat!(dest::AbstractData, bc, us)
79-
@inbounds begin
80-
tidx = thread_index()
81-
if tidx ≤ get_N(us)
82-
n = size(dest)
83-
I = kernel_indexes(tidx, n)
84-
dest[I] = bc[I]
85-
end
86-
end
87-
return nothing
88-
end
89-
9011
function cuda_copyto!(dest::AbstractData, bc)
9112
(_, _, Nv, _, Nh) = DataLayouts.universal_size(dest)
9213
us = DataLayouts.UniversalSize(dest)
9314
if Nv > 0 && Nh > 0
94-
nitems = prod(DataLayouts.universal_size(dest))
95-
auto_launch!(knl_copyto_flat!, (dest, bc, us), nitems; auto = true)
15+
args = (dest, bc, us)
16+
threads = threads_via_occupancy(knl_copyto!, args)
17+
n_max_threads = min(threads, get_N(us))
18+
p = partition(dest, n_max_threads)
19+
auto_launch!(
20+
knl_copyto!,
21+
args;
22+
threads_s = p.threads,
23+
blocks_s = p.blocks,
24+
)
9625
end
9726
return dest
9827
end
9928

100-
# TODO: can we use CUDA's luanch configuration for all data layouts?
101-
# Currently, it seems to have a slight performance degradation.
10229
#! format: off
103-
# Base.copyto!(dest::IJFH{S, Nij}, bc::DataLayouts.BroadcastedUnionIJFH{S, Nij, Nh}, ::ToCUDA) where {S, Nij, Nh} = cuda_copyto!(dest, bc)
30+
Base.copyto!(dest::IJFH{S, Nij}, bc::DataLayouts.BroadcastedUnionIJFH{S, Nij, Nh}, ::ToCUDA) where {S, Nij, Nh} = cuda_copyto!(dest, bc)
10431
Base.copyto!(dest::IFH{S, Ni, Nh}, bc::DataLayouts.BroadcastedUnionIFH{S, Ni, Nh}, ::ToCUDA) where {S, Ni, Nh} = cuda_copyto!(dest, bc)
10532
Base.copyto!(dest::IJF{S, Nij}, bc::DataLayouts.BroadcastedUnionIJF{S, Nij}, ::ToCUDA) where {S, Nij} = cuda_copyto!(dest, bc)
10633
Base.copyto!(dest::IF{S, Ni}, bc::DataLayouts.BroadcastedUnionIF{S, Ni}, ::ToCUDA) where {S, Ni} = cuda_copyto!(dest, bc)
10734
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)
35+
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)
36+
Base.copyto!(dest::VF{S, Nv}, bc::DataLayouts.BroadcastedUnionVF{S, Nv}, ::ToCUDA) where {S, Nv} = cuda_copyto!(dest, bc)
37+
Base.copyto!(dest::DataF{S}, bc::DataLayouts.BroadcastedUnionDataF{S}, ::ToCUDA) where {S} = cuda_copyto!(dest, bc)
11138
#! format: on

‎ext/cuda/data_layouts_fill.jl

+16-21
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,27 @@
1-
function knl_fill_flat!(dest::AbstractData, val, us)
2-
@inbounds begin
3-
tidx = thread_index()
4-
if tidx ≤ get_N(us)
5-
n = size(dest)
6-
I = kernel_indexes(tidx, n)
7-
@inbounds dest[I] = val
8-
end
1+
function knl_fill!(dest, val, us)
2+
I = universal_index(dest)
3+
if is_valid_index(dest, I, us)
4+
@inbounds dest[I] = val
95
end
106
return nothing
117
end
128

13-
function cuda_fill!(dest::AbstractData, val)
9+
function cuda_fill!(dest::AbstractData, bc)
1410
(_, _, Nv, _, Nh) = DataLayouts.universal_size(dest)
1511
us = DataLayouts.UniversalSize(dest)
1612
if Nv > 0 && Nh > 0
17-
nitems = prod(DataLayouts.universal_size(dest))
18-
auto_launch!(knl_fill_flat!, (dest, val, us), nitems; auto = true)
13+
args = (dest, bc, us)
14+
threads = threads_via_occupancy(knl_fill!, args)
15+
n_max_threads = min(threads, get_N(us))
16+
p = partition(dest, n_max_threads)
17+
auto_launch!(
18+
knl_fill!,
19+
args;
20+
threads_s = p.threads,
21+
blocks_s = p.blocks,
22+
)
1923
end
2024
return dest
2125
end
2226

23-
#! format: off
24-
Base.fill!(dest::IJFH{<:Any, <:Any}, val, ::ToCUDA) = cuda_fill!(dest, val)
25-
Base.fill!(dest::IFH{<:Any, <:Any}, val, ::ToCUDA) = cuda_fill!(dest, val)
26-
Base.fill!(dest::IJF{<:Any, <:Any}, val, ::ToCUDA) = cuda_fill!(dest, val)
27-
Base.fill!(dest::IF{<:Any, <:Any}, val, ::ToCUDA) = cuda_fill!(dest, val)
28-
Base.fill!(dest::VIFH{<:Any, <:Any, <:Any}, val, ::ToCUDA) = cuda_fill!(dest, val)
29-
Base.fill!(dest::VIJFH{<:Any, <:Any, <:Any}, val, ::ToCUDA) = cuda_fill!(dest, val)
30-
Base.fill!(dest::VF{<:Any, <:Any}, val, ::ToCUDA) = cuda_fill!(dest, val)
31-
Base.fill!(dest::DataF{<:Any}, val, ::ToCUDA) = cuda_fill!(dest, val)
32-
#! format: on
27+
Base.fill!(dest::AbstractData, val, ::ToCUDA) = cuda_fill!(dest, val)

0 commit comments

Comments
 (0)