Files
composable_kernel/composable_kernel/include/utility/math.hpp
Chao Liu 611944080a refactor
[ROCm/composable_kernel commit: 21f7e9f103]
2019-06-19 17:43:56 -05:00

108 lines
2.1 KiB
C++

#ifndef CK_MATH_HPP
#define CK_MATH_HPP
#include "config.hpp"
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; }
};
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;
}
};
template <class T>
__host__ __device__ constexpr T integer_divide_ceil(T a, T b)
{
static_assert(is_same<T, index_t>{} || is_same<T, int>{}, "wrong type");
return (a + b - 1) / b;
}
template <class T>
__host__ __device__ constexpr T integer_least_multiple(T a, T b)
{
static_assert(is_same<T, index_t>{} || is_same<T, int>{}, "wrong type");
return b * integer_divide_ceil(a, b);
}
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;
}
// this is WRONG
// TODO: implement least common multiple properly, instead of calling max()
template <class T, class... Ts>
__host__ __device__ constexpr T lcm(T x, Ts... xs)
{
return max(x, xs...);
}
} // namespace math
} // namspace ck
#endif