sequence.hpp 29 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
// SPDX-License-Identifier: MIT
arai713's avatar
arai713 committed
2
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
Chao Liu's avatar
Chao Liu committed
3

4
#pragma once
5

6
7
8
9
#ifdef _HIPCC_RTC_
#define CK_CODE_GEN_RTC
#endif

10
#ifndef __HIPCC_RTC__
arai713's avatar
arai713 committed
11
#ifndef CK_CODE_GEN_RTC
12
#include <ostream>
13
#endif
14
#endif
15

16
17
18
19
#include "ck/utility/integral_constant.hpp"
#include "ck/utility/type.hpp"
#include "ck/utility/functional.hpp"
#include "ck/utility/math.hpp"
20

21
22
namespace ck {

Chao Liu's avatar
Chao Liu committed
23
24
25
template <index_t, index_t, index_t>
struct static_for;

26
27
28
template <index_t...>
struct Sequence;

Chao Liu's avatar
Chao Liu committed
29
template <typename Seq, index_t I>
30
31
struct sequence_split;

Chao Liu's avatar
Chao Liu committed
32
template <typename>
33
struct sequence_reverse;
Chao Liu's avatar
Chao Liu committed
34

Chao Liu's avatar
Chao Liu committed
35
template <typename>
Chao Liu's avatar
Chao Liu committed
36
37
struct sequence_map_inverse;

Chao Liu's avatar
Chao Liu committed
38
template <typename>
39
40
41
42
43
struct is_valid_sequence_map;

template <index_t I, index_t... Is>
__host__ __device__ constexpr auto sequence_pop_front(Sequence<I, Is...>);

Chao Liu's avatar
Chao Liu committed
44
template <typename Seq>
45
46
__host__ __device__ constexpr auto sequence_pop_back(Seq);

Chao Liu's avatar
Chao Liu committed
47
template <index_t... Is>
48
49
struct Sequence
{
Chao Liu's avatar
Chao Liu committed
50
51
    using Type      = Sequence;
    using data_type = index_t;
52

53
    static constexpr index_t mSize = sizeof...(Is);
54

Chao Liu's avatar
Chao Liu committed
55
    __host__ __device__ static constexpr auto Size() { return Number<mSize>{}; }
56

Chao Liu's avatar
Chao Liu committed
57
58
59
    __host__ __device__ static constexpr auto GetSize() { return Size(); }

    __host__ __device__ static constexpr index_t At(index_t I)
60
    {
Chao Liu's avatar
Chao Liu committed
61
62
63
64
65
66
        // the last dummy element is to prevent compiler complain about empty array, when mSize = 0
        const index_t mData[mSize + 1] = {Is..., 0};
        return mData[I];
    }

    template <index_t I>
Chao Liu's avatar
Chao Liu committed
67
    __host__ __device__ static constexpr auto At(Number<I>)
Chao Liu's avatar
Chao Liu committed
68
    {
Chao Liu's avatar
Chao Liu committed
69
70
        static_assert(I < mSize, "wrong! I too large");

Chao Liu's avatar
Chao Liu committed
71
        return Number<At(I)>{};
Chao Liu's avatar
Chao Liu committed
72
73
    }

Chao Liu's avatar
Chao Liu committed
74
    template <index_t I>
Chao Liu's avatar
Chao Liu committed
75
    __host__ __device__ static constexpr auto Get(Number<I>)
Chao Liu's avatar
Chao Liu committed
76
    {
Chao Liu's avatar
Chao Liu committed
77
        return At(Number<I>{});
78
79
    }

Chao Liu's avatar
Chao Liu committed
80
81
82
83
84
    template <typename I>
    __host__ __device__ constexpr auto operator[](I i) const
    {
        return At(i);
    }
Chao Liu's avatar
Chao Liu committed
85

86
    template <index_t... IRs>
87
    __host__ __device__ static constexpr auto ReorderGivenNew2Old(Sequence<IRs...> /*new2old*/)
88
    {
Chao Liu's avatar
Chao Liu committed
89
        static_assert(sizeof...(Is) == sizeof...(IRs),
Chao Liu's avatar
Chao Liu committed
90
                      "wrong! reorder map should have the same size as Sequence to be rerodered");
Chao Liu's avatar
Chao Liu committed
91

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

Chao Liu's avatar
Chao Liu committed
94
        return Sequence<Type::At(Number<IRs>{})...>{};
95
96
    }

Chao Liu's avatar
Chao Liu committed
97
    // MapOld2New is Sequence<...>
Chao Liu's avatar
Chao Liu committed
98
    template <typename MapOld2New>
Chao Liu's avatar
Chao Liu committed
99
100
    __host__ __device__ static constexpr auto ReorderGivenOld2New(MapOld2New)
    {
Chao Liu's avatar
Chao Liu committed
101
        static_assert(MapOld2New::Size() == Size(),
Chao Liu's avatar
Chao Liu committed
102
103
104
105
106
107
108
                      "wrong! reorder map should have the same size as Sequence to be rerodered");

        static_assert(is_valid_sequence_map<MapOld2New>::value, "wrong! invalid reorder map");

        return ReorderGivenNew2Old(typename sequence_map_inverse<MapOld2New>::type{});
    }

109
110
111
112
    __host__ __device__ static constexpr auto Reverse()
    {
        return typename sequence_reverse<Type>::type{};
    }
Chao Liu's avatar
Chao Liu committed
113

Chao Liu's avatar
Chao Liu committed
114
    __host__ __device__ static constexpr auto Front()
115
    {
Chao Liu's avatar
Chao Liu committed
116
        static_assert(mSize > 0, "wrong!");
Chao Liu's avatar
Chao Liu committed
117
        return At(Number<0>{});
118
    }
119

Chao Liu's avatar
Chao Liu committed
120
    __host__ __device__ static constexpr auto Back()
121
    {
Chao Liu's avatar
Chao Liu committed
122
        static_assert(mSize > 0, "wrong!");
Chao Liu's avatar
Chao Liu committed
123
        return At(Number<mSize - 1>{});
124
    }
125

126
    __host__ __device__ static constexpr auto PopFront() { return sequence_pop_front(Type{}); }
Chao Liu's avatar
Chao Liu committed
127

128
    __host__ __device__ static constexpr auto PopBack() { return sequence_pop_back(Type{}); }
Chao Liu's avatar
Chao Liu committed
129
130
131

