From a00979e065a133a8bb1fb808ae96bd5ac216cee9 Mon Sep 17 00:00:00 2001 From: w4ffl35 <25737761+w4ffl35@users.noreply.github.com> Date: Wed, 14 Oct 2026 05:12:22 -0600 Subject: [PATCH] Bind STDP to graph Affine/Linear weights; refresh README training scope (#48, #49) Add nir_stdp_graph_bind/step/free (stdp_graph.h): snapshots a named node's [out][in] weights into an owned STDP instance, transposing across the pre-major/out-major layout gap, and writes each step back through nir_graph_set_weights. Conv2d/Conv1d fail with a named error pending a flattening decision; the header documents the on-chip local-learning flow. test_stdp_graph covers the Affine round-trip (weight delta plus changed checkpoint bytes), bit-exact equivalence with standalone nir_stdp_step, Linear binding, rejects, and non-finite spike handling. README Scope now marks training IO (rng/encode/data/checkpoint) done and Out of scope links the remaining training follow-ups (#39-#47). --- CMakeLists.txt | 1 + README.md | 33 +++- include/nir_c_runtime/stdp_graph.h | 69 ++++++++ src/stdp_graph.c | 137 +++++++++++++++ tests/CMakeLists.txt | 1 + tests/test_stdp_graph.c | 272 +++++++++++++++++++++++++++++ 6 files changed, 507 insertions(+), 6 deletions(-) create mode 100644 include/nir_c_runtime/stdp_graph.h create mode 100644 src/stdp_graph.c create mode 100644 tests/test_stdp_graph.c diff --git a/CMakeLists.txt b/CMakeLists.txt index 162e141..cf8540e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -18,6 +18,7 @@ add_library(nir_c_runtime src/runtime.c src/fixed_event.c src/train.c + src/stdp_graph.c src/rng.c src/encode.c src/data.c diff --git a/README.md b/README.md index a8fdc87..75e8d4f 100644 --- a/README.md +++ b/README.md @@ -77,15 +77,36 @@ target-specific kernels remain future optimization work. - [x] C-side training math (`train.h`): named surrogate gradients (fast_sigmoid/atan, no silent default) with Threshold/LI/Affine/Linear backward ops, softmax cross-entropy, SGD/Adam optimizers over raw - weight buffers, and a local STDP rule. Dataset loading, RNG/encoding, - checkpoints, and a graph-level backprop-through-time tape remain - future work (see the training epic). + weight buffers, and a local STDP rule with a graph-bound adapter + (`stdp_graph.h`) that trains a loaded Affine/Linear node's weights + from recorded pre/post spikes. +- [x] C-side training IO: seeded RNG (`rng.h`), rate/latency/delta + spike encoders (`encode.h`), IDX MNIST dataset loading (`data.h`), + and versioned checkpoint save/load with weight accessors and + summary export (`checkpoint.h`), all covered by C tests. ## Out of scope (for now) -- Full training pipelines. The loss/optimizer/STDP primitives above are - supported; data loading, spike encoding, checkpointing, and graph-level - backprop-through-time are not built yet. +- Full training pipelines. The training math and training IO above are + supported; still missing are the graph-level backprop-through-time tape + ([#39](https://github.com/Capsize-Games/nir-c-runtime/issues/39)), + backward ops for spatial/linear + ([#40](https://github.com/Capsize-Games/nir-c-runtime/issues/40)) and + spiking-neuron + ([#41](https://github.com/Capsize-Games/nir-c-runtime/issues/41)) nodes, + mini-batch training + ([#42](https://github.com/Capsize-Games/nir-c-runtime/issues/42)), + optimizer-state checkpoint resume + ([#43](https://github.com/Capsize-Games/nir-c-runtime/issues/43)), and an + end-to-end C training driver with a worked example + ([#44](https://github.com/Capsize-Games/nir-c-runtime/issues/44)). + Weight initializers + ([#45](https://github.com/Capsize-Games/nir-c-runtime/issues/45)), + learning-rate schedules + ([#46](https://github.com/Capsize-Games/nir-c-runtime/issues/46)), and + additional loss functions + ([#47](https://github.com/Capsize-Games/nir-c-runtime/issues/47)) are filed + as follow-ups. - General Curlee/JOE graph execution. The narrow, allocation-free event ABI is integrated as the reference path for [curlee#316](https://github.com/w4ffl35/curlee/issues/316); loading arbitrary diff --git a/include/nir_c_runtime/stdp_graph.h b/include/nir_c_runtime/stdp_graph.h new file mode 100644 index 0000000..e80589f --- /dev/null +++ b/include/nir_c_runtime/stdp_graph.h @@ -0,0 +1,69 @@ +/* Graph-bound STDP: apply the train.h local-learning rule to a + * loaded graph's weights (#48). + * + * `nir_stdp_t` keeps its own weight matrix in [in][out] pre-major + * layout, while graph Affine/Linear nodes store [out][in] + * out-major weights -- so manual wiring through + * `nir_graph_get_weights`/`nir_graph_set_weights` alone would + * silently transpose the matrix. This adapter is the documented, + * tested path instead: it binds to one named node, loads that + * node's current weights (transposed) into an owned STDP instance, + * and writes every step's update back through + * `nir_graph_set_weights` (biases are left unchanged). + * + * Only Affine and Linear nodes bind. Conv2d/Conv1d carry 3D/4D + * kernels with no agreed flattening into the STDP rule's 2D + * pre/post matrix, so they fail with a named error until that + * mapping is decided. + * + * Intended on-chip local-learning flow: run the graph with + * `nir_runtime_run`, record each step's input spikes to the bound + * node (pre) and output spikes from the node or its downstream + * neuron (post), pass those vectors to `nir_stdp_graph_step`, and + * persist the trained weights with `nir_checkpoint_save`. The + * binding owns the STDP-side copy: bind snapshots the graph's + * weights, and each step overwrites the graph's weights with the + * update. The graph must outlive the binding. + */ +#ifndef NIR_C_RUNTIME_STDP_GRAPH_H +#define NIR_C_RUNTIME_STDP_GRAPH_H + +#include + +#include "nir_c_runtime/graph.h" +#include "nir_c_runtime/train.h" + +typedef struct { + nir_stdp_t stdp; + nir_graph_t *graph; /* borrowed; must outlive the binding */ + char *node; /* bound node name, malloc'd */ + float *staging; /* [out][in] write-back buffer */ + size_t in_features; + size_t out_features; +} nir_stdp_graph_t; + +/* Bind `node` (Affine or Linear) in `graph`, snapshotting its + * current weights. Hyperparameters follow `nir_stdp_init` (taus + * finite and > 0, w_min <= w_max); the dimensions come from the + * node. Returns 0 and sets *error on unknown/unweighted/ + * non-2D nodes and on invalid hyperparameters. */ +int nir_stdp_graph_bind( + nir_stdp_graph_t *b, nir_graph_t *graph, const char *node, + float tau_pre, float tau_post, float a_plus, float a_minus, + float w_min, float w_max, char **error +); + +/* One STDP step from this step's pre/post spike vectors (lengths + * in_features/out_features), written back to the graph through + * `nir_graph_set_weights`. Returns 0 on NULL or non-finite + * spikes. */ +int nir_stdp_graph_step( + nir_stdp_graph_t *b, const float *pre_spikes, + const float *post_spikes +); + +/* Free the STDP state, node name, and staging buffer, and zero + * `b` out. Safe on an already-freed or zeroed binding. */ +void nir_stdp_graph_free(nir_stdp_graph_t *b); + +#endif /* NIR_C_RUNTIME_STDP_GRAPH_H */ diff --git a/src/stdp_graph.c b/src/stdp_graph.c new file mode 100644 index 0000000..d9b8c1d --- /dev/null +++ b/src/stdp_graph.c @@ -0,0 +1,137 @@ +#include "nir_c_runtime/stdp_graph.h" + +#include +#include + +#include "nir_c_runtime/checkpoint.h" + +static void set_error(char **error, const char *message) { + if (error == NULL) return; + *error = NULL; + if (message == NULL) return; + size_t n = strlen(message) + 1; + char *copy = (char *)malloc(n); + if (copy != NULL) memcpy(copy, message, n); + *error = copy; +} + +/* Out-major [o][i] -> pre-major [i][o]. */ +static void to_premajor( + const float *src, float *dst, size_t in, size_t out +) { + for (size_t i = 0; i < in; i++) { + for (size_t o = 0; o < out; o++) { + dst[i * out + o] = src[o * in + i]; + } + } +} + +/* Pre-major [i][o] -> out-major [o][i]. */ +static void to_outmajor( + const float *src, float *dst, size_t in, size_t out +) { + for (size_t o = 0; o < out; o++) { + for (size_t i = 0; i < in; i++) { + dst[o * in + i] = src[i * out + o]; + } + } +} + +int nir_stdp_graph_bind( + nir_stdp_graph_t *b, nir_graph_t *graph, const char *node, + float tau_pre, float tau_post, float a_plus, float a_minus, + float w_min, float w_max, char **error +) { + if (error != NULL) *error = NULL; + if (b == NULL || graph == NULL || node == NULL) { + set_error(error, "STDP bind needs a binding, graph, and node"); + return 0; + } + memset(b, 0, sizeof(*b)); + size_t index = nir_graph_find(graph, node); + if (index == (size_t)-1) { + set_error(error, "STDP bind names an unknown node"); + return 0; + } + const nir_node_t *found = &graph->nodes[index]; + size_t in = 0; + size_t out = 0; + if (found->kind == NIR_NODE_AFFINE) { + in = found->params.affine.in_features; + out = found->params.affine.out_features; + } else if (found->kind == NIR_NODE_LINEAR) { + in = found->params.linear.in_features; + out = found->params.linear.out_features; + } else { + set_error(error, "STDP bind needs an Affine or Linear node"); + return 0; + } + if (!nir_stdp_init( + &b->stdp, in, out, tau_pre, tau_post, a_plus, a_minus, + w_min, w_max, error + )) { + memset(b, 0, sizeof(*b)); + return 0; + } + b->node = (char *)malloc(strlen(node) + 1); + b->staging = (float *)calloc(in * out, sizeof(float)); + if (b->node == NULL || b->staging == NULL) { + nir_stdp_graph_free(b); + set_error(error, "out of memory"); + return 0; + } + strcpy(b->node, node); + const float *weight = NULL; + const float *bias = NULL; + size_t wn = 0; + size_t bn = 0; + if (!nir_graph_get_weights( + graph, node, &weight, &wn, &bias, &bn + ) || + wn != in * out) { + nir_stdp_graph_free(b); + set_error(error, "STDP bind cannot read the node weights"); + return 0; + } + /* Staging doubles as the pre-major load buffer; set_weight + * copies it into the STDP instance (and rejects non-finite + * graph weights). */ + to_premajor(weight, b->staging, in, out); + if (!nir_stdp_set_weight(&b->stdp, b->staging)) { + nir_stdp_graph_free(b); + set_error(error, "STDP bind found non-finite node weights"); + return 0; + } + b->graph = graph; + b->in_features = in; + b->out_features = out; + return 1; +} + +int nir_stdp_graph_step( + nir_stdp_graph_t *b, const float *pre_spikes, + const float *post_spikes +) { + if (b == NULL || b->graph == NULL || b->node == NULL || + b->staging == NULL) { + return 0; + } + if (!nir_stdp_step(&b->stdp, pre_spikes, post_spikes)) { + return 0; + } + to_outmajor( + b->stdp.weight, b->staging, b->in_features, b->out_features + ); + return nir_graph_set_weights( + b->graph, b->node, b->staging, b->in_features * b->out_features, + NULL, 0, NULL + ); +} + +void nir_stdp_graph_free(nir_stdp_graph_t *b) { + if (b == NULL) return; + nir_stdp_free(&b->stdp); + free(b->node); + free(b->staging); + memset(b, 0, sizeof(*b)); +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 574811f..0cfd331 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -16,6 +16,7 @@ nir_add_test(test_workspace) nir_add_test(test_tensor) nir_add_test(test_error_format) nir_add_test(test_train) +nir_add_test(test_stdp_graph) nir_add_test(test_rng) nir_add_test(test_encode) nir_add_test(test_data) diff --git a/tests/test_stdp_graph.c b/tests/test_stdp_graph.c new file mode 100644 index 0000000..2c0d57b --- /dev/null +++ b/tests/test_stdp_graph.c @@ -0,0 +1,272 @@ +/* Graph-bound STDP (#48): the adapter trains a loaded + * Affine/Linear node's weights from pre/post spike trains, matching + * standalone nir_stdp_step exactly and moving checkpoint bytes. + */ +#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/stdp_graph.h" +#include "nir_c_runtime/train.h" +#include "test_util.h" + +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; +} + +/* fc2 in lif_linear_graph.json is Affine with out=2, in=3. */ +static void test_affine_roundtrip(void) { + nir_graph_t graph = load_fixture("fixtures/lif_linear_graph.json"); + char *error = NULL; + char *before = nir_checkpoint_dumps(&graph, NULL, &error); + CHECK(error == NULL && before != NULL); + + const float *w = NULL; + const float *bias = NULL; + size_t wn = 0; + size_t bn = 0; + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &bias, &bn)); + CHECK(wn == 6 && bn == 2); + float initial[6]; + float initial_bias[2]; + memcpy(initial, w, sizeof(initial)); + memcpy(initial_bias, bias, sizeof(initial_bias)); + + nir_stdp_graph_t bound; + memset(&bound, 0, sizeof(bound)); + CHECK(nir_stdp_graph_bind( + &bound, &graph, "fc2", 2.0f, 2.0f, 0.01f, 0.012f, 0.0f, + 1.0f, &error + )); + CHECK(error == NULL); + float pre1[] = {1.0f, 0.0f, 0.0f}; + float post1[] = {0.0f, 0.0f}; + CHECK(nir_stdp_graph_step(&bound, pre1, post1)); + float pre2[] = {0.0f, 0.0f, 0.0f}; + float post2[] = {1.0f, 0.0f}; + CHECK(nir_stdp_graph_step(&bound, pre2, post2)); + + /* Only w[out=0][in=0] potentiates: a_plus * exp(-1/tau). */ + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &bias, &bn)); + CHECK_NEAR(w[0] - initial[0], 0.01 * exp(-0.5), 1e-7); + for (size_t i = 1; i < 6; i++) CHECK(w[i] == initial[i]); + CHECK(memcmp(bias, initial_bias, sizeof(initial_bias)) == 0); + + char *after = nir_checkpoint_dumps(&graph, NULL, &error); + CHECK(error == NULL && after != NULL); + CHECK(strcmp(before, after) != 0); + + nir_stdp_graph_free(&bound); + free(before); + free(after); + nir_graph_free(&graph); +} + +/* The adapter matches standalone nir_stdp_step bit-exactly on the + * same spike trains (weights modulo the [out][in] transpose, plus + * the eligibility traces). */ +static void test_equivalence(void) { + nir_graph_t graph = load_fixture("fixtures/lif_linear_graph.json"); + char *error = NULL; + nir_stdp_graph_t bound; + memset(&bound, 0, sizeof(bound)); + CHECK(nir_stdp_graph_bind( + &bound, &graph, "fc2", 3.0f, 4.0f, 0.02f, 0.015f, -1.0f, + 1.0f, &error + )); + CHECK(error == NULL); + + nir_stdp_t solo; + CHECK(nir_stdp_init( + &solo, 3, 2, 3.0f, 4.0f, 0.02f, 0.015f, -1.0f, 1.0f, + &error + )); + CHECK(error == NULL); + CHECK(nir_stdp_set_weight(&solo, bound.stdp.weight)); + + float pre[][3] = { + {1.0f, 0.0f, 1.0f}, {0.0f, 1.0f, 0.0f}, {1.0f, 1.0f, 0.0f}, + {0.0f, 0.0f, 1.0f}, {1.0f, 0.0f, 0.0f}, {0.0f, 1.0f, 1.0f}, + }; + float post[][2] = { + {0.0f, 1.0f}, {1.0f, 0.0f}, {0.0f, 0.0f}, + {1.0f, 1.0f}, {0.0f, 1.0f}, {1.0f, 0.0f}, + }; + for (size_t t = 0; t < 6; t++) { + CHECK(nir_stdp_graph_step(&bound, pre[t], post[t])); + CHECK(nir_stdp_step(&solo, pre[t], post[t])); + } + for (size_t i = 0; i < 6; i++) { + CHECK(bound.stdp.weight[i] == solo.weight[i]); + } + for (size_t i = 0; i < 3; i++) { + CHECK(bound.stdp.pre_trace[i] == solo.pre_trace[i]); + } + for (size_t o = 0; o < 2; o++) { + CHECK(bound.stdp.post_trace[o] == solo.post_trace[o]); + } + const float *w = NULL; + const float *bias = NULL; + size_t wn = 0; + size_t bn = 0; + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &bias, &bn)); + for (size_t i = 0; i < 3; i++) { + for (size_t o = 0; o < 2; o++) { + CHECK(w[o * 3 + i] == solo.weight[i * 2 + o]); + } + } + + nir_stdp_free(&solo); + nir_stdp_graph_free(&bound); + nir_graph_free(&graph); +} + +/* fc1 is Linear (in=4, out=3, no bias): same wiring, no bias. */ +static void test_linear(void) { + nir_graph_t graph = load_fixture("fixtures/lif_linear_graph.json"); + char *error = NULL; + const float *w = NULL; + const float *bias = NULL; + size_t wn = 0; + size_t bn = 0; + CHECK(nir_graph_get_weights(&graph, "fc1", &w, &wn, &bias, &bn)); + CHECK(wn == 12 && bias == NULL && bn == 0); + float initial[12]; + memcpy(initial, w, sizeof(initial)); + + nir_stdp_graph_t bound; + memset(&bound, 0, sizeof(bound)); + CHECK(nir_stdp_graph_bind( + &bound, &graph, "fc1", 2.0f, 2.0f, 0.05f, 0.05f, -2.0f, + 2.0f, &error + )); + CHECK(error == NULL); + float pre1[] = {0.0f, 1.0f, 0.0f, 1.0f}; + float post1[] = {1.0f, 0.0f, 0.0f}; + CHECK(nir_stdp_graph_step(&bound, pre1, post1)); + float pre2[] = {1.0f, 0.0f, 1.0f, 0.0f}; + float post2[] = {0.0f, 1.0f, 1.0f}; + CHECK(nir_stdp_graph_step(&bound, pre2, post2)); + + CHECK(nir_graph_get_weights(&graph, "fc1", &w, &wn, &bias, &bn)); + CHECK(memcmp(w, initial, sizeof(initial)) != 0); + for (size_t i = 0; i < 12; i++) CHECK(isfinite(w[i])); + + nir_stdp_graph_free(&bound); + nir_graph_free(&graph); +} + +static void test_rejects(void) { + nir_graph_t graph = load_fixture("fixtures/lif_linear_graph.json"); + nir_graph_t conv = load_fixture("fixtures/conv_net_graph.json"); + char *error = NULL; + nir_stdp_graph_t bound; + memset(&bound, 0, sizeof(bound)); + + CHECK(!nir_stdp_graph_bind( + &bound, &graph, "nope", 2.0f, 2.0f, 0.01f, 0.012f, 0.0f, + 1.0f, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_stdp_graph_bind( + &bound, &graph, "lif1", 2.0f, 2.0f, 0.01f, 0.012f, 0.0f, + 1.0f, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_stdp_graph_bind( + &bound, &conv, "conv1", 2.0f, 2.0f, 0.01f, 0.012f, 0.0f, + 1.0f, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_stdp_graph_bind( + NULL, &graph, "fc2", 2.0f, 2.0f, 0.01f, 0.012f, 0.0f, + 1.0f, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_stdp_graph_bind( + &bound, NULL, "fc2", 2.0f, 2.0f, 0.01f, 0.012f, 0.0f, + 1.0f, NULL + )); + CHECK(!nir_stdp_graph_bind( + &bound, &graph, NULL, 2.0f, 2.0f, 0.01f, 0.012f, 0.0f, + 1.0f, NULL + )); + CHECK(!nir_stdp_graph_bind( + &bound, &graph, "fc2", 0.0f, 2.0f, 0.01f, 0.012f, 0.0f, + 1.0f, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + + /* A zeroed binding refuses to step; free is safe on it. */ + float pre[] = {1.0f, 0.0f, 0.0f}; + float post[] = {0.0f, 1.0f}; + CHECK(!nir_stdp_graph_step(&bound, pre, post)); + CHECK(!nir_stdp_graph_step(NULL, pre, post)); + nir_stdp_graph_free(&bound); + nir_stdp_graph_free(NULL); + + nir_graph_free(&graph); + nir_graph_free(&conv); +} + +static void test_step_rejects_nonfinite(void) { + nir_graph_t graph = load_fixture("fixtures/lif_linear_graph.json"); + char *error = NULL; + nir_stdp_graph_t bound; + memset(&bound, 0, sizeof(bound)); + CHECK(nir_stdp_graph_bind( + &bound, &graph, "fc2", 2.0f, 2.0f, 0.01f, 0.012f, 0.0f, + 1.0f, &error + )); + CHECK(error == NULL); + const float *w = NULL; + const float *bias = NULL; + size_t wn = 0; + size_t bn = 0; + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &bias, &bn)); + float initial[6]; + memcpy(initial, w, sizeof(initial)); + + float bad[] = {1.0f, (float)INFINITY, 0.0f}; + float ok[] = {0.0f, 1.0f}; + CHECK(!nir_stdp_graph_step(&bound, bad, ok)); + CHECK(!nir_stdp_graph_step(&bound, NULL, ok)); + CHECK(!nir_stdp_graph_step(&bound, ok, NULL)); + CHECK(nir_graph_get_weights(&graph, "fc2", &w, &wn, &bias, &bn)); + CHECK(memcmp(w, initial, sizeof(initial)) == 0); + + nir_stdp_graph_free(&bound); + nir_graph_free(&graph); +} + +int main(void) { + test_affine_roundtrip(); + test_equivalence(); + test_linear(); + test_rejects(); + test_step_rejects_nonfinite(); + return 0; +}