Use WebGPU for color quantization

typescript/package.json

@@ -45,15 +45,23 @@
   "scripts": {
     "build": "wireit",
     "build:ts": "wireit",
-    "prepublishOnly": "npm run build"
+    "prepublishOnly": "npm run build",
+    "test": "npm run build && jasmine-browser-runner runSpecs",
+    "test:serve": "npm run build && jasmine-browser-runner serve"
   },
   "wireit": {
     "build": {
-      "dependencies": ["build:ts"]
+      "dependencies": [
+        "build:ts"
+      ]
     },
     "build:ts": {
       "command": "tsc --pretty",
-      "files": ["tsconfig.json", "**/*.ts", "!**/*.d.ts"],
+      "files": [
+        "tsconfig.json",
+        "**/*.ts",
+        "!**/*.d.ts"
+      ],
       "output": [
         ".tsbuildinfo",
         "**/*.js",
@@ -66,10 +74,13 @@
     }
   },
   "devDependencies": {
-    "@types/jasmine": "^3.10.3",
+    "@types/jasmine": "^3.10.18",
     "@types/node": "^18.7.17",
+    "@webgpu/types": "^0.1.54",
     "jasmine": "^4.0.2",
-    "typescript": "^4.5.5",
+    "jasmine-browser-runner": "^2.5.0",
+    "jasmine-core": "^5.6.0",
+    "typescript": "^4.9.5",
     "wireit": "^0.9.5"
   }
 }

typescript/palettes/palettes_test.ts

@@ -17,8 +17,6 @@
 
 import 'jasmine';
 
-import {Hct} from '../hct/hct.js';
-
 import {CorePalette} from './core_palette.js';
 import {TonalPalette} from './tonal_palette.js';
 

typescript/quantize-webgpu/kmeans/index.ts

@@ -0,0 +1,84 @@
+import { setupCompute } from './pipelines/compute.js';
+
+export async function extractDominantColorsKMeansGPU(
+    device: GPUDevice,
+    pixels: number[],
+    K: number,
+    initialCentroidsBuffer: GPUBuffer | null = null
+): Promise<GPUBuffer> {
+    const MAX_ITERATIONS = 256;
+    const CONVERGENCE_EPS = 0.01;
+    const CONVERGENCE_CHECK = 8;
+
+    const {
+        colorCount,
+        centroidsBuffer,
+        centroidsDeltaBuffer,
+        assignPipeline,
+        updatePipeline,
+        computeBindGroup
+    } = await setupCompute(device, pixels, K);
+
+    const stagingCentroidsDeltaBuffer = device.createBuffer({
+        label: 'centroids-delta-staging',
+        size: K * Float32Array.BYTES_PER_ELEMENT,
+        usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
+    });
+
+    let encoder = device.createCommandEncoder();
+
+    if (initialCentroidsBuffer) {
+        encoder.copyBufferToBuffer(
+            initialCentroidsBuffer, 0,
+            centroidsBuffer, 0,
+            3 * K * Float32Array.BYTES_PER_ELEMENT
+        );
+    } else {
+        const centroids = new Float32Array(3 * K);
+        for (let i = 0; i < 3 * K; i++) {
+            centroids[i] = Math.random();
+        }
+        device.queue.writeBuffer(centroidsBuffer, 0, centroids);
+    }
+
+    for (let i = 0; i < MAX_ITERATIONS; i++) {
+        const assignPass = encoder.beginComputePass();
+        assignPass.setPipeline(assignPipeline);
+        assignPass.setBindGroup(0, computeBindGroup);
+        assignPass.dispatchWorkgroups(Math.ceil(colorCount / 256));
+        assignPass.end();
+
+        const updatePass = encoder.beginComputePass();
+        updatePass.setPipeline(updatePipeline);
+        updatePass.setBindGroup(0, computeBindGroup);
+        updatePass.dispatchWorkgroups(Math.ceil(K / 16));
+        updatePass.end();
+
+        if (i !== 0 && i % CONVERGENCE_CHECK === 0) {
+            encoder.copyBufferToBuffer(
+                centroidsDeltaBuffer, 0,
+                stagingCentroidsDeltaBuffer, 0,
+                K * Float32Array.BYTES_PER_ELEMENT
+            );
+
+            const commandBuffer = encoder.finish();
+            device.queue.submit([commandBuffer]);
+            encoder = device.createCommandEncoder();
+
+            await stagingCentroidsDeltaBuffer.mapAsync(GPUMapMode.READ, 0, K * Float32Array.BYTES_PER_ELEMENT);
+            const centroidsDeltaData = new Float32Array(stagingCentroidsDeltaBuffer.getMappedRange());
+            const deltaSum = centroidsDeltaData.reduce((acc, val) => acc + val, 0);
+            stagingCentroidsDeltaBuffer.unmap();
+            if (deltaSum < CONVERGENCE_EPS) {
+                console.log(`Convergence reached at iteration ${i}`);
+                break;
+            }
+        }
+    }
+
+    device.queue.submit([encoder.finish()]);
+    await device.queue.onSubmittedWorkDone();
+
+    return centroidsBuffer;
+}
+

