functional.hip.hpp 3.66 KB
Newer Older
1
#pragma once
Chao Liu's avatar
Chao Liu committed
2
#include "integral_constant.hip.hpp"
3

Chao Liu's avatar
Chao Liu committed
4
5
6
struct forwarder
{
    template <typename T>
Chao Liu's avatar
Chao Liu committed
7
    __host__ __device__ constexpr T&& operator()(T&& x) const
Chao Liu's avatar
Chao Liu committed
8
    {
Chao Liu's avatar
Chao Liu committed
9
        return static_cast<T&&>(x);
Chao Liu's avatar
Chao Liu committed
10
11
12
    }
};

Chao Liu's avatar
Chao Liu committed
13
14
15
16
17
18
19
20
#if 0
template<class F>
__host__ __device__ constexpr auto unpacker(F f)
{
    return [=](auto xs_array){ f(xs...); };
}
#endif

Chao Liu's avatar
Chao Liu committed
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
// Emulate compile time if statement for C++14
//   Get the idea from
//   "https://baptiste-wicht.com/posts/2015/07/simulate-static_if-with-c11c14.html"
// TODO: use if constexpr, when C++17 is supported
template <bool Predicate>
struct static_if
{
};

template <>
struct static_if<true>
{
    using Type = static_if<true>;

    template <class F>
    __host__ __device__ constexpr auto operator()(F f) const
    {
        // This is a trick for compiler:
        //   Pass forwarder to lambda "f" as "auto" argument, and maks sure "f" will use it,
Chao Liu's avatar
Chao Liu committed
40
41
        //   this will make "f" a generic lambda, so that "f" won't be compiled until being
        //   instantiated here
Chao Liu's avatar
Chao Liu committed
42
43
44
45
46
        f(forwarder{});
        return Type{};
    }

    template <class F>
47
    __host__ __device__ static constexpr auto Else(F)
Chao Liu's avatar
Chao Liu committed
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    {
        return Type{};
    }
};

template <>
struct static_if<false>
{
    using Type = static_if<false>;

    template <class F>
    __host__ __device__ constexpr auto operator()(F) const
    {
        return Type{};
    }

    template <class F>
65
    __host__ __device__ static constexpr auto Else(F f)
Chao Liu's avatar
Chao Liu committed
66
67
68
    {
        // This is a trick for compiler:
        //   Pass forwarder to lambda "f" as "auto" argument, and maks sure "f" will use it,
Chao Liu's avatar
Chao Liu committed
69
70
        //   this will make "f" a generic lambda, so that "f" won't be compiled until being
        //   instantiated here
Chao Liu's avatar
Chao Liu committed
71
72
73
74
        f(forwarder{});
        return Type{};
    }
};
75
76
template <index_t Iter, index_t Remaining, index_t Increment>
struct static_for_impl
77
78
{
    template <class F>
Chao Liu's avatar
Chao Liu committed
79
    constexpr __host__ __device__ void operator()(F f) const
80
    {
81
82
        static_assert(Remaining % Increment == 0, "wrong! Remaining % Increment != 0");
        static_assert(Increment <= Remaining, "will go out-of-range");
83

84
85
        f(Number<Iter>{});
        static_for_impl<Iter + Increment, Remaining - Increment, Increment>{}(f);
86
87
88
    }
};

89
90
91
92
template <index_t Iter, index_t Increment>
struct static_for_impl<Iter, 0, Increment>
{
    template <class F>
Chao Liu's avatar
Chao Liu committed
93
    constexpr __host__ __device__ void operator()(F) const
94
    {
95
        // no work left, just return
96
97
98
99
        return;
    }
};

Chao Liu's avatar
Chao Liu committed
100
// F signature: F(Number<Iter>)
101
102
template <index_t NBegin, index_t NEnd, index_t Increment>
struct static_for
103
104
{
    template <class F>
Chao Liu's avatar
Chao Liu committed
105
    constexpr __host__ __device__ void operator()(F f) const
106
    {
Chao Liu's avatar
Chao Liu committed
107
108
        static_assert(NBegin <= NEnd, "wrongs! should have NBegin <= NEnd");

109
110
        static_assert((NEnd - NBegin) % Increment == 0,
                      "Wrong! should satisfy (NEnd - NBegin) % Increment == 0");
Chao Liu's avatar
Chao Liu committed
111

Chao Liu's avatar
Chao Liu committed
112
#if 0
Chao Liu's avatar
Chao Liu committed
113
        static_if<(NBegin < NEnd)>{}(
Chao Liu's avatar
Chao Liu committed
114
            [&](auto fwd) { static_for_impl<NBegin, NEnd - NBegin, fwd(Increment)>{}(f); });
Chao Liu's avatar
Chao Liu committed
115
116
117
#else
        static_for_impl<NBegin, NEnd - NBegin, Increment>{}(f);
#endif
118
119
120
    }
};

Chao Liu's avatar
Chao Liu committed
121
template <index_t NLoop>
122
123
struct static_const_reduce_n
{
Chao Liu's avatar
Chao Liu committed
124
    // signature of F: F(Number<I>)
125
126
127
128
129
130
    template <class F, class Reduce>
    __host__ __device__ constexpr auto operator()(F f, Reduce r) const
    {
        static_assert(NLoop > 1, "out-of-range");

        constexpr auto a = f(Number<NLoop - 1>{});
131
        auto b = static_const_reduce_n<NLoop - 1>{}(f, r); // TODO: cannot use constexpr here, weird
132
133
134
135
136
137
138
139
140
141
142
143
144
        return r(a, b);
    }
};

template <>
struct static_const_reduce_n<1>
{
    template <class F, class Reduce>
    __host__ __device__ constexpr auto operator()(F f, Reduce) const
    {
        return f(Number<0>{});
    }
};