Kokkos Core Kernels Package  Version of the Day
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Groups Pages
Kokkos_StaticCrsGraph.hpp
1 //@HEADER
2 // ************************************************************************
3 //
4 // Kokkos v. 4.0
5 // Copyright (2022) National Technology & Engineering
6 // Solutions of Sandia, LLC (NTESS).
7 //
8 // Under the terms of Contract DE-NA0003525 with NTESS,
9 // the U.S. Government retains certain rights in this software.
10 //
11 // Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12 // See https://kokkos.org/LICENSE for license information.
13 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14 //
15 //@HEADER
16 
17 #ifndef KOKKOS_STATICCRSGRAPH_HPP
18 #define KOKKOS_STATICCRSGRAPH_HPP
19 #ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
20 #define KOKKOS_IMPL_PUBLIC_INCLUDE
21 #define KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_STATICCRSGRAPH
22 #endif
23 
24 #include <string>
25 #include <vector>
26 
27 #include <Kokkos_View.hpp>
28 #include <Kokkos_Parallel.hpp>
29 #include <Kokkos_Parallel_Reduce.hpp>
30 
31 namespace Kokkos {
32 
33 namespace Impl {
34 template <class RowOffsetsType, class RowBlockOffsetsType>
35 struct StaticCrsGraphBalancerFunctor {
36  using int_type = typename RowOffsetsType::non_const_value_type;
37  RowOffsetsType row_offsets;
38  RowBlockOffsetsType row_block_offsets;
39 
40  int_type cost_per_row, num_blocks;
41 
42  StaticCrsGraphBalancerFunctor(RowOffsetsType row_offsets_,
43  RowBlockOffsetsType row_block_offsets_,
44  int_type cost_per_row_, int_type num_blocks_)
45  : row_offsets(row_offsets_),
46  row_block_offsets(row_block_offsets_),
47  cost_per_row(cost_per_row_),
48  num_blocks(num_blocks_) {}
49 
50  KOKKOS_INLINE_FUNCTION
51  void operator()(const int_type& iRow) const {
52  const int_type num_rows = row_offsets.extent(0) - 1;
53  const int_type num_entries = row_offsets(num_rows);
54  const int_type total_cost = num_entries + num_rows * cost_per_row;
55 
56  const double cost_per_workset = 1.0 * total_cost / num_blocks;
57 
58  const int_type row_cost =
59  row_offsets(iRow + 1) - row_offsets(iRow) + cost_per_row;
60 
61  int_type count = row_offsets(iRow + 1) + cost_per_row * iRow;
62 
63  if (iRow == num_rows - 1) row_block_offsets(num_blocks) = num_rows;
64 
65  if (true) {
66  int_type current_block =
67  (count - row_cost - cost_per_row) / cost_per_workset;
68  int_type end_block = count / cost_per_workset;
69 
70  // Handle some corner cases for the last two blocks.
71  if (current_block >= num_blocks - 2) {
72  if ((current_block == num_blocks - 2) &&
73  (count >= (current_block + 1) * cost_per_workset)) {
74  int_type row = iRow;
75  int_type cc = count - row_cost - cost_per_row;
76  int_type block = cc / cost_per_workset;
77  while ((block > 0) && (block == current_block)) {
78  cc = row_offsets(row) + row * cost_per_row;
79  block = cc / cost_per_workset;
80  row--;
81  }
82  if ((count - cc - row_cost - cost_per_row) <
83  num_entries - row_offsets(iRow + 1)) {
84  row_block_offsets(current_block + 1) = iRow + 1;
85  } else {
86  row_block_offsets(current_block + 1) = iRow;
87  }
88  }
89  } else {
90  if ((count >= (current_block + 1) * cost_per_workset) ||
91  (iRow + 2 == int_type(row_offsets.extent(0)))) {
92  if (end_block > current_block + 1) {
93  int_type num_block = end_block - current_block;
94  row_block_offsets(current_block + 1) = iRow;
95  for (int_type block = current_block + 2; block <= end_block;
96  block++)
97  if ((block < current_block + 2 + (num_block - 1) / 2))
98  row_block_offsets(block) = iRow;
99  else
100  row_block_offsets(block) = iRow + 1;
101  } else {
102  row_block_offsets(current_block + 1) = iRow + 1;
103  }
104  }
105  }
106  }
107  }
108 };
109 } // namespace Impl
110 
146 template <class GraphType>
149  using ordinal_type = const typename GraphType::data_type;
150 
151  private:
153  ordinal_type* colidx_;
161  const ordinal_type stride_;
162 
163  public:
171  KOKKOS_INLINE_FUNCTION
172  GraphRowViewConst(ordinal_type* const colidx_in, const ordinal_type& stride,
173  const ordinal_type& count)
174  : colidx_(colidx_in), stride_(stride), length(count) {}
175 
188  template <class OffsetType>
189  KOKKOS_INLINE_FUNCTION GraphRowViewConst(
190  const typename GraphType::entries_type& colidx_in,
191  const ordinal_type& stride, const ordinal_type& count,
192  const OffsetType& idx,
193  const std::enable_if_t<std::is_integral<OffsetType>::value, int>& = 0)
194  : colidx_(&colidx_in(idx)), stride_(stride), length(count) {}
195 
207 
213  KOKKOS_INLINE_FUNCTION
214  ordinal_type& colidx(const ordinal_type& i) const {
215  return colidx_[i * stride_];
216  }
217 
219  KOKKOS_INLINE_FUNCTION
220  ordinal_type& operator()(const ordinal_type& i) const { return colidx(i); }
221 };
222 
256 template <class DataType, class Arg1Type, class Arg2Type = void,
257  class Arg3Type = void,
258  typename SizeType = typename ViewTraits<DataType*, Arg1Type, Arg2Type,
259  Arg3Type>::size_type>
261  private:
263 
264  public:
265  using data_type = DataType;
266  using array_layout = typename traits::array_layout;
267  using execution_space = typename traits::execution_space;
268  using device_type = typename traits::device_type;
269  using memory_traits = typename traits::memory_traits;
270  using size_type = SizeType;
271 
272  using staticcrsgraph_type =
274  using HostMirror = StaticCrsGraph<data_type, array_layout,
275  typename traits::host_mirror_space,
276  memory_traits, size_type>;
277 
278  using row_map_type =
280  using entries_type =
282  using row_block_type =
284 
285  entries_type entries;
286  row_map_type row_map;
287  row_block_type row_block_offsets;
288 
290  KOKKOS_INLINE_FUNCTION
291  StaticCrsGraph() : entries(), row_map(), row_block_offsets() {}
292 
294  KOKKOS_INLINE_FUNCTION
296  : entries(rhs.entries),
297  row_map(rhs.row_map),
298  row_block_offsets(rhs.row_block_offsets) {}
299 
300  template <class EntriesType, class RowMapType>
301  KOKKOS_INLINE_FUNCTION StaticCrsGraph(const EntriesType& entries_,
302  const RowMapType& row_map_)
303  : entries(entries_), row_map(row_map_), row_block_offsets() {}
304 
309  KOKKOS_INLINE_FUNCTION
311  entries = rhs.entries;
312  row_map = rhs.row_map;
313  row_block_offsets = rhs.row_block_offsets;
314  return *this;
315  }
316 
320  KOKKOS_DEFAULTED_FUNCTION
321  ~StaticCrsGraph() = default;
322 
325  KOKKOS_INLINE_FUNCTION
326  size_type numRows() const {
327  return (row_map.extent(0) != 0)
328  ? row_map.extent(0) - static_cast<size_type>(1)
329  : static_cast<size_type>(0);
330  }
331 
332  KOKKOS_INLINE_FUNCTION constexpr bool is_allocated() const {
333  return (row_map.is_allocated() && entries.is_allocated());
334  }
335 
354  KOKKOS_INLINE_FUNCTION
355  GraphRowViewConst<StaticCrsGraph> rowConst(const data_type i) const {
356  const size_type start = row_map(i);
357  // count is guaranteed to fit in ordinal_type, as long as no row
358  // has duplicate entries.
359  const data_type count = static_cast<data_type>(row_map(i + 1) - start);
360 
361  if (count == 0) {
362  return GraphRowViewConst<StaticCrsGraph>(nullptr, 1, 0);
363  } else {
364  return GraphRowViewConst<StaticCrsGraph>(entries, 1, count, start);
365  }
366  }
367 
371  void create_block_partitioning(size_type num_blocks,
372  size_type fix_cost_per_row = 4) {
374  "StatisCrsGraph::load_balance_offsets", num_blocks + 1);
375 
376  Impl::StaticCrsGraphBalancerFunctor<
378  partitioner(row_map, block_offsets, fix_cost_per_row, num_blocks);
379 
380  Kokkos::parallel_for("Kokkos::StaticCrsGraph::create_block_partitioning",
382  partitioner);
383  typename device_type::execution_space().fence(
384  "Kokkos::StaticCrsGraph::create_block_partitioning:: fence after "
385  "partition");
386 
387  row_block_offsets = block_offsets;
388  }
389 };
390 
391 //----------------------------------------------------------------------------
392 
393 template <class StaticCrsGraphType, class InputSizeType>
394 typename StaticCrsGraphType::staticcrsgraph_type create_staticcrsgraph(
395  const std::string& label, const std::vector<InputSizeType>& input);
396 
397 template <class StaticCrsGraphType, class InputSizeType>
398 typename StaticCrsGraphType::staticcrsgraph_type create_staticcrsgraph(
399  const std::string& label,
400  const std::vector<std::vector<InputSizeType> >& input);
401 
402 //----------------------------------------------------------------------------
403 
404 template <class DataType, class Arg1Type, class Arg2Type, class Arg3Type,
405  typename SizeType>
406 typename StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
407  SizeType>::HostMirror
408 create_mirror_view(const StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
409  SizeType>& input);
410 
411 template <class DataType, class Arg1Type, class Arg2Type, class Arg3Type,
412  typename SizeType>
413 typename StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
414  SizeType>::HostMirror
415 create_mirror(const StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
416  SizeType>& input);
417 
418 } // namespace Kokkos
419 
420 //----------------------------------------------------------------------------
421 //----------------------------------------------------------------------------
422 
423 #include <impl/Kokkos_StaticCrsGraph_factory.hpp>
424 
425 //----------------------------------------------------------------------------
426 //----------------------------------------------------------------------------
427 
428 namespace Kokkos {
429 namespace Impl {
430 
431 template <class GraphType>
432 struct StaticCrsGraphMaximumEntry {
433  using execution_space = typename GraphType::execution_space;
434  using value_type = typename GraphType::data_type;
435 
436  const typename GraphType::entries_type entries;
437 
438  StaticCrsGraphMaximumEntry(const GraphType& graph) : entries(graph.entries) {}
439 
440  KOKKOS_INLINE_FUNCTION
441  void operator()(const unsigned i, value_type& update) const {
442  if (update < entries(i)) update = entries(i);
443  }
444 
445  KOKKOS_INLINE_FUNCTION
446  void init(value_type& update) const { update = 0; }
447 
448  KOKKOS_INLINE_FUNCTION
449  void join(value_type& update, const value_type& input) const {
450  if (update < input) update = input;
451  }
452 };
453 
454 } // namespace Impl
455 
456 template <class DataType, class Arg1Type, class Arg2Type, class Arg3Type,
457  typename SizeType>
458 DataType maximum_entry(const StaticCrsGraph<DataType, Arg1Type, Arg2Type,
459  Arg3Type, SizeType>& graph) {
460  using GraphType =
461  StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type, SizeType>;
462  using FunctorType = Impl::StaticCrsGraphMaximumEntry<GraphType>;
463 
464  DataType result = 0;
465  Kokkos::parallel_reduce("Kokkos::maximum_entry", graph.entries.extent(0),
466  FunctorType(graph), result);
467  return result;
468 }
469 
470 } // namespace Kokkos
471 
472 //----------------------------------------------------------------------------
473 //----------------------------------------------------------------------------
474 
475 #ifdef KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_STATICCRSGRAPH
476 #undef KOKKOS_IMPL_PUBLIC_INCLUDE
477 #undef KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_STATICCRSGRAPH
478 #endif
479 #endif /* #ifndef KOKKOS_CRSARRAY_HPP */
KOKKOS_INLINE_FUNCTION GraphRowViewConst< StaticCrsGraph > rowConst(const data_type i) const
Return a const view of row i of the graph.
KOKKOS_INLINE_FUNCTION StaticCrsGraph()
Construct an empty view.
const ordinal_type length
Number of entries in the row.
KOKKOS_INLINE_FUNCTION StaticCrsGraph(const StaticCrsGraph &rhs)
Copy constructor (shallow copy).
KOKKOS_INLINE_FUNCTION GraphRowViewConst(ordinal_type *const colidx_in, const ordinal_type &stride, const ordinal_type &count)
Constructor.
KOKKOS_INLINE_FUNCTION ordinal_type & colidx(const ordinal_type &i) const
(Const) reference to the column index of entry i in this row of the sparse matrix.
KOKKOS_DEFAULTED_FUNCTION ~StaticCrsGraph()=default
Destroy this view of the array. If the last view then allocated memory is deallocated.
const typename GraphType::data_type ordinal_type
The type of the column indices in the row.
KOKKOS_INLINE_FUNCTION ordinal_type & operator()(const ordinal_type &i) const
An alias for colidx.
void create_block_partitioning(size_type num_blocks, size_type fix_cost_per_row=4)
Create a row partitioning into a given number of blocks balancing non-zeros + a fixed cost per row...
Declaration of parallel operators.
KOKKOS_INLINE_FUNCTION size_type numRows() const
Return number of rows in the graph.
Execution policy for work over a range of an integral type.
Compressed row storage array.
Traits class for accessing attributes of a View.
KOKKOS_INLINE_FUNCTION GraphRowViewConst(const typename GraphType::entries_type &colidx_in, const ordinal_type &stride, const ordinal_type &count, const OffsetType &idx, const std::enable_if_t< std::is_integral< OffsetType >::value, int > &=0)
Constructor with offset into colidx array.
View of a row of a sparse graph.
KOKKOS_INLINE_FUNCTION StaticCrsGraph & operator=(const StaticCrsGraph &rhs)
Assign to a view of the rhs array. If the old view is the last view then allocated memory is dealloca...