Skip to content

Commit b41eb05

Browse files
charleskawczynskicharleskawczynski
committed
Implement data-specific cartesian index
Extend data-specific cartesian index
1 parent db5dae0 commit b41eb05

File tree

8 files changed

+248
-91
lines changed

8 files changed

+248
-91
lines changed

‎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: DataSpecificCartesianIndex, array_size
21+
import ClimaCore.DataLayouts: has_uniform_datalayouts
2022

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

‎ext/cuda/data_layouts_copyto.jl

Lines changed: 28 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,88 +1,5 @@
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))
13-
@inbounds dest[I] = src[I]
14-
end
15-
return nothing
16-
end
17-
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-
dest;
28-
threads_s = (Nij, Nij),
29-
blocks_s = (Nh, 1),
30-
)
31-
end
32-
return dest
33-
end
34-
35-
function Base.copyto!(
36-
dest::VIJFH{S, Nv, Nij, Nh},
37-
bc::DataLayouts.BroadcastedUnionVIJFH{S, Nv, Nij, Nh},
38-
::ToCUDA,
39-
) where {S, Nv, Nij, Nh}
40-
if Nv > 0 && Nh > 0
41-
Nv_per_block = min(Nv, fld(256, Nij * Nij))
42-
Nv_blocks = cld(Nv, Nv_per_block)
43-
auto_launch!(
44-
knl_copyto!,
45-
(dest, bc),
46-
dest;
47-
threads_s = (Nij, Nij, Nv_per_block),
48-
blocks_s = (Nh, Nv_blocks),
49-
)
50-
end
51-
return dest
52-
end
53-
54-
function Base.copyto!(
55-
dest::VF{S, Nv},
56-
bc::DataLayouts.BroadcastedUnionVF{S, Nv},
57-
::ToCUDA,
58-
) where {S, Nv}
59-
if Nv > 0
60-
auto_launch!(
61-
knl_copyto!,
62-
(dest, bc),
63-
dest;
64-
threads_s = (1, 1),
65-
blocks_s = (1, Nv),
66-
)
67-
end
68-
return dest
69-
end
70-
71-
function Base.copyto!(
72-
dest::DataF{S},
73-
bc::DataLayouts.BroadcastedUnionDataF{S},
74-
::ToCUDA,
75-
) where {S}
76-
auto_launch!(
77-
knl_copyto!,
78-
(dest, bc),
79-
dest;
80-
threads_s = (1, 1),
81-
blocks_s = (1, 1),
82-
)
83-
return dest
84-
end
85-
863
import ClimaCore.DataLayouts: isascalar
874
function knl_copyto_flat!(dest::AbstractData, bc, us)
885
@inbounds begin
@@ -96,24 +13,46 @@ function knl_copyto_flat!(dest::AbstractData, bc, us)
9613
return nothing
9714
end
9815

16+
function knl_copyto_flat_specialized!(dest::AbstractData, bc, us)
17+
@inbounds begin
18+
tidx = thread_index()
19+
if tidx ≤ get_N(us)
20+
n = array_size(dest)
21+
CIS = CartesianIndices(map(x -> Base.OneTo(x), n))
22+
I = DataSpecificCartesianIndex(CIS[tidx])
23+
dest[I] = bc[I]
24+
end
25+
end
26+
return nothing
27+
end
28+
9929
function cuda_copyto!(dest::AbstractData, bc)
10030
(_, _, Nv, Nh) = DataLayouts.universal_size(dest)
101-
us = DataLayouts.UniversalSize(dest)
10231
if Nv > 0 && Nh > 0
103-
auto_launch!(knl_copyto_flat!, (dest, bc, us), dest; auto = true)
32+
us = DataLayouts.UniversalSize(dest)
33+
if has_uniform_datalayouts(bc)
34+
auto_launch!(
35+
knl_copyto_flat_specialized!,
36+
(dest, bc, us),
37+
dest;
38+
auto = true,
39+
)
40+
else
41+
auto_launch!(knl_copyto_flat!, (dest, bc, us), dest; auto = true)
42+
end
10443
end
10544
return dest
10645
end
10746

