From 7871593cc85307c156badb6a3322fe2e05b07f8d Mon Sep 17 00:00:00 2001 From: Yibo Cai Date: Mon, 8 Dec 2025 19:03:04 +0800 Subject: [PATCH] [cpu] Implement all gather/reduce for arm64 cpu (#12527) --- python/pyproject_cpu.toml | 4 +- sgl-kernel/csrc/cpu/aarch64/shm.h | 127 +++++++++++++ sgl-kernel/csrc/cpu/shm.cpp | 297 +----------------------------- sgl-kernel/csrc/cpu/shm.h | 1 - sgl-kernel/csrc/cpu/x86_64/shm.h | 280 ++++++++++++++++++++++++++++ 5 files changed, 419 insertions(+), 290 deletions(-) create mode 100644 sgl-kernel/csrc/cpu/aarch64/shm.h create mode 100644 sgl-kernel/csrc/cpu/x86_64/shm.h diff --git a/python/pyproject_cpu.toml b/python/pyproject_cpu.toml index a537bc360..eb0fe75c5 100644 --- a/python/pyproject_cpu.toml +++ b/python/pyproject_cpu.toml @@ -23,13 +23,13 @@ dependencies = [ "build", "compressed-tensors", "datasets", - "decord", + "decord; platform_machine == 'x86_64'", "einops", "fastapi", "gguf", "hf_transfer", "huggingface_hub", - "intel-openmp", + "intel-openmp; platform_machine == 'x86_64'", "interegular", "llguidance>=0.7.11,<0.8.0", "modelscope", diff --git a/sgl-kernel/csrc/cpu/aarch64/shm.h b/sgl-kernel/csrc/cpu/aarch64/shm.h new file mode 100644 index 000000000..06b21ce00 --- /dev/null +++ b/sgl-kernel/csrc/cpu/aarch64/shm.h @@ -0,0 +1,127 @@ +#pragma once +#include + +#define VECTOR_LENGTH_IN_BYTES 16 + +__attribute__((target("+bf16"))) inline float32x4x2_t cvt_bf16_to_fp32(const bfloat16x8_t src) { + float32x4x2_t y; + y.val[0] = vcvtq_low_f32_bf16(src); + y.val[1] = vcvtq_high_f32_bf16(src); + return y; +} + +__attribute__((target("+bf16"))) inline bfloat16x8_t cvt_fp32_to_bf16(const float32x4x2_t src) { + return vcvtq_high_bf16_f32(vcvtq_low_bf16_f32(src.val[0]), src.val[1]); +} + +__attribute__((target("+bf16"))) inline void +reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) { + const int element_size = 2; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + float32x4x2_t inout_val = cvt_bf16_to_fp32(vld1q_bf16((const bfloat16_t*)(buffers[0] + i))); + for (int j = 1; j < world_size; j++) { + const float32x4x2_t in_val = cvt_bf16_to_fp32(vld1q_bf16((const bfloat16_t*)(buffers[j] + i))); + inout_val.val[0] = vaddq_f32(inout_val.val[0], in_val.val[0]); + inout_val.val[1] = vaddq_f32(inout_val.val[1], in_val.val[1]); + } + vst1q_bf16((bfloat16_t*)(to_buffer + i), cvt_fp32_to_bf16(inout_val)); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { + val += vcvtah_f32_bf16(*(bfloat16_t*)(buffers[j] + i)); + } + *(bfloat16_t*)(to_buffer + i) = vcvth_bf16_f32(val); + remain_elements--; + i += element_size; + } +} + +inline void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) { + const int element_size = 2; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + float16x8_t inout_val = vld1q_f16((const float16_t*)(buffers[0] + i)); + for (int j = 1; j < world_size; j++) { + const float16x8_t in_val = vld1q_f16((const float16_t*)(buffers[j] + i)); + inout_val = vaddq_f16(inout_val, in_val); + } + vst1q_f16((float16_t*)(to_buffer + i), inout_val); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float16_t val = 0.0f; + for (int j = 0; j < world_size; j++) { + val = vaddh_f16(val, *(float16_t*)(buffers[j] + i)); + } + *(float16_t*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +inline void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) { + const int element_size = 4; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + float32x4_t inout_val = vld1q_f32((const float*)(buffers[0] + i)); + for (int j = 1; j < world_size; j++) { + const float32x4_t in_val = vld1q_f32((const float*)(buffers[j] + i)); + inout_val = vaddq_f32(inout_val, in_val); + } + vst1q_f32((float32_t*)(to_buffer + i), inout_val); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { + val += *(float*)(buffers[j] + i); + } + *(float*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +inline void parallel_memcpy(void* to, void* from, size_t n_bytes) { + auto aligned_bytes = n_bytes - (n_bytes % VECTOR_LENGTH_IN_BYTES); + // process aligned part +#pragma omp parallel for + for (size_t i = 0; i < aligned_bytes; i += VECTOR_LENGTH_IN_BYTES) { + const uint8x16_t val = vld1q_u8((uint8_t*)from + i); + vst1q_u8((uint8_t*)to + i, val); + } + + // process remaining part + for (size_t i = aligned_bytes; i < n_bytes; i++) { + *((uint8_t*)to + i) = *((uint8_t*)from + i); + } +} + +#undef VECTOR_LENGTH_IN_BYTES diff --git a/sgl-kernel/csrc/cpu/shm.cpp b/sgl-kernel/csrc/cpu/shm.cpp index 1bf65a9b2..89d643aed 100644 --- a/sgl-kernel/csrc/cpu/shm.cpp +++ b/sgl-kernel/csrc/cpu/shm.cpp @@ -1,9 +1,15 @@ #include "shm.h" +#if defined(__x86_64__) +#include "x86_64/shm.h" +#elif defined(__aarch64__) +#include "aarch64/shm.h" +#else +#error "unsupported architecture" +#endif #include #include #include -#include #include #include #include @@ -100,51 +106,6 @@ void wait_buffer_state_until_2(int index, enum coll_state state0, enum coll_stat } } -__m512 cvt_bf16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); -inline __m512 cvt_bf16_to_fp32(const __m256i src) { - auto y = _mm512_cvtepu16_epi32(src); - return _mm512_castsi512_ps(_mm512_bslli_epi128(y, 2)); -} - -inline __m256i cvt_fp32_to_bf16(const __m512 src) __attribute__((target("avx512bw"))); -inline __m256i cvt_fp32_to_bf16(const __m512 src) { - __m512i value = _mm512_castps_si512(src); - __m512i nan = _mm512_set1_epi32(0xffff); - auto mask_value = _mm512_cmp_ps_mask(src, src, _CMP_ORD_Q); - __m512i ones = _mm512_set1_epi32(0x1); - __m512i vec_bias = _mm512_set1_epi32(0x7fff); - // uint32_t lsb = (input >> 16) & 1; - auto t_value = _mm512_and_si512(_mm512_srli_epi32(value, 16), ones); - // uint32_t rounding_bias = 0x7fff + lsb; - t_value = _mm512_add_epi32(t_value, vec_bias); - // input += rounding_bias; - t_value = _mm512_add_epi32(t_value, value); - // input = input >> 16; - t_value = _mm512_srli_epi32(t_value, 16); - // Check NaN before converting back to bf16 - t_value = _mm512_mask_blend_epi32(mask_value, nan, t_value); - return _mm512_cvtusepi32_epi16(t_value); -} - -__m512 cvt_fp16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); -inline __m512 cvt_fp16_to_fp32(const __m256i src) { - return _mm512_cvtph_ps(src); -} - -inline __m256i cvt_fp32_to_fp16(const __m512 src) __attribute__((target("avx512bw"))); -inline __m256i cvt_fp32_to_fp16(const __m512 src) { - return _mm512_cvtps_ph(src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); -} - -void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) - __attribute__((target("avx512bw"))); - -void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) - __attribute__((target("avx512bw"))); - -void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) - __attribute__((target("avx512bw"))); - void reduce_all_buffers( int start_elements, int num_elements, @@ -154,241 +115,19 @@ void reduce_all_buffers( char** buffers) { switch (scalar_type) { case c10::ScalarType::BFloat16: - reduce_bf16_buffers(start_elements, num_elements, to_buffer, buffers); + reduce_bf16_buffers(start_elements, num_elements, to_buffer, buffers, world_size); break; case c10::ScalarType::Half: - reduce_fp16_buffers(start_elements, num_elements, to_buffer, buffers); + reduce_fp16_buffers(start_elements, num_elements, to_buffer, buffers, world_size); break; case c10::ScalarType::Float: - reduce_fp32_buffers(start_elements, num_elements, to_buffer, buffers); + reduce_fp32_buffers(start_elements, num_elements, to_buffer, buffers, world_size); break; default: assert(!"Should not get here"); } } -#define CVT_ADD_BF16(x) \ - do { \ - auto in##x##_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ - inout_val = _mm512_add_ps(inout_val, in##x##_val); \ - } while (0) - -// Reduce functions down below use vectorized algorithm, the number of bytes -// processed each iteration depends on vector length. 256bit vector ==> 32 -// bytes, 512bit vector ==> 64 bytes If you change implementation of -// reduce_bf16_buffers, etc. , check whether this number needs to be changed -#define VECTOR_LENGTH_IN_BYTES 32 - -void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) { - const int element_size = 2; - const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; - int main_elements = num_elements - (num_elements % vector_length); - int remain_elements = num_elements % vector_length; - - // process aligned part -#pragma omp parallel for - for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; - i += VECTOR_LENGTH_IN_BYTES) { - auto inout_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); - switch (world_size) { - case 16: - CVT_ADD_BF16(15); - case 15: - CVT_ADD_BF16(14); - case 14: - CVT_ADD_BF16(13); - case 13: - CVT_ADD_BF16(12); - case 12: - CVT_ADD_BF16(11); - case 11: - CVT_ADD_BF16(10); - case 10: - CVT_ADD_BF16(9); - case 9: - CVT_ADD_BF16(8); - case 8: - CVT_ADD_BF16(7); - case 7: - CVT_ADD_BF16(6); - case 6: - CVT_ADD_BF16(5); - case 5: - CVT_ADD_BF16(4); - case 4: - CVT_ADD_BF16(3); - case 3: - CVT_ADD_BF16(2); - case 2: - CVT_ADD_BF16(1); - case 1: - break; - default: - for (int j = 1; j < world_size; j++) { - auto in_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); - inout_val = _mm512_add_ps(inout_val, in_val); - } - } - _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_bf16(inout_val)); - } - - // process remaining part - int i = (start_elements + main_elements) * element_size; - while (remain_elements > 0) { - float val = 0.0f; - for (int j = 0; j < world_size; j++) { - val += *(at::BFloat16*)(buffers[j] + i); - } - *(at::BFloat16*)(to_buffer + i) = val; - remain_elements--; - i += element_size; - } -} - -#define CVT_ADD_FP16(x) \ - do { \ - auto in##x##_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ - inout_val = _mm512_add_ps(inout_val, in##x##_val); \ - } while (0) - -void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) { - const int element_size = 2; - const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; - int main_elements = num_elements - (num_elements % vector_length); - int remain_elements = num_elements % vector_length; - - // process aligned part -#pragma omp parallel for - for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; - i += VECTOR_LENGTH_IN_BYTES) { - auto inout_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); - switch (world_size) { - case 16: - CVT_ADD_FP16(15); - case 15: - CVT_ADD_FP16(14); - case 14: - CVT_ADD_FP16(13); - case 13: - CVT_ADD_FP16(12); - case 12: - CVT_ADD_FP16(11); - case 11: - CVT_ADD_FP16(10); - case 10: - CVT_ADD_FP16(9); - case 9: - CVT_ADD_FP16(8); - case 8: - CVT_ADD_FP16(7); - case 7: - CVT_ADD_FP16(6); - case 6: - CVT_ADD_FP16(5); - case 5: - CVT_ADD_FP16(4); - case 4: - CVT_ADD_FP16(3); - case 3: - CVT_ADD_FP16(2); - case 2: - CVT_ADD_FP16(1); - case 1: - break; - default: - for (int j = 1; j < world_size; j++) { - auto in_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); - inout_val = _mm512_add_ps(inout_val, in_val); - } - } - _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_fp16(inout_val)); - } - - // process remaining part - int i = (start_elements + main_elements) * element_size; - while (remain_elements > 0) { - float val = 0.0f; - for (int j = 0; j < world_size; j++) { - val += *(at::Half*)(buffers[j] + i); - } - *(at::Half*)(to_buffer + i) = val; - remain_elements--; - i += element_size; - } -} - -#define CVT_ADD_F32(x) \ - do { \ - auto in##x##_val = _mm256_loadu_ps((float*)(buffers[x] + i)); \ - inout_val = _mm256_add_ps(inout_val, in##x##_val); \ - } while (0) - -void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) { - const int element_size = 4; - const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; - int main_elements = num_elements - (num_elements % vector_length); - int remain_elements = num_elements % vector_length; - - // process aligned part -#pragma omp parallel for - for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; - i += VECTOR_LENGTH_IN_BYTES) { - auto inout_val = _mm256_loadu_ps((float*)(buffers[0] + i)); - switch (world_size) { - case 16: - CVT_ADD_F32(15); - case 15: - CVT_ADD_F32(14); - case 14: - CVT_ADD_F32(13); - case 13: - CVT_ADD_F32(12); - case 12: - CVT_ADD_F32(11); - case 11: - CVT_ADD_F32(10); - case 10: - CVT_ADD_F32(9); - case 9: - CVT_ADD_F32(8); - case 8: - CVT_ADD_F32(7); - case 7: - CVT_ADD_F32(6); - case 6: - CVT_ADD_F32(5); - case 5: - CVT_ADD_F32(4); - case 4: - CVT_ADD_F32(3); - case 3: - CVT_ADD_F32(2); - case 2: - CVT_ADD_F32(1); - case 1: - break; - default: - for (int j = 1; j < world_size; j++) { - auto in_val = _mm256_loadu_ps((float*)(buffers[j] + i)); - inout_val = _mm256_add_ps(inout_val, in_val); - } - } - _mm256_storeu_ps((float*)(to_buffer + i), inout_val); - } - - // process remaining part - int i = (start_elements + main_elements) * element_size; - while (remain_elements > 0) { - float val = 0.0f; - for (int j = 0; j < world_size; j++) { - val += *(float*)(buffers[j] + i); - } - *(float*)(to_buffer + i) = val; - remain_elements--; - i += element_size; - } -} - static bool is_initialized = false; static int world_rank; @@ -443,22 +182,6 @@ void shm_initialize(int size, int rank, const char* addr_string, const char* por } } -static void parallel_memcpy(void* to, void* from, size_t n_bytes) __attribute__((target("avx512bw"))); -static void parallel_memcpy(void* to, void* from, size_t n_bytes) { - auto aligned_bytes = n_bytes - (n_bytes % VECTOR_LENGTH_IN_BYTES); - // process aligned part -#pragma omp parallel for - for (size_t i = 0; i < aligned_bytes; i += VECTOR_LENGTH_IN_BYTES) { - auto val = _mm256_loadu_si256((__m256i*)((char*)from + i)); - _mm256_storeu_si256((__m256i*)((char*)to + i), val); - } - - // process remaining part - for (size_t i = aligned_bytes; i < n_bytes; i++) { - *((char*)to + i) = *((char*)from + i); - } -} - #define positive_mod(num, mod) ((((num) % (mod)) + (mod)) % (mod)) #define rank_mod(rank) positive_mod(rank, world_size) size_t slice_size(size_t chunk_el, int slice_idx) { diff --git a/sgl-kernel/csrc/cpu/shm.h b/sgl-kernel/csrc/cpu/shm.h index 3e903972c..590fa4cbb 100644 --- a/sgl-kernel/csrc/cpu/shm.h +++ b/sgl-kernel/csrc/cpu/shm.h @@ -4,7 +4,6 @@ #ifndef __SHM_COLLECTIVES__ #define __SHM_COLLECTIVES__ -#define VECTOR_LENGTH_IN_BYTES 32 void shm_initialize(int size, int rank, const char* addr_string, const char* port_string); 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); diff --git a/sgl-kernel/csrc/cpu/x86_64/shm.h b/sgl-kernel/csrc/cpu/x86_64/shm.h new file mode 100644 index 000000000..83d45eaf0 --- /dev/null +++ b/sgl-kernel/csrc/cpu/x86_64/shm.h @@ -0,0 +1,280 @@ +#pragma once +#include + +// Reduce functions down below use vectorized algorithm, the number of bytes +// processed each iteration depends on vector length. 256bit vector ==> 32 +// bytes, 512bit vector ==> 64 bytes If you change implementation of +// reduce_bf16_buffers, etc. , check whether this number needs to be changed +#define VECTOR_LENGTH_IN_BYTES 32 + +inline __m512 cvt_bf16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); +inline __m512 cvt_bf16_to_fp32(const __m256i src) { + auto y = _mm512_cvtepu16_epi32(src); + return _mm512_castsi512_ps(_mm512_bslli_epi128(y, 2)); +} + +inline __m256i cvt_fp32_to_bf16(const __m512 src) __attribute__((target("avx512bw"))); +inline __m256i cvt_fp32_to_bf16(const __m512 src) { + __m512i value = _mm512_castps_si512(src); + __m512i nan = _mm512_set1_epi32(0xffff); + auto mask_value = _mm512_cmp_ps_mask(src, src, _CMP_ORD_Q); + __m512i ones = _mm512_set1_epi32(0x1); + __m512i vec_bias = _mm512_set1_epi32(0x7fff); + // uint32_t lsb = (input >> 16) & 1; + auto t_value = _mm512_and_si512(_mm512_srli_epi32(value, 16), ones); + // uint32_t rounding_bias = 0x7fff + lsb; + t_value = _mm512_add_epi32(t_value, vec_bias); + // input += rounding_bias; + t_value = _mm512_add_epi32(t_value, value); + // input = input >> 16; + t_value = _mm512_srli_epi32(t_value, 16); + // Check NaN before converting back to bf16 + t_value = _mm512_mask_blend_epi32(mask_value, nan, t_value); + return _mm512_cvtusepi32_epi16(t_value); +} + +inline __m512 cvt_fp16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); +inline __m512 cvt_fp16_to_fp32(const __m256i src) { + return _mm512_cvtph_ps(src); +} + +inline __m256i cvt_fp32_to_fp16(const __m512 src) __attribute__((target("avx512bw"))); +inline __m256i cvt_fp32_to_fp16(const __m512 src) { + return _mm512_cvtps_ph(src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); +} + +#define CVT_ADD_BF16(x) \ + do { \ + auto in##x##_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ + inout_val = _mm512_add_ps(inout_val, in##x##_val); \ + } while (0) + +__attribute__((target("avx512bw"))) inline void +reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) { + const int element_size = 2; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + auto inout_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); + switch (world_size) { + case 16: + CVT_ADD_BF16(15); + case 15: + CVT_ADD_BF16(14); + case 14: + CVT_ADD_BF16(13); + case 13: + CVT_ADD_BF16(12); + case 12: + CVT_ADD_BF16(11); + case 11: + CVT_ADD_BF16(10); + case 10: + CVT_ADD_BF16(9); + case 9: + CVT_ADD_BF16(8); + case 8: + CVT_ADD_BF16(7); + case 7: + CVT_ADD_BF16(6); + case 6: + CVT_ADD_BF16(5); + case 5: + CVT_ADD_BF16(4); + case 4: + CVT_ADD_BF16(3); + case 3: + CVT_ADD_BF16(2); + case 2: + CVT_ADD_BF16(1); + case 1: + break; + default: + for (int j = 1; j < world_size; j++) { + auto in_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); + inout_val = _mm512_add_ps(inout_val, in_val); + } + } + _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_bf16(inout_val)); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { + val += *(at::BFloat16*)(buffers[j] + i); + } + *(at::BFloat16*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +#define CVT_ADD_FP16(x) \ + do { \ + auto in##x##_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ + inout_val = _mm512_add_ps(inout_val, in##x##_val); \ + } while (0) + +__attribute__((target("avx512bw"))) inline void +reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) { + const int element_size = 2; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + auto inout_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); + switch (world_size) { + case 16: + CVT_ADD_FP16(15); + case 15: + CVT_ADD_FP16(14); + case 14: + CVT_ADD_FP16(13); + case 13: + CVT_ADD_FP16(12); + case 12: + CVT_ADD_FP16(11); + case 11: + CVT_ADD_FP16(10); + case 10: + CVT_ADD_FP16(9); + case 9: + CVT_ADD_FP16(8); + case 8: + CVT_ADD_FP16(7); + case 7: + CVT_ADD_FP16(6); + case 6: + CVT_ADD_FP16(5); + case 5: + CVT_ADD_FP16(4); + case 4: + CVT_ADD_FP16(3); + case 3: + CVT_ADD_FP16(2); + case 2: + CVT_ADD_FP16(1); + case 1: + break; + default: + for (int j = 1; j < world_size; j++) { + auto in_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); + inout_val = _mm512_add_ps(inout_val, in_val); + } + } + _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_fp16(inout_val)); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { + val += *(at::Half*)(buffers[j] + i); + } + *(at::Half*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +#define CVT_ADD_F32(x) \ + do { \ + auto in##x##_val = _mm256_loadu_ps((float*)(buffers[x] + i)); \ + inout_val = _mm256_add_ps(inout_val, in##x##_val); \ + } while (0) + +__attribute__((target("avx512bw"))) inline void +reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) { + const int element_size = 4; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + auto inout_val = _mm256_loadu_ps((float*)(buffers[0] + i)); + switch (world_size) { + case 16: + CVT_ADD_F32(15); + case 15: + CVT_ADD_F32(14); + case 14: + CVT_ADD_F32(13); + case 13: + CVT_ADD_F32(12); + case 12: + CVT_ADD_F32(11); + case 11: + CVT_ADD_F32(10); + case 10: + CVT_ADD_F32(9); + case 9: + CVT_ADD_F32(8); + case 8: + CVT_ADD_F32(7); + case 7: + CVT_ADD_F32(6); + case 6: + CVT_ADD_F32(5); + case 5: + CVT_ADD_F32(4); + case 4: + CVT_ADD_F32(3); + case 3: + CVT_ADD_F32(2); + case 2: + CVT_ADD_F32(1); + case 1: + break; + default: + for (int j = 1; j < world_size; j++) { + auto in_val = _mm256_loadu_ps((float*)(buffers[j] + i)); + inout_val = _mm256_add_ps(inout_val, in_val); + } + } + _mm256_storeu_ps((float*)(to_buffer + i), inout_val); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { + val += *(float*)(buffers[j] + i); + } + *(float*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +__attribute__((target("avx512bw"))) inline void parallel_memcpy(void* to, void* from, size_t n_bytes) { + auto aligned_bytes = n_bytes - (n_bytes % VECTOR_LENGTH_IN_BYTES); + // process aligned part +#pragma omp parallel for + for (size_t i = 0; i < aligned_bytes; i += VECTOR_LENGTH_IN_BYTES) { + auto val = _mm256_loadu_si256((__m256i*)((char*)from + i)); + _mm256_storeu_si256((__m256i*)((char*)to + i), val); + } + + // process remaining part + for (size_t i = aligned_bytes; i < n_bytes; i++) { + *((char*)to + i) = *((char*)from + i); + } +} + +#undef VECTOR_LENGTH_IN_BYTES