241 lines · plain
1//===----------------------------------------------------------------------===//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 __OMPX_H10#define __OMPX_H11 12#if (defined(__NVPTX__) || defined(__AMDGPU__))13#include <gpuintrin.h>14#define __OMPX_TARGET_IS_GPU15#endif16 17typedef unsigned long uint64_t;18typedef unsigned int uint32_t;19 20static inline uint32_t __warpSize(void) {21#ifdef __OMPX_TARGET_IS_GPU22 return __gpu_num_lanes();23#else24 __builtin_trap();25#endif26}27 28#ifdef __cplusplus29extern "C" {30#endif31 32int omp_get_ancestor_thread_num(int);33int omp_get_team_size(int);34 35#ifdef __cplusplus36}37#endif38 39/// Target kernel language extensions40///41/// These extensions exist for the host to allow fallback implementations,42/// however, they cannot be arbitrarily composed with OpenMP. If the rules of43/// the kernel language are followed, the host fallbacks should behave as44/// expected since the kernel is represented as 3 sequential outer loops, one45/// for each grid dimension, and three (nested) parallel loops, one for each46/// block dimension. This fallback is not supposed to be optimal and should be47/// configurable by the user.48///49///{50 51#ifdef __cplusplus52extern "C" {53#endif54 55enum {56 ompx_relaxed = __ATOMIC_RELAXED,57 ompx_aquire = __ATOMIC_ACQUIRE,58 ompx_release = __ATOMIC_RELEASE,59 ompx_acq_rel = __ATOMIC_ACQ_REL,60 ompx_seq_cst = __ATOMIC_SEQ_CST,61};62 63enum {64 ompx_dim_x = 0,65 ompx_dim_y = 1,66 ompx_dim_z = 2,67};68 69// TODO: The following implementation is for host fallback. We need to disable70// generation of host fallback in kernel language mode.71#pragma omp begin declare variant match(device = {kind(cpu)})72 73/// ompx_{thread,block}_{id,dim}74///{75#define _TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_C(NAME, VALUE) \76 static inline int ompx_##NAME(int Dim) { return VALUE; }77 78_TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_C(thread_id,79 omp_get_ancestor_thread_num(Dim + 1))80_TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_C(block_dim, omp_get_team_size(Dim + 1))81_TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_C(block_id, 0)82_TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_C(grid_dim, 1)83#undef _TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_C84///}85 86/// ompx_{sync_block}_{,divergent}87///{88#define _TGT_KERNEL_LANGUAGE_HOST_IMPL_SYNC_C(RETTY, NAME, ARGS, BODY) \89 static inline RETTY ompx_##NAME(ARGS) { BODY; }90 91_TGT_KERNEL_LANGUAGE_HOST_IMPL_SYNC_C(void, sync_block, int Ordering,92 _Pragma("omp barrier"))93_TGT_KERNEL_LANGUAGE_HOST_IMPL_SYNC_C(void, sync_block_acq_rel, void,94 ompx_sync_block(ompx_acq_rel))95_TGT_KERNEL_LANGUAGE_HOST_IMPL_SYNC_C(void, sync_block_divergent, int Ordering,96 ompx_sync_block(Ordering))97#undef _TGT_KERNEL_LANGUAGE_HOST_IMPL_SYNC_C98///}99 100static inline uint64_t ompx_ballot_sync(uint64_t mask, int pred) {101 __builtin_trap();102}103 104/// ompx_shfl_down_sync_{i,f,l,d}105///{106#define _TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC_HOST_IMPL(TYPE, TY) \107 static inline TYPE ompx_shfl_down_sync_##TY(uint64_t mask, TYPE var, \108 unsigned delta, int width) { \109 __builtin_trap(); \110 }111 112_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC_HOST_IMPL(int, i)113_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC_HOST_IMPL(float, f)114_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC_HOST_IMPL(long, l)115_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC_HOST_IMPL(double, d)116 117#undef _TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC_HOST_IMPL118///}119 120#pragma omp end declare variant121 122/// ompx_{sync_block}_{,divergent}123///{124#define _TGT_KERNEL_LANGUAGE_DECL_SYNC_C(RETTY, NAME, ARGS) \125 RETTY ompx_##NAME(ARGS);126 127_TGT_KERNEL_LANGUAGE_DECL_SYNC_C(void, sync_block, int Ordering)128_TGT_KERNEL_LANGUAGE_DECL_SYNC_C(void, sync_block_acq_rel, void)129_TGT_KERNEL_LANGUAGE_DECL_SYNC_C(void, sync_block_divergent, int Ordering)130#undef _TGT_KERNEL_LANGUAGE_DECL_SYNC_C131///}132 133/// ompx_{thread,block}_{id,dim}_{x,y,z}134///{135#define _TGT_KERNEL_LANGUAGE_DECL_GRID_C(NAME) \136 int ompx_##NAME(int Dim); \137 static inline int ompx_##NAME##_x() { return ompx_##NAME(ompx_dim_x); } \138 static inline int ompx_##NAME##_y() { return ompx_##NAME(ompx_dim_y); } \139 static inline int ompx_##NAME##_z() { return ompx_##NAME(ompx_dim_z); }140 141_TGT_KERNEL_LANGUAGE_DECL_GRID_C(thread_id)142_TGT_KERNEL_LANGUAGE_DECL_GRID_C(block_dim)143_TGT_KERNEL_LANGUAGE_DECL_GRID_C(block_id)144_TGT_KERNEL_LANGUAGE_DECL_GRID_C(grid_dim)145#undef _TGT_KERNEL_LANGUAGE_DECL_GRID_C146///}147 148uint64_t ompx_ballot_sync(uint64_t mask, int pred);149 150/// ompx_shfl_down_sync_{i,f,l,d}151///{152#define _TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(TYPE, TY) \153 TYPE ompx_shfl_down_sync_##TY(uint64_t mask, TYPE var, unsigned delta, \154 int width);155 156_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(int, i)157_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(float, f)158_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(long, l)159_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(double, d)160 161#undef _TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC162///}163 164#ifdef __cplusplus165}166#endif167 168#ifdef __cplusplus169 170namespace ompx {171 172enum {173 dim_x = ompx_dim_x,174 dim_y = ompx_dim_y,175 dim_z = ompx_dim_z,176};177 178enum {179 relaxed = ompx_relaxed ,180 aquire = ompx_aquire,181 release = ompx_release,182 acc_rel = ompx_acq_rel,183 seq_cst = ompx_seq_cst,184};185 186/// ompx::{thread,block}_{id,dim}_{,x,y,z}187///{188#define _TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_CXX(NAME) \189 static inline int NAME(int Dim) noexcept { return ompx_##NAME(Dim); } \190 static inline int NAME##_x() noexcept { return NAME(ompx_dim_x); } \191 static inline int NAME##_y() noexcept { return NAME(ompx_dim_y); } \192 static inline int NAME##_z() noexcept { return NAME(ompx_dim_z); }193 194_TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_CXX(thread_id)195_TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_CXX(block_dim)196_TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_CXX(block_id)197_TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_CXX(grid_dim)198#undef _TGT_KERNEL_LANGUAGE_HOST_IMPL_GRID_CXX199///}200 201/// ompx_{sync_block}_{,divergent}202///{203#define _TGT_KERNEL_LANGUAGE_HOST_IMPL_SYNC_CXX(RETTY, NAME, ARGS, CALL_ARGS) \204 static inline RETTY NAME(ARGS) { \205 return ompx_##NAME(CALL_ARGS); \206 }207 208_TGT_KERNEL_LANGUAGE_HOST_IMPL_SYNC_CXX(void, sync_block, int Ordering = acc_rel,209 Ordering)210_TGT_KERNEL_LANGUAGE_HOST_IMPL_SYNC_CXX(void, sync_block_divergent,211 int Ordering = acc_rel, Ordering)212#undef _TGT_KERNEL_LANGUAGE_HOST_IMPL_SYNC_CXX213///}214 215static inline uint64_t ballot_sync(uint64_t mask, int pred) {216 return ompx_ballot_sync(mask, pred);217}218 219/// shfl_down_sync220///{221#define _TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(TYPE, TY) \222 static inline TYPE shfl_down_sync(uint64_t mask, TYPE var, unsigned delta, \223 int width = __warpSize()) { \224 return ompx_shfl_down_sync_##TY(mask, var, delta, width); \225 }226 227_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(int, i)228_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(float, f)229_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(long, l)230_TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC(double, d)231 232#undef _TGT_KERNEL_LANGUAGE_SHFL_DOWN_SYNC233///}234 235} // namespace ompx236#endif237 238///}239 240#endif /* __OMPX_H */241