Skip to content

Commit d259da6

Browse files
committed
Address reviewer comments [perf]
1 parent b20fad1 commit d259da6

11 files changed

Lines changed: 223 additions & 353 deletions

File tree

ext/cuda/scopes.jl

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,6 @@ function check_device_assumptions()
2020
return nothing
2121
end
2222

23-
@inline x_component((; x, y, z)) =
24-
isone(y) && isone(z) ? x :
25-
throw(ArgumentError("y and z dimensions in launch configuration are not supported"))
26-
2723
DataLayouts.DataScope(::Type{<:CUDA.CuArray}) = ThisHost()
2824
DataLayouts.DataScope(::Type{<:CUDA.CuDeviceArray{<:Any, <:Any, A}}) where {A} =
2925
A == CUDA.AS.Local ? DataLayouts.ThisThread() :
@@ -49,12 +45,15 @@ DataLayouts.scoped_array(::ThisHost, ::Type{T}, dims) where {T} =
4945
5046
[`DataScope`](@ref) that represents all available threads on a GPU. Operations
5147
that require synchronizations or array allocations are not supported.
48+
49+
NOTE: This assumes that kernels are always launched with one-dimensional grids.
50+
Support for multidimensional grids may be added in a future release.
5251
"""
5352
struct ThisKernel <: DataLayouts.DataScope end
5453

5554
DataLayouts.partition(::ThisKernel) = ThisBlock()
56-
DataLayouts.num_partitions(::ThisKernel) = x_component(CUDA.gridDim())
57-
DataLayouts.partition_rank(::ThisKernel) = x_component(CUDA.blockIdx())
55+
DataLayouts.num_partitions(::ThisKernel) = CUDA.gridDim().x
56+
DataLayouts.partition_rank(::ThisKernel) = CUDA.blockIdx().x
5857

5958
"""
6059
ThisCooperativeGroup
@@ -72,12 +71,15 @@ abstract type ThisCooperativeGroup <: DataLayouts.DataScope end
7271
7372
[`DataScope`](@ref) that represents one thread block of [`ThisKernel`](@ref).
7473
Operations that require dynamically-sized array allocations are not supported.
74+
75+
NOTE: This assumes that kernels are always launched with one-dimensional blocks.
76+
Support for multidimensional blocks may be added in a future release.
7577
"""
7678
struct ThisBlock <: ThisCooperativeGroup end
7779

7880
DataLayouts.partition(::ThisBlock) = ThisWarp()
79-
DataLayouts.num_threads(::ThisBlock) = x_component(CUDA.blockDim())
80-
DataLayouts.thread_rank(::ThisBlock) = x_component(CUDA.threadIdx())
81+
DataLayouts.num_threads(::ThisBlock) = CUDA.blockDim().x
82+
DataLayouts.thread_rank(::ThisBlock) = CUDA.threadIdx().x
8183
DataLayouts.synchronize(::ThisBlock) = CUDA.sync_threads()
8284
DataLayouts.scoped_static_array(::ThisBlock, ::Type{T}, dims) where {T} =
8385
CUDA.CuStaticSharedArray(T, dims)
@@ -102,10 +104,10 @@ DataLayouts.partition(::ThisSubBlock{N}) where {N} =
102104
N < 4 ? DataLayouts.ThisThread() : ThisSubBlock{N ÷ 2}()
103105
DataLayouts.num_threads(::ThisSubBlock{N}) where {N} = N
104106
DataLayouts.thread_rank(::ThisSubBlock{N}) where {N} =
105-
N > THREADS_PER_WARP ? (x_component(CUDA.threadIdx()) - 1) % N + 1 :
107+
N > THREADS_PER_WARP ? (DataLayouts.thread_rank(ThisBlock()) - 1) % N + 1 :
106108
N < THREADS_PER_WARP ? (CUDA.laneid() - 1) % N + 1 : CUDA.laneid()
107109
DataLayouts.synchronize(::ThisSubBlock{N}) where {N} =
108-
N > THREADS_PER_WARP ? CUDA.sync_threads() : CUDA.sync_warp()
110+
N > THREADS_PER_WARP ? DataLayouts.synchronize(ThisBlock()) : CUDA.sync_warp()
109111

110112
# Assign threads in a sub-block one slice of an array shared across their block.
111113
function DataLayouts.scoped_static_array(scope::ThisSubBlock, ::Type{T}, dims) where {T}

src/DataLayouts/DataLayouts.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ include("scopes.jl")
2828
2929
An `N`-dimensional `AbstractArray` containing values of type `T`, stored in a
3030
parent array of type `A` whose memory layout is determined by the layout's type.
31-
Every value is identified by up to four indicesa vertical level `v`, horizontal
31+
Every value can be identified by four indices: a vertical level `v`, horizontal
3232
quadrature points `i` and `j`, and a horizontal element `h`. The components of
3333
each value are optionally stored along a hidden field axis `F` of the parent
3434
array, leading to a hybrid of the traditional "array-of-structs" (`F = 1`) and
@@ -47,8 +47,8 @@ Several layouts are available, named after the order of their parent axes:
4747
- [`VIJHWithF`](@ref) generalizes `VIJFH` and `VIJHF` to any `F` axis
4848
position, with `F = nothing` removing the axis altogether
4949
- [`VIH1`](@ref) and [`IH1JH2`](@ref) store vertical and horizontal planes of
50-
interpolated data for plotting, whose indices `ih1` and `jh2` combine `i` and
51-
`j` with corresponding `h1` and `h2` components of `h` in rectangular domains
50+
interpolated data for plotting, whose `ih1` and `jh2` indices combine `i` and
51+
`j` with `h1` and `h2` (orthogonal components of `h` in rectangular domains)
5252
5353
```julia-repl
5454
julia> data = VIJFH{Tuple{Int64, Float64, Int128}, 10, 5, 5, nothing}(Array{Int64}, 20);

src/DataLayouts/bitcast_struct.jl

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,6 @@
22
field_expr(f, value_expr) = :(Core.getfield($value_expr, $f))
33
field_expr(f, array_expr, index_expr) =
44
:(@inbounds $array_expr[struct_index($f, $array_expr, $index_expr...)])
5-
# Flat/linear variant (see get_struct_linear): read field f of the value stored
6-
# at linear position `base` using a single linear load into the array's flat
7-
# storage. `base` and `fstride` are already fully constant-folded affine
8-
# expressions of the layout's static shape params, and `$(f - 1)` is a literal,
9-
# so `array[base + (f - 1) * fstride]` lowers to one arrayref with a precomputed
10-
# offset -- no Cartesian reindex, no div/rem, no per-dimension bounds math.
11-
field_expr(f, array_expr, base_expr, fstride_expr) =
12-
:(@inbounds $array_expr[$base_expr + $(f - 1) * $fstride_expr])
135

146
# Keep array element read instructions separate unless the full value is needed
157
full_value_expr(@nospecialize(S), value_expr) = value_expr
@@ -163,22 +155,3 @@ For more information about `reinterpret` and padding, see the following:
163155
S = NTuple{num_indices, eltype(array)}
164156
return Expr(:block, :@inline, bitcast_struct_expr(T, S, :array, :index))
165157
end
166-
167-
# Flat/linear analogue of the array method (see get_struct_linear): reads the
168-
# `num_indices` entries at linear positions `base, base + fstride, ...` using the
169-
# flat `field_expr(f, array, base, fstride)`. Reuses bitcast_struct_expr so the
170-
# type-reconstruction machinery is shared with the Cartesian path.
171-
@generated function bitcast_struct_linear(
172-
::Type{T},
173-
array,
174-
::Val{num_indices},
175-
base,
176-
fstride,
177-
) where {T, num_indices}
178-
S = NTuple{num_indices, eltype(array)}
179-
return Expr(
180-
:block,
181-
:@inline,
182-
bitcast_struct_expr(T, S, :array, :base, :fstride),
183-
)
184-
end

src/DataLayouts/broadcast.jl

Lines changed: 19 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -7,27 +7,23 @@
77
struct DataStyle{N, D <: DataLayout{<:Any, N}} <: Broadcast.AbstractArrayStyle{N} end
88
DataStyle(::Type{D}) where {D} = DataStyle{ndims(D), layout_type(D)}()
99

10+
Base.ndims(::DataStyle{N}) where {N} = N
11+
1012
Broadcast.BroadcastStyle(::Type{D}) where {D <: DataLayout} = DataStyle(D)
1113

12-
# For styles with equal typenames but different dimensionalities, Base's fallback for
13-
# AbstractArrayStyle calls typeof(style)(Val(N)); DataStyle needs the layout type D, so
14-
# it cannot define that constructor and bypasses the fallback instead.
15-
Broadcast.BroadcastStyle(::DataStyle{<:Any, D1}, ::DataStyle{<:Any, D2}) where {D1, D2} =
16-
D1 == D2 || iszero(ndims(D2)) ? DataStyle(D1) :
17-
iszero(ndims(D1)) ? DataStyle(D2) : Broadcast.Unknown()
14+
# For styles with equal typenames but different dimensionalities, Base's
15+
# fallback for AbstractArrayStyle calls typeof(style)(Val(N)). DataStyle needs a
16+
# layout type D in addition to the dimensionality, so it bypasses the fallback.
17+
Broadcast.BroadcastStyle(style1::DataStyle, style2::DataStyle) =
18+
style1 == style2 || iszero(ndims(style2)) ? style1 :
19+
iszero(ndims(style1)) ? style2 : Broadcast.Unknown()
1820

1921
# Pass scalar values in Tuples of length 1 or in 0-dimensional AbstractArrays.
2022
# Add DefaultArrayStyle{0} and DataStyle{0} methods to avoid ambiguities.
2123
Broadcast.BroadcastStyle(style::DataStyle, ::Broadcast.Style{Tuple}) = style
2224
Broadcast.BroadcastStyle(style::DataStyle, ::Broadcast.AbstractArrayStyle{0}) = style
2325
Broadcast.BroadcastStyle(style::DataStyle, ::Broadcast.DefaultArrayStyle{0}) = style
24-
# The unconstrained typevar D2 is required for this method to cover the full
25-
# intersection of the two ambiguous signatures above (typeintersect drops the
26-
# DataLayout bound on the second style's layout type).
27-
Broadcast.BroadcastStyle(
28-
style::DataStyle{N, D1},
29-
::DataStyle{0, D2},
30-
) where {N, D1 <: DataLayout{<:Any, N}, D2} = style
26+
Broadcast.BroadcastStyle(style::DataStyle, ::DataStyle{0}) = style
3127

3228
# Enable automatic nested broadcasting over supported types of iterators.
3329
@inline Broadcast.broadcastable(data::DataLayout) =
@@ -42,17 +38,19 @@ A [`DataStyle`](@ref) broadcast expression whose [`layout_type`](@ref) is `D`.
4238
"""
4339
const LazyDataLayout{D} = Broadcast.Broadcasted{<:DataStyle{<:Any, D}}
4440

