10 #ifndef STOKHOS_COO_PRODUCT_TENSOR_HPP
11 #define STOKHOS_COO_PRODUCT_TENSOR_HPP
13 #include <type_traits>
15 #include "Kokkos_Core.hpp"
35 template<
typename ValueType,
class ExecutionSpace,
bool PackIndex >
40 template<
typename ValueType,
class ExecutionSpace >
45 typedef typename execution_space::size_type
size_type;
56 typedef Kokkos::View< value_type[], execution_space >
vec_type;
71 size_type ijk = ii | (jj << bits) | (kk << 2*bits);
75 KOKKOS_INLINE_FUNCTION
77 i = ijk & mask; ijk >>= bits;
99 m_coord( rhs.m_coord ),
100 m_value( rhs.m_value ),
102 m_flops( rhs.m_flops ) {}
107 m_coord = rhs.m_coord;
108 m_value = rhs.m_value;
110 m_flops = rhs.m_flops;
115 KOKKOS_INLINE_FUNCTION
119 KOKKOS_INLINE_FUNCTION
123 KOKKOS_INLINE_FUNCTION
126 unpack(m_coord(entry), i, j, k);
130 KOKKOS_INLINE_FUNCTION
132 {
return m_value(entry); }
135 KOKKOS_INLINE_FUNCTION
139 KOKKOS_INLINE_FUNCTION
142 template <
typename OrdinalType>
149 typedef typename Cijk_type::i_iterator i_iterator;
150 typedef typename Cijk_type::ik_iterator k_iterator;
151 typedef typename Cijk_type::ikj_iterator j_iterator;
155 for (i_iterator i_it = Cijk.
i_begin(); i_it!=Cijk.
i_end(); ++i_it) {
156 for (k_iterator k_it = Cijk.
k_begin(i_it); k_it != Cijk.
k_end(i_it);
158 OrdinalType k = index(k_it);
159 for (j_iterator j_it = Cijk.
j_begin(k_it); j_it != Cijk.
j_end(k_it);
161 OrdinalType
j = index(j_it);
170 #if defined( KOKKOS_ENABLE_CUDA )
171 enum { Align = std::is_same<ExecutionSpace,Kokkos::Cuda>::value ? 32 : 1 };
176 entry_count = (entry_count+Align-1) & ~(Align-1);
183 tensor.m_dim = basis.
size();
184 tensor.m_flops = 5*entry_count + tensor.m_dim;
187 typename coord_array_type::HostMirror
189 typename value_array_type::HostMirror
193 OrdinalType i=0,
j=0, k=0;
194 for (i_iterator i_it = Cijk.
i_begin(); i_it!=Cijk.
i_end(); ++i_it) {
196 for (k_iterator k_it = Cijk.
k_begin(i_it); k_it != Cijk.
k_end(i_it);
199 for (j_iterator j_it = Cijk.
j_begin(k_it); j_it != Cijk.
j_end(k_it);
202 ValueType c = Stokhos::value(j_it);
204 host_value(n) = (
j != k) ? c : 0.5*c;
205 host_coord(n) = pack(i,
j,k);
211 for (; n < entry_count; ++n) {
213 host_coord(n) = pack(i,
j,k);
223 void print(std::ostream& os)
const {
225 typename coord_array_type::HostMirror
227 typename value_array_type::HostMirror
232 os <<
"CooProductTensor: dim = "
233 << dimension() <<
", entry_count = "
234 << num_entry << std::endl
235 <<
"Entries: i j k : cijk" << std::endl;
238 unpack(host_coord(l), i, j, k);
239 ValueType
cijk = host_value(l);
240 os <<
"\t " << l <<
" : " << i <<
" " << j <<
" " << k <<
" = " << cijk
242 if (l > 0 && l % 32 == 0) os << std::endl;
249 template<
typename ValueType,
class ExecutionSpace>
259 typedef Kokkos::View< value_type[], execution_space >
vec_type;
282 m_coord( rhs.m_coord ),
283 m_value( rhs.m_value ),
285 m_flops( rhs.m_flops ) {}
290 m_coord = rhs.m_coord;
291 m_value = rhs.m_value;
293 m_flops = rhs.m_flops;
298 KOKKOS_INLINE_FUNCTION
302 KOKKOS_INLINE_FUNCTION
306 KOKKOS_INLINE_FUNCTION
309 i = m_coord(entry,0);
310 j = m_coord(entry,1);
311 k = m_coord(entry,2);
315 KOKKOS_INLINE_FUNCTION
317 {
return m_value( entry ); }
320 KOKKOS_INLINE_FUNCTION
324 KOKKOS_INLINE_FUNCTION
327 template <
typename OrdinalType>
334 typedef typename Cijk_type::i_iterator i_iterator;
335 typedef typename Cijk_type::ik_iterator k_iterator;
336 typedef typename Cijk_type::ikj_iterator j_iterator;
340 for (i_iterator i_it = Cijk.
i_begin(); i_it!=Cijk.
i_end(); ++i_it) {
341 for (k_iterator k_it = Cijk.
k_begin(i_it); k_it != Cijk.
k_end(i_it);
343 OrdinalType k = index(k_it);
344 for (j_iterator j_it = Cijk.
j_begin(k_it); j_it != Cijk.
j_end(k_it);
346 OrdinalType
j = index(j_it);
355 #if defined( KOKKOS_ENABLE_CUDA )
356 enum { Align = std::is_same<ExecutionSpace,Kokkos::Cuda>::value ? 32 : 1 };
361 entry_count = (entry_count+Align-1) & ~(Align-1);
368 tensor.m_dim = basis.
size();
369 tensor.m_flops = 5*entry_count + tensor.m_dim;
372 typename coord_array_type::HostMirror
374 typename value_array_type::HostMirror
379 OrdinalType i=0,
j=0, k=0;
380 for (i_iterator i_it = Cijk.
i_begin(); i_it!=Cijk.
i_end(); ++i_it) {
382 for (k_iterator k_it = Cijk.
k_begin(i_it); k_it != Cijk.
k_end(i_it);
385 for (j_iterator j_it = Cijk.
j_begin(k_it); j_it != Cijk.
j_end(k_it);
388 ValueType c = Stokhos::value(j_it);
390 host_value(n) = (
j != k) ? c : 0.5*c;
399 for (; n < entry_count; ++n) {
413 void print(std::ostream& os)
const {
415 typename coord_array_type::HostMirror
417 typename value_array_type::HostMirror
422 os <<
"CooProductTensor: dim = "
423 << dimension() <<
", entry_count = "
424 << num_entry << std::endl
425 <<
"Entries: i j k : cijk" << std::endl;
430 ValueType
cijk = host_value(l);
431 os <<
"\t " << l <<
" : " << i <<
" " << j <<
" " << k <<
" = " << cijk
433 if (l > 0 && l % 32 == 0) os << std::endl;
438 template<
class Device,
bool Pack,
typename OrdinalType,
typename ValueType >
439 CooProductTensor<ValueType, Device, Pack>
446 basis, Cijk, params );
449 template <
typename ValueType,
typename Device,
bool Pack>
457 template<
typename ValueType,
typename Device,
bool Pack >
465 template<
typename MatrixValue ,
typename VectorValue >
466 KOKKOS_INLINE_FUNCTION
468 const MatrixValue *
const a,
469 const VectorValue *
const x,
470 VectorValue *
const y )
472 const size_type nEntry = tensor.entry_count();
474 VectorValue
val = 0.0, carry_val = 0.0;
475 for (
size_type entry = 0 ; entry < nEntry ; ++entry ) {
476 tensor.coord(entry, i,
j, k);
477 val = tensor.value(entry) * ( a[
j] * x[k] + a[k] * x[
j] );
481 y[i_prev] += carry_val;
489 KOKKOS_INLINE_FUNCTION
491 {
return tensor.dimension(); }
493 KOKKOS_INLINE_FUNCTION
495 {
return tensor.dimension(); }
KOKKOS_INLINE_FUNCTION void coord(const size_type entry, size_type &i, size_type &j, size_type &k) const
Get (i,j,k) coordinates of an entry.
execution_space::size_type size_type
Kokkos::View< value_type[], execution_space > value_array_type
k_iterator k_begin() const
Iterator pointing to first k entry.
KOKKOS_INLINE_FUNCTION size_type dimension() const
Dimension of the tensor.
Data structure storing a sparse 3-tensor C(i,j,k) in a a compressed format.
static KOKKOS_INLINE_FUNCTION size_type matrix_size(const tensor_type &tensor)
KOKKOS_INLINE_FUNCTION size_type dimension() const
Dimension of the tensor.
Kokkos::View< size_type[][3], execution_space > coord_array_type
CooProductTensor(const CooProductTensor &rhs)
kj_iterator j_begin(const k_iterator &k) const
Iterator pointing to first j entry for given k.
KOKKOS_INLINE_FUNCTION size_type entry_count() const
Number of sparse entries.
static CooProductTensor create(const Stokhos::ProductBasis< OrdinalType, ValueType > &basis, const Stokhos::Sparse3Tensor< OrdinalType, ValueType > &Cijk, const Teuchos::ParameterList ¶ms=Teuchos::ParameterList())
kj_iterator j_end(const k_iterator &k) const
Iterator pointing to last j entry for given k.
CooProductTensor(const CooProductTensor &rhs)
KOKKOS_INLINE_FUNCTION size_type num_non_zeros() const
Number of non-zero's.
KOKKOS_INLINE_FUNCTION size_type entry_count() const
Number of sparse entries.
std::ostream & operator<<(std::ostream &os, const ProductContainer< coeff_type > &vec)
Kokkos::View< value_type[], execution_space > vec_type
Kokkos::View< size_type[], execution_space > coord_array_type
static size_type pack(const size_type i, const size_type j, const size_type k)
void deep_copy(const Stokhos::CrsMatrix< ValueType, DstDevice, Layout > &dst, const Stokhos::CrsMatrix< ValueType, SrcDevice, Layout > &src)
KOKKOS_INLINE_FUNCTION size_type num_non_zeros() const
Number of non-zero's.
void print(std::ostream &os) const
k_iterator k_end() const
Iterator pointing to last k entry.
KOKKOS_INLINE_FUNCTION constexpr std::enable_if< is_view_uq_pce< view_type >::value, typename CijkType< view_type >::type >::type cijk(const view_type &view)
Stokhos::Sparse3Tensor< int, double > Cijk_type
execution_space::size_type size_type
void print(std::ostream &os) const
static CooProductTensor create(const Stokhos::ProductBasis< OrdinalType, ValueType > &basis, const Stokhos::Sparse3Tensor< OrdinalType, ValueType > &Cijk, const Teuchos::ParameterList ¶ms=Teuchos::ParameterList())
Kokkos::View< value_type[], execution_space > vec_type
KOKKOS_INLINE_FUNCTION void coord(const size_type entry, size_type &i, size_type &j, size_type &k) const
Get (i,j,k) coordinates of an entry.
KOKKOS_INLINE_FUNCTION size_type num_flops() const
Number flop's per multiply-add.
ExecutionSpace execution_space
KOKKOS_INLINE_FUNCTION size_type num_flops() const
Number flop's per multiply-add.
KOKKOS_INLINE_FUNCTION const value_type & value(const size_type entry) const
Value of an entry.
#define TEUCHOS_ASSERT(assertion_test)
static KOKKOS_INLINE_FUNCTION size_type vector_size(const tensor_type &tensor)
KOKKOS_INLINE_FUNCTION void unpack(size_type ijk, size_type &i, size_type &j, size_type &k) const
kji_iterator i_begin(const kj_iterator &j) const
Iterator pointing to first i entry for given j and k.
Sparse product tensor using 'COO'-like storage format.
Kokkos::View< value_type[], execution_space > value_array_type
kji_iterator i_end(const kj_iterator &j) const
Iterator pointing to last i entry for given j and k.
CooProductTensor< ValueType, Device, Pack > create_coo_product_tensor(const Stokhos::ProductBasis< OrdinalType, ValueType > &basis, const Stokhos::Sparse3Tensor< OrdinalType, ValueType > &Cijk, const Teuchos::ParameterList ¶ms=Teuchos::ParameterList())
CooProductTensor< ValueType, Device, Pack > tensor_type
virtual ordinal_type size() const =0
Return total size of basis.
KOKKOS_INLINE_FUNCTION const value_type & value(const size_type entry) const
Value of an entry.
Device::size_type size_type
Stokhos::CrsMatrix< ValueType, Device, Layout >::HostMirror create_mirror_view(const Stokhos::CrsMatrix< ValueType, Device, Layout > &A)
static KOKKOS_INLINE_FUNCTION void apply(const tensor_type &tensor, const MatrixValue *const a, const VectorValue *const x, VectorValue *const y)
ExecutionSpace execution_space