Repository: microsoft/DeepSpeed
File: csrc/cpu/comm/shm.cpp (C++)

Code before:
        case c10::ScalarType::Half:
            reduce_fp16_buffers(start_elements, num_elements, to_buffer, buffers);
            break;
        case c10::ScalarType::Float:
            reduce_fp32_buffers(start_elements, num_elements, to_buffer, buffers);
            break;
        default: assert(!"Should not get here");
    }
}

#define CVT_ADD_BF16(x)                                                 \
    do {                                                                \
        auto in##x##_val = CVT_BF16_TO_FP32(VLOAD_U16(buffers[x] + i)); \
        inout_val = VADD_F32_2VL(inout_val, in##x##_val);               \
    } while (0)

void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers)
{
    const int element_size = 2;
#if TARGET_RISCV
    size_t vl = __riscv_vsetvl_e16m1(num_elements);
    vector_length_in_bytes = vl * element_size;
#elif TARGET_ARM
    const int vl = full_precision_elements_in_fixed_vector;
    vector_length_in_bytes = vl * element_size;
#else
    const int vl = vector_length_in_bytes / element_size;
#endif
    int main_elements = num_elements - (num_elements % vl);
    int remain_elements = num_elements % vl;

    // 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(VLOAD_U16(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);

Review comment (suggestion):
Suggest to add a comment after `#else` to mention it is for x86_64.

```
#else // x86_64
```

It looks not very obvious that the `#else` part is for x86_64 instruction set.

Code after:
        case c10::ScalarType::Half:
            reduce_fp16_buffers(start_elements, num_elements, to_buffer, buffers);
            break;
        case c10::ScalarType::Float:
            reduce_fp32_buffers(start_elements, num_elements, to_buffer, buffers);
            break;
        default: assert(!"Should not get here");
    }
}

#define CVT_ADD_BF16(x)                                                 \
    do {                                                                \
        auto in##x##_val = CVT_BF16_TO_FP32(VLOAD_U16(buffers[x] + i)); \
        inout_val = VADD_F32_2VL(inout_val, in##x##_val);               \
    } while (0)

void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers)
{
    const int element_size = 2;
#if TARGET_RISCV
    size_t vl = __riscv_vsetvl_e16m1(num_elements);
    vector_length_in_bytes = vl * element_size;
#elif TARGET_ARM
    const int vl = full_precision_elements_in_fixed_vector;
    vector_length_in_bytes = vl * element_size;
#else  // x86_64
    const int vl = vector_length_in_bytes / element_size;
#endif
    int main_elements = num_elements - (num_elements % vl);
    int remain_elements = num_elements % vl;

    // 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(VLOAD_U16(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);