Array.hpp 11.1 KB
Newer Older
1
2
3
#ifndef CK_ARRAY_HPP
#define CK_ARRAY_HPP

Chao Liu's avatar
Chao Liu committed
4
5
#include "Sequence.hpp"
#include "functional2.hpp"
6

7
8
namespace ck {

Chao Liu's avatar
Chao Liu committed
9
template <typename TData, index_t NSize>
10
11
struct Array
{
Chao Liu's avatar
Chao Liu committed
12
    using type      = Array<TData, NSize>;
Chao Liu's avatar
Chao Liu committed
13
    using data_type = TData;
14

Chao Liu's avatar
Chao Liu committed
15
    index_t mData[NSize];
16

Chao Liu's avatar
Chao Liu committed
17
    __host__ __device__ explicit constexpr Array() {}
18

Chao Liu's avatar
Chao Liu committed
19
20
21
    template <typename X, typename... Xs>
    __host__ __device__ explicit constexpr Array(X x, Xs... xs)
        : mData{static_cast<TData>(x), static_cast<TData>(xs)...}
22
    {
Chao Liu's avatar
Chao Liu committed
23
        static_assert(sizeof...(Xs) + 1 == NSize, "wrong! size");
24
25
    }

Chao Liu's avatar
Chao Liu committed
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
#if 0
    template <typename T>
    __host__ __device__ explicit constexpr Array(const T& x)
    {
        static_assert(T::Size() == NSize, "wrong! size");

        static_for<0, NSize, 1>{}([&](auto i){
           mData[i] = x.At(i); 
        })
    }
#endif

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

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

Chao Liu's avatar
Chao Liu committed
42
    template <index_t I>
Chao Liu's avatar
Chao Liu committed
43
    __host__ __device__ constexpr const TData& At(Number<I>) const
Chao Liu's avatar
Chao Liu committed
44
    {
Chao Liu's avatar
Chao Liu committed
45
46
        static_assert(I < NSize, "wrong!");

Chao Liu's avatar
Chao Liu committed
47
48
49
50
        return mData[I];
    }

    template <index_t I>
Chao Liu's avatar
Chao Liu committed
51
    __host__ __device__ constexpr TData& At(Number<I>)
Chao Liu's avatar
Chao Liu committed
52
    {
Chao Liu's avatar
Chao Liu committed
53
54
        static_assert(I < NSize, "wrong!");

Chao Liu's avatar
Chao Liu committed
55
56
57
        return mData[I];
    }

Chao Liu's avatar
Chao Liu committed
58
    __host__ __device__ constexpr const TData& At(index_t i) const { return mData[i]; }
59

Chao Liu's avatar
Chao Liu committed
60
61
62
63
    __host__ __device__ constexpr TData& At(index_t i) { return mData[i]; }

    template <typename I>
    __host__ __device__ constexpr const TData& operator[](I i) const
Chao Liu's avatar
Chao Liu committed
64
    {
Chao Liu's avatar
Chao Liu committed
65
66
        return At(i);
    }
Chao Liu's avatar
Chao Liu committed
67

Chao Liu's avatar
Chao Liu committed
68
69
70
71
    template <typename I>
    __host__ __device__ constexpr TData& operator()(I i)
    {
        return At(i);
Chao Liu's avatar
Chao Liu committed
72
73
    }

Chao Liu's avatar
Chao Liu committed
74
75
76
77
78
79
80
    template <typename T>
    __host__ __device__ constexpr type& operator=(const T& x)
    {
        static_for<0, Size(), 1>{}([&](auto i) { operator()(i) = x[i]; });

        return *this;
    }
Chao Liu's avatar
Chao Liu committed
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

    struct lambda_PushBack // emulate constexpr lambda
    {
        const Array<TData, NSize>& old_array;
        Array<TData, NSize + 1>& new_array;

        __host__ __device__ constexpr lambda_PushBack(const Array<TData, NSize>& old_array_,
                                                      Array<TData, NSize + 1>& new_array_)
            : old_array(old_array_), new_array(new_array_)
        {
        }

