Last active
April 3, 2022 13:08
-
-
Save bellbind/8318c595fb1464bfefa2827b0b55667a to your computer and use it in GitHub Desktop.
[WebGPU] sampling from bitmap image by compute shaders (for Chrome-100)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
<!doctype html> | |
<html> | |
<head> | |
<!-- IMPORTANT: The current Chrome requires some origin-trial token in <meta>. | |
To register origins at the last "WebGPU REGISTER" in https://developer.chrome.com/origintrials/ | |
This token is for a Web Origin "http://localhost:8000" (maybe expired at Mar 31, 2022) | |
--> | |
<meta http-equiv="origin-trial" | |
content="AkIL+/THBoi1QEsWbX5SOuMpL6+KGAXKrZE5Bz6yHTuijzvKz2MznuLqE+MH4YSqRi/v1fDK/6JyFzgibTTeNAsAAABJeyJvcmlnaW4iOiJodHRwOi8vbG9jYWxob3N0OjgwMDAiLCJmZWF0dXJlIjoiV2ViR1BVIiwiZXhwaXJ5IjoxNjUyODMxOTk5fQ==" /> | |
<meta http-equiv="origin-trial" | |
content="Akv07qcAop5MFaZYxJtHHjUuM8eV3GpbHkTeuhZo/4wsNjYnQ7GSGJyo7hRVZvpvyjYwilbJ8KbFVchI4O1DpA0AAABQeyJvcmlnaW4iOiJodHRwczovL2dpc3QuZ2l0aGFjay5jb206NDQzIiwiZmVhdHVyZSI6IldlYkdQVSIsImV4cGlyeSI6MTY1MjgzMTk5OX0=" /> | |
<script src="./main.js" type="module"></script> | |
<style>@media(prefers-color-scheme: dark){:root {color-scheme: dark;}}</style> | |
<link rel="icon" href="data:image/x-icon;," /> | |
</head> | |
<body> | |
<h1>(Notice: The origin-trial token in this page will be expired at May 15, 2022)</h1> | |
<canvas style="width: 80vmin; height: 80vmin; border: solid;" id="canvas"></canvas> | |
</body> | |
</html> |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// Image texture example for WebGPU API for Chrome-100: https://www.w3.org/TR/webgpu/ | |
const adapter = await navigator.gpu.requestAdapter(); | |
const device = await adapter.requestDevice(); | |
// arranged from: https://github.com/austinEng/webgpu-samples/tree/main/src/sample/particles | |
// prepare image | |
const bitmapSize = 512; | |
const emoji = "🎎"; | |
//const emoji = "🌄"; | |
console.log(emoji); | |
const img = new Image(); | |
img.src = "data:image/svg+xml," + encodeURIComponent(` | |
<svg xmlns="http://www.w3.org/2000/svg" width="${bitmapSize}" height="${bitmapSize}"> | |
<text x="0" y="${bitmapSize - bitmapSize / 8}" font-size="${bitmapSize}">${emoji}</text> | |
</svg> | |
`); | |
await img.decode(); | |
const bitmap = await createImageBitmap(img); | |
//[A. generate mipmap] copy bitmap alpha into prob mipmap | |
// - bitmap alpha value as probability | |
// - mipmap levels from original texture [width, height] to [1, 1]; step by [1/2, 1/2] e.g. mip level count of 8x8 texture = 4 | |
// - each smaller mipmap (r, g, b, a) as | |
// 4x larger mipmap (top-left, top-left + top-right, top-left + top-right + bottom-left, top-left + top-right + bottom-left + bottom-right) probabilities | |
// (0 <= prob <= 1) | |
// 1. copy alpha value in bitmap into the initial buffer | |
const ws = 256; | |
const copyAlphaWgsl = ` | |
struct Buf { | |
alphas: array<f32>; | |
}; | |
@group(0) @binding(0) var<storage, write> alpha_out: Buf; | |
@group(0) @binding(1) var bitmap: texture_2d<f32>; | |
let ws = ${ws}; | |
let width = ${bitmapSize}u; | |
@stage(compute) @workgroup_size(ws) fn copy_alpha(@builtin(global_invocation_id) giid: vec3<u32>) { | |
alpha_out.alphas[giid.y * width + giid.x] = textureLoad(bitmap, vec2<i32>(giid.xy), 0).a; | |
} | |
`; | |
const copyAlphaShader = device.createShaderModule({code: copyAlphaWgsl}); | |
// 2. shrink half&half from bufferA to bufferB, then write shrinked value into texture mipmap | |
const shrinkMipmapWgsl = ` | |
struct Buf { | |
alphas: array<f32>; | |
}; | |
@group(0) @binding(0) var<storage, read> alpha_in: Buf; | |
@group(0) @binding(1) var<storage, write> alpha_out: Buf; | |
@group(0) @binding(2) var mipmap: texture_storage_2d<rgba8unorm, write>; | |
let ws = ${ws}; | |
let width = ${bitmapSize}u; | |
@stage(compute) @workgroup_size(ws) fn shrink_mipmap(@builtin(global_invocation_id) giid: vec3<u32>) { | |
if (!all(giid.xy < vec2<u32>(textureDimensions(mipmap)))) {return;} | |
let dst = giid.y * width + giid.x; | |
let offs = 2u * giid.y * width + 2u * giid.x; | |
let tl = alpha_in.alphas[offs]; | |
let tr = alpha_in.alphas[offs + 1u]; | |
let bl = alpha_in.alphas[offs + width]; | |
let br = alpha_in.alphas[offs + width + 1u]; | |
let total = tl + tr + bl + br; | |
alpha_out.alphas[dst] = total / 4.0; | |
if (total == 0.0) { | |
textureStore(mipmap, vec2<i32>(giid.xy), vec4<f32>(0.0, 0.0, 0.0, 0.0)); | |
} else { | |
textureStore(mipmap, vec2<i32>(giid.xy), vec4<f32>(tl, tl + tr, tl + tr + bl, tl + tr + bl + br) / total); | |
} | |
} | |
`; | |
const shrinkMipmapShader = device.createShaderModule({code: shrinkMipmapWgsl}); | |
// texture and buffers | |
const mipLevelCount = Math.log2(bitmapSize) + 1; | |
const texture = device.createTexture({ | |
size: [bitmapSize, bitmapSize], mipLevelCount, format: "rgba8unorm", | |
usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.STORAGE_BINDING | GPUTextureUsage.COPY_DST | GPUTextureUsage.RENDER_ATTACHMENT, | |
}); | |
device.queue.copyExternalImageToTexture({source: bitmap}, {texture: texture}, [bitmapSize, bitmapSize]); | |
const alphasBuffer0 = device.createBuffer({size: bitmapSize * bitmapSize * Float32Array.BYTES_PER_ELEMENT, usage: GPUBufferUsage.STORAGE}); | |
const alphasBuffer1 = device.createBuffer({size: bitmapSize * bitmapSize * Float32Array.BYTES_PER_ELEMENT, usage: GPUBufferUsage.STORAGE}); | |
// pipelines | |
const copyAlphaPipeline = device.createComputePipeline({ | |
compute: {module: copyAlphaShader, entryPoint: "copy_alpha"} | |
}); | |
const shrinkMipmapPipeline = device.createComputePipeline({ | |
compute: {module: shrinkMipmapShader, entryPoint: "shrink_mipmap"} | |
}); | |
// command encoder | |
const initEncoder = device.createCommandEncoder(); | |
{//copy alpha | |
const bindGroup = device.createBindGroup({ | |
layout: copyAlphaPipeline.getBindGroupLayout(0), | |
entries: [ | |
{binding: 0, resource: {buffer: alphasBuffer0}}, | |
{binding: 1, resource: texture.createView({format: "rgba8unorm", dimension: "2d", baseMipLevel: 0, mipLevelCount: 1})}, | |
] | |
}); | |
const passEncoder = initEncoder.beginComputePass(); | |
passEncoder.setPipeline(copyAlphaPipeline); | |
passEncoder.setBindGroup(0, bindGroup); | |
passEncoder.dispatch(bitmapSize / ws, bitmapSize); | |
passEncoder.end(); | |
} | |
for (let level = 1; level < mipLevelCount; level++) {//shrink mipmap | |
const destSize = bitmapSize >> level; | |
const [alphaIn, alphaOut] = (level % 2 === 1) ? [alphasBuffer0, alphasBuffer1] : [alphasBuffer1, alphasBuffer0]; | |
const bindGroup = device.createBindGroup({ | |
layout: shrinkMipmapPipeline.getBindGroupLayout(0), | |
entries: [ | |
{binding: 0, resource: {buffer: alphaIn}}, | |
{binding: 1, resource: {buffer: alphaOut}}, | |
{binding: 2, resource: texture.createView({format: "rgba8unorm", dimension: "2d", baseMipLevel: level, mipLevelCount: 1})}, | |
] | |
}); | |
const passEncoder = initEncoder.beginComputePass(); | |
passEncoder.setPipeline(shrinkMipmapPipeline); | |
passEncoder.setBindGroup(0, bindGroup); | |
passEncoder.dispatch(Math.ceil(destSize / ws), destSize); | |
passEncoder.end(); | |
} | |
device.queue.submit([initEncoder.finish()]); | |
//[B. sampling particle from prob mipmap] particle generator | |
const generateParticlesWgsl = ` | |
struct Particle { | |
rgba: vec4<f32>; | |
uv: vec2<f32>; | |
// pad_for_align16: array<f32, 2>; | |
}; | |
struct Particles { | |
list: array<Particle>; | |
}; | |
@group(0) @binding(0) var<storage, read_write> particles: Particles; | |
@group(0) @binding(1) var texture: texture_2d<f32>; | |
var<private> rand_seed : vec2<f32>; | |
fn rand() -> f32 { | |
rand_seed.x = fract(cos(dot(rand_seed, vec2<f32>(23.14077926, 232.61690225))) * 136.8168); | |
rand_seed.y = fract(cos(dot(rand_seed, vec2<f32>(54.47856553, 345.84153136))) * 534.7645); | |
return rand_seed.y; | |
} | |
fn born() -> Particle { | |
var pos = vec2<i32>(0, 0); | |
for (var level = textureNumLevels(texture) - 1; level > 0; level = level - 1) { | |
let r = rand(); | |
let probs = textureLoad(texture, pos, level); | |
if (r < probs.r) { | |
pos = vec2<i32>(pos.x * 2, pos.y * 2); | |
} else if (r < probs.g) { | |
pos = vec2<i32>(pos.x * 2 + 1, pos.y * 2); | |
} else if (r < probs.b) { | |
pos = vec2<i32>(pos.x * 2, pos.y * 2 + 1); | |
} else { | |
pos = vec2<i32>(pos.x * 2 + 1, pos.y * 2 + 1); | |
} | |
} | |
let uv = vec2<f32>(pos) / vec2<f32>(textureDimensions(texture)); | |
let rgba = textureLoad(texture, pos, 0); | |
return Particle(rgba, uv); | |
} | |
let ws = ${ws}; | |
@stage(compute) @workgroup_size(ws) fn generate_particles(@builtin(global_invocation_id) giid: vec3<u32>) { | |
rand_seed = vec2<f32>(giid.xy); | |
particles.list[giid.x] = born(); | |
} | |
`; | |
const generateParticlesShader = device.createShaderModule({code: generateParticlesWgsl}); | |
const particleSize = (4 + 2 + 2/*as pading for vec4<f32>'s align 16*/) * Float32Array.BYTES_PER_ELEMENT ; | |
const wsCount = 64; | |
const particleCount = ws * wsCount; | |
const particlesBuffer = device.createBuffer({size: particleSize * particleCount, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.VERTEX}); | |
const generateParticlesPipeline = device.createComputePipeline({ | |
compute: {module: generateParticlesShader, entryPoint: "generate_particles"} | |
}); | |
const particlesEncoder = device.createCommandEncoder(); | |
{ | |
const bindGroup = device.createBindGroup({ | |
layout: generateParticlesPipeline.getBindGroupLayout(0), | |
entries: [ | |
{binding: 0, resource: {buffer: particlesBuffer}}, | |
{binding: 1, resource: texture.createView()}, | |
]}); | |
const passEncoder = particlesEncoder.beginComputePass(); | |
passEncoder.setPipeline(generateParticlesPipeline); | |
passEncoder.setBindGroup(0, bindGroup); | |
passEncoder.dispatch(wsCount); | |
passEncoder.end(); | |
} | |
device.queue.submit([particlesEncoder.finish()]); | |
//[C. render particle] as regular polygons | |
const shape = 4; | |
const perR = 64; | |
const renderWgsl = ` | |
struct Particle { | |
@location(0) rgba: vec4<f32>; | |
@location(1) uv: vec2<f32>; | |
}; | |
struct InOut { | |
@builtin(position) pos: vec4<f32>; | |
@location(0) color: vec4<f32>; | |
}; | |
@stage(vertex) fn vmain(@builtin(vertex_index) i: u32, particle: Particle) -> InOut { | |
let center = vec2<f32>(2.0 * (particle.uv.x - 0.5), -2.0 * (particle.uv.y - 0.5)); | |
let vid = i % 3u; | |
let tid = (i - vid) / 3u; | |
if (vid == 0u) { | |
return InOut(vec4<f32>(center, 0.0, 1.0), particle.rgba); | |
} | |
let t = radians(f32(tid + vid - 1u) / f32(${shape}) * 360.0); | |
let r = 1.0 / f32(${perR}); | |
return InOut(vec4<f32>(center + vec2<f32>(-sin(t), cos(t)) * r, 0.0, 1.0), particle.rgba); | |
} | |
@stage(fragment) fn fmain(io: InOut) -> @location(0) vec4<f32> { | |
return io.color; | |
}; | |
`; | |
const renderShader = device.createShaderModule({code: renderWgsl}); | |
const stride = { | |
stepMode: "instance", | |
arrayStride: particleSize, | |
attributes: [ | |
{shaderLocation: 0, offset: 0, format: "float32x4"}, | |
{shaderLocation: 1, offset: 4 * Float32Array.BYTES_PER_ELEMENT, format: "float32x2"}, | |
] | |
}; | |
// gpu config for canvas | |
const canvas = document.getElementById("canvas"); | |
const gpu = canvas.getContext("webgpu"); | |
const format = gpu.getPreferredFormat(adapter); | |
gpu.configure({device, format, size: [canvas.width, canvas.height]}); | |
const renderPipeline = device.createRenderPipeline({ | |
primitive: {topology: "triangle-list", cullMode: "back"}, | |
vertex: {module: renderShader, entryPoint: "vmain", buffers: [stride]}, | |
fragment: {module: renderShader, entryPoint: "fmain", targets: [{format}]}, | |
}); | |
{ | |
const view = gpu.getCurrentTexture().createView(); | |
const clearValue = {r: 0, g: 0, b: 0, a: 1}; | |
const renderPass = {colorAttachments: [{view, loadOp: "clear", clearValue, loadValue: clearValue, storeOp: "store"}]}; | |
const commandEncoder = device.createCommandEncoder(); | |
const passEncoder = commandEncoder.beginRenderPass(renderPass); | |
passEncoder.setPipeline(renderPipeline); | |
passEncoder.setVertexBuffer(0, particlesBuffer); | |
passEncoder.draw(3 * shape, particleCount); | |
passEncoder.end(); | |
device.queue.submit([commandEncoder.finish()]); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
demo for chrome-100: https://gist.githack.com/bellbind/8318c595fb1464bfefa2827b0b55667a/raw/index.html