45-
# Optimize axes(::LazyDataLayout) with statically inferrable sizes. Dynamic sizes avoid
46-
# Base's combine_axes, whose dimension-mismatch error path cannot compile in GPU kernels:
47-
# assuming broadcast-compatible arguments, each dimension takes the non-singleton axis.
41+
# Optimize axes(::LazyDataLayout) with statically inferrable axes when possible.
4842
@inline Broadcast._axes(bc::LazyDataLayout, ::Nothing) =
4943
has_inferred_size(bc) ? unrolled_map(Base.OneTo, inferred_size(bc)) :
50-
unrolled_reduce(broadcast_axes, unrolled_map(axes, bc.args))
51-
@inline broadcast_axes(axes1::Tuple, axes2::Tuple) =
52-
isempty(axes1) || isempty(axes2) ? (isempty(axes1) ? axes2 : axes1) :
44+
unrolled_reduce(stable_combine_axes, unrolled_map(axes, bc.args))
45+
46+
# Instead of Base's combine_axes, whose DimensionMismatch error cannot compile
47+
# on GPUs, use a version that always selects the first non-singleton dimension.
48+
@inline stable_combine_axes(axes1::Tuple, axes2::Tuple) =
49+
isempty(axes2) ? axes1 :
50+
isempty(axes1) ? axes2 :
5351
(
54-
isone(length(first(axes1))) ? first(axes2) : first(axes1),
55-
broadcast_axes(Base.tail(axes1), Base.tail(axes2))...,
52+
isone(length(first(axes2))) ? first(axes1) : first(axes2),
53+
stable_combine_axes(Base.tail(axes1), Base.tail(axes2))...,
5654
)
5755

