1//===-- Shared memory RPC client / server interface -------------*- 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// This file implements a remote procedure call mechanism to communicate between
10// heterogeneous devices that can share an address space atomically. We provide
11// a client and a server to facilitate the remote call. The client makes
12// requests to the server using a shared communication channel. We use separate
13// atomic signals to indicate which side, the client or the server is in
14// ownership of the buffer.
15//
16//===----------------------------------------------------------------------===//
17
18#ifndef LLVM_LIBC_SHARED_RPC_H
19#define LLVM_LIBC_SHARED_RPC_H
20
21#include "rpc_util.h"
22
23/// Use scoped atomic variants if they are available for the target.
24#if !__has_builtin(__scoped_atomic_load_n)
25#ifdef _MSC_VER // MSVC atomic support.
26#include <intrin.h>
27#define __scoped_atomic_load_n(src, ord, scp) \
28 __iso_volatile_load32((const volatile int32_t *)(src))
29#define __scoped_atomic_store_n(dst, src, ord, scp) \
30 (sizeof(*(dst)) == 4 \
31 ? __iso_volatile_store32((volatile int32_t *)(dst), (__int32)(src)) \
32 : __iso_volatile_store64((volatile int64_t *)(dst), (__int64)(src)))
33#define __scoped_atomic_fetch_or(src, val, ord, scp) \
34 _InterlockedOr((volatile long *)(src), (long)(val))
35#define __scoped_atomic_fetch_and(src, val, ord, scp) \
36 _InterlockedAnd((volatile long *)(src), (long)(val))
37#else // GNU atomic support.
38#define __scoped_atomic_load_n(src, ord, scp) __atomic_load_n(src, ord)
39#define __scoped_atomic_store_n(dst, src, ord, scp) \
40 __atomic_store_n(dst, src, ord)
41#define __scoped_atomic_fetch_or(src, val, ord, scp) \
42 __atomic_fetch_or(src, val, ord)
43#define __scoped_atomic_fetch_and(src, val, ord, scp) \
44 __atomic_fetch_and(src, val, ord)
45#define __scoped_atomic_fetch_add(src, val, ord, scp) \
46 __atomic_fetch_add(src, val, ord)
47#define __scoped_atomic_fetch_sub(src, val, ord, scp) \
48 __atomic_fetch_sub(src, val, ord)
49#endif
50#endif
51#if !__has_builtin(__scoped_atomic_thread_fence)
52#ifdef _MSC_VER
53#define __scoped_atomic_thread_fence(ord, scp) _ReadWriteBarrier()
54#else
55#define __scoped_atomic_thread_fence(ord, scp) __atomic_thread_fence(ord)
56#endif
57#endif
58
59namespace rpc {
60
61/// Generic codes that can be used when implementing the server.
62enum RPCStatus {
63 RPC_SUCCESS = 0x0,
64 RPC_ERROR = 0x1000,
65 RPC_UNHANDLED_OPCODE = 0x1001,
66};
67
68/// A fixed size channel used to communicate between the RPC client and server.
69struct Buffer {
70 uint64_t data[8];
71};
72static_assert(sizeof(Buffer) == 64, "Buffer size mismatch");
73
74/// A target specific struct containing a doorbell to wake the server-side
75/// thread.
76struct alignas(64) Doorbell {
77 RPC_GLOBAL uint64_t *value;
78 RPC_GLOBAL uint64_t *mailbox;
79 uint32_t event_id;
80};
81
82/// The information associated with a packet. This indicates which operations to
83/// perform and which threads are active in the slots.
84struct Header {
85 uint64_t mask;
86 uint32_t opcode;
87};
88
89/// The maximum number of parallel ports that the RPC interface can support.
90/// This should be greater than the expected hardware's occupancy to ensure the
91/// interface is non-blocking.
92constexpr static uint64_t MAX_PORT_COUNT = 16384;
93
94/// A common process used to synchronize communication between a client and a
95/// server. The process contains a read-only inbox and a write-only outbox used
96/// for signaling ownership of the shared buffer between both sides. We assign
97/// ownership of the buffer to the client if the inbox and outbox bits match,
98/// otherwise it is owned by the server.
99///
100/// This process is designed to allow the client and the server to exchange data
101/// using a fixed size packet in a mostly arbitrary order using the 'send' and
102/// 'recv' operations. The following restrictions to this scheme apply:
103/// - The client will always start with a 'send' operation.
104/// - The server will always start with a 'recv' operation.
105/// - Every 'send' or 'recv' call is mirrored by the other process.
106template <bool Invert> struct Process {
107 RPC_ATTRS Process() = default;
108 RPC_ATTRS Process(const Process &) = delete;
109 RPC_ATTRS Process &operator=(const Process &) = delete;
110 RPC_ATTRS Process(Process &&) = default;
111 RPC_ATTRS Process &operator=(Process &&) = default;
112 RPC_ATTRS ~Process() = default;
113
114 const uint32_t port_count = 0;
115 RPC_GLOBAL Doorbell *const doorbell = nullptr;
116 RPC_GLOBAL const uint32_t *const inbox = nullptr;
117 RPC_GLOBAL uint32_t *const outbox = nullptr;
118 RPC_GLOBAL Header *const header = nullptr;
119 RPC_GLOBAL Buffer *const packet = nullptr;
120
121 static constexpr uint64_t NUM_BITS_IN_WORD = sizeof(uint32_t) * 8;
122 uint32_t lock[MAX_PORT_COUNT / NUM_BITS_IN_WORD] = {0};
123
124 // The buffer is supplied by the host unqualified, so the casts below are
125 // C-style as they need to change the address space.
126 RPC_ATTRS Process(uint32_t port_count, void *buffer)
127 : port_count(port_count),
128 doorbell((RPC_GLOBAL Doorbell *)advance(ptr: buffer, bytes: doorbell_offset())),
129 inbox((RPC_GLOBAL uint32_t *)advance(ptr: buffer, bytes: inbox_offset(port_count))),
130 outbox(
131 (RPC_GLOBAL uint32_t *)advance(ptr: buffer, bytes: outbox_offset(port_count))),
132 header((RPC_GLOBAL Header *)advance(ptr: buffer, bytes: header_offset(port_count))),
133 packet(
134 (RPC_GLOBAL Buffer *)advance(ptr: buffer, bytes: buffer_offset(port_count))) {}
135
136 /// Allocate a memory buffer sufficient to store the following equivalent
137 /// representation in memory.
138 ///
139 /// struct Equivalent {
140 /// Doorbell doorbell;
141 /// Atomic<uint32_t> primary[port_count];
142 /// Atomic<uint32_t> secondary[port_count];
143 /// Header header[port_count];
144 /// Buffer packet[port_count][lane_size];
145 /// };
146 RPC_ATTRS static constexpr uint64_t allocation_size(uint32_t port_count,
147 uint32_t lane_size) {
148 return buffer_offset(port_count) + buffer_bytes(port_count, lane_size);
149 }
150
151 /// Ring the doorbell if the protocol was configured with one.
152 RPC_ATTRS void notify(uint64_t lane_mask) const {
153#ifndef _MSC_VER
154 if (!doorbell->value)
155 return;
156
157 uint32_t event_id = rpc::broadcast_value(lane_mask, x: doorbell->event_id);
158 if (rpc::is_first_lane(lane_mask)) {
159 // The interrupt is optional and is skipped if there is no mailbox.
160 if (!__scoped_atomic_fetch_add(doorbell->value, 1UL, __ATOMIC_RELAXED,
161 __MEMORY_SCOPE_SYSTEM) &&
162 doorbell->mailbox) {
163 __scoped_atomic_store_n(doorbell->mailbox,
164 static_cast<uint64_t>(doorbell->event_id),
165 __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM);
166 __scoped_atomic_thread_fence(__ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM);
167 signal_interrupt(event_id);
168 }
169 }
170#endif
171 }
172
173 /// Decrement the doorbell signal if the protocol is using one.
174 RPC_ATTRS void finish(uint64_t lane_mask) const {
175#ifndef _MSC_VER
176 if (!doorbell->value)
177 return;
178
179 if (rpc::is_first_lane(lane_mask))
180 __scoped_atomic_fetch_sub(doorbell->value, 1UL, __ATOMIC_RELAXED,
181 __MEMORY_SCOPE_SYSTEM);
182#endif
183 }
184
185 /// Retrieve the inbox state from memory shared between processes.
186 RPC_ATTRS uint32_t load_inbox(uint64_t lane_mask, uint32_t index) const {
187 return rpc::broadcast_value(
188 lane_mask, x: __scoped_atomic_load_n(&inbox[index], __ATOMIC_RELAXED,
189 __MEMORY_SCOPE_SYSTEM));
190 }
191
192 /// Retrieve the outbox state from memory shared between processes.
193 RPC_ATTRS uint32_t load_outbox(uint64_t lane_mask, uint32_t index) const {
194 return rpc::broadcast_value(
195 lane_mask, x: __scoped_atomic_load_n(&outbox[index], __ATOMIC_RELAXED,
196 __MEMORY_SCOPE_SYSTEM));
197 }
198
199 /// Signal to the other process that this one is finished with the buffer.
200 /// Equivalent to loading outbox followed by store of the inverted value
201 /// The outbox is write only by this warp and tracking the value locally is
202 /// cheaper than calling load_outbox to get the value to store.
203 RPC_ATTRS uint32_t invert_outbox(uint64_t lane_mask, uint32_t index,
204 uint32_t current_outbox) {
205 uint32_t inverted_outbox = !current_outbox;
206 rpc::sync_lane(lane_mask);
207 __scoped_atomic_thread_fence(__ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM);
208 if (rpc::is_first_lane(lane_mask))
209 __scoped_atomic_store_n(&outbox[index], inverted_outbox, __ATOMIC_RELAXED,
210 __MEMORY_SCOPE_SYSTEM);
211 return inverted_outbox;
212 }
213
214 /// Given the current outbox and inbox values, wait until the inbox changes
215 /// to indicate that this thread owns the buffer element.
216 RPC_ATTRS void wait_for_ownership(uint64_t lane_mask, uint32_t index,
217 uint32_t out, uint32_t in) {
218 while (buffer_unavailable(in, out)) {
219 sleep_briefly();
220 in = load_inbox(lane_mask, index);
221 }
222 __scoped_atomic_thread_fence(__ATOMIC_ACQUIRE, __MEMORY_SCOPE_SYSTEM);
223 }
224
225 /// The packet is a linearly allocated array of buffers used to communicate
226 /// with the other process. This function returns the appropriate slot in this
227 /// array such that the process can operate on an entire warp or wavefront.
228 RPC_ATTRS RPC_GLOBAL Buffer *get_packet(uint32_t index, uint32_t lane_size) {
229 return &packet[index * lane_size];
230 }
231
232 /// Determines if this process needs to wait for ownership of the buffer. We
233 /// invert the condition on one of the processes to indicate that if one
234 /// process owns the buffer then the other does not.
235 RPC_ATTRS static bool buffer_unavailable(uint32_t in, uint32_t out) {
236 bool cond = in != out;
237 return Invert ? !cond : cond;
238 }
239
240 /// Attempts to claim the lock at this index under the given execution mask.
241 /// Returns true on success.
242 RPC_ATTRS bool try_lock(uint64_t lane_mask, uint32_t index) {
243 bool claimed = false;
244 if (rpc::is_first_lane(lane_mask))
245 claimed = !set_nth(bits: lock, index);
246 claimed = rpc::broadcast_value(lane_mask, x: claimed);
247
248 // Do not move any reads past the point we obtain the lock.
249 if (claimed)
250 __scoped_atomic_thread_fence(__ATOMIC_ACQUIRE, __MEMORY_SCOPE_DEVICE);
251 return claimed;
252 }
253
254 /// Unlock the lock at index. We need a lane sync to keep this function
255 /// convergent, otherwise the compiler will sink the store and deadlock.
256 RPC_ATTRS void unlock(uint64_t lane_mask, uint32_t index) {
257 // Do not move any writes past the unlock.
258 __scoped_atomic_thread_fence(__ATOMIC_RELEASE, __MEMORY_SCOPE_DEVICE);
259
260 if (rpc::is_first_lane(lane_mask))
261 clear_nth(bits: lock, index);
262 rpc::sync_lane(lane_mask);
263 }
264
265 /// Number of bytes to allocate for an inbox or outbox.
266 RPC_ATTRS static constexpr uint64_t mailbox_bytes(uint32_t port_count) {
267 return port_count * sizeof(uint32_t);
268 }
269
270 /// Number of bytes to allocate for the buffer containing the packets.
271 RPC_ATTRS static constexpr uint64_t buffer_bytes(uint32_t port_count,
272 uint32_t lane_size) {
273 return port_count * lane_size * sizeof(Buffer);
274 }
275
276 /// The offset to the doorbell interface.
277 RPC_ATTRS static constexpr uint64_t doorbell_offset() { return 0; }
278
279 /// Offset of the inbox in memory. This is the same as the outbox if inverted.
280 RPC_ATTRS static constexpr uint64_t inbox_offset(uint32_t port_count) {
281 return sizeof(Doorbell) + (Invert ? mailbox_bytes(port_count) : 0);
282 }
283
284 /// Offset of the outbox in memory. This is the same as the inbox if inverted.
285 RPC_ATTRS static constexpr uint64_t outbox_offset(uint32_t port_count) {
286 return sizeof(Doorbell) + (Invert ? 0 : mailbox_bytes(port_count));
287 }
288
289 /// Offset of the header containing the opcode and mask after the mailboxes.
290 RPC_ATTRS static constexpr uint64_t header_offset(uint32_t port_count) {
291 return align_up(val: sizeof(Doorbell) + 2 * mailbox_bytes(port_count),
292 align: alignof(Header));
293 }
294
295 /// Offset of the buffer containing the packets after the inbox and outbox.
296 RPC_ATTRS static constexpr uint64_t buffer_offset(uint32_t port_count) {
297 return align_up(val: header_offset(port_count) + port_count * sizeof(Header),
298 align: alignof(Buffer));
299 }
300
301 /// Conditionally set the n-th bit in the atomic bitfield.
302 RPC_ATTRS static constexpr uint32_t set_nth(uint32_t *bits, uint32_t index) {
303 uint32_t slot = index / NUM_BITS_IN_WORD;
304 uint32_t bit = index % NUM_BITS_IN_WORD;
305 return __scoped_atomic_fetch_or(&bits[slot], 1u << bit, __ATOMIC_RELAXED,
306 __MEMORY_SCOPE_DEVICE) &
307 (1u << bit);
308 }
309
310 /// Conditionally clear the n-th bit in the atomic bitfield.
311 RPC_ATTRS static constexpr uint32_t clear_nth(uint32_t *bits,
312 uint32_t index) {
313 uint32_t slot = index / NUM_BITS_IN_WORD;
314 uint32_t bit = index % NUM_BITS_IN_WORD;
315 return __scoped_atomic_fetch_and(&bits[slot], ~0u ^ (1u << bit),
316 __ATOMIC_RELAXED, __MEMORY_SCOPE_DEVICE) &
317 (1u << bit);
318 }
319};
320
321/// Invokes a function across every active buffer across the total lane size.
322template <typename F>
323RPC_ATTRS static void invoke_rpc(F &&fn, uint32_t lane_size, uint64_t lane_mask,
324 RPC_GLOBAL Buffer *slot) {
325 if constexpr (is_process_gpu()) {
326 fn(&slot[rpc::get_lane_id()], rpc::get_lane_id());
327 } else {
328 for (uint32_t i = 0; i < lane_size; i += rpc::get_num_lanes())
329 if (lane_mask & (1ul << i))
330 fn(&slot[i], i);
331 }
332}
333
334/// The port provides the interface to communicate between the multiple
335/// processes. A port is conceptually an index into the memory provided by the
336/// underlying process that is guarded by a lock bit.
337template <bool T> struct Port {
338 RPC_ATTRS Port(Process<T> &process, uint64_t lane_mask, uint32_t lane_size,
339 uint32_t index, uint32_t out)
340 : process(process), lane_mask(lane_mask), lane_size(lane_size),
341 index(index), out(out), receive(false), owns_buffer(true) {}
342 RPC_ATTRS ~Port() { close(); }
343
344private:
345 RPC_ATTRS Port(const Port &) = delete;
346 RPC_ATTRS Port &operator=(const Port &) = delete;
347 RPC_ATTRS Port(Port &&) = delete;
348 RPC_ATTRS Port &operator=(Port &&) = delete;
349
350 friend struct Client;
351 friend struct Server;
352 friend struct rpc::optional<Port<T>>;
353
354public:
355 template <typename U> RPC_ATTRS void recv(U use);
356 template <typename F> RPC_ATTRS void send(F fill);
357 template <typename F, typename U> RPC_ATTRS void send_and_recv(F fill, U use);
358 template <typename W> RPC_ATTRS void recv_and_send(W work);
359 RPC_ATTRS void send_n(const void *const *src, uint64_t *size);
360 RPC_ATTRS void send_n(const void *src, uint64_t size);
361 template <typename A>
362 RPC_ATTRS void recv_n(void **dst, uint64_t *size, A &&alloc);
363
364 template <typename Ty> RPC_ATTRS void send_n(const Ty *src);
365 template <typename Ty> RPC_ATTRS void recv_n(Ty *dst);
366
367 RPC_ATTRS uint32_t get_opcode() const { return process.header[index].opcode; }
368
369 RPC_ATTRS uint32_t get_index() const { return index; }
370
371 RPC_ATTRS uint64_t get_lane_mask() const {
372 if constexpr (T)
373 return process.header[index].mask;
374 return lane_mask;
375 }
376
377private:
378 RPC_ATTRS void close() {
379 // Wait for all lanes to finish using the port.
380 rpc::sync_lane(lane_mask);
381
382 // The server is passive, if it owns the buffer when it closes we need to
383 // give ownership back to the client.
384 if (owns_buffer && T)
385 out = process.invert_outbox(lane_mask, index, out);
386 process.unlock(lane_mask, index);
387 if constexpr (T)
388 process.finish(lane_mask);
389 }
390
391 Process<T> &process;
392 uint64_t lane_mask;
393 uint32_t lane_size;
394 uint32_t index;
395 uint32_t out;
396 bool receive;
397 bool owns_buffer;
398};
399
400/// The RPC client used to make requests to the server.
401struct Client {
402 RPC_ATTRS Client() = default;
403 RPC_ATTRS Client(const Client &) = delete;
404 RPC_ATTRS Client &operator=(const Client &) = delete;
405 RPC_ATTRS ~Client() = default;
406
407 RPC_ATTRS Client(uint32_t port_count, void *buffer)
408 : process(port_count, buffer) {}
409
410 using Port = rpc::Port<false>;
411 template <uint32_t opcode> RPC_ATTRS Port open();
412
413private:
414 Process<false> process;
415};
416
417/// The RPC server used to respond to the client.
418struct Server {
419 RPC_ATTRS Server() = default;
420 RPC_ATTRS Server(const Server &) = delete;
421 RPC_ATTRS Server &operator=(const Server &) = delete;
422 RPC_ATTRS ~Server() = default;
423
424 RPC_ATTRS Server(uint32_t port_count, void *buffer)
425 : process(port_count, buffer) {}
426
427 using Port = rpc::Port<true>;
428 RPC_ATTRS rpc::optional<Port> try_open(uint32_t lane_size,
429 uint32_t start = 0);
430
431 RPC_ATTRS static constexpr uint64_t allocation_size(uint32_t lane_size,
432 uint32_t port_count) {
433 return Process<true>::allocation_size(port_count, lane_size);
434 }
435
436 RPC_ATTRS static constexpr uint64_t doorbell_offset() {
437 return Process<true>::doorbell_offset();
438 }
439
440private:
441 Process<true> process;
442};
443
444/// Applies \p fill to the shared buffer and initiates a send operation.
445template <bool T> template <typename F> RPC_ATTRS void Port<T>::send(F fill) {
446 uint32_t in = owns_buffer ? out ^ T : process.load_inbox(lane_mask, index);
447
448 // We need to wait until we own the buffer before sending.
449 process.wait_for_ownership(lane_mask, index, out, in);
450
451 // Apply the \p fill function to initialize the buffer and release the memory.
452 invoke_rpc(fill, lane_size, get_lane_mask(),
453 process.get_packet(index, lane_size));
454 out = process.invert_outbox(lane_mask, index, out);
455 owns_buffer = false;
456 receive = false;
457}
458
459/// Applies \p use to the shared buffer and acknowledges the send.
460template <bool T> template <typename U> RPC_ATTRS void Port<T>::recv(U use) {
461 // We only exchange ownership of the buffer during a receive if we are waiting
462 // for a previous receive to finish.
463 if (receive) {
464 out = process.invert_outbox(lane_mask, index, out);
465 owns_buffer = false;
466 }
467
468 uint32_t in = owns_buffer ? out ^ T : process.load_inbox(lane_mask, index);
469
470 // We need to wait until we own the buffer before receiving.
471 process.wait_for_ownership(lane_mask, index, out, in);
472
473 // Apply the \p use function to read the memory out of the buffer.
474 invoke_rpc(use, lane_size, get_lane_mask(),
475 process.get_packet(index, lane_size));
476 receive = true;
477 owns_buffer = true;
478}
479
480/// Combines a send and receive into a single function.
481template <bool T>
482template <typename F, typename U>
483RPC_ATTRS void Port<T>::send_and_recv(F fill, U use) {
484 send(fill);
485 recv(use);
486}
487
488/// Combines a receive and send operation into a single function. The \p work
489/// function modifies the buffer in-place and the send is only used to initiate
490/// the copy back.
491template <bool T>
492template <typename W>
493RPC_ATTRS void Port<T>::recv_and_send(W work) {
494 recv(work);
495 send([](RPC_GLOBAL Buffer *, uint32_t) { /* no-op */ });
496}
497
498/// Helper routine to simplify the interface when sending from the GPU using
499/// thread private pointers to the underlying value.
500template <bool T>
501RPC_ATTRS void Port<T>::send_n(const void *src, uint64_t size) {
502 const void **src_ptr = &src;
503 uint64_t *size_ptr = &size;
504 send_n(src_ptr, size_ptr);
505}
506
507/// Sends an arbitrarily sized data buffer \p src across the shared channel in
508/// multiples of the packet length.
509template <bool T>
510RPC_ATTRS void Port<T>::send_n(const void *const *src, uint64_t *size) {
511 constexpr uint64_t BUFFER_SIZE = sizeof(Buffer::data);
512 constexpr uint64_t FIRST_CHUNK = BUFFER_SIZE - sizeof(uint64_t);
513 uint64_t num_sends = 0;
514 send([&](RPC_GLOBAL Buffer *buffer, uint32_t id) {
515 reinterpret_cast<uint64_t *>(buffer->data)[0] = lane_value(val: size, id);
516 num_sends = is_process_gpu() ? lane_value(val: size, id)
517 : rpc::max(a: lane_value(val: size, id), b: num_sends);
518 uint64_t len =
519 lane_value(val: size, id) > FIRST_CHUNK ? FIRST_CHUNK : lane_value(val: size, id);
520 rpc_memcpy(dst: &buffer->data[1], src: lane_value(val: src, id), count: len);
521 });
522 uint64_t idx = FIRST_CHUNK;
523 uint64_t mask = process.header[index].mask;
524 while (rpc::ballot(lane_mask: mask, x: idx < num_sends && num_sends > FIRST_CHUNK)) {
525 send([=](RPC_GLOBAL Buffer *buffer, uint32_t id) {
526 uint64_t len = lane_value(val: size, id) - idx > BUFFER_SIZE
527 ? BUFFER_SIZE
528 : lane_value(val: size, id) - idx;
529 if (idx < lane_value(val: size, id))
530 rpc_memcpy(dst: buffer->data, src: advance(ptr: lane_value(val: src, id), bytes: idx), count: len);
531 });
532 idx += BUFFER_SIZE;
533 }
534}
535
536/// Receives an arbitrarily sized data buffer across the shared channel in
537/// multiples of the packet length. The \p alloc function is called with the
538/// size of the data so that we can initialize the size of the \p dst buffer.
539template <bool T>
540template <typename A>
541RPC_ATTRS void Port<T>::recv_n(void **dst, uint64_t *size, A &&alloc) {
542 constexpr uint64_t BUFFER_SIZE = sizeof(Buffer::data);
543 constexpr uint64_t FIRST_CHUNK = BUFFER_SIZE - sizeof(uint64_t);
544 uint64_t num_recvs = 0;
545 recv([&](RPC_GLOBAL Buffer *buffer, uint32_t id) {
546 lane_value(val: size, id) = reinterpret_cast<uint64_t *>(buffer->data)[0];
547 lane_value(val: dst, id) =
548 reinterpret_cast<uint8_t *>(alloc(lane_value(val: size, id)));
549 num_recvs = is_process_gpu() ? lane_value(val: size, id)
550 : rpc::max(a: lane_value(val: size, id), b: num_recvs);
551 uint64_t len =
552 lane_value(val: size, id) > FIRST_CHUNK ? FIRST_CHUNK : lane_value(val: size, id);
553 rpc_memcpy(dst: lane_value(val: dst, id), src: &buffer->data[1], count: len);
554 });
555 uint64_t idx = FIRST_CHUNK;
556 uint64_t mask = process.header[index].mask;
557 while (rpc::ballot(lane_mask: mask, x: idx < num_recvs && num_recvs > FIRST_CHUNK)) {
558 recv([=](RPC_GLOBAL Buffer *buffer, uint32_t id) {
559 uint64_t len = lane_value(val: size, id) - idx > BUFFER_SIZE
560 ? BUFFER_SIZE
561 : lane_value(val: size, id) - idx;
562 if (idx < lane_value(val: size, id))
563 rpc_memcpy(dst: advance(ptr: lane_value(val: dst, id), bytes: idx), src: buffer->data, count: len);
564 });
565 idx += BUFFER_SIZE;
566 }
567}
568
569/// Simplified version of `send_n` where the size is a known constant.
570template <bool T>
571template <typename Ty>
572RPC_ATTRS void Port<T>::send_n(const Ty *src) {
573 for (uint64_t idx = 0; idx < sizeof(Ty); idx += sizeof(Buffer::data)) {
574 const uint64_t bytes = rpc::min(a: sizeof(Ty) - idx, b: sizeof(Buffer::data));
575 send([&](RPC_GLOBAL Buffer *buffer, uint32_t id) {
576 rpc_memcpy(buffer->data, advance(&lane_value(src, id), idx), bytes);
577 });
578 }
579}
580
581/// Simplified version of `recv_n` where the size is a known constant.
582template <bool T>
583template <typename Ty>
584RPC_ATTRS void Port<T>::recv_n(Ty *dst) {
585 for (uint64_t idx = 0; idx < sizeof(Ty); idx += sizeof(Buffer::data)) {
586 const uint64_t bytes = rpc::min(a: sizeof(Ty) - idx, b: sizeof(Buffer::data));
587 recv([&](RPC_GLOBAL Buffer *buffer, uint32_t id) {
588 rpc_memcpy(advance(&lane_value(dst, id), idx), buffer->data, bytes);
589 });
590 }
591}
592
593/// Continually attempts to open a port to use as the client. The client can
594/// only open a port if we find an index that is in a valid sending state. That
595/// is, there are send operations pending that haven't been serviced on this
596/// port. Each port instance uses an associated \p opcode to tell the server
597/// what to do. The Client interface provides the appropriate lane size to the
598/// port using the platform's returned value.
599template <uint32_t opcode> RPC_ATTRS Client::Port Client::open() {
600 // Repeatedly perform a naive linear scan for a port that can be opened to
601 // send data.
602 for (uint32_t index = 0;; ++index) {
603 // Start from the beginning if we run out of ports to check.
604 if (index >= process.port_count)
605 index = 0;
606
607 // Attempt to acquire the lock on this index. Under NVIDIA's ITS the lanes
608 // may reconverge with differing index values, ensure they are convergent.
609 uint64_t lane_mask = rpc::get_lane_mask();
610 index = rpc::broadcast_value(lane_mask, x: index);
611 if (!process.try_lock(lane_mask, index))
612 continue;
613
614 uint32_t in = process.load_inbox(lane_mask, index);
615 uint32_t out = process.load_outbox(lane_mask, index);
616
617 // Once we acquire the index we need to check if we are in a valid sending
618 // state.
619 if (process.buffer_unavailable(in, out)) {
620 process.unlock(lane_mask, index);
621 continue;
622 }
623
624 if (rpc::is_first_lane(lane_mask)) {
625 process.header[index].opcode = opcode;
626 process.header[index].mask = lane_mask;
627 }
628 rpc::sync_lane(lane_mask);
629
630 process.notify(lane_mask);
631 return Port(process, lane_mask, rpc::get_num_lanes(), index, out);
632 }
633}
634
635/// Attempts to open a port to use as the server. The server can only open a
636/// port if it has a pending receive operation
637RPC_ATTRS rpc::optional<typename Server::Port>
638Server::try_open(uint32_t lane_size, uint32_t start) {
639 if (rpc::get_lane_id() >= lane_size)
640 return rpc::nullopt;
641
642 // Perform a naive linear scan for a port that has a pending request.
643 for (uint32_t index = start; index < process.port_count; ++index) {
644 uint64_t lane_mask = rpc::get_lane_mask();
645 uint32_t in = process.load_inbox(lane_mask, index);
646 uint32_t out = process.load_outbox(lane_mask, index);
647
648 // The server is passive, if there is no work pending don't bother
649 // opening a port.
650 if (process.buffer_unavailable(in, out))
651 continue;
652
653 // Attempt to acquire the lock on this index.
654 if (!process.try_lock(lane_mask, index))
655 continue;
656
657 in = process.load_inbox(lane_mask, index);
658 out = process.load_outbox(lane_mask, index);
659
660 if (process.buffer_unavailable(in, out)) {
661 process.unlock(lane_mask, index);
662 continue;
663 }
664
665 return rpc::optional<Port>(rpc::in_place, process, lane_mask, lane_size,
666 index, out);
667 }
668 return rpc::nullopt;
669}
670
671#if !__has_builtin(__scoped_atomic_load_n)
672#undef __scoped_atomic_load_n
673#undef __scoped_atomic_store_n
674#undef __scoped_atomic_fetch_or
675#undef __scoped_atomic_fetch_and
676#undef __scoped_atomic_fetch_add
677#undef __scoped_atomic_fetch_sub
678#endif
679#if !__has_builtin(__scoped_atomic_thread_fence)
680#undef __scoped_atomic_thread_fence
681#endif
682
683} // namespace rpc
684
685#endif // LLVM_LIBC_SHARED_RPC_H
686