Stokhos Package Browser (Single Doxygen Collection)  Version of the Day
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
Kokkos_Blas1_MP_Vector.hpp
Go to the documentation of this file.
1 // @HEADER
2 // *****************************************************************************
3 // Stokhos Package
4 //
5 // Copyright 2009 NTESS and the Stokhos contributors.
6 // SPDX-License-Identifier: BSD-3-Clause
7 // *****************************************************************************
8 // @HEADER
9 
10 #ifndef KOKKOS_BLAS1_MP_VECTOR_HPP
11 #define KOKKOS_BLAS1_MP_VECTOR_HPP
12 
13 #include "Sacado_ConfigDefs.h"
14 
15 #include "Sacado_MP_Vector.hpp"
18 #include "KokkosBlas.hpp"
19 
20 //----------------------------------------------------------------------------
21 // Specializations of Kokkos Vector/MultiVector math functions
22 //----------------------------------------------------------------------------
23 
24 namespace KokkosBlas {
25 
26 #if defined(HAVE_STOKHOS_ENSEMBLE_REDUCT)
27 
28 template <typename XD, typename ... XP,
29  typename YD, typename ... YP>
30 typename std::enable_if<
31  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value &&
32  Kokkos::is_view_mp_vector< Kokkos::View<YD,YP...> >::value,
33  typename Kokkos::Details::InnerProductSpaceTraits<
34  typename Kokkos::View<XD,XP...>::non_const_value_type >::dot_type
35  >::type
36 dot(const Kokkos::View<XD,XP...>& x,
37  const Kokkos::View<YD,YP...>& y)
38 {
39  typedef Kokkos::View<XD,XP...> XVector;
40  typedef Kokkos::View<YD,YP...> YVector;
41 
42  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
43  typename Kokkos::FlatArrayType<YVector>::type y_flat = y;
44 
45  return dot( x_flat, y_flat );
46 }
47 
48 template <typename RV,
49  typename XD, typename ... XP,
50  typename YD, typename ... YP>
51 typename std::enable_if<
52  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value &&
53  Kokkos::is_view_mp_vector< Kokkos::View<YD,YP...> >::value >::type
54 dot(const RV& r,
55  const Kokkos::View<XD,XP...>& x,
56  const Kokkos::View<YD,YP...>& y)
57 {
58  typedef Kokkos::View<XD,XP...> XVector;
59  typedef Kokkos::View<YD,YP...> YVector;
60 
61  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
62  typename Kokkos::FlatArrayType<YVector>::type y_flat = y;
63 
64  dot( r, x_flat, y_flat );
65 }
66 
67 template <typename XD, typename ... XP>
68 typename std::enable_if<
69  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value >::type
70 fill(const Kokkos::View<XD,XP...>& x,
71  const typename Kokkos::View<XD,XP...>::non_const_value_type& val) {
72  typedef Kokkos::View<XD,XP...> XVector;
73 
74  // Use the existing fill() implementation if we can
75  if (Sacado::is_constant(val)) {
76  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
77  fill( x_flat, val.coeff(0) );
78  }
79  else {
80  Kokkos::deep_copy(x, val);
81  }
82 }
83 
84 template <typename RV,
85  typename XD, typename ... XP>
86 typename std::enable_if<
87  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value >::type
89  const RV& r,
90  const Kokkos::View<XD,XP...>& x)
91 {
92  typedef Kokkos::View<XD,XP...> XVector;
93 
94  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
95 
96  nrm2_squared( r, x_flat );
97 }
98 
99 template <typename RV,
100  typename XD, typename ... XP>
101 typename std::enable_if<
102  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value >::type
103 nrm1(
104  const RV& r,
105  const Kokkos::View<XD,XP...>& x)
106 {
107  typedef Kokkos::View<XD,XP...> XVector;
108 
109  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
110 
111  nrm1( r, x_flat );
112 }
113 
114 template <typename RV,
115  typename XD, typename ... XP>
116 typename std::enable_if<
117  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value >::type
118 nrmInf(
119  const RV& r,
120  const Kokkos::View<XD,XP...>& x)
121 {
122  typedef Kokkos::View<XD,XP...> XVector;
123 
124  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
125 
126  nrmInf( r, x_flat );
127 }
128 
129 template <typename AV,
130  typename XD, typename ... XP,
131  typename BV,
132  typename YD, typename ... YP>
133 typename std::enable_if<
134  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value &&
135  Kokkos::is_view_mp_vector< Kokkos::View<YD,YP...> >::value >::type
136 axpby(const AV& a,
137  const Kokkos::View<XD,XP...>& x,
138  const BV& b,
139  const Kokkos::View<YD,YP...>& y)
140 {
141  typedef Kokkos::View<XD,XP...> XVector;
142  typedef Kokkos::View<YD,YP...> YVector;
143 
144  if (!Sacado::is_constant(a) || !Sacado::is_constant(b)) {
145  Kokkos::Impl::raise_error("axpby not implemented for non-constant a or b");
146  }
147 
148  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
149  typename Kokkos::FlatArrayType<YVector>::type y_flat = y;
150  auto aa = Sacado::Value<AV>::eval(a);
151  auto bb = Sacado::Value<BV>::eval(b);
152  axpby( aa, x_flat, bb, y_flat );
153 }
154 
155 // Currently not handling scal() when AV is a view
156 
157 template <typename RD, typename ... RP,
158  typename XD, typename ... XP>
159 typename std::enable_if<
160  Kokkos::is_view_mp_vector< Kokkos::View<RD,RP...> >::value &&
161  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value >::type
162 scal(const Kokkos::View<RD,RP...>& r,
163  const typename Kokkos::View<XD,XP...>::non_const_value_type& a,
164  const Kokkos::View<XD,XP...>& x)
165 {
166  typedef Kokkos::View<RD,RP...> RVector;
167  typedef Kokkos::View<XD,XP...> XVector;
168 
169  if (!Sacado::is_constant(a)) {
170  Kokkos::Impl::raise_error("scal not implemented for non-constant a");
171  }
172 
173  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
174  typename Kokkos::FlatArrayType<RVector>::type r_flat = r;
175  scal( r_flat, a.coeff(0), x_flat );
176 }
177 
178 // abs -- can't do this one by flattening. Hold out for refactoring of scalar
179 // types in Kokkos
180 
181 // We have a special verision of update for scalar alpha/beta/gamma since it
182 // is used in TrilinosCouplings CG solve (even though Tpetra doesn't).
183 template <typename XD, typename ... XP,
184  typename YD, typename ... YP,
185  typename ZD, typename ... ZP>
186 typename std::enable_if<
187  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value &&
188  Kokkos::is_view_mp_vector< Kokkos::View<YD,YP...> >::value &&
189  Kokkos::is_view_mp_vector< Kokkos::View<ZD,ZP...> >::value >::type
190 update(
191  const typename Kokkos::View<XD,XP...>::array_type::non_const_value_type& alpha,
192  const Kokkos::View<XD,XP...>& x,
193  const typename Kokkos::View<YD,YP...>::array_type::non_const_value_type& beta,
194  const Kokkos::View<YD,YP...>& y,
195  const typename Kokkos::View<ZD,ZP...>::array_type::non_const_value_type& gamma,
196  const Kokkos::View<ZD,ZP...>& z)
197 {
198  typedef Kokkos::View<XD,XP...> XVector;
199  typedef Kokkos::View<YD,YP...> YVector;
200  typedef Kokkos::View<ZD,ZP...> ZVector;
201 
202  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
203  typename Kokkos::FlatArrayType<YVector>::type y_flat = y;
204  typename Kokkos::FlatArrayType<ZVector>::type z_flat = z;
205 
206  update( alpha, x_flat, beta, y_flat, gamma, z_flat);
207 
208 }
209 
210 template <typename XD, typename ... XP,
211  typename YD, typename ... YP,
212  typename ZD, typename ... ZP>
213 typename std::enable_if<
214  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value &&
215  Kokkos::is_view_mp_vector< Kokkos::View<YD,YP...> >::value &&
216  Kokkos::is_view_mp_vector< Kokkos::View<ZD,ZP...> >::value >::type
217 update(
218  const typename Kokkos::View<XD,XP...>::non_const_value_type& alpha,
219  const Kokkos::View<XD,XP...>& x,
220  const typename Kokkos::View<YD,YP...>::non_const_value_type& beta,
221  const Kokkos::View<YD,YP...>& y,
222  const typename Kokkos::View<ZD,ZP...>::non_const_value_type& gamma,
223  const Kokkos::View<ZD,ZP...>& z)
224 {
225  if (!Sacado::is_constant(alpha) || !Sacado::is_constant(beta) ||
226  !Sacado::is_constant(gamma)) {
228  "update not implemented for non-constant alpha, beta, gamma");
229  }
230 
231  update( alpha.coeff(0), x, beta.coeff(0), y, gamma.coeff(0), z );
232 }
233 
234 template <typename RD, typename ... RP,
235  typename XD, typename ... XP>
236 typename std::enable_if<
237  Kokkos::is_view_mp_vector< Kokkos::View<RD,RP...> >::value &&
238  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value >::type
239 reciprocal(
240  const Kokkos::View<RD,RP...>& r,
241  const Kokkos::View<XD,XP...>& x)
242 {
243  typedef Kokkos::View<RD,RP...> RVector;
244  typedef Kokkos::View<XD,XP...> XVector;
245 
246  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
247  typename Kokkos::FlatArrayType<RVector>::type r_flat = r;
248  reciprocal( r_flat, x_flat );
249 }
250 
251 template <typename RD, typename ... RP,
252  typename XD, typename ... XP>
253 typename std::enable_if<
254  Kokkos::is_view_mp_vector< Kokkos::View<RD,RP...> >::value &&
255  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value >::type
256 sum(
257  const Kokkos::View<RD,RP...>& r,
258  const Kokkos::View<XD,XP...>& x)
259 {
260  typedef Kokkos::View<RD,RP...> RVector;
261  typedef Kokkos::View<XD,XP...> XVector;
262 
263  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
264  typename Kokkos::FlatArrayType<RVector>::type r_flat = r;
265  sum( r_flat, x_flat );
266 }
267 
268 template <typename RD, typename ... RP,
269  typename XD, typename ... XP,
270  typename WD, typename ... WP>
271 typename std::enable_if<
272  Kokkos::is_view_mp_vector< Kokkos::View<RD,RP...> >::value &&
273  Kokkos::is_view_mp_vector< Kokkos::View<XD,XP...> >::value &&
274  Kokkos::is_view_mp_vector< Kokkos::View<WD,WP...> >::value >::type
276  const Kokkos::View<RD,RP...>& r,
277  const Kokkos::View<XD,XP...>& x,
278  const Kokkos::View<WD,WP...>& w)
279 {
280  typedef Kokkos::View<RD,RP...> RVector;
281  typedef Kokkos::View<XD,XP...> XVector;
282  typedef Kokkos::View<WD,WP...> WVector;
283 
284  typename Kokkos::FlatArrayType<XVector>::type x_flat = x;
285  typename Kokkos::FlatArrayType<RVector>::type r_flat = r;
286  typename Kokkos::FlatArrayType<WVector>::type w_flat = w;
287  nrm2w_squared( r_flat, x_flat, w_flat );
288 }
289 
290 template <typename CD, typename ... CP,
291  typename AD, typename ... AP,
292  typename BD, typename ... BP>
293 typename std::enable_if<
294  Kokkos::is_view_mp_vector< Kokkos::View<CD,CP...> >::value &&
295  Kokkos::is_view_mp_vector< Kokkos::View<AD,AP...> >::value &&
296  Kokkos::is_view_mp_vector< Kokkos::View<BD,BP...> >::value >::type
297 mult(
298  const typename Kokkos::View<CD,CP...>::const_value_type& c,
299  const Kokkos::View<CD,CP...>& C,
300  const typename Kokkos::View<AD,AP...>::const_value_type& ab,
301  const Kokkos::View<AD,AP...>& A,
302  const Kokkos::View<BD,BP...>& B)
303 {
304  if (!Sacado::is_constant(c) || !Sacado::is_constant(ab)) {
305  Kokkos::Impl::raise_error("mult not implemented for non-constant c, ab");
306  }
307 
308  typedef Kokkos::View<CD,CP...> CVector;
309  typedef Kokkos::View<AD,AP...> AVector;
310  typedef Kokkos::View<BD,BP...> BVector;
311 
312  typename Kokkos::FlatArrayType<CVector>::type C_flat = C;
313  typename Kokkos::FlatArrayType<AVector>::type A_flat = A;
314  typename Kokkos::FlatArrayType<BVector>::type B_flat = B;
315  mult( c.coeff(0), C_flat, ab.coeff(0), A_flat, B_flat );
316 }
317 
318 #endif
319 
320 } // namespace KokkosBlas
321 
322 #endif /* #ifndef KOKKOS_MV_MP_VECTOR_HPP */
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value >::type fill(const Kokkos::View< XD, XP...> &x, const typename Kokkos::View< XD, XP...>::non_const_value_type &val)
KOKKOS_INLINE_FUNCTION void raise_error(const char *msg)
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value &&Kokkos::is_view_uq_pce< Kokkos::View< YD, YP...> >::value, typename Kokkos::Details::InnerProductSpaceTraits< typename Kokkos::View< XD, XP...>::non_const_value_type >::dot_type >::type dot(const Kokkos::View< XD, XP...> &x, const Kokkos::View< YD, YP...> &y)
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value >::type nrm1(const RV &r, const Kokkos::View< XD, XP...> &x)
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< RD, RP...> >::value &&Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value &&Kokkos::is_view_uq_pce< Kokkos::View< WD, WP...> >::value >::type nrm2w_squared(const Kokkos::View< RD, RP...> &r, const Kokkos::View< XD, XP...> &x, const Kokkos::View< WD, WP...> &w)
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value >::type nrm2_squared(const RV &r, const Kokkos::View< XD, XP...> &x)
KOKKOS_INLINE_FUNCTION bool is_constant(const T &x)
void deep_copy(const Stokhos::CrsMatrix< ValueType, DstDevice, Layout > &dst, const Stokhos::CrsMatrix< ValueType, SrcDevice, Layout > &src)
expr val()
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value &&Kokkos::is_view_uq_pce< Kokkos::View< YD, YP...> >::value &&Kokkos::is_view_uq_pce< Kokkos::View< ZD, ZP...> >::value >::type update(const typename Kokkos::View< XD, XP...>::array_type::non_const_value_type &alpha, const Kokkos::View< XD, XP...> &x, const typename Kokkos::View< YD, YP...>::array_type::non_const_value_type &beta, const Kokkos::View< YD, YP...> &y, const typename Kokkos::View< ZD, ZP...>::array_type::non_const_value_type &gamma, const Kokkos::View< ZD, ZP...> &z)
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value >::type nrmInf(const RV &r, const Kokkos::View< XD, XP...> &x)
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value &&Kokkos::is_view_uq_pce< Kokkos::View< YD, YP...> >::value >::type axpby(const AV &a, const Kokkos::View< XD, XP...> &x, const BV &b, const Kokkos::View< YD, YP...> &y)
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< RD, RP...> >::value &&Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value >::type scal(const Kokkos::View< RD, RP...> &r, const typename Kokkos::View< XD, XP...>::non_const_value_type &a, const Kokkos::View< XD, XP...> &x)
std::enable_if< Kokkos::is_view_uq_pce< Kokkos::View< RD, RP...> >::value &&Kokkos::is_view_uq_pce< Kokkos::View< XD, XP...> >::value >::type sum(const Kokkos::View< RD, RP...> &r, const Kokkos::View< XD, XP...> &x)