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

// cfc_step.cu — one CfC time step (Hasani 2022 closed-form recurrence)
// + batched variants + 3D transpose helper.

// Transpose a [N1, N2, N3] row-major tensor to [N2, N1, N3]. Used by
// the batched PerceptionTrainer to convert Mamba2's [B, K, hidden]
// per-step output into the [K, B, hidden] layout that makes per-K
// slot access contiguous in the trainer's K loop (and back again for
// grad_h_enriched_seq feeding Mamba2's backward).
//
// Launches one block per (i, j) pair, blockDim.x threads cover N3.
extern "C" __global__ void transpose_3d_swap_01(
    const float* __restrict__ in,    // [N1, N2, N3]
    float* __restrict__ out,         // [N2, N1, N3]
    int N1,
    int N2,
    int N3
) {
    int i = blockIdx.x;
    int j = blockIdx.y;
    int k = blockIdx.z * blockDim.x + threadIdx.x;
    if (i >= N1 || j >= N2 || k >= N3) return;
    out[((long long)j * N1 + i) * N3 + k] = in[((long long)i * N2 + j) * N3 + k];
}

//
//   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]
//   τ_eff    = max(tau[i], 1e-6)
//   decay[i] = exp(-dt / τ_eff)
//   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_decay[i]    = grad_h_new[i] * (h_old[i] - s[i])
//   d_pre[i]      = grad_h_new[i] * (1 - decay) * (1 - s^2)
//   grad_tau[i]  += d_decay[i] * decay[i] * dt / τ_eff²   (when tau > eps)
//   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 trained — grad_tau receives a `+=` accumulation (callers must
// pre-zero across the per-position K loop).
//
// 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_tau,             // [n_hid] — accumulated via +=
    float* __restrict__ grad_h_old,           // [n_hid]
    float* __restrict__ grad_x                // [n_in] — set to nullptr to skip
) {
    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 tau_eps = 1e-6f;
    const float tau_eff = fmaxf(tau[i], tau_eps);
    const float decay = expf(-dt_s / tau_eff);
    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();

    // Parameter-grad writes use `+=`: callers may invoke this kernel
    // multiple times per training step (per-position supervision) and
    // need the gradients summed across calls. Callers MUST pre-zero
    // grad_b / grad_w_in / grad_w_rec / grad_tau at the start of each step.
    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_tau accumulation. Zero contribution if tau is clamped to eps
    // (max() inactive on this side) — strict d(max)/d(tau) = 0 there.
    if (tau[i] > tau_eps) {
        const float d_decay = dh * (h_old[i] - s);
        grad_tau[i] += d_decay * decay * dt_s / (tau_eff * tau_eff);
    }

    // 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;

    // grad_x[k] = sum_i d_pre[i] * W_in[i, k]. Each thread i contributes to
    // n_in different output positions — different threads write to different
    // grad_x slots if we partition by i, but multiple i's contribute to each k.
    // Avoid atomicAdd: rely on thread 0 to do the n_in × n_hid sum
    // sequentially (n_in is small — typically 32 or 128, this is one-time per
    // backward call). Discriminated by nullptr → skip path for back-compat.
    if (grad_x != nullptr && i == 0) {
        for (int k = 0; k < n_in; ++k) {
            float gx = 0.0f;
            for (int j = 0; j < n_hid; ++j) {
                gx += sd_pre[j] * w_in[j * n_in + k];
            }
            grad_x[k] = gx;
        }
    }
}


// ─── Batched variants ────────────────────────────────────────────────
//
// Phase B (2026-05-17): block-per-batch refactor. Launch contract:
//   grid=(n_batch, 1, 1), block=(n_hid, 1, 1).
// Each block handles ONE batch's n_hid channels in parallel.
//
// Param-grad writes go through per-batch scratch buffers
// (grad_w_in_scratch [B, n_hid, n_in] etc.). Thread (bi, i) is the
// sole writer to its slice of each scratch tensor; the K-loop's 64
// invocations accumulate via += within the same (bi, i, *) slice.
// After the K-loop the caller runs `reduce_axis0` to sum B → final
// grad buffer (OVERWRITE semantics). No atomicAdd per
// feedback_no_atomicadd.md.
//
// Shared memory budget per block: sd_pre[n_hid] + sdecay[n_hid] =
// 2 * n_hid * 4 bytes. At n_hid=128 → 1 KiB. Far below 48 KiB limit,
// no cudaFuncSetAttribute needed.

