{
const { context, device, format } = await gpu.init(SIZE, SIZE)
const sampler = gpu.sampler(device, {
minFilter: 'linear',
magFilter: 'linear',
mipmapFilter: 'linear',
})
const texture = gpu.canvasTexture(device, sample.canvas)
const WORKGROUP_SIZE = SIZE
const WORKGROUP_DISPATCH_COUNT = 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]}`
}