F32-VSQRT for Hexagon - Implement xnn_rsqrt_f32 for simd/f32-hvx.h use Carmack approximation - Add portable simd.c.in template for hvx 34.65x faster than scalar_approx_rsqrt on Samsung S23 xnn_f32_vrsqrt_ukernel__hvx_u32/N:32640/real_time 24841 ns xnn_f32_vrsqrt_ukernel__hvx_u64/N:32640/real_time 13128 ns xnn_f32_vrsqrt_ukernel__hvx_u128/N:32640/real_time 8642 ns xnn_f32_vrsqrt_ukernel__scalar_rsqrt_u1/N:32640/real_time 2293708 ns xnn_f32_vrsqrt_ukernel__scalar_rsqrt_u2/N:32640/real_time 2310264 ns xnn_f32_vrsqrt_ukernel__scalar_rsqrt_u4/N:32640/real_time 2200541 ns xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u1/N:32640/real_time 790721 ns xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u2/N:32640/real_time 444373 ns xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u4/N:32640/real_time 299529 ns xnn_f32_vsqrt_ukernel__scalar_sqrt_u1/N:32640/real_time 1123451 ns xnn_f32_vsqrt_ukernel__scalar_sqrt_u2/N:32640/real_time 1217916 ns xnn_f32_vsqrt_ukernel__scalar_sqrt_u4/N:32640/real_time 1206876 ns PiperOrigin-RevId: 758132404
diff --git a/cmake/gen/avx2_microkernels.cmake b/cmake/gen/avx2_microkernels.cmake index 3eb82a5..5f4798a 100644 --- a/cmake/gen/avx2_microkernels.cmake +++ b/cmake/gen/avx2_microkernels.cmake
@@ -27,6 +27,7 @@ src/f32-raddstoreexpminusmax/gen/f32-raddstoreexpminusmax-avx2-rr2-p5-u32-acc2.c src/f32-velu/gen/f32-velu-avx2-rr1-lut4-p4-perm-u32.c src/f32-vlog/gen/f32-vlog-avx2-rational-3-3-div.c + src/f32-vrsqrt/gen/f32-vrsqrt-avx2-approx-rsqrt.c src/f32-vsigmoid/gen/f32-vsigmoid-avx2-rr1-p5-div-u16.c src/qd8-f16-qb4w-gemm/gen/qd8-f16-qb4w-gemm-1x8c8-minmax-avx2.c src/qd8-f16-qb4w-gemm/gen/qd8-f16-qb4w-gemm-3x8c8-minmax-avx2.c
diff --git a/cmake/gen/avx_microkernels.cmake b/cmake/gen/avx_microkernels.cmake index a5c29cb..7979c6c 100644 --- a/cmake/gen/avx_microkernels.cmake +++ b/cmake/gen/avx_microkernels.cmake
@@ -73,6 +73,7 @@ src/f32-vrnd/gen/f32-vrndne-avx-u16.c src/f32-vrnd/gen/f32-vrndu-avx-u16.c src/f32-vrnd/gen/f32-vrndz-avx-u16.c + src/f32-vrsqrt/gen/f32-vrsqrt-avx-approx-rsqrt.c src/f32-vrsqrt/gen/f32-vrsqrt-avx-rsqrt-u16.c src/f32-vsigmoid/gen/f32-vsigmoid-avx-rr2-p5-nr2-u16.c src/f32-vsin/gen/f32-vsin-avx-rational-5-4-div.c
diff --git a/cmake/gen/hvx_microkernels.cmake b/cmake/gen/hvx_microkernels.cmake index 7e91b5a..f895f6a 100644 --- a/cmake/gen/hvx_microkernels.cmake +++ b/cmake/gen/hvx_microkernels.cmake
@@ -222,6 +222,9 @@ src/f32-vrnd/gen/f32-vrndu-hvx-u64.c src/f32-vrnd/gen/f32-vrndz-hvx-u32.c src/f32-vrnd/gen/f32-vrndz-hvx-u64.c + src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u32.c + src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u64.c + src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u128.c src/f32-vtanh/gen/f32-vtanh-hvx-rational-9-8-nr.c src/qs8-packw/gen/qs8-packw-x96c4-gemm-gio-hvx.c src/qs8-packw/gen/qs8-packw-x96c4-gemm-goi-hvx.c
diff --git a/cmake/gen/scalar_microkernels.cmake b/cmake/gen/scalar_microkernels.cmake index f580547..8f296db 100644 --- a/cmake/gen/scalar_microkernels.cmake +++ b/cmake/gen/scalar_microkernels.cmake
@@ -121,8 +121,8 @@ src/f32-vrnd/gen/f32-vrndne-scalar-libm-u1.c src/f32-vrnd/gen/f32-vrndu-scalar-libm-u1.c src/f32-vrnd/gen/f32-vrndz-scalar-libm-u1.c + src/f32-vrsqrt/gen/f32-vrsqrt-scalar-approx-rsqrt.c src/f32-vrsqrt/gen/f32-vrsqrt-scalar-rsqrt-u1.c - src/f32-vrsqrt/gen/f32-vrsqrt-scalar-rsqrt-u4.c src/f32-vsigmoid/gen/f32-vsigmoid-scalar-rr2-lut64-p2-div-u2.c src/f32-vsin/gen/f32-vsin-scalar-rational-5-4-div.c src/f32-vsqrt/gen/f32-vsqrt-scalar-sqrt-u1.c @@ -498,6 +498,7 @@ src/f32-vrnd/gen/f32-vrndz-scalar-libm-u2.c src/f32-vrnd/gen/f32-vrndz-scalar-libm-u4.c src/f32-vrsqrt/gen/f32-vrsqrt-scalar-rsqrt-u2.c + src/f32-vrsqrt/gen/f32-vrsqrt-scalar-rsqrt-u4.c src/f32-vsigmoid/gen/f32-vsigmoid-scalar-rr2-lut64-p2-div-u1.c src/f32-vsigmoid/gen/f32-vsigmoid-scalar-rr2-lut64-p2-div-u4.c src/f32-vsigmoid/gen/f32-vsigmoid-scalar-rr2-lut2048-p1-div-u1.c
diff --git a/cmake/gen/sse2_microkernels.cmake b/cmake/gen/sse2_microkernels.cmake index 95896f5..ea85906 100644 --- a/cmake/gen/sse2_microkernels.cmake +++ b/cmake/gen/sse2_microkernels.cmake
@@ -42,6 +42,7 @@ src/f32-vrnd/gen/f32-vrndne-sse2-u8.c src/f32-vrnd/gen/f32-vrndu-sse2-u8.c src/f32-vrnd/gen/f32-vrndz-sse2-u8.c + src/f32-vrsqrt/gen/f32-vrsqrt-sse2-approx-rsqrt.c src/f32-vsigmoid/gen/f32-vsigmoid-sse2-rr2-lut64-p2-div-u8.c src/f32-vsin/gen/f32-vsin-sse2-rational-5-4-div.c src/f32-vtanh/gen/f32-vtanh-sse2-rational-9-8-div.c
diff --git a/cmake/gen/sse_microkernels.cmake b/cmake/gen/sse_microkernels.cmake index 17b7017..158dae6 100644 --- a/cmake/gen/sse_microkernels.cmake +++ b/cmake/gen/sse_microkernels.cmake
@@ -166,6 +166,11 @@ src/f32-rminmax/gen/f32-rminmax-sse-u8-acc2.c src/f32-rminmax/gen/f32-rminmax-sse-u12-acc3.c src/f32-rminmax/gen/f32-rminmax-sse-u16-acc2.c + src/f32-rsum/gen/f32-rsum-sse-u4.c + src/f32-rsum/gen/f32-rsum-sse-u8-acc2.c + src/f32-rsum/gen/f32-rsum-sse-u12-acc3.c + src/f32-rsum/gen/f32-rsum-sse-u16-acc2.c + src/f32-rsum/gen/f32-rsum-sse-u16-acc4.c src/f32-spmm/gen/f32-spmm-4x1-minmax-sse.c src/f32-spmm/gen/f32-spmm-8x1-minmax-sse.c src/f32-spmm/gen/f32-spmm-16x1-minmax-sse.c
diff --git a/gen/avx2_microkernels.bzl b/gen/avx2_microkernels.bzl index ef8ea72..7785af9 100644 --- a/gen/avx2_microkernels.bzl +++ b/gen/avx2_microkernels.bzl
@@ -23,6 +23,7 @@ "src/f32-raddstoreexpminusmax/gen/f32-raddstoreexpminusmax-avx2-rr2-p5-u32-acc2.c", "src/f32-velu/gen/f32-velu-avx2-rr1-lut4-p4-perm-u32.c", "src/f32-vlog/gen/f32-vlog-avx2-rational-3-3-div.c", + "src/f32-vrsqrt/gen/f32-vrsqrt-avx2-approx-rsqrt.c", "src/f32-vsigmoid/gen/f32-vsigmoid-avx2-rr1-p5-div-u16.c", "src/qd8-f16-qb4w-gemm/gen/qd8-f16-qb4w-gemm-1x8c8-minmax-avx2.c", "src/qd8-f16-qb4w-gemm/gen/qd8-f16-qb4w-gemm-3x8c8-minmax-avx2.c",
diff --git a/gen/avx_microkernels.bzl b/gen/avx_microkernels.bzl index 08fd93d..cc36c59 100644 --- a/gen/avx_microkernels.bzl +++ b/gen/avx_microkernels.bzl
@@ -69,6 +69,7 @@ "src/f32-vrnd/gen/f32-vrndne-avx-u16.c", "src/f32-vrnd/gen/f32-vrndu-avx-u16.c", "src/f32-vrnd/gen/f32-vrndz-avx-u16.c", + "src/f32-vrsqrt/gen/f32-vrsqrt-avx-approx-rsqrt.c", "src/f32-vrsqrt/gen/f32-vrsqrt-avx-rsqrt-u16.c", "src/f32-vsigmoid/gen/f32-vsigmoid-avx-rr2-p5-nr2-u16.c", "src/f32-vsin/gen/f32-vsin-avx-rational-5-4-div.c",
diff --git a/gen/hvx_microkernels.bzl b/gen/hvx_microkernels.bzl index 7d7f8d1..1502004 100644 --- a/gen/hvx_microkernels.bzl +++ b/gen/hvx_microkernels.bzl
@@ -219,6 +219,9 @@ "src/f32-vrnd/gen/f32-vrndu-hvx-u64.c", "src/f32-vrnd/gen/f32-vrndz-hvx-u32.c", "src/f32-vrnd/gen/f32-vrndz-hvx-u64.c", + "src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u32.c", + "src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u64.c", + "src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u128.c", "src/f32-vtanh/gen/f32-vtanh-hvx-rational-9-8-nr.c", "src/qs8-packw/gen/qs8-packw-x96c4-gemm-gio-hvx.c", "src/qs8-packw/gen/qs8-packw-x96c4-gemm-goi-hvx.c",
diff --git a/gen/scalar_microkernels.bzl b/gen/scalar_microkernels.bzl index 6f30c8a..40515c6 100644 --- a/gen/scalar_microkernels.bzl +++ b/gen/scalar_microkernels.bzl
@@ -117,8 +117,8 @@ "src/f32-vrnd/gen/f32-vrndne-scalar-libm-u1.c", "src/f32-vrnd/gen/f32-vrndu-scalar-libm-u1.c", "src/f32-vrnd/gen/f32-vrndz-scalar-libm-u1.c", + "src/f32-vrsqrt/gen/f32-vrsqrt-scalar-approx-rsqrt.c", "src/f32-vrsqrt/gen/f32-vrsqrt-scalar-rsqrt-u1.c", - "src/f32-vrsqrt/gen/f32-vrsqrt-scalar-rsqrt-u4.c", "src/f32-vsigmoid/gen/f32-vsigmoid-scalar-rr2-lut64-p2-div-u2.c", "src/f32-vsin/gen/f32-vsin-scalar-rational-5-4-div.c", "src/f32-vsqrt/gen/f32-vsqrt-scalar-sqrt-u1.c", @@ -495,6 +495,7 @@ "src/f32-vrnd/gen/f32-vrndz-scalar-libm-u2.c", "src/f32-vrnd/gen/f32-vrndz-scalar-libm-u4.c", "src/f32-vrsqrt/gen/f32-vrsqrt-scalar-rsqrt-u2.c", + "src/f32-vrsqrt/gen/f32-vrsqrt-scalar-rsqrt-u4.c", "src/f32-vsigmoid/gen/f32-vsigmoid-scalar-rr2-lut64-p2-div-u1.c", "src/f32-vsigmoid/gen/f32-vsigmoid-scalar-rr2-lut64-p2-div-u4.c", "src/f32-vsigmoid/gen/f32-vsigmoid-scalar-rr2-lut2048-p1-div-u1.c",
diff --git a/gen/sse2_microkernels.bzl b/gen/sse2_microkernels.bzl index 8088064..39136ed 100644 --- a/gen/sse2_microkernels.bzl +++ b/gen/sse2_microkernels.bzl
@@ -38,6 +38,7 @@ "src/f32-vrnd/gen/f32-vrndne-sse2-u8.c", "src/f32-vrnd/gen/f32-vrndu-sse2-u8.c", "src/f32-vrnd/gen/f32-vrndz-sse2-u8.c", + "src/f32-vrsqrt/gen/f32-vrsqrt-sse2-approx-rsqrt.c", "src/f32-vsigmoid/gen/f32-vsigmoid-sse2-rr2-lut64-p2-div-u8.c", "src/f32-vsin/gen/f32-vsin-sse2-rational-5-4-div.c", "src/f32-vtanh/gen/f32-vtanh-sse2-rational-9-8-div.c",
diff --git a/gen/sse_microkernels.bzl b/gen/sse_microkernels.bzl index 4579cb2..3f83f34 100644 --- a/gen/sse_microkernels.bzl +++ b/gen/sse_microkernels.bzl
@@ -163,6 +163,11 @@ "src/f32-rminmax/gen/f32-rminmax-sse-u8-acc2.c", "src/f32-rminmax/gen/f32-rminmax-sse-u12-acc3.c", "src/f32-rminmax/gen/f32-rminmax-sse-u16-acc2.c", + "src/f32-rsum/gen/f32-rsum-sse-u4.c", + "src/f32-rsum/gen/f32-rsum-sse-u8-acc2.c", + "src/f32-rsum/gen/f32-rsum-sse-u12-acc3.c", + "src/f32-rsum/gen/f32-rsum-sse-u16-acc2.c", + "src/f32-rsum/gen/f32-rsum-sse-u16-acc4.c", "src/f32-spmm/gen/f32-spmm-4x1-minmax-sse.c", "src/f32-spmm/gen/f32-spmm-8x1-minmax-sse.c", "src/f32-spmm/gen/f32-spmm-16x1-minmax-sse.c",
diff --git a/scripts/generate-f32-vrsqrt.sh b/scripts/generate-f32-vrsqrt.sh index 564b0bc..102de55 100755 --- a/scripts/generate-f32-vrsqrt.sh +++ b/scripts/generate-f32-vrsqrt.sh
@@ -38,4 +38,15 @@ tools/xngen src/f32-vrsqrt/avx512f-rsqrt.c.in -D BATCH_TILE=32 -o src/f32-vrsqrt/gen/f32-vrsqrt-avx512f-rsqrt-u32.c & tools/xngen src/f32-vrsqrt/avx512f-rsqrt.c.in -D BATCH_TILE=64 -o src/f32-vrsqrt/gen/f32-vrsqrt-avx512f-rsqrt-u64.c & +################################## Hexagon HVX ################################# +tools/xngen src/f32-vrsqrt/simd.c.in -D BATCH_TILE=32 -D ARCH=hvx -o src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u32.c & +tools/xngen src/f32-vrsqrt/simd.c.in -D BATCH_TILE=64 -D ARCH=hvx -o src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u64.c & +tools/xngen src/f32-vrsqrt/simd.c.in -D BATCH_TILE=128 -D ARCH=hvx -o src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u128.c & + +##################################### SIMD ##################################### +tools/xngen src/f32-vrsqrt/simd-approx-rsqrt.c.in -D ARCH=scalar -D BATCH_TILES=1,2,4,8 -o src/f32-vrsqrt/gen/f32-vrsqrt-scalar-approx-rsqrt.c & +tools/xngen src/f32-vrsqrt/simd-approx-rsqrt.c.in -D ARCH=sse2 -D BATCH_TILES=4,8,16 -o src/f32-vrsqrt/gen/f32-vrsqrt-sse2-approx-rsqrt.c & +tools/xngen src/f32-vrsqrt/simd-approx-rsqrt.c.in -D ARCH=avx -D BATCH_TILES=8,16,32 -o src/f32-vrsqrt/gen/f32-vrsqrt-avx-approx-rsqrt.c & +tools/xngen src/f32-vrsqrt/simd-approx-rsqrt.c.in -D ARCH=avx2 -D BATCH_TILES=8,16,32 -o src/f32-vrsqrt/gen/f32-vrsqrt-avx2-approx-rsqrt.c & + wait
diff --git a/src/configs/unary-elementwise-config.c b/src/configs/unary-elementwise-config.c index 1af028c..ecd31b4 100644 --- a/src/configs/unary-elementwise-config.c +++ b/src/configs/unary-elementwise-config.c
@@ -1379,28 +1379,26 @@ f32_rsqrt_config.ukernel = (xnn_vunary_ukernel_fn) xnn_f32_vrsqrt_ukernel__scalar_rsqrt_u1; } #elif XNN_ARCH_X86 || XNN_ARCH_X86_64 + // TODO: b/404616456 - These kernels are faster, but produce numerically inconsistent results + // xnn_f32_vrsqrt_ukernel__avx512f_rsqrt_u32; + // xnn_f32_vrsqrt_ukernel__fma3_rsqrt_u16; + // xnn_f32_vrsqrt_ukernel__avx_rsqrt_u16; + // xnn_f32_vrsqrt_ukernel__sse_rsqrt_u8; const struct xnn_hardware_config* hardware_config = xnn_init_hardware_config(); assert(hardware_config != NULL); - #if XNN_ENABLE_AVX512F - // TODO: b/404616456 - This kernel should only be enabled if we allow - // numerically inconsistent results across microarchitectures. - if (false && !XNN_PLATFORM_MOBILE && hardware_config->use_x86_avx512f) { - f32_rsqrt_config.ukernel = (xnn_vunary_ukernel_fn) xnn_f32_vrsqrt_ukernel__avx512f_rsqrt_u32; - } else - #endif - if (!XNN_PLATFORM_MOBILE && hardware_config->use_x86_fma3) { - f32_rsqrt_config.ukernel = (xnn_vunary_ukernel_fn) xnn_f32_vrsqrt_ukernel__fma3_rsqrt_u16; + if (hardware_config->use_x86_avx2) { + f32_rsqrt_config.ukernel = (xnn_vunary_ukernel_fn) xnn_f32_vrsqrt_ukernel__avx2_approx_rsqrt_u16; } else if (hardware_config->use_x86_avx) { - f32_rsqrt_config.ukernel = (xnn_vunary_ukernel_fn) xnn_f32_vrsqrt_ukernel__avx_rsqrt_u16; + f32_rsqrt_config.ukernel = (xnn_vunary_ukernel_fn) xnn_f32_vrsqrt_ukernel__avx_approx_rsqrt_u16; } else { - f32_rsqrt_config.ukernel = (xnn_vunary_ukernel_fn) xnn_f32_vrsqrt_ukernel__sse_rsqrt_u8; + f32_rsqrt_config.ukernel = (xnn_vunary_ukernel_fn) xnn_f32_vrsqrt_ukernel__sse2_approx_rsqrt_u8; } #elif XNN_ARCH_RISCV && XNN_ENABLE_RISCV_VECTOR const struct xnn_hardware_config* hardware_config = xnn_init_hardware_config(); f32_rsqrt_config.ukernel = (xnn_vunary_ukernel_fn) xnn_f32_vrsqrt_ukernel__rvv_rsqrt_u4v; #else f32_rsqrt_config.ukernel = - (xnn_vunary_ukernel_fn)xnn_f32_vrsqrt_ukernel__scalar_rsqrt_u4; + (xnn_vunary_ukernel_fn)xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u4; #endif }
diff --git a/src/f32-vrsqrt/f32-vrsqrt.h b/src/f32-vrsqrt/f32-vrsqrt.h index 412cf7f..0b7646f 100644 --- a/src/f32-vrsqrt/f32-vrsqrt.h +++ b/src/f32-vrsqrt/f32-vrsqrt.h
@@ -19,6 +19,10 @@ XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__scalar_rsqrt_u1, 1, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__scalar_rsqrt_u2, 2, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__scalar_rsqrt_u4, 4, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u1, 1, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u2, 2, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u4, 4, false, float, struct xnn_f32_default_params, NULL) + #if XNN_ARCH_RISCV && XNN_ENABLE_RISCV_VECTOR XNN_UKERNEL_WITH_PARAMS(xnn_arch_riscv_vector, xnn_f32_vrsqrt_ukernel__rvv_rsqrt_u1v, 1, true, float, struct xnn_f32_default_params, NULL) @@ -36,12 +40,21 @@ XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__sse_rsqrt_u4, 4, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__sse_rsqrt_u8, 8, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__sse_rsqrt_u16, 16, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__sse2_approx_rsqrt_u4, 4, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__sse2_approx_rsqrt_u8, 8, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__sse2_approx_rsqrt_u16, 16, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx, xnn_f32_vrsqrt_ukernel__avx_rsqrt_u8, 8, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx, xnn_f32_vrsqrt_ukernel__avx_rsqrt_u16, 16, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx, xnn_f32_vrsqrt_ukernel__avx_rsqrt_u32, 32, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx, xnn_f32_vrsqrt_ukernel__avx_approx_rsqrt_u8, 8, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx, xnn_f32_vrsqrt_ukernel__avx_approx_rsqrt_u16, 16, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx, xnn_f32_vrsqrt_ukernel__avx_approx_rsqrt_u32, 32, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_fma3, xnn_f32_vrsqrt_ukernel__fma3_rsqrt_u8, 8, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_fma3, xnn_f32_vrsqrt_ukernel__fma3_rsqrt_u16, 16, false, float, struct xnn_f32_default_params, NULL) XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_fma3, xnn_f32_vrsqrt_ukernel__fma3_rsqrt_u32, 32, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx2, xnn_f32_vrsqrt_ukernel__avx2_approx_rsqrt_u8, 8, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx2, xnn_f32_vrsqrt_ukernel__avx2_approx_rsqrt_u16, 16, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx2, xnn_f32_vrsqrt_ukernel__avx2_approx_rsqrt_u32, 32, false, float, struct xnn_f32_default_params, NULL) #endif // XNN_ARCH_X86 || XNN_ARCH_X86_64 #if XNN_ENABLE_AVX512F && (XNN_ARCH_X86 || XNN_ARCH_X86_64) @@ -50,6 +63,11 @@ XNN_UKERNEL_WITH_PARAMS(xnn_arch_x86_avx512f, xnn_f32_vrsqrt_ukernel__avx512f_rsqrt_u64, 64, false, float, struct xnn_f32_default_params, NULL) #endif // XNN_ARCH_X86 || XNN_ARCH_X86_64 +#if XNN_ARCH_HEXAGON && XNN_ENABLE_HVX +XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__hvx_u32, 32, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__hvx_u64, 64, false, float, struct xnn_f32_default_params, NULL) +XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_vrsqrt_ukernel__hvx_u128, 128, false, float, struct xnn_f32_default_params, NULL) +#endif // XNN_ARCH_HEXAGON && XNN_ENABLE_HVX #ifdef XNN_DEFINED_UKERNEL_WITH_PARAMS #undef XNN_DEFINED_UKERNEL_WITH_PARAMS
diff --git a/src/f32-vrsqrt/gen/f32-vrsqrt-avx-approx-rsqrt.c b/src/f32-vrsqrt/gen/f32-vrsqrt-avx-approx-rsqrt.c new file mode 100644 index 0000000..fbe8524 --- /dev/null +++ b/src/f32-vrsqrt/gen/f32-vrsqrt-avx-approx-rsqrt.c
@@ -0,0 +1,351 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/f32-vrsqrt/simd-approx-rsqrt.c.in +// Generator: tools/xngen +// +// Copyright 2025 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 <assert.h> +#include <stddef.h> +#include <stdint.h> + +#include "src/xnnpack/simd/f32-avx.h" + +#include "src/xnnpack/common.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vunary.h" + +// In the following, we use Newton-Raphson steps on the equation +// $x^{-2} - a$, which expands to: +// +// $$x_{k+1} = 0.5 * x_k * (3.0 - a * x_k^2)$$ +// +// So we do the following steps: +// +// 1. t0 = x_k +// 2. t1 = t0 * t0 (x_k^2) +// 3. t2 = a * t1 (a * x_k^2) +// 4. t3 = 3.0 - t2 (3.0 - a * x_k^2) +// 5. t4 = 0.5 * t0 (0.5 * x_k) +// 6. y = t3 * t4 (0.5 * x_k * (3.0 - a * x_k^2)) +// +// Where $x_k$ is the original approximation and `y` contains the improved +// approximation $x_{k+1}$. + +// rsqrtps and the like do not produce consistent results across +// microarchitectures, so we use this famous (and portable) approximation +// instead: https://en.wikipedia.org/wiki/Fast_inverse_square_root +// This is not faster than _*_rsqrt(14)_ps, but it is numerically consistent +// across all CPUs. +static XNN_INLINE xnn_simd_f32_t +approx_reciprocal_sqrt_f32(xnn_simd_f32_t a) { + const int32_t magic = 0x5F375A86; + const __m128i vmagic = _mm_set1_epi32(magic); + __m128i a_lo = _mm256_castsi256_si128(_mm256_castps_si256(a)); + __m128i a_hi = _mm256_extractf128_si256(_mm256_castps_si256(a), 1); + __m128i res_lo = _mm_sub_epi32(vmagic, _mm_srai_epi32(a_lo, 1)); + __m128i res_hi = _mm_sub_epi32(vmagic, _mm_srai_epi32(a_hi, 1)); + return _mm256_castsi256_ps( + _mm256_insertf128_si256(_mm256_castsi128_si256(res_lo), res_hi, 1)); +} + + +void xnn_f32_vrsqrt_ukernel__avx_approx_rsqrt_u8( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 8); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 8; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 8; + } + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +} + +void xnn_f32_vrsqrt_ukernel__avx_approx_rsqrt_u16( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 8); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= 16 * sizeof(float); batch -= 16 * sizeof(float)) { + xnn_simd_f32_t vx_0 = xnn_loadu_f32(input + 0 * xnn_simd_size_f32); + xnn_simd_f32_t vx_1 = xnn_loadu_f32(input + 1 * xnn_simd_size_f32); + input += 16; + + xnn_simd_f32_t vy_0 = approx_reciprocal_sqrt_f32(vx_0); + xnn_simd_f32_t vy_1 = approx_reciprocal_sqrt_f32(vx_1); + + xnn_simd_f32_t vt1_0, vt2_0, vt3_0, vt4_0; + xnn_simd_f32_t vt1_1, vt2_1, vt3_1, vt4_1; + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + + xnn_storeu_f32(output + 0 * xnn_simd_size_f32, vy_0); + xnn_storeu_f32(output + 1 * xnn_simd_size_f32, vy_1); + output += 16; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 8; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 8; + } + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +} + +void xnn_f32_vrsqrt_ukernel__avx_approx_rsqrt_u32( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 8); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= 32 * sizeof(float); batch -= 32 * sizeof(float)) { + xnn_simd_f32_t vx_0 = xnn_loadu_f32(input + 0 * xnn_simd_size_f32); + xnn_simd_f32_t vx_1 = xnn_loadu_f32(input + 1 * xnn_simd_size_f32); + xnn_simd_f32_t vx_2 = xnn_loadu_f32(input + 2 * xnn_simd_size_f32); + xnn_simd_f32_t vx_3 = xnn_loadu_f32(input + 3 * xnn_simd_size_f32); + input += 32; + + xnn_simd_f32_t vy_0 = approx_reciprocal_sqrt_f32(vx_0); + xnn_simd_f32_t vy_1 = approx_reciprocal_sqrt_f32(vx_1); + xnn_simd_f32_t vy_2 = approx_reciprocal_sqrt_f32(vx_2); + xnn_simd_f32_t vy_3 = approx_reciprocal_sqrt_f32(vx_3); + + xnn_simd_f32_t vt1_0, vt2_0, vt3_0, vt4_0; + xnn_simd_f32_t vt1_1, vt2_1, vt3_1, vt4_1; + xnn_simd_f32_t vt1_2, vt2_2, vt3_2, vt4_2; + xnn_simd_f32_t vt1_3, vt2_3, vt3_3, vt4_3; + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + + xnn_storeu_f32(output + 0 * xnn_simd_size_f32, vy_0); + xnn_storeu_f32(output + 1 * xnn_simd_size_f32, vy_1); + xnn_storeu_f32(output + 2 * xnn_simd_size_f32, vy_2); + xnn_storeu_f32(output + 3 * xnn_simd_size_f32, vy_3); + output += 32; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 8; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 8; + } + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +}
diff --git a/src/f32-vrsqrt/gen/f32-vrsqrt-avx2-approx-rsqrt.c b/src/f32-vrsqrt/gen/f32-vrsqrt-avx2-approx-rsqrt.c new file mode 100644 index 0000000..3f605a3 --- /dev/null +++ b/src/f32-vrsqrt/gen/f32-vrsqrt-avx2-approx-rsqrt.c
@@ -0,0 +1,347 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/f32-vrsqrt/simd-approx-rsqrt.c.in +// Generator: tools/xngen +// +// Copyright 2025 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 <assert.h> +#include <stddef.h> +#include <stdint.h> + +#include "src/xnnpack/simd/f32-avx2.h" + +#include "src/xnnpack/common.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vunary.h" + +// In the following, we use Newton-Raphson steps on the equation +// $x^{-2} - a$, which expands to: +// +// $$x_{k+1} = 0.5 * x_k * (3.0 - a * x_k^2)$$ +// +// So we do the following steps: +// +// 1. t0 = x_k +// 2. t1 = t0 * t0 (x_k^2) +// 3. t2 = a * t1 (a * x_k^2) +// 4. t3 = 3.0 - t2 (3.0 - a * x_k^2) +// 5. t4 = 0.5 * t0 (0.5 * x_k) +// 6. y = t3 * t4 (0.5 * x_k * (3.0 - a * x_k^2)) +// +// Where $x_k$ is the original approximation and `y` contains the improved +// approximation $x_{k+1}$. + +// rsqrtps and the like do not produce consistent results across +// microarchitectures, so we use this famous (and portable) approximation +// instead: https://en.wikipedia.org/wiki/Fast_inverse_square_root +// This is not faster than _*_rsqrt(14)_ps, but it is numerically consistent +// across all CPUs. +static XNN_INLINE xnn_simd_f32_t +approx_reciprocal_sqrt_f32(xnn_simd_f32_t a) { + const int32_t magic = 0x5F375A86; + const __m256i vmagic = _mm256_set1_epi32(magic); + return _mm256_castsi256_ps( + _mm256_sub_epi32(vmagic, _mm256_srai_epi32(_mm256_castps_si256(a), 1))); +} + + +void xnn_f32_vrsqrt_ukernel__avx2_approx_rsqrt_u8( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 8); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 8; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 8; + } + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +} + +void xnn_f32_vrsqrt_ukernel__avx2_approx_rsqrt_u16( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 8); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= 16 * sizeof(float); batch -= 16 * sizeof(float)) { + xnn_simd_f32_t vx_0 = xnn_loadu_f32(input + 0 * xnn_simd_size_f32); + xnn_simd_f32_t vx_1 = xnn_loadu_f32(input + 1 * xnn_simd_size_f32); + input += 16; + + xnn_simd_f32_t vy_0 = approx_reciprocal_sqrt_f32(vx_0); + xnn_simd_f32_t vy_1 = approx_reciprocal_sqrt_f32(vx_1); + + xnn_simd_f32_t vt1_0, vt2_0, vt3_0, vt4_0; + xnn_simd_f32_t vt1_1, vt2_1, vt3_1, vt4_1; + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + + xnn_storeu_f32(output + 0 * xnn_simd_size_f32, vy_0); + xnn_storeu_f32(output + 1 * xnn_simd_size_f32, vy_1); + output += 16; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 8; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 8; + } + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +} + +void xnn_f32_vrsqrt_ukernel__avx2_approx_rsqrt_u32( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 8); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= 32 * sizeof(float); batch -= 32 * sizeof(float)) { + xnn_simd_f32_t vx_0 = xnn_loadu_f32(input + 0 * xnn_simd_size_f32); + xnn_simd_f32_t vx_1 = xnn_loadu_f32(input + 1 * xnn_simd_size_f32); + xnn_simd_f32_t vx_2 = xnn_loadu_f32(input + 2 * xnn_simd_size_f32); + xnn_simd_f32_t vx_3 = xnn_loadu_f32(input + 3 * xnn_simd_size_f32); + input += 32; + + xnn_simd_f32_t vy_0 = approx_reciprocal_sqrt_f32(vx_0); + xnn_simd_f32_t vy_1 = approx_reciprocal_sqrt_f32(vx_1); + xnn_simd_f32_t vy_2 = approx_reciprocal_sqrt_f32(vx_2); + xnn_simd_f32_t vy_3 = approx_reciprocal_sqrt_f32(vx_3); + + xnn_simd_f32_t vt1_0, vt2_0, vt3_0, vt4_0; + xnn_simd_f32_t vt1_1, vt2_1, vt3_1, vt4_1; + xnn_simd_f32_t vt1_2, vt2_2, vt3_2, vt4_2; + xnn_simd_f32_t vt1_3, vt2_3, vt3_3, vt4_3; + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + + xnn_storeu_f32(output + 0 * xnn_simd_size_f32, vy_0); + xnn_storeu_f32(output + 1 * xnn_simd_size_f32, vy_1); + xnn_storeu_f32(output + 2 * xnn_simd_size_f32, vy_2); + xnn_storeu_f32(output + 3 * xnn_simd_size_f32, vy_3); + output += 32; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 8; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 8; + } + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +}
diff --git a/src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u128.c b/src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u128.c new file mode 100644 index 0000000..1def458 --- /dev/null +++ b/src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u128.c
@@ -0,0 +1,118 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/f32-vrsqrt/simd.c.in +// Generator: tools/xngen +// +// Copyright 2025 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 <assert.h> + +#include "src/xnnpack/simd/f32-hvx.h" + +#include "src/xnnpack/common.h" +#include "src/xnnpack/vunary.h" + +// In the following, we do a single Newton-Raphson step on the equation +// $x^{-2} - a$, which expands to: +// +// $$x_{k+1} = 0.5 * x_k * (3.0 - a * x_k^2)$$ +// +// So we do the following steps: +// +// 1. t0 = x_k +// 2. t1 = t0 * t0 (x_k^2) +// 3. t3 = a * t1 - 3.0 (a * x_k^2 - 3.0) +// 4. t4 = 0.5 * t0 (-0.5 * x_k) +// 5. y = t3 * t4 ((-0.5 * x_k) * (a * x_k^2 - 3.0)) +// +// Where $x_k$ is the original 12-bit approximation and `y` contains the final +// 24-bit approximation $x_{k+1}$. + + +void xnn_f32_vrsqrt_ukernel__hvx_u128( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + + // Constants for the Newton-Raphson iteration. + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vneg_half, -0.5f); + + for (; batch >= 128 * sizeof(float); batch -= 128 * sizeof(float)) { + const xnn_simd_f32_t vx0 = xnn_loadu_f32(input + 0); + const xnn_simd_f32_t vx1 = xnn_loadu_f32(input + 32); + const xnn_simd_f32_t vx2 = xnn_loadu_f32(input + 64); + const xnn_simd_f32_t vx3 = xnn_loadu_f32(input + 96); + input += 128; + + // Generate the initial approximation. + const xnn_simd_f32_t vt0_0 = xnn_rsqrt_f32(vx0); + const xnn_simd_f32_t vt0_1 = xnn_rsqrt_f32(vx1); + const xnn_simd_f32_t vt0_2 = xnn_rsqrt_f32(vx2); + const xnn_simd_f32_t vt0_3 = xnn_rsqrt_f32(vx3); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1_0 = xnn_mul_f32(vt0_0, vt0_0); + const xnn_simd_f32_t vt1_1 = xnn_mul_f32(vt0_1, vt0_1); + const xnn_simd_f32_t vt1_2 = xnn_mul_f32(vt0_2, vt0_2); + const xnn_simd_f32_t vt1_3 = xnn_mul_f32(vt0_3, vt0_3); + const xnn_simd_f32_t vt3_0 = xnn_fmsub_f32(vx0, vt1_0, vthree); + const xnn_simd_f32_t vt3_1 = xnn_fmsub_f32(vx1, vt1_1, vthree); + const xnn_simd_f32_t vt3_2 = xnn_fmsub_f32(vx2, vt1_2, vthree); + const xnn_simd_f32_t vt3_3 = xnn_fmsub_f32(vx3, vt1_3, vthree); + const xnn_simd_f32_t vt4_0 = xnn_mul_f32(vneg_half, vt0_0); + const xnn_simd_f32_t vt4_1 = xnn_mul_f32(vneg_half, vt0_1); + const xnn_simd_f32_t vt4_2 = xnn_mul_f32(vneg_half, vt0_2); + const xnn_simd_f32_t vt4_3 = xnn_mul_f32(vneg_half, vt0_3); + const xnn_simd_f32_t vy0 = xnn_mul_f32(vt3_0, vt4_0); + const xnn_simd_f32_t vy1 = xnn_mul_f32(vt3_1, vt4_1); + const xnn_simd_f32_t vy2 = xnn_mul_f32(vt3_2, vt4_2); + const xnn_simd_f32_t vy3 = xnn_mul_f32(vt3_3, vt4_3); + + // Store the results. + xnn_storeu_f32(output + 0, vy0); + xnn_storeu_f32(output + 32, vy1); + xnn_storeu_f32(output + 64, vy2); + xnn_storeu_f32(output + 96, vy3); + output += 128; + } + for (; batch >= 32 * sizeof(float); batch -= 32 * sizeof(float)) { + const xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 32; + + // Generate the initial approximation. + const xnn_simd_f32_t vt0 = xnn_rsqrt_f32(vx); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1 = xnn_mul_f32(vt0, vt0); + const xnn_simd_f32_t vt3 = xnn_fmsub_f32(vx, vt1, vthree); + const xnn_simd_f32_t vt4 = xnn_mul_f32(vneg_half, vt0); + const xnn_simd_f32_t vy = xnn_mul_f32(vt3, vt4); + + xnn_storeu_f32(output, vy); + output += 32; + } + if XNN_UNLIKELY(batch != 0) { + const xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + // Generate the initial 12-bit approximation. + const xnn_simd_f32_t vt0 = xnn_rsqrt_f32(vx); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1 = xnn_mul_f32(vt0, vt0); + const xnn_simd_f32_t vt3 = xnn_fmsub_f32(vx, vt1, vthree); + const xnn_simd_f32_t vt4 = xnn_mul_f32(vneg_half, vt0); + xnn_simd_f32_t vy = xnn_mul_f32(vt3, vt4); + + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +}
diff --git a/src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u32.c b/src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u32.c new file mode 100644 index 0000000..b68435e --- /dev/null +++ b/src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u32.c
@@ -0,0 +1,80 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/f32-vrsqrt/simd.c.in +// Generator: tools/xngen +// +// Copyright 2025 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 <assert.h> + +#include "src/xnnpack/simd/f32-hvx.h" + +#include "src/xnnpack/common.h" +#include "src/xnnpack/vunary.h" + +// In the following, we do a single Newton-Raphson step on the equation +// $x^{-2} - a$, which expands to: +// +// $$x_{k+1} = 0.5 * x_k * (3.0 - a * x_k^2)$$ +// +// So we do the following steps: +// +// 1. t0 = x_k +// 2. t1 = t0 * t0 (x_k^2) +// 3. t3 = a * t1 - 3.0 (a * x_k^2 - 3.0) +// 4. t4 = 0.5 * t0 (-0.5 * x_k) +// 5. y = t3 * t4 ((-0.5 * x_k) * (a * x_k^2 - 3.0)) +// +// Where $x_k$ is the original 12-bit approximation and `y` contains the final +// 24-bit approximation $x_{k+1}$. + + +void xnn_f32_vrsqrt_ukernel__hvx_u32( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + + // Constants for the Newton-Raphson iteration. + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vneg_half, -0.5f); + + for (; batch >= 32 * sizeof(float); batch -= 32 * sizeof(float)) { + const xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 32; + + // Generate the initial approximation. + const xnn_simd_f32_t vt0 = xnn_rsqrt_f32(vx); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1 = xnn_mul_f32(vt0, vt0); + const xnn_simd_f32_t vt3 = xnn_fmsub_f32(vx, vt1, vthree); + const xnn_simd_f32_t vt4 = xnn_mul_f32(vneg_half, vt0); + const xnn_simd_f32_t vy = xnn_mul_f32(vt3, vt4); + + xnn_storeu_f32(output, vy); + output += 32; + } + if XNN_UNLIKELY(batch != 0) { + const xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + // Generate the initial 12-bit approximation. + const xnn_simd_f32_t vt0 = xnn_rsqrt_f32(vx); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1 = xnn_mul_f32(vt0, vt0); + const xnn_simd_f32_t vt3 = xnn_fmsub_f32(vx, vt1, vthree); + const xnn_simd_f32_t vt4 = xnn_mul_f32(vneg_half, vt0); + xnn_simd_f32_t vy = xnn_mul_f32(vt3, vt4); + + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +}
diff --git a/src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u64.c b/src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u64.c new file mode 100644 index 0000000..eaec285 --- /dev/null +++ b/src/f32-vrsqrt/gen/f32-vrsqrt-hvx-u64.c
@@ -0,0 +1,104 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/f32-vrsqrt/simd.c.in +// Generator: tools/xngen +// +// Copyright 2025 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 <assert.h> + +#include "src/xnnpack/simd/f32-hvx.h" + +#include "src/xnnpack/common.h" +#include "src/xnnpack/vunary.h" + +// In the following, we do a single Newton-Raphson step on the equation +// $x^{-2} - a$, which expands to: +// +// $$x_{k+1} = 0.5 * x_k * (3.0 - a * x_k^2)$$ +// +// So we do the following steps: +// +// 1. t0 = x_k +// 2. t1 = t0 * t0 (x_k^2) +// 3. t3 = a * t1 - 3.0 (a * x_k^2 - 3.0) +// 4. t4 = 0.5 * t0 (-0.5 * x_k) +// 5. y = t3 * t4 ((-0.5 * x_k) * (a * x_k^2 - 3.0)) +// +// Where $x_k$ is the original 12-bit approximation and `y` contains the final +// 24-bit approximation $x_{k+1}$. + + +void xnn_f32_vrsqrt_ukernel__hvx_u64( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + + // Constants for the Newton-Raphson iteration. + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vneg_half, -0.5f); + + for (; batch >= 64 * sizeof(float); batch -= 64 * sizeof(float)) { + const xnn_simd_f32_t vx0 = xnn_loadu_f32(input + 0); + const xnn_simd_f32_t vx1 = xnn_loadu_f32(input + 32); + input += 64; + + // Generate the initial approximation. + const xnn_simd_f32_t vt0_0 = xnn_rsqrt_f32(vx0); + const xnn_simd_f32_t vt0_1 = xnn_rsqrt_f32(vx1); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1_0 = xnn_mul_f32(vt0_0, vt0_0); + const xnn_simd_f32_t vt1_1 = xnn_mul_f32(vt0_1, vt0_1); + const xnn_simd_f32_t vt3_0 = xnn_fmsub_f32(vx0, vt1_0, vthree); + const xnn_simd_f32_t vt3_1 = xnn_fmsub_f32(vx1, vt1_1, vthree); + const xnn_simd_f32_t vt4_0 = xnn_mul_f32(vneg_half, vt0_0); + const xnn_simd_f32_t vt4_1 = xnn_mul_f32(vneg_half, vt0_1); + const xnn_simd_f32_t vy0 = xnn_mul_f32(vt3_0, vt4_0); + const xnn_simd_f32_t vy1 = xnn_mul_f32(vt3_1, vt4_1); + + // Store the results. + xnn_storeu_f32(output + 0, vy0); + xnn_storeu_f32(output + 32, vy1); + output += 64; + } + for (; batch >= 32 * sizeof(float); batch -= 32 * sizeof(float)) { + const xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 32; + + // Generate the initial approximation. + const xnn_simd_f32_t vt0 = xnn_rsqrt_f32(vx); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1 = xnn_mul_f32(vt0, vt0); + const xnn_simd_f32_t vt3 = xnn_fmsub_f32(vx, vt1, vthree); + const xnn_simd_f32_t vt4 = xnn_mul_f32(vneg_half, vt0); + const xnn_simd_f32_t vy = xnn_mul_f32(vt3, vt4); + + xnn_storeu_f32(output, vy); + output += 32; + } + if XNN_UNLIKELY(batch != 0) { + const xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + // Generate the initial 12-bit approximation. + const xnn_simd_f32_t vt0 = xnn_rsqrt_f32(vx); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1 = xnn_mul_f32(vt0, vt0); + const xnn_simd_f32_t vt3 = xnn_fmsub_f32(vx, vt1, vthree); + const xnn_simd_f32_t vt4 = xnn_mul_f32(vneg_half, vt0); + xnn_simd_f32_t vy = xnn_mul_f32(vt3, vt4); + + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +}
diff --git a/src/f32-vrsqrt/gen/f32-vrsqrt-scalar-approx-rsqrt.c b/src/f32-vrsqrt/gen/f32-vrsqrt-scalar-approx-rsqrt.c new file mode 100644 index 0000000..dc169ca --- /dev/null +++ b/src/f32-vrsqrt/gen/f32-vrsqrt-scalar-approx-rsqrt.c
@@ -0,0 +1,455 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/f32-vrsqrt/simd-approx-rsqrt.c.in +// Generator: tools/xngen +// +// Copyright 2025 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 <assert.h> +#include <stddef.h> +#include <stdint.h> + +#include "src/xnnpack/simd/f32-scalar.h" + +#include "src/xnnpack/common.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vunary.h" + +// In the following, we use Newton-Raphson steps on the equation +// $x^{-2} - a$, which expands to: +// +// $$x_{k+1} = 0.5 * x_k * (3.0 - a * x_k^2)$$ +// +// So we do the following steps: +// +// 1. t0 = x_k +// 2. t1 = t0 * t0 (x_k^2) +// 3. t2 = a * t1 (a * x_k^2) +// 4. t3 = 3.0 - t2 (3.0 - a * x_k^2) +// 5. t4 = 0.5 * t0 (0.5 * x_k) +// 6. y = t3 * t4 (0.5 * x_k * (3.0 - a * x_k^2)) +// +// Where $x_k$ is the original approximation and `y` contains the improved +// approximation $x_{k+1}$. + +// rsqrtps and the like do not produce consistent results across +// microarchitectures, so we use this famous (and portable) approximation +// instead: https://en.wikipedia.org/wiki/Fast_inverse_square_root +// This is not faster than _*_rsqrt(14)_ps, but it is numerically consistent +// across all CPUs. +static XNN_INLINE xnn_simd_f32_t +approx_reciprocal_sqrt_f32(xnn_simd_f32_t a) { + const int32_t magic = 0x5F375A86; + int32_t a_int; + memcpy(&a_int, &a, 4); + a_int = magic - (a_int >> 1); + memcpy(&a, &a_int, 4); + return a; +} + + +void xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u1( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 1); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 1; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 1; + } +} + +void xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u2( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 1); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= 2 * sizeof(float); batch -= 2 * sizeof(float)) { + xnn_simd_f32_t vx_0 = xnn_loadu_f32(input + 0 * xnn_simd_size_f32); + xnn_simd_f32_t vx_1 = xnn_loadu_f32(input + 1 * xnn_simd_size_f32); + input += 2; + + xnn_simd_f32_t vy_0 = approx_reciprocal_sqrt_f32(vx_0); + xnn_simd_f32_t vy_1 = approx_reciprocal_sqrt_f32(vx_1); + + xnn_simd_f32_t vt1_0, vt2_0, vt3_0, vt4_0; + xnn_simd_f32_t vt1_1, vt2_1, vt3_1, vt4_1; + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + + xnn_storeu_f32(output + 0 * xnn_simd_size_f32, vy_0); + xnn_storeu_f32(output + 1 * xnn_simd_size_f32, vy_1); + output += 2; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 1; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 1; + } +} + +void xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u4( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 1); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= 4 * sizeof(float); batch -= 4 * sizeof(float)) { + xnn_simd_f32_t vx_0 = xnn_loadu_f32(input + 0 * xnn_simd_size_f32); + xnn_simd_f32_t vx_1 = xnn_loadu_f32(input + 1 * xnn_simd_size_f32); + xnn_simd_f32_t vx_2 = xnn_loadu_f32(input + 2 * xnn_simd_size_f32); + xnn_simd_f32_t vx_3 = xnn_loadu_f32(input + 3 * xnn_simd_size_f32); + input += 4; + + xnn_simd_f32_t vy_0 = approx_reciprocal_sqrt_f32(vx_0); + xnn_simd_f32_t vy_1 = approx_reciprocal_sqrt_f32(vx_1); + xnn_simd_f32_t vy_2 = approx_reciprocal_sqrt_f32(vx_2); + xnn_simd_f32_t vy_3 = approx_reciprocal_sqrt_f32(vx_3); + + xnn_simd_f32_t vt1_0, vt2_0, vt3_0, vt4_0; + xnn_simd_f32_t vt1_1, vt2_1, vt3_1, vt4_1; + xnn_simd_f32_t vt1_2, vt2_2, vt3_2, vt4_2; + xnn_simd_f32_t vt1_3, vt2_3, vt3_3, vt4_3; + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + + xnn_storeu_f32(output + 0 * xnn_simd_size_f32, vy_0); + xnn_storeu_f32(output + 1 * xnn_simd_size_f32, vy_1); + xnn_storeu_f32(output + 2 * xnn_simd_size_f32, vy_2); + xnn_storeu_f32(output + 3 * xnn_simd_size_f32, vy_3); + output += 4; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 1; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 1; + } +} + +void xnn_f32_vrsqrt_ukernel__scalar_approx_rsqrt_u8( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 1); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= 8 * sizeof(float); batch -= 8 * sizeof(float)) { + xnn_simd_f32_t vx_0 = xnn_loadu_f32(input + 0 * xnn_simd_size_f32); + xnn_simd_f32_t vx_1 = xnn_loadu_f32(input + 1 * xnn_simd_size_f32); + xnn_simd_f32_t vx_2 = xnn_loadu_f32(input + 2 * xnn_simd_size_f32); + xnn_simd_f32_t vx_3 = xnn_loadu_f32(input + 3 * xnn_simd_size_f32); + xnn_simd_f32_t vx_4 = xnn_loadu_f32(input + 4 * xnn_simd_size_f32); + xnn_simd_f32_t vx_5 = xnn_loadu_f32(input + 5 * xnn_simd_size_f32); + xnn_simd_f32_t vx_6 = xnn_loadu_f32(input + 6 * xnn_simd_size_f32); + xnn_simd_f32_t vx_7 = xnn_loadu_f32(input + 7 * xnn_simd_size_f32); + input += 8; + + xnn_simd_f32_t vy_0 = approx_reciprocal_sqrt_f32(vx_0); + xnn_simd_f32_t vy_1 = approx_reciprocal_sqrt_f32(vx_1); + xnn_simd_f32_t vy_2 = approx_reciprocal_sqrt_f32(vx_2); + xnn_simd_f32_t vy_3 = approx_reciprocal_sqrt_f32(vx_3); + xnn_simd_f32_t vy_4 = approx_reciprocal_sqrt_f32(vx_4); + xnn_simd_f32_t vy_5 = approx_reciprocal_sqrt_f32(vx_5); + xnn_simd_f32_t vy_6 = approx_reciprocal_sqrt_f32(vx_6); + xnn_simd_f32_t vy_7 = approx_reciprocal_sqrt_f32(vx_7); + + xnn_simd_f32_t vt1_0, vt2_0, vt3_0, vt4_0; + xnn_simd_f32_t vt1_1, vt2_1, vt3_1, vt4_1; + xnn_simd_f32_t vt1_2, vt2_2, vt3_2, vt4_2; + xnn_simd_f32_t vt1_3, vt2_3, vt3_3, vt4_3; + xnn_simd_f32_t vt1_4, vt2_4, vt3_4, vt4_4; + xnn_simd_f32_t vt1_5, vt2_5, vt3_5, vt4_5; + xnn_simd_f32_t vt1_6, vt2_6, vt3_6, vt4_6; + xnn_simd_f32_t vt1_7, vt2_7, vt3_7, vt4_7; + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt1_4 = xnn_mul_f32(vy_4, vy_4); + vt1_5 = xnn_mul_f32(vy_5, vy_5); + vt1_6 = xnn_mul_f32(vy_6, vy_6); + vt1_7 = xnn_mul_f32(vy_7, vy_7); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt2_4 = xnn_mul_f32(vx_4, vt1_4); + vt2_5 = xnn_mul_f32(vx_5, vt1_5); + vt2_6 = xnn_mul_f32(vx_6, vt1_6); + vt2_7 = xnn_mul_f32(vx_7, vt1_7); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt3_4 = xnn_sub_f32(vthree, vt2_4); + vt3_5 = xnn_sub_f32(vthree, vt2_5); + vt3_6 = xnn_sub_f32(vthree, vt2_6); + vt3_7 = xnn_sub_f32(vthree, vt2_7); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vt4_4 = xnn_mul_f32(vhalf, vy_4); + vt4_5 = xnn_mul_f32(vhalf, vy_5); + vt4_6 = xnn_mul_f32(vhalf, vy_6); + vt4_7 = xnn_mul_f32(vhalf, vy_7); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + vy_4 = xnn_mul_f32(vt3_4, vt4_4); + vy_5 = xnn_mul_f32(vt3_5, vt4_5); + vy_6 = xnn_mul_f32(vt3_6, vt4_6); + vy_7 = xnn_mul_f32(vt3_7, vt4_7); + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt1_4 = xnn_mul_f32(vy_4, vy_4); + vt1_5 = xnn_mul_f32(vy_5, vy_5); + vt1_6 = xnn_mul_f32(vy_6, vy_6); + vt1_7 = xnn_mul_f32(vy_7, vy_7); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt2_4 = xnn_mul_f32(vx_4, vt1_4); + vt2_5 = xnn_mul_f32(vx_5, vt1_5); + vt2_6 = xnn_mul_f32(vx_6, vt1_6); + vt2_7 = xnn_mul_f32(vx_7, vt1_7); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt3_4 = xnn_sub_f32(vthree, vt2_4); + vt3_5 = xnn_sub_f32(vthree, vt2_5); + vt3_6 = xnn_sub_f32(vthree, vt2_6); + vt3_7 = xnn_sub_f32(vthree, vt2_7); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vt4_4 = xnn_mul_f32(vhalf, vy_4); + vt4_5 = xnn_mul_f32(vhalf, vy_5); + vt4_6 = xnn_mul_f32(vhalf, vy_6); + vt4_7 = xnn_mul_f32(vhalf, vy_7); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + vy_4 = xnn_mul_f32(vt3_4, vt4_4); + vy_5 = xnn_mul_f32(vt3_5, vt4_5); + vy_6 = xnn_mul_f32(vt3_6, vt4_6); + vy_7 = xnn_mul_f32(vt3_7, vt4_7); + + xnn_storeu_f32(output + 0 * xnn_simd_size_f32, vy_0); + xnn_storeu_f32(output + 1 * xnn_simd_size_f32, vy_1); + xnn_storeu_f32(output + 2 * xnn_simd_size_f32, vy_2); + xnn_storeu_f32(output + 3 * xnn_simd_size_f32, vy_3); + xnn_storeu_f32(output + 4 * xnn_simd_size_f32, vy_4); + xnn_storeu_f32(output + 5 * xnn_simd_size_f32, vy_5); + xnn_storeu_f32(output + 6 * xnn_simd_size_f32, vy_6); + xnn_storeu_f32(output + 7 * xnn_simd_size_f32, vy_7); + output += 8; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 1; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 1; + } +}
diff --git a/src/f32-vrsqrt/gen/f32-vrsqrt-sse2-approx-rsqrt.c b/src/f32-vrsqrt/gen/f32-vrsqrt-sse2-approx-rsqrt.c new file mode 100644 index 0000000..9a84b95 --- /dev/null +++ b/src/f32-vrsqrt/gen/f32-vrsqrt-sse2-approx-rsqrt.c
@@ -0,0 +1,347 @@ +// clang-format off +// Auto-generated file. Do not edit! +// Template: src/f32-vrsqrt/simd-approx-rsqrt.c.in +// Generator: tools/xngen +// +// Copyright 2025 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 <assert.h> +#include <stddef.h> +#include <stdint.h> + +#include "src/xnnpack/simd/f32-sse2.h" + +#include "src/xnnpack/common.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vunary.h" + +// In the following, we use Newton-Raphson steps on the equation +// $x^{-2} - a$, which expands to: +// +// $$x_{k+1} = 0.5 * x_k * (3.0 - a * x_k^2)$$ +// +// So we do the following steps: +// +// 1. t0 = x_k +// 2. t1 = t0 * t0 (x_k^2) +// 3. t2 = a * t1 (a * x_k^2) +// 4. t3 = 3.0 - t2 (3.0 - a * x_k^2) +// 5. t4 = 0.5 * t0 (0.5 * x_k) +// 6. y = t3 * t4 (0.5 * x_k * (3.0 - a * x_k^2)) +// +// Where $x_k$ is the original approximation and `y` contains the improved +// approximation $x_{k+1}$. + +// rsqrtps and the like do not produce consistent results across +// microarchitectures, so we use this famous (and portable) approximation +// instead: https://en.wikipedia.org/wiki/Fast_inverse_square_root +// This is not faster than _*_rsqrt(14)_ps, but it is numerically consistent +// across all CPUs. +static XNN_INLINE xnn_simd_f32_t +approx_reciprocal_sqrt_f32(xnn_simd_f32_t a) { + const int32_t magic = 0x5F375A86; + const __m128i vmagic = _mm_set1_epi32(magic); + return _mm_castsi128_ps( + _mm_sub_epi32(vmagic, _mm_srai_epi32(_mm_castps_si128(a), 1))); +} + + +void xnn_f32_vrsqrt_ukernel__sse2_approx_rsqrt_u4( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 4); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 4; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 4; + } + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +} + +void xnn_f32_vrsqrt_ukernel__sse2_approx_rsqrt_u8( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 4); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= 8 * sizeof(float); batch -= 8 * sizeof(float)) { + xnn_simd_f32_t vx_0 = xnn_loadu_f32(input + 0 * xnn_simd_size_f32); + xnn_simd_f32_t vx_1 = xnn_loadu_f32(input + 1 * xnn_simd_size_f32); + input += 8; + + xnn_simd_f32_t vy_0 = approx_reciprocal_sqrt_f32(vx_0); + xnn_simd_f32_t vy_1 = approx_reciprocal_sqrt_f32(vx_1); + + xnn_simd_f32_t vt1_0, vt2_0, vt3_0, vt4_0; + xnn_simd_f32_t vt1_1, vt2_1, vt3_1, vt4_1; + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + + xnn_storeu_f32(output + 0 * xnn_simd_size_f32, vy_0); + xnn_storeu_f32(output + 1 * xnn_simd_size_f32, vy_1); + output += 8; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 4; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 4; + } + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +} + +void xnn_f32_vrsqrt_ukernel__sse2_approx_rsqrt_u16( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == 4); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + for (; batch >= 16 * sizeof(float); batch -= 16 * sizeof(float)) { + xnn_simd_f32_t vx_0 = xnn_loadu_f32(input + 0 * xnn_simd_size_f32); + xnn_simd_f32_t vx_1 = xnn_loadu_f32(input + 1 * xnn_simd_size_f32); + xnn_simd_f32_t vx_2 = xnn_loadu_f32(input + 2 * xnn_simd_size_f32); + xnn_simd_f32_t vx_3 = xnn_loadu_f32(input + 3 * xnn_simd_size_f32); + input += 16; + + xnn_simd_f32_t vy_0 = approx_reciprocal_sqrt_f32(vx_0); + xnn_simd_f32_t vy_1 = approx_reciprocal_sqrt_f32(vx_1); + xnn_simd_f32_t vy_2 = approx_reciprocal_sqrt_f32(vx_2); + xnn_simd_f32_t vy_3 = approx_reciprocal_sqrt_f32(vx_3); + + xnn_simd_f32_t vt1_0, vt2_0, vt3_0, vt4_0; + xnn_simd_f32_t vt1_1, vt2_1, vt3_1, vt4_1; + xnn_simd_f32_t vt1_2, vt2_2, vt3_2, vt4_2; + xnn_simd_f32_t vt1_3, vt2_3, vt3_3, vt4_3; + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + // Do a Newton-Raphson step as described above. + vt1_0 = xnn_mul_f32(vy_0, vy_0); + vt1_1 = xnn_mul_f32(vy_1, vy_1); + vt1_2 = xnn_mul_f32(vy_2, vy_2); + vt1_3 = xnn_mul_f32(vy_3, vy_3); + vt2_0 = xnn_mul_f32(vx_0, vt1_0); + vt2_1 = xnn_mul_f32(vx_1, vt1_1); + vt2_2 = xnn_mul_f32(vx_2, vt1_2); + vt2_3 = xnn_mul_f32(vx_3, vt1_3); + vt3_0 = xnn_sub_f32(vthree, vt2_0); + vt3_1 = xnn_sub_f32(vthree, vt2_1); + vt3_2 = xnn_sub_f32(vthree, vt2_2); + vt3_3 = xnn_sub_f32(vthree, vt2_3); + vt4_0 = xnn_mul_f32(vhalf, vy_0); + vt4_1 = xnn_mul_f32(vhalf, vy_1); + vt4_2 = xnn_mul_f32(vhalf, vy_2); + vt4_3 = xnn_mul_f32(vhalf, vy_3); + vy_0 = xnn_mul_f32(vt3_0, vt4_0); + vy_1 = xnn_mul_f32(vt3_1, vt4_1); + vy_2 = xnn_mul_f32(vt3_2, vt4_2); + vy_3 = xnn_mul_f32(vt3_3, vt4_3); + + xnn_storeu_f32(output + 0 * xnn_simd_size_f32, vy_0); + xnn_storeu_f32(output + 1 * xnn_simd_size_f32, vy_1); + xnn_storeu_f32(output + 2 * xnn_simd_size_f32, vy_2); + xnn_storeu_f32(output + 3 * xnn_simd_size_f32, vy_3); + output += 16; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 4; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += 4; + } + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +}
diff --git a/src/f32-vrsqrt/simd-approx-rsqrt.c.in b/src/f32-vrsqrt/simd-approx-rsqrt.c.in new file mode 100644 index 0000000..bc055c5 --- /dev/null +++ b/src/f32-vrsqrt/simd-approx-rsqrt.c.in
@@ -0,0 +1,159 @@ +// Copyright 2025 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. + +$BATCH_TILES = tuple(int(bt) for bt in BATCH_TILES.split(",")) +$SIMD_SIZE = BATCH_TILES[0] +#include <assert.h> +#include <stddef.h> +#include <stdint.h> + +#include "src/xnnpack/simd/f32-${ARCH}.h" + +#include "src/xnnpack/common.h" +#include "src/xnnpack/microparams.h" +#include "src/xnnpack/vunary.h" + +// In the following, we use Newton-Raphson steps on the equation +// $x^{-2} - a$, which expands to: +// +// $$x_{k+1} = 0.5 * x_k * (3.0 - a * x_k^2)$$ +// +// So we do the following steps: +// +// 1. t0 = x_k +// 2. t1 = t0 * t0 (x_k^2) +// 3. t2 = a * t1 (a * x_k^2) +// 4. t3 = 3.0 - t2 (3.0 - a * x_k^2) +// 5. t4 = 0.5 * t0 (0.5 * x_k) +// 6. y = t3 * t4 (0.5 * x_k * (3.0 - a * x_k^2)) +// +// Where $x_k$ is the original approximation and `y` contains the improved +// approximation $x_{k+1}$. +$NEWTON_RHAPSON_ITERATIONS = 2 + +// rsqrtps and the like do not produce consistent results across +// microarchitectures, so we use this famous (and portable) approximation +// instead: https://en.wikipedia.org/wiki/Fast_inverse_square_root +// This is not faster than _*_rsqrt(14)_ps, but it is numerically consistent +// across all CPUs. +static XNN_INLINE xnn_simd_f32_t +approx_reciprocal_sqrt_f32(xnn_simd_f32_t a) { + const int32_t magic = 0x5F375A86; + $if ARCH == "sse2": + const __m128i vmagic = _mm_set1_epi32(magic); + return _mm_castsi128_ps( + _mm_sub_epi32(vmagic, _mm_srai_epi32(_mm_castps_si128(a), 1))); + $elif ARCH == "avx": + const __m128i vmagic = _mm_set1_epi32(magic); + __m128i a_lo = _mm256_castsi256_si128(_mm256_castps_si256(a)); + __m128i a_hi = _mm256_extractf128_si256(_mm256_castps_si256(a), 1); + __m128i res_lo = _mm_sub_epi32(vmagic, _mm_srai_epi32(a_lo, 1)); + __m128i res_hi = _mm_sub_epi32(vmagic, _mm_srai_epi32(a_hi, 1)); + return _mm256_castsi256_ps( + _mm256_insertf128_si256(_mm256_castsi128_si256(res_lo), res_hi, 1)); + $elif ARCH == "avx2": + const __m256i vmagic = _mm256_set1_epi32(magic); + return _mm256_castsi256_ps( + _mm256_sub_epi32(vmagic, _mm256_srai_epi32(_mm256_castps_si256(a), 1))); + $elif ARCH == "avx512f": + const __m512i vmagic = _mm512_set1_epi32(magic); + return _mm512_castsi512_ps( + _mm512_sub_epi32(vmagic, _mm512_srai_epi32(_mm512_castps_si512(a), 1))); + $elif ARCH == "scalar": + int32_t a_int; + memcpy(&a_int, &a, 4); + a_int = magic - (a_int >> 1); + memcpy(&a, &a_int, 4); + return a; +} + +$for BATCH_TILE in BATCH_TILES: + $assert BATCH_TILE % SIMD_SIZE == 0 + $assert BATCH_TILE >= SIMD_SIZE + $SIMD_TILE = BATCH_TILE // SIMD_SIZE + + void xnn_f32_vrsqrt_ukernel__${ARCH}_approx_rsqrt_u${BATCH_TILE}( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params unused_params[restrict XNN_MIN_ELEMENTS(1)]) + { + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + assert(xnn_simd_size_f32 == ${SIMD_SIZE}); + + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vhalf, 0.5f); + + // Register pressure isn't much of an issue in these kernels, but for some + // reason this is necessary to avoid up to *~10x* slower code. + // Stack alignment...? + XNN_FORCE_REALIZATION(vthree); + XNN_FORCE_REALIZATION(vhalf); + + $if SIMD_TILE > 1: + for (; batch >= ${BATCH_TILE} * sizeof(float); batch -= ${BATCH_TILE} * sizeof(float)) { + $for N in range(SIMD_TILE): + xnn_simd_f32_t vx_${N} = xnn_loadu_f32(input + ${N} * xnn_simd_size_f32); + input += ${BATCH_TILE}; + + $for N in range(SIMD_TILE): + xnn_simd_f32_t vy_${N} = approx_reciprocal_sqrt_f32(vx_${N}); + + $for N in range(SIMD_TILE): + xnn_simd_f32_t vt1_${N}, vt2_${N}, vt3_${N}, vt4_${N}; + $for i in range(NEWTON_RHAPSON_ITERATIONS): + // Do a Newton-Raphson step as described above. + $for N in range(SIMD_TILE): + vt1_${N} = xnn_mul_f32(vy_${N}, vy_${N}); + $for N in range(SIMD_TILE): + vt2_${N} = xnn_mul_f32(vx_${N}, vt1_${N}); + $for N in range(SIMD_TILE): + vt3_${N} = xnn_sub_f32(vthree, vt2_${N}); + $for N in range(SIMD_TILE): + vt4_${N} = xnn_mul_f32(vhalf, vy_${N}); + $for N in range(SIMD_TILE): + vy_${N} = xnn_mul_f32(vt3_${N}, vt4_${N}); + + $for N in range(SIMD_TILE): + xnn_storeu_f32(output + ${N} * xnn_simd_size_f32, vy_${N}); + output += ${BATCH_TILE}; + } + for (; batch >= xnn_simd_bytes_f32; batch -= xnn_simd_bytes_f32) { + xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += ${SIMD_SIZE}; + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + $for i in range(NEWTON_RHAPSON_ITERATIONS): + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_storeu_f32(output, vy); + output += ${SIMD_SIZE}; + } + $if SIMD_SIZE > 1: + if XNN_UNLIKELY(batch != 0) { + xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + xnn_simd_f32_t vy = approx_reciprocal_sqrt_f32(vx); + + xnn_simd_f32_t vt1, vt2, vt3, vt4; + $for i in range(NEWTON_RHAPSON_ITERATIONS): + // Do a Newton-Raphson step as described above. + vt1 = xnn_mul_f32(vy, vy); + vt2 = xnn_mul_f32(vx, vt1); + vt3 = xnn_sub_f32(vthree, vt2); + vt4 = xnn_mul_f32(vhalf, vy); + vy = xnn_mul_f32(vt3, vt4); + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } + }
diff --git a/src/f32-vrsqrt/simd.c.in b/src/f32-vrsqrt/simd.c.in new file mode 100644 index 0000000..7eca932 --- /dev/null +++ b/src/f32-vrsqrt/simd.c.in
@@ -0,0 +1,103 @@ +// Copyright 2025 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 % 32 == 0 +$assert BATCH_TILE >= 32 +$SIMD_TILE = BATCH_TILE // 32 +#include <assert.h> + +#include "src/xnnpack/simd/f32-${ARCH}.h" + +#include "src/xnnpack/common.h" +#include "src/xnnpack/vunary.h" + +// In the following, we do a single Newton-Raphson step on the equation +// $x^{-2} - a$, which expands to: +// +// $$x_{k+1} = 0.5 * x_k * (3.0 - a * x_k^2)$$ +// +// So we do the following steps: +// +// 1. t0 = x_k +// 2. t1 = t0 * t0 (x_k^2) +// 3. t3 = a * t1 - 3.0 (a * x_k^2 - 3.0) +// 4. t4 = 0.5 * t0 (-0.5 * x_k) +// 5. y = t3 * t4 ((-0.5 * x_k) * (a * x_k^2 - 3.0)) +// +// Where $x_k$ is the original 12-bit approximation and `y` contains the final +// 24-bit approximation $x_{k+1}$. + + +void xnn_f32_vrsqrt_ukernel__${ARCH}_u${BATCH_TILE}( + size_t batch, + const float* input, + float* output, + const struct xnn_f32_default_params params[restrict XNN_MIN_ELEMENTS(1)]) +{ + assert(batch != 0); + assert(batch % sizeof(float) == 0); + assert(input != NULL); + assert(output != NULL); + + // Constants for the Newton-Raphson iteration. + XNN_SIMD_CONST_F32(vthree, 3.0f); + XNN_SIMD_CONST_F32(vneg_half, -0.5f); + + $if BATCH_TILE > 32: + for (; batch >= ${BATCH_TILE} * sizeof(float); batch -= ${BATCH_TILE} * sizeof(float)) { + $for N in range(SIMD_TILE): + const xnn_simd_f32_t vx${N} = xnn_loadu_f32(input + ${N * 32}); + input += ${BATCH_TILE}; + + // Generate the initial approximation. + $for N in range(SIMD_TILE): + const xnn_simd_f32_t vt0_${N} = xnn_rsqrt_f32(vx${N}); + + // Do a single Newton-Raphson step as described above. + $for N in range(SIMD_TILE): + const xnn_simd_f32_t vt1_${N} = xnn_mul_f32(vt0_${N}, vt0_${N}); + $for N in range(SIMD_TILE): + const xnn_simd_f32_t vt3_${N} = xnn_fmsub_f32(vx${N}, vt1_${N}, vthree); + $for N in range(SIMD_TILE): + const xnn_simd_f32_t vt4_${N} = xnn_mul_f32(vneg_half, vt0_${N}); + $for N in range(SIMD_TILE): + const xnn_simd_f32_t vy${N} = xnn_mul_f32(vt3_${N}, vt4_${N}); + + // Store the results. + $for N in range(SIMD_TILE): + xnn_storeu_f32(output + ${N * 32}, vy${N}); + output += ${BATCH_TILE}; + } + for (; batch >= 32 * sizeof(float); batch -= 32 * sizeof(float)) { + const xnn_simd_f32_t vx = xnn_loadu_f32(input); + input += 32; + + // Generate the initial approximation. + const xnn_simd_f32_t vt0 = xnn_rsqrt_f32(vx); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1 = xnn_mul_f32(vt0, vt0); + const xnn_simd_f32_t vt3 = xnn_fmsub_f32(vx, vt1, vthree); + const xnn_simd_f32_t vt4 = xnn_mul_f32(vneg_half, vt0); + const xnn_simd_f32_t vy = xnn_mul_f32(vt3, vt4); + + xnn_storeu_f32(output, vy); + output += 32; + } + if XNN_UNLIKELY(batch != 0) { + const xnn_simd_f32_t vx = xnn_load_tail_f32(input, batch >> XNN_LOG2_SIZEOF_FLOAT); + + // Generate the initial 12-bit approximation. + const xnn_simd_f32_t vt0 = xnn_rsqrt_f32(vx); + + // Do a single Newton-Raphson step as described above. + const xnn_simd_f32_t vt1 = xnn_mul_f32(vt0, vt0); + const xnn_simd_f32_t vt3 = xnn_fmsub_f32(vx, vt1, vthree); + const xnn_simd_f32_t vt4 = xnn_mul_f32(vneg_half, vt0); + xnn_simd_f32_t vy = xnn_mul_f32(vt3, vt4); + + xnn_store_tail_f32(output, vy, batch >> XNN_LOG2_SIZEOF_FLOAT); + } +}
diff --git a/src/xnnpack/simd/f32-hvx.h b/src/xnnpack/simd/f32-hvx.h index 3b3e3a4..fd2e34e 100644 --- a/src/xnnpack/simd/f32-hvx.h +++ b/src/xnnpack/simd/f32-hvx.h
@@ -266,6 +266,11 @@ #define XNN_SIMD_HAVE_RSQRT_F32 0 #define XNN_SIMD_NUM_RCP_ITER_F32 2 +static XNN_INLINE xnn_simd_f32_t xnn_rsqrt_f32(xnn_simd_f32_t a) { + const HVX_Vector vmagic = Q6_V_vsplat_R(0x5F375A86); + return Q6_Vw_vsub_VwVw(vmagic, Q6_Vw_vasr_VwR(a, 1)); +} + #define XNN_SIMD_HAVE_REDUCE_MAX_F32 1 static XNN_INLINE float xnn_reduce_max_f32(xnn_simd_f32_t v) { v = Q6_Vsf_vmax_VsfVsf(v, Q6_V_vror_VR(v, 64));