Kokkos Core Kernels Package  Version of the Day
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Groups Pages
Kokkos_Crs.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_IMPL_PUBLIC_INCLUDE
18 #include <Kokkos_Macros.hpp>
19 static_assert(false,
20  "Including non-public Kokkos header files is not allowed.");
21 #endif
22 #ifndef KOKKOS_CRS_HPP
23 #define KOKKOS_CRS_HPP
24 
25 #include <Kokkos_View.hpp>
26 #include <Kokkos_CopyViews.hpp>
27 
28 namespace Kokkos {
29 
60 template <class DataType, class Arg1Type, class Arg2Type = void,
61  typename SizeType = typename ViewTraits<DataType*, Arg1Type, Arg2Type,
62  void>::size_type>
63 class Crs {
64  protected:
66 
67  public:
68  using data_type = DataType;
69  using array_layout = typename traits::array_layout;
70  using execution_space = typename traits::execution_space;
71  using memory_space = typename traits::memory_space;
72  using device_type = typename traits::device_type;
73  using size_type = SizeType;
74 
76  using HostMirror =
80 
81  row_map_type row_map;
82  entries_type entries;
83 
84  /*
85  * Default Constructors, operators and destructor
86  */
87  KOKKOS_DEFAULTED_FUNCTION Crs() = default;
88  KOKKOS_DEFAULTED_FUNCTION Crs(Crs const&) = default;
89  KOKKOS_DEFAULTED_FUNCTION Crs(Crs&&) = default;
90  KOKKOS_DEFAULTED_FUNCTION Crs& operator=(Crs const&) = default;
91  KOKKOS_DEFAULTED_FUNCTION Crs& operator=(Crs&&) = default;
92  KOKKOS_DEFAULTED_FUNCTION ~Crs() = default;
93 
98  template <class EntriesType, class RowMapType>
99  KOKKOS_INLINE_FUNCTION Crs(const RowMapType& row_map_,
100  const EntriesType& entries_)
101  : row_map(row_map_), entries(entries_) {}
102 
105  KOKKOS_INLINE_FUNCTION
106  size_type numRows() const {
107  return (row_map.extent(0) != 0)
108  ? row_map.extent(0) - static_cast<size_type>(1)
109  : static_cast<size_type>(0);
110  }
111 };
112 
113 /*--------------------------------------------------------------------------*/
114 
115 template <class OutCounts, class DataType, class Arg1Type, class Arg2Type,
116  class SizeType>
117 void get_crs_transpose_counts(
118  OutCounts& out, Crs<DataType, Arg1Type, Arg2Type, SizeType> const& in,
119  std::string const& name = "transpose_counts");
120 
121 template <class OutCounts, class InCrs>
122 typename OutCounts::value_type get_crs_row_map_from_counts(
123  OutCounts& out, InCrs const& in, std::string const& name = "row_map");
124 
125 template <class DataType, class Arg1Type, class Arg2Type, class SizeType>
126 void transpose_crs(Crs<DataType, Arg1Type, Arg2Type, SizeType>& out,
127  Crs<DataType, Arg1Type, Arg2Type, SizeType> const& in);
128 
129 } // namespace Kokkos
130 
131 /*--------------------------------------------------------------------------*/
132 
133 /*--------------------------------------------------------------------------*/
134 
135 namespace Kokkos {
136 namespace Impl {
137 
138 template <class InCrs, class OutCounts>
139 class GetCrsTransposeCounts {
140  public:
141  using execution_space = typename InCrs::execution_space;
142  using self_type = GetCrsTransposeCounts<InCrs, OutCounts>;
143  using index_type = typename InCrs::size_type;
144 
145  private:
146  InCrs in;
147  OutCounts out;
148 
149  public:
150  KOKKOS_INLINE_FUNCTION
151  void operator()(index_type i) const { atomic_increment(&out[in.entries(i)]); }
152  GetCrsTransposeCounts(InCrs const& arg_in, OutCounts const& arg_out)
153  : in(arg_in), out(arg_out) {
154  using policy_type = RangePolicy<index_type, execution_space>;
156  const closure_type closure(*this,
157  policy_type(0, index_type(in.entries.size())));
158  closure.execute();
159  execution_space().fence(
160  "Kokkos::Impl::GetCrsTransposeCounts::GetCrsTransposeCounts: fence "
161  "after functor execution");
162  }
163 };
164 
165 template <class InCounts, class OutRowMap>
166 class CrsRowMapFromCounts {
167  public:
168  using execution_space = typename InCounts::execution_space;
169  using value_type = typename OutRowMap::value_type;
170  using index_type = typename InCounts::size_type;
171  using last_value_type =
173 
174  private:
175  InCounts m_in;
176  OutRowMap m_out;
177  last_value_type m_last_value;
178 
179  public:
180  KOKKOS_INLINE_FUNCTION
181  void operator()(index_type i, value_type& update, bool final_pass) const {
182  if (i < static_cast<index_type>(m_in.size())) {
183  update += m_in(i);
184  if (final_pass) m_out(i + 1) = update;
185  } else if (final_pass) {
186  m_out(0) = 0;
187  m_last_value() = update;
188  }
189  }
190  KOKKOS_INLINE_FUNCTION
191  void init(value_type& update) const { update = 0; }
192  KOKKOS_INLINE_FUNCTION
193  void join(value_type& update, const value_type& input) const {
194  update += input;
195  }
196  using self_type = CrsRowMapFromCounts<InCounts, OutRowMap>;
197  CrsRowMapFromCounts(InCounts const& arg_in, OutRowMap const& arg_out)
198  : m_in(arg_in), m_out(arg_out), m_last_value("last_value") {}
199  value_type execute() {
200  using policy_type = RangePolicy<index_type, execution_space>;
202  closure_type closure(*this, policy_type(0, m_in.size() + 1));
203  closure.execute();
204  auto last_value =
205  Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, m_last_value);
206  return last_value();
207  }
208 };
209 
210 template <class InCrs, class OutCrs>
211 class FillCrsTransposeEntries {
212  public:
213  using execution_space = typename InCrs::execution_space;
214  using memory_space = typename InCrs::memory_space;
215  using value_type = typename OutCrs::entries_type::value_type;
216  using index_type = typename InCrs::size_type;
217 
218  private:
219  using counters_type = View<index_type*, memory_space>;
220  InCrs in;
221  OutCrs out;
222  counters_type counters;
223 
224  public:
225  KOKKOS_INLINE_FUNCTION
226  void operator()(index_type i) const {
227  auto begin = in.row_map(i);
228  auto end = in.row_map(i + 1);
229  for (auto j = begin; j < end; ++j) {
230  auto ti = in.entries(j);
231  auto tbegin = out.row_map(ti);
232  auto tj = atomic_fetch_add(&counters(ti), 1);
233  out.entries(tbegin + tj) = i;
234  }
235  }
236  using self_type = FillCrsTransposeEntries<InCrs, OutCrs>;
237  FillCrsTransposeEntries(InCrs const& arg_in, OutCrs const& arg_out)
238  : in(arg_in), out(arg_out), counters("counters", arg_out.numRows()) {
239  using policy_type = RangePolicy<index_type, execution_space>;
241  const closure_type closure(*this, policy_type(0, index_type(in.numRows())));
242  closure.execute();
243  execution_space().fence(
244  "Kokkos::Impl::FillCrsTransposeEntries::FillCrsTransposeEntries: fence "
245  "after functor execution");
246  }
247 };
248 
249 } // namespace Impl
250 } // namespace Kokkos
251 
252 /*--------------------------------------------------------------------------*/
253 
254 /*--------------------------------------------------------------------------*/
255 
256 namespace Kokkos {
257 
258 template <class OutCounts, class DataType, class Arg1Type, class Arg2Type,
259  class SizeType>
260 void get_crs_transpose_counts(
261  OutCounts& out, Crs<DataType, Arg1Type, Arg2Type, SizeType> const& in,
262  std::string const& name) {
263  using InCrs = Crs<DataType, Arg1Type, Arg2Type, SizeType>;
264  out = OutCounts(name, in.numRows());
265  Kokkos::Impl::GetCrsTransposeCounts<InCrs, OutCounts> functor(in, out);
266 }
267 
268 template <class OutRowMap, class InCounts>
269 typename OutRowMap::value_type get_crs_row_map_from_counts(
270  OutRowMap& out, InCounts const& in, std::string const& name) {
271  out = OutRowMap(view_alloc(WithoutInitializing, name), in.size() + 1);
272  Kokkos::Impl::CrsRowMapFromCounts<InCounts, OutRowMap> functor(in, out);
273  return functor.execute();
274 }
275 
276 template <class DataType, class Arg1Type, class Arg2Type, class SizeType>
277 void transpose_crs(Crs<DataType, Arg1Type, Arg2Type, SizeType>& out,
278  Crs<DataType, Arg1Type, Arg2Type, SizeType> const& in) {
279  using crs_type = Crs<DataType, Arg1Type, Arg2Type, SizeType>;
280  using memory_space = typename crs_type::memory_space;
281  using counts_type = View<SizeType*, memory_space>;
282  {
283  counts_type counts;
284  Kokkos::get_crs_transpose_counts(counts, in);
285  Kokkos::get_crs_row_map_from_counts(out.row_map, counts,
286  "tranpose_row_map");
287  }
288  out.entries = decltype(out.entries)("transpose_entries", in.entries.size());
289  Kokkos::Impl::FillCrsTransposeEntries<crs_type, crs_type> entries_functor(
290  in, out);
291 }
292 
293 template <class CrsType, class Functor,
294  class ExecutionSpace = typename CrsType::execution_space>
295 struct CountAndFillBase;
296 
297 template <class CrsType, class Functor, class ExecutionSpace>
298 struct CountAndFillBase {
299  using data_type = typename CrsType::data_type;
300  using size_type = typename CrsType::size_type;
301  using row_map_type = typename CrsType::row_map_type;
302  using counts_type = row_map_type;
303  CrsType m_crs;
304  Functor m_functor;
305  counts_type m_counts;
306  struct Count {};
307  KOKKOS_FUNCTION void operator()(Count, size_type i) const {
308  m_counts(i) = m_functor(i, nullptr);
309  }
310  struct Fill {};
311  KOKKOS_FUNCTION void operator()(Fill, size_type i) const {
312  auto j = m_crs.row_map(i);
313  /* we don't want to access entries(entries.size()), even if its just to get
314  its address and never use it. this can happen when row (i) is empty and
315  all rows after it are also empty. we could compare to row_map(i + 1), but
316  that is a read from global memory, whereas dimension_0() should be part
317  of the View in registers (or constant memory) */
318  data_type* fill = (j == static_cast<decltype(j)>(m_crs.entries.extent(0)))
319  ? nullptr
320  : (&(m_crs.entries(j)));
321  m_functor(i, fill);
322  }
323  CountAndFillBase(CrsType& crs, Functor const& f) : m_crs(crs), m_functor(f) {}
324 };
325 
326 template <class CrsType, class Functor>
327 struct CountAndFill : public CountAndFillBase<CrsType, Functor> {
328  using base_type = CountAndFillBase<CrsType, Functor>;
329  using typename base_type::Count;
330  using typename base_type::counts_type;
331  using typename base_type::data_type;
332  using typename base_type::Fill;
333  using typename base_type::size_type;
334  using entries_type = typename CrsType::entries_type;
335  using self_type = CountAndFill<CrsType, Functor>;
336  CountAndFill(CrsType& crs, size_type nrows, Functor const& f)
337  : base_type(crs, f) {
338  using execution_space = typename CrsType::execution_space;
339  this->m_counts = counts_type("counts", nrows);
340  {
341  using count_policy_type = RangePolicy<size_type, execution_space, Count>;
342  using count_closure_type =
344  const count_closure_type closure(*this, count_policy_type(0, nrows));
345  closure.execute();
346  }
347  auto nentries = Kokkos::get_crs_row_map_from_counts(this->m_crs.row_map,
348  this->m_counts);
349  this->m_counts = counts_type();
350  this->m_crs.entries = entries_type("entries", nentries);
351  {
352  using fill_policy_type = RangePolicy<size_type, execution_space, Fill>;
353  using fill_closure_type =
355  const fill_closure_type closure(*this, fill_policy_type(0, nrows));
356  closure.execute();
357  }
358  crs = this->m_crs;
359  }
360 };
361 
362 template <class CrsType, class Functor>
363 void count_and_fill_crs(CrsType& crs, typename CrsType::size_type nrows,
364  Functor const& f) {
365  Kokkos::CountAndFill<CrsType, Functor>(crs, nrows, f);
366 }
367 
368 } // namespace Kokkos
369 
370 #endif /* #define KOKKOS_CRS_HPP */
Implementation detail of parallel_scan.
Memory management for host memory.
Compressed row storage array.
Definition: Kokkos_Crs.hpp:63
KOKKOS_INLINE_FUNCTION size_type numRows() const
Return number of rows in the graph.
Definition: Kokkos_Crs.hpp:106
KOKKOS_INLINE_FUNCTION Crs(const RowMapType &row_map_, const EntriesType &entries_)
Assign to a view of the rhs array. If the old view is the last view then allocated memory is dealloca...
Definition: Kokkos_Crs.hpp:99
Implementation of the ParallelFor operator that has a partial specialization for the device...
Traits class for accessing attributes of a View.