@rlabs-inc/sparse 0.1.0

package/src/gpu.ts

@@ -0,0 +1,382 @@
+// ============================================================================
+// SPARSE - GPU Array API
+// High-level TypeScript API for GPU compute
+// Sparsity-first. Indices are READ-ONLY, never corrupted.
+// ============================================================================
+
+import { ptr, toArrayBuffer, type Pointer } from 'bun:ffi'
+import { symbols, DataType } from './ffi'
+
+// ============================================================================
+// TYPES
+// ============================================================================
+
+export type BufferPtr = Pointer
+export type ContextPtr = Pointer
+
+export interface GPUDeviceInfo {
+  name: string
+  memory: bigint
+}
+
+// ============================================================================
+// GPU CONTEXT
+// ============================================================================
+
+let _context: ContextPtr | null = null
+
+/**
+ * Initialize the GPU context (called automatically on first use)
+ */
+export function init(): ContextPtr {
+  if (_context) return _context
+  _context = symbols.sparse_init()
+  if (!_context) {
+    throw new Error('Failed to initialize Metal GPU context')
+  }
+  return _context
+}
+
+/**
+ * Get device info
+ */
+export function deviceInfo(): GPUDeviceInfo {
+  const ctx = init()
+  return {
+    name: symbols.sparse_device_name(ctx) as unknown as string,
+    memory: symbols.sparse_device_memory(ctx),
+  }
+}
+
+/**
+ * Synchronize - wait for all GPU operations to complete
+ */
+export function sync(): void {
+  if (_context) {
+    symbols.sparse_sync(_context)
+  }
+}
+
+/**
+ * Cleanup GPU resources (call when done)
+ */
+export function cleanup(): void {
+  if (_context) {
+    symbols.sparse_cleanup(_context)
+    _context = null
+  }
+}
+
+// ============================================================================
+// GPU ARRAY CLASS
+// ============================================================================
+
+export class GPUArray {
+  private _ptr: BufferPtr
+  private _count: number
+  private _dtype: number
+  private _freed = false
+
+  constructor(ptr: BufferPtr, count: number, dtype: number) {
+    this._ptr = ptr
+    this._count = count
+    this._dtype = dtype
+  }
+
+  get ptr(): BufferPtr {
+    if (this._freed) throw new Error('GPUArray has been freed')
+    return this._ptr
+  }
+
+  get count(): number {
+    return this._count
+  }
+
+  get dtype(): number {
+    return this._dtype
+  }
+
+  get isFloat(): boolean {
+    return this._dtype === DataType.FLOAT32
+  }
+
+  get isBool(): boolean {
+    return this._dtype === DataType.UINT32
+  }
+
+  /**
+   * Read data back to CPU as Float32Array
+   */
+  toFloat32Array(): Float32Array {
+    const result = new Float32Array(this._count)
+    symbols.sparse_to_float(this._ptr, ptr(result), this._count)
+    return result
+  }
+
+  /**
+   * Read data back to CPU as Uint32Array
+   */
+  toUint32Array(): Uint32Array {
+    const result = new Uint32Array(this._count)
+    symbols.sparse_to_uint(this._ptr, ptr(result), this._count)
+    return result
+  }
+
+  /**
+   * Read data back as regular array
+   */
+  toArray(): number[] {
+    if (this._dtype === DataType.FLOAT32) {
+      return Array.from(this.toFloat32Array())
+    } else {
+      return Array.from(this.toUint32Array())
+    }
+  }
+
+  /**
+   * Free GPU memory (called automatically by GC, but can be called manually)
+   */
+  free(): void {
+    if (!this._freed) {
+      symbols.sparse_buffer_free(this._ptr)
+      this._freed = true
+    }
+  }
+}
+
+// ============================================================================
+// ARRAY CREATION
+// ============================================================================
+
+/**
+ * Create a GPU array filled with zeros
+ */
+export function zeros(count: number, dtype: number = DataType.FLOAT32): GPUArray {
+  const ctx = init()
+  const bufPtr = symbols.sparse_zeros(ctx, count, dtype)
+  if (!bufPtr) throw new Error('Failed to create zeros buffer')
+  return new GPUArray(bufPtr, count, dtype)
+}
+
+/**
+ * Create a GPU array filled with a value
+ */
+export function full(count: number, value: number): GPUArray {
+  const ctx = init()
+  const bufPtr = symbols.sparse_full(ctx, count, value)
+  if (!bufPtr) throw new Error('Failed to create full buffer')
+  return new GPUArray(bufPtr, count, DataType.FLOAT32)
+}
+
+/**
+ * Create a GPU array from Float32Array or number[]
+ */
+export function array(data: Float32Array | number[]): GPUArray {
+  const ctx = init()
+  const float32 = data instanceof Float32Array ? data : new Float32Array(data)
+  const bufPtr = symbols.sparse_from_float(ctx, ptr(float32), float32.length)
+  if (!bufPtr) throw new Error('Failed to create array buffer')
+  return new GPUArray(bufPtr, float32.length, DataType.FLOAT32)
+}
+
+/**
+ * Create a GPU array of uint32 indices
+ */
+export function indices(data: Uint32Array | number[]): GPUArray {
+  const ctx = init()
+  const uint32 = data instanceof Uint32Array ? data : new Uint32Array(data)
+  const bufPtr = symbols.sparse_from_uint(ctx, ptr(uint32), uint32.length)
+  if (!bufPtr) throw new Error('Failed to create indices buffer')
+  return new GPUArray(bufPtr, uint32.length, DataType.UINT32)
+}
+
+// ============================================================================
+// CORE OPERATIONS
+// ============================================================================
+
+/**
+ * Scatter-add: target[indices[i]] += values[i]
+ * CRITICAL: This is atomic and thread-safe. Indices are READ-ONLY, never corrupted.
+ * This is the core operation for sparse synapse transmission.
+ */
+export function scatterAdd(
+  target: GPUArray,
+  idxs: GPUArray,
+  values: GPUArray
+): void {
+  const ctx = init()
+  symbols.sparse_scatter_add(ctx, target.ptr, idxs.ptr, values.ptr, idxs.count)
+}
+
+/**
+ * Gather: result[i] = source[indices[i]]
+ */
+export function gather(source: GPUArray, idxs: GPUArray): GPUArray {
+  const ctx = init()
+  const bufPtr = symbols.sparse_gather(ctx, source.ptr, idxs.ptr, idxs.count)
+  if (!bufPtr) throw new Error('Failed to gather')
+  return new GPUArray(bufPtr, idxs.count, DataType.FLOAT32)
+}
+
+/**
+ * Gather boolean: result[i] = source[indices[i]] (for boolean arrays)
+ */
+export function gatherBool(source: GPUArray, idxs: GPUArray): GPUArray {
+  const ctx = init()
+  const bufPtr = symbols.sparse_gather_bool(ctx, source.ptr, idxs.ptr, idxs.count)
+  if (!bufPtr) throw new Error('Failed to gather bool')
+  return new GPUArray(bufPtr, idxs.count, DataType.UINT32)
+}
+
+// ============================================================================
+// ELEMENT-WISE OPERATIONS
+// ============================================================================
+
+/**
+ * Add two arrays or add scalar
+ */
+export function add(a: GPUArray, b: GPUArray | number): GPUArray {
+  const ctx = init()
+  let bufPtr: BufferPtr
+
+  if (typeof b === 'number') {
+    bufPtr = symbols.sparse_add_scalar(ctx, a.ptr, b)
+  } else {
+    bufPtr = symbols.sparse_add(ctx, a.ptr, b.ptr)
+  }
+
+  if (!bufPtr) throw new Error('Failed to add')
+  return new GPUArray(bufPtr, a.count, DataType.FLOAT32)
+}
+
+/**
+ * Subtract: a - b
+ */
+export function subtract(a: GPUArray, b: GPUArray | number): GPUArray {
+  if (typeof b === 'number') {
+    return add(a, -b)
+  } else {
+    // a - b = a + (-1 * b)
+    const negB = multiply(b, -1)
+    const result = add(a, negB)
+    negB.free()
+    return result
+  }
+}
+
+/**
+ * Multiply two arrays or multiply by scalar
+ */
+export function multiply(a: GPUArray, b: GPUArray | number): GPUArray {
+  const ctx = init()
+  let bufPtr: BufferPtr
+
+  if (typeof b === 'number') {
+    bufPtr = symbols.sparse_multiply_scalar(ctx, a.ptr, b)
+  } else {
+    bufPtr = symbols.sparse_multiply(ctx, a.ptr, b.ptr)
+  }
+
+  if (!bufPtr) throw new Error('Failed to multiply')
+  return new GPUArray(bufPtr, a.count, DataType.FLOAT32)
+}
+
+/**
+ * Square each element
+ */
+export function square(a: GPUArray): GPUArray {
+  const ctx = init()
+  const bufPtr = symbols.sparse_square(ctx, a.ptr)
+  if (!bufPtr) throw new Error('Failed to square')
+  return new GPUArray(bufPtr, a.count, DataType.FLOAT32)
+}
+
+// ============================================================================
+// CONDITIONAL OPERATIONS
+// ============================================================================
+
+/**
+ * Compare: result[i] = (a[i] >= threshold) ? 1 : 0
+ */
+export function greaterEqual(a: GPUArray, threshold: number): GPUArray {
+  const ctx = init()
+  const bufPtr = symbols.sparse_greater_equal(ctx, a.ptr, threshold)
+  if (!bufPtr) throw new Error('Failed to compare')
+  return new GPUArray(bufPtr, a.count, DataType.UINT32)
+}
+
+/**
+ * Where: result[i] = condition[i] ? ifTrue[i] : ifFalse[i]
+ * Can also use scalar values for ifTrue/ifFalse
+ */
+export function where(
+  condition: GPUArray,
+  ifTrue: GPUArray | number,
+  ifFalse: GPUArray | number
+): GPUArray {
+  const ctx = init()
+  let bufPtr: BufferPtr
+
+  if (typeof ifTrue === 'number' && typeof ifFalse === 'number') {
+    bufPtr = symbols.sparse_where_scalar(ctx, condition.ptr, ifTrue, ifFalse)
+  } else if (ifTrue instanceof GPUArray && ifFalse instanceof GPUArray) {
+    bufPtr = symbols.sparse_where(ctx, condition.ptr, ifTrue.ptr, ifFalse.ptr)
+  } else if (typeof ifTrue === 'number' && ifFalse instanceof GPUArray) {
+    // Mixed: scalar ifTrue, array ifFalse -> create temp array
+    const ifTrueArr = full(condition.count, ifTrue)
+    bufPtr = symbols.sparse_where(ctx, condition.ptr, ifTrueArr.ptr, ifFalse.ptr)
+    ifTrueArr.free()
+  } else if (ifTrue instanceof GPUArray && typeof ifFalse === 'number') {
+    // Mixed: array ifTrue, scalar ifFalse -> create temp array
+    const ifFalseArr = full(condition.count, ifFalse)
+    bufPtr = symbols.sparse_where(ctx, condition.ptr, ifTrue.ptr, ifFalseArr.ptr)
+    ifFalseArr.free()
+  } else {
+    throw new Error('where() invalid argument types')
+  }
+
+  if (!bufPtr) throw new Error('Failed to where')
+  return new GPUArray(bufPtr, condition.count, DataType.FLOAT32)
+}
+
+// ============================================================================
+// REDUCTION OPERATIONS
+// ============================================================================
+
+/**
+ * Sum all elements
+ */
+export function sum(a: GPUArray): number {
+  const ctx = init()
+  if (a.dtype === DataType.FLOAT32) {
+    return symbols.sparse_sum(ctx, a.ptr)
+  } else {
+    return symbols.sparse_sum_bool(ctx, a.ptr)
+  }
+}
+
+// ============================================================================
+// RANDOM OPERATIONS
+// ============================================================================
+
+export const random = {
+  /**
+   * Random uniform values in [low, high)
+   */
+  uniform(count: number, low = 0, high = 1): GPUArray {
+    const ctx = init()
+    const bufPtr = symbols.sparse_random_uniform(ctx, count, low, high)
+    if (!bufPtr) throw new Error('Failed to create random uniform')
+    return new GPUArray(bufPtr, count, DataType.FLOAT32)
+  },
+
+  /**
+   * Random normal values with mean and std
+   */
+  normal(count: number, mean = 0, std = 1): GPUArray {
+    const ctx = init()
+    const bufPtr = symbols.sparse_random_normal(ctx, count, mean, std)
+    if (!bufPtr) throw new Error('Failed to create random normal')
+    return new GPUArray(bufPtr, count, DataType.FLOAT32)
+  },
+}

