Digital Drosophila Archive
26 lines
1.4 KiB
C++
26 lines
1.4 KiB
C++
#include <torch/extension.h>
|
|
|
|
void state_step_cuda(torch::Tensor v, torch::Tensor g, torch::Tensor refractory,
|
|
torch::Tensor arrival, torch::Tensor spikes,
|
|
torch::Tensor spike_counts, double g_decay, double alpha,
|
|
double rest, double reset, double threshold,
|
|
int64_t refractory_steps);
|
|
void force_step_cuda(torch::Tensor v, torch::Tensor refractory,
|
|
torch::Tensor spikes, torch::Tensor spike_counts,
|
|
torch::Tensor targets, torch::Tensor forced,
|
|
double reset, int64_t refractory_steps);
|
|
void fill_edges_cuda(torch::Tensor indptr, torch::Tensor sources,
|
|
torch::Tensor output_offsets, torch::Tensor output);
|
|
void segmented_sum_cuda(torch::Tensor sorted_edge_ids,
|
|
torch::Tensor segment_starts,
|
|
torch::Tensor segment_targets,
|
|
torch::Tensor weights,
|
|
torch::Tensor destination);
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("state_step", &state_step_cuda, "Canonical float32 LIF update (CUDA)");
|
|
m.def("force_step", &force_step_cuda, "Canonical external spike injection (CUDA)");
|
|
m.def("fill_edges", &fill_edges_cuda, "Source/edge ordered CSR expansion (CUDA)");
|
|
m.def("segmented_sum", &segmented_sum_cuda, "Fixed-order sequential segmented float32 sum (CUDA)");
|
|
}
|