@rlabs-inc/sparse 0.1.0

package/src/test-scatter-loop.ts

@@ -0,0 +1,226 @@
+// ============================================================================
+// TEST: Scatter-Add in Loop
+// This is the EXACT pattern that crashed node-mlx.
+// If this works, we've solved the problem.
+// ============================================================================
+
+import * as gpu from './gpu'
+
+console.log('=== SPARSE: Testing Scatter-Add in Loop ===\n')
+
+// Get device info
+const info = gpu.deviceInfo()
+console.log('Device:', info.name)
+console.log('Memory:', (Number(info.memory) / 1024 / 1024 / 1024).toFixed(1), 'GB')
+console.log()
+
+// Test 1: Basic scatter-add
+console.log('Test 1: Basic scatter-add')
+{
+  const arr = gpu.zeros(10)
+  const idxs = gpu.indices([0, 2, 5])
+  const vals = gpu.array([1, 2, 3])
+
+  gpu.scatterAdd(arr, idxs, vals)
+
+  console.log('  Expected: [1, 0, 2, 0, 0, 3, 0, 0, 0, 0]')
+  console.log('  Got:     ', arr.toArray())
+  console.log()
+
+  arr.free()
+  idxs.free()
+  vals.free()
+}
+
+// Test 2: Scatter-add with duplicate indices (critical!)
+console.log('Test 2: Scatter-add with DUPLICATE indices')
+{
+  const arr = gpu.zeros(5)
+  const idxs = gpu.indices([0, 0, 0, 2, 2])
+  const vals = gpu.array([1, 2, 3, 0.5, 0.5])
+
+  gpu.scatterAdd(arr, idxs, vals)
+
+  console.log('  Expected: [6, 0, 1, 0, 0] (1+2+3=6, 0.5+0.5=1)')
+  console.log('  Got:     ', arr.toArray())
+  console.log()
+
+  arr.free()
+  idxs.free()
+  vals.free()
+}
+
+// Test 3: THE CRITICAL TEST - Scatter-add in a loop with REUSED indices
+console.log('Test 3: Scatter-add in LOOP with REUSED indices')
+console.log('  (This is what crashed node-mlx!)')
+{
+  // Create indices ONCE, reuse in every iteration
+  const idxs = gpu.indices([0, 1, 2])
+  const vals = gpu.array([1, 2, 3])
+
+  for (let i = 0; i < 10; i++) {
+    const arr = gpu.zeros(10)
+    gpu.scatterAdd(arr, idxs, vals)
+    console.log(`  Iteration ${i}: ${arr.toArray().slice(0, 5).join(', ')}...`)
+    arr.free()
+  }
+
+  console.log('  ALL 10 ITERATIONS COMPLETED!')
+  console.log()
+
+  idxs.free()
+  vals.free()
+}
+
+// Test 4: Large-scale performance
+console.log('Test 4: Large-scale scatter-add (1M neurons, 100K synapses)')
+{
+  const numNeurons = 1_000_000
+  const numSynapses = 100_000
+
+  // Create indices once
+  const idxData = new Uint32Array(numSynapses)
+  for (let i = 0; i < numSynapses; i++) {
+    idxData[i] = Math.floor(Math.random() * numNeurons)
+  }
+  const idxs = gpu.indices(idxData)
+  const vals = gpu.random.uniform(numSynapses, 0, 1)
+
+  console.log('  Running 10 iterations...')
+  const start = performance.now()
+
+  for (let i = 0; i < 10; i++) {
+    const arr = gpu.zeros(numNeurons)
+    gpu.scatterAdd(arr, idxs, vals)
+    arr.free()
+  }
+
+  const elapsed = performance.now() - start
+  console.log(`  Time: ${elapsed.toFixed(2)}ms (${(elapsed / 10).toFixed(2)}ms per iteration)`)
+  console.log()
+
+  idxs.free()
+  vals.free()
+}
+
+// Test 5: Full Izhikevich pattern
+console.log('Test 5: Full Izhikevich neuron simulation pattern')
+{
+  const numNeurons = 500
+  const numSynapses = 5000
+  const numSteps = 20
+
+  // Neuron state
+  let v = gpu.full(numNeurons, -65)
+  let u = gpu.multiply(v, 0.2)
+
+  // Static connectivity - created ONCE, reused every timestep
+  const preIdxData = new Uint32Array(numSynapses)
+  const postIdxData = new Uint32Array(numSynapses)
+  for (let i = 0; i < numSynapses; i++) {
+    preIdxData[i] = Math.floor(Math.random() * numNeurons)
+    postIdxData[i] = Math.floor(Math.random() * numNeurons)
+  }
+  const preIndices = gpu.indices(preIdxData)
+  const postIndices = gpu.indices(postIdxData)
+  const weights = gpu.random.uniform(numSynapses, 0, 0.5)
+
+  // Input
+  const inputIdxData = new Uint32Array(50)
+  for (let i = 0; i < 50; i++) inputIdxData[i] = i
+  const inputIndices = gpu.indices(inputIdxData)
+  const inputValues = gpu.full(50, 15)
+
+  console.log('  Simulating', numSteps, 'timesteps...')
+  const start = performance.now()
+
+  for (let t = 0; t < numSteps; t++) {
+    // 1. Detect firing
+    const fired = gpu.greaterEqual(v, 30)
+
+    // 2. Sparse transmission
+    const preFired = gpu.gatherBool(fired, preIndices)
+    const zerosArr = gpu.zeros(numSynapses)
+    const contribution = gpu.where(preFired, weights, zerosArr)
+
+    // 3. Scatter-add currents
+    let current = gpu.zeros(numNeurons)
+    gpu.scatterAdd(current, postIndices, contribution)
+    gpu.scatterAdd(current, inputIndices, inputValues)
+
+    // 4. Izhikevich update
+    const vSq = gpu.square(v)
+    const vSq004 = gpu.multiply(vSq, 0.04)
+    const v5 = gpu.multiply(v, 5)
+    const sum1 = gpu.add(vSq004, v5)
+    const sum2 = gpu.add(sum1, 140)
+    const negU = gpu.multiply(u, -1)
+    const sum3 = gpu.add(sum2, negU)
+    const dv = gpu.add(sum3, current)
+
+    const bv = gpu.multiply(v, 0.2)
+    const diff = gpu.subtract(bv, u)
+    const du = gpu.multiply(diff, 0.02)
+
+    const newV = gpu.add(v, dv)
+    const newU = gpu.add(u, du)
+
+    // 5. Reset spiked neurons
+    const vReset = gpu.where(fired, -65, newV)
+    const uBump = gpu.add(u, 8)
+    const uReset = gpu.where(fired, uBump, newU)
+
+    // Count spikes
+    const spikeCount = gpu.sum(fired)
+
+    if (t % 5 === 0) {
+      console.log(`    Step ${t}: ${spikeCount} spikes`)
+    }
+
+    // Cleanup intermediates
+    fired.free()
+    preFired.free()
+    zerosArr.free()
+    contribution.free()
+    current.free()
+    vSq.free()
+    vSq004.free()
+    v5.free()
+    sum1.free()
+    sum2.free()
+    negU.free()
+    sum3.free()
+    dv.free()
+    bv.free()
+    diff.free()
+    du.free()
+    newV.free()
+    newU.free()
+    uBump.free()
+
+    // Swap state
+    v.free()
+    u.free()
+    v = vReset
+    u = uReset
+  }
+
+  const elapsed = performance.now() - start
+  console.log(`  Time: ${elapsed.toFixed(2)}ms (${(elapsed / numSteps).toFixed(2)}ms per step)`)
+  console.log('  SIMULATION COMPLETE!')
+  console.log()
+
+  // Cleanup
+  v.free()
+  u.free()
+  preIndices.free()
+  postIndices.free()
+  weights.free()
+  inputIndices.free()
+  inputValues.free()
+}
+
+// Cleanup
+gpu.cleanup()
+
+console.log('=== ALL TESTS PASSED ===')

