Sequence.hip.hpp 2.51 KB
Newer Older
1
2
3
4
#pragma once
#include "constant_integral.hip.hpp"
#include "functional.hip.hpp"

Chao Liu's avatar
Chao Liu committed
5
template <index_t... Is>
6
7
8
9
struct Sequence
{
    using Type = Sequence<Is...>;

Chao Liu's avatar
Chao Liu committed
10
    static constexpr index_t nDim = sizeof...(Is);
11

Chao Liu's avatar
Chao Liu committed
12
    const index_t mData[nDim] = {Is...};
13

Chao Liu's avatar
Chao Liu committed
14
15
    template <index_t I>
    __host__ __device__ constexpr index_t Get(Number<I>) const
16
17
18
19
20
    {
        return mData[I];
    }

    // this is ugly, only for nDIm = 4
Chao Liu's avatar
Chao Liu committed
21
    template <index_t I0, index_t I1, index_t I2, index_t I3>
22
23
24
25
26
27
    __host__ __device__ constexpr auto ReorderByGetNewFromOld(Sequence<I0, I1, I2, I3>) const
    {
        static_assert(nDim == 4, "nDim != 4");

        constexpr auto old_sequence = Type{};

Chao Liu's avatar
Chao Liu committed
28
29
30
31
        constexpr index_t NR0 = old_sequence.mData[I0];
        constexpr index_t NR1 = old_sequence.mData[I1];
        constexpr index_t NR2 = old_sequence.mData[I2];
        constexpr index_t NR3 = old_sequence.mData[I3];
32
33
34
35

        return Sequence<NR0, NR1, NR2, NR3>{};
    }

Chao Liu's avatar
Chao Liu committed
36
    template <index_t I0, index_t I1, index_t I2, index_t I3>
37
38
39
40
41
42
43
    __host__ __device__ constexpr auto ReorderByPutOldToNew(Sequence<I0, I1, I2, I3>) const
    {
        // don't know how to implement this
        printf("Sequence::ReorderByPutOldToNew not implemented");
        assert(false);
    }

Chao Liu's avatar
Chao Liu committed
44
    template <index_t I>
45
46
47
48
49
50
51
52
53
54
55
56
57
58
    __host__ __device__ constexpr auto PushBack(Number<I>) const
    {
        return Sequence<Is..., I>{};
    }

    __host__ __device__ constexpr auto PopBack() const;

    template <class F>
    __host__ __device__ constexpr auto Transform(F f) const
    {
        return Sequence<f(Is)...>{};
    }
};

Chao Liu's avatar
Chao Liu committed
59
template <index_t... Is, index_t I>
60
61
62
63
64
65
__host__ __device__ constexpr auto sequence_pop_back(Sequence<Is..., I>)
{
    static_assert(sizeof...(Is) >= 1, "empty Sequence!");
    return Sequence<Is...>{};
}

Chao Liu's avatar
Chao Liu committed
66
template <class F, index_t... Xs, index_t... Ys>
67
68
69
70
71
72
73
__host__ __device__ constexpr auto sequence_sequence_op(Sequence<Xs...>, Sequence<Ys...>, F f)
{
    static_assert(Sequence<Xs...>::nDim == Sequence<Ys...>::nDim, "Dim not the same");

    return Sequence<f(Xs, Ys)...>{};
}

Chao Liu's avatar
Chao Liu committed
74
template <index_t... Xs, index_t... Ys>
75
76
77
78
__host__ __device__ constexpr auto sequence_sequence_add(Sequence<Xs...>, Sequence<Ys...>)
{
    struct add
    {
Chao Liu's avatar
Chao Liu committed
79
        __host__ __device__ constexpr index_t operator()(index_t x, index_t y) const
80
81
82
83
84
85
86
87
        {
            return x + y;
        }
    };

    return sequence_sequence_op(Sequence<Xs...>{}, Sequence<Ys...>{}, add{});
}

Chao Liu's avatar
Chao Liu committed
88
template <index_t... Is>
89
90
91
92
__host__ __device__ constexpr auto Sequence<Is...>::PopBack() const
{
    return sequence_pop_back(Type{});
}