        template <index_t I>
        __host__ __device__ constexpr void operator()(Number<I>) const
        {
Chao Liu's avatar
Chao Liu committed
96
            new_array(Number<I>{}) = old_array[I];
Chao Liu's avatar
Chao Liu committed
97
98
99
        }
    };

Chao Liu's avatar
Chao Liu committed
100
    __host__ __device__ constexpr auto PushBack(TData x) const
101
102
103
    {
        Array<TData, NSize + 1> new_array;

Chao Liu's avatar
Chao Liu committed
104
        static_for<0, NSize, 1>{}(lambda_PushBack(*this, new_array));
105

Chao Liu's avatar
Chao Liu committed
106
        new_array(Number<NSize>{}) = x;
107
108
109

        return new_array;
    }
110
};
111

Chao Liu's avatar
Chao Liu committed
112
// Arr: Array
Chao Liu's avatar
Chao Liu committed
113
// Picks: Sequence<...>
Chao Liu's avatar
Chao Liu committed
114
template <typename Arr, typename Picks>
Chao Liu's avatar
Chao Liu committed
115
struct ArrayElementPicker
Chao Liu's avatar
Chao Liu committed
116
{
Chao Liu's avatar
Chao Liu committed
117
    using type      = ArrayElementPicker;
Chao Liu's avatar
Chao Liu committed
118
119
    using data_type = typename Arr::data_type;

Chao Liu's avatar
Chao Liu committed
120
121
122
    __host__ __device__ constexpr ArrayElementPicker() = delete;

    __host__ __device__ explicit constexpr ArrayElementPicker(Arr& array) : mArray{array}
Chao Liu's avatar
Chao Liu committed
123
124
125
126
    {
        constexpr index_t imax =
            accumulate_on_sequence(Picks{}, math::maxer<index_t>{}, Number<0>{});

Chao Liu's avatar
Chao Liu committed
127
        static_assert(imax < Arr::Size(), "wrong! exceeding # array element");
Chao Liu's avatar
Chao Liu committed
128
129
    }

Chao Liu's avatar
Chao Liu committed
130
    __host__ __device__ static constexpr auto Size() { return Picks::Size(); }
Chao Liu's avatar
Chao Liu committed
131
132

    template <index_t I>
Chao Liu's avatar
Chao Liu committed
133
    __host__ __device__ constexpr const data_type& At(Number<I>) const
Chao Liu's avatar
Chao Liu committed
134
    {
Chao Liu's avatar
Chao Liu committed
135
136
137
138
        static_assert(I < Size(), "wrong!");

        constexpr auto IP = Picks{}[I];
        return mArray[IP];
Chao Liu's avatar
Chao Liu committed
139
140
    }

Chao Liu's avatar
Chao Liu committed
141
142
    template <index_t I>
    __host__ __device__ constexpr data_type& At(Number<I>)
Chao Liu's avatar
Chao Liu committed
143
    {
Chao Liu's avatar
Chao Liu committed
144
145
146
147
        static_assert(I < Size(), "wrong!");

        constexpr auto IP = Picks{}[I];
        return mArray(IP);
Chao Liu's avatar
Chao Liu committed
148
149
    }

Chao Liu's avatar
Chao Liu committed
150
151
    template <typename I>
    __host__ __device__ constexpr const data_type& operator[](I i) const
Chao Liu's avatar
Chao Liu committed
152
    {
Chao Liu's avatar
Chao Liu committed
153
        return At(i);
Chao Liu's avatar
Chao Liu committed
154
155
    }

Chao Liu's avatar
Chao Liu committed
156
157
    template <typename I>
    __host__ __device__ constexpr data_type& operator()(I i)
Chao Liu's avatar
Chao Liu committed
158
    {
Chao Liu's avatar
Chao Liu committed
159
        return At(i);
Chao Liu's avatar
Chao Liu committed
160
161
    }

Chao Liu's avatar
Chao Liu committed
162
163
164
165
166
167
168
169
170
    template <typename T>
    __host__ __device__ constexpr type& operator=(const T& a)
    {
        static_for<0, Size(), 1>{}([&](auto i) { operator()(i) = a[i]; });

        return *this;
    }

    Arr& mArray;
Chao Liu's avatar
Chao Liu committed
171
172
};

Chao Liu's avatar
Chao Liu committed
173
template <typename Arr, typename Picks>
Chao Liu's avatar
Chao Liu committed
174
175
176
177
178
__host__ __device__ constexpr auto pick_array_element(Arr& a, Picks)
{
    return ArrayElementPicker<Arr, Picks>(a);
}

Chao Liu's avatar
Chao Liu committed
179
180
181
182
183
184
185
186
187
188
189
190
#if 1
template <typename T>
__host__ __device__ constexpr auto to_array(const T& x)
{
    Array<typename T::data_type, T::Size()> y;

    static_for<0, T::Size(), 1>{}([&](auto i) { y.At(i) = x.At(i); });

    return y;
}
#endif

Chao Liu's avatar
Chao Liu committed
191
192
193
194
195
196
template <index_t... Is>
__host__ __device__ constexpr auto sequence2array(Sequence<Is...>)
{
    return Array<index_t, sizeof...(Is)>{Is...};
}

Chao Liu's avatar
Chao Liu committed
197
template <typename TData, index_t NSize>
Chao Liu's avatar
Chao Liu committed
198
199
__host__ __device__ constexpr auto make_zero_array()
{
Chao Liu's avatar
Chao Liu committed
200
    constexpr auto zero_sequence = typename uniform_sequence_gen<NSize, 0>::type{};
Chao Liu's avatar
Chao Liu committed
201
202
    constexpr auto zero_array    = sequence2array(zero_sequence);
    return zero_array;
Chao Liu's avatar
Chao Liu committed
203
204
}

Chao Liu's avatar
Chao Liu committed
205
template <typename TData, index_t NSize, index_t... IRs>
Chao Liu's avatar
Chao Liu committed
206
__host__ __device__ constexpr auto reorder_array_given_new2old(const Array<TData, NSize>& old_array,
Chao Liu's avatar
Chao Liu committed
207
                                                               Sequence<IRs...> /*new2old*/)
208
209
210
{
    static_assert(NSize == sizeof...(IRs), "NSize not consistent");

Chao Liu's avatar
Chao Liu committed
211
    static_assert(is_valid_sequence_map<Sequence<IRs...>>{}, "wrong! invalid reorder map");
212

213
    return Array<TData, NSize>{old_array[IRs]...};
214
215
}

Chao Liu's avatar
Chao Liu committed
216
template <typename TData, index_t NSize, typename MapOld2New>
Chao Liu's avatar
Chao Liu committed
217
struct lambda_reorder_array_given_old2new
Chao Liu's avatar
Chao Liu committed
218
{
Chao Liu's avatar
Chao Liu committed
219
220
    const Array<TData, NSize>& old_array;
    Array<TData, NSize>& new_array;
Chao Liu's avatar
Chao Liu committed
221

Chao Liu's avatar
Chao Liu committed
222
223
224
    __host__ __device__ constexpr lambda_reorder_array_given_old2new(
        const Array<TData, NSize>& old_array_, Array<TData, NSize>& new_array_)
        : old_array(old_array_), new_array(new_array_)
Chao Liu's avatar
Chao Liu committed
225
226
227
228
229
230
    {
    }

    template <index_t IOldDim>
    __host__ __device__ constexpr void operator()(Number<IOldDim>) const
    {
Chao Liu's avatar
Chao Liu committed
231
        TData old_data = old_array[IOldDim];
Chao Liu's avatar
Chao Liu committed
232

Chao Liu's avatar
Chao Liu committed
233
        constexpr index_t INewDim = MapOld2New::At(Number<IOldDim>{});
Chao Liu's avatar
Chao Liu committed
234

Chao Liu's avatar
Chao Liu committed
235
        new_array(Number<INewDim>{}) = old_data;
Chao Liu's avatar
Chao Liu committed
236
237
238
    }
};

Chao Liu's avatar
Chao Liu committed
239
template <typename TData, index_t NSize, index_t... IRs>
Chao Liu's avatar
Chao Liu committed
240
__host__ __device__ constexpr auto reorder_array_given_old2new(const Array<TData, NSize>& old_array,
Chao Liu's avatar
Chao Liu committed
241
                                                               Sequence<IRs...> /*old2new*/)
Chao Liu's avatar
Chao Liu committed
242
243
244
245
246
{
    Array<TData, NSize> new_array;

    static_assert(NSize == sizeof...(IRs), "NSize not consistent");

Chao Liu's avatar
Chao Liu committed
247
248
    static_assert(is_valid_sequence_map<Sequence<IRs...>>::value, "wrong! invalid reorder map");

Chao Liu's avatar
Chao Liu committed
249
    static_for<0, NSize, 1>{}(
Chao Liu's avatar
Chao Liu committed
250
        lambda_reorder_array_given_old2new<TData, NSize, Sequence<IRs...>>(old_array, new_array));
Chao Liu's avatar
Chao Liu committed
251
252
253

    return new_array;
}
Chao Liu's avatar
Chao Liu committed
254

Chao Liu's avatar
Chao Liu committed
255
template <typename TData, index_t NSize, typename ExtractSeq>
Chao Liu's avatar
Chao Liu committed
256
__host__ __device__ constexpr auto extract_array(const Array<TData, NSize>& old_array, ExtractSeq)
257
258
259
260
261
262
263
{
    Array<TData, ExtractSeq::GetSize()> new_array;

    constexpr index_t new_size = ExtractSeq::GetSize();

    static_assert(new_size <= NSize, "wrong! too many extract");

Chao Liu's avatar
Chao Liu committed
264
    static_for<0, new_size, 1>{}([&](auto I) { new_array(I) = old_array[ExtractSeq::At(I)]; });
265
266
267
268

    return new_array;
}

Chao Liu's avatar
Chao Liu committed
269
270
template <typename F, typename X, typename Y, typename Z> // emulate constepxr lambda for array
// math
Chao Liu's avatar
Chao Liu committed
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
struct lambda_array_math
{
    const F& f;
    const X& x;
    const Y& y;
    Z& z;

