277 lines · c
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-exception6//7//===----------------------------------------------------------------------===//8 9#ifndef LLVM_LIBC_SHARED_RPC_UTIL_H10#define LLVM_LIBC_SHARED_RPC_UTIL_H11 12#include <stddef.h>13#include <stdint.h>14 15#if (defined(__NVPTX__) || defined(__AMDGPU__)) && \16 !((defined(__CUDA__) && !defined(__CUDA_ARCH__)) || \17 (defined(__HIP__) && !defined(__HIP_DEVICE_COMPILE__)))18#include <gpuintrin.h>19#define RPC_TARGET_IS_GPU20#endif21 22// Workaround for missing __has_builtin in < GCC 10.23#ifndef __has_builtin24#define __has_builtin(x) 025#endif26 27#ifndef RPC_ATTRS28#if defined(__CUDA__) || defined(__HIP__)29#define RPC_ATTRS __attribute__((host, device)) inline30#else31#define RPC_ATTRS inline32#endif33#endif34 35namespace rpc {36 37template <typename T> struct type_identity {38 using type = T;39};40 41template <class T, T v> struct type_constant {42 static inline constexpr T value = v;43};44 45template <class T> struct remove_reference : type_identity<T> {};46template <class T> struct remove_reference<T &> : type_identity<T> {};47template <class T> struct remove_reference<T &&> : type_identity<T> {};48 49template <class T> struct is_const : type_constant<bool, false> {};50template <class T> struct is_const<const T> : type_constant<bool, true> {};51 52/// Freestanding implementation of std::move.53template <class T>54RPC_ATTRS constexpr typename remove_reference<T>::type &&move(T &&t) {55 return static_cast<typename remove_reference<T>::type &&>(t);56}57 58/// Freestanding implementation of std::forward.59template <typename T>60RPC_ATTRS constexpr T &&forward(typename remove_reference<T>::type &value) {61 return static_cast<T &&>(value);62}63template <typename T>64RPC_ATTRS constexpr T &&forward(typename remove_reference<T>::type &&value) {65 return static_cast<T &&>(value);66}67 68struct in_place_t {69 RPC_ATTRS explicit in_place_t() = default;70};71 72struct nullopt_t {73 RPC_ATTRS constexpr explicit nullopt_t() = default;74};75 76constexpr inline in_place_t in_place{};77constexpr inline nullopt_t nullopt{};78 79/// Freestanding and minimal implementation of std::optional.80template <typename T> class optional {81 template <typename U> struct OptionalStorage {82 union {83 char empty;84 U stored_value;85 };86 87 bool in_use = false;88 89 RPC_ATTRS ~OptionalStorage() { reset(); }90 91 RPC_ATTRS constexpr OptionalStorage() : empty() {}92 93 template <typename... Args>94 RPC_ATTRS constexpr explicit OptionalStorage(in_place_t, Args &&...args)95 : stored_value(forward<Args>(args)...) {}96 97 RPC_ATTRS constexpr void reset() {98 if (in_use)99 stored_value.~U();100 in_use = false;101 }102 };103 104 OptionalStorage<T> storage;105 106public:107 RPC_ATTRS constexpr optional() = default;108 RPC_ATTRS constexpr optional(nullopt_t) {}109 110 RPC_ATTRS constexpr optional(const T &t) : storage(in_place, t) {111 storage.in_use = true;112 }113 RPC_ATTRS constexpr optional(const optional &) = default;114 115 RPC_ATTRS constexpr optional(T &&t) : storage(in_place, move(t)) {116 storage.in_use = true;117 }118 RPC_ATTRS constexpr optional(optional &&O) = default;119 120 RPC_ATTRS constexpr optional &operator=(T &&t) {121 storage = move(t);122 return *this;123 }124 RPC_ATTRS constexpr optional &operator=(optional &&) = default;125 126 RPC_ATTRS constexpr optional &operator=(const T &t) {127 storage = t;128 return *this;129 }130 RPC_ATTRS constexpr optional &operator=(const optional &) = default;131 132 RPC_ATTRS constexpr void reset() { storage.reset(); }133 134 RPC_ATTRS constexpr const T &value() const & { return storage.stored_value; }135 136 RPC_ATTRS constexpr T &value() & { return storage.stored_value; }137 138 RPC_ATTRS constexpr explicit operator bool() const { return storage.in_use; }139 RPC_ATTRS constexpr bool has_value() const { return storage.in_use; }140 RPC_ATTRS constexpr const T *operator->() const {141 return &storage.stored_value;142 }143 RPC_ATTRS constexpr T *operator->() { return &storage.stored_value; }144 RPC_ATTRS constexpr const T &operator*() const & {145 return storage.stored_value;146 }147 RPC_ATTRS constexpr T &operator*() & { return storage.stored_value; }148 149 RPC_ATTRS constexpr T &&value() && { return move(storage.stored_value); }150 RPC_ATTRS constexpr T &&operator*() && { return move(storage.stored_value); }151};152 153/// Suspend the thread briefly to assist the thread scheduler during busy loops.154RPC_ATTRS void sleep_briefly() {155#if __has_builtin(__nvvm_reflect)156 if (__nvvm_reflect("__CUDA_ARCH") >= 700)157 asm("nanosleep.u32 64;" ::: "memory");158#elif __has_builtin(__builtin_amdgcn_s_sleep)159 __builtin_amdgcn_s_sleep(2);160#elif __has_builtin(__builtin_ia32_pause)161 __builtin_ia32_pause();162#elif __has_builtin(__builtin_arm_isb)163 __builtin_arm_isb(0xf);164#else165 // Simply do nothing if sleeping isn't supported on this platform.166#endif167}168 169/// Conditional to indicate if this process is running on the GPU.170RPC_ATTRS constexpr bool is_process_gpu() {171#ifdef RPC_TARGET_IS_GPU172 return true;173#else174 return false;175#endif176}177 178/// Wait for all lanes in the group to complete.179RPC_ATTRS void sync_lane([[maybe_unused]] uint64_t lane_mask) {180#ifdef RPC_TARGET_IS_GPU181 return __gpu_sync_lane(lane_mask);182#endif183}184 185/// Copies the value from the first active thread to the rest.186RPC_ATTRS uint32_t broadcast_value([[maybe_unused]] uint64_t lane_mask,187 uint32_t x) {188#ifdef RPC_TARGET_IS_GPU189 return __gpu_read_first_lane_u32(lane_mask, x);190#else191 return x;192#endif193}194 195/// Returns the number lanes that participate in the RPC interface.196RPC_ATTRS uint32_t get_num_lanes() {197#ifdef RPC_TARGET_IS_GPU198 return __gpu_num_lanes();199#else200 return 1;201#endif202}203 204/// Returns the id of the thread inside of an AMD wavefront executing together.205RPC_ATTRS uint64_t get_lane_mask() {206#ifdef RPC_TARGET_IS_GPU207 return __gpu_lane_mask();208#else209 return 1;210#endif211}212 213/// Returns the id of the thread inside of an AMD wavefront executing together.214RPC_ATTRS uint32_t get_lane_id() {215#ifdef RPC_TARGET_IS_GPU216 return __gpu_lane_id();217#else218 return 0;219#endif220}221 222/// Conditional that is only true for a single thread in a lane.223RPC_ATTRS bool is_first_lane([[maybe_unused]] uint64_t lane_mask) {224#ifdef RPC_TARGET_IS_GPU225 return __gpu_is_first_in_lane(lane_mask);226#else227 return true;228#endif229}230 231/// Returns a bitmask of threads in the current lane for which \p x is true.232RPC_ATTRS uint64_t ballot([[maybe_unused]] uint64_t lane_mask, bool x) {233#ifdef RPC_TARGET_IS_GPU234 return __gpu_ballot(lane_mask, x);235#else236 return x;237#endif238}239 240/// Return \p val aligned "upwards" according to \p align.241template <typename V, typename A>242RPC_ATTRS constexpr V align_up(V val, A align) {243 return ((val + V(align) - 1) / V(align)) * V(align);244}245 246/// Utility to provide a unified interface between the CPU and GPU's memory247/// model. On the GPU stack variables are always private to a lane so we can248/// simply use the variable passed in. On the CPU we need to allocate enough249/// space for the whole lane and index into it.250template <typename V> RPC_ATTRS V &lane_value(V *val, uint32_t id) {251 if constexpr (is_process_gpu())252 return *val;253 return val[id];254}255 256/// Advance the \p p by \p bytes.257template <typename T, typename U> RPC_ATTRS T *advance(T *ptr, U bytes) {258 if constexpr (is_const<T>::value)259 return reinterpret_cast<T *>(reinterpret_cast<const uint8_t *>(ptr) +260 bytes);261 else262 return reinterpret_cast<T *>(reinterpret_cast<uint8_t *>(ptr) + bytes);263}264 265/// Wrapper around the optimal memory copy implementation for the target.266RPC_ATTRS void rpc_memcpy(void *dst, const void *src, size_t count) {267 __builtin_memcpy(dst, src, count);268}269 270template <class T> RPC_ATTRS constexpr const T &max(const T &a, const T &b) {271 return (a < b) ? b : a;272}273 274} // namespace rpc275 276#endif // LLVM_LIBC_SHARED_RPC_UTIL_H277