xref: /petsc/src/sys/objects/device/impls/cupm/kernels.hpp (revision 2ea277ce7ccf004135b03702db4a237d7886e1ba) !
16d54fb17SJacob Faibussowitsch #ifndef PETSC_DEVICE_CUPM_KERNELS_HPP
26d54fb17SJacob Faibussowitsch #define PETSC_DEVICE_CUPM_KERNELS_HPP
36d54fb17SJacob Faibussowitsch 
46d54fb17SJacob Faibussowitsch #include <petscdevice_cupm.h>
56d54fb17SJacob Faibussowitsch 
66d54fb17SJacob Faibussowitsch #if defined(__cplusplus)
76d54fb17SJacob Faibussowitsch 
86d54fb17SJacob Faibussowitsch namespace Petsc
96d54fb17SJacob Faibussowitsch {
106d54fb17SJacob Faibussowitsch 
116d54fb17SJacob Faibussowitsch namespace device
126d54fb17SJacob Faibussowitsch {
136d54fb17SJacob Faibussowitsch 
146d54fb17SJacob Faibussowitsch namespace cupm
156d54fb17SJacob Faibussowitsch {
166d54fb17SJacob Faibussowitsch 
176d54fb17SJacob Faibussowitsch namespace kernels
186d54fb17SJacob Faibussowitsch {
196d54fb17SJacob Faibussowitsch 
206d54fb17SJacob Faibussowitsch namespace util
216d54fb17SJacob Faibussowitsch {
226d54fb17SJacob Faibussowitsch 
236d54fb17SJacob Faibussowitsch template <typename SizeType, typename T>
246d54fb17SJacob Faibussowitsch PETSC_DEVICE_INLINE_DECL static void grid_stride_1D(const SizeType size, T &&func) noexcept
256d54fb17SJacob Faibussowitsch {
266d54fb17SJacob Faibussowitsch   for (SizeType i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += blockDim.x * gridDim.x) func(i);
276d54fb17SJacob Faibussowitsch   return;
286d54fb17SJacob Faibussowitsch }
296d54fb17SJacob Faibussowitsch 
306d54fb17SJacob Faibussowitsch } // namespace util
316d54fb17SJacob Faibussowitsch 
326d54fb17SJacob Faibussowitsch } // namespace kernels
336d54fb17SJacob Faibussowitsch 
3490585354SJacob Faibussowitsch namespace functors
3590585354SJacob Faibussowitsch {
3690585354SJacob Faibussowitsch 
3790585354SJacob Faibussowitsch template <typename T>
3890585354SJacob Faibussowitsch class plus_equals {
3990585354SJacob Faibussowitsch public:
4090585354SJacob Faibussowitsch   using value_type = T;
4190585354SJacob Faibussowitsch 
4290585354SJacob Faibussowitsch   PETSC_HOSTDEVICE_DECL constexpr explicit plus_equals(value_type v = value_type{}) noexcept : v_{std::move(v)} { }
4390585354SJacob Faibussowitsch 
4490585354SJacob Faibussowitsch   PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL constexpr value_type operator()(const value_type &val) const noexcept { return val + v_; }
4590585354SJacob Faibussowitsch 
4690585354SJacob Faibussowitsch private:
4790585354SJacob Faibussowitsch   value_type v_;
4890585354SJacob Faibussowitsch };
4990585354SJacob Faibussowitsch 
50*2ea277ceSJacob Faibussowitsch template <typename T>
51*2ea277ceSJacob Faibussowitsch class times_equals {
52*2ea277ceSJacob Faibussowitsch public:
53*2ea277ceSJacob Faibussowitsch   using value_type = T;
54*2ea277ceSJacob Faibussowitsch 
55*2ea277ceSJacob Faibussowitsch   PETSC_HOSTDEVICE_DECL constexpr explicit times_equals(value_type v = value_type{}) noexcept : v_{std::move(v)} { }
56*2ea277ceSJacob Faibussowitsch 
57*2ea277ceSJacob Faibussowitsch   PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL constexpr value_type operator()(const value_type &val) const noexcept { return val * v_; }
58*2ea277ceSJacob Faibussowitsch 
59*2ea277ceSJacob Faibussowitsch private:
60*2ea277ceSJacob Faibussowitsch   value_type v_;
61*2ea277ceSJacob Faibussowitsch };
62*2ea277ceSJacob Faibussowitsch 
6390585354SJacob Faibussowitsch namespace
6490585354SJacob Faibussowitsch {
6590585354SJacob Faibussowitsch 
6690585354SJacob Faibussowitsch template <typename T>
6790585354SJacob Faibussowitsch PETSC_HOSTDEVICE_INLINE_DECL constexpr plus_equals<T> make_plus_equals(const T &v) noexcept
6890585354SJacob Faibussowitsch {
6990585354SJacob Faibussowitsch   return plus_equals<T>{v};
7090585354SJacob Faibussowitsch }
7190585354SJacob Faibussowitsch 
72*2ea277ceSJacob Faibussowitsch template <typename T>
73*2ea277ceSJacob Faibussowitsch PETSC_HOSTDEVICE_INLINE_DECL constexpr times_equals<T> make_times_equals(const T &v) noexcept
74*2ea277ceSJacob Faibussowitsch {
75*2ea277ceSJacob Faibussowitsch   return times_equals<T>{v};
76*2ea277ceSJacob Faibussowitsch }
77*2ea277ceSJacob Faibussowitsch 
7890585354SJacob Faibussowitsch } // anonymous namespace
7990585354SJacob Faibussowitsch 
8090585354SJacob Faibussowitsch } // namespace functors
8190585354SJacob Faibussowitsch 
826d54fb17SJacob Faibussowitsch } // namespace cupm
836d54fb17SJacob Faibussowitsch 
846d54fb17SJacob Faibussowitsch } // namespace device
856d54fb17SJacob Faibussowitsch 
866d54fb17SJacob Faibussowitsch } // namespace Petsc
876d54fb17SJacob Faibussowitsch 
886d54fb17SJacob Faibussowitsch #endif // __cplusplus
896d54fb17SJacob Faibussowitsch 
906d54fb17SJacob Faibussowitsch #endif // PETSC_DEVICE_CUPM_KERNELS_HPP
91