KOKKOS: optimize pair_style pace/kk on GPUs and add CPU-backend kernels - #54
Closed
stanmoore1 wants to merge 12 commits into
Closed
KOKKOS: optimize pair_style pace/kk on GPUs and add CPU-backend kernels#54stanmoore1 wants to merge 12 commits into
stanmoore1 wants to merge 12 commits into
Conversation
…s (Stage 0) Introduce the GPU performance-tuning infrastructure for the Kokkos ACE (ML-PACE) pair style by replacing the single hard-coded team size (32 for every GPU TeamPolicy kernel) with named, per-kernel constants that are selected per backend (CUDA/HIP/SYCL vs host): team_size_compute_neigh, team_size_compute_radial, team_size_compute_ai, team_size_compute_derivative This mirrors the per-kernel tuning constants used by the Kokkos SNAP pair style and is the hook for a later empirical per-architecture tuning sweep. The constants intentionally reproduce the previous behaviour (team size 32 on all GPU kernels), so this change is performance- and bitwise-neutral. Verified: builds cleanly with the Kokkos Serial backend, and the copper example (examples/PACKAGES/pace/in.pace.product) produces identical energies, pressures and consistency checks (v_delenergy, v_delpress) to the pre-change baseline.
… (Stage 1) Replace the interleaved complex arrays A_sph and weights (Kokkos Views of SNAComplex) with separate real/imaginary double arrays (A_sph_re/A_sph_im and weights_re/weights_im). The atom index is the innermost (stride-1) dimension, so the per-neighbor atomic accumulation in ComputeAi and the adjoint accumulation in ComputeWeights now coalesce across the atom index on GPUs, instead of being strided by sizeof(complex) when touching the .re and .im halves of an interleaved complex element. This mirrors the ulisttot_re/ulisttot_im layout used by the heavily optimized Kokkos SNAP pair style and is the prerequisite for eliminating these atomics entirely in a follow-on change. Readers in ConjugateAi and ComputeDerivative reconstruct a complex value from the two arrays (using complex(re,-im) where the conjugate was taken). The algorithm is unchanged. This change is isolated to PairPACEKokkos; PairPACEExtrapolationKokkos is an independent class and is untouched. Verified: builds cleanly with the Kokkos Serial backend. The Kokkos compute kernels were exercised directly on the Serial backend (via a temporary test-only bypass of the host fallback, not included in this commit) and reproduce the serial reference evaluator bit-for-bit on the copper example (examples/PACKAGES/pace/in.pace.product): identical energies and pressures over 100 steps, with v_delenergy = 0 and v_delpress at the 1e-11 round-off level.
The m=0, m=1, and m>1 branches of ComputeDerivative each carried an identical inner loop over the radial functions n. Factor that loop into a single KOKKOS_INLINE_FUNCTION helper (compute_derivative_radial) shared by all three branches, with a wscale argument folding in the factor-of-2 used for the m>0 (-m) cases. The helper also hoists the zero-weight test ahead of the fr/dfr reads and the complex products, so channels with a zero weight no longer pay for the radial memory traffic and the Y_DR/grad_phi multiplies before being skipped. Behavior preserving: a 50-step fcc Cu recursive-ACE run matches the non-Kokkos PACE reference bit-for-bit (temperature, energy, pressure).
General GPU hygiene: pow() is a ~hundreds-of-cycles transcendental on the
device, so replace it where the exponent is a small integer.
- cutoff_func_poly: pow(x,3)/pow(x,4)/pow(x,5) -> explicit products
(runs per atom in the embedding inner cutoff and the ZBL transition).
- Fexp: pow(w*a,3) -> cube.
Mathematically identical (integer powers), so results are bit-identical.
Zero weights_re/weights_im/weights_rank1 with a first-touch pass in ComputeFS (one thread per atom, run just before ComputeWeights) instead of a full-array deep_copy every chunk.
…ComputeRho Eliminate the full (both-sign m) complex A array and the ConjugateAi kernel that expanded it from the half-basis A_sph. ComputeRho now reads A_sph directly, applying A(l,-p) = (-1)^p * conj(A(l,p)) for negative-m factors. This halves the A working set of the top kernel (better L2 reuse), removes ConjugateAi's full-A write traffic each chunk, and frees the A allocation. Cost is a d_idx_sph lookup + sign branch per A read in the product chain. The product-chain scratch (A_list/A_forward_prod) and dB_flatten are kept as develop's interleaved-complex arrays (the real/imag split of those per-thread scratch arrays gave no speedup, so it is not included here). Removes the A array (and its grow/memory_usage), and the TagPairPACEConjugateAi kernel, tag, declaration and launch. Validated on the Serial backend: energy/pressure bit-identical to reference, forces vs CPU PairPACE 1.3e-12.
In the spline radial path, calcSplines wrote the rnl functions into the d_values/d_derivatives scratch buffers, which were then copied into fr/dfr with an index "transpose". Because fr (LayoutLeft, l fastest of the last two dims) and d_values share the exact same flat layout (func = n*(lmax+1)+l), that copy was an identity copy - a full extra round-trip of the radial functions per neighbor (write d_values, read d_values, write fr). Store fr/dfr as 3D arrays (natom, maxneigh, (lmax+1)*nradmax) indexed by n*(lmax+1)+l so the rnl spline writes them directly, removing the buffer and the copy (~2x less radial-function store traffic in ComputeRadial). The consumers (accumulate_A_sph, compute_derivative_radial) and the direct evaluator just use the flattened index; reads stay coalesced (atom innermost). d_values/d_derivatives remain for the single-function hard-core spline. Validated on the Serial backend: spline path energy/pressure bit-identical and forces vs CPU PairPACE 1.3e-12; direct path unchanged (2.5e-9, spline interpolation error).
… ComputeWeights
dB_flatten was the largest scratch array (natom x idx_ms_combs_max x rankmax
complex), written in ComputeRho (top kernel, 29%) and read in ComputeWeights
(15%). Eliminate it:
- ComputeRho no longer needs the leave-one-out products, so it drops the
backward product pass and the A_list/A_forward_prod global scratch too,
collapsing to a single running register product B = prod(A) (read from
A_sph via the shared read_A conjugate-symmetry helper) followed by the
rho(p) accumulation.
- ComputeWeights recomputes each leave-one-out product dB[t] = prod_{s!=t} A_s
on the fly from A_sph (O(rank^2) cheap, cached reads; no array, no
register product-chain which had regressed).
Net: three global scratch arrays removed (A_list, A_forward_prod, dB_flatten)
and the ComputeRho backward pass deleted, trading the largest scratch array's
store+load traffic for recompute in ComputeWeights. Same memory-for-compute
trade as the A-array elimination.
Validated on the Serial backend: energy bit-identical, forces vs CPU PairPACE
1.3e-12 (pressure differs at 3e-9, product-order round-off).
Also add a const qualifier to SNAComplex::real_part_product(const complex&)
in kokkos_type.h. Making B a mutable running product here (it was a const
copy of A_forward_prod before) means B.real_part_product(d_ctildes(...)) is
now called on a non-const complex; with a real_type argument that is
implicitly convertible to complex, both real_part_product overloads become
viable and Clang (HIP/CUDA device pass) rejects the call as ambiguous.
Neither overload mutates the object, so the const makes the const-argument
overload the unique best match and is otherwise a no-op.
ComputeRadial and ComputeAi run over the identical flat (atom,neighbor) team layout, and ComputeAi already consumes the radial outputs (gr/fr/cr) that ComputeRadial produced. Call evaluate_splines() at the top of ComputeAi so each thread computes the radial functions for its own (ii,jj) slot and reads them back in the same kernel. This removes one full kernel launch per chunk and the global round-trip of the rank-1 radial arrays between the two kernels. Each (ii,jj) is handled by a single thread, so there is no cross-thread dependency between the write and the read. dfr/dgr/dcr are still stored for ComputeDerivative, so this is a pure launch/traffic saving with no algorithmic change. Validated on the fcc-Cu product deck (in.pace.product) by running the GPU kernels on the Serial backend: energy matches the CPU PairPACE reference to ~13 significant digits and per-atom forces to 1.3e-12.
The ComputeNeigh kernel in pair pace/kk and pace/extrapolation/kk caches the short neighbor list in level-0 (on-chip shared) team scratch memory, sized team_size*maxneigh*sizeof(int). On GPUs shared memory is a very limited resource, so runs with many neighbors and/or many atomic species could abort with "Requested too much scratch memory on level 0" (CUDA) or "could not find a valid team size" (HIP). See lammps#5063. Query the maximum available level-0 scratch from Kokkos via TeamPolicy::scratch_size_max(0) and transparently fall back to level-1 (global memory) scratch when the request does not fit, printing a warning the first time. The limit is queried rather than hard-coded (e.g. 48 KiB), so larger shared-memory limits such as the opt-in >48 KiB shared memory in newer Kokkos are picked up automatically. Add a "neigh" pair_style keyword (auto|shared|global) so the user can override the automatic choice, and document it. (cherry picked from commit 1fef170)
pair_style pace/kk never executed a single Kokkos kernel on a CPU. In a CPU-only build LMPDeviceType and LMPHostType are the same type, so host_flag was always true and both init_style() and compute() returned early: compute() synced to the host and called PairPACE::compute(), i.e. the external ML-PACE library evaluator, and init_style() refused to run with more than one thread. That is the delegating "fake port" pattern the KOKKOS package rules forbid, and it also meant pace/kk could not use more than one CPU core. Remove both early returns so the Kokkos kernels run on the Serial and OpenMP backends. init_style() now reaches the neighbor list request and the copy_pertype()/copy_splines()/copy_tilde() setup on host, which it previously skipped entirely. Turn host_flag into a static constexpr so the host/device branches are resolved at compile time and the unused kernel set is not instantiated, matching pair_snap_kokkos. Guard the pace/kk/host style registration with LMP_KOKKOS_GPU for the same reason, and drop the now misleading "on the GPU" from the recursive-evaluator error message. Verified against the non-Kokkos product evaluator on examples/PACKAGES/pace/in.pace.product: identical E_pair at every printed step of a 100 step NVE run, and per-atom forces agreeing to 1e-12 eV/A (1e-11 relative), for both 1 and 4 OpenMP threads. (cherry picked from commit 73b6ce5)
Squashes the pace/kk CPU optimization series (7 commits) onto the GPU
optimization work already on this branch, keeping each backend's own data
layout rather than forcing one on both.
The host kernels give one thread a whole atom and loop its neighbors
inside, so nothing else writes that atom's accumulators and the adds need
no atomics; the basis-function work is fused (rho -> F(rho) -> weights in
one pass), flattened onto precomputed gather offsets with the View strides
hoisted into raw pointers, batched across eight atoms, and specialized on
the density count at compile time. Dispatch is "if constexpr (host_flag)",
so only one kernel set is instantiated per backend.
Dual layout. The two backends want opposite things from the same arrays,
so both are declared and only the live one is allocated (see grow()):
- host: interleaved complex A / A_sph / weights, so re and im land in the
same cache line for a thread walking one atom, plus the full A array
expanded once per step by ConjugateAi and the stored dB_flatten.
- device: split re/im A_sph and weights with the atom index innermost,
which is what a warp scattering across atoms needs to coalesce, with
the -m entries applied on the fly in read_A() and the leave-one-out
products recomputed rather than stored.
fr/dfr stay single: both backends share the flat (n*(lmax+1)+l) array the
rnl spline writes directly, with the host gathers reindexed onto it.
Transposing it to n-innermost for the host was measured and was a net loss.
Radial evaluation stays fused into ComputeAi on the device (this branch
dropped the separate ComputeRadial kernel); the host keeps ComputeRadialCPU,
which walks a whole atom and prefetches the next neighbor's spline rows.
Weight zeroing is per-backend too: first-touch inside ComputeFS on the
device, a bulk deep_copy on the host, where per-atom scalar stores lose.
Validated on the fcc-Cu product deck against the CPU PairPACE reference,
both branches of the compile-time switch exercised on the Serial backend:
per-atom forces match to 9.75e-13 on the host path and 9.80e-13 with the
device path forced on.
Performance note: on 2048 atoms the host path runs 1.25x the reference
evaluator (14.9 s vs 18.6 s). The standalone CPU branch reaches 1.70x, so
roughly two thirds of that gain is not yet recovered here; the cause is not
kernel dispatch, the NDENSITY specializations, the weight zeroing or the
fr ordering, all of which were checked. The device path is unchanged in
kind from the GPU series and still needs a CUDA build to re-confirm.
Folded in (previously separate follow-up commits on the branch) so this
commit builds and runs correctly on a real GPU backend on its own:
- grow() tested A_sph_re.extent(0), but A_sph_re is only allocated on the
device path; on a host backend its extent stays 0 so the whole per-atom
allocation block was freed and re-allocated (and re-zeroed, re-faulted)
every step. Test A_rank1 instead, which both paths allocate.
- The explicit instantiation template class PairPACEKokkos<LMPDeviceType>
(and, in a GPU build, PairPACEKokkos<LMPHostType>) instantiates every
member, so the device compiler pass type-checks the CPU-only kernels.
"if constexpr (host_flag)" does not help the LMPHostType instantiation,
where host_flag is always true. Guard the six CPU kernel bodies that
reference host-only helpers or the LayoutRight assumption with an
"#ifndef LMP_KK_DEVICE_COMPILE" preprocessor check, which is false for
both instantiations in the device pass, so the bodies compile away
there. This clears two device-pass blockers: the reference to the plain
__host__ function pace_batched_derivative() from a __host__ __device__
context, and the LayoutRight static_assert firing when t_ace_4c is
LayoutLeft on the device.
- compute_ai_one() took the address of A_sph on the shared path the device
executes, but grow() never allocates A_sph there; guard it under
"if constexpr (!NEED_ATOMICS)" with a null base and zero strides.
- pace_batched_derivative() and its call sites hard-coded double; retyped
to KK_FLOAT so LMP_KOKKOS_SINGLE_SINGLE / SINGLE_DOUBLE builds work.
- team-size constants keyed off the build config (KOKKOS_ENABLE_CUDA and
friends) rather than host_flag, so in a GPU build the pace/kk/host CPU
backend launched with GPU team sizes and 32x the team scratch it needs.
Key them off host_flag.
- memory_usage() did not count the host-only A, A_sph, weights and
dB_flatten arrays; calcSplines() still declared wl2/wl3/w2l1/w3l2, dead
since the Horner rewrite.
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
1 participant
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.
Summary
Optimizes the KOKKOS version of
pair_style pace(ACE) on GPU backends, adds native CPU-backend kernels so thatpace/kkno longer falls back to the non-accelerated evaluator on host builds, and fixes an abort that occurred when the short neighbor list exceeded the available team shared memory.Measured on the fcc Cu
productdeck, relative to the first commit in this series:No new files, no new styles; 12 commits, each one a self-contained change that was benchmarked individually.
Related Issue(s)
Addresses lammps#5063 (
pace/kkaborting with "Requested too much scratch memory on level 0" on CUDA, or "could not find a valid team size" on HIP).Author(s)
Stan Moore, Sandia National Laboratories
Licensing
By submitting this pull request, I agree, that my contribution will be included in LAMMPS and redistributed under either the GNU General Public License version 2 (GPL v2) or the GNU Lesser General Public License version 2.1 (LGPL v2.1).
Artificial Intelligence (AI) Tools Usage
AI tools were used for this pull request and the disclosure below replaces the default statement in the template.
The great majority of the code changes were generated using Claude Code (Anthropic). This includes the GPU optimization sequence, the CPU-backend kernels and the dual host/device layout that integrates them, and the accompanying documentation updates. The human author directed the work, selected which changes to keep, reviewed the resulting code, ran all of the hardware benchmarking on H100 and MI300A, and contributed fixes on top of the generated code (including the
grow()re-allocation fix and the HIP/CUDA compile fixes). Several AI-generated changes were measured, found to be regressions, and dropped before submission (see Implementation Notes).Backward Compatibility
Mostly preserved; three user-visible changes:
neigh auto|shared|globalkeyword is added topair_style paceandpair_style pace/extrapolation. It is only recognized by the KOKKOS styles; the non-accelerated styles stop with an "unknown keyword" error if it is given. The defaultautoreproduces the previous behavior except that it falls back to level-1 scratch instead of aborting.pace/kknow runs the KOKKOS kernels instead of deferring to the non-accelerated evaluator. Results agree with the previous behavior to round-off. The previous restriction thatpace/kkcould only be run on a single CPU thread is lifted.recursiveevaluator still has no KOKKOS implementation. On host backends it now falls back to the non-accelerated evaluator with a warning instead of being rejected; on GPU backends it remains an error, as before.No changes are required to existing input scripts.
Implementation Notes
GPU optimizations (commits 1-9). The largest wins come from changing what is stored rather than how it is computed:
A_sph/weightsas separate re/im arraysComputeDerivativeradial looppow()for integer exponentsComputeFSfirst-touchdeep_copyAarray, readA_sphvia conjugate symmetryConjugateAipass on devicefr/dfrdirectlydB_flattenComputeWeightsinstead of storedComputeRadialintoComputeAiThe same ordering holds on both vendors, and the peak is at the same commit on each.
Shared-memory fallback (commit 10).
ComputeNeighcaches the short neighbor list in level-0 (on-chip) team scratch sizedteam_size*maxneigh*sizeof(int). The maximum available level-0 scratch is now queried from Kokkos viaTeamPolicy::scratch_size_max(0)rather than assuming a fixed limit, and the kernel transparently falls back to level-1 (global) scratch when the request does not fit, warning once. Theneighkeyword allows overriding the automatic choice.CPU-backend kernels (commits 11-12).
host_flagbecomes a compile-timestatic constexpr, so the host and device kernel sets are selected withif constexprand only one is ever instantiated (the same pattern aspair_snap_kokkos). The host kernels give one thread an entire atom and loop its neighbors inside, which removes every atomic and keeps the atom's accumulators as the innermost working set; the basis-function work is fused (rho -> F(rho) -> weights in one pass), flattened onto precomputed gather offsets, batched across eight atoms, and specialized on the density count at compile time.Because the two backends want opposite things from the same data, both layouts are kept and only the live one is allocated: the host uses interleaved-complex
A/A_sph/weights(re and im in the same cache line for a thread walking one atom), the device uses the split re/im arrays with the atom index innermost (coalescing for a warp scattering across atoms).fr/dfrare shared.Changes that were tried and dropped. Tiling
ComputeRho/ComputeWeightsas anMDRangePolicyand adding aLaunchBoundsoccupancy knob were both benchmarked and removed:LaunchBoundswas neutral to -1% and theMDRangeconversion cost 16% on MI300A (and ~0.5% on H100). An optional spline-free "direct" Chebyshev radial path and anO(rank)prefix/suffix rewrite of the leave-one-out products were also implemented, measured, and dropped (the first for adding ~150 lines for an unmeasurable gain, the second as a large regression caused by the added register pressure).Correctness. Verified against the non-accelerated evaluator on the fcc Cu
productdeck, with both branches of the compile-time host/device switch exercised: per-atom forces agree to 9.8e-13 eV/A and energies to ~13 significant digits. All threeneighkeyword values give identical forces, an unknown value is rejected, and every commit in the series builds individually.kokkos_type.hgains a one-word fix:SNAComplex::real_part_product(const complex &)is nowconst, which resolves an overload-ambiguity warning that GCC emitted at every call site.Post Submission Checklist
Further Information, Files, and Links
The
neighkeyword is documented indoc/src/pair_pace.rstwith a.. versionadded:: TBDdirective.Still outstanding at the time of submission: validation of the new host kernels under Kokkos OpenMP with more than one thread is in progress. The kernels are written one-atom-per-thread specifically so that no atomics are needed, but the threaded path has not yet been measured or verified here, and reviewers should treat that as unconfirmed until it is reported in this pull request.
Generated by Claude Code