diff --git a/cmake/gen/neonv8_microkernels.cmake b/cmake/gen/neonv8_microkernels.cmake index 133eee7b154..895784e4384 100644 --- a/cmake/gen/neonv8_microkernels.cmake +++ b/cmake/gen/neonv8_microkernels.cmake @@ -10,6 +10,7 @@ SET(PROD_NEONV8_MICROKERNEL_SRCS + src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u32.c src/f32-qs8-vcvt/gen/f32-qs8-vcvt-neonv8-u32.c src/f32-qu8-vcvt/gen/f32-qu8-vcvt-neonv8-u32.c src/f32-vrnd/gen/f32-vrndd-neonv8-u8.c @@ -29,6 +30,9 @@ SET(PROD_NEONV8_MICROKERNEL_SRCS src/qs8-qc8w-igemm/gen/qs8-qc8w-igemm-2x8c2s4-minmax-fp32-neonv8-mlal.c) SET(NON_PROD_NEONV8_MICROKERNEL_SRCS + src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u8.c + src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u16.c + src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u24.c src/f32-qs8-vcvt/gen/f32-qs8-vcvt-neonv8-u8.c src/f32-qs8-vcvt/gen/f32-qs8-vcvt-neonv8-u16.c src/f32-qs8-vcvt/gen/f32-qs8-vcvt-neonv8-u24.c diff --git a/gen/neonv8_microkernels.bzl b/gen/neonv8_microkernels.bzl index 9ca3621f3a9..364de5bbc6b 100644 --- a/gen/neonv8_microkernels.bzl +++ b/gen/neonv8_microkernels.bzl @@ -6,6 +6,7 @@ # PROD_NEONV8_MICROKERNEL_SRCS = [ + "src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u32.c", "src/f32-qs8-vcvt/gen/f32-qs8-vcvt-neonv8-u32.c", "src/f32-qu8-vcvt/gen/f32-qu8-vcvt-neonv8-u32.c", "src/f32-vrnd/gen/f32-vrndd-neonv8-u8.c", @@ -26,6 +27,9 @@ PROD_NEONV8_MICROKERNEL_SRCS = [ ] NON_PROD_NEONV8_MICROKERNEL_SRCS = [ + "src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u8.c", + "src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u16.c", + "src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u24.c", "src/f32-qs8-vcvt/gen/f32-qs8-vcvt-neonv8-u8.c", "src/f32-qs8-vcvt/gen/f32-qs8-vcvt-neonv8-u16.c", "src/f32-qs8-vcvt/gen/f32-qs8-vcvt-neonv8-u24.c", diff --git a/scripts/generate-f32-qs8-vcvt.sh b/scripts/generate-f32-qs8-vcvt.sh index 8fc02a5f458..9a07e706580 100755 --- a/scripts/generate-f32-qs8-vcvt.sh +++ b/scripts/generate-f32-qs8-vcvt.sh @@ -5,6 +5,11 @@ # LICENSE file in the root directory of this source tree. ################################## ARM NEON ################################### +tools/xngen src/bf16-qs8-vcvt/neonv8.c.in -D BATCH_TILE=8 -o src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u8.c & +tools/xngen src/bf16-qs8-vcvt/neonv8.c.in -D BATCH_TILE=16 -o src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u16.c & +tools/xngen src/bf16-qs8-vcvt/neonv8.c.in -D BATCH_TILE=24 -o src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u24.c & +tools/xngen src/bf16-qs8-vcvt/neonv8.c.in -D BATCH_TILE=32 -o src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u32.c & + tools/xngen src/f16-qs8-vcvt/neonfp16arith.c.in -D BATCH_TILE=8 -D DATATYPE=QS8 -o src/f16-qs8-vcvt/gen/f16-qs8-vcvt-neonfp16arith-u8.c & tools/xngen src/f16-qs8-vcvt/neonfp16arith.c.in -D BATCH_TILE=16 -D DATATYPE=QS8 -o src/f16-qs8-vcvt/gen/f16-qs8-vcvt-neonfp16arith-u16.c & tools/xngen src/f16-qs8-vcvt/neonfp16arith.c.in -D BATCH_TILE=24 -D DATATYPE=QS8 -o src/f16-qs8-vcvt/gen/f16-qs8-vcvt-neonfp16arith-u24.c & diff --git a/src/bf16-qs8-vcvt/bf16-qs8-vcvt.inc b/src/bf16-qs8-vcvt/bf16-qs8-vcvt.inc index 8d89e9c9b84..3c470e3f8f9 100644 --- a/src/bf16-qs8-vcvt/bf16-qs8-vcvt.inc +++ b/src/bf16-qs8-vcvt/bf16-qs8-vcvt.inc @@ -9,6 +9,13 @@ #define XNN_DEFINED_QUANTIZED #endif +#if XNN_ARCH_ARM || XNN_ARCH_ARM64 +XNN_UKERNEL(xnn_arch_arm_neon_v8, xnn_bf16_qs8_vcvt_ukernel__neonv8_u8, 8, false, xnn_bfloat16, XNN_QUANTIZED(int8_t), struct xnn_bf16_qs8_cvt_params, xnn_init_bf16_qs8_cvt_scalar_params) +XNN_UKERNEL(xnn_arch_arm_neon_v8, xnn_bf16_qs8_vcvt_ukernel__neonv8_u16, 16, false, xnn_bfloat16, XNN_QUANTIZED(int8_t), struct xnn_bf16_qs8_cvt_params, xnn_init_bf16_qs8_cvt_scalar_params) +XNN_UKERNEL(xnn_arch_arm_neon_v8, xnn_bf16_qs8_vcvt_ukernel__neonv8_u24, 24, false, xnn_bfloat16, XNN_QUANTIZED(int8_t), struct xnn_bf16_qs8_cvt_params, xnn_init_bf16_qs8_cvt_scalar_params) +XNN_UKERNEL(xnn_arch_arm_neon_v8, xnn_bf16_qs8_vcvt_ukernel__neonv8_u32, 32, false, xnn_bfloat16, XNN_QUANTIZED(int8_t), struct xnn_bf16_qs8_cvt_params, xnn_init_bf16_qs8_cvt_scalar_params) +#endif // XNN_ARCH_ARM || XNN_ARCH_ARM64 + #if XNN_ENABLE_AVX512SKX && (XNN_ARCH_X86 || XNN_ARCH_X86_64) XNN_UKERNEL(xnn_arch_x86_avx512skx, xnn_bf16_qs8_vcvt_ukernel__avx512skx_u32, 32, false, xnn_bfloat16, XNN_QUANTIZED(int8_t), struct xnn_bf16_qs8_cvt_params, xnn_init_bf16_qs8_cvt_scalar_params) XNN_UKERNEL(xnn_arch_x86_avx512skx, xnn_bf16_qs8_vcvt_ukernel__avx512skx_u64, 64, false, xnn_bfloat16, XNN_QUANTIZED(int8_t), struct xnn_bf16_qs8_cvt_params, xnn_init_bf16_qs8_cvt_scalar_params) diff --git a/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u16.c b/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u16.c new file mode 100644 index 00000000000..074821f3142 --- /dev/null +++ b/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u16.c @@ -0,0 +1,124 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/bf16-qs8-vcvt/neonv8.c.in +// Generator: tools/xngen +// +// Copyright 2026 Google LLC +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include +#include +#include +#include +#include + +#include "src/xnnpack/common.h" +#include "src/xnnpack/intrinsics-polyfill.h" +#include "src/xnnpack/math.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vcvt.h" + + +void xnn_bf16_qs8_vcvt_ukernel__neonv8_u16( + size_t batch, + const xnn_bfloat16* input, + int8_t* output, + const struct xnn_bf16_qs8_cvt_params* restrict params) XNN_OOB_READS +{ + assert(batch != 0); + assert(batch % sizeof(xnn_bfloat16) == 0); + assert(input != NULL); + assert(output != NULL); + + const uint16_t* i = (const uint16_t*) input; + // Match the scalar path when the reciprocal underflows in BF16. + const float32x4_t vscale = vdupq_n_f32( + math_max_f32(FLT_MIN, xnn_bfloat16_to_float(params->scalar.scale))); + const int16x8_t voutput_zero_point = + vdupq_n_s16(params->scalar.output_zero_point); + for (; batch >= 16 * sizeof(xnn_bfloat16); batch -= 16 * sizeof(xnn_bfloat16)) { + const uint16x8_t vbf0 = vld1q_u16(i); i += 8; + const uint16x8_t vbf1 = vld1q_u16(i); i += 8; + + float32x4_t vx0 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf0), 16)); + float32x4_t vx1 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf0), 16)); + float32x4_t vx2 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf1), 16)); + float32x4_t vx3 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf1), 16)); + + vx0 = vmulq_f32(vx0, vscale); + vx1 = vmulq_f32(vx1, vscale); + vx2 = vmulq_f32(vx2, vscale); + vx3 = vmulq_f32(vx3, vscale); + + const int32x4_t vacc0 = vcvtnq_s32_f32(vx0); + const int32x4_t vacc1 = vcvtnq_s32_f32(vx1); + const int32x4_t vacc2 = vcvtnq_s32_f32(vx2); + const int32x4_t vacc3 = vcvtnq_s32_f32(vx3); + + int16x8_t vacc01 = vcombine_s16(vqmovn_s32(vacc0), vqmovn_s32(vacc1)); + int16x8_t vacc23 = vcombine_s16(vqmovn_s32(vacc2), vqmovn_s32(vacc3)); + + vacc01 = vqaddq_s16(vacc01, voutput_zero_point); + vacc23 = vqaddq_s16(vacc23, voutput_zero_point); + + const int8x16_t vy0 = vcombine_s8(vqmovn_s16(vacc01), vqmovn_s16(vacc23)); + + vst1q_s8(output, vy0); output += 16; + } + for (; batch >= 8 * sizeof(xnn_bfloat16); batch -= 8 * sizeof(xnn_bfloat16)) { + const uint16x8_t vbf = vld1q_u16(i); i += 8; + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + const int8x8_t vy = vqmovn_s16(vacc); + vst1_s8(output, vy); output += 8; + } + if XNN_UNLIKELY(batch != 0) { + assert(batch >= 1 * sizeof(xnn_bfloat16)); + assert(batch <= 7 * sizeof(xnn_bfloat16)); + const uint16x8_t vbf = vld1q_u16(i); + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + int8x8_t vy = vqmovn_s16(vacc); + if (batch & (4 * sizeof(xnn_bfloat16))) { + vst1_lane_u32((void*) output, vreinterpret_u32_s8(vy), 0); output += 4; + vy = vext_s8(vy, vy, 4); + } + if (batch & (2 * sizeof(xnn_bfloat16))) { + vst1_lane_u16((void*) output, vreinterpret_u16_s8(vy), 0); output += 2; + vy = vext_s8(vy, vy, 2); + } + if (batch & (1 * sizeof(xnn_bfloat16))) { + vst1_lane_s8(output, vy, 0); + } + } +} diff --git a/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u24.c b/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u24.c new file mode 100644 index 00000000000..2477455a32b --- /dev/null +++ b/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u24.c @@ -0,0 +1,135 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/bf16-qs8-vcvt/neonv8.c.in +// Generator: tools/xngen +// +// Copyright 2026 Google LLC +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include +#include +#include +#include +#include + +#include "src/xnnpack/common.h" +#include "src/xnnpack/intrinsics-polyfill.h" +#include "src/xnnpack/math.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vcvt.h" + + +void xnn_bf16_qs8_vcvt_ukernel__neonv8_u24( + size_t batch, + const xnn_bfloat16* input, + int8_t* output, + const struct xnn_bf16_qs8_cvt_params* restrict params) XNN_OOB_READS +{ + assert(batch != 0); + assert(batch % sizeof(xnn_bfloat16) == 0); + assert(input != NULL); + assert(output != NULL); + + const uint16_t* i = (const uint16_t*) input; + // Match the scalar path when the reciprocal underflows in BF16. + const float32x4_t vscale = vdupq_n_f32( + math_max_f32(FLT_MIN, xnn_bfloat16_to_float(params->scalar.scale))); + const int16x8_t voutput_zero_point = + vdupq_n_s16(params->scalar.output_zero_point); + for (; batch >= 24 * sizeof(xnn_bfloat16); batch -= 24 * sizeof(xnn_bfloat16)) { + const uint16x8_t vbf0 = vld1q_u16(i); i += 8; + const uint16x8_t vbf1 = vld1q_u16(i); i += 8; + const uint16x8_t vbf2 = vld1q_u16(i); i += 8; + + float32x4_t vx0 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf0), 16)); + float32x4_t vx1 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf0), 16)); + float32x4_t vx2 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf1), 16)); + float32x4_t vx3 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf1), 16)); + float32x4_t vx4 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf2), 16)); + float32x4_t vx5 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf2), 16)); + + vx0 = vmulq_f32(vx0, vscale); + vx1 = vmulq_f32(vx1, vscale); + vx2 = vmulq_f32(vx2, vscale); + vx3 = vmulq_f32(vx3, vscale); + vx4 = vmulq_f32(vx4, vscale); + vx5 = vmulq_f32(vx5, vscale); + + const int32x4_t vacc0 = vcvtnq_s32_f32(vx0); + const int32x4_t vacc1 = vcvtnq_s32_f32(vx1); + const int32x4_t vacc2 = vcvtnq_s32_f32(vx2); + const int32x4_t vacc3 = vcvtnq_s32_f32(vx3); + const int32x4_t vacc4 = vcvtnq_s32_f32(vx4); + const int32x4_t vacc5 = vcvtnq_s32_f32(vx5); + + int16x8_t vacc01 = vcombine_s16(vqmovn_s32(vacc0), vqmovn_s32(vacc1)); + int16x8_t vacc23 = vcombine_s16(vqmovn_s32(vacc2), vqmovn_s32(vacc3)); + int16x8_t vacc45 = vcombine_s16(vqmovn_s32(vacc4), vqmovn_s32(vacc5)); + + vacc01 = vqaddq_s16(vacc01, voutput_zero_point); + vacc23 = vqaddq_s16(vacc23, voutput_zero_point); + vacc45 = vqaddq_s16(vacc45, voutput_zero_point); + + const int8x16_t vy0 = vcombine_s8(vqmovn_s16(vacc01), vqmovn_s16(vacc23)); + const int8x8_t vy2 = vqmovn_s16(vacc45); + + vst1q_s8(output, vy0); output += 16; + vst1_s8(output, vy2); output += 8; + } + for (; batch >= 8 * sizeof(xnn_bfloat16); batch -= 8 * sizeof(xnn_bfloat16)) { + const uint16x8_t vbf = vld1q_u16(i); i += 8; + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + const int8x8_t vy = vqmovn_s16(vacc); + vst1_s8(output, vy); output += 8; + } + if XNN_UNLIKELY(batch != 0) { + assert(batch >= 1 * sizeof(xnn_bfloat16)); + assert(batch <= 7 * sizeof(xnn_bfloat16)); + const uint16x8_t vbf = vld1q_u16(i); + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + int8x8_t vy = vqmovn_s16(vacc); + if (batch & (4 * sizeof(xnn_bfloat16))) { + vst1_lane_u32((void*) output, vreinterpret_u32_s8(vy), 0); output += 4; + vy = vext_s8(vy, vy, 4); + } + if (batch & (2 * sizeof(xnn_bfloat16))) { + vst1_lane_u16((void*) output, vreinterpret_u16_s8(vy), 0); output += 2; + vy = vext_s8(vy, vy, 2); + } + if (batch & (1 * sizeof(xnn_bfloat16))) { + vst1_lane_s8(output, vy, 0); + } + } +} diff --git a/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u32.c b/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u32.c new file mode 100644 index 00000000000..e4293acaad9 --- /dev/null +++ b/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u32.c @@ -0,0 +1,144 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/bf16-qs8-vcvt/neonv8.c.in +// Generator: tools/xngen +// +// Copyright 2026 Google LLC +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include +#include +#include +#include +#include + +#include "src/xnnpack/common.h" +#include "src/xnnpack/intrinsics-polyfill.h" +#include "src/xnnpack/math.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vcvt.h" + + +void xnn_bf16_qs8_vcvt_ukernel__neonv8_u32( + size_t batch, + const xnn_bfloat16* input, + int8_t* output, + const struct xnn_bf16_qs8_cvt_params* restrict params) XNN_OOB_READS +{ + assert(batch != 0); + assert(batch % sizeof(xnn_bfloat16) == 0); + assert(input != NULL); + assert(output != NULL); + + const uint16_t* i = (const uint16_t*) input; + // Match the scalar path when the reciprocal underflows in BF16. + const float32x4_t vscale = vdupq_n_f32( + math_max_f32(FLT_MIN, xnn_bfloat16_to_float(params->scalar.scale))); + const int16x8_t voutput_zero_point = + vdupq_n_s16(params->scalar.output_zero_point); + for (; batch >= 32 * sizeof(xnn_bfloat16); batch -= 32 * sizeof(xnn_bfloat16)) { + const uint16x8_t vbf0 = vld1q_u16(i); i += 8; + const uint16x8_t vbf1 = vld1q_u16(i); i += 8; + const uint16x8_t vbf2 = vld1q_u16(i); i += 8; + const uint16x8_t vbf3 = vld1q_u16(i); i += 8; + + float32x4_t vx0 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf0), 16)); + float32x4_t vx1 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf0), 16)); + float32x4_t vx2 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf1), 16)); + float32x4_t vx3 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf1), 16)); + float32x4_t vx4 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf2), 16)); + float32x4_t vx5 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf2), 16)); + float32x4_t vx6 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf3), 16)); + float32x4_t vx7 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf3), 16)); + + vx0 = vmulq_f32(vx0, vscale); + vx1 = vmulq_f32(vx1, vscale); + vx2 = vmulq_f32(vx2, vscale); + vx3 = vmulq_f32(vx3, vscale); + vx4 = vmulq_f32(vx4, vscale); + vx5 = vmulq_f32(vx5, vscale); + vx6 = vmulq_f32(vx6, vscale); + vx7 = vmulq_f32(vx7, vscale); + + const int32x4_t vacc0 = vcvtnq_s32_f32(vx0); + const int32x4_t vacc1 = vcvtnq_s32_f32(vx1); + const int32x4_t vacc2 = vcvtnq_s32_f32(vx2); + const int32x4_t vacc3 = vcvtnq_s32_f32(vx3); + const int32x4_t vacc4 = vcvtnq_s32_f32(vx4); + const int32x4_t vacc5 = vcvtnq_s32_f32(vx5); + const int32x4_t vacc6 = vcvtnq_s32_f32(vx6); + const int32x4_t vacc7 = vcvtnq_s32_f32(vx7); + + int16x8_t vacc01 = vcombine_s16(vqmovn_s32(vacc0), vqmovn_s32(vacc1)); + int16x8_t vacc23 = vcombine_s16(vqmovn_s32(vacc2), vqmovn_s32(vacc3)); + int16x8_t vacc45 = vcombine_s16(vqmovn_s32(vacc4), vqmovn_s32(vacc5)); + int16x8_t vacc67 = vcombine_s16(vqmovn_s32(vacc6), vqmovn_s32(vacc7)); + + vacc01 = vqaddq_s16(vacc01, voutput_zero_point); + vacc23 = vqaddq_s16(vacc23, voutput_zero_point); + vacc45 = vqaddq_s16(vacc45, voutput_zero_point); + vacc67 = vqaddq_s16(vacc67, voutput_zero_point); + + const int8x16_t vy0 = vcombine_s8(vqmovn_s16(vacc01), vqmovn_s16(vacc23)); + const int8x16_t vy2 = vcombine_s8(vqmovn_s16(vacc45), vqmovn_s16(vacc67)); + + vst1q_s8(output, vy0); output += 16; + vst1q_s8(output, vy2); output += 16; + } + for (; batch >= 8 * sizeof(xnn_bfloat16); batch -= 8 * sizeof(xnn_bfloat16)) { + const uint16x8_t vbf = vld1q_u16(i); i += 8; + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + const int8x8_t vy = vqmovn_s16(vacc); + vst1_s8(output, vy); output += 8; + } + if XNN_UNLIKELY(batch != 0) { + assert(batch >= 1 * sizeof(xnn_bfloat16)); + assert(batch <= 7 * sizeof(xnn_bfloat16)); + const uint16x8_t vbf = vld1q_u16(i); + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + int8x8_t vy = vqmovn_s16(vacc); + if (batch & (4 * sizeof(xnn_bfloat16))) { + vst1_lane_u32((void*) output, vreinterpret_u32_s8(vy), 0); output += 4; + vy = vext_s8(vy, vy, 4); + } + if (batch & (2 * sizeof(xnn_bfloat16))) { + vst1_lane_u16((void*) output, vreinterpret_u16_s8(vy), 0); output += 2; + vy = vext_s8(vy, vy, 2); + } + if (batch & (1 * sizeof(xnn_bfloat16))) { + vst1_lane_s8(output, vy, 0); + } + } +} diff --git a/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u8.c b/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u8.c new file mode 100644 index 00000000000..ef8295207c3 --- /dev/null +++ b/src/bf16-qs8-vcvt/gen/bf16-qs8-vcvt-neonv8-u8.c @@ -0,0 +1,95 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/bf16-qs8-vcvt/neonv8.c.in +// Generator: tools/xngen +// +// Copyright 2026 Google LLC +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include +#include +#include +#include +#include + +#include "src/xnnpack/common.h" +#include "src/xnnpack/intrinsics-polyfill.h" +#include "src/xnnpack/math.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vcvt.h" + + +void xnn_bf16_qs8_vcvt_ukernel__neonv8_u8( + size_t batch, + const xnn_bfloat16* input, + int8_t* output, + const struct xnn_bf16_qs8_cvt_params* restrict params) XNN_OOB_READS +{ + assert(batch != 0); + assert(batch % sizeof(xnn_bfloat16) == 0); + assert(input != NULL); + assert(output != NULL); + + const uint16_t* i = (const uint16_t*) input; + // Match the scalar path when the reciprocal underflows in BF16. + const float32x4_t vscale = vdupq_n_f32( + math_max_f32(FLT_MIN, xnn_bfloat16_to_float(params->scalar.scale))); + const int16x8_t voutput_zero_point = + vdupq_n_s16(params->scalar.output_zero_point); + for (; batch >= 8 * sizeof(xnn_bfloat16); batch -= 8 * sizeof(xnn_bfloat16)) { + const uint16x8_t vbf = vld1q_u16(i); i += 8; + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + const int8x8_t vy = vqmovn_s16(vacc); + vst1_s8(output, vy); output += 8; + } + if XNN_UNLIKELY(batch != 0) { + assert(batch >= 1 * sizeof(xnn_bfloat16)); + assert(batch <= 7 * sizeof(xnn_bfloat16)); + const uint16x8_t vbf = vld1q_u16(i); + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + int8x8_t vy = vqmovn_s16(vacc); + if (batch & (4 * sizeof(xnn_bfloat16))) { + vst1_lane_u32((void*) output, vreinterpret_u32_s8(vy), 0); output += 4; + vy = vext_s8(vy, vy, 4); + } + if (batch & (2 * sizeof(xnn_bfloat16))) { + vst1_lane_u16((void*) output, vreinterpret_u16_s8(vy), 0); output += 2; + vy = vext_s8(vy, vy, 2); + } + if (batch & (1 * sizeof(xnn_bfloat16))) { + vst1_lane_s8(output, vy, 0); + } + } +} diff --git a/src/bf16-qs8-vcvt/neonv8.c.in b/src/bf16-qs8-vcvt/neonv8.c.in new file mode 100644 index 00000000000..53c92cfbd7d --- /dev/null +++ b/src/bf16-qs8-vcvt/neonv8.c.in @@ -0,0 +1,126 @@ +// Copyright 2026 Google LLC +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +$assert BATCH_TILE % 8 == 0 +$assert BATCH_TILE >= 8 +$SIMD_TILE = BATCH_TILE // 8 +#include +#include +#include +#include +#include + +#include "src/xnnpack/common.h" +#include "src/xnnpack/intrinsics-polyfill.h" +#include "src/xnnpack/math.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vcvt.h" + + +void xnn_bf16_qs8_vcvt_ukernel__neonv8_u${BATCH_TILE}( + size_t batch, + const xnn_bfloat16* input, + int8_t* output, + const struct xnn_bf16_qs8_cvt_params* restrict params) XNN_OOB_READS +{ + assert(batch != 0); + assert(batch % sizeof(xnn_bfloat16) == 0); + assert(input != NULL); + assert(output != NULL); + + const uint16_t* i = (const uint16_t*) input; + // Match the scalar path when the reciprocal underflows in BF16. + const float32x4_t vscale = vdupq_n_f32( + math_max_f32(FLT_MIN, xnn_bfloat16_to_float(params->scalar.scale))); + const int16x8_t voutput_zero_point = + vdupq_n_s16(params->scalar.output_zero_point); + $if BATCH_TILE > 8: + for (; batch >= ${BATCH_TILE} * sizeof(xnn_bfloat16); batch -= ${BATCH_TILE} * sizeof(xnn_bfloat16)) { + $for N in range(SIMD_TILE): + const uint16x8_t vbf${N} = vld1q_u16(i); i += 8; + + $for N in range(SIMD_TILE): + float32x4_t vx${2*N} = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf${N}), 16)); + float32x4_t vx${2*N+1} = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf${N}), 16)); + + $for N in range(2 * SIMD_TILE): + vx${N} = vmulq_f32(vx${N}, vscale); + + $for N in range(2 * SIMD_TILE): + const int32x4_t vacc${N} = vcvtnq_s32_f32(vx${N}); + + $for N in range(SIMD_TILE): + int16x8_t vacc${2*N}${2*N+1} = vcombine_s16(vqmovn_s32(vacc${2*N}), vqmovn_s32(vacc${2*N+1})); + + $for N in range(SIMD_TILE): + vacc${2*N}${2*N+1} = vqaddq_s16(vacc${2*N}${2*N+1}, voutput_zero_point); + + $for N in range(0, SIMD_TILE, 2): + $if N + 1 < SIMD_TILE: + const int8x16_t vy${N} = vcombine_s8(vqmovn_s16(vacc${2*N}${2*N+1}), vqmovn_s16(vacc${2*N+2}${2*N+3})); + $else: + const int8x8_t vy${N} = vqmovn_s16(vacc${2*N}${2*N+1}); + + $for N in range(0, SIMD_TILE, 2): + $if N + 1 < SIMD_TILE: + vst1q_s8(output, vy${N}); output += 16; + $else: + vst1_s8(output, vy${N}); output += 8; + } + for (; batch >= 8 * sizeof(xnn_bfloat16); batch -= 8 * sizeof(xnn_bfloat16)) { + const uint16x8_t vbf = vld1q_u16(i); i += 8; + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + const int8x8_t vy = vqmovn_s16(vacc); + vst1_s8(output, vy); output += 8; + } + if XNN_UNLIKELY(batch != 0) { + assert(batch >= 1 * sizeof(xnn_bfloat16)); + assert(batch <= 7 * sizeof(xnn_bfloat16)); + const uint16x8_t vbf = vld1q_u16(i); + + float32x4_t vx_lo = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(vbf), 16)); + float32x4_t vx_hi = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(vbf), 16)); + + vx_lo = vmulq_f32(vx_lo, vscale); + vx_hi = vmulq_f32(vx_hi, vscale); + + const int32x4_t vacc_lo = vcvtnq_s32_f32(vx_lo); + const int32x4_t vacc_hi = vcvtnq_s32_f32(vx_hi); + + int16x8_t vacc = + vcombine_s16(vqmovn_s32(vacc_lo), vqmovn_s32(vacc_hi)); + vacc = vqaddq_s16(vacc, voutput_zero_point); + + int8x8_t vy = vqmovn_s16(vacc); + if (batch & (4 * sizeof(xnn_bfloat16))) { + vst1_lane_u32((void*) output, vreinterpret_u32_s8(vy), 0); output += 4; + vy = vext_s8(vy, vy, 4); + } + if (batch & (2 * sizeof(xnn_bfloat16))) { + vst1_lane_u16((void*) output, vreinterpret_u16_s8(vy), 0); output += 2; + vy = vext_s8(vy, vy, 2); + } + if (batch & (1 * sizeof(xnn_bfloat16))) { + vst1_lane_s8(output, vy, 0); + } + } +} diff --git a/src/configs/gemm-config.c b/src/configs/gemm-config.c index 1765aa10ab4..3c4a08765b2 100644 --- a/src/configs/gemm-config.c +++ b/src/configs/gemm-config.c @@ -2582,6 +2582,7 @@ static void init_qd8_f16_qb4w_gemm_config(void) { static void init_qd8_bf16_qb4w_gemm_config(void) { // Common parameters. + qd8_bf16_qb4w_gemm_config.arch = 0; qd8_bf16_qb4w_gemm_config.log2_input_element_size = XNN_LOG2_SIZEOF_INT8_T; qd8_bf16_qb4w_gemm_config.log2_filter_element_size = XNN_LOG2_SIZEOF_UINT8_T; qd8_bf16_qb4w_gemm_config.log2_filter_element_bit_size = XNN_LOG2_BIT_SIZEOF_INT4; @@ -2595,6 +2596,8 @@ static void init_qd8_bf16_qb4w_gemm_config(void) { qd8_bf16_qb4w_gemm_config.init.bf16_qb4w = xnn_init_bf16_qb4w_minmax_scalar_params; qd8_bf16_qb4w_gemm_config.mr = 4; qd8_bf16_qb4w_gemm_config.nr = 4; + qd8_bf16_qb4w_gemm_config.log2_kr = 0; + qd8_bf16_qb4w_gemm_config.log2_sr = 0; qd8_bf16_qb4w_gemm_config.planes = 2; // The i8mm kernels convert fp32->bf16 with `vcvt_bf16_f32`, so they require both @@ -2608,6 +2611,7 @@ static void init_qd8_bf16_qb4w_gemm_config(void) { (hardware_config->arch_flags & xnn_arch_arm_neon_i8mm) && (hardware_config->arch_flags & xnn_arch_arm_neon_bf16)) { #if XNN_ARCH_ARM64 && XNN_ENABLE_ARM_I8MM && XNN_ENABLE_ARM_BF16 + qd8_bf16_qb4w_gemm_config.arch = xnn_arch_arm_neon_i8mm; qd8_bf16_qb4w_gemm_config.minmax.dqgemm[XNN_MR_TO_INDEX(1)] = XNN_INIT_HMP_DQGEMM_UKERNEL(xnn_qd8_bf16_qb4w_gemm_minmax_ukernel_1x16c8__neoni8mmbf16); qd8_bf16_qb4w_gemm_config.minmax.dqgemm[XNN_MR_TO_INDEX(4)] = XNN_INIT_HMP_DQGEMM_UKERNEL(xnn_qd8_bf16_qb4w_gemm_minmax_ukernel_4x16c8__neoni8mmbf16); qd8_bf16_qb4w_gemm_config.mr = 4; @@ -2618,6 +2622,7 @@ static void init_qd8_bf16_qb4w_gemm_config(void) { } else if (XNN_ENABLE_ARM_DOTPROD && (hardware_config->arch_flags & xnn_arch_arm_neon_dot)) { #if XNN_ENABLE_ARM_DOTPROD + qd8_bf16_qb4w_gemm_config.arch = xnn_arch_arm_neon_dot; qd8_bf16_qb4w_gemm_config.minmax.dqgemm[XNN_MR_TO_INDEX(1)] = XNN_INIT_HMP_DQGEMM_UKERNEL(xnn_qd8_bf16_qb4w_gemm_minmax_ukernel_1x16c4__neondot); qd8_bf16_qb4w_gemm_config.minmax.dqgemm[XNN_MR_TO_INDEX(4)] = XNN_INIT_HMP_DQGEMM_UKERNEL(xnn_qd8_bf16_qb4w_gemm_minmax_ukernel_4x16c4__neondot); qd8_bf16_qb4w_gemm_config.mr = 4; diff --git a/src/configs/unary-elementwise-config.c b/src/configs/unary-elementwise-config.c index d9f57e0c71a..b741f5cf5f2 100644 --- a/src/configs/unary-elementwise-config.c +++ b/src/configs/unary-elementwise-config.c @@ -1295,7 +1295,23 @@ static void init_f16_tanh_config(void) { } static void init_bf16_to_qs8_cvt_config(void) { - #if XNN_ARCH_X86 || XNN_ARCH_X86_64 + #if XNN_ARCH_ARM + const struct xnn_hardware_config* hardware_config = xnn_init_hardware_config(); + assert(hardware_config != NULL); + if (hardware_config->arch_flags & xnn_arch_arm_neon_v8) { + bf16_to_qs8_cvt_config.ukernel = XNN_INIT_UNARY_UKERNEL(xnn_bf16_qs8_vcvt_ukernel__neonv8_u32); + bf16_to_qs8_cvt_config.element_tile = 32; + bf16_to_qs8_cvt_config.init = (xnn_init_unary_uparams_fn) xnn_init_bf16_qs8_cvt_scalar_params; + } else { + bf16_to_qs8_cvt_config.ukernel = XNN_INIT_UNARY_UKERNEL(xnn_bf16_qs8_vcvt_ukernel__scalar_imagic_u4); + bf16_to_qs8_cvt_config.element_tile = 4; + bf16_to_qs8_cvt_config.init = (xnn_init_unary_uparams_fn) xnn_init_bf16_qs8_cvt_scalar_params; + } + #elif XNN_ARCH_ARM64 + bf16_to_qs8_cvt_config.ukernel = XNN_INIT_UNARY_UKERNEL(xnn_bf16_qs8_vcvt_ukernel__neonv8_u32); + bf16_to_qs8_cvt_config.element_tile = 32; + bf16_to_qs8_cvt_config.init = (xnn_init_unary_uparams_fn) xnn_init_bf16_qs8_cvt_scalar_params; + #elif XNN_ARCH_X86 || XNN_ARCH_X86_64 const struct xnn_hardware_config* hardware_config = xnn_init_hardware_config(); assert(hardware_config != NULL); (void) hardware_config; // May be unused. diff --git a/src/subgraph/rewrites/cvt_to_fp32.cc b/src/subgraph/rewrites/cvt_to_fp32.cc index d74a59e0b02..9d68a189dfe 100644 --- a/src/subgraph/rewrites/cvt_to_fp32.cc +++ b/src/subgraph/rewrites/cvt_to_fp32.cc @@ -277,16 +277,33 @@ OpAction GetOpActionFp16(const xnn_subgraph_t subgraph, const xnn_node& node) { OpAction GetOpActionBf16(const xnn_subgraph_t subgraph, const xnn_node& node) { switch (node.type) { case xnn_node_type_fully_connected: { - // Fully-connected with a blockwise int4 filter and a bf16 output has a - // native fused-bf16 GEMM. Only keep it native when the unsigned - // qdu8_bf16_qb4w path is available: on that hardware the activations are - // dynamically quantized to qduint8 by convert_gemm_to_qduint8. Elsewhere, - // fall through to the fp32 GEMM + f32->bf16 convert lowering (which always - // works). const xnn_value& filter = subgraph->values[node.inputs[1]]; if (filter.datatype == xnn_datatype_qbint4) { - if (xnn_init_qdu8_bf16_qb4w_gemm_config() != nullptr) { - return OpAction::kTransparent; + const xnn_value& input = subgraph->values[node.inputs[0]]; + const xnn_value& output = subgraph->values[node.outputs[0]]; + const bool inline_lhs_packing = node.flags & XNN_FLAG_INLINE_LHS_PACKING; + const xnn_datatype packed_input_datatype = + inline_lhs_packing ? node.packed_input_datatype : input.datatype; + switch (packed_input_datatype) { + case xnn_datatype_qduint8: + if (xnn_init_qdu8_bf16_qb4w_gemm_config() != nullptr) { + return OpAction::kTransparent; + } + break; + case xnn_datatype_qdint8: { + const xnn_gemm_config* config = + xnn_init_qd8_bf16_qb4w_gemm_config(); + if (config != nullptr && config->arch != 0 && + output.datatype == xnn_datatype_bf16 && + filter.quantization.scale_type == xnn_datatype_bf16 && + (!inline_lhs_packing || + input.datatype == xnn_datatype_bf16)) { + return OpAction::kTransparent; + } + break; + } + default: + break; } } break; diff --git a/test/bf16-qs8-vcvt.cc b/test/bf16-qs8-vcvt.cc index c5aa1c37814..0e10c974476 100644 --- a/test/bf16-qs8-vcvt.cc +++ b/test/bf16-qs8-vcvt.cc @@ -34,6 +34,14 @@ TEST(ukernel, output_zero_point) { \ TestOutputZeroPoint( \ arch_flags, batch_tile, ukernel, init_params); \ + } \ + TEST(ukernel, output_saturation) { \ + TestOutputSaturation( \ + arch_flags, batch_tile, ukernel, init_params); \ + } \ + TEST(ukernel, output_overflow) { \ + TestOutputOverflow( \ + arch_flags, batch_tile, ukernel, init_params); \ } #include "src/bf16-qs8-vcvt/bf16-qs8-vcvt.inc" #undef XNN_UKERNEL diff --git a/test/subgraph/bf16-to-fp32-fallback.cc b/test/subgraph/bf16-to-fp32-fallback.cc index 6e83eb593e0..08ffe1d208e 100644 --- a/test/subgraph/bf16-to-fp32-fallback.cc +++ b/test/subgraph/bf16-to-fp32-fallback.cc @@ -3,6 +3,7 @@ // This source code is licensed under the BSD-style license found in the // LICENSE file in the root directory of this source tree. +#include #include #include #include @@ -49,12 +50,108 @@ namespace { using XnnTensor = litert::tensor::Tensor; using litert::tensor::BuildXnnpackGraph; -using litert::tensor::OwningCpuBuffer; using litert::tensor::Type; using litert::tensor::XnnpackGraph; using testing::Eq; using xnnpack::IsIsomorphicTo; +struct InlineQd8Bf16Qb4wFullyConnected { + std::unique_ptr graph; + xnn_subgraph_t subgraph = nullptr; + uint32_t output_id = XNN_INVALID_VALUE_ID; + // BF16 1.0 scales backing the patched BF16 scale pointer below. The Tensor + // API only produces FP16 blockwise scales. + std::vector bf16_scales = {0x3F80, 0x3F80}; + + void Build(enum xnn_datatype input_datatype, + enum xnn_datatype output_datatype = xnn_datatype_bf16, + enum xnn_datatype scale_datatype = xnn_datatype_bf16) { + // The Tensor backend only emits a valid qd8_qb4w graph for FP32 inputs + // (it inserts a dynamic-quantize convert to qdint8). A BF16 input would + // reach xnn_define_fully_connected as BF16, which is rejected with + // xnn_status_invalid_parameter. Always build with FP32, then patch the + // input datatype to the requested type below. + XnnTensor input({.type = Type::kFP32, .shape = {3, 32}}); + // 2x32 int4 weights (64 nibbles, 32 bytes). {-8, -8} packs to 0x88, + // which decodes to all-zero weights with zero_point=8. + std::vector weights_data( + 32, litert::tensor::int4_t{-8, -8}); + auto quantization = std::make_shared( + std::vector{1.0f, 1.0f}, std::vector{8, 8}, + /*block_size=*/32, /*quantized_dimension=*/0); + XnnTensor weights({.type = Type::kI4, + .shape = {2, 32}, + .buffer = std::move(weights_data), + .quantization = std::move(quantization)}); + XnnTensor output = FullyConnected(input, weights); + LRT_TENSOR_ASSERT_OK_AND_ASSIGN(graph, BuildXnnpackGraph({output})); + subgraph = graph->subgraph(); + + struct xnn_node* fully_connected = nullptr; + for (size_t i = 0; i < subgraph->num_nodes; ++i) { + if (subgraph->nodes[i].type == xnn_node_type_fully_connected) { + fully_connected = &subgraph->nodes[i]; + break; + } + } + ASSERT_NE(fully_connected, nullptr); + output_id = fully_connected->outputs[0]; + + // The Tensor API always emits FP16 blockwise scales; patch to BF16 when + // requested. Only the scale_type enum matters for the fallback decision + // under test. + xnn_value& filter = subgraph->values[fully_connected->inputs[1]]; + ASSERT_EQ(filter.datatype, xnn_datatype_qbint4); + if (scale_datatype == xnn_datatype_bf16) { + filter.quantization.scale_type = xnn_datatype_bf16; + filter.quantization.blockwise_scale.bf16_scale = + reinterpret_cast(bf16_scales.data()); + } else { + filter.quantization.scale_type = scale_datatype; + } + + // FullyConnected infers the output type from the FP32 input; override when + // the test needs a different output datatype. + if (output_datatype != xnn_datatype_fp32) { + xnn_value& out = subgraph->values[output_id]; + out.datatype = output_datatype; + out.size = xnn_tensor_get_size(&out); + } + + // Match the graph state produced by packed-LHS fusion while retaining the + // original input datatype used by the inline packer. The Tensor API + // backend inserted a convert to qdint8; rewire the FC to consume the + // original input like fusion does. + bool rewired = false; + for (size_t i = 0; i < subgraph->num_nodes; ++i) { + struct xnn_node& node = subgraph->nodes[i]; + if (&node != fully_connected && node.num_outputs > 0 && + node.outputs[0] == fully_connected->inputs[0]) { + fully_connected->inputs[0] = node.inputs[0]; + rewired = true; + break; + } + } + ASSERT_TRUE(rewired); + if (input_datatype != xnn_datatype_fp32) { + xnn_value& in = subgraph->values[fully_connected->inputs[0]]; + in.datatype = input_datatype; + in.size = xnn_tensor_get_size(&in); + } + fully_connected->flags |= XNN_FLAG_INLINE_LHS_PACKING; + fully_connected->packed_input_datatype = xnn_datatype_qdint8; + } +}; + +const struct xnn_node* FindFullyConnected(const xnn_subgraph_t subgraph) { + for (size_t i = 0; i < subgraph->num_nodes; i++) { + if (subgraph->nodes[i].type == xnn_node_type_fully_connected) { + return &subgraph->nodes[i]; + } + } + return nullptr; +} + class Bf16ToFp32FallbackTest : public testing::Test { public: void SetUp() override { @@ -270,4 +367,116 @@ TEST_F(Bf16ToFp32FallbackTest, BinaryRewrite) { EXPECT_THAT(graph, IsIsomorphicTo(expected_graph)); } +class Bf16ToFp32OptimizedQd8Test : public testing::Test { + public: + void SetUp() override { +#if XNN_ARCH_ARM || XNN_ARCH_ARM64 +#if XNN_ENABLE_ARM_DOTPROD + mock_config_.arch_flags |= xnn_arch_arm_neon_dot; +#else + GTEST_SKIP(); +#endif +#else + GTEST_SKIP(); +#endif + xnn_set_hardware_config(&mock_config_); + } + + void TearDown() override { xnn_reset_hardware_config(); } + + xnn_hardware_config mock_config_{}; +}; + +TEST_F(Bf16ToFp32OptimizedQd8Test, KeepsBf16InputNative) { + InlineQd8Bf16Qb4wFullyConnected builder; + builder.Build(xnn_datatype_bf16); + + ASSERT_THAT(xnn_subgraph_fallback_from_bf16_to_fp32( + builder.subgraph, /*optimization_flags=*/0), + Eq(xnn_status_success)); + + const struct xnn_node* fully_connected = + FindFullyConnected(builder.subgraph); + ASSERT_NE(fully_connected, nullptr); + EXPECT_EQ(builder.subgraph->values[fully_connected->inputs[0]].datatype, + xnn_datatype_bf16); + EXPECT_EQ(builder.subgraph->values[fully_connected->outputs[0]].datatype, + xnn_datatype_bf16); + EXPECT_EQ(fully_connected->outputs[0], builder.output_id); +} + +TEST_F(Bf16ToFp32OptimizedQd8Test, RewritesFp32InputForSafety) { + InlineQd8Bf16Qb4wFullyConnected builder; + builder.Build(xnn_datatype_fp32); + + ASSERT_THAT(xnn_subgraph_fallback_from_bf16_to_fp32( + builder.subgraph, /*optimization_flags=*/0), + Eq(xnn_status_success)); + + const struct xnn_node* fully_connected = + FindFullyConnected(builder.subgraph); + ASSERT_NE(fully_connected, nullptr); + EXPECT_EQ(builder.subgraph->values[fully_connected->inputs[0]].datatype, + xnn_datatype_fp32); + EXPECT_EQ(builder.subgraph->values[fully_connected->outputs[0]].datatype, + xnn_datatype_fp32); + EXPECT_EQ(builder.subgraph->values[builder.output_id].datatype, + xnn_datatype_bf16); + EXPECT_NE(fully_connected->outputs[0], builder.output_id); +} + +TEST_F(Bf16ToFp32OptimizedQd8Test, RewritesFp32OutputForSafety) { + InlineQd8Bf16Qb4wFullyConnected builder; + builder.Build(xnn_datatype_bf16, xnn_datatype_fp32); + + ASSERT_THAT(xnn_subgraph_fallback_from_bf16_to_fp32( + builder.subgraph, /*optimization_flags=*/0), + Eq(xnn_status_success)); + + const struct xnn_node* fully_connected = + FindFullyConnected(builder.subgraph); + ASSERT_NE(fully_connected, nullptr); + EXPECT_EQ(builder.subgraph->values[fully_connected->inputs[0]].datatype, + xnn_datatype_fp32); + EXPECT_EQ(builder.subgraph->values[fully_connected->outputs[0]].datatype, + xnn_datatype_fp32); +} + +TEST_F(Bf16ToFp32OptimizedQd8Test, RewritesFp16WeightScalesForSafety) { + InlineQd8Bf16Qb4wFullyConnected builder; + builder.Build(xnn_datatype_bf16, xnn_datatype_bf16, xnn_datatype_fp16); + + ASSERT_THAT(xnn_subgraph_fallback_from_bf16_to_fp32( + builder.subgraph, /*optimization_flags=*/0), + Eq(xnn_status_success)); + + const struct xnn_node* fully_connected = + FindFullyConnected(builder.subgraph); + ASSERT_NE(fully_connected, nullptr); + EXPECT_EQ(builder.subgraph->values[fully_connected->inputs[0]].datatype, + xnn_datatype_fp32); + EXPECT_EQ(builder.subgraph->values[fully_connected->outputs[0]].datatype, + xnn_datatype_fp32); +} + +TEST_F(Bf16ToFp32FallbackTest, RewritesScalarQd8Fallback) { + InlineQd8Bf16Qb4wFullyConnected builder; + builder.Build(xnn_datatype_bf16); + + ASSERT_THAT(xnn_subgraph_fallback_from_bf16_to_fp32( + builder.subgraph, /*optimization_flags=*/0), + Eq(xnn_status_success)); + + const struct xnn_node* fully_connected = + FindFullyConnected(builder.subgraph); + ASSERT_NE(fully_connected, nullptr); + EXPECT_EQ(builder.subgraph->values[fully_connected->inputs[0]].datatype, + xnn_datatype_fp32); + EXPECT_EQ(builder.subgraph->values[fully_connected->outputs[0]].datatype, + xnn_datatype_fp32); + EXPECT_EQ(builder.subgraph->values[builder.output_id].datatype, + xnn_datatype_bf16); + EXPECT_NE(fully_connected->outputs[0], builder.output_id); +} + } // namespace