// Block-per-batch (Phase B): grid=(n_batch, 1, 1), block=(n_hid, 1, 1).
// Each block handles ONE batch's n_hid channels in parallel. Removes
// the internal batch loop — frees ~31 SMs per launch at B=32.
extern "C" __global__ void cfc_step_batched(
    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_batch, n_in]
    const float* __restrict__ h_old,   // [n_batch, n_hid]
    float dt_s,
    int n_in,
    int n_hid,
    int n_batch,
    float* __restrict__ h_new          // [n_batch, n_hid]
) {
    // Dynamic shared layout (caller supplies (n_in + n_hid) * 4 bytes):
    //   s_x     [n_in]   — staged x_b row (read by every channel thread)
    //   s_h_old [n_hid]  — staged h_old_b row (same)
    extern __shared__ float smem_fwd[];
    float* s_x     = smem_fwd;
    float* s_h_old = s_x + n_in;

    int bi  = blockIdx.x;
    int tid = threadIdx.x;
    if (bi >= n_batch) return;

    // Cooperative staging: 128× redundant DRAM reads per row → 1× per row
    // shared across all channel threads. Both rows fit since n_in = n_hid
    // = HIDDEN.
    if (tid < n_in)  s_x[tid]     = x[(long long)bi * n_in + tid];
    if (tid < n_hid) s_h_old[tid] = h_old[(long long)bi * n_hid + tid];
    __syncthreads();

    if (tid >= n_hid) return;
    const int i = tid;
    const float bias_i = b[i];
    const float decay  = expf(-dt_s / fmaxf(tau[i], 1e-6f));
    const float one_minus_decay = 1.0f - decay;

    float pre = bias_i;
    for (int k = 0; k < n_in;  ++k) pre += w_in[i * n_in + k]  * s_x[k];
    for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * s_h_old[k];
    h_new[(long long)bi * n_hid + i] =
        s_h_old[i] * decay + one_minus_decay * tanhf(pre);
}


