2d:null 2e:["$","title",null,{"children":"WebGPU Experiments | Compute Data Pipeline | Mike Tahani | Observable"}] 33:T1c4b,{ const { context, adapter, device } = await gpu.init(1, 1) // ================================================== // shaders // ================================================== // for now, this just returns a simple transformation. in the future, // it could mimic the d3-scale api. const scaleShaderModule = device.createShaderModule({ label: 'scale shader module', code: ` @group(0) @binding(0) var data: array; // should come up with a shader architecture that allows easy composition fn scale(value: f32) -> f32 { return sin(value); } @compute @workgroup_size(1) fn cs( @builtin(global_invocation_id) id: vec3 ) { let i = id.x; data[i] = scale(data[i]); } ` }) // 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 scalePipeline = device.createComputePipeline({ label: 'scale compute pipeline', layout: 'auto', compute: { module: scaleShaderModule, entryPoint: 'cs', }, }) 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 lets us read gpu-derived values // on the cpu const extentResultBuffer = device.createBuffer({ label: 'extent dequant result buffer', size: extentDequantBuffer.size, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, }) const dataResultBuffer = device.createBuffer({ label: 'data result buffer', size: dataBuffer.size, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, }) // ================================================== // bind groups // ================================================== const scaleBindGroup = device.createBindGroup({ label: 'scale bind group', layout: scalePipeline.getBindGroupLayout(0), entries: [ { binding: 0, resource: { buffer: dataBuffer } }, ], }) 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(scalePipeline) pass.setBindGroup(0, scaleBindGroup) pass.dispatchWorkgroups(input.length) 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(dataBuffer, 0, dataResultBuffer, 0, dataResultBuffer.size) encoder.copyBufferToBuffer(extentDequantBuffer, 0, extentResultBuffer, 0, extentResultBuffer.size) device.queue.submit([encoder.finish()]) // ================================================== // gpu -> cpu // ================================================== const scaled = await gpu.readBuffer(dataResultBuffer) const extent = await gpu.readBuffer(extentResultBuffer) return htl.html`${[ Inputs.table([['inputs', ...input], ['scaled', ...scaled]]), Inputs.table([{ '(scaled) min': extent[0], '(scaled) max': extent[1] }], { columns: ['(scaled) min', '(scaled) max'] }) ]}` }34:[] 31:{"id":"f4f0cf0dd0c89321","slug":"webgpu-experiments-compute-data-pipeline","title":"WebGPU Experiments | Compute Data Pipeline","thumbnail":"a569904e52c3216830f9539d6eeb7d1b026599c0a35a10e6a01ea6fad95d9b87","access_level":"public","create_time":"$D2024-04-02T12:28:31.406Z","update_time":"$D2024-06-12T02:13:59.485Z","view_time":null,"first_public_time":"$D2024-04-02T14:16:31.549Z","has_unpublished_edits":null,"fork":{"id":"ad1d28f6637d0898","slug":"webgpu-experiments-quantized-compute","title":"WebGPU Experiments | Quantized compute","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":350,"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":1,"collection":{"id":"b12e2b165ab97944","slug":"webgpu-basics","title":"WebGPU Basics","owner":{"id":"f258f685ba4dff9a","login":"miketahani","name":"Mike Tahani"}},"type":"notebook","body":{"title":"WebGPU Experiments | Compute Data Pipeline","cells":[{"id":0,"value":"# WebGPU Experiments | Compute Data Pipeline\n\nCompute pipeline for working with data, with the goal of creating a library similar to D3 on the GPU... like a G3.js of sorts\n\n⚠️ Unoptimized! 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() - 0.5) * Math.PI), \n reduce: () => Float32Array.from({ length: 10 }).fill(0).map(() => (Math.random() - 0.5) * Math.PI)\n }\n)","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","output":"$undefined","database":"$undefined"},{"id":269,"value":"Below, see the entire pipeline: `inputs` which are fed to the scaling shader (`scaled`), which are then fed into the extent shader which calculates the bounds of the new data over two shader passes.","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":7,"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"