Neko 1.99.5
A portable framework for high-order spectral element flow simulations
Loading...
Searching...
No Matches
gs_nvshmem_kernels.h
Go to the documentation of this file.
1/*
2 Copyright (c) 2024-2025, The Neko Authors
3 All rights reserved.
4
5 Redistribution and use in source and binary forms, with or without
6 modification, are permitted provided that the following conditions
7 are met:
8
9 * Redistributions of source code must retain the above copyright
10 notice, this list of conditions and the following disclaimer.
11
12 * Redistributions in binary form must reproduce the above
13 copyright notice, this list of conditions and the following
14 disclaimer in the documentation and/or other materials provided
15 with the distribution.
16
17 * Neither the name of the authors nor the names of its
18 contributors may be used to endorse or promote products derived
19 from this software without specific prior written permission.
20
21 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
22 "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
23 LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
24 FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
25 COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
26 INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
27 BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
28 LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
29 CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
30 LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
31 ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
32 POSSIBILITY OF SUCH DAMAGE.
33*/
34
35#ifndef __GS_NVSHMEM_KERNELS__
36#define __GS_NVSHMEM_KERNELS__
37
38#include <nvshmemx.h>
39
40/*
41 * Fused pack-and-push kernels with rank-indexed signaling.
42 *
43 * Signaling uses two symmetric arrays of pe_size slots, indexed by the
44 * REMOTE PE's rank (so peer lists need not be uniform in length or aligned
45 * in order across ranks, and the symmetric allocations are collective-safe):
46 * - doneSig = &done_sig[my_rank] on the destination: set to iter by our
47 * put_signal when our slab has landed there.
48 * - readySlot = &ready_sig[destRank] locally: set to iter by the
49 * destination once it has consumed our round-iter slab.
50 * The round counter iter advances once per gs op (lockstep across ranks),
51 * and all waits use CMP_GE, so no cross-rank counter matching is needed.
52 *
53 * These are SINGLE-BLOCK kernels (launched with one block): the pack is a
54 * block-stride loop, so the block-local __syncthreads() suffices to order
55 * the pack before the put. A multi-block pack would race with the push --
56 * there is no grid-wide barrier between pack and push. Single-block
57 * transfers were also found to perform best.
58 *
59 * The ready wait comes BEFORE the pack: the destination posts ready(iter-1)
60 * only after consuming our round iter-1 slab, which implies our previous
61 * (non-blocking) put has fully drained src -- so the pack below can safely
62 * overwrite it. Packing before this wait would race with a still-in-flight
63 * nbi put on proxy-based transports.
64 */
65
66template< typename T >
67__global__ void pack_pushShmemKernel(const T * __restrict__ u,
68 T * dest,
69 T * __restrict__ src,
70 const int * __restrict__ dof,
71 const int destRank,
72 const int n,
73 uint64_t iter,
74 uint64_t * doneSig,
75 uint64_t * readySlot);
76
77template<>
78__global__ void pack_pushShmemKernel(const float * __restrict__ u,
79 float * dest,
80 float * __restrict__ src,
81 const int * __restrict__ dof,
82 const int destRank,
83 const int n,
84 uint64_t iter,
85 uint64_t * doneSig,
86 uint64_t * readySlot)
87{
88
89 /* Wait until destRank has consumed our previous round (see note above) */
90 if (threadIdx.x == 0) {
91 nvshmem_signal_wait_until(readySlot, NVSHMEM_CMP_GE, iter - 1);
92 }
93 __syncthreads();
94
95 /* Pack with a block-stride loop (single-block kernel) */
96 for (int j = threadIdx.x; j < n; j += blockDim.x) {
97 src[j] = u[dof[j]-1];
98 }
99 __syncthreads();
100
101 /* Push data and set done_sig[my_rank] = iter on the destination */
102 nvshmemx_float_put_signal_nbi_block(dest, src, n,
103 doneSig, iter,
104 NVSHMEM_SIGNAL_SET, destRank);
105}
106
107template<>
108__global__ void pack_pushShmemKernel(const double * __restrict__ u,
109 double * dest,
110 double * __restrict__ src,
111 const int * __restrict__ dof,
112 const int destRank,
113 const int n,
114 uint64_t iter,
115 uint64_t * doneSig,
116 uint64_t * readySlot)
117{
118
119 /* Wait until destRank has consumed our previous round (see note above) */
120 if (threadIdx.x == 0) {
121 nvshmem_signal_wait_until(readySlot, NVSHMEM_CMP_GE, iter - 1);
122 }
123 __syncthreads();
124
125 /* Pack with a block-stride loop (single-block kernel) */
126 for (int j = threadIdx.x; j < n; j += blockDim.x) {
127 src[j] = u[dof[j]-1];
128 }
129 __syncthreads();
130
131 /* Push data and set done_sig[my_rank] = iter on the destination */
132 nvshmemx_double_put_signal_nbi_block(dest, src, n,
133 doneSig, iter,
134 NVSHMEM_SIGNAL_SET, destRank);
135}
136
137/* Wait until the slab from a recv peer has landed (doneSlot is our local
138 done_sig[src] slot, set by the peer's put_signal) */
139__global__ void pushShmemKernelWait(uint64_t iter,
140 uint64_t *doneSlot)
141{
142 if (blockIdx.x==0 && threadIdx.x == 0) {
143 nvshmem_signal_wait_until(doneSlot, NVSHMEM_CMP_GE, iter);
144 }
145}
146
147/* Post our ready signal to the peer we receive from: sets
148 ready_sig[my_rank] = iter on srcRank, allowing it to overwrite its send
149 slab for the next round. Launched after the unpack on the same stream. */
150__global__ void postReadyShmemKernel(uint64_t *readySlot,
151 uint64_t iter,
152 const int srcRank)
153{
154 if (blockIdx.x==0 && threadIdx.x == 0) {
155 nvshmemx_signal_op(readySlot, iter, NVSHMEM_SIGNAL_SET, srcRank);
156 }
157}
158
159/*
160 * Fused nc-component pack-and-push (single-block, rank-indexed signaling,
161 * see the scalar kernels above). u is the compact shared buffer,
162 * component-outer with per-component stride ns; src/dest are interleaved
163 * (nc per position). The pushed payload is nc*n elements.
164 */
165template< typename T >
166__global__ void pack_pushShmemKernel_vec(const T * __restrict__ u,
167 T * dest,
168 T * __restrict__ src,
169 const int * __restrict__ dof,
170 const int destRank,
171 const int n,
172 const int nc,
173 const int ns,
174 uint64_t iter,
175 uint64_t * doneSig,
176 uint64_t * readySlot);
177
178template<>
179__global__ void pack_pushShmemKernel_vec(const float * __restrict__ u,
180 float * dest,
181 float * __restrict__ src,
182 const int * __restrict__ dof,
183 const int destRank,
184 const int n,
185 const int nc,
186 const int ns,
187 uint64_t iter,
188 uint64_t * doneSig,
189 uint64_t * readySlot)
190{
191
192 /* Wait until destRank has consumed our previous round */
193 if (threadIdx.x == 0) {
194 nvshmem_signal_wait_until(readySlot, NVSHMEM_CMP_GE, iter - 1);
195 }
196 __syncthreads();
197
198 /* Pack with a block-stride loop (single-block kernel) */
199 for (int j = threadIdx.x; j < n; j += blockDim.x) {
200 const int idx = dof[j] - 1;
201 for (int c = 0; c < nc; c++)
202 src[nc*j + c] = u[c*ns + idx];
203 }
204 __syncthreads();
205
206 /* Push data and set done_sig[my_rank] = iter on the destination */
207 nvshmemx_float_put_signal_nbi_block(dest, src, (size_t) nc * n,
208 doneSig, iter,
209 NVSHMEM_SIGNAL_SET, destRank);
210}
211
212template<>
213__global__ void pack_pushShmemKernel_vec(const double * __restrict__ u,
214 double * dest,
215 double * __restrict__ src,
216 const int * __restrict__ dof,
217 const int destRank,
218 const int n,
219 const int nc,
220 const int ns,
221 uint64_t iter,
222 uint64_t * doneSig,
223 uint64_t * readySlot)
224{
225
226 /* Wait until destRank has consumed our previous round */
227 if (threadIdx.x == 0) {
228 nvshmem_signal_wait_until(readySlot, NVSHMEM_CMP_GE, iter - 1);
229 }
230 __syncthreads();
231
232 /* Pack with a block-stride loop (single-block kernel) */
233 for (int j = threadIdx.x; j < n; j += blockDim.x) {
234 const int idx = dof[j] - 1;
235 for (int c = 0; c < nc; c++)
236 src[nc*j + c] = u[c*ns + idx];
237 }
238 __syncthreads();
239
240 /* Push data and set done_sig[my_rank] = iter on the destination */
241 nvshmemx_double_put_signal_nbi_block(dest, src, (size_t) nc * n,
242 doneSig, iter,
243 NVSHMEM_SIGNAL_SET, destRank);
244}
245
246#endif
ale_add_kinematics_kernel
__global__ void ale_add_kinematics_kernel(const int n, T *__restrict__ wx, T *__restrict__ wy, T *__restrict__ wz, const T *__restrict__ x_ref, const T *__restrict__ y_ref, const T *__restrict__ z_ref, const T *__restrict__ phi, const T *__restrict__ x, const T *__restrict__ y, const T *__restrict__ z, const kinematics_params_t kin_params)
Definition ale_kinematics_kernel.h:52
u
__global__ void T *__restrict__ T *__restrict__ const T *__restrict__ u
Definition ax_helm_full_kernel.h:46
j
const int j
Definition ax_helm_full_kernel.h:94
__syncthreads
__syncthreads()
pack_pushShmemKernel_vec
__global__ void pack_pushShmemKernel_vec(const T *__restrict__ u, T *dest, T *__restrict__ src, const int *__restrict__ dof, const int destRank, const int n, const int nc, const int ns, uint64_t iter, uint64_t *doneSig, uint64_t *readySlot)
pushShmemKernelWait
__global__ void pushShmemKernelWait(uint64_t iter, uint64_t *doneSlot)
Definition gs_nvshmem_kernels.h:139
pack_pushShmemKernel
__global__ void pack_pushShmemKernel(const T *__restrict__ u, T *dest, T *__restrict__ src, const int *__restrict__ dof, const int destRank, const int n, uint64_t iter, uint64_t *doneSig, uint64_t *readySlot)
postReadyShmemKernel
__global__ void postReadyShmemKernel(uint64_t *readySlot, uint64_t iter, const int srcRank)
Definition gs_nvshmem_kernels.h:150