// Block-per-batch backward.
//   grid=(n_batch, 1, 1)  block=(n_hid, 1, 1)
//   shared mem: 2 * n_hid * 4 bytes (sd_pre + sdecay; one row each)
//
// Param-grad scratch tensors hold per-batch slices. Thread (bi, i) is
// sole writer to its slice of each scratch tensor for every K-iteration
// the K-loop calls this kernel — no race even with B parallel blocks.
// Caller zeroes the scratch buffers once per training step; this kernel
// uses += to accumulate across the K-loop's 64 invocations.
//
// Outputs grad_h_old and grad_x are per-batch indexed; each block bi is
// sole writer to its [bi, :] slice (overwrite, single launch).
extern "C" __global__ void cfc_step_backward_batched(
    const float* __restrict__ w_in,                  // [n_hid, n_in]
    const float* __restrict__ w_rec,                 // [n_hid, n_hid]
    const float* __restrict__ b,                     // [n_hid]
    const float* __restrict__ tau,                   // [n_hid]
    const float* __restrict__ x,                     // [n_batch, n_in]
    const float* __restrict__ h_old,                 // [n_batch, n_hid]
    const float* __restrict__ grad_h_new,            // [n_batch, n_hid]
    float dt_s,
    int n_in,
    int n_hid,
    int n_batch,
    float* __restrict__ grad_w_in_scratch,           // [n_batch, n_hid, n_in]    accum +=
    float* __restrict__ grad_w_rec_scratch,          // [n_batch, n_hid, n_hid]   accum +=
    float* __restrict__ grad_b_scratch,              // [n_batch, n_hid]          accum +=
    float* __restrict__ grad_tau_scratch,            // [n_batch, n_hid]          accum +=
    float* __restrict__ grad_h_old,                  // [n_batch, n_hid]          overwrite
    float* __restrict__ grad_x                       // [n_batch, n_in]           overwrite (nullptr OK)
) {
    extern __shared__ float smem[];
    // Shared layout — caller passes (2*n_hid + n_in + n_hid) * 4 bytes:
    //   sd_pre  [n_hid]   — d_pre_b per channel (set in Pass 0, consumed Pass 1..3)
    //   sdecay  [n_hid]   — reserved for future passes / kept for ABI symmetry
    //   s_x     [n_in]    — staged x_b row (eliminates 128× redundant DRAM reads)
    //   s_h_old [n_hid]   — staged h_old_b row (same)
    float* sd_pre  = smem;
    float* sdecay  = sd_pre + n_hid;
    float* s_x     = sdecay + n_hid;
    float* s_h_old = s_x + n_in;

    int bi  = blockIdx.x;
    int tid = threadIdx.x;
    // block_dim is >= max(n_hid, n_in); each thread serves both roles
    // across passes. Most paths gate on whether (tid < n_hid) / (tid < n_in).
    if (bi >= n_batch) return;

    // Cooperative staging — coalesced loads, single sync.
    if (tid < n_in)  s_x[tid]     = x[(long long)bi * n_in + tid];
    if (tid < n_hid) s_h_old[tid] = h_old[(long long)bi * n_hid + tid];
    __syncthreads();

    // Pass 0: compute pre[i], d_pre[i], decay[i]; write per-channel scratch.
    //   Thread tid serves the channel role i = tid. Reads from shared.
    float dh_b = 0.0f;
    float decay_i = 0.0f;
    float s_i = 0.0f;
    if (tid < n_hid) {
        const int i = tid;
        const float bias_i = b[i];
        const float tau_eps = 1e-6f;
        const float tau_i_raw = tau[i];
        const float tau_eff = fmaxf(tau_i_raw, tau_eps);
        decay_i = expf(-dt_s / tau_eff);
        sdecay[i] = decay_i;
        dh_b = grad_h_new[(long long)bi * n_hid + i];

        float pre = bias_i;
        for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * s_x[k];
        for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * s_h_old[k];
        s_i = tanhf(pre);
        const float d_pre_b = dh_b * (1.0f - decay_i) * (1.0f - s_i * s_i);
        sd_pre[i] = d_pre_b;

        // Per-channel scratch (one float per thread, no coalescing concern).
        grad_b_scratch[(long long)bi * n_hid + i] += d_pre_b;

        // Branchless tau-grad — gates contribution by (tau > eps) without
        // warp divergence.
        const float d_decay_b = dh_b * (s_h_old[i] - s_i);
        const float gate     = (tau_i_raw > tau_eps) ? 1.0f : 0.0f;
        const float contrib  = d_decay_b * decay_i * dt_s / (tau_eff * tau_eff);
        grad_tau_scratch[(long long)bi * n_hid + i] += gate * contrib;
    }
    __syncthreads();

    // Pass 1: grad_w_in[bi, i, k] += d_pre[i] * x[k]
    //         grad_w_rec[bi, i, k] += d_pre[i] * h_old[k]
    //   Thread tid serves the column role k = tid. For each i, every warp
    //   thread writes to grad_w_*[i, tid] → COALESCED 128B/transaction.
    //   Prior layout (thread = i) strided cross-thread writes by n_in/n_hid
    //   = 512B, paying 32 separate transactions per warp instruction.
    if (tid < n_in) {
        const int k = tid;
        const float x_k     = s_x[k];
        const float h_old_k = s_h_old[k];   // valid because n_in == n_hid here
        for (int i = 0; i < n_hid; ++i) {
            const float d_pre_i = sd_pre[i];
            grad_w_in_scratch[((long long)bi * n_hid + i) * n_in  + k]
                += d_pre_i * x_k;
            grad_w_rec_scratch[((long long)bi * n_hid + i) * n_hid + k]
                += d_pre_i * h_old_k;
        }
    }

    // Pass 2: grad_h_old[bi, i] = sum_j sd_pre[j] * W_rec[j, i] + dh_b * decay.
    //   Thread tid is back to channel role i = tid; dh_b and decay_i
    //   stayed in registers from Pass 0.
    if (tid < n_hid) {
        const int i = tid;
        float gh = dh_b * decay_i;
        for (int j = 0; j < n_hid; ++j) {
            gh += sd_pre[j] * w_rec[j * n_hid + i];
        }
        grad_h_old[(long long)bi * n_hid + i] = gh;
    }

    // Pass 3: grad_x[bi, k] = sum_j sd_pre[j] * W_in[j, k]
    //   Thread tid serves column role k = tid. Sequential along j, sd_pre
    //   in shared. One coalesced write per thread (grad_x layout matches).
    if (grad_x != nullptr && tid < n_in) {
        const int k = tid;
        float gx = 0.0f;
        for (int j = 0; j < n_hid; ++j) {
            gx += sd_pre[j] * w_in[j * n_in + k];
        }
        grad_x[(long long)bi * n_in + k] = gx;
    }
}