172 lines
7.8 KiB
C++
172 lines
7.8 KiB
C++
// CUDA Driver/NVRTC support for a restricted integer AltAILIF/Delta KNP backend.
|
|
#pragma once
|
|
#include <cuda.h>
|
|
#include <nvrtc.h>
|
|
#include <cstdint>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace knp_altai_cuda {
|
|
inline void checked(CUresult code) {
|
|
if (code != CUDA_SUCCESS) {
|
|
const char *message = nullptr;
|
|
cuGetErrorString(code, &message);
|
|
throw std::runtime_error(std::string("AltAI CUDA: ") + (message ? message : "driver error"));
|
|
}
|
|
}
|
|
struct NeuronState {
|
|
float potential, pre;
|
|
std::int64_t activity;
|
|
int threshold, negative, leak, reset;
|
|
unsigned flags;
|
|
};
|
|
static_assert(sizeof(NeuronState) == 40, "CUDA neuron ABI");
|
|
inline const char *kernel_source = R"CUDA(
|
|
struct NeuronState { float potential, pre; long long activity; int threshold, negative, leak, reset; unsigned flags; };
|
|
extern "C" __global__ void population(NeuronState *neurons, long long *pending, long long *absolute,
|
|
long long *events, unsigned char *spikes, int *error,
|
|
int n, int slots, unsigned long long step) {
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n) return;
|
|
NeuronState state = neurons[i];
|
|
float v = (state.flags & 16) ? (float)state.reset : state.potential;
|
|
state.pre = v;
|
|
unsigned long long index = (step % slots) * (unsigned long long)n + i;
|
|
long long bound = (long long)(v < 0 ? -v : v) + absolute[index] + (state.leak < 0 ? -state.leak : state.leak);
|
|
if (bound >= 16777216LL) { atomicExch(error, 1); spikes[i] = 0; return; }
|
|
if (events[index]) v += (float)pending[index];
|
|
pending[index] = absolute[index] = events[index] = 0;
|
|
v += state.leak * (((state.flags & 4) && v < 0) ? -1 : 1);
|
|
if (state.activity > 0) --state.activity;
|
|
else if (state.activity < 0) ++state.activity;
|
|
if (state.activity == 0) state.activity = 9223372036854775807LL;
|
|
bool fired = false, reset = false;
|
|
if (v >= state.threshold) {
|
|
fired = state.activity > 0;
|
|
if (state.flags & 1) v -= state.threshold;
|
|
if (state.flags & 2) { v = (float)state.reset; reset = true; }
|
|
}
|
|
if (v <= -(float)state.negative && !reset) {
|
|
if (state.flags & 8) v = -(float)state.negative;
|
|
else if (state.flags & 2) v = (float)(-state.reset);
|
|
else if (state.flags & 1) v += state.negative;
|
|
}
|
|
state.potential = v;
|
|
neurons[i] = state;
|
|
spikes[i] = fired;
|
|
}
|
|
extern "C" __global__ void projection(const unsigned *rows, const unsigned *sources,
|
|
const int *weights, const unsigned *delays, const int *counts,
|
|
long long *pending, long long *absolute, long long *events,
|
|
int n, int slots, unsigned long long step) {
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n) return;
|
|
for (unsigned s = rows[i]; s < rows[i + 1]; ++s) {
|
|
int count = counts[sources[s]];
|
|
if (!count) continue;
|
|
long long impact = (long long)weights[s] * count;
|
|
unsigned long long index = ((step + delays[s]) % slots) * (unsigned long long)n + i;
|
|
pending[index] += impact;
|
|
absolute[index] += impact < 0 ? -impact : impact;
|
|
events[index] += count;
|
|
}
|
|
}
|
|
|
|
extern "C" __global__ void projection_wide(const unsigned *rows, const unsigned *sources,
|
|
const int *weights, const unsigned *delays, const int *counts,
|
|
long long *pending, long long *absolute, long long *events,
|
|
int n, int slots, unsigned long long step) {
|
|
int i = blockIdx.x, t = threadIdx.x;
|
|
if (i >= n) return;
|
|
__shared__ long long sums[256], absums[256], event_sums[256];
|
|
for (int delay=1;delay<slots;++delay) {
|
|
long long value=0, ab=0, ev=0;
|
|
for (unsigned s=rows[i]+t;s<rows[i+1];s+=256) {
|
|
if (delays[s] != (unsigned)delay) continue;
|
|
int count=counts[sources[s]];
|
|
long long impact=(long long)weights[s]*count;
|
|
value+=impact;ab+=impact<0?-impact:impact;ev+=count;
|
|
}
|
|
sums[t]=value;absums[t]=ab;event_sums[t]=ev;
|
|
__syncthreads();
|
|
for (int shift=128;shift;shift/=2) {
|
|
if(t<shift) {sums[t]+=sums[t+shift];absums[t]+=absums[t+shift];event_sums[t]+=event_sums[t+shift];}
|
|
__syncthreads();
|
|
}
|
|
if(t==0) {
|
|
unsigned long long index=((step+delay)%slots)*(unsigned long long)n+i;
|
|
pending[index]+=sums[0];absolute[index]+=absums[0];events[index]+=event_sums[0];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
}
|
|
)CUDA";
|
|
|
|
class Runtime {
|
|
public:
|
|
CUcontext context = nullptr;
|
|
CUdevice device = 0;
|
|
CUmodule module = nullptr;
|
|
CUfunction population_kernel = nullptr, projection_kernel = nullptr, wide_kernel = nullptr;
|
|
std::vector<CUdeviceptr> allocations;
|
|
Runtime() = default;
|
|
Runtime(const Runtime &) = delete;
|
|
Runtime &operator=(const Runtime &) = delete;
|
|
~Runtime() {
|
|
if (context) {
|
|
cuCtxSetCurrent(context);
|
|
cuCtxSynchronize();
|
|
for (auto pointer : allocations) cuMemFree(pointer);
|
|
if (module) cuModuleUnload(module);
|
|
cuDevicePrimaryCtxRelease(device);
|
|
}
|
|
}
|
|
void initialize(int index) {
|
|
checked(cuInit(0));
|
|
checked(cuDeviceGet(&device, index));
|
|
checked(cuDevicePrimaryCtxRetain(&context, device));
|
|
current();
|
|
int major = 0, minor = 0;
|
|
checked(cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device));
|
|
checked(cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device));
|
|
std::string architecture = "--gpu-architecture=compute_" + std::to_string(major) + std::to_string(minor);
|
|
nvrtcProgram program = nullptr;
|
|
auto code = nvrtcCreateProgram(&program, kernel_source, "knp_altai.cu", 0, nullptr, nullptr);
|
|
if (code != NVRTC_SUCCESS) throw std::runtime_error(nvrtcGetErrorString(code));
|
|
const char *options[] = {architecture.c_str(), "--std=c++11", "--fmad=false"};
|
|
code = nvrtcCompileProgram(program, 3, options);
|
|
if (code != NVRTC_SUCCESS) {
|
|
size_t length = 0;
|
|
nvrtcGetProgramLogSize(program, &length);
|
|
std::string log(length, '\0');
|
|
nvrtcGetProgramLog(program, log.data());
|
|
nvrtcDestroyProgram(&program);
|
|
throw std::runtime_error("NVRTC: " + log);
|
|
}
|
|
size_t length = 0;
|
|
auto result = nvrtcGetPTXSize(program, &length);
|
|
if (result != NVRTC_SUCCESS) { nvrtcDestroyProgram(&program); throw std::runtime_error(nvrtcGetErrorString(result)); }
|
|
std::vector<char> ptx(length);
|
|
result = nvrtcGetPTX(program, ptx.data());
|
|
nvrtcDestroyProgram(&program);
|
|
if (result != NVRTC_SUCCESS) throw std::runtime_error(nvrtcGetErrorString(result));
|
|
checked(cuModuleLoadData(&module, ptx.data()));
|
|
checked(cuModuleGetFunction(&population_kernel, module, "population"));
|
|
checked(cuModuleGetFunction(&projection_kernel, module, "projection"));
|
|
checked(cuModuleGetFunction(&wide_kernel, module, "projection_wide"));
|
|
}
|
|
void current() const { checked(cuCtxSetCurrent(context)); }
|
|
CUdeviceptr allocate(size_t bytes, const void *data = nullptr) {
|
|
CUdeviceptr result = 0;
|
|
checked(cuMemAlloc(&result, bytes ? bytes : 4));
|
|
allocations.push_back(result);
|
|
if (data && bytes) checked(cuMemcpyHtoD(result, data, bytes));
|
|
else checked(cuMemsetD8(result, 0, bytes ? bytes : 4));
|
|
return result;
|
|
}
|
|
void launch(CUfunction function, int n, void **arguments) {
|
|
if (n) checked(cuLaunchKernel(function, (n + 255) / 256, 1, 1, 256, 1, 1, 0, nullptr, arguments, nullptr));
|
|
}
|
|
};
|
|
} // namespace knp_altai_cuda
|