@manycore/aholo-splat-transform 1.2.9 → 1.2.11

package/CHANGELOG.md

@@ -1,5 +1,16 @@
 # ChangeLOG
 
+## 1.2.11
+
+- 移除`webgpu`依赖，改为自主编译版本
+- 优化`k-means`计算，采用`vec4`对齐，简化欧几里得距离计算。(性能提升约50%)
+- 优化`clusterAverage`，减少js端overhead
+- native线程池共享
+
+## 1.2.10
+
+- 增加`darwin-arm64`的预编译版本
+
 ## 1.2.9
 
 - 改进编码风格开启`verbatimModuleSyntax` & `isolatedModules`，迁移至`OXC`

package/README.md

@@ -8,7 +8,7 @@ A 3DGS modifier used by aholo
 - system
   - windows: windows 22H2+, x86_64, D3D12 or Vulkan compatible GPU(When use GPU features, dedicated GPU for better performance)
   - linux: x86_64, glibc >= 2.34, libstdc++ >= 3.4.30, Vulkan compatible GPU(When use GPU features, dedicated GPU for better performance)
-  - osx: Not Supported
+  - osx: apple silicon ARM64 only.
 
 ## Usage
 

package/dist/file/sog.js

@@ -3,7 +3,7 @@ import { Buffer } from 'node:buffer';
 import { decodeWebP, encodeWebP, WebPLosslessProfile } from '../native/index.js';
 import { ColIdx } from '../SplatData.js';
 import { SH_C0, SH_MAPS, NUM_F_REST_TO_SH_DEGREE } from '../constant.js';
-import { getOrCreateDevice, kmeans, logger, mortonSort, quantize1d, isUrl, extractFromRootDir, clamp, } from '../utils/index.js';
+import { getOrCreateDevice, kMeans, logger, mortonSort, quantize1d, isUrl, extractFromRootDir, clamp, } from '../utils/index.js';
 const ZIP_MAGIC = 0x04034b50;
 const PERM_TABLE = [
     // original quat idx ---> actual storage idx
@@ -485,7 +485,7 @@ export class SogFile {
             const device = await getOrCreateDevice();
             logger.info(`SOG SH${shDegree} k-means with clusters=${paletteSize} iterations=${this.iterations}`);
             logger.time(`SOG SH${shDegree} k-means`);
-            const { centroids, labels } = await kmeans(shDataTable, paletteSize, this.iterations, device);
+            const { centroids, labels } = await kMeans(shDataTable, paletteSize, this.iterations, device);
             logger.timeEnd(`SOG SH${shDegree} k-means`);
             const codebook = quantize1d(centroids);
             // write centroids

package/dist/native/index.d.ts

@@ -51,4 +51,4 @@ export declare function decodeAVIFBatched(inputs: Array<Uint8Array | Buffer>): {
     width: number;
     height: number;
 }[];
-export declare function clusterAverage(dataTable: Float32Array[], clusters: Uint32Array[], output: Float32Array[]): void;
+export declare function clusterAverage(dataTable: Float32Array[], labels: Uint32Array, k: number, output: Float32Array[]): void;

package/dist/native/index.js

@@ -1,26 +1,36 @@
 import { createRequire } from 'node:module';
 import { SplatData } from '../SplatData.js';
 import { Buffer } from 'node:buffer';
-import { isMusl } from './utils.js';
-import p from '../../package.json' with { type: 'json' };
+import { getNativePackageName } from './utils.js';
 const getModule = (function () {
     let m = undefined;
     const require = createRequire(import.meta.url);
-    let runtime = undefined;
-    if (process.platform === 'win32') {
-        runtime = 'msvc';
-    }
-    else if (process.platform === 'linux') {
-        runtime = isMusl() ? 'musl' : 'gnu';
-    }
-    const binaryPackage = `${p.name}-${process.platform}-${process.arch}${runtime ? `-${runtime}` : ''}`;
+    const binaryModule = getNativePackageName() + '/splat-transform.node';
     return function () {
         if (!m) {
-            m = require(binaryPackage);
+            m = require(binaryModule);
         }
         return m;
     };
 })();
+const [defaultThreadPool, smallThreadPool] = (function () {
+    let defaultTheadPool;
+    let smallTheadPool;
+    return [
+        function () {
+            if (!defaultTheadPool) {
+                defaultTheadPool = new (getModule().ThreadPool)();
+            }
+            return defaultTheadPool;
+        },
+        function () {
+            if (!smallTheadPool) {
+                smallTheadPool = new (getModule().ThreadPool)(4);
+            }
+            return smallTheadPool;
+        },
+    ];
+})();
 export function generateLod(splat, levelParameters, blockPrecision, minSize, maxStep) {
     if (splat.counts === 0) {
         return {
@@ -44,7 +54,7 @@ export function generateLod(splat, levelParameters, blockPrecision, minSize, max
             parameters[i * 2 + 1] = scaleBoost;
         }
     }
-    const { blockBoxes, blockRefs, gaussianCount, data } = getModule().generate_lod(inputBuffers, splat.shCounts, buffer, blockPrecision, minSize, maxStep);
+    const { blockBoxes, blockRefs, gaussianCount, data } = getModule().generate_lod(inputBuffers, splat.shCounts, buffer, blockPrecision, minSize, maxStep, defaultThreadPool());
     const blockView = new Float32Array(blockBoxes.buffer, blockBoxes.byteOffset, blockBoxes.byteLength / 4);
     const blockRefsView = new Uint32Array(blockRefs.buffer, blockRefs.byteOffset, blockRefs.byteLength / 4);
     const blockCount = blockView.length / 6;
@@ -108,15 +118,15 @@ export function encodeAVIFBatched(inputs) {
     return getModule().avif_encode_rgba_batched(inputs.map(i => ({
         ...i,
         data: i.data instanceof Buffer ? i.data : Buffer.from(i.data.buffer, i.data.byteOffset, i.data.byteLength),
-    })));
+    })), smallThreadPool());
 }
 export function decodeAVIF(data) {
     const buffer = data instanceof Buffer ? data : Buffer.from(data.buffer, data.byteOffset, data.byteLength);
     return getModule().avif_decode_rgba(buffer);
 }
 export function decodeAVIFBatched(inputs) {
-    return getModule().avif_decode_rgba_batched(inputs.map(i => (i instanceof Buffer ? i : Buffer.from(i.buffer, i.byteOffset, i.byteLength))));
+    return getModule().avif_decode_rgba_batched(inputs.map(i => (i instanceof Buffer ? i : Buffer.from(i.buffer, i.byteOffset, i.byteLength))), smallThreadPool());
 }
