Skip to content

Parallelize ClimaCore sparse matrix assembly + polygon intersections#110

Merged
asinghvi17 merged 7 commits into
mainfrom
as/parallelize-sparsematrix
Jun 23, 2026
Merged

Parallelize ClimaCore sparse matrix assembly + polygon intersections#110
asinghvi17 merged 7 commits into
mainfrom
as/parallelize-sparsematrix

Conversation

@asinghvi17

@asinghvi17 asinghvi17 commented Jun 4, 2026

Copy link
Copy Markdown
Member

Adds a new interface for the intersection operator, so that ClimaCore only needs to define its intersection operator, not the full sparse matrix assembly.

Also adds some benchmarks for the ClimaCore regridder construction, as well as performance optimisation / improvements there.

@asinghvi17 asinghvi17 marked this pull request as ready for review June 4, 2026 14:37
@akshaysridhar akshaysridhar self-requested a review June 5, 2026 16:45
@asinghvi17 asinghvi17 changed the title Parallelize ClimaCore sparse matrix ops Parallelize ClimaCore sparse matrix assembly + polygon intersections Jun 6, 2026
asinghvi17 and others added 7 commits June 10, 2026 20:27
The standard regridder's parallel COO assembly is now driven by one trait
plus two hooks on the intersection operator, each with a default that
reproduces today's area-based weight:

- `IntersectionReturnStyle(op)`: `OutOfPlaceSingleResult` (default;
  `op(src, dst) -> area`, the driver stores the triplet) vs `InPlace`
  (`op(rows, cols, vals, item, src_tree, dst_tree)` pushes its own COO).
- `work_items(op, candidate_pairs)`: the units of parallel work
  (default: one candidate pair each).
- `output_matrix_size(op, src_tree, dst_tree)`: the assembled shape
  (default: dst-cells x src-cells).

`intersection_areas` collapses to a single method (`threaded::BoolsAsTypes`)
that resolves the return-style trait once and threads it through the shared
`_parallel_coo` / `_assemble_chunk` / `_run_and_store!` skeleton; the removed
`compute_intersection_areas` becomes the `OutOfPlaceSingleResult` branch.
`DefaultIntersectionOperator` uses every default, so the standard path is
behaviorally unchanged.

The interface is `@public` so downstream packages can define their own
parallel-assembling operators.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Both spectral-element assemblers previously built their sparse matrices in a
hand-written serial loop. They now plug into the shared parallel
`intersection_areas` driver as two `InPlace` operators:

- `SEToFVIntersectionOperator`: per (elem, cell) pair, pushes an Nq^2 block
  of B-weights per intersection polygon (matrix N_fv x N_nodes).
- `FVToSEIntersectionOperator`: overrides `work_items` to make the element
  (not the candidate pair) the unit of work, since the per-element local
  mass-matrix solve needs all of an element's cells at once
  (matrix N_nodes x N_fv).

Each entrypoint shrinks to constructing an operator and calling
`intersection_areas`; the private `_get_candidate_pairs` and the two assembly
loops are gone. Assembly at h_elem=32 / 8 threads is ~5.5x faster.

