NVIDIA_cuda-samples

1
/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
2
 *
3
 * Redistribution and use in source and binary forms, with or without
4
 * modification, are permitted provided that the following conditions
5
 * are met:
6
 *  * Redistributions of source code must retain the above copyright
7
 *    notice, this list of conditions and the following disclaimer.
8
 *  * Redistributions in binary form must reproduce the above copyright
9
 *    notice, this list of conditions and the following disclaimer in the
10
 *    documentation and/or other materials provided with the distribution.
11
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
12
 *    contributors may be used to endorse or promote products derived
13
 *    from this software without specific prior written permission.
14
 *
15
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
16
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
19
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
23
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26
 */
27

28
#include "helper_multiprocess.h"
29
#include <cstdlib>
30
#include <string>
31

32
int sharedMemoryCreate(const char *name, size_t sz, sharedMemoryInfo *info) {
33
#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
34
  info->size = sz;
35
  info->shmHandle = CreateFileMapping(INVALID_HANDLE_VALUE, NULL,
36
                                      PAGE_READWRITE, 0, (DWORD)sz, name);
37
  if (info->shmHandle == 0) {
38
    return GetLastError();
39
  }
40

41
  info->addr = MapViewOfFile(info->shmHandle, FILE_MAP_ALL_ACCESS, 0, 0, sz);
42
  if (info->addr == NULL) {
43
    return GetLastError();
44
  }
45

46
  return 0;
47
#else
48
  int status = 0;
49

50
  info->size = sz;
51

52
  info->shmFd = shm_open(name, O_RDWR | O_CREAT, 0777);
53
  if (info->shmFd < 0) {
54
    return errno;
55
  }
56

57
  status = ftruncate(info->shmFd, sz);
58
  if (status != 0) {
59
    return status;
60
  }
61

62
  info->addr = mmap(0, sz, PROT_READ | PROT_WRITE, MAP_SHARED, info->shmFd, 0);
63
  if (info->addr == NULL) {
64
    return errno;
65
  }
66

67
  return 0;
68
#endif
69
}
70

71
int sharedMemoryOpen(const char *name, size_t sz, sharedMemoryInfo *info) {
72
#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
73
  info->size = sz;
74

75
  info->shmHandle = OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, name);
76
  if (info->shmHandle == 0) {
77
    return GetLastError();
78
  }
79

80
  info->addr = MapViewOfFile(info->shmHandle, FILE_MAP_ALL_ACCESS, 0, 0, sz);
81
  if (info->addr == NULL) {
82
    return GetLastError();
83
  }
84

85
  return 0;
86
#else
87
  info->size = sz;
88

89
  info->shmFd = shm_open(name, O_RDWR, 0777);
90
  if (info->shmFd < 0) {
91
    return errno;
92
  }
93

94
  info->addr = mmap(0, sz, PROT_READ | PROT_WRITE, MAP_SHARED, info->shmFd, 0);
95
  if (info->addr == NULL) {
96
    return errno;
97
  }
98

99
  return 0;
100
#endif
101
}
102

103
void sharedMemoryClose(sharedMemoryInfo *info) {
104
#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
105
  if (info->addr) {
106
    UnmapViewOfFile(info->addr);
107
  }
108
  if (info->shmHandle) {
109
    CloseHandle(info->shmHandle);
110
  }
111
#else
112
  if (info->addr) {
113
    munmap(info->addr, info->size);
114
  }
115
  if (info->shmFd) {
116
    close(info->shmFd);
117
  }
118
#endif
119
}
120

121
int spawnProcess(Process *process, const char *app, char *const *args) {
122
#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
123
  STARTUPINFO si = {0};
124
  BOOL status;
125
  size_t arglen = 0;
126
  size_t argIdx = 0;
127
  std::string arg_string;
128
  memset(process, 0, sizeof(*process));
129

130
  while (*args) {
131
    arg_string.append(*args).append(1, ' ');
132
    args++;
133
  }
134

135
  status = CreateProcess(app, LPSTR(arg_string.c_str()), NULL, NULL, FALSE, 0,
136
                         NULL, NULL, &si, process);
137

138
  return status ? 0 : GetLastError();
139
#else
140
  *process = fork();
141
  if (*process == 0) {
142
    if (0 > execvp(app, args)) {
143
      return errno;
144
    }
145
  } else if (*process < 0) {
146
    return errno;
147
  }
148
  return 0;
149
#endif
150
}
151

