163 lines
7.3 KiB
C++
163 lines
7.3 KiB
C++
#include "gpu_stdp_bridge.h"
|
|
#include "gpu_stdp_types.h"
|
|
#include <knp/backends/cpu-library/impl/populations/training/stdp.h>
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <cstdlib>
|
|
#include <iostream>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <stdexcept>
|
|
|
|
namespace knp_gpu_experiment
|
|
{
|
|
using Params = knp::synapse_traits::synapse_parameters<Synapse>;
|
|
using NeuronParams = knp::neuron_traits::neuron_parameters<Neuron>;
|
|
struct Context
|
|
{
|
|
Population* population;
|
|
std::vector<std::reference_wrapper<Params>> synapses;
|
|
std::vector<std::size_t> offsets;
|
|
std::vector<GpuStdpSynapse> state;
|
|
std::vector<GpuStdpNeuron> neurons;
|
|
std::vector<std::uint64_t> last;
|
|
std::vector<float> weights;
|
|
GpuStdpContext* gpu = nullptr;
|
|
~Context() { gpu_stdp_destroy(gpu); }
|
|
};
|
|
static thread_local std::map<const Population*, std::unique_ptr<Context>> contexts;
|
|
static thread_local std::size_t checked_steps = 0, gpu_steps = 0;
|
|
static thread_local double max_error = 0;
|
|
|
|
static GpuStdpSynapse pack(const Params& p)
|
|
{
|
|
const auto& r = p.rule_;
|
|
return {r.synaptic_resource_, r.w_min_, r.w_max_, r.d_u_, p.weight_,
|
|
r.dopamine_plasticity_period_, p.delay_, r.has_contributed_, r.had_hebbian_update_};
|
|
}
|
|
static GpuStdpNeuron pack(const NeuronParams& n, bool spike)
|
|
{
|
|
return {n.additional_threshold_, n.dopamine_value_, n.synapse_sum_threshold_coefficient_,
|
|
n.free_synaptic_resource_, n.synaptic_resource_threshold_, n.stability_,
|
|
n.stability_change_parameter_, n.stability_change_at_isi_, n.d_h_,
|
|
n.last_step_, n.last_spike_step_, n.first_isi_spike_, n.isi_max_,
|
|
n.dopamine_plasticity_time_, n.resource_drain_coefficient_,
|
|
static_cast<int>(n.isi_status_), n.is_being_forced_, spike};
|
|
}
|
|
static void unpack(const GpuStdpNeuron& s, NeuronParams& n)
|
|
{
|
|
n.additional_threshold_ = s.additional;
|
|
n.free_synaptic_resource_ = s.free_resource;
|
|
n.stability_ = s.stability;
|
|
n.last_step_ = s.last_step; n.last_spike_step_ = s.last_spike; n.first_isi_spike_ = s.first_isi;
|
|
n.isi_status_ = static_cast<knp::neuron_traits::ISIPeriodType>(s.status);
|
|
}
|
|
static void unpack(const GpuStdpSynapse& s, Params& p)
|
|
{
|
|
p.weight_ = s.weight;
|
|
p.rule_.synaptic_resource_ = s.resource;
|
|
p.rule_.has_contributed_ = s.contributed;
|
|
p.rule_.had_hebbian_update_ = s.hebbian;
|
|
// last_spike_step_ remains authoritative on CPU (the projection phase updates it).
|
|
}
|
|
static bool near(double a, double b)
|
|
{
|
|
if (a == b) return true;
|
|
const double error = std::abs(a - b);
|
|
max_error = std::max(max_error, error);
|
|
return std::isfinite(a) && std::isfinite(b) && error <= 1e-4 + 1e-6 * std::max(std::abs(a), std::abs(b));
|
|
}
|
|
static void compare(const GpuStdpNeuron& actual, const GpuStdpNeuron& expected)
|
|
{
|
|
if (!near(actual.additional, expected.additional) || !near(actual.free_resource, expected.free_resource) ||
|
|
!near(actual.stability, expected.stability) || actual.last_step != expected.last_step ||
|
|
actual.last_spike != expected.last_spike || actual.first_isi != expected.first_isi ||
|
|
actual.status != expected.status)
|
|
throw std::runtime_error("CUDA STDP neuron-state parity failed");
|
|
}
|
|
static void synchronize(Context& c)
|
|
{
|
|
gpu_stdp_download(c.gpu, c.state.data());
|
|
for (std::size_t j = 0; j < c.synapses.size(); ++j) unpack(c.state[j], c.synapses[j]);
|
|
}
|
|
void flush()
|
|
{
|
|
for (auto& entry : contexts) synchronize(*entry.second);
|
|
if (!contexts.empty())
|
|
std::cerr << "CUDA_STDP gpu_steps=" << gpu_steps << " checked_steps=" << checked_steps
|
|
<< " max_state_abs_error=" << max_error << '\n';
|
|
}
|
|
void clear()
|
|
{
|
|
contexts.clear(); checked_steps = gpu_steps = 0; max_error = 0;
|
|
}
|
|
void train(Population& population, Projections& projections,
|
|
const knp::core::messaging::SpikeMessage& message, knp::core::Step step, Reference reference)
|
|
{
|
|
if (projections.empty()) { reference(population, projections, message, step); return; }
|
|
bool active = !message.neuron_indexes_.empty();
|
|
for (const auto& n : population)
|
|
active |= n.dopamine_value_ > 0 || (n.dopamine_value_ < 0 && n.isi_status_ != knp::neuron_traits::ISIPeriodType::is_forced)
|
|
|| (!(step - n.last_step_ <= n.isi_max_ && n.isi_status_ != knp::neuron_traits::ISIPeriodType::is_forced)
|
|
&& std::fabs(n.free_synaptic_resource_) >= n.synaptic_resource_threshold_);
|
|
if (!active) return;
|
|
auto& owned = contexts[&population];
|
|
if (!owned)
|
|
{
|
|
owned = std::make_unique<Context>();
|
|
auto& c = *owned;
|
|
c.population = &population;
|
|
c.offsets.push_back(0);
|
|
for (std::size_t i = 0; i < population.size(); ++i)
|
|
{
|
|
const auto& synapses = knp::backends::cpu::populations::impl::training::stdp::
|
|
get_all_connected_synapses<Synapse>(projections, i);
|
|
c.synapses.insert(c.synapses.end(), synapses.begin(), synapses.end());
|
|
c.offsets.push_back(c.synapses.size());
|
|
}
|
|
if (c.synapses.empty()) { owned.reset(); reference(population, projections, message, step); return; }
|
|
c.neurons.resize(population.size()); c.state.reserve(c.synapses.size());
|
|
for (auto& s : c.synapses) c.state.push_back(pack(s));
|
|
c.last.resize(c.synapses.size()); c.weights.resize(c.synapses.size());
|
|
for (std::size_t j = 0; j < c.synapses.size(); ++j) c.weights[j] = c.synapses[j].get().weight_;
|
|
c.gpu = gpu_stdp_create(c.state.data(), c.weights.data(), c.offsets.data(), population.size(), c.synapses.size());
|
|
}
|
|
auto& c = *owned;
|
|
if (c.neurons.size() != population.size()) throw std::runtime_error("CUDA STDP requires fixed topology");
|
|
const bool verify = std::getenv("KNP_CUDA_VERIFY") != nullptr;
|
|
if (verify) synchronize(c);
|
|
std::vector<bool> spiked(population.size());
|
|
for (auto i : message.neuron_indexes_) spiked.at(i) = true;
|
|
for (std::size_t i = 0; i < population.size(); ++i) c.neurons[i] = pack(population[i], spiked[i]);
|
|
for (std::size_t j = 0; j < c.synapses.size(); ++j) c.last[j] = c.synapses[j].get().rule_.last_spike_step_;
|
|
std::vector<GpuStdpNeuron> expected_neurons;
|
|
std::vector<GpuStdpSynapse> expected_synapses;
|
|
if (verify)
|
|
{
|
|
reference(population, projections, message, step);
|
|
for (std::size_t i = 0; i < population.size(); ++i) expected_neurons.push_back(pack(population[i], spiked[i]));
|
|
for (auto& s : c.synapses) expected_synapses.push_back(pack(s));
|
|
}
|
|
gpu_stdp_step(c.gpu, c.neurons.data(), c.last.data(), c.weights.data(), step);
|
|
++gpu_steps;
|
|
for (std::size_t i = 0; i < population.size(); ++i)
|
|
{
|
|
if (verify) compare(c.neurons[i], expected_neurons[i]);
|
|
unpack(c.neurons[i], population[i]);
|
|
}
|
|
for (std::size_t j = 0; j < c.synapses.size(); ++j) c.synapses[j].get().weight_ = c.weights[j];
|
|
if (verify)
|
|
{
|
|
gpu_stdp_download(c.gpu, c.state.data());
|
|
for (std::size_t j = 0; j < c.synapses.size(); ++j)
|
|
{
|
|
const auto& s = c.state[j]; const auto& e = expected_synapses[j];
|
|
if (!near(s.resource, e.resource) || !near(s.weight, e.weight) ||
|
|
s.contributed != e.contributed || s.hebbian != e.hebbian)
|
|
throw std::runtime_error("CUDA STDP synapse-state parity failed at " + std::to_string(step));
|
|
unpack(s, c.synapses[j]);
|
|
}
|
|
++checked_steps;
|
|
}
|
|
}
|
|
}
|