Adding support for Kokkos kernel params autotuner

[ndtrung81]

Summary

This PR adds a performance tuner for the GPU kernel launch parameters for the KOKKOS package. The tuner, if enabled, will scan through the possible values of the team sizes and vector sizes, measures the overall performance within a specified interval, and find the combination that gives the best performance. The optimal combination will be used for the pair_compute_neighlist kernel(defined inpair_kokkos.h`) from that point on.

This PR is built upon the features provided by #4573. Based on the findings of that PR, the most influential kernel parameters are pair/team/size and threads/per/atom, that is, teamsize and vectorsize, respectively, of the pair style kernels. That's the motivation for scanning over these 2 parameters in the /kk pair styles in this PR. Bond styles have a single parameter to tune (bond_block_size).

Related Issue(s)

N/A

Author(s)

Trung Nguyen (U Chicago)

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

By submitting this pull request, I confirm that I did NOT use any AI tools to generate
all or parts of the code and modifications in this pull request.

Backward Compatibility

Should maintain backward compatiblity

Implementation Notes

It would be more efficient if the tuner can be refactored out of individual pair styles, to avoid the repetition for each /kk pair style.

Post Submission Checklist

• The feature or features in this pull request is complete
• Licensing information is complete
• Corresponding author information is complete
• The source code follows the LAMMPS formatting guidelines
• Suitable new documentation files and/or updates to the existing docs are included
• The added/updated documentation is integrated and tested with the documentation build system
• The feature has been verified to work with the conventional build system
• The feature has been verified to work with the CMake based build system
• Suitable tests have been added to the unittest tree.
• A package specific README file has been included or updated
• One or more example input decks are included

Further Information, Files, and Links

[Copilot]

Pull request overview

This pull request adds an autotuning feature for GPU kernel parameters in the KOKKOS package. The autotuner scans through combinations of team sizes (64-512 in steps of 32) and vector sizes (threads per atom: 1, 2, 4, 8, 16, 32) to find the optimal configuration that gives the best performance. The feature is disabled by default and can be enabled via the package kokkos auto/tuning command with a specified interval for performance measurement.

Key changes:

• Adds new TuneKokkos class that manages the scanning and performance measurement logic
• Integrates autotuner into three pair styles: pair_lj_cut, pair_lj_cut_coul_long, and pair_lj_charmm_coul_long
• Adds auto/tuning keyword to the package kokkos command with interval parameter

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
src/KOKKOS/tune_kokkos.h	Header for new TuneKokkos class with performance tracking and parameter management
src/KOKKOS/tune_kokkos.cpp	Implementation of autotuner logic including parameter scanning, timing, and optimal selection
src/KOKKOS/pair_lj_cut_kokkos.h	Adds tuner member pointer to PairLJCutKokkos class
src/KOKKOS/pair_lj_cut_kokkos.cpp	Integrates autotuner in compute and init_style methods
src/KOKKOS/pair_lj_cut_coul_long_kokkos.h	Adds tuner member pointer to PairLJCutCoulLongKokkos class
src/KOKKOS/pair_lj_cut_coul_long_kokkos.cpp	Integrates autotuner in compute and init_style methods, adds missing copymode reset
src/KOKKOS/pair_lj_charmm_coul_long_kokkos.h	Adds tuner member pointer to PairLJCharmmCoulLongKokkos class
src/KOKKOS/pair_lj_charmm_coul_long_kokkos.cpp	Integrates autotuner in compute and init_style methods
src/KOKKOS/kokkos.h	Adds autotuning flag to KokkosLMP class
src/KOKKOS/kokkos.cpp	Implements auto/tuning command parsing and disables on non-GPU builds
cmake/Modules/Packages/KOKKOS.cmake	Adds tune_kokkos.cpp to KOKKOS package sources for CMake build
doc/src/package.rst	Documents the new auto/tuning keyword for package kokkos command

[ndtrung81]

Preliminary results from the autotuner:

• in.lj: 1.5X speedup vs. using default values:
  perf-lj.pdf

• in.rhodo: 1.3X speedup vs. using default values:
  perf-rhodo.pdf

[sjplimp]

@ndtrung81 I really like this new feature - thanks for this PR. Made several comments on the doc info for the autotuning option. Am ready to approve for merging once you respond.

It would also be very good to add a couple of your perf plots for autotuning to the Benchmark web page.

[ndtrung81]

@sjplimp thank you for the feedback! I have updated the doc page with your suggestion to clarify the case when nevery is 0 and mention the defaults at the bottom of the page.

Attached is the plot showing the performance with and without auto/tuning for in.lj and in.rhodo with nevery = 100, nsamples = 5, mode = median with 1 MPI on 1 A100 PCIe GPU, the KOKKOS package was built with single-precision.

Do you think they are suitable for the Benchmark page?

[sjplimp]

@ndtrung81 I think these should be added to the package command doc page, as well as the Benchmark page. Talk to @stanmoore1 about the Benchmark page.

[sjplimp]

@stanmoore1 I also think the Benchmark page could be updated to report Benchmark timings that reflect the increased autotuning timesteps rates, where they make a difference,

[sjplimp]

@ndtrung81 Have my other comments on the doc page been addressed? If so, then I will approve this PR for merging,

[sjplimp]

@akohlmey per email, you can merge this when Trung completes his final (small) edits

[ndtrung81]

@akohlmey I have tried to address Steve's comments and edits. The doc page is good from my point of view. I will assign it to you, unless you need me to make more changes. Thanks!

[akohlmey]

I will assign it to you, unless you need me to make more changes. Thanks!

@ndtrung81 I made a few cosmetic changes to package.rst and am now running this one more through code review. A few minor findings:

1. get_current_vector_size() indexes team_sizes instead of vector_sizes — live for the GENERIC dpd/dipole integrations that call it, yielding an invalid vector length.
2. utils::logmesg in kokkos.cpp has 3 {} placeholders but 2 arguments (and the first should be the interval, not perf_nsamples) — throws when autotuning is enabled.
3. if (tuner) delete tuner; is redundant. You can just do delete tuner;
4. package.rst default-settings line still lists auto/tuning = disabled.
5. Double blank line around the #else // LMP_KOKKOS_GPU not defined block.
6. Constructor accepts interval == 0 / nsamples == 0 (< 0 checks); not reachable as a live div-by-zero since a tuner is only created when autotuning > 0, but <= 0 would be clearer defensively.
7. Side-effecting getter: get_optimal_combination_idx() mutates opt_perf and, in MEDIAN mode, std::sorts the performance rows in place — yet it's called from regular_performance_check() purely for logging. Make it const/non-mutating, or compute the median on a copy.
8. Repetition inside the tuner: allocate(), set_team_size_values(), set_vector_size_values() each duplicate the ncombinations/memory->create block.

Please check if these are valid concerns.

[ndtrung81]

@akohlmey thanks for the code review. These are valid concerns.

1. This is a bug, I need to index vector_sizes for get_current_vector_size() from the current combination index.
2. Missing autotuning for the first placeholder.
3. Will simplify this.
4. This is intentional because auto/tuning is by default disabled. What is your suggestion?
5. Will remove the double blank lines.
6. Will check of <= 0 for interval and nsamples.
7. Will make this change.
8. The duplicated memory->create() blocks (followed by destroy()) in allocate(), set_team_size_values(), set_vector_size_values() are because the two set functions may be called after allocate() is already called in the tuner's constructor. A Pair class may want to override the scanned values of team sizes or vector sizes.

[akohlmey]

@akohlmey thanks for the code review. These are valid concerns.

[...]

4. This is intentional because auto/tuning is by default disabled. What is your suggestion?

This was the result of the AI review. I was too tired to double check everything.
I think that disabling auto-tuning by default is the right call. We had a similar
thing for KSPACE (benchmarking the FFT time percentage and auto-tuning it) that
is now both disabled by default.

[ndtrung81]

@akohlmey I am done with addressing the concerns from code review.

[akohlmey]

@akohlmey I am done with addressing the concerns from code review.

Then let me quickly merge it before somebody get any more ideas for changes...