251 lines
14 KiB
C++
251 lines
14 KiB
C++
#include "altai_cuda_backend.h"
|
|
#include "altai_cuda_runtime.h"
|
|
#include <knp/meta/stringify.h>
|
|
#include <knp/meta/variant_helpers.h>
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cmath>
|
|
#include <cstdlib>
|
|
#include <iostream>
|
|
#include <limits>
|
|
|
|
namespace knp_altai_cuda {
|
|
using Neuron = knp::neuron_traits::AltAILIF;
|
|
using Synapse = knp::synapse_traits::DeltaSynapse;
|
|
using Population = knp::core::Population<Neuron>;
|
|
using Projection = knp::core::Projection<Synapse>;
|
|
using Spike = knp::core::messaging::SpikeMessage;
|
|
using Impact = knp::core::messaging::SynapticImpactMessage;
|
|
using NeuronList = boost::mp11::mp_list<Neuron>;
|
|
using SynapseList = boost::mp11::mp_list<Synapse>;
|
|
using PopulationList = boost::mp11::mp_list<Population>;
|
|
using ProjectionList = boost::mp11::mp_list<Projection>;
|
|
|
|
class Device final : public knp::core::Device {
|
|
public:
|
|
int index;
|
|
explicit Device(int i) : index(i) {
|
|
std::array<std::uint8_t,16> bytes{};
|
|
bytes[0]=0xA1;bytes[1]=0x7A;bytes[12]=(i>>24)&255;bytes[13]=(i>>16)&255;bytes[14]=(i>>8)&255;bytes[15]=i&255;
|
|
base_.uid_=knp::core::UID(bytes);
|
|
}
|
|
knp::core::DeviceType get_type() const override { return knp::core::DeviceType::GPU; }
|
|
const std::string get_name() const override {
|
|
CUdevice device;checked(cuDeviceGet(&device,index));
|
|
char name[256];checked(cuDeviceGetName(name,sizeof(name),device));
|
|
return std::string("CUDA ")+name;
|
|
}
|
|
float get_power() const override { return std::numeric_limits<float>::quiet_NaN(); }
|
|
};
|
|
struct Backend::State {
|
|
struct Pop {
|
|
Population *host;
|
|
int n,slots=2;
|
|
CUdeviceptr neurons=0,pending=0,absolute=0,events=0,spikes=0;
|
|
std::vector<NeuronState> states;
|
|
std::vector<unsigned char> flags;
|
|
};
|
|
struct Proj {
|
|
Projection *host;
|
|
size_t target;
|
|
int sources=0,max_row=0;
|
|
CUdeviceptr rows=0,source=0,weights=0,delays=0,counts=0;
|
|
std::vector<int> mask;
|
|
};
|
|
Runtime runtime;
|
|
std::vector<Pop> pops;
|
|
std::vector<Proj> projections;
|
|
CUdeviceptr error=0;
|
|
};
|
|
Backend::Backend() {
|
|
if (const char *v=std::getenv("KNP_CUDA_DEVICE")) {
|
|
std::string value(v);size_t end=0;device_index_=std::stoi(value,&end);
|
|
if (end!=value.size() || device_index_<0) throw std::runtime_error("Invalid KNP_CUDA_DEVICE");
|
|
}
|
|
}
|
|
Backend::~Backend() = default;
|
|
std::shared_ptr<knp::core::Backend> Backend::create() { return std::make_shared<Backend>(); }
|
|
std::vector<std::string> Backend::get_supported_neurons() const {
|
|
return knp::meta::get_supported_type_names<knp::neuron_traits::AllNeurons,NeuronList>(knp::neuron_traits::neurons_names);
|
|
}
|
|
std::vector<std::string> Backend::get_supported_synapses() const {
|
|
return knp::meta::get_supported_type_names<knp::synapse_traits::AllSynapses,SynapseList>(knp::synapse_traits::synapses_names);
|
|
}
|
|
std::vector<size_t> Backend::get_supported_population_indexes() const {
|
|
return knp::meta::get_supported_type_indexes<knp::core::AllPopulations,PopulationList>();
|
|
}
|
|
std::vector<size_t> Backend::get_supported_projection_indexes() const {
|
|
return knp::meta::get_supported_type_indexes<knp::core::AllProjections,ProjectionList>();
|
|
}
|
|
std::vector<std::unique_ptr<knp::core::Device>> Backend::get_devices() const {
|
|
checked(cuInit(0));int count=0;checked(cuDeviceGetCount(&count));
|
|
std::vector<std::unique_ptr<knp::core::Device>> result;
|
|
for(int i=0;i<count;++i) result.push_back(std::make_unique<Device>(i));
|
|
return result;
|
|
}
|
|
void Backend::select_devices(const std::set<knp::core::UID> &uids) {
|
|
if(state_ || uids.size()!=1) throw std::runtime_error("Select exactly one CUDA device before initialization");
|
|
for(auto &device:get_devices()) if(uids.count(device->get_uid())) {select_device(std::move(device));return;}
|
|
throw std::runtime_error("Unknown CUDA device UID");
|
|
}
|
|
void Backend::select_device(std::unique_ptr<knp::core::Device> &&device) {
|
|
auto *gpu=dynamic_cast<Device*>(device.get());
|
|
if(state_ || !gpu) throw std::runtime_error("Select a device from this CUDA backend before initialization");
|
|
device_index_=gpu->index;knp::core::Backend::select_device(std::move(device));
|
|
}
|
|
void Backend::load_all_populations(const std::vector<knp::core::AllPopulationsVariant> &populations) {
|
|
if(state_) throw std::runtime_error("CUDA topology is immutable after initialization; create a new backend");
|
|
for(const auto &p:populations) if(!std::holds_alternative<Population>(p))
|
|
throw std::runtime_error("CUDA backend supports integer AltAILIF inference only");
|
|
Base::load_all_populations(populations);
|
|
}
|
|
void Backend::load_all_projections(const std::vector<knp::core::AllProjectionsVariant> &projections) {
|
|
if(state_) throw std::runtime_error("CUDA topology is immutable after initialization; create a new backend");
|
|
for(const auto &p:projections) if(!std::holds_alternative<Projection>(p))
|
|
throw std::runtime_error("CUDA backend supports DeltaSynapse only; no STDP");
|
|
Base::load_all_projections(projections);
|
|
}
|
|
void Backend::remove_populations(const std::vector<knp::core::UID> &) {throw std::runtime_error("CUDA topology mutation unsupported");}
|
|
void Backend::remove_projections(const std::vector<knp::core::UID> &) {throw std::runtime_error("CUDA topology mutation unsupported");}
|
|
void Backend::start_learning() {throw std::runtime_error("CUDA AltAI inference backend has no plasticity");}
|
|
void Backend::_init() {
|
|
if(state_) {state_->runtime.current();return;}
|
|
auto state=std::make_unique<State>();
|
|
state->runtime.initialize(device_index_);
|
|
auto &runtime=state->runtime;
|
|
state->error=runtime.allocate(sizeof(int));
|
|
for(auto it=Base::begin_populations();it!=Base::end_populations();++it) {
|
|
auto &population=std::get<Population>(*it);
|
|
if(population.size()>10000000) throw std::runtime_error("CUDA population size limit exceeded");
|
|
State::Pop p{};p.host=&population;p.n=static_cast<int>(population.size());
|
|
for(auto &neuron:population) {
|
|
if(!std::isfinite(neuron.potential_) || std::round(neuron.potential_)!=neuron.potential_ ||
|
|
std::abs(neuron.potential_)>=16777216.F || neuron.additional_threshold_!=0)
|
|
throw std::runtime_error("CUDA requires integer finite potential and zero additional threshold");
|
|
unsigned flags=(neuron.is_diff_?1:0)|(neuron.is_reset_?2:0)|(neuron.leak_rev_?4:0)|
|
|
(neuron.saturate_?8:0)|(neuron.do_not_save_?16:0);
|
|
p.states.push_back({neuron.potential_,neuron.pre_impact_potential_,neuron.activity_time_,
|
|
neuron.activation_threshold_,neuron.negative_activation_threshold_,neuron.potential_leak_,
|
|
neuron.potential_reset_value_,flags});
|
|
}
|
|
p.flags.resize(p.n);state->pops.push_back(std::move(p));
|
|
}
|
|
std::vector<std::vector<std::uint64_t>> absolute_bound;
|
|
for(const auto &p:state->pops) absolute_bound.emplace_back(p.n,0);
|
|
for(auto it=Base::begin_projections();it!=Base::end_projections();++it) {
|
|
auto &projection=std::get<Projection>(it->arg_);
|
|
size_t target=state->pops.size();
|
|
for(size_t i=0;i<state->pops.size();++i) if(state->pops[i].host->get_uid()==projection.get_postsynaptic()) target=i;
|
|
if(target==state->pops.size()) throw std::runtime_error("Projection target is not loaded");
|
|
auto &population=state->pops[target];
|
|
if(projection.size()>std::numeric_limits<std::uint32_t>::max()) throw std::runtime_error("CUDA projection too large");
|
|
std::vector<unsigned> rows(population.n+1,0),sources(projection.size()),delays(projection.size());
|
|
std::vector<int> weights(projection.size());
|
|
int source_count=0;
|
|
for(const auto &synapse:projection) {
|
|
const auto ¶meters=std::get<knp::core::synapse_data>(synapse);
|
|
auto source=std::get<knp::core::source_neuron_id>(synapse),post=std::get<knp::core::target_neuron_id>(synapse);
|
|
if(source>=10000000 || post>=static_cast<size_t>(population.n)) throw std::runtime_error("CUDA synapse index outside limits");
|
|
if(!std::isfinite(parameters.weight_) || std::round(parameters.weight_)!=parameters.weight_ ||
|
|
std::abs(parameters.weight_)>65535 || (parameters.weight_==0 && std::signbit(parameters.weight_)))
|
|
throw std::runtime_error("CUDA requires finite integer weights with abs <=65535; negative zero unsupported");
|
|
if(parameters.delay_<1 || parameters.delay_>64) throw std::runtime_error("CUDA supports delays 1..64");
|
|
using Output=knp::synapse_traits::OutputType;
|
|
if(parameters.output_type_!=Output::EXCITATORY && parameters.output_type_!=Output::INHIBITORY_CURRENT)
|
|
throw std::runtime_error("CUDA supports excitatory and inhibitory-current Delta only");
|
|
auto &bound=absolute_bound[target][post];
|
|
bound+=static_cast<std::uint64_t>(std::abs(parameters.weight_));
|
|
if(bound>static_cast<std::uint64_t>(std::numeric_limits<std::int64_t>::max())/(1024*65))
|
|
throw std::runtime_error("CUDA integer queue accumulation bound exceeded");
|
|
++rows[post+1];source_count=std::max(source_count,static_cast<int>(source+1));
|
|
population.slots=std::max(population.slots,static_cast<int>(parameters.delay_+1));
|
|
}
|
|
if(projection.get_presynaptic()) {
|
|
auto pre=std::find_if(state->pops.begin(),state->pops.end(),[&](const auto &p){return p.host->get_uid()==projection.get_presynaptic();});
|
|
if(pre==state->pops.end() || source_count>pre->n) throw std::runtime_error("CUDA invalid presynaptic population/index");
|
|
}
|
|
int max_row=0;
|
|
for(int i=0;i<population.n;++i) {max_row=std::max(max_row,static_cast<int>(rows[i+1]));rows[i+1]+=rows[i];}
|
|
auto cursor=rows;
|
|
for(const auto &synapse:projection) {
|
|
const auto ¶meters=std::get<knp::core::synapse_data>(synapse);
|
|
auto post=std::get<knp::core::target_neuron_id>(synapse);auto index=cursor[post]++;
|
|
sources[index]=std::get<knp::core::source_neuron_id>(synapse);
|
|
weights[index]=static_cast<int>(parameters.weight_)*(parameters.output_type_==knp::synapse_traits::OutputType::INHIBITORY_CURRENT?-1:1);
|
|
delays[index]=parameters.delay_;
|
|
}
|
|
State::Proj p{};p.host=&projection;p.target=target;p.sources=source_count;p.max_row=max_row;p.mask.resize(source_count);
|
|
p.rows=runtime.allocate(rows.size()*sizeof(unsigned),rows.data());
|
|
p.source=runtime.allocate(sources.size()*sizeof(unsigned),sources.data());
|
|
p.weights=runtime.allocate(weights.size()*sizeof(int),weights.data());
|
|
p.delays=runtime.allocate(delays.size()*sizeof(unsigned),delays.data());
|
|
p.counts=runtime.allocate(p.mask.size()*sizeof(int));state->projections.push_back(std::move(p));
|
|
}
|
|
for(auto &p:state->pops) {
|
|
p.neurons=runtime.allocate(p.states.size()*sizeof(NeuronState),p.states.data());
|
|
auto bytes=static_cast<size_t>(p.n)*p.slots*sizeof(std::int64_t);
|
|
p.pending=runtime.allocate(bytes);p.absolute=runtime.allocate(bytes);p.events=runtime.allocate(bytes);
|
|
p.spikes=runtime.allocate(p.flags.size());
|
|
}
|
|
Base::_init();
|
|
state_=std::move(state);
|
|
char name[256];checked(cuDeviceGetName(name,sizeof(name),state_->runtime.device));
|
|
std::cerr<<"KNP AltAI CUDA integer inference initialized on "<<name<<"\n";
|
|
}
|
|
void Backend::_step() {
|
|
if(!state_) throw std::runtime_error("CUDA backend not initialized");
|
|
auto &s=*state_;auto &runtime=s.runtime;runtime.current();
|
|
auto step=static_cast<unsigned long long>(get_step());
|
|
auto &bus=get_message_bus();auto &endpoint=get_message_endpoint();
|
|
bus.route_messages();endpoint.receive_all_messages();
|
|
for(auto &p:s.pops) {
|
|
if(!endpoint.unload_messages<Impact>(p.host->get_uid()).empty())
|
|
throw std::runtime_error("Direct external impact messages unsupported by CUDA integer backend");
|
|
void *arguments[]={&p.neurons,&p.pending,&p.absolute,&p.events,&p.spikes,&s.error,&p.n,&p.slots,&step};
|
|
runtime.launch(runtime.population_kernel,p.n,arguments);
|
|
if(p.n) checked(cuMemcpyDtoH(p.flags.data(),p.spikes,p.flags.size()));
|
|
Spike message{{p.host->get_uid(),get_step()},{}};
|
|
for(int i=0;i<p.n;++i) if(p.flags[i]) message.neuron_indexes_.push_back(i);
|
|
if(!message.neuron_indexes_.empty()) endpoint.send_message(message);
|
|
}
|
|
int error=0;checked(cuMemcpyDtoH(&error,s.error,sizeof(error)));
|
|
if(error) throw std::runtime_error("CUDA exactness guard: potential + absolute incoming impacts reaches 2^24");
|
|
bus.route_messages();endpoint.receive_all_messages();
|
|
for(auto &p:s.projections) {
|
|
auto messages=endpoint.unload_messages<Spike>(p.host->get_uid());
|
|
if(messages.empty()) continue;
|
|
std::fill(p.mask.begin(),p.mask.end(),0);
|
|
bool active=false;
|
|
for(const auto &message:messages) for(auto index:message.neuron_indexes_) {
|
|
// Spikes beyond this projection's used source range have no connections.
|
|
if(index>=p.mask.size()) continue;
|
|
if(++p.mask[index]>1024) throw std::runtime_error("CUDA input multiplicity exceeds 1024 per projection/step");
|
|
active=true;
|
|
}
|
|
if(!active) continue;
|
|
checked(cuMemcpyHtoD(p.counts,p.mask.data(),p.mask.size()*sizeof(int)));
|
|
auto &post=s.pops[p.target];
|
|
void *arguments[]={&p.rows,&p.source,&p.weights,&p.delays,&p.counts,&post.pending,&post.absolute,&post.events,&post.n,&post.slots,&step};
|
|
if(p.max_row>512 && post.n) checked(cuLaunchKernel(runtime.wide_kernel,post.n,1,1,256,1,1,0,nullptr,arguments,nullptr));
|
|
else runtime.launch(runtime.projection_kernel,post.n,arguments);
|
|
}
|
|
bus.route_messages();endpoint.receive_all_messages();
|
|
(void)gad_step();
|
|
}
|
|
void Backend::synchronize() {
|
|
if(!state_) return;
|
|
state_->runtime.current();
|
|
for(auto &p:state_->pops) {
|
|
if(p.n) checked(cuMemcpyDtoH(p.states.data(),p.neurons,p.states.size()*sizeof(NeuronState)));
|
|
for(int i=0;i<p.n;++i) {
|
|
(*p.host)[i].potential_=p.states[i].potential;(*p.host)[i].pre_impact_potential_=p.states[i].pre;
|
|
(*p.host)[i].activity_time_=p.states[i].activity;
|
|
}
|
|
}
|
|
}
|
|
knp::core::Backend::DataRanges Backend::get_network_data() const {
|
|
const_cast<Backend*>(this)->synchronize();return Base::get_network_data();
|
|
}
|
|
void Backend::_uninit() {synchronize();Base::_uninit();}
|
|
BOOST_DLL_ALIAS(knp_altai_cuda::Backend::create, create_knp_backend)
|
|
}
|