152
int waitProcess(Process *process) {
153
#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
154
  DWORD exitCode;
155
  WaitForSingleObject(process->hProcess, INFINITE);
156
  GetExitCodeProcess(process->hProcess, &exitCode);
157
  CloseHandle(process->hProcess);
158
  CloseHandle(process->hThread);
159
  return (int)exitCode;
160
#else
161
  int status = 0;
162
  do {
163
    if (0 > waitpid(*process, &status, 0)) {
164
      return errno;
165
    }
166
  } while (!WIFEXITED(status));
167
  return WEXITSTATUS(status);
168
#endif
169
}
170

171
#if defined(__linux__) || defined(__QNX__)
172
int ipcCreateSocket(ipcHandle *&handle, const char *name,
173
                    const std::vector<Process> &processes) {
174
  int server_fd;
175
  struct sockaddr_un servaddr;
176

177
  handle = new ipcHandle;
178
  memset(handle, 0, sizeof(*handle));
179
  handle->socket = -1;
180
  handle->socketName = NULL;
181

182
  // Creating socket file descriptor
183
  if ((server_fd = socket(AF_UNIX, SOCK_DGRAM, 0)) == 0) {
184
    perror("IPC failure: Socket creation failed");
185
    return -1;
186
  }
187

188
  unlink(name);
189
  bzero(&servaddr, sizeof(servaddr));
190
  servaddr.sun_family = AF_UNIX;
191

192
  size_t len = strlen(name);
193
  if (len > (sizeof(servaddr.sun_path) - 1)) {
194
    perror("IPC failure: Cannot bind provided name to socket. Name too large");
195
    return -1;
196
  }
197

198
  strncpy(servaddr.sun_path, name, len);
199

200
  if (bind(server_fd, (struct sockaddr *)&servaddr, SUN_LEN(&servaddr)) < 0) {
201
    perror("IPC failure: Binding socket failed");
202
    return -1;
203
  }
204

205
  handle->socketName = new char[strlen(name) + 1];
206
  strcpy(handle->socketName, name);
207
  handle->socket = server_fd;
208
  return 0;
209
}
210

211
int ipcOpenSocket(ipcHandle *&handle) {
212
  int sock = 0;
213
  struct sockaddr_un cliaddr;
214

215
  handle = new ipcHandle;
216
  memset(handle, 0, sizeof(*handle));
217

218
  if ((sock = socket(AF_UNIX, SOCK_DGRAM, 0)) < 0) {
219
    perror("IPC failure:Socket creation error");
220
    return -1;
221
  }
222

223
  bzero(&cliaddr, sizeof(cliaddr));
224
  cliaddr.sun_family = AF_UNIX;
225
  char temp[10];
226

227
  // Create unique name for the socket.
228
  sprintf(temp, "%u", getpid());
229

230
  strcpy(cliaddr.sun_path, temp);
231
  if (bind(sock, (struct sockaddr *)&cliaddr, sizeof(cliaddr)) < 0) {
232
    perror("IPC failure: Binding socket failed");
233
    return -1;
234
  }
235

236
  handle->socket = sock;
237
  handle->socketName = new char[strlen(temp) + 1];
238
  strcpy(handle->socketName, temp);
239

240
  return 0;
241
}
242

243
int ipcCloseSocket(ipcHandle *handle) {
244
  if (!handle) {
245
    return -1;
246
  }
247

248
  if (handle->socketName) {
249
    unlink(handle->socketName);
250
    delete[] handle->socketName;
251
  }
252
  close(handle->socket);
253
  delete handle;
254
  return 0;
255
}
256

