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

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/utility/statically_indexed_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/statically_indexed_array_multi_index.hpp Source File
statically_indexed_array_multi_index.hpp
Go to the documentation of this file.
1 // SPDX-License-Identifier: MIT
2 // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
3 
4 #ifndef CK_STATICALLY_INDEXED_ARRAY_MULTI_INDEX_HPP
5 #define CK_STATICALLY_INDEXED_ARRAY_MULTI_INDEX_HPP
6 
7 #include "common_header.hpp"
8 #include "ck/utility/math_v2.hpp"
9 
10 namespace ck {
11 
12 template <index_t N>
13 using MultiIndex = StaticallyIndexedArray<index_t, N>;
14 
15 template <typename... Xs>
16 __host__ __device__ constexpr auto make_multi_index(Xs&&... xs)
17 {
18  return make_statically_indexed_array<index_t>(index_t{xs}...);
19 }
20 
21 template <index_t NSize>
22 __host__ __device__ constexpr auto make_zero_multi_index()
23 {
24  return unpack([](auto... xs) { return make_multi_index(xs...); },
26 }
27 
28 template <typename T>
29 __host__ __device__ constexpr auto to_multi_index(const T& x)
30 {
31  return unpack([](auto... ys) { return make_multi_index(ys...); }, x);
32 }
33 
34 // Here should use MultiIndex<NSize>, instead of Tuple<Ys...>, although the former
35 // is the alias of the latter. This is because compiler cannot infer the NSize if
36 // using MultiIndex<NSize>
37 // TODO: how to fix this?
38 template <typename... Ys,
39  typename X,
41 __host__ __device__ constexpr auto operator+=(Tuple<Ys...>& y, const X& x)
42 {
43  static_assert(X::Size() == sizeof...(Ys), "wrong! size not the same");
44  constexpr index_t NSize = sizeof...(Ys);
45  static_for<0, NSize, 1>{}([&](auto i) { y(i) += x[i]; });
46  return y;
47 }
48 
49 template <typename... Ys,
50  typename X,
52 __host__ __device__ constexpr auto operator-=(Tuple<Ys...>& y, const X& x)
53 {
54  static_assert(X::Size() == sizeof...(Ys), "wrong! size not the same");
55  constexpr index_t NSize = sizeof...(Ys);
56  static_for<0, NSize, 1>{}([&](auto i) { y(i) -= x[i]; });
57  return y;
58 }
59 
60 template <typename... Xs,
61  typename Y,
63 __host__ __device__ constexpr auto operator+(const Tuple<Xs...>& x, const Y& y)
64 {
65  static_assert(Y::Size() == sizeof...(Xs), "wrong! size not the same");
66  constexpr index_t NSize = sizeof...(Xs);
67 
68  Tuple<Xs...> r;
69  static_for<0, NSize, 1>{}([&](auto i) { r(i) = x[i] + y[i]; });
70  return r;
71 }
72 
73 template <typename... Xs,
74  typename Y,
76 __host__ __device__ constexpr auto operator-(const Tuple<Xs...>& x, const Y& y)
77 {
78  static_assert(Y::Size() == sizeof...(Xs), "wrong! size not the same");
79  constexpr index_t NSize = sizeof...(Xs);
80 
81  Tuple<Xs...> r;
82  static_for<0, NSize, 1>{}([&](auto i) { r(i) = x[i] - y[i]; });
83  return r;
84 }
85 
86 template <typename... Xs,
87  typename Y,
89 __host__ __device__ constexpr auto operator*(const Tuple<Xs...>& x, const Y& y)
90 {
91  static_assert(Y::Size() == sizeof...(Xs), "wrong! size not the same");
92  constexpr index_t NSize = sizeof...(Xs);
93 
94  Tuple<Xs...> r;
95  static_for<0, NSize, 1>{}([&](auto i) { r(i) = x[i] * y[i]; });
96  return r;
97 }
98 
99 // MultiIndex = scalar * MultiIndex
100 template <typename... Xs,
101  typename Y,
103 __host__ __device__ constexpr auto operator*(Y a, const Tuple<Xs...>& x)
104 {
105  constexpr index_t NSize = sizeof...(Xs);
106 
107  Tuple<Xs...> r;
108  static_for<0, NSize, 1>{}([&](auto i) { r(i) = a * x[i]; });
109  return r;
110 }
111 
112 // MultiIndex = MultiIndex * scalar
113 template <typename... Xs,
114  typename Y,
116 __host__ __device__ constexpr auto operator*(const Tuple<Xs...>& x, Y a)
117 {
118  return a * x;
119 }
120 
121 namespace mathext {
122 
123 template <typename... Xs>
124 __host__ __device__ constexpr auto exp(const Tuple<Xs...>& x)
125 {
126  constexpr index_t NSize = sizeof...(Xs);
127 
128  Tuple<Xs...> r;
129  static_for<0, NSize, 1>{}([&](auto i) { r(i) = math::exp(x[i]); });
130  return r;
131 }
132 
133 template <typename... Xs, typename Y>
134 __host__ __device__ constexpr auto max(const Tuple<Xs...>& x, const Y& y)
135 {
136  static_assert(Y::Size() == sizeof...(Xs), "wrong! size not the same");
137  constexpr index_t NSize = sizeof...(Xs);
138 
139  Tuple<Xs...> r;
140  static_for<0, NSize, 1>{}([&](auto i) { r(i) = math::max(x[i], y[i]); });
141  return r;
142 }
143 
144 } // namespace mathext
145 
146 template <typename... Xs>
147 __host__ __device__ void print_multi_index(const Tuple<Xs...>& x)
148 {
149  printf("{");
150  printf("MultiIndex, ");
151  printf("size %d,", index_t{sizeof...(Xs)});
152  static_for<0, sizeof...(Xs), 1>{}(
153  [&](auto i) { printf("%d ", static_cast<index_t>(x.At(i))); });
154  printf("}");
155 }
156 
157 } // namespace ck
158 #endif
__host__ T exp(T x)
Definition: math_v2.hpp:391
__host__ constexpr __device__ T max(T x)
Definition: math.hpp:84
__host__ constexpr __device__ auto max(const Tuple< Xs... > &x, const Y &y)
Definition: statically_indexed_array_multi_index.hpp:134
__host__ constexpr __device__ auto exp(const Tuple< Xs... > &x)
Definition: statically_indexed_array_multi_index.hpp:124
Definition: ck.hpp:267
__host__ constexpr __device__ auto make_multi_index(Xs &&... xs)
Definition: array_multi_index.hpp:15
__host__ constexpr __device__ auto operator+(const MultiIndex< NSize > &a, const T &b)
Definition: array_multi_index.hpp:50
__host__ constexpr __device__ auto operator+=(MultiIndex< NSize > &y, const X &x)
Definition: array_multi_index.hpp:34
__host__ constexpr __device__ auto operator-=(MultiIndex< NSize > &y, const X &x)
Definition: array_multi_index.hpp:42
__host__ constexpr __device__ auto to_multi_index(const T &x)
Definition: array_multi_index.hpp:28
__host__ constexpr __device__ auto operator-(const MultiIndex< NSize > &a, const T &b)
Definition: array_multi_index.hpp:60
__host__ constexpr __device__ auto operator*(const MultiIndex< NSize > &a, const T &b)
Definition: array_multi_index.hpp:70
__host__ constexpr __device__ auto unpack(F &&f, X &&x)
Definition: functional4.hpp:46
int32_t index_t
Definition: ck.hpp:298
__host__ constexpr __device__ auto make_zero_multi_index()
Definition: array_multi_index.hpp:21
Array< index_t, N > MultiIndex
Definition: array_multi_index.hpp:12
__host__ __device__ void print_multi_index(const Tuple< Xs... > &x)
Definition: statically_indexed_array_multi_index.hpp:147
const GenericPointer< typename T::ValueType > T2 value
Definition: pointer.h:1350
const GenericPointer< typename T::ValueType > T2 T::AllocatorType & a
Definition: pointer.h:1249
Definition: tuple.hpp:117
Definition: type.hpp:187
Definition: functional2.hpp:33
typename sequence_gen< NSize, F >::type type
Definition: sequence.hpp:295