math.hpp 5.36 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
// SPDX-License-Identifier: MIT
Illia Silin's avatar
Illia Silin committed
2
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
Chao Liu's avatar
Chao Liu committed
3

Chao Liu's avatar
Chao Liu committed
4
#pragma once
Chao Liu's avatar
Chao Liu committed
5

Chao Liu's avatar
Chao Liu committed
6
#include "ck/ck.hpp"
Chao Liu's avatar
Chao Liu committed
7
#include "integral_constant.hpp"
Chao Liu's avatar
Chao Liu committed
8
#include "number.hpp"
Chao Liu's avatar
Chao Liu committed
9
#include "type.hpp"
Umang Yadav's avatar
Umang Yadav committed
10
#include "enable_if.hpp"
Chao Liu's avatar
Chao Liu committed
11
12
13
14

namespace ck {
namespace math {

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

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

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

Umang Yadav's avatar
Umang Yadav committed
33
34
35
36
37
38
39
struct multiplies
{
    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
40
41
};

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

Umang Yadav's avatar
Umang Yadav committed
48
49
50
51
template <typename 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
52
53
};

Umang Yadav's avatar
Umang Yadav committed
54
55
56
57
58
59
template <typename 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");
Chao Liu's avatar
Chao Liu committed
60

Umang Yadav's avatar
Umang Yadav committed
61
62
        return (a + b - Number<1>{}) / b;
    }
Chao Liu's avatar
Chao Liu committed
63
64
};

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

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

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

Umang Yadav's avatar
Umang Yadav committed
83
84
85
86
87
template <typename T>
__host__ __device__ constexpr T max(T x)
{
    return x;
}
Chao Liu's avatar
Chao Liu committed
88

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

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

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

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

Umang Yadav's avatar
Umang Yadav committed
112
    return max(x, max(ys...));
Chao Liu's avatar
Chao Liu committed
113
114
}

Umang Yadav's avatar
Umang Yadav committed
115
116
117
118
119
template <typename T>
__host__ __device__ constexpr T min(T x)
{
    return x;
}
Chao Liu's avatar
Chao Liu committed
120

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

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

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

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

Umang Yadav's avatar
Umang Yadav committed
144
    return min(x, min(ys...));
Chao Liu's avatar
Chao Liu committed
145
146
}

rocking5566's avatar
rocking5566 committed
147
template <typename T>
Umang Yadav's avatar
Umang Yadav committed
148
149
150
__host__ __device__ constexpr T clamp(const T& x, const T& lowerbound, const T& upperbound)
{
    return min(max(x, lowerbound), upperbound);
rocking5566's avatar
rocking5566 committed
151
152
}

153
// disallow implicit type casting
Umang Yadav's avatar
Umang Yadav committed
154
155
template <typename T>
__device__ T exp(T x);
156

157
158
// TODO: add f16 support using v_exp_f16

Umang Yadav's avatar
Umang Yadav committed
159
160
161
162
163
template <>
__device__ float exp<float>(float x)
{
    return __expf(x);
}
164

Umang Yadav's avatar
Umang Yadav committed
165
166
167
168
169
template <>
__device__ double exp<double>(double x)
{
    return exp(x);
}
170

Umang Yadav's avatar
Umang Yadav committed
171
static inline __host__ float exp(float x) { return ::expf(x); }
172

Umang Yadav's avatar
Umang Yadav committed
173
static inline __host__ double exp(double x) { return std::exp(x); }
174

175
// greatest common divisor, aka highest common factor
Umang Yadav's avatar
Umang Yadav committed
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
__host__ __device__ constexpr index_t gcd(index_t x, index_t y)
{
    if(x < 0)
    {
        return gcd(-x, y);
    }
    else if(y < 0)
    {
        return gcd(x, -y);
    }
    else if(x == y || x == 0)
    {
        return y;
    }
    else if(y == 0)
    {
        return x;
    }
    else if(x > y)
    {
        return gcd(x % y, y);
    }
    else
    {
        return gcd(x, y % x);
    }
Chao Liu's avatar
Chao Liu committed
202
203
204
}

template <index_t X, index_t Y>
Umang Yadav's avatar
Umang Yadav committed
205
206
207
__host__ __device__ constexpr auto gcd(Number<X>, Number<Y>)
{
    constexpr auto r = gcd(X, Y);
Chao Liu's avatar
Chao Liu committed
208

Umang Yadav's avatar
Umang Yadav committed
209
    return Number<r>{};
Chao Liu's avatar
Chao Liu committed
210
211
}

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

// least common multiple
Chao Liu's avatar
Chao Liu committed
219
template <typename X, typename Y>
Umang Yadav's avatar
Umang Yadav committed
220
221
222
__host__ __device__ constexpr auto lcm(X x, Y y)
{
    return (x * y) / gcd(x, y);
Chao Liu's avatar
Chao Liu committed
223
224
}

Umang Yadav's avatar
Umang Yadav committed
225
226
227
228
template <typename X, typename... Ys, typename enable_if<sizeof...(Ys) >= 2, bool>::type = false>
__host__ __device__ constexpr auto lcm(X x, Ys... ys)
{
    return lcm(x, lcm(ys...));
Chao Liu's avatar
Chao Liu committed
229
230
}

Umang Yadav's avatar
Umang Yadav committed
231
232
233
234
template <typename T>
struct equal
{
    __host__ __device__ constexpr bool operator()(T x, T y) const { return x == y; }
Chao Liu's avatar
Chao Liu committed
235
236
};

Umang Yadav's avatar
Umang Yadav committed
237
238
239
240
template <typename T>
struct less
{
    __host__ __device__ constexpr bool operator()(T x, T y) const { return x < y; }
Chao Liu's avatar
Chao Liu committed
241
242
};

243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
template <index_t X>
__host__ __device__ constexpr auto next_power_of_two()
{
    // TODO: X need to be 2 ~ 0x7fffffff. 0, 1, or larger than 0x7fffffff will compile fail
    constexpr index_t Y = 1 << (32 - __builtin_clz(X - 1));
    return Y;
}

template <index_t X>
__host__ __device__ constexpr auto next_power_of_two(Number<X> x)
{
    // TODO: X need to be 2 ~ 0x7fffffff. 0, 1, or larger than 0x7fffffff will compile fail
    constexpr index_t Y = 1 << (32 - __builtin_clz(x.value - 1));
    return Number<Y>{};
}

Chao Liu's avatar
Chao Liu committed
259
} // namespace math
Chao Liu's avatar
Chao Liu committed
260
} // namespace ck