2025-05-22 02:17:10 +08:00
|
|
|
#include <torch/all.h>
|
2025-04-09 00:37:46 +08:00
|
|
|
|
|
|
|
|
#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
|
|
|
|
|
|
|
|
|
|
#ifndef __SHM_COLLECTIVES__
|
|
|
|
|
#define __SHM_COLLECTIVES__
|
|
|
|
|
#define VECTOR_LENGTH_IN_BYTES 32
|
2025-06-10 16:08:15 +08:00
|
|
|
void shm_initialize(int size, int rank, const char* addr_string, const char* port_string);
|
2025-04-09 00:37:46 +08:00
|
|
|
void all_reduce_outer_loop(torch::Tensor& data, size_t numel, int data_size);
|
|
|
|
|
torch::Tensor& all_gather(torch::Tensor& result, torch::Tensor& data, int dim, size_t numel, int data_size);
|
|
|
|
|
#endif
|