    template <index_t... Xs>
    __host__ __device__ static constexpr auto PushFront(Sequence<Xs...>)
132
    {
Chao Liu's avatar
Chao Liu committed
133
        return Sequence<Xs..., Is...>{};
134
135
    }

Chao Liu's avatar
Chao Liu committed
136
137
    template <index_t... Xs>
    __host__ __device__ static constexpr auto PushFront(Number<Xs>...)
138
    {
Chao Liu's avatar
Chao Liu committed
139
        return Sequence<Xs..., Is...>{};
140
141
    }

Chao Liu's avatar
Chao Liu committed
142
143
144
145
146
    template <index_t... Xs>
    __host__ __device__ static constexpr auto PushBack(Sequence<Xs...>)
    {
        return Sequence<Is..., Xs...>{};
    }
147

Chao Liu's avatar
Chao Liu committed
148
    template <index_t... Xs>
Chao Liu's avatar
Chao Liu committed
149
    __host__ __device__ static constexpr auto PushBack(Number<Xs>...)
150
    {
Chao Liu's avatar
Chao Liu committed
151
152
        return Sequence<Is..., Xs...>{};
    }
Chao Liu's avatar
Chao Liu committed
153

Chao Liu's avatar
Chao Liu committed
154
    template <index_t... Ns>
155
    __host__ __device__ static constexpr auto Extract(Number<Ns>...)
Chao Liu's avatar
Chao Liu committed
156
    {
Chao Liu's avatar
Chao Liu committed
157
        return Sequence<Type::At(Number<Ns>{})...>{};
Chao Liu's avatar
Chao Liu committed
158
    }
Chao Liu's avatar
Chao Liu committed
159

Chao Liu's avatar
Chao Liu committed
160
    template <index_t... Ns>
161
    __host__ __device__ static constexpr auto Extract(Sequence<Ns...>)
Chao Liu's avatar
Chao Liu committed
162
    {
Chao Liu's avatar
Chao Liu committed
163
        return Sequence<Type::At(Number<Ns>{})...>{};
Chao Liu's avatar
Chao Liu committed
164
    }
165
166

    template <index_t I, index_t X>
167
168
    __host__ __device__ static constexpr auto Modify(Number<I>, Number<X>)
    {
Chao Liu's avatar
Chao Liu committed
169
        static_assert(I < Size(), "wrong!");
170
171

        using seq_split          = sequence_split<Type, I>;
Chao Liu's avatar
Chao Liu committed
172
173
        constexpr auto seq_left  = typename seq_split::left_type{};
        constexpr auto seq_right = typename seq_split::right_type{}.PopFront();
174
175
176

        return seq_left.PushBack(Number<X>{}).PushBack(seq_right);
    }
Chao Liu's avatar
Chao Liu committed
177

Chao Liu's avatar
Chao Liu committed
178
    template <typename F>
Chao Liu's avatar
Chao Liu committed
179
180
181
182
    __host__ __device__ static constexpr auto Transform(F f)
    {
        return Sequence<f(Is)...>{};
    }
Chao Liu's avatar
Chao Liu committed
183
184
185
186
187
188
189
190

