math.hpp 4.44 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
4
#ifndef CK_MATH_HPP
#define CK_MATH_HPP

#include "config.hpp"
Chao Liu's avatar
Chao Liu committed
5
#include "integral_constant.hpp"
Chao Liu's avatar
Chao Liu committed
6
#include "number.hpp"
Chao Liu's avatar
Chao Liu committed
7
#include "type.hpp"
Chao Liu's avatar
Chao Liu committed
8
#include "enable_if.hpp"
Chao Liu's avatar
Chao Liu committed
9
10
11
12

namespace ck {
namespace math {

zjing14's avatar
zjing14 committed
13
template <typename T, T s>
Chao Liu's avatar
Chao Liu committed
14
15
16
17
18
struct scales
{
    __host__ __device__ constexpr T operator()(T a) const { return s * a; }
};

zjing14's avatar
zjing14 committed
19
template <typename T>
Chao Liu's avatar
Chao Liu committed
20
21
22
23
24
struct plus
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a + b; }
};

zjing14's avatar
zjing14 committed
25
template <typename T>
Chao Liu's avatar
Chao Liu committed
26
27
28
29
30
31
struct minus
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a - b; }
};

struct multiplies
Chao Liu's avatar
Chao Liu committed
32
33
34
35
36
37
38
39
{
    template <typename A, typename B>
    __host__ __device__ constexpr auto operator()(const A& a, const B& b) const
    {
        return a * b;
    }
};

zjing14's avatar
zjing14 committed
40
template <typename T>
Chao Liu's avatar
Chao Liu committed
41
struct maximize
Chao Liu's avatar
Chao Liu committed
42
43
44
45
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a >= b ? a : b; }
};

zjing14's avatar
zjing14 committed
46
template <typename T>
Chao Liu's avatar
Chao Liu committed
47
48
49
50
51
struct minimize
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a <= b ? a : b; }
};

zjing14's avatar
zjing14 committed
52
template <typename T>
Chao Liu's avatar
Chao Liu committed
53
54
55
56
57
58
struct integer_divide_ceiler
{
    __host__ __device__ constexpr T operator()(T a, T b) const
    {
        static_assert(is_same<T, index_t>{} || is_same<T, int>{}, "wrong type");

zjing14's avatar
zjing14 committed
59
        return (a + b - Number<1>{}) / b;
Chao Liu's avatar
Chao Liu committed
60
61
62
    }
};

zjing14's avatar
zjing14 committed
63
template <typename X, typename Y>
64
65
66
67
68
__host__ __device__ constexpr auto integer_divide_floor(X x, Y y)
{
    return x / y;
}

zjing14's avatar
zjing14 committed
69
template <typename X, typename Y>
Chao Liu's avatar
Chao Liu committed
70
__host__ __device__ constexpr auto integer_divide_ceil(X x, Y y)
Chao Liu's avatar
Chao Liu committed
71
{
72
    return (x + y - Number<1>{}) / y;
Chao Liu's avatar
Chao Liu committed
73
74
}

zjing14's avatar
zjing14 committed
75
template <typename X, typename Y>
Chao Liu's avatar
Chao Liu committed
76
__host__ __device__ constexpr auto integer_least_multiple(X x, Y y)
Chao Liu's avatar
Chao Liu committed
77
{
Chao Liu's avatar
Chao Liu committed
78
    return y * integer_divide_ceil(x, y);
Chao Liu's avatar
Chao Liu committed
79
80
}

zjing14's avatar
zjing14 committed
81
template <typename T>
Chao Liu's avatar
Chao Liu committed
82
83
84
85
86
__host__ __device__ constexpr T max(T x)
{
    return x;
}

zjing14's avatar
zjing14 committed
87
88
template <typename T>
__host__ __device__ constexpr T max(T x, T y)
Chao Liu's avatar
Chao Liu committed
89
{
zjing14's avatar
zjing14 committed
90
91
    return x > y ? x : y;
}
Chao Liu's avatar
Chao Liu committed
92

zjing14's avatar
zjing14 committed
93
94
95
96
97
template <index_t X>
__host__ __device__ constexpr index_t max(Number<X>, index_t y)
{
    return X > y ? X : y;
}
Chao Liu's avatar
Chao Liu committed
98

zjing14's avatar
zjing14 committed
99
100
101
102
103
template <index_t Y>
__host__ __device__ constexpr index_t max(index_t x, Number<Y>)
{
    return x > Y ? x : Y;
}
Chao Liu's avatar
Chao Liu committed
104

zjing14's avatar
zjing14 committed
105
106
107
108
109
110
template <typename X, typename... Ys>
__host__ __device__ constexpr auto max(X x, Ys... ys)
{
    static_assert(sizeof...(Ys) > 0, "not enough argument");

    return max(x, max(ys...));
Chao Liu's avatar
Chao Liu committed
111
112
}

zjing14's avatar
zjing14 committed
113
template <typename T>
Chao Liu's avatar
Chao Liu committed
114
115
116
117
118
__host__ __device__ constexpr T min(T x)
{
    return x;
}

