mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-03-22 08:07:38 +00:00
66 lines
1.9 KiB
C++
66 lines
1.9 KiB
C++
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#ifndef CK_FUNCTIONAL4_HPP
|
|
#define CK_FUNCTIONAL4_HPP
|
|
|
|
#include "sequence.hpp"
|
|
#include "tuple.hpp"
|
|
#include "array.hpp"
|
|
|
|
namespace ck {
|
|
|
|
namespace detail {
|
|
|
|
template <typename Indices>
|
|
struct unpack_impl;
|
|
|
|
template <index_t... Is>
|
|
struct unpack_impl<Sequence<Is...>>
|
|
{
|
|
template <typename F, typename X>
|
|
__host__ __device__ constexpr auto operator()(F&& f, X&& x) const
|
|
{
|
|
return ck::forward<F>(f)(ck::forward<X>(x).At(Number<Is>{})...);
|
|
}
|
|
};
|
|
|
|
template <typename Seq0, typename Seq1>
|
|
struct unpack2_impl;
|
|
|
|
// TODO: remove this, after properly implementing unpack that takes any number of containers
|
|
template <index_t... Is, index_t... Js>
|
|
struct unpack2_impl<Sequence<Is...>, Sequence<Js...>>
|
|
{
|
|
template <typename F, typename X, typename Y>
|
|
__host__ __device__ constexpr auto operator()(F&& f, X&& x, Y&& y) const
|
|
{
|
|
return ck::forward<F>(f)(ck::forward<X>(x).At(Number<Is>{})...,
|
|
ck::forward<Y>(y).At(Number<Js>{})...);
|
|
}
|
|
};
|
|
|
|
} // namespace detail
|
|
|
|
template <typename F, typename X>
|
|
__host__ __device__ constexpr auto unpack(F&& f, X&& x)
|
|
{
|
|
using X_ = remove_reference_t<X>;
|
|
return detail::unpack_impl<typename arithmetic_sequence_gen<0, X_::Size(), 1>::type>{}(
|
|
ck::forward<F>(f), ck::forward<X>(x));
|
|
}
|
|
|
|
// TODO: properly implement unpack that takes any number of containers
|
|
template <typename F, typename X, typename Y>
|
|
__host__ __device__ constexpr auto unpack2(F&& f, X&& x, Y&& y)
|
|
{
|
|
using X_ = remove_reference_t<X>;
|
|
using Y_ = remove_reference_t<Y>;
|
|
return detail::unpack2_impl<typename arithmetic_sequence_gen<0, X_::Size(), 1>::type,
|
|
typename arithmetic_sequence_gen<0, Y_::Size(), 1>::type>{}(
|
|
ck::forward<F>(f), ck::forward<X>(x), ck::forward<Y>(y));
|
|
}
|
|
|
|
} // namespace ck
|
|
#endif
|