Skip to content

Extract perspective transform logic into a separate function #7034

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Disservin wants to merge 2 commits into
base: master
Choose a base branch
from
+156 −154
Open

Extract perspective transform logic into a separate function #7034

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0a326e0
Cleanup nnue_feature_transformer.h
Disservin Aug 4, 2026
8501f66
Merge branch 'master' into cleanup-nnue_feature_transformer.h
Disservin Aug 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
310 changes: 156 additions & 154 deletions src/nnue/nnue_feature_transformer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,8 +226,6 @@ class FeatureTransformer {
OutputType* output,
int bucket,
[[maybe_unused]] NNZInfo<OutputDimensions>& nnzInfo) const {

using namespace SIMD;
accumulatorStack.evaluate(pos, *this, cache);
const auto& accumulatorState = accumulatorStack.latest();

Expand All @@ -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<i16, HalfDimensions>& accumulation,
OutputType* output,
IndexType perspective,
[[maybe_unused]] NNZInfo<OutputDimensions>& 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<const vec_t*>(&(accumulation[perspectives[p]][0]));
const vec_t* in1 =
reinterpret_cast<const vec_t*>(&(accumulation[perspectives[p]][HalfDimensions / 2]));
vec_t* out = reinterpret_cast<vec_t*>(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<const vec_t*>(&accumulation[0]);
const vec_t* in1 = reinterpret_cast<const vec_t*>(&accumulation[HalfDimensions / 2]);
vec_t* out = reinterpret_cast<vec_t*>(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<vec_t>(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<vec_t>(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<int>(perspectives[p])][j + 0];
BiasType sum1 =
accumulation[static_cast<int>(perspectives[p])][j + HalfDimensions / 2];

sum0 = std::clamp<BiasType>(sum0, 0, FtMaxVal);
sum1 = std::clamp<BiasType>(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<OutputType>(unsigned(sum0 * sum1) / 512);
}
sum0 = std::clamp<BiasType>(sum0, 0, FtMaxVal);
sum1 = std::clamp<BiasType>(sum1, 0, FtMaxVal);

#endif
output[offset + j] = static_cast<OutputType>(unsigned(sum0 * sum1) / 512);
}

return psqt;
} // end of function transform()
#endif
}

public:
alignas(CacheLineSize) BiasesArray biases;
alignas(CacheLineSize) WeightArray weights;

Expand Down
Loading