    __host__ __device__ constexpr lambda_array_math(const F& f_, const X& x_, const Y& y_, Z& z_)
        : f(f_), x(x_), y(y_), z(z_)
    {
    }

    template <index_t IDim_>
    __host__ __device__ constexpr void operator()(Number<IDim_>) const
    {
        constexpr auto IDim = Number<IDim_>{};
Chao Liu's avatar
Chao Liu committed
287
        z(IDim)             = f(x[IDim], y[IDim]);
Chao Liu's avatar
Chao Liu committed
288
289
290
    }
};

291
// Array = Array + Array
Chao Liu's avatar
Chao Liu committed
292
template <typename TData, index_t NSize>
Chao Liu's avatar
Chao Liu committed
293
__host__ __device__ constexpr auto operator+(Array<TData, NSize> a, Array<TData, NSize> b)
Chao Liu's avatar
Chao Liu committed
294
295
296
{
    Array<TData, NSize> result;

297
    auto f = math::plus<index_t>{};
Chao Liu's avatar
Chao Liu committed
298

Chao Liu's avatar
Chao Liu committed
299
300
301
    static_for<0, NSize, 1>{}(
        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
            f, a, b, result));
Chao Liu's avatar
Chao Liu committed
302
303
304

    return result;
}
Chao Liu's avatar
Chao Liu committed
305

