functional4.hpp 1.8 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
#ifndef CK_FUNCTIONAL4_HPP
#define CK_FUNCTIONAL4_HPP

Chao Liu's avatar
Chao Liu committed
4
#include "sequence.hpp"
Chao Liu's avatar
Chao Liu committed
5
#include "tuple.hpp"
Chao Liu's avatar
Chao Liu committed
6
#include "array.hpp"
Chao Liu's avatar
Chao Liu committed
7
8
9
10
11
12
13
14
15
16
17
18

namespace ck {

namespace detail {

template <typename Indices>
struct unpack_impl;

template <index_t... Is>
struct unpack_impl<Sequence<Is...>>
{
    template <typename F, typename X>
Chao Liu's avatar
Chao Liu committed
19
    __host__ __device__ constexpr auto operator()(F&& f, X&& x) const
Chao Liu's avatar
Chao Liu committed
20
    {
Chao Liu's avatar
Chao Liu committed
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
        return std::forward<F>(f)(std::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 std::forward<F>(f)(std::forward<X>(x).At(Number<Is>{})...,
                                  std::forward<Y>(y).At(Number<Js>{})...);
Chao Liu's avatar
Chao Liu committed
37
38
39
40
41
42
    }
};

} // namespace detail

template <typename F, typename X>
Chao Liu's avatar
Chao Liu committed
43
44
45
46
47
48
49
50
51
52
__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>{}(
        std::forward<F>(f), std::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)
Chao Liu's avatar
Chao Liu committed
53
{
Chao Liu's avatar
Chao Liu committed
54
55
56
57
58
    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>{}(
        std::forward<F>(f), std::forward<X>(x), std::forward<Y>(y));
Chao Liu's avatar
Chao Liu committed
59
60
61
62
}

} // namespace ck
#endif