2d:null 2e:["$","title",null,{"children":"WebGPU Experiments | Quantized compute | Mike Tahani | Observable"}] 33:T15fa,{ const { context, adapter, device } = await gpu.init(1, 1) // ================================================== // shaders // ================================================== // this shader uses atomics which only work with u32 and i32. to // get the extent of float values, we have to calculate them using // int types, output them in a buffer, then pass that buffer // to another shader to convert it to f32s in a subsequent shader pass. const extentShaderModule = device.createShaderModule({ label: 'extent quantization compute module', code: ` @group(0) @binding(0) var data: array; @group(0) @binding(1) var quantized_extent: array>; const QUANTIZE_FACTOR = 32768.0; fn updateExtent(index: u32, value: f32) { let quantizedValue = i32(value * QUANTIZE_FACTOR); if (index == 0) { atomicStore(&quantized_extent[0], quantizedValue); atomicStore(&quantized_extent[1], quantizedValue); } atomicMin(&quantized_extent[0], quantizedValue); atomicMax(&quantized_extent[1], quantizedValue); } @compute @workgroup_size(1) fn cs( @builtin(global_invocation_id) id: vec3 ) { let i = id.x; updateExtent(i, data[i]); } `, }) const extentDequantShaderModule = device.createShaderModule({ label: 'extent dequantization compute module', code: ` @group(0) @binding(0) var quantized_extent: array; @group(0) @binding(1) var extent: array; const DEQUANTIZE_FACTOR = 1.0 / 32768.0; fn getExtent(index: u32) -> vec2f { // Loads the quantized normal values into a vector and dequantizes them. return vec2f(f32(quantized_extent[0]), f32(quantized_extent[1])) * DEQUANTIZE_FACTOR; } @compute @workgroup_size(1) fn cs( @builtin(global_invocation_id) id: vec3 ) { let i = id.x; let ext = getExtent(i); extent = array(ext.x, ext.y); } `, }) // ================================================== // pipelines // ================================================== const extentPipeline = device.createComputePipeline({ label: 'extent quant compute pipeline', layout: 'auto', compute: { module: extentShaderModule, entryPoint: 'cs', }, }) const extentDequantPipeline = device.createComputePipeline({ label: 'extent dequant compute pipeline', layout: 'auto', compute: { module: extentDequantShaderModule, entryPoint: 'cs', }, }) // ================================================== // buffers // ================================================== const dataBuffer = device.createBuffer({ label: 'data buffer', size: input.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(dataBuffer, 0, input); // this buffer contains quantized u32 values that represent f32s. // we will convert them to f32s in a subsequent shader pass. const extentQuantBuffer = device.createBuffer({ size: 2 * 4, // min and max * 4 bytes per (u32), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST, }) // this contains f32 values after they've been converted in the // dequantization shader const extentDequantBuffer = device.createBuffer({ size: 2 * 4, // min and max * 4 bytes per (f32), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST, }) // mapped extent results buffer that lets us read gpu-derived values // on the cpu const resultBuffer = device.createBuffer({ label: 'extent dequant result buffer', size: extentDequantBuffer.size, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, }) // ================================================== // bind groups // ================================================== const extentBindGroup = device.createBindGroup({ label: 'extent quant bind group', layout: extentPipeline.getBindGroupLayout(0), entries: [ { binding: 0, resource: { buffer: dataBuffer } }, { binding: 1, resource: { buffer: extentQuantBuffer } }, ], }) const extentDequantBindGroup = device.createBindGroup({ label: 'extent dequant bind group', layout: extentDequantPipeline.getBindGroupLayout(0), entries: [ { binding: 0, resource: { buffer: extentQuantBuffer } }, { binding: 1, resource: { buffer: extentDequantBuffer } }, ], }) // ================================================== // render // ================================================== const encoder = device.createCommandEncoder({ label: 'encoder' }) const pass = encoder.beginComputePass({ label: 'compute pass' }) pass.setPipeline(extentPipeline) pass.setBindGroup(0, extentBindGroup) pass.dispatchWorkgroups(input.length) pass.setPipeline(extentDequantPipeline) pass.setBindGroup(0, extentDequantBindGroup) pass.dispatchWorkgroups(input.length) pass.end() encoder.copyBufferToBuffer(extentDequantBuffer, 0, resultBuffer, 0, resultBuffer.size) device.queue.submit([encoder.finish()]) // ================================================== // gpu -> cpu // ================================================== const output = await gpu.readBuffer(resultBuffer) return JSON.stringify({ input, min_and_max_output_from_compute_shader: output }, null, 2) }34:[] 31:{"id":"ad1d28f6637d0898","slug":"webgpu-experiments-quantized-compute","title":"WebGPU Experiments | Quantized compute","thumbnail":"15e099ea1ecc3610e6c822ace3bd27c2adfd29f5225c29c0cc0d2554618e16f4","access_level":"public","create_time":"$D2024-04-01T22:30:31.693Z","update_time":"$D2024-11-10T19:36:04.620Z","view_time":null,"first_public_time":"$D2024-04-01T23:36:42.622Z","has_unpublished_edits":null,"fork":{"id":"49b3f6114b11d8e3","slug":"webgpu-experiments-multi-pipeline","title":"WebGPU Experiments | Multi-pipeline","owner":{"id":"f258f685ba4dff9a","login":"miketahani","name":"Mike Tahani"}},"stars":0,"starred":false,"owner":{"type":"team","id":"f258f685ba4dff9a","login":"miketahani","name":"Mike Tahani","avatar_url":"https://avatars-next.observableusercontent.com/eaac9d07c89480b6348b8d4e15a8ada6cd440e3b8e1ba1b1a330f055aa368db8","tier":"starter_2024"},"version":329,"publish_version":null,"role":null,"authored":false,"authors":[{"type":"individual","id":"02c6b10033eb7a61","login":"miketahani","name":"Mike Tahani","avatar_url":"https://avatars-next.observableusercontent.com/8341f220c6016b9edcb3cf852d6526ce60b354ab1ec9cd907f76cfc54425f28e","tier":null}],"members":[{"type":"individual","id":"02c6b10033eb7a61","login":"miketahani","name":"Mike Tahani","avatar_url":"https://avatars-next.observableusercontent.com/8341f220c6016b9edcb3cf852d6526ce60b354ab1ec9cd907f76cfc54425f28e","tier":null,"role":"editor"},{"type":"team","id":"f258f685ba4dff9a","login":"miketahani","name":"Mike Tahani","avatar_url":"https://avatars-next.observableusercontent.com/eaac9d07c89480b6348b8d4e15a8ada6cd440e3b8e1ba1b1a330f055aa368db8","tier":"starter_2024","role":"editor"}],"comments":[],"custom_thumbnail":false,"trashed":false,"lapsed":null,"subscribed":null,"forks":2,"collection":{"id":"b12e2b165ab97944","slug":"webgpu-basics","title":"WebGPU Basics","owner":{"id":"f258f685ba4dff9a","login":"miketahani","name":"Mike Tahani"}},"type":"notebook","body":{"title":"WebGPU Experiments | Quantized compute","cells":[{"id":0,"value":"# WebGPU Experiments | Quantized compute\n\nUses atomics to prevent parallelization issues, but atomics only work with `u32` and `i32`, so we store float types as ints (quantization) and convert them back to floats later (dequantization).\n\nhttps://toji.dev/webgpu-best-practices/compute-vertex-data#synchronizing-data-access-with-atomics\n\n⚠️ Many inefficiencies here since this is only for illustrative/learning purposes. Use Chrome (https://caniuse.com/webgpu)","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","output":"","database":"$undefined"},{"id":254,"value":"viewof input = Inputs.button(\n \"New Inputs\", {\n value: Float32Array.from({ length: 10 }).fill(0).map(() => Math.random()), \n reduce: () => Float32Array.from({ length: 10 }).fill(0).map(() => Math.random())\n }\n)","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","output":"$undefined","database":"$undefined"},{"id":269,"value":"Below, see the input values and the outputted min and max of those values from the shader pipeline","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","output":"","database":"$undefined"},{"id":19,"value":"$33","pinned":true,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","output":"$undefined","database":"$undefined"},{"id":289,"value":"gpu = ({\n init: async (width = 512, height = 512) => {\n const canvas = document.createElement('canvas');\n canvas.width = width;\n canvas.height = height;\n const context = canvas.getContext('webgpu');\n\n const adapter = await navigator.gpu.requestAdapter();\n const device = await adapter.requestDevice();\n const format = navigator.gpu.getPreferredCanvasFormat();\n context.configure({ device, format });\n \n return { context, adapter, device }\n },\n readBuffer: async (buffer, ArrayType = Float32Array) => {\n await buffer.mapAsync(GPUMapMode.READ)\n const result = new ArrayType(buffer.getMappedRange().slice())\n buffer.unmap()\n return result\n }\n})","pinned":true,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","output":"$undefined","database":"$undefined"}],"files":"$W34","stdlib":"1","titling":"auto"}} 32:"$31:body"