sequence_helper.hpp 930 Bytes
Newer Older
Chao Liu's avatar
Chao Liu committed
1
// SPDX-License-Identifier: MIT
Illia Silin's avatar
Illia Silin committed
2
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
Chao Liu's avatar
Chao Liu committed
3

4
#pragma once
Chao Liu's avatar
Chao Liu committed
5

6
#include "ck/utility/tuple.hpp"
Chao Liu's avatar
Chao Liu committed
7
8
9

namespace ck {

10
11
12
13
14
15
template <index_t... Is>
__host__ __device__ constexpr auto make_sequence(Number<Is>...)
{
    return Sequence<Is...>{};
}

Chao Liu's avatar
Chao Liu committed
16
17
// F() returns index_t
// F use default constructor
Chao Liu's avatar
Chao Liu committed
18
19
20
21
22
23
template <typename F, index_t N>
__host__ __device__ constexpr auto generate_sequence(F, Number<N>)
{
    return typename sequence_gen<N, F>::type{};
}

Chao Liu's avatar
Chao Liu committed
24
// F() returns Number<>
25
26
27
28
29
30
31
template <typename F, index_t N>
__host__ __device__ constexpr auto generate_sequence_v2(F&& f, Number<N>)
{
    return unpack([&f](auto&&... xs) { return make_sequence(f(xs)...); },
                  typename arithmetic_sequence_gen<0, N, 1>::type{});
}

32
33
34
35
36
37
template <index_t... Is>
__host__ __device__ constexpr auto to_sequence(Tuple<Number<Is>...>)
{
    return Sequence<Is...>{};
}

Chao Liu's avatar
Chao Liu committed
38
} // namespace ck