// Copyright 2020 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 $assert RR_STEPS in [1, 2] $assert DIV_ALGO in ["div", "nr2"] $ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" $SIMD_TILE = BATCH_TILE // 8 #include #include #include #include static const int32_t mask_table[14] = {-1, -1, -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0}; void xnn_f32_sigmoid_ukernel__avx_rr${RR_STEPS}_p5_${DIV_ALGO}_x${BATCH_TILE}( size_t n, const float* x, float* y, const void* params) { assert(n % sizeof(float) == 0); const __m256 vsign_mask = _mm256_set1_ps(-0.0f); const __m256 vmagic_bias = _mm256_set1_ps(0x1.8000FEp23f); const __m256 vlog2e = _mm256_set1_ps(0x1.715476p0f); $if RR_STEPS == 1: const __m256 vminus_ln2 = _mm256_set1_ps(-0x1.62E43p-1f); $else: const __m256 vminus_ln2_hi = _mm256_set1_ps(-0x1.62E43p-1f); const __m256 vminus_ln2_lo = _mm256_set1_ps(0x1.05C61p-29f); const __m256 vc5 = _mm256_set1_ps(0x1.0F9F9Cp-7f); const __m256 vc4 = _mm256_set1_ps(0x1.573A1Ap-5f); const __m256 vc3 = _mm256_set1_ps(0x1.555A80p-3f); const __m256 vc2 = _mm256_set1_ps(0x1.FFFDC6p-2f); const __m256 vc1 = _mm256_set1_ps(0x1.FFFFF6p-1f); const __m256 vone = _mm256_set1_ps(1.0f); $if DIV_ALGO == "nr2": const __m256 vtwo = _mm256_set1_ps(2.0f); const __m256 vdenorm_cutoff = _mm256_set1_ps(-0x1.5D589Ep+6f); $if BATCH_TILE > 8: for (; n >= ${BATCH_TILE} * sizeof(float); n -= ${BATCH_TILE} * sizeof(float)) { const __m256 vx${ABC[0]} = _mm256_loadu_ps(x); $for N in range(1, SIMD_TILE): const __m256 vx${ABC[N]} = _mm256_loadu_ps(x + ${N * 8}); x += ${BATCH_TILE}; $for N in range(SIMD_TILE): const __m256 vz${ABC[N]} = _mm256_or_ps(vx${ABC[N]}, vsign_mask); $for N in range(SIMD_TILE): __m256 vn${ABC[N]} = _mm256_add_ps(_mm256_mul_ps(vz${ABC[N]}, vlog2e), vmagic_bias); $for N in range(SIMD_TILE): const __m128 vs_lo${ABC[N]} = _mm_castsi128_ps(_mm_slli_epi32(_mm_castps_si128(_mm256_castps256_ps128(vn${ABC[N]})), 23)); const __m128 vs_hi${ABC[N]} = _mm_castsi128_ps(_mm_slli_epi32(_mm_castps_si128(_mm256_extractf128_ps(vn${ABC[N]}, 1)), 23)); const __m256 vs${ABC[N]} = _mm256_insertf128_ps(_mm256_castps128_ps256(vs_lo${ABC[N]}), vs_hi${ABC[N]}, 1); $for N in range(SIMD_TILE): vn${ABC[N]} = _mm256_sub_ps(vn${ABC[N]}, vmagic_bias); $if RR_STEPS == 1: $for N in range(SIMD_TILE): __m256 vt${ABC[N]} = _mm256_add_ps(_mm256_mul_ps(vn${ABC[N]}, vminus_ln2), vz${ABC[N]}); $else: $for N in range(SIMD_TILE): __m256 vt${ABC[N]} = _mm256_add_ps(_mm256_mul_ps(vn${ABC[N]}, vminus_ln2_hi), vz${ABC[N]}); $for N in range(SIMD_TILE): vt${ABC[N]} = _mm256_add_ps(_mm256_mul_ps(vn${ABC[N]}, vminus_ln2_lo), vt${ABC[N]}); $for N in range(SIMD_TILE): __m256 vp${ABC[N]} = _mm256_add_ps(_mm256_mul_ps(vc5, vt${ABC[N]}), vc4); $for N in range(SIMD_TILE): vp${ABC[N]} = _mm256_add_ps(_mm256_mul_ps(vp${ABC[N]}, vt${ABC[N]}), vc3); $for N in range(SIMD_TILE): vp${ABC[N]} = _mm256_add_ps(_mm256_mul_ps(vp${ABC[N]}, vt${ABC[N]}), vc2); $for N in range(SIMD_TILE): vp${ABC[N]} = _mm256_add_ps(_mm256_mul_ps(vp${ABC[N]}, vt${ABC[N]}), vc1); $for N in range(SIMD_TILE): vt${ABC[N]} = _mm256_mul_ps(vt${ABC[N]}, vs${ABC[N]}); $for N in range(SIMD_TILE): const __m256 ve${ABC[N]} = _mm256_add_ps(_mm256_mul_ps(vt${ABC[N]}, vp${ABC[N]}), vs${ABC[N]}); $for N in range(SIMD_TILE): const __m256 vd${ABC[N]} = _mm256_add_ps(ve${ABC[N]}, vone); $if DIV_ALGO == "div": $for N in range(SIMD_TILE): __m256 vf${ABC[N]} = _mm256_div_ps(ve${ABC[N]}, vd${ABC[N]}); $else: $for N in range(SIMD_TILE): __m256 vr${ABC[N]} = _mm256_rcp_ps(vd${ABC[N]}); $for N in range(SIMD_TILE): vr${ABC[N]} = _mm256_mul_ps(vr${ABC[N]}, _mm256_sub_ps(vtwo, _mm256_mul_ps(vr${ABC[N]}, vd${ABC[N]}))); vr${ABC[N]} = _mm256_mul_ps(vr${ABC[N]}, _mm256_sub_ps(vtwo, _mm256_mul_ps(vr${ABC[N]}, vd${ABC[N]}))); $for N in range(SIMD_TILE): __m256 vf${ABC[N]} = _mm256_mul_ps(ve${ABC[N]}, vr${ABC[N]}); $for N in range(SIMD_TILE): vf${ABC[N]} = _mm256_andnot_ps(_mm256_cmp_ps(vz${ABC[N]}, vdenorm_cutoff, _CMP_LT_OS), vf${ABC[N]}); $for N in range(SIMD_TILE): vf${ABC[N]} = _mm256_blendv_ps(_mm256_sub_ps(vone, vf${ABC[N]}), vf${ABC[N]}, vx${ABC[N]}); _mm256_storeu_ps(y, vf${ABC[0]}); $for N in range(1, SIMD_TILE): _mm256_storeu_ps(y + ${N * 8}, vf${ABC[N]}); y += ${BATCH_TILE}; } for (; n >= 8 * sizeof(float); n -= 8 * sizeof(float)) { const __m256 vx = _mm256_loadu_ps(x); x += 8; const __m256 vz = _mm256_or_ps(vx, vsign_mask); __m256 vn = _mm256_add_ps(_mm256_mul_ps(vz, vlog2e), vmagic_bias); const __m128 vs_lo = _mm_castsi128_ps(_mm_slli_epi32(_mm_castps_si128(_mm256_castps256_ps128(vn)), 23)); const __m128 vs_hi = _mm_castsi128_ps(_mm_slli_epi32(_mm_castps_si128(_mm256_extractf128_ps(vn, 1)), 23)); const __m256 vs = _mm256_insertf128_ps(_mm256_castps128_ps256(vs_lo), vs_hi, 1); vn = _mm256_sub_ps(vn, vmagic_bias); $if RR_STEPS == 1: __m256 vt = _mm256_add_ps(_mm256_mul_ps(vn, vminus_ln2), vz); $else: __m256 vt = _mm256_add_ps(_mm256_mul_ps(vn, vminus_ln2_hi), vz); vt = _mm256_add_ps(_mm256_mul_ps(vn, vminus_ln2_lo), vt); __m256 vp = _mm256_add_ps(_mm256_mul_ps(vc5, vt), vc4); vp = _mm256_add_ps(_mm256_mul_ps(vp, vt), vc3); vp = _mm256_add_ps(_mm256_mul_ps(vp, vt), vc2); vp = _mm256_add_ps(_mm256_mul_ps(vp, vt), vc1); vt = _mm256_mul_ps(vt, vs); const __m256 ve = _mm256_add_ps(_mm256_mul_ps(vt, vp), vs); const __m256 vd = _mm256_add_ps(ve, vone); $if DIV_ALGO == "div": __m256 vf = _mm256_div_ps(ve, vd); $else: __m256 vr = _mm256_rcp_ps(vd); vr = _mm256_mul_ps(vr, _mm256_sub_ps(vtwo, _mm256_mul_ps(vr, vd))); vr = _mm256_mul_ps(vr, _mm256_sub_ps(vtwo, _mm256_mul_ps(vr, vd))); __m256 vf = _mm256_mul_ps(ve, vr); vf = _mm256_andnot_ps(_mm256_cmp_ps(vz, vdenorm_cutoff, _CMP_LT_OS), vf); vf = _mm256_blendv_ps(_mm256_sub_ps(vone, vf), vf, vx); _mm256_storeu_ps(y, vf); y += 8; } if XNN_UNLIKELY(n != 0) { assert(n >= 1 * sizeof(float)); assert(n <= 7 * sizeof(float)); __m256i vmask = _mm256_loadu_si256((const __m256i*) ((uintptr_t) &mask_table[7] - n)); const __m256 vx = _mm256_maskload_ps(x, vmask); const __m256 vz = _mm256_or_ps(vx, vsign_mask); __m256 vn = _mm256_add_ps(_mm256_mul_ps(vz, vlog2e), vmagic_bias); const __m128 vs_lo = _mm_castsi128_ps(_mm_slli_epi32(_mm_castps_si128(_mm256_castps256_ps128(vn)), 23)); const __m128 vs_hi = _mm_castsi128_ps(_mm_slli_epi32(_mm_castps_si128(_mm256_extractf128_ps(vn, 1)), 23)); const __m256 vs = _mm256_insertf128_ps(_mm256_castps128_ps256(vs_lo), vs_hi, 1); vn = _mm256_sub_ps(vn, vmagic_bias); $if RR_STEPS == 1: __m256 vt = _mm256_add_ps(_mm256_mul_ps(vn, vminus_ln2), vz); $else: __m256 vt = _mm256_add_ps(_mm256_mul_ps(vn, vminus_ln2_hi), vz); vt = _mm256_add_ps(_mm256_mul_ps(vn, vminus_ln2_lo), vt); __m256 vp = _mm256_add_ps(_mm256_mul_ps(vc5, vt), vc4); vp = _mm256_add_ps(_mm256_mul_ps(vp, vt), vc3); vp = _mm256_add_ps(_mm256_mul_ps(vp, vt), vc2); vp = _mm256_add_ps(_mm256_mul_ps(vp, vt), vc1); vt = _mm256_mul_ps(vt, vs); const __m256 ve = _mm256_add_ps(_mm256_mul_ps(vt, vp), vs); const __m256 vd = _mm256_add_ps(ve, vone); $if DIV_ALGO == "div": __m256 vf = _mm256_div_ps(ve, vd); $else: __m256 vr = _mm256_rcp_ps(vd); vr = _mm256_mul_ps(vr, _mm256_sub_ps(vtwo, _mm256_mul_ps(vr, vd))); vr = _mm256_mul_ps(vr, _mm256_sub_ps(vtwo, _mm256_mul_ps(vr, vd))); __m256 vf = _mm256_mul_ps(ve, vr); vf = _mm256_andnot_ps(_mm256_cmp_ps(vz, vdenorm_cutoff, _CMP_LT_OS), vf); vf = _mm256_blendv_ps(_mm256_sub_ps(vone, vf), vf, vx); // _mm256_maskstore_ps(y, vmask, vf) could be used here, but triggers msan failures (probably an msan bug). __m128 vf_lo = _mm256_castps256_ps128(vf); if (n & (4 * sizeof(float))) { _mm_storeu_ps(y, vf_lo); vf_lo = _mm256_extractf128_ps(vf, 1); y += 4; } if (n & (2 * sizeof(float))) { _mm_storel_pi((__m64*) y, vf_lo); vf_lo = _mm_movehl_ps(vf_lo, vf_lo); y += 2; } if (n & (1 * sizeof(float))) { _mm_store_ss(y, vf_lo); } } }