math.hpp 3.69 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
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

namespace ck {
namespace math {

template <class T, T s>
struct scales
{
    __host__ __device__ constexpr T operator()(T a) const { return s * a; }
};

template <class T>
struct plus
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a + b; }
};

template <class T>
struct minus
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a - b; }
};

template <class T>
struct multiplies
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a * b; }
};

Chao Liu's avatar
Chao Liu committed
36
37
38
39
40
41
42
43
44
struct multiplies_v2
{
    template <typename A, typename B>
    __host__ __device__ constexpr auto operator()(const A& a, const B& b) const
    {
        return a * b;
    }
};

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

Chao Liu's avatar
Chao Liu committed
51
52
53
54
55
56
template <class T>
struct minimize
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a <= b ? a : b; }
};

Chao Liu's avatar
Chao Liu committed
57
58
59
60
61
62
63
64
65
66
67
template <class T>
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");

        return (a + b - 1) / b;
    }
};

68
69
70
71
72
73
template <class X, class Y>
__host__ __device__ constexpr auto integer_divide_floor(X x, Y y)
{
    return x / y;
}

Chao Liu's avatar
Chao Liu committed
74
75
template <class X, class Y>
__host__ __device__ constexpr auto integer_divide_ceil(X x, Y y)
Chao Liu's avatar
Chao Liu committed
76
{
Chao Liu's avatar
Chao Liu committed
77
    return (x + y - 1) / y;
Chao Liu's avatar
Chao Liu committed
78
79
}

Chao Liu's avatar
Chao Liu committed
80
81
template <class X, class Y>
__host__ __device__ constexpr auto integer_least_multiple(X x, Y y)
Chao Liu's avatar
Chao Liu committed
82
{
Chao Liu's avatar
Chao Liu committed
83
    return y * integer_divide_ceil(x, y);
Chao Liu's avatar
Chao Liu committed
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
}

template <class T>
__host__ __device__ constexpr T max(T x)
{
    return x;
}

template <class T, class... Ts>
__host__ __device__ constexpr T max(T x, Ts... xs)
{
    static_assert(sizeof...(xs) > 0, "not enough argument");

    auto y = max(xs...);

    static_assert(is_same<decltype(y), T>{}, "not the same type");

    return x > y ? x : y;
}

template <class T>
__host__ __device__ constexpr T min(T x)
{
    return x;
}

template <class T, class... Ts>
__host__ __device__ constexpr T min(T x, Ts... xs)
{
    static_assert(sizeof...(xs) > 0, "not enough argument");

    auto y = min(xs...);

    static_assert(is_same<decltype(y), T>{}, "not the same type");

    return x < y ? x : y;
}

122
// greatest common divisor, aka highest common factor
Chao Liu's avatar
Chao Liu committed
123
__host__ __device__ constexpr index_t gcd(index_t x, index_t y)
Chao Liu's avatar
Chao Liu committed
124
{
Chao Liu's avatar
Chao Liu committed
125
    if(x == y || x == 0)
Chao Liu's avatar
Chao Liu committed
126
127
128
    {
        return y;
    }
Chao Liu's avatar
Chao Liu committed
129
    else if(y == 0)
Chao Liu's avatar
Chao Liu committed
130
131
132
    {
        return x;
    }
Chao Liu's avatar
Chao Liu committed
133
    else if(x > y)
Chao Liu's avatar
Chao Liu committed
134
    {
Chao Liu's avatar
Chao Liu committed
135
        return gcd(x - y, y);
Chao Liu's avatar
Chao Liu committed
136
    }
Chao Liu's avatar
Chao Liu committed
137
    else
Chao Liu's avatar
Chao Liu committed
138
    {
Chao Liu's avatar
Chao Liu committed
139
        return gcd(x, y - x);
Chao Liu's avatar
Chao Liu committed
140
141
142
143
    }
}

template <index_t X, index_t Y>
144
__host__ __device__ constexpr auto gcd(Number<X>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
145
{
Chao Liu's avatar
Chao Liu committed
146
147
148
    constexpr auto r = gcd(X, Y);

    return Number<r>{};
Chao Liu's avatar
Chao Liu committed
149
150
}

Chao Liu's avatar
Chao Liu committed
151
152
153
template <typename X,
          typename... Ys,
          typename std::enable_if<sizeof...(Ys) >= 2, bool>::type = false>
154
__host__ __device__ constexpr auto gcd(X x, Ys... ys)
Chao Liu's avatar
Chao Liu committed
155
{
156
    return gcd(x, ys...);
Chao Liu's avatar
Chao Liu committed
157
158
159
}

// least common multiple
Chao Liu's avatar
Chao Liu committed
160
161
template <typename X, typename Y>
__host__ __device__ constexpr auto lcm(X x, Y y)
Chao Liu's avatar
Chao Liu committed
162
{
163
    return (x * y) / gcd(x, y);
Chao Liu's avatar
Chao Liu committed
164
165
}

Chao Liu's avatar
Chao Liu committed
166
167
168
template <typename X,
          typename... Ys,
          typename std::enable_if<sizeof...(Ys) >= 2, bool>::type = false>
Chao Liu's avatar
Chao Liu committed
169
__host__ __device__ constexpr auto lcm(X x, Ys... ys)
Chao Liu's avatar
Chao Liu committed
170
{
Chao Liu's avatar
Chao Liu committed
171
    return lcm(x, lcm(ys...));
Chao Liu's avatar
Chao Liu committed
172
173
}

Chao Liu's avatar
Chao Liu committed
174
175
176
177
178
179
180
181
182
183
184
185
template <class T>
struct equal
{
    __host__ __device__ constexpr bool operator()(T x, T y) const { return x == y; }
};

template <class T>
struct less
{
    __host__ __device__ constexpr bool operator()(T x, T y) const { return x < y; }
};

Chao Liu's avatar
Chao Liu committed
186
} // namespace math
Chao Liu's avatar
Chao Liu committed
187
} // namespace ck
Chao Liu's avatar
Chao Liu committed
188
189

#endif