libc/src/__support/GPU/amdgpu/utils.h

   1 //===-------------- AMDGPU implementation of GPU utils ----------*- 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_SRC_SUPPORT_GPU_AMDGPU_IO_H
  10 #define LLVM_LIBC_SRC_SUPPORT_GPU_AMDGPU_IO_H
  11
  12 #include "src/__support/common.h"
  13 #include "src/__support/macros/config.h"
  14
  15 #include <stdint.h>
  16
  17 namespace __llvm_libc {
  18 namespace gpu {
  19
  20 /// The number of threads that execute in lock-step in a lane.
  21 constexpr const uint64_t LANE_SIZE = __AMDGCN_WAVEFRONT_SIZE;
  22
  23 /// Returns the number of workgroups in the 'x' dimension of the grid.
  24 LIBC_INLINE uint32_t get_num_blocks_x() {
  25   return __builtin_amdgcn_grid_size_x() / __builtin_amdgcn_workgroup_size_x();
  26 }
  27
  28 /// Returns the number of workgroups in the 'y' dimension of the grid.
  29 LIBC_INLINE uint32_t get_num_blocks_y() {
  30   return __builtin_amdgcn_grid_size_y() / __builtin_amdgcn_workgroup_size_y();
  31 }
  32
  33 /// Returns the number of workgroups in the 'z' dimension of the grid.
  34 LIBC_INLINE uint32_t get_num_blocks_z() {
  35   return __builtin_amdgcn_grid_size_z() / __builtin_amdgcn_workgroup_size_z();
  36 }
  37
  38 /// Returns the total number of workgruops in the grid.
  39 LIBC_INLINE uint64_t get_num_blocks() {
  40   return get_num_blocks_x() * get_num_blocks_y() * get_num_blocks_z();
  41 }
  42
  43 /// Returns the 'x' dimension of the current AMD workgroup's id.
  44 LIBC_INLINE uint32_t get_block_id_x() {
  45   return __builtin_amdgcn_workgroup_id_x();
  46 }
  47
  48 /// Returns the 'y' dimension of the current AMD workgroup's id.
  49 LIBC_INLINE uint32_t get_block_id_y() {
  50   return __builtin_amdgcn_workgroup_id_y();
  51 }
  52
  53 /// Returns the 'z' dimension of the current AMD workgroup's id.
  54 LIBC_INLINE uint32_t get_block_id_z() {
  55   return __builtin_amdgcn_workgroup_id_z();
  56 }
  57
  58 /// Returns the absolute id of the AMD workgroup.
  59 LIBC_INLINE uint64_t get_block_id() {
  60   return get_block_id_x() + get_num_blocks_x() * get_block_id_y() +
  61          get_num_blocks_x() * get_num_blocks_y() * get_block_id_z();
  62 }
  63
  64 /// Returns the number of workitems in the 'x' dimension.
  65 LIBC_INLINE uint32_t get_num_threads_x() {
  66   return __builtin_amdgcn_workgroup_size_x();
  67 }
  68
  69 /// Returns the number of workitems in the 'y' dimension.
  70 LIBC_INLINE uint32_t get_num_threads_y() {
  71   return __builtin_amdgcn_workgroup_size_y();
  72 }
  73
  74 /// Returns the number of workitems in the 'z' dimension.
  75 LIBC_INLINE uint32_t get_num_threads_z() {
  76   return __builtin_amdgcn_workgroup_size_z();
  77 }
  78
  79 /// Returns the total number of workitems in the workgroup.
  80 LIBC_INLINE uint64_t get_num_threads() {
  81   return get_num_threads_x() * get_num_threads_y() * get_num_threads_z();
  82 }
  83
  84 /// Returns the 'x' dimension id of the workitem in the current AMD workgroup.
  85 LIBC_INLINE uint32_t get_thread_id_x() {
  86   return __builtin_amdgcn_workitem_id_x();
  87 }
  88
  89 /// Returns the 'y' dimension id of the workitem in the current AMD workgroup.
  90 LIBC_INLINE uint32_t get_thread_id_y() {
  91   return __builtin_amdgcn_workitem_id_y();
  92 }
  93
  94 /// Returns the 'z' dimension id of the workitem in the current AMD workgroup.
  95 LIBC_INLINE uint32_t get_thread_id_z() {
  96   return __builtin_amdgcn_workitem_id_z();
  97 }
  98
  99 /// Returns the absolute id of the thread in the current AMD workgroup.
 100 LIBC_INLINE uint64_t get_thread_id() {
 101   return get_thread_id_x() + get_num_threads_x() * get_thread_id_y() +
 102          get_num_threads_x() * get_num_threads_y() * get_thread_id_z();
 103 }
 104
 105 /// Returns the size of an AMD wavefront. Either 32 or 64 depending on hardware.
 106 LIBC_INLINE uint32_t get_lane_size() { return LANE_SIZE; }
 107
 108 /// Returns the id of the thread inside of an AMD wavefront executing together.
 109 [[clang::convergent]] LIBC_INLINE uint32_t get_lane_id() {
 110   if constexpr (LANE_SIZE == 64)
 111     return __builtin_amdgcn_mbcnt_hi(~0u, __builtin_amdgcn_mbcnt_lo(~0u, 0u));
 112   else
 113     return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
 114 }
 115
 116 /// Returns the bit-mask of active threads in the current wavefront.
 117 [[clang::convergent]] LIBC_INLINE uint64_t get_lane_mask() {
 118   return __builtin_amdgcn_read_exec();
 119 }
 120
 121 /// Copies the value from the first active thread in the wavefront to the rest.
 122 [[clang::convergent]] LIBC_INLINE uint32_t broadcast_value(uint32_t x) {
 123   return __builtin_amdgcn_readfirstlane(x);
 124 }
 125
 126 /// Returns a bitmask of threads in the current lane for which \p x is true.
 127 [[clang::convergent]] LIBC_INLINE uint64_t ballot(uint64_t lane_mask, bool x) {
 128   // the lane_mask & gives the nvptx semantics when lane_mask is a subset of
 129   // the active threads
 130   if constexpr (LANE_SIZE == 64) {
 131     return lane_mask & __builtin_amdgcn_ballot_w64(x);
 132   } else {
 133     return lane_mask & __builtin_amdgcn_ballot_w32(x);
 134   }
 135 }
 136
 137 /// Waits for all the threads in the block to converge and issues a fence.
 138 [[clang::convergent]] LIBC_INLINE void sync_threads() {
 139   __builtin_amdgcn_s_barrier();
 140   __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
 141 }
 142
 143 /// Wait for all threads in the wavefront to converge, this is a noop on AMDGPU.
 144 [[clang::convergent]] LIBC_INLINE void sync_lane(uint64_t) {
 145   __builtin_amdgcn_wave_barrier();
 146 }
 147
 148 /// Returns the current value of the GPU's processor clock.
 149 /// NOTE: The RDNA3 and RDNA2 architectures use a 20-bit cycle cycle counter.
 150 LIBC_INLINE uint64_t processor_clock() {
 151   if constexpr (LIBC_HAS_BUILTIN(__builtin_amdgcn_s_memtime))
 152     return __builtin_amdgcn_s_memtime();
 153   else if constexpr (LIBC_HAS_BUILTIN(__builtin_readcyclecounter))
 154     return __builtin_readcyclecounter();
 155   else
 156     return 0;
 157 }
 158
 159 /// Returns a fixed-frequency timestamp. The actual frequency is dependent on
 160 /// the card and can only be queried via the driver.
 161 LIBC_INLINE uint64_t fixed_frequency_clock() {
 162   if constexpr (LIBC_HAS_BUILTIN(__builtin_amdgcn_s_sendmsg_rtnl))
 163     return __builtin_amdgcn_s_sendmsg_rtnl(0x83);
 164   else if constexpr (LIBC_HAS_BUILTIN(__builtin_amdgcn_s_memrealtime))
 165     return __builtin_amdgcn_s_memrealtime();
 166   else if constexpr (LIBC_HAS_BUILTIN(__builtin_amdgcn_s_memtime))
 167     return __builtin_amdgcn_s_memtime();
 168   else
 169     return 0;
 170 }
 171
 172 } // namespace gpu
 173 } // namespace __llvm_libc
 174
 175 #endif