A toy MPI wrapper for Kokkos.
- Requires c++20
- Requires c++23 for std::mdspan
macOS standard toolchain has incomplete std::mdspan
support: std::layout_stride
is not implemented as of xcode 15.3 release candidates
At sandia, with MPICH and the VPN enabled, you may need to do this before running any tests:
export FI_PROVIDER=tcp
mkdir -p build && cd build
cmake .. \
-DCMAKE_CXX_COMPILER=mpicxx \
-DKokkosComm_ENABLE_MDSPAN=ON \
-DKokkosComm_USE_KOKKOS_MDSPAN=ON \
-DKokkos_DIR=/path/to/kokkos-install/lib/cmake/Kokkos
make
ctest
cwpearson.github.io/kokkos-mpi/
https://www.sphinx-doc.org/en/master/usage/domains/cpp.html
Kokkos::View |
mdspan |
|
---|---|---|
MPI_Isend | x | x |
MPI_Recv | x | x |
MPI_Send | x | |
MPI_Reduce |
- host data
mdspan
- device data
mdspan
- first pass could be a MpiDatatypePacker which just constructs an MPI Datatype matching the mdspan and hands it off to MPI to deal with the non-contiguous data
- second pass would be to somehow associate a Kokkos memory space with the
mdspan
so we know how to allocate intermediate packing buffers
- use
Kokkos::deep_copy
to handle packing and unpacking of non-contiguousKokkos::View
- When non-contiguous views are passed to an MPI function, a temporary contiguous view of matching extent is allocated, and
Kokkos::deep_copy
is used to pack the data.
- When non-contiguous views are passed to an MPI function, a temporary contiguous view of matching extent is allocated, and
- "Immediate" functions (e.g.
isend
) return aKokkosComm::Req
, which can bewait()
-ed to block until the input view can be reused.Req
also manages the lifetimes of any intermediate views needed for packing the data, releasing those views whenwait()
is complete. -
KokkosComm::Traits
is specialized forKokkos::View
andmdspan
- whether
View
needs to be packed or not - what
pack
does forView
- what
unpack
does forView
- spans (distance between beginning of first byte and end of last byte)
- whether
- macOS xcode 15.3 doesn't support
std::mdspan
, so we usekokkos/mdspan
. - Pluggable packing strategies
- This would probably be a template parameter on the interface, which would be specialized to actually implement the various MPI operations
- Constructing matching MPI datatype and sending
- Packing into a contiguous buffer and sending
- How to handle discriminate between mdspans of host or device data, for packing
- Alternatively, construct an MPI datatype matching the non-contiguous sp
- MPI threaded-ness and Kokkos backends (Serial with multiple instances, Threads, etc)
- Are there circumstances in which we can fuse packing into another kernel?
- A better pack/unpack interface
- Maybe a
PackTraits<View>
where users can specializePackTraits
for any types they want to handle - Also, could introduce a runtime packing argument to the various functions, like a pack tag
- Maybe a
- More convenient collective wrappers
- Outer dimension has destination rank?
- Custom reductions?
test_2dhalo.cpp
: a 2d halo exchangetest_sendrecv.cpp
: ping-pong between ranks 0 and 1
shopt -s globstar
clang-format-8 -i {src,unit_tests,perf_tests}/**/*.[ch]pp