Skip to content

Compress threat indexing, update feature set name to Full Threats v2 #398

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
sscg13 wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Compress threat indexing, update feature set name to Full Threats v2 #398

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
fba3059
compress indices, match SF
sscg13 Feb 24, 2026
60b256e
use right feature set in data loader bench
sscg13 Feb 24, 2026
6e39609
we hate resolving merge conflicts
sscg13 Feb 25, 2026
34aa5ca
update modules
sscg13 Feb 28, 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
4 changes: 2 additions & 2 deletions model/modules/features/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,15 @@
from dataclasses import dataclass

from .halfka_v2_hm import HalfKav2Hm
from .full_threats import FullThreats
from .full_threatsv2 import FullThreatsv2

import tyro
from typing import Annotated


_FEATURES: dict[str, type] = {
"HalfKAv2_hm^": HalfKav2Hm,
"Full_Threats^": FullThreats,
"Full_Threatsv2^": FullThreatsv2,
}


Expand Down
21 changes: 11 additions & 10 deletions model/modules/features/full_threats.py → model/modules/features/full_threatsv2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,24 +7,24 @@
from .halfka_v2_hm import InverseKingBuckets, _halfka_idx


class FullThreats(DoubleFeatureTransformer):
HASH = 0x8F234CB8
FEATURE_NAME = "Full_Threats^"
INPUT_FEATURE_NAME = "Full_Threats"
class FullThreatsv2(DoubleFeatureTransformer):
HASH = 0x8F234CB9
FEATURE_NAME = "Full_Threatsv2^"
INPUT_FEATURE_NAME = "Full_Threatsv2"
MAX_ACTIVE_FEATURES = 128 + 32

NUM_SQ = 64
NUM_PT = 12
NUM_PLANES = NUM_SQ * NUM_PT # 768
NUM_BUCKETS = NUM_SQ // 2 # 32

NUM_THREAT_FEATURES = 60144
NUM_THREAT_FEATURES = 53564
NUM_PSQ_FEATURES = NUM_PLANES * NUM_BUCKETS # 24,576
NUM_INPUTS = NUM_THREAT_FEATURES + NUM_PSQ_FEATURES # 84,720
NUM_INPUTS = NUM_THREAT_FEATURES + NUM_PSQ_FEATURES # 78,140
NUM_INPUTS_VIRTUAL = NUM_PLANES # 768

# Export size: threats + 11-piece-type PSQ (704 * 32 = 22,528)
NUM_REAL_FEATURES = NUM_THREAT_FEATURES + 704 * 32 # 82,672
NUM_REAL_FEATURES = NUM_THREAT_FEATURES + 704 * 32 # 76,092

def __init__(self, num_outputs: int):
super().__init__(self.NUM_INPUTS, num_outputs)
Expand Down Expand Up @@ -186,17 +186,18 @@ def halfka_psqts() -> list[int]:
chess.QUEEN: 2538,
}

num_threats = 53564
num_psq = 768 * 32 # 24,576
num_total = 60_144 + num_psq
num_total = num_threats + num_psq
values = [0] * num_total