package/src/index.ts

@@ -0,0 +1,7 @@
+// ============================================================================
+// SPARSE - Sparsity-First GPU Compute
+// Minimal. Elegant. For biological neural simulation, not transformers.
+// ============================================================================
+
+export * from './gpu'
+export { DataType } from './ffi'

package/src/test-debug-spikes.ts

@@ -0,0 +1,70 @@
+// Debug: Why no spikes?
+
+import * as gpu from './gpu'
+
+console.log('=== Debug: Izhikevich Dynamics ===\n')
+
+const numNeurons = 10
+
+// Initial state
+let v = gpu.full(numNeurons, -65)
+let u = gpu.multiply(v, 0.2)
+
+// Constant input to all neurons
+const input = gpu.full(numNeurons, 15)
+
+console.log('Initial:')
+console.log('  v:', v.toArray())
+console.log('  u:', u.toArray())
+console.log()
+
+for (let t = 0; t < 15; t++) {
+  // Izhikevich equations
+  // dv = 0.04*v^2 + 5*v + 140 - u + I
+  const vSq = gpu.square(v)
+  const term1 = gpu.multiply(vSq, 0.04)  // 0.04*v^2
+  const term2 = gpu.multiply(v, 5)        // 5*v
+  const sum1 = gpu.add(term1, term2)
+  const sum2 = gpu.add(sum1, 140)
+  const sum3 = gpu.subtract(sum2, u)
+  const dv = gpu.add(sum3, input)
+
+  // du = a * (b*v - u) where a=0.02, b=0.2
+  const bv = gpu.multiply(v, 0.2)
+  const diff = gpu.subtract(bv, u)
+  const du = gpu.multiply(diff, 0.02)
+
+  // Update
+  const newV = gpu.add(v, dv)
+  const newU = gpu.add(u, du)
+
+  // Check for spikes
+  const fired = gpu.greaterEqual(newV, 30)
+  const spikeCount = gpu.sum(fired)
+
+  console.log(`Step ${t}:`)
+  console.log('  v:', newV.toArray().map(x => x.toFixed(1)))
+  console.log('  dv:', dv.toArray().map(x => x.toFixed(1)))
+  console.log('  spikes:', spikeCount)
+
+  // Reset if spiked
+  const vReset = gpu.where(fired, -65, newV)
+  const uBump = gpu.add(newU, 8)
+  const uReset = gpu.where(fired, uBump, newU)
+
+  // Cleanup
+  vSq.free(); term1.free(); term2.free(); sum1.free(); sum2.free(); sum3.free(); dv.free()
+  bv.free(); diff.free(); du.free()
+  newV.free(); newU.free(); fired.free(); uBump.free()
+  v.free(); u.free()
+
+  v = vReset
+  u = uReset
+
+  if (spikeCount > 0) {
+    console.log('\n  >>> SPIKE DETECTED! <<<\n')
+  }
+  console.log()
+}
+
+gpu.cleanup()

