cwenzi/neuroflow-cpp

1/*2 * NeuroFlow GPU Kernels (HIP/ROCm)3 * 编译: hipcc -shared -fPIC -O3 -o libneuroflow_gpu.so neuroflow_gpu.cpp --offload-arch=gfx90a4 * 用法: Python ctypes加载5 */6 7#include <hip/hip_runtime.h>8#include <hipblas.h>9#include <cmath>10#include <cstdio>11#include <cstdlib>12 13#define CHECK_HIP(cmd) do { \14    hipError_t e = cmd; \15    if (e != hipSuccess) { \16        fprintf(stderr, "HIP error %s:%d '%s'(%d)\n", __FILE__, __LINE__, hipGetErrorString(e), e); \17        exit(1); \18    } \19} while(0)20 21// ═══════════════════════════════════════════22// GPU Kernels23// ═══════════════════════════════════════════24 25// ReLU: y = max(x, 0)26__global__ void relu_kernel(float* y, const float* x, int n) {27    int i = blockIdx.x * blockDim.x + threadIdx.x;28    if (i < n) y[i] = fmaxf(x[i], 0.0f);29}30 31// Sigmoid: y = 1/(1+exp(-x))32__global__ void sigmoid_kernel(float* y, const float* x, int n) {33    int i = blockIdx.x * blockDim.x + threadIdx.x;34    if (i < n) y[i] = 1.0f / (1.0f + expf(-x[i]));35}36 37// L2 normalize rows (in-place)38__global__ void l2_normalize_rows_kernel(float* x, int rows, int cols) {39    int row = blockIdx.x;40    if (row >= rows) return;41    float sum_sq = 0.0f;42    for (int j = 0; j < cols; j++) {43        float v = x[row * cols + j];44        sum_sq += v * v;45    }46    float norm = sqrtf(sum_sq) + 1e-8f;47    for (int j = 0; j < cols; j++) {48        x[row * cols + j] /= norm;49    }50}51 52// Masked noise: y = x * mask + noise53__global__ void mask_noise_kernel(float* y, const float* x, const float* mask, int n, float noise_std, unsigned long seed) {54    int i = blockIdx.x * blockDim.x + threadIdx.x;55    if (i >= n) return;56    // Simple LCG random57    unsigned long s = seed + i * 2654435761UL;58    float r = (float)(s & 0xFFFFFF) / 0xFFFFFF - 0.5f;59    y[i] = x[i] * mask[i] + r * noise_std;60}61 62// MSE loss gradient: grad = 2*(pred - target)/N63__global__ void mse_grad_kernel(float* grad, const float* pred, const float* target, int n) {64    int i = blockIdx.x * blockDim.x + threadIdx.x;65    if (i < n) grad[i] = 2.0f * (pred[i] - target[i]) / static_cast<float>(n);66}67 68// SGD update: w -= lr * (grad + wd * w)69__global__ void sgd_update_kernel(float* w, const float* grad, int n, float lr, float wd) {70    int i = blockIdx.x * blockDim.x + threadIdx.x;71    if (i < n) w[i] -= lr * (grad[i] + wd * w[i]);72}73 74// Masked top-K selection (SAE): keep only top K75__global__ void topk_mask_kernel(float* y, int n, int k, int offset) {76    // Each block processes one row77    extern __shared__ float shared[];78    int tid = threadIdx.x;79    int row = blockIdx.x;80    81    // Load row into shared memory82    if (tid < k * 2) {  // double buffer83        int idx = row * n + offset + tid;84        shared[tid] = (tid < n - offset) ? fabsf(y[idx]) : -1.0f;85    }86    __syncthreads();87    88    // Simple bitonic sort for top-k (k is small, e.g., 65)89    // ... simplified: just threshold90    if (tid == 0) {91        // Copy to local, sort, find threshold92        float local[256];93        int len = min(n - offset, 256);94        for (int j = 0; j < len; j++) {95            int idx = row * n + offset + j;96            local[j] = fabsf(y[idx]);97        }98        // Simple bubble sort for top-k99        for (int j = 0; j < k && j < len; j++) {100            int max_idx = j;101            for (int m = j+1; m < len; m++) {102                if (local[m] > local[max_idx]) max_idx = m;103            }104            float tmp = local[j];105            local[j] = local[max_idx];106            local[max_idx] = tmp;107        }108        float threshold = (k <= len) ? local[k-1] : 0.0f;109        // Apply mask110        for (int j = 0; j < len; j++) {111            int idx = row * n + offset + j;112            if (fabsf(y[idx]) < threshold && fabsf(y[idx]) > 0) {113                y[idx] = 0.0f;114            }115        }116    }117}118 119// ═══════════════════════════════════════════120// Python-callable C functions121// ═══════════════════════════════════════════122 123extern "C" {124 125// GPU info126int gpu_get_count() {127    int count;128    CHECK_HIP(hipGetDeviceCount(&count));129    return count;130}131 132void gpu_get_name(int id, char* name, int max_len) {133    hipDeviceProp_t prop;134    CHECK_HIP(hipGetDeviceProperties(&prop, id));135    strncpy(name, prop.name, max_len);136}137 138// Init139void gpu_init(int device_id) {140    CHECK_HIP(hipSetDevice(device_id));141    printf("[GPU] Initialized device %d\n", device_id);142}143 144// Allocate / Free145float* gpu_alloc(int n) {146    float* ptr;147    CHECK_HIP(hipMalloc(&ptr, n * sizeof(float)));148    return ptr;149}150 151void gpu_free(float* ptr) {152    CHECK_HIP(hipFree(ptr));153}154 155// Copy H→D, D→H156void gpu_memcpy_htod(float* d_ptr, const float* h_ptr, int n) {157    CHECK_HIP(hipMemcpy(d_ptr, h_ptr, n * sizeof(float), hipMemcpyHostToDevice));158}159 160void gpu_memcpy_dtoh(float* h_ptr, const float* d_ptr, int n) {161    CHECK_HIP(hipMemcpy(h_ptr, d_ptr, n * sizeof(float), hipMemcpyDeviceToHost));162}163 164// ReLU165void gpu_relu(float* d_y, const float* d_x, int n) {166    int block = 256;167    int grid = (n + block - 1) / block;168    relu_kernel<<<grid, block>>>(d_y, d_x, n);169    CHECK_HIP(hipGetLastError());170    CHECK_HIP(hipDeviceSynchronize());171}172 173// Sigmoid174void gpu_sigmoid(float* d_y, const float* d_x, int n) {175    int block = 256;176    int grid = (n + block - 1) / block;177    sigmoid_kernel<<<grid, block>>>(d_y, d_x, n);178    CHECK_HIP(hipGetLastError());179    CHECK_HIP(hipDeviceSynchronize());180}181 182// L2 normalize rows183void gpu_l2_norm_rows(float* d_x, int rows, int cols) {184    l2_normalize_rows_kernel<<<rows, 1>>>(d_x, rows, cols);185    CHECK_HIP(hipGetLastError());186    CHECK_HIP(hipDeviceSynchronize());187}188 189// Mask + Noise190void gpu_mask_noise(float* d_y, const float* d_x, const float* d_mask, int n, float noise_std, unsigned long seed) {191    int block = 256;192    int grid = (n + block - 1) / block;193    mask_noise_kernel<<<grid, block>>>(d_y, d_x, d_mask, n, noise_std, seed);194    CHECK_HIP(hipGetLastError());195    CHECK_HIP(hipDeviceSynchronize());196}197 198// Matrix multiply: C = A @ B  (using hipBLAS)199// A: [M, K], B: [K, N], C: [M, N]200void gpu_matmul(float* d_C, const float* d_A, const float* d_B, int M, int N, int K) {201    static hipblasHandle_t handle = nullptr;202    if (!handle) hipblasCreate(&handle);203    204    float alpha = 1.0f, beta = 0.0f;205    hipblasSgemm(handle, HIPBLAS_OP_N, HIPBLAS_OP_N,206                 N, M, K,207                 &alpha, d_B, N, d_A, K,208                 &beta, d_C, N);209    CHECK_HIP(hipDeviceSynchronize());210}211 212// SGD update213void gpu_sgd_update(float* d_w, const float* d_grad, int n, float lr, float wd) {214    int block = 256;215    int grid = (n + block - 1) / block;216    sgd_update_kernel<<<grid, block>>>(d_w, d_grad, n, lr, wd);217    CHECK_HIP(hipGetLastError());218    CHECK_HIP(hipDeviceSynchronize());219}220 221// Fill with zeros222void gpu_fill_zero(float* d_ptr, int n) {223    CHECK_HIP(hipMemset(d_ptr, 0, n * sizeof(float)));224}225 226// Print GPU memory info227void gpu_mem_info() {228    size_t free, total;229    CHECK_HIP(hipMemGetInfo(&free, &total));230    printf("[GPU] Memory: %.1f GB free / %.1f GB total\n", 231           free / 1e9, total / 1e9);232}233 234} // extern "C"235