Skip to content

Instantly share code, notes, and snippets.

@bellbind
Last active April 3, 2022 13:08
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save bellbind/8318c595fb1464bfefa2827b0b55667a to your computer and use it in GitHub Desktop.
Save bellbind/8318c595fb1464bfefa2827b0b55667a to your computer and use it in GitHub Desktop.
[WebGPU] sampling from bitmap image by compute shaders (for Chrome-100)
<!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>
// 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()]);
}
@bellbind
Copy link
Author

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment