From 31fc3176367d0c329c875affa6c15701255e8d23 Mon Sep 17 00:00:00 2001 From: w4ffl35 <25737761+w4ffl35@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:46:08 -0600 Subject: [PATCH] Add C training-IO: RNG, encoders, IDX loader, checkpoints Closes #27, closes #28, closes #29, closes #30. - src/rng: splitmix64 seeded RNG, portable C11, fixed test vectors. - src/encode: rate/latency (linear+log)/delta/random/repeat codings ported from SpikeEncoder/spikegen semantics; goldens from the stdlib-only tools/encode_golden.py, bit-exact in test_encode. - src/data: IDX MNIST reader with v/255 normalization matching the torch transform; tiny committed fixtures from tools/make_mnist_fixture.py. - src/checkpoint: JSON writer, version-stamped spikeforge-c-checkpoint v1 save/load, weight get/set accessors, and graph_summary export for the reverse bridge; tools/patch_nir_weights.py merges a checkpoint into a Python-side summary. Tests: 34/34 ctest pass, including test_rng, test_encode, test_data, test_checkpoint; ASan/UBSan build also 34/34. --- CMakeLists.txt | 4 + include/nir_c_runtime/checkpoint.h | 123 +++ include/nir_c_runtime/data.h | 47 ++ include/nir_c_runtime/encode.h | 87 +++ include/nir_c_runtime/rng.h | 30 + src/checkpoint.c | 1047 ++++++++++++++++++++++++++ src/data.c | 142 ++++ src/encode.c | 217 ++++++ src/rng.c | 23 + tests/CMakeLists.txt | 4 + tests/fixtures/encode_golden.json | 347 +++++++++ tests/fixtures/mnist_tiny_images.idx | Bin 0 -> 28 bytes tests/fixtures/mnist_tiny_labels.idx | Bin 0 -> 11 bytes tests/test_checkpoint.c | 268 +++++++ tests/test_data.c | 111 +++ tests/test_encode.c | 176 +++++ tests/test_rng.c | 57 ++ tools/encode_golden.py | 232 ++++++ tools/make_mnist_fixture.py | 45 ++ tools/patch_nir_weights.py | 102 +++ 20 files changed, 3062 insertions(+) create mode 100644 include/nir_c_runtime/checkpoint.h create mode 100644 include/nir_c_runtime/data.h create mode 100644 include/nir_c_runtime/encode.h create mode 100644 include/nir_c_runtime/rng.h create mode 100644 src/checkpoint.c create mode 100644 src/data.c create mode 100644 src/encode.c create mode 100644 src/rng.c create mode 100644 tests/fixtures/encode_golden.json create mode 100644 tests/fixtures/mnist_tiny_images.idx create mode 100644 tests/fixtures/mnist_tiny_labels.idx create mode 100644 tests/test_checkpoint.c create mode 100644 tests/test_data.c create mode 100644 tests/test_encode.c create mode 100644 tests/test_rng.c create mode 100644 tools/encode_golden.py create mode 100644 tools/make_mnist_fixture.py create mode 100644 tools/patch_nir_weights.py diff --git a/CMakeLists.txt b/CMakeLists.txt index a335e79..0b62d74 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -17,6 +17,10 @@ add_library(nir_c_runtime src/ops.c src/runtime.c src/fixed_event.c + src/rng.c + src/encode.c + src/data.c + src/checkpoint.c ) target_include_directories(nir_c_runtime PUBLIC include) target_link_libraries(nir_c_runtime PUBLIC m) diff --git a/include/nir_c_runtime/checkpoint.h b/include/nir_c_runtime/checkpoint.h new file mode 100644 index 0000000..6de7d4c --- /dev/null +++ b/include/nir_c_runtime/checkpoint.h @@ -0,0 +1,123 @@ +/* Torch-free checkpoint save/load plus the reverse bridge (#29, #30). + * + * Checkpoints are version-stamped JSON (mirroring the + * `spikeforge-nir-graph` v1 envelope discipline): + * + * {"format": "spikeforge-c-checkpoint", "version": 1, + * "topology": "conv_net", "seed": 7, "step": 100, + * "nodes": [{"name": "fc", "kind": "Affine", + * "shape": [4, 8], "weight": [...], "bias": [...]}]} + * + * Weights are flat row-major float arrays with an explicit shape; + * floats print with %.9g, which round-trips float32 exactly. Unknown + * format versions fail with a named error, as do missing nodes, + * kind/shape mismatches, and non-finite weights -- never a silent + * reinterpretation. + * + * The JSON writer here is new code (the existing json module stays + * parse-only, owned by another worker). The weight setter mutates + * the graph's public param buffers in place, which is also the + * mutation path a future C optimizer would use. + * + * `nir_graph_export_summary` is the C half of the reverse bridge: + * it re-emits a loaded graph as `graph_summary`-shaped JSON with + * current (possibly C-trained) weights, which `nir_graph_load` + * accepts and `tools/patch_nir_weights.py` can merge into a + * Python-side summary. + */ +#ifndef NIR_C_RUNTIME_CHECKPOINT_H +#define NIR_C_RUNTIME_CHECKPOINT_H + +#include +#include + +#include "nir_c_runtime/graph.h" + +#define NIR_CHECKPOINT_FORMAT "spikeforge-c-checkpoint" +#define NIR_CHECKPOINT_VERSION 1 + +/* Growing JSON text builder. */ +typedef struct nir_json_writer nir_json_writer_t; + +nir_json_writer_t *nir_json_writer_new(void); +void nir_json_writer_free(nir_json_writer_t *writer); + +/* The finished document; valid until the next call or free. NULL + * after any failed call. */ +const char *nir_json_writer_data(const nir_json_writer_t *writer); +size_t nir_json_writer_length(const nir_json_writer_t *writer); + +int nir_json_begin_object(nir_json_writer_t *writer); +int nir_json_begin_array(nir_json_writer_t *writer); +/* Close the innermost container. Fails when unbalanced. */ +int nir_json_end(nir_json_writer_t *writer); +int nir_json_key(nir_json_writer_t *writer, const char *key); +int nir_json_string(nir_json_writer_t *writer, const char *value); +int nir_json_float(nir_json_writer_t *writer, float value); +int nir_json_double(nir_json_writer_t *writer, double value); +int nir_json_int(nir_json_writer_t *writer, long long value); +int nir_json_bool(nir_json_writer_t *writer, int value); +int nir_json_null(nir_json_writer_t *writer); + +typedef struct { + char *topology; /* malloc'd, NULL when the file omits it */ + uint64_t seed; + unsigned long step; +} nir_checkpoint_meta_t; + +void nir_checkpoint_meta_free(nir_checkpoint_meta_t *meta); + +/* Serialize `graph`'s weights plus `meta` (NULL for weights only) + * to a malloc'd JSON document. */ +char *nir_checkpoint_dumps( + const nir_graph_t *graph, const nir_checkpoint_meta_t *meta, + char **error +); + +/* Write the same document to `path`. */ +int nir_checkpoint_save( + const char *path, const nir_graph_t *graph, + const nir_checkpoint_meta_t *meta, char **error +); + +/* Load weights from a checkpoint document into an already-loaded + * graph: every weighted node must be present with matching kind + * and shape. Fills `meta` (NULL to skip; caller frees with + * `nir_checkpoint_meta_free`). */ +int nir_checkpoint_loads( + const char *text, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + char **error +); + +/* Load weights from the checkpoint file at `path`. */ +int nir_checkpoint_load( + const char *path, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + char **error +); + +/* Point `*weight`/`*bias` at node `name`'s buffers with element + * counts (`*bias` NULL, `*bias_count` 0 when the node has no bias). + * Returns 0 for unknown nodes and nodes without weights. */ +int nir_graph_get_weights( + const nir_graph_t *graph, const char *name, const float **weight, + size_t *weight_count, const float **bias, size_t *bias_count +); + +/* Copy `weight_count`/`bias_count` floats into node `name`'s + * buffers. Weight counts must match exactly; pass `bias` NULL to + * leave the bias unchanged (nodes without a bias reject bias + * data). Returns 0 and sets *error on any mismatch. */ +int nir_graph_set_weights( + nir_graph_t *graph, const char *name, const float *weight, + size_t weight_count, const float *bias, size_t bias_count, + char **error +); + +/* Re-emit `graph` as `graph_summary`-shaped JSON (nodes with kind + * params, edges with delayed flags) carrying the current weights. + * Returns a malloc'd document, or NULL with *error set. */ +char *nir_graph_export_summary( + const nir_graph_t *graph, char **error +); + +#endif /* NIR_C_RUNTIME_CHECKPOINT_H */ diff --git a/include/nir_c_runtime/data.h b/include/nir_c_runtime/data.h new file mode 100644 index 0000000..185072c --- /dev/null +++ b/include/nir_c_runtime/data.h @@ -0,0 +1,47 @@ +/* Torch-free dataset loading (#27). + * + * Minimal IDX-format MNIST reader (no new dependencies): the magic + * numbers, big-endian headers, and row-major u8 pixels of the + * original MNIST files. Normalization matches spikeforge's torch + * transform exactly (`datasets.transform`: Grayscale, Resize, + * ToTensor, Normalize((0,), (1,))) -- ToTensor scales bytes by + * 1/255 and the Normalize is an identity -- so C decodes each + * pixel as v / 255.0f. Multi-byte samples land in `nir_tensor_t` + * frames (batch-1, one tensor per sample); batching stays the + * caller's loop until the tensor ABI grows a batch axis. + */ +#ifndef NIR_C_RUNTIME_DATA_H +#define NIR_C_RUNTIME_DATA_H + +#include +#include + +#include "nir_c_runtime/tensor.h" + +/* Decode `n` IDX pixels into normalized floats (v / 255). */ +void nir_data_normalize(const uint8_t *pixels, float *out, size_t n); + +/* Read an IDX image file (magic 2051) into a malloc'd row-major + * buffer of `*count * *rows * *cols` bytes. Dimensions are capped + * (count <= 2^24, rows/cols <= 1024) and the file must hold exactly + * the declared payload. Returns 0 and sets *error on failure. */ +int nir_idx_read_images( + const char *path, uint8_t **pixels, size_t *count, size_t *rows, + size_t *cols, char **error +); + +/* Read an IDX label file (magic 2049) into a malloc'd buffer of + * `*count` bytes. Returns 0 and sets *error on failure. */ +int nir_idx_read_labels( + const char *path, uint8_t **labels, size_t *count, char **error +); + +/* Normalize sample `index` of an image buffer into a fresh (C=1, H, + * W) tensor. Returns a tensor with data == NULL on bad sizes or + * allocation failure. */ +nir_tensor_t nir_data_frame( + const uint8_t *pixels, size_t count, size_t rows, size_t cols, + size_t index +); + +#endif /* NIR_C_RUNTIME_DATA_H */ diff --git a/include/nir_c_runtime/encode.h b/include/nir_c_runtime/encode.h new file mode 100644 index 0000000..2778719 --- /dev/null +++ b/include/nir_c_runtime/encode.h @@ -0,0 +1,87 @@ +/* Spike encoders ported from spikeforge (#28). + * + * Mirrors `SpikeEncoder` (spikeforge/encoding/spike_encoder.py) on top + * of snnTorch's `spikegen` semantics: + * + * - rate: clamp01(gain * x) Bernoulli sampling, i.e. `spikegen.rate` + * with gain=1 (`rate_conv`: clamp then `torch.bernoulli`). A spike + * fires iff a uniform draw u satisfies u < p. The draws come from + * `nir_rng_t` (splitmix64), NOT torch's MT19937, so bit equality + * with torch holds only given the same uniform stream; the rule + * above is the documented equivalence contract (draw k = t * n + + * i, consumed t-major, i-minor, one `nir_rng_next_f32` per draw). + * - latency: first-spike-time code, `spikegen.latency` with the + * encoder's tau/threshold/clip/normalize/linear flags. Linear and + * log branches both implemented; spike times round half-to-even + * (`rintf`), matching `torch.round`. + * - delta: `spikegen.delta` over a linspace(0, 1, steps) ramp times + * each pixel (padding=False, off_spike=False), threshold given as + * a fraction (`SpikeEncoder` divides its percent by 100). + * - random: the noise baseline, `rate_conv(rand * scale)` with two + * RNG draws per element (magnitude, then threshold). + * - repeat: the legacy raw path (`_repeat_pixels`): normalised + * pixels repeated across steps, no spikes. + * + * Outputs are time-major: out[t * n + i] is pixel i at step t, + * matching the encoder's [T, B=1, N] layout flattened per sample. + */ +#ifndef NIR_C_RUNTIME_ENCODE_H +#define NIR_C_RUNTIME_ENCODE_H + +#include +#include + +#include "nir_c_runtime/rng.h" + +#define NIR_LATENCY_EPSILON 1e-7f + +typedef struct { + size_t num_steps; + float tau; + float threshold; + int clip; + int normalize; + int linear; + size_t first_spike_time; + float epsilon; +} nir_latency_params_t; + +/* Bernoulli rate code of `x[0..n)` over `num_steps`. Consumes one RNG + * draw per (t, i). Returns 0 and sets *error on bad sizes. */ +int nir_encode_rate( + const float *x, size_t n, size_t num_steps, float gain, + nir_rng_t *rng, uint8_t *out, char **error +); + +/* First-spike-time latency code of `x[0..n)`. Every input must lie + * in [0, 1] (the torch side raises otherwise). Returns 0 and sets + * *error on out-of-range input, bad parameters, or a spike time + * that overflows `num_steps` (mirroring spikegen's no-bypass + * raise). */ +int nir_encode_latency( + const float *x, size_t n, const nir_latency_params_t *params, + uint8_t *out, char **error +); + +/* Delta code of `x[0..n)` over `num_steps` with a fractional + * change threshold. Returns 0 and sets *error on bad sizes. */ +int nir_encode_delta( + const float *x, size_t n, size_t num_steps, float threshold, + uint8_t *out, char **error +); + +/* Noise baseline: rate-coded uniform noise, input ignored. + * Returns 0 and sets *error on bad sizes. */ +int nir_encode_random( + size_t n, size_t num_steps, float scale, nir_rng_t *rng, + uint8_t *out, char **error +); + +/* Repeat `x[0..n)` across `num_steps` into `out[t * n + i]`. + * Returns 0 and sets *error on bad sizes. */ +int nir_encode_repeat( + const float *x, size_t n, size_t num_steps, float *out, + char **error +); + +#endif /* NIR_C_RUNTIME_ENCODE_H */ diff --git a/include/nir_c_runtime/rng.h b/include/nir_c_runtime/rng.h new file mode 100644 index 0000000..3050c6c --- /dev/null +++ b/include/nir_c_runtime/rng.h @@ -0,0 +1,30 @@ +/* Deterministic seeded RNG for training-side C code (#28). + * + * splitmix64 (Steele, Lea, O'Neill 2014): portable C11, no platform + * entropy dependency, so embedded targets keep working. The stream + * for a seed is fixed by the algorithm below -- identical across + * runs and platforms -- which is what makes seeded spike encoding + * reproducible. This is a sampler, not a CSPRNG. + */ +#ifndef NIR_C_RUNTIME_RNG_H +#define NIR_C_RUNTIME_RNG_H + +#include + +typedef struct { + uint64_t state; +} nir_rng_t; + +/* Prime the generator with `seed`. */ +void nir_rng_seed(nir_rng_t *rng, uint64_t seed); + +/* Return the next 64 bits of the stream. */ +uint64_t nir_rng_next_u64(nir_rng_t *rng); + +/* Return a uniform double in [0, 1) from the top 53 bits of one draw. */ +double nir_rng_next_f64(nir_rng_t *rng); + +/* Return a uniform float in [0, 1) from the top 24 bits of one draw. */ +float nir_rng_next_f32(nir_rng_t *rng); + +#endif /* NIR_C_RUNTIME_RNG_H */ diff --git a/src/checkpoint.c b/src/checkpoint.c new file mode 100644 index 0000000..9a887b3 --- /dev/null +++ b/src/checkpoint.c @@ -0,0 +1,1047 @@ +#include "nir_c_runtime/checkpoint.h" + +#include +#include +#include +#include + +#include "nir_c_runtime/json.h" + +struct nir_json_writer { + char *text; + size_t length; + size_t capacity; + char *stack; /* '{' / '[' per open container */ + size_t depth; + size_t stack_capacity; + int need_comma; + int failed; +}; + +nir_json_writer_t *nir_json_writer_new(void) { + nir_json_writer_t *writer = + (nir_json_writer_t *)calloc(1, sizeof(*writer)); + return writer; +} + +void nir_json_writer_free(nir_json_writer_t *writer) { + if (writer == NULL) { + return; + } + free(writer->text); + free(writer->stack); + free(writer); +} + +const char *nir_json_writer_data(const nir_json_writer_t *writer) { + if (writer == NULL || writer->failed || writer->text == NULL) { + return NULL; + } + return writer->text; +} + +size_t nir_json_writer_length(const nir_json_writer_t *writer) { + if (writer == NULL || writer->failed) { + return 0; + } + return writer->length; +} + +static int writer_reserve(nir_json_writer_t *w, size_t extra) { + if (w->failed) { + return 0; + } + if (w->length + extra + 1 <= w->capacity) { + return 1; + } + size_t capacity = w->capacity == 0 ? 256 : w->capacity; + while (w->length + extra + 1 > capacity) { + capacity *= 2; + } + char *text = (char *)realloc(w->text, capacity); + if (text == NULL) { + w->failed = 1; + return 0; + } + w->text = text; + w->capacity = capacity; + return 1; +} + +static int writer_put(nir_json_writer_t *w, const char *s, size_t n) { + if (!writer_reserve(w, n)) { + return 0; + } + memcpy(w->text + w->length, s, n); + w->length += n; + w->text[w->length] = '\0'; + return 1; +} + +static int writer_byte(nir_json_writer_t *w, char c) { + return writer_put(w, &c, 1); +} + +/* Emit a separator (or the object colon handling lives in key). */ +static int writer_value(nir_json_writer_t *w) { + if (w->failed) { + return 0; + } + if (w->need_comma && !writer_byte(w, ',')) { + return 0; + } + w->need_comma = 1; + return 1; +} + +static int writer_push(nir_json_writer_t *w, char kind, char open) { + if (!writer_value(w) || !writer_byte(w, open)) { + return 0; + } + if (w->depth == w->stack_capacity) { + size_t capacity = w->stack_capacity == 0 ? 8 : w->stack_capacity * 2; + char *stack = (char *)realloc(w->stack, capacity); + if (stack == NULL) { + w->failed = 1; + return 0; + } + w->stack = stack; + w->stack_capacity = capacity; + } + w->stack[w->depth++] = kind; + w->need_comma = 0; + return 1; +} + +int nir_json_begin_object(nir_json_writer_t *writer) { + return writer != NULL && writer_push(writer, '{', '{'); +} + +int nir_json_begin_array(nir_json_writer_t *writer) { + return writer != NULL && writer_push(writer, '[', '['); +} + +int nir_json_end(nir_json_writer_t *writer) { + if (writer == NULL || writer->failed || writer->depth == 0) { + if (writer != NULL) { + writer->failed = 1; + } + return 0; + } + char kind = writer->stack[--writer->depth]; + writer->need_comma = 1; + return writer_byte(writer, kind == '{' ? '}' : ']'); +} + +int nir_json_key(nir_json_writer_t *writer, const char *key) { + if (writer == NULL || key == NULL || writer->failed || + writer->depth == 0 || writer->stack[writer->depth - 1] != '{') { + if (writer != NULL) { + writer->failed = 1; + } + return 0; + } + if (writer->need_comma && !writer_byte(writer, ',')) { + return 0; + } + writer->need_comma = 0; + if (!nir_json_string(writer, key)) { + return 0; + } + writer->need_comma = 0; + return writer_byte(writer, ':'); +} + +int nir_json_string(nir_json_writer_t *writer, const char *value) { + static const char *escapes[32] = { + "\\u0000", "\\u0001", "\\u0002", "\\u0003", + "\\u0004", "\\u0005", "\\u0006", "\\u0007", + "\\b", "\\t", "\\n", "\\u000b", + "\\f", "\\r", "\\u000e", "\\u000f", + "\\u0010", "\\u0011", "\\u0012", "\\u0013", + "\\u0014", "\\u0015", "\\u0016", "\\u0017", + "\\u0018", "\\u0019", "\\u001a", "\\u001b", + "\\u001c", "\\u001d", "\\u001e", "\\u001f", + }; + if (writer == NULL || value == NULL || !writer_value(writer) || + !writer_byte(writer, '"')) { + return 0; + } + for (const unsigned char *p = (const unsigned char *)value; *p; p++) { + if (*p == '"' || *p == '\\') { + if (!writer_byte(writer, '\\') || !writer_byte(writer, *p)) { + return 0; + } + } else if (*p < 32) { + const char *esc = escapes[*p]; + if (!writer_put(writer, esc, strlen(esc))) { + return 0; + } + } else if (!writer_byte(writer, *p)) { + return 0; + } + } + writer->need_comma = 1; + return writer_byte(writer, '"'); +} + +static int writer_number(nir_json_writer_t *w, const char *text) { + return writer_value(w) && writer_put(w, text, strlen(text)); +} + +int nir_json_float(nir_json_writer_t *writer, float value) { + char text[32]; + if (writer == NULL) { + return 0; + } + if (!isfinite(value)) { + return nir_json_null(writer); + } + snprintf(text, sizeof(text), "%.9g", value); + return writer_number(writer, text); +} + +int nir_json_double(nir_json_writer_t *writer, double value) { + char text[32]; + if (writer == NULL) { + return 0; + } + if (!isfinite(value)) { + return nir_json_null(writer); + } + snprintf(text, sizeof(text), "%.17g", value); + return writer_number(writer, text); +} + +int nir_json_int(nir_json_writer_t *writer, long long value) { + char text[32]; + if (writer == NULL) { + return 0; + } + snprintf(text, sizeof(text), "%lld", value); + return writer_number(writer, text); +} + +int nir_json_bool(nir_json_writer_t *writer, int value) { + if (writer == NULL) { + return 0; + } + return writer_number(writer, value ? "true" : "false"); +} + +int nir_json_null(nir_json_writer_t *writer) { + if (writer == NULL) { + return 0; + } + return writer_number(writer, "null"); +} + +static void set_error(char **error, const char *message) { + if (error == NULL || *error != NULL) { + return; + } + size_t length = strlen(message) + 1; + *error = (char *)malloc(length); + if (*error != NULL) { + memcpy(*error, message, length); + } +} + +void nir_checkpoint_meta_free(nir_checkpoint_meta_t *meta) { + if (meta == NULL) { + return; + } + free(meta->topology); + memset(meta, 0, sizeof(*meta)); +} + +/* Row-major element counts per weighted kind; 0 weight means the + * kind carries no weights. */ +static int weight_shapes( + const nir_node_t *node, const float **weight, size_t *wn, + const float **bias, size_t *bn +) { + switch (node->kind) { + case NIR_NODE_CONV2D: { + const nir_conv2d_params_t *p = &node->params.conv2d; + *weight = p->weight; + *wn = p->out_channels * p->in_channels * p->kernel_h * + p->kernel_w; + *bias = p->bias; + *bn = p->bias != NULL ? p->out_channels : 0; + return 1; + } + case NIR_NODE_AFFINE: { + const nir_affine_params_t *p = &node->params.affine; + *weight = p->weight; + *wn = p->out_features * p->in_features; + *bias = p->bias; + *bn = p->bias != NULL ? p->out_features : 0; + return 1; + } + case NIR_NODE_CONV1D: { + const nir_conv1d_params_t *p = &node->params.conv1d; + *weight = p->weight; + *wn = p->out_channels * p->in_channels * p->kernel; + *bias = p->bias; + *bn = p->bias != NULL ? p->out_channels : 0; + return 1; + } + case NIR_NODE_LINEAR: { + const nir_linear_params_t *p = &node->params.linear; + *weight = p->weight; + *wn = p->out_features * p->in_features; + *bias = NULL; + *bn = 0; + return 1; + } + default: + return 0; + } +} + +int nir_graph_get_weights( + const nir_graph_t *graph, const char *name, const float **weight, + size_t *weight_count, const float **bias, size_t *bias_count +) { + if (graph == NULL || name == NULL || weight == NULL || + weight_count == NULL || bias == NULL || bias_count == NULL) { + return 0; + } + size_t index = nir_graph_find(graph, name); + if (index == (size_t)-1) { + return 0; + } + return weight_shapes( + &graph->nodes[index], weight, weight_count, bias, bias_count + ); +} + +int nir_graph_set_weights( + nir_graph_t *graph, const char *name, const float *weight, + size_t weight_count, const float *bias, size_t bias_count, + char **error +) { + if (graph == NULL || name == NULL) { + set_error(error, "weight update needs a graph and node name"); + return 0; + } + size_t index = nir_graph_find(graph, name); + if (index == (size_t)-1) { + set_error(error, "weight update names an unknown node"); + return 0; + } + const float *dst_w = NULL; + const float *dst_b = NULL; + size_t want_w = 0; + size_t want_b = 0; + if (!weight_shapes( + &graph->nodes[index], &dst_w, &want_w, &dst_b, &want_b + )) { + set_error(error, "weight update targets an unweighted node"); + return 0; + } + if (weight == NULL || weight_count != want_w) { + set_error(error, "weight update has the wrong weight count"); + return 0; + } + if (bias == NULL) { + if (bias_count != 0) { + set_error(error, "weight update has an inconsistent bias"); + return 0; + } + } else if (dst_b == NULL || bias_count != want_b) { + set_error(error, "weight update has the wrong bias count"); + return 0; + } + memcpy((float *)dst_w, weight, want_w * sizeof(float)); + if (bias != NULL) { + memcpy((float *)dst_b, bias, want_b * sizeof(float)); + } + return 1; +} + +static const char *kind_name(nir_node_kind_t kind) { + switch (kind) { + case NIR_NODE_INPUT: return "Input"; + case NIR_NODE_OUTPUT: return "Output"; + case NIR_NODE_CONV2D: return "Conv2d"; + case NIR_NODE_AFFINE: return "Affine"; + case NIR_NODE_FLATTEN: return "Flatten"; + case NIR_NODE_AVGPOOL2D: return "AvgPool2d"; + case NIR_NODE_SUMPOOL2D: return "SumPool2d"; + case NIR_NODE_LI: return "LI"; + case NIR_NODE_THRESHOLD: return "Threshold"; + case NIR_NODE_SCALE: return "Scale"; + case NIR_NODE_DELAY: return "Delay"; + case NIR_NODE_LIF: return "LIF"; + case NIR_NODE_IF: return "IF"; + case NIR_NODE_CUBALIF: return "CubaLIF"; + case NIR_NODE_CONV1D: return "Conv1d"; + case NIR_NODE_LINEAR: return "Linear"; + default: return NULL; + } +} + +static int write_floats( + nir_json_writer_t *w, const float *values, size_t n +) { + if (!nir_json_begin_array(w)) { + return 0; + } + for (size_t i = 0; i < n; i++) { + if (!isfinite(values[i]) || !nir_json_float(w, values[i])) { + return 0; + } + } + return nir_json_end(w); +} + +/* Explicit shape vector per weighted kind, matching the loader's + * dimension order. */ +static int write_shape(nir_json_writer_t *w, const nir_node_t *node) { + if (!nir_json_key(w, "shape") || !nir_json_begin_array(w)) { + return 0; + } + long long dims[4]; + size_t ndim = 0; + switch (node->kind) { + case NIR_NODE_CONV2D: { + const nir_conv2d_params_t *p = &node->params.conv2d; + dims[0] = (long long)p->out_channels; + dims[1] = (long long)p->in_channels; + dims[2] = (long long)p->kernel_h; + dims[3] = (long long)p->kernel_w; + ndim = 4; + break; + } + case NIR_NODE_AFFINE: { + const nir_affine_params_t *p = &node->params.affine; + dims[0] = (long long)p->out_features; + dims[1] = (long long)p->in_features; + ndim = 2; + break; + } + case NIR_NODE_CONV1D: { + const nir_conv1d_params_t *p = &node->params.conv1d; + dims[0] = (long long)p->out_channels; + dims[1] = (long long)p->in_channels; + dims[2] = (long long)p->kernel; + ndim = 3; + break; + } + case NIR_NODE_LINEAR: { + const nir_linear_params_t *p = &node->params.linear; + dims[0] = (long long)p->out_features; + dims[1] = (long long)p->in_features; + ndim = 2; + break; + } + default: + return 0; + } + for (size_t i = 0; i < ndim; i++) { + if (!nir_json_int(w, dims[i])) { + return 0; + } + } + return nir_json_end(w); +} + +static int write_node(nir_json_writer_t *w, const nir_node_t *node) { + const float *weight = NULL; + const float *bias = NULL; + size_t wn = 0; + size_t bn = 0; + if (!weight_shapes(node, &weight, &wn, &bias, &bn)) { + return 0; + } + return nir_json_begin_object(w) && nir_json_key(w, "name") && + nir_json_string(w, node->name) && nir_json_key(w, "kind") && + nir_json_string(w, kind_name(node->kind)) && + write_shape(w, node) && nir_json_key(w, "weight") && + write_floats(w, weight, wn) && nir_json_key(w, "bias") && + (bias != NULL ? write_floats(w, bias, bn) + : nir_json_null(w)) && + nir_json_end(w); +} + +char *nir_checkpoint_dumps( + const nir_graph_t *graph, const nir_checkpoint_meta_t *meta, + char **error +) { + if (graph == NULL) { + set_error(error, "checkpoint save needs a graph"); + return NULL; + } + nir_json_writer_t *w = nir_json_writer_new(); + int ok = w != NULL && nir_json_begin_object(w) && + nir_json_key(w, "format") && + nir_json_string(w, NIR_CHECKPOINT_FORMAT) && + nir_json_key(w, "version") && + nir_json_int(w, NIR_CHECKPOINT_VERSION); + if (ok && meta != NULL) { + if (meta->topology != NULL) { + ok = nir_json_key(w, "topology") && + nir_json_string(w, meta->topology); + } + ok = ok && nir_json_key(w, "seed") && + nir_json_double(w, (double)meta->seed) && + nir_json_key(w, "step") && + nir_json_int(w, (long long)meta->step); + } + ok = ok && nir_json_key(w, "nodes") && nir_json_begin_array(w); + for (size_t i = 0; ok && i < graph->node_count; i++) { + const nir_node_t *node = &graph->nodes[i]; + const float *weight = NULL; + const float *bias = NULL; + size_t wn = 0; + size_t bn = 0; + if (weight_shapes(node, &weight, &wn, &bias, &bn)) { + ok = write_node(w, node); + } + } + ok = ok && nir_json_end(w) && nir_json_end(w); + char *text = NULL; + if (ok && nir_json_writer_data(w) != NULL) { + size_t length = nir_json_writer_length(w) + 1; + text = (char *)malloc(length); + if (text != NULL) { + memcpy(text, nir_json_writer_data(w), length); + } + } + nir_json_writer_free(w); + if (text == NULL) { + set_error(error, "checkpoint weights must be finite"); + } + return text; +} + +int nir_checkpoint_save( + const char *path, const nir_graph_t *graph, + const nir_checkpoint_meta_t *meta, char **error +) { + if (path == NULL) { + set_error(error, "checkpoint save needs a path"); + return 0; + } + char *text = nir_checkpoint_dumps(graph, meta, error); + if (text == NULL) { + return 0; + } + FILE *f = fopen(path, "wb"); + int ok = 0; + if (f != NULL) { + size_t length = strlen(text); + ok = fwrite(text, 1, length, f) == length; + ok = fclose(f) == 0 && ok; + } + free(text); + if (!ok) { + set_error(error, "cannot write checkpoint file"); + } + return ok; +} + +/* Read a flat numeric array into `out[0..n)`; every element must be + * a finite number. */ +static int read_floats( + const json_value_t *array, float *out, size_t n +) { + if (array == NULL || array->type != JSON_ARRAY || + json_array_length(array) != n) { + return 0; + } + for (size_t i = 0; i < n; i++) { + const json_value_t *item = json_array_get(array, i); + if (item == NULL || item->type != JSON_NUMBER || + !isfinite(item->as.number)) { + return 0; + } + out[i] = (float)item->as.number; + } + return 1; +} + +static int shape_matches(const nir_node_t *node, const json_value_t *shape) { + long long want[4]; + size_t ndim = 0; + switch (node->kind) { + case NIR_NODE_CONV2D: { + const nir_conv2d_params_t *p = &node->params.conv2d; + want[0] = (long long)p->out_channels; + want[1] = (long long)p->in_channels; + want[2] = (long long)p->kernel_h; + want[3] = (long long)p->kernel_w; + ndim = 4; + break; + } + case NIR_NODE_AFFINE: { + const nir_affine_params_t *p = &node->params.affine; + want[0] = (long long)p->out_features; + want[1] = (long long)p->in_features; + ndim = 2; + break; + } + case NIR_NODE_CONV1D: { + const nir_conv1d_params_t *p = &node->params.conv1d; + want[0] = (long long)p->out_channels; + want[1] = (long long)p->in_channels; + want[2] = (long long)p->kernel; + ndim = 3; + break; + } + case NIR_NODE_LINEAR: { + const nir_linear_params_t *p = &node->params.linear; + want[0] = (long long)p->out_features; + want[1] = (long long)p->in_features; + ndim = 2; + break; + } + default: + return 0; + } + if (shape == NULL || shape->type != JSON_ARRAY || + json_array_length(shape) != ndim) { + return 0; + } + for (size_t i = 0; i < ndim; i++) { + const json_value_t *dim = json_array_get(shape, i); + if (dim == NULL || dim->type != JSON_NUMBER || + dim->as.number != (double)want[i]) { + return 0; + } + } + return 1; +} + +static int load_node_weights( + nir_graph_t *graph, const json_value_t *entry, char **error +) { + const char *name = json_as_string(json_object_get(entry, "name")); + const char *kind = json_as_string(json_object_get(entry, "kind")); + if (name == NULL || kind == NULL) { + set_error(error, "checkpoint node is missing name or kind"); + return 0; + } + size_t index = nir_graph_find(graph, name); + if (index == (size_t)-1) { + set_error(error, "checkpoint names an unknown node"); + return 0; + } + nir_node_t *node = &graph->nodes[index]; + const char *want_kind = kind_name(node->kind); + if (want_kind == NULL || strcmp(kind, want_kind) != 0) { + set_error(error, "checkpoint node kind does not match graph"); + return 0; + } + const float *dst_w = NULL; + const float *dst_b = NULL; + size_t wn = 0; + size_t bn = 0; + if (!weight_shapes(node, &dst_w, &wn, &dst_b, &bn) || + !shape_matches(node, json_object_get(entry, "shape"))) { + set_error(error, "checkpoint node shape does not match graph"); + return 0; + } + if (!read_floats(json_object_get(entry, "weight"), (float *)dst_w, wn)) { + set_error(error, "checkpoint node weight is not a flat array"); + return 0; + } + const json_value_t *bias = json_object_get(entry, "bias"); + if (bias == NULL || bias->type == JSON_NULL) { + if (dst_b != NULL) { + set_error(error, "checkpoint drops a node bias"); + return 0; + } + return 1; + } + if (dst_b == NULL || !read_floats(bias, (float *)dst_b, bn)) { + set_error(error, "checkpoint node bias does not match graph"); + return 0; + } + return 1; +} + +static char *read_whole_file(const char *path) { + FILE *f = fopen(path, "rb"); + if (f == NULL) { + return NULL; + } + fseek(f, 0, SEEK_END); + long size = ftell(f); + fseek(f, 0, SEEK_SET); + if (size < 0) { + fclose(f); + return NULL; + } + char *text = (char *)malloc((size_t)size + 1); + if (text == NULL || (size > 0 && + fread(text, 1, (size_t)size, f) != (size_t)size)) { + free(text); + fclose(f); + return NULL; + } + text[size] = '\0'; + fclose(f); + return text; +} + +int nir_checkpoint_loads( + const char *text, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + char **error +) { + if (text == NULL || graph == NULL) { + set_error(error, "checkpoint load needs text and a graph"); + return 0; + } + char *parse_error = NULL; + json_value_t *root = json_parse(text, &parse_error); + if (root == NULL) { + free(parse_error); + set_error(error, "checkpoint is not valid JSON"); + return 0; + } + const char *format = json_as_string(json_object_get(root, "format")); + const json_value_t *version = json_object_get(root, "version"); + int ok = format != NULL && strcmp(format, NIR_CHECKPOINT_FORMAT) == 0; + if (!ok) { + set_error(error, "checkpoint has an unknown format marker"); + } else if (version == NULL || version->type != JSON_NUMBER || + version->as.number != (double)NIR_CHECKPOINT_VERSION) { + set_error(error, "checkpoint has an unsupported version"); + ok = 0; + } + const json_value_t *nodes = json_object_get(root, "nodes"); + if (ok && (nodes == NULL || nodes->type != JSON_ARRAY)) { + set_error(error, "checkpoint has no node array"); + ok = 0; + } + unsigned char *seen = NULL; + if (ok) { + seen = (unsigned char *)calloc( + graph->node_count == 0 ? 1 : graph->node_count, 1 + ); + ok = seen != NULL; + } + for (size_t i = 0; ok && i < json_array_length(nodes); i++) { + const json_value_t *entry = json_array_get(nodes, i); + const char *name = json_as_string(json_object_get(entry, "name")); + ok = load_node_weights(graph, entry, error); + if (ok) { + seen[nir_graph_find(graph, name)] = 1; + } + } + /* Every weighted node must be covered: a partial resume would + * silently mix stale and fresh weights. */ + for (size_t i = 0; ok && i < graph->node_count; i++) { + const float *w = NULL; + const float *b = NULL; + size_t wn = 0; + size_t bn = 0; + if (weight_shapes(&graph->nodes[i], &w, &wn, &b, &bn) && + !seen[i]) { + set_error(error, "checkpoint lacks weights for a node"); + ok = 0; + } + } + if (ok && meta != NULL) { + memset(meta, 0, sizeof(*meta)); + const char *topology = json_as_string( + json_object_get(root, "topology") + ); + if (topology != NULL) { + meta->topology = (char *)malloc(strlen(topology) + 1); + if (meta->topology == NULL) { + set_error(error, "out of memory"); + ok = 0; + } else { + strcpy(meta->topology, topology); + } + } + meta->seed = (uint64_t)json_as_number( + json_object_get(root, "seed"), 0.0 + ); + meta->step = (unsigned long)json_as_number( + json_object_get(root, "step"), 0.0 + ); + } + free(seen); + json_free(root); + return ok; +} + +int nir_checkpoint_load( + const char *path, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + char **error +) { + if (path == NULL) { + set_error(error, "checkpoint load needs a path"); + return 0; + } + char *text = read_whole_file(path); + if (text == NULL) { + set_error(error, "cannot read checkpoint file"); + return 0; + } + int ok = nir_checkpoint_loads(text, graph, meta, error); + free(text); + return ok; +} + +/* Emit `weight` (flat row-major, `wn` elements) as a nested array + * with per-level sizes `dims[0..ndim)`. */ +static int write_nested( + nir_json_writer_t *w, const float *weight, const size_t *dims, + size_t ndim +) { + if (ndim == 1) { + if (!nir_json_begin_array(w)) { + return 0; + } + for (size_t i = 0; i < dims[0]; i++) { + if (!isfinite(weight[i]) || !nir_json_float(w, weight[i])) { + return 0; + } + } + return nir_json_end(w); + } + if (!nir_json_begin_array(w)) { + return 0; + } + size_t stride = 1; + for (size_t i = 1; i < ndim; i++) { + stride *= dims[i]; + } + for (size_t i = 0; i < dims[0]; i++) { + if (!write_nested(w, weight + i * stride, dims + 1, ndim - 1)) { + return 0; + } + } + return nir_json_end(w); +} + +static int write_pair(nir_json_writer_t *w, const char *key, int a, int b) { + return nir_json_key(w, key) && nir_json_begin_array(w) && + nir_json_int(w, a) && nir_json_int(w, b) && nir_json_end(w); +} + +static int write_scalar( + nir_json_writer_t *w, const char *key, float value +) { + return nir_json_key(w, key) && nir_json_float(w, value); +} + +static int export_weighted(nir_json_writer_t *w, const nir_node_t *node) { + const float *weight = NULL; + const float *bias = NULL; + size_t wn = 0; + size_t bn = 0; + size_t dims[4]; + size_t ndim = 0; + if (!weight_shapes(node, &weight, &wn, &bias, &bn)) { + return 0; + } + switch (node->kind) { + case NIR_NODE_CONV2D: { + const nir_conv2d_params_t *p = &node->params.conv2d; + dims[0] = p->out_channels; + dims[1] = p->in_channels; + dims[2] = (size_t)p->kernel_h; + dims[3] = (size_t)p->kernel_w; + ndim = 4; + break; + } + case NIR_NODE_AFFINE: { + const nir_affine_params_t *p = &node->params.affine; + dims[0] = p->out_features; + dims[1] = p->in_features; + ndim = 2; + break; + } + case NIR_NODE_CONV1D: { + const nir_conv1d_params_t *p = &node->params.conv1d; + dims[0] = p->out_channels; + dims[1] = p->in_channels; + dims[2] = p->kernel; + ndim = 3; + break; + } + case NIR_NODE_LINEAR: { + const nir_linear_params_t *p = &node->params.linear; + dims[0] = p->out_features; + dims[1] = p->in_features; + ndim = 2; + break; + } + default: + return 0; + } + int ok = nir_json_key(w, "weight") && + write_nested(w, weight, dims, ndim) && nir_json_key(w, "bias"); + if (!ok) { + return 0; + } + if (node->kind == NIR_NODE_LINEAR) { + ok = nir_json_null(w); + } else if (bias != NULL) { + ok = write_floats(w, bias, bn); + } else { + ok = nir_json_null(w); + } + return ok; +} + +static int export_conv2d(nir_json_writer_t *w, const nir_node_t *node) { + const nir_conv2d_params_t *p = &node->params.conv2d; + return export_weighted(w, node) && + write_pair(w, "stride", p->stride_h, p->stride_w) && + write_pair(w, "padding", p->padding_h, p->padding_w) && + write_pair(w, "dilation", p->dilation_h, p->dilation_w); +} + +static int export_conv1d(nir_json_writer_t *w, const nir_node_t *node) { + const nir_conv1d_params_t *p = &node->params.conv1d; + return export_weighted(w, node) && + nir_json_key(w, "stride") && nir_json_int(w, p->stride) && + nir_json_key(w, "padding") && nir_json_int(w, p->padding) && + nir_json_key(w, "dilation") && nir_json_int(w, p->dilation); +} + +static int export_pool2d(nir_json_writer_t *w, const nir_node_t *node) { + const nir_pool2d_params_t *p = &node->params.pool2d; + return write_pair(w, "kernel_size", p->kernel_h, p->kernel_w) && + write_pair(w, "stride", p->stride_h, p->stride_w) && + write_pair(w, "padding", p->padding_h, p->padding_w); +} + +static int export_lif(nir_json_writer_t *w, const nir_node_t *node) { + const nir_lif_params_t *p = &node->params.lif; + return write_scalar(w, "tau", p->tau) && + write_scalar(w, "r", p->r) && + write_scalar(w, "v_leak", p->v_leak) && + write_scalar(w, "v_threshold", p->v_threshold) && + write_scalar(w, "v_reset", p->v_reset); +} + +static int export_if(nir_json_writer_t *w, const nir_node_t *node) { + const nir_if_params_t *p = &node->params.if_; + return write_scalar(w, "r", p->r) && + write_scalar(w, "v_threshold", p->v_threshold) && + write_scalar(w, "v_reset", p->v_reset); +} + +static int export_cuba(nir_json_writer_t *w, const nir_node_t *node) { + const nir_cuba_params_t *p = &node->params.cuba; + return write_scalar(w, "tau_syn", p->tau_syn) && + write_scalar(w, "tau_mem", p->tau_mem) && + write_scalar(w, "r", p->r) && + write_scalar(w, "v_leak", p->v_leak) && + write_scalar(w, "v_threshold", p->v_threshold) && + write_scalar(w, "v_reset", p->v_reset) && + write_scalar(w, "w_in", p->w_in); +} + +static int export_params(nir_json_writer_t *w, const nir_node_t *node) { + if (!nir_json_key(w, "params") || !nir_json_begin_object(w)) { + return 0; + } + int ok = 1; + switch (node->kind) { + case NIR_NODE_CONV2D: + ok = export_conv2d(w, node); + break; + case NIR_NODE_CONV1D: + ok = export_conv1d(w, node); + break; + case NIR_NODE_AFFINE: + case NIR_NODE_LINEAR: + ok = export_weighted(w, node); + break; + case NIR_NODE_AVGPOOL2D: + case NIR_NODE_SUMPOOL2D: + ok = export_pool2d(w, node); + break; + case NIR_NODE_LI: { + const nir_li_params_t *p = &node->params.li; + ok = write_scalar(w, "tau", p->tau) && + write_scalar(w, "r", p->r) && + write_scalar(w, "v_leak", p->v_leak); + break; + } + case NIR_NODE_LIF: + ok = export_lif(w, node); + break; + case NIR_NODE_IF: + ok = export_if(w, node); + break; + case NIR_NODE_CUBALIF: + ok = export_cuba(w, node); + break; + case NIR_NODE_THRESHOLD: + ok = write_scalar( + w, "threshold", node->params.threshold.threshold + ); + break; + case NIR_NODE_SCALE: + ok = write_scalar(w, "scale", node->params.scale.scale); + break; + default: + break; /* Input/Output/Flatten/Delay carry no params */ + } + return ok && nir_json_end(w); +} + +char *nir_graph_export_summary( + const nir_graph_t *graph, char **error +) { + if (graph == NULL) { + set_error(error, "summary export needs a graph"); + return NULL; + } + nir_json_writer_t *w = nir_json_writer_new(); + int ok = w != NULL && nir_json_begin_object(w) && + nir_json_key(w, "nodes") && nir_json_begin_array(w); + for (size_t i = 0; ok && i < graph->node_count; i++) { + const nir_node_t *node = &graph->nodes[i]; + const char *kind = kind_name(node->kind); + ok = kind != NULL && nir_json_begin_object(w) && + nir_json_key(w, "name") && + nir_json_string(w, node->name) && nir_json_key(w, "kind") && + nir_json_string(w, kind) && export_params(w, node) && + nir_json_end(w); + } + ok = ok && nir_json_end(w) && nir_json_key(w, "edges") && + nir_json_begin_array(w); + for (size_t i = 0; ok && i < graph->edge_count; i++) { + size_t src = graph->edges[i].source_index; + size_t dst = graph->edges[i].target_index; + if (src >= graph->node_count || dst >= graph->node_count) { + ok = 0; + break; + } + ok = nir_json_begin_object(w) && nir_json_key(w, "source") && + nir_json_string(w, graph->nodes[src].name) && + nir_json_key(w, "target") && + nir_json_string(w, graph->nodes[dst].name) && + nir_json_key(w, "delayed") && + nir_json_bool( + w, graph->nodes[dst].kind == NIR_NODE_DELAY + ) && + nir_json_end(w); + } + ok = ok && nir_json_end(w) && nir_json_end(w); + char *text = NULL; + if (ok && nir_json_writer_data(w) != NULL) { + size_t length = nir_json_writer_length(w) + 1; + text = (char *)malloc(length); + if (text != NULL) { + memcpy(text, nir_json_writer_data(w), length); + } + } + nir_json_writer_free(w); + if (text == NULL) { + set_error(error, "summary export failed to serialize"); + } + return text; +} diff --git a/src/data.c b/src/data.c new file mode 100644 index 0000000..45a2ea5 --- /dev/null +++ b/src/data.c @@ -0,0 +1,142 @@ +#include "nir_c_runtime/data.h" + +#include +#include +#include + +#define IDX_MAGIC_IMAGES 2051u +#define IDX_MAGIC_LABELS 2049u +#define IDX_MAX_COUNT ((size_t)1 << 24) +#define IDX_MAX_DIM 1024u + +void nir_data_normalize(const uint8_t *pixels, float *out, size_t n) { + for (size_t i = 0; i < n; i++) { + out[i] = (float)pixels[i] / 255.0f; + } +} + +static void set_error(char **error, const char *message) { + if (error == NULL || *error != NULL) { + return; + } + size_t length = strlen(message) + 1; + *error = (char *)malloc(length); + if (*error != NULL) { + memcpy(*error, message, length); + } +} + +static int read_u32be(FILE *f, uint32_t *out) { + unsigned char bytes[4]; + if (fread(bytes, 1, 4, f) != 4) { + return 0; + } + *out = ((uint32_t)bytes[0] << 24) | ((uint32_t)bytes[1] << 16) | + ((uint32_t)bytes[2] << 8) | (uint32_t)bytes[3]; + return 1; +} + +/* Read the whole payload after the header; the file must end + * exactly at the declared size (no truncation, no trailing junk). */ +static uint8_t *read_payload(FILE *f, size_t size) { + uint8_t *buffer = (uint8_t *)malloc(size == 0 ? 1 : size); + if (buffer == NULL) { + return NULL; + } + if (size > 0 && fread(buffer, 1, size, f) != size) { + free(buffer); + return NULL; + } + if (fgetc(f) != EOF) { + free(buffer); + return NULL; + } + return buffer; +} + +int nir_idx_read_images( + const char *path, uint8_t **pixels, size_t *count, size_t *rows, + size_t *cols, char **error +) { + if (path == NULL || pixels == NULL || count == NULL || rows == NULL || + cols == NULL) { + set_error(error, "IDX image read needs out-parameters"); + return 0; + } + FILE *f = fopen(path, "rb"); + if (f == NULL) { + set_error(error, "cannot open IDX image file"); + return 0; + } + uint32_t magic = 0, n = 0, h = 0, w = 0; + int ok = read_u32be(f, &magic) && read_u32be(f, &n) && + read_u32be(f, &h) && read_u32be(f, &w); + if (!ok || magic != IDX_MAGIC_IMAGES || n == 0 || n > IDX_MAX_COUNT || + h == 0 || h > IDX_MAX_DIM || w == 0 || w > IDX_MAX_DIM) { + fclose(f); + set_error(error, "bad IDX image header"); + return 0; + } + size_t total = (size_t)n * (size_t)h * (size_t)w; + uint8_t *buffer = read_payload(f, total); + fclose(f); + if (buffer == NULL) { + set_error(error, "IDX image payload is truncated or ragged"); + return 0; + } + *pixels = buffer; + *count = n; + *rows = h; + *cols = w; + return 1; +} + +int nir_idx_read_labels( + const char *path, uint8_t **labels, size_t *count, char **error +) { + if (path == NULL || labels == NULL || count == NULL) { + set_error(error, "IDX label read needs out-parameters"); + return 0; + } + FILE *f = fopen(path, "rb"); + if (f == NULL) { + set_error(error, "cannot open IDX label file"); + return 0; + } + uint32_t magic = 0, n = 0; + int ok = read_u32be(f, &magic) && read_u32be(f, &n); + if (!ok || magic != IDX_MAGIC_LABELS || n == 0 || + n > IDX_MAX_COUNT) { + fclose(f); + set_error(error, "bad IDX label header"); + return 0; + } + uint8_t *buffer = read_payload(f, n); + fclose(f); + if (buffer == NULL) { + set_error(error, "IDX label payload is truncated or ragged"); + return 0; + } + *labels = buffer; + *count = n; + return 1; +} + +nir_tensor_t nir_data_frame( + const uint8_t *pixels, size_t count, size_t rows, size_t cols, + size_t index +) { + nir_tensor_t empty; + memset(&empty, 0, sizeof(empty)); + if (pixels == NULL || count == 0 || rows == 0 || cols == 0 || + index >= count) { + return empty; + } + size_t shape[3] = {1, rows, cols}; + nir_tensor_t frame = nir_tensor_alloc(shape, 3); + if (frame.data == NULL) { + return frame; + } + nir_data_normalize(pixels + index * rows * cols, frame.data, frame.size); + return frame; +} diff --git a/src/encode.c b/src/encode.c new file mode 100644 index 0000000..db8e978 --- /dev/null +++ b/src/encode.c @@ -0,0 +1,217 @@ +#include "nir_c_runtime/encode.h" + +#include +#include +#include + +static void set_error(char **error, const char *message) { + if (error == NULL || *error != NULL) { + return; + } + size_t length = strlen(message) + 1; + *error = (char *)malloc(length); + if (*error != NULL) { + memcpy(*error, message, length); + } +} + +static float clamp01(float v) { + if (v < 0.0f) { + return 0.0f; + } + if (v > 1.0f) { + return 1.0f; + } + return v; +} + +int nir_encode_rate( + const float *x, size_t n, size_t num_steps, float gain, + nir_rng_t *rng, uint8_t *out, char **error +) { + if (x == NULL || rng == NULL || out == NULL || n == 0 || + num_steps == 0 || !isfinite(gain)) { + set_error(error, "rate encode needs inputs, steps, and a gain"); + return 0; + } + for (size_t t = 0; t < num_steps; t++) { + for (size_t i = 0; i < n; i++) { + float p = clamp01(gain * x[i]); + out[t * n + i] = nir_rng_next_f32(rng) < p ? 1 : 0; + } + } + return 1; +} + +static int latency_param_ok( + const nir_latency_params_t *p, char **error +) { + if (p == NULL || p->num_steps == 0 || !(p->tau > 0.0f) || + p->first_spike_time >= p->num_steps) { + set_error(error, "latency encode needs steps, tau, spike time"); + return 0; + } + if (p->linear) { + if (!(p->threshold >= 0.0f) || !(p->threshold <= 1.0f)) { + set_error(error, "linear latency threshold must lie in [0, 1]"); + return 0; + } + } else if (!(p->threshold > 0.0f) || !(p->threshold < 1.0f)) { + set_error(error, "log latency threshold must lie in (0, 1)"); + return 0; + } + return 1; +} + +/* Linear branch: st = min(-tau*(x-1), -tau*(thr-1)) + first, with + * tau replaced by steps-1-first when normalizing. */ +static void latency_linear( + const float *x, size_t n, const nir_latency_params_t *p, float *st +) { + float tau = p->normalize + ? (float)(p->num_steps - 1 - p->first_spike_time) + : p->tau; + float cap = -tau * (p->threshold - 1.0f); + for (size_t i = 0; i < n; i++) { + float time = -tau * (x[i] - 1.0f); + st[i] = (time < cap ? time : cap) + (float)p->first_spike_time; + } +} + +/* Log branch: st = tau*ln(xc/(xc-thr)) (+ normalize rescale), with + * inputs saturated at thr+eps. */ +static void latency_log( + const float *x, size_t n, const nir_latency_params_t *p, float *st +) { + float floor = p->threshold + p->epsilon; + float first = (float)p->first_spike_time; + float max_shift = 0.0f; + for (size_t i = 0; i < n; i++) { + float xc = x[i] < floor ? floor : x[i]; + st[i] = p->tau * logf(xc / (xc - p->threshold)) + first; + if (st[i] - first > max_shift) { + max_shift = st[i] - first; + } + } + if (p->normalize) { + float span = (float)(p->num_steps - 1) - first; + for (size_t i = 0; i < n; i++) { + st[i] = max_shift > 0.0f + ? (st[i] - first) * span / max_shift + first + : first; + } + } +} + +int nir_encode_latency( + const float *x, size_t n, const nir_latency_params_t *params, + uint8_t *out, char **error +) { + if (x == NULL || out == NULL || n == 0 || + !latency_param_ok(params, error)) { + return 0; + } + for (size_t i = 0; i < n; i++) { + if (!(x[i] >= 0.0f) || !(x[i] <= 1.0f)) { + set_error(error, "latency input must lie in [0, 1]"); + return 0; + } + } + float *st = (float *)malloc(n * sizeof(float)); + if (st == NULL) { + set_error(error, "out of memory"); + return 0; + } + if (params->linear) { + latency_linear(x, n, params, st); + } else { + latency_log(x, n, params, st); + } + /* spikegen raises when the latest spike overflows num_steps + * (SpikeEncoder never bypasses); check before scattering. */ + double latest = 0.0; + for (size_t i = 0; i < n; i++) { + double rounded = rint((double)st[i]); + if (rounded > latest) { + latest = rounded; + } + } + if (latest > (double)(params->num_steps - 1)) { + free(st); + set_error(error, "latency spike time exceeds num_steps"); + return 0; + } + memset(out, 0, params->num_steps * n); + for (size_t i = 0; i < n; i++) { + if (params->clip && x[i] < params->threshold) { + continue; + } + size_t t = (size_t)rint((double)st[i]); + if (t > params->num_steps - 1) { + continue; /* rm_idx: drop out-of-range spikes */ + } + out[t * n + i] = 1; + } + free(st); + return 1; +} + +int nir_encode_delta( + const float *x, size_t n, size_t num_steps, float threshold, + uint8_t *out, char **error +) { + if (x == NULL || out == NULL || n == 0 || num_steps == 0 || + !isfinite(threshold)) { + set_error(error, "delta encode needs inputs and steps"); + return 0; + } + /* linspace(0, 1, steps): t/(steps-1), or 0 for a single step. */ + for (size_t t = 0; t < num_steps; t++) { + float ramp = num_steps > 1 + ? (float)t / (float)(num_steps - 1) + : 0.0f; + float prev = num_steps > 1 && t > 0 + ? (float)(t - 1) / (float)(num_steps - 1) + : 0.0f; + for (size_t i = 0; i < n; i++) { + /* Separate statements: one expression could contract. */ + float level = ramp * x[i]; + float before = prev * x[i]; + out[t * n + i] = (level - before) >= threshold ? 1 : 0; + } + } + return 1; +} + +int nir_encode_random( + size_t n, size_t num_steps, float scale, nir_rng_t *rng, + uint8_t *out, char **error +) { + if (rng == NULL || out == NULL || n == 0 || num_steps == 0 || + !isfinite(scale)) { + set_error(error, "random encode needs steps, scale, and an RNG"); + return 0; + } + for (size_t t = 0; t < num_steps; t++) { + for (size_t i = 0; i < n; i++) { + float magnitude = nir_rng_next_f32(rng); + float p = clamp01(scale * magnitude); + out[t * n + i] = nir_rng_next_f32(rng) < p ? 1 : 0; + } + } + return 1; +} + +int nir_encode_repeat( + const float *x, size_t n, size_t num_steps, float *out, + char **error +) { + if (x == NULL || out == NULL || n == 0 || num_steps == 0) { + set_error(error, "repeat encode needs inputs and steps"); + return 0; + } + for (size_t t = 0; t < num_steps; t++) { + memcpy(&out[t * n], x, n * sizeof(float)); + } + return 1; +} diff --git a/src/rng.c b/src/rng.c new file mode 100644 index 0000000..48d9642 --- /dev/null +++ b/src/rng.c @@ -0,0 +1,23 @@ +#include "nir_c_runtime/rng.h" + +void nir_rng_seed(nir_rng_t *rng, uint64_t seed) { + rng->state = seed; +} + +uint64_t nir_rng_next_u64(nir_rng_t *rng) { + uint64_t z = rng->state + UINT64_C(0x9E3779B97F4A7C15); + rng->state = z; + z = (z ^ (z >> 30)) * UINT64_C(0xBF58476D1CE4E5B9); + z = (z ^ (z >> 27)) * UINT64_C(0x94D049BB133111EB); + return z ^ (z >> 31); +} + +double nir_rng_next_f64(nir_rng_t *rng) { + uint64_t bits = nir_rng_next_u64(rng) >> 11; + return (double)bits * (1.0 / 9007199254740992.0); +} + +float nir_rng_next_f32(nir_rng_t *rng) { + uint32_t bits = (uint32_t)(nir_rng_next_u64(rng) >> 40); + return (float)bits * (1.0f / 16777216.0f); +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 4e9fbe4..83f2955 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -15,6 +15,10 @@ nir_add_test(test_quantized) nir_add_test(test_workspace) nir_add_test(test_tensor) nir_add_test(test_error_format) +nir_add_test(test_rng) +nir_add_test(test_encode) +nir_add_test(test_data) +nir_add_test(test_checkpoint) # OOM sweep includes src/json.c directly, so it must not link the library # (that would duplicate the json_* symbols); it only needs the headers. diff --git a/tests/fixtures/encode_golden.json b/tests/fixtures/encode_golden.json new file mode 100644 index 0000000..2b6f511 --- /dev/null +++ b/tests/fixtures/encode_golden.json @@ -0,0 +1,347 @@ +{ + "rate": { + "x": [ + 0.0, + 0.25, + 0.5, + 0.75, + 1.0 + ], + "steps": 8, + "gain": 1.0, + "seed": 12345, + "bits": [ + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 1, + 1, + 0, + 0, + 1, + 1, + 1 + ] + }, + "latency_linear": { + "x": [ + 0.0, + 0.1, + 0.55, + 0.9, + 1.0 + ], + "steps": 10, + "tau": 5.0, + "threshold": 0.01, + "clip": false, + "normalize": true, + "linear": true, + "bits": [ + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0 + ] + }, + "latency_linear_clip": { + "x": [ + 0.0, + 0.005, + 0.55, + 1.0 + ], + "steps": 10, + "tau": 5.0, + "threshold": 0.01, + "clip": true, + "normalize": true, + "linear": true, + "bits": [ + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + "latency_linear_plain": { + "x": [ + 0.2, + 0.6, + 1.0 + ], + "steps": 12, + "tau": 2.0, + "threshold": 0.01, + "clip": false, + "normalize": false, + "linear": true, + "first_spike_time": 1, + "bits": [ + 0, + 0, + 0, + 0, + 0, + 1, + 0, + 1, + 0, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + "latency_log": { + "x": [ + 0.05, + 0.3, + 0.7, + 1.0 + ], + "steps": 10, + "tau": 5.0, + "threshold": 0.01, + "normalize": true, + "linear": false, + "bits": [ + 0, + 0, + 0, + 1, + 0, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0 + ] + }, + "delta": { + "x": [ + 0.0, + 0.3, + 1.0 + ], + "steps": 6, + "threshold": 0.04, + "bits": [ + 0, + 0, + 0, + 0, + 1, + 1, + 0, + 1, + 1, + 0, + 1, + 1, + 0, + 1, + 1, + 0, + 1, + 1 + ] + }, + "random": { + "n": 4, + "steps": 6, + "scale": 0.5, + "seed": 999, + "bits": [ + 0, + 0, + 0, + 0, + 0, + 1, + 0, + 1, + 0, + 0, + 0, + 1, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 0, + 1, + 0, + 1, + 0 + ] + } +} diff --git a/tests/fixtures/mnist_tiny_images.idx b/tests/fixtures/mnist_tiny_images.idx new file mode 100644 index 0000000000000000000000000000000000000000..97d0d34115cc14d01d87d0027f5c4de3cb683b94 GIT binary patch literal 28 dcmZQz;9zE8U;tqzFb-Yy|Np8`21X`k762op1Q`GT literal 0 HcmV?d00001 diff --git a/tests/fixtures/mnist_tiny_labels.idx b/tests/fixtures/mnist_tiny_labels.idx new file mode 100644 index 0000000000000000000000000000000000000000..a69706da6dea8403c42d6a107b8bb44626ce2b9d GIT binary patch literal 11 ScmZQz;9z86U|?ov<^%u$ivS=1 literal 0 HcmV?d00001 diff --git a/tests/test_checkpoint.c b/tests/test_checkpoint.c new file mode 100644 index 0000000..3e7038e --- /dev/null +++ b/tests/test_checkpoint.c @@ -0,0 +1,268 @@ +/* Checkpoint round-trips, weight accessors, summary export (#29). */ +#include +#include +#include + +#include "nir_c_runtime/checkpoint.h" +#include "nir_c_runtime/graph.h" +#include "nir_c_runtime/json.h" +#include "nir_c_runtime/runtime.h" +#include "test_util.h" + +static void check_writer(void) { + nir_json_writer_t *w = nir_json_writer_new(); + CHECK(w != NULL); + CHECK(nir_json_begin_object(w)); + CHECK(nir_json_key(w, "name")); + CHECK(nir_json_string(w, "a\"b\\c\nd")); + CHECK(nir_json_key(w, "w")); + CHECK(nir_json_begin_array(w)); + CHECK(nir_json_float(w, 0.5f)); + CHECK(nir_json_double(w, 0.1)); + CHECK(nir_json_int(w, -7)); + CHECK(nir_json_bool(w, 1)); + CHECK(nir_json_null(w)); + CHECK(nir_json_end(w)); + CHECK(nir_json_end(w)); + const char *text = nir_json_writer_data(w); + CHECK(text != NULL); + CHECK(strcmp(text, + "{\"name\":\"a\\\"b\\\\c\\nd\",\"w\":[0.5,0.10000000000000001," + "-7,true,null]}") == 0); + CHECK(nir_json_writer_length(w) == strlen(text)); + CHECK(!nir_json_end(w)); /* unbalanced close fails */ + CHECK(nir_json_writer_data(w) == NULL); + nir_json_writer_free(w); + + nir_json_writer_t *bad = nir_json_writer_new(); + CHECK(!nir_json_key(bad, "k")); /* key outside an object fails */ + CHECK(!nir_json_end(bad)); /* end with nothing open fails */ + nir_json_writer_free(bad); +} + +static nir_graph_t load_fixture(const char *path) { + nir_graph_t graph; + memset(&graph, 0, sizeof(graph)); + char *text = read_file(path); + char *error = NULL; + json_value_t *json = json_parse(text, &error); + CHECK(error == NULL); + CHECK(nir_graph_load(json, &graph, &error)); + CHECK(error == NULL); + json_free(json); + free(text); + return graph; +} + +static void check_weights_equal( + const nir_graph_t *a, const nir_graph_t *b +) { + CHECK(a->node_count == b->node_count); + for (size_t i = 0; i < a->node_count; i++) { + const float *wa = NULL; + const float *wb = NULL; + const float *ba = NULL; + const float *bb = NULL; + size_t wna = 0; + size_t wnb = 0; + size_t bna = 0; + size_t bnb = 0; + int ha = nir_graph_get_weights( + a, a->nodes[i].name, &wa, &wna, &ba, &bna + ); + int hb = nir_graph_get_weights( + b, b->nodes[i].name, &wb, &wnb, &bb, &bnb + ); + CHECK(ha == hb); + if (!ha) { + continue; + } + CHECK(wna == wnb && bna == bnb); + CHECK(memcmp(wa, wb, wna * sizeof(float)) == 0); + if (bna > 0) { + CHECK(memcmp(ba, bb, bna * sizeof(float)) == 0); + } + } +} + +static nir_tensor_t run_readout(nir_graph_t *graph) { + nir_runtime_t runtime; + char *error = NULL; + CHECK(nir_runtime_init(&runtime, graph, &error)); + size_t shape[1] = {4}; + nir_tensor_t frames[6]; + for (int t = 0; t < 6; t++) { + frames[t] = nir_tensor_alloc(shape, 1); + for (size_t i = 0; i < frames[t].size; i++) { + frames[t].data[i] = (float)(t + 1) * 0.1f; + } + } + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(nir_runtime_run(&runtime, frames, 6, &readout, &error)); + CHECK(error == NULL); + for (int t = 0; t < 6; t++) { + nir_tensor_free(&frames[t]); + } + nir_runtime_free(&runtime); + return readout; +} + +int main(void) { + check_writer(); + + char *error = NULL; + nir_graph_t graph = load_fixture("fixtures/lif_linear_graph.json"); + + /* Weight accessors: get, set, and every mismatch shape. */ + const float *w = NULL; + const float *b = NULL; + size_t wn = 0; + size_t bn = 0; + CHECK(nir_graph_get_weights(&graph, "fc1", &w, &wn, &b, &bn)); + CHECK(wn > 0 && b == NULL && bn == 0); /* Linear has no bias */ + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &b, &bn)); + CHECK(wn > 0 && b != NULL && bn > 0); /* Affine has a bias */ + CHECK(!nir_graph_get_weights(&graph, "lif1", &w, &wn, &b, &bn)); + CHECK(!nir_graph_get_weights(&graph, "nope", &w, &wn, &b, &bn)); + + float *fc1_before = (float *)malloc(wn * sizeof(float)); + CHECK(fc1_before != NULL); + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &b, &bn)); + memcpy(fc1_before, w, wn * sizeof(float)); + float *zeros = (float *)calloc(wn > bn ? wn : bn, sizeof(float)); + CHECK(zeros != NULL); + CHECK(nir_graph_set_weights(&graph, "fc2", zeros, wn, zeros, bn, + &error)); + CHECK(error == NULL); + CHECK(!nir_graph_set_weights(&graph, "fc2", zeros, wn + 1, NULL, 0, + &error)); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_graph_set_weights(&graph, "fc1", zeros, 1, zeros, 1, + &error)); /* Linear rejects bias data */ + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_graph_set_weights(&graph, "lif1", zeros, 1, NULL, 0, + &error)); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_graph_set_weights(&graph, "nope", zeros, 1, NULL, 0, + &error)); + CHECK(error != NULL); + free(error); + error = NULL; + + /* Save with metadata, clobber, reload: bit-exact resume. */ + nir_checkpoint_meta_t meta; + memset(&meta, 0, sizeof(meta)); + meta.topology = "lif_linear"; + meta.seed = 42; + meta.step = 100; + /* Restore fc2 first so the checkpoint holds trained weights. */ + CHECK(nir_graph_set_weights(&graph, "fc2", fc1_before, wn, b, bn, + &error)); + char *dump = nir_checkpoint_dumps(&graph, &meta, &error); + CHECK(error == NULL); + CHECK(dump != NULL); + CHECK(strstr(dump, "\"format\":\"spikeforge-c-checkpoint\"") != NULL); + CHECK(strstr(dump, "\"version\":1") != NULL); + CHECK(strstr(dump, "\"topology\":\"lif_linear\"") != NULL); + CHECK(nir_graph_set_weights(&graph, "fc2", zeros, wn, zeros, bn, + &error)); + nir_checkpoint_meta_t back; + memset(&back, 0, sizeof(back)); + CHECK(nir_checkpoint_loads(dump, &graph, &back, &error)); + CHECK(error == NULL); + CHECK(back.topology != NULL && strcmp(back.topology, "lif_linear") == 0); + CHECK(back.seed == 42 && back.step == 100); + nir_checkpoint_meta_free(&back); + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &b, &bn)); + CHECK(memcmp(w, fc1_before, wn * sizeof(float)) == 0); + + /* File round-trip through TMPDIR (never the source tree). */ + const char *tmp = getenv("TMPDIR"); + char path[512]; + sprintf(path, "%s/nir-checkpoint-test.json", + tmp != NULL ? tmp : "/tmp"); + CHECK(nir_checkpoint_save(path, &graph, &meta, &error)); + CHECK(error == NULL); + CHECK(nir_graph_set_weights(&graph, "fc2", zeros, wn, zeros, bn, + &error)); + memset(&back, 0, sizeof(back)); + CHECK(nir_checkpoint_load(path, &graph, &back, &error)); + CHECK(error == NULL); + CHECK(back.step == 100); + nir_checkpoint_meta_free(&back); + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &b, &bn)); + CHECK(memcmp(w, fc1_before, wn * sizeof(float)) == 0); + free(fc1_before); + free(zeros); + + /* Rejects: bad version, bad format, missing node, bad shape. */ + char hacked[65536]; + strcpy(hacked, dump); + char *ver = strstr(hacked, "\"version\":1"); + CHECK(ver != NULL); + ver[strlen("\"version\":")] = '9'; + CHECK(!nir_checkpoint_loads(hacked, &graph, NULL, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_checkpoint_loads("{\"format\":\"nope\",\"version\":1," + "\"nodes\":[]}", &graph, NULL, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_checkpoint_loads("{\"format\":" + "\"spikeforge-c-checkpoint\",\"version\":1,\"nodes\":[]}", + &graph, NULL, &error)); /* drops every weighted node */ + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_checkpoint_loads("not json", &graph, NULL, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + + /* Non-finite weights refuse to serialize. */ + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &b, &bn)); + float *mutable = (float *)w; + float saved = mutable[0]; + mutable[0] = INFINITY; + CHECK(nir_checkpoint_dumps(&graph, NULL, &error) == NULL); + CHECK(error != NULL); + free(error); + error = NULL; + mutable[0] = saved; + + /* Reverse bridge: export, reload, identical weights + readout. */ + nir_tensor_t before = run_readout(&graph); + char *summary = nir_graph_export_summary(&graph, &error); + CHECK(error == NULL); + CHECK(summary != NULL); + json_value_t *summary_json = json_parse(summary, &error); + CHECK(error == NULL); + nir_graph_t again; + memset(&again, 0, sizeof(again)); + CHECK(nir_graph_load(summary_json, &again, &error)); + CHECK(error == NULL); + check_weights_equal(&graph, &again); + nir_tensor_t after = run_readout(&again); + CHECK(before.size == after.size); + for (size_t i = 0; i < before.size; i++) { + CHECK(before.data[i] == after.data[i]); + } + nir_tensor_free(&before); + nir_tensor_free(&after); + nir_graph_free(&again); + json_free(summary_json); + free(summary); + free(dump); + + nir_graph_free(&graph); + return 0; +} diff --git a/tests/test_data.c b/tests/test_data.c new file mode 100644 index 0000000..5805b75 --- /dev/null +++ b/tests/test_data.c @@ -0,0 +1,111 @@ +/* Torch-free IDX loading: exact bytes, torch-matching normalize. */ +#include +#include + +#include "nir_c_runtime/data.h" +#include "test_util.h" + +static void write_temp( + const char *name, const unsigned char *bytes, size_t n, char *path +) { + const char *tmp = getenv("TMPDIR"); + if (tmp == NULL) { + tmp = "/tmp"; + } + sprintf(path, "%s/nir-data-test-%s", tmp, name); + FILE *f = fopen(path, "wb"); + CHECK(f != NULL); + CHECK(fwrite(bytes, 1, n, f) == n); + CHECK(fclose(f) == 0); +} + +int main(void) { + char *error = NULL; + uint8_t *pixels = NULL; + uint8_t *labels = NULL; + size_t count = 0; + size_t rows = 0; + size_t cols = 0; + size_t nlabels = 0; + + CHECK(nir_idx_read_images( + "fixtures/mnist_tiny_images.idx", &pixels, &count, &rows, &cols, + &error + )); + CHECK(error == NULL); + CHECK(count == 3 && rows == 2 && cols == 2); + /* tools/make_mnist_fixture.py wrote these exact bytes. */ + const uint8_t want_px[12] = { + 0, 85, 170, 255, 255, 170, 85, 0, 1, 2, 3, 4, + }; + CHECK(memcmp(pixels, want_px, sizeof(want_px)) == 0); + + CHECK(nir_idx_read_labels( + "fixtures/mnist_tiny_labels.idx", &labels, &nlabels, &error + )); + CHECK(error == NULL); + CHECK(nlabels == 3); + CHECK(labels[0] == 7 && labels[1] == 3 && labels[2] == 9); + + /* Normalize matches ToTensor+Normalize((0,),(1,)): v / 255. */ + float norm[12]; + nir_data_normalize(pixels, norm, 12); + CHECK(norm[0] == 0.0f); + CHECK_NEAR(norm[1], 85.0 / 255.0, 1e-7); + CHECK_NEAR(norm[2], 170.0 / 255.0, 1e-7); + CHECK(norm[3] == 1.0f); + + /* One sample lands in a (1, H, W) tensor frame. */ + nir_tensor_t frame = nir_data_frame(pixels, count, rows, cols, 2); + CHECK(frame.data != NULL); + CHECK(frame.ndim == 3); + CHECK(frame.shape[0] == 1 && frame.shape[1] == 2 && + frame.shape[2] == 2); + CHECK_NEAR(frame.data[0], 1.0 / 255.0, 1e-7); + CHECK_NEAR(frame.data[3], 4.0 / 255.0, 1e-7); + nir_tensor_free(&frame); + nir_tensor_t bad = nir_data_frame(pixels, count, rows, cols, 3); + CHECK(bad.data == NULL); + free(pixels); + free(labels); + + /* Malformed inputs fail with named errors, not crashes. */ + char path[512]; + const unsigned char bad_magic[16] = { + 0, 0, 8, 2, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, + }; + write_temp("bad-magic.idx", bad_magic, sizeof(bad_magic), path); + pixels = NULL; + CHECK(!nir_idx_read_images(path, &pixels, &count, &rows, &cols, + &error)); + CHECK(error != NULL); + free(error); + error = NULL; + + /* Declares one 1x1 image but the payload is empty. */ + const unsigned char truncated[16] = { + 0, 0, 8, 3, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, + }; + write_temp("truncated.idx", truncated, sizeof(truncated), path); + CHECK(!nir_idx_read_images(path, &pixels, &count, &rows, &cols, + &error)); + CHECK(error != NULL); + free(error); + error = NULL; + + /* Trailing junk after the declared payload is rejected. */ + const unsigned char ragged[18] = { + 0, 0, 8, 1, 0, 0, 0, 1, 9, 9, + }; + write_temp("ragged.idx", ragged, 10, path); + labels = NULL; + CHECK(!nir_idx_read_labels(path, &labels, &nlabels, &error)); + CHECK(error != NULL); + free(error); + + CHECK(!nir_idx_read_images( + "fixtures/does-not-exist.idx", &pixels, &count, &rows, &cols, + NULL + )); + return 0; +} diff --git a/tests/test_encode.c b/tests/test_encode.c new file mode 100644 index 0000000..8f67e2b --- /dev/null +++ b/tests/test_encode.c @@ -0,0 +1,176 @@ +/* Spike encoders vs committed goldens (tools/encode_golden.py). */ +#include +#include + +#include "nir_c_runtime/encode.h" +#include "nir_c_runtime/json.h" +#include "test_util.h" + +static float *load_floats(const json_value_t *array, size_t *n) { + size_t count = json_array_length(array); + float *out = (float *)malloc(count * sizeof(float)); + CHECK(out != NULL); + for (size_t i = 0; i < count; i++) { + out[i] = (float)json_as_number(json_array_get(array, i), 0.0); + } + *n = count; + return out; +} + +static void check_bits( + const json_value_t *entry, const uint8_t *got, size_t n +) { + const json_value_t *bits = json_object_get(entry, "bits"); + CHECK(json_array_length(bits) == n); + for (size_t i = 0; i < n; i++) { + int want = (int)json_as_number(json_array_get(bits, i), -1); + if (got[i] != want) { + fprintf( + stderr, "bit %zu: got %u want %d\n", i, got[i], want + ); + exit(1); + } + } +} + +static size_t num_steps_of(const json_value_t *entry) { + return (size_t)json_as_number(json_object_get(entry, "steps"), 0); +} + +int main(void) { + char *text = read_file("fixtures/encode_golden.json"); + char *error = NULL; + json_value_t *goldens = json_parse(text, &error); + CHECK(error == NULL); + CHECK(goldens != NULL); + + /* Rate: Bernoulli rule against the same RNG stream. */ + const json_value_t *rate = json_object_get(goldens, "rate"); + size_t n = 0; + float *x = load_floats(json_object_get(rate, "x"), &n); + size_t steps = num_steps_of(rate); + uint8_t *out = (uint8_t *)malloc(steps * n); + CHECK(out != NULL); + nir_rng_t rng; + nir_rng_seed(&rng, 12345); + CHECK(nir_encode_rate(x, n, steps, 1.0f, &rng, out, &error)); + CHECK(error == NULL); + check_bits(rate, out, steps * n); + /* Deterministic: the same seed reproduces the train. */ + nir_rng_seed(&rng, 12345); + uint8_t *again = (uint8_t *)malloc(steps * n); + CHECK(nir_encode_rate(x, n, steps, 1.0f, &rng, again, &error)); + CHECK(memcmp(out, again, steps * n) == 0); + free(again); + free(x); + + /* Latency: linear, clipped, unnormalized, and log branches. */ + const char *cases[4] = { + "latency_linear", "latency_linear_clip", "latency_linear_plain", + "latency_log", + }; + for (int c = 0; c < 4; c++) { + const json_value_t *entry = json_object_get(goldens, cases[c]); + x = load_floats(json_object_get(entry, "x"), &n); + steps = num_steps_of(entry); + free(out); + out = (uint8_t *)malloc(steps * n); + nir_latency_params_t params; + memset(¶ms, 0, sizeof(params)); + params.num_steps = steps; + params.tau = (float)json_as_number( + json_object_get(entry, "tau"), 5.0 + ); + params.threshold = (float)json_as_number( + json_object_get(entry, "threshold"), 0.01 + ); + const json_value_t *clip = json_object_get(entry, "clip"); + params.clip = clip != NULL && clip->type == JSON_BOOL && + clip->as.boolean; + const json_value_t *norm = json_object_get(entry, "normalize"); + params.normalize = norm != NULL && norm->type == JSON_BOOL && + norm->as.boolean; + const json_value_t *lin = json_object_get(entry, "linear"); + params.linear = lin != NULL && lin->type == JSON_BOOL && + lin->as.boolean; + params.first_spike_time = (size_t)json_as_number( + json_object_get(entry, "first_spike_time"), 0 + ); + params.epsilon = NIR_LATENCY_EPSILON; + CHECK(nir_encode_latency(x, n, ¶ms, out, &error)); + CHECK(error == NULL); + check_bits(entry, out, steps * n); + free(x); + } + + /* Delta over the 0..1 ramp. */ + const json_value_t *delta = json_object_get(goldens, "delta"); + x = load_floats(json_object_get(delta, "x"), &n); + steps = num_steps_of(delta); + free(out); + out = (uint8_t *)malloc(steps * n); + float thr = (float)json_as_number( + json_object_get(delta, "threshold"), 0 + ); + CHECK(nir_encode_delta(x, n, steps, thr, out, &error)); + CHECK(error == NULL); + check_bits(delta, out, steps * n); + free(x); + + /* Random noise baseline. */ + const json_value_t *random = json_object_get(goldens, "random"); + n = (size_t)json_as_number(json_object_get(random, "n"), 0); + steps = num_steps_of(random); + free(out); + out = (uint8_t *)malloc(steps * n); + nir_rng_seed(&rng, 999); + float scale = (float)json_as_number( + json_object_get(random, "scale"), 0 + ); + CHECK(nir_encode_random(n, steps, scale, &rng, out, &error)); + CHECK(error == NULL); + check_bits(random, out, steps * n); + free(out); + + /* Repeat: raw pixels across steps, no spikes involved. */ + float px[3] = {0.25f, 0.5f, 1.0f}; + float rep[12]; + CHECK(nir_encode_repeat(px, 3, 4, rep, &error)); + for (int t = 0; t < 4; t++) { + CHECK(rep[t * 3] == 0.25f && rep[t * 3 + 1] == 0.5f && + rep[t * 3 + 2] == 1.0f); + } + + /* Error paths mirror the torch side's raises. */ + uint8_t tiny[8]; + float bad[2] = {0.5f, 1.5f}; + nir_latency_params_t lp; + memset(&lp, 0, sizeof(lp)); + lp.num_steps = 4; + lp.tau = 5.0f; + lp.threshold = 0.01f; + lp.normalize = 1; + lp.linear = 1; + lp.epsilon = NIR_LATENCY_EPSILON; + CHECK(!nir_encode_latency(bad, 2, &lp, tiny, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + /* Unnormalized tau=100 overflows 4 steps: spike time is 99. */ + lp.normalize = 0; + lp.tau = 100.0f; + float ok[1] = {0.0f}; + CHECK(!nir_encode_latency(ok, 1, &lp, tiny, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + nir_rng_seed(&rng, 0); + CHECK(!nir_encode_rate(NULL, 0, 4, 1.0f, &rng, tiny, NULL)); + CHECK(!nir_encode_delta(px, 3, 0, 0.04f, tiny, NULL)); + CHECK(!nir_encode_random(3, 4, 0.5f, NULL, tiny, NULL)); + CHECK(!nir_encode_repeat(px, 3, 0, rep, NULL)); + + json_free(goldens); + free(text); + return 0; +} diff --git a/tests/test_rng.c b/tests/test_rng.c new file mode 100644 index 0000000..f50172a --- /dev/null +++ b/tests/test_rng.c @@ -0,0 +1,57 @@ +/* Seeded RNG: known splitmix64 vectors plus determinism. */ +#include + +#include "nir_c_runtime/rng.h" +#include "test_util.h" + +static void check_stream(uint64_t seed, const uint64_t want[4]) { + nir_rng_t rng; + nir_rng_seed(&rng, seed); + for (int i = 0; i < 4; i++) { + CHECK(nir_rng_next_u64(&rng) == want[i]); + } +} + +int main(void) { + /* Reference vectors (splitmix64, independently computed). */ + const uint64_t seed0[4] = { + 0xe220a8397b1dcdafULL, 0x6e789e6aa1b965f4ULL, + 0x06c45d188009454fULL, 0xf88bb8a8724c81ecULL, + }; + const uint64_t seed1[4] = { + 0x910a2dec89025cc1ULL, 0xbeeb8da1658eec67ULL, + 0xf893a2eefb32555eULL, 0x71c18690ee42c90bULL, + }; + const uint64_t seed_big[4] = { + 0x22118258a9d111a0ULL, 0x346edce5f713f8edULL, + 0x1e9a57bc80e6721dULL, 0x2d160e7e5c3f42caULL, + }; + check_stream(0, seed0); + check_stream(1, seed1); + check_stream(12345, seed_big); + + /* Same seed, same stream, across instances and float widths. */ + nir_rng_t a; + nir_rng_t b; + nir_rng_seed(&a, 999); + nir_rng_seed(&b, 999); + for (int i = 0; i < 1000; i++) { + CHECK(nir_rng_next_u64(&a) == nir_rng_next_u64(&b)); + } + nir_rng_seed(&a, 7); + nir_rng_seed(&b, 7); + for (int i = 0; i < 1000; i++) { + double x = nir_rng_next_f64(&a); + double y = nir_rng_next_f64(&b); + CHECK(x == y && x >= 0.0 && x < 1.0); + float fx = nir_rng_next_f32(&a); + float fy = nir_rng_next_f32(&b); + CHECK(fx == fy && fx >= 0.0f && fx < 1.0f); + } + + /* Distinct seeds diverge immediately. */ + nir_rng_seed(&a, 0); + nir_rng_seed(&b, 1); + CHECK(nir_rng_next_u64(&a) != nir_rng_next_u64(&b)); + return 0; +} diff --git a/tools/encode_golden.py b/tools/encode_golden.py new file mode 100644 index 0000000..cacc1f7 --- /dev/null +++ b/tools/encode_golden.py @@ -0,0 +1,232 @@ +"""Emit golden spike trains for the C encoder test. + +Stdlib-only: needs no torch, so it runs anywhere. It re-implements +the documented coding rules in float32 (via struct) on top of a pure +Python splitmix64, and writes tests/fixtures/encode_golden.json, +which tests/test_encode.c compares against bit-exactly. + +Rule sources (see include/nir_c_runtime/encode.h): +- rate: spike iff u < clamp01(gain*x), one splitmix64-derived f32 + draw per (t, i), t-major -- the spikegen.rate/rate_conv Bernoulli + rule. torch's own stream is MT19937, so these goldens pin the + documented rule, not torch's stream (the equivalence contract the + C header states). +- latency: spikegen.latency first-spike times, linear and log + branches, rounded half-to-even like torch.round. +- delta: spikegen.delta over linspace(0, 1, steps) * pixel. +- random: rate_conv(rand * scale) with two draws per element. + +The script refuses to emit goldens whose latency spike times sit +within 1e-4 of a .5 rounding boundary, so C/Python libm +differences cannot flip a bit. +""" + +import json +import math +import struct +from pathlib import Path + +OUT = Path(__file__).resolve().parent.parent / "tests" / "fixtures" +MASK = (1 << 64) - 1 + + +def f32(value: float) -> float: + """Round a float to float32 precision.""" + return struct.unpack("f", struct.pack("f", value))[0] + + +class SplitMix64: + """Pure-Python splitmix64 matching src/rng.c.""" + + def __init__(self, seed: int) -> None: + self.state = seed & MASK + + def next_u64(self) -> int: + """Return the next 64 bits of the stream.""" + self.state = (self.state + 0x9E3779B97F4A7C15) & MASK + value = self.state + value = ((value ^ (value >> 30)) * 0xBF58476D1CE4E5B9) & MASK + value = ((value ^ (value >> 27)) * 0x94D049BB133111EB) & MASK + return value ^ (value >> 31) + + def next_f32(self) -> float: + """Return a uniform float32 in [0, 1).""" + bits = (self.next_u64() >> 40) & 0xFFFFFF + return f32(bits * (1.0 / 16777216.0)) + + +def clamp01(value: float) -> float: + """Clamp a float32 value to [0, 1].""" + return min(1.0, max(0.0, value)) + + +def golden_rate(x, steps, gain, seed): + """Bernoulli rate code bits, t-major.""" + rng = SplitMix64(seed) + n = len(x) + out = [0] * (steps * n) + for time in range(steps): + for i, pixel in enumerate(x): + prob = clamp01(f32(gain * pixel)) + out[time * n + i] = 1 if rng.next_f32() < prob else 0 + return out + + +def check_margin(times, label): + """Reject spike times too close to a rounding boundary.""" + for time in times: + frac = abs(time - math.floor(time + 0.5)) + assert abs(frac - 0.5) > 1e-4, f"{label}: {time} near boundary" + + +def golden_latency_linear(x, steps, tau, thr, clip, norm, first=0): + """Linear first-spike-time code bits, t-major.""" + plain = tau if not norm else float(steps - 1 - first) + cap = f32(-plain * f32(thr - 1.0)) + times = [] + for pixel in x: + time = f32(-plain * f32(pixel - 1.0)) + times.append(f32(min(time, cap) + first)) + check_margin(times, "latency-linear") + return _scatter(times, x, steps, thr, clip) + + +def golden_latency_log(x, steps, tau, thr, norm, eps=1e-7): + """Logarithmic first-spike-time code bits, t-major.""" + floor = thr + eps + times = [] + for pixel in x: + sat = max(pixel, floor) + times.append(f32(tau * math.log(sat / (sat - thr)))) + if norm: + peak = max(times) + span = float(steps - 1) + times = [f32(t * span / peak) for t in times] if peak > 0 else times + check_margin(times, "latency-log") + return _scatter(times, x, steps, thr, False) + + +def _scatter(times, x, steps, thr, clip): + """Place one spike per pixel at its rounded spike time.""" + n = len(x) + assert max(round(t) for t in times) <= steps - 1 + out = [0] * (steps * n) + for i, time in enumerate(times): + if clip and x[i] < thr: + continue + step = round(time) + if step <= steps - 1: + out[step * n + i] = 1 + return out + + +def golden_delta(x, steps, thr): + """Delta code bits over a 0..1 ramp, t-major.""" + n = len(x) + out = [0] * (steps * n) + for time in range(steps): + ramp = f32(time / (steps - 1)) if steps > 1 else 0.0 + prev = f32((time - 1) / (steps - 1)) if time > 0 else 0.0 + for i, pixel in enumerate(x): + change = f32(f32(ramp * pixel) - f32(prev * pixel)) + assert abs(change - thr) > 1e-6, "delta near boundary" + out[time * n + i] = 1 if change >= thr else 0 + return out + + +def golden_random(n, steps, scale, seed): + """Noise-baseline bits, t-major.""" + rng = SplitMix64(seed) + out = [0] * (steps * n) + for time in range(steps): + for i in range(n): + magnitude = rng.next_f32() + prob = clamp01(f32(scale * magnitude)) + out[time * n + i] = 1 if rng.next_f32() < prob else 0 + return out + + +def main() -> int: + """Write the golden JSON fixture.""" + goldens = { + "rate": { + "x": [0.0, 0.25, 0.5, 0.75, 1.0], + "steps": 8, + "gain": 1.0, + "seed": 12345, + "bits": golden_rate( + [0.0, 0.25, 0.5, 0.75, 1.0], 8, 1.0, 12345 + ), + }, + "latency_linear": { + "x": [0.0, 0.1, 0.55, 0.9, 1.0], + "steps": 10, + "tau": 5.0, + "threshold": 0.01, + "clip": False, + "normalize": True, + "linear": True, + "bits": golden_latency_linear( + [0.0, 0.1, 0.55, 0.9, 1.0], 10, 5.0, 0.01, + False, True, + ), + }, + "latency_linear_clip": { + "x": [0.0, 0.005, 0.55, 1.0], + "steps": 10, + "tau": 5.0, + "threshold": 0.01, + "clip": True, + "normalize": True, + "linear": True, + "bits": golden_latency_linear( + [0.0, 0.005, 0.55, 1.0], 10, 5.0, 0.01, True, True, + ), + }, + "latency_linear_plain": { + "x": [0.2, 0.6, 1.0], + "steps": 12, + "tau": 2.0, + "threshold": 0.01, + "clip": False, + "normalize": False, + "linear": True, + "first_spike_time": 1, + "bits": golden_latency_linear( + [0.2, 0.6, 1.0], 12, 2.0, 0.01, False, False, first=1, + ), + }, + "latency_log": { + "x": [0.05, 0.3, 0.7, 1.0], + "steps": 10, + "tau": 5.0, + "threshold": 0.01, + "normalize": True, + "linear": False, + "bits": golden_latency_log( + [0.05, 0.3, 0.7, 1.0], 10, 5.0, 0.01, True, + ), + }, + "delta": { + "x": [0.0, 0.3, 1.0], + "steps": 6, + "threshold": 0.04, + "bits": golden_delta([0.0, 0.3, 1.0], 6, 0.04), + }, + "random": { + "n": 4, + "steps": 6, + "scale": 0.5, + "seed": 999, + "bits": golden_random(4, 6, 0.5, 999), + }, + } + OUT.mkdir(parents=True, exist_ok=True) + with open(OUT / "encode_golden.json", "w", encoding="utf-8") as handle: + json.dump(goldens, handle, indent=2) + handle.write("\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/make_mnist_fixture.py b/tools/make_mnist_fixture.py new file mode 100644 index 0000000..b153c7c --- /dev/null +++ b/tools/make_mnist_fixture.py @@ -0,0 +1,45 @@ +"""Write the tiny committed IDX fixtures for the C data test. + +Stdlib-only (no torch, no datasets download): the pixels and labels +are hardcoded so regeneration is byte-stable. The C test +(tests/test_data.c) reads these with zero Python involvement. + +Layout: 3 images of 2x2 pixels plus 3 labels. +""" + +import struct +from pathlib import Path + +OUT_DIR = Path(__file__).resolve().parent.parent / "tests" / "fixtures" + +IMAGES = [ + [0, 85, 170, 255], + [255, 170, 85, 0], + [1, 2, 3, 4], +] +LABELS = [7, 3, 9] + + +def write_idx(path: Path, magic: int, dims: list, payload: bytes) -> None: + """Write one IDX file with a big-endian header.""" + header = struct.pack(">" + "I" * (1 + len(dims)), magic, *dims) + path.write_bytes(header + payload) + + +def main() -> int: + """Write the image and label fixtures.""" + OUT_DIR.mkdir(parents=True, exist_ok=True) + flat = bytes(px for image in IMAGES for px in image) + write_idx( + OUT_DIR / "mnist_tiny_images.idx", 2051, + [len(IMAGES), 2, 2], flat, + ) + write_idx( + OUT_DIR / "mnist_tiny_labels.idx", 2049, + [len(LABELS)], bytes(LABELS), + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/patch_nir_weights.py b/tools/patch_nir_weights.py new file mode 100644 index 0000000..f418418 --- /dev/null +++ b/tools/patch_nir_weights.py @@ -0,0 +1,102 @@ +"""Patch a graph_summary JSON with weights from a C checkpoint. + +Stdlib-only reverse-bridge helper (#30): reads a `graph_summary` +document (spikeforge.nir_bridge.exporter.graph_summary) and a +`spikeforge-c-checkpoint` file, and writes a summary whose weighted +nodes carry the checkpoint's flat row-major weights re-nested to the +node's shape. The output feeds the same consumers as any summary -- +notably a fresh `nir_graph_load`, or a NIR graph rebuild on the +Python side. + +Name, kind, and shape mismatches fail loudly; nothing is silently +reinterpreted. Nodes the checkpoint does not mention keep their +summary weights. + +Usage: + patch_nir_weights.py summary.json checkpoint.json -o out.json +""" + +import argparse +import json +import sys +from typing import Any + + +def nest(flat: list, shape: list) -> Any: + """Re-nest a flat row-major list to `shape`.""" + if len(shape) == 1: + if len(flat) != shape[0]: + raise ValueError(f"weight has {len(flat)} values, not {shape}") + return list(flat) + stride = 1 + for dim in shape[1:]: + stride *= dim + if len(flat) != shape[0] * stride: + raise ValueError(f"weight has {len(flat)} values, not {shape}") + return [ + nest(flat[i * stride:(i + 1) * stride], shape[1:]) + for i in range(shape[0]) + ] + + +def patch(summary: dict, checkpoint: dict) -> dict: + """Return a copy of `summary` with checkpoint weights applied.""" + if checkpoint.get("format") != "spikeforge-c-checkpoint": + raise ValueError("not a spikeforge-c-checkpoint document") + if checkpoint.get("version") != 1: + raise ValueError( + f"unsupported checkpoint version {checkpoint.get('version')!r}" + ) + nodes = {node["name"]: node for node in summary["nodes"]} + patched = 0 + for entry in checkpoint["nodes"]: + name = entry["name"] + if name not in nodes: + raise ValueError(f"checkpoint names unknown node {name!r}") + node = nodes[name] + if node["kind"] != entry["kind"]: + raise ValueError( + f"node {name!r}: kind {entry['kind']!r} != " + f"{node['kind']!r}" + ) + node["params"]["weight"] = nest(entry["weight"], entry["shape"]) + bias = entry.get("bias") + if bias is None: + if node["params"].get("bias") is not None and node["kind"] in ( + "Affine", + ): + raise ValueError(f"node {name!r}: checkpoint drops a bias") + node["params"]["bias"] = None + else: + node["params"]["bias"] = list(bias) + patched += 1 + if patched == 0: + raise ValueError("checkpoint carries no nodes") + return summary + + +def main(argv: list) -> int: + """Patch a summary file with a checkpoint file.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("summary", help="graph_summary JSON file") + parser.add_argument("checkpoint", help="checkpoint JSON file") + parser.add_argument("-o", "--output", required=True, + help="patched summary output path") + args = parser.parse_args(argv) + try: + with open(args.summary, encoding="utf-8") as handle: + summary = json.load(handle) + with open(args.checkpoint, encoding="utf-8") as handle: + checkpoint = json.load(handle) + result = patch(summary, checkpoint) + except (OSError, ValueError, KeyError) as error: + print(f"patch_nir_weights: {error}", file=sys.stderr) + return 1 + with open(args.output, "w", encoding="utf-8") as handle: + json.dump(result, handle, indent=2) + handle.write("\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:]))