Add C-side training math: surrogate hook, backward ops, loss, optimizers

CMakeLists.txt

@@ -17,6 +17,7 @@ add_library(nir_c_runtime
   src/ops.c
   src/runtime.c
   src/fixed_event.c
+  src/train.c
 )
 target_include_directories(nir_c_runtime PUBLIC include)
 target_link_libraries(nir_c_runtime PUBLIC m)

README.md

@@ -74,10 +74,18 @@ target-specific kernels remain future optimization work.
 - [x] Allocation-free fixed-point event-module ABI for bounded, single-output
       temporal detectors (`fixed_event.h`), deterministic guard-state/action
       helpers, and an 8,192-row exhaustive C/Curlee truth-table fingerprint.
+- [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).
 
 ## Out of scope (for now)
 
-- Training. This is inference-only.
+- 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.
 - 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

include/nir_c_runtime/ops.h

@@ -166,4 +166,65 @@ int nir_op_scale_into(
     nir_tensor_t *output
 );
 
+/* Training: single-step backward ops. The forward spike
+ * nonlinearities stay strict Heavisides; the backward pass below
+ * substitutes a smooth surrogate derivative, mirroring snnTorch's
+ * `surrogate` module (see spikeforge/introspection/surrogate.py).
+ * Names resolve through nir_surrogate_from_name in train.h, which
+ * rejects unknown names instead of silently defaulting. */
+
+typedef enum {
+    NIR_SURROGATE_FAST_SIGMOID,
+    NIR_SURROGATE_ATAN
+} nir_surrogate_kind_t;
+
+/* Surrogate derivative dS/dx at x. fast_sigmoid with slope s:
+ * s / (1 + |s*x|)^2 (snnTorch default s = 25). atan with
+ * steepness a: (a/2) / (1 + (pi*a*x/2)^2) (snnTorch default
+ * a = 2). Returns 0 unless `kind` is known and `slope` is
+ * finite and positive. */
+int nir_surrogate_grad(
+    nir_surrogate_kind_t kind, float slope, float x, float *out
+);
+
+/* Backward through Threshold: forward is the strict Heaviside,
+ * and grad_input[i] = grad_output[i] * surr(membrane[i] -
+ * threshold). `membrane` is the retained pre-threshold input.
+ * Shapes must match; `grad_input` may exactly alias
+ * `grad_output` (elementwise rule). */
+int nir_op_threshold_backward(
+    const nir_threshold_params_t *params, const nir_tensor_t *membrane,
+    const nir_tensor_t *grad_output, nir_surrogate_kind_t surrogate,
+    float slope, nir_tensor_t *grad_input
+);
+
+/* Backward through one LI step: with decay = exp(-1 / tau),
+ * grad_input = grad_output * r * (1 - decay) and grad_prev =
+ * grad_output * decay. `grad_prev` may be NULL when the
+ * previous-state gradient is unneeded. Either output may exactly
+ * alias `grad_output`; the two outputs must not overlap each
+ * other. */
+int nir_op_li_backward(
+    const nir_li_params_t *params, const nir_tensor_t *grad_output,
+    nir_tensor_t *grad_input, nir_tensor_t *grad_prev
+);
+
+/* Backward through y = Wx + b: grad_weight += outer(grad_output,
+ * input), grad_bias += grad_output, grad_input = W^T grad_output.
+ * Weight and bias gradients accumulate (the caller zeroes them);
+ * `grad_bias` may be NULL when the node has no bias. All three
+ * outputs must be disjoint from the inputs and from each other. */
+int nir_op_affine_backward(
+    const nir_affine_params_t *params, const nir_tensor_t *input,
+    const nir_tensor_t *grad_output, nir_tensor_t *grad_weight,
+    nir_tensor_t *grad_bias, nir_tensor_t *grad_input
+);
+
+/* Backward through y = Wx: as above without the bias term. */
+int nir_op_linear_backward(
+    const nir_linear_params_t *params, const nir_tensor_t *input,
+    const nir_tensor_t *grad_output, nir_tensor_t *grad_weight,
+    nir_tensor_t *grad_input
+);
+
 #endif /* NIR_C_RUNTIME_OPS_H */

include/nir_c_runtime/train.h

