// Tencent is pleased to support the open source community by making ncnn available. // // Copyright (C) 2022 THL A29 Limited, a Tencent company. All rights reserved. // // Licensed under the BSD 3-Clause License (the "License"); you may not use this file except // in compliance with the License. You may obtain a copy of the License at // // https://opensource.org/licenses/BSD-3-Clause // // Unless required by applicable law or agreed to in writing, software distributed // under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR // CONDITIONS OF ANY KIND, either express or implied. See the License for the // specific language governing permissions and limitations under the License. #if NCNN_RUNTIME_CPU && NCNN_ARM84BF16 && __aarch64__ && !__ARM_FEATURE_BF16_VECTOR_ARITHMETIC void cast_fp32_to_bf16_neon_bf16(const Mat& bottom_blob, Mat& top_blob, const Option& opt); void cast_bf16_to_fp32_neon_bf16(const Mat& bottom_blob, Mat& top_blob, const Option& opt); #endif static void cast_fp32_to_bf16_neon(const Mat& bottom_blob, Mat& top_blob, const Option& opt) { #if NCNN_RUNTIME_CPU && NCNN_ARM84BF16 && __aarch64__ && !__ARM_FEATURE_BF16_VECTOR_ARITHMETIC if (ncnn::cpu_support_arm_bf16()) { cast_fp32_to_bf16_neon_bf16(bottom_blob, top_blob, opt); return; } #endif const int w = bottom_blob.w; const int h = bottom_blob.h; const int d = bottom_blob.d; const int channels = bottom_blob.c; const int elempack = bottom_blob.elempack; const int size = w * h * d * elempack; #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < channels; q++) { const float* ptr = bottom_blob.channel(q); #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC __bf16* outptr = top_blob.channel(q); #else unsigned short* outptr = top_blob.channel(q); #endif int i = 0; #if __ARM_NEON for (; i + 15 < size; i += 16) { #if __aarch64__ #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC asm volatile( "prfm pldl1keep, [%0, #512] \n" "ld1 {v0.4s, v1.4s, v2.4s, v3.4s}, [%0], #64 \n" "bfcvtn v0.4h, v0.4s \n" "bfcvtn2 v0.8h, v1.4s \n" "bfcvtn v1.4h, v2.4s \n" "bfcvtn2 v1.8h, v3.4s \n" "st1 {v0.8h, v1.8h}, [%1], #32 \n" : "=r"(ptr), // %0 "=r"(outptr) // %1 : "0"(ptr), "1"(outptr) : "memory", "v0", "v1"); #else // __ARM_FEATURE_BF16_VECTOR_ARITHMETIC asm volatile( "prfm pldl1keep, [%0, #512] \n" "ld1 {v0.4s, v1.4s, v2.4s, v3.4s}, [%0], #64 \n" "shrn v0.4h, v0.4s, #16 \n" "shrn v1.4h, v1.4s, #16 \n" "shrn v2.4h, v2.4s, #16 \n" "shrn v3.4h, v3.4s, #16 \n" "st1 {v0.4h, v1.4h, v2.4h, v3.4h}, [%1], #32 \n" : "=r"(ptr), // %0 "=r"(outptr) // %1 : "0"(ptr), "1"(outptr) : "memory", "v0", "v1", "v2", "v3"); #endif // __ARM_FEATURE_BF16_VECTOR_ARITHMETIC #else // __aarch64__ asm volatile( "pld [%0, #512] \n" "vldm %0!, {d0-d7} \n" "vshrn.u32 d0, q0, #16 \n" "vshrn.u32 d1, q1, #16 \n" "vshrn.u32 d2, q2, #16 \n" "vshrn.u32 d3, q3, #16 \n" "vst1.u16 {d0-d3}, [%1 :128]! \n" : "=r"(ptr), // %0 "=r"(outptr) // %1 : "0"(ptr), "1"(outptr) : "memory", "q0", "q1", "q2", "q3"); #endif // __aarch64__ } for (; i + 7 < size; i += 8) { float32x4_t _p0_fp32 = vld1q_f32(ptr); float32x4_t _p1_fp32 = vld1q_f32(ptr + 4); #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC bfloat16x4_t _p0_bf16 = vcvt_bf16_f32(_p0_fp32); bfloat16x4_t _p1_bf16 = vcvt_bf16_f32(_p1_fp32); bfloat16x8_t _p_bf16 = vcombine_bf16(_p0_bf16, _p1_bf16); vst1q_bf16(outptr, _p_bf16); #else uint16x4_t _p0_bf16 = float2bfloat(_p0_fp32); uint16x4_t _p1_bf16 = float2bfloat(_p1_fp32); uint16x8_t _p_bf16 = vcombine_u16(_p0_bf16, _p1_bf16); vst1q_u16(outptr, _p_bf16); #endif ptr += 8; outptr += 8; } for (; i + 3 < size; i += 4) { float32x4_t _p_fp32 = vld1q_f32(ptr); #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC bfloat16x4_t _p_bf16 = vcvt_bf16_f32(_p_fp32); vst1_bf16(outptr, _p_bf16); #else uint16x4_t _p_bf16 = float2bfloat(_p_fp32); vst1_u16(outptr, _p_bf16); #endif ptr += 4; outptr += 4; } #endif for (; i < size; i++) { #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC asm volatile( "ldr s0, [%0], #4 \n" "bfcvt h0, s0 \n" "str h0, [%1], #2 \n" : "=r"(ptr), // %0 "=r"(outptr) // %1 : "0"(ptr), "1"(outptr) : "memory", "s0"); // because intrinsic cause ndk clang crash // *outptr++ = vcvth_bf16_f32(*ptr++); #else *outptr++ = float32_to_bfloat16(*ptr++); #endif } } } static void cast_bf16_to_fp32_neon(const Mat& bottom_blob, Mat& top_blob, const Option& opt) { #if NCNN_RUNTIME_CPU && NCNN_ARM84BF16 && __aarch64__ && !__ARM_FEATURE_BF16_VECTOR_ARITHMETIC if (ncnn::cpu_support_arm_bf16()) { cast_bf16_to_fp32_neon_bf16(bottom_blob, top_blob, opt); return; } #endif const int w = bottom_blob.w; const int h = bottom_blob.h; const int d = bottom_blob.d; const int channels = bottom_blob.c; const int elempack = bottom_blob.elempack; const int size = w * h * d * elempack; #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < channels; q++) { #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC const __bf16* ptr = bottom_blob.channel(q); #else const unsigned short* ptr = bottom_blob.channel(q); #endif float* outptr = top_blob.channel(q); int i = 0; #if __ARM_NEON for (; i + 15 < size; i += 16) { #if __aarch64__ asm volatile( "prfm pldl1keep, [%0, #256] \n" "ld1 {v0.4h, v1.4h, v2.4h, v3.4h}, [%0], #32 \n" "shll v0.4s, v0.4h, #16 \n" "shll v1.4s, v1.4h, #16 \n" "shll v2.4s, v2.4h, #16 \n" "shll v3.4s, v3.4h, #16 \n" "st1 {v0.4s, v1.4s, v2.4s, v3.4s}, [%1], #64 \n" : "=r"(ptr), // %0 "=r"(outptr) // %1 : "0"(ptr), "1"(outptr) : "memory", "v0", "v1", "v2", "v3"); #else // __aarch64__ asm volatile( "pld [%0, #256] \n" "vld1.u16 {d4-d7}, [%0 :128]! \n" "vshll.u16 q0, d4, #16 \n" "vshll.u16 q1, d5, #16 \n" "vshll.u16 q2, d6, #16 \n" "vshll.u16 q3, d7, #16 \n" "vstm %1!, {d0-d7} \n" : "=r"(ptr), // %0 "=r"(outptr) // %1 : "0"(ptr), "1"(outptr) : "memory", "q0", "q1", "q2", "q3"); #endif // __aarch64__ } for (; i + 7 < size; i += 8) { #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC bfloat16x8_t _p_bf16 = vld1q_bf16(ptr); float32x4_t _p0_fp32 = vcvt_f32_bf16(vget_low_bf16(_p_bf16)); float32x4_t _p1_fp32 = vcvt_f32_bf16(vget_high_bf16(_p_bf16)); #else uint16x8_t _p_bf16 = vld1q_u16(ptr); float32x4_t _p0_fp32 = bfloat2float(vget_low_u16(_p_bf16)); float32x4_t _p1_fp32 = bfloat2float(vget_high_u16(_p_bf16)); #endif vst1q_f32(outptr, _p0_fp32); vst1q_f32(outptr + 4, _p1_fp32); ptr += 8; outptr += 8; } for (; i + 3 < size; i += 4) { #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC bfloat16x4_t _p_bf16 = vld1_bf16(ptr); float32x4_t _p_fp32 = vcvt_f32_bf16(_p_bf16); #else uint16x4_t _p_bf16 = vld1_u16(ptr); float32x4_t _p_fp32 = bfloat2float(_p_bf16); #endif vst1q_f32(outptr, _p_fp32); ptr += 4; outptr += 4; } #endif for (; i < size; i++) { #if __ARM_FEATURE_BF16_VECTOR_ARITHMETIC *outptr++ = vcvtah_f32_bf16(*ptr++); #else *outptr++ = bfloat16_to_float32(*ptr++); #endif } } }