sequence_helper.hpp 1006 Bytes
Newer Older
Umang Yadav's avatar
Umang Yadav committed
1
2
3

#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Weverything"
Chao Liu's avatar
Chao Liu committed
4
// SPDX-License-Identifier: MIT
Illia Silin's avatar
Illia Silin committed
5
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
Chao Liu's avatar
Chao Liu committed
6

7
#pragma once
Chao Liu's avatar
Chao Liu committed
8

9
#include "ck/utility/tuple.hpp"
Chao Liu's avatar
Chao Liu committed
10
11
12

namespace ck {

13
14
15
16
17
18
19
template <index_t... Is>
__host__ __device__ constexpr auto make_sequence(Number<Is>...)
{
    return Sequence<Is...>{};
}

// F returns index_t
Chao Liu's avatar
Chao Liu committed
20
21
22
23
24
25
template <typename F, index_t N>
__host__ __device__ constexpr auto generate_sequence(F, Number<N>)
{
    return typename sequence_gen<N, F>::type{};
}

26
27
28
29
30
31
32
33
// F returns Number<>
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{});
}

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

Chao Liu's avatar
Chao Liu committed
40
} // namespace ck
Umang Yadav's avatar
Umang Yadav committed
41
42

#pragma clang diagnostic pop