10847
# TODO: can we use CUDA's luanch configuration for all data layouts?
10948
# Currently, it seems to have a slight performance degradation.
11049
#! format: off
111-
# Base.copyto!(dest::IJFH{S, Nij}, bc::DataLayouts.BroadcastedUnionIJFH{S, Nij, Nh}, ::ToCUDA) where {S, Nij, Nh} = cuda_copyto!(dest, bc)
50+
Base.copyto!(dest::IJFH{S, Nij}, bc::DataLayouts.BroadcastedUnionIJFH{S, Nij, Nh}, ::ToCUDA) where {S, Nij, Nh} = cuda_copyto!(dest, bc)
11251
Base.copyto!(dest::IFH{S, Ni, Nh}, bc::DataLayouts.BroadcastedUnionIFH{S, Ni, Nh}, ::ToCUDA) where {S, Ni, Nh} = cuda_copyto!(dest, bc)
11352
Base.copyto!(dest::IJF{S, Nij}, bc::DataLayouts.BroadcastedUnionIJF{S, Nij}, ::ToCUDA) where {S, Nij} = cuda_copyto!(dest, bc)
11453
Base.copyto!(dest::IF{S, Ni}, bc::DataLayouts.BroadcastedUnionIF{S, Ni}, ::ToCUDA) where {S, Ni} = cuda_copyto!(dest, bc)
11554
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)
116-
# 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)
117-
# Base.copyto!(dest::VF{S, Nv}, bc::DataLayouts.BroadcastedUnionVF{S, Nv}, ::ToCUDA) where {S, Nv} = cuda_copyto!(dest, bc)
118-
# Base.copyto!(dest::DataF{S}, bc::DataLayouts.BroadcastedUnionDataF{S}, ::ToCUDA) where {S} = cuda_copyto!(dest, bc)
55+
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)
56+
Base.copyto!(dest::VF{S, Nv}, bc::DataLayouts.BroadcastedUnionVF{S, Nv}, ::ToCUDA) where {S, Nv} = cuda_copyto!(dest, bc)
57+
Base.copyto!(dest::DataF{S}, bc::DataLayouts.BroadcastedUnionDataF{S}, ::ToCUDA) where {S} = cuda_copyto!(dest, bc)
11958
#! format: on

‎ext/cuda/data_layouts_fill.jl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,9 @@ function knl_fill_flat!(dest::AbstractData, val, us)
22
@inbounds begin
33
tidx = thread_index()
44
if tidx ≤ get_N(us)
5-
n = size(dest)
6-
I = kernel_indexes(tidx, n)
5+
n = array_size(dest)
6+
CIS = CartesianIndices(map(x -> Base.OneTo(x), n))
7+
I = DataSpecificCartesianIndex(CIS[tidx])
78
@inbounds dest[I] = val
89
end
910
end

‎src/DataLayouts/DataLayouts.jl

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,13 @@ include("struct.jl")
4949

5050
abstract type AbstractData{S} end
5151

52+
abstract type AbstractDataSpecificCartesianIndex{N} <:
53+
Base.AbstractCartesianIndex{N} end
54+
55+
struct DataSpecificCartesianIndex{N} <: AbstractDataSpecificCartesianIndex{N}
56+
I::CartesianIndex{N}
57+
end
58+
5259
Base.size(data::AbstractData, i::Integer) = size(data)[i]
5360