306
// Array = Array - Array
Chao Liu's avatar
Chao Liu committed
307
template <typename TData, index_t NSize>
308
309
310
311
__host__ __device__ constexpr auto operator-(Array<TData, NSize> a, Array<TData, NSize> b)
{
    Array<TData, NSize> result;

312
    auto f = math::minus<index_t>{};
313

Chao Liu's avatar
Chao Liu committed
314
315
316
    static_for<0, NSize, 1>{}(
        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
            f, a, b, result));
317
318
319
320

    return result;
}

Chao Liu's avatar
Chao Liu committed
321
// Array += Array
Chao Liu's avatar
Chao Liu committed
322
template <typename TData, index_t NSize>
Chao Liu's avatar
Chao Liu committed
323
324
325
326
327
328
329
__host__ __device__ constexpr auto operator+=(Array<TData, NSize>& a, Array<TData, NSize> b)
{
    a = a + b;
    return a;
}

// Array -= Array
Chao Liu's avatar
Chao Liu committed
330
template <typename TData, index_t NSize>
Chao Liu's avatar
Chao Liu committed
331
332
333
334
335
__host__ __device__ constexpr auto operator-=(Array<TData, NSize>& a, Array<TData, NSize> b)
{
    a = a - b;
    return a;
}
336
// Array = Array + Sequence
Chao Liu's avatar
Chao Liu committed
337
template <typename TData, index_t NSize, index_t... Is>
338
339
340
341
342
343
__host__ __device__ constexpr auto operator+(Array<TData, NSize> a, Sequence<Is...> b)
{
    static_assert(sizeof...(Is) == NSize, "wrong! size not the same");

    Array<TData, NSize> result;

344
    auto f = math::plus<index_t>{};
345

Chao Liu's avatar
Chao Liu committed
346
347
348
    static_for<0, NSize, 1>{}(
        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
            f, a, b, result));