    __host__ __device__ static void Print()
    {
        printf("{");
        printf("size %d, ", index_t{Size()});
        static_for<0, Size(), 1>{}([&](auto i) { printf("%d ", At(i).value); });
        printf("}");
    }
191
192
};

Chao Liu's avatar
Chao Liu committed
193
// merge sequence
Chao Liu's avatar
Chao Liu committed
194
195
196
197
198
template <typename Seq, typename... Seqs>
struct sequence_merge
{
    using type = typename sequence_merge<Seq, typename sequence_merge<Seqs...>::type>::type;
};
Chao Liu's avatar
Chao Liu committed
199

Chao Liu's avatar
Chao Liu committed
200
201
202
template <index_t... Xs, index_t... Ys>
struct sequence_merge<Sequence<Xs...>, Sequence<Ys...>>
{
Chao Liu's avatar
Chao Liu committed
203
    using type = Sequence<Xs..., Ys...>;
Chao Liu's avatar
Chao Liu committed
204
};
Chao Liu's avatar
Chao Liu committed
205

Chao Liu's avatar
Chao Liu committed
206
207
208
209
210
211
template <typename Seq>
struct sequence_merge<Seq>
{
    using type = Seq;
};

Chao Liu's avatar
Chao Liu committed
212
// generate sequence
Chao Liu's avatar
Chao Liu committed
213
214
template <index_t NSize, typename F>
struct sequence_gen
Chao Liu's avatar
Chao Liu committed
215
{
Chao Liu's avatar
Chao Liu committed
216
217
218
219
220
221
    template <index_t IBegin, index_t NRemain, typename G>
    struct sequence_gen_impl
    {
        static constexpr index_t NRemainLeft  = NRemain / 2;
        static constexpr index_t NRemainRight = NRemain - NRemainLeft;
        static constexpr index_t IMiddle      = IBegin + NRemainLeft;
Chao Liu's avatar
Chao Liu committed
222

Chao Liu's avatar
Chao Liu committed
223
224
225
226
        using type = typename sequence_merge<
            typename sequence_gen_impl<IBegin, NRemainLeft, G>::type,
            typename sequence_gen_impl<IMiddle, NRemainRight, G>::type>::type;
    };
Chao Liu's avatar
Chao Liu committed
227

Chao Liu's avatar
Chao Liu committed
228
229
230
231
232
233
    template <index_t I, typename G>
    struct sequence_gen_impl<I, 1, G>
    {
        static constexpr index_t Is = G{}(Number<I>{});
        using type                  = Sequence<Is>;
    };
Chao Liu's avatar
Chao Liu committed
234

Chao Liu's avatar
Chao Liu committed
235
236
237
238
239
    template <index_t I, typename G>
    struct sequence_gen_impl<I, 0, G>
    {
        using type = Sequence<>;
    };
Chao Liu's avatar
Chao Liu committed
240

Chao Liu's avatar
Chao Liu committed
241
242
243
244
    using type = typename sequence_gen_impl<0, NSize, F>::type;
};

// arithmetic sequence
Chao Liu's avatar
Chao Liu committed
245
template <index_t IBegin, index_t IEnd, index_t Increment>
246
struct arithmetic_sequence_gen
Chao Liu's avatar
Chao Liu committed
247
{
Chao Liu's avatar
Chao Liu committed
248
249
250
251
252
253
254
255
    struct F
    {
        __host__ __device__ constexpr index_t operator()(index_t i) const
        {
            return i * Increment + IBegin;
        }
    };

256
257
258
259
260
261
262
    using type0 = typename sequence_gen<(IEnd - IBegin) / Increment, F>::type;
    using type1 = Sequence<>;

    static constexpr bool kHasContent =
        (Increment > 0 && IBegin < IEnd) || (Increment < 0 && IBegin > IEnd);

    using type = typename conditional<kHasContent, type0, type1>::type;
Chao Liu's avatar
Chao Liu committed
263
264
265
266
267
268
};

// uniform sequence
template <index_t NSize, index_t I>
struct uniform_sequence_gen
{
Chao Liu's avatar
Chao Liu committed
269
    struct F
Chao Liu's avatar
Chao Liu committed
270
271
272
273
    {
        __host__ __device__ constexpr index_t operator()(index_t) const { return I; }
    };

Chao Liu's avatar
Chao Liu committed
274
    using type = typename sequence_gen<NSize, F>::type;
Chao Liu's avatar
Chao Liu committed
275
276
277
};

// reverse inclusive scan (with init) sequence
Chao Liu's avatar
Chao Liu committed
278
template <typename, typename, index_t>
Chao Liu's avatar
Chao Liu committed
279
struct sequence_reverse_inclusive_scan;
Chao Liu's avatar
Chao Liu committed
280

Chao Liu's avatar
Chao Liu committed
281
template <index_t I, index_t... Is, typename Reduce, index_t Init>
Chao Liu's avatar
Chao Liu committed
282
struct sequence_reverse_inclusive_scan<Sequence<I, Is...>, Reduce, Init>
Chao Liu's avatar
Chao Liu committed
283
{
Chao Liu's avatar
Chao Liu committed
284
    using old_scan = typename sequence_reverse_inclusive_scan<Sequence<Is...>, Reduce, Init>::type;
Chao Liu's avatar
Chao Liu committed
285
286
287

    static constexpr index_t new_reduce = Reduce{}(I, old_scan{}.Front());

Chao Liu's avatar
Chao Liu committed
288
    using type = typename sequence_merge<Sequence<new_reduce>, old_scan>::type;
Chao Liu's avatar
Chao Liu committed
289
290
};

Chao Liu's avatar
Chao Liu committed
291
template <index_t I, typename Reduce, index_t Init>
Chao Liu's avatar
Chao Liu committed
292
struct sequence_reverse_inclusive_scan<Sequence<I>, Reduce, Init>
Chao Liu's avatar
Chao Liu committed
293
{
Chao Liu's avatar
Chao Liu committed
294
    using type = Sequence<Reduce{}(I, Init)>;
Chao Liu's avatar
Chao Liu committed
295
296
};

Chao Liu's avatar
Chao Liu committed
297
template <typename Reduce, index_t Init>
Chao Liu's avatar
Chao Liu committed
298
struct sequence_reverse_inclusive_scan<Sequence<>, Reduce, Init>
Chao Liu's avatar
Chao Liu committed
299
{
Chao Liu's avatar
Chao Liu committed
300
    using type = Sequence<>;
Chao Liu's avatar
Chao Liu committed
301
302
};

Chao Liu's avatar
Chao Liu committed
303
// split sequence
Chao Liu's avatar
Chao Liu committed
304
template <typename Seq, index_t I>
Chao Liu's avatar
Chao Liu committed
305
306
struct sequence_split
{
Chao Liu's avatar
Chao Liu committed
307
    static constexpr index_t NSize = Seq{}.Size();
Chao Liu's avatar
Chao Liu committed
308

Chao Liu's avatar
Chao Liu committed
309
310
    using range0 = typename arithmetic_sequence_gen<0, I, 1>::type;
    using range1 = typename arithmetic_sequence_gen<I, NSize, 1>::type;
Chao Liu's avatar
Chao Liu committed
311

Chao Liu's avatar
Chao Liu committed
312
313
    using left_type  = decltype(Seq::Extract(range0{}));
    using right_type = decltype(Seq::Extract(range1{}));
Chao Liu's avatar
Chao Liu committed
314
315
};

Chao Liu's avatar
Chao Liu committed
316
// reverse sequence
Chao Liu's avatar
Chao Liu committed
317
template <typename Seq>
Chao Liu's avatar
Chao Liu committed
318
319
struct sequence_reverse
{
Chao Liu's avatar
Chao Liu committed
320
    static constexpr index_t NSize = Seq{}.Size();
Chao Liu's avatar
Chao Liu committed
321
322

