Skip to content

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
mainfrom
loop-signature-padding-loopcontext
Open

Unify the e-ph loop argument lists, drop the batched dummy padding, and name LoopContext construction#14
jaemolihm wants to merge 4 commits into
mainfrom
loop-signature-padding-loopcontext

Conversation

@jaemolihm

Copy link
Copy Markdown
Owner

Motivation

Three TODO.md "## GPU" items, all readability debt in the two hottest loops in the package. PR #13
split use_gpu into backend + batched but left the loops themselves as they were:

  • the per-point and batched loops take divergent argument lists, so a reader comparing them cannot
    tell shared data from path-specific data;
  • both batched loops pad a partial final batch with duplicated valid data so the batched kernels
    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 = 0 sentinel) — flagged as a semantic trap in the
    2026-07-17 architecture review;
  • LoopContext is built from positional constructors disambiguated only by whether argument 3 is an
    Integer outer index or a UnitRange batch.

Pure refactor: no physics, no intended numerical change. bte_scattering_increments is untouched
and 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,). backend becomes a shared positional in both loops of all three drivers,
removing three hardcoded backend = CPUBackend() constants whose correctness a reader previously had
to 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 positional
list appears exactly once, which is what let a stale "backend is positional #9" claim drift into a
file 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_batch trailing-prefix view. This deletes the outer-k tail fill and its
fill!(P_mk, 1), and the outer-q tail fill including the wtk = 0 sentinel — the payload width is
now the only guarantee, identically at both sites.

One driver-side change was needed: get_eph_Rq_to_kq_batched! asserted its workspace sizes with
equality, so a shared max-width workspace threw on a partial batch. Its sibling
get_eph_kR_to_kq_batched! already solved this (>= assertions plus internal
view(ws.g, :, 1:nq)), so this mirrors that — making the two siblings consistent rather than adding a
convention. No consumer needed a dense fallback and no driver signature was widened:
view(::CuArray, :, :, 1:n) is a derived CuArray, not a SubArray, and reshaping keeps that.

C — named LoopContext construction. LoopContext(backend, ::BatchedMode; outer_index, batch, n_batch_max) is keyword-only, and the per-k context is visibly derived via with_outer_index. All
six 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-k
batch 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 EPDataKBatched padding paragraph, and README_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

suite result
runtests.jl, no CUDA 22519 / 22519 (3m58s)
runtests.jl, with CUDA (A100) 22636 / 22636 (2m19s)
test_gpu.jl on A100-80GB all testsets, incl. new GPU partial k-batch 5/5
test_calculator_contract.jl 27/27, incl. new LoopContext named construction 7/7
boltzmann/test_gpu_boltzmann_calculator.jl 8/8 incl. the GPU arm
MigdalEliashberg.jl/test/runtests.jl 41/41, both GPU testsets

Every reference deviation is bit-identical to the 1901141 baseline — outer-k CPU+batched vs
CPU+per-point 8.027231620421635e-16, outer-q 3.0427801012277363e-16, outer-q CPU-vs-GPU
8.114080269940629e-16, BTE Sᵢ 1.3380159109184602e-15, Sₒ 6.11269303778053e-16. The bit-exact
assertions still hold, including ME's cgb.g2 == cgf.g2 at nk = 173 with nk_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 generic
batched_gemm!/eigen_batched methods are per-slice mul!/syev!), and eigen_batched is safe on
both paths because the deleted padding was a duplicate of column nk_batch, leaving the set of
distinct matrices — and any heevjBatched! sweep count — unchanged. But cuBLAS selects its algorithm
from N non-monotonically at small M, so N-narrowing is reduction-order-relevant: at the outer-q
shape (M=48, K=700), N 64→32 differs by 3.34e-15, 32→16 by 3.95e-15, 20→6 by 3.97e-15. The
widths 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 = 66 at nk_outer_batch_max = 20 still gives 4 tiles
(20/20/20/6), so a nonzero Sᵢ tile_offset and a partial tail both still occur, with the
precondition guard intact. with_outer_index(ctx_batch, ik) reproduces the old 5-field context
exactly. A new unit test pins the >=/prefix-view workspace change directly (bitwise equal to the
independent ws === nothing fresh-allocation branch at nmodes = 4, with a permutation control
proving the arm non-vacuous).

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.
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.

1 participant