Archive MaleCNS research sources, provenance, bibliography and immutable data inventory

2026-09-29 03:38:32 +08:00
commit 3981628ff8
682 changed files with 161475 additions and 0 deletions
@@ -0,0 +1,152 @@
+#include <torch/extension.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace {
+
+__global__ void state_kernel(float* v, float* g, int* refractory,
+                             float* arrival, unsigned char* spikes,
+                             int* spike_counts, int n, float g_decay,
+                             float alpha, float rest, float reset,
+                             float threshold, int refractory_steps) {
+  int i = blockIdx.x * blockDim.x + threadIdx.x;
+  if (i >= n) return;
+  // Phase 4B executes mul_ and add_ as separate kernels: prevent contraction.
+  float gi = __fadd_rn(__fmul_rn(g[i], g_decay), 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) {
+    // Match the separate float32 operations used by the eager reference.
+    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] = refractory_steps;
+    spike_counts[i] += 1;
+  }
+  v[i] = vi;
+  spikes[i] = fired ? 1 : 0;
+}
+
+__global__ void force_kernel(float* v, int* refractory,
+                             unsigned char* spikes, int* spike_counts,
+                             const int64_t* targets,
+                             const unsigned char* forced, int count,
+                             float reset, int refractory_steps) {
+  int j = blockIdx.x * blockDim.x + threadIdx.x;
+  if (j >= count || !forced[j]) return;
+  int64_t i = targets[j];
+  // A natural spike already recorded at this target represents the same event.
+  if (spikes[i]) return;
+  if (refractory[i] == 0) {
+    spikes[i] = 1;
+    v[i] = reset;
+    refractory[i] = refractory_steps;
+    spike_counts[i] += 1;
+  }
+}
+
+__global__ void fanout_kernel(const unsigned char* spikes,
+                              const int64_t* indptr, const int* indices,
+                              const float* weights, float* destination,
+                              int n) {
+  int src = blockIdx.x * blockDim.x + threadIdx.x;
+  if (src >= n || !spikes[src]) return;
+  int64_t start = indptr[src];
+  int64_t stop = indptr[src + 1];
+  for (int64_t e = start; e < stop; ++e) {
+    atomicAdd(destination + indices[e], weights[e]);
+  }
+}
+
+__global__ void fill_edges_kernel(const int64_t* indptr,
+                                  const int64_t* sources,
+                                  const int64_t* output_offsets,
+                                  int64_t* output, int source_count) {
+  int j = blockIdx.x * blockDim.x + threadIdx.x;
+  if (j >= source_count) return;
+  int64_t src = sources[j];
+  int64_t begin = indptr[src];
+  int64_t degree = indptr[src + 1] - begin;
+  int64_t out = output_offsets[j];
+  for (int64_t k = 0; k < degree; ++k) output[out + k] = begin + k;
+}
+
+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");
+}
+}  // namespace
+
+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) {
+  check_cuda(v); check_cuda(g); check_cuda(refractory); check_cuda(arrival);
+  check_cuda(spikes); check_cuda(spike_counts);
+  TORCH_CHECK(v.scalar_type() == torch::kFloat32 && g.scalar_type() == torch::kFloat32,
+              "v/g must be float32");
+  TORCH_CHECK(refractory.scalar_type() == torch::kInt32 && spike_counts.scalar_type() == torch::kInt32,
+              "refractory/counts must be int32");
+  TORCH_CHECK(spikes.scalar_type() == torch::kUInt8, "spikes must be uint8");
+  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>(),
+      spike_counts.data_ptr<int>(), n, (float)g_decay, (float)alpha,
+      (float)rest, (float)reset, (float)threshold, (int)refractory_steps);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+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) {
+  check_cuda(targets); check_cuda(forced);
+  int count = targets.numel();
+  force_kernel<<<(count + 255) / 256, 256>>>(
+      v.data_ptr<float>(), refractory.data_ptr<int>(),
+      spikes.data_ptr<unsigned char>(), spike_counts.data_ptr<int>(),
+      targets.data_ptr<int64_t>(), forced.data_ptr<unsigned char>(), count,
+      (float)reset, (int)refractory_steps);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+void fanout_cuda(torch::Tensor spikes, torch::Tensor indptr,
+                 torch::Tensor indices, torch::Tensor weights,
+                 torch::Tensor destination) {
+  check_cuda(spikes); check_cuda(indptr); check_cuda(indices);
+  check_cuda(weights); check_cuda(destination);
+  int n = spikes.numel();
+  fanout_kernel<<<(n + 255) / 256, 256>>>(
+      spikes.data_ptr<unsigned char>(), indptr.data_ptr<int64_t>(),
+      indices.data_ptr<int>(), weights.data_ptr<float>(),
+      destination.data_ptr<float>(), n);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+void fill_edges_cuda(torch::Tensor indptr, torch::Tensor sources,
+                     torch::Tensor output_offsets, torch::Tensor output) {
+  check_cuda(indptr); check_cuda(sources); check_cuda(output_offsets); check_cuda(output);
+  TORCH_CHECK(indptr.scalar_type() == torch::kInt64 && sources.scalar_type() == torch::kInt64 &&
+              output_offsets.scalar_type() == torch::kInt64 && output.scalar_type() == torch::kInt64,
+              "fill_edges tensors must be int64");
+  int count = sources.numel();
+  fill_edges_kernel<<<(count + 255) / 256, 256>>>(
+      indptr.data_ptr<int64_t>(), sources.data_ptr<int64_t>(),
+      output_offsets.data_ptr<int64_t>(), output.data_ptr<int64_t>(), count);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}