    using seq_split = sequence_split<Seq, NSize / 2>;
Chao Liu's avatar
Chao Liu committed
323
    using type      = typename sequence_merge<
Chao Liu's avatar
Chao Liu committed
324
325
        typename sequence_reverse<typename seq_split::right_type>::type,
        typename sequence_reverse<typename seq_split::left_type>::type>::type;
Chao Liu's avatar
Chao Liu committed
326
327
328
329
330
};

template <index_t I>
struct sequence_reverse<Sequence<I>>
{
Chao Liu's avatar
Chao Liu committed
331
    using type = Sequence<I>;
Chao Liu's avatar
Chao Liu committed
332
333
334
335
336
};

template <index_t I0, index_t I1>
struct sequence_reverse<Sequence<I0, I1>>
{
Chao Liu's avatar
Chao Liu committed
337
    using type = Sequence<I1, I0>;
Chao Liu's avatar
Chao Liu committed
338
};
Chao Liu's avatar
Chao Liu committed
339

Chao Liu's avatar
Chao Liu committed
340
#if 1
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
template <typename Reduce, typename Seq, typename... Seqs>
struct sequence_reduce
{
    using type = typename sequence_reduce<Reduce,
                                          Seq,
                                          typename sequence_reduce<Reduce, Seqs...>::type>::type;
};

template <typename Reduce, index_t... Xs, index_t... Ys>
struct sequence_reduce<Reduce, Sequence<Xs...>, Sequence<Ys...>>
{
    using type = Sequence<Reduce{}(Xs, Ys)...>;
};

template <typename Reduce, typename Seq>
struct sequence_reduce<Reduce, Seq>
{
    using type = Seq;
};
#endif

Chao Liu's avatar
Chao Liu committed
362
363
template <typename Values, typename Ids, typename Compare>
struct sequence_sort_impl
Chao Liu's avatar
Chao Liu committed
364
{
Chao Liu's avatar
Chao Liu committed
365
366
367
368
369
370
371
    template <typename LeftValues,
              typename LeftIds,
              typename RightValues,
              typename RightIds,
              typename MergedValues,
              typename MergedIds,
              typename Comp>
Chao Liu's avatar
Chao Liu committed
372
373
    struct sorted_sequence_merge_impl
    {
Chao Liu's avatar
Chao Liu committed
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
        static constexpr bool choose_left = LeftValues::Front() < RightValues::Front();

        static constexpr index_t chosen_value =
            choose_left ? LeftValues::Front() : RightValues::Front();
        static constexpr index_t chosen_id = choose_left ? LeftIds::Front() : RightIds::Front();

        using new_merged_values = decltype(MergedValues::PushBack(Number<chosen_value>{}));
        using new_merged_ids    = decltype(MergedIds::PushBack(Number<chosen_id>{}));

        using new_left_values =
            typename conditional<choose_left, decltype(LeftValues::PopFront()), LeftValues>::type;
        using new_left_ids =
            typename conditional<choose_left, decltype(LeftIds::PopFront()), LeftIds>::type;

        using new_right_values =
            typename conditional<choose_left, RightValues, decltype(RightValues::PopFront())>::type;
        using new_right_ids =
            typename conditional<choose_left, RightIds, decltype(RightIds::PopFront())>::type;

        using merge = sorted_sequence_merge_impl<new_left_values,
                                                 new_left_ids,
                                                 new_right_values,
                                                 new_right_ids,
                                                 new_merged_values,
                                                 new_merged_ids,
                                                 Comp>;
        // this is output
        using merged_values = typename merge::merged_values;
        using merged_ids    = typename merge::merged_ids;
Chao Liu's avatar
Chao Liu committed
403
404
    };

Chao Liu's avatar
Chao Liu committed
405
406
407
408
409
410
411
412
413
414
415
416
    template <typename LeftValues,
              typename LeftIds,
              typename MergedValues,
              typename MergedIds,
              typename Comp>
    struct sorted_sequence_merge_impl<LeftValues,
                                      LeftIds,
                                      Sequence<>,
                                      Sequence<>,
                                      MergedValues,
                                      MergedIds,
                                      Comp>
Chao Liu's avatar
Chao Liu committed
417
    {
Chao Liu's avatar
Chao Liu committed
418
419
        using merged_values = typename sequence_merge<MergedValues, LeftValues>::type;
        using merged_ids    = typename sequence_merge<MergedIds, LeftIds>::type;
Chao Liu's avatar
Chao Liu committed
420
421
    };

Chao Liu's avatar
Chao Liu committed
422
423
424
425
426
427
428
429
430
431
432
433
    template <typename RightValues,
              typename RightIds,
              typename MergedValues,
              typename MergedIds,
              typename Comp>
    struct sorted_sequence_merge_impl<Sequence<>,
                                      Sequence<>,
                                      RightValues,
                                      RightIds,
                                      MergedValues,
                                      MergedIds,
                                      Comp>
Chao Liu's avatar
Chao Liu committed
434
    {
Chao Liu's avatar
Chao Liu committed
435
436
        using merged_values = typename sequence_merge<MergedValues, RightValues>::type;
        using merged_ids    = typename sequence_merge<MergedIds, RightIds>::type;
Chao Liu's avatar
Chao Liu committed
437
438
    };

Chao Liu's avatar
Chao Liu committed
439
440
441
442
443
    template <typename LeftValues,
              typename LeftIds,
              typename RightValues,
              typename RightIds,
              typename Comp>
Chao Liu's avatar
Chao Liu committed
444
445
    struct sorted_sequence_merge
    {
Chao Liu's avatar
Chao Liu committed
446
447
448
449
450
451
452
453
454
455
        using merge = sorted_sequence_merge_impl<LeftValues,
                                                 LeftIds,
                                                 RightValues,
                                                 RightIds,
                                                 Sequence<>,
                                                 Sequence<>,
                                                 Comp>;

        using merged_values = typename merge::merged_values;
        using merged_ids    = typename merge::merged_ids;
Chao Liu's avatar
Chao Liu committed
456
457
    };

Chao Liu's avatar
Chao Liu committed
458
459
460
461
    static constexpr index_t nsize = Values::Size();

    using split_unsorted_values = sequence_split<Values, nsize / 2>;
    using split_unsorted_ids    = sequence_split<Ids, nsize / 2>;
Chao Liu's avatar
Chao Liu committed
462

Chao Liu's avatar
Chao Liu committed
463
464
465
466
467
    using left_unsorted_values = typename split_unsorted_values::left_type;
    using left_unsorted_ids    = typename split_unsorted_ids::left_type;
    using left_sort          = sequence_sort_impl<left_unsorted_values, left_unsorted_ids, Compare>;
    using left_sorted_values = typename left_sort::sorted_values;
    using left_sorted_ids    = typename left_sort::sorted_ids;
Chao Liu's avatar
Chao Liu committed
468

Chao Liu's avatar
Chao Liu committed
469
470
471
472
473
474
475
476
477
478
479
480
481
482
    using right_unsorted_values = typename split_unsorted_values::right_type;
    using right_unsorted_ids    = typename split_unsorted_ids::right_type;
    using right_sort = sequence_sort_impl<right_unsorted_values, right_unsorted_ids, Compare>;
    using right_sorted_values = typename right_sort::sorted_values;
    using right_sorted_ids    = typename right_sort::sorted_ids;

    using merged_sorted = sorted_sequence_merge<left_sorted_values,
                                                left_sorted_ids,
                                                right_sorted_values,
                                                right_sorted_ids,
                                                Compare>;

    using sorted_values = typename merged_sorted::merged_values;
    using sorted_ids    = typename merged_sorted::merged_ids;
Chao Liu's avatar
Chao Liu committed
483
484
};

Chao Liu's avatar
Chao Liu committed
485
486
template <index_t ValueX, index_t ValueY, index_t IdX, index_t IdY, typename Compare>
struct sequence_sort_impl<Sequence<ValueX, ValueY>, Sequence<IdX, IdY>, Compare>
Chao Liu's avatar
Chao Liu committed
487
{
Chao Liu's avatar
Chao Liu committed
488
489
490
491
492
493
    static constexpr bool choose_x = Compare{}(ValueX, ValueY);

    using sorted_values =
        typename conditional<choose_x, Sequence<ValueX, ValueY>, Sequence<ValueY, ValueX>>::type;
    using sorted_ids = typename conditional<choose_x, Sequence<IdX, IdY>, Sequence<IdY, IdX>>::type;
};
Chao Liu's avatar
Chao Liu committed
494

Chao Liu's avatar
Chao Liu committed
495
496
497
498
499
template <index_t Value, index_t Id, typename Compare>
struct sequence_sort_impl<Sequence<Value>, Sequence<Id>, Compare>
{
    using sorted_values = Sequence<Value>;
    using sorted_ids    = Sequence<Id>;
Chao Liu's avatar
Chao Liu committed
500
501
};

502
503
504
505
506
507
508
template <typename Compare>
struct sequence_sort_impl<Sequence<>, Sequence<>, Compare>
{
    using sorted_values = Sequence<>;
    using sorted_ids    = Sequence<>;
};

Chao Liu's avatar
Chao Liu committed
509
510
template <typename Values, typename Compare>
struct sequence_sort
Chao Liu's avatar
Chao Liu committed
511
{
Chao Liu's avatar
Chao Liu committed
512
513
514
515
516
517
    using unsorted_ids = typename arithmetic_sequence_gen<0, Values::Size(), 1>::type;
    using sort         = sequence_sort_impl<Values, unsorted_ids, Compare>;

    // this is output
    using type                = typename sort::sorted_values;
    using sorted2unsorted_map = typename sort::sorted_ids;
Chao Liu's avatar
Chao Liu committed
518
519
};

Chao Liu's avatar
Chao Liu committed
520
template <typename Values, typename Less, typename Equal>
Chao Liu's avatar
Chao Liu committed
521
522
struct sequence_unique_sort
{
Chao Liu's avatar
Chao Liu committed
523
524
525
526
527
    template <typename RemainValues,
              typename RemainIds,
              typename UniquifiedValues,
              typename UniquifiedIds,
              typename Eq>
Chao Liu's avatar
Chao Liu committed
528
529
    struct sorted_sequence_uniquify_impl
    {
Chao Liu's avatar
Chao Liu committed
530
531
532
533
534
535
536
537
538
539
540
541
        static constexpr index_t current_value = RemainValues::Front();
        static constexpr index_t current_id    = RemainIds::Front();

