205 lines
8.0 KiB
C
205 lines
8.0 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* fd_byte_dataset — byte-level mmap'd next-byte-prediction loader.
|
|
* See fd_byte_dataset.h for the contract.
|
|
*/
|
|
#include "fd_byte_dataset.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <fcntl.h>
|
|
#include <unistd.h>
|
|
#include <errno.h>
|
|
#include <sys/mman.h>
|
|
#include <sys/stat.h>
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* splitmix64 — cheap, well-mixed PRNG for deriving round keys. */
|
|
/* ------------------------------------------------------------------ */
|
|
static inline uint64_t fd_splitmix64(uint64_t x) {
|
|
x += 0x9E3779B97F4A7C15ULL;
|
|
x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL;
|
|
x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL;
|
|
return x ^ (x >> 31);
|
|
}
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* Cycle-walking Feistel permutation over [0, n_windows). */
|
|
/* */
|
|
/* We build a balanced Feistel network over 2*hbits bits (a domain of */
|
|
/* size 2^(2*hbits) >= n_windows). Feistel with a bijective round is a */
|
|
/* bijection on that power-of-two domain; to get a bijection on the */
|
|
/* exact range [0, n) we "cycle-walk": re-encrypt until the output */
|
|
/* falls below n. This yields a true permutation of [0, n) for any n, */
|
|
/* with no O(n) shuffle table, and is fully determined by the keys. */
|
|
/* ------------------------------------------------------------------ */
|
|
|
|
#define FD_FEISTEL_ROUNDS 4
|
|
|
|
typedef struct {
|
|
uint32_t hbits; /* bits per half */
|
|
uint64_t hmask; /* (1<<hbits)-1 */
|
|
uint64_t domain; /* 1 << (2*hbits) */
|
|
uint64_t key[FD_FEISTEL_ROUNDS]; /* per-round keys */
|
|
} fd_feistel;
|
|
|
|
static void fd_feistel_init(fd_feistel *f, uint64_t n, uint64_t seed, uint64_t epoch) {
|
|
/* smallest even bit-width whose domain covers n (n>=1). */
|
|
uint32_t bits = 1;
|
|
while ((1ULL << bits) < n) bits++;
|
|
if (bits & 1u) bits++; /* make it even so halves are equal */
|
|
if (bits < 2) bits = 2;
|
|
f->hbits = bits / 2u;
|
|
f->hmask = (1ULL << f->hbits) - 1ULL;
|
|
f->domain = 1ULL << bits;
|
|
uint64_t s = fd_splitmix64(seed ^ 0xD1B54A32D192ED03ULL);
|
|
s = fd_splitmix64(s ^ epoch);
|
|
for (int r = 0; r < FD_FEISTEL_ROUNDS; r++) {
|
|
s = fd_splitmix64(s + 0x9E3779B97F4A7C15ULL);
|
|
f->key[r] = s;
|
|
}
|
|
}
|
|
|
|
/* One pass of the Feistel network over the power-of-two domain. */
|
|
static uint64_t fd_feistel_encrypt(const fd_feistel *f, uint64_t x) {
|
|
uint64_t L = (x >> f->hbits) & f->hmask;
|
|
uint64_t R = x & f->hmask;
|
|
for (int r = 0; r < FD_FEISTEL_ROUNDS; r++) {
|
|
uint64_t Fr = fd_splitmix64((R + f->key[r]) ^ (uint64_t)r) & f->hmask;
|
|
uint64_t nL = R;
|
|
uint64_t nR = L ^ Fr;
|
|
L = nL;
|
|
R = nR;
|
|
}
|
|
return (L << f->hbits) | R;
|
|
}
|
|
|
|
/* Permute x in [0, n) -> a distinct value in [0, n). */
|
|
static uint64_t fd_feistel_permute(const fd_feistel *f, uint64_t n, uint64_t x) {
|
|
uint64_t y = fd_feistel_encrypt(f, x);
|
|
/* cycle-walk: guaranteed to land in [0,n) within finitely many steps
|
|
* because the map is a bijection on the (>= n) power-of-two domain. */
|
|
while (y >= n) y = fd_feistel_encrypt(f, y);
|
|
return y;
|
|
}
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* Corpus */
|
|
/* ------------------------------------------------------------------ */
|
|
|
|
int fd_byte_corpus_open(fd_byte_corpus *c, const char *path) {
|
|
if (!c || !path) { errno = EINVAL; return -1; }
|
|
c->fd = -1; c->data = NULL; c->size = 0;
|
|
|
|
int fd = open(path, O_RDONLY);
|
|
if (fd < 0) return -1;
|
|
|
|
struct stat st;
|
|
if (fstat(fd, &st) != 0 || st.st_size <= 0) { close(fd); return -1; }
|
|
|
|
void *m = mmap(NULL, (size_t)st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
|
|
if (m == MAP_FAILED) { close(fd); return -1; }
|
|
|
|
/* Advise the kernel we will stream/random-access; harmless if unsupported. */
|
|
#ifdef MADV_RANDOM
|
|
madvise(m, (size_t)st.st_size, MADV_RANDOM);
|
|
#endif
|
|
|
|
c->fd = fd;
|
|
c->data = (const uint8_t *)m;
|
|
c->size = (size_t)st.st_size;
|
|
return 0;
|
|
}
|
|
|
|
void fd_byte_corpus_close(fd_byte_corpus *c) {
|
|
if (!c) return;
|
|
if (c->data && c->size) munmap((void *)c->data, c->size);
|
|
if (c->fd >= 0) close(c->fd);
|
|
c->fd = -1; c->data = NULL; c->size = 0;
|
|
}
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* Splits */
|
|
/* ------------------------------------------------------------------ */
|
|
|
|
int fd_byte_split_init(fd_byte_split *s, const fd_byte_corpus *c,
|
|
size_t offset, size_t length, size_t seq_len) {
|
|
if (!s || !c || !c->data) return -1;
|
|
if (seq_len == 0) return -1;
|
|
if (offset > c->size || length > c->size - offset) return -1; /* no overflow */
|
|
if (length <= seq_len) return -1; /* need seq_len+1 */
|
|
|
|
s->base = c->data + offset;
|
|
s->offset = offset;
|
|
s->length = length;
|
|
s->seq_len = seq_len;
|
|
s->n_windows = length - seq_len; /* last start = length-seq_len-1 */
|
|
return 0;
|
|
}
|
|
|
|
int fd_byte_split_enwik8_train(fd_byte_split *s, const fd_byte_corpus *c, size_t seq_len) {
|
|
if (!c || c->size != FD_ENWIK8_SIZE) return -1;
|
|
return fd_byte_split_init(s, c, FD_ENWIK8_TRAIN_OFF, FD_ENWIK8_TRAIN_LEN, seq_len);
|
|
}
|
|
int fd_byte_split_enwik8_val(fd_byte_split *s, const fd_byte_corpus *c, size_t seq_len) {
|
|
if (!c || c->size != FD_ENWIK8_SIZE) return -1;
|
|
return fd_byte_split_init(s, c, FD_ENWIK8_VAL_OFF, FD_ENWIK8_VAL_LEN, seq_len);
|
|
}
|
|
int fd_byte_split_enwik8_test(fd_byte_split *s, const fd_byte_corpus *c, size_t seq_len) {
|
|
if (!c || c->size != FD_ENWIK8_SIZE) return -1;
|
|
return fd_byte_split_init(s, c, FD_ENWIK8_TEST_OFF, FD_ENWIK8_TEST_LEN, seq_len);
|
|
}
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* Sampling */
|
|
/* ------------------------------------------------------------------ */
|
|
|
|
size_t fd_byte_window_pos(const fd_byte_split *s,
|
|
uint64_t seed, uint64_t epoch, uint64_t ordinal) {
|
|
if (seed == FD_BYTE_SEED_UNSET) {
|
|
fprintf(stderr,
|
|
"fd_byte_dataset: FATAL — window sampled with seed=0 (unseeded). "
|
|
"The reproducible A/B requires an explicit nonzero seed.\n");
|
|
abort();
|
|
}
|
|
fd_feistel f;
|
|
fd_feistel_init(&f, s->n_windows, seed, epoch);
|
|
return (size_t)fd_feistel_permute(&f, s->n_windows, ordinal % s->n_windows);
|
|
}
|
|
|
|
int fd_byte_get_batch(const fd_byte_split *s,
|
|
uint64_t seed, uint64_t epoch, uint64_t step,
|
|
uint32_t rank, uint32_t num_ranks, uint32_t batch,
|
|
int *inputs, int *targets) {
|
|
if (!s || !s->base || !inputs || !targets) return -1;
|
|
if (num_ranks == 0 || batch == 0 || rank >= num_ranks) return -1;
|
|
if (seed == FD_BYTE_SEED_UNSET) {
|
|
fprintf(stderr,
|
|
"fd_byte_dataset: FATAL — get_batch called with seed=0 (unseeded). "
|
|
"The reproducible A/B requires an explicit nonzero seed.\n");
|
|
abort();
|
|
}
|
|
|
|
/* One Feistel context for the whole (seed, epoch) permutation. */
|
|
fd_feistel f;
|
|
fd_feistel_init(&f, s->n_windows, seed, epoch);
|
|
|
|
const uint64_t seq = (uint64_t)s->seq_len;
|
|
const uint64_t nwin = (uint64_t)s->n_windows;
|
|
const uint64_t base_ord = (step * (uint64_t)num_ranks + (uint64_t)rank) * (uint64_t)batch;
|
|
|
|
for (uint32_t b = 0; b < batch; b++) {
|
|
uint64_t ordinal = (base_ord + b) % nwin;
|
|
uint64_t p = fd_feistel_permute(&f, nwin, ordinal); /* window start byte */
|
|
const uint8_t *w = s->base + p;
|
|
int *in = inputs + (size_t)b * s->seq_len;
|
|
int *tg = targets + (size_t)b * s->seq_len;
|
|
for (uint64_t j = 0; j < seq; j++) {
|
|
in[j] = (int)w[j];
|
|
tg[j] = (int)w[j + 1]; /* p+seq <= n_windows-1+seq = length-1: in range */
|
|
}
|
|
}
|
|
return 0;
|
|
}
|