| 1 | #include "cpu/cpu_types.hpp" |
| 2 | |
| 3 | #include <fcntl.h> |
| 4 | #include <sys/mman.h> |
| 5 | #include <sys/stat.h> |
| 6 | #include <unistd.h> |
| 7 | |
| 8 | #if defined(__aarch64__) || defined(__powerpc64__) |
| 9 | #include <atomic> |
| 10 | #endif |
| 11 | |
| 12 | namespace { |
| 13 | #define MAX_SHM_RANK_NUM 8 |
| 14 | #define PER_THREAD_SHM_BUFFER_BYTES (4 * 1024 * 1024) |
| 15 | static_assert(PER_THREAD_SHM_BUFFER_BYTES % 2 == 0); |
| 16 | #define PER_THREAD_SHM_BUFFER_OFFSET (PER_THREAD_SHM_BUFFER_BYTES >> 1) |
| 17 | #define MIN_THREAD_PROCESS_SIZE (256) |
| 18 | #define MAX_P2P_SEND_TENSOR_NUM 8 |
| 19 | |
| 20 | template <typename scalar_t> |
| 21 | struct KernelVecType { |
| 22 | using scalar_vec_t = void; |
| 23 | }; |
| 24 | |
| 25 | template <> |
| 26 | struct KernelVecType<float> { |
| 27 | using scalar_vec_t = vec_op::FP32Vec16; |
| 28 | }; |
| 29 | |
| 30 | template <> |
| 31 | struct KernelVecType<c10::BFloat16> { |
| 32 | using scalar_vec_t = vec_op::BF16Vec16; |
| 33 | }; |
| 34 | |
| 35 | template <> |
| 36 | struct KernelVecType<c10::Half> { |
| 37 | using scalar_vec_t = vec_op::FP16Vec16; |
| 38 | }; |
| 39 | |
| 40 | struct ThreadSHMContext { |
| 41 | #if defined(__aarch64__) || defined(__powerpc64__) |
| 42 | // memory model is weaker on AArch64, so we use atomic variables for |
| 43 | // consumer (load-acquire) and producer (store-release) to make sure |
| 44 | // that a stamp cannot be ready before the corresponding data is ready. |
| 45 | std::atomic<char> _curr_thread_stamp[2]; |
| 46 | std::atomic<char> _ready_thread_stamp[2]; |
| 47 | static_assert(std::atomic<char>::is_always_lock_free); |
| 48 | #else |
| 49 | volatile char _curr_thread_stamp[2]; |
| 50 | volatile char _ready_thread_stamp[2]; |
| 51 | #endif // __aarch64__ |
| 52 | int local_stamp_buffer_idx; |
| 53 | int remote_stamp_buffer_idx; |
| 54 | int thread_id; |
| 55 | int thread_num; |
| 56 | int rank; |
| 57 | int group_size; |
| 58 | size_t _spinning_count; |
| 59 | int swizzled_ranks[MAX_SHM_RANK_NUM]; |
| 60 | void* thread_shm_ptrs[MAX_SHM_RANK_NUM]; |
| 61 | ThreadSHMContext* shm_contexts[MAX_SHM_RANK_NUM]; |
| 62 | size_t _thread_buffer_mask[2]; |
| 63 | char _padding2[40]; |
| 64 | |
| 65 | ThreadSHMContext(const int thread_id, const int thread_num, const int rank, |
| 66 | const int group_size, void* thread_shm_ptr) |
| 67 | : local_stamp_buffer_idx(0), |
| 68 | remote_stamp_buffer_idx(0), |
| 69 | thread_id(thread_id), |
| 70 | thread_num(thread_num), |
| 71 | rank(rank), |
| 72 | group_size(group_size), |
| 73 | _spinning_count(0) { |
| 74 | static_assert(sizeof(ThreadSHMContext) % 64 == 0); |
| 75 | TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM); |
| 76 | TORCH_CHECK((size_t)this % 64 == 0); |
| 77 | TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0); |
| 78 | #if defined(__aarch64__) || defined(__powerpc64__) |
| 79 | _curr_thread_stamp[0].store(1, std::memory_order_relaxed); |
| 80 | _curr_thread_stamp[1].store(1, std::memory_order_relaxed); |
| 81 | _ready_thread_stamp[0].store(0, std::memory_order_relaxed); |
| 82 | _ready_thread_stamp[1].store(0, std::memory_order_relaxed); |
| 83 | #else |
| 84 | _curr_thread_stamp[0] = 1; |
| 85 | _curr_thread_stamp[1] = 1; |
| 86 | _ready_thread_stamp[0] = 0; |
| 87 | _ready_thread_stamp[1] = 0; |
| 88 | #endif // __aarch64__ |
| 89 | _thread_buffer_mask[0] = 0; |
| 90 | _thread_buffer_mask[1] = 0; |
| 91 | for (int i = 0; i < MAX_SHM_RANK_NUM; ++i) { |
| 92 | shm_contexts[i] = nullptr; |
| 93 | thread_shm_ptrs[i] = nullptr; |
| 94 | swizzled_ranks[i] = (i + rank) % group_size; |
| 95 | } |
| 96 | set_context(rank, this, thread_shm_ptr); |
| 97 | } |
| 98 | |
| 99 | void set_stamp_buffer_idx(int local, int remote) { |
| 100 | local_stamp_buffer_idx = local; |
| 101 | remote_stamp_buffer_idx = remote; |
| 102 | } |
| 103 | |
| 104 | void set_context(int rank, ThreadSHMContext* ptr, void* thread_shm_ptr) { |
| 105 | TORCH_CHECK(rank < MAX_SHM_RANK_NUM); |
| 106 | TORCH_CHECK(ptr); |
| 107 | TORCH_CHECK(thread_shm_ptr); |
| 108 | TORCH_CHECK_EQ(ptr->thread_num, thread_num); |
| 109 | TORCH_CHECK_EQ(ptr->thread_id, thread_id); |
| 110 | shm_contexts[rank] = ptr; |
| 111 | thread_shm_ptrs[rank] = thread_shm_ptr; |
| 112 | } |
| 113 | |
| 114 | template <typename T> |
| 115 | T* get_thread_shm_ptr(int rank) { |
| 116 | return reinterpret_cast<T*>( |
| 117 | reinterpret_cast<int8_t*>(thread_shm_ptrs[rank]) + |
| 118 | (PER_THREAD_SHM_BUFFER_OFFSET & |
| 119 | _thread_buffer_mask[local_stamp_buffer_idx])); |
| 120 | } |
| 121 | |
| 122 | void next_buffer() { |
| 123 | _thread_buffer_mask[local_stamp_buffer_idx] ^= 0xFFFFFFFFFFFFFFFF; |
| 124 | } |
| 125 | |
| 126 | char get_curr_stamp(int idx) const { |
| 127 | #if defined(__aarch64__) || defined(__powerpc64__) |
| 128 | return _curr_thread_stamp[idx].load(std::memory_order_acquire); |
| 129 | #else |
| 130 | return _curr_thread_stamp[idx]; |
| 131 | #endif // __aarch64__ |
| 132 | } |
| 133 | |
| 134 | char get_ready_stamp(int idx) const { |
| 135 | #if defined(__aarch64__) || defined(__powerpc64__) |
| 136 | return _ready_thread_stamp[idx].load(std::memory_order_acquire); |
| 137 | #else |
| 138 | return _ready_thread_stamp[idx]; |
| 139 | #endif // __aarch64__ |
| 140 | } |
| 141 | |
| 142 | void next_stamp() { |
| 143 | #if defined(__aarch64__) || defined(__powerpc64__) |
| 144 | _curr_thread_stamp[local_stamp_buffer_idx].fetch_add( |
| 145 | 1, std::memory_order_release); |
| 146 | #else |
| 147 | _mm_mfence(); |
| 148 | _curr_thread_stamp[local_stamp_buffer_idx] += 1; |
| 149 | #endif // __aarch64__ |
| 150 | } |
| 151 | |
| 152 | void commit_ready_stamp() { |
| 153 | #if defined(__aarch64__) || defined(__powerpc64__) |
| 154 | _ready_thread_stamp[local_stamp_buffer_idx].store( |
| 155 | _curr_thread_stamp[local_stamp_buffer_idx].load( |
| 156 | std::memory_order_relaxed), |
| 157 | std::memory_order_release); |
| 158 | #else |
| 159 | _mm_mfence(); |
| 160 | _ready_thread_stamp[local_stamp_buffer_idx] = |
| 161 | _curr_thread_stamp[local_stamp_buffer_idx]; |
| 162 | #endif // __aarch64__ |
| 163 | } |
| 164 | |
| 165 | int get_swizzled_rank(int idx) { return swizzled_ranks[idx]; } |
| 166 | |
| 167 | template <typename Cond> |
| 168 | void wait_for_all(Cond&& cond) { |
| 169 | for (int idx = 1; idx < group_size; ++idx) { |
| 170 | int rank = get_swizzled_rank(idx); |
| 171 | wait_for_one(rank, std::forward<Cond>(cond)); |
| 172 | } |
| 173 | } |
| 174 | |
| 175 | template <typename Cond> |
| 176 | void wait_for_one(int rank, Cond&& cond) { |
| 177 | ThreadSHMContext* rank_ctx = shm_contexts[rank]; |
| 178 | for (;;) { |
| 179 | char local_curr_stamp = get_curr_stamp(local_stamp_buffer_idx); |
| 180 | char local_ready_stamp = get_ready_stamp(local_stamp_buffer_idx); |
| 181 | char rank_curr_stamp = rank_ctx->get_curr_stamp(remote_stamp_buffer_idx); |
| 182 | char rank_ready_stamp = |
| 183 | rank_ctx->get_ready_stamp(remote_stamp_buffer_idx); |
| 184 | if (cond(local_curr_stamp, local_ready_stamp, rank_curr_stamp, |
| 185 | rank_ready_stamp)) { |
| 186 | break; |
| 187 | } |
| 188 | ++_spinning_count; |
| 189 | #if defined(__aarch64__) |
| 190 | __asm__ __volatile__("yield"); |
| 191 | #elif defined(__powerpc64__) |
| 192 | __asm__ __volatile__("or 1,1,1"); |
| 193 | #else |
| 194 | _mm_pause(); |
| 195 | #endif // __aarch64__ |
| 196 | } |
| 197 | } |
| 198 | |
| 199 | static bool check_no_buffer_conflict(char local_curr_stamp, |
| 200 | char local_ready_stamp, |
| 201 | char rank_curr_stamp, |
| 202 | char rank_ready_stamp) { |
| 203 | char temp = rank_curr_stamp + 2; |
| 204 | return local_curr_stamp != temp; |
| 205 | } |
| 206 | |
| 207 | static bool check_stamp_ready(char local_curr_stamp, char local_ready_stamp, |
| 208 | char rank_curr_stamp, char rank_ready_stamp) { |
| 209 | char temp = local_curr_stamp + 1; |
| 210 | return (local_curr_stamp == rank_ready_stamp) || (temp == rank_ready_stamp); |
| 211 | } |
| 212 | |
| 213 | std::string to_string() const { |
| 214 | std::stringstream ss; |
| 215 | ss << "SHMContext:"; |
| 216 | ss << "\nrank: " << rank; |
| 217 | ss << "\ngroup_size: " << group_size; |
| 218 | ss << "\nthread_num: " << thread_num; |
| 219 | ss << "\nthread_id: " << thread_id; |
| 220 | |
| 221 | ss << "\nshm_ctx_stat_loop_seq: ["; |
| 222 | for (int i = 0; i < group_size; ++i) { |
| 223 | ss << swizzled_ranks[i] << ", "; |
| 224 | } |
| 225 | ss << "]"; |
| 226 | |
| 227 | ss << "\nshm_contexts: ["; |
| 228 | for (int i = 0; i < group_size; ++i) { |
| 229 | if (shm_contexts[i]) { |
| 230 | ss << shm_contexts[i]->rank << ", "; |
| 231 | } |
| 232 | } |
| 233 | ss << "]"; |
| 234 | |
| 235 | return ss.str(); |
| 236 | } |
| 237 | }; |
| 238 | |
| 239 | class SHMManager { |
| 240 | public: |
| 241 | explicit SHMManager(const std::string& name, const int rank, |
| 242 | const int group_size, const int thread_num) |
| 243 | : _rank(rank), |
| 244 | _group_size(group_size), |
| 245 | _thread_num(thread_num), |
| 246 | _shm_names({""}), |
| 247 | _shared_mem_ptrs({nullptr}), |
| 248 | _shm_ctx(nullptr) { |
| 249 | _shm_names[rank] = get_shm_name(name, rank); |
| 250 | _shared_mem_ptrs[rank] = init_shm(rank); |
| 251 | _shm_ctx = reinterpret_cast<ThreadSHMContext*>(_shared_mem_ptrs[rank]); |
| 252 | |
| 253 | for (int i = 0; i < _thread_num; ++i) { |
| 254 | ThreadSHMContext* ctx = new (_shm_ctx + i) |
| 255 | ThreadSHMContext(i, _thread_num, _rank, _group_size, |
| 256 | compute_thread_shm_ptr(_shm_ctx, i)); |
| 257 | } |
| 258 | } |
| 259 | |
| 260 | void join(const std::string& name) { |
| 261 | for (int rank_idx = 0; rank_idx < _group_size; ++rank_idx) { |
| 262 | if (rank_idx != _rank) { |
| 263 | TORCH_CHECK(_shm_names[rank_idx].empty()); |
| 264 | TORCH_CHECK(_shared_mem_ptrs[rank_idx] == nullptr); |
| 265 | _shm_names[rank_idx] = get_shm_name(name, rank_idx); |
| 266 | _shared_mem_ptrs[rank_idx] = init_shm(rank_idx); |
| 267 | ThreadSHMContext* target_ctx = |
| 268 | reinterpret_cast<ThreadSHMContext*>(_shared_mem_ptrs[rank_idx]); |
| 269 | for (int thread_idx = 0; thread_idx < _thread_num; ++thread_idx) { |
| 270 | _shm_ctx[thread_idx].set_context( |
| 271 | rank_idx, target_ctx + thread_idx, |
| 272 | compute_thread_shm_ptr(target_ctx, thread_idx)); |
| 273 | } |
| 274 | } |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | ~SHMManager() { destroy_shm(); } |
| 279 | |
| 280 | ThreadSHMContext* get_shm_ctx() const { return _shm_ctx; } |
| 281 | |
| 282 | static std::string get_shm_name(const std::string& name, int rank) { |
| 283 | return name + "_" + std::to_string(rank); |
| 284 | } |
| 285 | |
| 286 | static int64_t create_singleton_instance(const std::string& name, |
| 287 | const int group_size, const int rank, |
| 288 | const int thread_num) { |
| 289 | std::lock_guard<std::mutex> guard(SingletonInstancesLock); |
| 290 | SingletonInstances.emplace_back( |
| 291 | std::make_unique<SHMManager>(name, rank, group_size, thread_num)); |
| 292 | return static_cast<int64_t>(SingletonInstances.size() - 1); |
| 293 | } |
| 294 | |
| 295 | static SHMManager* get_singleton_instance(int64_t handle) { |
| 296 | return SingletonInstances[handle].get(); |
| 297 | } |
| 298 | |
| 299 | protected: |
| 300 | static std::vector<std::unique_ptr<SHMManager>> SingletonInstances; |
| 301 | static std::mutex SingletonInstancesLock; |
| 302 | |
| 303 | private: |
| 304 | static size_t round_to_alignment(size_t num) { |
| 305 | return ((num + 63) / 64) * 64; |
| 306 | } |
| 307 | |
| 308 | int8_t* compute_thread_shm_ptr(ThreadSHMContext* ctx, int thread_id) { |
| 309 | int8_t* thread_shm_ptr = |
| 310 | reinterpret_cast<int8_t*>(ctx) + |
| 311 | round_to_alignment(_thread_num * sizeof(ThreadSHMContext)); |
| 312 | return thread_shm_ptr + |
| 313 | thread_id * round_to_alignment(PER_THREAD_SHM_BUFFER_BYTES); |
| 314 | } |
| 315 | |
| 316 | size_t compute_shm_size() { |
| 317 | const size_t rounded_rank_buffer_size = |
| 318 | round_to_alignment(PER_THREAD_SHM_BUFFER_BYTES) * _thread_num; |
| 319 | const size_t rounded_thread_shm_ctx_size = |
| 320 | round_to_alignment(_thread_num * sizeof(ThreadSHMContext)); |
| 321 | const size_t shm_size = |
| 322 | rounded_thread_shm_ctx_size + rounded_rank_buffer_size; |
| 323 | return shm_size; |
| 324 | } |
| 325 | |
| 326 | void* init_shm(int target_rank) { |
| 327 | const std::string& shm_name = _shm_names[target_rank]; |
| 328 | const int local_rank = _rank; |
| 329 | const size_t shm_size = compute_shm_size(); |
| 330 | |
| 331 | int fd = -1; |
| 332 | if (local_rank == target_rank) { |
| 333 | fd = shm_open(shm_name.c_str(), O_CREAT | O_EXCL | O_RDWR, |
| 334 | S_IRUSR | S_IWUSR); |
| 335 | |
| 336 | if (fd == -1) |
| 337 | TORCH_CHECK(false, "create shm in SHMManager failed. errno: " + |
| 338 | std::to_string(errno)); |
| 339 | |
| 340 | if (ftruncate(fd, shm_size) == -1) |
| 341 | TORCH_CHECK(false, "ftruncate in SHMManager failed. errno: " + |
| 342 | std::to_string(errno)); |
| 343 | } else { |
| 344 | fd = shm_open(shm_name.c_str(), O_RDWR, S_IRUSR | S_IWUSR); |
| 345 | |
| 346 | if (fd == -1) |
| 347 | TORCH_CHECK(false, "open shm in SHMManager failed. errno: " + |
| 348 | std::to_string(errno)); |
| 349 | } |
| 350 | |
| 351 | void* shm_ptr = mmap(nullptr, shm_size, PROT_READ | PROT_WRITE, |
| 352 | MAP_SHARED | MAP_POPULATE, fd, 0); |
| 353 | |
| 354 | if (shm_ptr == MAP_FAILED) { |
| 355 | TORCH_CHECK(false, |
| 356 | "mmap in SHMManager failed. errno: " + std::to_string(errno)); |
| 357 | } |
| 358 | |
| 359 | if (close(fd) != 0) { |
| 360 | TORCH_CHECK( |
| 361 | false, "close in SHMManager failed. errno: " + std::to_string(errno)); |
| 362 | } |
| 363 | |
| 364 | TORCH_CHECK((size_t)shm_ptr % 64 == 0); |
| 365 | |
| 366 | return shm_ptr; |
| 367 | } |
| 368 | |
| 369 | void destroy_shm() { |
| 370 | std::stringstream ss; |
| 371 | ss << "local rank " << _rank << ": ["; |
| 372 | for (int thread_id = 0; thread_id < _thread_num; ++thread_id) { |
| 373 | ss << _shm_ctx[thread_id]._spinning_count << ", "; |
| 374 | } |
| 375 | ss << "]\n"; |
| 376 | |
| 377 | for (int i = 0; i < MAX_SHM_RANK_NUM; ++i) { |
| 378 | if (_shared_mem_ptrs[i] != nullptr) { |
| 379 | munmap(_shared_mem_ptrs[i], compute_shm_size()); |
| 380 | } |
| 381 | |
| 382 | if (!_shm_names[i].empty()) { |
| 383 | shm_unlink(_shm_names[i].c_str()); |
| 384 | } |
| 385 | } |
| 386 | } |
| 387 | |
| 388 | int _rank; |
| 389 | int _group_size; |
| 390 | int _thread_num; |
| 391 | std::array<std::string, MAX_SHM_RANK_NUM> _shm_names; |
| 392 | std::array<void*, MAX_SHM_RANK_NUM> _shared_mem_ptrs; |
| 393 | ThreadSHMContext* _shm_ctx; |
| 394 | }; |
| 395 | |
| 396 | namespace shm_cc_ops { |
| 397 | template <typename scalar_t, typename F> |
| 398 | void shm_cc_loop(ThreadSHMContext* ctx, int64_t elem_num, F&& inner_func) { |
| 399 | int thread_num = ctx->thread_num; |
| 400 | int64_t total_bytes = elem_num * sizeof(scalar_t); |
| 401 | int64_t total_units_num = |
| 402 | (total_bytes + MIN_THREAD_PROCESS_SIZE - 1) / MIN_THREAD_PROCESS_SIZE; |
| 403 | int64_t per_thread_units_num = |
| 404 | (total_units_num + thread_num - 1) / thread_num; |
| 405 | int64_t per_unit_elem_num = MIN_THREAD_PROCESS_SIZE / sizeof(scalar_t); |
| 406 | int64_t max_per_thread_iteration_elem_num = |
| 407 | (PER_THREAD_SHM_BUFFER_BYTES >> 1) / |
| 408 | sizeof(scalar_t); // Note: double buffer |
| 409 | int64_t per_thread_elem_num = per_unit_elem_num * per_thread_units_num; |
| 410 | |
| 411 | #pragma omp parallel for schedule(static, 1) |
| 412 | for (int i = 0; i < thread_num; ++i) { |
| 413 | int64_t offset = i * per_thread_elem_num; |
| 414 | int64_t end = std::min(elem_num, offset + per_thread_elem_num); |
| 415 | int64_t curr_elem_num = |
| 416 | std::min(max_per_thread_iteration_elem_num, end - offset); |
| 417 | ThreadSHMContext* thread_ctx = ctx + i; |
| 418 | bool fast_mode = ((end - offset) <= max_per_thread_iteration_elem_num); |
| 419 | |
| 420 | while (curr_elem_num > 0) { |
| 421 | inner_func(thread_ctx, offset, curr_elem_num, fast_mode); |
| 422 | |
| 423 | thread_ctx->next_stamp(); |
| 424 | thread_ctx->next_buffer(); |
| 425 | offset += max_per_thread_iteration_elem_num; |
| 426 | curr_elem_num = std::min(max_per_thread_iteration_elem_num, end - offset); |
| 427 | } |
| 428 | } |
| 429 | } |
| 430 | |
| 431 | void reset_threads_stamp_buffer_idx(ThreadSHMContext* ctx, int local, |
| 432 | int remote) { |
| 433 | int thread_num = ctx->thread_num; |
| 434 | for (int i = 0; i < thread_num; ++i) { |
| 435 | ThreadSHMContext* thread_ctx = ctx + i; |
| 436 | thread_ctx->set_stamp_buffer_idx(local, remote); |
| 437 | } |
| 438 | } |
| 439 | }; // namespace shm_cc_ops |
| 440 | |
| 441 | namespace shm_cc_ops { |
| 442 | |
| 443 | void memcpy_from_shm(void* dst, void* src, const int64_t bytes) { |
| 444 | const int64_t aligned_bytes = ((bytes >> 6) << 6); // 64 bytes aligned |
| 445 | int64_t i = 0; |
| 446 | #pragma GCC unroll 4 |
| 447 | for (; i < aligned_bytes; i += 64) { |
| 448 | vec_op::INT8Vec64 data( |
| 449 | true, (int8_t*)src + i); // stream loading shm to avoid caching |
| 450 | data.save((int8_t*)dst + i); |
| 451 | } |
| 452 | if (aligned_bytes < bytes) { |
| 453 | vec_op::INT8Vec64 data(true, (int8_t*)src + aligned_bytes); |
| 454 | data.save((int8_t*)dst + aligned_bytes, bytes - aligned_bytes); |
| 455 | } |
| 456 | } |
| 457 | |
| 458 | void memcpy_to_shm(void* dst, void* src, const int64_t bytes) { |
| 459 | #pragma GCC unroll 4 |
| 460 | for (int64_t i = 0; i < bytes; i += 64) { |
| 461 | vec_op::INT8Vec64 data((int8_t*)src + i); |
| 462 | data.nt_save((int8_t*)dst + i); |
| 463 | } |
| 464 | } |
| 465 | |
| 466 | void memcpy(void* dst, void* src, const int64_t bytes) { |
| 467 | const int64_t aligned_bytes = ((bytes >> 6) << 6); // 64 bytes aligned |
| 468 | int64_t i = 0; |
| 469 | #pragma GCC unroll 4 |
| 470 | for (; i < aligned_bytes; i += 64) { |
| 471 | vec_op::INT8Vec64 data((int8_t*)src + i); |
| 472 | data.save((int8_t*)dst + i); |
| 473 | } |
| 474 | if (aligned_bytes < bytes) { |
| 475 | vec_op::INT8Vec64 data((int8_t*)src + aligned_bytes); |
| 476 | data.save((int8_t*)dst + aligned_bytes, bytes - aligned_bytes); |
| 477 | } |
| 478 | } |
| 479 | |
| 480 | template <typename scalar_t, int RANKS> |
| 481 | void all_reduce_sum_impl(ThreadSHMContext* ctx, scalar_t* data, |
| 482 | size_t elem_num) { |
| 483 | CPU_KERNEL_GUARD_IN(all_reduce_sum_impl) |
| 484 | using vec_t = typename KernelVecType<scalar_t>::scalar_vec_t; |
| 485 | constexpr int64_t vec_elem_num = vec_t::get_elem_num(); |
| 486 | const int worldsize = ctx->group_size; |
| 487 | |
| 488 | shm_cc_ops::shm_cc_loop<scalar_t>( |
| 489 | ctx, elem_num, |
| 490 | [&](ThreadSHMContext* thread_ctx, int64_t data_offset, |
| 491 | int64_t data_elem_num, bool fast_mode) { |
| 492 | int rank = thread_ctx->rank; |
| 493 | scalar_t* thread_shm_ptr = |
| 494 | thread_ctx->get_thread_shm_ptr<scalar_t>(rank); |
| 495 | scalar_t* thread_data_ptr = data + data_offset; |
| 496 | int64_t thread_data_elem_num = data_elem_num * sizeof(scalar_t); |
| 497 | |
| 498 | scalar_t* remote_data_ptrs[RANKS - 1]; |
| 499 | vec_op::unroll_loop<int, RANKS - 1>([&](int idx) { |
| 500 | remote_data_ptrs[idx] = thread_ctx->get_thread_shm_ptr<scalar_t>( |
| 501 | thread_ctx->get_swizzled_rank(idx + 1)); |
| 502 | }); |
| 503 | |
| 504 | if (!fast_mode) { |
| 505 | thread_ctx->wait_for_all(ThreadSHMContext::check_no_buffer_conflict); |
| 506 | } |
| 507 | |
| 508 | shm_cc_ops::memcpy_to_shm(thread_shm_ptr, thread_data_ptr, |
| 509 | thread_data_elem_num); |
| 510 | thread_ctx->commit_ready_stamp(); |
| 511 | int64_t aligned_data_elem_num = |
| 512 | (data_elem_num / vec_elem_num) * vec_elem_num; |
| 513 | int64_t i = 0; |
| 514 | thread_ctx->wait_for_all(ThreadSHMContext::check_stamp_ready); |
| 515 | #pragma GCC unroll 4 |
| 516 | for (; i < aligned_data_elem_num; i += vec_elem_num) { |
| 517 | vec_t local_data(thread_data_ptr + i); // load from cache |
| 518 | vec_op::FP32Vec16 local_data_fp32(local_data); |
| 519 | vec_op::unroll_loop<int, RANKS - 1>([&](int idx) { |
| 520 | vec_t remote_data( |
| 521 | true, remote_data_ptrs[idx] + i); // stream load from shm |
| 522 | vec_op::FP32Vec16 remote_data_fp32(remote_data); |
| 523 | local_data_fp32 = local_data_fp32 + remote_data_fp32; // sum reduce |
| 524 | }); |
| 525 | vec_t reduced_data(local_data_fp32); |
| 526 | reduced_data.save(thread_data_ptr + i); |
| 527 | } |
| 528 | |
| 529 | if (i < data_elem_num) { |
| 530 | vec_t local_data(thread_data_ptr + i); // load from cache |
| 531 | vec_op::FP32Vec16 local_data_fp32(local_data); |
| 532 | vec_op::unroll_loop<int, RANKS - 1>([&](int idx) { |
| 533 | vec_t remote_data( |
| 534 | true, remote_data_ptrs[idx] + i); // stream load from shm |
| 535 | vec_op::FP32Vec16 remote_data_fp32(remote_data); |
| 536 | local_data_fp32 = local_data_fp32 + remote_data_fp32; // sum reduce |
| 537 | }); |
| 538 | vec_t reduced_data(local_data_fp32); |
| 539 | reduced_data.save(thread_data_ptr + i, |
| 540 | data_elem_num - aligned_data_elem_num); |
| 541 | } |
| 542 | }); |
| 543 | |
| 544 | return; |
| 545 | } |
| 546 | }; // namespace shm_cc_ops |
| 547 | |
| 548 | std::vector<std::unique_ptr<SHMManager>> SHMManager::SingletonInstances = {}; |
| 549 | std::mutex SHMManager::SingletonInstancesLock = {}; |
| 550 | |
| 551 | template <typename scalar_t> |
| 552 | void shm_allreduce_sum(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num) { |
| 553 | switch (ctx->group_size) { |
| 554 | case 2: |
| 555 | shm_cc_ops::all_reduce_sum_impl<scalar_t, 2>(ctx, data, elem_num); |
| 556 | break; |
| 557 | case 3: |
| 558 | shm_cc_ops::all_reduce_sum_impl<scalar_t, 3>(ctx, data, elem_num); |
| 559 | break; |
| 560 | case 4: |
| 561 | shm_cc_ops::all_reduce_sum_impl<scalar_t, 4>(ctx, data, elem_num); |
| 562 | break; |
| 563 | case 8: |
| 564 | shm_cc_ops::all_reduce_sum_impl<scalar_t, 8>(ctx, data, elem_num); |
| 565 | break; |
| 566 | default: |
| 567 | TORCH_CHECK(false, |
| 568 | "Invalid world size: " + std::to_string(ctx->group_size)); |
| 569 | } |
| 570 | } |
| 571 | |
| 572 | template <typename scalar_t> |
| 573 | void shm_gather_impl(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num, |
| 574 | scalar_t** outputs, const int dst) { |
| 575 | CPU_KERNEL_GUARD_IN(shm_gather_impl) |
| 576 | const int worldsize = ctx->group_size; |
| 577 | TORCH_CHECK_LT(dst, worldsize); |
| 578 | shm_cc_ops::shm_cc_loop<scalar_t>( |
| 579 | ctx, elem_num, |
| 580 | [&](ThreadSHMContext* thread_ctx, int64_t data_offset, |
| 581 | int64_t data_elem_num, bool fast_mode) { |
| 582 | int rank = thread_ctx->rank; |
| 583 | scalar_t* thread_shm_ptr = |
| 584 | thread_ctx->get_thread_shm_ptr<scalar_t>(rank); |
| 585 | |
| 586 | if (!fast_mode) { |
| 587 | thread_ctx->wait_for_all(ThreadSHMContext::check_no_buffer_conflict); |
| 588 | } |
| 589 | |
| 590 | shm_cc_ops::memcpy(thread_shm_ptr, data + data_offset, |
| 591 | data_elem_num * sizeof(scalar_t)); |
| 592 | thread_ctx->commit_ready_stamp(); |
| 593 | if (rank == dst) { |
| 594 | shm_cc_ops::memcpy(outputs[rank] + data_offset, data + data_offset, |
| 595 | data_elem_num * sizeof(scalar_t)); |
| 596 | for (int i = 1; i < worldsize; ++i) { |
| 597 | int src_rank = thread_ctx->get_swizzled_rank(i); |
| 598 | scalar_t* src_ptr = |
| 599 | thread_ctx->get_thread_shm_ptr<scalar_t>(src_rank); // shm |
| 600 | scalar_t* dst_ptr = outputs[src_rank] + data_offset; |
| 601 | thread_ctx->wait_for_one(src_rank, |
| 602 | ThreadSHMContext::check_stamp_ready); |
| 603 | shm_cc_ops::memcpy(dst_ptr, src_ptr, |
| 604 | data_elem_num * sizeof(scalar_t)); |
| 605 | } |
| 606 | } |
| 607 | }); |
| 608 | |
| 609 | return; |
| 610 | } |
| 611 | |
| 612 | struct MemPiece { |
| 613 | void* ptr; |
| 614 | int64_t size; |
| 615 | |
| 616 | template <typename T> |
| 617 | T* data_ptr() { |
| 618 | return reinterpret_cast<T*>(ptr); |
| 619 | } |
| 620 | }; |
| 621 | |
| 622 | struct TensorListMeta { |
| 623 | int64_t tensor_bytes[MAX_P2P_SEND_TENSOR_NUM]; |
| 624 | torch::ScalarType tensor_types[MAX_P2P_SEND_TENSOR_NUM]; |
| 625 | int64_t tensor_num; |
| 626 | int64_t total_bytes; |
| 627 | |
| 628 | TensorListMeta() : tensor_num(0), total_bytes(0) { |
| 629 | static_assert(sizeof(TensorListMeta) % 64 == 0); |
| 630 | static_assert(sizeof(TensorListMeta) < |
| 631 | MIN_THREAD_PROCESS_SIZE); // To ensure the metadata always |
| 632 | // hold by the thread 0 |
| 633 | for (int i = 0; i < MAX_P2P_SEND_TENSOR_NUM; ++i) { |
| 634 | tensor_bytes[i] = 0; |
| 635 | tensor_ptrs[i] = nullptr; |
| 636 | tensor_types[i] = torch::ScalarType::Undefined; |
| 637 | } |
| 638 | } |
| 639 | |
| 640 | // For send and recv |
| 641 | void bind_tensor_list(std::vector<torch::Tensor>& tensor_list) { |
| 642 | TORCH_CHECK(tensor_types[0] == torch::ScalarType::Undefined, |
| 643 | "Re-bind TensorListMeta is not allowed.") |
| 644 | TORCH_CHECK_LE(tensor_list.size(), MAX_P2P_SEND_TENSOR_NUM); |
| 645 | tensor_num = tensor_list.size(); |
| 646 | int64_t bytes_sum = 0; |
| 647 | for (int i = 0; i < tensor_list.size(); ++i) { |
| 648 | torch::Tensor& t = tensor_list[i]; |
| 649 | TORCH_CHECK(t.is_contiguous()); |
| 650 | tensor_bytes[i] = t.nbytes(); |
| 651 | tensor_types[i] = t.scalar_type(); |
| 652 | tensor_ptrs[i] = t.data_ptr(); |
| 653 | bytes_sum += t.nbytes(); |
| 654 | } |
| 655 | total_bytes = bytes_sum; |
| 656 | } |
| 657 | |
| 658 | // For recv |
| 659 | std::vector<torch::Tensor> generate_tensor_list() { |
| 660 | std::vector<torch::Tensor> tensor_list; |
| 661 | tensor_list.reserve(tensor_num); |
| 662 | |
| 663 | for (int i = 0; i < tensor_num; ++i) { |
| 664 | int64_t bytes = tensor_bytes[i]; |
| 665 | auto type = tensor_types[i]; |
| 666 | int64_t elem_bytes = torch::elementSize(type); |
| 667 | |
| 668 | TORCH_CHECK_EQ(bytes % elem_bytes, 0); |
| 669 | int64_t elem_num = bytes / elem_bytes; |
| 670 | auto options = torch::TensorOptions().dtype(type).device(torch::kCPU); |
| 671 | tensor_list.emplace_back(torch::empty({elem_num}, options)); |
| 672 | } |
| 673 | return tensor_list; |
| 674 | } |
| 675 | |
| 676 | MemPiece get_data(int64_t offset) { |
| 677 | for (int i = 0; i < tensor_num; ++i) { |
| 678 | if (offset < tensor_bytes[i]) { |
| 679 | return {reinterpret_cast<int8_t*>(tensor_ptrs[i]) + offset, |
| 680 | tensor_bytes[i] - offset}; |
| 681 | } |
| 682 | offset -= tensor_bytes[i]; |
| 683 | } |
| 684 | return {nullptr, 0}; |
| 685 | } |
| 686 | |
| 687 | private: |
| 688 | void* tensor_ptrs[MAX_P2P_SEND_TENSOR_NUM]; |
| 689 | int8_t _padding[40]; |
| 690 | }; |
| 691 | |
| 692 | void shm_send_tensor_list_impl(ThreadSHMContext* ctx, int64_t dst, |
| 693 | const std::vector<torch::Tensor>& tensor_list) { |
| 694 | CPU_KERNEL_GUARD_IN(shm_send_tensor_list_impl) |
| 695 | std::vector<torch::Tensor> tensor_list_with_metadata; |
| 696 | tensor_list_with_metadata.reserve(1 + tensor_list.size()); |
| 697 | |
| 698 | auto options = torch::TensorOptions().dtype(torch::kInt8).device(torch::kCPU); |
| 699 | tensor_list_with_metadata.emplace_back( |
| 700 | torch::empty({sizeof(TensorListMeta)}, options)); |
| 701 | tensor_list_with_metadata.insert(tensor_list_with_metadata.end(), |
| 702 | tensor_list.begin(), tensor_list.end()); |
| 703 | |
| 704 | torch::Tensor& metadata_tensor = tensor_list_with_metadata[0]; |
| 705 | TORCH_CHECK_EQ(metadata_tensor.nbytes(), sizeof(TensorListMeta)); |
| 706 | |
| 707 | TensorListMeta* metadata = new (metadata_tensor.data_ptr()) TensorListMeta(); |
| 708 | metadata->bind_tensor_list(tensor_list_with_metadata); |
| 709 | |
| 710 | shm_cc_ops::reset_threads_stamp_buffer_idx(ctx, 0, 1); |
| 711 | shm_cc_ops::shm_cc_loop<int8_t>( |
| 712 | ctx, metadata->total_bytes, |
| 713 | [&](ThreadSHMContext* thread_ctx, int64_t data_offset, |
| 714 | int64_t data_elem_num, bool fast_mode) { |
| 715 | int rank = thread_ctx->rank; |
| 716 | int64_t curr_shm_offset = 0; |
| 717 | thread_ctx->wait_for_one(dst, |
| 718 | ThreadSHMContext::check_no_buffer_conflict); |
| 719 | while (curr_shm_offset < data_elem_num) { |
| 720 | MemPiece frag = metadata->get_data(data_offset + curr_shm_offset); |
| 721 | frag.size = std::min(frag.size, data_elem_num - curr_shm_offset); |
| 722 | shm_cc_ops::memcpy( |
| 723 | thread_ctx->get_thread_shm_ptr<int8_t>(rank) + curr_shm_offset, |
| 724 | frag.ptr, frag.size); |
| 725 | curr_shm_offset += frag.size; |
| 726 | } |
| 727 | thread_ctx->commit_ready_stamp(); |
| 728 | }); |
| 729 | } |
| 730 | |
| 731 | std::vector<torch::Tensor> shm_recv_tensor_list_impl(ThreadSHMContext* ctx, |
| 732 | int64_t src) { |
| 733 | CPU_KERNEL_GUARD_IN(shm_recv_tensor_list_impl) |
| 734 | auto options = torch::TensorOptions().dtype(torch::kInt8).device(torch::kCPU); |
| 735 | torch::Tensor metadata_tensor = |
| 736 | torch::empty({sizeof(TensorListMeta)}, options); |
| 737 | |
| 738 | shm_cc_ops::reset_threads_stamp_buffer_idx(ctx, 1, 0); |
| 739 | ctx->wait_for_one(src, ThreadSHMContext::check_stamp_ready); |
| 740 | shm_cc_ops::memcpy(metadata_tensor.data_ptr(), |
| 741 | ctx->get_thread_shm_ptr<void>(src), |
| 742 | sizeof(TensorListMeta)); |
| 743 | TensorListMeta* src_metadata = |
| 744 | reinterpret_cast<TensorListMeta*>(metadata_tensor.data_ptr()); |
| 745 | std::vector<torch::Tensor> tensor_list_with_metadata = |
| 746 | src_metadata->generate_tensor_list(); |
| 747 | |
| 748 | TensorListMeta metadata; |
| 749 | metadata.bind_tensor_list(tensor_list_with_metadata); |
| 750 | TORCH_CHECK_EQ(metadata.tensor_num, src_metadata->tensor_num); |
| 751 | TORCH_CHECK_EQ(metadata.total_bytes, src_metadata->total_bytes); |
| 752 | |
| 753 | shm_cc_ops::shm_cc_loop<int8_t>( |
| 754 | ctx, metadata.total_bytes, |
| 755 | [&](ThreadSHMContext* thread_ctx, int64_t data_offset, |
| 756 | int64_t data_elem_num, bool fast_mode) { |
| 757 | thread_ctx->wait_for_one(src, ThreadSHMContext::check_stamp_ready); |
| 758 | int64_t curr_shm_offset = 0; |
| 759 | while (curr_shm_offset < data_elem_num) { |
| 760 | MemPiece frag = metadata.get_data(data_offset + curr_shm_offset); |
| 761 | frag.size = std::min(frag.size, data_elem_num - curr_shm_offset); |
| 762 | shm_cc_ops::memcpy( |
| 763 | frag.ptr, |
| 764 | thread_ctx->get_thread_shm_ptr<int8_t>(src) + curr_shm_offset, |
| 765 | frag.size); |
| 766 | curr_shm_offset += frag.size; |
| 767 | } |
| 768 | }); |
| 769 | |
| 770 | std::vector<torch::Tensor> tensor_list; |
| 771 | tensor_list.reserve(metadata.tensor_num - 1); |
| 772 | tensor_list.insert(tensor_list.begin(), tensor_list_with_metadata.begin() + 1, |
| 773 | tensor_list_with_metadata.end()); |
| 774 | |
| 775 | return tensor_list; |
| 776 | } |
| 777 | } // namespace |
| 778 | |
| 779 | void shm_gather(int64_t handle, torch::Tensor& data, |
| 780 | const std::optional<std::vector<torch::Tensor>>& outputs, |
| 781 | int64_t dst) { |
| 782 | TORCH_CHECK(data.is_contiguous()) |
| 783 | VLLM_DISPATCH_FLOATING_TYPES(data.scalar_type(), "shm_gather_impl", [&] { |
| 784 | CPU_KERNEL_GUARD_IN(shm_gather_impl) |
| 785 | |
| 786 | if (outputs.has_value()) { |
| 787 | TORCH_CHECK_LE(outputs->size(), MAX_SHM_RANK_NUM); |
| 788 | scalar_t* output_ptrs[MAX_SHM_RANK_NUM] = {nullptr}; |
| 789 | for (int i = 0; i < outputs->size(); ++i) { |
| 790 | output_ptrs[i] = outputs->at(i).data_ptr<scalar_t>(); |
| 791 | } |
| 792 | shm_gather_impl(SHMManager::get_singleton_instance(handle)->get_shm_ctx(), |
| 793 | data.data_ptr<scalar_t>(), data.numel(), output_ptrs, |
| 794 | dst); |
| 795 | } else { |
| 796 | shm_gather_impl(SHMManager::get_singleton_instance(handle)->get_shm_ctx(), |
| 797 | data.data_ptr<scalar_t>(), data.numel(), (scalar_t**)(0), |
| 798 | dst); |
| 799 | } |
| 800 | |
| 801 | CPU_KERNEL_GUARD_OUT(shm_gather_impl) |
| 802 | }); |
| 803 | } |
| 804 | |
| 805 | void shm_all_gather(int64_t handle, const torch::Tensor& data, |
| 806 | torch::Tensor& output) { |
| 807 | TORCH_CHECK(data.is_contiguous()) |
| 808 | TORCH_CHECK(output.is_contiguous()) |
| 809 | |
| 810 | const int64_t input_elem_num = data.numel(); |
| 811 | const int64_t output_elem_num = output.numel(); |
| 812 | TORCH_CHECK_EQ(output_elem_num % input_elem_num, 0); |
| 813 | const int world_size = output_elem_num / input_elem_num; |
| 814 | |
| 815 | VLLM_DISPATCH_FLOATING_TYPES(data.scalar_type(), "shm_all_gather_impl", [&] { |
| 816 | CPU_KERNEL_GUARD_IN(shm_all_gather_impl) |
| 817 | auto ctx = SHMManager::get_singleton_instance(handle)->get_shm_ctx(); |
| 818 | TORCH_CHECK_EQ(ctx->group_size, world_size); |
| 819 | |
| 820 | scalar_t* output_ptrs[MAX_SHM_RANK_NUM] = {nullptr}; |
| 821 | for (int i = 0; i < world_size; ++i) { |
| 822 | output_ptrs[i] = output.data_ptr<scalar_t>() + i * input_elem_num; |
| 823 | } |
| 824 | shm_gather_impl(ctx, data.data_ptr<scalar_t>(), data.numel(), output_ptrs, |
| 825 | ctx->rank); |
| 826 | CPU_KERNEL_GUARD_OUT(shm_all_gather_impl) |
| 827 | }); |
| 828 | } |
| 829 | |
| 830 | void shm_allreduce(int64_t handle, torch::Tensor& data) { |
| 831 | TORCH_CHECK(data.is_contiguous()) |
| 832 | VLLM_DISPATCH_FLOATING_TYPES(data.scalar_type(), "shm_allreduce_sum", [&] { |
| 833 | CPU_KERNEL_GUARD_IN(shm_allreduce_sum) |
| 834 | shm_allreduce_sum(SHMManager::get_singleton_instance(handle)->get_shm_ctx(), |
| 835 | data.data_ptr<scalar_t>(), data.numel()); |
| 836 | CPU_KERNEL_GUARD_OUT(shm_allreduce_sum) |
| 837 | }); |
| 838 | } |
| 839 | |
| 840 | void shm_send_tensor_list(int64_t handle, |
| 841 | const std::vector<torch::Tensor>& tensor_list, |
| 842 | int64_t dst) { |
| 843 | CPU_KERNEL_GUARD_IN(shm_send_tensor_list) |
| 844 | shm_send_tensor_list_impl( |
| 845 | SHMManager::get_singleton_instance(handle)->get_shm_ctx(), dst, |
| 846 | tensor_list); |
| 847 | CPU_KERNEL_GUARD_OUT(shm_send_tensor_list) |
| 848 | } |
| 849 | |
| 850 | std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src) { |
| 851 | CPU_KERNEL_GUARD_IN(shm_recv_tensor_list) |
| 852 | auto tensor_list = shm_recv_tensor_list_impl( |
| 853 | SHMManager::get_singleton_instance(handle)->get_shm_ctx(), src); |
| 854 | CPU_KERNEL_GUARD_OUT(shm_recv_tensor_list) |
| 855 | return tensor_list; |
| 856 | } |
| 857 | |
| 858 | int64_t init_shm_manager(const std::string& name, const int64_t group_size, |
| 859 | const int64_t rank, const int64_t thread_num) { |
| 860 | return SHMManager::create_singleton_instance(name, group_size, rank, |
| 861 | thread_num); |
| 862 | } |
| 863 | |
| 864 | std::string join_shm_manager(int64_t handle, const std::string& name) { |
| 865 | auto shm_manager = SHMManager::get_singleton_instance(handle); |
| 866 | TORCH_CHECK(shm_manager); |
| 867 | shm_manager->join(name); |
| 868 | return shm_manager->get_shm_ctx()->to_string(); |
| 869 | } |