Writers.hpp

Go to the documentation of this file.

#pragma once
#include "ExpertRouter.hpp"     // For ExpertWeights struct
#include "FeatureExtractor.hpp" // For ChessInput struct
#include "WDLNormalizer.hpp"    // For WDLOutput struct
#include <array>
#include <cstdint>
#include <fstream>
#include <iostream>
#include <memory>
#include <string>
#include <vector>
 
struct alignas(64) PackedChessInput {
  // Planes 0-29 are categorical (pieces, moves, threats). 
  // Pack each 64-square plane into a SINGLE 64-bit integer.
  // Size: 30 * 8 bytes = 240 bytes
  uint64_t bitboards[30]; 
 
  // Planes 30-31 are continuous (King distance), keep as uint8_t
  // Size: 2 * 64 bytes = 128 bytes
  uint8_t continuous_layers[2][64]; 
 
  // Global features
  // Size: 16 * 4 bytes = 64 bytes
  float global[16]; 
};
 
class DatasetWriter {
private:
  std::string filename;
  size_t batch_size;
 
  // Buffer stores the packed structs to save memory
  std::vector<PackedChessInput> buffer;
 
public:
  DatasetWriter(const std::string &fname, size_t batch_capacity = 1000);
  ~DatasetWriter();
  void add(const ChessInput &features);
  void flush();
};
 
class LabelWriter {
private:
  std::string filename;
  std::vector<float> buffer;
  size_t batch_size;
 
public:
  LabelWriter(const std::string &fname, size_t batch_capacity = 1000);
  ~LabelWriter();
  void add(float label);
  void flush();
};

class WDLWriter {
private:
  std::string filename;
  std::vector<float> buffer; // Flattened [w, d, l, m, w, d, l, m, ...]
  size_t batch_size;
 
public:
  WDLWriter(const std::string &fname, size_t batch_capacity = 1000);
  ~WDLWriter();

  void add(const WDLOutput &wdl);
 
  void flush();
};

class ExpertWeightsWriter {
private:
  std::string filename;
  std::vector<float> buffer; // Flat buffer of all weights
  size_t batch_size;
  static constexpr int NUM_EXPERTS = ExpertRouter::NUM_EXPERTS;
 
public:
  ExpertWeightsWriter(const std::string &fname, size_t batch_capacity = 1000);
  ~ExpertWeightsWriter();

  void add(const ExpertRouter::ExpertWeights &weights);

  void add(const float *weights);
 
  void flush();
};

class ExpertDatasetWriter {
private:
  static constexpr int NUM_EXPERTS = 6;
  static constexpr float AUX_GATE_THRESHOLD = 0.1f;
 
  std::array<std::unique_ptr<DatasetWriter>, NUM_EXPERTS> feature_writers;
  std::array<std::unique_ptr<WDLWriter>, NUM_EXPERTS> label_writers;
  std::array<size_t, NUM_EXPERTS> counts;
 
public:
  ExpertDatasetWriter(const std::string &output_dir,
                      size_t batch_capacity = 10000);
  ~ExpertDatasetWriter();

  void add(const ChessInput &features, const WDLOutput &wdl,
           const ExpertRouter::ExpertWeights &weights);

  void flush();

  void print_stats() const;
};