Digital Drosophila Archive
88 lines
4.9 KiB
Plaintext
88 lines
4.9 KiB
Plaintext
#include <torch/extension.h>
|
|
#include <c10/cuda/CUDAException.h>
|
|
#include <cuda.h>
|
|
#include <cuda_runtime.h>
|
|
#include <cstdint>
|
|
|
|
namespace {
|
|
void check_cuda(const torch::Tensor& x) {
|
|
TORCH_CHECK(x.is_cuda(), "all tensors must be CUDA tensors");
|
|
TORCH_CHECK(x.is_contiguous(), "all tensors must be contiguous");
|
|
}
|
|
|
|
__global__ void state_kernel(float* v, float* g, int* refractory,
|
|
float* arrival, unsigned char* spikes,
|
|
int* counts, int n, float gd, float alpha,
|
|
float rest, float reset, float threshold,
|
|
int ref_steps) {
|
|
int i=blockIdx.x*blockDim.x+threadIdx.x;
|
|
if(i>=n) return;
|
|
float gi=__fadd_rn(__fmul_rn(g[i],gd),arrival[i]);
|
|
g[i]=gi; arrival[i]=0.0f;
|
|
int r=refractory[i]; r=r>0?r-1:0; refractory[i]=r;
|
|
float vi=v[i]; bool available=r==0;
|
|
if(available){
|
|
float leak=__fsub_rn(vi,rest);
|
|
float drive=__fsub_rn(gi,leak);
|
|
vi=__fadd_rn(vi,__fmul_rn(alpha,drive));
|
|
} else vi=reset;
|
|
bool fired=available && vi>=threshold;
|
|
if(fired){vi=reset;refractory[i]=ref_steps;counts[i]+=1;}
|
|
v[i]=vi;spikes[i]=fired?1:0;
|
|
}
|
|
|
|
__global__ void force_kernel(float* v,int* refractory,unsigned char* spikes,
|
|
int* counts,const int64_t* targets,
|
|
const unsigned char* forced,int n,float reset,
|
|
int ref_steps){
|
|
int j=blockIdx.x*blockDim.x+threadIdx.x;
|
|
if(j>=n||!forced[j]) return;
|
|
int64_t i=targets[j]; if(spikes[i]||refractory[i]!=0) return;
|
|
spikes[i]=1;v[i]=reset;refractory[i]=ref_steps;counts[i]+=1;
|
|
}
|
|
|
|
__global__ void fill_edges_kernel(const int64_t* indptr,const int64_t* sources,
|
|
const int64_t* offsets,int64_t* output,int n){
|
|
int j=blockIdx.x*blockDim.x+threadIdx.x;
|
|
if(j>=n) return;
|
|
int64_t src=sources[j],begin=indptr[src],degree=indptr[src+1]-begin,out=offsets[j];
|
|
for(int64_t k=0;k<degree;++k) output[out+k]=begin+k;
|
|
}
|
|
|
|
// Exactly one CUDA thread owns a target segment. The loop order is the
|
|
// canonical target -> source -> edge_id order supplied by the host.
|
|
__global__ void segmented_sum_kernel(const int64_t* edges,const int64_t* starts,
|
|
const int* targets,const float* weights,
|
|
float* destination,int segments,
|
|
int64_t edge_count){
|
|
int s=blockIdx.x*blockDim.x+threadIdx.x;
|
|
if(s>=segments) return;
|
|
int64_t begin=starts[s],end=(s+1<segments)?starts[s+1]:edge_count;
|
|
float acc=0.0f;
|
|
for(int64_t j=begin;j<end;++j) acc=__fadd_rn(acc,weights[edges[j]]);
|
|
destination[targets[s]]=__fadd_rn(destination[targets[s]],acc);
|
|
}
|
|
}
|
|
|
|
void state_step_cuda(torch::Tensor v,torch::Tensor g,torch::Tensor refractory,
|
|
torch::Tensor arrival,torch::Tensor spikes,torch::Tensor counts,
|
|
double gd,double alpha,double rest,double reset,double threshold,
|
|
int64_t ref_steps){
|
|
check_cuda(v);check_cuda(g);check_cuda(refractory);check_cuda(arrival);check_cuda(spikes);check_cuda(counts);
|
|
int n=v.numel();state_kernel<<<(n+255)/256,256>>>(v.data_ptr<float>(),g.data_ptr<float>(),refractory.data_ptr<int>(),arrival.data_ptr<float>(),spikes.data_ptr<unsigned char>(),counts.data_ptr<int>(),n,(float)gd,(float)alpha,(float)rest,(float)reset,(float)threshold,(int)ref_steps);C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|
|
void force_step_cuda(torch::Tensor v,torch::Tensor refractory,torch::Tensor spikes,
|
|
torch::Tensor counts,torch::Tensor targets,torch::Tensor forced,
|
|
double reset,int64_t ref_steps){
|
|
check_cuda(targets);check_cuda(forced);int n=targets.numel();force_kernel<<<(n+255)/256,256>>>(v.data_ptr<float>(),refractory.data_ptr<int>(),spikes.data_ptr<unsigned char>(),counts.data_ptr<int>(),targets.data_ptr<int64_t>(),forced.data_ptr<unsigned char>(),n,(float)reset,(int)ref_steps);C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|
|
void fill_edges_cuda(torch::Tensor indptr,torch::Tensor sources,torch::Tensor offsets,torch::Tensor output){
|
|
check_cuda(indptr);check_cuda(sources);check_cuda(offsets);check_cuda(output);int n=sources.numel();fill_edges_kernel<<<(n+255)/256,256>>>(indptr.data_ptr<int64_t>(),sources.data_ptr<int64_t>(),offsets.data_ptr<int64_t>(),output.data_ptr<int64_t>(),n);C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|
|
void segmented_sum_cuda(torch::Tensor edges,torch::Tensor starts,torch::Tensor targets,
|
|
torch::Tensor weights,torch::Tensor destination){
|
|
check_cuda(edges);check_cuda(starts);check_cuda(targets);check_cuda(weights);check_cuda(destination);
|
|
TORCH_CHECK(edges.scalar_type()==torch::kInt64&&starts.scalar_type()==torch::kInt64&&targets.scalar_type()==torch::kInt32,"index dtype mismatch");
|
|
int n=targets.numel();if(n) segmented_sum_kernel<<<(n+255)/256,256>>>(edges.data_ptr<int64_t>(),starts.data_ptr<int64_t>(),targets.data_ptr<int>(),weights.data_ptr<float>(),destination.data_ptr<float>(),n,edges.numel());C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|