/* 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 #include #include /* 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; }