@@ -0,0 +1,133 @@
+/* C-side training math: surrogate-name resolution, cross-entropy
+ * loss, SGD/Adam optimizers over raw weight buffers, and a local
+ * STDP rule ported from spikeforge's STDPSynapse.
+ *
+ * This mirrors the torch training step (cross_entropy on logits,
+ * then a torch.optim.Adam step) for single-step topologies whose
+ * backward pass the caller composes from the layer backward ops in
+ * ops.h. A graph-level tape retaining per-step activations for
+ * backprop-through-time is future work; optimizers here operate on
+ * caller-provided float buffers (e.g. a node's weight array).
+ */
+#ifndef NIR_C_RUNTIME_TRAIN_H
+#define NIR_C_RUNTIME_TRAIN_H
+
+#include <stddef.h>
+
+#include "nir_c_runtime/ops.h"
+#include "nir_c_runtime/tensor.h"
+
+/* Resolve "fast_sigmoid" or "atan" to a surrogate kind. Any other
+ * name (or NULL) fails with a malloc'd error naming it, mirroring
+ * spikeforge's resolve_surrogate raising KeyError -- there is no
+ * silent default surrogate. */
+int nir_surrogate_from_name(
+    const char *name, nir_surrogate_kind_t *out, char **error
+);
+
+/* Softmax cross-entropy for one logit vector, matching
+ * torch.nn.functional.cross_entropy on a single sample: loss =
+ * -log(softmax(logits)[target]), grad = softmax(logits) with
+ * grad[target] -= 1. Batch reduction is the caller's job. */
+int nir_cross_entropy(
+    const nir_tensor_t *logits, size_t target, float *loss,
+    nir_tensor_t *grad
+);
+
+/* SGD with optional momentum and L2 weight decay, matching
+ * torch.optim.SGD with dampening=0 and nesterov=False:
+ *   d = grad + weight_decay * w
+ *   v = momentum * v + d
+ *   w -= lr * v
+ */
+typedef struct {
+    float lr;
+    float momentum;
+    float weight_decay;
+    float *velocity;
+    size_t count;
+} nir_sgd_t;
+
+/* lr must be finite and >= 0, momentum and weight_decay finite
+ * and >= 0, count nonzero. The velocity buffer starts at zero. */
+int nir_sgd_init(
+    nir_sgd_t *opt, size_t count, float lr, float momentum,
+    float weight_decay, char **error
+);
+int nir_sgd_step(
+    nir_sgd_t *opt, float *weights, const float *grad, size_t count
+);
+void nir_sgd_free(nir_sgd_t *opt);
+
+/* Adam with bias correction, matching torch.optim.Adam with
+ * amsgrad=False:
+ *   d = grad + weight_decay * w
+ *   m = beta1 * m + (1 - beta1) * d
+ *   v = beta2 * v + (1 - beta2) * d^2
+ *   w -= lr * (m / (1 - beta1^t)) /
+ *       (sqrt(v / (1 - beta2^t)) + eps)
+ * where t counts steps from 1.
+ */
+typedef struct {
+    float lr;
+    float beta1;
+    float beta2;
+    float eps;
+    float weight_decay;
+    float *m;
+    float *v;
+    size_t count;
+    unsigned long step;
+} nir_adam_t;
+
+/* lr must be finite and >= 0, betas finite in [0, 1), eps finite
+ * and > 0, weight_decay finite and >= 0, count nonzero. */
+int nir_adam_init(
+    nir_adam_t *opt, size_t count, float lr, float beta1,
+    float beta2, float eps, float weight_decay, char **error
+);
+int nir_adam_step(
+    nir_adam_t *opt, float *weights, const float *grad, size_t count
+);
+void nir_adam_free(nir_adam_t *opt);
+
+/* Pair-based STDP with pre/post eligibility traces, ported from
+ * spikeforge.memory.STDPSynapse (Song/Abbott rule). Weights are
+ * [in_features][out_features], row-major pre-major, matching the
+ * torch [in, out] layout and the spikes @ weight forward. */
+typedef struct {
+    size_t in_features;
+    size_t out_features;
+    float tau_pre;
+    float tau_post;
+    float a_plus;
+    float a_minus;
+    float w_min;
+    float w_max;
+    float *weight;
+    float *pre_trace;
+    float *post_trace;
+} nir_stdp_t;
+
+/* taus must be finite and > 0, w_min <= w_max, both sizes nonzero.
+ * Weights and traces start at zero. */
+int nir_stdp_init(
+    nir_stdp_t *s, size_t in_features, size_t out_features,
+    float tau_pre, float tau_post, float a_plus, float a_minus,
+    float w_min, float w_max, char **error
+);
+
+/* Overwrite the weight matrix (in*out floats, pre-major). */
+int nir_stdp_set_weight(nir_stdp_t *s, const float *weight);
+
+/* One STDP step: decay both traces first (a spike reads the
+ * *prior* step's trace), potentiate on post spikes landing on
+ * lingering pre trace, depress on pre spikes landing on lingering
+ * post trace, clamp to [w_min, w_max], then add this step's
+ * spikes to the traces. */
+int nir_stdp_step(
+    nir_stdp_t *s, const float *pre_spikes, const float *post_spikes
+);
+void nir_stdp_free(nir_stdp_t *s);
+
+#endif /* NIR_C_RUNTIME_TRAIN_H */