1//===-- Shared memory RPC client / server utilities -------------*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#ifndef LLVM_LIBC_SHARED_RPC_UTIL_H
10#define LLVM_LIBC_SHARED_RPC_UTIL_H
11
12#include <stddef.h>
13#include <stdint.h>
14
15#if (defined(__NVPTX__) || defined(__AMDGPU__) || defined(__SPIRV__)) && \
16 !((defined(__CUDA__) && !defined(__CUDA_ARCH__)) || \
17 (defined(__HIP__) && !defined(__HIP_DEVICE_COMPILE__)))
18#include <gpuintrin.h>
19#define RPC_TARGET_IS_GPU
20#endif
21
22// Workaround for missing __has_builtin in < GCC 10.
23#ifndef __has_builtin
24#define __has_builtin(x) 0
25#endif
26
27#ifndef RPC_ATTRS
28#if defined(__CUDA__) || defined(__HIP__)
29#define RPC_ATTRS __attribute__((host, device)) inline
30#else
31#define RPC_ATTRS inline
32#endif
33#endif
34
35#ifndef RPC_GLOBAL
36#ifdef RPC_TARGET_IS_GPU
37#define RPC_GLOBAL __gpu_global
38#else
39#define RPC_GLOBAL
40#endif
41#endif
42
43namespace rpc {
44
45template <typename T> struct type_identity {
46 using type = T;
47};
48
49template <typename T, T v> struct type_constant {
50 static inline constexpr T value = v;
51};
52
53/// Freestanding type trait helpers.
54template <typename T> struct remove_cv : type_identity<T> {};
55template <typename T> struct remove_cv<const T> : type_identity<T> {};
56template <typename T> using remove_cv_t = typename remove_cv<T>::type;
57
58template <typename T> struct remove_pointer : type_identity<T> {};
59template <typename T> struct remove_pointer<T *> : type_identity<T> {};
60template <typename T> using remove_pointer_t = typename remove_pointer<T>::type;
61
62template <typename T> struct remove_const : type_identity<T> {};
63template <typename T> struct remove_const<const T> : type_identity<T> {};
64template <typename T> using remove_const_t = typename remove_const<T>::type;
65
66template <typename T> struct remove_reference : type_identity<T> {};
67template <typename T> struct remove_reference<T &> : type_identity<T> {};
68template <typename T> struct remove_reference<T &&> : type_identity<T> {};
69template <typename T>
70using remove_reference_t = typename remove_reference<T>::type;
71
72template <typename T> struct is_const : type_constant<bool, false> {};
73template <typename T> struct is_const<const T> : type_constant<bool, true> {};
74template <typename T> inline constexpr bool is_const_v = is_const<T>::value;
75
76template <typename T> struct is_pointer : type_constant<bool, false> {};
77template <typename T> struct is_pointer<T *> : type_constant<bool, true> {};
78template <typename T>
79struct is_pointer<T *const> : type_constant<bool, true> {};
80template <typename T> inline constexpr bool is_pointer_v = is_pointer<T>::value;
81
82template <typename T, typename U>
83struct is_same : type_constant<bool, false> {};
84template <typename T> struct is_same<T, T> : type_constant<bool, true> {};
85template <typename T, typename U>
86inline constexpr bool is_same_v = is_same<T, U>::value;
87
88template <typename T> struct is_void : type_constant<bool, false> {};
89template <> struct is_void<void> : type_constant<bool, true> {};
90template <typename T> inline constexpr bool is_void_v = is_void<T>::value;
91
92// Scary trait that can change within a TU, use with caution.
93template <typename...> using void_t = void;
94template <typename T, typename = void>
95struct is_complete : type_constant<bool, false> {};
96template <typename T>
97struct is_complete<T, void_t<decltype(sizeof(T))>> : type_constant<bool, true> {
98};
99template <typename T>
100inline constexpr bool is_complete_v = is_complete<T>::value;
101
102template <typename T>
103struct is_trivially_copyable
104 : public type_constant<bool, __is_trivially_copyable(T)> {};
105template <typename T>
106inline constexpr bool is_trivially_copyable_v = is_trivially_copyable<T>::value;
107
108template <typename T, typename... Args>
109struct is_trivially_constructible
110 : type_constant<bool, __is_trivially_constructible(T, Args...)> {};
111template <typename T, typename... Args>
112inline constexpr bool is_trivially_constructible_v =
113 is_trivially_constructible<T>::value;
114
115/// Tag type to indicate an array of elements being passed through RPC.
116template <typename T> struct span {
117 T *data;
118 uint64_t size;
119 RPC_ATTRS operator T *() const { return data; }
120};
121
122template <typename T> struct is_span : type_constant<bool, false> {};
123template <typename T> struct is_span<span<T>> : type_constant<bool, true> {};
124template <typename T> inline constexpr bool is_span_v = is_span<T>::value;
125
126template <typename T> struct remove_span : type_identity<T> {};
127template <typename T> struct remove_span<span<T>> : type_identity<T *> {};
128template <typename T> using remove_span_t = typename remove_span<T>::type;
129
130template <bool B, typename T, typename F>
131struct conditional : type_identity<T> {};
132template <typename T, typename F>
133struct conditional<false, T, F> : type_identity<F> {};
134template <bool B, typename T, typename F>
135using conditional_t = typename conditional<B, T, F>::type;
136
137/// Freestanding implementation of std::move.
138template <typename T>
139RPC_ATTRS constexpr typename remove_reference<T>::type &&move(T &&t) {
140 return static_cast<typename remove_reference<T>::type &&>(t);
141}
142
143/// Freestanding integer sequence.
144template <typename T, T... Ints> struct integer_sequence {
145 template <T Next> using append = integer_sequence<T, Ints..., Next>;
146};
147
148namespace detail {
149template <typename T, int N> struct make_integer_sequence {
150 using type =
151 typename make_integer_sequence<T, N - 1>::type::template append<N>;
152};
153template <typename T> struct make_integer_sequence<T, -1> {
154 using type = integer_sequence<T>;
155};
156} // namespace detail
157
158template <uint64_t... Ints>
159using index_sequence = integer_sequence<uint64_t, Ints...>;
160template <int N>
161using make_index_sequence =
162 typename detail::make_integer_sequence<uint64_t, N - 1>::type;
163template <typename... Ts>
164using index_sequence_for = make_index_sequence<sizeof...(Ts)>;
165
166/// Freestanding implementation of std::forward.
167template <typename T>
168RPC_ATTRS constexpr T &&forward(typename remove_reference<T>::type &value) {
169 return static_cast<T &&>(value);
170}
171template <typename T>
172RPC_ATTRS constexpr T &&forward(typename remove_reference<T>::type &&value) {
173 return static_cast<T &&>(value);
174}
175
176struct in_place_t {
177 RPC_ATTRS explicit in_place_t() = default;
178};
179
180struct nullopt_t {
181 RPC_ATTRS constexpr explicit nullopt_t() = default;
182};
183
184constexpr inline in_place_t in_place{};
185constexpr inline nullopt_t nullopt{};
186
187/// Freestanding and minimal implementation of std::optional.
188template <typename T> struct optional {
189 template <typename U> struct OptionalStorage {
190 union {
191 char empty;
192 U stored_value;
193 };
194
195 bool in_use = false;
196
197 RPC_ATTRS ~OptionalStorage() { reset(); }
198
199 RPC_ATTRS constexpr OptionalStorage() : empty() {}
200
201 template <typename... Args>
202 RPC_ATTRS constexpr explicit OptionalStorage(in_place_t, Args &&...args)
203 : stored_value(forward<Args>(args)...) {
204 in_use = true;
205 }
206
207 RPC_ATTRS constexpr void reset() {
208 if (in_use)
209 stored_value.~U();
210 in_use = false;
211 }
212 };
213
214 OptionalStorage<T> storage;
215
216public:
217 RPC_ATTRS constexpr optional() = default;
218 RPC_ATTRS constexpr optional(nullopt_t) {}
219
220 RPC_ATTRS constexpr optional(const T &t) : storage(in_place, t) {}
221 RPC_ATTRS constexpr optional(const optional &) = default;
222
223 RPC_ATTRS constexpr optional(T &&t) : storage(in_place, move(t)) {}
224 RPC_ATTRS constexpr optional(optional &&O) = default;
225
226 template <typename... Args>
227 RPC_ATTRS constexpr optional(in_place_t, Args &&...args)
228 : storage(in_place, forward<Args>(args)...) {}
229
230 RPC_ATTRS constexpr optional &operator=(T &&t) {
231 storage = move(t);
232 return *this;
233 }
234 RPC_ATTRS constexpr optional &operator=(optional &&) = default;
235
236 RPC_ATTRS constexpr optional &operator=(const T &t) {
237 storage = t;
238 return *this;
239 }
240 RPC_ATTRS constexpr optional &operator=(const optional &) = default;
241
242 RPC_ATTRS constexpr void reset() { storage.reset(); }
243
244 RPC_ATTRS constexpr const T &value() const & { return storage.stored_value; }
245
246 RPC_ATTRS constexpr T &value() & { return storage.stored_value; }
247
248 RPC_ATTRS constexpr explicit operator bool() const { return storage.in_use; }
249 RPC_ATTRS constexpr bool has_value() const { return storage.in_use; }
250 RPC_ATTRS constexpr const T *operator->() const {
251 return &storage.stored_value;
252 }
253 RPC_ATTRS constexpr T *operator->() { return &storage.stored_value; }
254 RPC_ATTRS constexpr const T &operator*() const & {
255 return storage.stored_value;
256 }
257 RPC_ATTRS constexpr T &operator*() & { return storage.stored_value; }
258
259 RPC_ATTRS constexpr T &&value() && { return move(storage.stored_value); }
260 RPC_ATTRS constexpr T &&operator*() && { return move(storage.stored_value); }
261};
262
263/// Minimal array type.
264template <typename T, uint64_t N> struct array {
265 T elems[N];
266
267 RPC_ATTRS constexpr T *data() { return elems; }
268 RPC_ATTRS constexpr const T *data() const { return elems; }
269 RPC_ATTRS static constexpr uint64_t size() { return N; }
270
271 RPC_ATTRS constexpr T &operator[](uint64_t i) { return elems[i]; }
272 RPC_ATTRS constexpr const T &operator[](uint64_t i) const { return elems[i]; }
273};
274
275/// Minimal tuple type.
276template <typename... Ts> struct tuple;
277template <> struct tuple<> {};
278
279template <typename Head, typename... Tail>
280struct tuple<Head, Tail...> : tuple<Tail...> {
281 Head head;
282
283 RPC_ATTRS constexpr tuple() = default;
284
285 template <typename OHead, typename... OTail>
286 RPC_ATTRS constexpr tuple &operator=(const tuple<OHead, OTail...> &other) {
287 head = other.get_head();
288 this->get_tail() = other.get_tail();
289 return *this;
290 }
291
292 RPC_ATTRS constexpr tuple(const Head &h, const Tail &...t)
293 : tuple<Tail...>(t...), head(h) {}
294
295 RPC_ATTRS constexpr Head &get_head() { return head; }
296 RPC_ATTRS constexpr const Head &get_head() const { return head; }
297
298 RPC_ATTRS constexpr tuple<Tail...> &get_tail() { return *this; }
299 RPC_ATTRS constexpr const tuple<Tail...> &get_tail() const { return *this; }
300};
301
302template <size_t Idx, typename T> struct tuple_element;
303template <size_t Idx, typename Head, typename... Tail>
304struct tuple_element<Idx, tuple<Head, Tail...>>
305 : tuple_element<Idx - 1, tuple<Tail...>> {};
306template <typename Head, typename... Tail>
307struct tuple_element<0, tuple<Head, Tail...>> {
308 using type = remove_cv_t<remove_reference_t<Head>>;
309};
310template <size_t Idx, typename T>
311using tuple_element_t = typename tuple_element<Idx, T>::type;
312
313template <uint64_t Idx, typename Head, typename... Tail>
314RPC_ATTRS constexpr auto &get(tuple<Head, Tail...> &t) {
315 if constexpr (Idx == 0)
316 return t.get_head();
317 else
318 return get<Idx - 1>(t.get_tail());
319}
320template <uint64_t Idx, typename Head, typename... Tail>
321RPC_ATTRS constexpr const auto &get(const tuple<Head, Tail...> &t) {
322 if constexpr (Idx == 0)
323 return t.get_head();
324 else
325 return get<Idx - 1>(t.get_tail());
326}
327
328namespace detail {
329template <typename F, typename Tuple, uint64_t... Is>
330RPC_ATTRS auto apply(F &&f, Tuple &&t, index_sequence<Is...>) {
331 return f(get<Is>(static_cast<Tuple &&>(t))...);
332}
333} // namespace detail
334
335template <typename F, typename... Ts>
336RPC_ATTRS auto apply(F &&f, tuple<Ts...> &t) {
337 return detail::apply(static_cast<F &&>(f), t,
338 make_index_sequence<sizeof...(Ts)>{});
339}
340
341/// Suspend the thread briefly to assist the thread scheduler during busy loops.
342RPC_ATTRS void sleep_briefly() {
343#if __has_builtin(__nvvm_reflect)
344 if (__nvvm_reflect("__CUDA_ARCH") >= 700)
345 asm("nanosleep.u32 64;" ::: "memory");
346#elif __has_builtin(__builtin_amdgcn_s_sleep)
347 __builtin_amdgcn_s_sleep(2);
348#elif __has_builtin(__builtin_ia32_pause)
349 __builtin_ia32_pause();
350#elif __has_builtin(__builtin_arm_isb)
351 __builtin_arm_isb(0xf);
352#else
353 // Simply do nothing if sleeping isn't supported on this platform.
354#endif
355}
356
357/// Conditional to indicate if this process is running on the GPU.
358RPC_ATTRS constexpr bool is_process_gpu() {
359#ifdef RPC_TARGET_IS_GPU
360 return true;
361#else
362 return false;
363#endif
364}
365
366/// Wait for all lanes in the group to complete.
367RPC_ATTRS void sync_lane([[maybe_unused]] uint64_t lane_mask) {
368#ifdef RPC_TARGET_IS_GPU
369 return __gpu_sync_lane(lane_mask);
370#endif
371}
372
373/// Copies the value from the first active thread to the rest.
374RPC_ATTRS uint32_t broadcast_value([[maybe_unused]] uint64_t lane_mask,
375 uint32_t x) {
376#ifdef RPC_TARGET_IS_GPU
377 return __gpu_read_first_lane_u32(lane_mask, x);
378#else
379 return x;
380#endif
381}
382
383/// Returns the number lanes that participate in the RPC interface.
384RPC_ATTRS uint32_t get_num_lanes() {
385#ifdef RPC_TARGET_IS_GPU
386 return __gpu_num_lanes();
387#else
388 return 1;
389#endif
390}
391
392/// Returns a bitmask of the currently active lanes.
393RPC_ATTRS uint64_t get_lane_mask() {
394#ifdef RPC_TARGET_IS_GPU
395 return __gpu_lane_mask();
396#else
397 return 1;
398#endif
399}
400
401/// Returns the id of the thread inside of an AMD wavefront executing together.
402RPC_ATTRS uint32_t get_lane_id() {
403#ifdef RPC_TARGET_IS_GPU
404 return __gpu_lane_id();
405#else
406 return 0;
407#endif
408}
409
410/// Conditional that is only true for a single thread in a lane.
411RPC_ATTRS bool is_first_lane([[maybe_unused]] uint64_t lane_mask) {
412#ifdef RPC_TARGET_IS_GPU
413 return __gpu_is_first_in_lane(lane_mask);
414#else
415 return true;
416#endif
417}
418
419/// Returns a bitmask of threads in the current lane for which \p x is true.
420RPC_ATTRS uint64_t ballot([[maybe_unused]] uint64_t lane_mask, bool x) {
421#ifdef RPC_TARGET_IS_GPU
422 return __gpu_ballot(lane_mask, x);
423#else
424 return x;
425#endif
426}
427
428/// Signal an interrupt from the device to wake the server. Only supported for
429/// AMDGPU targets currently.
430RPC_ATTRS void signal_interrupt([[maybe_unused]] uint32_t event_id) {
431#ifdef __AMDGPU__
432 constexpr uint32_t MSG_INTERRUPT = 1;
433 __builtin_amdgcn_s_sendmsg(MSG_INTERRUPT, event_id);
434#endif
435}
436
437/// Return \p val aligned "upwards" according to \p align.
438template <typename V, typename A>
439RPC_ATTRS constexpr V align_up(V val, A align) {
440 return ((val + V(align) - 1) / V(align)) * V(align);
441}
442
443/// Utility to provide a unified interface between the CPU and GPU's memory
444/// model. On the GPU stack variables are always private to a lane so we can
445/// simply use the variable passed in. On the CPU we need to allocate enough
446/// space for the whole lane and index into it.
447template <typename V> RPC_ATTRS V &lane_value(V *val, uint32_t id) {
448 if constexpr (is_process_gpu())
449 return *val;
450 return val[id];
451}
452
453/// Advance the \p p by \p bytes.
454template <typename T, typename U> RPC_ATTRS T *advance(T *ptr, U bytes) {
455 if constexpr (is_const<T>::value)
456 return reinterpret_cast<T *>(reinterpret_cast<const uint8_t *>(ptr) +
457 bytes);
458 else
459 return reinterpret_cast<T *>(reinterpret_cast<uint8_t *>(ptr) + bytes);
460}
461
462/// Wrapper around the optimal memory copy implementation for the target.
463template <typename D, typename S>
464RPC_ATTRS void rpc_memcpy(D *dst, S *src, uint64_t count) {
465#if __has_builtin(__builtin_memcpy)
466 if (count)
467 __builtin_memcpy(dst, src, count);
468#else
469 // The casts are C-style as they may need to change the address space.
470 for (uint64_t i = 0; i < count; ++i)
471 ((uint8_t *)dst)[i] = ((const uint8_t *)src)[i];
472#endif
473}
474
475/// Minimal string length function.
476template <typename T> RPC_ATTRS constexpr uint64_t string_length(const T *s) {
477 const T *end = s;
478 for (; *end != '\0'; ++end)
479 ;
480 return static_cast<uint64_t>(end - s + 1);
481}
482
483/// Helper for dealing with function pointers and lambda types.
484template <typename> struct function_traits;
485template <typename R, typename... Args> struct function_traits<R (*)(Args...)> {
486 using return_type = R;
487 using arg_types = rpc::tuple<Args...>;
488 static constexpr uint64_t ARITY = sizeof...(Args);
489};
490template <typename R, typename... Args>
491struct function_traits<R (*)(Args...) noexcept> {
492 using return_type = R;
493 using arg_types = rpc::tuple<Args...>;
494 static constexpr uint64_t ARITY = sizeof...(Args);
495};
496template <typename T> T &&declval();
497template <typename T>
498struct function_traits
499 : function_traits<decltype(+declval<rpc::remove_reference_t<T>>())> {};
500
501template <typename T, typename U>
502RPC_ATTRS constexpr T max(const T &a, const U &b) {
503 return (a < b) ? b : a;
504}
505
506template <typename T, typename U>
507RPC_ATTRS constexpr T min(const T &a, const U &b) {
508 return (a < b) ? a : b;
509}
510
511} // namespace rpc
512
513#endif // LLVM_LIBC_SHARED_RPC_UTIL_H
514