Digital Drosophila Archive
214 lines
10 KiB
Plaintext
214 lines
10 KiB
Plaintext
#include <torch/extension.h>
|
|
#include <ATen/cuda/CUDAContext.h>
|
|
#include <c10/cuda/CUDAGuard.h>
|
|
#include <c10/cuda/CUDAException.h>
|
|
#include <cuda.h>
|
|
#include <cuda_runtime.h>
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <limits>
|
|
|
|
namespace {
|
|
constexpr int BLOCK = 128;
|
|
|
|
void check_tensor(const torch::Tensor& x, const torch::Tensor& like,
|
|
at::ScalarType dtype, const char* name) {
|
|
TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.dim() == 1,
|
|
name, " must be a contiguous CUDA vector");
|
|
TORCH_CHECK(x.device() == like.device(), name, " is on a different CUDA device");
|
|
TORCH_CHECK(x.scalar_type() == dtype, name, " has the wrong dtype");
|
|
}
|
|
void check_state(const torch::Tensor& g, const torch::Tensor& ref) {
|
|
TORCH_CHECK(g.scalar_type() == torch::kFloat32 || g.scalar_type() == torch::kFloat64,
|
|
"state must be float32 or float64");
|
|
check_tensor(g, g, g.scalar_type(), "g");
|
|
check_tensor(ref, g, torch::kInt32, "ref");
|
|
TORCH_CHECK(g.numel() == ref.numel(), "g/ref length mismatch");
|
|
TORCH_CHECK(g.numel() <= std::numeric_limits<int>::max(), "too many neurons");
|
|
}
|
|
void index_range(const torch::Tensor& x, int64_t upper, const char* name) {
|
|
if (x.numel()) {
|
|
TORCH_CHECK(x.min().item<int64_t>() >= 0 && x.max().item<int64_t>() < upper,
|
|
name, " index out of bounds");
|
|
}
|
|
}
|
|
|
|
template <typename T> __device__ __forceinline__ T add_rn(T a, T b);
|
|
template <> __device__ __forceinline__ float add_rn(float a, float b) { return __fadd_rn(a, b); }
|
|
template <> __device__ __forceinline__ double add_rn(double a, double b) { return __dadd_rn(a, b); }
|
|
template <typename T> __device__ __forceinline__ T mul_rn(T a, T b);
|
|
template <> __device__ __forceinline__ float mul_rn(float a, float b) { return __fmul_rn(a, b); }
|
|
template <> __device__ __forceinline__ double mul_rn(double a, double b) { return __dmul_rn(a, b); }
|
|
|
|
template <typename T>
|
|
__global__ void state_kernel(T* v, T* g, int32_t* ref, const uint8_t* requests,
|
|
uint8_t* spikes, int32_t* counts, int n, T gd, T c0,
|
|
T cg, T cv, T threshold, T reset, T seed_volt, int ref_steps) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n) return;
|
|
int r = ref[i];
|
|
r = r > 0 ? r - 1 : 0;
|
|
const bool available = r == 0;
|
|
T vi = v[i];
|
|
if (available) {
|
|
const T oldg = g[i];
|
|
const T oldv = vi;
|
|
// Match NumPy float-array expression evaluation: each operation rounds in T.
|
|
g[i] = mul_rn(gd, oldg);
|
|
const T driven = mul_rn(cg, oldg);
|
|
const T leaked = mul_rn(cv, oldv);
|
|
vi = add_rn(c0, add_rn(driven, leaked));
|
|
if (requests[i]) vi = seed_volt;
|
|
}
|
|
const bool fired = available && vi > threshold;
|
|
if (fired) {
|
|
vi = reset;
|
|
g[i] = T(0);
|
|
r = ref_steps;
|
|
counts[i] += 1;
|
|
}
|
|
// Nonavailable neurons retain V/g. Resetting fired cells before arrivals is
|
|
// equivalent to reset-after-arrivals because all their arrivals are rejected.
|
|
v[i] = vi;
|
|
ref[i] = r;
|
|
spikes[i] = fired ? 1 : 0;
|
|
}
|
|
|
|
__global__ void fill_kernel(const int64_t* indptr, const int64_t* sources,
|
|
const int64_t* offsets, int64_t* output, int n) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n) return;
|
|
const int64_t source = sources[i];
|
|
const int64_t begin = indptr[source];
|
|
const int64_t end = indptr[source + 1];
|
|
for (int64_t j = begin; j < end; ++j) output[offsets[i] + j - begin] = j;
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void arrival_kernel(const int64_t* edges, const int64_t* starts,
|
|
const int32_t* targets, const T* weights, T* g,
|
|
const int32_t* ref, int segments, int64_t edge_count) {
|
|
const int segment = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (segment >= segments) return;
|
|
const int target = targets[segment];
|
|
if (ref[target] != 0) return;
|
|
const int64_t end = segment + 1 < segments ? starts[segment + 1] : edge_count;
|
|
// Start at existing g, never sum weights separately: cancellation can differ.
|
|
T accumulator = g[target];
|
|
for (int64_t j = starts[segment]; j < end; ++j)
|
|
accumulator = add_rn(accumulator, weights[edges[j]]);
|
|
g[target] = accumulator;
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void external_kernel(const T* external, T* g, const int32_t* ref, int n) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i < n && ref[i] == 0 && external[i] != T(0))
|
|
g[i] = add_rn(g[i], external[i]);
|
|
}
|
|
} // namespace
|
|
|
|
void author_state_cuda(torch::Tensor v, torch::Tensor g, torch::Tensor ref,
|
|
torch::Tensor requests, torch::Tensor spikes, torch::Tensor counts,
|
|
double gd, double c0, double cg, double cv, double threshold,
|
|
double reset, double seed_volt, int64_t refractory_steps) {
|
|
check_state(g, ref);
|
|
const c10::cuda::CUDAGuard guard(g.device());
|
|
check_tensor(v, g, g.scalar_type(), "v");
|
|
check_tensor(requests, g, torch::kUInt8, "requests");
|
|
check_tensor(spikes, g, torch::kUInt8, "spikes");
|
|
check_tensor(counts, g, torch::kInt32, "counts");
|
|
const int64_t n = g.numel();
|
|
TORCH_CHECK(v.numel() == n && requests.numel() == n && spikes.numel() == n && counts.numel() == n,
|
|
"state vector length mismatch");
|
|
TORCH_CHECK(refractory_steps > 0 && refractory_steps <= std::numeric_limits<int>::max(),
|
|
"invalid refractory length");
|
|
for (double x : {gd, c0, cg, cv, threshold, reset, seed_volt})
|
|
TORCH_CHECK(std::isfinite(x), "nonfinite coefficient");
|
|
if (!n) return;
|
|
const auto stream = at::cuda::getCurrentCUDAStream(g.get_device());
|
|
AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "author_state", [&] {
|
|
state_kernel<scalar_t><<<(n + BLOCK - 1) / BLOCK, BLOCK, 0, stream>>>(
|
|
v.data_ptr<scalar_t>(), g.data_ptr<scalar_t>(), ref.data_ptr<int32_t>(),
|
|
requests.data_ptr<uint8_t>(), spikes.data_ptr<uint8_t>(), counts.data_ptr<int32_t>(),
|
|
static_cast<int>(n), static_cast<scalar_t>(gd), static_cast<scalar_t>(c0),
|
|
static_cast<scalar_t>(cg), static_cast<scalar_t>(cv), static_cast<scalar_t>(threshold),
|
|
static_cast<scalar_t>(reset), static_cast<scalar_t>(seed_volt), static_cast<int>(refractory_steps));
|
|
});
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|
|
|
|
void author_fill_edges_cuda(torch::Tensor indptr, torch::Tensor sources,
|
|
torch::Tensor offsets, torch::Tensor output) {
|
|
check_tensor(indptr, indptr, torch::kInt64, "indptr");
|
|
const c10::cuda::CUDAGuard guard(indptr.device());
|
|
check_tensor(sources, indptr, torch::kInt64, "sources");
|
|
check_tensor(offsets, indptr, torch::kInt64, "offsets");
|
|
check_tensor(output, indptr, torch::kInt64, "output");
|
|
TORCH_CHECK(indptr.numel() >= 1 && sources.numel() == offsets.numel(), "CSR shape mismatch");
|
|
TORCH_CHECK(sources.numel() <= std::numeric_limits<int>::max(), "too many sources");
|
|
TORCH_CHECK(indptr[0].item<int64_t>() == 0, "CSR indptr must begin at zero");
|
|
if (indptr.numel() > 1)
|
|
TORCH_CHECK((indptr.slice(0, 1) >= indptr.slice(0, 0, -1)).all().item<bool>(), "CSR indptr not monotonic");
|
|
index_range(sources, indptr.numel() - 1, "source");
|
|
if (!sources.numel()) {
|
|
TORCH_CHECK(output.numel() == 0, "nonempty expansion for no active sources");
|
|
return;
|
|
}
|
|
const auto degrees = indptr.index_select(0, sources + 1) - indptr.index_select(0, sources);
|
|
const auto expected_offsets = degrees.cumsum(0) - degrees;
|
|
TORCH_CHECK(torch::equal(offsets, expected_offsets), "incorrect output offsets");
|
|
TORCH_CHECK(degrees.sum().item<int64_t>() == output.numel(), "incorrect expansion length");
|
|
const auto stream = at::cuda::getCurrentCUDAStream(indptr.get_device());
|
|
fill_kernel<<<(sources.numel() + BLOCK - 1) / BLOCK, BLOCK, 0, stream>>>(
|
|
indptr.data_ptr<int64_t>(), sources.data_ptr<int64_t>(), offsets.data_ptr<int64_t>(),
|
|
output.data_ptr<int64_t>(), static_cast<int>(sources.numel()));
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|
|
|
|
void author_arrivals_cuda(torch::Tensor edges, torch::Tensor starts, torch::Tensor targets,
|
|
torch::Tensor weights, torch::Tensor g, torch::Tensor ref) {
|
|
check_state(g, ref);
|
|
const c10::cuda::CUDAGuard guard(g.device());
|
|
check_tensor(edges, g, torch::kInt64, "edges");
|
|
check_tensor(starts, g, torch::kInt64, "starts");
|
|
check_tensor(targets, g, torch::kInt32, "targets");
|
|
check_tensor(weights, g, g.scalar_type(), "weights");
|
|
TORCH_CHECK(starts.numel() == targets.numel(), "segment count mismatch");
|
|
TORCH_CHECK(starts.numel() <= std::numeric_limits<int>::max(), "too many target segments");
|
|
if (!starts.numel()) {
|
|
TORCH_CHECK(!edges.numel(), "edges without target segments");
|
|
return;
|
|
}
|
|
index_range(edges, weights.numel(), "edge");
|
|
index_range(starts, edges.numel(), "segment start");
|
|
index_range(targets, g.numel(), "target");
|
|
TORCH_CHECK(starts[0].item<int64_t>() == 0, "first segment must start at zero");
|
|
if (starts.numel() > 1) {
|
|
TORCH_CHECK((starts.slice(0, 1) > starts.slice(0, 0, -1)).all().item<bool>(), "segment starts not increasing");
|
|
TORCH_CHECK((targets.slice(0, 1) > targets.slice(0, 0, -1)).all().item<bool>(), "targets must be unique and increasing");
|
|
}
|
|
const auto stream = at::cuda::getCurrentCUDAStream(g.get_device());
|
|
AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "author_arrivals", [&] {
|
|
arrival_kernel<scalar_t><<<(starts.numel() + BLOCK - 1) / BLOCK, BLOCK, 0, stream>>>(
|
|
edges.data_ptr<int64_t>(), starts.data_ptr<int64_t>(), targets.data_ptr<int32_t>(),
|
|
weights.data_ptr<scalar_t>(), g.data_ptr<scalar_t>(), ref.data_ptr<int32_t>(),
|
|
static_cast<int>(starts.numel()), edges.numel());
|
|
});
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|
|
|
|
void author_external_cuda(torch::Tensor external, torch::Tensor g, torch::Tensor ref) {
|
|
check_state(g, ref);
|
|
const c10::cuda::CUDAGuard guard(g.device());
|
|
check_tensor(external, g, g.scalar_type(), "external");
|
|
TORCH_CHECK(external.numel() == g.numel(), "external/state length mismatch");
|
|
if (!g.numel()) return;
|
|
const auto stream = at::cuda::getCurrentCUDAStream(g.get_device());
|
|
AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "author_external", [&] {
|
|
external_kernel<scalar_t><<<(g.numel() + BLOCK - 1) / BLOCK, BLOCK, 0, stream>>>(
|
|
external.data_ptr<scalar_t>(), g.data_ptr<scalar_t>(), ref.data_ptr<int32_t>(), static_cast<int>(g.numel()));
|
|
});
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|