zjing14's avatar
zjing14 committed
119
120
121
122
123
124
125
126
template <typename T>
__host__ __device__ constexpr T min(T x, T y)
{
    return x < y ? x : y;
}

template <index_t X>
__host__ __device__ constexpr index_t min(Number<X>, index_t y)
Chao Liu's avatar
Chao Liu committed
127
{
zjing14's avatar
zjing14 committed
128
129
    return X < y ? X : y;
}
Chao Liu's avatar
Chao Liu committed
130

zjing14's avatar
zjing14 committed
131
132
133
134
135
template <index_t Y>
__host__ __device__ constexpr index_t min(index_t x, Number<Y>)
{
    return x < Y ? x : Y;
}
Chao Liu's avatar
Chao Liu committed
136

zjing14's avatar
zjing14 committed
137
138
139
140
template <typename X, typename... Ys>
__host__ __device__ constexpr auto min(X x, Ys... ys)
{
    static_assert(sizeof...(Ys) > 0, "not enough argument");
Chao Liu's avatar
Chao Liu committed
141

zjing14's avatar
zjing14 committed
142
    return min(x, min(ys...));
Chao Liu's avatar
Chao Liu committed
143
144
}

145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
// disallow implicit type casting
template <typename T>
__device__ T exp(T x);

template <>
__device__ float exp<float>(float x)
{
    return __expf(x);
}

template <>
__device__ double exp<double>(double x)
{
    return exp(x);
}

161
// greatest common divisor, aka highest common factor
Chao Liu's avatar
Chao Liu committed
162
__host__ __device__ constexpr index_t gcd(index_t x, index_t y)
Chao Liu's avatar
Chao Liu committed
163
{
164
165
166
167
168
169
170
171
172
    if(x < 0)
    {
        return gcd(-x, y);
    }
    else if(y < 0)
    {
        return gcd(x, -y);
    }
    else if(x == y || x == 0)
Chao Liu's avatar
Chao Liu committed
173
174
175
    {
        return y;
    }
Chao Liu's avatar
Chao Liu committed
176
    else if(y == 0)
Chao Liu's avatar
Chao Liu committed
177
178
179
    {
        return x;
    }
Chao Liu's avatar
Chao Liu committed
180
    else if(x > y)
Chao Liu's avatar
Chao Liu committed
181
    {
182
        return gcd(x % y, y);
Chao Liu's avatar
Chao Liu committed
183
    }
Chao Liu's avatar
Chao Liu committed
184
    else
Chao Liu's avatar
Chao Liu committed
185
    {
186
        return gcd(x, y % x);
Chao Liu's avatar
Chao Liu committed
187
188
189
190
    }
}

template <index_t X, index_t Y>
191
__host__ __device__ constexpr auto gcd(Number<X>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
192
{
Chao Liu's avatar
Chao Liu committed
193
194
195
    constexpr auto r = gcd(X, Y);

    return Number<r>{};
Chao Liu's avatar
Chao Liu committed
196
197
}

Chao Liu's avatar
Chao Liu committed
198
template <typename X, typename... Ys, typename enable_if<sizeof...(Ys) >= 2, bool>::type = false>
199
__host__ __device__ constexpr auto gcd(X x, Ys... ys)
Chao Liu's avatar
Chao Liu committed
200
{
201
    return gcd(x, gcd(ys...));
Chao Liu's avatar
Chao Liu committed
202
203
204
}

// least common multiple
Chao Liu's avatar
Chao Liu committed
205
206
template <typename X, typename Y>
__host__ __device__ constexpr auto lcm(X x, Y y)
Chao Liu's avatar
Chao Liu committed
207
{
208
    return (x * y) / gcd(x, y);
Chao Liu's avatar
Chao Liu committed
209
210
}

Chao Liu's avatar
Chao Liu committed
211
template <typename X, typename... Ys, typename enable_if<sizeof...(Ys) >= 2, bool>::type = false>
Chao Liu's avatar
Chao Liu committed
212
__host__ __device__ constexpr auto lcm(X x, Ys... ys)
Chao Liu's avatar
Chao Liu committed
213
{
Chao Liu's avatar
Chao Liu committed
214
    return lcm(x, lcm(ys...));
Chao Liu's avatar
Chao Liu committed
215
216
}

zjing14's avatar
zjing14 committed
217
template <typename T>
Chao Liu's avatar
Chao Liu committed
218
219
220
221
222
struct equal
{
    __host__ __device__ constexpr bool operator()(T x, T y) const { return x == y; }
};

zjing14's avatar
zjing14 committed
223
template <typename T>
Chao Liu's avatar
Chao Liu committed
224
225
226
227
228
struct less
{
    __host__ __device__ constexpr bool operator()(T x, T y) const { return x < y; }
};

Chao Liu's avatar
Chao Liu committed
229
} // namespace math
Chao Liu's avatar
Chao Liu committed
230
} // namespace ck
Chao Liu's avatar
Chao Liu committed
231
232

#endif