package/src/test-stress.ts

@@ -0,0 +1,160 @@
+// ============================================================================
+// STRESS TEST: Let's see what the M1 Max 64GB can handle!
+// ============================================================================
+
+import * as gpu from './gpu'
+
+console.log('=== STRESS TEST: Frying the M1 Max ===\n')
+
+const info = gpu.deviceInfo()
+console.log('Device:', info.name)
+console.log('Memory:', (Number(info.memory) / 1024 / 1024 / 1024).toFixed(1), 'GB')
+console.log()
+
+// Test configurations - let's scale up!
+const tests = [
+  { neurons: 1_000, synapses: 10_000, steps: 100 },
+  { neurons: 10_000, synapses: 100_000, steps: 100 },
+  { neurons: 100_000, synapses: 1_000_000, steps: 100 },
+  { neurons: 1_000_000, synapses: 10_000_000, steps: 50 },
+  { neurons: 10_000_000, synapses: 100_000_000, steps: 10 },  // 10M neurons, 100M synapses!
+]
+
+for (const { neurons, synapses, steps } of tests) {
+  console.log(`\n${'='.repeat(60)}`)
+  console.log(`Testing: ${neurons.toLocaleString()} neurons, ${synapses.toLocaleString()} synapses`)
+  console.log(`${'='.repeat(60)}`)
+
+  try {
+    // Create neuron state
+    console.log('  Allocating neuron state...')
+    const allocStart = performance.now()
+
+    let v = gpu.full(neurons, -65)
+    let u = gpu.multiply(v, 0.2)
+
+    // Create connectivity
+    console.log('  Creating connectivity...')
+    const preIdxData = new Uint32Array(synapses)
+    const postIdxData = new Uint32Array(synapses)
+    for (let i = 0; i < synapses; i++) {
+      preIdxData[i] = Math.floor(Math.random() * neurons)
+      postIdxData[i] = Math.floor(Math.random() * neurons)
+    }
+
+    const preIndices = gpu.indices(preIdxData)
+    const postIndices = gpu.indices(postIdxData)
+    const weights = gpu.random.uniform(synapses, 0, 0.5)
+
+    // Input to 10% of neurons
+    const numInput = Math.floor(neurons * 0.1)
+    const inputIdxData = new Uint32Array(numInput)
+    for (let i = 0; i < numInput; i++) inputIdxData[i] = i
+    const inputIndices = gpu.indices(inputIdxData)
+    const inputValues = gpu.full(numInput, 15)
+
+    const allocTime = performance.now() - allocStart
+    console.log(`  Allocation time: ${allocTime.toFixed(0)}ms`)
+
+    // Simulate
+    console.log(`  Running ${steps} timesteps...`)
+    const simStart = performance.now()
+    let totalSpikes = 0
+
+    for (let t = 0; t < steps; t++) {
+      // 1. Detect firing
+      const fired = gpu.greaterEqual(v, 30)
+
+      // 2. Sparse transmission
+      const preFired = gpu.gatherBool(fired, preIndices)
+      const zerosArr = gpu.zeros(synapses)
+      const contribution = gpu.where(preFired, weights, zerosArr)
+
+      // 3. Scatter-add currents
+      let current = gpu.zeros(neurons)
+      gpu.scatterAdd(current, postIndices, contribution)
+      gpu.scatterAdd(current, inputIndices, inputValues)
+
+      // 4. Izhikevich update (simplified for speed)
+      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 bv = gpu.multiply(v, 0.2)
+      const diff = gpu.subtract(bv, u)
+      const du = gpu.multiply(diff, 0.02)
+
+      const newV = gpu.add(v, dv)
+      const newU = gpu.add(u, du)
+
+      // 5. Reset
+      const vReset = gpu.where(fired, -65, newV)
+      const uBump = gpu.add(newU, 8)
+      const uReset = gpu.where(fired, uBump, newU)
+
+      // Count spikes
+      const spikeCount = gpu.sum(fired)
+      totalSpikes += spikeCount
+
+      // Cleanup intermediates
+      fired.free()
+      preFired.free()
+      zerosArr.free()
+      contribution.free()
+      current.free()
+      vSq.free()
+      term1.free()
+      term2.free()
+      sum1.free()
+      sum2.free()
+      sum3.free()
+      dv.free()
+      bv.free()
+      diff.free()
+      du.free()
+      newV.free()
+      newU.free()
+      uBump.free()
+
+      v.free()
+      u.free()
+      v = vReset
+      u = uReset
+
+      if (t % Math.max(1, Math.floor(steps / 5)) === 0) {
+        process.stdout.write(`    Step ${t}/${steps} (${spikeCount.toLocaleString()} spikes)\n`)
+      }
+    }
+
+    const simTime = performance.now() - simStart
+    const msPerStep = simTime / steps
+
+    console.log()
+    console.log(`  ✅ COMPLETED!`)
+    console.log(`  Total time: ${simTime.toFixed(0)}ms`)
+    console.log(`  Per step: ${msPerStep.toFixed(2)}ms`)
+    console.log(`  Total spikes: ${totalSpikes.toLocaleString()}`)
+    console.log(`  Neurons/ms: ${(neurons / msPerStep).toLocaleString()}`)
+    console.log(`  Synapses/ms: ${(synapses / msPerStep).toLocaleString()}`)
+
+    // Cleanup
+    v.free()
+    u.free()
+    preIndices.free()
+    postIndices.free()
+    weights.free()
+    inputIndices.free()
+    inputValues.free()
+
+  } catch (error) {
+    console.log(`  ❌ FAILED: ${error}`)
+    break
+  }
+}
+
+gpu.cleanup()
+console.log('\n=== STRESS TEST COMPLETE ===')

package/src/test-webgpu.ts

@@ -0,0 +1,31 @@
+/**
+ * Quick test to check WebGPU availability in Bun
+ */
+
+// Check if WebGPU is available
+console.log('Checking WebGPU support in Bun...')
+console.log()
+
+// @ts-ignore - checking global availability
+const hasGPU = typeof navigator !== 'undefined' && 'gpu' in navigator
+
+console.log('navigator.gpu available:', hasGPU)
+
+if (hasGPU) {
+  // @ts-ignore
+  const adapter = await navigator.gpu.requestAdapter()
+  console.log('Adapter:', adapter)
+
+  if (adapter) {
+    const device = await adapter.requestDevice()
+    console.log('Device:', device)
+    console.log('Device limits:', device.limits)
+  }
+} else {
+  console.log()
+  console.log('WebGPU not available in Bun runtime.')
+  console.log('Options:')
+  console.log('  1. Use wgpu-native bindings (Rust → FFI)')
+  console.log('  2. Use Metal directly via FFI')
+  console.log('  3. Use Dawn (Google\'s WebGPU implementation) via FFI')
+}