feat: add reduction ops conversion for MPI + cleanup impl

src/KokkosComm/nccl/allreduce.hpp

@@ -9,6 +9,7 @@
 #include <KokkosComm/concepts.hpp>
 #include <KokkosComm/traits.hpp>
 #include <KokkosComm/datatype.hpp>
+#include <KokkosComm/reduction_op.hpp>
 #include "nccl_space.hpp"
 
 #include "impl/pack_traits.hpp"
@@ -49,7 +50,7 @@ namespace Impl {
 template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp>
 struct AllReduce<SendView, RecvView, RedOp, Kokkos::Cuda, NcclSpace> {
   static auto execute(Handle<Kokkos::Cuda, NcclSpace> &h, const SendView sv, RecvView rv) -> Req<NcclSpace> {
-    return nccl::allreduce(h.space(), sv, rv, nccl::Impl::reduction_op_v<RedOp>, h.comm());
+    return nccl::allreduce(h.space(), sv, rv, reduction_op<NcclSpace, RedOp>(), h.comm());
   }
 };
 

src/KokkosComm/nccl/reduce.hpp

@@ -9,6 +9,7 @@
 #include <KokkosComm/concepts.hpp>
 #include <KokkosComm/traits.hpp>
 #include <KokkosComm/datatype.hpp>
+#include <KokkosComm/reduction_op.hpp>
 #include "nccl_space.hpp"
 
 #include <KokkosComm/impl/contiguous.hpp>
@@ -70,7 +71,7 @@ namespace Impl {
 template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp>
 struct Reduce<SendView, RecvView, RedOp, Kokkos::Cuda, NcclSpace> {
   static auto execute(Handle<Kokkos::Cuda, NcclSpace> &h, const SendView sv, RecvView rv, int root) -> Req<NcclSpace> {
-    return nccl::reduce(h.space(), sv, rv, nccl::Impl::reduction_op_v<RedOp>, root, h.rank(), h.comm());
+    return nccl::reduce(h.space(), sv, rv, reduction_op<NcclSpace, RedOp>(), root, h.rank(), h.comm());
   }
 };
 

src/KokkosComm/reduction_op.hpp

@@ -1,107 +1,120 @@
-//@HEADER
-// ************************************************************************
-//
-//                        Kokkos v. 4.0
-//       Copyright (2022) National Technology & Engineering
-//               Solutions of Sandia, LLC (NTESS).
-//
-// Under the terms of Contract DE-NA0003525 with NTESS,
-// the U.S. Government retains certain rights in this software.
-//
-// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
-// See https://kokkos.org/LICENSE for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//@HEADER
+// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
 
 #pragma once
 
+#include <type_traits>
+
 #include <Kokkos_Core.hpp>
 #ifdef KOKKOSCOMM_ENABLE_NCCL
 #include <nccl.h>
 #endif
 
 #include <KokkosComm/concepts.hpp>
+#include "mpi/mpi_space.hpp"
+#ifdef KOKKOSCOMM_ENABLE_NCCL
+#include "nccl/nccl_space.hpp"
+#endif
 
 namespace KokkosComm {
+namespace {
+
+// clang-format off
+#define DECL_REDUCTION_OP_FOR(operator) \
+  struct operator {};                   \
+  template <> struct Impl::is_reduction_operator<operator> : public std::true_type {}
+// clang-format on
+
+}  // namespace
+
+DECL_REDUCTION_OP_FOR(BAnd);
+DECL_REDUCTION_OP_FOR(BOr);
+DECL_REDUCTION_OP_FOR(BXor);
+DECL_REDUCTION_OP_FOR(LAnd);
+DECL_REDUCTION_OP_FOR(LOr);
+DECL_REDUCTION_OP_FOR(LXor);
+DECL_REDUCTION_OP_FOR(Max);
+DECL_REDUCTION_OP_FOR(MaxLoc);
+DECL_REDUCTION_OP_FOR(Min);
+DECL_REDUCTION_OP_FOR(MinLoc);
+DECL_REDUCTION_OP_FOR(Sum);
+DECL_REDUCTION_OP_FOR(Prod);
+DECL_REDUCTION_OP_FOR(Average);
+
+namespace Impl {
+
+template <ReductionOperator RO>
+constexpr auto mpi_reduction_op() -> MPI_Op {
+  if constexpr (std::is_same_v<RO, BAnd>) {
+    return MPI_BAND;
+  } else if constexpr (std::is_same_v<RO, BOr>) {
+    return MPI_BOR;
+  } else if constexpr (std::is_same_v<RO, BXor>) {
+    return MPI_BXOR;
+  } else if constexpr (std::is_same_v<RO, LAnd>) {
+    return MPI_LAND;
+  } else if constexpr (std::is_same_v<RO, LOr>) {
+    return MPI_LOR;
+  } else if constexpr (std::is_same_v<RO, LXor>) {
+    return MPI_LXOR;
+  } else if constexpr (std::is_same_v<RO, Max>) {
+    return MPI_MAX;
+  } else if constexpr (std::is_same_v<RO, MaxLoc>) {
+    return MPI_MAXLOC;
+  } else if constexpr (std::is_same_v<RO, Min>) {
+    return MPI_MIN;
+  } else if constexpr (std::is_same_v<RO, MinLoc>) {
+    return MPI_MINLOC;
+  } else if constexpr (std::is_same_v<RO, Sum>) {
+    return MPI_SUM;
+  } else if constexpr (std::is_same_v<RO, Prod>) {
+    return MPI_PROD;
+  } else {
+    static_assert(std::is_void_v<RO>, "KokkosComm::Impl::mpi_reduction_op: operator not implemented");
+    return MPI_SUM;  // unreachable
+  }
+}
 
-struct BAnd {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::BAnd> : public std::true_type {};
-
-struct BOr {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::BOr> : public std::true_type {};
-
-struct LAnd {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::LAnd> : public std::true_type {};
-
-struct LOr {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::LOr> : public std::true_type {};
-
-struct Max {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::Max> : public std::true_type {};
-
-struct MaxLoc {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::MaxLoc> : public std::true_type {};
-
-struct Min {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::Min> : public std::true_type {};
-
-struct MinLoc {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::MinLoc> : public std::true_type {};
-
-struct MinMax {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::MinMax> : public std::true_type {};
-
-struct MinMaxLoc {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::MinMaxLoc> : public std::true_type {};
-
-struct Sum {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::Sum> : public std::true_type {};
-
-struct Prod {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::Prod> : public std::true_type {};
-
-struct Average {};
-template <>
-struct KokkosComm::Impl::is_reduction_operator<KokkosComm::Average> : public std::true_type {};
-
-#ifdef KOKKOSCOMM_ENABLE_NCCL
-namespace Experimental::nccl::Impl {
-
-template <ReductionOperator RedOp>
-constexpr auto reduction_op() -> ncclRedOp_t {
-  if constexpr (std::is_same_v<RedOp, KokkosComm::Max>) {
-    return ncclMax;
-  } else if constexpr (std::is_same_v<RedOp, KokkosComm::Min>) {
-    return ncclMin;
-  } else if constexpr (std::is_same_v<RedOp, KokkosComm::Sum>) {
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+template <ReductionOperator RO>
+constexpr auto nccl_reduction_op() -> ncclRedOp_t {
+  if constexpr (std::is_same_v<RO, Sum>) {
     return ncclSum;
-  } else if constexpr (std::is_same_v<RedOp, KokkosComm::Prod>) {
+  } else if constexpr (std::is_same_v<RO, Prod>) {
     return ncclProd;
-  } else if constexpr (std::is_same_v<RedOp, KokkosComm::Average>) {
+  } else if constexpr (std::is_same_v<RO, Min>) {
+    return ncclMin;
+  } else if constexpr (std::is_same_v<RO, Max>) {
+    return ncclMax;
+  } else if constexpr (std::is_same_v<RO, Average>) {
     return ncclAvg;
   } else {
-    static_assert(std::is_void_v<RedOp>, "nccl::reduction_op: operator not implemented");
-    return ncclMax;  // unreachable
+    static_assert(std::is_void_v<RO>, "KokkosComm::Impl::nccl_reduction_op: operator not implemented");
+    return ncclSum;  // unreachable
   }
 }
+#endif
 
-template <ReductionOperator RedOp>
-inline constexpr ncclRedOp_t reduction_op_v = reduction_op<RedOp>();
+}  // namespace Impl
 
-}  // namespace Experimental::nccl::Impl
+template <CommunicationSpace CS, ReductionOperator RO>
+[[nodiscard]] constexpr auto reduction_op() -> typename CS::reduction_op_type {
+  if constexpr (std::is_same_v<CS, MpiSpace>) {
+    return Impl::mpi_reduction_op<RO>();
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+  } else if constexpr (std::is_same_v<CS, Experimental::NcclSpace>) {
+    return Impl::nccl_reduction_op<RO>();
 #endif
+  } else {
+    static_assert(std::is_void_v<CS>,
+                  "KokkosComm::reduction_op: conversion not implemented for this communication space");
+    return Impl::mpi_reduction_op<RO>();  // unreachable
+  }
+}
+
+template <CommunicationSpace CS, ReductionOperator RO>
+[[nodiscard]] constexpr auto reduction_op_for(RO&&) -> typename CS::reduction_op_type {
+  return reduction_op<CS, std::remove_cvref_t<RO>>();
+}
 
 }  // namespace KokkosComm