diff --git a/CMakeLists.txt b/CMakeLists.txt index 4ab1aa7..cabae61 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -27,6 +27,7 @@ add_library(nir_c_runtime src/data.c src/checkpoint.c src/minibatch.c + src/tape.c ) target_include_directories(nir_c_runtime PUBLIC include) target_link_libraries(nir_c_runtime PUBLIC m) diff --git a/include/nir_c_runtime/tape.h b/include/nir_c_runtime/tape.h new file mode 100644 index 0000000..0fcb580 --- /dev/null +++ b/include/nir_c_runtime/tape.h @@ -0,0 +1,169 @@ +/* Graph-level backprop-through-time tape (#39). + * + * The layer backward ops in ops.h cover one step of one node; the + * caller would have to hand-compose them, and only for single-step + * topologies. The tape closes that gap for a whole exported graph: + * it runs a multi-step forward pass that retains every step's + * per-node summed inputs, values, and neuron state, then walks the + * graph backward in reverse topological order, unrolling Delay and + * neuron state through time and accumulating parameter gradients + * per node. + * + * The forward pass reuses the runtime's Delay-breaking scheduler + * (an internal nir_runtime_t, never stepped, supplies the exact + * order and cycle validation) and evaluates the same public layer + * ops, so its readout matches nir_runtime_run bit for bit. One + * deliberate difference: every tape forward starts from unprimed + * temporal state (no retained membranes or delay buffers across + * forwards), since a training step pairs one fresh forward with + * one backward. + * + * Backward semantics: + * + * - The Output node's time-averaging backpropagates as grad/T + * into every step's output value. + * - Delay nodes route their step-t value gradient to the + * contributing sources' step-(t-1) value gradients, mirroring + * the forward buffer refresh. + * - LI nodes unroll exactly: each step's membrane gradient is + * the consumer gradient plus the next step's previous-state + * gradient. + * - LIF/IF/CubaLIF follow the spiking-neuron convention in + * ops.h: the forward hard reset carries zero gradient, so no + * credit flows through the stored membrane across steps, and + * the spike path carries the surrogate gradient within each + * step. The tape wires the previous-state gradients into + * grad_stored faithfully (where shapes agree); the ops ignore + * them. CubaLIF's synaptic-current recurrence has no gradient + * input slot on nir_op_cuba_backward, so cross-step synaptic + * credit is truncated as well. + * + * The tape owns all retained tensors and gradient buffers; the + * graph (and its parameter arrays) is borrowed and must outlive + * the tape. Parameter gradients are exposed read-only for the + * caller to feed nir_sgd_step/nir_adam_step. Inference APIs are + * untouched: this module only adds new symbols. + */ +#ifndef NIR_C_RUNTIME_TAPE_H +#define NIR_C_RUNTIME_TAPE_H + +#include + +#include "nir_c_runtime/graph.h" +#include "nir_c_runtime/ops.h" +#include "nir_c_runtime/runtime.h" +#include "nir_c_runtime/tensor.h" + +typedef struct { + const nir_graph_t *graph; /* borrowed; must outlive the tape */ + /* Scheduler reuse: supplies the Delay-breaking topological + * order and cycle validation. Never stepped. */ + nir_runtime_t sched; + size_t num_steps; + nir_surrogate_kind_t surrogate; + float slope; + /* Retained forward pass, indexed [step * node_count + node]: + * summed inputs, node values, post-step neuron state (LI, + * LIF, IF, CubaLIF), and CubaLIF synaptic current. Slots a + * node never fills stay empty (data == NULL); an unprimed + * Delay's step-0 value is one such slot, mirroring the + * runtime leaving it unset. */ + nir_tensor_t *inputs; + nir_tensor_t *values; + nir_tensor_t *stored; + nir_tensor_t *syn; + /* Backward buffers, same indexing: dL/dvalue per step and + * node, plus dL/dstored-state routed from the next step for + * LI/LIF/IF/CubaLIF nodes. */ + nir_tensor_t *grads; + nir_tensor_t *state_grads; + /* Accumulated parameter gradients, indexed by node: + * grad_weight for Affine/Linear/Conv2d/Conv1d (Conv weights + * are flat vectors in params->weight order, as documented + * on the layer backward ops), grad_bias for biased nodes, + * grad_scale[i] for Scale nodes. Unweighted nodes keep an + * empty grad_weight (data == NULL); unbiased nodes an empty + * grad_bias; non-Scale nodes a zero grad_scale. */ + nir_tensor_t *grad_weight; + nir_tensor_t *grad_bias; + float *grad_scale; + int forwarded; /* a forward pass is currently retained */ + int have_grads; /* a backward pass has populated the grads */ +} nir_tape_t; + +/* Bind `graph` for `num_steps`-step runs: validate the graph + * through the shared scheduler (unschedulable graphs fail here + * with the runtime's cycle error), allocate the retained and + * gradient buffers, and size the parameter-gradient slots from + * the node parameters. `graph` must outlive `tape`. + * `num_steps` must be nonzero; the surrogate/slope pair must + * pass nir_surrogate_grad's validation (known kind, finite + * positive slope). Returns 0 and sets *error on failure. */ +int nir_tape_init( + nir_tape_t *tape, const nir_graph_t *graph, size_t num_steps, + nir_surrogate_kind_t surrogate, float slope, char **error +); + +/* Free every retained and gradient buffer plus the internal + * scheduler, and zero `tape` out. Safe on an already-freed or + * zeroed tape. */ +void nir_tape_free(nir_tape_t *tape); + +/* Run `frames[0..num_steps)` through the graph from unprimed + * temporal state, retaining every step's per-node inputs, + * values, and neuron state, and return the time-averaged + * Output-node value in `readout_out` (caller frees it with + * nir_tensor_free), matching nir_runtime_run on a fresh + * runtime. Any previously retained pass (and its gradients) + * is discarded. Returns 0 and sets *error on failure, + * including the runtime's own input/shape errors. */ +int nir_tape_forward( + nir_tape_t *tape, const nir_tensor_t *frames, + nir_tensor_t *readout_out, char **error +); + +/* Backpropagate `grad_readout` (dL/dreadout, shaped exactly as + * the retained readout) through the retained pass: seed each + * step's output gradient with grad/T, walk steps newest to + * oldest and nodes in reverse topological order, unroll Delay + * and neuron state across steps, and accumulate the parameter + * gradients. Requires a retained forward pass; any previous + * gradients are zeroed first, so backward may run again after + * a new forward (or with a new grad_readout). Returns 0 and + * sets *error on failure. */ +int nir_tape_backward( + nir_tape_t *tape, const nir_tensor_t *grad_readout, + char **error +); + +/* Read-only parameter gradients after a successful backward. + * The weight gradient for an Affine/Linear/Conv2d/Conv1d + * node, or NULL for any other kind, an out-of-range index, + * or a tape with no gradients. */ +const nir_tensor_t *nir_tape_weight_grad( + const nir_tape_t *tape, size_t node_index +); + +/* The bias gradient for a biased Affine/Conv2d/Conv1d node, + * else NULL under the same conditions as above. */ +const nir_tensor_t *nir_tape_bias_grad( + const nir_tape_t *tape, size_t node_index +); + +/* The accumulated scale gradient for a Scale node + * (sum over steps of sum(grad_output * input), per + * nir_op_scale_backward). Writes *out and returns 1, or + * returns 0 for any other kind, an out-of-range index, a + * tape with no gradients, or a NULL out. */ +int nir_tape_scale_grad( + const nir_tape_t *tape, size_t node_index, float *out +); + +/* dL/dframes[step]: the retained input node's value gradient + * at that step, or NULL for an out-of-range step or a tape + * with no gradients. */ +const nir_tensor_t *nir_tape_input_grad( + const nir_tape_t *tape, size_t step +); + +#endif /* NIR_C_RUNTIME_TAPE_H */ diff --git a/src/tape.c b/src/tape.c new file mode 100644 index 0000000..b50355e --- /dev/null +++ b/src/tape.c @@ -0,0 +1,1166 @@ +#include "nir_c_runtime/tape.h" + +#include +#include +#include + +/* The message is always data: copied verbatim, never interpreted + * as a printf format string. */ +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); + } +} + +/* Join a caller-supplied name into a message by concatenation, so the + * name can never act as a printf conversion. */ +static void set_error_arg( + char **error, const char *prefix, const char *arg, const char *suffix +) { + if (error == NULL || *error != NULL) { + return; + } + if (arg == NULL) { + arg = ""; + } + size_t prefix_length = strlen(prefix); + size_t arg_length = strlen(arg); + size_t suffix_length = strlen(suffix); + *error = (char *)malloc(prefix_length + arg_length + suffix_length + 1); + if (*error != NULL) { + memcpy(*error, prefix, prefix_length); + memcpy(*error + prefix_length, arg, arg_length); + memcpy( + *error + prefix_length + arg_length, suffix, suffix_length + 1 + ); + } +} + +static int checked_mul(size_t a, size_t b, size_t *out) { + if (b != 0 && a > SIZE_MAX / b) { + return 0; + } + *out = a * b; + return 1; +} + +/* A Delay source feeds its consumers from the previous step's + * buffer, which only exists from step 1 on: every tape forward + * starts from unprimed temporal state. */ +static int source_contributes( + const nir_graph_t *graph, size_t source, size_t step +) { + return graph->nodes[source].kind != NIR_NODE_DELAY || step > 0; +} + +static int tensor_valid(const nir_tensor_t *tensor) { + if (tensor == NULL || tensor->data == NULL || + tensor->ndim > NIR_TENSOR_MAX_DIMS) { + return 0; + } + size_t count = 1; + for (size_t i = 0; i < tensor->ndim; i++) { + if (tensor->shape[i] != 0 && count > SIZE_MAX / tensor->shape[i]) { + return 0; + } + count *= tensor->shape[i]; + } + return count == tensor->size && tensor->capacity >= tensor->size; +} + +static int tensor_same_shape( + const nir_tensor_t *a, const nir_tensor_t *b +) { + if (a->ndim != b->ndim) { + return 0; + } + for (size_t i = 0; i < a->ndim; i++) { + if (a->shape[i] != b->shape[i]) { + return 0; + } + } + return 1; +} + +static void free_slots(nir_tensor_t *slots, size_t count) { + if (slots == NULL) { + return; + } + for (size_t i = 0; i < count; i++) { + nir_tensor_free(&slots[i]); + } +} + +/* Retained-slot count. Safe: num_steps is only stored after the + * checked multiplication in nir_tape_init succeeds. */ +static size_t tape_slots(const nir_tape_t *tape) { + if (tape->graph == NULL) { + return 0; + } + return tape->num_steps * tape->graph->node_count; +} + +static void clear_retained(nir_tape_t *tape) { + size_t slots = tape_slots(tape); + free_slots(tape->inputs, slots); + free_slots(tape->values, slots); + free_slots(tape->stored, slots); + free_slots(tape->syn, slots); + free_slots(tape->grads, slots); + free_slots(tape->state_grads, slots); + tape->forwarded = 0; + tape->have_grads = 0; +} + +/* Parameter-gradient element counts for one node: weight_count + * is 0 for unweighted kinds, bias_count is 0 for unbiased + * ones. Fails only on arithmetic overflow. */ +static int grad_counts( + const nir_node_t *node, size_t *weight_count, size_t *bias_count +) { + size_t w = 0, b = 0; + switch (node->kind) { + case NIR_NODE_AFFINE: + if (!checked_mul( + node->params.affine.out_features, + node->params.affine.in_features, &w)) { + return 0; + } + if (node->params.affine.bias != NULL) { + b = node->params.affine.out_features; + } + break; + case NIR_NODE_LINEAR: + if (!checked_mul( + node->params.linear.out_features, + node->params.linear.in_features, &w)) { + return 0; + } + break; + case NIR_NODE_CONV2D: { + const nir_conv2d_params_t *p = &node->params.conv2d; + w = p->out_channels; + if (!checked_mul(w, p->in_channels, &w) || + !checked_mul(w, p->kernel_h, &w) || + !checked_mul(w, p->kernel_w, &w)) { + return 0; + } + if (p->bias != NULL) { + b = p->out_channels; + } + break; + } + case NIR_NODE_CONV1D: { + const nir_conv1d_params_t *p = &node->params.conv1d; + w = p->out_channels; + if (!checked_mul(w, p->in_channels, &w) || + !checked_mul(w, p->kernel, &w)) { + return 0; + } + if (p->bias != NULL) { + b = p->out_channels; + } + break; + } + default: + break; + } + *weight_count = w; + *bias_count = b; + return 1; +} + +static int alloc_param_grads(nir_tape_t *tape, char **error) { + size_t n = tape->graph->node_count; + for (size_t i = 0; i < n; i++) { + size_t w = 0, b = 0; + if (!grad_counts(&tape->graph->nodes[i], &w, &b)) { + set_error(error, "tape parameter block is too large"); + return 0; + } + if (w > 0) { + size_t shape[1] = {w}; + tape->grad_weight[i] = nir_tensor_alloc(shape, 1); + if (tape->grad_weight[i].data == NULL) { + set_error(error, "out of memory"); + return 0; + } + } + if (b > 0) { + size_t shape[1] = {b}; + tape->grad_bias[i] = nir_tensor_alloc(shape, 1); + if (tape->grad_bias[i].data == NULL) { + set_error(error, "out of memory"); + return 0; + } + } + } + return 1; +} + +int nir_tape_init( + nir_tape_t *tape, const nir_graph_t *graph, size_t num_steps, + nir_surrogate_kind_t surrogate, float slope, char **error +) { + float probe; + if (tape == NULL) { + set_error(error, "invalid tape arguments"); + return 0; + } + memset(tape, 0, sizeof(*tape)); + if (graph == NULL || num_steps == 0 || + !nir_surrogate_grad(surrogate, slope, 0.0f, &probe)) { + set_error(error, "invalid tape arguments"); + return 0; + } + /* Scheduler reuse: this validates the graph exactly as + * inference would (including the Delay-breaking cycle + * check) and computes the order the tape walks. */ + if (!nir_runtime_init(&tape->sched, graph, error)) { + memset(tape, 0, sizeof(*tape)); + return 0; + } + size_t slots = 0; + if (!checked_mul(num_steps, graph->node_count, &slots)) { + nir_tape_free(tape); + set_error(error, "tape run is too large"); + return 0; + } + tape->graph = graph; + tape->num_steps = num_steps; + tape->surrogate = surrogate; + tape->slope = slope; + size_t n = graph->node_count; + tape->inputs = (nir_tensor_t *)calloc( + slots == 0 ? 1 : slots, sizeof(nir_tensor_t) + ); + tape->values = (nir_tensor_t *)calloc( + slots == 0 ? 1 : slots, sizeof(nir_tensor_t) + ); + tape->stored = (nir_tensor_t *)calloc( + slots == 0 ? 1 : slots, sizeof(nir_tensor_t) + ); + tape->syn = (nir_tensor_t *)calloc( + slots == 0 ? 1 : slots, sizeof(nir_tensor_t) + ); + tape->grads = (nir_tensor_t *)calloc( + slots == 0 ? 1 : slots, sizeof(nir_tensor_t) + ); + tape->state_grads = (nir_tensor_t *)calloc( + slots == 0 ? 1 : slots, sizeof(nir_tensor_t) + ); + tape->grad_weight = (nir_tensor_t *)calloc( + n == 0 ? 1 : n, sizeof(nir_tensor_t) + ); + tape->grad_bias = (nir_tensor_t *)calloc( + n == 0 ? 1 : n, sizeof(nir_tensor_t) + ); + tape->grad_scale = (float *)calloc(n == 0 ? 1 : n, sizeof(float)); + if (tape->inputs == NULL || tape->values == NULL || + tape->stored == NULL || tape->syn == NULL || + tape->grads == NULL || tape->state_grads == NULL || + tape->grad_weight == NULL || tape->grad_bias == NULL || + tape->grad_scale == NULL) { + nir_tape_free(tape); + set_error(error, "out of memory"); + return 0; + } + if (!alloc_param_grads(tape, error)) { + nir_tape_free(tape); + return 0; + } + return 1; +} + +void nir_tape_free(nir_tape_t *tape) { + if (tape == NULL) { + return; + } + clear_retained(tape); + size_t n = tape->graph != NULL ? tape->graph->node_count : 0; + free_slots(tape->grad_weight, n); + free_slots(tape->grad_bias, n); + free(tape->inputs); + free(tape->values); + free(tape->stored); + free(tape->syn); + free(tape->grads); + free(tape->state_grads); + free(tape->grad_weight); + free(tape->grad_bias); + free(tape->grad_scale); + nir_runtime_free(&tape->sched); + memset(tape, 0, sizeof(*tape)); +} + +/* Sum the retained step-`step` values feeding `target`, mirroring + * runtime.c's sum_incoming: an unprimed Delay source (any Delay + * at step 0) contributes nothing, and a target with no + * contributing source sums to a single-element zero. */ +static int sum_incoming( + nir_tape_t *tape, size_t step, size_t target, nir_tensor_t *out, + char **error +) { + const nir_graph_t *graph = tape->graph; + size_t n = graph->node_count; + int have_any = 0, found_edge = 0; + memset(out, 0, sizeof(*out)); + for (size_t e = 0; e < graph->edge_count; e++) { + if (graph->edges[e].target_index != target) { + continue; + } + found_edge = 1; + size_t src = graph->edges[e].source_index; + if (!source_contributes(graph, src, step)) { + continue; + } + const nir_tensor_t *value = &tape->values[step * n + src]; + if (!have_any) { + *out = nir_tensor_clone(value); + if (out->data == NULL) { + set_error(error, "out of memory or invalid node tensor"); + return 0; + } + have_any = 1; + } else if (!nir_tensor_add_inplace(out, value)) { + nir_tensor_free(out); + set_error(error, "shape mismatch summing node inputs"); + return 0; + } + } + if (!found_edge) { + set_error_arg(error, "node ", graph->nodes[target].name, + " received no inputs"); + return 0; + } + if (!have_any) { + size_t one[1] = {1}; + *out = nir_tensor_alloc(one, 1); /* zero-initialized scalar */ + if (out->data == NULL) { + set_error(error, "out of memory"); + return 0; + } + } + return 1; +} + +static nir_tensor_t alloc_like(const nir_tensor_t *input) { + if (input == NULL) { + return (nir_tensor_t){0}; + } + return nir_tensor_alloc(input->shape, input->ndim); +} + +/* The previous step's retained membrane for a stateful node, or + * NULL on the first step (treated as zeros, matching the + * runtime's unprimed state). */ +static const nir_tensor_t *prev_state( + nir_tape_t *tape, size_t step, size_t index +) { + if (step == 0) { + return NULL; + } + size_t n = tape->graph->node_count; + return &tape->stored[(step - 1) * n + index]; +} + +static int apply_li( + nir_tape_t *tape, size_t step, size_t index, + const nir_tensor_t *input, char **error +) { + const nir_node_t *node = &tape->graph->nodes[index]; + size_t n = tape->graph->node_count; + nir_tensor_t membrane = nir_op_li( + &node->params.li, input, prev_state(tape, step, index) + ); + if (membrane.data == NULL) { + set_error(error, "node operation failed or out of memory"); + return 0; + } + nir_tensor_t next = nir_tensor_clone(&membrane); + if (next.data == NULL) { + nir_tensor_free(&membrane); + set_error(error, "node operation failed or out of memory"); + return 0; + } + tape->values[step * n + index] = membrane; + tape->stored[step * n + index] = next; + return 1; +} + +static int apply_spiking( + nir_tape_t *tape, size_t step, size_t index, + const nir_tensor_t *input, char **error +) { + const nir_node_t *node = &tape->graph->nodes[index]; + size_t n = tape->graph->node_count; + const nir_tensor_t *prev = prev_state(tape, step, index); + nir_tensor_t spike = alloc_like(input); + nir_tensor_t stored = alloc_like(input); + int ok = spike.data != NULL && stored.data != NULL; + if (ok && node->kind == NIR_NODE_LIF) { + ok = nir_op_lif_into( + &node->params.lif, input, prev, &spike, &stored + ); + } else if (ok) { + ok = nir_op_if_into( + &node->params.if_, input, prev, &spike, &stored + ); + } + if (!ok) { + nir_tensor_free(&spike); + nir_tensor_free(&stored); + set_error(error, "node operation failed or out of memory"); + return 0; + } + tape->values[step * n + index] = spike; + tape->stored[step * n + index] = stored; + return 1; +} + +static int apply_cuba( + nir_tape_t *tape, size_t step, size_t index, + const nir_tensor_t *input, char **error +) { + const nir_node_t *node = &tape->graph->nodes[index]; + size_t n = tape->graph->node_count; + const nir_tensor_t *syn_prev = NULL; + const nir_tensor_t *mem_prev = prev_state(tape, step, index); + if (step > 0) { + syn_prev = &tape->syn[(step - 1) * n + index]; + } + nir_tensor_t spike = alloc_like(input); + nir_tensor_t syn = alloc_like(input); + nir_tensor_t stored = alloc_like(input); + int ok = spike.data != NULL && syn.data != NULL && + stored.data != NULL && nir_op_cuba_into( + &node->params.cuba, input, syn_prev, mem_prev, + &spike, &syn, &stored + ); + if (!ok) { + nir_tensor_free(&spike); + nir_tensor_free(&syn); + nir_tensor_free(&stored); + set_error(error, "node operation failed or out of memory"); + return 0; + } + tape->values[step * n + index] = spike; + tape->syn[step * n + index] = syn; + tape->stored[step * n + index] = stored; + return 1; +} + +/* Evaluate one non-Input, non-Delay node from its retained + * summed input, retaining the value (and post-step state for + * stateful kinds). The op sequence mirrors runtime.c's + * apply_node exactly. */ +static int apply_node( + nir_tape_t *tape, size_t step, size_t index, + const nir_tensor_t *input, char **error +) { + const nir_node_t *node = &tape->graph->nodes[index]; + size_t n = tape->graph->node_count; + nir_tensor_t *value = &tape->values[step * n + index]; + switch (node->kind) { + case NIR_NODE_CONV2D: + *value = nir_op_conv2d(&node->params.conv2d, input); + break; + case NIR_NODE_CONV1D: + *value = nir_op_conv1d(&node->params.conv1d, input); + break; + case NIR_NODE_AFFINE: + *value = nir_op_affine(&node->params.affine, input); + break; + case NIR_NODE_LINEAR: + *value = nir_op_linear(&node->params.linear, input); + break; + case NIR_NODE_FLATTEN: + *value = nir_op_flatten(input); + break; + case NIR_NODE_AVGPOOL2D: + *value = nir_op_avgpool2d(&node->params.pool2d, input); + break; + case NIR_NODE_SUMPOOL2D: + *value = nir_op_sumpool2d(&node->params.pool2d, input); + break; + case NIR_NODE_LI: + return apply_li(tape, step, index, input, error); + case NIR_NODE_LIF: + case NIR_NODE_IF: + return apply_spiking(tape, step, index, input, error); + case NIR_NODE_CUBALIF: + return apply_cuba(tape, step, index, input, error); + case NIR_NODE_THRESHOLD: + *value = nir_op_threshold(&node->params.threshold, input); + break; + case NIR_NODE_SCALE: + *value = nir_op_scale(&node->params.scale, input); + break; + case NIR_NODE_OUTPUT: + *value = nir_tensor_clone(input); /* pass-through */ + break; + default: /* Input and Delay never reach here. */ + break; + } + if (value->data == NULL) { + set_error(error, "node operation failed or out of memory"); + return 0; + } + return 1; +} + +static int run_step( + nir_tape_t *tape, const nir_tensor_t *frame, size_t step, + char **error +) { + const nir_graph_t *graph = tape->graph; + size_t n = graph->node_count; + tape->values[step * n + graph->input_index] = + nir_tensor_clone(frame); + if (tape->values[step * n + graph->input_index].data == NULL) { + set_error(error, "invalid input tensor or out of memory"); + return 0; + } + /* A Delay's value is the previous step's retained summed + * input -- the same buffer the runtime refreshes after each + * step. At step 0 the slot stays empty (unprimed). */ + if (step > 0) { + for (size_t i = 0; i < n; i++) { + if (graph->nodes[i].kind != NIR_NODE_DELAY) { + continue; + } + tape->values[step * n + i] = nir_tensor_clone( + &tape->inputs[(step - 1) * n + i] + ); + if (tape->values[step * n + i].data == NULL) { + set_error(error, "invalid delay state or out of memory"); + return 0; + } + } + } + for (size_t o = 0; o < n; o++) { + size_t index = tape->sched.order[o]; + nir_node_kind_t kind = graph->nodes[index].kind; + if (kind == NIR_NODE_INPUT || kind == NIR_NODE_DELAY) { + continue; + } + if (!sum_incoming( + tape, step, index, + &tape->inputs[step * n + index], error) || + !apply_node( + tape, step, index, + &tape->inputs[step * n + index], error)) { + return 0; + } + } + for (size_t i = 0; i < n; i++) { + if (graph->nodes[i].kind != NIR_NODE_DELAY) { + continue; + } + if (!sum_incoming( + tape, step, i, &tape->inputs[step * n + i], + error)) { + return 0; + } + } + return 1; +} + +/* Zeroed gradient buffers shaped like the retained values and + * states they differentiate. */ +static int alloc_grad_buffers(nir_tape_t *tape, char **error) { + size_t slots = tape_slots(tape); + for (size_t s = 0; s < slots; s++) { + if (tape->values[s].data != NULL) { + tape->grads[s] = nir_tensor_alloc( + tape->values[s].shape, tape->values[s].ndim + ); + if (tape->grads[s].data == NULL) { + set_error(error, "out of memory"); + return 0; + } + } + if (tape->stored[s].data != NULL) { + tape->state_grads[s] = nir_tensor_alloc( + tape->stored[s].shape, tape->stored[s].ndim + ); + if (tape->state_grads[s].data == NULL) { + set_error(error, "out of memory"); + return 0; + } + } + } + return 1; +} + +int nir_tape_forward( + nir_tape_t *tape, const nir_tensor_t *frames, + nir_tensor_t *readout_out, char **error +) { + if (tape == NULL || tape->graph == NULL || frames == NULL || + readout_out == NULL) { + set_error(error, "invalid tape arguments"); + return 0; + } + clear_retained(tape); + const nir_graph_t *graph = tape->graph; + size_t n = graph->node_count; + for (size_t t = 0; t < tape->num_steps; t++) { + if (!run_step(tape, &frames[t], t, error)) { + clear_retained(tape); + return 0; + } + } + nir_tensor_t total = nir_tensor_clone( + &tape->values[graph->output_index] + ); + if (total.data == NULL) { + clear_retained(tape); + set_error(error, "out of memory or invalid output tensor"); + return 0; + } + for (size_t t = 1; t < tape->num_steps; t++) { + if (!nir_tensor_add_inplace( + &total, &tape->values[t * n + graph->output_index])) { + nir_tensor_free(&total); + clear_retained(tape); + set_error(error, "runtime output shape changed"); + return 0; + } + } + for (size_t i = 0; i < total.size; i++) { + total.data[i] /= (float)tape->num_steps; + } + if (!alloc_grad_buffers(tape, error)) { + nir_tensor_free(&total); + clear_retained(tape); + return 0; + } + *readout_out = total; + tape->forwarded = 1; + return 1; +} + +/* Accumulate a consumer's input gradient into a source's value + * gradient. Equal sizes add pointwise; a single-element source + * (a Delay holding the zero scalar) folds the broadcast sum, + * mirroring the forward broadcast add. Any other mismatch is + * unreachable after a successful forward. */ +static int accumulate_grad(nir_tensor_t *acc, const nir_tensor_t *g) { + if (!tensor_valid(acc) || !tensor_valid(g)) { + return 0; + } + if (acc->size == g->size) { + for (size_t i = 0; i < acc->size; i++) { + acc->data[i] += g->data[i]; + } + return 1; + } + if (acc->size == 1) { + double sum = 0.0; + for (size_t i = 0; i < g->size; i++) { + sum += (double)g->data[i]; + } + acc->data[0] += (float)sum; + return 1; + } + return 0; +} + +/* Route `grad` (dL/d of `target`'s summed input) to every source + * that contributed to that sum at step `src_step`. */ +static int route_to_sources( + nir_tape_t *tape, size_t target, size_t src_step, + const nir_tensor_t *grad, char **error +) { + const nir_graph_t *graph = tape->graph; + size_t n = graph->node_count; + for (size_t e = 0; e < graph->edge_count; e++) { + if (graph->edges[e].target_index != target) { + continue; + } + size_t src = graph->edges[e].source_index; + if (!source_contributes(graph, src, src_step)) { + continue; + } + if (!accumulate_grad( + &tape->grads[src_step * n + src], grad)) { + set_error(error, "gradient shape mismatch routing node inputs"); + return 0; + } + } + return 1; +} + +/* The layer backward ops overwrite their grad_input, so one + * reusable scratch tensor serves every node; it only grows. */ +static int scratch_for(nir_tensor_t *scratch, size_t need) { + if (scratch->data != NULL && scratch->capacity >= need) { + return 1; + } + nir_tensor_free(scratch); + size_t shape[1] = {need}; + *scratch = nir_tensor_alloc(shape, 1); + return scratch->data != NULL; +} + +static const char *backward_error_for(nir_node_kind_t kind) { + switch (kind) { + case NIR_NODE_AFFINE: + return "affine node backward pass failed"; + case NIR_NODE_LINEAR: + return "linear node backward pass failed"; + case NIR_NODE_CONV2D: + return "conv2d node backward pass failed"; + case NIR_NODE_CONV1D: + return "conv1d node backward pass failed"; + case NIR_NODE_FLATTEN: + return "flatten node backward pass failed"; + case NIR_NODE_AVGPOOL2D: + return "avgpool2d node backward pass failed"; + case NIR_NODE_SUMPOOL2D: + return "sumpool2d node backward pass failed"; + case NIR_NODE_THRESHOLD: + return "threshold node backward pass failed"; + case NIR_NODE_SCALE: + return "scale node backward pass failed"; + case NIR_NODE_LI: + return "LI node backward pass failed"; + case NIR_NODE_LIF: + return "LIF node backward pass failed"; + case NIR_NODE_IF: + return "IF node backward pass failed"; + case NIR_NODE_CUBALIF: + return "CubaLIF node backward pass failed"; + default: + return "node backward pass failed"; + } +} + +/* This step's routed stored-state gradient, for the spiking + * ops' grad_stored tape wiring -- or NULL when the slot is + * empty or shape-shifted (defensive; the ops ignore it under + * the reset convention either way). */ +static const nir_tensor_t *stored_grad_or_null( + nir_tape_t *tape, size_t step, size_t index +) { + size_t n = tape->graph->node_count; + const nir_tensor_t *sg = &tape->state_grads[step * n + index]; + const nir_tensor_t *input = &tape->inputs[step * n + index]; + if (sg->data != NULL && tensor_same_shape(sg, input)) { + return sg; + } + return NULL; +} + +static int backward_spiking( + nir_tape_t *tape, size_t step, size_t index, + nir_tensor_t *scratch, char **error +) { + const nir_node_t *node = &tape->graph->nodes[index]; + size_t n = tape->graph->node_count; + const nir_tensor_t *input = &tape->inputs[step * n + index]; + const nir_tensor_t *gv = &tape->grads[step * n + index]; + const nir_tensor_t *v_prev = + step > 0 ? &tape->stored[(step - 1) * n + index] : NULL; + nir_tensor_t *gprev = step > 0 ? + &tape->state_grads[(step - 1) * n + index] : NULL; + const nir_tensor_t *gstored = + stored_grad_or_null(tape, step, index); + if (!scratch_for(scratch, input->size)) { + set_error(error, "out of memory"); + return 0; + } + int ok = 0; + if (node->kind == NIR_NODE_LIF) { + ok = nir_op_lif_backward( + &node->params.lif, input, v_prev, gv, gstored, + tape->surrogate, tape->slope, scratch, gprev + ); + } else if (node->kind == NIR_NODE_IF) { + ok = nir_op_if_backward( + &node->params.if_, input, v_prev, gv, gstored, + tape->surrogate, tape->slope, scratch, gprev + ); + } else { + const nir_tensor_t *syn_prev = step > 0 ? + &tape->syn[(step - 1) * n + index] : NULL; + /* nir_op_cuba_backward has no gradient input for the + * synaptic current, so cross-step synaptic credit is + * truncated here (see tape.h). */ + ok = nir_op_cuba_backward( + &node->params.cuba, input, syn_prev, v_prev, gv, + gstored, tape->surrogate, tape->slope, scratch, + NULL, gprev + ); + } + if (!ok) { + set_error(error, backward_error_for(node->kind)); + return 0; + } + return route_to_sources(tape, index, step, scratch, error); +} + +static int backward_li( + nir_tape_t *tape, size_t step, size_t index, + nir_tensor_t *scratch, char **error +) { + const nir_node_t *node = &tape->graph->nodes[index]; + size_t n = tape->graph->node_count; + nir_tensor_t *gv = &tape->grads[step * n + index]; + nir_tensor_t *sg = &tape->state_grads[step * n + index]; + nir_tensor_t *sprev = step > 0 ? + &tape->state_grads[(step - 1) * n + index] : NULL; + /* The membrane gradient is the consumer gradient plus the + * next step's previous-state gradient; both are final by + * now (reverse step order), so fold in place. */ + if (!nir_tensor_add_inplace(gv, sg) || + !scratch_for(scratch, gv->size) || + !nir_op_li_backward(&node->params.li, gv, scratch, sprev)) { + set_error(error, backward_error_for(node->kind)); + return 0; + } + return route_to_sources(tape, index, step, scratch, error); +} + +static int backward_affine( + nir_tape_t *tape, size_t step, size_t index, + nir_tensor_t *scratch, char **error +) { + const nir_node_t *node = &tape->graph->nodes[index]; + size_t n = tape->graph->node_count; + const nir_tensor_t *input = &tape->inputs[step * n + index]; + const nir_tensor_t *gv = &tape->grads[step * n + index]; + int ok = 0; + if (!scratch_for(scratch, input->size)) { + set_error(error, "out of memory"); + return 0; + } + if (node->kind == NIR_NODE_AFFINE) { + nir_tensor_t *gb = node->params.affine.bias != NULL ? + &tape->grad_bias[index] : NULL; + ok = nir_op_affine_backward( + &node->params.affine, input, gv, + &tape->grad_weight[index], gb, scratch + ); + } else { + ok = nir_op_linear_backward( + &node->params.linear, input, gv, + &tape->grad_weight[index], scratch + ); + } + if (!ok) { + set_error(error, backward_error_for(node->kind)); + return 0; + } + return route_to_sources(tape, index, step, scratch, error); +} + +static int backward_conv( + nir_tape_t *tape, size_t step, size_t index, + nir_tensor_t *scratch, char **error +) { + const nir_node_t *node = &tape->graph->nodes[index]; + size_t n = tape->graph->node_count; + const nir_tensor_t *input = &tape->inputs[step * n + index]; + const nir_tensor_t *gv = &tape->grads[step * n + index]; + int ok = 0; + if (!scratch_for(scratch, input->size)) { + set_error(error, "out of memory"); + return 0; + } + if (node->kind == NIR_NODE_CONV2D) { + nir_tensor_t *gb = node->params.conv2d.bias != NULL ? + &tape->grad_bias[index] : NULL; + ok = nir_op_conv2d_backward( + &node->params.conv2d, input, gv, + &tape->grad_weight[index], gb, scratch + ); + } else { + nir_tensor_t *gb = node->params.conv1d.bias != NULL ? + &tape->grad_bias[index] : NULL; + ok = nir_op_conv1d_backward( + &node->params.conv1d, input, gv, + &tape->grad_weight[index], gb, scratch + ); + } + if (!ok) { + set_error(error, backward_error_for(node->kind)); + return 0; + } + return route_to_sources(tape, index, step, scratch, error); +} + +static int backward_stateless( + nir_tape_t *tape, size_t step, size_t index, + nir_tensor_t *scratch, char **error +) { + const nir_node_t *node = &tape->graph->nodes[index]; + size_t n = tape->graph->node_count; + const nir_tensor_t *input = &tape->inputs[step * n + index]; + const nir_tensor_t *gv = &tape->grads[step * n + index]; + int ok = 0; + if (!scratch_for(scratch, input->size)) { + set_error(error, "out of memory"); + return 0; + } + switch (node->kind) { + case NIR_NODE_FLATTEN: + ok = nir_op_flatten_backward(input, gv, scratch); + break; + case NIR_NODE_AVGPOOL2D: + ok = nir_op_avgpool2d_backward( + &node->params.pool2d, input, gv, scratch + ); + break; + case NIR_NODE_SUMPOOL2D: + ok = nir_op_sumpool2d_backward( + &node->params.pool2d, input, gv, scratch + ); + break; + case NIR_NODE_THRESHOLD: + ok = nir_op_threshold_backward( + &node->params.threshold, input, gv, + tape->surrogate, tape->slope, scratch + ); + break; + case NIR_NODE_SCALE: + ok = nir_op_scale_backward( + &node->params.scale, input, gv, scratch, + &tape->grad_scale[index] + ); + break; + default: + break; + } + if (!ok) { + set_error(error, backward_error_for(node->kind)); + return 0; + } + return route_to_sources(tape, index, step, scratch, error); +} + +static int backward_node( + nir_tape_t *tape, size_t step, size_t index, + nir_tensor_t *scratch, char **error +) { + nir_node_kind_t kind = tape->graph->nodes[index].kind; + switch (kind) { + case NIR_NODE_INPUT: + return 1; + case NIR_NODE_DELAY: + if (step == 0) { + return 1; + } + /* The step-t value is the previous step's buffer: + * route its gradient to the step-(t-1) sources. */ + return route_to_sources( + tape, index, step - 1, + &tape->grads[ + step * tape->graph->node_count + index + ], + error + ); + case NIR_NODE_OUTPUT: + return route_to_sources( + tape, index, step, + &tape->grads[ + step * tape->graph->node_count + index + ], + error + ); + case NIR_NODE_AFFINE: + case NIR_NODE_LINEAR: + return backward_affine(tape, step, index, scratch, error); + case NIR_NODE_CONV2D: + case NIR_NODE_CONV1D: + return backward_conv(tape, step, index, scratch, error); + case NIR_NODE_FLATTEN: + case NIR_NODE_AVGPOOL2D: + case NIR_NODE_SUMPOOL2D: + case NIR_NODE_THRESHOLD: + case NIR_NODE_SCALE: + return backward_stateless( + tape, step, index, scratch, error + ); + case NIR_NODE_LI: + return backward_li(tape, step, index, scratch, error); + case NIR_NODE_LIF: + case NIR_NODE_IF: + case NIR_NODE_CUBALIF: + return backward_spiking( + tape, step, index, scratch, error + ); + default: + set_error(error, "tape encountered an invalid node kind"); + return 0; + } +} + +static void zero_tensor_data(nir_tensor_t *tensor) { + if (tensor->data != NULL) { + memset(tensor->data, 0, tensor->size * sizeof(float)); + } +} + +static void zero_backward_buffers(nir_tape_t *tape) { + size_t slots = tape_slots(tape); + for (size_t s = 0; s < slots; s++) { + zero_tensor_data(&tape->grads[s]); + zero_tensor_data(&tape->state_grads[s]); + } + for (size_t i = 0; i < tape->graph->node_count; i++) { + zero_tensor_data(&tape->grad_weight[i]); + zero_tensor_data(&tape->grad_bias[i]); + } + memset( + tape->grad_scale, 0, + tape->graph->node_count * sizeof(float) + ); +} + +/* Seed one step's output-value gradient with grad/T. A + * single-element step output (a step whose output summed only + * unprimed-delay scalars) folds the broadcast sum, mirroring + * the forward averaging. */ +static int seed_output_step( + nir_tensor_t *g, const nir_tensor_t *grad_readout, float inv, + char **error +) { + if (!tensor_valid(g)) { + set_error(error, "readout gradient shape mismatch"); + return 0; + } + if (g->size == grad_readout->size) { + for (size_t k = 0; k < g->size; k++) { + g->data[k] += grad_readout->data[k] * inv; + } + return 1; + } + if (g->size == 1) { + double sum = 0.0; + for (size_t k = 0; k < grad_readout->size; k++) { + sum += (double)grad_readout->data[k]; + } + g->data[0] += (float)sum * inv; + return 1; + } + set_error(error, "readout gradient shape mismatch"); + return 0; +} + +int nir_tape_backward( + nir_tape_t *tape, const nir_tensor_t *grad_readout, + char **error +) { + if (tape == NULL || tape->graph == NULL || + grad_readout == NULL) { + set_error(error, "invalid tape arguments"); + return 0; + } + if (!tape->forwarded) { + set_error(error, "tape has no retained forward pass"); + return 0; + } + const nir_graph_t *graph = tape->graph; + size_t n = graph->node_count; + if (!tensor_valid(grad_readout) || !tensor_same_shape( + grad_readout, &tape->values[graph->output_index])) { + set_error(error, "readout gradient shape mismatch"); + return 0; + } + zero_backward_buffers(tape); + tape->have_grads = 0; + float inv = 1.0f / (float)tape->num_steps; + for (size_t t = 0; t < tape->num_steps; t++) { + if (!seed_output_step( + &tape->grads[t * n + graph->output_index], + grad_readout, inv, error)) { + return 0; + } + } + nir_tensor_t scratch = {0}; + for (size_t ti = tape->num_steps; ti-- > 0;) { + for (size_t oi = n; oi-- > 0;) { + if (!backward_node( + tape, ti, tape->sched.order[oi], &scratch, + error)) { + nir_tensor_free(&scratch); + return 0; + } + } + } + nir_tensor_free(&scratch); + tape->have_grads = 1; + return 1; +} + +static int is_weighted(nir_node_kind_t kind) { + return kind == NIR_NODE_AFFINE || kind == NIR_NODE_LINEAR || + kind == NIR_NODE_CONV2D || kind == NIR_NODE_CONV1D; +} + +static int node_has_bias(const nir_node_t *node) { + switch (node->kind) { + case NIR_NODE_AFFINE: + return node->params.affine.bias != NULL; + case NIR_NODE_CONV2D: + return node->params.conv2d.bias != NULL; + case NIR_NODE_CONV1D: + return node->params.conv1d.bias != NULL; + default: + return 0; + } +} + +const nir_tensor_t *nir_tape_weight_grad( + const nir_tape_t *tape, size_t node_index +) { + if (tape == NULL || tape->graph == NULL || !tape->have_grads || + node_index >= tape->graph->node_count || + !is_weighted(tape->graph->nodes[node_index].kind)) { + return NULL; + } + return &tape->grad_weight[node_index]; +} + +const nir_tensor_t *nir_tape_bias_grad( + const nir_tape_t *tape, size_t node_index +) { + if (tape == NULL || tape->graph == NULL || !tape->have_grads || + node_index >= tape->graph->node_count || + !node_has_bias(&tape->graph->nodes[node_index])) { + return NULL; + } + return &tape->grad_bias[node_index]; +} + +int nir_tape_scale_grad( + const nir_tape_t *tape, size_t node_index, float *out +) { + if (tape == NULL || tape->graph == NULL || out == NULL || + !tape->have_grads || + node_index >= tape->graph->node_count || + tape->graph->nodes[node_index].kind != NIR_NODE_SCALE) { + return 0; + } + *out = tape->grad_scale[node_index]; + return 1; +} + +const nir_tensor_t *nir_tape_input_grad( + const nir_tape_t *tape, size_t step +) { + if (tape == NULL || tape->graph == NULL || !tape->have_grads || + step >= tape->num_steps) { + return NULL; + } + return &tape->grads[ + step * tape->graph->node_count + tape->graph->input_index + ]; +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 38ae4a9..770627a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -27,6 +27,7 @@ nir_add_test(test_encode) nir_add_test(test_data) nir_add_test(test_checkpoint) nir_add_test(test_minibatch) +nir_add_test(test_tape) # 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/tape_conv1d_graph.json b/tests/fixtures/tape_conv1d_graph.json new file mode 100644 index 0000000..b832e67 --- /dev/null +++ b/tests/fixtures/tape_conv1d_graph.json @@ -0,0 +1,81 @@ +{ + "nodes": [ + { + "name": "input", + "kind": "Input", + "params": {} + }, + { + "name": "conv", + "kind": "Conv1d", + "params": { + "weight": [ + [ + [0.5, -0.25], + [0.75, 0.125] + ], + [ + [-0.5, 0.5], + [0.25, -0.75] + ], + [ + [0.125, 0.375], + [-0.625, 0.875] + ] + ], + "bias": [0.1, -0.1, 0.0], + "stride": 1, + "padding": 0, + "dilation": 1 + } + }, + { + "name": "fl", + "kind": "Flatten", + "params": {} + }, + { + "name": "lin", + "kind": "Linear", + "params": { + "weight": [ + [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, -0.75, -0.5], + [-0.25, 0.0, 0.25, 0.5, 0.75, -0.75, -0.5, -0.25, 0.0], + [0.25, 0.5, 0.75, -0.75, -0.5, -0.25, 0.0, 0.25, 0.5], + [0.75, -0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, -0.75] + ] + } + }, + { + "name": "sc", + "kind": "Scale", + "params": { + "scale": 0.5 + } + }, + { + "name": "fc", + "kind": "Affine", + "params": { + "weight": [ + [-1.0, -0.5, 0.0, 0.5], + [1.0, -1.0, -0.5, 0.0] + ], + "bias": [0.25, -0.25] + } + }, + { + "name": "output", + "kind": "Output", + "params": {} + } + ], + "edges": [ + {"source": "input", "target": "conv", "delayed": false}, + {"source": "conv", "target": "fl", "delayed": false}, + {"source": "fl", "target": "lin", "delayed": false}, + {"source": "lin", "target": "sc", "delayed": false}, + {"source": "sc", "target": "fc", "delayed": false}, + {"source": "fc", "target": "output", "delayed": false} + ] +} diff --git a/tests/fixtures/tape_conv_graph.json b/tests/fixtures/tape_conv_graph.json new file mode 100644 index 0000000..2e95e59 --- /dev/null +++ b/tests/fixtures/tape_conv_graph.json @@ -0,0 +1,80 @@ +{ + "nodes": [ + { + "name": "input", + "kind": "Input", + "params": {} + }, + { + "name": "conv", + "kind": "Conv2d", + "params": { + "weight": [ + [ + [ + [0.5, -0.25], + [0.75, 0.125] + ] + ], + [ + [ + [-0.5, 0.5], + [0.25, -0.75] + ] + ] + ], + "bias": [0.1, -0.1], + "stride": [1, 1], + "padding": [0, 0], + "dilation": [1, 1] + } + }, + { + "name": "avg", + "kind": "AvgPool2d", + "params": { + "kernel_size": [2, 2], + "stride": [2, 2], + "padding": [0, 0] + } + }, + { + "name": "sum", + "kind": "SumPool2d", + "params": { + "kernel_size": [2, 2], + "stride": [2, 2], + "padding": [0, 0] + } + }, + { + "name": "fl", + "kind": "Flatten", + "params": {} + }, + { + "name": "fc", + "kind": "Affine", + "params": { + "weight": [ + [0.4, 0.3], + [0.1, -0.4] + ], + "bias": [0.0, 0.05] + } + }, + { + "name": "output", + "kind": "Output", + "params": {} + } + ], + "edges": [ + {"source": "input", "target": "conv", "delayed": false}, + {"source": "conv", "target": "avg", "delayed": false}, + {"source": "avg", "target": "sum", "delayed": false}, + {"source": "sum", "target": "fl", "delayed": false}, + {"source": "fl", "target": "fc", "delayed": false}, + {"source": "fc", "target": "output", "delayed": false} + ] +} diff --git a/tests/fixtures/tape_delay_li_graph.json b/tests/fixtures/tape_delay_li_graph.json new file mode 100644 index 0000000..e6f9721 --- /dev/null +++ b/tests/fixtures/tape_delay_li_graph.json @@ -0,0 +1,69 @@ +{ + "nodes": [ + { + "name": "input", + "kind": "Input", + "params": {} + }, + { + "name": "fc1", + "kind": "Affine", + "params": { + "weight": [ + [0.5, -0.25], + [0.75, 0.125], + [-0.5, 0.5] + ], + "bias": [0.1, -0.1, 0.0] + } + }, + { + "name": "li1", + "kind": "LI", + "params": { + "tau": 2.0, + "r": 1.0, + "v_leak": 0.0 + } + }, + { + "name": "dly", + "kind": "Delay", + "params": {} + }, + { + "name": "fc2", + "kind": "Affine", + "params": { + "weight": [ + [0.4, 0.3, -0.2], + [0.1, -0.4, 0.6] + ], + "bias": [0.0, 0.05] + } + }, + { + "name": "li2", + "kind": "LI", + "params": { + "tau": 3.0, + "r": 0.5, + "v_leak": 0.1 + } + }, + { + "name": "output", + "kind": "Output", + "params": {} + } + ], + "edges": [ + {"source": "input", "target": "fc1", "delayed": false}, + {"source": "fc1", "target": "li1", "delayed": false}, + {"source": "li1", "target": "fc2", "delayed": false}, + {"source": "li1", "target": "dly", "delayed": false}, + {"source": "dly", "target": "fc2", "delayed": false}, + {"source": "fc2", "target": "li2", "delayed": false}, + {"source": "li2", "target": "output", "delayed": false} + ] +} diff --git a/tests/fixtures/tape_if_graph.json b/tests/fixtures/tape_if_graph.json new file mode 100644 index 0000000..0987145 --- /dev/null +++ b/tests/fixtures/tape_if_graph.json @@ -0,0 +1,51 @@ +{ + "nodes": [ + { + "name": "input", + "kind": "Input", + "params": {} + }, + { + "name": "fc1", + "kind": "Affine", + "params": { + "weight": [ + [0.5, -0.5], + [0.25, 0.75] + ], + "bias": [0.0, 0.1] + } + }, + { + "name": "neu", + "kind": "IF", + "params": { + "r": 2.0, + "v_threshold": 1.0, + "v_reset": 0.0 + } + }, + { + "name": "fc2", + "kind": "Affine", + "params": { + "weight": [ + [1.0, -1.0], + [-0.5, 0.5] + ], + "bias": [0.1, -0.1] + } + }, + { + "name": "output", + "kind": "Output", + "params": {} + } + ], + "edges": [ + {"source": "input", "target": "fc1", "delayed": false}, + {"source": "fc1", "target": "neu", "delayed": false}, + {"source": "neu", "target": "fc2", "delayed": false}, + {"source": "fc2", "target": "output", "delayed": false} + ] +} diff --git a/tests/fixtures/tape_scalar_delay_graph.json b/tests/fixtures/tape_scalar_delay_graph.json new file mode 100644 index 0000000..761fc92 --- /dev/null +++ b/tests/fixtures/tape_scalar_delay_graph.json @@ -0,0 +1,32 @@ +{ + "nodes": [ + { + "name": "input", + "kind": "Input", + "params": {} + }, + { + "name": "dly", + "kind": "Delay", + "params": {} + }, + { + "name": "fc1", + "kind": "Affine", + "params": { + "weight": [[2.0]], + "bias": [0.5] + } + }, + { + "name": "output", + "kind": "Output", + "params": {} + } + ], + "edges": [ + {"source": "input", "target": "dly", "delayed": false}, + {"source": "dly", "target": "fc1", "delayed": false}, + {"source": "fc1", "target": "output", "delayed": false} + ] +} diff --git a/tests/fixtures/tape_spike_graph.json b/tests/fixtures/tape_spike_graph.json new file mode 100644 index 0000000..723a98e --- /dev/null +++ b/tests/fixtures/tape_spike_graph.json @@ -0,0 +1,86 @@ +{ + "nodes": [ + { + "name": "input", + "kind": "Input", + "params": {} + }, + { + "name": "fc1", + "kind": "Affine", + "params": { + "weight": [ + [0.5, -0.25], + [0.75, 0.125], + [-0.5, 0.5] + ], + "bias": [0.1, -0.1, 0.0] + } + }, + { + "name": "thr", + "kind": "Threshold", + "params": { + "threshold": 0.5 + } + }, + { + "name": "lif", + "kind": "LIF", + "params": { + "tau": 2.0, + "r": 1.0, + "v_leak": 0.0, + "v_threshold": 1.0, + "v_reset": 0.0 + } + }, + { + "name": "ifn", + "kind": "IF", + "params": { + "r": 1.0, + "v_threshold": 1.0, + "v_reset": 0.0 + } + }, + { + "name": "cuba", + "kind": "CubaLIF", + "params": { + "tau_syn": 2.0, + "tau_mem": 3.0, + "r": 1.0, + "v_leak": 0.0, + "v_threshold": 1.0, + "v_reset": 0.0, + "w_in": 1.0 + } + }, + { + "name": "fc2", + "kind": "Affine", + "params": { + "weight": [ + [0.4, 0.3, -0.2], + [0.1, -0.4, 0.6] + ], + "bias": [0.0, 0.05] + } + }, + { + "name": "output", + "kind": "Output", + "params": {} + } + ], + "edges": [ + {"source": "input", "target": "fc1", "delayed": false}, + {"source": "fc1", "target": "thr", "delayed": false}, + {"source": "thr", "target": "lif", "delayed": false}, + {"source": "lif", "target": "ifn", "delayed": false}, + {"source": "ifn", "target": "cuba", "delayed": false}, + {"source": "cuba", "target": "fc2", "delayed": false}, + {"source": "fc2", "target": "output", "delayed": false} + ] +} diff --git a/tests/fixtures/tape_step_graph.json b/tests/fixtures/tape_step_graph.json new file mode 100644 index 0000000..72adf11 --- /dev/null +++ b/tests/fixtures/tape_step_graph.json @@ -0,0 +1,69 @@ +{ + "nodes": [ + { + "name": "input", + "kind": "Input", + "params": {} + }, + { + "name": "fc1", + "kind": "Affine", + "params": { + "weight": [ + [0.5, -0.5], + [0.25, 0.75] + ], + "bias": [0.0, 0.1] + } + }, + { + "name": "thr", + "kind": "Threshold", + "params": { + "threshold": 0.5 + } + }, + { + "name": "lif", + "kind": "LIF", + "params": { + "tau": 2.0, + "r": 3.0, + "v_leak": 0.0, + "v_threshold": 1.0, + "v_reset": 0.0 + } + }, + { + "name": "sc", + "kind": "Scale", + "params": { + "scale": 2.0 + } + }, + { + "name": "fc2", + "kind": "Affine", + "params": { + "weight": [ + [1.0, -1.0], + [-0.5, 0.5] + ], + "bias": [0.1, -0.1] + } + }, + { + "name": "output", + "kind": "Output", + "params": {} + } + ], + "edges": [ + {"source": "input", "target": "fc1", "delayed": false}, + {"source": "fc1", "target": "thr", "delayed": false}, + {"source": "thr", "target": "lif", "delayed": false}, + {"source": "lif", "target": "sc", "delayed": false}, + {"source": "sc", "target": "fc2", "delayed": false}, + {"source": "fc2", "target": "output", "delayed": false} + ] +} diff --git a/tests/test_tape.c b/tests/test_tape.c new file mode 100644 index 0000000..51a42cf --- /dev/null +++ b/tests/test_tape.c @@ -0,0 +1,1005 @@ +/* Graph-level BPTT tape (#39). + * + * Forward retention matches nir_runtime_run bit for bit (small + * Delay+LI graph plus the real exported conv_net with Delay + * reset cycles). Multi-step gradients match central finite + * differences on smooth topologies covering every stateless + * kind (Affine, Linear, Conv2d, Conv1d, Flatten, AvgPool2d, + * SumPool2d, Scale), LI recurrence, and Delay edges. The + * single-step spiking chains (Threshold, LIF, IF, Scale) match + * float64 closed forms of the documented surrogate rules, and + * a scalar Delay graph proves credit flows into earlier steps + * (its last frame has exactly zero gradient). + */ +#include +#include +#include + +#include "nir_c_runtime/graph.h" +#include "nir_c_runtime/json.h" +#include "nir_c_runtime/loss.h" +#include "nir_c_runtime/runtime.h" +#include "nir_c_runtime/tape.h" +#include "nir_c_runtime/train.h" +#include "test_util.h" + +static nir_graph_t load_graph(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 nir_tensor_t make_vector(const float *values, size_t n) { + size_t shape[1] = {n}; + nir_tensor_t t = nir_tensor_alloc(shape, 1); + CHECK(t.data != NULL); + memcpy(t.data, values, n * sizeof(float)); + return t; +} + +static void check_vector( + const nir_tensor_t *t, const double *want, size_t n, double tol +) { + CHECK(t->size == n); + for (size_t i = 0; i < n; i++) { + CHECK_NEAR(t->data[i], want[i], tol); + } +} + +static void check_finite(const nir_tensor_t *t) { + CHECK(t != NULL && t->data != NULL); + for (size_t i = 0; i < t->size; i++) { + CHECK(isfinite(t->data[i])); + } +} + +static unsigned lcg_next(unsigned *state) { + *state = *state * 1103515245u + 12345u; + return (*state >> 16) & 0x7fffu; +} + +static double fast_sigmoid(double slope, double x) { + double d = 1.0 + fabs(slope * x); + return slope / (d * d); +} + +/* One fresh-runtime MSE evaluation: the finite-difference + * oracle. A new runtime per call keeps temporal state + * unprimed, matching the tape's fresh forward. */ +static float run_loss( + nir_graph_t *graph, nir_tensor_t *frames, size_t steps, + const nir_tensor_t *target +) { + char *error = NULL; + nir_runtime_t runtime; + CHECK(nir_runtime_init(&runtime, graph, &error)); + CHECK(error == NULL); + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(nir_runtime_run(&runtime, frames, steps, &readout, &error)); + CHECK(error == NULL); + size_t shape[1] = {readout.size}; + nir_tensor_t grad = nir_tensor_alloc(shape, 1); + CHECK(grad.data != NULL); + float loss = 0.0f; + CHECK(nir_mse_loss(&readout, target, &loss, &grad)); + nir_tensor_free(&readout); + nir_tensor_free(&grad); + nir_runtime_free(&runtime); + return loss; +} + +/* Flat parameter views of a weighted node for FD perturbation: + * weight count (0 when unweighted), plus the bias view (NULL/0 + * when the node has no bias). */ +static void node_param_views( + nir_node_t *node, float **w, size_t *nw, float **b, size_t *nb +) { + *w = NULL; + *nw = 0; + *b = NULL; + *nb = 0; + switch (node->kind) { + case NIR_NODE_AFFINE: + *w = node->params.affine.weight; + *nw = node->params.affine.out_features * + node->params.affine.in_features; + *b = node->params.affine.bias; + if (*b != NULL) { + *nb = node->params.affine.out_features; + } + break; + case NIR_NODE_LINEAR: + *w = node->params.linear.weight; + *nw = node->params.linear.out_features * + node->params.linear.in_features; + break; + case NIR_NODE_CONV2D: + *w = node->params.conv2d.weight; + *nw = node->params.conv2d.out_channels * + node->params.conv2d.in_channels * + node->params.conv2d.kernel_h * + node->params.conv2d.kernel_w; + *b = node->params.conv2d.bias; + if (*b != NULL) { + *nb = node->params.conv2d.out_channels; + } + break; + case NIR_NODE_CONV1D: + *w = node->params.conv1d.weight; + *nw = node->params.conv1d.out_channels * + node->params.conv1d.in_channels * + node->params.conv1d.kernel; + *b = node->params.conv1d.bias; + if (*b != NULL) { + *nb = node->params.conv1d.out_channels; + } + break; + default: + break; + } +} + +static void check_fd_buffer( + nir_graph_t *graph, nir_tensor_t *frames, size_t steps, + const nir_tensor_t *target, float *buf, const float *analytic, + size_t count, double eps, double tol +) { + for (size_t k = 0; k < count; k++) { + float save = buf[k]; + buf[k] = (float)(save + eps); + float up = run_loss(graph, frames, steps, target); + buf[k] = (float)(save - eps); + float down = run_loss(graph, frames, steps, target); + buf[k] = save; + CHECK_NEAR(analytic[k], (up - down) / (2.0 * eps), tol); + } +} + +/* Central-difference check of every parameter gradient, every + * Scale gradient, and every input-frame gradient against the + * runtime forward. */ +static void check_fd( + nir_graph_t *graph, nir_tape_t *tape, nir_tensor_t *frames, + size_t steps, const nir_tensor_t *target, double eps, double tol +) { + for (size_t i = 0; i < graph->node_count; i++) { + nir_node_t *node = &graph->nodes[i]; + if (node->kind == NIR_NODE_SCALE) { + float got = 0.0f; + CHECK(nir_tape_scale_grad(tape, i, &got)); + float save = node->params.scale.scale; + node->params.scale.scale = (float)(save + eps); + float up = run_loss(graph, frames, steps, target); + node->params.scale.scale = (float)(save - eps); + float down = run_loss(graph, frames, steps, target); + node->params.scale.scale = save; + CHECK_NEAR(got, (up - down) / (2.0 * eps), tol); + continue; + } + float *w = NULL, *b = NULL; + size_t nw = 0, nb = 0; + node_param_views(node, &w, &nw, &b, &nb); + if (nw == 0) { + continue; + } + const nir_tensor_t *gw = nir_tape_weight_grad(tape, i); + CHECK(gw != NULL && gw->size == nw); + check_fd_buffer( + graph, frames, steps, target, w, gw->data, nw, eps, + tol + ); + if (nb > 0) { + const nir_tensor_t *gb = nir_tape_bias_grad(tape, i); + CHECK(gb != NULL && gb->size == nb); + check_fd_buffer( + graph, frames, steps, target, b, gb->data, nb, + eps, tol + ); + } + } + for (size_t t = 0; t < steps; t++) { + const nir_tensor_t *gi = nir_tape_input_grad(tape, t); + CHECK(gi != NULL && gi->size == frames[t].size); + check_fd_buffer( + graph, frames, steps, target, frames[t].data, + gi->data, frames[t].size, eps, tol + ); + } +} + +/* Run one graph through a fresh runtime and through the tape, + * and require bit-identical readouts. */ +static void check_forward_exact( + nir_graph_t *graph, nir_tensor_t *frames, size_t steps +) { + char *error = NULL; + nir_runtime_t runtime; + CHECK(nir_runtime_init(&runtime, graph, &error)); + CHECK(error == NULL); + nir_tensor_t want; + memset(&want, 0, sizeof(want)); + CHECK(nir_runtime_run(&runtime, frames, steps, &want, &error)); + CHECK(error == NULL); + nir_runtime_free(&runtime); + nir_tape_t tape; + memset(&tape, 0, sizeof(tape)); + CHECK(nir_tape_init( + &tape, graph, steps, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error == NULL); + nir_tensor_t got; + memset(&got, 0, sizeof(got)); + CHECK(nir_tape_forward(&tape, frames, &got, &error)); + CHECK(error == NULL); + CHECK(got.size == want.size); + for (size_t i = 0; i < got.size; i++) { + CHECK(got.data[i] == want.data[i]); + } + nir_tensor_free(&want); + nir_tensor_free(&got); + nir_tape_free(&tape); +} + +static void test_forward_matches_runtime(void) { + nir_graph_t graph = load_graph( + "fixtures/tape_delay_li_graph.json" + ); + float fdat[4][2] = { + {0.5f, -0.25f}, + {-0.75f, 0.5f}, + {0.25f, 0.75f}, + {-0.125f, -0.5f}, + }; + nir_tensor_t frames[4]; + for (size_t t = 0; t < 4; t++) { + frames[t] = make_vector(fdat[t], 2); + } + check_forward_exact(&graph, frames, 4); + for (size_t t = 0; t < 4; t++) { + nir_tensor_free(&frames[t]); + } + nir_graph_free(&graph); +} + +/* The real exported conv_net, whose Threshold/Delay/Scale reset + * loops are genuine Delay feedback cycles. */ +static void test_forward_conv_net(void) { + char *graph_text = read_file("fixtures/conv_net_graph.json"); + char *io_text = read_file("fixtures/conv_net_io.json"); + char *error = NULL; + json_value_t *graph_json = json_parse(graph_text, &error); + CHECK(error == NULL); + json_value_t *io_json = json_parse(io_text, &error); + CHECK(error == NULL); + nir_graph_t graph; + memset(&graph, 0, sizeof(graph)); + CHECK(nir_graph_load(graph_json, &graph, &error)); + CHECK(error == NULL); + const json_value_t *shape_json = json_object_get(io_json, "shape"); + size_t steps = (size_t)json_as_number( + json_array_get(shape_json, 0), 0 + ); + size_t channels = (size_t)json_as_number( + json_array_get(shape_json, 1), 0 + ); + size_t height = (size_t)json_as_number( + json_array_get(shape_json, 2), 0 + ); + size_t width = (size_t)json_as_number( + json_array_get(shape_json, 3), 0 + ); + nir_tensor_t *frames = (nir_tensor_t *)calloc( + steps, sizeof(nir_tensor_t) + ); + CHECK(frames != NULL); + size_t shape[3] = {channels, height, width}; + const json_value_t *input_json = json_object_get(io_json, "input"); + for (size_t t = 0; t < steps; t++) { + frames[t] = nir_tensor_alloc(shape, 3); + CHECK(frames[t].data != NULL); + const json_value_t *frame = json_array_get(input_json, t); + size_t index = 0; + for (size_t c = 0; c < channels; c++) { + const json_value_t *row_c = json_array_get(frame, c); + for (size_t h = 0; h < height; h++) { + const json_value_t *row_h = json_array_get(row_c, h); + for (size_t w = 0; w < width; w++) { + frames[t].data[index++] = (float)json_as_number( + json_array_get(row_h, w), 0.0 + ); + } + } + } + } + check_forward_exact(&graph, frames, steps); + for (size_t t = 0; t < steps; t++) { + nir_tensor_free(&frames[t]); + } + free(frames); + nir_graph_free(&graph); + json_free(graph_json); + json_free(io_json); + free(graph_text); + free(io_text); +} + +static void test_fd_delay_li(void) { + nir_graph_t graph = load_graph( + "fixtures/tape_delay_li_graph.json" + ); + size_t steps = 4; + float fdat[4][2] = { + {0.5f, -0.25f}, + {-0.75f, 0.5f}, + {0.25f, 0.75f}, + {-0.125f, -0.5f}, + }; + nir_tensor_t frames[4]; + for (size_t t = 0; t < steps; t++) { + frames[t] = make_vector(fdat[t], 2); + } + float tv[] = {0.5f, -0.5f}; + nir_tensor_t target = make_vector(tv, 2); + char *error = NULL; + nir_tape_t tape; + memset(&tape, 0, sizeof(tape)); + CHECK(nir_tape_init( + &tape, &graph, steps, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error == NULL); + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(nir_tape_forward(&tape, frames, &readout, &error)); + CHECK(error == NULL); + size_t gshape[1] = {readout.size}; + nir_tensor_t dloss = nir_tensor_alloc(gshape, 1); + CHECK(dloss.data != NULL); + float loss = 0.0f; + CHECK(nir_mse_loss(&readout, &target, &loss, &dloss)); + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + /* Float32 central differences at eps=1e-3 agree to ~1e-5 + * here; 1e-3 separates rounding from formula errors. */ + check_fd(&graph, &tape, frames, steps, &target, 1e-3, 1e-3); + /* A second backward from the same retained pass + * reproduces every gradient bit for bit: buffers are + * re-zeroed, never accumulated into. */ + size_t fc1 = nir_graph_find(&graph, "fc1"); + size_t fc2 = nir_graph_find(&graph, "fc2"); + float snap_w1[6], snap_b1[3], snap_w2[6], snap_b2[2]; + memcpy(snap_w1, nir_tape_weight_grad(&tape, fc1)->data, + sizeof(snap_w1)); + memcpy(snap_b1, nir_tape_bias_grad(&tape, fc1)->data, + sizeof(snap_b1)); + memcpy(snap_w2, nir_tape_weight_grad(&tape, fc2)->data, + sizeof(snap_w2)); + memcpy(snap_b2, nir_tape_bias_grad(&tape, fc2)->data, + sizeof(snap_b2)); + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + CHECK(memcmp(snap_w1, nir_tape_weight_grad(&tape, fc1)->data, + sizeof(snap_w1)) == 0); + CHECK(memcmp(snap_b1, nir_tape_bias_grad(&tape, fc1)->data, + sizeof(snap_b1)) == 0); + CHECK(memcmp(snap_w2, nir_tape_weight_grad(&tape, fc2)->data, + sizeof(snap_w2)) == 0); + CHECK(memcmp(snap_b2, nir_tape_bias_grad(&tape, fc2)->data, + sizeof(snap_b2)) == 0); + for (size_t t = 0; t < steps; t++) { + nir_tensor_free(&frames[t]); + } + nir_tensor_free(&target); + nir_tensor_free(&readout); + nir_tensor_free(&dloss); + nir_tape_free(&tape); + nir_graph_free(&graph); +} + +static void test_delay_flow(void) { + nir_graph_t graph = load_graph( + "fixtures/tape_scalar_delay_graph.json" + ); + size_t steps = 4; + double f[4] = {0.25, -0.5, 0.75, 1.0}; + nir_tensor_t frames[4]; + for (size_t t = 0; t < steps; t++) { + float v = (float)f[t]; + frames[t] = make_vector(&v, 1); + } + float tv[] = {1.0f}; + nir_tensor_t target = make_vector(tv, 1); + char *error = NULL; + nir_tape_t tape; + memset(&tape, 0, sizeof(tape)); + CHECK(nir_tape_init( + &tape, &graph, steps, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error == NULL); + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(nir_tape_forward(&tape, frames, &readout, &error)); + CHECK(error == NULL); + size_t gshape[1] = {1}; + nir_tensor_t dloss = nir_tensor_alloc(gshape, 1); + CHECK(dloss.data != NULL); + float loss = 0.0f; + CHECK(nir_mse_loss(&readout, &target, &loss, &dloss)); + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + /* Closed form: step 0 reads the unprimed Delay as 0, later + * steps read the previous frame through it. */ + double w = 2.0, b = 0.5, tgt = 1.0; + double o0 = b; + double o1 = w * f[0] + b; + double o2 = w * f[1] + b; + double o3 = w * f[2] + b; + double ro = (o0 + o1 + o2 + o3) / 4.0; + double want_loss = (ro - tgt) * (ro - tgt); + double g = 2.0 * (ro - tgt); + CHECK_NEAR(readout.data[0], ro, 1e-6); + CHECK_NEAR(loss, want_loss, 1e-6); + size_t fc = nir_graph_find(&graph, "fc1"); + const nir_tensor_t *gw = nir_tape_weight_grad(&tape, fc); + const nir_tensor_t *gb = nir_tape_bias_grad(&tape, fc); + CHECK(gw != NULL && gb != NULL); + CHECK_NEAR(gw->data[0], g / 4.0 * (f[0] + f[1] + f[2]), 1e-6); + CHECK_NEAR(gb->data[0], g, 1e-6); + for (size_t t = 0; t < 3; t++) { + const nir_tensor_t *gi = nir_tape_input_grad(&tape, t); + CHECK(gi != NULL); + /* The only path from these frames to the loss crosses + * the Delay edge into a later step. */ + CHECK_NEAR(gi->data[0], g / 4.0 * w, 1e-6); + } + /* The last frame feeds only the Delay buffer, which no + * later step reads: exactly zero gradient. */ + const nir_tensor_t *glast = nir_tape_input_grad(&tape, 3); + CHECK(glast != NULL && glast->data[0] == 0.0f); + for (size_t t = 0; t < steps; t++) { + nir_tensor_free(&frames[t]); + } + nir_tensor_free(&target); + nir_tensor_free(&readout); + nir_tensor_free(&dloss); + nir_tape_free(&tape); + nir_graph_free(&graph); +} + +/* One-step Affine/Threshold/LIF/Scale/Affine chain against a + * float64 closed form of the documented surrogate rules. */ +static void test_step_lif(void) { + nir_graph_t graph = load_graph("fixtures/tape_step_graph.json"); + float xv[] = {1.5f, -0.25f}; + nir_tensor_t frame = make_vector(xv, 2); + float tv[] = {0.3f, 0.7f}; + nir_tensor_t target = make_vector(tv, 2); + char *error = NULL; + nir_tape_t tape; + memset(&tape, 0, sizeof(tape)); + CHECK(nir_tape_init( + &tape, &graph, 1, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error == NULL); + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(nir_tape_forward(&tape, &frame, &readout, &error)); + CHECK(error == NULL); + nir_tensor_t dloss = make_vector(tv, 2); + float loss = 0.0f; + CHECK(nir_mse_loss(&readout, &target, &loss, &dloss)); + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + double w1[2][2] = {{0.5, -0.5}, {0.25, 0.75}}; + double b1[2] = {0.0, 0.1}; + double w2[2][2] = {{1.0, -1.0}, {-0.5, 0.5}}; + double b2[2] = {0.1, -0.1}; + double x[2] = {1.5, -0.25}; + double tgt[2] = {0.3, 0.7}; + double hp[2], h[2], v[2], s[2], c[2], lg[2]; + for (size_t o = 0; o < 2; o++) { + hp[o] = w1[o][0] * x[0] + w1[o][1] * x[1] + b1[o]; + h[o] = hp[o] > 0.5 ? 1.0 : 0.0; + } + double decay = exp(-1.0 / 2.0); + for (size_t o = 0; o < 2; o++) { + v[o] = 3.0 * (1.0 - decay) * h[o]; + s[o] = v[o] > 1.0 ? 1.0 : 0.0; + c[o] = 2.0 * s[o]; + } + CHECK(s[0] == 1.0 && s[1] == 0.0); + for (size_t o = 0; o < 2; o++) { + lg[o] = w2[o][0] * c[0] + w2[o][1] * c[1] + b2[o]; + } + double dl[2] = {lg[0] - tgt[0], lg[1] - tgt[1]}; + double dc[2], ds[2], g[2], dh[2], dhp[2]; + double gscale = 0.0; + for (size_t i = 0; i < 2; i++) { + dc[i] = w2[0][i] * dl[0] + w2[1][i] * dl[1]; + gscale += dc[i] * s[i]; + ds[i] = dc[i] * 2.0; + g[i] = ds[i] * fast_sigmoid(25.0, v[i] - 1.0); + dh[i] = g[i] * 3.0 * (1.0 - decay); + dhp[i] = dh[i] * fast_sigmoid(25.0, hp[i] - 0.5); + } + double want_w1[4], want_w2[4], want_dx[2]; + for (size_t o = 0; o < 2; o++) { + for (size_t i = 0; i < 2; i++) { + want_w1[o * 2 + i] = dhp[o] * x[i]; + want_w2[o * 2 + i] = dl[o] * c[i]; + } + } + for (size_t i = 0; i < 2; i++) { + want_dx[i] = w1[0][i] * dhp[0] + w1[1][i] * dhp[1]; + } + size_t fc1 = nir_graph_find(&graph, "fc1"); + size_t fc2 = nir_graph_find(&graph, "fc2"); + size_t sc = nir_graph_find(&graph, "sc"); + check_vector(nir_tape_weight_grad(&tape, fc1), want_w1, 4, 1e-5); + check_vector(nir_tape_bias_grad(&tape, fc1), dhp, 2, 1e-5); + check_vector(nir_tape_weight_grad(&tape, fc2), want_w2, 4, 1e-5); + check_vector(nir_tape_bias_grad(&tape, fc2), dl, 2, 1e-5); + float got_scale = 0.0f; + CHECK(nir_tape_scale_grad(&tape, sc, &got_scale)); + CHECK_NEAR(got_scale, gscale, 1e-5); + check_vector(nir_tape_input_grad(&tape, 0), want_dx, 2, 1e-5); + nir_tensor_free(&frame); + nir_tensor_free(&target); + nir_tensor_free(&readout); + nir_tensor_free(&dloss); + nir_tape_free(&tape); + nir_graph_free(&graph); +} + +static void test_step_if(void) { + nir_graph_t graph = load_graph("fixtures/tape_if_graph.json"); + float xv[] = {1.5f, -0.25f}; + nir_tensor_t frame = make_vector(xv, 2); + float tv[] = {0.3f, 0.7f}; + nir_tensor_t target = make_vector(tv, 2); + char *error = NULL; + nir_tape_t tape; + memset(&tape, 0, sizeof(tape)); + CHECK(nir_tape_init( + &tape, &graph, 1, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error == NULL); + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(nir_tape_forward(&tape, &frame, &readout, &error)); + CHECK(error == NULL); + nir_tensor_t dloss = make_vector(tv, 2); + float loss = 0.0f; + CHECK(nir_mse_loss(&readout, &target, &loss, &dloss)); + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + double w1[2][2] = {{0.5, -0.5}, {0.25, 0.75}}; + double b1[2] = {0.0, 0.1}; + double w2[2][2] = {{1.0, -1.0}, {-0.5, 0.5}}; + double b2[2] = {0.1, -0.1}; + double x[2] = {1.5, -0.25}; + double tgt[2] = {0.3, 0.7}; + double hp[2], m[2], s[2], lg[2]; + for (size_t o = 0; o < 2; o++) { + hp[o] = w1[o][0] * x[0] + w1[o][1] * x[1] + b1[o]; + m[o] = 2.0 * hp[o]; + s[o] = m[o] > 1.0 ? 1.0 : 0.0; + } + CHECK(s[0] == 1.0 && s[1] == 0.0); + for (size_t o = 0; o < 2; o++) { + lg[o] = w2[o][0] * s[0] + w2[o][1] * s[1] + b2[o]; + } + double dl[2] = {lg[0] - tgt[0], lg[1] - tgt[1]}; + double ds[2], g[2], dh[2]; + for (size_t i = 0; i < 2; i++) { + ds[i] = w2[0][i] * dl[0] + w2[1][i] * dl[1]; + g[i] = ds[i] * fast_sigmoid(25.0, m[i] - 1.0); + dh[i] = g[i] * 2.0; + } + double want_w1[4], want_w2[4], want_dx[2]; + for (size_t o = 0; o < 2; o++) { + for (size_t i = 0; i < 2; i++) { + want_w1[o * 2 + i] = dh[o] * x[i]; + want_w2[o * 2 + i] = dl[o] * s[i]; + } + } + for (size_t i = 0; i < 2; i++) { + want_dx[i] = w1[0][i] * dh[0] + w1[1][i] * dh[1]; + } + size_t fc1 = nir_graph_find(&graph, "fc1"); + size_t fc2 = nir_graph_find(&graph, "fc2"); + check_vector(nir_tape_weight_grad(&tape, fc1), want_w1, 4, 1e-5); + check_vector(nir_tape_bias_grad(&tape, fc1), dh, 2, 1e-5); + check_vector(nir_tape_weight_grad(&tape, fc2), want_w2, 4, 1e-5); + check_vector(nir_tape_bias_grad(&tape, fc2), dl, 2, 1e-5); + check_vector(nir_tape_input_grad(&tape, 0), want_dx, 2, 1e-5); + nir_tensor_free(&frame); + nir_tensor_free(&target); + nir_tensor_free(&readout); + nir_tensor_free(&dloss); + nir_tape_free(&tape); + nir_graph_free(&graph); +} + +static void fill_lcg(nir_tensor_t *t, unsigned *state) { + for (size_t i = 0; i < t->size; i++) { + t->data[i] = + ((float)(int)(lcg_next(state) % 2001u) - 1000.0f) / + 2000.0f; + } +} + +static void test_fd_conv(void) { + nir_graph_t graph = load_graph( + "fixtures/tape_conv_graph.json" + ); + size_t steps = 2; + size_t shape[3] = {1, 5, 5}; + nir_tensor_t frames[2]; + unsigned state = 0xC0FFEEu; + for (size_t t = 0; t < steps; t++) { + frames[t] = nir_tensor_alloc(shape, 3); + CHECK(frames[t].data != NULL); + fill_lcg(&frames[t], &state); + } + float tv[] = {0.3f, -0.2f}; + nir_tensor_t target = make_vector(tv, 2); + char *error = NULL; + nir_tape_t tape; + memset(&tape, 0, sizeof(tape)); + CHECK(nir_tape_init( + &tape, &graph, steps, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error == NULL); + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(nir_tape_forward(&tape, frames, &readout, &error)); + CHECK(error == NULL); + size_t gshape[1] = {readout.size}; + nir_tensor_t dloss = nir_tensor_alloc(gshape, 1); + CHECK(dloss.data != NULL); + float loss = 0.0f; + CHECK(nir_mse_loss(&readout, &target, &loss, &dloss)); + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + check_fd(&graph, &tape, frames, steps, &target, 1e-3, 1e-3); + for (size_t t = 0; t < steps; t++) { + nir_tensor_free(&frames[t]); + } + nir_tensor_free(&target); + nir_tensor_free(&readout); + nir_tensor_free(&dloss); + nir_tape_free(&tape); + nir_graph_free(&graph); +} + +static void test_fd_conv1d(void) { + nir_graph_t graph = load_graph( + "fixtures/tape_conv1d_graph.json" + ); + size_t steps = 2; + size_t shape[2] = {2, 4}; + nir_tensor_t frames[2]; + unsigned state = 0x1DECAFEu; + for (size_t t = 0; t < steps; t++) { + frames[t] = nir_tensor_alloc(shape, 2); + CHECK(frames[t].data != NULL); + fill_lcg(&frames[t], &state); + } + float tv[] = {-0.4f, 0.6f}; + nir_tensor_t target = make_vector(tv, 2); + char *error = NULL; + nir_tape_t tape; + memset(&tape, 0, sizeof(tape)); + CHECK(nir_tape_init( + &tape, &graph, steps, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error == NULL); + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(nir_tape_forward(&tape, frames, &readout, &error)); + CHECK(error == NULL); + size_t gshape[1] = {readout.size}; + nir_tensor_t dloss = nir_tensor_alloc(gshape, 1); + CHECK(dloss.data != NULL); + float loss = 0.0f; + CHECK(nir_mse_loss(&readout, &target, &loss, &dloss)); + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + check_fd(&graph, &tape, frames, steps, &target, 1e-3, 1e-3); + for (size_t t = 0; t < steps; t++) { + nir_tensor_free(&frames[t]); + } + nir_tensor_free(&target); + nir_tensor_free(&readout); + nir_tensor_free(&dloss); + nir_tape_free(&tape); + nir_graph_free(&graph); +} + +/* Multi-step Threshold/LIF/IF/CubaLIF chain: spikes are not + * finite-differenceable, so this pins successful wiring (no + * crash, finite gradients, deterministic re-runs) while the + * single-step tests above pin exactness. */ +static void test_spike_smoke(void) { + nir_graph_t graph = load_graph( + "fixtures/tape_spike_graph.json" + ); + size_t steps = 4; + float fdat[4][2] = { + {0.5f, -0.25f}, + {-0.75f, 0.5f}, + {0.25f, 0.75f}, + {-0.125f, -0.5f}, + }; + nir_tensor_t frames[4]; + for (size_t t = 0; t < steps; t++) { + frames[t] = make_vector(fdat[t], 2); + } + float tv[] = {0.5f, -0.5f}; + nir_tensor_t target = make_vector(tv, 2); + char *error = NULL; + nir_surrogate_kind_t kind; + CHECK(nir_surrogate_from_name("atan", &kind, &error)); + CHECK(error == NULL && kind == NIR_SURROGATE_ATAN); + nir_tape_t tape; + memset(&tape, 0, sizeof(tape)); + CHECK(nir_tape_init(&tape, &graph, steps, kind, 2.0f, &error)); + CHECK(error == NULL); + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(nir_tape_forward(&tape, frames, &readout, &error)); + CHECK(error == NULL); + check_finite(&readout); + size_t gshape[1] = {readout.size}; + nir_tensor_t dloss = nir_tensor_alloc(gshape, 1); + CHECK(dloss.data != NULL); + float loss = 0.0f; + CHECK(nir_mse_loss(&readout, &target, &loss, &dloss)); + CHECK(isfinite(loss)); + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + size_t fc1 = nir_graph_find(&graph, "fc1"); + size_t fc2 = nir_graph_find(&graph, "fc2"); + check_finite(nir_tape_weight_grad(&tape, fc1)); + check_finite(nir_tape_bias_grad(&tape, fc1)); + check_finite(nir_tape_weight_grad(&tape, fc2)); + check_finite(nir_tape_bias_grad(&tape, fc2)); + for (size_t t = 0; t < steps; t++) { + check_finite(nir_tape_input_grad(&tape, t)); + } + float snap1[6], snap2[6]; + memcpy(snap1, nir_tape_weight_grad(&tape, fc1)->data, + sizeof(snap1)); + memcpy(snap2, nir_tape_weight_grad(&tape, fc2)->data, + sizeof(snap2)); + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + CHECK(memcmp(snap1, nir_tape_weight_grad(&tape, fc1)->data, + sizeof(snap1)) == 0); + CHECK(memcmp(snap2, nir_tape_weight_grad(&tape, fc2)->data, + sizeof(snap2)) == 0); + for (size_t t = 0; t < steps; t++) { + nir_tensor_free(&frames[t]); + } + nir_tensor_free(&target); + nir_tensor_free(&readout); + nir_tensor_free(&dloss); + nir_tape_free(&tape); + nir_graph_free(&graph); +} + +static void test_errors(void) { + nir_graph_t graph = load_graph( + "fixtures/tape_delay_li_graph.json" + ); + char *error = NULL; + nir_tape_t tape; + memset(&tape, 0, sizeof(tape)); + CHECK(!nir_tape_init( + NULL, &graph, 4, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_tape_init( + &tape, NULL, 4, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + nir_tape_free(&tape); + CHECK(!nir_tape_init( + &tape, &graph, 0, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + nir_tape_free(&tape); + const float bad_slopes[] = {0.0f, -2.0f, (float)INFINITY, + (float)NAN}; + for (size_t i = 0; i < 4; i++) { + CHECK(!nir_tape_init( + &tape, &graph, 4, NIR_SURROGATE_ATAN, bad_slopes[i], + &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + nir_tape_free(&tape); + } + CHECK(!nir_tape_init( + &tape, &graph, 4, (nir_surrogate_kind_t)99, 2.0f, + &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + nir_tape_free(&tape); + /* A cycle no Delay breaks fails init, like the runtime. */ + const char *cyclic = + "{\"nodes\":[" + "{\"name\":\"input\",\"kind\":\"Input\",\"params\":{}}," + "{\"name\":\"fc\",\"kind\":\"Affine\"," + "\"params\":{\"weight\":[[1.0]],\"bias\":[0.0]}}," + "{\"name\":\"output\",\"kind\":\"Output\",\"params\":{}}]," + "\"edges\":[" + "{\"source\":\"input\",\"target\":\"fc\"}," + "{\"source\":\"fc\",\"target\":\"fc\"}," + "{\"source\":\"fc\",\"target\":\"output\"}]}"; + json_value_t *cyclic_json = json_parse(cyclic, &error); + CHECK(error == NULL); + nir_graph_t loop; + memset(&loop, 0, sizeof(loop)); + CHECK(nir_graph_load(cyclic_json, &loop, &error)); + CHECK(error == NULL); + CHECK(!nir_tape_init( + &tape, &loop, 2, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error != NULL && strstr(error, "Delay") != NULL); + free(error); + error = NULL; + nir_tape_free(&tape); + json_free(cyclic_json); + nir_graph_free(&loop); + /* Sequencing on a live tape. */ + CHECK(nir_tape_init( + &tape, &graph, 4, NIR_SURROGATE_FAST_SIGMOID, 25.0f, + &error + )); + CHECK(error == NULL); + float fv[] = {0.5f, -0.25f}; + nir_tensor_t frame = make_vector(fv, 2); + nir_tensor_t readout; + memset(&readout, 0, sizeof(readout)); + CHECK(!nir_tape_forward(&tape, NULL, &readout, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_tape_forward(&tape, &frame, NULL, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(!nir_tape_forward(NULL, &frame, &readout, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + size_t fc1 = nir_graph_find(&graph, "fc1"); + size_t li1 = nir_graph_find(&graph, "li1"); + CHECK(!nir_tape_backward(&tape, &frame, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + CHECK(nir_tape_weight_grad(&tape, fc1) == NULL); + CHECK(nir_tape_bias_grad(&tape, fc1) == NULL); + float scale_out = 0.0f; + CHECK(!nir_tape_scale_grad(&tape, fc1, &scale_out)); + CHECK(nir_tape_input_grad(&tape, 0) == NULL); + /* Misshapen frames fail the forward and retain + * nothing: backward still refuses afterwards. */ + float bv[] = {1.0f, 2.0f, 3.0f}; + nir_tensor_t badframes[4]; + for (size_t t = 0; t < 4; t++) { + badframes[t] = make_vector(bv, 3); + } + CHECK(!nir_tape_forward(&tape, badframes, &readout, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + for (size_t t = 0; t < 4; t++) { + nir_tensor_free(&badframes[t]); + } + CHECK(!nir_tape_backward(&tape, &frame, &error)); + free(error); + error = NULL; + nir_tensor_t frames[4]; + for (size_t t = 0; t < 4; t++) { + frames[t] = make_vector(fv, 2); + } + CHECK(nir_tape_forward(&tape, frames, &readout, &error)); + CHECK(error == NULL); + /* Still no gradients until backward runs. */ + CHECK(nir_tape_weight_grad(&tape, fc1) == NULL); + CHECK(nir_tape_input_grad(&tape, 0) == NULL); + CHECK(!nir_tape_backward(&tape, NULL, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + float wv[] = {1.0f, 2.0f, 3.0f}; + nir_tensor_t wrong = make_vector(wv, 3); + CHECK(!nir_tape_backward(&tape, &wrong, &error)); + CHECK(error != NULL); + free(error); + error = NULL; + size_t rshape[1] = {readout.size}; + nir_tensor_t dloss = nir_tensor_alloc(rshape, 1); + CHECK(dloss.data != NULL); + for (size_t i = 0; i < dloss.size; i++) { + dloss.data[i] = 1.0f; + } + CHECK(nir_tape_backward(&tape, &dloss, &error)); + CHECK(error == NULL); + /* Accessor kind and range guards. */ + CHECK(nir_tape_weight_grad(&tape, li1) == NULL); + CHECK(nir_tape_bias_grad(&tape, li1) == NULL); + CHECK(nir_tape_weight_grad(&tape, graph.node_count) == NULL); + CHECK(nir_tape_bias_grad(&tape, graph.node_count) == NULL); + CHECK(nir_tape_weight_grad(NULL, fc1) == NULL); + CHECK(!nir_tape_scale_grad(&tape, fc1, &scale_out)); + CHECK(!nir_tape_scale_grad(&tape, graph.node_count, &scale_out)); + CHECK(!nir_tape_scale_grad(&tape, fc1, NULL)); + CHECK(nir_tape_input_grad(&tape, 4) == NULL); + CHECK(nir_tape_input_grad(NULL, 0) == NULL); + CHECK(nir_tape_input_grad(&tape, 0) != NULL); + /* A re-forward discards the old gradients. */ + nir_tensor_free(&readout); + CHECK(nir_tape_forward(&tape, frames, &readout, &error)); + CHECK(error == NULL); + CHECK(nir_tape_weight_grad(&tape, fc1) == NULL); + nir_tensor_free(&frame); + nir_tensor_free(&wrong); + for (size_t t = 0; t < 4; t++) { + nir_tensor_free(&frames[t]); + } + nir_tensor_free(&readout); + nir_tensor_free(&dloss); + nir_tape_free(&tape); + nir_tape_free(&tape); + nir_tape_free(NULL); + nir_graph_free(&graph); + memset(&tape, 0, sizeof(tape)); + nir_tape_free(&tape); +} + +int main(void) { + test_forward_matches_runtime(); + test_forward_conv_net(); + test_fd_delay_li(); + test_delay_flow(); + test_step_lif(); + test_step_if(); + test_fd_conv(); + test_fd_conv1d(); + test_spike_smoke(); + test_errors(); + return 0; +}