Extra kwargs (e.g. `normalize`) are absorbed by the entrypoint signatures and
not forwarded, matching the previous behavior.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
On a small cubed sphere, assert the threaded and serial intersection matrices
are exactly equal for both SE->FV and FV->SE and both regular and Gilbert
(space-filling-curve) element orderings. There are no cross-pair duplicate
(row,col) entries in either direction, so chunking cannot change any sum,
making this a tight data-race / ordering regression check.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add a "Customizing the weights" section to how_it_works.md (with
OutOfPlaceSingleResult and InPlace worked examples), render the seven new
public symbols in the API reference, and update the CLAUDE.md architecture
notes to describe the trait + two hooks and the ClimaCore operators.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
`GO.area(::Spherical, polygon)` infers `Any` (GeometryOps v0.1.40 boxes a
multi-method local closure — JuliaGeo/GeometryOps.jl#407), so
spherical_triangle_area returned `Any` and boxed every quadrature weight in
accumulate_principled_b (~4000 heap allocs per intersection polygon). Add a
`::Float64` barrier on the area call, and build the ring container with a
static `SA[...]` to drop its per-triangle allocation.

Value-identical (ClimaCore suite green). SE→FV assembly: ~1624 → ~45
allocs/pair (~36x), 380 MiB → 48 MB; FV→SE benefits via the shared accumulator.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Tighten the docstrings introduced in this branch (the intersection-operator
interface in intersection_areas.jl and the ClimaCore SE/FV extension) without
dropping any content — all PDF equation references and @ref links preserved.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@asinghvi17 asinghvi17 force-pushed the as/parallelize-sparsematrix branch from 27f8110 to 9b66683 Compare June 11, 2026 03:28
@akshaysridhar akshaysridhar requested a review from hgpeterson June 11, 2026 16:12
@hgpeterson

Copy link
Copy Markdown
Collaborator

I've looked over the changes and verified that tests pass locally (including examples/low_level/climacore_exp.jl). Most of the PR goes over my head, but I get the general idea of parallelizing some of these operations. I question whether this added complexity in the code is really necessary, though. Running the benchmarks in bench on a single thread takes ~5 seconds for the slowest case while on 8 threads the slowest case is ~1.5 seconds. Are there any use cases where these computations need to occur more than once? From what I understand, we only need to do them once at the beginning of the simulation.

@hgpeterson

Copy link
Copy Markdown
Collaborator

BTW: I also ran a /code-review in claude and it had 10 comments. They were mostly small stuff from what I could tell, but happy to share them all if desired. The one it flagged as most important was:


src/regridder/intersection_areas.jl:161

Summary: _parallel_coo's threaded branch crashes with MethodError: reducing over an empty collection when items is empty, because ChunkSplitters.chunks([]; n=...) yields 0 partitions so all_results = [] and reduce(vcat, getindex.([], 1)) has no init.

Failure scenario: Verified live: ConservativeRegridding.intersection_areas(GO.Planar(), GOCore.True(), dst_tree, src_tree) for two non-overlapping grids throws this error, while False() correctly returns an empty sparse matrix. Since Regridder(dst, src) defaults to threaded = True(), any Regridder construction between grids/spaces with zero candidate pairs (disjoint domains, a regional patch, or — newly via this PR — an SE element set that doesn't overlap an FV grid) crashes instead of producing an empty/all-zero matrix. The old SE↔FV code built sparse(rows,cols,vals,...) directly and handled this fine regardless of threaded; this PR routes those constructors through the same vulnerable shared path for the first time.


Should be an easy fix, but would we ever call Regridder for a case with no overlaps?

@asinghvi17

asinghvi17 commented Jun 13, 2026

Copy link
Copy Markdown
Member Author

Hmm for that one it's certainly possible that one might try to construct a reusable regridder on two non-overlapping areas (though it would be a bit dumb). The fix should be simple enough that we can just do it there.

In general the reason for complexity is twofold:

  1. Polygon intersection (and the other things required for spectral element) is the heaviest part of the workload, so parallelizing that is always nice. Especially for the large (e.g. two 3600x1800 grids) cases, this becomes quite essential to having a reasonable runtime, since those can take 200-500 seconds.
  2. This allows the user to specify exactly what workflow they want to use, so e.g. for an intersection-grid for flux regridding, you can hook in to all the infrastructure and it just works, without having to reimplement the parallelism etc. Especially if you might want to collect extra information, having these hooks becomes very useful.

@hgpeterson

Copy link
Copy Markdown
Collaborator

Hi @asinghvi17, sorry for the late reply; was at a conference.

I can see how very high resolution setups would benefit from this change. If that's a high priority right now, then I'm okay with you merging this in as it doesn't seem to affect our current ClimaCoupler runs.

@asinghvi17

Copy link
Copy Markdown
Member Author

will fix docs in a follow up PR

@asinghvi17 asinghvi17 merged commit f9fe362 into main Jun 23, 2026
2 of 4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants