char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Ryan Crpc: support apple RDMA as an RPC transport (#26421)b114b4739

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
22.1 KiB739 linesraw
1#include "transport.h"
2#include "ggml-impl.h"
3
4#ifdef _WIN32
5#  define WIN32_LEAN_AND_MEAN
6#  ifndef NOMINMAX
7#     define NOMINMAX
8#  endif
9#  include <windows.h>
10#  include <winsock2.h>
11#else
12#  include <arpa/inet.h>
13#  include <sys/socket.h>
14#  include <sys/types.h>
15#  include <netinet/in.h>
16#  include <netinet/tcp.h>
17#  include <netdb.h>
18#  include <unistd.h>
19#endif
20#include <cstdlib>
21#include <cstring>
22#include <mutex>
23#include <optional>
24
25#ifdef GGML_RPC_RDMA
26#  include <infiniband/verbs.h>
27#  include <array>
28#  include <time.h>
29#  ifndef _WIN32
30#    include <poll.h>
31#  endif
32#  ifdef GGML_RPC_RDMA_APPLE
33#    include "transport-apple.h"
34#  endif
35#endif // GGML_RPC_RDMA
36
37#ifdef _WIN32
38typedef SOCKET sockfd_t;
39using ssize_t = __int64;
40#else
41typedef int sockfd_t;
42#endif
43
44static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG");
45
46#define LOG_DBG(...) \
47    do { if (RPC_DEBUG) GGML_LOG_DEBUG(__VA_ARGS__); } while (0)
48
49#ifdef GGML_RPC_RDMA
50static constexpr size_t RDMA_GID_SIZE = 16;            // RoCE GID / IB GID is always 16 bytes
51using rdma_gid_t = std::array<uint8_t, RDMA_GID_SIZE>;
52#endif // GGML_RPC_RDMA
53
54#if defined(GGML_RPC_RDMA) && !defined(GGML_RPC_RDMA_APPLE)
55static constexpr size_t RDMA_CHUNK    = 256 * 1024;   // 256 KiB per send/recv (fits default 8 MiB memlock)
56static constexpr int    RDMA_RX_DEPTH = 24;            // pre-posted recv ring: 24 × 256 KiB = 6 MiB
57
58struct rdma_conn {
59    struct ibv_context * ctx = nullptr;
60    struct ibv_pd * pd  = nullptr;
61    struct ibv_cq * scq = nullptr;   // send completions
62    struct ibv_cq * rcq = nullptr;   // recv completions
63    struct ibv_qp * qp  = nullptr;
64
65    void          * tx_buf = nullptr;
66    struct ibv_mr * tx_mr  = nullptr;
67
68    void          * rx_buf = nullptr; // RDMA_RX_DEPTH × RDMA_CHUNK contiguous
69    struct ibv_mr * rx_mr  = nullptr;
70    int             rx_head = 0;
71
72    uint32_t        max_inline = 0;
73
74    uint8_t * rx_slot(int i) const {
75        return static_cast<uint8_t *>(rx_buf) + static_cast<size_t>(i) * RDMA_CHUNK;
76    }
77
78    bool post_rx(int i) {
79        struct ibv_sge sge = {};
80        sge.addr   = (uintptr_t)rx_slot(i);
81        sge.length = RDMA_CHUNK;
82        sge.lkey   = rx_mr->lkey;
83        struct ibv_recv_wr wr = {}, * bad = nullptr;
84        wr.wr_id   = (uint64_t)i;
85        wr.sg_list = &sge;
86        wr.num_sge = 1;
87        return ibv_post_recv(qp, &wr, &bad) == 0;
88    }
89
90    ~rdma_conn() {
91        if (tx_mr) ibv_dereg_mr(tx_mr);
92        if (rx_mr) ibv_dereg_mr(rx_mr);
93        free(tx_buf);
94        free(rx_buf);
95        if (qp)  ibv_destroy_qp(qp);
96        if (scq) ibv_destroy_cq(scq);
97        if (rcq) ibv_destroy_cq(rcq);
98        if (pd)  ibv_dealloc_pd(pd);
99        if (ctx) ibv_close_device(ctx);
100    }
101};
102
103// Local RDMA parameters captured during the probe phase and later consumed
104// by rdma_activate() after the remote side's caps arrive via HELLO.
105struct rdma_local_info {
106    uint32_t qpn     = 0;
107    uint32_t psn     = 0;
108    uint8_t  gid[RDMA_GID_SIZE] = {};
109    uint8_t  ib_port = 0;
110    int      gid_idx = 0;
111    enum ibv_mtu path_mtu = IBV_MTU_1024;
112};
113
114struct rdma_caps {
115    uint32_t qpn;
116    uint32_t psn;
117    uint8_t  gid[RDMA_GID_SIZE];
118};
119
120static_assert(sizeof(rdma_caps) == RPC_CONN_CAPS_SIZE, "rdma_caps must match conn_caps size");
121
122#endif // GGML_RPC_RDMA && !GGML_RPC_RDMA_APPLE
123
124struct socket_t::impl {
125    impl(sockfd_t fd) : use_rdma(false), fd(fd) {}
126    ~impl();
127    bool send_data(const void * data, size_t size);
128    bool recv_data(void * data, size_t size);
129    bool flush();
130    void get_caps(uint8_t * local_caps);
131    void update_caps(const uint8_t * remote_caps);
132
133#ifdef GGML_RPC_RDMA
134    std::optional<rdma_gid_t> rdma_build_target_gid();
135
136#  ifdef GGML_RPC_RDMA_APPLE
137    std::unique_ptr<apple_rdma> rdma;
138#  else
139    bool rdma_probe();
140    bool rdma_send(const void * data, size_t size);
141    bool rdma_recv(void * data, size_t size);
142    bool tcp_peer_closed();
143    bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid);
144    bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc);
145
146    std::unique_ptr<rdma_conn> rdma;
147    rdma_local_info            rdma_local = {};
148#  endif
149#endif // GGML_RPC_RDMA
150    bool     use_rdma;
151    sockfd_t fd;
152};
153
154socket_t::impl::~impl() {
155#ifdef GGML_RPC_RDMA
156    rdma.reset();
157#endif // GGML_RPC_RDMA
158    LOG_DBG("[%s] closing socket %d\n", __func__, this->fd);
159#ifdef _WIN32
160    if (fd != INVALID_SOCKET) closesocket(this->fd);
161#else
162    if (fd >= 0) close(this->fd);
163#endif
164}
165
166#ifdef GGML_RPC_RDMA
167
168// Build a RoCE GID-shaped 16-byte target from a TCP socket's local address.
169// Used to match the socket's local IP against the kernel's GID table so that
170// a single memcmp handles IPv4, IPv4-mapped IPv6, and native IPv6 uniformly:
171//   AF_INET                -> ::ffff:a.b.c.d  (bytes 10-11 = 0xff, last 4 = IPv4)
172//   AF_INET6 (IPv4-mapped) -> ::ffff:a.b.c.d  (already in GID shape)
173//   AF_INET6 (native v6)   -> the 16-byte IPv6 address as-is
174// Returns std::nullopt on unsupported family or getsockname failure.
175std::optional<rdma_gid_t> socket_t::impl::rdma_build_target_gid() {
176    sockaddr_storage addr = {};
177    socklen_t addr_len = sizeof(addr);
178    if (getsockname(fd, reinterpret_cast<sockaddr *>(&addr), &addr_len) != 0) {
179        return std::nullopt;
180    }
181    rdma_gid_t target = {};
182    if (addr.ss_family == AF_INET) {
183        const auto * a = reinterpret_cast<const sockaddr_in *>(&addr);
184        target[10] = 0xff;
185        target[11] = 0xff;
186        memcpy(&target[12], &a->sin_addr, 4);
187        return target;
188    }
189    if (addr.ss_family == AF_INET6) {
190        const auto * a = reinterpret_cast<const sockaddr_in6 *>(&addr);
191        memcpy(target.data(), &a->sin6_addr, RDMA_GID_SIZE);
192        return target;
193    }
194    return std::nullopt;
195}
196
197#ifndef GGML_RPC_RDMA_APPLE
198
199bool socket_t::impl::tcp_peer_closed() {
200    if (fd < 0) return false;
201#ifndef _WIN32
202    struct pollfd pfd = { fd, POLLIN | POLLRDHUP, 0 };
203    int r = poll(&pfd, 1, 0);
204    return r > 0 && (pfd.revents & (POLLHUP | POLLERR | POLLRDHUP));
205#else
206    return false;
207#endif
208}
209
210bool socket_t::impl::rdma_probe() {
211    const char * dev_env = std::getenv("GGML_RDMA_DEV");
212    const char * gid_env = std::getenv("GGML_RDMA_GID");
213
214    auto target_gid = rdma_build_target_gid();
215    if (!target_gid) {
216        return false;
217    }
218
219    const uint8_t ib_port = 1;
220    int num_devs = 0;
221    ibv_device ** devs = ibv_get_device_list(&num_devs);
222    if (!devs || num_devs == 0) return false;
223
224    ibv_context * ibctx = nullptr;
225    const char * matched_dev = nullptr;
226    int gid_idx = gid_env ? atoi(gid_env) : -1;
227    int gid_version = IBV_GID_TYPE_IB;  // 0 = unknown/IB
228
229    for (int d = 0; d < num_devs; d++) {
230        const char * dn = ibv_get_device_name(devs[d]);
231        if (dev_env && strcmp(dev_env, dn) != 0) continue;
232
233        ibv_context * ctx = ibv_open_device(devs[d]);
234        if (!ctx) continue;
235
236        ibv_port_attr pa;
237        if (ibv_query_port(ctx, ib_port, &pa) != 0) { ibv_close_device(ctx); continue; }
238
239        int found_gid = gid_idx;
240        int found_version = IBV_GID_TYPE_IB;
241        if (found_gid < 0) {
242            // Find a GID on this port whose bytes equal the local TCP address
243            // (IPv4 or IPv6). Prefer RoCE v2 (UDP/IP, L3-routable) over v1
244            // (raw Ethernet, same-L2 only) so silent hangs on L3-routed paths
245            // are avoided. ibv_query_gid_ex returns gid+type in one call.
246            int v2_idx = -1;
247            int v1_idx = -1;
248            for (int i = 0; i < pa.gid_tbl_len; i++) {
249                ibv_gid_entry entry = {};
250                if (ibv_query_gid_ex(ctx, ib_port, i, &entry, 0) != 0) continue;
251                if (memcmp(entry.gid.raw, target_gid->data(), RDMA_GID_SIZE) != 0) continue;
252                if (entry.gid_type == IBV_GID_TYPE_ROCE_V2 && v2_idx < 0) {
253                    v2_idx = i;
254                } else if (entry.gid_type == IBV_GID_TYPE_ROCE_V1 && v1_idx < 0) {
255                    v1_idx = i;
256                }
257            }
258            if (v2_idx >= 0) {
259                found_gid = v2_idx;
260                found_version = IBV_GID_TYPE_ROCE_V2;
261            } else if (v1_idx >= 0) {
262                found_gid = v1_idx;
263                found_version = IBV_GID_TYPE_ROCE_V1;
264            }
265        } else {
266            // Explicit GID index from GGML_RDMA_GID — fetch its type for logging.
267            ibv_gid_entry entry = {};
268            if (ibv_query_gid_ex(ctx, ib_port, found_gid, &entry, 0) == 0) {
269                found_version = entry.gid_type;
270            }
271        }
272        if (found_gid >= 0) {
273            ibctx = ctx;
274            gid_idx = found_gid;
275            gid_version = found_version;
276            matched_dev = dn;
277            rdma_local.path_mtu = pa.active_mtu;
278            break;
279        }
280        ibv_close_device(ctx);
281    }
282    ibv_free_device_list(devs);
283    if (!ibctx) return false;
284
285    rdma_local.ib_port = ib_port;
286    rdma_local.gid_idx = gid_idx;
287
288    rdma = std::make_unique<rdma_conn>();
289    rdma->ctx = ibctx;
290
291    rdma->pd = ibv_alloc_pd(ibctx);
292    if (!rdma->pd) return false;
293
294    rdma->scq = ibv_create_cq(ibctx, 16, nullptr, nullptr, 0);
295    rdma->rcq = ibv_create_cq(ibctx, RDMA_RX_DEPTH + 4, nullptr, nullptr, 0);
296    if (!rdma->scq || !rdma->rcq) return false;
297
298    ibv_qp_init_attr qia = {};
299    qia.send_cq = rdma->scq;
300    qia.recv_cq = rdma->rcq;
301    qia.qp_type = IBV_QPT_RC;
302    qia.cap.max_send_wr     = 4;
303    qia.cap.max_recv_wr     = RDMA_RX_DEPTH + 4;
304    qia.cap.max_send_sge    = 1;
305    qia.cap.max_recv_sge    = 1;
306    qia.cap.max_inline_data = 256;
307
308    rdma->qp = ibv_create_qp(rdma->pd, &qia);
309    if (!rdma->qp) return false;
310    rdma->max_inline = qia.cap.max_inline_data;
311
312    rdma->tx_buf = aligned_alloc(4096, RDMA_CHUNK);
313    rdma->rx_buf = aligned_alloc(4096, static_cast<size_t>(RDMA_RX_DEPTH) * RDMA_CHUNK);
314    if (!rdma->tx_buf || !rdma->rx_buf) return false;
315
316    rdma->tx_mr = ibv_reg_mr(rdma->pd, rdma->tx_buf, RDMA_CHUNK, IBV_ACCESS_LOCAL_WRITE);
317    rdma->rx_mr = ibv_reg_mr(rdma->pd, rdma->rx_buf, static_cast<size_t>(RDMA_RX_DEPTH) * RDMA_CHUNK,
318                           IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE);
319    if (!rdma->tx_mr || !rdma->rx_mr) return false;
320
321    ibv_gid local_gid;
322    if (ibv_query_gid(ibctx, ib_port, gid_idx, &local_gid) != 0) return false;
323
324    rdma_local.qpn = rdma->qp->qp_num;
325    rdma_local.psn = rdma->qp->qp_num & 0xffffff;
326    memcpy(&rdma_local.gid, &local_gid, RDMA_GID_SIZE);
327
328    const char * ver_str = "";
329    if (gid_version == IBV_GID_TYPE_ROCE_V2) {
330        ver_str = " RoCEv2";
331    } else if (gid_version == IBV_GID_TYPE_ROCE_V1) {
332        ver_str = " RoCEv1";
333    }
334    GGML_LOG_INFO("RDMA probed: dev=%s gid=%d%s qpn=%u inline=%u\n",
335                  matched_dev, gid_idx, ver_str, rdma_local.qpn, rdma->max_inline);
336    return true;
337}
338
339// Phase 2: Given remote QPN/PSN/GID, transition QP: RESET->INIT->pre-post->RTR->RTS.
340// On success, the connection is live and ready for rdma_send/rdma_recv.
341bool socket_t::impl::rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid) {
342    // RESET -> INIT
343    {
344        struct ibv_qp_attr a = {};
345        a.qp_state        = IBV_QPS_INIT;
346        a.port_num        = rdma_local.ib_port;
347        a.pkey_index      = 0;
348        a.qp_access_flags = IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_LOCAL_WRITE;
349        if (ibv_modify_qp(rdma->qp, &a,
350                IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS) != 0) {
351            return false;
352        }
353    }
354
355    for (int i = 0; i < RDMA_RX_DEPTH; i++) {
356        if (!rdma->post_rx(i)) return false;
357    }
358
359    // INIT -> RTR
360    {
361        struct ibv_qp_attr a = {};
362        a.qp_state           = IBV_QPS_RTR;
363        a.path_mtu           = rdma_local.path_mtu;
364        a.dest_qp_num        = remote_qpn;
365        a.rq_psn             = remote_psn;
366        a.max_dest_rd_atomic = 1;
367        a.min_rnr_timer      = 1;
368        a.ah_attr.is_global  = 1;
369        memcpy(&a.ah_attr.grh.dgid, remote_gid, RDMA_GID_SIZE);
370        a.ah_attr.grh.hop_limit  = 1;
371        a.ah_attr.grh.sgid_index = rdma_local.gid_idx;
372        a.ah_attr.dlid       = 0;
373        a.ah_attr.port_num   = rdma_local.ib_port;
374        if (ibv_modify_qp(rdma->qp, &a,
375                IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN |
376                IBV_QP_RQ_PSN | IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER) != 0) {
377            return false;
378        }
379    }
380
381    // RTR -> RTS
382    {
383        struct ibv_qp_attr a = {};
384        a.qp_state     = IBV_QPS_RTS;
385        a.timeout      = 14;
386        a.retry_cnt    = 7;
387        a.rnr_retry    = 7;
388        a.sq_psn       = rdma_local.psn;
389        a.max_rd_atomic = 1;
390        if (ibv_modify_qp(rdma->qp, &a,
391                IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT | IBV_QP_RNR_RETRY |
392                IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC) != 0) {
393            return false;
394        }
395    }
396
397    GGML_LOG_INFO("RDMA activated: qpn=%u->%u mtu=%d rx_depth=%d\n",
398                  rdma_local.qpn, remote_qpn, 128 << rdma_local.path_mtu, RDMA_RX_DEPTH);
399    return true;
400}
401
402bool socket_t::impl::rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc) {
403    for (uint64_t s = 0; ; s++) {
404        int n = ibv_poll_cq(cq, 1, wc);
405        if (n > 0) {
406            if (wc->status != IBV_WC_SUCCESS) {
407                GGML_LOG_ERROR("RDMA CQ wc error: status=%d (%s) vendor_err=0x%x\n",
408                    wc->status, ibv_wc_status_str(wc->status), wc->vendor_err);
409            }
410            return wc->status == IBV_WC_SUCCESS;
411        }
412        if (n < 0) return false;
413        if ((s & 0xFFFFF) == 0 && s > 0) {
414            if (tcp_peer_closed()) {
415                return false;
416            }
417        }
418    }
419}
420
421bool socket_t::impl::rdma_send(const void * data, size_t size) {
422    rdma_conn * c = rdma.get();
423    const uint8_t * src = (const uint8_t *)data;
424    size_t rem = size;
425    while (rem > 0) {
426        size_t chunk = std::min(rem, RDMA_CHUNK);
427
428        struct ibv_sge sge = {};
429        struct ibv_send_wr wr = {}, * bad = nullptr;
430        wr.opcode  = IBV_WR_SEND;
431        wr.sg_list = &sge;
432        wr.num_sge = 1;
433
434        if (chunk <= c->max_inline) {
435            sge.addr   = (uintptr_t)src;
436            sge.length = chunk;
437            wr.send_flags = IBV_SEND_SIGNALED | IBV_SEND_INLINE;
438        } else {
439            memcpy(c->tx_buf, src, chunk);
440            sge.addr   = (uintptr_t)c->tx_buf;
441            sge.length = chunk;
442            sge.lkey   = c->tx_mr->lkey;
443            wr.send_flags = IBV_SEND_SIGNALED;
444        }
445
446        if (ibv_post_send(c->qp, &wr, &bad) != 0) return false;
447        struct ibv_wc wc;
448        if (!rdma_poll(c->scq, &wc)) return false;
449
450        src += chunk;
451        rem -= chunk;
452    }
453    return true;
454}
455
456bool socket_t::impl::rdma_recv(void * data, size_t size) {
457    rdma_conn * c = rdma.get();
458    uint8_t * dst = (uint8_t *)data;
459    size_t rem = size;
460    while (rem > 0) {
461        struct ibv_wc wc;
462        if (!rdma_poll(c->rcq, &wc)) return false;
463
464        int slot = (int)wc.wr_id;
465        size_t got = wc.byte_len;
466        memcpy(dst, c->rx_slot(slot), got);
467
468        if (!c->post_rx(slot)) return false;
469
470        dst += got;
471        rem -= got;
472    }
473    return true;
474}
475
476#endif // !GGML_RPC_RDMA_APPLE (Linux RC transport)
477
478#endif // GGML_RPC_RDMA
479
480bool socket_t::impl::send_data(const void * data, size_t size) {
481#ifdef GGML_RPC_RDMA_APPLE
482    if (use_rdma) {
483        return rdma->send(data, size);
484    }
485#elif defined(GGML_RPC_RDMA)
486    if (use_rdma) {
487        return rdma_send(data, size);
488    }
489#endif
490    size_t bytes_sent = 0;
491    while (bytes_sent < size) {
492        size_t size_to_send = std::min(size - bytes_sent, MAX_CHUNK_SIZE);
493        ssize_t n = send(fd, (const char *)data + bytes_sent, size_to_send, 0);
494        if (n < 0) {
495            GGML_LOG_ERROR("send failed (bytes_sent=%zu, size_to_send=%zu)\n",
496                           bytes_sent, size_to_send);
497            return false;
498        }
499        bytes_sent += (size_t)n;
500    }
501    return true;
502}
503
504bool socket_t::impl::recv_data(void * data, size_t size) {
505#ifdef GGML_RPC_RDMA_APPLE
506    if (use_rdma) {
507        return rdma->recv(data, size);
508    }
509#elif defined(GGML_RPC_RDMA)
510    if (use_rdma) {
511        return rdma_recv(data, size);
512    }
513#endif
514    size_t bytes_recv = 0;
515    while (bytes_recv < size) {
516        size_t size_to_recv = std::min(size - bytes_recv, MAX_CHUNK_SIZE);
517        ssize_t n = recv(fd, (char *)data + bytes_recv, size_to_recv, 0);
518        if (n < 0) {
519            GGML_LOG_ERROR("recv failed (bytes_recv=%zu, size_to_recv=%zu)\n",
520                           bytes_recv, size_to_recv);
521            return false;
522        }
523        if (n == 0) {
524            LOG_DBG("recv returned 0 (peer closed?)\n");
525            return false;
526        }
527        bytes_recv += (size_t)n;
528    }
529    return true;
530}
531
532void socket_t::impl::get_caps(uint8_t * local_caps) {
533    memset(local_caps, 0, RPC_CONN_CAPS_SIZE);
534#ifdef GGML_RPC_RDMA
535    if (std::getenv("GGML_RPC_NO_RDMA")) {
536        return;
537    }
538#  ifdef GGML_RPC_RDMA_APPLE
539    auto target_gid = rdma_build_target_gid();
540    if (target_gid) {
541        rdma = apple_rdma::probe(fd, target_gid->data(), local_caps);
542    }
543#  else
544    rdma_local = {};
545    if (rdma_probe()) {
546        rdma_caps rc = {};
547        rc.qpn = rdma_local.qpn;
548        rc.psn = rdma_local.psn;
549        memcpy(rc.gid, rdma_local.gid, RDMA_GID_SIZE);
550        memcpy(local_caps, &rc, sizeof(rc));
551    } else {
552        rdma.reset();
553    }
554#  endif
555#endif // GGML_RPC_RDMA
556}
557
558void socket_t::impl::update_caps(const uint8_t * remote_caps) {
559#ifdef GGML_RPC_RDMA
560    // a peer that has no RDMA advertises all-zero caps and takes no further part
561    // in the negotiation, so drop to TCP without reporting a failure
562    bool remote_rdma = false;
563    for (size_t i = 0; i < RPC_CONN_CAPS_SIZE; i++) {
564        remote_rdma |= remote_caps[i] != 0;
565    }
566    if (!rdma || !remote_rdma) {
567        rdma.reset();
568        return;
569    }
570#  ifdef GGML_RPC_RDMA_APPLE
571    bool activated = rdma->activate(remote_caps);
572#  else
573    rdma_caps rc = {};
574    memcpy(&rc, remote_caps, sizeof(rc));
575    bool activated = rdma_activate(rc.qpn, rc.psn, rc.gid);
576#  endif
577    if (activated) {
578        use_rdma = true;
579    } else {
580        GGML_LOG_ERROR("RDMA activate failed, staying on TCP\n");
581        rdma.reset();
582    }
583#else
584    (void)remote_caps;
585#endif // GGML_RPC_RDMA
586}
587
588bool socket_t::impl::flush() {
589#ifdef GGML_RPC_RDMA_APPLE
590    if (use_rdma) {
591        return rdma->flush();
592    }
593#endif
594    return true;
595}
596
597/////////////////////////////////////////////////////////////////////////////
598
599socket_t::socket_t(std::unique_ptr<impl> p) : pimpl(std::move(p)) {}
600
601socket_t::~socket_t() = default;
602
603bool socket_t::send_data(const void * data, size_t size) {
604    return pimpl->send_data(data, size);
605}
606
607bool socket_t::recv_data(void * data, size_t size) {
608    return pimpl->recv_data(data, size);
609}
610
611bool socket_t::flush() {
612    return pimpl->flush();
613}
614
615void socket_t::get_caps(uint8_t * local_caps) {
616    return pimpl->get_caps(local_caps);
617}
618
619void socket_t::update_caps(const uint8_t * remote_caps) {
620    return pimpl->update_caps(remote_caps);
621}
622
623static bool is_valid_fd(sockfd_t sockfd) {
624#ifdef _WIN32
625    return sockfd != INVALID_SOCKET;
626#else
627    return sockfd >= 0;
628#endif
629}
630
631static bool set_no_delay(sockfd_t sockfd) {
632    int flag = 1;
633    // set TCP_NODELAY to disable Nagle's algorithm
634    int ret = setsockopt(sockfd, IPPROTO_TCP, TCP_NODELAY, (char *)&flag, sizeof(int));
635    return ret == 0;
636}
637
638static bool set_reuse_addr(sockfd_t sockfd) {
639    int flag = 1;
640    int ret = setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, (char *)&flag, sizeof(int));
641    return ret == 0;
642}
643
644socket_ptr socket_t::accept() {
645    auto client_socket_fd = ::accept(pimpl->fd, NULL, NULL);
646    if (!is_valid_fd(client_socket_fd)) {
647        return nullptr;
648    }
649    if (!set_no_delay(client_socket_fd)) {
650        GGML_LOG_ERROR("Failed to set TCP_NODELAY\n");
651        return nullptr;
652    }
653    return socket_ptr(new socket_t(std::make_unique<impl>(client_socket_fd)));
654}
655
656socket_ptr socket_t::create_server(const char * host, int port) {
657    auto sockfd = socket(AF_INET, SOCK_STREAM, 0);
658    if (!is_valid_fd(sockfd)) {
659        return nullptr;
660    }
661    if (!set_reuse_addr(sockfd)) {
662        GGML_LOG_ERROR("Failed to set SO_REUSEADDR\n");
663        return nullptr;
664    }
665    if (inet_addr(host) == INADDR_NONE) {
666        GGML_LOG_ERROR("Invalid host address: %s\n", host);
667        return nullptr;
668    }
669    struct sockaddr_in serv_addr;
670    serv_addr.sin_family = AF_INET;
671    serv_addr.sin_addr.s_addr = inet_addr(host);
672    serv_addr.sin_port = htons(port);
673
674    if (bind(sockfd, (struct sockaddr *) &serv_addr, sizeof(serv_addr)) < 0) {
675        return nullptr;
676    }
677    if (listen(sockfd, 1) < 0) {
678        return nullptr;
679    }
680    return socket_ptr(new socket_t(std::make_unique<impl>(sockfd)));
681}
682
683socket_ptr socket_t::connect(const char * host, int port) {
684    auto sockfd = socket(AF_INET, SOCK_STREAM, 0);
685    if (!is_valid_fd(sockfd)) {
686        return nullptr;
687    }
688    if (!set_no_delay(sockfd)) {
689        GGML_LOG_ERROR("Failed to set TCP_NODELAY\n");
690        return nullptr;
691    }
692    struct sockaddr_in addr;
693    addr.sin_family = AF_INET;
694    addr.sin_port = htons(port);
695    struct hostent * server = gethostbyname(host);
696    if (server == NULL) {
697        GGML_LOG_ERROR("Cannot resolve host '%s'\n", host);
698        return nullptr;
699    }
700    memcpy(&addr.sin_addr.s_addr, server->h_addr, server->h_length);
701    if (::connect(sockfd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
702        return nullptr;
703    }
704    return socket_ptr(new socket_t(std::make_unique<impl>(sockfd)));
705}
706
707#ifdef _WIN32
708static std::mutex g_rpc_transport_mu;
709static bool g_rpc_transport_wsa_started = false;
710#endif
711
712bool rpc_transport_init() {
713#ifdef _WIN32
714    std::lock_guard<std::mutex> lock(g_rpc_transport_mu);
715    if (g_rpc_transport_wsa_started) {
716        return true;
717    }
718    WSADATA wsaData;
719    int res = WSAStartup(MAKEWORD(2, 2), &wsaData);
720    if (res != 0) {
721        return false;
722    }
723    g_rpc_transport_wsa_started = true;
724    return true;
725#else
726    return true;
727#endif
728}
729
730void rpc_transport_shutdown() {
731#ifdef _WIN32
732    std::lock_guard<std::mutex> lock(g_rpc_transport_mu);
733    if (!g_rpc_transport_wsa_started) {
734        return;
735    }
736    WSACleanup();
737    g_rpc_transport_wsa_started = false;
738#endif
739}