CliMA / CliMA/ClimaCore.jl

Parallelize `materialize!(::FieldNameDict,::FieldNameDict)`

Open
#1,659 2 comments 0 reactions 0 assignees View on GitHub
enhancement performance
Dominant language
Julia
Stars
117
Forks
19
Avg merge
3d 4h
Merged PRs (30d)
41

Description

It's occurred to me that [`@. cache.b₂′ = b₂ - A₂₁ * x₁`](https://github.com/CliMA/ClimaCore.jl/blob/30b823c4fe9592df2b3a516d8de16d1ad0e7fd0d/src/MatrixFields/field_matrix_solver.jl#L311) lowers to [this](https://github.com/CliMA/ClimaCore.jl/blob/30b823c4fe9592df2b3a516d8de16d1ad0e7fd0d/src/MatrixFields/field_name_dict.jl#L508-L527):

```julia
function Base.Broadcast.materialize!(
dest::FieldNameDict,
vector_or_matrix::FieldNameDict,
)
!is_lazy(dest) || error("Cannot materialize into a lazy FieldNameDict")
is_subset_that_covers_set(keys(vector_or_matrix), keys(dest)) || error(
"Broadcast result and destination keys are incompatible: \
$(set_string(keys(vector_or_matrix))) vs. $(set_string(keys(dest)))",
) # It is not always the case that keys(vector_or_matrix) == keys(dest).
foreach(keys(vector_or_matrix)) do key
entry = vector_or_matrix[key]
if dest[key] isa UniformScaling
dest[key] == entry || error("UniformScaling is immutable")
elseif entry isa UniformScaling
dest[key] .= (entry,)
else
dest[key] .= entry
end
end
end
```
So, there are even more kernel launches than I realized in `run_field_matrix_solver!`. We should fuse this loop into a single kernel launch, and parallelize if possible.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with materialize!(::FieldNameDict, ::FieldNameDict) in src/MatrixFields/field_name_dict.jl, then inspect its use from run_field_matrix_solver! in src/MatrixFields/field_matrix_solver.jl. Determine how the loop can be fused into one kernel launch and parallelized; done means the materialization no longer causes the extra per-entry launches while preserving the existing key and UniformScaling behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
hpc, performance
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.