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

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

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/utility/container_element_picker.hpp Source File
container_element_picker.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_CONTAINER_ELEMENT_PICKER_HPP
5 #define CK_CONTAINER_ELEMENT_PICKER_HPP
6 
7 #include "functional2.hpp"
8 #include "sequence.hpp"
9 
10 namespace ck {
11 
12 // Arr: Array or StaticallyIndexedArray
13 // Picks: Sequence<...>
14 template <typename Arr, typename Picks>
16 {
18 #if 0
19  using data_type = typename Arr::data_type;
20 #endif
21 
22  __host__ __device__ constexpr ContainerElementPicker() = delete;
23 
24  __host__ __device__ constexpr ContainerElementPicker(Arr& array) : mArray{array}
25  {
26  constexpr index_t imax =
28 
29  static_assert(imax < Arr::Size(), "wrong! exceeding # array element");
30  }
31 
32  __host__ __device__ static constexpr auto Size() { return Picks::Size(); }
33 
34  template <index_t I>
35  __host__ __device__ constexpr const auto& At(Number<I> i) const
36  {
37  static_assert(I < Size(), "wrong!");
38 
39  constexpr auto IP = Picks{}[i];
40  return mArray[IP];
41  }
42 
43  template <index_t I>
44  __host__ __device__ constexpr auto& At(Number<I> i)
45  {
46  static_assert(I < Size(), "wrong!");
47 
48  constexpr auto IP = Picks{}[i];
49  return mArray(IP);
50  }
51 
52  template <index_t I>
53  __host__ __device__ constexpr const auto& operator[](Number<I> i) const
54  {
55  return At(i);
56  }
57 
58  template <index_t I>
59  __host__ __device__ constexpr auto& operator()(Number<I> i)
60  {
61  return At(i);
62  }
63 
64  template <typename T>
65  __host__ __device__ constexpr auto operator=(const T& a)
66  {
67  static_assert(T::Size() == Size(), "wrong! size not the same");
68 
69  static_for<0, Size(), 1>{}([&](auto i) { operator()(i) = a[i]; });
70 
71  return *this;
72  }
73 
74  private:
75  Arr& mArray;
76 };
77 
78 // Arr: Array or StaticallyIndexedArray
79 // Picks: Sequence<...>
80 template <typename Arr, typename Picks>
82 {
84 #if 0
85  using data_type = typename Arr::data_type;
86 #endif
87 
88  __host__ __device__ constexpr ConstantContainerElementPicker() = delete;
89 
90  __host__ __device__ constexpr ConstantContainerElementPicker(const Arr& array) : mArray{array}
91  {
92  constexpr index_t imax =
94 
95  static_assert(imax < Arr::Size(), "wrong! exceeding # array element");
96  }
97 
98  __host__ __device__ static constexpr auto Size() { return Picks::Size(); }
99 
100  template <index_t I>
101  __host__ __device__ constexpr const auto& At(Number<I> i) const
102  {
103  static_assert(I < Size(), "wrong!");
104 
105  constexpr auto IP = Picks{}[i];
106  return mArray[IP];
107  }
108 
109  template <index_t I>
110  __host__ __device__ constexpr const auto& operator[](Number<I> i) const
111  {
112  return At(i);
113  }
114 
115  private:
116  const Arr& mArray;
117 };
118 
119 template <typename Arr, typename Picks, typename X>
120 __host__ __device__ constexpr auto operator+=(ContainerElementPicker<Arr, Picks>& y, const X& x)
121 {
123  constexpr index_t nsize = Y::Size();
124 
125  static_assert(nsize == X::Size(), "wrong! size not the same");
126 
127  static_for<0, nsize, 1>{}([&](auto i) { y(i) += x[i]; });
128 
129  return y;
130 }
131 
132 template <typename Arr, typename Picks, typename X>
133 __host__ __device__ constexpr auto operator-=(ContainerElementPicker<Arr, Picks>& y, const X& x)
134 {
136  constexpr index_t nsize = Y::Size();
137 
138  static_assert(nsize == X::Size(), "wrong! size not the same");
139 
140  static_for<0, nsize, 1>{}([&](auto i) { y(i) -= x[i]; });
141 
142  return y;
143 }
144 
145 template <typename Arr, typename Picks>
146 __host__ __device__ constexpr auto pick_container_element(Arr& a, Picks)
147 {
149 }
150 
151 template <typename Arr, typename Picks>
152 __host__ __device__ constexpr auto pick_container_element(const Arr& a, Picks)
153 {
155 }
156 
157 } // namespace ck
158 #endif
Definition: ck.hpp:267
__host__ constexpr __device__ auto operator+=(MultiIndex< NSize > &y, const X &x)
Definition: array_multi_index.hpp:34
__host__ constexpr __device__ index_t reduce_on_sequence(Seq, Reduce f, Number< Init >)
Definition: sequence.hpp:884
__host__ constexpr __device__ auto operator-=(MultiIndex< NSize > &y, const X &x)
Definition: array_multi_index.hpp:42
__host__ constexpr __device__ auto pick_container_element(Arr &a, Picks)
Definition: container_element_picker.hpp:146
int32_t index_t
Definition: ck.hpp:298
const GenericPointer< typename T::ValueType > T2 T::AllocatorType & a
Definition: pointer.h:1249
Definition: container_element_picker.hpp:82
__host__ constexpr __device__ const auto & operator[](Number< I > i) const
Definition: container_element_picker.hpp:110
__host__ static constexpr __device__ auto Size()
Definition: container_element_picker.hpp:98
__host__ constexpr __device__ const auto & At(Number< I > i) const
Definition: container_element_picker.hpp:101
__host__ constexpr __device__ ConstantContainerElementPicker(const Arr &array)
Definition: container_element_picker.hpp:90
__host__ constexpr __device__ ConstantContainerElementPicker()=delete
Definition: container_element_picker.hpp:16
__host__ constexpr __device__ auto & operator()(Number< I > i)
Definition: container_element_picker.hpp:59
__host__ constexpr __device__ auto operator=(const T &a)
Definition: container_element_picker.hpp:65
__host__ static constexpr __device__ auto Size()
Definition: container_element_picker.hpp:32
__host__ constexpr __device__ ContainerElementPicker(Arr &array)
Definition: container_element_picker.hpp:24
__host__ constexpr __device__ ContainerElementPicker()=delete
__host__ constexpr __device__ auto & At(Number< I > i)
Definition: container_element_picker.hpp:44
__host__ constexpr __device__ const auto & operator[](Number< I > i) const
Definition: container_element_picker.hpp:53
__host__ constexpr __device__ const auto & At(Number< I > i) const
Definition: container_element_picker.hpp:35
Definition: integral_constant.hpp:20
Definition: math.hpp:44
Definition: functional2.hpp:33