package/src/test-limits.ts

@@ -0,0 +1,140 @@
+// ============================================================================
+// LIMIT TEST: Find the M1 Max 64GB ceiling
+// ============================================================================
+
+import * as gpu from './gpu'
+
+console.log('=== LIMIT TEST: Finding the Ceiling ===\n')
+
+const info = gpu.deviceInfo()
+console.log('Device:', info.name)
+console.log('Memory:', (Number(info.memory) / 1024 / 1024 / 1024).toFixed(1), 'GB')
+console.log()
+
+// Memory estimation:
+// - 1 neuron: 2 floats (v, u) = 8 bytes
+// - 1 synapse: 2 uint32 (pre, post) + 1 float (weight) = 12 bytes
+// - Plus intermediate buffers during computation (~5x overhead)
+
+// Let's scale up aggressively
+const tests = [
+  { neurons: 10_000_000, synapses: 100_000_000 },    // 10M neurons, 100M synapses (worked!)
+  { neurons: 20_000_000, synapses: 200_000_000 },    // 20M neurons, 200M synapses
+  { neurons: 50_000_000, synapses: 500_000_000 },    // 50M neurons, 500M synapses
+  { neurons: 100_000_000, synapses: 1_000_000_000 }, // 100M neurons, 1B synapses!
+  { neurons: 200_000_000, synapses: 2_000_000_000 }, // 200M neurons, 2B synapses!
+]
+
+function formatBytes(bytes: number): string {
+  const gb = bytes / 1024 / 1024 / 1024
+  return `${gb.toFixed(2)} GB`
+}
+
+function estimateMemory(neurons: number, synapses: number): number {
+  // Neuron state: v, u (2 floats per neuron)
+  const neuronBytes = neurons * 2 * 4
+
+  // Synapse data: preIdx, postIdx (2 uint32) + weight (1 float)
+  const synapseBytes = synapses * (2 * 4 + 4)
+
+  // Intermediate buffers during computation (rough estimate: ~3x neuron + 2x synapse)
+  const overhead = neuronBytes * 3 + synapseBytes * 2
+
+  return neuronBytes + synapseBytes + overhead
+}
+
+for (const { neurons, synapses } of tests) {
+  console.log(`\n${'='.repeat(70)}`)
+  console.log(`Testing: ${(neurons / 1_000_000).toFixed(0)}M neurons, ${(synapses / 1_000_000).toFixed(0)}M synapses`)
+  console.log(`Estimated memory: ${formatBytes(estimateMemory(neurons, synapses))}`)
+  console.log(`${'='.repeat(70)}`)
+
+  try {
+    // Create neuron state
+    console.log('  [1/4] Allocating neuron state...')
+    const t1 = performance.now()
+    let v = gpu.full(neurons, -65)
+    let u = gpu.multiply(v, 0.2)
+    console.log(`        Done in ${(performance.now() - t1).toFixed(0)}ms`)
+
+    // Create connectivity - this is the big one
+    console.log('  [2/4] Creating pre-synaptic indices...')
+    const t2 = performance.now()
+    const preIdxData = new Uint32Array(synapses)
+    for (let i = 0; i < synapses; i++) {
+      preIdxData[i] = Math.floor(Math.random() * neurons)
+    }
+    const preIndices = gpu.indices(preIdxData)
+    console.log(`        Done in ${(performance.now() - t2).toFixed(0)}ms`)
+
+    console.log('  [3/4] Creating post-synaptic indices...')
+    const t3 = performance.now()
+    const postIdxData = new Uint32Array(synapses)
+    for (let i = 0; i < synapses; i++) {
+      postIdxData[i] = Math.floor(Math.random() * neurons)
+    }
+    const postIndices = gpu.indices(postIdxData)
+    console.log(`        Done in ${(performance.now() - t3).toFixed(0)}ms`)
+
+    console.log('  [4/4] Creating synaptic weights...')
+    const t4 = performance.now()
+    const weights = gpu.random.uniform(synapses, 0, 0.5)
+    console.log(`        Done in ${(performance.now() - t4).toFixed(0)}ms`)
+
+    // Run ONE timestep to verify it works
+    console.log('  Running 1 timestep...')
+    const simStart = performance.now()
+
+    // Detect firing
+    const fired = gpu.greaterEqual(v, 30)
+    const spikeCount = gpu.sum(fired)
+
+    // Sparse transmission
+    const preFired = gpu.gatherBool(fired, preIndices)
+    const zerosArr = gpu.zeros(synapses)
+    const contribution = gpu.where(preFired, weights, zerosArr)
+
+    // Scatter-add
+    let current = gpu.zeros(neurons)
+    gpu.scatterAdd(current, postIndices, contribution)
+
+    // Izhikevich
+    const vSq = gpu.square(v)
+    const term1 = gpu.multiply(vSq, 0.04)
+    const term2 = gpu.multiply(v, 5)
+    const sum1 = gpu.add(term1, term2)
+    const sum2 = gpu.add(sum1, 140)
+    const sum3 = gpu.subtract(sum2, u)
+    const dv = gpu.add(sum3, current)
+    const newV = gpu.add(v, dv)
+
+    const simTime = performance.now() - simStart
+
+    console.log()
+    console.log(`  ✅ SUCCESS!`)
+    console.log(`  1 timestep: ${simTime.toFixed(0)}ms`)
+    console.log(`  Spikes: ${spikeCount.toLocaleString()}`)
+    console.log(`  Throughput: ${((neurons + synapses) / simTime / 1000).toFixed(1)}M ops/ms`)
+
+    // Cleanup
+    v.free(); u.free()
+    preIndices.free(); postIndices.free(); weights.free()
+    fired.free(); preFired.free(); zerosArr.free(); contribution.free()
+    current.free(); vSq.free(); term1.free(); term2.free()
+    sum1.free(); sum2.free(); sum3.free(); dv.free(); newV.free()
+
+    // Let JS GC handle the arrays
+
+  } catch (error: any) {
+    console.log()
+    console.log(`  ❌ HIT THE CEILING!`)
+    console.log(`  Error: ${error.message || error}`)
+    console.log()
+    console.log(`  Maximum sustainable scale:`)
+    console.log(`    Previous test passed - that's your limit!`)
+    break
+  }
+}
+
+gpu.cleanup()
+console.log('\n=== LIMIT TEST COMPLETE ===')