math.hpp 3.55 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
46
47
48
49
50
template <class T>
struct maxer
{
    __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
57
58
59
60
61
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;
    }
};

62
63
64
65
66
67
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
68
69
template <class X, class Y>
__host__ __device__ constexpr auto integer_divide_ceil(X x, Y y)
Chao Liu's avatar
Chao Liu committed
70
{
Chao Liu's avatar
Chao Liu committed
71
    return (x + y - 1) / y;
Chao Liu's avatar
Chao Liu committed
72
73
}

Chao Liu's avatar
Chao Liu committed
74
75
template <class X, class Y>
__host__ __device__ constexpr auto integer_least_multiple(X x, Y y)
Chao Liu's avatar
Chao Liu committed
76
{
Chao Liu's avatar
Chao Liu committed
77
    return y * integer_divide_ceil(x, y);
Chao Liu's avatar
Chao Liu committed
78
79
80
81
82
83
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
}

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;
}

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

template <index_t X, index_t Y>
138
__host__ __device__ constexpr auto gcd(Number<X>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
139
{
Chao Liu's avatar
Chao Liu committed
140
141
142
    constexpr auto r = gcd(X, Y);

    return Number<r>{};
Chao Liu's avatar
Chao Liu committed
143
144
}

Chao Liu's avatar
Chao Liu committed
145
146
147
template <typename X,
          typename... Ys,
          typename std::enable_if<sizeof...(Ys) >= 2, bool>::type = false>
148
__host__ __device__ constexpr auto gcd(X x, Ys... ys)
Chao Liu's avatar
Chao Liu committed
149
{
150
    return gcd(x, ys...);
Chao Liu's avatar
Chao Liu committed
151
152
153
}

// least common multiple
Chao Liu's avatar
Chao Liu committed
154
155
template <typename X, typename Y>
__host__ __device__ constexpr auto lcm(X x, Y y)
Chao Liu's avatar
Chao Liu committed
156
{
157
    return (x * y) / gcd(x, y);
Chao Liu's avatar
Chao Liu committed
158
159
}

Chao Liu's avatar
Chao Liu committed
160
161
162
template <typename X,
          typename... Ys,
          typename std::enable_if<sizeof...(Ys) >= 2, bool>::type = false>
Chao Liu's avatar
Chao Liu committed
163
__host__ __device__ constexpr auto lcm(X x, Ys... ys)
Chao Liu's avatar
Chao Liu committed
164
{
Chao Liu's avatar
Chao Liu committed
165
    return lcm(x, lcm(ys...));
Chao Liu's avatar
Chao Liu committed
166
167
}

Chao Liu's avatar
Chao Liu committed
168
169
170
171
172
173
174
175
176
177
178
179
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
180
} // namespace math
Chao Liu's avatar
Chao Liu committed
181
} // namespace ck
Chao Liu's avatar
Chao Liu committed
182
183

#endif