deepspeedai--deepspeed
wehub-resource-sync
107 行
3.9 KiB
C
107 行
3.9 KiB
C
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
// NOTE:
|
|
// This shared-memory implementation targets AArch64 CPUs.
|
|
// Minimum supported architecture is ARMv8-A with NEON (Advanced SIMD) support.
|
|
// Systems without NEON are not supported.
|
|
|
|
#include <arm_neon.h>
|
|
#include <stddef.h>
|
|
#include <stdint.h>
|
|
#include <cmath>
|
|
|
|
// 128 bits = 16 bytes -> fits 8 fp16/bf16 or 4 fp32 elements.
|
|
static int vector_length_in_bytes = 16;
|
|
// When widening fp16/bf16 -> fp32, 4 elements fit in one 128-bit register.
|
|
// Using 8 would require two 128-bit registers, so limit to 4.
|
|
static constexpr int full_precision_elements_in_fixed_vector = 4;
|
|
|
|
static inline float32x4_t cvt_bf16_to_fp32(const uint16x4_t input)
|
|
{
|
|
// Zero-extend 16-bit to 32-bit and shift left by 16 bits
|
|
// BF16 has the same exponent/sign bits as FP32, just missing lower mantissa bits
|
|
uint32x4_t result_32 = vshll_n_u16(input, 16);
|
|
return vreinterpretq_f32_u32(result_32);
|
|
}
|
|
|
|
static inline float32x4_t cvt_fp16_to_fp32(float16x4_t input)
|
|
{
|
|
// Converts 4 FP16 values to 4 FP32 values
|
|
return vcvt_f32_f16(input);
|
|
}
|
|
|
|
// While converting fp32 to fp16, before truncating lsb, it should be rounded to nearest even and
|
|
// Converts 4 float32 -> 4 bfloat16 with round-to-nearest-even (RNE) and NaN handling
|
|
static inline uint16x4_t cvt_fp32_to_bf16(float32x4_t src)
|
|
{
|
|
// Reinterpret float32 bits as uint32
|
|
uint32x4_t u32 = vreinterpretq_u32_f32(src);
|
|
|
|
const uint32x4_t ones = vdupq_n_u32(0x1);
|
|
const uint32x4_t vec_bias =
|
|
vdupq_n_u32(0x7FFF); // one less than half of the dropped bits range
|
|
const uint16x4_t nan_bf16 = vdup_n_u16(0xFFFF);
|
|
|
|
// RNE: lsb = (input >> 16) & 1
|
|
uint32x4_t lsb = vandq_u32(vshrq_n_u32(u32, 16), ones);
|
|
|
|
// rounding_bias = 0x7FFF + lsb, lsb can be 0 or 1.
|
|
uint32x4_t bias = vaddq_u32(vec_bias, lsb);
|
|
|
|
// input += rounding_bias
|
|
u32 = vaddq_u32(u32, bias);
|
|
|
|
// >> 16 to get bfloat16
|
|
// vshrq_n_u32 - keeps 32 bit width after shift
|
|
// vshrn_n_u32 - keeps 16 bits width after shift
|
|
uint16x4_t bf16 = vshrn_n_u32(u32, 16);
|
|
|
|
// vmvnq_u32 is bitwise NOT
|
|
// NaN mask: ~(src == src) -> 1 if NaN
|
|
// for normal num, ~(src == src) -> 0
|
|
uint32x4_t isnan = vmvnq_u32(vceqq_f32(src, src));
|
|
|
|
// Select nan_bf16 if isnan (use 16-bit mask)
|
|
uint16x4_t mask = vreinterpret_u16_u32(vget_low_u32(isnan));
|
|
return vbsl_u16(mask, nan_bf16, bf16);
|
|
}
|
|
|
|
// fp32 and fp16 are IEEE formats.
|
|
// converting fp32 to fp16 is handled by vcvt_f16_f32 internally without arbitrarily truncating the
|
|
// lsb but rounds to nearest.
|
|
static inline float16x4_t cvt_fp32_to_fp16(float32x4_t input)
|
|
{
|
|
// Converts 4 FP32 values to 4 FP16 values with rounding
|
|
return vcvt_f16_f32(input);
|
|
}
|
|
|
|
// Reduce functions down below use vectorized algorithm, the number of bytes processed each
|
|
// iteration depends on vector length. 128bit vector ==> 16 bytes. sticking to NEON 128 bit
|
|
|
|
void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers);
|
|
void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers);
|
|
void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers);
|
|
|
|
void parallel_memcpy(void* to, void* from, size_t n_bytes);
|
|
|
|
#define VLOAD_U8(X) vld1q_u8((uint8_t*)(X))
|
|
#define VLOAD_U16(X) vld1_u16((uint16_t*)(X))
|
|
#define VLOAD_F16(X) vld1_f16((float16_t*)(X))
|
|
#define VLOAD_F32(X) vld1q_f32((float32_t*)(X))
|
|
|
|
#define VSTORE_U8(A, B) vst1q_u8((uint8_t*)(A), B)
|
|
#define VSTORE_U16(A, B) vst1_u16((uint16_t*)(A), B)
|
|
#define VSTORE_F16(A, B) vst1_f16((float16_t*)(A), B) // fp16 supported from armv8.2-a+fp16
|
|
#define VSTORE_F32(A, B) vst1q_f32((float32_t*)(A), B)
|
|
|
|
#define VADD_F32(A, B) vaddq_f32(A, B)
|
|
#define VADD_F32_2VL(A, B) vaddq_f32(A, B)
|
|
|
|
#define CVT_BF16_TO_FP32(X) cvt_bf16_to_fp32(X)
|
|
#define CVT_FP16_TO_FP32(X) cvt_fp16_to_fp32(X)
|
|
#define CVT_FP32_TO_BF16(X) cvt_fp32_to_bf16(X)
|
|
#define CVT_FP32_TO_FP16(X) cvt_fp32_to_fp16(X)
|