5856
# Make ndims support nested broadcasts whose axes have not been instantiated.

src/DataLayouts/deprecated.jl

Lines changed: 19 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,56 +1,34 @@
11
# Backwards-compatibility shims for the pre-rewrite DataLayouts API.
22
#
3-
# The PR replaces the per-layout struct hierarchy (AbstractData, IJFH, VIJFH,
4-
# ...) with the single DataLayout type family. Downstream packages (ClimaAtmos,
5-
# ClimaCoupler, user code) that haven't migrated yet can keep referring to the
6-
# old names through the aliases below.
7-
#
83
# These are not deprecated with warnings — they're plain aliases — so old code
94
# continues to type-check without noise. Remove this file when all downstream
105
# consumers have migrated to the new names.
11-
#
12-
# Old names whose semantics changed incompatibly (e.g. concrete layout types
13-
# whose parent arrays gained a leading vertical axis, or the *Singleton
14-
# dispatch tokens) are intentionally not aliased: resolving them to a subtly
15-
# different meaning would be worse than an UndefVarError.
166

17-
# Old: AbstractData{S} with S the element type. New: DataLayout{T, N, F, S, A}
18-
# with T the element type, so parametrized uses like AbstractData{<:Real}
19-
# still work through this alias.
20-
const AbstractData{S} = DataLayout{S}
7+
export AbstractData, IJFH, IJHF
218