257
int ipcRecvShareableHandle(ipcHandle *handle, ShareableHandle *shHandle) {
258
  struct msghdr msg = {0};
259
  struct iovec iov[1];
260
  struct cmsghdr cm;
261

262
  // Union to guarantee alignment requirements for control array
263
  union {
264
    struct cmsghdr cm;
265
    // This will not work on QNX as QNX CMSG_SPACE calls __cmsg_alignbytes
266
    // And __cmsg_alignbytes is a runtime function instead of compile-time macros
267
    // char control[CMSG_SPACE(sizeof(int))]
268
    char* control;
269
  } control_un;
270

271
  size_t sizeof_control = CMSG_SPACE(sizeof(int)) * sizeof(char);
272
  control_un.control = (char*) malloc(sizeof_control);
273
  struct cmsghdr *cmptr;
274
  ssize_t n;
275
  int receivedfd;
276
  char dummy_buffer[1];
277
  ssize_t sendResult;
278
  msg.msg_control = control_un.control;
279
  msg.msg_controllen = sizeof_control;
280

281
  iov[0].iov_base = (void *)dummy_buffer;
282
  iov[0].iov_len = sizeof(dummy_buffer);
283

284
  msg.msg_iov = iov;
285
  msg.msg_iovlen = 1;
286
  if ((n = recvmsg(handle->socket, &msg, 0)) <= 0) {
287
    perror("IPC failure: Receiving data over socket failed");
288
    free(control_un.control);
289
    return -1;
290
  }
291

292
  if (((cmptr = CMSG_FIRSTHDR(&msg)) != NULL) &&
293
      (cmptr->cmsg_len == CMSG_LEN(sizeof(int)))) {
294
    if ((cmptr->cmsg_level != SOL_SOCKET) || (cmptr->cmsg_type != SCM_RIGHTS)) {
295
      free(control_un.control);
296
      return -1;
297
    }
298

299
    memmove(&receivedfd, CMSG_DATA(cmptr), sizeof(receivedfd));
300
    *(int *)shHandle = receivedfd;
301
  } else {
302
    free(control_un.control);
303
    return -1;
304
  }
305

306
  free(control_un.control);
307
  return 0;
308
}
309

310
int ipcRecvDataFromClient(ipcHandle *serverHandle, void *data, size_t size) {
311
  ssize_t readResult;
312
  struct sockaddr_un cliaddr;
313
  socklen_t len = sizeof(cliaddr);
314

315
  readResult = recvfrom(serverHandle->socket, data, size, 0,
316
                        (struct sockaddr *)&cliaddr, &len);
317
  if (readResult == -1) {
318
    perror("IPC failure: Receiving data over socket failed");
319
    return -1;
320
  }
321
  return 0;
322
}
323

324
int ipcSendDataToServer(ipcHandle *handle, const char *serverName,
325
                        const void *data, size_t size) {
326
  ssize_t sendResult;
327
  struct sockaddr_un serveraddr;
328

329
  bzero(&serveraddr, sizeof(serveraddr));
330
  serveraddr.sun_family = AF_UNIX;
331
  strncpy(serveraddr.sun_path, serverName, sizeof(serveraddr.sun_path) - 1);
332

333
  sendResult = sendto(handle->socket, data, size, 0,
334
                      (struct sockaddr *)&serveraddr, sizeof(serveraddr));
335
  if (sendResult <= 0) {
336
    perror("IPC failure: Sending data over socket failed");
337
  }
338

339
  return 0;
340
}
341

