Conversation
…mpute_zdiff_gradp Replaces O(nedges * nlev^2) nested jk/jk1 search loops with O(nedges * nlev * log(nlev)) np.searchsorted on reversed z_ifc slices. Eliminates per-iteration boolean array allocation and array_ns.where calls. Phase 1: batch searchsorted for all z_me[jk] per edge per side. Phase 2: single searchsorted for z_aux2, applied via boolean mask. Largest grid (exclaim_ch_r04b09_dsl): 2.95s -> 0.93s (3.2x)
…zdiff_gradp Removes the (nedges, 2, nlev) vertidx_gradp allocation and the final vertoffset_gradp subtraction step. Computes vertoffset_gradp offsets directly in the per-edge loop. Initializes vertoffset_gradp as zeros instead of broadcasting jk_field.
- Pre-compute z_ifc_asc = z_ifc[:, ::-1] to avoid per-edge reversal - Extract e2c_0, e2c_1 column arrays to avoid per-edge indexing - Unroll for side in range(2) into explicit side 0 and side 1 - Hoist jk_field_slice and z_me_slice out of side processing
The zdiff_gradp computation is now fast enough (0.92s vs 2.95s baseline on largest grid) that the factory test no longer needs the cpu_only guard. The vwind_impl_wgt speed concern is also resolved.
Replaces 62K per-edge searchsorted calls with 4 batched 2D calls using the row-offset trick. Detects CuPy vs NumPy at runtime. All data stays on-device (no per-iteration D2H transfers). Largest grid (exclaim_ch_r04b09_dsl): 2.95s -> 0.74s (4.0x)
Replaces _get_xp / _cp / _np with the array_ns returned by data_alloc.array_namespace(), which already provides backend- agnostic searchsorted, clip, where, take_along_axis, etc.
- Remove unnecessary z_me_m masking (invalid jk values are masked out by valid_jk after the search) - Remove .copy() from z_ifc_asc (fancy indexing creates new arrays) - Compute fill_high from z_ifc directly (fewer elements than z_ifc_e0) - Remove zdiff_gradp exchange calls from function, delegate to factory via do_exchange=True
The z_me halo exchange inside compute_zdiff_gradp was redundant: z_mc already has valid halo values from the factory dependency chain (Z_MC provider exchanges its halo). Since z_me is computed locally from z_mc[e2c], it inherits correct halo values without an additional exchange. - Remove exchange parameter from compute_zdiff_gradp signature - Remove functools.partial wrapping in factory registration - Remove dims and decomposition imports (no longer needed) - Update unit test call site
Contributor
Author
|
cscs-ci run distributed |
Contributor
Author
|
cscs-ci run distributed |
…n, and endpoint fallback Implements iter-001 design D1-D6: - New compute_zdiff_gradp_v2 next to existing baseline. - _batched_searchsorted_v2 keeps reductions on-device and casts to float64. - Host-side slicing at horizontal_start / horizontal_start_1 replaces device masks. - Global fill bounds over all searched arrays (D3). - Phase-2-only endpoint fallback where pos_aux >= nlev+1-fi (D6). - Gated validation via ICON4PY_VALIDATE_ZDIFF_GRADP env var (D4). Tests: - Datatest now runs both baseline and v2 variants with validation enabled. - Added endpoint-forcing unit test asserting v2 matches main-semantics golden and baseline diverges, run with validation ON and OFF.
… round-robin benchmark, NaN test fix, hoist invariant [R30-R34]
… scan for cupy compatibility
- Add cupy first-match-scan fast path variant compute_zdiff_gradp_exact_v3. - Numpy path delegates to compute_zdiff_gradp_exact_v2 unchanged. - Add _first_match_scan_reference numpy oracle and boundary-case unit test. - Wire exact_v3 into existing datatest and nan-validation parametrizations.
- Add numpy-exact3 / cupy-exact3 variants to Tier A and Tier B. - Add Tier A internal consistency checks: reference scan vs numpy-exact2, cupy-exact3 vs numpy-exact2. - Extend endpoint forcing and weird-input probes with exact_v3 paths. - Clean up pre-existing ruff findings in the temporary benchmark file.
… and route dispatch to exact fallback [skip ci]
…codegen blocked) [skip ci]
…d positions [skip ci]
…][R64][F5][F6] [skip ci]
…st path [skip ci]
…el compile test [skip ci]
…implementation - Single compute_zdiff_gradp function using array_namespace (numpy/cupy), no backend dispatch, no raw kernels, no validation bundle, no globals. - Implements main's bracket predicate and jk_start carry exactly. - Correct for all finite inputs: non-monotone z_ifc, interior ties, non-increasing z_me, nlev=1. - Remove all variant functions, dispatch, exact* helpers, and LAST_*_PATH globals. - Delete compute_zdiff_gradp_gt4py.py. - Update tests to compare only the new function against _main_reference and the serialized datatest golden.
… variant Adds a premise-free, performance-oriented alternative to compute_zdiff_gradp in the same module. Cell 0 keeps the existing batched first-match; cell 1 and phase-2 cells use a Hillis-Steele doubling suffix-min successor table with a gather-scan carry. Edges are processed in memory-capped chunks. - _successor_table builds (E, nq, nlev) first-match tables. - _carry_gather_scan advances the carried lower bound via table lookup. - _zdiff_gradp_chunk_size caps per-chunk successor-table memory to ~256 MiB. - compute_zdiff_gradp_fast shares the same signature as compute_zdiff_gradp. - Tests are parametrized over both functions and pass against _main_reference.
When queries is 1-D (phase 2, constant z_aux2 per edge), the bracket (upper >= q[:,None]) & (q[:,None] >= lower) is the same for every jk. Compute it once before the loop instead of nlev times.
…unction Replace premise-free sequential-carry implementation with a single compute_zdiff_gradp function using batched searchsorted. Recovers the per-edge offset trick from the old v2 (2e5c5a7) and uses side='right' to match the reference tie handling. Adds a Hillis-Steele scan for cell 1 to reproduce the reference jk_start carry when E3 is violated. Cites vertical.py:558 and vertical.py:625 for the z_ifc monotonicity proof.
…ation Drop compute_zdiff_gradp_fast parametrization. P2 (zero-thickness) now passes with the searchsorted function, so keep it as a regular test. P4 (non-monotone z_ifc) remains xfail because it violates E1, which searchsorted requires.
side='right' caused a cupy CPU fallback (4723ms vs 17ms). Reverted to side='left' (GPU-accelerated) and added a uniform vectorized where that corrects the off-by-one when the query exactly equals a z_ifc boundary (side='left' returns jk+1, Fortran first-match takes jk). Also fixed the _batched_searchsorted offset to prevent cross-row overlap when fill values place queries outside the row's own range (nlev=1, fi=0).
The tie correction (take_along_axis gather of (491520, 60) int64 indices into z_ifc, x4 searchsorted calls) cost 8x: 27s numpy / 5s GPU vs 3.4s / 17ms without. Detecting a tie costs the same as correcting it (both need the gather), so there is no cheap middle ground. The premise (z_me strictly between z_ifc boundaries, no exact ties) holds for all production VerticalGridConfig: z_ifc layers are separated by millions to billions of ULPs (enforced by _check_and_correct_layer_thickness, vertical.py:625), so each midpoint z_mc[c,j] = 0.5*(z_ifc[c,j]+z_ifc[c,j+1]) is a distinct float strictly between its boundaries, and z_me (a convex combination of two such midpoints) is strictly between them as well. The module docstring now states this argument with citations. P1 (interior tie, forces z_me == z_ifc boundary by direct assignment) is xfailed with the same reasoning as P2/P3/P4: synthetic input that the grid builder does not produce.
…r reduction (R66)
Contributor
Author
|
cscs-ci run default;MODEL_SUBPACKAGES=common:driver;SESSIONS=model;MODEL_SUBSETS=datatest |
Contributor
Author
|
cscs-ci run default;MODEL_SUBPACKAGES=common:driver;SESSIONS=model;MODEL_SUBSETS=datatest;LEVELS=integration |
…long sentences (R71, R72) [skip ci]
|
When developing, you can test your changes on CSCS CI before merge with the You can pass options to override pipeline variables, for example:
Available options are:
For each option, See The Merging Once your PR is approved and ready for merging, add it to the merge queue. The Optional Tests To run benchmarks you can use:
For more detailed information please look at CI in the EXCLAIM universe. |
Contributor
There was a problem hiding this comment.
Pull request overview
Optimizes pressure-gradient metrics for GPU execution through batched array operations and shared edge-height computation.
Changes:
- Replaces nested host loops with batched
searchsorted. - Adds shared
z_meand selective halo exchanges. - Expands regression, randomized, and edge-case tests.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
compute_zdiff_gradp.py |
Implements vectorized bracket searches. |
metric_fields.py |
Reuses precomputed z_me. |
metrics_factory.py |
Registers z_me and selective exchange. |
states/factory.py |
Adds exchange-field filtering. |
test_compute_zdiff_gradp.py |
Adds reference and probe tests. |
test_metric_fields.py |
Updates flat-index test inputs. |
test_metrics_factory.py |
Enables GPU factory testing. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| + 1.0 | ||
| ) | ||
| r = max_num * array_ns.arange(m, dtype=array_ns.float64)[:, None] | ||
| p = array_ns.searchsorted((a + r).ravel(), (v + r).ravel()).reshape(v.shape) |
| array_ns.expand_dims(z_me, axis=1)[horizontal_start:, :, :] | ||
| - z_mc[e2c][horizontal_start:, :, :] | ||
| zdiff_gradp[hs:, :, :] = array_ns.expand_dims(z_me, axis=1)[hs:, :, :] - z_mc[e2c][hs:, :, :] | ||
| vertoffset_gradp = array_ns.zeros((nedges, 2, nlev), dtype=gtx.int32) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR makes the
compute_zdiff_gradpmetrics kernel fast on GPU by replacing the per-element host access in nested Python loops with a single batchedsearchsortedover all edge rows. The implementation uses the array namespace (numpy/cupy), so the same code runs on CPU and GPU without backend dispatch.compute_zdiff_gradpnow takes a precomputedz_meedge-level field and returnszdiff_gradpandvertoffset_gradp. It reverses each cell column to ascending order, masks levels aboveflat_idx, and offsets every row by a float64 constant large enough to keep cross-row queries separate. One flattenedsearchsortedcall finds the bracket index for all(edge, level)queries at once; a clip to[flat_idx, nlev-1]gives the samejk_startas the sequential reference loop.The change relies on two grid invariants documented in the module docstring:
z_ifcis strictly decreasing in the vertical (enforced by the grid builder), andz_meis non-increasing per edge. Under those invariants the searchsorted result is non-decreasing injkand matches the reference first-match scan. Synthetic probe inputs that violate the invariants are markedxfailwith the reason recorded in the test file.z_meis extracted into its own registeredProgramFieldProviderwithdo_exchange=True, socompute_wgtfac_e,compute_flat_max_idx, andcompute_zdiff_gradpconsume a single exchangedz_meinstead of computing it inline.compute_wgtfac_enow acceptsz_medirectly and no longer performs an exchange. The field providers instates/factory.pygained an optionalexchange_fieldsfilter so thecompute_zdiff_gradpprovider can exchange onlyzdiff_gradpand not its inputs.Tests in
model/common/tests/common/metrics/unit_tests/test_compute_zdiff_gradp.pywere updated: the existing datatest now runs against the new signature, a deterministic element-by-element reference mirrors the main-branch semantics, random small and full-scale cases compare the vectorized result against that reference, and probe inputs exercise the documented invariants.Benchmark at full scale (491520 edges x 60 levels, GH200, minimum over repetitions):
🤖 Opened by an agent on behalf of @msimberg.