@rlabs-inc/sparse 0.1.0

package/native/sparse.metal

@@ -0,0 +1,215 @@
+// ============================================================================
+// SPARSE - Metal Compute Shaders for Biological Neural Simulation
+// Sparsity-first GPU compute. No transformer bloat.
+// ============================================================================
+
+#include <metal_stdlib>
+using namespace metal;
+
+// ============================================================================
+// SCATTER-ADD KERNEL
+// The core operation for sparse synapse transmission.
+// Thread-safe atomic add - indices are READ-ONLY, never corrupted.
+// ============================================================================
+
+kernel void scatter_add(
+    device float* target [[buffer(0)]],           // Target array to add into
+    device const uint* indices [[buffer(1)]],     // Indices (READ-ONLY!)
+    device const float* values [[buffer(2)]],     // Values to add
+    constant uint& count [[buffer(3)]],           // Number of elements
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+
+    uint idx = indices[gid];
+    float val = values[gid];
+
+    // Atomic add - thread-safe, no corruption
+    atomic_fetch_add_explicit(
+        (device atomic_float*)&target[idx],
+        val,
+        memory_order_relaxed
+    );
+}
+
+// ============================================================================
+// ELEMENT-WISE OPERATIONS
+// ============================================================================
+
+kernel void add_arrays(
+    device const float* a [[buffer(0)]],
+    device const float* b [[buffer(1)]],
+    device float* result [[buffer(2)]],
+    constant uint& count [[buffer(3)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    result[gid] = a[gid] + b[gid];
+}
+
+kernel void add_scalar(
+    device const float* a [[buffer(0)]],
+    constant float& scalar [[buffer(1)]],
+    device float* result [[buffer(2)]],
+    constant uint& count [[buffer(3)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    result[gid] = a[gid] + scalar;
+}
+
+kernel void multiply_arrays(
+    device const float* a [[buffer(0)]],
+    device const float* b [[buffer(1)]],
+    device float* result [[buffer(2)]],
+    constant uint& count [[buffer(3)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    result[gid] = a[gid] * b[gid];
+}
+
+kernel void multiply_scalar(
+    device const float* a [[buffer(0)]],
+    constant float& scalar [[buffer(1)]],
+    device float* result [[buffer(2)]],
+    constant uint& count [[buffer(3)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    result[gid] = a[gid] * scalar;
+}
+
+kernel void square(
+    device const float* a [[buffer(0)]],
+    device float* result [[buffer(1)]],
+    constant uint& count [[buffer(2)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    float val = a[gid];
+    result[gid] = val * val;
+}
+
+// ============================================================================
+// CONDITIONAL OPERATIONS
+// ============================================================================
+
+kernel void greater_equal(
+    device const float* a [[buffer(0)]],
+    constant float& threshold [[buffer(1)]],
+    device uint* result [[buffer(2)]],  // Boolean as uint (0 or 1)
+    constant uint& count [[buffer(3)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    result[gid] = (a[gid] >= threshold) ? 1 : 0;
+}
+
+kernel void where_select(
+    device const uint* condition [[buffer(0)]],   // Boolean mask
+    device const float* if_true [[buffer(1)]],
+    device const float* if_false [[buffer(2)]],
+    device float* result [[buffer(3)]],
+    constant uint& count [[buffer(4)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    result[gid] = condition[gid] ? if_true[gid] : if_false[gid];
+}
+
+kernel void where_scalar(
+    device const uint* condition [[buffer(0)]],   // Boolean mask
+    constant float& if_true [[buffer(1)]],
+    constant float& if_false [[buffer(2)]],
+    device float* result [[buffer(3)]],
+    constant uint& count [[buffer(4)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    result[gid] = condition[gid] ? if_true : if_false;
+}
+
+// ============================================================================
+// REDUCTION OPERATIONS
+// ============================================================================
+
+kernel void sum_reduce(
+    device const float* input [[buffer(0)]],
+    device float* output [[buffer(1)]],
+    constant uint& count [[buffer(2)]],
+    threadgroup float* shared [[threadgroup(0)]],
+    uint gid [[thread_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint group_size [[threads_per_threadgroup]]
+) {
+    // Load into shared memory
+    shared[lid] = (gid < count) ? input[gid] : 0.0f;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Parallel reduction in shared memory
+    for (uint stride = group_size / 2; stride > 0; stride /= 2) {
+        if (lid < stride) {
+            shared[lid] += shared[lid + stride];
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    // Write result
+    if (lid == 0) {
+        atomic_fetch_add_explicit(
+            (device atomic_float*)output,
+            shared[0],
+            memory_order_relaxed
+        );
+    }
+}
+
+// ============================================================================
+// GATHER OPERATION (for sparse indexing)
+// ============================================================================
+
+kernel void gather(
+    device const float* source [[buffer(0)]],
+    device const uint* indices [[buffer(1)]],
+    device float* result [[buffer(2)]],
+    constant uint& count [[buffer(3)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    result[gid] = source[indices[gid]];
+}
+
+kernel void gather_bool(
+    device const uint* source [[buffer(0)]],    // Boolean array as uint
+    device const uint* indices [[buffer(1)]],
+    device uint* result [[buffer(2)]],
+    constant uint& count [[buffer(3)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    result[gid] = source[indices[gid]];
+}
+
+// ============================================================================
+// FILL OPERATIONS
+// ============================================================================
+
+kernel void fill_float(
+    device float* arr [[buffer(0)]],
+    constant float& value [[buffer(1)]],
+    constant uint& count [[buffer(2)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    arr[gid] = value;
+}
+
+kernel void fill_zeros(
+    device float* arr [[buffer(0)]],
+    constant uint& count [[buffer(1)]],
+    uint gid [[thread_position_in_grid]]
+) {
+    if (gid >= count) return;
+    arr[gid] = 0.0f;
+}

package/package.json

@@ -0,0 +1,38 @@
+{
+  "name": "@rlabs-inc/sparse",
+  "version": "0.1.0",
+  "description": "Sparsity-first GPU compute for biological neural simulation",
+  "type": "module",
+  "main": "src/index.ts",
+  "types": "src/index.ts",
+  "files": [
+    "src",
+    "native",
+    "libsparse.dylib",
+    "CLAUDE.md"
+  ],
+  "scripts": {
+    "test": "bun test",
+    "build": "bun build src/index.ts --outdir dist --target node",
+    "build:native": "cd native && make"
+  },
+  "os": ["darwin"],
+  "cpu": ["arm64", "x64"],
+  "keywords": [
+    "gpu",
+    "metal",
+    "sparse",
+    "neural",
+    "simulation",
+    "scatter-add",
+    "apple-silicon"
+  ],
+  "author": "RLabs Inc",
+  "license": "MIT",
+  "dependencies": {
+    "@rlabs-inc/signals": "^1.0.0"
+  },
+  "devDependencies": {
+    "@types/bun": "latest"
+  }
+}

package/src/ffi.ts

@@ -0,0 +1,156 @@
+// ============================================================================
+// SPARSE - FFI Bindings
+// TypeScript bindings to the native Metal library via Bun FFI
+// ============================================================================
+
+import { dlopen, FFIType, ptr, suffix, type Pointer } from 'bun:ffi'
+import { join, dirname } from 'path'
+import { fileURLToPath } from 'url'
+
+// Get the directory of this file to find the native library
+const __dirname = dirname(fileURLToPath(import.meta.url))
+const libPath = join(__dirname, '..', `libsparse.${suffix}`)
+
+// Data types
+export const DataType = {
+  FLOAT32: 0,
+  UINT32: 1,
+} as const
+
+export type DataType = typeof DataType[keyof typeof DataType]
+
+// Load the native library
+const lib = dlopen(libPath, {
+  // Context management
+  sparse_init: {
+    returns: FFIType.ptr,
+    args: [],
+  },
+  sparse_cleanup: {
+    returns: FFIType.void,
+    args: [FFIType.ptr],
+  },
+  sparse_sync: {
+    returns: FFIType.void,
+    args: [FFIType.ptr],
+  },
+  sparse_device_name: {
+    returns: FFIType.cstring,
+    args: [FFIType.ptr],
+  },
+  sparse_device_memory: {
+    returns: FFIType.u64,
+    args: [FFIType.ptr],
+  },
+
+  // Buffer management
+  sparse_zeros: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.u32, FFIType.i32],
+  },
+  sparse_full: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.u32, FFIType.f32],
+  },
+  sparse_from_float: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.u32],
+  },
+  sparse_from_uint: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.u32],
+  },
+  sparse_to_float: {
+    returns: FFIType.void,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.u32],
+  },
+  sparse_to_uint: {
+    returns: FFIType.void,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.u32],
+  },
+  sparse_buffer_count: {
+    returns: FFIType.u32,
+    args: [FFIType.ptr],
+  },
+  sparse_buffer_dtype: {
+    returns: FFIType.i32,
+    args: [FFIType.ptr],
+  },
+  sparse_buffer_free: {
+    returns: FFIType.void,
+    args: [FFIType.ptr],
+  },
+
+  // Core operations
+  sparse_scatter_add: {
+    returns: FFIType.void,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.ptr, FFIType.ptr, FFIType.u32],
+  },
+  sparse_gather: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.ptr, FFIType.u32],
+  },
+  sparse_gather_bool: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.ptr, FFIType.u32],
+  },
+
+  // Element-wise operations
+  sparse_add: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.ptr],
+  },
+  sparse_add_scalar: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.f32],
+  },
+  sparse_multiply: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.ptr],
+  },
+  sparse_multiply_scalar: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.f32],
+  },
+  sparse_square: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr],
+  },
+
+  // Conditional operations
+  sparse_greater_equal: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.f32],
+  },
+  sparse_where: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.ptr, FFIType.ptr],
+  },
+  sparse_where_scalar: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.ptr, FFIType.f32, FFIType.f32],
+  },
+
+  // Reduction operations
+  sparse_sum: {
+    returns: FFIType.f32,
+    args: [FFIType.ptr, FFIType.ptr],
+  },
+  sparse_sum_bool: {
+    returns: FFIType.u32,
+    args: [FFIType.ptr, FFIType.ptr],
+  },
+
+  // Random operations
+  sparse_random_uniform: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.u32, FFIType.f32, FFIType.f32],
+  },
+  sparse_random_normal: {
+    returns: FFIType.ptr,
+    args: [FFIType.ptr, FFIType.u32, FFIType.f32, FFIType.f32],
+  },
+})
+
+export const symbols = lib.symbols
+export type Symbols = typeof symbols