src/webgpu.ts

type GpuBuffer = {
  destroy(): void
  mapAsync(mode: number): Promise<void>
  getMappedRange(): ArrayBuffer
  unmap(): void
}

type GpuComputePass = {
  setPipeline(pipeline: unknown): void
  setBindGroup(index: number, group: unknown): void
  dispatchWorkgroups(count: number): void
  end(): void
}

type GpuDevice = {
  limits: { maxStorageBufferBindingSize: number; maxBufferSize: number }
  queue: { writeBuffer(buffer: GpuBuffer, offset: number, data: BufferSource): void; submit(commands: unknown[]): void }
  createBuffer(descriptor: { size: number; usage: number; mappedAtCreation?: boolean }): GpuBuffer
  createShaderModule(descriptor: { code: string }): unknown
  createComputePipeline(descriptor: { layout: string; compute: { module: unknown; entryPoint: string } }): unknown
  createBindGroup(descriptor: { layout: unknown; entries: Array<{ binding: number; resource: { buffer: GpuBuffer } }> }): unknown
  createCommandEncoder(): {
    beginComputePass(): GpuComputePass
    copyBufferToBuffer(source: GpuBuffer, sourceOffset: number, destination: GpuBuffer, destinationOffset: number, size: number): void
    finish(): unknown
  }
  destroy(): void
  lost: Promise<{ message: string }>
}

type GpuAdapter = { requestDevice(): Promise<GpuDevice> }
type GpuNavigator = { requestAdapter(): Promise<GpuAdapter | null> }

const MAP_READ = 0x0001
const STORAGE = 0x0080
const COPY_SRC = 0x0004
const COPY_DST = 0x0008

const VECTOR_ADD_SHADER = `
@group(0) @binding(0) var<storage, read> left: array<f32>;
@group(0) @binding(1) var<storage, read> right: array<f32>;
@group(0) @binding(2) var<storage, read_write> output: array<f32>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
  let index = id.x;
  if (index < arrayLength(&output)) {
    output[index] = left[index] + right[index];
  }
}
`

function vectorInputs(start: number, end: number) {
  const count = end - start
  const left = new Float32Array(count)
  const right = new Float32Array(count)
  for (let i = 0; i < count; i++) {
    const n = start + i
    left[i] = ((n * 17) % 100_003) / 100_003
    right[i] = ((n * 29 + 7) % 100_019) / 100_019
  }
  return { left, right }
}

export function runVectorAddCpu(start: number, end: number) {
  const { left, right } = vectorInputs(start, end)
  const output = new Float32Array(left.length)
  for (let i = 0; i < output.length; i++) output[i] = left[i] + right[i]
  return output
}

/** Portable f32 checkpoint representation, bounded to 2.67 MB at 500k values. */
export function encodeVectorAddOutput(values: Float32Array) {
  const bytes = new Uint8Array(values.length * 4)
  const view = new DataView(bytes.buffer)
  for (let i = 0; i < values.length; i++) view.setFloat32(i * 4, values[i], true)
  let binary = ''
  for (let offset = 0; offset < bytes.length; offset += 0x8000) binary += String.fromCharCode(...bytes.subarray(offset, offset + 0x8000))
  return { encoding: 'f32-le-base64', elements: values.length, data: btoa(binary) }
}

export async function digestVectorAddResult(values: Float32Array) {
  const bytes = values.slice().buffer
  const result = await crypto.subtle.digest('SHA-256', bytes)
  return Array.from(new Uint8Array(result), byte => byte.toString(16).padStart(2, '0')).join('')
}

/** Creates a bounded WebGPU vector-add worker, or null when WebGPU is unavailable. */
export async function createVectorAddWorker() {
  const gpu = (navigator as Navigator & { gpu?: GpuNavigator }).gpu
  if (!gpu) return null
  const adapter = await gpu.requestAdapter()
  if (!adapter) return null
  const device = await adapter.requestDevice()
  let disposed = false
  let lost = false
  void device.lost.then(() => { lost = true })
  const module = device.createShaderModule({ code: VECTOR_ADD_SHADER })
  const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } })

  return {
    async run(start: number, end: number) {
      if (disposed || lost) throw new Error('WebGPU device was lost. Reconnect this node and retry the job.')
      const count = end - start
      if (!Number.isSafeInteger(count) || count <= 0 || count > 500_000) throw new Error('WebGPU batch exceeds the safe 500,000-element limit.')
      const byteLength = count * Float32Array.BYTES_PER_ELEMENT
      if (byteLength > device.limits.maxStorageBufferBindingSize || byteLength > device.limits.maxBufferSize) {
        throw new Error('This GPU cannot fit one batch within its storage-buffer limits.')
      }

      const { left, right } = vectorInputs(start, end)

      const usage = STORAGE | COPY_DST
      const leftBuffer = device.createBuffer({ size: byteLength, usage })
      const rightBuffer = device.createBuffer({ size: byteLength, usage })
      const outputBuffer = device.createBuffer({ size: byteLength, usage: STORAGE | COPY_SRC })
      const readbackBuffer = device.createBuffer({ size: byteLength, usage: MAP_READ | COPY_DST })
      try {
        device.queue.writeBuffer(leftBuffer, 0, left)
        device.queue.writeBuffer(rightBuffer, 0, right)
        const bindGroup = device.createBindGroup({
          layout: (pipeline as { getBindGroupLayout(index: number): unknown }).getBindGroupLayout(0),
          entries: [
            { binding: 0, resource: { buffer: leftBuffer } },
            { binding: 1, resource: { buffer: rightBuffer } },
            { binding: 2, resource: { buffer: outputBuffer } },
          ],
        })
        const encoder = device.createCommandEncoder()
        const pass = encoder.beginComputePass()
        pass.setPipeline(pipeline)
        pass.setBindGroup(0, bindGroup)
        pass.dispatchWorkgroups(Math.ceil(count / 64))
        pass.end()
        encoder.copyBufferToBuffer(outputBuffer, 0, readbackBuffer, 0, byteLength)
        device.queue.submit([encoder.finish()])
        await readbackBuffer.mapAsync(MAP_READ)
        const result = new Float32Array(readbackBuffer.getMappedRange().slice(0))
        readbackBuffer.unmap()

        for (let i = 0; i < count; i++) {
          const expected = Math.fround(left[i] + right[i])
          if (!Number.isFinite(result[i]) || result[i] !== expected) throw new Error('WebGPU batch validation failed; result was not checkpointed.')
        }
        return result
      } finally {
        leftBuffer.destroy()
        rightBuffer.destroy()
        outputBuffer.destroy()
        readbackBuffer.destroy()
      }
    },
    close() {
      if (!disposed) { disposed = true; device.destroy() }
    },
  }
}