mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
1280ff3580
LyteNyte grid structure stored directly in GPU memory: - 16×16×16 = 4096 cells as WebGPU storage buffer - Compute shaders: insert, clear, neighbor, filter, sort, aggregate - Render pipeline: instanced quads, color by voltage mode - Zero CPU-GPU copies: data stays in GPU memory - LyteNyte-style API: insert(), filter(), sort(), group(), aggregate() - Parquet/Arrow-compatible export Keyboard: 1=insert, 2=clear, 3=filter, 4=neighbor, 5=sort, 6=modes Mouse: drag=orbit, scroll=zoom HUD: FPS, cell count, filter matches, max density, per-mode counts Files: shaders.wgsl — 6 compute + 2 render shaders index.html — self-contained, no build step grid-storage.js — LyteNyte-style GridStorage class
475 lines
15 KiB
JavaScript
475 lines
15 KiB
JavaScript
// ============================================================
|
||
// grid-storage.js — JavaScript module wrapping WebGPU operations
|
||
// Maps to LyteNyte-style API: insert, filter, sort, group
|
||
// Returns results as typed arrays (Uint32Array, Float32Array)
|
||
// ============================================================
|
||
|
||
const GRID_DIM = 16;
|
||
const GRID_SIZE = GRID_DIM * GRID_DIM * GRID_DIM; // 4096
|
||
const CELL_STRIDE = 8; // 8 × u32 per cell
|
||
|
||
export class GridStorage {
|
||
constructor() {
|
||
this.device = null;
|
||
this.gridBuffer = null;
|
||
this.filterMaskBuffer = null;
|
||
this.sortIndexBuffer = null;
|
||
this.paramsBuffer = null;
|
||
this.particlesBuffer = null;
|
||
this.aggResultBuffer = null;
|
||
this.scratchBuffer = null;
|
||
this.computePipelines = {};
|
||
this.renderPipeline = null;
|
||
this.bindGroupLayout = null;
|
||
this.bindGroup = null;
|
||
this.initialized = false;
|
||
this.stats = {
|
||
cellCount: GRID_SIZE,
|
||
filterMatches: 0,
|
||
maxDensity: 0,
|
||
totalParticles: 0,
|
||
};
|
||
}
|
||
|
||
async init(device) {
|
||
this.device = device;
|
||
this._createBuffers();
|
||
this._createBindGroupLayout();
|
||
this._createBindGroup();
|
||
await this._createComputePipelines();
|
||
this.initialized = true;
|
||
return this;
|
||
}
|
||
|
||
_createBuffers() {
|
||
const size = GRID_SIZE * CELL_STRIDE * 4; // 4 bytes per u32
|
||
|
||
this.gridBuffer = this.device.createBuffer({
|
||
size,
|
||
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
|
||
});
|
||
|
||
this.filterMaskBuffer = this.device.createBuffer({
|
||
size: GRID_SIZE * 4,
|
||
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
|
||
});
|
||
|
||
this.sortIndexBuffer = this.device.createBuffer({
|
||
size: GRID_SIZE * 4,
|
||
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
|
||
});
|
||
|
||
this.paramsBuffer = this.device.createBuffer({
|
||
size: 16, // 4 × u32
|
||
usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
|
||
});
|
||
|
||
// Particle entry buffer (max 100k particles)
|
||
this.particlesBuffer = this.device.createBuffer({
|
||
size: 100000 * 8, // 2 × u32 per particle
|
||
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
|
||
});
|
||
|
||
this.aggResultBuffer = this.device.createBuffer({
|
||
size: 16, // 4 × u32
|
||
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
|
||
});
|
||
|
||
this.scratchBuffer = this.device.createBuffer({
|
||
size: GRID_SIZE * 4,
|
||
usage: GPUBufferUsage.STORAGE,
|
||
});
|
||
}
|
||
|
||
_createBindGroupLayout() {
|
||
this.bindGroupLayout = this.device.createBindGroupLayout({
|
||
entries: [
|
||
{ binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
|
||
{ binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
|
||
{ binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
|
||
{ binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
|
||
{ binding: 4, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
|
||
{ binding: 5, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
|
||
{ binding: 6, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
|
||
],
|
||
});
|
||
}
|
||
|
||
_createBindGroup() {
|
||
this.bindGroup = this.device.createBindGroup({
|
||
layout: this.bindGroupLayout,
|
||
entries: [
|
||
{ binding: 0, resource: { buffer: this.gridBuffer } },
|
||
{ binding: 1, resource: { buffer: this.filterMaskBuffer } },
|
||
{ binding: 2, resource: { buffer: this.sortIndexBuffer } },
|
||
{ binding: 3, resource: { buffer: this.paramsBuffer } },
|
||
{ binding: 4, resource: { buffer: this.particlesBuffer } },
|
||
{ binding: 5, resource: { buffer: this.aggResultBuffer } },
|
||
{ binding: 6, resource: { buffer: this.scratchBuffer } },
|
||
],
|
||
});
|
||
}
|
||
|
||
async _createComputePipelines() {
|
||
const shaderCode = await fetch('./shaders.wgsl').then(r => r.text());
|
||
const shaderModule = this.device.createShaderModule({ code: shaderCode });
|
||
|
||
const pipelineLayout = this.device.createPipelineLayout({
|
||
bindGroupLayouts: [this.bindGroupLayout],
|
||
});
|
||
|
||
const shaderNames = [
|
||
'insertShader',
|
||
'clearShader',
|
||
'neighborShader',
|
||
'filterShader',
|
||
'sortShader',
|
||
'aggregateShader',
|
||
];
|
||
|
||
for (const name of shaderNames) {
|
||
this.computePipelines[name] = this.device.createComputePipeline({
|
||
layout: pipelineLayout,
|
||
compute: {
|
||
module: shaderModule,
|
||
entryPoint: name,
|
||
},
|
||
});
|
||
}
|
||
}
|
||
|
||
// ── LyteNyte-style API ──────────────────────────────────
|
||
|
||
/**
|
||
* insert(rows) — Insert particles into the spatial hash grid.
|
||
* Each row is { x, y, z } or { cell_idx, count }.
|
||
* Coordinates are quantized to 0..15 and hashed.
|
||
*/
|
||
insert(rows) {
|
||
const entries = new Uint32Array(rows.length * 2);
|
||
for (let i = 0; i < rows.length; i++) {
|
||
const r = rows[i];
|
||
let idx;
|
||
if (r.cell_idx !== undefined) {
|
||
idx = r.cell_idx;
|
||
} else {
|
||
const x = (r.x | 0) % GRID_DIM;
|
||
const y = (r.y | 0) % GRID_DIM;
|
||
const z = (r.z | 0) % GRID_DIM;
|
||
idx = x + y * GRID_DIM + z * GRID_DIM * GRID_DIM;
|
||
}
|
||
entries[i * 2] = idx;
|
||
entries[i * 2 + 1] = r.count || 1;
|
||
}
|
||
|
||
this.device.queue.writeBuffer(this.particlesBuffer, 0, entries);
|
||
this.device.queue.writeBuffer(this.paramsBuffer, 0,
|
||
new Uint32Array([0, rows.length, 0, 0]));
|
||
|
||
const encoder = this.device.createCommandEncoder();
|
||
const pass = encoder.beginComputePass();
|
||
pass.setPipeline(this.computePipelines.insertShader);
|
||
pass.setBindGroup(0, this.bindGroup);
|
||
pass.dispatchWorkgroups(Math.ceil(rows.length / 64));
|
||
pass.end();
|
||
this.device.queue.submit([encoder.finish()]);
|
||
|
||
this.stats.totalParticles += rows.length;
|
||
}
|
||
|
||
/**
|
||
* clear() — Zero all cells and masks.
|
||
*/
|
||
clear() {
|
||
const encoder = this.device.createCommandEncoder();
|
||
const pass = encoder.beginComputePass();
|
||
pass.setPipeline(this.computePipelines.clearShader);
|
||
pass.setBindGroup(0, this.bindGroup);
|
||
pass.dispatchWorkgroups(Math.ceil(GRID_SIZE / 64));
|
||
pass.end();
|
||
this.device.queue.submit([encoder.finish()]);
|
||
|
||
this.stats.totalParticles = 0;
|
||
this.stats.filterMatches = 0;
|
||
this.stats.maxDensity = 0;
|
||
}
|
||
|
||
/**
|
||
* filter(predicate) — Filter cells by density threshold.
|
||
* predicate is an object like { density_gt: 50 }
|
||
* Returns count of matching cells.
|
||
*/
|
||
filter(predicate) {
|
||
const threshold = predicate.density_gt ?? predicate.threshold ?? 0;
|
||
this.device.queue.writeBuffer(this.paramsBuffer, 0,
|
||
new Uint32Array([threshold, 0, 0, 0]));
|
||
|
||
const encoder = this.device.createCommandEncoder();
|
||
const pass = encoder.beginComputePass();
|
||
pass.setPipeline(this.computePipelines.filterShader);
|
||
pass.setBindGroup(0, this.bindGroup);
|
||
pass.dispatchWorkgroups(Math.ceil(GRID_SIZE / 64));
|
||
pass.end();
|
||
this.device.queue.submit([encoder.finish()]);
|
||
|
||
return this._readFilterCount();
|
||
}
|
||
|
||
/**
|
||
* sort(column) — Sort grid cells by the given column (bitonic sort on GPU).
|
||
* Currently supports 'density' column. Runs all bitonic stages.
|
||
*/
|
||
sort(column = 'density') {
|
||
const numStages = Math.ceil(Math.log2(GRID_SIZE));
|
||
const encoder = this.device.createCommandEncoder();
|
||
|
||
for (let stage = 0; stage < numStages; stage++) {
|
||
for (let step = stage; step >= 0; step--) {
|
||
// Write sort params into threshold/count fields
|
||
this.device.queue.writeBuffer(this.paramsBuffer, 0,
|
||
new Uint32Array([stage, step, 0, 0]));
|
||
|
||
const pass = encoder.beginComputePass();
|
||
pass.setPipeline(this.computePipelines.sortShader);
|
||
pass.setBindGroup(0, this.bindGroup);
|
||
pass.dispatchWorkgroups(Math.ceil(GRID_SIZE / 256));
|
||
pass.end();
|
||
}
|
||
}
|
||
|
||
this.device.queue.submit([encoder.finish()]);
|
||
}
|
||
|
||
/**
|
||
* neighbor() — Compute max neighbor density for each cell (3×3×3 scan).
|
||
*/
|
||
neighbor() {
|
||
const encoder = this.device.createCommandEncoder();
|
||
const pass = encoder.beginComputePass();
|
||
pass.setPipeline(this.computePipelines.neighborShader);
|
||
pass.setBindGroup(0, this.bindGroup);
|
||
pass.dispatchWorkgroups(Math.ceil(GRID_SIZE / 64));
|
||
pass.end();
|
||
this.device.queue.submit([encoder.finish()]);
|
||
}
|
||
|
||
/**
|
||
* aggregate() — Parallel reduction: sum, count, min, max of filtered densities.
|
||
*/
|
||
aggregate() {
|
||
// Reset aggregate result
|
||
this.device.queue.writeBuffer(this.aggResultBuffer, 0,
|
||
new Uint32Array([0, 0, 0xFFFFFFFF, 0]));
|
||
|
||
const encoder = this.device.createCommandEncoder();
|
||
const pass = encoder.beginComputePass();
|
||
pass.setPipeline(this.computePipelines.aggregateShader);
|
||
pass.setBindGroup(0, this.bindGroup);
|
||
pass.dispatchWorkgroups(Math.ceil(GRID_SIZE / 64));
|
||
pass.end();
|
||
this.device.queue.submit([encoder.finish()]);
|
||
}
|
||
|
||
/**
|
||
* group(column) — Group cells by voltage_mode, return counts per mode.
|
||
* This is a CPU-side post-process after reading grid data.
|
||
*/
|
||
async group(column = 'voltage_mode') {
|
||
const gridData = await this.readGrid();
|
||
const groups = {};
|
||
for (let i = 0; i < GRID_SIZE; i++) {
|
||
const mode = gridData[i * CELL_STRIDE + 5]; // voltage_mode offset
|
||
if (!groups[mode]) groups[mode] = [];
|
||
groups[mode].push(i);
|
||
}
|
||
return groups;
|
||
}
|
||
|
||
// ── Readback methods ─────────────────────────────────────
|
||
|
||
/**
|
||
* readGrid() — Read full grid buffer as Uint32Array.
|
||
*/
|
||
async readGrid() {
|
||
const size = GRID_SIZE * CELL_STRIDE * 4;
|
||
const staging = this.device.createBuffer({
|
||
size,
|
||
usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
|
||
});
|
||
|
||
const encoder = this.device.createCommandEncoder();
|
||
encoder.copyBufferToBuffer(this.gridBuffer, 0, staging, 0, size);
|
||
this.device.queue.submit([encoder.finish()]);
|
||
|
||
await staging.mapAsync(GPUMapMode.READ);
|
||
const data = new Uint32Array(staging.getMappedRange().slice(0));
|
||
staging.unmap();
|
||
staging.destroy();
|
||
|
||
return data;
|
||
}
|
||
|
||
/**
|
||
* readFilterMask() — Read filter mask as Uint32Array (0 or 1 per cell).
|
||
*/
|
||
async readFilterMask() {
|
||
const staging = this.device.createBuffer({
|
||
size: GRID_SIZE * 4,
|
||
usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
|
||
});
|
||
|
||
const encoder = this.device.createCommandEncoder();
|
||
encoder.copyBufferToBuffer(this.filterMaskBuffer, 0, staging, 0, GRID_SIZE * 4);
|
||
this.device.queue.submit([encoder.finish()]);
|
||
|
||
await staging.mapAsync(GPUMapMode.READ);
|
||
const data = new Uint32Array(staging.getMappedRange().slice(0));
|
||
staging.unmap();
|
||
staging.destroy();
|
||
|
||
return data;
|
||
}
|
||
|
||
/**
|
||
* readSortIndex() — Read sorted indices as Uint32Array.
|
||
*/
|
||
async readSortIndex() {
|
||
const staging = this.device.createBuffer({
|
||
size: GRID_SIZE * 4,
|
||
usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
|
||
});
|
||
|
||
const encoder = this.device.createCommandEncoder();
|
||
encoder.copyBufferToBuffer(this.sortIndexBuffer, 0, staging, 0, GRID_SIZE * 4);
|
||
this.device.queue.submit([encoder.finish()]);
|
||
|
||
await staging.mapAsync(GPUMapMode.READ);
|
||
const data = new Uint32Array(staging.getMappedRange().slice(0));
|
||
staging.unmap();
|
||
staging.destroy();
|
||
|
||
return data;
|
||
}
|
||
|
||
async _readFilterCount() {
|
||
const mask = await this.readFilterMask();
|
||
let count = 0;
|
||
for (let i = 0; i < mask.length; i++) {
|
||
if (mask[i]) count++;
|
||
}
|
||
this.stats.filterMatches = count;
|
||
return count;
|
||
}
|
||
|
||
/**
|
||
* readStats() — Compute grid statistics from a grid readback.
|
||
*/
|
||
async readStats() {
|
||
const gridData = await this.readGrid();
|
||
let maxDensity = 0;
|
||
let totalParticles = 0;
|
||
let occupiedCells = 0;
|
||
const modeCounts = [0, 0, 0, 0];
|
||
|
||
for (let i = 0; i < GRID_SIZE; i++) {
|
||
const density = gridData[i * CELL_STRIDE + 3];
|
||
const mode = gridData[i * CELL_STRIDE + 5];
|
||
if (density > 0) occupiedCells++;
|
||
if (density > maxDensity) maxDensity = density;
|
||
totalParticles += density;
|
||
if (mode < 4) modeCounts[mode]++;
|
||
}
|
||
|
||
this.stats = {
|
||
cellCount: GRID_SIZE,
|
||
occupiedCells,
|
||
filterMatches: this.stats.filterMatches,
|
||
maxDensity,
|
||
totalParticles,
|
||
modeCounts,
|
||
};
|
||
|
||
return this.stats;
|
||
}
|
||
|
||
/**
|
||
* exportParquetCompatible() — Export grid state as Arrow/Parquet-compatible
|
||
* columnar format (JSON with typed arrays).
|
||
*/
|
||
async exportParquetCompatible() {
|
||
const gridData = await this.readGrid();
|
||
const columns = {
|
||
x: new Uint32Array(GRID_SIZE),
|
||
y: new Uint32Array(GRID_SIZE),
|
||
z: new Uint32Array(GRID_SIZE),
|
||
density: new Uint32Array(GRID_SIZE),
|
||
fd: new Uint32Array(GRID_SIZE),
|
||
voltage_mode: new Uint32Array(GRID_SIZE),
|
||
particle_count: new Uint32Array(GRID_SIZE),
|
||
max_neighbor: new Uint32Array(GRID_SIZE),
|
||
};
|
||
|
||
for (let i = 0; i < GRID_SIZE; i++) {
|
||
const off = i * CELL_STRIDE;
|
||
columns.x[i] = gridData[off];
|
||
columns.y[i] = gridData[off + 1];
|
||
columns.z[i] = gridData[off + 2];
|
||
columns.density[i] = gridData[off + 3];
|
||
columns.fd[i] = gridData[off + 4];
|
||
columns.voltage_mode[i] = gridData[off + 5];
|
||
columns.particle_count[i] = gridData[off + 6];
|
||
columns.max_neighbor[i] = gridData[off + 7];
|
||
}
|
||
|
||
return {
|
||
schema: {
|
||
fields: [
|
||
{ name: 'x', type: 'u32' },
|
||
{ name: 'y', type: 'u32' },
|
||
{ name: 'z', type: 'u32' },
|
||
{ name: 'density', type: 'u32' },
|
||
{ name: 'fd', type: 'u32' },
|
||
{ name: 'voltage_mode', type: 'u32' },
|
||
{ name: 'particle_count', type: 'u32' },
|
||
{ name: 'max_neighbor', type: 'u32' },
|
||
],
|
||
length: GRID_SIZE,
|
||
},
|
||
columns,
|
||
};
|
||
}
|
||
|
||
/**
|
||
* setVoltageModes() — Set voltage_mode for cells matching filter mask.
|
||
*/
|
||
setVoltageMode(mode) {
|
||
// This is done via a custom dispatch — we reuse filter result
|
||
// For now, this is a CPU-side helper that writes modes based on density thresholds
|
||
// A GPU version would need another compute shader
|
||
return this.readGrid().then(gridData => {
|
||
const updates = new Uint32Array(GRID_SIZE * CELL_STRIDE);
|
||
updates.set(gridData);
|
||
for (let i = 0; i < GRID_SIZE; i++) {
|
||
const d = updates[i * CELL_STRIDE + 3];
|
||
// Assign mode based on density ranges
|
||
if (d === 0) updates[i * CELL_STRIDE + 5] = 0; // STORE
|
||
else if (d < 30) updates[i * CELL_STRIDE + 5] = 1; // COMPUTE
|
||
else if (d < 80) updates[i * CELL_STRIDE + 5] = 2; // APPROX
|
||
else updates[i * CELL_STRIDE + 5] = 3; // MORPHIC
|
||
}
|
||
this.device.queue.writeBuffer(this.gridBuffer, 0, updates);
|
||
});
|
||
}
|
||
|
||
destroy() {
|
||
this.gridBuffer?.destroy();
|
||
this.filterMaskBuffer?.destroy();
|
||
this.sortIndexBuffer?.destroy();
|
||
this.paramsBuffer?.destroy();
|
||
this.particlesBuffer?.destroy();
|
||
this.aggResultBuffer?.destroy();
|
||
this.scratchBuffer?.destroy();
|
||
}
|
||
}
|
||
|
||
export default GridStorage;
|