Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
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>
19static_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
28namespace Kokkos {
29
60template <class DataType, class Arg1Type, class Arg2Type = void,
61 typename SizeType = typename ViewTraits<DataType*, Arg1Type, Arg2Type,
62 void>::size_type>
63class 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
115template <class OutCounts, class DataType, class Arg1Type, class Arg2Type,
116 class SizeType>
117void get_crs_transpose_counts(
118 OutCounts& out, Crs<DataType, Arg1Type, Arg2Type, SizeType> const& in,
119 std::string const& name = "transpose_counts");
120
121template <class OutCounts, class InCrs>
122typename OutCounts::value_type get_crs_row_map_from_counts(
123 OutCounts& out, InCrs const& in, std::string const& name = "row_map");
124
125template <class DataType, class Arg1Type, class Arg2Type, class SizeType>
126void 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
135namespace Kokkos {
136namespace Impl {
137
138template <class InCrs, class OutCounts>
139class 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
165template <class InCounts, class OutRowMap>
166class 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
210template <class InCrs, class OutCrs>
211class 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
256namespace Kokkos {
257
258template <class OutCounts, class DataType, class Arg1Type, class Arg2Type,
259 class SizeType>
260void 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
268template <class OutRowMap, class InCounts>
269typename 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
276template <class DataType, class Arg1Type, class Arg2Type, class SizeType>
277void 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
293template <class CrsType, class Functor,
294 class ExecutionSpace = typename CrsType::execution_space>
295struct CountAndFillBase;
296
297template <class CrsType, class Functor, class ExecutionSpace>
298struct 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
326template <class CrsType, class Functor>
327struct 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
362template <class CrsType, class Functor>
363void 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 */
Compressed row storage array.
Definition: Kokkos_Crs.hpp:63
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
KOKKOS_INLINE_FUNCTION size_type numRows() const
Return number of rows in the graph.
Definition: Kokkos_Crs.hpp:106
Memory management for host memory.
Implementation of the ParallelFor operator that has a partial specialization for the device.
Implementation detail of parallel_scan.
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
Traits class for accessing attributes of a View.