        static constexpr bool is_unique_value = (current_value != UniquifiedValues::Back());

        using new_remain_values = decltype(RemainValues::PopFront());
        using new_remain_ids    = decltype(RemainIds::PopFront());

        using new_uniquified_values =
            typename conditional<is_unique_value,
                                 decltype(UniquifiedValues::PushBack(Number<current_value>{})),
                                 UniquifiedValues>::type;
Chao Liu's avatar
Chao Liu committed
542

Chao Liu's avatar
Chao Liu committed
543
544
545
546
547
548
549
550
551
552
553
554
555
556
        using new_uniquified_ids =
            typename conditional<is_unique_value,
                                 decltype(UniquifiedIds::PushBack(Number<current_id>{})),
                                 UniquifiedIds>::type;

        using uniquify = sorted_sequence_uniquify_impl<new_remain_values,
                                                       new_remain_ids,
                                                       new_uniquified_values,
                                                       new_uniquified_ids,
                                                       Eq>;

        // this is output
        using uniquified_values = typename uniquify::uniquified_values;
        using uniquified_ids    = typename uniquify::uniquified_ids;
Chao Liu's avatar
Chao Liu committed
557
558
    };

Chao Liu's avatar
Chao Liu committed
559
560
561
562
563
564
    template <typename UniquifiedValues, typename UniquifiedIds, typename Eq>
    struct sorted_sequence_uniquify_impl<Sequence<>,
                                         Sequence<>,
                                         UniquifiedValues,
                                         UniquifiedIds,
                                         Eq>
Chao Liu's avatar
Chao Liu committed
565
    {
Chao Liu's avatar
Chao Liu committed
566
567
        using uniquified_values = UniquifiedValues;
        using uniquified_ids    = UniquifiedIds;
Chao Liu's avatar
Chao Liu committed
568
569
    };

Chao Liu's avatar
Chao Liu committed
570
    template <typename SortedValues, typename SortedIds, typename Eq>
Chao Liu's avatar
Chao Liu committed
571
572
    struct sorted_sequence_uniquify
    {
Chao Liu's avatar
Chao Liu committed
573
574
575
576
577
578
579
580
        using uniquify = sorted_sequence_uniquify_impl<decltype(SortedValues::PopFront()),
                                                       decltype(SortedIds::PopFront()),
                                                       Sequence<SortedValues::Front()>,
                                                       Sequence<SortedIds::Front()>,
                                                       Eq>;

