codekingpro/portable-devtools
114k
1//===----------------------------------------------------------------------===//2//3// Part of libcu++, the C++ Standard Library for your entire system,4// under the Apache License v2.0 with LLVM Exceptions.5// See https://llvm.org/LICENSE.txt for license information.6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception7// SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES.8//9//===----------------------------------------------------------------------===//10 11#ifndef _CUDA_BARRIER12#define _CUDA_BARRIER13 14#include "std/barrier"15 16// Forward-declare CUtensorMap for use in cp_async_bulk_tensor_* PTX wrapping17// functions. These functions take a pointer to CUtensorMap, so do not need to18// know its size. This type is defined in cuda.h (driver API) as:19//20// typedef struct CUtensorMap_st { [ .. snip .. ] } CUtensorMap;21//22// We need to forward-declare both CUtensorMap_st (the struct) and CUtensorMap23// (the typedef):24struct CUtensorMap_st;25typedef struct CUtensorMap_st CUtensorMap;26 27_LIBCUDACXX_BEGIN_NAMESPACE_CUDA_DEVICE_EXPERIMENTAL28 29// Experimental exposure of TMA PTX:30//31// - cp_async_bulk_global_to_shared32// - cp_async_bulk_shared_to_global33// - cp_async_bulk_tensor_{1,2,3,4,5}d_global_to_shared34// - cp_async_bulk_tensor_{1,2,3,4,5}d_shared_to_global35// - fence_proxy_async_shared_cta36// - cp_async_bulk_commit_group37// - cp_async_bulk_wait_group_read<0, …, 7>38 39// These PTX wrappers are only available when the code is compiled compute40// capability 9.0 and above. The check for (!defined(__CUDA_MINIMUM_ARCH__)) is41// necessary to prevent cudafe from ripping out the device functions before42// device compilation begins.43#ifdef __cccl_lib_experimental_ctk12_cp_async_exposure44 45// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk46inline _LIBCUDACXX_DEVICE47void cp_async_bulk_global_to_shared(void *__dest, const void *__src, _CUDA_VSTD::uint32_t __size, ::cuda::barrier<::cuda::thread_scope_block> &__bar)48{49 _LIBCUDACXX_DEBUG_ASSERT(__size % 16 == 0, "Size must be multiple of 16.");50 _LIBCUDACXX_DEBUG_ASSERT(__isShared(__dest), "Destination must be shared memory address.");51 _LIBCUDACXX_DEBUG_ASSERT(__isGlobal(__src), "Source must be global memory address.");52 53 asm volatile(54 "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n"55 :56 : "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__dest))),57 "l"(static_cast<_CUDA_VSTD::uint64_t>(__cvta_generic_to_global(__src))),58 "r"(__size),59 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(::cuda::device::barrier_native_handle(__bar))))60 : "memory");61}62 63 64// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk65inline _LIBCUDACXX_DEVICE66void cp_async_bulk_shared_to_global(void *__dest, const void * __src, _CUDA_VSTD::uint32_t __size)67{68 _LIBCUDACXX_DEBUG_ASSERT(__size % 16 == 0, "Size must be multiple of 16.");69 _LIBCUDACXX_DEBUG_ASSERT(__isGlobal(__dest), "Destination must be global memory address.");70 _LIBCUDACXX_DEBUG_ASSERT(__isShared(__src), "Source must be shared memory address.");71 72 asm volatile(73 "cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;\n"74 :75 : "l"(static_cast<_CUDA_VSTD::uint64_t>(__cvta_generic_to_global(__dest))),76 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__src))),77 "r"(__size)78 : "memory");79}80 81// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor82inline _LIBCUDACXX_DEVICE83void cp_async_bulk_tensor_1d_global_to_shared(84 void *__dest, const CUtensorMap *__tensor_map , int __c0, ::cuda::barrier<::cuda::thread_scope_block> &__bar)85{86 asm volatile(87 "cp.async.bulk.tensor.1d.shared::cluster.global.tile.mbarrier::complete_tx::bytes "88 "[%0], [%1, {%2}], [%3];\n"89 :90 : "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__dest))),91 "l"(__tensor_map),92 "r"(__c0),93 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(::cuda::device::barrier_native_handle(__bar))))94 : "memory");95}96 97// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor98inline _LIBCUDACXX_DEVICE99void cp_async_bulk_tensor_2d_global_to_shared(100 void *__dest, const CUtensorMap *__tensor_map , int __c0, int __c1, ::cuda::barrier<::cuda::thread_scope_block> &__bar)101{102 asm volatile(103 "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes "104 "[%0], [%1, {%2, %3}], [%4];\n"105 :106 : "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__dest))),107 "l"(__tensor_map),108 "r"(__c0),109 "r"(__c1),110 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(::cuda::device::barrier_native_handle(__bar))))111 : "memory");112}113 114// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor115inline _LIBCUDACXX_DEVICE116void cp_async_bulk_tensor_3d_global_to_shared(117 void *__dest, const CUtensorMap *__tensor_map, int __c0, int __c1, int __c2, ::cuda::barrier<::cuda::thread_scope_block> &__bar)118{119 asm volatile(120 "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes "121 "[%0], [%1, {%2, %3, %4}], [%5];\n"122 :123 : "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__dest))),124 "l"(__tensor_map),125 "r"(__c0),126 "r"(__c1),127 "r"(__c2),128 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(::cuda::device::barrier_native_handle(__bar))))129 : "memory");130}131 132// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor133inline _LIBCUDACXX_DEVICE134void cp_async_bulk_tensor_4d_global_to_shared(135 void *__dest, const CUtensorMap *__tensor_map , int __c0, int __c1, int __c2, int __c3, ::cuda::barrier<::cuda::thread_scope_block> &__bar)136{137 asm volatile(138 "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes "139 "[%0], [%1, {%2, %3, %4, %5}], [%6];\n"140 :141 : "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__dest))),142 "l"(__tensor_map),143 "r"(__c0),144 "r"(__c1),145 "r"(__c2),146 "r"(__c3),147 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(::cuda::device::barrier_native_handle(__bar))))148 : "memory");149}150 151// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor152inline _LIBCUDACXX_DEVICE153void cp_async_bulk_tensor_5d_global_to_shared(154 void *__dest, const CUtensorMap *__tensor_map , int __c0, int __c1, int __c2, int __c3, int __c4, ::cuda::barrier<::cuda::thread_scope_block> &__bar)155{156 asm volatile(157 "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes "158 "[%0], [%1, {%2, %3, %4, %5, %6}], [%7];\n"159 :160 : "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__dest))),161 "l"(__tensor_map),162 "r"(__c0),163 "r"(__c1),164 "r"(__c2),165 "r"(__c3),166 "r"(__c4),167 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(::cuda::device::barrier_native_handle(__bar))))168 : "memory");169}170 171// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor172inline _LIBCUDACXX_DEVICE173void cp_async_bulk_tensor_1d_shared_to_global(174 const CUtensorMap *__tensor_map, int __c0, const void *__src)175{176 asm volatile(177 "cp.async.bulk.tensor.1d.global.shared::cta.tile.bulk_group "178 "[%0, {%1}], [%2];\n"179 :180 : "l"(__tensor_map),181 "r"(__c0),182 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__src)))183 : "memory");184}185 186// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor187inline _LIBCUDACXX_DEVICE188void cp_async_bulk_tensor_2d_shared_to_global(189 const CUtensorMap *__tensor_map, int __c0, int __c1, const void *__src)190{191 asm volatile(192 "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group "193 "[%0, {%1, %2}], [%3];\n"194 :195 : "l"(__tensor_map),196 "r"(__c0),197 "r"(__c1),198 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__src)))199 : "memory");200}201 202// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor203inline _LIBCUDACXX_DEVICE204void cp_async_bulk_tensor_3d_shared_to_global(205 const CUtensorMap *__tensor_map, int __c0, int __c1, int __c2, const void *__src)206{207 asm volatile(208 "cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group "209 "[%0, {%1, %2, %3}], [%4];\n"210 :211 : "l"(__tensor_map),212 "r"(__c0),213 "r"(__c1),214 "r"(__c2),215 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__src)))216 : "memory");217}218 219// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor220inline _LIBCUDACXX_DEVICE221void cp_async_bulk_tensor_4d_shared_to_global(222 const CUtensorMap *__tensor_map, int __c0, int __c1, int __c2, int __c3, const void *__src)223{224 asm volatile(225 "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group "226 "[%0, {%1, %2, %3, %4}], [%5];\n"227 :228 : "l"(__tensor_map),229 "r"(__c0),230 "r"(__c1),231 "r"(__c2),232 "r"(__c3),233 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__src)))234 : "memory");235}236 237// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor238inline _LIBCUDACXX_DEVICE239void cp_async_bulk_tensor_5d_shared_to_global(240 const CUtensorMap *__tensor_map, int __c0, int __c1, int __c2, int __c3, int __c4, const void *__src)241{242 asm volatile(243 "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group "244 "[%0, {%1, %2, %3, %4, %5}], [%6];\n"245 :246 : "l"(__tensor_map),247 "r"(__c0),248 "r"(__c1),249 "r"(__c2),250 "r"(__c3),251 "r"(__c4),252 "r"(static_cast<_CUDA_VSTD::uint32_t>(__cvta_generic_to_shared(__src)))253 : "memory");254}255 256// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-membar257inline _LIBCUDACXX_DEVICE258void fence_proxy_async_shared_cta() {259 asm volatile("fence.proxy.async.shared::cta; \n":::"memory");260}261 262// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-commit-group263inline _LIBCUDACXX_DEVICE264void cp_async_bulk_commit_group()265{266 asm volatile("cp.async.bulk.commit_group;\n" ::: "memory");267}268 269// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-wait-group270template <int n_prior>271inline _LIBCUDACXX_DEVICE272void cp_async_bulk_wait_group_read()273{274 static_assert(n_prior <= 63, "cp_async_bulk_wait_group_read: waiting for more than 63 groups is not supported.");275 asm volatile("cp.async.bulk.wait_group.read %0; \n"276 :277 : "n"(n_prior)278 : "memory");279}280 281#endif // __cccl_lib_experimental_ctk12_cp_async_exposure282 283_LIBCUDACXX_END_NAMESPACE_CUDA_DEVICE_EXPERIMENTAL284 285#endif // _CUDA_BARRIER286