MU Library
call_main.hpp
Go to the documentation of this file.
1 // SPDX-License-Identifier: Apache-2.0
2 // Copyright 2024 XCENA Inc.
3 
4 #pragma once
5 
6 #include <cstdint>
7 #include <cstring>
8 
9 #include "call_main_template.hpp"
10 #include "mu/ndarray/ndarray.hpp"
11 namespace mu
12 {
13 
14 bool isNormalMode(uint32_t mode);
15 
16 template <typename Function>
17 struct function_traits : public function_traits_helper<typename add_pointer<Function>::type>
18 {
19 };
20 
21 template <typename F, unsigned Arity = function_traits<F>::arity>
22 class CallHelper;
23 
24 namespace detail
25 {
26 
27 template <typename T>
28 T loadArgumentSlot(const void* address)
29 {
30  std::remove_cv_t<T> value;
31  std::memcpy(&value, address, sizeof(value));
32  return value;
33 }
34 
35 template <typename Pointer>
36 Pointer unpackKernelPointer(bool isNDArrayValue, const char* buffer)
37 {
38  // The driver passes &taskCmd.param[i] in each argv slot. For NDArray,
39  // param[i] points to the per-task descriptor, whose data() is the slice.
40  auto parameterSlot = loadArgumentSlot<const void*>(buffer);
41  if (isNDArrayValue)
42  {
43  auto ndarray = loadArgumentSlot<const mu::NDArray<char>*>(parameterSlot);
44  return reinterpret_cast<Pointer>(const_cast<char*>(ndarray->data()));
45  }
46  return loadArgumentSlot<Pointer>(parameterSlot);
47 }
48 
49 template <typename Pointer>
50 void storeKernelPointer(bool isNDArrayValue, char* buffer, Pointer value)
51 {
52  auto parameterSlot = loadArgumentSlot<void*>(buffer);
53  if (isNDArrayValue)
54  {
55  auto ndarray = loadArgumentSlot<mu::NDArray<char>*>(parameterSlot);
56  ndarray->setData(reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(value)));
57  }
58  else
59  {
60  std::memcpy(parameterSlot, &value, sizeof(value));
61  }
62 }
63 
64 template <typename Param>
65 decltype(auto) unpackKernelArgument(uint32_t mode, bool isNDArrayValue, char* buffer)
66 {
67  if constexpr (directional_pointer_traits<Param>::value)
68  {
69  using adapter = directional_pointer_traits<Param>;
70  using wire_type = typename adapter::wire_type;
71  return adapter::fromWire(unpackKernelPointer<wire_type>(isNDArrayValue, buffer));
72  }
73  else if constexpr (is_pointer<Param>::value)
74  {
75  return unpackKernelPointer<Param>(isNDArrayValue, buffer);
76  }
77  else
78  {
79  Param* value = *reinterpret_cast<Param**>(buffer);
80  if constexpr (is_ndarray_v<Param>)
81  {
82  void* argumentAddress = *reinterpret_cast<void**>(buffer);
83  value = *reinterpret_cast<Param**>(argumentAddress);
84  }
85  else if (!isNormalMode(mode))
86  {
87  void* argumentAddress = *reinterpret_cast<void**>(buffer);
88  if (isNDArrayValue)
89  {
90  auto ndarray = reinterpret_cast<mu::NDArray<char>*>(argumentAddress);
91  void* dataPointer = ndarray->data();
92  value = *reinterpret_cast<Param**>(dataPointer);
93  }
94  else
95  {
96  value = *reinterpret_cast<Param**>(argumentAddress);
97  }
98  }
99  return static_cast<Param&>(*value);
100  }
101 }
102 
103 } // namespace detail
104 
105 #define GET_ARG(num) \
106  static_assert(detail::supported_directional_parameter_v<typename function_traits<F>::arg##num##_type>, \
107  "Directional kernel pointer parameters must be passed by value, not by reference"); \
108  using remove_reference##num##_t = typename remove_reference<typename function_traits<F>::arg##num##_type>::type; \
109  decltype(auto) a##num = detail::unpackKernelArgument<remove_reference##num##_t>( \
110  mode, (ndarrayValueBit & (1ull << (num - 1))) != 0, buffer); \
111  buffer += sizeof(void*);
112 
113 #define PRIMITIVE_CAT(a, ...) a##__VA_ARGS__
114 #define ARGS(x) PRIMITIVE_CAT(ARGS_, x)
115 #define ARGS_0
116 #define ARGS_1 GET_ARG(1);
117 #define ARGS_2 ARGS_1 GET_ARG(2);
118 #define ARGS_3 ARGS_2 GET_ARG(3);
119 #define ARGS_4 ARGS_3 GET_ARG(4);
120 #define ARGS_5 ARGS_4 GET_ARG(5);
121 #define ARGS_6 ARGS_5 GET_ARG(6);
122 #define ARGS_7 ARGS_6 GET_ARG(7);
123 #define ARGS_8 ARGS_7 GET_ARG(8);
124 #define ARGS_9 ARGS_8 GET_ARG(9);
125 
126 // Preserve legacy mutable pointer-reference updates between invocations.
127 // Ordinary pointer values and directional wrappers have no writeback here.
128 #define STORE_ARG(num) \
129  if constexpr (std::is_lvalue_reference_v<typename function_traits<F>::arg##num##_type> && \
130  is_pointer<remove_reference##num##_t>::value && !std::is_const_v<remove_reference##num##_t>) \
131  { \
132  detail::storeKernelPointer((ndarrayValueBit & (1ull << (num - 1))) != 0, \
133  argumentStart + (num - 1) * sizeof(void*), a##num); \
134  }
135 #define STORES(x) PRIMITIVE_CAT(STORES_, x)
136 #define STORES_0
137 #define STORES_1 STORE_ARG(1);
138 #define STORES_2 STORES_1 STORE_ARG(2);
139 #define STORES_3 STORES_2 STORE_ARG(3);
140 #define STORES_4 STORES_3 STORE_ARG(4);
141 #define STORES_5 STORES_4 STORE_ARG(5);
142 #define STORES_6 STORES_5 STORE_ARG(6);
143 #define STORES_7 STORES_6 STORE_ARG(7);
144 #define STORES_8 STORES_7 STORE_ARG(8);
145 #define STORES_9 STORES_8 STORE_ARG(9);
146 
147 #define CALL_PARAMS(x) PRIMITIVE_CAT(CALL_PARAMS_, x)
148 #define CALL_PARAMS_0
149 #define CALL_PARAMS_1 a1
150 #define CALL_PARAMS_2 CALL_PARAMS_1, a2
151 #define CALL_PARAMS_3 CALL_PARAMS_2, a3
152 #define CALL_PARAMS_4 CALL_PARAMS_3, a4
153 #define CALL_PARAMS_5 CALL_PARAMS_4, a5
154 #define CALL_PARAMS_6 CALL_PARAMS_5, a6
155 #define CALL_PARAMS_7 CALL_PARAMS_6, a7
156 #define CALL_PARAMS_8 CALL_PARAMS_7, a8
157 #define CALL_PARAMS_9 CALL_PARAMS_8, a9
158 #define CALL_FUNC(x) fp(CALL_PARAMS(x))
159 
160 #define CALL_HELPER(_NUM) \
161  template <typename F> \
162  class CallHelper<F, _NUM> \
163  { \
164  public: \
165  template <typename FP> \
166  static void invoke(FP fp, [[maybe_unused]] uint32_t mode, [[maybe_unused]] uint32_t ndarrayValueBit, \
167  [[maybe_unused]] char* buffer) \
168  { \
169  [[maybe_unused]] char* argumentStart = buffer; \
170  ARGS(_NUM); \
171  CALL_FUNC(_NUM); \
172  STORES(_NUM); \
173  } \
174  }
175 
176 CALL_HELPER(0);
177 CALL_HELPER(1);
178 CALL_HELPER(2);
179 CALL_HELPER(3);
180 CALL_HELPER(4);
181 CALL_HELPER(5);
182 CALL_HELPER(6);
183 CALL_HELPER(7);
184 CALL_HELPER(8);
185 CALL_HELPER(9);
186 
187 #undef GET_ARG
188 #undef PRIMITIVE_CAT
189 #undef ARGS
190 #undef ARGS_0
191 #undef ARGS_1
192 #undef ARGS_2
193 #undef ARGS_3
194 #undef ARGS_4
195 #undef ARGS_5
196 #undef ARGS_6
197 #undef ARGS_7
198 #undef ARGS_8
199 #undef ARGS_9
200 #undef STORE_ARG
201 #undef STORES
202 #undef STORES_0
203 #undef STORES_1
204 #undef STORES_2
205 #undef STORES_3
206 #undef STORES_4
207 #undef STORES_5
208 #undef STORES_6
209 #undef STORES_7
210 #undef STORES_8
211 #undef STORES_9
212 #undef CALL_PARAMS
213 #undef CALL_PARAMS_0
214 #undef CALL_PARAMS_1
215 #undef CALL_PARAMS_2
216 #undef CALL_PARAMS_3
217 #undef CALL_PARAMS_4
218 #undef CALL_PARAMS_5
219 #undef CALL_PARAMS_6
220 #undef CALL_PARAMS_7
221 #undef CALL_PARAMS_8
222 #undef CALL_PARAMS_9
223 #undef CALL_FUNC
224 #undef CALL_HELPER
225 
226 template <typename FP>
227 inline void callMain(FP func, uint32_t mode, uint32_t ndarrayValueBit, char* buffer)
228 {
229  CallHelper<typename remove_pointer<FP>::type>::invoke(func, mode, ndarrayValueBit, buffer);
230 }
231 
232 #undef GET_ARG
233 
234 } // namespace mu
call_main_template.hpp
mu::NDArray
NDArray is a class that can be used in both host and device code. NDArray is automatically divided in...
Definition: ndarray.hpp:23
mu::NDArray::data
T * data()
Definition: ndarray.hpp:79
mu::detail::unpackKernelPointer
Pointer unpackKernelPointer(bool isNDArrayValue, const char *buffer)
Definition: call_main.hpp:36
mu::detail::loadArgumentSlot
T loadArgumentSlot(const void *address)
Definition: call_main.hpp:28
mu::detail::storeKernelPointer
void storeKernelPointer(bool isNDArrayValue, char *buffer, Pointer value)
Definition: call_main.hpp:50
mu::detail::unpackKernelArgument
decltype(auto) unpackKernelArgument(uint32_t mode, bool isNDArrayValue, char *buffer)
Definition: call_main.hpp:65
mu
Definition: access.hpp:10
mu::callMain
void callMain(FP func, uint32_t mode, uint32_t ndarrayValueBit, char *buffer)
Definition: call_main.hpp:227
mu::CALL_HELPER
CALL_HELPER(0)
mu::isNormalMode
bool isNormalMode(uint32_t mode)
mu::detail::directional_pointer_traits_impl
Definition: call_main_template.hpp:273
mu::function_traits_helper
Definition: call_main_template.hpp:16
mu::function_traits
Definition: call_main.hpp:18
mu::is_pointer
Definition: call_main_template.hpp:204