-export function clusterAverage(dataTable, clusters, output) {
-    return getModule().cluster_average(dataTable.map(t => Buffer.from(t.buffer, t.byteOffset, t.byteLength)), clusters.map(t => Buffer.from(t.buffer, t.byteOffset, t.byteLength)), output.map(t => Buffer.from(t.buffer, t.byteOffset, t.byteLength)));
+export function clusterAverage(dataTable, labels, k, output) {
+    return getModule().cluster_average(dataTable.map(t => Buffer.from(t.buffer, t.byteOffset, t.byteLength)), Buffer.from(labels.buffer, labels.byteOffset, labels.byteLength), k, output.map(t => Buffer.from(t.buffer, t.byteOffset, t.byteLength)), defaultThreadPool());
 }

package/dist/native/utils.d.ts

@@ -1 +1 @@
-export declare function isMusl(): boolean;
+export declare function getNativePackageName(): string;

package/dist/native/utils.js

@@ -1,6 +1,7 @@
 import { readFileSync } from 'node:fs';
 import child_process from 'node:child_process';
-export function isMusl() {
+import p from '../../package.json' with { type: 'json' };
+function isMusl() {
     let musl = false;
     if (process.platform === 'linux') {
         musl = isMuslFromFilesystem();
@@ -52,3 +53,13 @@ function isMuslFromChildProcess() {
         return false;
     }
 }
+export function getNativePackageName() {
+    let runtime = undefined;
+    if (process.platform === 'win32') {
+        runtime = 'msvc';
+    }
+    else if (process.platform === 'linux') {
+        runtime = isMusl() ? 'musl' : 'gnu';
+    }
+    return `${p.name}-${process.platform}-${process.arch}${runtime ? `-${runtime}` : ''}`;
+}

package/dist/utils/index.d.ts

@@ -15,7 +15,7 @@ export * from './StreamChunkDecoder.js';
 export * from './math.js';
 export * from './sh-rotate.js';
 export * from './splat.js';
-export * from './k-means.js';
+export * from './k-means/index.js';
 export * from './quantize-1d.js';
 export * from './webgpu.js';
 export * from './voxel/common.js';

package/dist/utils/index.js

@@ -89,7 +89,7 @@ export * from './StreamChunkDecoder.js';
 export * from './math.js';
 export * from './sh-rotate.js';
 export * from './splat.js';
-export * from './k-means.js';
+export * from './k-means/index.js';
 export * from './quantize-1d.js';
 export * from './webgpu.js';
 export * from './voxel/common.js';

package/dist/utils/k-means/GpuClustering.d.ts

@@ -0,0 +1,16 @@
+export default class GpuClustering {
+    private device;
+    private numPoints;
+    private numCentroids;
+    private batchSize;
+    private resource;
+    private numBatches;
+    private concurrencyBatches;
+    private concurrencyRuns;
+    private workgroupSize;
+    private pointStride;
+    private vecColumns;
+    constructor(device: GPUDevice, numPoints: number, numColumns: number, numCentroids: number);
+    execute(points: Float32Array[], centroids: Float32Array[], labels: Uint32Array): Promise<void>;
+    destroy(): void;
+}

package/dist/utils/k-means/GpuClustering.js

@@ -0,0 +1,287 @@
+import { logger } from '../index.js';
+const chunkSize = 128;
+const workgroupSize = 128;
+function clusterWgsl(vecColumns) {
+    return /* wgsl */ `
+struct Uniforms {
+    numPoints: u32,
+    numCentroids: u32
+};
+
+@group(0) @binding(0) var<uniform> uniforms: Uniforms;
+@group(0) @binding(1) var<storage, read> points: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> centroids: array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read> centroidSq: array<f32>;
+@group(0) @binding(4) var<storage, read_write> results: array<u32>;
+
+const vecColumns = ${vecColumns}u;
+const chunkSize = ${chunkSize}u;
+const workgroupSize = ${workgroupSize}u;
+var<workgroup> sharedChunk: array<vec4<f32>, vecColumns * chunkSize>;
+var<workgroup> sharedSq: array<f32, chunkSize>;
+
+fn calcDistance(point: array<vec4<f32>, vecColumns>, centroid: u32) -> f32 {
+    let ci = centroid * vecColumns;
+    var result = sharedSq[centroid];
+
+    for (var i = 0u; i < vecColumns; i++) {
+        // euclid distance simplify
+        // (centroid - point) ^ 2 = centroid ^ 2 - 2 * dot(centroid, point) + point ^ 2
+        // point ^ 2 omitted, for same point find nearest centroid is not necessary
+        result -= 2.0 * dot(point[i], sharedChunk[ci + i]);
+    }
+
+    return result;
+}
+
+@compute @workgroup_size(workgroupSize)
+fn main(
+    @builtin(local_invocation_index) localId : u32,
+    @builtin(global_invocation_id) globalId: vec3u,
+    @builtin(num_workgroups) numWorkgroups: vec3u
+) {
+    let pointIndex = globalId.x + globalId.y * numWorkgroups.x * workgroupSize;
+
+    var point: array<vec4<f32>, vecColumns>;
+    if (pointIndex < uniforms.numPoints) {
+        for (var i = 0u; i < vecColumns; i++) {
+            point[i] = points[pointIndex * vecColumns + i];
+        }
+    }
+
+    var mind = 1000000.0;
+    var mini = 0u;
+
+    let numChunks = u32(ceil(f32(uniforms.numCentroids) / f32(chunkSize)));
+    for (var i = 0u; i < numChunks; i++) {
+        let chunkToLoad = min(chunkSize, uniforms.numCentroids - i * chunkSize);
+        for (var row = localId; row < chunkToLoad; row += workgroupSize) {
+            let srcRow = i * chunkSize + row;
+            let dst = row * vecColumns;
+            let src = srcRow * vecColumns;
+
+            for (var c = 0u; c < vecColumns; c++) {
+                sharedChunk[dst + c] = centroids[src + c];
+            }
+            sharedSq[row] = centroidSq[srcRow];
+        }
+
+        workgroupBarrier();
+
+        if (pointIndex < uniforms.numPoints) {
+            let thisChunkSize = min(chunkSize, uniforms.numCentroids - i * chunkSize);
+            for (var c = 0u; c < thisChunkSize; c++) {
+                let d = calcDistance(point, c);
+                if (d < mind) {
+                    mind = d;
+                    mini = i * chunkSize + c;
+                }
+            }
+        }
+
+        workgroupBarrier();
+    }
+
+    if (pointIndex < uniforms.numPoints) {
+        results[pointIndex] = mini;
+    }
+}
+`;
+}
+function packVec4Data(result, dataTable, numRows, rowOffset, vecColumns, norms) {
+    const numColumns = dataTable.length;
+    const stride = vecColumns * 4;
+    for (let r = 0; r < numRows; r++) {
+        const dst = r * stride;
+        let norm = 0.0;
+        for (let c = 0; c < numColumns; c++) {
+            const v = dataTable[c][rowOffset + r];
+            result[dst + c] = v;
+            if (norms) {
+                norm += v * v;
+            }
+        }
+        for (let c = numColumns; c < stride; c++) {
+            result[dst + c] = 0.0;
+        }
+        if (norms) {
+            norms[r] = norm;
+        }
+    }
+}
+const MAX_CONCURRENCY_BATCHES = 10;
+export default class GpuClustering {
+    constructor(device, numPoints, numColumns, numCentroids) {
+        this.device = device;
+        this.numPoints = numPoints;
+        this.numCentroids = numCentroids;
+        this.vecColumns = Math.ceil(numColumns / 4);
+        this.workgroupSize = workgroupSize;
+        this.pointStride = this.vecColumns * 4;
+        const storageLimit = Math.min(device.limits.maxBufferSize, device.limits.maxStorageBufferBindingSize);
+        const workgroupsPerBatch = Math.max(1, Math.min(device.limits.maxComputeWorkgroupsPerDimension, // device dispatch limit
+        Math.floor(storageLimit / (this.pointStride * this.workgroupSize * 4)), // point storage limit
+        Math.ceil(numPoints / this.workgroupSize)));
+        this.batchSize = workgroupsPerBatch * this.workgroupSize;
+        this.numBatches = Math.ceil(numPoints / this.batchSize);
+        this.concurrencyBatches = Math.min(MAX_CONCURRENCY_BATCHES, this.numBatches);
+        this.concurrencyRuns = Math.ceil(this.numBatches / this.concurrencyBatches);
+        const shader = device.createShaderModule({
+            code: clusterWgsl(this.vecColumns),
+        });
+        const pipeline = device.createComputePipeline({
+            layout: 'auto',
+            compute: {
+                module: shader,
+                entryPoint: 'main',
+            },
+        });
+        const pointsBackBuffer = new Float32Array(this.pointStride * this.batchSize);
+        const centroidsBackBuffer = new Float32Array(this.pointStride * numCentroids);
+        const centroidSqBackBuffer = new Float32Array(numCentroids);
+        const uniformBackBuffer = new Uint32Array([0, numCentroids]);
+        const pointsBuffers = [];
+        const centroidsBuffer = device.createBuffer({
+            size: centroidsBackBuffer.byteLength,
+            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE,
+        });
+        const centroidSqBuffer = device.createBuffer({
+            size: centroidSqBackBuffer.byteLength,
+            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE,
+        });
+        const uniformBuffer = device.createBuffer({
+            size: 256 * this.concurrencyBatches,
+            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM,
+        });
+        const resultBuffer = device.createBuffer({
+            size: this.concurrencyBatches * this.batchSize * 4,
+            usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE,
+        });
+        const resultReadBackBuffer = device.createBuffer({
+            size: this.concurrencyBatches * this.batchSize * 4,
+            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
+        });
+        const layout = pipeline.getBindGroupLayout(0);
+        const bindGroups = [];
+        for (let i = 0; i < this.concurrencyBatches; i++) {
+            const pointsBuffer = device.createBuffer({
+                size: pointsBackBuffer.byteLength,
+                usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE,
+            });
+            pointsBuffers.push(pointsBuffer);
+            const entries = [
+                {
+                    binding: 0,
+                    resource: {
+                        buffer: uniformBuffer,
+                        offset: i * 256,
+                        size: 8,
+                    },
+                },
+                {
+                    binding: 1,
+                    resource: pointsBuffer,
+                },
+                {
+                    binding: 2,
+                    resource: centroidsBuffer,
+                },
+                {
+                    binding: 3,
+                    resource: centroidSqBuffer,
+                },
+                {
+                    binding: 4,
+                    resource: {
+                        buffer: resultBuffer,
+                        offset: i * this.batchSize * 4,
+                        size: this.batchSize * 4,
+                    },
+                },
+            ];
+            bindGroups.push(device.createBindGroup({
+                layout,
+                entries,
+            }));
+        }
+        const gpuBuffers = {
+            uniform: uniformBuffer,
+            points: pointsBuffers,
+            centroids: centroidsBuffer,
+            centroidSq: centroidSqBuffer,
+            result: resultBuffer,
+            resultReadBack: resultReadBackBuffer,
+        };
+        const backBuffers = {
+            uniform: uniformBackBuffer,
+            points: pointsBackBuffer,
+            centroids: centroidsBackBuffer,
+            centroidSq: centroidSqBackBuffer,
+        };
+        this.resource = {
+            pipeline,
+            bindGroups,
+            gpuBuffers,
+            backBuffers,
+            uploadedBatches: [],
+        };
+        logger.info(`GPU k-means kernel bootstrapped with batch ${workgroupsPerBatch}*${this.workgroupSize}*${this.numBatches}, concurrency: ${this.concurrencyBatches}, runs: ${this.concurrencyRuns}`);
+    }
+    async execute(points, centroids, labels) {
+        const { device, numPoints, numCentroids, numBatches, batchSize, resource, concurrencyBatches, concurrencyRuns, pointStride, vecColumns, workgroupSize, } = this;
+        // upload centroid data to gpu
+        packVec4Data(resource.backBuffers.centroids, centroids, numCentroids, 0, vecColumns, resource.backBuffers.centroidSq);
+        device.queue.writeBuffer(resource.gpuBuffers.centroids, 0, resource.backBuffers.centroids.buffer);
+        device.queue.writeBuffer(resource.gpuBuffers.centroidSq, 0, resource.backBuffers.centroidSq.buffer);
+        for (let i = 0; i < concurrencyRuns; i++) {
+            const batchStart = i * concurrencyBatches;
+            let resultCount = 0;
+            for (let j = 0; j < concurrencyBatches; j++) {
+                const batchIndex = batchStart + j;
+                if (batchIndex >= numBatches) {
+                    break;
+                }
+                const currentBatchSize = Math.min(numPoints - batchIndex * batchSize, batchSize);
+                resultCount += currentBatchSize;
+                // write this batch of point data to gpu
+                if (resource.uploadedBatches[j] !== batchIndex) {
+                    packVec4Data(resource.backBuffers.points, points, currentBatchSize, batchIndex * batchSize, vecColumns);
+                    device.queue.writeBuffer(resource.gpuBuffers.points[j], 0, resource.backBuffers.points.buffer, 0, pointStride * currentBatchSize * 4);
+                    resource.backBuffers.uniform[0] = currentBatchSize;
+                    device.queue.writeBuffer(resource.gpuBuffers.uniform, 256 * j, resource.backBuffers.uniform.buffer, 0, 8);
+                    resource.uploadedBatches[j] = batchIndex;
+                }
+            }
+            const encoder = device.createCommandEncoder();
+            const computePass = encoder.beginComputePass();
+            computePass.setPipeline(resource.pipeline);
+            for (let j = 0; j < concurrencyBatches; j++) {
+                const batchIndex = batchStart + j;
+                if (batchIndex >= numBatches) {
+                    break;
+                }
+                const currentBatchSize = Math.min(numPoints - batchIndex * batchSize, batchSize);
+                const groups = Math.ceil(currentBatchSize / workgroupSize);
+                computePass.setBindGroup(0, resource.bindGroups[j]);
+                computePass.dispatchWorkgroups(groups);
+            }
+            computePass.end();
+            encoder.copyBufferToBuffer(resource.gpuBuffers.result, 0, resource.gpuBuffers.resultReadBack, 0, resultCount * 4);
+            device.queue.submit([encoder.finish()]);
+            await resource.gpuBuffers.resultReadBack.mapAsync(GPUMapMode.READ);
+            const mapped = resource.gpuBuffers.resultReadBack.getMappedRange();
+            labels.set(new Uint32Array(mapped, 0, resultCount), batchStart * batchSize);
+            resource.gpuBuffers.resultReadBack.unmap();
+        }
+    }
+    destroy() {
+        this.resource.gpuBuffers.uniform.destroy();
+        this.resource.gpuBuffers.centroids.destroy();
+        this.resource.gpuBuffers.centroidSq?.destroy();
+        this.resource.gpuBuffers.result.destroy();
+        this.resource.gpuBuffers.resultReadBack.destroy();
+        for (const buffer of this.resource.gpuBuffers.points) {
+            buffer.destroy();
+        }
+    }
+}

package/dist/utils/{k-means.d.ts → k-means/index.d.ts}

@@ -1,4 +1,4 @@
-export declare function kmeans(points: Float32Array[], k: number, iterations: number, device: GPUDevice): Promise<{
+export declare function kMeans(points: Float32Array[], k: number, iterations: number, device: GPUDevice): Promise<{
     centroids: Float32Array<ArrayBufferLike>[];
     labels: Uint32Array<ArrayBuffer>;
 }>;

package/dist/utils/k-means/index.js

@@ -0,0 +1,71 @@
+import { clusterAverage } from '../../native/index.js';
+import { logger } from '../index.js';
+import GpuClustering from './GpuClustering.js';
+// in the 1d case we use quantile-based initialization for better handling of skewed data
+function initializeCentroids1D(data, centroids) {
+    const n = data.length;
+    const k = centroids.length;
+    // Sort data to compute quantiles
+    const sorted = Float32Array.from(data).sort((a, b) => a - b);
+    for (let i = 0; i < k; ++i) {
+        // Place centroid at the center of its expected cluster region
+        const quantile = (2 * i + 1) / (2 * k);
+        const index = Math.min(Math.floor(quantile * n), n - 1);
+        centroids[i] = sorted[index];
+    }
+}
+// use floyd's algorithm to pick m unique random indices from 0..n-1
+function pickRandomIndices(n, m) {
+    const chosen = new Set();
+    for (let j = n - m; j < n; j++) {
+        const t = Math.floor(Math.random() * (j + 1));
+        chosen.add(chosen.has(t) ? j : t);
+    }
+    return [...chosen];
+}
+function initializeCentroids(dataTable, centroids) {
+    const indices = pickRandomIndices(dataTable[0].length, centroids[0].length);
+    for (let i = 0; i < centroids[0].length; i++) {
+        for (let j = 0; j < dataTable.length; j++) {
+            centroids[j][i] = dataTable[j][indices[i]];
+        }
+    }
+}
+// https://github.com/playcanvas/splat-transform/blob/main/src/lib/spatial/k-means.ts
+export async function kMeans(points, k, iterations, device) {
+    const numRows = points.length > 0 ? points[0].length : 0;
+    if (numRows < k) {
+        return {
+            centroids: points,
+            // use a typed array here so downstream code can rely on
+            // labels supporting subarray(), even in this early-return
+            // path used for very small datasets.
+            labels: new Uint32Array(numRows).map((_, i) => i),
+        };
+    }
+    const centroids = points.map(_ => new Float32Array(k));
+    if (points.length === 1) {
+        initializeCentroids1D(points[0], centroids[0]);
+    }
+    else {
+        initializeCentroids(points, centroids);
+    }
+    const gpuClustering = new GpuClustering(device, numRows, points.length, k);
+    const labels = new Uint32Array(numRows);
+    let converged = false;
+    let steps = 0;
+    while (!converged) {
+        logger.info(`kmeans iteration ${steps + 1}`);
+        await gpuClustering.execute(points, centroids, labels);
+        clusterAverage(points, labels, k, centroids);
+        steps++;
+        if (steps >= iterations) {
+            converged = true;
+        }
+    }
+    gpuClustering.destroy();
+    return {
+        centroids,
+        labels,
+    };
+}

package/dist/utils/webgpu.js

@@ -1,11 +1,12 @@
 /// <reference types="@webgpu/types" />
 import { createRequire } from 'node:module';
 import { logger } from './index.js';
+import { getNativePackageName } from '../native/utils.js';
 const getModule = (function () {
     let m = undefined;
     return function () {
         if (!m) {
-            m = createRequire(import.meta.url)('webgpu');
+            m = createRequire(import.meta.url)(getNativePackageName() + '/dawn.node');
             Object.assign(globalThis, m.globals);
         }
         return m;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@manycore/aholo-splat-transform",
-  "version": "1.2.9",
+  "version": "1.2.11",
   "description": "Transform & filter Gaussian splats file",
   "author": "egs",
   "repository": "https://github.com/manycoretech/aholo-viewer.git",
@@ -27,11 +27,11 @@
   "dependencies": {
     "commander": "^14.0.2",
     "fflate": "^0.8.2",
-    "tslib": "^2.8.1",
-    "webgpu": "^0.4.0"
+    "tslib": "^2.8.1"
   },
   "optionalDependencies": {
-    "@manycore/aholo-splat-transform-linux-x64-gnu": "1.2.9",
-    "@manycore/aholo-splat-transform-win32-x64-msvc": "1.2.9"
+    "@manycore/aholo-splat-transform-darwin-arm64": "1.2.11",
+    "@manycore/aholo-splat-transform-linux-x64-gnu": "1.2.11",
+    "@manycore/aholo-splat-transform-win32-x64-msvc": "1.2.11"
   }
 }

package/dist/utils/k-means.js

@@ -1,340 +0,0 @@
-import { clusterAverage } from '../native/index.js';
-import { logger } from './index.js';
-// in the 1d case we use quantile-based initialization for better handling of skewed data
-function initializeCentroids1D(data, centroids) {
-    const n = data.length;
-    const k = centroids.length;
-    // Sort data to compute quantiles
-    const sorted = Float32Array.from(data).sort((a, b) => a - b);
-    for (let i = 0; i < k; ++i) {
-        // Place centroid at the center of its expected cluster region
-        const quantile = (2 * i + 1) / (2 * k);
-        const index = Math.min(Math.floor(quantile * n), n - 1);
-        centroids[i] = sorted[index];
-    }
-}
-// use floyd's algorithm to pick m unique random indices from 0..n-1
-function pickRandomIndices(n, m) {
-    const chosen = new Set();
-    for (let j = n - m; j < n; j++) {
-        const t = Math.floor(Math.random() * (j + 1));
-        chosen.add(chosen.has(t) ? j : t);
-    }
-    return [...chosen];
-}
-function initializeCentroids(dataTable, centroids) {
-    const indices = pickRandomIndices(dataTable[0].length, centroids[0].length);
-    for (let i = 0; i < centroids[0].length; i++) {
-        for (let j = 0; j < dataTable.length; j++) {
-            centroids[j][i] = dataTable[j][indices[i]];
-        }
-    }
-}
-const chunkSize = 128;
-const workgroupSize = 64;
-function clusterWgsl(numColumns) {
-    return /* wgsl */ `
-struct Uniforms {
-    numPoints: u32,
-    numCentroids: u32
-};
-
-@group(0) @binding(0) var<uniform> uniforms: Uniforms;
-@group(0) @binding(1) var<storage, read> points: array<f32>;
-@group(0) @binding(2) var<storage, read> centroids: array<f32>;
-@group(0) @binding(3) var<storage, read_write> results: array<u32>;
-
-const numColumns = ${numColumns};   // number of columns in the points and centroids tables
-const chunkSize = ${chunkSize}u;             // must be a multiple of 64
-const workgroupSize = ${workgroupSize}u;
-var<workgroup> sharedChunk: array<f32, numColumns * chunkSize>;
-
-// calculate the squared distance between the point and centroid
-fn calcDistanceSqr(point: array<f32, numColumns>, centroid: u32) -> f32 {
-    var result = 0.0;
-
-    var ci = centroid * numColumns;
-
-    for (var i = 0u; i < numColumns; i++) {
-        let v = f32(point[i] - sharedChunk[ci+i]);
-        result += v * v;
-    }
-
-    return result;
-}
-
-@compute @workgroup_size(workgroupSize)
-fn main(
-    @builtin(local_invocation_index) local_id : u32,
-    @builtin(global_invocation_id) global_id: vec3u,
-    @builtin(num_workgroups) num_workgroups: vec3u
-) {
-    // calculate row index for this thread point
-    let pointIndex = global_id.x + global_id.y * num_workgroups.x * workgroupSize;
-
-    // copy the point data from global memory
-    var point: array<f32, numColumns>;
-    if (pointIndex < uniforms.numPoints) {
-        for (var i = 0u; i < numColumns; i++) {
-            point[i] = points[pointIndex * numColumns + i];
-        }
-    }
-
-    var mind = 1000000.0;
-    var mini = 0u;
-
-    // work through the list of centroids in shared memory chunks
-    let numChunks = u32(ceil(f32(uniforms.numCentroids) / f32(chunkSize)));
-    for (var i = 0u; i < numChunks; i++) {
-
-        // copy this thread's slice of the centroid shared chunk data
-        let dstRow = local_id * (chunkSize / workgroupSize);
-        let srcRow = min(uniforms.numCentroids, i * chunkSize + local_id * chunkSize / workgroupSize);
-        let numRows = min(uniforms.numCentroids, srcRow + chunkSize / workgroupSize) - srcRow;
-
-        var dst = dstRow * numColumns;
-        var src = srcRow * numColumns;
-
-        for (var c = 0u; c < numRows * numColumns; c++) {
-            sharedChunk[dst + c] = centroids[src + c];
-        }
-
-        // wait for all threads to finish writing their part of centroids shared memory buffer
-        workgroupBarrier();
-
-        // loop over the next chunk of centroids finding the closest
-        if (pointIndex < uniforms.numPoints) {
-            let thisChunkSize = min(chunkSize, uniforms.numCentroids - i * chunkSize);
-            for (var c = 0u; c < thisChunkSize; c++) {
-                let d = calcDistanceSqr(point, c);
-                if (d < mind) {
-                    mind = d;
-                    mini = i * chunkSize + c;
-                }
-            }
-        }
-
-        // next loop will overwrite the shared memory, so wait
-        workgroupBarrier();
-    }
-
-    if (pointIndex < uniforms.numPoints) {
-        results[pointIndex] = mini;
-    }
-}
-`;
-}
-function interleaveData(result, dataTable, numRows, rowOffset) {
-    const numColumns = dataTable.length;
-    for (let c = 0; c < numColumns; ++c) {
-        const column = dataTable[c];
-        for (let r = 0; r < numRows; ++r) {
-            result[r * numColumns + c] = column[rowOffset + r];
-        }
-    }
-}
-const MAX_CONCURRENCY_BATCHES = 10;
-class GpuClustering {
-    constructor(device, numPoints, numColumns, numCentroids) {
-        this.device = device;
-        this.numPoints = numPoints;
-        this.numColumns = numColumns;
-        this.numCentroids = numCentroids;
-        const workgroupsPerBatch = Math.min(device.limits.maxComputeWorkgroupsPerDimension, // device dispatch limit
-        Math.floor(device.limits.maxBufferSize / (numColumns * workgroupSize * 4)), // point storage limit
-        Math.ceil(numPoints / workgroupSize));
-        this.batchSize = workgroupsPerBatch * workgroupSize;
-        this.numBatches = Math.ceil(numPoints / this.batchSize);
-        this.concurrencyBatches = Math.min(MAX_CONCURRENCY_BATCHES, this.numBatches);
-        this.concurrencyRuns = Math.ceil(this.numBatches / this.concurrencyBatches);
-        const shader = device.createShaderModule({
-            code: clusterWgsl(numColumns),
-        });
-        const pipeline = device.createComputePipeline({
-            layout: 'auto',
-            compute: {
-                module: shader,
-                entryPoint: 'main',
-            },
-        });
-        const pointsBackBuffer = new Float32Array(numColumns * this.batchSize);
-        const centroidsBackBuffer = new Float32Array(numColumns * numCentroids);
-        const uniformBackBuffer = new Uint32Array([0, numCentroids]);
-        const pointsBuffers = [];
-        const centroidsBuffer = device.createBuffer({
-            size: centroidsBackBuffer.byteLength,
-            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE,
-        });
-        const uniformBuffer = device.createBuffer({
-            size: 256 * this.concurrencyBatches,
-            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM,
-        });
-        const resultBuffer = device.createBuffer({
-            size: this.concurrencyBatches * this.batchSize * 4,
-            usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE,
-        });
-        const resultReadBackBuffer = device.createBuffer({
-            size: this.concurrencyBatches * this.batchSize * 4,
-            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
-        });
-        const layout = pipeline.getBindGroupLayout(0);
-        const bindGroups = [];
-        for (let i = 0; i < this.concurrencyBatches; i++) {
-            const pointsBuffer = device.createBuffer({
-                size: pointsBackBuffer.byteLength,
-                usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE,
-            });
-            pointsBuffers.push(pointsBuffer);
-            bindGroups.push(device.createBindGroup({
-                layout,
-                entries: [
-                    {
-                        binding: 0,
-                        resource: {
-                            buffer: uniformBuffer,
-                            offset: i * 256,
-                            size: 8,
-                        },
-                    },
-                    {
-                        binding: 1,
-                        resource: pointsBuffer,
-                    },
-                    {
-                        binding: 2,
-                        resource: centroidsBuffer,
-                    },
-                    {
-                        binding: 3,
-                        resource: {
-                            buffer: resultBuffer,
-                            offset: i * this.batchSize * 4,
-                            size: this.batchSize * 4,
-                        },
-                    },
-                ],
-            }));
-        }
-        this.resource = {
-            pipeline,
-            bindGroups,
-            gpuBuffers: {
-                uniform: uniformBuffer,
-                points: pointsBuffers,
-                centroids: centroidsBuffer,
-                result: resultBuffer,
-                resultReadBack: resultReadBackBuffer,
-            },
-            backBuffers: {
-                uniform: uniformBackBuffer,
-                points: pointsBackBuffer,
-                centroids: centroidsBackBuffer,
-            },
-            uploadedBatches: [],
-        };
-        logger.info(`GPU k-means kernel bootstrapped with batch ${workgroupsPerBatch}*${workgroupSize}*${this.numBatches}, concurrency: ${this.concurrencyBatches}, runs: ${this.concurrencyRuns}`);
-    }
-    async execute(points, centroids, labels) {
-        const { device, numPoints, numColumns, numCentroids, numBatches, batchSize, resource, concurrencyBatches, concurrencyRuns, } = this;
-        // upload centroid data to gpu
-        interleaveData(resource.backBuffers.centroids, centroids, numCentroids, 0);
-        device.queue.writeBuffer(resource.gpuBuffers.centroids, 0, resource.backBuffers.centroids.buffer);
-        for (let i = 0; i < concurrencyRuns; i++) {
-            const batchStart = i * concurrencyBatches;
-            let resultCount = 0;
-            for (let j = 0; j < concurrencyBatches; j++) {
-                const batchIndex = batchStart + j;
-                if (batchIndex >= numBatches) {
-                    break;
-                }
-                const currentBatchSize = Math.min(numPoints - batchIndex * batchSize, batchSize);
-                resultCount += currentBatchSize;
-                // write this batch of point data to gpu
-                if (resource.uploadedBatches[j] !== batchIndex) {
-                    interleaveData(resource.backBuffers.points, points, currentBatchSize, batchIndex * batchSize);
-                    device.queue.writeBuffer(resource.gpuBuffers.points[j], 0, resource.backBuffers.points.buffer, 0, numColumns * currentBatchSize * 4);
-                    resource.backBuffers.uniform[0] = currentBatchSize;
-                    device.queue.writeBuffer(resource.gpuBuffers.uniform, 256 * j, resource.backBuffers.uniform.buffer, 0, 8);
-                    resource.uploadedBatches[j] = batchIndex;
-                }
-            }
-            const encoder = device.createCommandEncoder();
-            const computePass = encoder.beginComputePass();
-            computePass.setPipeline(resource.pipeline);
-            for (let j = 0; j < concurrencyBatches; j++) {
-                const batchIndex = batchStart + j;
-                if (batchIndex >= numBatches) {
-                    break;
-                }
-                const currentBatchSize = Math.min(numPoints - batchIndex * batchSize, batchSize);
-                const groups = Math.ceil(currentBatchSize / workgroupSize);
-                computePass.setBindGroup(0, resource.bindGroups[j]);
-                computePass.dispatchWorkgroups(groups);
-            }
-            computePass.end();
-            encoder.copyBufferToBuffer(resource.gpuBuffers.result, 0, resource.gpuBuffers.resultReadBack, 0, resultCount * 4);
-            device.queue.submit([encoder.finish()]);
-            await resource.gpuBuffers.resultReadBack.mapAsync(GPUMapMode.READ);
-            const mapped = resource.gpuBuffers.resultReadBack.getMappedRange();
-            labels.set(new Uint32Array(mapped, 0, resultCount), batchStart * batchSize);
-            resource.gpuBuffers.resultReadBack.unmap();
-        }
-    }
-    destroy() {
-        this.resource.gpuBuffers.uniform.destroy();
-        this.resource.gpuBuffers.centroids.destroy();
-        this.resource.gpuBuffers.result.destroy();
-        this.resource.gpuBuffers.resultReadBack.destroy();
-        for (const buffer of this.resource.gpuBuffers.points) {
-            buffer.destroy();
-        }
-    }
-}
-function groupLabels(labels, k) {
-    const clusters = [];
-    for (let i = 0; i < k; ++i) {
-        clusters[i] = [];
-    }
-    for (let i = 0; i < labels.length; ++i) {
-        clusters[labels[i]].push(i);
-    }
-    return clusters.map(c => new Uint32Array(c));
-}
-// https://github.com/playcanvas/splat-transform/blob/main/src/lib/spatial/k-means.ts
-export async function kmeans(points, k, iterations, device) {
-    const numRows = points.length > 0 ? points[0].length : 0;
-    if (numRows < k) {
-        return {
-            centroids: points,
-            // use a typed array here so downstream code can rely on
-            // labels supporting subarray(), even in this early-return
-            // path used for very small datasets.
-            labels: new Uint32Array(numRows).map((_, i) => i),
-        };
-    }
-    const centroids = points.map(_ => new Float32Array(k));
-    if (points.length === 1) {
-        initializeCentroids1D(points[0], centroids[0]);
-    }
-    else {
-        initializeCentroids(points, centroids);
-    }
-    const gpuClustering = new GpuClustering(device, numRows, points.length, k);
-    const labels = new Uint32Array(numRows);
-    let converged = false;
-    let steps = 0;
-    while (!converged) {
-        logger.info(`kmeans iteration ${steps + 1}`);
-        await gpuClustering.execute(points, centroids, labels);
-        clusterAverage(points, groupLabels(labels, k), centroids);
-        steps++;
-        if (steps >= iterations) {
-            converged = true;
-        }
-    }
-    gpuClustering.destroy();
-    return {
-        centroids,
-        labels,
-    };
-}