"git@developer.sourcefind.cn:modelzoo/alphafold2_jax.git" did not exist on "98caef21efa959e44ed01ab33cfb15ab04a39418"
math.hpp 4.34 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
struct minus
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a - b; }
};

zjing14's avatar
zjing14 committed
31
template <typename T>
Chao Liu's avatar
Chao Liu committed
32
33
34
35
36
struct multiplies
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a * b; }
};

Chao Liu's avatar
Chao Liu committed
37
38
39
40
41
42
43
44
45
struct multiplies_v2
{
    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
46
template <typename T>
Chao Liu's avatar
Chao Liu committed
47
struct maximize
Chao Liu's avatar
Chao Liu committed
48
49
50
51
{
    __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
struct minimize
{
    __host__ __device__ constexpr T operator()(T a, T b) const { return a <= b ? a : b; }
};

zjing14's avatar
zjing14 committed
58
template <typename T>
Chao Liu's avatar
Chao Liu committed
59
60
61
62
63
64
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
65
        return (a + b - Number<1>{}) / b;
Chao Liu's avatar
Chao Liu committed
66
67
68
    }
};

zjing14's avatar
zjing14 committed
69
template <typename X, typename Y>
70
71
72
73
74
__host__ __device__ constexpr auto integer_divide_floor(X x, Y y)
{
    return x / y;
}

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

zjing14's avatar
zjing14 committed
81
template <typename X, typename Y>
Chao Liu's avatar
Chao Liu committed
82
__host__ __device__ constexpr auto integer_least_multiple(X x, Y y)
Chao Liu's avatar
Chao Liu committed
83
{
Chao Liu's avatar
Chao Liu committed
84
    return y * integer_divide_ceil(x, y);
Chao Liu's avatar
Chao Liu committed
85
86
}

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

zjing14's avatar
zjing14 committed
93
94
template <typename T>
__host__ __device__ constexpr T max(T x, T y)
Chao Liu's avatar
Chao Liu committed
95
{
zjing14's avatar
zjing14 committed
96
97
    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 X>
__host__ __device__ constexpr index_t max(Number<X>, index_t y)
{
    return X > y ? X : y;
}
Chao Liu's avatar
Chao Liu committed
104

zjing14's avatar
zjing14 committed
105
106
107
108
109
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
110

zjing14's avatar
zjing14 committed
111
112
113
114
115
116
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
117
118
}

zjing14's avatar
zjing14 committed
119
template <typename T>
Chao Liu's avatar
Chao Liu committed
120
121
122
123
124
__host__ __device__ constexpr T min(T x)
{
    return x;
}

zjing14's avatar
zjing14 committed
125
126
127
128
129
130
131
132
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
133
{
zjing14's avatar
zjing14 committed
134
135
    return X < y ? X : y;
}
Chao Liu's avatar
Chao Liu committed
136

zjing14's avatar
zjing14 committed
137
138
139
140
141
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
142

zjing14's avatar
zjing14 committed
143
144
145
146
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
147

zjing14's avatar
zjing14 committed
148
    return min(x, min(ys...));
Chao Liu's avatar
Chao Liu committed
149
150
}

151
// greatest common divisor, aka highest common factor
Chao Liu's avatar
Chao Liu committed
152
__host__ __device__ constexpr index_t gcd(index_t x, index_t y)
Chao Liu's avatar
Chao Liu committed
153
{
154
155
156
157
158
159
160
161
162
    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
163
164
165
    {
        return y;
    }
Chao Liu's avatar
Chao Liu committed
166
    else if(y == 0)
Chao Liu's avatar
Chao Liu committed
167
168
169
    {
        return x;
    }
Chao Liu's avatar
Chao Liu committed
170
    else if(x > y)
Chao Liu's avatar
Chao Liu committed
171
    {
172
        return gcd(x % y, y);
Chao Liu's avatar
Chao Liu committed
173
    }
Chao Liu's avatar
Chao Liu committed
174
    else
Chao Liu's avatar
Chao Liu committed
175
    {
176
        return gcd(x, y % x);
Chao Liu's avatar
Chao Liu committed
177
178
179
180
    }
}

template <index_t X, index_t Y>
181
__host__ __device__ constexpr auto gcd(Number<X>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
182
{
Chao Liu's avatar
Chao Liu committed
183
184
185
    constexpr auto r = gcd(X, Y);

    return Number<r>{};
Chao Liu's avatar
Chao Liu committed
186
187
}

Chao Liu's avatar
Chao Liu committed
188
template <typename X, typename... Ys, typename enable_if<sizeof...(Ys) >= 2, bool>::type = false>
189
__host__ __device__ constexpr auto gcd(X x, Ys... ys)
Chao Liu's avatar
Chao Liu committed
190
{
191
    return gcd(x, gcd(ys...));
Chao Liu's avatar
Chao Liu committed
192
193
194
}

// least common multiple
Chao Liu's avatar
Chao Liu committed
195
196
template <typename X, typename Y>
__host__ __device__ constexpr auto lcm(X x, Y y)
Chao Liu's avatar
Chao Liu committed
197
{
198
    return (x * y) / gcd(x, y);
Chao Liu's avatar
Chao Liu committed
199
200
}

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

zjing14's avatar
zjing14 committed
207
template <typename T>
Chao Liu's avatar
Chao Liu committed
208
209
210
211
212
struct equal
{
    __host__ __device__ constexpr bool operator()(T x, T y) const { return x == y; }
};

zjing14's avatar
zjing14 committed
213
template <typename T>
Chao Liu's avatar
Chao Liu committed
214
215
216
217
218
struct less
{
    __host__ __device__ constexpr bool operator()(T x, T y) const { return x < y; }
};

Chao Liu's avatar
Chao Liu committed
219
} // namespace math
Chao Liu's avatar
Chao Liu committed
220
} // namespace ck
Chao Liu's avatar
Chao Liu committed
221
222

#endif