diff --git a/src/nnue/nnue_feature_transformer.h b/src/nnue/nnue_feature_transformer.h index 3edd763e2b4..297000de5e9 100644 --- a/src/nnue/nnue_feature_transformer.h +++ b/src/nnue/nnue_feature_transformer.h @@ -226,8 +226,6 @@ class FeatureTransformer { OutputType* output, int bucket, [[maybe_unused]] NNZInfo& nnzInfo) const { - - using namespace SIMD; accumulatorStack.evaluate(pos, *this, cache); const auto& accumulatorState = accumulatorStack.latest(); @@ -240,194 +238,198 @@ class FeatureTransformer { const auto& accumulation = accumulatorState.accumulation; for (IndexType p = 0; p < 2; ++p) - { - const IndexType offset = (HalfDimensions / 2) * p; + transform_perspective(accumulation[perspectives[p]], output, p, nnzInfo); + + return psqt; + } + + private: + static void transform_perspective(const std::array& accumulation, + OutputType* output, + IndexType perspective, + [[maybe_unused]] NNZInfo& nnzInfo) { + + using namespace SIMD; + const IndexType offset = (HalfDimensions / 2) * perspective; #if defined(VECTOR) - [[maybe_unused]] auto cursor = nnzInfo.make_cursor(p); - - constexpr IndexType OutputChunkSize = MaxChunkSize; - static_assert((HalfDimensions / 2) % OutputChunkSize == 0); - constexpr IndexType NumOutputChunks = HalfDimensions / 2 / OutputChunkSize; - - [[maybe_unused]] const vec_t Zero = vec_zero(); - [[maybe_unused]] const vec_t FtMax = vec_set_16(FtMaxVal); - [[maybe_unused]] constexpr int shift = 7; - - const vec_t* in0 = reinterpret_cast(&(accumulation[perspectives[p]][0])); - const vec_t* in1 = - reinterpret_cast(&(accumulation[perspectives[p]][HalfDimensions / 2])); - vec_t* out = reinterpret_cast(output + offset); - - // Per the NNUE architecture, here we want to multiply pairs of - // clipped elements and divide the product by 128. To do this, - // we can naively perform min/max operation to clip each of the - // four int16 vectors, mullo pairs together, then pack them into - // one int8 vector. However, there exists a faster way. - - // The idea here is to use the implicit clipping from packus to - // save us two vec_max_16 instructions. This clipping works due - // to the fact that any int16 integer below zero will be zeroed - // on packus. - - // Consider the case where the second element is negative. - // If we do standard clipping, that element will be zero, which - // means our pairwise product is zero. If we perform packus and - // remove the lower-side clip for the second element, then our - // product before packus will be negative, and is zeroed on pack. - // The two operations produce equivalent results, but the second - // one (using packus) saves one max operation per pair. - - // But here we run into a problem: mullo does not preserve the - // sign of the multiplication. We can get around this by doing - // mulhi, which keeps the sign. But that requires an additional - // tweak. - - // mulhi cuts off the last 16 bits of the resulting product, - // which is the same as performing a rightward shift of 16 bits. - // We can use this to our advantage. Recall that we want to - // divide the final product by 128, which is equivalent to a - // 7-bit right shift. Intuitively, if we shift the clipped - // value left by 9, and perform mulhi, which shifts the product - // right by 16 bits, then we will net a right shift of 7 bits. - // However, this won't work as intended. Since we clip the - // values to have a maximum value of 127, shifting it by 9 bits - // might occupy the signed bit, resulting in some positive - // values being interpreted as negative after the shift. - - // There is a way, however, to get around this limitation. When - // loading the network, scale accumulator weights and biases by - // 2. To get the same pairwise multiplication result as before, - // we need to divide the product by 128 * 2 * 2 = 512, which - // amounts to a right shift of 9 bits. So now we only have to - // shift left by 7 bits, perform mulhi (shifts right by 16 bits) - // and net a 9 bit right shift. Since we scaled everything by - // two, the values are clipped at 127 * 2 = 254, which occupies - // 8 bits. Shifting it by 7 bits left will no longer occupy the - // signed bit, so we are safe. - - for (IndexType j = 0; j < NumOutputChunks; j += 2) + [[maybe_unused]] auto cursor = nnzInfo.make_cursor(perspective); + + constexpr IndexType OutputChunkSize = MaxChunkSize; + static_assert((HalfDimensions / 2) % OutputChunkSize == 0); + constexpr IndexType NumOutputChunks = HalfDimensions / 2 / OutputChunkSize; + + [[maybe_unused]] const vec_t Zero = vec_zero(); + [[maybe_unused]] const vec_t FtMax = vec_set_16(FtMaxVal); + [[maybe_unused]] constexpr int shift = 7; + + const vec_t* in0 = reinterpret_cast(&accumulation[0]); + const vec_t* in1 = reinterpret_cast(&accumulation[HalfDimensions / 2]); + vec_t* out = reinterpret_cast(output + offset); + + // Per the NNUE architecture, here we want to multiply pairs of + // clipped elements and divide the product by 128. To do this, + // we can naively perform min/max operation to clip each of the + // four int16 vectors, mullo pairs together, then pack them into + // one int8 vector. However, there exists a faster way. + + // The idea here is to use the implicit clipping from packus to + // save us two vec_max_16 instructions. This clipping works due + // to the fact that any int16 integer below zero will be zeroed + // on packus. + + // Consider the case where the second element is negative. + // If we do standard clipping, that element will be zero, which + // means our pairwise product is zero. If we perform packus and + // remove the lower-side clip for the second element, then our + // product before packus will be negative, and is zeroed on pack. + // The two operations produce equivalent results, but the second + // one (using packus) saves one max operation per pair. + + // But here we run into a problem: mullo does not preserve the + // sign of the multiplication. We can get around this by doing + // mulhi, which keeps the sign. But that requires an additional + // tweak. + + // mulhi cuts off the last 16 bits of the resulting product, + // which is the same as performing a rightward shift of 16 bits. + // We can use this to our advantage. Recall that we want to + // divide the final product by 128, which is equivalent to a + // 7-bit right shift. Intuitively, if we shift the clipped + // value left by 9, and perform mulhi, which shifts the product + // right by 16 bits, then we will net a right shift of 7 bits. + // However, this won't work as intended. Since we clip the + // values to have a maximum value of 127, shifting it by 9 bits + // might occupy the signed bit, resulting in some positive + // values being interpreted as negative after the shift. + + // There is a way, however, to get around this limitation. When + // loading the network, scale accumulator weights and biases by + // 2. To get the same pairwise multiplication result as before, + // we need to divide the product by 128 * 2 * 2 = 512, which + // amounts to a right shift of 9 bits. So now we only have to + // shift left by 7 bits, perform mulhi (shifts right by 16 bits) + // and net a 9 bit right shift. Since we scaled everything by + // two, the values are clipped at 127 * 2 = 254, which occupies + // 8 bits. Shifting it by 7 bits left will no longer occupy the + // signed bit, so we are safe. + + for (IndexType j = 0; j < NumOutputChunks; j += 2) + { + vec_t packed[2]; + for (IndexType k = 0; k < 2; ++k) { - vec_t packed[2]; - for (IndexType k = 0; k < 2; ++k) - { - const IndexType i = (j + k) * 2; + const IndexType i = (j + k) * 2; - vec_t acc0a = in0[i + 0]; - vec_t acc0b = in0[i + 1]; - vec_t acc1a = in1[i + 0]; - vec_t acc1b = in1[i + 1]; + vec_t acc0a = in0[i + 0]; + vec_t acc0b = in0[i + 1]; + vec_t acc1a = in1[i + 0]; + vec_t acc1b = in1[i + 1]; - static_assert(FtMaxVal == 255); + static_assert(FtMaxVal == 255); #if defined(USE_NEON) - uint16x8_t mul0 = vmull_u8(vqmovun_s16(acc0a), vqmovun_s16(acc1a)); - uint16x8_t mul1 = vmull_u8(vqmovun_s16(acc0b), vqmovun_s16(acc1b)); + uint16x8_t mul0 = vmull_u8(vqmovun_s16(acc0a), vqmovun_s16(acc1a)); + uint16x8_t mul1 = vmull_u8(vqmovun_s16(acc0b), vqmovun_s16(acc1b)); - uint8x16x2_t uzp = - vuzpq_u8(vreinterpretq_u8_u16(mul0), vreinterpretq_u8_u16(mul1)); - uint8x16_t pab = vshrq_n_u8(uzp.val[1], 1); - vec_t result = reinterpret_cast(pab); + uint8x16x2_t uzp = vuzpq_u8(vreinterpretq_u8_u16(mul0), vreinterpretq_u8_u16(mul1)); + uint8x16_t pab = vshrq_n_u8(uzp.val[1], 1); + vec_t result = reinterpret_cast(pab); #elif defined(USE_LSX) || defined(USE_LASX) - vec_t pa = vec_packus_16(acc0a, acc0b); - vec_t pb = vec_packus_16(acc1a, acc1b); + vec_t pa = vec_packus_16(acc0a, acc0b); + vec_t pb = vec_packus_16(acc1a, acc1b); - vec_t hi = vec_mulhi_8(pa, pb); - vec_t result = vec_srli_8(hi, 1); + vec_t hi = vec_mulhi_8(pa, pb); + vec_t result = vec_srli_8(hi, 1); #elif defined(__wasm__) - // _mm_mulhi_epi16 is lowered to 32-bit multiplies, so we take - // a similar approach as the NEON path. - vec_t mul0 = vec_packus_16(acc0a, acc0b); - vec_t mul1 = vec_packus_16(acc1a, acc1b); - - vec_t low = wasm_u16x8_extmul_low_u8x16(mul0, mul1); - vec_t hi = wasm_u16x8_extmul_high_u8x16(mul0, mul1); - - // equivalent to vuzp2_u8 - vec_t merged = wasm_i8x16_shuffle(low, hi, 1, 3, 5, 7, 9, 11, 13, 15, 17, 19, - 21, 23, 25, 27, 29, 31); - vec_t result = wasm_u8x16_shr(merged, 1); + // _mm_mulhi_epi16 is lowered to 32-bit multiplies, so we take + // a similar approach as the NEON path. + vec_t mul0 = vec_packus_16(acc0a, acc0b); + vec_t mul1 = vec_packus_16(acc1a, acc1b); + + vec_t low = wasm_u16x8_extmul_low_u8x16(mul0, mul1); + vec_t hi = wasm_u16x8_extmul_high_u8x16(mul0, mul1); + + // equivalent to vuzp2_u8 + vec_t merged = wasm_i8x16_shuffle(low, hi, 1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, + 23, 25, 27, 29, 31); + vec_t result = wasm_u8x16_shr(merged, 1); #else - vec_t sum0a = vec_slli_16(vec_max_16(vec_min_16(acc0a, FtMax), Zero), shift); - vec_t sum0b = vec_slli_16(vec_max_16(vec_min_16(acc0b, FtMax), Zero), shift); - vec_t sum1a = vec_min_16(acc1a, FtMax); - vec_t sum1b = vec_min_16(acc1b, FtMax); + vec_t sum0a = vec_slli_16(vec_max_16(vec_min_16(acc0a, FtMax), Zero), shift); + vec_t sum0b = vec_slli_16(vec_max_16(vec_min_16(acc0b, FtMax), Zero), shift); + vec_t sum1a = vec_min_16(acc1a, FtMax); + vec_t sum1b = vec_min_16(acc1b, FtMax); - vec_t pa = vec_mulhi_16(sum0a, sum1a); - vec_t pb = vec_mulhi_16(sum0b, sum1b); + vec_t pa = vec_mulhi_16(sum0a, sum1a); + vec_t pb = vec_mulhi_16(sum0b, sum1b); - vec_t result = vec_packus_16(pa, pb); + vec_t result = vec_packus_16(pa, pb); #endif - packed[k] = out[j + k] = result; - } - - cursor.record2(packed[0], packed[1]); + packed[k] = out[j + k] = result; } + cursor.record2(packed[0], packed[1]); + } + #elif defined(USE_RVV) - usize j = 0; - usize VL = __riscv_vsetvlmax_e8m1(); - vuint8m1_t vid8; - vuint16m2_t vid16; - if (VL <= 256) - vid8 = __riscv_vid_v_u8m1(VL); - else - vid16 = __riscv_vid_v_u16m2(VL); - const auto& accp = accumulation[perspectives[p]]; + usize j = 0; + usize VL = __riscv_vsetvlmax_e8m1(); + vuint8m1_t vid8; + vuint16m2_t vid16; + if (VL <= 256) + vid8 = __riscv_vid_v_u8m1(VL); + else + vid16 = __riscv_vid_v_u16m2(VL); - for (usize vl; j < HalfDimensions / 2; j += vl) - { - vl = __riscv_vsetvl_e16m2(HalfDimensions / 2 - j); + for (usize vl; j < HalfDimensions / 2; j += vl) + { + vl = __riscv_vsetvl_e16m2(HalfDimensions / 2 - j); - vint16m2_t acc0 = __riscv_vle16_v_i16m2(&accp[j], vl); - vint16m2_t acc1 = __riscv_vle16_v_i16m2(&accp[j + HalfDimensions / 2], vl); + vint16m2_t acc0 = __riscv_vle16_v_i16m2(&accumulation[j], vl); + vint16m2_t acc1 = __riscv_vle16_v_i16m2(&accumulation[j + HalfDimensions / 2], vl); - acc0 = __riscv_vmax(acc0, 0, vl); - acc1 = __riscv_vmax(acc1, 0, vl); + acc0 = __riscv_vmax(acc0, 0, vl); + acc1 = __riscv_vmax(acc1, 0, vl); - vuint8m1_t pa = __riscv_vnclipu(__riscv_vreinterpret_u16m2(acc0), 0, 0, vl); - vuint8m1_t pb = __riscv_vnclipu(__riscv_vreinterpret_u16m2(acc1), 0, 0, vl); + vuint8m1_t pa = __riscv_vnclipu(__riscv_vreinterpret_u16m2(acc0), 0, 0, vl); + vuint8m1_t pb = __riscv_vnclipu(__riscv_vreinterpret_u16m2(acc1), 0, 0, vl); - vuint8m1_t hi = __riscv_vmulhu(pa, pb, vl); - vuint8m1_t result = __riscv_vsrl(hi, 1, vl); + vuint8m1_t hi = __riscv_vmulhu(pa, pb, vl); + vuint8m1_t result = __riscv_vsrl(hi, 1, vl); - __riscv_vse8(&output[offset + j], result, vl); + __riscv_vse8(&output[offset + j], result, vl); - vbool8_t m = __riscv_vmsne(result, 0, vl); - usize cnt = __riscv_vcpop(m, vl); - vuint16m2_t vidx; - if (VL <= 256) - vidx = __riscv_vzext_vf2(__riscv_vcompress(vid8, m, vl), cnt); - else - vidx = __riscv_vcompress(vid16, m, vl); - __riscv_vse16(&nnzInfo.nnz[nnzInfo.count], __riscv_vadd(vidx, offset + j, cnt), - cnt); - nnzInfo.count += cnt; - } + vbool8_t m = __riscv_vmsne(result, 0, vl); + usize cnt = __riscv_vcpop(m, vl); + vuint16m2_t vidx; + if (VL <= 256) + vidx = __riscv_vzext_vf2(__riscv_vcompress(vid8, m, vl), cnt); + else + vidx = __riscv_vcompress(vid16, m, vl); + __riscv_vse16(&nnzInfo.nnz[nnzInfo.count], __riscv_vadd(vidx, offset + j, cnt), cnt); + nnzInfo.count += cnt; + } #else - for (IndexType j = 0; j < HalfDimensions / 2; ++j) - { - BiasType sum0 = accumulation[static_cast(perspectives[p])][j + 0]; - BiasType sum1 = - accumulation[static_cast(perspectives[p])][j + HalfDimensions / 2]; - - sum0 = std::clamp(sum0, 0, FtMaxVal); - sum1 = std::clamp(sum1, 0, FtMaxVal); + for (IndexType j = 0; j < HalfDimensions / 2; ++j) + { + BiasType sum0 = accumulation[j]; + BiasType sum1 = accumulation[j + HalfDimensions / 2]; - output[offset + j] = static_cast(unsigned(sum0 * sum1) / 512); - } + sum0 = std::clamp(sum0, 0, FtMaxVal); + sum1 = std::clamp(sum1, 0, FtMaxVal); -#endif + output[offset + j] = static_cast(unsigned(sum0 * sum1) / 512); } - return psqt; - } // end of function transform() +#endif + } + public: alignas(CacheLineSize) BiasesArray biases; alignas(CacheLineSize) WeightArray weights;