chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[cpu] Implement all gather/reduce for arm64 cpu (#12527)

Browse Source
This commit is contained in:
Yibo Cai
2025-12-08 19:03:04 +08:00
committed by GitHub
parent 4a62a0e3cd
commit 7871593cc8
5 changed files with 419 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/pyproject_cpu.toml
View File

@@ -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",

127
sgl-kernel/csrc/cpu/aarch64/shm.h Normal file
View File

@@ -0,0 +1,127 @@
#pragma once
#include <arm_neon.h>
#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

297
sgl-kernel/csrc/cpu/shm.cpp
View File

@@ -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 <ATen/ATen.h>
#include <errno.h>
#include <fcntl.h>
#include <immintrin.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
@@ -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) {

1
sgl-kernel/csrc/cpu/shm.h
View File

@@ -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);

280
sgl-kernel/csrc/cpu/x86_64/shm.h Normal file
View File

@@ -0,0 +1,280 @@
#pragma once
#include <immintrin.h>
// 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