342
int ipcSendShareableHandle(ipcHandle *handle,
343
                           const std::vector<ShareableHandle> &shareableHandles,
344
                           Process process, int data) {
345
  struct msghdr msg;
346
  struct iovec iov[1];
347

348
  union {
349
    struct cmsghdr cm;
350
    char* control;
351
  } control_un;
352

353
  size_t sizeof_control = CMSG_SPACE(sizeof(int)) * sizeof(char);
354
  control_un.control = (char*) malloc(sizeof_control);
355

356
  struct cmsghdr *cmptr;
357
  ssize_t readResult;
358
  struct sockaddr_un cliaddr;
359
  socklen_t len = sizeof(cliaddr);
360

361
  // Construct client address to send this SHareable handle to
362
  bzero(&cliaddr, sizeof(cliaddr));
363
  cliaddr.sun_family = AF_UNIX;
364
  char temp[10];
365
  sprintf(temp, "%u", process);
366
  strcpy(cliaddr.sun_path, temp);
367
  len = sizeof(cliaddr);
368

369
  // Send corresponding shareable handle to the client
370
  int sendfd = (int)shareableHandles[data];
371

372
  msg.msg_control = control_un.control;
373
  msg.msg_controllen = sizeof_control;
374

375
  cmptr = CMSG_FIRSTHDR(&msg);
376
  cmptr->cmsg_len = CMSG_LEN(sizeof(int));
377
  cmptr->cmsg_level = SOL_SOCKET;
378
  cmptr->cmsg_type = SCM_RIGHTS;
379

380
  memmove(CMSG_DATA(cmptr), &sendfd, sizeof(sendfd));
381

382
  msg.msg_name = (void *)&cliaddr;
383
  msg.msg_namelen = sizeof(struct sockaddr_un);
384

385
  iov[0].iov_base = (void *)"";
386
  iov[0].iov_len = 1;
387
  msg.msg_iov = iov;
388
  msg.msg_iovlen = 1;
389

390
  ssize_t sendResult = sendmsg(handle->socket, &msg, 0);
391
  if (sendResult <= 0) {
392
    perror("IPC failure: Sending data over socket failed");
393
    free(control_un.control);
394
    return -1;
395
  }
396

397
  free(control_un.control);
398
  return 0;
399
}
400

401
int ipcSendShareableHandles(
402
    ipcHandle *handle, const std::vector<ShareableHandle> &shareableHandles,
403
    const std::vector<Process> &processes) {
404
  // Send all shareable handles to every single process.
405
  for (int i = 0; i < shareableHandles.size(); i++) {
406
    for (int j = 0; j < processes.size(); j++) {
407
      checkIpcErrors(
408
          ipcSendShareableHandle(handle, shareableHandles, processes[j], i));
409
    }
410
  }
411
  return 0;
412
}
413

414
int ipcRecvShareableHandles(ipcHandle *handle,
415
                            std::vector<ShareableHandle> &shareableHandles) {
416
  for (int i = 0; i < shareableHandles.size(); i++) {
417
    checkIpcErrors(ipcRecvShareableHandle(handle, &shareableHandles[i]));
418
  }
419
  return 0;
420
}
421

422
int ipcCloseShareableHandle(ShareableHandle shHandle) {
423
  return close(shHandle);
424
}
425

426
#elif defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
427
// Generic name to build individual Mailslot names by appending process ids.
428
LPTSTR SlotName = (LPTSTR)TEXT("\\\\.\\mailslot\\sample_mailslot_");
429

430
int ipcCreateSocket(ipcHandle *&handle, const char *name,
431
                    const std::vector<Process> &processes) {
432
  handle = new ipcHandle;
433
  handle->hMailslot.resize(processes.size());
434

435
  // Open Mailslots of all clients and store respective handles.
436
  for (int i = 0; i < handle->hMailslot.size(); ++i) {
437
    std::basic_string<TCHAR> childSlotName(SlotName);
438
    char tempBuf[20];
439
    _itoa_s(processes[i].dwProcessId, tempBuf, 10);
440
    childSlotName += TEXT(tempBuf);
441

442
    HANDLE hFile =
443
        CreateFile(TEXT(childSlotName.c_str()), GENERIC_WRITE, FILE_SHARE_READ,
444
                   (LPSECURITY_ATTRIBUTES)NULL, OPEN_EXISTING,
445
                   FILE_ATTRIBUTE_NORMAL, (HANDLE)NULL);
446
    if (hFile == INVALID_HANDLE_VALUE) {
447
      printf("IPC failure: Opening Mailslot by CreateFile failed with %d\n",
448
             GetLastError());
449
      return -1;
450
    }
451
    handle->hMailslot[i] = hFile;
452
  }
453
  return 0;
454
}
455

