Skip to content

Commit 01dc842

Browse files
cwpearsoncwpearson
committed
Initial attempt at comm/space wrapping and noncontig views
1 parent fbaf029 commit 01dc842

5 files changed

+203
-0
lines changed

‎src/KokkosComm_config.hpp

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
#include "mpi/KokkosComm_mpi.hpp"
20+
#include "KokkosComm_concepts.hpp"
21+
22+
namespace KokkosComm {
23+
24+
using GenericTransport = ::KokkosComm::Mpi;
25+
using SpecialTransport = ::KokkosComm::Mpi;
26+
template <KokkosExecutionSpace ExecSpace>
27+
using Handle = Mpi::Handle<ExecSpace>;
28+
29+
} // namespace KokkosComm

‎src/KokkosComm_plan.hpp

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
namespace KokkosComm {
20+
21+
template <KokkosExecutionSpace ExecSpace, typename CommFunc>
22+
Mpi::Handle<ExecSpace> plan(const ExecSpace &space, MPI_Comm comm, CommFunc f) {
23+
return Mpi::Plan<ExecSpace, CommFunc>(space, comm, f).handle();
24+
}
25+
26+
} // namespace KokkosComm

‎src/KokkosComm_point_to_point.hpp

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
#include "KokkosComm_api.hpp"
20+
#include "KokkosComm_concepts.hpp"
21+
#include "KokkosComm_config.hpp"
22+
#include "KokkosComm_plan.hpp"
23+
24+
namespace KokkosComm {
25+
26+
template <typename Handle, KokkosView RecvView>
27+
void irecv(Handle &h, RecvView &rv, int src, int tag) {
28+
if constexpr (Impl::api_avail_v<SpecialTransport, Impl::Api::Irecv>) {
29+
SpecialTransport::irecv(h, rv, src, tag);
30+
} else {
31+
GenericTransport::irecv(h, rv, src, tag);
32+
}
33+
}
34+
35+
template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
36+
KokkosComm::Handle<ExecSpace> irecv(const ExecSpace &space, const RecvView &rv, int dest, int tag, MPI_Comm comm) {
37+
using MyHandle = KokkosComm::Handle<ExecSpace>;
38+
return KokkosComm::plan(space, comm, [=](MyHandle &handle) { KokkosComm::irecv(handle, rv, dest, tag); });
39+
}
40+
41+
template <typename Handle, KokkosView SendView>
42+
void isend(Handle &h, SendView &sv, int src, int tag) {
43+
if constexpr (Impl::api_avail_v<SpecialTransport, Impl::Api::Isend>) {
44+
SpecialTransport::isend(h, sv, src, tag);
45+
} else {
46+
GenericTransport::isend(h, sv, src, tag);
47+
}
48+
}
49+
50+
template <KokkosExecutionSpace ExecSpace, KokkosView SendView>
51+
KokkosComm::Handle<ExecSpace> isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
52+
using MyHandle = KokkosComm::Handle<ExecSpace>;
53+
return KokkosComm::plan(space, comm, [=](MyHandle &handle) { KokkosComm::isend(handle, sv, dest, tag); });
54+
}
55+
56+
} // namespace KokkosComm

‎src/impl/KokkosComm_api.hpp

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
namespace KokkosComm::Impl {
20+
21+
enum class Api { Irecv, Isend };
22+
23+
// catch-all: no transports implement any APIs
24+
template <typename Transport, Impl::Api API>
25+
struct api_avail : public std::false_type {};
26+
27+
template <typename Transport, Impl::Api API>
28+
constexpr bool api_avail_v = api_avail<Transport, API>::value;
29+
30+
} // namespace KokkosComm::Impl

‎src/impl/KokkosComm_contiguous.hpp

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
#include <string>
20+
21+
#include <Kokkos_Core.hpp>
22+
23+
#include "KokkosComm_concepts.hpp"
24+
25+
namespace KokkosComm::Impl {
26+
27+
template <KokkosView View>
28+
struct contiguous_view {
29+
using type = Kokkos::View<typename View::non_const_data_type, Kokkos::LayoutRight, typename View::memory_space>;
30+
};
31+
32+
template <KokkosView View>
33+
using contiguous_view_t = contiguous_view<View>::type;
34+
35+
template <KokkosView View, KokkosExecutionSpace Space>
36+
auto allocate_contiguous_for(const Space &space, const std::string &label, View &v) {
37+
using non_const_packed_view_type = contiguous_view_t<View>;
38+
39+
if constexpr (KokkosComm::rank<View>() == 1) {
40+
return non_const_packed_view_type(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, label), v.extent(0));
41+
} else if constexpr (KokkosComm::rank<View>() == 2) {
42+
return non_const_packed_view_type(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, label), v.extent(0),
43+
v.extent(1));
44+
} else {
45+
static_assert(std::is_void_v<View>, "allocate_contiguous_for for views > rank 2 not implemented");
46+
}
47+
}
48+
49+
template <KokkosExecutionSpace Space, KokkosView DstView, KokkosView SrcView>
50+
auto resize_contiguous_for(const Space &space, DstView &out, const SrcView &in) {
51+
static_assert(DstView::rank == SrcView::rank, "");
52+
53+
if constexpr (KokkosComm::rank<DstView>() == 1) {
54+
Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0));
55+
} else if constexpr (KokkosComm::rank<DstView>() == 2) {
56+
Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0), in.extent(1));
57+
} else {
58+
static_assert(std::is_void_v<DstView>, "realloc_contiguous_for for views > rank 2 not implemented");
59+
}
60+
}
61+
62+
} // namespace KokkosComm::Impl

0 commit comments

Comments
 (0)