foxhunt/crates/ml-alpha/cuda/cfc_step.cu

// cfc_step.cu — one CfC time step (Hasani 2022 closed-form recurrence).
//
//   pre[i]   = sum_k w_in[i,k] * x[k] + sum_k w_rec[i,k] * h_old[k] + b[i]
//   decay[i] = exp(-dt_s / max(tau[i], 1e-6))
//   h_new[i] = h_old[i] * decay[i] + (1 - decay[i]) * tanh(pre[i])
//
// Launched as one block of n_hid threads (or grid of blocks if n_hid > 128).
// Each thread computes one hidden unit; reductions are intra-thread (each
// thread does its own dot product). No atomicAdd. Per
// `feedback_no_atomicadd.md` and the GPU/CPU contract.

extern "C" __global__ void cfc_step(
    const float* __restrict__ w_in,    // [n_hid, n_in], row-major
    const float* __restrict__ w_rec,   // [n_hid, n_hid], row-major
    const float* __restrict__ b,       // [n_hid]
    const float* __restrict__ tau,     // [n_hid]
    const float* __restrict__ x,       // [n_in]
    const float* __restrict__ h_old,   // [n_hid]
    float dt_s,
    int n_in,
    int n_hid,
    float* __restrict__ h_new          // [n_hid]
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n_hid) return;

    float pre = b[i];
    for (int k = 0; k < n_in; ++k) {
        pre += w_in[i * n_in + k] * x[k];
    }
    for (int k = 0; k < n_hid; ++k) {
        pre += w_rec[i * n_hid + k] * h_old[k];
    }
    const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f));
    h_new[i] = h_old[i] * decay + (1.0f - decay) * tanhf(pre);
}

// Backward through ONE cfc_step (truncated BPTT, K=1).
//
// Forward (for reference):
//   pre[i]   = b[i] + sum_k W_in[i,k] x[k] + sum_k W_rec[i,k] h_old[k]
//   decay[i] = exp(-dt / max(tau[i], 1e-6))
//   s[i]     = tanh(pre[i])
//   h_new[i] = h_old[i] * decay + (1 - decay) * s
//
// Given grad_h_new[i] = dL/dh_new[i], compute:
//   d_pre[i]      = grad_h_new[i] * (1 - decay) * (1 - s^2)
//   grad_b[i]    += d_pre[i]
//   grad_W_in[i,k]  += d_pre[i] * x[k]
//   grad_W_rec[i,k] += d_pre[i] * h_old[k]
//   grad_h_old[i] += grad_h_new[i] * decay + sum_j d_pre[j] * W_rec[j, i]
//
// tau is held frozen (structurally log-uniform initialised per spec
// Section 2); no grad_tau path in this version.
//
// One thread per hidden unit (128 threads in one block). Block tree-
// reduce not needed — each thread owns its own grad_W rows and grad_h
// component; grad_h_old accumulates via shared `d_pre` reads.

extern "C" __global__ void cfc_step_backward(
    const float* __restrict__ w_in,           // [n_hid, n_in]
    const float* __restrict__ w_rec,          // [n_hid, n_hid]
    const float* __restrict__ b,
    const float* __restrict__ tau,
    const float* __restrict__ x,              // [n_in]
    const float* __restrict__ h_old,          // [n_hid]
    const float* __restrict__ grad_h_new,     // [n_hid]
    float dt_s,
    int n_in,
    int n_hid,
    float* __restrict__ grad_w_in,            // [n_hid, n_in]
    float* __restrict__ grad_w_rec,           // [n_hid, n_hid]
    float* __restrict__ grad_b,               // [n_hid]
    float* __restrict__ grad_h_old            // [n_hid]
) {
    extern __shared__ float smem[];
    float* sd_pre = smem;          // [n_hid]
    float* sdecay = sd_pre + n_hid; // [n_hid]

    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n_hid) return;

    // Recompute forward pre + tanh to derive d_pre.
    float pre = b[i];
    for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * x[k];
    for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * h_old[k];
    const float s = tanhf(pre);
    const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f));
    const float dh = grad_h_new[i];
    const float d_pre = dh * (1.0f - decay) * (1.0f - s * s);

    sd_pre[i] = d_pre;
    sdecay[i] = decay;
    __syncthreads();

    grad_b[i] = d_pre;
    for (int k = 0; k < n_in; ++k) {
        grad_w_in[i * n_in + k] = d_pre * x[k];
    }
    for (int k = 0; k < n_hid; ++k) {
        grad_w_rec[i * n_hid + k] = d_pre * h_old[k];
    }

    // grad_h_old[i] = dh * decay + sum_j d_pre[j] * W_rec[j, i]
    float gh = dh * decay;
    for (int j = 0; j < n_hid; ++j) {
        gh += sd_pre[j] * w_rec[j * n_hid + i];
    }
    grad_h_old[i] = gh;
}