Public
Edited
May 7, 2024
Insert cell
Insert cell
{
const { context, device, format } = await gpu.init(SIZE, SIZE)

// input image
const sampler = gpu.sampler(device, {
minFilter: 'linear',
magFilter: 'linear',
mipmapFilter: 'linear',
})
const texture = gpu.canvasTexture(device, sample.canvas)

// blur
const WORKGROUP_SIZE = SIZE
const WORKGROUP_DISPATCH_COUNT = SIZE // Math.ceil((SIZE ** 2) / WORKGROUP_SIZE)
const blurPipeline = device.createComputePipeline({
layout: 'auto',
compute: {
module: device.createShaderModule({
code: `
struct Uniforms {
direction: u32,
}

@group(0) @binding(0) var tex_in: texture_2d<f32>;
@group(0) @binding(1) var tex_out: texture_storage_2d<rgba8unorm, write>;
@group(0) @binding(2) var<uniform> uniforms: Uniforms;

const offsets = array<u32, 5>(0, 1, 2, 3, 4); // k = 5
const weight = array<f32, 5>(0.2270270270, 0.1945945946, 0.1216216216, 0.0540540541, 0.0162162162);

fn blur(image_tex: texture_2d<f32>, pos: vec2u, direction: vec2u) -> vec4<f32> {
let size = textureDimensions(image_tex);
var color = textureLoad(image_tex, pos, 0) * weight[0];
for (var i = 1; i < 5; i++) {
let offset = vec2(offsets[i]) * direction;
let o1 = pos.xy + offset;
if (all(o1 < size)) {
color += textureLoad(image_tex, o1, 0) * weight[i];
}
let o2 = pos.xy - offset;
if (all(o2 < size)) {
color += textureLoad(image_tex, o2, 0) * weight[i];
}
}
return color;
}

const h = vec2u(0, 1);
const v = vec2u(1, 0);

@compute
@workgroup_size(${WORKGROUP_SIZE})
fn cs(
@builtin(workgroup_id) workgroup_id : vec3<u32>,
@builtin(local_invocation_id) local_invocation_id : vec3<u32>,
) {
let pos = vec2u(workgroup_id.x, local_invocation_id.x);
let size = textureDimensions(tex_in, 0);
if (all(pos < size)) {
// let color = textureSampleLevel(tex_in, tex_sampler, vec2f(0), 0);
// \`select\` is like a ternary with (false, true, condition) call sig
let direction = select(h, v, bool(uniforms.direction));
let color = blur(tex_in, pos, direction);
textureStore(tex_out, pos, color);
}
}
`
})
}
})

const [tex0, tex1] = util.arr(2, () =>
gpu.createTexture(device, {
width: sample.canvas.width,
height: sample.canvas.height
})
)

const [buf0, buf1] = util.arr(2, b => {
const buf = device.createBuffer({
label: `buf${b}`,
size: 4,
usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
})
device.queue.writeBuffer(buf, 0, new Uint32Array([b]))
return buf
})

// bind groups
// const computeConstants = device.createBindGroup({
// layout: blurPipeline.getBindGroupLayout(0),
// entries: [
// { binding: 0, resource: sampler }
// ]
// })

const computeLayoutIdx = 0
const computeInitial = device.createBindGroup({
layout: blurPipeline.getBindGroupLayout(computeLayoutIdx),
entries: [
{ binding: 0, resource: texture.createView() },
{ binding: 1, resource: tex0.createView() },
{ binding: 2, resource: { buffer: buf0 } }
],
})
const computeA = device.createBindGroup({
layout: blurPipeline.getBindGroupLayout(computeLayoutIdx),
entries: [
{ binding: 0, resource: tex0.createView() },
{ binding: 1, resource: tex1.createView() },
{ binding: 2, resource: { buffer: buf1 } }
],
})
const computeB = device.createBindGroup({
layout: blurPipeline.getBindGroupLayout(computeLayoutIdx),
entries: [
{ binding: 0, resource: tex1.createView() },
{ binding: 1, resource: tex0.createView() },
{ binding: 2, resource: { buffer: buf0 } }
],
})
// for rendering the final blurred output
const module = device.createShaderModule({
code: `
struct VertexOut {
@builtin(position) position: vec4f,
@location(0) texcoord: vec2f,
}

@group(0) @binding(0) var tex_sampler: sampler;
@group(0) @binding(1) var tex: texture_2d<f32>;
@vertex
fn vs(
@builtin(vertex_index) vertexIndex : u32
) -> VertexOut {
let quad = array(
vec2f(0.0, 0.0),
vec2f(1.0, 0.0),
vec2f(0.0, 1.0),
vec2f(0.0, 1.0),
vec2f(1.0, 0.0),
vec2f(1.0, 1.0),
);
let texcoord = quad[vertexIndex];
return VertexOut(
vec4f((texcoord - 0.5) * 2, 0, 1),
texcoord
);
}
@fragment
fn fs(fsIn: VertexOut) -> @location(0) vec4f {
return textureSample(tex, tex_sampler, fsIn.texcoord);
}
`
})
const pipeline = device.createRenderPipeline({
layout: 'auto',
vertex: {
module,
entryPoint: 'vs',
},
fragment: {
module,
entryPoint: 'fs',
targets: [{ format }],
},
})
const result = device.createBindGroup({
layout: pipeline.getBindGroupLayout(0),
entries: [
{ binding: 0, resource: sampler },
{ binding: 1, resource: tex1.createView() },
],
})
function render() {
const encoder = device.createCommandEncoder();
const computePass = encoder.beginComputePass()
computePass.setPipeline(blurPipeline)

// computePass.setBindGroup(0, computeConstants)
computePass.setBindGroup(computeLayoutIdx, computeInitial)
computePass.dispatchWorkgroups(WORKGROUP_DISPATCH_COUNT)

computePass.setBindGroup(computeLayoutIdx, computeA)
computePass.dispatchWorkgroups(WORKGROUP_DISPATCH_COUNT)

for (let i = 0; i < 4; i++) {
computePass.setBindGroup(computeLayoutIdx, computeB)
computePass.dispatchWorkgroups(WORKGROUP_DISPATCH_COUNT)
computePass.setBindGroup(computeLayoutIdx, computeA)
computePass.dispatchWorkgroups(WORKGROUP_DISPATCH_COUNT)
}

computePass.end()
const pass = encoder.beginRenderPass({
colorAttachments: [
{
view: context.getCurrentTexture().createView(),
clearValue: [0, 0, 0, 1],
loadOp: 'clear',
storeOp: 'store',
},
],
});
pass.setPipeline(pipeline);
pass.setBindGroup(0, result);
pass.draw(6);
pass.end();
device.queue.submit([encoder.finish()]);
}
render()

return htl.html`${[sample.canvas, context.canvas]}`
}
Insert cell
Insert cell
Insert cell
sample = {
const canvas = document.createElement('canvas');
canvas.width = SIZE;
canvas.height = SIZE;
const context = canvas.getContext('2d');

// context.lineWidth = 3
// context.strokeStyle = '#fff'
for (let i = 0; i < 20; i++) {
context.fillStyle = `hsl(${util.rand(360)}, 100%, 60%)`
context.beginPath()
const r = util.rand(SIZE / 8)
const m = r * 1.1
context.arc(util.rand(m, SIZE - m), util.rand(m, SIZE - m), r, 0, util.TWO_PI)
context.fill()
// context.stroke()
}

return context
}
Insert cell
gpu = {
const init = async (width = 512, height = 512) => {
const canvas = document.createElement('canvas');
canvas.width = width
canvas.height = height
const context = canvas.getContext('webgpu');

const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();
const format = navigator.gpu.getPreferredCanvasFormat();
context.configure({ device, format, alphaMode: 'premultiplied', });
return { context, adapter, device, format }
}
const sampler = (device, options = {}) => {
return device.createSampler({
magFilter: 'linear',
minFilter: 'linear',
...options
})
}

const createTexture = (
device,
{
width,
height,
format = 'rgba8unorm',
usage =
GPUTextureUsage.COPY_DST |
GPUTextureUsage.STORAGE_BINDING |
GPUTextureUsage.TEXTURE_BINDING,
}
) => {
return device.createTexture({
format,
size: [width, height],
usage
})
}
const canvasTexture = (
device,
canvas,
{
format = 'rgba8unorm',
flipY = true,
usage =
GPUTextureUsage.TEXTURE_BINDING |
GPUTextureUsage.COPY_DST |
GPUTextureUsage.RENDER_ATTACHMENT
} = {}
) => {
const texture = device.createTexture({
format,
size: [canvas.width, canvas.height],
usage
})
device.queue.copyExternalImageToTexture(
{ source: canvas, flipY },
{ texture },
[canvas.width, canvas.height]
)
return texture
}
return {
init,
sampler,
canvasTexture,
createTexture
}
}
Insert cell
util = ({
TWO_PI: 2 * Math.PI,
rand: (min = 1, max) => {
if (max === undefined) {
return Math.random() * min
}
return Math.random() * (max - min) + min
},
// random signed
rands: () => (Math.random() - 0.5) * 2,
arr: (size, callback) => {
const arr = new Array(size)
if (typeof callback !== 'function') {
return arr.fill(callback)
}
return arr.fill(null).map((_, i) => callback(i))
},
})
Insert cell

Purpose-built for displays of data

Observable is your go-to platform for exploring data and creating expressive data visualizations. Use reactive JavaScript notebooks for prototyping and a collaborative canvas for visual data exploration and dashboard creation.
Learn more