diff --git a/README_GPU.md b/README_GPU.md index cfa7fc5..caeacb6 100644 --- a/README_GPU.md +++ b/README_GPU.md @@ -196,7 +196,9 @@ payload, named for which momentum is the outer loop and which is batched on the `eph_batched_bytes_per_point(calc, EPDataKBatched)` for the loop's memory-adaptive batch sizing. The k side is streamed per k-batch (host-staged, no whole-grid device stack), and the payload is trimmed to the batch's actual width — a consumer reads its own size from any field - (e.g. `size(eps, 4)`) and never sees a padded tail (the outer-k convention). + (e.g. `size(eps, 4)`) and never sees a padded tail (the outer-k convention). Both loops consume + their max-width staging buffers through contiguous trailing-prefix views, so a partial final batch + computes only its own columns; nothing is padded with duplicated data. - Both loops fold their device-buffer byte accounting into `src/calculator/eph_device_staging.jl`: `_outer_{k,q}_staging_bytes(…)` return the loop's `(per_point, committed)` device-byte counts, and `plan_batch(backend, per_point, committed, cap; …)` turns those into the memory-adaptive batch @@ -254,10 +256,6 @@ unchanged). No window handling is needed in the calculator beyond addressing its benchmarked and validated bit-identical, but the gain is small on the GPU (CUDA's pool already recycles device buffers), so it is deferred. Best done together with the calculator loop, where one workspace allocated at loop setup is reused across all (k, q). -- **View-instead-of-fill for the outer-q staging buffers.** `run_eph_over_q_and_k`'s per-batch - staging fills padded per-k buffers (`Uk_batch`, …). Passing `nk_batch` and taking width-`nk` views - into the padded buffers (instead of filling) would avoid the copy. A hot-path change with no - correctness component, so benchmark it on its own before adopting. - **Energy window and long-range/polar on the GPU** — left on the CPU per-k path for now. - **MPI / multi-GPU** for the GPU loop — not in this foundation. - **Backend as a type parameter instead of a backend object (future).** The `use_gpu` keyword is @@ -278,7 +276,7 @@ unchanged). No window handling is needed in the calculator beyond addressing its - **A QR-based batched eigensolve (future).** The batched eigensolve uses `CUSOLVER.heevjBatched!` (Jacobi). A QR-based `HEEV` (e.g. via cuSolverDx) may be faster for the small matrices here; worth evaluating, but not in this PR. Accuracy is not a motivation — Jacobi is already at - machine precision. Note that `cusolverDnXsyevBatched` is *not* the answer: it is 2–3× faster per + machine precision. Note that `cusolverDnXsyevBatched` is *not* the answer: it is 2-3× faster per matrix but needs ~1.05 MB of workspace per matrix (65× heevj), which caps one call at ~72k matrices on an 80 GB A100 and competes with the device-resident e-ph tiles. - **Parametrize `Model` over its `WannierObject` array type (future).** Widening `WannierObject` diff --git a/TODO.md b/TODO.md index 2bbe848..3bd2925 100644 --- a/TODO.md +++ b/TODO.md @@ -14,22 +14,24 @@ https://github.com/JuliaMolSim/DFTK.jl/blob/master/src/architecture.jl and follow the pattern in https://docs.dftk.org/stable/developer/gpu_computations/ . -- [ ] Remove all the `nk_batch+1:nk_batch_max` dummy padding fills (e.g. - `src/calculator/run_eph_over_k_and_kq.jl` ~L702) by passing the actual `nk_batch` through to the - batched interpolation/kernels instead of padding partial tails with duplicated valid data to run on - dense `nk_batch_max`-sized arrays. +- [x] Remove all the `nk_batch+1:nk_batch_max` dummy padding fills — done at both sites + (`run_eph_over_k_and_kq.jl`, `run_eph_over_q_and_k.jl`). The batched interpolation/eigensolve/kernels + are handed contiguous width-`nk_batch` trailing-prefix views of the max-width staging buffers, so a + partial final batch computes only its own columns. `get_eph_Rq_to_kq_batched!` now asserts its + workspace sizes with `>=` and prefix-views them internally, as its sibling + `get_eph_kR_to_kq_batched!` already did. - [ ] JML should review `TiledDeviceOutput` (the Sᵢ tiling machinery: `tile_begin!` / `tile_download!` / `tile_offset` / `tile_length`, used by `BoltzmannCalculator`'s batched path). -- [ ] Unify the CPU/GPU e-ph loop signatures (`_loop_eph_over_k_and_kq` vs - `_loop_eph_over_k_and_kq_batched`, `src/calculator/run_eph_over_k_and_kq.jl`). Decided direction (JML, - PR #9): keep **two functions with different names**, but give them the **identical positional - argument list** and push the path-specific data into **keyword arguments** — CPU kwargs - `(epstates, ep_ekpRs, epmat, ep_ekpR_obj, dyn_threads, epmat_R, epobj_ekpR_R, ep_ekpR_Rs)`, GPU - kwargs `(epmat_dev, backend)`. Constraint (`_setup`/`_loop` Core.Box rule): the CPU `_loop` must - destructure any NamedTuple into locals at the top before `@threads`, never index it inside the - threaded closure. +- [x] Unify the per-point/batched e-ph loop signatures — done for both driver families. Each pair + takes the identical positional list of what BOTH need (outer-k: `(model, kpts, qpts, kqpts, + el_k_save, el_kq_save, ph_save, precompute_ph, backend)`; outer-q: `(model, kpts, qpts, el_k_save, + ph_save, eph_buffers, backend)`) and its own path-specific data as individual keyword arguments + (per-point: the host interpolators/channels, resp. the precomputed k+q states; batched: + `epmat_dev`, resp. `el_ham_dev`). `backend` is the last shared positional rather than a batched-only + kwarg, so no `_loop_*` hardcodes `CPUBackend()` any more (`run_eph_over_k_and_q` too, which has no + batched twin). The convention is stated once, above `_loop_eph_over_k_and_kq`. - [x] ~~Reconsider whether `backend` should be built inside `_loop_eph_over_k_and_kq_batched` rather than in `_setup_eph_over_k_and_kq`.~~ Subsumed: `backend` is now a user-facing driver keyword, so it @@ -43,11 +45,8 @@ which is exactly right now that batched-on-`CPUBackend` is a supported configuration. The real defect was the flag: `calc.on_gpu = backend isa GPUBackend` inferred the loop shape from the backend, which builds the wrong buffers under CPU+batched. It is now `calc.batched = mode isa - BatchedMode`, from the `mode::LoopMode` keyword the drivers pass to `setup_calculator!`. - -- [ ] Clean up the `LoopContext` construction at the batch/per-k scope (deferred out of the - `backend`/`batched` PR as an independent follow-up) - (`src/calculator/run_eph_over_k_and_kq.jl` ~L713/L728). The batch-scope `ctx_batch` and per-k - `ctx_k` are built from positional constructors that are disambiguated by whether the argument is an - `Integer` outer index (`ik`) or a `UnitRange` batch (`iks_batch`) — flagged as flaky in review. - Consider a clearer, explicitly-named construction API for the two scopes. + BatchedMode`, from the positional `mode::LoopMode` the drivers pass to `setup_calculator!`. + +- [x] Clean up the `LoopContext` construction at the batch/per-k scope — done. The `BatchedMode` + convenience constructor is keyword-only (`batch` / `outer_index` / required `n_batch_max`), so + argument 3 is never positional there, and the per-k context is `with_outer_index(ctx_batch, ik)`. diff --git a/src/EPData.jl b/src/EPData.jl index 4b62a25..4a62add 100644 --- a/src/EPData.jl +++ b/src/EPData.jl @@ -86,10 +86,13 @@ batch of matrices takes the plural name `eps`; `ek`/`ekq`/`uk`/`ukq`/`wtk` are n they are (the batch is just their trailing axis). All fields are trimmed to the batch's actual width `nk` (the outer-k convention): the final partial -batch has `nk < n_batch_max`, and the loop hands width-`nk` views into its internally padded staging -buffers (a trailing-prefix device view, so it stays contiguous). A consumer therefore reads its own -size from any field (e.g. `size(eps, 4)`) and never sees a padded tail. (The loop still zeros the -internal `wtk` padding as defense-in-depth, but that padding is not exposed here.) +batch has `nk < n_batch_max`, and the loop hands width-`nk` views into its max-width staging buffers +(a trailing-prefix device view, so it stays contiguous). A consumer therefore reads its own size from +any field (e.g. `size(eps, 4)`) and never sees a padded tail. There is no padded tail internally +either — the loop computes only the batch's own columns — so the payload's width is the *only* +guarantee: a consumer that sized itself from `ctx.n_batch_max` instead of `size(eps, 4)` would read +stale data from the previous batch (uninitialised if the run has only one, partial batch), which is +worse than uninitialised because it is finite and plausible. Fields (`m` = k+q band, `n` = k band, `k` = batch column): - `eps` :: `(nw, nw, nmodes, nk)` — eigenbasis e-ph matrices. Out-of-window bands are already zeroed diff --git a/src/calculator/AbstractCalculator.jl b/src/calculator/AbstractCalculator.jl index 2754ff0..2772a72 100644 --- a/src/calculator/AbstractCalculator.jl +++ b/src/calculator/AbstractCalculator.jl @@ -93,8 +93,21 @@ Fields: - `mode` :: `SingleMode()` (per-(k, q) host loop) or `BatchedMode()` (device-batched loop). - `outer_index` :: current outer index (`ik` for outer-k loops, `iq` for the outer-q loop); `0` at batch scope. -- `batch` :: outer-iteration range of the current batch (`1:0` on the CPU paths). -- `n_batch_max` :: loop batch cap, for device-buffer sizing. +- `batch` :: outer-iteration range of the current batch (`1:0` when the loop does not batch its + outer axis — the per-point loops, and the batched outer-q loop, which batches the INNER k axis). +- `n_batch_max` :: loop batch cap, for device-buffer sizing. Which axis it caps follows the loop: the + OUTER-k width in the batched outer-k loop, the INNER-k width in the batched outer-q loop. A + consumer that reads it as an outer-axis width (as `TiledDeviceOutput` does, pairing it with + `ctx.batch`) is therefore outer-k-only. + +The two batched forms and the per-k derivation are constructed by name, never positionally: + +```julia +LoopContext(backend, SingleMode(), ik) # per-point +LoopContext(backend, BatchedMode(); batch = iks, n_batch_max = nk_max) # outer-k, batch scope +LoopContext(backend, BatchedMode(); outer_index = iq, n_batch_max = nk_max) # outer-q, one q +with_outer_index(ctx_batch, ik) # one k of a batch +``` """ struct LoopContext{BT <: AbstractBackend, MT <: LoopMode} backend :: BT @@ -109,10 +122,24 @@ end LoopContext(backend::AbstractBackend, ::SingleMode, outer_index::Integer) = LoopContext(backend, SingleMode(), outer_index, 1:0, 0) -# BatchedMode context at batch scope (device loops): there is no single outer index spanning the whole -# batch, so `outer_index = 0` is the "no single outer index — use `batch`" sentinel. -LoopContext(backend::AbstractBackend, ::BatchedMode, batch::UnitRange, n_batch_max::Integer) = - LoopContext(backend, BatchedMode(), 0, batch, n_batch_max) +# BatchedMode context, keyword-only so the two scopes a batched loop needs are told apart by NAME +# rather than by whether argument 3 is an `Integer` or a `UnitRange`: +# * batch scope (outer-k loop): pass `batch`; `outer_index = 0` is the "no single outer index — use +# `batch`" sentinel. +# * one outer iteration (outer-q loop): pass `outer_index`; `batch = 1:0` means "the outer axis is +# not batched here" (that loop batches its inner k axis). +LoopContext(backend::AbstractBackend, ::BatchedMode; outer_index::Integer = 0, + batch::UnitRange = 1:0, n_batch_max::Integer) = + LoopContext(backend, BatchedMode(), outer_index, batch, n_batch_max) + +""" + with_outer_index(ctx::LoopContext, outer_index::Integer) -> LoopContext + +The same context at one outer index, all other fields kept. The batched outer-k loop uses it to +derive its per-k context from the batch-scope one, which is exactly what the two contexts differ by. +""" +with_outer_index(ctx::LoopContext, outer_index::Integer) = + LoopContext(ctx.backend, ctx.mode, outer_index, ctx.batch, ctx.n_batch_max) # ============================================================================= @@ -246,7 +273,7 @@ if VERSION >= v"1.11.0-DEV.469" "postprocess_calculator!, calculator_begin!, calculator_end!, " * "OuterKLoop, OuterQLoop, OuterIteration, OuterIterationBatch, " * "AbstractElPhPayload, EPData, EPDataQBatched, EPDataKBatched, " * - "LoopContext, SingleMode, BatchedMode, LoopMode, " * + "LoopContext, with_outer_index, SingleMode, BatchedMode, LoopMode, " * "AbstractBackend, CPUBackend, GPUBackend, gpu_backend, alloc, free_bytes, synchronize, " * "batched_gemm!, eph_window_scatter!, bte_window_accumulate!, " * "eph_batched_bytes_per_point, allowed_eph_phonon_basis, " * diff --git a/src/calculator/run_eph_over_k_and_kq.jl b/src/calculator/run_eph_over_k_and_kq.jl index d0a464d..b74992f 100644 --- a/src/calculator/run_eph_over_k_and_kq.jl +++ b/src/calculator/run_eph_over_k_and_kq.jl @@ -125,8 +125,8 @@ function run_eph_over_k_and_kq( _loop_eph_over_k_and_kq_batched(model, setup.kpts, setup.qpts, setup.kqpts, setup.el_k_save, setup.el_kq_save, - setup.ph_save, setup.precompute_ph, - setup.epmat_dev, setup.backend; + setup.ph_save, setup.precompute_ph, setup.backend; + setup.epmat_dev, calculators, energy_conservation, screening_params, progress_print_step, nq_batch_max, nk_outer_batch_max, symmetry, verbosity, @@ -135,10 +135,10 @@ function run_eph_over_k_and_kq( _loop_eph_over_k_and_kq(model, setup.kpts, setup.qpts, setup.kqpts, setup.el_k_save, setup.el_kq_save, - setup.ph_save, setup.precompute_ph, + setup.ph_save, setup.precompute_ph, setup.backend; setup.epstates, setup.ep_ekpRs, setup.epmat, setup.ep_ekpR_obj, setup.dyn_threads, - setup.epmat_R, setup.epobj_ekpR_R, setup.ep_ekpR_Rs; + setup.epmat_R, setup.epobj_ekpR_R, setup.ep_ekpR_Rs, calculators, skip_eph, energy_conservation, screening_params, progress_print_step, nchunks_threads, @@ -303,14 +303,30 @@ function _setup_eph_over_k_and_kq( end +# CANONICAL statement of the `_loop_*` argument convention; the other three `_loop_*` functions point +# here instead of restating it. +# +# A driver family's per-point and batched loop shapes take the IDENTICAL positional list — exactly what +# BOTH shapes need, in the same order — and each takes only its own path-specific data as keyword +# arguments. So the two call sites read as two shapes of one call rather than two unrelated functions. +# For this family that list is +# (model, kpts, qpts, kqpts, el_k_save, el_kq_save, ph_save, precompute_ph, backend) +# and this shape's kwargs are the host interpolators / per-thread channels (the batched shape's is +# `epmat_dev`). `backend` is shared rather than batched-only: both shapes use it for the same thing, +# the `LoopContext`. +# +# Path-specific data stays INDIVIDUAL kwargs rather than one bundled NamedTuple, so every name is a +# typed argument of the lowered body method — which is what keeps the `@threads` closure below free of +# `Core.Box` without a destructure-before-`@threads` rule. function _loop_eph_over_k_and_kq( model :: Model{FT}, kpts, qpts, kqpts, el_k_save, el_kq_save, ph_save, precompute_ph, + backend; epstates, ep_ekpRs, epmat, ep_ekpR_obj, dyn_threads, - epmat_R, epobj_ekpR_R, ep_ekpR_Rs; + epmat_R, epobj_ekpR_R, ep_ekpR_Rs, calculators = [], skip_eph = false, energy_conservation = (:None, 0.0), @@ -323,7 +339,6 @@ function _loop_eph_over_k_and_kq( (; nw, nmodes) = model nk = kpts.n - backend = CPUBackend() for ik in 1:nk if mod(ik, progress_print_step) == 0 && mpi_isroot() @@ -525,6 +540,9 @@ end # batched over k-batches and q-batches with device staging — differs from the per-(k,q) loop, # not because it holds any device-specific code. # +# Positional list identical to `_loop_eph_over_k_and_kq`'s (the argument convention is stated in full +# in the comment above that function); this shape's only path-specific kwarg is `epmat_dev`. +# # Backend: `GPUBackend` is the production configuration. Because nothing here is device-specific, # the loop also runs on `CPUBackend` (`batched = true`), which is a VALIDATION configuration only: # the k-batch loop is serial, `batched_gemm!` degrades to a `mul!` loop, and `plan_batch` returns @@ -586,7 +604,8 @@ function _loop_eph_over_k_and_kq_batched( kpts, qpts, kqpts, el_k_save, el_kq_save, ph_save, precompute_ph, - epmat_dev, backend; + backend; + epmat_dev, calculators = [], energy_conservation = (:None, 0.0), screening_params = nothing, @@ -686,10 +705,11 @@ function _loop_eph_over_k_and_kq_batched( end # ----- persistent workspace (allocated once, reused across all (k, q)) ----- - # All device staging is sized to the full batch and used as plain CuArrays (not - # batch-sliced views), so the batched drivers' reshape/cuBLAS calls stay on dense arrays. - # Every buffer below comes from `alloc(backend, ...)`, so the backend is the single authority on - # where "device" is. + # All device staging is sized to the MAX batch width and consumed through contiguous + # trailing-prefix views of width `nk_batch` / `nq_batch`, so a partial final batch computes only + # its own columns and every argument the batched drivers see is still a dense array (such a view + # of a device array is a device array, not a `SubArray`). Every buffer below comes from + # `alloc(backend, ...)`, so the backend is the single authority on where "device" is. # RR->kR over a batch of `nk_batch_max` outer-k at once: one batched kernel per batch instead of one # launch-bound single-k call per k. `ep_ekpR_all` holds g(k, R_ep) for the whole batch; the inner @@ -738,7 +758,7 @@ function _loop_eph_over_k_and_kq_batched( # Grid coordinates as (3 × n) real device matrices, uploaded once. Both phase builds below read # them directly, so nothing on the phase path is staged on the host or copied H2D inside the loop. # Negated once here rather than conjugating the phase tile every batch: the convention needs - # conj(exp(2πi R_p·x_k)) = exp(2πi R_p·(−x_k)), and the two are bitwise identical (FP negation is + # conj(exp(2πi R_p·x_k)) = exp(2πi R_p·(-x_k)), and the two are bitwise identical (FP negation is # exact, and `cispi` is exactly symmetric). This is the only consumer of the k coordinates. # Out-of-place on purpose: on `CPUBackend` `_kpoints_to_device_matrix` returns a view onto # `kpts.vectors`, so negating in place would corrupt the k-points. @@ -746,7 +766,7 @@ function _loop_eph_over_k_and_kq_batched( xkq_dev = _kpoints_to_device_matrix(backend, kqpts) # The two Fourier phase matrices of the k+q convention (see `get_eph_RR_to_kR_batched!`): - # P_mk[ip, k] = exp(2πi R_p · (−x_k)) — folded into g(k, R_ep) once per outer-k batch + # P_mk[ip, k] = exp(2πi R_p · (-x_k)) — folded into g(k, R_ep) once per outer-k batch # P_kq[ip, j] = exp(2πi R_p · x_{k+q_j}) — the kR->kq phase, INDEPENDENT of the outer k # so one built P_kq tile serves every k of the batch. That reuse factor is `nk_batch`, i.e. # `nk_outer_batch_max`: lowering that cap shrinks this saving proportionally. @@ -756,11 +776,6 @@ function _loop_eph_over_k_and_kq_batched( irvecp_mat = _irvec_to_device_matrix(model.epmat.irvec_next, epmat_dev, FT) P_mk = alloc(backend, Complex{FT}, nr_ep, nk_batch_max) P_kq = alloc(backend, Complex{FT}, nr_ep, nq_batch_max) - # Defensive: only columns 1:nk_batch are rewritten per batch, so a partial final batch leaves the - # tail columns holding whatever the previous batch wrote. Nothing reads them — the k loop runs - # `1:nk_batch` — and 1 is the identity of the convention multiply, so the padded (never-read) - # slice of ep_ekpR_all stays meaningful whether or not it has been written yet. - fill!(P_mk, 1) # `iq` index staging for one (k, q-tile). iqs_batch = Vector{Int}(undef, nq_batch_max) @@ -793,9 +808,11 @@ function _loop_eph_over_k_and_kq_batched( kend = min(kstart + nk_batch_max - 1, nk) iks_batch = kstart:kend nk_batch = length(iks_batch) + rng_k = 1:nk_batch # this batch's columns within the nk_batch_max-sized staging buffers - # Stack U(k) and the k list for this outer-k batch (pad the partial tail with valid - # duplicated data so the batched RR->kR runs on dense `nk_batch_max`-sized arrays). + # Stack U(k) and the k list for this outer-k batch. A partial final batch fills only its own + # `nk_batch` columns: everything below runs at width `nk_batch` via prefix views into the + # max-width buffers, the same convention the q axis of this loop already uses. for (ik_ind, ik) in enumerate(iks_batch) # k-side window projection: the `nbandk_max` contiguous eigenvector columns around this # k's in-window range (all nw columns when full-band). The window selection itself @@ -804,11 +821,9 @@ function _loop_eph_over_k_and_kq_batched( @views uks_host[:, :, ik_ind] .= el_k_save[ik].u_full[:, nb0+1:nb0+nbandk_max] ks_batch[ik_ind] = kpts.vectors[ik] end - for ik_ind in (nk_batch+1):nk_batch_max - @views uks_host[:, :, ik_ind] .= uks_host[:, :, nk_batch] - ks_batch[ik_ind] = ks_batch[nk_batch] - end - copyto!(uks_dev, uks_host) + # H2D of just this batch's columns (5-arg contiguous copy — copying a host<->device SubArray + # view instead would fall back to scalar indexing; same reason as the `iqs_batch` copy below). + copyto!(uks_dev, 1, uks_host, 1, nw * nbandk_max * nk_batch) if mpi_isroot() && div(kend, progress_print_step) > div(kstart - 1, progress_print_step) @info "$(now()) ik = $kstart:$kend / $nk" @@ -816,15 +831,17 @@ function _loop_eph_over_k_and_kq_batched( end # One batched RR->kR over the whole batch: g(k, R_ep) for all k in the batch, stored in the - # k+q convention (multiplied by P_mk, the phase at −x_k) so the kR->kq phase below is - # k-independent. - @views fourier_phase!(P_mk[:, 1:nk_batch], irvecp_mat, mxk_dev[:, iks_batch]) - get_eph_RR_to_kR_batched!(ep_ekpR_all, itp_epmat, ks_batch, uks_dev; - additional_phase = P_mk) + # k+q convention (multiplied by P_mk, the phase at -x_k) so the kR->kq phase below is + # k-independent. The driver derives the batch width from `size(uks, 3)` and asserts that the + # other three arguments agree, so handing it four width-`nk_batch` views is within its contract. + @views fourier_phase!(P_mk[:, rng_k], irvecp_mat, mxk_dev[:, iks_batch]) + get_eph_RR_to_kR_batched!(view(ep_ekpR_all, :, :, rng_k), itp_epmat, + view(ks_batch, rng_k), view(uks_dev, :, :, rng_k); + additional_phase = view(P_mk, :, rng_k)) # Outer-batch-resident calculators (re)point/zero their per-batch device buffer here, before # this batch's scatters; no-op (default hooks) for calculators that hold their whole output. - ctx_batch = LoopContext(backend, BatchedMode(), iks_batch, nk_batch_max) + ctx_batch = LoopContext(backend, BatchedMode(); batch = iks_batch, n_batch_max = nk_batch_max) foreach(c -> calculator_begin!(c, OuterIterationBatch(), ctx_batch), calculators) qstart = 1 @@ -870,7 +887,7 @@ function _loop_eph_over_k_and_kq_batched( # Hand the tile's e-ph matrix (still on the device) to each calculator, which forms # g2 / scatters it on the device; no D2H of the e-ph matrix here. - ctx_k = LoopContext(backend, BatchedMode(), ik, iks_batch, nk_batch_max) + ctx_k = with_outer_index(ctx_batch, ik) payload = EPDataQBatched( view(epkq_dev, :, :, :, rng_q), view(g2_dev, :, :, :, rng_q), view(ωq_dev, :, rng_q), ik, ikqs_used, ibandk_offsets[ik]) diff --git a/src/calculator/run_eph_over_k_and_q.jl b/src/calculator/run_eph_over_k_and_q.jl index fdb7fd4..2450c14 100644 --- a/src/calculator/run_eph_over_k_and_q.jl +++ b/src/calculator/run_eph_over_k_and_q.jl @@ -83,7 +83,7 @@ function run_eph_over_k_and_q( _loop_eph_over_k_and_q(model, setup.kpts, setup.qpts, setup.kqpts, setup.el_k_save, setup.el_kq_save, setup.ph_save, - setup.precompute_el_kq, + setup.precompute_el_kq, setup.backend, setup.epstates, setup.ep_ekpRs, setup.epmat, setup.ep_ekpR_obj, setup.ham_threads, setup.vel_threads, setup.pos_threads; calculators, skip_eph, window_kq, @@ -222,16 +222,21 @@ function _setup_eph_over_k_and_q( precompute_el_kq, nband_max, epstates, ep_ekpRs, epmat, ep_ekpR_obj, ham_threads, vel_threads, pos_threads, + backend, iband_min, iband_max, ) end +# `backend` is a positional argument here too, as in the two loop pairs that have a batched twin +# (run_eph_over_k_and_kq.jl / run_eph_over_q_and_k.jl), so no `_loop_*` hardcodes `CPUBackend()`. It +# sits last among the shared arguments, before this shape's own data. This driver has no batched twin, +# so the rest of its argument list is not unified with anything. function _loop_eph_over_k_and_q( model :: Model{FT}, kpts, qpts, kqpts, el_k_save, el_kq_save, ph_save, - precompute_el_kq, + precompute_el_kq, backend, epstates, ep_ekpRs, epmat, ep_ekpR_obj, ham_threads, vel_threads, pos_threads; calculators = [], @@ -246,7 +251,6 @@ function _loop_eph_over_k_and_q( nk = kpts.n nq = qpts.n - backend = CPUBackend() for ik in 1:nk if mod(ik, progress_print_step) == 0 && mpi_isroot() diff --git a/src/calculator/run_eph_over_q_and_k.jl b/src/calculator/run_eph_over_q_and_k.jl index b8f83d0..3fdb24d 100644 --- a/src/calculator/run_eph_over_q_and_k.jl +++ b/src/calculator/run_eph_over_q_and_k.jl @@ -116,17 +116,19 @@ function run_eph_over_q_and_k( "the batched path does not support precompute_el_kq (k+q states are eigensolved in the loop).")) _loop_eph_over_q_and_k_batched(model, setup.kpts, setup.qpts, - setup.el_k_save, setup.ph_save, setup.eph_buffers, - setup.el_ham_dev, setup.backend; + setup.el_k_save, setup.ph_save, + setup.eph_buffers, setup.backend; + setup.el_ham_dev, calculators, skip_eph, window_kq, energy_conservation, screening_params, progress_print_step, eph_phonon_basis, verbosity, nk_batch_max, ) else _loop_eph_over_q_and_k(model, - setup.kpts, setup.qpts, setup.kqpts, - setup.el_k_save, setup.el_kq_save, setup.ph_save, - setup.precompute_el_kq, setup.eph_buffers; + setup.kpts, setup.qpts, + setup.el_k_save, setup.ph_save, + setup.eph_buffers, setup.backend; + setup.kqpts, setup.el_kq_save, setup.precompute_el_kq, calculators, skip_eph, window_kq, energy_conservation, screening_params, progress_print_step, nchunks_threads, @@ -277,12 +279,17 @@ function _setup_eph_over_q_and_k( end +# Same argument convention as `_loop_eph_over_k_and_kq` (stated in full in the comment above it). This +# family's shared positional list is +# (model, kpts, qpts, el_k_save, ph_save, eph_buffers, backend) +# and this shape's kwargs are the precomputed k+q data (the batched shape eigensolves k+q in the loop +# instead, and its kwarg is `el_ham_dev`). function _loop_eph_over_q_and_k( model :: Model{FT}, - kpts, qpts, kqpts, - el_k_save, el_kq_save, ph_save, - precompute_el_kq, - eph_buffers :: EphOuterQLoopBuffers{FT}; + kpts, qpts, + el_k_save, ph_save, + eph_buffers :: EphOuterQLoopBuffers{FT}, backend; + kqpts, el_kq_save, precompute_el_kq, calculators = [], skip_eph = false, window_kq = (-Inf, Inf), @@ -297,7 +304,6 @@ function _loop_eph_over_q_and_k( (; epstates, ep_eRpq_obj, ep_eRpqs, epmat, ham_threads, vel_threads) = eph_buffers nk = kpts.n nq = qpts.n - backend = CPUBackend() for iq in 1:nq if verbosity > 0 && mod(iq, progress_print_step) == 0 && mpi_isroot() @@ -420,6 +426,9 @@ end # `mul!`-loop `batched_gemm!`, and `plan_batch` returns the `nk_batch_max` cap verbatim). The `_dev` # suffix means "on `backend`", which is the host there. # +# Positional list identical to `_loop_eph_over_q_and_k`'s (see the comment above it); this shape's only +# path-specific kwarg is `el_ham_dev`. +# # Scope (asserted below; the per-point path handles the rest): no screening, # energy_conservation = (:None, 0.0), skip_eph = false, and every calculator supports the # `EPDataKBatched` payload. Windows are supported via eigenvector-column masking (out-of-window @@ -429,8 +438,8 @@ function _loop_eph_over_q_and_k_batched( model :: Model{FT}, kpts, qpts, el_k_save, ph_save, - eph_buffers :: EphOuterQLoopBuffers{FT}, - el_ham_dev, backend; + eph_buffers :: EphOuterQLoopBuffers{FT}, backend; + el_ham_dev, calculators = [], skip_eph = false, window_kq = (-Inf, Inf), @@ -497,9 +506,15 @@ function _loop_eph_over_q_and_k_batched( itp_ep_eRpq = BatchedWannierInterpolator(ep_eRpq_dev; batch_size = nk_batch_max) # ----- persistent per-batch device workspace (sized to nk_batch_max) ----- + # Consumed through contiguous trailing-prefix views of width `nk_batch`, so a partial final batch + # computes only its own columns and every argument the batched drivers see is still a dense array + # (such a view of a device array is a device array, not a `SubArray`). ep_ws = RqToKQWorkspace(ep_eRpq_dev.op_r, ndata_eRpq, nw, nw, nmodes, nk_batch_max) ep_batch = alloc(backend, Complex{FT}, nw, nw, nmodes, nk_batch_max) Hkq_flat = alloc(backend, Complex{FT}, nw * nw, nk_batch_max) + # The k+q eigensolve wants H as an (nw, nw, k) stack; reshape the dense parent once here so the + # per-batch argument is a plain prefix view of a dense array. + Hkq3 = reshape(Hkq_flat, nw, nw, nk_batch_max) Uk_batch = alloc(backend, Complex{FT}, nw, nw, nk_batch_max) Ukq_batch = alloc(backend, Complex{FT}, nw, nw, nk_batch_max) ek_batch = alloc(backend, FT, nw, nk_batch_max) @@ -537,21 +552,20 @@ function _loop_eph_over_q_and_k_batched( # Per-q calculator begin: allocate (first q) + zero the device accumulator (OuterIteration # bracket, same as the CPU loop; ctx carries backend + n_batch_max for the device buffer). - ctx_q = LoopContext(backend, BatchedMode(), iq, 1:0, nk_batch_max) + ctx_q = LoopContext(backend, BatchedMode(); outer_index = iq, n_batch_max = nk_batch_max) foreach(c -> calculator_begin!(c, OuterIteration(), ctx_q), calculators) for kstart in 1:nk_batch_max:nk kend = min(kstart + nk_batch_max - 1, nk) iks_batch = kstart:kend nk_batch = length(iks_batch) + rng_k = 1:nk_batch # this batch's columns within the nk_batch_max-sized buffers # k / k+q lists + k-side data staged on the host, then uploaded (streaming). `Uk` is # zero-padded outside each k's in-window range `rng`, so the full nw×nw Rq→kq rotation - # reproduces the CPU's windowed rotation with zeros outside. The tail (partial final - # batch) is padded with the last valid k so the batched Fourier / eigensolve see finite - # data; its weight is zeroed so — unlike the outer-k loop, where padded duplicates scatter - # to a unique in-window index harmlessly — a padded k column cannot double-count into the - # q-summed χ (the calculator multiplies by `wtk`). + # reproduces the CPU's windowed rotation with zeros outside. A partial final batch fills + # only its own `nk_batch` columns: everything below runs at width `nk_batch` via prefix + # views into the max-width buffers, so there is no padded tail to neutralise. for (ik_ind, ik) in enumerate(iks_batch) el = el_k_save[ik] ks_batch[ik_ind] = kpts.vectors[ik] @@ -561,44 +575,43 @@ function _loop_eph_over_q_and_k_batched( ek_host[:, ik_ind] .= el.e_full wtk_host[ik_ind] = kpts.weights[ik] end - for ik_ind in (nk_batch + 1):nk_batch_max - ks_batch[ik_ind] = ks_batch[nk_batch] - kqs_batch[ik_ind] = kqs_batch[nk_batch] - @views Uk_host[:, :, ik_ind] .= Uk_host[:, :, nk_batch] - @views ek_host[:, ik_ind] .= ek_host[:, nk_batch] - wtk_host[ik_ind] = 0 - end - copyto!(Uk_batch, Uk_host) - copyto!(ek_batch, ek_host) - copyto!(wtk_batch, wtk_host) + # H2D of just this batch's columns (5-arg contiguous copies — copying a host<->device + # SubArray view instead would fall back to scalar indexing). + copyto!(Uk_batch, 1, Uk_host, 1, nw * nw * nk_batch) + copyto!(ek_batch, 1, ek_host, 1, nw * nk_batch) + copyto!(wtk_batch, 1, wtk_host, 1, nk_batch) + + # The width-`nk_batch` prefix views of the staging buffers, named once so the drivers below + # and the payload provably see the same arrays (a trailing-prefix view of a device array is + # itself a device array, so the extension kernels take these directly). + ep_v = view(ep_batch, :, :, :, rng_k) + Uk_v = view(Uk_batch, :, :, rng_k) + Ukq_v = view(Ukq_batch, :, :, rng_k) + ks_v = view(ks_batch, rng_k) # k+q eigensolve on the device (batched). No gauge fixing needed: χ is gauge-invariant. # TODO: `eigen_batched` allocates (E, U) each batch; an in-place variant into ek/Ukq # scratch would remove the per-batch allocation. - get_fourier_batched!(Hkq_flat, itp_el_ham, kqs_batch) - Ekq, Ukq = eigen_batched(reshape(Hkq_flat, nw, nw, nk_batch_max)) # (nw,·), (nw,nw,·) + get_fourier_batched!(view(Hkq_flat, :, rng_k), itp_el_ham, view(kqs_batch, rng_k)) + Ekq, Ukq = eigen_batched(view(Hkq3, :, :, rng_k)) # (nw, nk_batch), (nw, nw, nk_batch) # k+q window mask: zero eigenvector COLUMNS m outside [wmin, wmax] (Ekq[m,k]). This # zeroes ep_kq[m,·] and every k+q-side matrix element for out-of-window m, so those # (m,n) pairs contribute exactly 0 — reproducing the CPU's `for m in el_kq.rng` loop. - mask_kq = (Ekq .>= wmin) .& (Ekq .<= wmax) # (nw, ·) Bool - Ukq_batch .= Ukq .* reshape(mask_kq, 1, nw, nk_batch_max) + mask_kq = (Ekq .>= wmin) .& (Ekq .<= wmax) # (nw, nk_batch) Bool + Ukq_v .= Ukq .* reshape(mask_kq, 1, nw, nk_batch) # Batched Rq→kq e-ph interpolation: ep_batch[m,n,ν,k] = Ukq(k)' * g(k) * Uk(k). - get_eph_Rq_to_kq_batched!(ep_batch, itp_ep_eRpq, ks_batch, Uk_batch, Ukq_batch; ws = ep_ws) - - use_polar_eph && add_eph_dipole_batched!(ep_batch, coeffs_dev, Ukq_batch, Uk_batch, mmats_batch) - - # Hand the calculator width-`nk_batch` views (the outer-k convention): the internal - # staging stays padded to `nk_batch_max` for the dense batched eigensolve / Fourier, but - # the payload is trimmed so an unweighted reduction cannot count padded columns. (The - # internal `wtk` zero-padding above is kept as defense-in-depth.) A trailing-prefix view - # of a device array is contiguous, so the extension kernels take these directly. - rng_k = 1:nk_batch - payload = EPDataKBatched( - view(ep_batch, :, :, :, rng_k), view(ek_batch, :, rng_k), view(Ekq, :, rng_k), - view(Uk_batch, :, :, rng_k), view(Ukq_batch, :, :, rng_k), view(wtk_batch, rng_k), - view(ks_batch, rng_k), iq) + get_eph_Rq_to_kq_batched!(ep_v, itp_ep_eRpq, ks_v, Uk_v, Ukq_v; ws = ep_ws) + + use_polar_eph && add_eph_dipole_batched!(ep_v, coeffs_dev, Ukq_v, Uk_v, + view(mmats_batch, :, :, rng_k)) + + # A consumer reads its own size from any payload field (`size(eps, 4)`); there is no padded + # tail anywhere. `Ekq`/`Ukq` are already exactly this batch's width; they are viewed anyway + # so `ek`/`ekq` (which share one type parameter) stay the same type. + payload = EPDataKBatched(ep_v, view(ek_batch, :, rng_k), view(Ekq, :, rng_k), + Uk_v, Ukq_v, view(wtk_batch, rng_k), ks_v, iq) foreach(c -> run_calculator!(c, payload, ctx_q), calculators) end # k batch diff --git a/src/wannier_to_bloch_batched.jl b/src/wannier_to_bloch_batched.jl index f61dacc..d70944e 100644 --- a/src/wannier_to_bloch_batched.jl +++ b/src/wannier_to_bloch_batched.jl @@ -308,8 +308,9 @@ The parent is the electron-Wannier / phonon-Bloch object (`op_r` `(nw^2*nmodes, `ep_kq_all[m,n,ν,k] = Σ_{iw,jw} conj(ukqs[iw,m,k]) · g[iw,jw,ν,k] · uks[jw,n,k]` is applied as two `batched_gemm!`s (`ukq(k)'` on the left over batch `k`, `uk(k)` on the right over batch `(ν,k)`). -Pass an [`RqToKQWorkspace`](@ref) (sized for this `nk`) as `ws` to reuse the `g`/`tmp`/`uk_rep` -scratch across calls instead of allocating it each call — the per-q hot loop does this. +Pass an [`RqToKQWorkspace`](@ref) as `ws` (sized for at least this `nk`) to reuse the `g`/`tmp`/`uk_rep` +scratch across calls instead of allocating it each call — the per-q hot loop does this, sizing `ws` +for the max batch width and passing `nk <=` that for a partial final batch. Full-band only: like [`get_eph_RR_to_kR_batched!`](@ref), all `nk` k-points must share the same `nbandk`/`nbandkq` (energy windows are handled by callers with masks). @@ -330,10 +331,16 @@ function get_eph_Rq_to_kq_batched!(ep_kq_all::AbstractArray{Complex{T},4}, tmp = similar(parent.op_r, Complex{T}, nbandkq, nw * nmodes, nk) uk_rep = similar(parent.op_r, Complex{T}, nw, nbandk, nmodes * nk) else - g, tmp, uk_rep = ws.g, ws.tmp, ws.uk_rep - @assert size(g) == (parent.ndata, nk) - @assert size(tmp) == (nbandkq, nw * nmodes, nk) - @assert size(uk_rep) == (nw, nbandk, nmodes * nk) + # `ws` is sized for the max batch width; use the first `nk` columns (a partial final batch + # passes nk < capacity), so the whole loop runs without padding the batch back up. Same + # convention as `get_eph_kR_to_kq_batched!` above. + @assert size(ws.g, 1) == parent.ndata && size(ws.g, 2) >= nk + @assert size(ws.tmp, 1) == nbandkq && size(ws.tmp, 2) == nw * nmodes && size(ws.tmp, 3) >= nk + @assert size(ws.uk_rep, 1) == nw && size(ws.uk_rep, 2) == nbandk && + size(ws.uk_rep, 3) >= nmodes * nk + g = view(ws.g, :, 1:nk) + tmp = view(ws.tmp, :, :, 1:nk) + uk_rep = view(ws.uk_rep, :, :, 1:nmodes*nk) end # Fourier over R_el at every k -> g(k) in (nw, nw, nmodes, nk); index legend g[iw, jw, ν, k] diff --git a/test/test_calculator_contract.jl b/test/test_calculator_contract.jl index e1c73b2..ec50207 100644 --- a/test/test_calculator_contract.jl +++ b/test/test_calculator_contract.jl @@ -3,7 +3,7 @@ using ElectronPhonon using ElectronPhonon: AbstractCalculator, OuterKLoop, OuterQLoop, EPData, EPDataQBatched, supports, LoopContext, SingleMode, BatchedMode, CPUBackend, GPUBackend, AbstractBackend, OuterIteration, OuterIterationBatch, - calculator_begin!, calculator_end!, to_device + calculator_begin!, calculator_end!, to_device, with_outer_index # Stage-2 calculator-contract checks (CPU-only): the `supports` trait, the fail-early payload checks # the drivers do at entry, the `calculators`-as-kwarg change, and the screening-disabled error. @@ -79,9 +79,11 @@ ElectronPhonon.calculator_begin!(c::_ModeDispatchCalc, ::OuterIterationBatch, :: (push!(c.fired, (:batch, :batched)); c) @testset "loop-mode bracket dispatch (DECISION-6)" begin - # `LoopContext` carries the backend first, the mode second. + # `LoopContext` carries the backend first, the mode second. `ctx_pt` deliberately uses the + # field-wise constructor: it probes a SingleMode context with a nonzero `n_batch_max`, a + # combination production never builds (the SingleMode convenience forces `1:0, 0`). ctx_pt = LoopContext(CPUBackend(), SingleMode(), 1, 1:0, 4) - ctx_bt = LoopContext(CPUBackend(), BatchedMode(), 0, 1:4, 4) + ctx_bt = LoopContext(CPUBackend(), BatchedMode(); batch = 1:4, n_batch_max = 4) @test ctx_pt isa LoopContext{CPUBackend, SingleMode} @test ctx_bt isa LoopContext{CPUBackend, BatchedMode} # The backend-first order keeps the partial annotation `LoopContext{<:GPUBackend}` valid (any mode). @@ -116,3 +118,29 @@ ElectronPhonon.calculator_begin!(c::_ModeDispatchCalc, ::OuterIterationBatch, :: v = [1.0, 2.0, 3.0] @test to_device(CPUBackend(), v) === v end + +@testset "LoopContext named construction" begin + # The three forms the drivers build, each by name: argument 3 is never positional in BatchedMode, + # so an `Integer` outer index and a `UnitRange` batch can no longer be confused. + ctx_pt = LoopContext(CPUBackend(), SingleMode(), 7) + @test (ctx_pt.outer_index, ctx_pt.batch, ctx_pt.n_batch_max) == (7, 1:0, 0) + + # Batch scope (batched outer-k): no single outer index spans the batch, hence the `0` sentinel. + ctx_bt = LoopContext(CPUBackend(), BatchedMode(); batch = 3:5, n_batch_max = 4) + @test (ctx_bt.outer_index, ctx_bt.batch, ctx_bt.n_batch_max) == (0, 3:5, 4) + + # One outer iteration of the batched outer-q loop: the outer axis is not batched there, so `batch` + # keeps its "no outer batch" default. + ctx_q = LoopContext(CPUBackend(), BatchedMode(); outer_index = 2, n_batch_max = 10) + @test (ctx_q.outer_index, ctx_q.batch, ctx_q.n_batch_max) == (2, 1:0, 10) + + # `with_outer_index` changes only the outer index. + ctx_k = with_outer_index(ctx_bt, 4) + @test ctx_k isa typeof(ctx_bt) + @test (ctx_k.backend, ctx_k.mode, ctx_k.batch, ctx_k.n_batch_max) == + (ctx_bt.backend, ctx_bt.mode, ctx_bt.batch, ctx_bt.n_batch_max) + @test ctx_k.outer_index == 4 + + # `n_batch_max` is required: a batched context with no batch cap is a construction error. + @test_throws UndefKeywordError LoopContext(CPUBackend(), BatchedMode(); batch = 1:4) +end diff --git a/test/test_gpu.jl b/test/test_gpu.jl index c5eedc0..ddc846f 100644 --- a/test/test_gpu.jl +++ b/test/test_gpu.jl @@ -221,7 +221,9 @@ ElectronPhonon.calculator_begin!(c::_OptInQCalc, ::ElectronPhonon.OuterIteration ElectronPhonon.calculator_end!(c::_OptInQCalc, ::ElectronPhonon.OuterIteration, ctx) = (c.flush += 1; nothing) @testset "outer-q batched calculator hook plumbing" begin - ctx = ElectronPhonon.LoopContext(ElectronPhonon.CPUBackend(), ElectronPhonon.SingleMode(), 1, 1:0, 4) + # The context the batched outer-q loop builds per q: one outer index, no outer batch. + ctx = ElectronPhonon.LoopContext(ElectronPhonon.CPUBackend(), ElectronPhonon.BatchedMode(); + outer_index = 1, n_batch_max = 4) # Default opts out; a calculator with no run_calculator! method for the payload is a MethodError. @test ElectronPhonon.supports(_PlainQCalc(), ElectronPhonon.EPDataKBatched) == false pl = ElectronPhonon.EPDataKBatched(nothing, nothing, nothing, nothing, nothing, nothing, nothing, 1) @@ -382,6 +384,65 @@ end end end +# Partial final K-batch (the q-axis analogue above is per-k). Both batched loops run a narrower final +# batch by handing the k-side drivers width-`m` contiguous prefix views of their max-width staging: +# `get_eph_RR_to_kR_batched!` takes the batch width from `size(uks, 3)` and asserts the rest, and +# `get_eph_Rq_to_kq_batched!` asserts its workspace with `>=` and prefix-views it internally. A +# width-`m` view run must reproduce the first `m` slices of the full-width run and leave the tail +# untouched. Backend-agnostic (`to_dev`/`arr_dev` as in `check_eph_batched`), so it runs without CUDA. +function check_eph_partial_view_k(to_dev, arr_dev; rtol) + nwe, nmodes, nr_el, nr_ep = 3, 4, 12, 9 + nband, nk, m = nwe, 9, 5 # view width m < full width nk + irvec_el = sort([Vec3(rand(-2:2, 3)...) for _ in 1:nr_el], by = x -> reverse(x)) + irvec_ep = sort([Vec3(rand(-2:2, 3)...) for _ in 1:nr_ep], by = x -> reverse(x)) + ks = [Vec3(rand(3)...) for _ in 1:nk] + uks = arr_dev(rand(ComplexF64, nwe, nband, nk)) + ukqs = arr_dev(rand(ComplexF64, nwe, nband, nk)) + + # (a) RR→kR, the outer-k loop's k-axis: output / k list / uks / additional_phase all prefix-viewed. + epmat_d = to_dev(WannierObject(irvec_el, rand(ComplexF64, nwe^2*nmodes*nr_ep, nr_el); + irvec_next = irvec_ep)) + itp_epmat = get_interpolator(epmat_d; fourier_mode="batched", batch_size=nk) + P = arr_dev(rand(ComplexF64, nr_ep, nk)) + ndata = nwe * nband * nmodes + full = arr_dev(zeros(ComplexF64, ndata, nr_ep, nk)) + get_eph_RR_to_kR_batched!(full, itp_epmat, ks, uks; additional_phase = P) + part = arr_dev(zeros(ComplexF64, ndata, nr_ep, nk)) + get_eph_RR_to_kR_batched!(view(part, :, :, 1:m), itp_epmat, view(ks, 1:m), + view(uks, :, :, 1:m); additional_phase = view(P, :, 1:m)) + @test isapprox(Array(view(part, :, :, 1:m)), Array(view(full, :, :, 1:m)); rtol) + @test all(Array(view(part, :, :, m+1:nk)) .== 0) # the narrow call wrote only its own columns + + # (b) Rq→kq, the outer-q loop's k-axis: same ONE max-width workspace serves the full and the + # narrow call, which is what the `>=` workspace assertions and internal prefix views are for. + eRpq_d = to_dev(WannierObject(irvec_el, rand(ComplexF64, nwe^2*nmodes, nr_el))) + itp_eRpq = get_interpolator(eRpq_d; fourier_mode="batched", batch_size=nk) + ws = ElectronPhonon.RqToKQWorkspace(eRpq_d.op_r, nwe^2*nmodes, nband, nband, nmodes, nk) + full_q = arr_dev(zeros(ComplexF64, nband, nband, nmodes, nk)) + get_eph_Rq_to_kq_batched!(full_q, itp_eRpq, ks, uks, ukqs; ws) + part_q = arr_dev(zeros(ComplexF64, nband, nband, nmodes, nk)) + get_eph_Rq_to_kq_batched!(view(part_q, :, :, :, 1:m), itp_eRpq, view(ks, 1:m), + view(uks, :, :, 1:m), view(ukqs, :, :, 1:m); ws) + @test isapprox(Array(view(part_q, :, :, :, 1:m)), Array(view(full_q, :, :, :, 1:m)); rtol) + @test all(Array(view(part_q, :, :, :, m+1:nk)) .== 0) +end + +@testset "partial k-batch (prefix views into max-width buffers, CPU)" begin + check_eph_partial_view_k(identity, identity; rtol=1e-10) +end + +@testset "GPU partial k-batch (prefix views into max-width buffers)" begin + if !GPU_AVAILABLE + @info "CUDA not available/functional — skipping GPU partial k-batch test" + else + # A trailing-prefix view of a CuArray is itself a CuArray, which is why the extension's + # `CuArray`-annotated eigensolve / cuBLAS methods still dispatch on these arguments. + @test view(CuArray(zeros(ComplexF64, 2, 3, 4)), :, :, 1:2) isa CuArray + check_eph_partial_view_k(obj -> to_device(ElectronPhonon.gpu_backend(), obj), CuArray; + rtol=1e-9) + end +end + @testset "eph_apply_rotations! rejects a non-dense g" begin # The two-GEMM rotation paths merge `g`'s band and mode axes with a `reshape`, so a strided `g`