typescript/quantize-webgpu/kmeans/pipelines/compute.ts

@@ -0,0 +1,145 @@
+import * as colorUtils from '../../../utils/color_utils.js';
+
+interface ComputeResult {
+    colorCount: number;
+    centroidsBuffer: GPUBuffer;
+    centroidsDeltaBuffer: GPUBuffer;
+    assignPipeline: GPUComputePipeline;
+    updatePipeline: GPUComputePipeline;
+    computeBindGroup: GPUBindGroup;
+}
+
+export async function setupCompute(
+    device: GPUDevice,
+    pixels: number[],
+    K: number
+): Promise<ComputeResult> {
+    const colorHistogram = new Map<number, number>();
+    for (let i = 0; i < pixels.length; i++) {
+        const pixel = pixels[i];
+        const r = colorUtils.redFromArgb(pixel);
+        const g = colorUtils.greenFromArgb(pixel);
+        const b = colorUtils.blueFromArgb(pixel);
+        const key = (r << 16) | (g << 8) | b;
+        colorHistogram.set(key, (colorHistogram.get(key) ?? 0) + 1);
+    }
+
+    const colorCount = colorHistogram.size;
+    const histogramArray = new Float32Array(colorCount * 4);
+    let i = 0;
+    for (const [key, count] of colorHistogram) {
+        const r = (key >> 16) & 0xFF;
+        const g = (key >> 8) & 0xFF;
+        const b = key & 0xFF;
+
+        histogramArray[i * 4] = r / 255;
+        histogramArray[i * 4 + 1] = g / 255;
+        histogramArray[i * 4 + 2] = b / 255;
+        histogramArray[i * 4 + 3] = count;
+        i++;
+    }
+
+    const countsUniformBuffer = device.createBuffer({
+        size: 2 * Uint32Array.BYTES_PER_ELEMENT,
+        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
+    });
+    device.queue.writeBuffer(countsUniformBuffer, 0, new Uint32Array([K, colorCount]));
+
+    const histogramBuffer = device.createBuffer({
+        label: 'histogram-compute',
+        size: histogramArray.byteLength,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
+    });
+    device.queue.writeBuffer(histogramBuffer, 0, histogramArray);
+
+    const centroidsBuffer = device.createBuffer({
+        label: 'centroids-compute',
+        size: 3 * K * Float32Array.BYTES_PER_ELEMENT,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC
+    });
+
+    const clustersBuffer = device.createBuffer({
+        label: 'clusters-compute',
+        size: colorCount * Uint32Array.BYTES_PER_ELEMENT,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
+    });
+
+    const centroidsDeltaBuffer = device.createBuffer({
+        label: 'centroids-delta-compute',
+        size: K * Float32Array.BYTES_PER_ELEMENT,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
+    });
+
+    const kUniformBuffer = device.createBuffer({
+        size: Uint32Array.BYTES_PER_ELEMENT,
+        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
+    });
+    device.queue.writeBuffer(kUniformBuffer, 0, new Uint32Array([K]));
+
+    const assignModule = device.createShaderModule({
+        code: await fetch(new URL('../shaders/assign.wgsl', import.meta.url).toString()).then(res => res.text())
+    });
+    const updateModule = device.createShaderModule({
+        code: await fetch(new URL('../shaders/update.wgsl', import.meta.url).toString()).then(res => res.text())
+    });
+
+    const computeBindGroupLayout = device.createBindGroupLayout({
+        entries: [{
+            binding: 0,
+            visibility: GPUShaderStage.COMPUTE,
+            buffer: { type: 'read-only-storage' }
+        },
+        {
+            binding: 1,
+            visibility: GPUShaderStage.COMPUTE,
+            buffer: { type: 'uniform' }
+        },
+        {
+            binding: 2,
+            visibility: GPUShaderStage.COMPUTE,
+            buffer: { type: 'storage' }
+        },
+        {
+            binding: 3,
+            visibility: GPUShaderStage.COMPUTE,
+            buffer: { type: 'storage' }
+        },
+        {
+            binding: 4,
+            visibility: GPUShaderStage.COMPUTE,
+            buffer: { type: 'storage' }
+        }]
+    });
+
+    const computeBindGroup = device.createBindGroup({
+        layout: computeBindGroupLayout,
+        entries: [
+            { binding: 0, resource: { buffer: histogramBuffer } },
+            { binding: 1, resource: { buffer: countsUniformBuffer } },
+            { binding: 2, resource: { buffer: centroidsBuffer } },
+            { binding: 3, resource: { buffer: clustersBuffer } },
+            { binding: 4, resource: { buffer: centroidsDeltaBuffer } }
+        ]
+    });
+    const computePipelineLayout = device.createPipelineLayout({
+        bindGroupLayouts: [computeBindGroupLayout]
+    });
+
+    const updatePipeline = device.createComputePipeline({
+        layout: computePipelineLayout,
+        compute: { module: updateModule }
+    });
+    const assignPipeline = device.createComputePipeline({
+        layout: computePipelineLayout,
+        compute: { module: assignModule }
+    });
+
+    return {
+        colorCount,
+        centroidsBuffer,
+        centroidsDeltaBuffer,
+        assignPipeline,
+        updatePipeline,
+        computeBindGroup
+    };
+} 

