Unify the e-ph loop argument lists, drop the batched dummy padding, and name LoopContext construction - #14
Open
jaemolihm wants to merge 4 commits into
Open
Unify the e-ph loop argument lists, drop the batched dummy padding, and name LoopContext construction#14jaemolihm wants to merge 4 commits into
jaemolihm wants to merge 4 commits into
Conversation
The two batched-scope convenience constructors were told apart only by whether argument 3 was an `Integer` outer index or a `UnitRange` batch, and the call sites hand-wrote the `1:0` / bare-`ik` fillers. Replace them with one keyword-only `BatchedMode` constructor (`outer_index` / `batch` / required `n_batch_max`) plus `with_outer_index`, so argument 3 is never positional in `BatchedMode` and the per-k context is visibly the batch context at one index. Also documents that `n_batch_max` means the outer-k batch width in the outer-k loop and the inner-k batch width in the outer-q loop, hence that a `TiledDeviceOutput`-style consumer is outer-k-only. No arithmetic changes: CPU+batched vs CPU+per-point stays bit-identical (Pb 4^3 g2 8.027e-16, Pb 6^3 Si 1.3380159109184602e-15 / So 6.11269303778053e-16, outer-q 3.043e-16).
Both driver families now have their two loop shapes take the identical positional argument list -
everything BOTH need - and take only their own path-specific data as individual keyword arguments:
outer-k: (model, kpts, qpts, kqpts, el_k_save, el_kq_save, ph_save, precompute_ph, backend)
per-point kwargs: epstates, ep_ekpRs, epmat, ep_ekpR_obj, dyn_threads,
epmat_R, epobj_ekpR_R, ep_ekpR_Rs; batched kwarg: epmat_dev
outer-q: (model, kpts, qpts, el_k_save, ph_save, eph_buffers, backend)
per-point kwargs: kqpts, el_kq_save, precompute_el_kq; batched kwarg: el_ham_dev
`backend` is a shared positional rather than a batched-only kwarg (both shapes use it for the same
thing, the LoopContext), so the per-point loops no longer hardcode `CPUBackend()`. Same change in
`run_eph_over_k_and_q`, which has no batched twin - `backend` sits in the same slot #9 there.
Individual kwargs, not a bundled NamedTuple, so every name stays a typed argument of the lowered body
method: checked that all five `_loop_*` bodies still lower with zero `Core.Box`, and that the only
`Any` slots in their inferred specializations are the pre-existing ones (`@threads`' `err`, the
untaken `covariant_derivative_of_g` branch's `tmp`, and outer-q's `mask_kq`).
No arithmetic changes. Full CUDA-free suite 22515/22515; CPU+batched vs CPU+per-point unchanged to
the last bit (Pb 4^3 outer-k g2 8.027231620421635e-16, outer-q 3.0427801012277363e-16; Pb 6^3
Si 1.3380159109184602e-15 / So 6.11269303778053e-16).
…ph loops
Both batched loops padded a partial final k-batch with duplicated valid data so the batched
interpolation/eigensolve ran on dense `nk_batch_max`-sized arrays. They now hand those drivers
contiguous width-`nk_batch` trailing-prefix views of the same max-width staging buffers - the
convention both loops already used on their q axis - so a partial batch computes only its own columns
and there is no padded tail to reason about.
Outer-k (`run_eph_over_k_and_kq`): the tail fill and the defensive `fill!(P_mk, 1)` are gone;
`get_eph_RR_to_kR_batched!` takes four width-`nk_batch` views (it derives the width from
`size(uks, 3)` and asserts the rest, so this is inside its existing contract); the H2D of `uks` is the
5-arg linear form (a host<->device SubArray copy would fall back to scalar indexing).
Outer-q (`run_eph_over_q_and_k`): the tail fill - including the `wtk = 0` sentinel that neutralised it
- is gone; the Fourier, eigensolve, window mask, Rq->kq and polar-dipole calls all run at `nk_batch`;
`Hkq_flat`'s (nw, nw, k) reshape is hoisted next to the buffer block so each per-batch argument is a
plain prefix view of a dense array; three 5-arg linear H2D copies.
Driver-side, the only change needed was in `get_eph_Rq_to_kq_batched!`, whose workspace assertions
were equalities and so rejected a partial batch: they are now `>=` with internal `view(ws.g, :, 1:nk)`
etc., mirroring what its sibling `get_eph_kR_to_kq_batched!` already did.
Docs moved with the code: the `EPDataKBatched` padding paragraph (the payload width is now the ONLY
width guarantee - there is no zeroed `wtk` tail behind it), README_GPU's outer-q bullet, and the
README_GPU "Deferred" entry that proposed exactly this change.
Verification. Item B is BIT-NEUTRAL, not round-off-visible: narrowing a batch changes the batched-GEMM
batch count and the GEMM's N, never the contraction length, so no per-element sum is reassociated.
Every reference number is unchanged to the last digit, so the two literals in
test_gpu_boltzmann_calculator.jl's comment were re-measured and left as they are:
outer-k CPU+batched vs CPU+per-point, Pb 4^3, nk_outer_batch_max=5/nq_batch_max=7 (13 k-batches,
last width 4; 10 q-tiles, last width 1) g2 8.027231620421635e-16
outer-q CPU+batched vs CPU+per-point, Pb 4^3, nk_batch_max=10 (7 batches, last width 4)
A 3.0427801012277363e-16
outer-q partial k-batch CPU vs GPU, nk_batch_max=10 A 8.114080269940629e-16
BTE Pb 6^3 CPU+batched vs CPU+per-point (nk=66, nk_outer_batch_max=20 -> 4 Si tiles at nonzero
tile_offset, last k-batch width 6) Si 1.3380159109184602e-15
So 6.11269303778053e-16
Bit-exact tests all still hold: GPU multi-q-tile `cg.omega_q == cg7.omega_q`, and ME's
full-resident-vs-block-resident `cgb.g2 == cgf.g2` at nk=173 / nk_outer_batch_max=4 (44 batches, last
width 1). Full CUDA-free suite 22519/22519; test_gpu.jl + BTE + MigdalEliashberg 41/41 (both GPU
testsets) green on an A100-80GB.
New unit coverage `check_eph_partial_view_k` (the k-axis analogue of the existing q-axis
`check_eph_partial_view`): full-width vs width-m prefix-view runs of `get_eph_RR_to_kR_batched!` and of
`get_eph_Rq_to_kq_batched!` against one shared max-width workspace, asserting both the first m slices
and that the tail was left untouched - on the CPU (no CUDA) as well as the GPU.
…onvention comment Review round 1 on `1901141..e0b3923`. Comments, docs and one local-hoist only - no behaviour change, and every reference number is unchanged (see below). F2 - `run_eph_over_k_and_q.jl` claimed `backend` is positional #9 "the same slot it occupies in the two loop pairs that have a batched twin". It is #7 in the outer-q pair, whose shared list has seven elements. Reworded to "the last shared positional"; the same false #9 removed from `TODO.md` bullet 4, which contradicted its own 7-element list two lines earlier. F5 - that drift happened because the item-A convention was explained in four in-source blocks with the positional list spelled out in three. The block above `_loop_eph_over_k_and_kq` is now marked CANONICAL and carries the convention plus the individual-kwargs/Core.Box rationale; the two batched blocks are one-line pointers to it, and the outer-q per-point block gives only its own (different) list. Each positional list now appears exactly once. F3 - the outer-q batch body repeated the same four view expressions 11 times, after which the payload comment had to assert in prose that they were the same width-`nk_batch` views the drivers were handed. Hoisted `ep_v` / `Uk_v` / `Ukq_v` / `ks_v` next to the H2D copies (the window mask writes through `Ukq_v` too) and dropped the prose assertion: the invariant is now an identity, so a future edit cannot narrow one range and not the others. The outer-k site reuses `rng_k` in four spots and is left alone. F4 - `src/EPData.jl` said a consumer sizing itself from `ctx.n_batch_max` would read "uninitialised" columns. For every batch after the first the tail holds the PREVIOUS batch's values - finite and plausible, i.e. worse than uninitialised. Reworded. Dash sweep of four pre-existing non-ASCII dashes (all present at `1901141`): U+2212 at `run_eph_over_k_and_kq.jl:754, 762, 827` and U+2013 at `README_GPU.md:279`. Em-dashes (U+2014, the repo's prose style) untouched. One U+2212 remains in `ext/ElectronPhononCUDAExt.jl:25` (a comment in a file this PR does not otherwise touch); `TODO.md:5` quotes both characters on purpose. F1 was a correction to the RECORD, not the code: the claim in `e0b3923`'s message that item B is bit-neutral is over-broad on the GPU. The reviewer refuted the reasoning on the A100 at the shape the outer-q loop uses (`get_fourier_batched!` -> `mul!`, M=48, K=700): narrowing N 64->32 gives max rel 3.34e-15, 32->16 3.95e-15, 20->6 3.97e-15, because cuBLAS selects its algorithm from N non-monotonically at small M. Nothing shipped is affected - the 10->4 width the tests hit and every outer-k RR->kR shape came back exactly 0.0, and the two literals in `test_gpu_boltzmann_calculator.jl`'s comment are CPU-only measurements on a CPU-only path. The correct statement, now in the plan's PROGRESS log: empirically bit-neutral on CPU by construction (OpenBLAS's K-blocking order is N-independent at every shape probed) and at the GPU batch widths tested, not guaranteed on GPU. `eigen_batched` is safe on both backends by a stronger argument: the deleted padding was a duplicate of column `nk_batch`, so the set of DISTINCT matrices - and any batch-level Jacobi sweep count in `heevjBatched!` - is unchanged. Verification: full CUDA-free suite 22519/22519; `test_gpu.jl`, `test_calculator_contract.jl`, BTE 8/8 and MigdalEliashberg 41/41 (both GPU testsets, incl. the bit-exact `cgb.g2 == cgf.g2`) all green. Reference deviations identical to `e0b3923`: outer-k 8.027231620421635e-16, outer-q 3.0427801012277363e-16, outer-q CPU-vs-GPU 8.114080269940629e-16, BTE Si 1.3380159109184602e-15 / So 6.11269303778053e-16.
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.
Motivation
Three
TODO.md"## GPU" items, all readability debt in the two hottest loops in the package. PR #13split
use_gpuintobackend+batchedbut left the loops themselves as they were:tell shared data from path-specific data;
run on dense
nk_batch_max-sized arrays, and the two sites neutralize the tail differently(outer-k trims to views, outer-q relied on a
wtk = 0sentinel) — flagged as a semantic trap in the2026-07-17 architecture review;
LoopContextis built from positional constructors disambiguated only by whether argument 3 is anIntegerouter index or aUnitRangebatch.Pure refactor: no physics, no intended numerical change.
bte_scattering_incrementsis untouchedand no
ext/file changes.Change
A — one positional list per loop pair. The per-point and batched loops now take identical
positional arguments, with path-specific data as keywords: per-point
(epstates, ep_ekpRs, epmat, ep_ekpR_obj, dyn_threads, epmat_R, epobj_ekpR_R, ep_ekpR_Rs),batched
(epmat_dev,).backendbecomes a shared positional in both loops of all three drivers,removing three hardcoded
backend = CPUBackend()constants whose correctness a reader previously hadto re-derive from the entry validation. The argument-convention comment is now stated once
(canonical block above
_loop_eph_over_k_and_kq) with one-line pointers elsewhere — each positionallist appears exactly once, which is what let a stale "
backendis positional #9" claim drift into afile where it is #7.
B — the dummy padding is gone from both sites. Partial tails are no longer computed at all: every
consumer takes a width-
nk_batchtrailing-prefix view. This deletes the outer-k tail fill and itsfill!(P_mk, 1), and the outer-q tail fill including thewtk = 0sentinel — the payload width isnow the only guarantee, identically at both sites.
One driver-side change was needed:
get_eph_Rq_to_kq_batched!asserted its workspace sizes withequality, so a shared max-width workspace threw on a partial batch. Its sibling
get_eph_kR_to_kq_batched!already solved this (>=assertions plus internalview(ws.g, :, 1:nq)), so this mirrors that — making the two siblings consistent rather than adding aconvention. No consumer needed a dense fallback and no driver signature was widened:
view(::CuArray, :, :, 1:n)is a derivedCuArray, not aSubArray, and reshaping keeps that.C — named
LoopContextconstruction.LoopContext(backend, ::BatchedMode; outer_index, batch, n_batch_max)is keyword-only, and the per-k context is visibly derived viawith_outer_index. Allsix driver sites use named forms; the field-wise constructor stays available for the unit tests that
deliberately probe field combinations production never builds.
n_batch_max's dual meaning (outer-kbatch width in one loop, inner-k in the other) is now documented, with the note that a
TiledDeviceOutput-style consumer is outer-k-only.Docs moved with the code: the
EPDataKBatchedpadding paragraph, andREADME_GPU.md's Deferred"View-instead-of-fill for the outer-q staging buffers … benchmark it on its own before adopting" entry
is deleted rather than left contradicting the code.
Test
runtests.jl, no CUDAruntests.jl, with CUDA (A100)test_gpu.jlon A100-80GBGPU partial k-batch5/5test_calculator_contract.jlLoopContext named construction7/7boltzmann/test_gpu_boltzmann_calculator.jlMigdalEliashberg.jl/test/runtests.jlEvery reference deviation is bit-identical to the
1901141baseline — outer-k CPU+batched vsCPU+per-point
8.027231620421635e-16, outer-q3.0427801012277363e-16, outer-q CPU-vs-GPU8.114080269940629e-16, BTESᵢ 1.3380159109184602e-15,Sₒ 6.11269303778053e-16. The bit-exactassertions still hold, including ME's
cgb.g2 == cgf.g2at nk = 173 withnk_outer_batch_max = 4(44 batches, last width 1 — the strongest partial-tail GPU arm in the stack).
Bit-neutrality is not a guarantee on the GPU, and the record says so. Narrowing a batch is
bit-neutral on the CPU by construction (OpenBLAS's K-blocking order is
N-independent; the genericbatched_gemm!/eigen_batchedmethods are per-slicemul!/syev!), andeigen_batchedis safe onboth paths because the deleted padding was a duplicate of column
nk_batch, leaving the set ofdistinct matrices — and any
heevjBatched!sweep count — unchanged. But cuBLAS selects its algorithmfrom
Nnon-monotonically at small M, so N-narrowing is reduction-order-relevant: at the outer-qshape (M=48, K=700), N 64→32 differs by
3.34e-15, 32→16 by3.95e-15, 20→6 by3.97e-15. Thewidths this PR's tests actually hit (10→4, and every outer-k RR→kR shape) come back exactly
0.0,which is why the measured deviations are unchanged — but do not add a bit-exact GPU regression across
a partial tail on the strength of it.
Coverage added by PR #13 is preserved:
nk = 66atnk_outer_batch_max = 20still gives 4 tiles(20/20/20/6), so a nonzero Sᵢ
tile_offsetand a partial tail both still occur, with theprecondition guard intact.
with_outer_index(ctx_batch, ik)reproduces the old 5-field contextexactly. A new unit test pins the
>=/prefix-view workspace change directly (bitwise equal to theindependent
ws === nothingfresh-allocation branch atnmodes = 4, with a permutation controlproving the arm non-vacuous).