5461
"""
@@ -1102,6 +1109,7 @@ function VIJFH{S, Nv, Nij, Nh}(
11021109
array::AbstractArray{T, 5},
11031110
) where {S, Nv, Nij, Nh, T}
11041111
check_basetype(T, S)
1112+
@assert size(array, 1) == Nv
11051113
@assert size(array, 2) == size(array, 3) == Nij
11061114
@assert size(array, 4) == typesize(T, S)
11071115
@assert size(array, 5) == Nh
@@ -1271,6 +1279,7 @@ function VIFH{S, Nv, Ni, Nh}(
12711279
array::AbstractArray{T, 4},
12721280
) where {S, Nv, Ni, Nh, T}
12731281
check_basetype(T, S)
1282+
@assert size(array, 1) == Nv
12741283
@assert size(array, 2) == Ni
12751284
@assert size(array, 3) == typesize(T, S)
12761285
@assert size(array, 4) == Nh
@@ -1568,6 +1577,16 @@ get_Nij(::IFH{S, Nij}) where {S, Nij} = Nij
15681577
get_Nij(::IJF{S, Nij}) where {S, Nij} = Nij
15691578
get_Nij(::IF{S, Nij}) where {S, Nij} = Nij
15701579

1580+
@inline field_dim(::IJKFVH) = 4
1581+
@inline field_dim(::IJFH) = 3
1582+
@inline field_dim(::IFH) = 2
1583+
@inline field_dim(::DataF) = 1
1584+
@inline field_dim(::IJF) = 3
1585+
@inline field_dim(::IF) = 2
1586+
@inline field_dim(::VF) = 2
1587+
@inline field_dim(::VIJFH) = 4
1588+
@inline field_dim(::VIFH) = 3
1589+
15711590
Base.ndims(data::AbstractData) = Base.ndims(typeof(data))
15721591
Base.ndims(::Type{T}) where {T <: AbstractData} =
15731592
Base.ndims(parent_array_type(T))
@@ -1641,5 +1660,7 @@ include("copyto.jl")
16411660
include("fused_copyto.jl")
16421661
include("fill.jl")
16431662
include("mapreduce.jl")
1663+
include("cartesian_index.jl")
1664+
include("has_uniform_datalayouts.jl")
16441665

16451666
end # module

‎src/DataLayouts/cartesian_index.jl

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
#! format: off
2+
# ============================================================ From Base.Broadcast
3+
@inline function Base.getindex(bc::Base.Broadcast.Broadcasted, I::DataSpecificCartesianIndex)
4+
@boundscheck checkbounds(bc, I)
5+
@inbounds _broadcast_getindex(bc, I)
6+
end
7+
8+
@inline Base.checkbounds(bc::Base.Broadcast.Broadcasted, I::DataSpecificCartesianIndex) =
9+
Base.checkbounds_indices(Bool, axes(bc), (I,)) || Base.throw_boundserror(bc, (I,))
10+
11+
Base.@propagate_inbounds _broadcast_getindex(A::Union{Ref,AbstractArray{<:Any,0},Number}, I) = A[] # Scalar-likes can just ignore all indices
12+
Base.@propagate_inbounds _broadcast_getindex(::Ref{Type{T}}, I) where {T} = T
13+
# Tuples are statically known to be singleton or vector-like
14+
Base.@propagate_inbounds _broadcast_getindex(A::Tuple{Any}, I) = A[1]
15+
Base.@propagate_inbounds _broadcast_getindex(A::Tuple, I) = A[I[1]]
16+
# Everything else falls back to dynamically dropping broadcasted indices based upon its axes
17+
# Base.@propagate_inbounds _broadcast_getindex(A, I) = A[Base.Broadcast.newindex(A, I)]
18+
Base.@propagate_inbounds _broadcast_getindex(A, I) = A[I]
19+
20+
# For Broadcasted
21+
Base.@propagate_inbounds function _broadcast_getindex(bc::Base.Broadcast.Broadcasted{<:Any,<:Any,<:Any,<:Any}, I)
22+
args = _getindex(bc.args, I)
23+
return _broadcast_getindex_evalf(bc.f, args...)
24+
end
25+
# Hack around losing Type{T} information in the final args tuple. Julia actually
26+
# knows (in `code_typed`) the _value_ of these types, statically displaying them,
27+
# but inference is currently skipping inferring the type of the types as they are
28+
# transiently placed in a tuple as the argument list is lispily constructed. These
29+
# additional methods recover type stability when a `Type` appears in one of the
30+
# first two arguments of a function.
31+
Base.@propagate_inbounds function _broadcast_getindex(bc::Base.Broadcast.Broadcasted{<:Any,<:Any,<:Any,<:Tuple{Ref{Type{T}},Vararg{Any}}}, I) where {T}
32+
args = _getindex(Base.tail(bc.args), I)
33+
return _broadcast_getindex_evalf(bc.f, T, args...)
34+
end
35+
Base.@propagate_inbounds function _broadcast_getindex(bc::Base.Broadcast.Broadcasted{<:Any,<:Any,<:Any,<:Tuple{Any,Ref{Type{T}},Vararg{Any}}}, I) where {T}
36+
arg1 = _broadcast_getindex(bc.args[1], I)
37+
args = _getindex(Base.tail(Base.tail(bc.args)), I)
38+
return _broadcast_getindex_evalf(bc.f, arg1, T, args...)
39+
end
40+
Base.@propagate_inbounds function _broadcast_getindex(bc::Base.Broadcast.Broadcasted{<:Any,<:Any,<:Any,<:Tuple{Ref{Type{T}},Ref{Type{S}},Vararg{Any}}}, I) where {T,S}
41+
args = _getindex(Base.tail(Base.tail(bc.args)), I)
42+
return _broadcast_getindex_evalf(bc.f, T, S, args...)
43+
end
44+
45+
# Utilities for _broadcast_getindex
46+
Base.@propagate_inbounds _getindex(args::Tuple, I) = (_broadcast_getindex(args[1], I), _getindex(Base.tail(args), I)...)
47+
Base.@propagate_inbounds _getindex(args::Tuple{Any}, I) = (_broadcast_getindex(args[1], I),)
48+
Base.@propagate_inbounds _getindex(args::Tuple{}, I) = ()
49+
50+
@inline _broadcast_getindex_evalf(f::Tf, args::Vararg{Any,N}) where {Tf,N} = f(args...) # not propagate_inbounds
51+
# ============================================================
52+
53+
#! format: on
54+
# Datalayouts
55+
@propagate_inbounds function Base.getindex(
56+
data::AbstractData{S},
57+
I::DataSpecificCartesianIndex,
58+
) where {S}
59+
@inbounds get_struct(parent(data), S, Val(field_dim(data)), I.I)
60+
end
61+
@propagate_inbounds function Base.setindex!(
62+
data::AbstractData{S},
63+
val,
64+
I::DataSpecificCartesianIndex,
65+
) where {S}
66+
@inbounds set_struct!(
67+
parent(data),
68+
convert(S, val),
69+
Val(field_dim(data)),
70+
I.I,
71+
)
72+
end
73+
74+
@inline array_size(::IJKFVH{S, Nij, Nk, Nv, Nh}) where {S, Nij, Nk, Nv, Nh} =
75+
(Nij, Nij, Nk, 1, Nv, Nh)
76+
@inline array_size(::IJFH{S, Nij, Nh}) where {S, Nij, Nh} = (Nij, Nij, 1, Nh)
77+
@inline array_size(::IFH{S, Ni, Nh}) where {S, Ni, Nh} = (Ni, 1, Nh)
78+
@inline array_size(::DataF{S}) where {S} = (1,)
79+
@inline array_size(::IJF{S, Nij}) where {S, Nij} = (Nij, Nij, 1)
80+
@inline array_size(::IF{S, Ni}) where {S, Ni} = (Ni, 1)
81+
@inline array_size(::VF{S, Nv}) where {S, Nv} = (Nv, 1)
82+
@inline array_size(::VIJFH{S, Nv, Nij, Nh}) where {S, Nv, Nij, Nh} =
83+
(Nv, Nij, Nij, 1, Nh)
84+
@inline array_size(::VIFH{S, Nv, Ni, Nh}) where {S, Nv, Ni, Nh} =
85+
(Nv, Ni, 1, Nh)

‎src/DataLayouts/has_uniform_datalayouts.jl

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
@inline function first_datalayout_in_bc(args::Tuple, rargs...)
2+
x1 = first_datalayout_in_bc(args[1], rargs...)
3+
x1 isa AbstractData && return x1
4+
return first_datalayout_in_bc(Base.tail(args), rargs...)
5+
end
6+
7+
@inline first_datalayout_in_bc(args::Tuple{Any}, rargs...) =
8+
first_datalayout_in_bc(args[1], rargs...)
9+
@inline first_datalayout_in_bc(args::Tuple{}, rargs...) = nothing
10+
@inline first_datalayout_in_bc(x) = nothing
11+
@inline first_datalayout_in_bc(x::AbstractData) = x
12+
13+
@inline first_datalayout_in_bc(bc::Base.Broadcast.Broadcasted) =
14+
first_datalayout_in_bc(bc.args)
15+
16+
@inline _has_uniform_datalayouts_args(
17+
truesofar,
18+
start::Any,
19+
args::Tuple,
20+
rargs...,
21+
) =
22+
truesofar &&
23+
_has_uniform_datalayouts(truesofar, start, args[1], rargs...) &&
24+
_has_uniform_datalayouts_args(truesofar, start, Base.tail(args), rargs...)
25+
26+
@inline _has_uniform_datalayouts_args(
27+
truesofar,
28+
start::Any,
29+
args::Tuple{Any},
30+
rargs...,
31+
) = truesofar && _has_uniform_datalayouts(truesofar, start, args[1], rargs...)
32+
@inline _has_uniform_datalayouts_args(
33+
truesofar,
34+
::Any,
35+
args::Tuple{},
36+
rargs...,
37+
) = truesofar
38+
39+
@inline function _has_uniform_datalayouts(
40+
truesofar,
41+
start::Any,
42+
bc::Base.Broadcast.Broadcasted,
43+
)
44+
return truesofar && _has_uniform_datalayouts_args(truesofar, start, bc.args)
45+
end
46+
for DL in (:IJKFVH, :IJFH, :IFH, :DataF, :IJF, :IF, :VF, :VIJFH, :VIFH)
47+
@eval begin
48+
@inline _has_uniform_datalayouts(truesofar, ::$(DL), ::$(DL)) = true
49+
end
50+
end
51+
@inline _has_uniform_datalayouts(truesofar, ::Any, x::AbstractData) = false
52+
@inline _has_uniform_datalayouts(truesofar, ::Any, x) = truesofar
53+
54+
# Find the first datalayout in the broadcast expression (BCE),
55+
# and compare against every other datalayout in the BCE
56+
@inline has_uniform_datalayouts(bc::Base.Broadcast.Broadcasted) =
57+
_has_uniform_datalayouts_args(true, first_datalayout_in_bc(bc), bc.args)
58+
59+
@inline has_uniform_datalayouts(bc::AbstractData) = true

0 commit comments

Comments
 (0)