Implemented kspace_modify nonblocking yes for Kokkos-based KSPACE

[hagertnl]

Summary

Implements kspace_modify nonblocking yes, which enables using non-blocking MPI_Isend during FFT data exchange. This requires more memory than the traditional MPI_Send (but the same as collective yes), and is optimal in many cases. I've included some scaling data below. This can result in performance benefits at medium scale, especially on smaller MD systems on larger node counts. Helps to overlap communication with unpacking of received data.

Author(s)

Nick Hagerty, ORNL, hagertynl@ornl.gov

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

These changes are fully backwards compatible, with no performance impacts. It is a new option.

Implementation Notes

At the bottom of the PR is performance data, using upstream LAMMPS (commit abba7fc) pt2pt as the reference (relative value 1.0), then implementing and re-testing with default (blocking MPI), non-blocking MPI, and collectives. This is normalized to upstream blocking MPI (higher is better). We used the Rhodo benchmark with a 4x4x4 replica per Frontier node (for ~2mil atoms per node) and 8x8x8 replica per Frontier node (~16mil atoms per node).

Most lines are 1.00 +/- 1%, with the exception of the non-blocking and the collectives as we scale out, which diverge up (good, 20% performance boost) and down (bad, 20+% worse than blocking pt2pt MPI), respectively.

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

[hagertnl]

@stanmoore1 you may be the one most interested in reviewing this PR, since your reviewed the collective yes in the past. I've initially written it with the isend keyword, but I'm not opposed to changing. Wanted to make sure things look good, then I'll write the docs for it.

Since this is about as much optimization of the comms as I can think of doing, it may also be a good time to copy these changes into non-Kokkos KSPACE. Let me know if you'd like that to be included in this PR.

[hagertnl]

I also verified that energy breakdowns match between all 3 communication implementations and the baseline from upstream.

[hagertnl]

There's several loops that I think could benefit from OpenMP threading, too, but I see very few OpenMP pragmas throughout the source code, outside of the OPENMP package. Is there a reason for that, or am I free to add OpenMP if I chose?

[stanmoore1]

@hagertnl nice work! This looks clean and should be good to go once you write the documentation. If you have time to port this to non-Kokkos KSPACE that would be great too. I don't think a typical user will know what isend means, so nonblocking yes or send/nonblock yes may be better. @sjplimp or @akohlmey may have an opinion.

There's several loops that I think could benefit from OpenMP threading, too, but I see very few OpenMP pragmas throughout the source code, outside of the OPENMP package. Is there a reason for that, or am I free to add OpenMP if I chose?

If you want to thread over loops in the KOKKOS package then you can add Kokkos parallel_for with a Host execution space which will give OpenMP if a user compiles with HIP + OpenMP Kokkos backends. Raw OpenMP threading needs to be contained in the OpenMP package to maintain the simplicity and readability of the vanilla code.

[hagertnl]

Thanks for the feedback! I agree that nonblocking would be more consistent with collective, so I'll swap keywords to that. Thanks for the context on OpenMP. I'll evaluate and do some profiling to see if it's worth it. I will see if I can get it ported to the non-Kokkos implementation as well.

[Copilot]

Pull request overview

This PR implements kspace_modify isend yes for Kokkos-based KSPACE styles (PPPM), enabling non-blocking MPI_Isend operations during FFT data exchange. This provides performance benefits at medium scale, especially for smaller MD systems on larger node counts by overlapping communication with data unpacking. The feature requires more memory than blocking MPI_Send (similar to collective mode) and is opt-in with no performance impact when disabled.

Key Changes

• Added isend_flag to KSpace base class for controlling non-blocking MPI behavior
• Modified Kokkos FFT3d and Remap classes to support MPI_Isend with proper buffer management
• Updated PPPM Kokkos implementation to propagate the isend flag through FFT/remap construction

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/kspace.h	Added isend_flag member variable to enable/disable MPI_Isend
src/kspace.cpp	Added command parsing for kspace_modify isend and flag initialization
src/KOKKOS/remap_kokkos.h	Added struct members for isend support (isend_reqs, send_bufloc, useisend flag)
src/KOKKOS/remap_kokkos.cpp	Implemented isend logic with buffer management, MPI_Isend calls, and MPI_Waitall synchronization
src/KOKKOS/pppm_kokkos.cpp	Propagated isend_flag to FFT3d and Remap constructors
src/KOKKOS/fft3d_kokkos.h	Updated function signatures to accept useisend parameter
src/KOKKOS/fft3d_kokkos.cpp	Propagated useisend parameter through FFT plan creation

