ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-24 23:05:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Re-use already-existing scales<> functor template

Browse Source
This commit is contained in:
Po Yen Chen
2024-04-09 08:06:38 +00:00
parent ad45cf8613
commit 5d0ebdbfe4
4 changed files with 31 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
include/ck_tile/core/numeric/math.hpp
View File

@@ -13,12 +13,37 @@
namespace ck_tile {
template <typename T, T s>
template <typename Scale, Scale lhs>
struct scales_c
{
template <typename Right>
CK_TILE_HOST_DEVICE constexpr auto operator()(const Right& rhs) const -> decltype(lhs * rhs)
{
return lhs * rhs;
}
};
template <typename Scale>
struct scales
{
CK_TILE_HOST_DEVICE constexpr T operator()(T a) const { return s * a; }
static_assert(std::is_copy_constructible_v<Scale>);
CK_TILE_HOST_DEVICE constexpr explicit scales(Scale lhs) : lhs_(lhs) {}
template <typename Right>
CK_TILE_HOST_DEVICE constexpr auto operator()(const Right& rhs) const
-> decltype(std::declval<const Scale&>() * rhs)
{
return lhs_ * rhs;
}
private:
Scale lhs_;
};
template <typename Scale>
__host__ __device__ scales(Scale)->scales<Scale>;
template <typename Left = void, typename Right = Left>
struct plus
{