our_knp/cats-dogs/gpu_stdp.cu

#include "gpu_stdp_types.h"
#include <cuda_runtime.h>
#include <cmath>
#include <stdexcept>
#include <string>

static void checked(cudaError_t status)
{
    if (status != cudaSuccess) throw std::runtime_error(std::string("AltAI CUDA: ") + cudaGetErrorString(status));
}
struct GpuStdpContext
{
    GpuStdpSynapse* synapses = nullptr;
    GpuStdpNeuron* neurons = nullptr;
    std::uint64_t* last = nullptr;
    std::size_t* offsets = nullptr;
    float* weights = nullptr;
    std::size_t n = 0, s = 0;
};

__device__ static float factor(float stability)
{
    return fminf(static_cast<float>(pow(2.0, static_cast<double>(-stability))), 1.0f);
}
__device__ static void recalculate(GpuStdpSynapse* synapses, float* weights, std::size_t lo, std::size_t hi)
{
    for (std::size_t j = lo; j < hi; ++j)
    {
        auto& s = synapses[j];
        const float resource = fmaxf(s.resource, 0.0f);
        const float diff = s.w_max - s.w_min;
        s.weight = s.w_min + diff * resource / (diff + resource);
        weights[j] = s.weight;
    }
}
__global__ static void plasticity(GpuStdpNeuron* neurons, GpuStdpSynapse* synapses,
                                  const std::uint64_t* last, const std::size_t* offsets,
                                  float* weights, std::size_t count, std::uint64_t step)
{
    const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= count) return;
    auto& n = neurons[i];
    const std::size_t lo = offsets[i], hi = offsets[i + 1];
    if (n.spiked)
    {
        n.last_spike = step;
        if (n.forced) n.status = 0;
        else
        {
            if (n.status == 3 || n.status == 0) { n.status = 1; n.first_isi = step; }
            else if (n.status == 1) { if (n.last_step - step < n.isi_max) n.status = 2; }
            else if (n.status == 2)
            {
                if (n.last_step - step >= n.isi_max || n.dopamine != 0) { n.status = 1; n.first_isi = step; }
            }
            n.last_step = step;
        }
        if (n.status == 1) n.stability -= n.stability_at_isi;
        n.additional = 0;
        for (std::size_t j = lo; j < hi; ++j)
        {
            auto& s = synapses[j];
            n.additional += s.weight * (s.weight > 0);
            auto begin = step - s.period;
            auto end = step;
            if (begin > end) { auto t = begin; begin = end; end = t; }
            const auto point = last[j] + s.delay - 1;
            const bool had_spike = begin <= point && point <= end;
            if (n.status != 2 || had_spike) s.contributed = had_spike;
        }
        n.additional *= n.coefficient;
        if (n.status != 2)
            for (std::size_t j = lo; j < hi; ++j) synapses[j].hebbian = 0;
        if (n.status != 0)
            for (std::size_t j = lo; j < hi; ++j)
            {
                auto& s = synapses[j];
                s.resource -= s.d_u;
                n.free_resource += s.d_u;
                if (s.contributed && !s.hebbian)
                {
                    const float delta = n.d_h * factor(n.stability);
                    s.resource += delta;
                    n.free_resource -= delta;
                    s.hebbian = 1;
                }
            }
        recalculate(synapses, weights, lo, hi);
    }
    if (n.dopamine > 0 || (n.dopamine < 0 && n.status != 0))
    {
        for (std::size_t j = lo; j < hi; ++j)
        {
            auto& s = synapses[j];
            if (step - n.last_spike <= n.dopamine_time && s.contributed)
            {
                const float delta = static_cast<float>(n.dopamine * factor(n.stability));
                s.resource += delta;
                n.free_resource -= delta;
            }
        }
        if (n.forced || n.dopamine < 0)
        {
            n.stability = static_cast<float>(n.stability - n.dopamine * n.stability_change);
            n.stability = fmaxf(n.stability, 0.0f);
        }
        else
        {
            const double difference = static_cast<double>(step - n.first_isi - n.isi_max);
            n.stability = static_cast<float>(n.stability + n.stability_change * n.dopamine *
                fmax(2.0 - fabs(difference) / n.isi_max, -1.0));
        }
        recalculate(synapses, weights, lo, hi);
    }
    if (step - n.last_step <= n.isi_max && n.status != 0) return;
    if (fabsf(n.free_resource) < n.resource_threshold) return;
    const float delta = n.free_resource / static_cast<float>((hi - lo) + n.drain);
    for (std::size_t j = lo; j < hi; ++j) synapses[j].resource += delta;
    n.free_resource = 0;
    recalculate(synapses, weights, lo, hi);
}

GpuStdpContext* gpu_stdp_create(const GpuStdpSynapse* synapses, const float* weights,
                              const std::size_t* offsets, std::size_t n, std::size_t s)
{
    auto* c = new GpuStdpContext;
    c->n = n; c->s = s;
    try
    {
        checked(cudaMalloc(&c->synapses, s * sizeof(GpuStdpSynapse)));
        checked(cudaMalloc(&c->neurons, n * sizeof(GpuStdpNeuron)));
        checked(cudaMalloc(&c->last, s * sizeof(std::uint64_t)));
        checked(cudaMalloc(&c->offsets, (n + 1) * sizeof(std::size_t)));
        checked(cudaMalloc(&c->weights, s * sizeof(float)));
        checked(cudaMemcpy(c->synapses, synapses, s * sizeof(GpuStdpSynapse), cudaMemcpyHostToDevice));
        checked(cudaMemcpy(c->offsets, offsets, (n + 1) * sizeof(std::size_t), cudaMemcpyHostToDevice));
        checked(cudaMemcpy(c->weights, weights, s * sizeof(float), cudaMemcpyHostToDevice));
    }
    catch (...) { gpu_stdp_destroy(c); throw; }
    return c;
}
void gpu_stdp_step(GpuStdpContext* c, GpuStdpNeuron* neurons, const std::uint64_t* last,
                   float* weights, std::uint64_t step)
{
    checked(cudaMemcpy(c->neurons, neurons, c->n * sizeof(GpuStdpNeuron), cudaMemcpyHostToDevice));
    checked(cudaMemcpy(c->last, last, c->s * sizeof(std::uint64_t), cudaMemcpyHostToDevice));
    plasticity<<<(c->n + 31) / 32, 32>>>(c->neurons, c->synapses, c->last, c->offsets, c->weights, c->n, step);
    checked(cudaGetLastError());
    checked(cudaMemcpy(neurons, c->neurons, c->n * sizeof(GpuStdpNeuron), cudaMemcpyDeviceToHost));
    checked(cudaMemcpy(weights, c->weights, c->s * sizeof(float), cudaMemcpyDeviceToHost));
}
void gpu_stdp_download(GpuStdpContext* c, GpuStdpSynapse* synapses)
{
    checked(cudaMemcpy(synapses, c->synapses, c->s * sizeof(GpuStdpSynapse), cudaMemcpyDeviceToHost));
}
void gpu_stdp_destroy(GpuStdpContext* c) noexcept
{
    if (!c) return;
    cudaFree(c->synapses); cudaFree(c->neurons); cudaFree(c->last); cudaFree(c->offsets); cudaFree(c->weights);
    delete c;
}