typescript/quantize-webgpu/kmeans/shaders/assign.wgsl

@@ -0,0 +1,40 @@
+struct Counts {
+    centroids: u32,
+    colors: u32
+};
+
+@group(0) @binding(0) var<storage> histogram: array<f32>;
+@group(0) @binding(1) var<uniform> counts: Counts;
+@group(0) @binding(2) var<storage, read_write> centroids: array<f32>;
+@group(0) @binding(3) var<storage, read_write> clusters: array<u32>;
+
+fn dist(a: vec3f, b: vec3f) -> f32 {
+    return pow((a.x - b.x), 2) + pow((a.y - b.y), 2) + pow((a.z - b.z), 2);
+}
+
+@compute @workgroup_size(256)
+fn cs(@builtin(global_invocation_id) id: vec3u) {
+    if (id.x >= counts.colors) {
+        return;
+    }
+
+    let pos = vec3f(histogram[id.x * 4], histogram[id.x * 4 + 1], histogram[id.x * 4 + 2]);
+    let count = histogram[id.x * 4 + 3];
+
+    var min_dist = -1.;
+    var closest = 0u;
+
+    for (var i = 0u; i < counts.centroids; i++) {
+        let centroid = vec3f(centroids[3*i], centroids[3*i + 1], centroids[3*i + 2]);
+        if (centroid.x == -1.0 || centroid.y == -1.0 || centroid.z == -1.0) {
+            continue;
+        }
+        let d = dist(pos, centroid);
+        if (min_dist == -1 || d < min_dist){
+            closest = i;
+            min_dist = d;
+        }
+    }
+
+    clusters[id.x] = closest;
+}

typescript/quantize-webgpu/kmeans/shaders/update.wgsl

@@ -0,0 +1,52 @@
+struct Counts {
+    centroids: u32,
+    colors: u32
+};
+
+@group(0) @binding(0) var<storage> histogram: array<f32>;
+@group(0) @binding(1) var<uniform> counts: Counts;
+@group(0) @binding(2) var<storage, read_write> centroids: array<f32>;
+@group(0) @binding(3) var<storage, read_write> clusters: array<u32>;
+@group(0) @binding(4) var<storage, read_write> centroids_delta: array<f32>;
+
+fn dist(a: vec3f, b: vec3f) -> f32 {
+    return sqrt(pow((a.x - b.x), 2) + pow((a.y - b.y), 2) + pow((a.z - b.z), 2));
+}
+
+@compute @workgroup_size(16)
+fn cs(@builtin(global_invocation_id) id: vec3u) {
+    let centroid = id.x;
+
+    if (centroid >= counts.centroids) {
+        return;
+    }
+
+    var sum = vec3f(0);
+    var count = 0u;
+
+    for (var i = 0u; i < counts.colors; i++) {
+        if (clusters[i] == centroid) {
+            let pixel = vec3f(histogram[i * 4], histogram[i * 4 + 1], histogram[i * 4 + 2]);
+            let pixel_count = u32(histogram[i * 4 + 3]);
+            sum += pixel * f32(pixel_count);
+            count += pixel_count;
+        }
+    }
+
+    if (count > 0u) {
+        let old_pos = vec3f(centroids[3*centroid], centroids[3*centroid + 1], centroids[3*centroid + 2]);
+        let new_pos = sum / f32(count);
+
+        centroids[3*centroid] = new_pos.x;
+        centroids[3*centroid + 1] = new_pos.y;
+        centroids[3*centroid + 2] = new_pos.z;
+
+        let d = dist(old_pos, new_pos);
+        centroids_delta[centroid] = d;
+    } else {
+        centroids[3*centroid] = -1.0;
+        centroids[3*centroid + 1] = -1.0;
+        centroids[3*centroid + 2] = -1.0;
+        centroids_delta[centroid] = 0.0;
+    }
+}