Can we make the LocalGeometry / IntervalMesh lazy?
- Dominant language
- Julia
- Stars
- 117
- Forks
- 19
- Avg merge
- 3d 4h
- Merged PRs (30d)
- 41
Description
The one frequent feedback tip we give to developers is: make small structs over large structs, especially if you want to make a field of them.
Unfortunately, we've not followed that tip ourselves, and the `LocalGeometry` struct is huge for `ExtrudedFiniteDifferenceSpace`s:
```julia
julia> DataLayouts.ncomponents(Spaces.local_geometry_data(space))
42
```
The reason that this is problem is that: the bandwidth of a cuda kernel depends on the number of struct fields accessed in the kernel. The more fields available, the larger the bandwidth can grow. If there is a 0-difference in the number of accesses between the large and small struct, then difference between the large and small structs is only about 9% (perhaps due to the difference in memory strides).
Here's a demo I was experimenting with:
```julia
#=
julia --project
include(joinpath("..", "perf", "struct_bandwidth.jl"))
=#
ENV["CLIMACOMMS_DEVICE"] = "CUDA"
using Revise
import ClimaComms
ClimaComms.@import_required_backends
import ClimaCore:Spaces, Fields
import ClimaCore
import ClimaCore.DataLayouts as DataLayouts
using BenchmarkTools;
@isdefined(TU) || include(joinpath(pkgdir(ClimaCore), "test", "TestUtilities", "TestUtilities.jl"));
import .TestUtilities as TU;
FT = Float64;
space = TU.CenterExtrudedFiniteDifferenceSpace(FT;zelem=63, helem=30);
abstract type AbstractStruct end
Base.@kwdef struct BigStruct{T} <: AbstractStruct
x1::T = 0
x2::T = 0
x3::T = 0
x4::T = 0
x5::T = 0
x6::T = 0
x7::T = 0
x8::T = 0
x9::T = 0
x10::T = 0
x11::T = 0
x12::T = 0
x13::T = 0
x14::T = 0
x15::T = 0
x16::T = 0
x17::T = 0
x18::T = 0
x19::T = 0
x20::T = 0
x21::T = 0
x22::T = 0
x23::T = 0
x24::T = 0
x25::T = 0
x26::T = 0
x27::T = 0
x28::T = 0
x29::T = 0
x30::T = 0
x31::T = 0
x32::T = 0
x33::T = 0
x34::T = 0
x35::T = 0
x36::T = 0
x37::T = 0
x38::T = 0
x39::T = 0
x40::T = 0
x41::T = 0
x42::T = 0
end;
Base.@kwdef struct SmallStruct{T} <: AbstractStruct
x1::T = 0
x2::T = 0
end;
@inline get_x1(s::BigStruct) = s.x1
@inline get_x2(s::BigStruct) = s.x2
@inline get_x3(s::BigStruct) = s.x3
@inline get_x4(s::BigStruct) = s.x4
@inline get_x5(s::BigStruct) = s.x5
@inline get_x6(s::BigStruct) = s.x6
@inline get_x7(s::BigStruct) = s.x7
@inline get_x1(s::SmallStruct) = s.x1
@inline get_x2(s::SmallStruct) = s.x2
@inline get_x3(s::SmallStruct) = s.x1 + 1
@inline get_x4(s::SmallStruct) = s.x2 + 2
@inline get_x5(s::SmallStruct) = s.x1 + 1 + 1
@inline get_x6(s::SmallStruct) = s.x2 + 2 + 2
@inline get_x7(s::SmallStruct) = s.x2 + 2 + 2 + 3
# # SmallStruct is lazy
# @inline foo(s::AbstractStruct) = get_x2(s) + get_x3(s) + get_x4(s) + get_x5(s) + get_x6(s) + get_x7(s)
# # Only difference is struct size
@inline foo(s::AbstractStruct) = s.x2 + 1
# Unused variables
# @inline function foo(s::AbstractStruct)
# y = get_x3(s) + get_x4(s) + get_x5(s) + get_x6(s) + get_x7(s)
# return get_x2(s)
# end
function kernel!(x, y)
@. x.x2 = foo(y)
return nothing
end
xb = Fields.Field(BigStruct{FT}, space);
yb = Fields.Field(BigStruct{FT}, space);
xs = Fields.Field(SmallStruct{FT}, space);
ys = Fields.Field(SmallStruct{FT}, space);
@show ClimaComms.device(xb)
kernel!(xb, yb);
kernel!(xs, ys);
println("\n--------------------------- kernel, BigStruct")
trial = BenchmarkTools.@benchmark ClimaComms.@cuda_sync $(ClimaComms.device(xb)) kernel!($xb, $yb);
show(stdout, MIME("text/plain"), trial);
println("\n--------------------------- kernel, SmallStruct")
trial = BenchmarkTools.@benchmark ClimaComms.@cuda_sync $(ClimaComms.device(xb)) kernel!($xs, $ys);
show(stdout, MIME("text/plain"), trial);
nothing
@show DataLayouts.ncomponents(Spaces.local_geometry_data(space)) # 42
```
One thing that we can do to better understand potential improvement would be to measure field accesses in the benchmarks (which we could do on the cpu). This will inform us how much better a lazy LocalGeometry may help us for our kernels.
The LocalGeometry used to be more lazy, however, this resulted in _more_ memory accesses (i.e., caching more variables can lead to decreases or increases in bandwidth per kernel).
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the benchmark setup in perf/struct_bandwidth.jl and the provided BigStruct/SmallStruct CUDA experiment, using TestUtilities.jl to reproduce the space. Measure field accesses on the CPU and compare bandwidth for eager versus lazy LocalGeometry or IntervalMesh; done means the measurements clarify whether laziness improves the relevant kernels.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- performance
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100