        using uniquified_values = typename uniquify::uniquified_values;
        using uniquified_ids    = typename uniquify::uniquified_ids;
Chao Liu's avatar
Chao Liu committed
581
582
    };

Chao Liu's avatar
Chao Liu committed
583
584
585
    using sort          = sequence_sort<Values, Less>;
    using sorted_values = typename sort::type;
    using sorted_ids    = typename sort::sorted2unsorted_map;
Chao Liu's avatar
Chao Liu committed
586

Chao Liu's avatar
Chao Liu committed
587
588
589
590
591
    using uniquify = sorted_sequence_uniquify<sorted_values, sorted_ids, Equal>;

    // this is output
    using type                = typename uniquify::uniquified_values;
    using sorted2unsorted_map = typename uniquify::uniquified_ids;
Chao Liu's avatar
Chao Liu committed
592
593
};

Chao Liu's avatar
Chao Liu committed
594
template <typename SeqMap>
Chao Liu's avatar
Chao Liu committed
595
596
struct is_valid_sequence_map : is_same<typename arithmetic_sequence_gen<0, SeqMap::Size(), 1>::type,
                                       typename sequence_sort<SeqMap, math::less<index_t>>::type>
Chao Liu's avatar
Chao Liu committed
597
598
{
};
599

Chao Liu's avatar
Chao Liu committed
600
601
template <typename SeqMap>
struct sequence_map_inverse
Chao Liu's avatar
Chao Liu committed
602
{
Chao Liu's avatar
Chao Liu committed
603
604
605
606
607
    template <typename X2Y, typename WorkingY2X, index_t XBegin, index_t XRemain>
    struct sequence_map_inverse_impl
    {
        static constexpr auto new_y2x =
            WorkingY2X::Modify(X2Y::At(Number<XBegin>{}), Number<XBegin>{});
Chao Liu's avatar
Chao Liu committed
608

Chao Liu's avatar
Chao Liu committed
609
610
611
612
        using type =
            typename sequence_map_inverse_impl<X2Y, decltype(new_y2x), XBegin + 1, XRemain - 1>::
                type;
    };
Chao Liu's avatar
Chao Liu committed
613

Chao Liu's avatar
Chao Liu committed
614
615
616
617
618
    template <typename X2Y, typename WorkingY2X, index_t XBegin>
    struct sequence_map_inverse_impl<X2Y, WorkingY2X, XBegin, 0>
    {
        using type = WorkingY2X;
    };
Chao Liu's avatar
Chao Liu committed
619
620

    using type =
Chao Liu's avatar
Chao Liu committed
621
622
        typename sequence_map_inverse_impl<SeqMap,
                                           typename uniform_sequence_gen<SeqMap::Size(), 0>::type,
Chao Liu's avatar
Chao Liu committed
623
                                           0,
Chao Liu's avatar
Chao Liu committed
624
                                           SeqMap::Size()>::type;
Chao Liu's avatar
Chao Liu committed
625
626
};

627
628
629
630
631
632
template <index_t... Xs, index_t... Ys>
__host__ __device__ constexpr bool operator==(Sequence<Xs...>, Sequence<Ys...>)
{
    return ((Xs == Ys) && ...);
}

Chao Liu's avatar
Chao Liu committed
633
template <index_t... Xs, index_t... Ys>
Chao Liu's avatar
Chao Liu committed
634
__host__ __device__ constexpr auto operator+(Sequence<Xs...>, Sequence<Ys...>)
Chao Liu's avatar
Chao Liu committed
635
636
637
638
639
640
641
{
    static_assert(sizeof...(Xs) == sizeof...(Ys), "wrong! inconsistent size");

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

template <index_t... Xs, index_t... Ys>
Chao Liu's avatar
Chao Liu committed
642
__host__ __device__ constexpr auto operator-(Sequence<Xs...>, Sequence<Ys...>)
Chao Liu's avatar
Chao Liu committed
643
644
645
646
647
648
649
{
    static_assert(sizeof...(Xs) == sizeof...(Ys), "wrong! inconsistent size");

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

template <index_t... Xs, index_t... Ys>
Chao Liu's avatar
Chao Liu committed
650
__host__ __device__ constexpr auto operator*(Sequence<Xs...>, Sequence<Ys...>)
Chao Liu's avatar
Chao Liu committed
651
652
653
654
655
656
657
{
    static_assert(sizeof...(Xs) == sizeof...(Ys), "wrong! inconsistent size");

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

template <index_t... Xs, index_t... Ys>
Chao Liu's avatar
Chao Liu committed
658
__host__ __device__ constexpr auto operator/(Sequence<Xs...>, Sequence<Ys...>)
Chao Liu's avatar
Chao Liu committed
659
660
661
662
663
664
665
{
    static_assert(sizeof...(Xs) == sizeof...(Ys), "wrong! inconsistent size");

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

template <index_t... Xs, index_t... Ys>
Chao Liu's avatar
Chao Liu committed
666
__host__ __device__ constexpr auto operator%(Sequence<Xs...>, Sequence<Ys...>)
Chao Liu's avatar
Chao Liu committed
667
668
669
670
671
672
673
{
    static_assert(sizeof...(Xs) == sizeof...(Ys), "wrong! inconsistent size");

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

template <index_t... Xs, index_t Y>
Chao Liu's avatar
Chao Liu committed
674
__host__ __device__ constexpr auto operator+(Sequence<Xs...>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
675
{
Chao Liu's avatar
Chao Liu committed
676
    return Sequence<(Xs + Y)...>{};
Chao Liu's avatar
Chao Liu committed
677
678
679
}

template <index_t... Xs, index_t Y>
Chao Liu's avatar
Chao Liu committed
680
__host__ __device__ constexpr auto operator-(Sequence<Xs...>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
681
{
Chao Liu's avatar
Chao Liu committed
682
    return Sequence<(Xs - Y)...>{};
Chao Liu's avatar
Chao Liu committed
683
684
685
}

template <index_t... Xs, index_t Y>
Chao Liu's avatar
Chao Liu committed
686
__host__ __device__ constexpr auto operator*(Sequence<Xs...>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
687
{
Chao Liu's avatar
Chao Liu committed
688
    return Sequence<(Xs * Y)...>{};
Chao Liu's avatar
Chao Liu committed
689
690
691
}

template <index_t... Xs, index_t Y>
Chao Liu's avatar
Chao Liu committed
692
__host__ __device__ constexpr auto operator/(Sequence<Xs...>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
693
{
Chao Liu's avatar
Chao Liu committed
694
    return Sequence<(Xs / Y)...>{};
Chao Liu's avatar
Chao Liu committed
695
696
697
}

template <index_t... Xs, index_t Y>
Chao Liu's avatar
Chao Liu committed
698
__host__ __device__ constexpr auto operator%(Sequence<Xs...>, Number<Y>)
Chao Liu's avatar
Chao Liu committed
699
{
Chao Liu's avatar
Chao Liu committed
700
    return Sequence<(Xs % Y)...>{};
Chao Liu's avatar
Chao Liu committed
701
702
}

Chao Liu's avatar
Chao Liu committed
703
704
template <index_t Y, index_t... Xs>
__host__ __device__ constexpr auto operator+(Number<Y>, Sequence<Xs...>)
Chao Liu's avatar
Chao Liu committed
705
{
Chao Liu's avatar
Chao Liu committed
706
    return Sequence<(Y + Xs)...>{};
Chao Liu's avatar
Chao Liu committed
707
708
}

Chao Liu's avatar
Chao Liu committed
709
710
template <index_t Y, index_t... Xs>
__host__ __device__ constexpr auto operator-(Number<Y>, Sequence<Xs...>)
Chao Liu's avatar
Chao Liu committed
711
{
Chao Liu's avatar
Chao Liu committed
712
    return Sequence<(Y - Xs)...>{};
Chao Liu's avatar
Chao Liu committed
713
714
}

Chao Liu's avatar
Chao Liu committed
715
716
template <index_t Y, index_t... Xs>
__host__ __device__ constexpr auto operator*(Number<Y>, Sequence<Xs...>)
Chao Liu's avatar
Chao Liu committed
717
{
Chao Liu's avatar
Chao Liu committed
718
    return Sequence<(Y * Xs)...>{};
Chao Liu's avatar
Chao Liu committed
719
720
}

Chao Liu's avatar
Chao Liu committed
721
722
template <index_t Y, index_t... Xs>
__host__ __device__ constexpr auto operator/(Number<Y>, Sequence<Xs...>)
Chao Liu's avatar
Chao Liu committed
723
{
Chao Liu's avatar
Chao Liu committed
724
    return Sequence<(Y / Xs)...>{};
Chao Liu's avatar
Chao Liu committed
725
726
}

Chao Liu's avatar
Chao Liu committed
727
728
template <index_t Y, index_t... Xs>
__host__ __device__ constexpr auto operator%(Number<Y>, Sequence<Xs...>)
Chao Liu's avatar
Chao Liu committed
729
{
Chao Liu's avatar
Chao Liu committed
730
    return Sequence<(Y % Xs)...>{};
Chao Liu's avatar
Chao Liu committed
731
732
}

733
734
735
736
737
738
template <index_t I, index_t... Is>
__host__ __device__ constexpr auto sequence_pop_front(Sequence<I, Is...>)
{
    return Sequence<Is...>{};
}

Chao Liu's avatar
Chao Liu committed
739
template <typename Seq>
Chao Liu's avatar
Chao Liu committed
740
__host__ __device__ constexpr auto sequence_pop_back(Seq)
741
{
Chao Liu's avatar
Chao Liu committed
742
    static_assert(Seq::Size() > 0, "wrong! cannot pop an empty Sequence!");
743
    return sequence_pop_front(Seq::Reverse()).Reverse();
744
}
745

Chao Liu's avatar
Chao Liu committed
746
747
748
749
750
751
template <typename... Seqs>
__host__ __device__ constexpr auto merge_sequences(Seqs...)
{
    return typename sequence_merge<Seqs...>::type{};
}

752
753
754
755
756
757
template <typename F, index_t... Xs>
__host__ __device__ constexpr auto transform_sequences(F f, Sequence<Xs...>)
{
    return Sequence<f(Xs)...>{};
}

Chao Liu's avatar
Chao Liu committed
758
template <typename F, index_t... Xs, index_t... Ys>
759
__host__ __device__ constexpr auto transform_sequences(F f, Sequence<Xs...>, Sequence<Ys...>)
760
{
761
    static_assert(Sequence<Xs...>::mSize == Sequence<Ys...>::mSize, "Dim not the same");
762
763
764
765

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

Chao Liu's avatar
Chao Liu committed
766
template <typename F, index_t... Xs, index_t... Ys, index_t... Zs>
767
768
769
770
771
772
773
774
775
776
__host__ __device__ constexpr auto
transform_sequences(F f, Sequence<Xs...>, Sequence<Ys...>, Sequence<Zs...>)
{
    static_assert(Sequence<Xs...>::mSize == Sequence<Ys...>::mSize &&
                      Sequence<Xs...>::mSize == Sequence<Zs...>::mSize,
                  "Dim not the same");

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

Chao Liu's avatar
Chao Liu committed
777
template <typename Seq, typename Reduce, index_t Init>
Chao Liu's avatar
Chao Liu committed
778
__host__ __device__ constexpr auto reverse_inclusive_scan_sequence(Seq, Reduce, Number<Init>)
779
{
Chao Liu's avatar
Chao Liu committed
780
    return typename sequence_reverse_inclusive_scan<Seq, Reduce, Init>::type{};
781
782
}

Chao Liu's avatar
Chao Liu committed
783
784
785
786
787
788
789
template <typename Seq, typename Reduce, index_t Init>
__host__ __device__ constexpr auto reverse_exclusive_scan_sequence(Seq, Reduce, Number<Init>)
{
    return reverse_inclusive_scan_sequence(Seq::PopFront(), Reduce{}, Number<Init>{})
        .PushBack(Number<Init>{});
}

Chao Liu's avatar
Chao Liu committed
790
template <typename Seq, typename Reduce, index_t Init>
Chao Liu's avatar
Chao Liu committed
791
__host__ __device__ constexpr auto inclusive_scan_sequence(Seq, Reduce, Number<Init>)
792
{
Chao Liu's avatar
Chao Liu committed
793
    return reverse_inclusive_scan_sequence(Seq{}.Reverse(), Reduce{}, Number<Init>{}).Reverse();
794
}
795

Chao Liu's avatar
Chao Liu committed
796
template <typename Seq, index_t... Is>
797
__host__ __device__ constexpr auto pick_sequence_elements_by_ids(Seq, Sequence<Is...> /* ids */)
Chao Liu's avatar
Chao Liu committed
798
799
800
801
{
    return Sequence<Seq::At(Number<Is>{})...>{};
}

Chao Liu's avatar
Chao Liu committed
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
#if 1
namespace detail {
template <typename WorkSeq, typename RemainSeq, typename RemainMask>
struct pick_sequence_elements_by_mask_impl
{
    using new_work_seq = typename conditional<RemainMask::Front(),
                                              decltype(WorkSeq::PushBack(RemainSeq::Front())),
                                              WorkSeq>::type;

    using type =
        typename pick_sequence_elements_by_mask_impl<new_work_seq,
                                                     decltype(RemainSeq::PopFront()),
                                                     decltype(RemainMask::PopFront())>::type;
};

template <typename WorkSeq>
struct pick_sequence_elements_by_mask_impl<WorkSeq, Sequence<>, Sequence<>>
{
    using type = WorkSeq;
};

} // namespace detail

825
826
827
template <typename Seq, typename Mask>
__host__ __device__ constexpr auto pick_sequence_elements_by_mask(Seq, Mask)
{
Chao Liu's avatar
Chao Liu committed
828
829
830
    static_assert(Seq::Size() == Mask::Size(), "wrong!");

    return typename detail::pick_sequence_elements_by_mask_impl<Sequence<>, Seq, Mask>::type{};
831
832
}

Chao Liu's avatar
Chao Liu committed
833
834
835
namespace detail {
template <typename WorkSeq, typename RemainValues, typename RemainIds>
struct modify_sequence_elements_by_ids_impl
Chao Liu's avatar
Chao Liu committed
836
{
Chao Liu's avatar
Chao Liu committed
837
    using new_work_seq = decltype(WorkSeq::Modify(RemainIds::Front(), RemainValues::Front()));
Chao Liu's avatar
Chao Liu committed
838

Chao Liu's avatar
Chao Liu committed
839
840
841
842
843
    using type =
        typename modify_sequence_elements_by_ids_impl<new_work_seq,
                                                      decltype(RemainValues::PopFront()),
                                                      decltype(RemainIds::PopFront())>::type;
};
Chao Liu's avatar
Chao Liu committed
844

Chao Liu's avatar
Chao Liu committed
845
846
847
848
template <typename WorkSeq>
struct modify_sequence_elements_by_ids_impl<WorkSeq, Sequence<>, Sequence<>>
{
    using type = WorkSeq;
Chao Liu's avatar
Chao Liu committed
849
};
Chao Liu's avatar
Chao Liu committed
850
851
852
853
854
855
856
857
858
859
} // namespace detail

template <typename Seq, typename Values, typename Ids>
__host__ __device__ constexpr auto modify_sequence_elements_by_ids(Seq, Values, Ids)
{
    static_assert(Values::Size() == Ids::Size() && Seq::Size() >= Values::Size(), "wrong!");

    return typename detail::modify_sequence_elements_by_ids_impl<Seq, Values, Ids>::type{};
}
#endif
Chao Liu's avatar
Chao Liu committed
860

Chao Liu's avatar
Chao Liu committed
861
template <typename Seq, typename Reduce, index_t Init>
Chao Liu's avatar
Chao Liu committed
862
__host__ __device__ constexpr index_t
Chao Liu's avatar
Chao Liu committed
863
reduce_on_sequence(Seq, Reduce f, Number<Init> /*initial_value*/)
Chao Liu's avatar
Chao Liu committed
864
865
866
{
    index_t result = Init;

Chao Liu's avatar
Chao Liu committed
867
868
869
870
    for(index_t i = 0; i < Seq::Size(); ++i)
    {
        result = f(result, Seq::At(i));
    }
Chao Liu's avatar
Chao Liu committed
871
872
873
874

    return result;
}

Chao Liu's avatar
Chao Liu committed
875
876
// TODO: a generic any_of for any container
template <typename Seq, typename F>
Chao Liu's avatar
Chao Liu committed
877
__host__ __device__ constexpr bool sequence_any_of(Seq, F f)
Chao Liu's avatar
Chao Liu committed
878
879
880
881
882
883
884
885
886
887
888
889
890
{
    bool flag = false;

    for(index_t i = 0; i < Seq::Size(); ++i)
    {
        flag = flag || f(Seq::At(i));
    }

    return flag;
}

// TODO: a generic all_of for any container
template <typename Seq, typename F>
Chao Liu's avatar
Chao Liu committed
891
__host__ __device__ constexpr bool sequence_all_of(Seq, F f)
Chao Liu's avatar
Chao Liu committed
892
893
894
895
896
897
898
899
900
901
902
{
    bool flag = true;

    for(index_t i = 0; i < Seq::Size(); ++i)
    {
        flag = flag && f(Seq::At(i));
    }

    return flag;
}

903
904
905
906
907
908
template <typename Sx, typename Sy>
using sequence_merge_t = typename sequence_merge<Sx, Sy>::type;

template <index_t NSize, index_t I>
using uniform_sequence_gen_t = typename uniform_sequence_gen<NSize, I>::type;

909
} // namespace ck
910

911
#ifndef __HIPCC_RTC__
arai713's avatar
arai713 committed
912
#ifndef CK_CODE_GEN_RTC
913
914
915
916
917
918
919
920
921
922
template <ck::index_t... Is>
std::ostream& operator<<(std::ostream& os, const ck::Sequence<Is...>)
{
    using S = ck::Sequence<Is...>;
    os << "{";
    ck::static_for<0, S::Size() - ck::Number<1>{}, 1>{}(
        [&](auto i) { os << S::At(i).value << ", "; });
    os << S::At(S::Size() - ck::Number<1>{}).value << "}";
    return os;
}
923
#endif
924
#endif