for ksq in range(64):
for s in range(64):
for pt, val in piece_values.items():
idxw = 60_144 + _halfka_idx(
idxw = num_threats + _halfka_idx(
True, ksq, s, chess.Piece(pt, chess.WHITE)
)
idxb = 60_144 + _halfka_idx(
idxb = num_threats + _halfka_idx(
True, ksq, s, chess.Piece(pt, chess.BLACK)
)
values[idxw] = val
Expand Down
2 changes: 1 addition & 1 deletion model/utils/serialize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from ..config import ModelConfig
from ..model import NNUEModel
from ..modules import BaseFeatureTransformer, get_feature_cls
from ..modules.features.full_threats import FullThreats
from ..modules.features.full_threatsv2 import FullThreatsv2
from ..quantize import QuantizationConfig


Expand Down
47 changes: 31 additions & 16 deletions training_data_loader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ struct HalfKAv2_hm {
}
};

constexpr int numvalidtargets[12] = {6, 6, 10, 10, 8, 8, 8, 8, 10, 10, 0, 0};
constexpr int numvalidtargets[12] = {5, 6, 9, 9, 7, 7, 7, 7, 9, 9, 0, 0};

using ThreatOffsetTable = std::array<std::array<int, 66>, 12>;

Expand Down Expand Up @@ -143,10 +143,10 @@ constexpr auto threatfeaturecalc = []() {

constexpr ThreatOffsetTable threatoffsets = threatfeaturecalc.table;
constexpr int threatfeatures = threatfeaturecalc.totalfeatures;
static_assert(threatfeatures == 60144);
static_assert(threatfeatures == 53564);

struct Full_Threats {
static constexpr std::string_view NAME = "Full_Threats";
struct Full_Threatsv2 {
static constexpr std::string_view NAME = "Full_Threatsv2";

static constexpr int SQUARE_NB = 64;
static constexpr int PIECE_NB = 12;
Expand Down Expand Up @@ -175,11 +175,11 @@ struct Full_Threats {
};

static constexpr int map[PIECE_TYPE_NB][PIECE_TYPE_NB] = {
{0, 1, -1, 2, -1, -1},
{0, 1, 2, 3, 4, -1},
{0, 1, 2, 3, -1, -1},
{0, 1, 2, 3, -1, -1},
{0, 1, 2, 3, 4, -1},
{ 0, 1, -1, 2, -1, -1},
{ 0, 1, 2, 3, 4, -1},
{ 0, 1, 2, 3, -1, -1},
{ 0, 1, 2, 3, -1, -1},
{ 0, 1, 2, 3, 4, -1},
{-1, -1, -1, -1, -1, -1}
};

Expand All @@ -196,7 +196,7 @@ struct Full_Threats {
// clang-format on

static constexpr int NUM_SQ = 64;
static constexpr int NUM_PT = 12;
static constexpr int NUM_PT = 11;
static constexpr int NUM_PLANES = NUM_SQ * NUM_PT;

static constexpr int NUM_THREAT_FEATURES = threatfeatures;
Expand Down Expand Up @@ -226,14 +226,29 @@ struct Full_Threats {
{
return -1;
}

if (enemy && (attkr.type() == attkd.type()) && (attkr.type() != PieceType::Pawn))
{
attkd = attkr;
std::swap(from, to);
}

int targetindex = int((attkr == attkd) ? Color::White : attkd.color())
* ((1 + numvalidtargets[(int) attkr])/ 2) + map[(int) attkr.type()][(int) attkd.type()];

if ((8 * (int)attkd.color() + (int)attkd.type()) > (8 + (int)attkr.type())
&& attkr != Piece(PieceType::Pawn, Color::Black))
{
targetindex--;
}


Bitboard attacks = (attkr.type() == PieceType::Pawn)
? bb::pawnAttacks(Bitboard::square(Square(from)), attkr.color())
: bb::detail::pseudoAttacks()[attkr.type()][Square(from)];
Bitboard upto = Bitboard::square(to);

return int(threatoffsets[(int) attkr][65]
+ (int(attkd.color()) * (numvalidtargets[(int) attkr] / 2)
+ map[(int) attkr.type()][(int) attkd.type()])
* threatoffsets[(int) attkr][64]
+ targetindex * threatoffsets[(int) attkr][64]
+ threatoffsets[(int) attkr][(int) from]
+ (Bitboard::fromBits((1ULL << (int) to) - 1) & attacks).count());
}
Expand Down Expand Up @@ -357,7 +372,7 @@ auto find_feature(std::string_view name) {

auto get_feature(std::string_view name) {
return find_feature<HalfKAv2_hm, //
Full_Threats //
Full_Threatsv2 //
>(name);
}

Expand Down Expand Up @@ -1081,7 +1096,7 @@ int main(int argc, char** argv) {
.pc_y2 = 2.0,
.pc_y3 = 1.0};
const DataloaderDDPConfig ddp_config = {.rank = 0, .world_size = 1};
auto stream = create_sparse_batch_stream("Full_Threats", concurrency, file_count, files,
auto stream = create_sparse_batch_stream("Full_Threatsv2", concurrency, file_count, files,
batch_size, cyclic, config, ddp_config);

auto t0 = std::chrono::high_resolution_clock::now();
Expand Down