/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/utility/array_multi_index.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/utility/array_multi_index.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/utility/array_multi_index.hpp Source File
array_multi_index.hpp
Go to the documentation of this file.
1 // SPDX-License-Identifier: MIT
2 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
3 
4 #ifndef CK_ARRAY_MULTI_INDEX_HPP
5 #define CK_ARRAY_MULTI_INDEX_HPP
6 
7 #include "common_header.hpp"
8 
9 namespace ck {
10 
11 template <index_t N>
12 using MultiIndex = Array<index_t, N>;
13 
14 template <typename... Xs>
15 __host__ __device__ constexpr auto make_multi_index(Xs&&... xs)
16 {
17  return make_array<index_t>(index_t{xs}...);
18 }
19 
20 template <index_t NSize>
21 __host__ __device__ constexpr auto make_zero_multi_index()
22 {
23  return unpack([](auto... xs) { return make_multi_index(xs...); },
24  typename uniform_sequence_gen<NSize, 0>::type{});
25 }
26 
27 template <typename T>
28 __host__ __device__ constexpr auto to_multi_index(const T& x)
29 {
30  return unpack([](auto... ys) { return make_multi_index(ys...); }, x);
31 }
32 
33 template <index_t NSize, typename X>
34 __host__ __device__ constexpr auto operator+=(MultiIndex<NSize>& y, const X& x)
35 {
36  static_assert(X::Size() == NSize, "wrong! size not the same");
37  static_for<0, NSize, 1>{}([&](auto i) { y(i) += x[i]; });
38  return y;
39 }
40 
41 template <index_t NSize, typename X>
42 __host__ __device__ constexpr auto operator-=(MultiIndex<NSize>& y, const X& x)
43 {
44  static_assert(X::Size() == NSize, "wrong! size not the same");
45  static_for<0, NSize, 1>{}([&](auto i) { y(i) -= x[i]; });
46  return y;
47 }
48 
49 template <index_t NSize, typename T>
50 __host__ __device__ constexpr auto operator+(const MultiIndex<NSize>& a, const T& b)
51 {
52  using type = MultiIndex<NSize>;
53  static_assert(T::Size() == NSize, "wrong! size not the same");
54  type r;
55  static_for<0, NSize, 1>{}([&](auto i) { r(i) = a[i] + b[i]; });
56  return r;
57 }
58 
59 template <index_t NSize, typename T>
60 __host__ __device__ constexpr auto operator-(const MultiIndex<NSize>& a, const T& b)
61 {
62  using type = MultiIndex<NSize>;
63  static_assert(T::Size() == NSize, "wrong! size not the same");
64  type r;
65  static_for<0, NSize, 1>{}([&](auto i) { r(i) = a[i] - b[i]; });
66  return r;
67 }
68 
69 template <index_t NSize, typename T>
70 __host__ __device__ constexpr auto operator*(const MultiIndex<NSize>& a, const T& b)
71 {
72  using type = MultiIndex<NSize>;
73  static_assert(T::Size() == NSize, "wrong! size not the same");
74  type r;
75  static_for<0, NSize, 1>{}([&](auto i) { r(i) = a[i] * b[i]; });
76  return r;
77 }
78 
79 } // namespace ck
80 #endif
common_header.hpp
ck
Definition: ck.hpp:267
ck::make_multi_index
__host__ constexpr __device__ auto make_multi_index(Xs &&... xs)
Definition: array_multi_index.hpp:15
ck::operator+
__host__ constexpr __device__ auto operator+(const MultiIndex< NSize > &a, const T &b)
Definition: array_multi_index.hpp:50
ck::operator+=
__host__ constexpr __device__ auto operator+=(MultiIndex< NSize > &y, const X &x)
Definition: array_multi_index.hpp:34
ck::operator-=
__host__ constexpr __device__ auto operator-=(MultiIndex< NSize > &y, const X &x)
Definition: array_multi_index.hpp:42
ck::to_multi_index
__host__ constexpr __device__ auto to_multi_index(const T &x)
Definition: array_multi_index.hpp:28
ck::operator-
__host__ constexpr __device__ auto operator-(const MultiIndex< NSize > &a, const T &b)
Definition: array_multi_index.hpp:60
ck::operator*
__host__ constexpr __device__ auto operator*(const MultiIndex< NSize > &a, const T &b)
Definition: array_multi_index.hpp:70
ck::unpack
__host__ constexpr __device__ auto unpack(F &&f, X &&x)
Definition: functional4.hpp:46
ck::index_t
int32_t index_t
Definition: ck.hpp:298
ck::make_zero_multi_index
__host__ constexpr __device__ auto make_zero_multi_index()
Definition: array_multi_index.hpp:21
a
const GenericPointer< typename T::ValueType > T2 T::AllocatorType & a
Definition: pointer.h:1249
ck::Array
Definition: array.hpp:14
ck::static_for
Definition: functional2.hpp:33
ck::uniform_sequence_gen::type
typename sequence_gen< NSize, F >::type type
Definition: sequence.hpp:295