| 22 | static bool all_ranks_local_p = false; |
| 23 | |
| 24 | void initialize(int size, int rank) |
| 25 | { |
| 26 | if (is_initialized) return; |
| 27 | |
| 28 | // Check whether all ranks is on the same physical machine. |
| 29 | // If true, we will use an SHM based low latency allreduce |
| 30 | |
| 31 | auto ls_string = std::getenv("LOCAL_SIZE"); |
| 32 | int ls = 0; |
| 33 | if (ls_string != NULL) { ls = std::stoi(std::getenv("LOCAL_SIZE")); } |
| 34 | |
| 35 | if (size >= 1 && size == ls) { all_ranks_local_p = true; } |
| 36 | |
| 37 | world_size = size; |
| 38 | world_rank = rank; |
| 39 | is_initialized = 1; |
| 40 | |
| 41 | auto addr_string = std::getenv("MASTER_ADDR"); |
| 42 | if (addr_string == NULL) { addr_string = ""; } |
| 43 | auto port_string = std::getenv("MASTER_PORT"); |
| 44 | if (port_string == NULL) { port_string = ""; } |
| 45 | |
| 46 | if (all_ranks_local_p) { shm_initialize(size, rank, addr_string, port_string); } |
| 47 | } |
| 48 | |
| 49 | void inference_all_reduce_(torch::Tensor& data, int op); |
| 50 |
nothing calls this directly
no test coverage detected