---

You can also share your feedback on Copilot code review for a chance to win a $100 gift card. Take the survey.

[Copilot]

Missing null check for isend_reqs allocation. If malloc fails, isend_reqs will be null but execution continues. Add a null check similar to line 384: if (plan->isend_reqs == nullptr) return nullptr;

Suggested change

	if (plan->useisend)
	plan->isend_reqs = (MPI_Request *) malloc(nsend*sizeof(MPI_Request));
	plan->send_bufloc = (int *) malloc(nsend*sizeof(int));
	if (plan->send_bufloc == nullptr) return nullptr;
	if (plan->useisend) {
	plan->isend_reqs = (MPI_Request *) malloc(nsend*sizeof(MPI_Request));
	if (plan->isend_reqs == nullptr) return nullptr;
	}
	plan->send_bufloc = (int *) malloc(nsend*sizeof(int));
	if (plan->send_bufloc == nullptr) return nullptr;

[hagertnl]

The suggested conditional is wrong (send_bufloc gets allocated even if !plan->useisend), but I can add in a check to make sure isend_reqs is not nullptr, too.

[hagertnl]

I've added the refactored comms backend to the non-Kokkos KSPACE, and am currently testing on Frontier at 56 cores per node. Compared to upstream, I think I am seeing a slight performance regression (<=5%) just comparing the default blocking pt2pt MPI of upstream to the new implementation, but I don't have all my testing data in yet. If there is a performance regression, I'll identify a specific case and do some profiling to see what might be taking longer.

[hagertnl]

Over the weekend, I had some benchmarking runs for the non-Kokkos implementation. Results in plots below, normalized speedups compared to upstream pt2pt (default settings). The 2 worst-scaling are the collectives ones (both upstream & pt2pt), but what's puzzling me is that the new collective implementation (which is what I used on Kokkos' remap and copied over here) is significantly slower than the original nested for loop. I suspect that the MPI_Allreduce is the bottleneck here, so I'll try re-implementing as a nested for loop to see if that helps performance. Otherwise, using non-blocking comms seems to be able to give up to around a 5-10% speedup in certain cases.

"Small" is 256K atoms per node, "large" is 2M atoms per node, using the Rhodo benchmark.

[hagertnl]

I replaced the Allreduce-based commringlist creation for collectives in non-Kokkos KSPACE with a C++ std::set method. This is implicitly what the previous code was doing, but using std::set removes a lot of unnecessary code, like for-loops to check if something exists in the commring already, and sorting. From what I've seen so far, this performs exactly the same as the pure-C method before, but with cleaner code. I want to do a quick check to see if it's beneficial for the Kokkos implementation, then I'll update docs and get the PR ready for merge.

[hagertnl]

The set-based implementation resulted in very good performance for the Kokkos side -- nearly double increase in timesteps/s for a small system on 512 nodes. Smaller boosts for other system sizes and node counts, but overall seems like a good change for both small and large systems. It looks like MPI_Allreduce (at least on cray-mpich) just isn't performant enough to benefit here.

Another interesting bit is that this set-based local comm ring build results in multiple different MPI groups (and multiple Alltoallv's across subsets of the processors), while the MPI_Allreduce approach created one single MPI group that handled all communication for all FFTs at once. I imagine multiple Alltoallv's would also be optimal. Fewer 0's, more concentrated communications.

Working on docs now.

[hagertnl]

Here are the updated performance plots. At 512 nodes, we still do generally see a pretty sharp drop-off, both relative to 64-node performance and relative to the baseline blocking pt2pt implementation, but in non-Kokkos implementation, it's equal or better than the existing implementation, and in the Kokkos implementation, it's much better than the existing Allreduce-based method.

Kokkos-based (both large & small):

Small non-kokkos system:

Large non-kokkos system:

[hagertnl]

I've been bothered by that purple line in the large CPU system, where the 8-node case shows a significant regression from upstream. I tracked this down to primarily the usage of plan->self, which handled whether send/recv's to self should be memcpy or MPI_Send. This was copied over from the Irecv/Send implementation when I re-did the Kokkos backend, and it's actually better to NOT do plan->self with collectives. So I removed that in both Kokkos & non-Kokkos case.

Beyond that, performance is quite variable, since I'm only doing 3 short runs. So the current performance is now in line with what we would've seen before.

[stanmoore1]

@hagertnl thanks so much for working on this. Is this PR ready to merge now, or do you have more work to do?

[hagertnl]

AFAIK, I'm finished! Docs should be updated for nonblocking yes/no, so should be good on that front.

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 4 comments.