forked from dftfeDevelopers/dftfe
-
Notifications
You must be signed in to change notification settings - Fork 0
/
MPICommunicatorP2PKernels.h
141 lines (130 loc) · 5.02 KB
/
MPICommunicatorP2PKernels.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
// ---------------------------------------------------------------------
//
// Copyright (c) 2017-2022 The Regents of the University of Michigan and DFT-FE
// authors.
//
// This file is part of the DFT-FE code.
//
// The DFT-FE code is free software; you can use it, redistribute
// it, and/or modify it under the terms of the GNU Lesser General
// Public License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
// The full text of the license can be found in the file LICENSE at
// the top level of the DFT-FE distribution.
//
// ---------------------------------------------------------------------
//
/*
* @author Sambit Das
*/
#ifndef dftfeMPICommunicatorP2PKernels_h
#define dftfeMPICommunicatorP2PKernels_h
#include <MemorySpaceType.h>
#include <MemoryStorage.h>
#include <TypeConfig.h>
#include <DataTypeOverloads.h>
#ifdef DFTFE_WITH_DEVICE
# include <DeviceTypeConfig.h>
#endif
namespace dftfe
{
namespace utils
{
template <typename ValueType, dftfe::utils::MemorySpace memorySpace>
class MPICommunicatorP2PKernels
{
public:
using SizeTypeVector = utils::MemoryStorage<size_type, memorySpace>;
/**
* @brief Function template for architecture adaptable gather kernel to send buffer
* @tparam ValueType the type of the number
* @tparam memorySpace
* @param[in] dataArray data array with locally owned entries
* @param[in] ownedLocalIndicesForTargetProcs
* @param[in] blockSize
* @param[out] sendBuffer
*/
template <typename ValueTypeComm>
static void
gatherLocallyOwnedEntriesSendBufferToTargetProcs(
const MemoryStorage<ValueType, memorySpace> &dataArray,
const SizeTypeVector &ownedLocalIndicesForTargetProcs,
const size_type blockSize,
MemoryStorage<ValueTypeComm, memorySpace> &sendBuffer);
/**
* @brief Function template for architecture adaptable accumlate kernel from recv buffer
* @tparam ValueType the type of the number
* @tparam memorySpace
* @param[in] recvBuffer
* @param[in] ownedLocalIndicesForTargetProcs
* @param[in] blockSize
* @param[out] dataArray
*/
template <typename ValueTypeComm>
static void
accumAddLocallyOwnedContrRecvBufferFromTargetProcs(
const MemoryStorage<ValueTypeComm, memorySpace> &recvBuffer,
const SizeTypeVector & ownedLocalIndicesForTargetProcs,
const size_type blockSize,
const size_type locallyOwnedSize,
const size_type ghostSize,
MemoryStorage<ValueType, memorySpace> &dataArray);
/**
* @brief Function template for copying type1 to type2
* @param[in] blockSize
* @param[in] type1Array
* @param[out] type2Array
*/
template <typename ValueType1, typename ValueType2>
static void
copyValueType1ArrToValueType2Arr(const size_type blockSize,
const ValueType1 *type1Array,
ValueType2 * type2Array);
};
#ifdef DFTFE_WITH_DEVICE
template <typename ValueType>
class MPICommunicatorP2PKernels<ValueType,
dftfe::utils::MemorySpace::DEVICE>
{
public:
template <typename ValueTypeComm>
static void
gatherLocallyOwnedEntriesSendBufferToTargetProcs(
const MemoryStorage<ValueType, dftfe::utils::MemorySpace::DEVICE>
&dataArray,
const MemoryStorage<size_type, dftfe::utils::MemorySpace::DEVICE>
& ownedLocalIndicesForTargetProcs,
const size_type blockSize,
MemoryStorage<ValueTypeComm, dftfe::utils::MemorySpace::DEVICE>
& sendBuffer,
dftfe::utils::deviceStream_t deviceCommStream);
template <typename ValueTypeComm>
static void
accumAddLocallyOwnedContrRecvBufferFromTargetProcs(
const MemoryStorage<ValueTypeComm, dftfe::utils::MemorySpace::DEVICE>
&recvBuffer,
const MemoryStorage<size_type, dftfe::utils::MemorySpace::DEVICE>
& ownedLocalIndicesForTargetProcs,
const size_type blockSize,
const size_type locallyOwnedSize,
const size_type ghostSize,
MemoryStorage<ValueType, dftfe::utils::MemorySpace::DEVICE> &dataArray,
dftfe::utils::deviceStream_t deviceCommStream);
/**
* @brief Function template for copying type1 to type2
* @param[in] blockSize
* @param[in] type1Array
* @param[out] type2Array
*/
template <typename ValueType1, typename ValueType2>
static void
copyValueType1ArrToValueType2Arr(
const size_type blockSize,
const ValueType1 * type1Array,
ValueType2 * type2Array,
dftfe::utils::deviceStream_t deviceCommStream);
};
#endif
} // namespace utils
} // namespace dftfe
#endif