Stokhos Package Browser (Single Doxygen Collection)  Version of the Day
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
Stokhos_CooProductTensor.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 STOKHOS_COO_PRODUCT_TENSOR_HPP
11 #define STOKHOS_COO_PRODUCT_TENSOR_HPP
12 
13 #include <type_traits>
14 
15 #include "Kokkos_Core.hpp"
16 
17 #include "Stokhos_Multiply.hpp"
18 #include "Stokhos_ProductBasis.hpp"
21 
22 //----------------------------------------------------------------------------
23 //----------------------------------------------------------------------------
24 
25 namespace Stokhos {
26 
35 template< typename ValueType, class ExecutionSpace, bool PackIndex >
37 
40 template< typename ValueType, class ExecutionSpace >
41 class CooProductTensor<ValueType,ExecutionSpace,true> {
42 public:
43 
44  typedef ExecutionSpace execution_space;
45  typedef typename execution_space::size_type size_type;
46  typedef ValueType value_type;
47 
48 private:
49 
50  // Number of bits available for each index
51  static const size_type bits = (sizeof(size_type)*8) / 3;
52 
53  // Mask for packing index
54  static const size_type mask = (1 << bits)-1;
55 
56  typedef Kokkos::View< value_type[], execution_space > vec_type;
57  typedef Kokkos::View< size_type[], execution_space > coord_array_type;
58  typedef Kokkos::View< value_type[], execution_space > value_array_type;
59 
64 
65  // Pack (i,j,k) into a single integer
66  static size_type
67  pack( const size_type i, const size_type j, const size_type k ) {
68  const size_type ii = i & mask;
69  const size_type jj = j & mask;
70  const size_type kk = k & mask;
71  size_type ijk = ii | (jj << bits) | (kk << 2*bits);
72  return ijk;
73  }
74 
75  KOKKOS_INLINE_FUNCTION
76  void unpack( size_type ijk, size_type& i, size_type& j, size_type& k ) const {
77  i = ijk & mask; ijk >>= bits;
78  j = ijk & mask;
79  k = ijk >> bits;
80  }
81 
82 public:
83 
85  static const size_type max_index = 1 << bits;
86 
87  inline
89 
90  inline
92  m_coord(),
93  m_value(),
94  m_dim(0),
95  m_flops(0) {}
96 
97  inline
99  m_coord( rhs.m_coord ),
100  m_value( rhs.m_value ),
101  m_dim( rhs.m_dim ),
102  m_flops( rhs.m_flops ) {}
103 
104  inline
105  CooProductTensor & operator = ( const CooProductTensor & rhs )
106  {
107  m_coord = rhs.m_coord;
108  m_value = rhs.m_value;
109  m_dim = rhs.m_dim;
110  m_flops = rhs.m_flops;
111  return *this;
112  }
113 
115  KOKKOS_INLINE_FUNCTION
116  size_type dimension() const { return m_dim; }
117 
119  KOKKOS_INLINE_FUNCTION
120  size_type entry_count() const { return m_coord.extent(0); }
121 
123  KOKKOS_INLINE_FUNCTION
124  void coord( const size_type entry,
125  size_type& i, size_type& j, size_type& k ) const {
126  unpack(m_coord(entry), i, j, k);
127  }
128 
130  KOKKOS_INLINE_FUNCTION
131  const value_type & value( const size_type entry ) const
132  { return m_value(entry); }
133 
135  KOKKOS_INLINE_FUNCTION
136  size_type num_non_zeros() const { return m_coord.extent(0); }
137 
139  KOKKOS_INLINE_FUNCTION
140  size_type num_flops() const { return m_flops; }
141 
142  template <typename OrdinalType>
143  static CooProductTensor
147  {
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;
152 
153  // Compute entry count
154  size_type entry_count = 0;
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);
157  ++k_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);
160  ++j_it) {
161  OrdinalType j = index(j_it);
162  if (j >= k) {
163  ++entry_count;
164  }
165  }
166  }
167  }
168 
169  // Align entry_count
170 #if defined( KOKKOS_ENABLE_CUDA )
171  enum { Align = std::is_same<ExecutionSpace,Kokkos::Cuda>::value ? 32 : 1 };
172 #else
173  enum { Align = 1 };
174 #endif
175 
176  entry_count = (entry_count+Align-1) & ~(Align-1);
177  TEUCHOS_ASSERT(entry_count % Align == 0);
178 
179  // Allocate tensor data
180  CooProductTensor tensor;
181  tensor.m_coord = coord_array_type( "tensor_coord", entry_count );
182  tensor.m_value = value_array_type( "tensor_value", entry_count );
183  tensor.m_dim = basis.size();
184  tensor.m_flops = 5*entry_count + tensor.m_dim;
185 
186  // Create mirror, is a view if is host memory
187  typename coord_array_type::HostMirror
188  host_coord = Kokkos::create_mirror_view( tensor.m_coord );
189  typename value_array_type::HostMirror
190  host_value = Kokkos::create_mirror_view( tensor.m_value );
191 
192  size_type n = 0;
193  OrdinalType i=0, j=0, k=0;
194  for (i_iterator i_it = Cijk.i_begin(); i_it!=Cijk.i_end(); ++i_it) {
195  i = index(i_it);
196  for (k_iterator k_it = Cijk.k_begin(i_it); k_it != Cijk.k_end(i_it);
197  ++k_it) {
198  k = index(k_it);
199  for (j_iterator j_it = Cijk.j_begin(k_it); j_it != Cijk.j_end(k_it);
200  ++j_it) {
201  j = index(j_it);
202  ValueType c = Stokhos::value(j_it);
203  if (j >= k) {
204  host_value(n) = (j != k) ? c : 0.5*c;
205  host_coord(n) = pack(i,j,k);
206  ++n;
207  }
208  }
209  }
210  }
211  for (; n < entry_count; ++n) {
212  host_value(n) = 0.0;
213  host_coord(n) = pack(i,j,k);
214  }
215 
216  // Copy data to device if necessary
217  Kokkos::deep_copy( tensor.m_coord, host_coord );
218  Kokkos::deep_copy( tensor.m_value, host_value );
219 
220  return tensor;
221  }
222 
223  void print(std::ostream& os) const {
224  size_type num_entry = entry_count();
225  typename coord_array_type::HostMirror
226  host_coord = Kokkos::create_mirror_view( m_coord );
227  typename value_array_type::HostMirror
228  host_value = Kokkos::create_mirror_view( m_value );
229  Kokkos::deep_copy( host_coord, m_coord );
230  Kokkos::deep_copy( host_value, m_value );
231 
232  os << "CooProductTensor: dim = "
233  << dimension() << ", entry_count = "
234  << num_entry << std::endl
235  << "Entries: i j k : cijk" << std::endl;
236  for (size_type l=0; l<num_entry; ++l) {
237  size_type i,j,k;
238  unpack(host_coord(l), i, j, k);
239  ValueType cijk = host_value(l);
240  os << "\t " << l << " : " << i << " " << j << " " << k << " = " << cijk
241  << std::endl;
242  if (l > 0 && l % 32 == 0) os << std::endl;
243  }
244  }
245 };
246 
249 template< typename ValueType, class ExecutionSpace>
250 class CooProductTensor<ValueType,ExecutionSpace,false> {
251 public:
252 
253  typedef ExecutionSpace execution_space;
254  typedef typename execution_space::size_type size_type;
255  typedef ValueType value_type;
256 
257 private:
258 
259  typedef Kokkos::View< value_type[], execution_space > vec_type;
260  typedef Kokkos::View< size_type[][3], execution_space > coord_array_type;
261  typedef Kokkos::View< value_type[], execution_space > value_array_type;
262 
267 
268 public:
269 
270  inline
272 
273  inline
275  m_coord(),
276  m_value(),
277  m_dim(0),
278  m_flops(0) {}
279 
280  inline
282  m_coord( rhs.m_coord ),
283  m_value( rhs.m_value ),
284  m_dim( rhs.m_dim ),
285  m_flops( rhs.m_flops ) {}
286 
287  inline
288  CooProductTensor & operator = ( const CooProductTensor & rhs )
289  {
290  m_coord = rhs.m_coord;
291  m_value = rhs.m_value;
292  m_dim = rhs.m_dim;
293  m_flops = rhs.m_flops;
294  return *this;
295  }
296 
298  KOKKOS_INLINE_FUNCTION
299  size_type dimension() const { return m_dim; }
300 
302  KOKKOS_INLINE_FUNCTION
303  size_type entry_count() const { return m_coord.extent(0); }
304 
306  KOKKOS_INLINE_FUNCTION
307  void coord( const size_type entry,
308  size_type& i, size_type& j, size_type& k ) const {
309  i = m_coord(entry,0);
310  j = m_coord(entry,1);
311  k = m_coord(entry,2);
312  }
313 
315  KOKKOS_INLINE_FUNCTION
316  const value_type & value( const size_type entry ) const
317  { return m_value( entry ); }
318 
320  KOKKOS_INLINE_FUNCTION
321  size_type num_non_zeros() const { return m_coord.extent(0); }
322 
324  KOKKOS_INLINE_FUNCTION
325  size_type num_flops() const { return m_flops; }
326 
327  template <typename OrdinalType>
328  static CooProductTensor
332  {
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;
337 
338  // Compute entry count
339  size_type entry_count = 0;
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);
342  ++k_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);
345  ++j_it) {
346  OrdinalType j = index(j_it);
347  if (j >= k) {
348  ++entry_count;
349  }
350  }
351  }
352  }
353 
354  // Align entry_count
355 #if defined( KOKKOS_ENABLE_CUDA )
356  enum { Align = std::is_same<ExecutionSpace,Kokkos::Cuda>::value ? 32 : 1 };
357 #else
358  enum { Align = 1 };
359 #endif
360 
361  entry_count = (entry_count+Align-1) & ~(Align-1);
362  TEUCHOS_ASSERT(entry_count % Align == 0);
363 
364  // Allocate tensor data
365  CooProductTensor tensor;
366  tensor.m_coord = coord_array_type( "tensor_coord", entry_count );
367  tensor.m_value = value_array_type( "tensor_value", entry_count );
368  tensor.m_dim = basis.size();
369  tensor.m_flops = 5*entry_count + tensor.m_dim;
370 
371  // Create mirror, is a view if is host memory
372  typename coord_array_type::HostMirror
373  host_coord = Kokkos::create_mirror_view( tensor.m_coord );
374  typename value_array_type::HostMirror
375  host_value = Kokkos::create_mirror_view( tensor.m_value );
376 
377  // Set entries
378  size_type n = 0;
379  OrdinalType i=0, j=0, k=0;
380  for (i_iterator i_it = Cijk.i_begin(); i_it!=Cijk.i_end(); ++i_it) {
381  i = index(i_it);
382  for (k_iterator k_it = Cijk.k_begin(i_it); k_it != Cijk.k_end(i_it);
383  ++k_it) {
384  k = index(k_it);
385  for (j_iterator j_it = Cijk.j_begin(k_it); j_it != Cijk.j_end(k_it);
386  ++j_it) {
387  j = index(j_it);
388  ValueType c = Stokhos::value(j_it);
389  if (j >= k) {
390  host_value(n) = (j != k) ? c : 0.5*c;
391  host_coord(n,0) = i;
392  host_coord(n,1) = j;
393  host_coord(n,2) = k;
394  ++n;
395  }
396  }
397  }
398  }
399  for (; n < entry_count; ++n) {
400  host_value(n) = 0.0;
401  host_coord(n,0) = i;
402  host_coord(n,1) = j;
403  host_coord(n,2) = k;
404  }
405 
406  // Copy data to device if necessary
407  Kokkos::deep_copy( tensor.m_coord, host_coord );
408  Kokkos::deep_copy( tensor.m_value, host_value );
409 
410  return tensor;
411  }
412 
413  void print(std::ostream& os) const {
414  size_type num_entry = entry_count();
415  typename coord_array_type::HostMirror
416  host_coord = Kokkos::create_mirror_view( m_coord );
417  typename value_array_type::HostMirror
418  host_value = Kokkos::create_mirror_view( m_value );
419  Kokkos::deep_copy( host_coord, m_coord );
420  Kokkos::deep_copy( host_value, m_value );
421 
422  os << "CooProductTensor: dim = "
423  << dimension() << ", entry_count = "
424  << num_entry << std::endl
425  << "Entries: i j k : cijk" << std::endl;
426  for (size_type l=0; l<num_entry; ++l) {
427  size_type i = host_coord(l,0);
428  size_type j = host_coord(l,1);
429  size_type k = host_coord(l,2);
430  ValueType cijk = host_value(l);
431  os << "\t " << l << " : " << i << " " << j << " " << k << " = " << cijk
432  << std::endl;
433  if (l > 0 && l % 32 == 0) os << std::endl;
434  }
435  }
436 };
437 
438 template< class Device, bool Pack, typename OrdinalType, typename ValueType >
439 CooProductTensor<ValueType, Device, Pack>
444 {
446  basis, Cijk, params );
447 }
448 
449 template <typename ValueType, typename Device, bool Pack>
450 std::ostream& operator << (
451  std::ostream& os, const CooProductTensor<ValueType,Device,Pack>& tensor)
452 {
453  tensor.print(os);
454  return os;
455 }
456 
457 template< typename ValueType, typename Device, bool Pack >
458 class BlockMultiply< CooProductTensor< ValueType, Device, Pack > >
459 {
460 public:
461 
462  typedef typename Device::size_type size_type;
464 
465  template< typename MatrixValue , typename VectorValue >
466  KOKKOS_INLINE_FUNCTION
467  static void apply( const tensor_type & tensor,
468  const MatrixValue * const a,
469  const VectorValue * const x,
470  VectorValue * const y )
471  {
472  const size_type nEntry = tensor.entry_count();
473  size_type i = 0, j = 0, k = 0, i_prev = -1;
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] );
478  if (i == i_prev)
479  carry_val += val;
480  else {
481  y[i_prev] += carry_val;
482  carry_val = val;
483  }
484  i_prev = i;
485  }
486  y[i] += carry_val;
487  }
488 
489  KOKKOS_INLINE_FUNCTION
490  static size_type matrix_size( const tensor_type & tensor )
491  { return tensor.dimension(); }
492 
493  KOKKOS_INLINE_FUNCTION
494  static size_type vector_size( const tensor_type & tensor )
495  { return tensor.dimension(); }
496 };
497 
498 } /* namespace Stokhos */
499 
500 //----------------------------------------------------------------------------
501 //----------------------------------------------------------------------------
502 
503 #endif /* #ifndef STOKHOS_COO_PRODUCT_TENSOR_HPP */
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::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
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 &params=Teuchos::ParameterList())
kj_iterator j_end(const k_iterator &k) const
Iterator pointing to last j entry for given k.
KOKKOS_INLINE_FUNCTION size_type num_non_zeros() const
Number of non-zero&#39;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< 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&#39;s.
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
expr val()
static CooProductTensor create(const Stokhos::ProductBasis< OrdinalType, ValueType > &basis, const Stokhos::Sparse3Tensor< OrdinalType, ValueType > &Cijk, const Teuchos::ParameterList &params=Teuchos::ParameterList())
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&#39;s per multiply-add.
KOKKOS_INLINE_FUNCTION size_type num_flops() const
Number flop&#39;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 &#39;COO&#39;-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 &params=Teuchos::ParameterList())
int n
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.
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)