feat(mpi): add non-blocking reduce (`mpi:ireduce`) by dssgabriel · Pull Request #198 · kokkos/kokkos-comm

dssgabriel · 2025-12-04T08:53:09Z

Prepare for exposure in "core" API.

Based on #187 and #188.

dssgabriel · 2025-12-18T13:57:31Z

I am adding a unit test on this because it fails on MPI + CUDA for a reason I don't understand yet.

dssgabriel · 2025-12-19T00:02:40Z

Honestly, I am completely lost on why this fails on OMPI + CUDA...

I reworked the MPI packing logic to use KokkosComm/impl/contiguous.hpp (there was a TODO comment for this for some time anyway), and I managed to get it to pass on my workstation with OMPI + OMP, but it always fails with CUDA. I can't figure out why, so I am abandoning for tonight. GDB debugging/backtrace did not help.

If someone has an idea about what is wrong with the test/the ireduce implementation, I'd be more than happy to receive some feedback.

cedricchevalier19 · 2026-01-15T16:56:59Z

@dssgabriel can you comment this issue and tell us what is needed to merge this PR.

dssgabriel · 2026-01-15T20:39:13Z

Okay, so the issue is that Open MPI does not implement non-blocking reduction operations with CUDA buffers.
This is referenced in several places:

Open MPI's documentation:
- https://docs.open-mpi.org/en/v5.0.x/tuning-apps/networking/cuda.html#which-mpi-apis-do-not-work-with-cuda-aware
- https://docs.open-mpi.org/en/v5.0.x/tuning-apps/networking/cuda.html#which-mpi-apis-do-not-work-with-cuda-aware-ucx (I think we are falling in this case?)
Open MPI issues:
- CUDA-aware Ireduce and Iallreduce operations for GPU tensors segfault open-mpi/ompi#9845
- coll collectives segfaults on CUDA buffers open-mpi/ompi#12045

Some work was done (started by @devreal 🙂), but it does not seem to have been picked up by anyone on the Open MPI team?
Initial work: open-mpi/ompi#12318, then split into:

Do we want to voice our needs and open an issue on the Open MPI repo too?

dssgabriel · 2026-01-15T20:41:04Z

I will mark this function as unimplemented if running the Open MPI impl.
This will be reversed once #205 is merged, when implementing host staging everywhere in the low-level MPI bindings.

dssgabriel · 2026-01-15T21:06:24Z

The failing jobs are expected since we error out on Open MPI. Do we merge like this @cedricchevalier19 @cwpearson ?

cedricchevalier19 · 2026-01-28T19:48:31Z

Can you add something to skip these tests with OpenMPI? I'll prefer if the CI to always pass on the main branch.

dssgabriel · 2026-01-30T13:02:22Z

It looks like the issue on the compile-only job is due to improper #includes in our headers. The #error PP directive ends up triggering a compilation failure in unrelated tests. I can't track down the exact issue, but I assume eliminating it would likely require moving some code around, which isn't ideal for this PR.

@cedricchevalier19 Do you want to merge like this, or should we attempt fixing this first?

cedricchevalier19 · 2026-01-30T16:47:47Z

I'll prefer not to merge something with broken tests.
We can try to fix it in this PR, and extract it in a separate PR then rebase.

Copilot

Pull request overview

This pull request adds support for non-blocking reduce operations (mpi::ireduce) to KokkosComm, preparing the functionality for exposure in the core API. The changes are based on two previous PRs: #187 (reduction operator conversion) and #188 (CommunicationSpace refactoring).

Changes:

Added new non-blocking reduce function KokkosComm::mpi::ireduce with support for both contiguous and non-contiguous views
Refactored MPI packer implementation to use a common contiguous view utility and moved from KokkosComm::Impl to KokkosComm::mpi::Impl namespace
Updated header organization to support both MPI and NCCL communication spaces simultaneously (changed from #elif to separate #if blocks)

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
unit_tests/mpi/test_ireduce.cpp	New test file for non-blocking reduce with contiguous 1D views
unit_tests/CMakeLists.txt	Adds test configuration for the new ireduce test
src/KokkosComm/mpi/reduce.hpp	Implements new `ireduce` function and updates existing `reduce` functions with namespace changes
src/KokkosComm/mpi/impl/packer.hpp	Refactors packer to use common contiguous view utilities and updates namespace
src/KokkosComm/mpi/impl/pack_traits.hpp	Updates namespace from `KokkosComm::Impl` to `KokkosComm::mpi::Impl`
src/KokkosComm/mpi/send.hpp	Updates PackTraits namespace references and adds label parameter to pack calls
src/KokkosComm/mpi/recv.hpp	Updates PackTraits namespace references
src/KokkosComm/mpi/isend.hpp	Updates PackTraits namespace references and adds label parameter to pack calls
src/KokkosComm/mpi/irecv.hpp	Updates PackTraits namespace references
src/KokkosComm/mpi/allreduce.hpp	Adds static_assert for unsupported Open MPI + CUDA/HIP configuration and formatting updates
src/KokkosComm/collective.hpp	Updates includes to add allreduce.hpp and reduce.hpp, removes redundant infrastructure includes
src/KokkosComm/point_to_point.hpp	Adds MPI isend/irecv includes and removes redundant NCCL infrastructure includes
src/KokkosComm/KokkosComm.hpp	Changes from mutually exclusive `#elif` to separate `#if` blocks to allow both MPI and NCCL simultaneously, adds NCCL collective includes
unit_tests/test_sendrecv.cpp	Refactors to use ExecSpace alias and explicit Handle construction
unit_tests/test_allreduce.cpp	Adds GTEST_SKIP guards for Open MPI + CUDA/HIP, formatting updates
unit_tests/test_allgather.cpp	Updates communication space names from `Mpi`/`Nccl` to `MpiSpace`/`NcclSpace`
unit_tests/mpi/test_reduce.cpp	Adds include for mpi.h and Kokkos_Core.hpp

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Needed to propagate changes to isend/irecv which were using our first attempt at doing non-contiguous view packing. Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

`static_assert(false)` always triggers a compilation failure, even in templated functions that are never instantiated, until [CWG2518/P2593](https://cplusplus.github.io/CWG/issues/2518.html). This is a workaround that allows us to have a type-dependent `false` value that can only be evaluated at template instantiation-time, thus avoiding spurious static assertions failures in functions that are never called. This workaround is only temporarily needed for disabling reduce/all-reduce support for Open MPI+CUDA/HIP combinations, which don't work, and will be removed once host staging is implemented for the MPI backend. Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

With non-contiguous data support. Disabled when MPI impl is Open MPI & Kokkos CUDA/HIP backend are active. Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

Skip when Open MPI + CUDA/HIP Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

Only disable `ireduce` and `iallreduce` when the Views are in CUDA or HIP execution spaces. Used `V::execution_space` instead of `V::memory_space` because there are multiple flavors of those, which MPI impls don't distinguish as far as support goes (it also makes for a big ugly boolean expression).

dssgabriel · 2026-02-05T19:35:32Z

@cedricchevalier19 This is (finally!) ready for review and merging.
Unfortunately there is no point in running the SNL CI since it's still with CMake 3.23 and #212 bumped it to 3.25.

cedricchevalier19

Add context for #ifdef

Co-authored-by: Cédric Chevalier <cedric.chevalier019@proton.me>

cedricchevalier19

Thanks a lot @dssgabriel!

dssgabriel requested review from cedricchevalier19 and cwpearson December 4, 2025 08:53

dssgabriel self-assigned this Dec 4, 2025

dssgabriel added C-enhancement Category: an enhancement A-mpi Area: KokkosComm MPI backend implementation labels Dec 4, 2025

dssgabriel mentioned this pull request Dec 4, 2025

feat(core): expose MPI-based reduce #199

Merged

cwpearson added the SNL-CI-APPROVAL Required to run SNL CI on non-SNL contributions label Dec 4, 2025

dssgabriel force-pushed the feature/mpi-ireduce branch from 91e723f to 3ecb363 Compare December 17, 2025 21:18

cwpearson added SNL-CI-APPROVAL Required to run SNL CI on non-SNL contributions and removed SNL-CI-APPROVAL Required to run SNL CI on non-SNL contributions labels Dec 17, 2025

dssgabriel force-pushed the feature/mpi-ireduce branch 3 times, most recently from b939894 to 04ba07a Compare December 18, 2025 13:56

cwpearson added SNL-CI-APPROVAL Required to run SNL CI on non-SNL contributions and removed SNL-CI-APPROVAL Required to run SNL CI on non-SNL contributions labels Dec 18, 2025

cwpearson added SNL-CI-APPROVAL Required to run SNL CI on non-SNL contributions and removed SNL-CI-APPROVAL Required to run SNL CI on non-SNL contributions labels Dec 19, 2025

cedricchevalier19 added this to the Version 0.1 milestone Jan 15, 2026

dssgabriel mentioned this pull request Jan 20, 2026

feat(mpi): implement host staging #208

Open

16 tasks

dssgabriel force-pushed the feature/mpi-ireduce branch from 4d878cc to e718bdc Compare January 30, 2026 11:20

dssgabriel force-pushed the feature/mpi-ireduce branch from 24ac107 to 9c9bf5a Compare February 2, 2026 23:03

dssgabriel force-pushed the feature/mpi-ireduce branch 2 times, most recently from 1f40de3 to 30d2133 Compare February 5, 2026 12:40

cedricchevalier19 requested a review from Copilot February 5, 2026 14:05

Copilot started reviewing on behalf of cedricchevalier19 February 5, 2026 14:05 View session

Copilot AI reviewed Feb 5, 2026

View reviewed changes

Comment thread src/KokkosComm/mpi/allreduce.hpp Outdated

Comment thread src/KokkosComm/mpi/impl/pack_traits.hpp Outdated

Comment thread src/KokkosComm/mpi/reduce.hpp

dssgabriel added 7 commits February 5, 2026 15:30

feat(mpi): use KC/contiguous to implement packing

4b556cc

Needed to propagate changes to isend/irecv which were using our first attempt at doing non-contiguous view packing. Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

feat(mpi): impl non-blocking ireduce binding

73d9a64

With non-contiguous data support. Disabled when MPI impl is Open MPI & Kokkos CUDA/HIP backend are active. Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

refactor(mpi): disable iallreduce for Open MPI + CUDA/HIP

4b99111

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

test(mpi): add ireduce test

2d7af02

Skip when Open MPI + CUDA/HIP Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

test(mpi): core allgather was broken with old comm space names

6538ae6

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

fix: header inclusion to allow both MPI & NCCL enabled together

5949734

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

dssgabriel force-pushed the feature/mpi-ireduce branch from aae9fcc to 5949734 Compare February 5, 2026 15:04

dssgabriel and others added 4 commits February 5, 2026 16:19

fix: wrong namespace + missing CMake for assert helper

bdfef94

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

Update src/KokkosComm/mpi/impl/pack_traits.hpp

962d0e2

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

fix: IWYU in MPI/NCCL bindings

7b1d8ec

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

refactor: remove helper header

a0a59ac

dssgabriel force-pushed the feature/mpi-ireduce branch from ae8a0b1 to dea249e Compare February 5, 2026 18:41

dssgabriel force-pushed the feature/mpi-ireduce branch from dea249e to 21e2227 Compare February 5, 2026 18:45

cedricchevalier19 mentioned this pull request Feb 10, 2026

Open MPI doesn't support non-blocking reduction colls on GPU #215

Open

cedricchevalier19 reviewed Feb 10, 2026

View reviewed changes

Comment thread unit_tests/test_allreduce.cpp

Comment thread unit_tests/test_allreduce.cpp

Comment thread src/KokkosComm/mpi/allreduce.hpp

Comment thread src/KokkosComm/mpi/reduce.hpp

dssgabriel and others added 2 commits February 10, 2026 11:02

chore: add FIXME notice to skipped test

924765b

Co-authored-by: Cédric Chevalier <cedric.chevalier019@proton.me>

chore: add FIXME notices addressing kokkos#215

e015c42

Co-authored-by: Cédric Chevalier <cedric.chevalier019@proton.me>

cedricchevalier19 approved these changes Feb 10, 2026

View reviewed changes

cedricchevalier19 merged commit a291c64 into kokkos:develop Feb 10, 2026
8 of 10 checks passed

dssgabriel deleted the feature/mpi-ireduce branch February 10, 2026 10:19

Uh oh!

Conversation

dssgabriel commented Dec 4, 2025

Uh oh!

dssgabriel commented Dec 18, 2025

Uh oh!

dssgabriel commented Dec 19, 2025

Uh oh!

cedricchevalier19 commented Jan 15, 2026

Uh oh!

dssgabriel commented Jan 15, 2026

Uh oh!

dssgabriel commented Jan 15, 2026

Uh oh!

dssgabriel commented Jan 15, 2026

Uh oh!

cedricchevalier19 commented Jan 28, 2026

Uh oh!

dssgabriel commented Jan 30, 2026

Uh oh!

cedricchevalier19 commented Jan 30, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dssgabriel commented Feb 5, 2026

Uh oh!

cedricchevalier19 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cedricchevalier19 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants