functional2.hpp 1.19 KB
Newer Older
1
2
3
#ifndef CK_FUNCTIONAL2_HPP
#define CK_FUNCTIONAL2_HPP

Chao Liu's avatar
Chao Liu committed
4
#include "functional.hpp"
Chao Liu's avatar
Chao Liu committed
5
#include "sequence.hpp"
6

7
8
namespace ck {

Chao Liu's avatar
Chao Liu committed
9
10
namespace detail {

Chao Liu's avatar
Chao Liu committed
11
12
template <class>
struct static_for_impl;
13

Chao Liu's avatar
Chao Liu committed
14
15
template <index_t... Is>
struct static_for_impl<Sequence<Is...>>
16
17
{
    template <class F>
Chao Liu's avatar
Chao Liu committed
18
    __host__ __device__ constexpr void operator()(F f) const
19
    {
Chao Liu's avatar
Chao Liu committed
20
        swallow{(f(Number<Is>{}), 0)...};
21
22
23
    }
};

Chao Liu's avatar
Chao Liu committed
24
25
} // namespace detail

Chao Liu's avatar
Chao Liu committed
26
27
28
// F signature: F(Number<Iter>)
template <index_t NBegin, index_t NEnd, index_t Increment>
struct static_for
29
{
30
    __host__ __device__ constexpr static_for()
31
    {
32
        static_assert(Increment != 0 && (NEnd - NBegin) % Increment == 0,
Chao Liu's avatar
Chao Liu committed
33
                      "Wrong! should satisfy (NEnd - NBegin) % Increment == 0");
34
        static_assert((Increment > 0 && NBegin <= NEnd) || (Increment < 0 && NBegin >= NEnd),
Chao Liu's avatar
Chao Liu committed
35
36
                      "wrongs! should (Increment > 0 && NBegin <= NEnd) || (Increment < 0 && "
                      "NBegin >= NEnd)");
37
    }
38

39
40
41
    template <class F>
    __host__ __device__ constexpr void operator()(F f) const
    {
Chao Liu's avatar
Chao Liu committed
42
43
        detail::static_for_impl<typename arithmetic_sequence_gen<NBegin, NEnd, Increment>::type>{}(
            f);
44
45
46
    }
};

47
48
} // namespace ck
#endif