Line data Source code
1 : //
2 : // Class BareField
3 : // A BareField consists of multple LFields and represents a field.
4 : //
5 : #include "Ippl.h"
6 :
7 : #include <Kokkos_ReductionIdentity.hpp>
8 : #include <cstdlib>
9 : #include <limits>
10 : #include <map>
11 : #include <utility>
12 :
13 : #include "Communicate/DataTypes.h"
14 :
15 : #include "Utility/Inform.h"
16 : #include "Utility/IpplInfo.h"
17 : namespace Kokkos {
18 : template <typename T, unsigned Dim>
19 : struct reduction_identity<ippl::Vector<T, Dim>> {
20 : KOKKOS_FORCEINLINE_FUNCTION static ippl::Vector<T, Dim> sum() {
21 56 : return ippl::Vector<T, Dim>(0);
22 : }
23 : KOKKOS_FORCEINLINE_FUNCTION static ippl::Vector<T, Dim> prod() {
24 : return ippl::Vector<T, Dim>(1);
25 : }
26 : KOKKOS_FORCEINLINE_FUNCTION static ippl::Vector<T, Dim> min() {
27 48 : return ippl::Vector<T, Dim>(std::numeric_limits<T>::infinity());
28 : }
29 : KOKKOS_FORCEINLINE_FUNCTION static ippl::Vector<T, Dim> max() {
30 48 : return ippl::Vector<T, Dim>(-std::numeric_limits<T>::infinity());
31 : }
32 : };
33 : } // namespace Kokkos
34 :
35 : namespace KokkosCorrection {
36 : template <typename Scalar, class Space = Kokkos::HostSpace>
37 : struct Max : Kokkos::Max<Scalar, Space> {
38 : using Super = Kokkos::Max<Scalar, Space>;
39 : using value_type = typename Super::value_type;
40 48 : KOKKOS_INLINE_FUNCTION Max(value_type& vref)
41 48 : : Super(vref) {}
42 : KOKKOS_INLINE_FUNCTION void join(value_type& dest, const value_type& src) const {
43 : using ippl::max;
44 : using Kokkos::max;
45 : dest = max(dest, src);
46 : }
47 : };
48 : template <typename Scalar, class Space = Kokkos::HostSpace>
49 : struct Min : Kokkos::Min<Scalar, Space> {
50 : using Super = Kokkos::Min<Scalar, Space>;
51 : using value_type = typename Super::value_type;
52 48 : KOKKOS_INLINE_FUNCTION Min(value_type& vref)
53 48 : : Super(vref) {}
54 : KOKKOS_INLINE_FUNCTION void join(value_type& dest, const value_type& src) const {
55 : using ippl::min;
56 : using Kokkos::min;
57 : dest = min(dest, src);
58 : }
59 : };
60 : template <typename Scalar, class Space = Kokkos::HostSpace>
61 : struct Sum : Kokkos::Sum<Scalar, Space> {
62 : using Super = Kokkos::Sum<Scalar, Space>;
63 : using value_type = typename Super::value_type;
64 168 : KOKKOS_INLINE_FUNCTION Sum(value_type& vref)
65 168 : : Super(vref) {}
66 : KOKKOS_INLINE_FUNCTION void join(value_type& dest, const value_type& src) const {
67 : dest += src;
68 : }
69 : };
70 : template <typename Scalar, class Space = Kokkos::HostSpace>
71 : struct Prod : Kokkos::Prod<Scalar, Space> {
72 : using Super = Kokkos::Prod<Scalar, Space>;
73 : using value_type = typename Super::value_type;
74 24 : KOKKOS_INLINE_FUNCTION Prod(value_type& vref)
75 24 : : Super(vref) {}
76 : KOKKOS_INLINE_FUNCTION void join(value_type& dest, const value_type& src) const {
77 : dest *= src;
78 : }
79 : };
80 : } // namespace KokkosCorrection
81 :
82 : namespace ippl {
83 : namespace detail {
84 : template <typename T, unsigned Dim, class... ViewArgs>
85 : struct isExpression<BareField<T, Dim, ViewArgs...>> : std::true_type {};
86 : } // namespace detail
87 :
88 : template <typename T, unsigned Dim, class... ViewArgs>
89 48 : BareField<T, Dim, ViewArgs...>::BareField()
90 48 : : nghost_m(1)
91 48 : , layout_m(nullptr) {}
92 :
93 : template <typename T, unsigned Dim, class... ViewArgs>
94 24 : BareField<T, Dim, ViewArgs...> BareField<T, Dim, ViewArgs...>::deepCopy() const {
95 24 : BareField<T, Dim, ViewArgs...> copy(*layout_m, nghost_m);
96 24 : Kokkos::deep_copy(copy.dview_m, dview_m);
97 24 : return copy;
98 0 : }
99 :
100 : template <typename T, unsigned Dim, class... ViewArgs>
101 1508 : BareField<T, Dim, ViewArgs...>::BareField(Layout_t& l, int nghost)
102 1508 : : nghost_m(nghost)
103 : // , owned_m(0)
104 1508 : , layout_m(&l) {
105 1508 : setup();
106 1508 : }
107 :
108 : template <typename T, unsigned Dim, class... ViewArgs>
109 48 : void BareField<T, Dim, ViewArgs...>::initialize(Layout_t& l, int nghost) {
110 48 : if (layout_m == 0) {
111 48 : layout_m = &l;
112 48 : nghost_m = nghost;
113 48 : setup();
114 : }
115 48 : }
116 :
117 : // ML
118 : template <typename T, unsigned Dim, class... ViewArgs>
119 96 : void BareField<T, Dim, ViewArgs...>::updateLayout(Layout_t& l, int nghost) {
120 : // std::cout << "Got in BareField::updateLayout()" << std::endl;
121 96 : layout_m = &l;
122 96 : nghost_m = nghost;
123 96 : setup();
124 96 : }
125 :
126 : template <typename T, unsigned Dim, class... ViewArgs>
127 1652 : void BareField<T, Dim, ViewArgs...>::setup() {
128 1652 : owned_m = layout_m->getLocalNDIndex();
129 :
130 4956 : auto resize = [&]<size_t... Idx>(const std::index_sequence<Idx...>&) {
131 1652 : this->resize((owned_m[Idx].length() + 2 * nghost_m)...);
132 : };
133 1652 : resize(std::make_index_sequence<Dim>{});
134 1652 : }
135 :
136 : template <typename T, unsigned Dim, class... ViewArgs>
137 : template <typename... Args>
138 1652 : void BareField<T, Dim, ViewArgs...>::resize(Args... args) {
139 1652 : Kokkos::resize(dview_m, args...);
140 1652 : }
141 :
142 : template <typename T, unsigned Dim, class... ViewArgs>
143 292 : void BareField<T, Dim, ViewArgs...>::fillHalo() {
144 292 : if (layout_m->comm.size() > 1) {
145 292 : halo_m.fillHalo(dview_m, layout_m);
146 : }
147 292 : if (layout_m->isAllPeriodic_m) {
148 : using Op = typename detail::HaloCells<T, Dim, ViewArgs...>::assign;
149 0 : halo_m.template applyPeriodicSerialDim<Op>(dview_m, layout_m, nghost_m);
150 : }
151 292 : }
152 :
153 : template <typename T, unsigned Dim, class... ViewArgs>
154 140 : void BareField<T, Dim, ViewArgs...>::accumulateHalo() {
155 140 : if (layout_m->comm.size() > 1) {
156 140 : halo_m.accumulateHalo(dview_m, layout_m);
157 : }
158 140 : if (layout_m->isAllPeriodic_m) {
159 : using Op = typename detail::HaloCells<T, Dim, ViewArgs...>::rhs_plus_assign;
160 72 : halo_m.template applyPeriodicSerialDim<Op>(dview_m, layout_m, nghost_m);
161 : }
162 140 : }
163 :
164 : template <typename T, unsigned Dim, class... ViewArgs>
165 8 : void BareField<T, Dim, ViewArgs...>::accumulateHalo_noghost(int nghost) {
166 8 : if (layout_m->comm.size() > 1) {
167 8 : halo_m.accumulateHalo_noghost(dview_m, layout_m, nghost);
168 : }
169 8 : }
170 :
171 : template <typename T, unsigned Dim, class... ViewArgs>
172 692 : BareField<T, Dim, ViewArgs...>& BareField<T, Dim, ViewArgs...>::operator=(T x) {
173 : using index_array_type = typename RangePolicy<Dim, execution_space>::index_array_type;
174 692 : ippl::parallel_for(
175 692 : "BareField::operator=(T)", getRangePolicy(dview_m),
176 46127374 : KOKKOS_CLASS_LAMBDA(const index_array_type& args) { apply(dview_m, args) = x; });
177 692 : return *this;
178 : }
179 :
180 : template <typename T, unsigned Dim, class... ViewArgs>
181 : template <typename E, size_t N>
182 528 : BareField<T, Dim, ViewArgs...>& BareField<T, Dim, ViewArgs...>::operator=(
183 : const detail::Expression<E, N>& expr) {
184 : using capture_type = detail::CapturedExpression<E, N>;
185 528 : capture_type expr_ = reinterpret_cast<const capture_type&>(expr);
186 : using index_array_type = typename RangePolicy<Dim, execution_space>::index_array_type;
187 528 : ippl::parallel_for(
188 528 : "BareField::operator=(const Expression&)", getRangePolicy(dview_m, nghost_m),
189 10042132 : KOKKOS_CLASS_LAMBDA(const index_array_type& args) {
190 5020538 : apply(dview_m, args) = apply(expr_, args);
191 : });
192 528 : return *this;
193 : }
194 :
195 : template <typename T, unsigned Dim, class... ViewArgs>
196 0 : void BareField<T, Dim, ViewArgs...>::write(std::ostream& out) const {
197 0 : Kokkos::fence();
198 0 : detail::write<T, Dim>(dview_m, out);
199 0 : }
200 :
201 : template <typename T, unsigned Dim, class... ViewArgs>
202 0 : void BareField<T, Dim, ViewArgs...>::write(Inform& inf) const {
203 0 : write(inf.getDestination());
204 0 : }
205 :
206 : #define DefineReduction(fun, name, op, MPI_Op) \
207 : template <typename T, unsigned Dim, class... ViewArgs> \
208 : T BareField<T, Dim, ViewArgs...>::name(int nghost) const { \
209 : PAssert_LE(nghost, nghost_m); \
210 : T temp = Kokkos::reduction_identity<T>::name(); \
211 : using index_array_type = typename RangePolicy<Dim, execution_space>::index_array_type; \
212 : ippl::parallel_reduce( \
213 : "fun", getRangePolicy(dview_m, nghost_m - nghost), \
214 : KOKKOS_CLASS_LAMBDA(const index_array_type& args, T& valL) { \
215 : T myVal = apply(dview_m, args); \
216 : op; \
217 : }, \
218 : KokkosCorrection::fun<T>(temp)); \
219 : T globaltemp = 0.0; \
220 : layout_m->comm.allreduce(temp, globaltemp, 1, MPI_Op<T>()); \
221 : return globaltemp; \
222 : }
223 :
224 1671288 : DefineReduction(Sum, sum, valL += myVal, std::plus)
225 477528 : DefineReduction(Max, max, using Kokkos::max; valL = max(valL, myVal), std::greater)
226 477528 : DefineReduction(Min, min, using Kokkos::min; valL = min(valL, myVal), std::less)
227 238752 : DefineReduction(Prod, prod, valL *= myVal, std::multiplies)
228 :
229 : } // namespace ippl
|
