249 lines
11 KiB
C
249 lines
11 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* ternary_optim.c — SR4 4-bit stochastic-rounded latent-accumulator optimizer.
|
|
* See ternary_optim.h for the public API and docs/HOW_IT_WORKS.md for the math.
|
|
*
|
|
* The core (init distribution, ternarize, per-matrix step normalization, 8.8
|
|
* fixed-point stochastic rounding, stratified Weyl low-discrepancy dither,
|
|
* SplitMix64 counter-hash, saturating add) is lifted BIT-FOR-BIT from the
|
|
* SR4 research prototype's `sr_*` core, decoupled from MNIST
|
|
* and generalized from the fixed 784-256-10 MLP to an arbitrary (rows x cols)
|
|
* matrix. No MNIST I/O, no fixed dims, no training loop.
|
|
*
|
|
* All stochastic-rounding arithmetic is pure integer (flag-immune): the numerics
|
|
* are bit-identical across -O levels and -ffast-math.
|
|
*/
|
|
|
|
#include "ternary_optim.h"
|
|
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
|
|
/* Stratified low-discrepancy dither constants (research prototype, verbatim).
|
|
LD_STRATUM_BITS: log2 of the stratum length in batches — within one stratum
|
|
each weight's dither walks the full Weyl orbit (h_i + LD_STRIDE*j) mod 256,
|
|
hitting every residue exactly once, so a constant frac f fires exactly f times
|
|
per stratum. LD_STRIDE must be odd (bijective mod 256); 159 ~ 256/phi is the
|
|
near-golden rotation (best equidistribution at every partial-window length). */
|
|
#define TO_LD_STRATUM_BITS 8
|
|
#define TO_LD_STRIDE 159u
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Opaque state
|
|
* ------------------------------------------------------------------------- */
|
|
struct ternary_optim {
|
|
int8_t *q; /* the ENTIRE per-weight state: int4 latent, stored
|
|
int8 clamped [-q_clamp, +q_clamp] (bring-up) */
|
|
int rows, cols, numel;
|
|
int64_t idx_base; /* global weight index base for the SR hash */
|
|
ternary_optim_config cfg;
|
|
};
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Internal xorshift RNG — used ONLY for the initial q distribution (the update
|
|
* path is stateless/counter-hash). Identical algorithm to the research prototype's sm_rand.
|
|
* ------------------------------------------------------------------------- */
|
|
static void to_sm_seed(uint64_t *state, uint64_t seed) { *state = seed ? seed : 1; }
|
|
static float to_sm_rand(uint64_t *state) {
|
|
uint64_t x = *state; x ^= x << 13; x ^= x >> 7; x ^= x << 17; *state = x;
|
|
return (float)(x >> 11) / (float)(1ULL << 53);
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* SplitMix64-style counter-based hash: pure integer, order-independent, derived
|
|
* from (seed, counter, global weight index) — deterministic per seed and
|
|
* flag-immune. Nothing is stored. (research prototype, verbatim.)
|
|
* ------------------------------------------------------------------------- */
|
|
static inline uint64_t sm64_mix(uint64_t x) {
|
|
x += 0x9E3779B97F4A7C15ULL;
|
|
x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL;
|
|
x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL;
|
|
return x ^ (x >> 31);
|
|
}
|
|
static inline uint64_t sr_hash(uint64_t seed, uint64_t batch_counter,
|
|
uint64_t global_idx) {
|
|
uint64_t x = sm64_mix(seed ^ 0xD1B54A32D192ED03ULL);
|
|
x = sm64_mix(x ^ batch_counter);
|
|
x = sm64_mix(x ^ global_idx);
|
|
return x;
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Config
|
|
* ------------------------------------------------------------------------- */
|
|
void ternary_optim_config_default(ternary_optim_config *cfg) {
|
|
if (!cfg) return;
|
|
cfg->target_step = 0.001f;
|
|
cfg->q_clamp = 7;
|
|
cfg->theta = 2;
|
|
cfg->use_dither = 1;
|
|
cfg->seed = 42;
|
|
cfg->idx_base = 0;
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Lifecycle
|
|
* ------------------------------------------------------------------------- */
|
|
ternary_optim *ternary_optim_create(int rows, int cols,
|
|
const ternary_optim_config *cfg_in) {
|
|
if (rows <= 0 || cols <= 0) return NULL;
|
|
/* guard against numel overflow */
|
|
if ((int64_t)rows * (int64_t)cols > (int64_t)0x7FFFFFFF) return NULL;
|
|
|
|
ternary_optim *opt = calloc(1, sizeof(*opt));
|
|
if (!opt) return NULL;
|
|
|
|
if (cfg_in) opt->cfg = *cfg_in;
|
|
else ternary_optim_config_default(&opt->cfg);
|
|
/* sanity floors — keep the mechanism well-defined */
|
|
if (opt->cfg.q_clamp < 1) opt->cfg.q_clamp = 1;
|
|
if (opt->cfg.theta < 1) opt->cfg.theta = 1;
|
|
|
|
opt->rows = rows; opt->cols = cols; opt->numel = rows * cols;
|
|
opt->idx_base = opt->cfg.idx_base;
|
|
opt->q = malloc((size_t)opt->numel);
|
|
if (!opt->q) { free(opt); return NULL; }
|
|
|
|
/* Init q in {-3, 0, +3} at {33.5%, 33%, 33.5%} — matches the reference
|
|
start stats (67% active, signs uniform); magnitude 3 = theta+1 gives
|
|
every active weight one LSB of dead-zone margin. Seeded from
|
|
(seed ^ idx_base) so distinct matrices get distinct inits. */
|
|
uint64_t rng; to_sm_seed(&rng, opt->cfg.seed ^ (uint64_t)opt->idx_base);
|
|
for (int i = 0; i < opt->numel; i++) {
|
|
float r = to_sm_rand(&rng);
|
|
opt->q[i] = (r < 0.335f) ? -3 : (r < 0.665f) ? 0 : 3;
|
|
}
|
|
return opt;
|
|
}
|
|
|
|
void ternary_optim_free(ternary_optim *opt) {
|
|
if (!opt) return;
|
|
free(opt->q);
|
|
free(opt);
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Forward: ternarize the accumulator. w = 0 if |q| < theta else sign(q).
|
|
* ------------------------------------------------------------------------- */
|
|
void ternary_optim_materialize(const ternary_optim *opt, float *out) {
|
|
if (!opt || !out) return;
|
|
const int theta = opt->cfg.theta;
|
|
for (int i = 0; i < opt->numel; i++) {
|
|
int q = opt->q[i];
|
|
out[i] = (q >= theta) ? 1.0f : (q <= -theta) ? -1.0f : 0.0f;
|
|
}
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* SR4 update: per batch, per matrix. Saturating stochastic-rounded integer add
|
|
* into q. (research prototype sr4_update, generalized: q_clamp / theta / dither from cfg,
|
|
* idx_base / seed from the handle; numerics otherwise byte-identical.)
|
|
* ------------------------------------------------------------------------- */
|
|
void ternary_optim_step(ternary_optim *opt, const float *grad_sum,
|
|
int batch_size, uint64_t batch_counter,
|
|
float target_step) {
|
|
if (!opt || !grad_sum || batch_size <= 0) return;
|
|
|
|
const int q_clamp = opt->cfg.q_clamp;
|
|
const uint64_t seed_u = opt->cfg.seed;
|
|
if (target_step <= 0.0f) target_step = opt->cfg.target_step;
|
|
|
|
float inv_bs = 1.0f / (float)batch_size;
|
|
|
|
/* Stratified LD dither coordinates (pure functions of batch_counter —
|
|
nothing stored): the stratum index re-keys every weight's phase hash once
|
|
per 2^LD_STRATUM_BITS batches; the Weyl rotation LD_STRIDE*j advances the
|
|
shared orbit one near-golden step per batch. */
|
|
uint64_t stratum = batch_counter >> TO_LD_STRATUM_BITS;
|
|
uint32_t j = (uint32_t)(batch_counter & ((1u << TO_LD_STRATUM_BITS) - 1u));
|
|
uint32_t weyl = (TO_LD_STRIDE * j) & 255u;
|
|
|
|
/* Per-matrix step scale: s = target_step / (mean|g| + 1e-12) — continuous
|
|
scalar recomputed each batch from raw stats, zero persistent bits. */
|
|
double abs_sum = 0.0;
|
|
for (int i = 0; i < opt->numel; i++)
|
|
abs_sum += fabsf(grad_sum[i] * inv_bs);
|
|
float mean_abs = (float)(abs_sum / opt->numel);
|
|
float s = target_step / (mean_abs + 1e-12f);
|
|
|
|
for (int i = 0; i < opt->numel; i++) {
|
|
float g = grad_sum[i] * inv_bs;
|
|
/* Grid-snap to 8.8 fixed point (the only float->discrete boundary).
|
|
Defensive clamp at +/-16 LSB before lrintf: any |step| >= 15 LSB
|
|
saturates q from every starting value regardless of the SR bit, so
|
|
this is semantics-preserving and guards lrintf against outlier
|
|
g/mean|g| ratios. */
|
|
float prop = -g * s * 256.0f;
|
|
if (prop > 4096.0f) prop = 4096.0f;
|
|
if (prop < -4096.0f) prop = -4096.0f;
|
|
int32_t step_fx = (int32_t)lrintf(prop);
|
|
|
|
/* Deterministic stochastic rounding, LD-dithered: floor to integer LSBs
|
|
plus a Bernoulli(frac/256) increment. Arithmetic shift gives floor
|
|
semantics for negative step_fx (frac stays in [0,255]). E[k+sr_bit] =
|
|
step_fx/256 exactly; within a stratum j -> u is a bijection on Z_256,
|
|
so the realized fire frequency tracks f/256 with O(1) discrepancy. */
|
|
int32_t k = step_fx >> 8;
|
|
int32_t f = step_fx & 255;
|
|
uint64_t h = sr_hash(seed_u, stratum,
|
|
(uint64_t)(opt->idx_base + i));
|
|
uint32_t u;
|
|
if (opt->cfg.use_dither)
|
|
u = ((uint32_t)(h & 255) + weyl) & 255u; /* stratified Weyl */
|
|
else
|
|
u = (uint32_t)(h & 255); /* plain i.i.d. */
|
|
int32_t sr_bit = ((int32_t)u < f) ? 1 : 0;
|
|
|
|
/* Saturating add — the rails are legitimate fully-committed states,
|
|
never wraparound. */
|
|
int32_t nq = (int32_t)opt->q[i] + k + sr_bit;
|
|
if (nq > q_clamp) nq = q_clamp;
|
|
if (nq < -q_clamp) nq = -q_clamp;
|
|
opt->q[i] = (int8_t)nq;
|
|
}
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Accessors
|
|
* ------------------------------------------------------------------------- */
|
|
int ternary_optim_numel(const ternary_optim *opt) { return opt ? opt->numel : 0; }
|
|
int ternary_optim_rows(const ternary_optim *opt) { return opt ? opt->rows : 0; }
|
|
int ternary_optim_cols(const ternary_optim *opt) { return opt ? opt->cols : 0; }
|
|
const int8_t *ternary_optim_q(const ternary_optim *opt) { return opt ? opt->q : NULL; }
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Optional per-column scales (research prototype apply_col_scale / col_dot_sums, already
|
|
* general in (rows, cols)).
|
|
* ------------------------------------------------------------------------- */
|
|
void ternary_apply_col_scale(float *Z, const float *scale, int rows, int cols) {
|
|
if (!Z || !scale) return;
|
|
for (int r = 0; r < rows; r++)
|
|
for (int c = 0; c < cols; c++)
|
|
Z[r * cols + c] *= scale[c];
|
|
}
|
|
|
|
void ternary_col_scale_grad(const float *D, const float *Z_raw, float *out,
|
|
int rows, int cols) {
|
|
if (!D || !Z_raw || !out) return;
|
|
memset(out, 0, (size_t)cols * sizeof(float));
|
|
for (int r = 0; r < rows; r++)
|
|
for (int c = 0; c < cols; c++)
|
|
out[c] += D[r * cols + c] * Z_raw[r * cols + c];
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* State serialization — STUBBED. TODO: persist packed q (4 b/w) + shape + cfg
|
|
* for checkpoint/resume. See README.md "Deferred / stubbed".
|
|
* ------------------------------------------------------------------------- */
|
|
int ternary_optim_save(const ternary_optim *opt, const char *path) {
|
|
(void)opt; (void)path;
|
|
return TERNARY_OPTIM_ENOTIMPL;
|
|
}
|
|
int ternary_optim_load(const char *path, ternary_optim **out) {
|
|
(void)path;
|
|
if (out) *out = NULL;
|
|
return TERNARY_OPTIM_ENOTIMPL;
|
|
}
|