349
350
351
352
353

    return result;
}

// Array = Array - Sequence
Chao Liu's avatar
Chao Liu committed
354
template <typename TData, index_t NSize, index_t... Is>
355
356
357
358
359
360
__host__ __device__ constexpr auto operator-(Array<TData, NSize> a, Sequence<Is...> b)
{
    static_assert(sizeof...(Is) == NSize, "wrong! size not the same");

    Array<TData, NSize> result;

361
    auto f = math::minus<index_t>{};
362

Chao Liu's avatar
Chao Liu committed
363
364
365
    static_for<0, NSize, 1>{}(
        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
            f, a, b, result));
366
367
368
369

    return result;
}

Chao Liu's avatar
Chao Liu committed
370
// Array = Array * Sequence
Chao Liu's avatar
Chao Liu committed
371
template <typename TData, index_t NSize, index_t... Is>
Chao Liu's avatar
Chao Liu committed
372
373
374
375
376
377
__host__ __device__ constexpr auto operator*(Array<TData, NSize> a, Sequence<Is...> b)
{
    static_assert(sizeof...(Is) == NSize, "wrong! size not the same");

    Array<TData, NSize> result;

378
    auto f = math::multiplies<index_t>{};
Chao Liu's avatar
Chao Liu committed
379

Chao Liu's avatar
Chao Liu committed
380
381
382
    static_for<0, NSize, 1>{}(
        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
            f, a, b, result));
Chao Liu's avatar
Chao Liu committed
383
384
385

    return result;
}
386

387
// Array = Sequence - Array
Chao Liu's avatar
Chao Liu committed
388
template <typename TData, index_t NSize, index_t... Is>
389
__host__ __device__ constexpr auto operator-(Sequence<Is...> a, Array<TData, NSize> b)
390
{
391
392
393
    static_assert(sizeof...(Is) == NSize, "wrong! size not the same");

    Array<TData, NSize> result;
394

395
    auto f = math::minus<index_t>{};
396

Chao Liu's avatar
Chao Liu committed
397
398
399
    static_for<0, NSize, 1>{}(
        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
            f, a, b, result));
400
401
402
403

    return result;
}

Chao Liu's avatar
Chao Liu committed
404
template <typename TData, index_t NSize, typename Reduce>
405
406
407
408
409
410
411
__host__ __device__ constexpr TData
accumulate_on_array(const Array<TData, NSize>& a, Reduce f, TData init)
{
    TData result = init;

    static_assert(NSize > 0, "wrong");

Chao Liu's avatar
Chao Liu committed
412
    static_for<0, NSize, 1>{}([&](auto I) { result = f(result, a[I]); });
413
414
415

    return result;
}
416

417
418
} // namespace ck
#endif