Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
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
31namespace Kokkos {
32
33namespace Impl {
34template <class RowOffsetsType, class RowBlockOffsetsType>
35struct 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
146template <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
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
256template <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
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
393template <class StaticCrsGraphType, class InputSizeType>
394typename StaticCrsGraphType::staticcrsgraph_type create_staticcrsgraph(
395 const std::string& label, const std::vector<InputSizeType>& input);
396
397template <class StaticCrsGraphType, class InputSizeType>
398typename StaticCrsGraphType::staticcrsgraph_type create_staticcrsgraph(
399 const std::string& label,
400 const std::vector<std::vector<InputSizeType> >& input);
401
402//----------------------------------------------------------------------------
403
404template <class DataType, class Arg1Type, class Arg2Type, class Arg3Type,
405 typename SizeType>
406typename StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
407 SizeType>::HostMirror
408create_mirror_view(const StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
409 SizeType>& input);
410
411template <class DataType, class Arg1Type, class Arg2Type, class Arg3Type,
412 typename SizeType>
413typename StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
414 SizeType>::HostMirror
415create_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
428namespace Kokkos {
429namespace Impl {
430
431template <class GraphType>
432struct 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
456template <class DataType, class Arg1Type, class Arg2Type, class Arg3Type,
457 typename SizeType>
458DataType 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 */
Declaration of parallel operators.
Execution policy for work over a range of an integral type.
Compressed row storage array.
KOKKOS_DEFAULTED_FUNCTION ~StaticCrsGraph()=default
Destroy this view of the array. If the last view then allocated memory is deallocated.
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.
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.
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...
KOKKOS_INLINE_FUNCTION StaticCrsGraph(const StaticCrsGraph &rhs)
Copy constructor (shallow copy).
KOKKOS_INLINE_FUNCTION size_type numRows() const
Return number of rows in the graph.
KOKKOS_INLINE_FUNCTION constexpr std::enable_if_t< std::is_integral< iType >::value, size_t > extent(const iType &r) const noexcept
rank() to be implemented
View of a row of a sparse graph.
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.
const ordinal_type length
Number of entries in the row.
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.
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.
Traits class for accessing attributes of a View.