456
int ipcOpenSocket(ipcHandle *&handle) {
457
  handle = new ipcHandle;
458
  HANDLE hSlot;
459

460
  std::basic_string<TCHAR> clientSlotName(SlotName);
461
  char tempBuf[20];
462
  _itoa_s(GetCurrentProcessId(), tempBuf, 10);
463
  clientSlotName += TEXT(tempBuf);
464

465
  hSlot = CreateMailslot((LPSTR)clientSlotName.c_str(), 0,
466
                         MAILSLOT_WAIT_FOREVER, (LPSECURITY_ATTRIBUTES)NULL);
467
  if (hSlot == INVALID_HANDLE_VALUE) {
468
    printf("IPC failure: CreateMailslot failed for client with %d\n",
469
           GetLastError());
470
    return -1;
471
  }
472

473
  handle->hMailslot.push_back(hSlot);
474
  return 0;
475
}
476

477
int ipcSendData(HANDLE mailslot, const void *data, size_t sz) {
478
  BOOL result;
479
  DWORD cbWritten;
480

481
  result = WriteFile(mailslot, data, (DWORD)sz, &cbWritten, (LPOVERLAPPED)NULL);
482
  if (!result) {
483
    printf("IPC failure: WriteFile failed with %d.\n", GetLastError());
484
    return -1;
485
  }
486
  return 0;
487
}
488

489
int ipcRecvData(ipcHandle *handle, void *data, size_t sz) {
490
  DWORD cbRead = 0;
491

492
  if (!ReadFile(handle->hMailslot[0], data, (DWORD)sz, &cbRead, NULL)) {
493
    printf("IPC failure: ReadFile failed with %d.\n", GetLastError());
494
    return -1;
495
  }
496

497
  if (sz != (size_t)cbRead) {
498
    printf(
499
        "IPC failure: ReadFile didn't receive the expected number of bytes\n");
500
    return -1;
501
  }
502

503
  return 0;
504
}
505

506
int ipcSendShareableHandles(
507
    ipcHandle *handle, const std::vector<ShareableHandle> &shareableHandles,
508
    const std::vector<Process> &processes) {
509
  // Send all shareable handles to every single process.
510
  for (int i = 0; i < processes.size(); i++) {
511
    HANDLE hProcess =
512
        OpenProcess(PROCESS_DUP_HANDLE, FALSE, processes[i].dwProcessId);
513
    if (hProcess == INVALID_HANDLE_VALUE) {
514
      printf("IPC failure: OpenProcess failed (%d)\n", GetLastError());
515
      return -1;
516
    }
517

518
    for (int j = 0; j < shareableHandles.size(); j++) {
519
      HANDLE hDup = INVALID_HANDLE_VALUE;
520
      // Duplicate the handle into the target process's space
521
      if (!DuplicateHandle(GetCurrentProcess(), shareableHandles[j], hProcess,
522
                           &hDup, 0, FALSE, DUPLICATE_SAME_ACCESS)) {
523
        printf("IPC failure: DuplicateHandle failed (%d)\n", GetLastError());
524
        return -1;
525
      }
526
      checkIpcErrors(ipcSendData(handle->hMailslot[i], &hDup, sizeof(hDup)));
527
    }
528
    CloseHandle(hProcess);
529
  }
530
  return 0;
531
}
532

533
int ipcRecvShareableHandles(ipcHandle *handle,
534
                            std::vector<ShareableHandle> &shareableHandles) {
535
  for (int i = 0; i < shareableHandles.size(); i++) {
536
    checkIpcErrors(
537
        ipcRecvData(handle, &shareableHandles[i], sizeof(shareableHandles[i])));
538
  }
539
  return 0;
540
}
541

542
int ipcCloseSocket(ipcHandle *handle) {
543
  for (int i = 0; i < handle->hMailslot.size(); i++) {
544
    CloseHandle(handle->hMailslot[i]);
545
  }
546
  delete handle;
547
  return 0;
548
}
549

550
int ipcCloseShareableHandle(ShareableHandle shHandle) {
551
  CloseHandle(shHandle);
552
  return 0;
553
}
554

555
#endif
556