22-
# Old: IJFH{S, Nij, A} and IJHF{S, Nij, A}, with 4-D parent arrays. New: 2D
23-
# fields store their values in VIJFH/VIJHF layouts with Nv = 1 and 5-D parent
24-
# arrays, so these aliases keep method signatures like
25-
# `remap!(target::IJFH{S, Nq}, ...)` (ClimaCoreTempestRemap) dispatching
26-
# correctly on 2D field values. Code that indexes into `parent(data)` with the
27-
# old 4-D shape still needs to be updated for the new parent-array layout.
28-
const IJFH{S, Nij} = VIJFH{S, 1, Nij, Nij, nothing}
29-
const IJHF{S, Nij} = VIJHF{S, 1, Nij, Nij, nothing}
9+
# Old: AbstractData{T}
10+
# New: DataLayout{T, N, F, S, A}
11+
const AbstractData{T} = DataLayout{T}
3012

31-
# These were exported on main, and packages like ClimaCoreTempestRemap access
32-
# them through `using ClimaCore.DataLayouts`.
33-
export AbstractData, IJFH, IJHF
13+
# Old: IJFH/IJHF layouts with 4-D parent arrays
14+
# New: VIJFH/VIJHF layouts with Nv = 1 and 5-D parent arrays
15+
# Used by ClimaCoreTempestRemap's remap!(target::IJFH{T, Nq}, ...).
16+
const IJFH{T, Nij} = VIJFH{T, 1, Nij, Nij, nothing}
17+
const IJHF{T, Nij} = VIJHF{T, 1, Nij, Nij, nothing}
3418

35-
# Old: getindex/setindex! with a 5-D "universal index" CartesianIndex(i, j, f,
36-
# v, h), whose f component was ignored because the entire value was read or
37-
# written. New: Cartesian indices have one coordinate per array dimension, in
38-
# the order (v, i, j, h). Packages like ClimaCoreTempestRemap still index full
39-
# layouts, slabs, and columns with universal indices (e.g.,
40-
# `slab(data, h)[CartesianIndex(i, j, 1, 1, 1)]`), so the methods below
41-
# translate 5-D indices from the old convention to the new one. They override
42-
# the generic AbstractArray interpretation of a `CartesianIndex{5}` as a 4-D
43-
# index with a trailing singleton coordinate, so new code must not add trailing
44-
# coordinates when indexing into 4-dimensional layouts. Remove these methods
45-
# together with the aliases above.
46-
@inline universal_index_shim(I::CartesianIndex{5}) =
47-
CartesianIndex(I[4], I[1], I[2], I[5])
19+
# Old: getindex/setindex! with a 5-D CartesianIndex(i, j, _, v, h)
20+
# New: getindex/setindex! with a 4-D CartesianIndex(v, i, j, h)
21+
# Used by ClimaCoreTempestRemap in slab(data, h)[CartesianIndex(i, j, 1, 1, 1)]
22+
# and similar functions. This overrides the generic AbstractArray interpretation
23+
# of a CartesianIndex{5}; until it is removed, new code should not add a
24+
# trailing singleton coordinate when indexing into VIJFH/VIJHF layouts.
4825
@propagate_inbounds Base.getindex(data::VIJHWithF, I::CartesianIndex{5}) =
49-
getindex(data, universal_index_shim(I))
26+
getindex(data, CartesianIndex(I[4], I[1], I[2], I[5]))
5027
@propagate_inbounds Base.setindex!(data::VIJHWithF, value, I::CartesianIndex{5}) =
51-
setindex!(data, value, universal_index_shim(I))
28+
setindex!(data, value, CartesianIndex(I[4], I[1], I[2], I[5]))
5229

53-
# Old: every component of a universal index was ignored for 0-dimensional data.
30+
# Old: single-point layouts ignore every component of a 5-D CartesianIndex
31+
# New: single-point layouts without any index
5432
@propagate_inbounds Base.getindex(data::DataF, I::CartesianIndex{5}) = data[]
5533
@propagate_inbounds Base.setindex!(data::DataF, value, I::CartesianIndex{5}) =
5634
setindex!(data, value)

0 commit comments

Comments
 (0)