"tools/collect_code_preds.py" did not exist on "2a271dbf60fe1c0917b4f1f7529a2421f33a569f"
functional2.hpp 1.16 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
#pragma once
2

Chao Liu's avatar
Chao Liu committed
3
4
#include "ck/utility/functional.hpp"
#include "ck/utility/sequence.hpp"
5

6
7
namespace ck {

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

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

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

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

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

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

46
} // namespace ck