math.hpp 3.21 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
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
42
43
44
45
46
47
48
49
50
51
52
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;
    }
};

53
54
55
56
57
58
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
59
60
template <class X, class Y>
__host__ __device__ constexpr auto integer_divide_ceil(X x, Y y)
Chao Liu's avatar
Chao Liu committed
61
{
Chao Liu's avatar
Chao Liu committed
62
    return (x + y - 1) / y;
Chao Liu's avatar
Chao Liu committed
63
64
}

Chao Liu's avatar
Chao Liu committed
65
66
template <class X, class Y>
__host__ __device__ constexpr auto integer_least_multiple(X x, Y y)
Chao Liu's avatar
Chao Liu committed
67
{
Chao Liu's avatar
Chao Liu committed
68
    return y * integer_divide_ceil(x, y);
Chao Liu's avatar
Chao Liu committed
69
70
71
72
73
74
75
76
77
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
}

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

107
// greatest common divisor, aka highest common factor
Chao Liu's avatar
Chao Liu committed
108
template <typename T>
109
__host__ __device__ constexpr T gcd(T x, T y)
Chao Liu's avatar
Chao Liu committed
110
{
Chao Liu's avatar
Chao Liu committed
111
    if(x == y || x == 0)
Chao Liu's avatar
Chao Liu committed
112
113
114
    {
        return y;
    }
Chao Liu's avatar
Chao Liu committed
115
    else if(y == 0)
Chao Liu's avatar
Chao Liu committed
116
117
118
    {
        return x;
    }
Chao Liu's avatar
Chao Liu committed
119
    else if(x > y)
Chao Liu's avatar
Chao Liu committed
120
    {
Chao Liu's avatar
Chao Liu committed
121
        return gcd(x - y, y);
Chao Liu's avatar
Chao Liu committed
122
    }
Chao Liu's avatar
Chao Liu committed
123
    else
Chao Liu's avatar
Chao Liu committed
124
    {
Chao Liu's avatar
Chao Liu committed
125
        return gcd(x, y - x);
Chao Liu's avatar
Chao Liu committed
126
127
128
129
    }
}

template <index_t X, index_t Y>
130
__host__ __device__ constexpr auto gcd(Number<X>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
131
{
132
    constexpr auto result = gcd(X, Y);
Chao Liu's avatar
Chao Liu committed
133
134
135
136
    return Number<result>{};
}

template <typename X, typename... Ys>
137
__host__ __device__ constexpr auto gcd(X x, Ys... ys)
Chao Liu's avatar
Chao Liu committed
138
{
139
    return gcd(x, ys...);
Chao Liu's avatar
Chao Liu committed
140
141
142
143
144
145
}

// least common multiple
template <typename T>
__host__ __device__ constexpr T lcm(T x, T y)
{
146
    return (x * y) / gcd(x, y);
Chao Liu's avatar
Chao Liu committed
147
148
}

Chao Liu's avatar
Chao Liu committed
149
150
template <typename X, typename... Ys>
__host__ __device__ constexpr auto lcm(X x, Ys... ys)
Chao Liu's avatar
Chao Liu committed
151
{
Chao Liu's avatar
Chao Liu committed
152
    return lcm(x, lcm(ys...));
Chao Liu's avatar
Chao Liu committed
153
154
}

Chao Liu's avatar
Chao Liu committed
155
156
157
158
159
160
161
162
163
164
165
166
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
167
168
169
170
} // namespace math
} // namspace ck

#endif