statically_indexed_array.hpp 2.35 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
4
5
6
7
8
9
10
#ifndef CK_STATICALLY_INDEXED_ARRAY_HPP
#define CK_STATICALLY_INDEXED_ARRAY_HPP

#include "functional2.hpp"
#include "sequence.hpp"
#include "tuple.hpp"

namespace ck {

namespace detail {
11
12
template <typename X, typename Y>
struct tuple_concat;
Chao Liu's avatar
Chao Liu committed
13

14
15
template <typename... Xs, typename... Ys>
struct tuple_concat<Tuple<Xs...>, Tuple<Ys...>>
Chao Liu's avatar
Chao Liu committed
16
{
17
18
    using type = Tuple<Xs..., Ys...>;
};
Chao Liu's avatar
Chao Liu committed
19

20
21
22
23
24
25
26
template <typename T, index_t N>
struct StaticallyIndexedArrayImpl
{
    using type =
        typename tuple_concat<typename StaticallyIndexedArrayImpl<T, N / 2>::type,
                              typename StaticallyIndexedArrayImpl<T, N - N / 2>::type>::type;
};
Chao Liu's avatar
Chao Liu committed
27

28
29
30
31
32
33
34
35
36
37
38
template <typename T>
struct StaticallyIndexedArrayImpl<T, 0>
{
    using type = Tuple<>;
};

template <typename T>
struct StaticallyIndexedArrayImpl<T, 1>
{
    using type = Tuple<T>;
};
Chao Liu's avatar
Chao Liu committed
39
40
} // namespace detail

41
42
template <typename T, index_t N>
using StaticallyIndexedArray = typename detail::StaticallyIndexedArrayImpl<T, N>::type;
Chao Liu's avatar
Chao Liu committed
43
44
45
46
47
48
49
50
51
52
53
54
55
56

template <typename X, typename... Xs>
__host__ __device__ constexpr auto make_statically_indexed_array(const X& x, const Xs&... xs)
{
    return StaticallyIndexedArray<X, sizeof...(Xs) + 1>(x, static_cast<X>(xs)...);
}

// make empty StaticallyIndexedArray
template <typename X>
__host__ __device__ constexpr auto make_statically_indexed_array()
{
    return StaticallyIndexedArray<X, 0>();
}

57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
template <typename T, index_t N>
struct StaticallyIndexedArray_v2
{
    __host__ __device__ constexpr StaticallyIndexedArray_v2() = default;

    __host__ __device__ static constexpr index_t Size() { return N; }

    // read access
    template <index_t I>
    __host__ __device__ constexpr const auto& At(Number<I>) const
    {
        static_assert(I < N, "wrong! out of range");

        return data_[I];
    }

    // write access
    template <index_t I>
    __host__ __device__ constexpr auto& At(Number<I>)
    {
        static_assert(I < N, "wrong! out of range");

        return data_[I];
    }

    // read access
    template <index_t I>
    __host__ __device__ constexpr const auto& operator[](Number<I> i) const
    {
        return At(i);
    }

    // write access
    template <index_t I>
    __host__ __device__ constexpr auto& operator()(Number<I> i)
    {
        return At(i);
    }

    __host__ __device__ static constexpr bool IsStaticBuffer() { return true; }

    T data_[N];
};

Chao Liu's avatar
Chao Liu committed
101
102
} // namespace ck
#endif