MCPcopy Create free account
hub / github.com/deepspeedai/DeepSpeed / shm_initialize

Function shm_initialize

csrc/cpu/comm/shm.cpp:340–394  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

338static int world_rank;
339
340void shm_initialize(int size, int rank, char* addr_string, char* port_string)
341{
342 if (is_initialized) return;
343 is_initialized = 1;
344
345 world_size = size;
346 world_rank = rank;
347
348 char shm_name_prefix[NAME_BUF_SIZE];
349 char shm_name[NAME_BUF_SIZE];
350 snprintf(shm_name_prefix,
351 NAME_BUF_SIZE,
352 "%s_%d_%s_%s",
353 SHM_BUFFER_NAME,
354 getuid(),
355 addr_string,
356 port_string);
357 // create shared workspace for SHM based allreduce
358 SharedData allreduce_buffer;
359 // allocate workspace_buf for current rank
360 struct allreduce_workspace* workspace_buf;
361 struct allreduce_workspace* workspace_buf_other;
362 workspace_buf = (struct allreduce_workspace*)malloc(sizeof(struct allreduce_workspace));
363 snprintf(shm_name, NAME_BUF_SIZE, "%s_%d", shm_name_prefix, rank);
364 shared_create(&allreduce_buffer, shm_name, workspace_buf, sizeof(struct allreduce_workspace));
365 workspace_buf = (struct allreduce_workspace*)allreduce_buffer.bytes;
366 workspace_buf->states[0] = coll_alt2_allreduce_naive__copy_in_done;
367 workspace_buf->states[1] = coll_begin;
368
369 // calloc used for defensive zero-init; the loop below writes every element before any read
370 workspace = (struct allreduce_workspace**)calloc(size, sizeof(struct allreduce_workspace*));
371 symmetric_buffer[0] = (char**)calloc(size, sizeof(char*));
372 symmetric_buffer[1] = (char**)calloc(size, sizeof(char*));
373 distributed_buffer[0] = (char**)calloc(size, sizeof(char*));
374 distributed_buffer[1] = (char**)calloc(size, sizeof(char*));
375
376 // map shm of all ranks
377 for (int i = 0; i < size; i++) {
378 if (i != rank) {
379 snprintf(shm_name, NAME_BUF_SIZE, "%s_%d", shm_name_prefix, i);
380 // printf("open %s, %d\n", shm_name, rank);
381 do {
382 shared_open(&allreduce_buffer, shm_name, sizeof(struct allreduce_workspace));
383 } while (allreduce_buffer.descriptor == -1 && errno == ENOENT);
384 workspace_buf_other = (struct allreduce_workspace*)allreduce_buffer.bytes;
385 workspace[i] = workspace_buf_other;
386 } else {
387 workspace[i] = workspace_buf;
388 }
389 symmetric_buffer[0][i] = workspace[i]->buffer + BUFFER0_OFFSET(0);
390 symmetric_buffer[1][i] = workspace[i]->buffer + BUFFER0_OFFSET(1);
391 distributed_buffer[0][i] = workspace[i]->buffer + BUFFER1_OFFSET(0);
392 distributed_buffer[1][i] = workspace[i]->buffer + BUFFER1_OFFSET(1);
393 }
394}
395
396void parallel_memcpy(void* to, void* from, size_t n_bytes)
397{

Callers 2

initializeFunction · 0.85
initializeFunction · 0.85

Calls 2

shared_createFunction · 0.85
shared_openFunction · 0.85

Tested by

no test coverage detected