Public
Edited
May 20, 2024
Insert cell
Insert cell
{
const { device, context } = await gpu.init(1024, 1024)

const DISPATCH_SIZE = [Math.ceil(1024 / 16), Math.ceil(1024 / 16), 1]

// Contents of texture don't matter, we're just displaying tiles
const texInput = gpu.canvasTexture(device, sample.canvas, { usage: gpu.Usage.StorageTexture })

const texOutput = device.createTexture({
size: [1024, 1024],
format: 'rgba8unorm',
usage:
GPUTextureUsage.STORAGE_BINDING |
GPUTextureUsage.COPY_SRC |
GPUTextureUsage.TEXTURE_BINDING
})

const shaderModule = device.createShaderModule({
code: `
@group(0) @binding(0) var texInput: texture_2d<f32>;
@group(0) @binding(1) var texOutput: texture_storage_2d<rgba8unorm, write>;

const tileSize: vec2u = vec2u(16, 16);
const blurRadius: i32 = 10;
const count: f32 = pow(f32(blurRadius) * 2 + 1, 2);

@compute
@workgroup_size(16, 16, 1)
fn cs(@builtin(global_invocation_id) global_id: vec3u) {
let texSize: vec2u = textureDimensions(texInput).xy;

let tileNW: vec2u = (global_id.xy / tileSize) * tileSize;
let localPos: vec2u = global_id.xy % tileSize;
let texPos: vec2u = tileNW + localPos;

if (texPos.x < texSize.x && texPos.y < texSize.y) {
var colorSum: vec4f = vec4f(0.0, 0.0, 0.0, 0.0);

for (var y = -blurRadius; y <= blurRadius; y++) {
for (var x = -blurRadius; x <= blurRadius; x++) {
let samplePos = vec2i(texPos) + vec2i(x, y);
if (
samplePos.x >= 0 && samplePos.x < i32(texSize.x) &&
samplePos.y >= 0 && samplePos.y < i32(texSize.y)
) {
colorSum = colorSum + textureLoad(texInput, samplePos, 0);
}
}
}
let color: vec4f = colorSum / count;
textureStore(texOutput, texPos, color);
}
}
`
})

const bindGroupLayout = device.createBindGroupLayout({
entries: [
{
binding: 0,
visibility: GPUShaderStage.COMPUTE,
texture: { format: 'rgba8unorm' }
},
{
binding: 1,
visibility: GPUShaderStage.COMPUTE | GPUShaderStage.FRAGMENT,
storageTexture: { format: 'rgba8unorm' }
}
]
})

const pipelineLayout = device.createPipelineLayout({
bindGroupLayouts: [bindGroupLayout]
})

const pipeline = device.createComputePipeline({
layout: pipelineLayout,
compute: {
module: shaderModule,
entryPoint: 'cs'
}
})

const bindGroup0 = device.createBindGroup({
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: texInput.createView() },
{ binding: 1, resource: texOutput.createView() }
]
})

const bindGroup1 = device.createBindGroup({
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: texOutput.createView() },
{ binding: 1, resource: texInput.createView() },
]
})

function render() {
const commandEncoder = device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(pipeline);

for (let i = 0; i < 3; i++) {
passEncoder.setBindGroup(0, bindGroup0);
passEncoder.dispatchWorkgroups(...DISPATCH_SIZE);

passEncoder.setBindGroup(0, bindGroup1);
passEncoder.dispatchWorkgroups(...DISPATCH_SIZE);
}
passEncoder.end();
device.queue.submit([commandEncoder.finish()]);
}

render()

const debug = new TextureRenderer(device, texOutput)
debug.render()

return debug.canvas
}
Insert cell
class TextureRenderer {
device = null
texture = null
context = null
sampler = null
pipeline = null
bindGroup = null
constructor(
device /* : GPUDevice */,
texture /* : GPUTexture */,
{
format = null,
alphaMode = 'premultiplied',
style = {},
} = {}
) {
format ??= gpu.format()
this.device = device
this.context = gpu.context(texture.width, texture.height);
this.context.configure({ device, format, alphaMode })
for (const key in style) {
this.context.canvas.style[key] = style[key]
}
const module = device.createShaderModule({
code: TextureRenderer.Shader
})

this.sampler = gpu.sampler(device, gpu.LinearSampler)

this.pipeline = device.createRenderPipeline({
layout: 'auto',
vertex: {
module,
entryPoint: 'vs',
},
fragment: {
module,
entryPoint: 'fs',
targets: [{ format }],
},
})

this.bindGroup = device.createBindGroup({
layout: this.pipeline.getBindGroupLayout(0),
entries: [
{ binding: 0, resource: this.sampler },
{ binding: 1, resource: texture.createView() },
],
})
}

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

static Shader = `
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(vs_out: VertexOut) -> @location(0) vec4f {
return textureSample(tex, tex_sampler, vs_out.texcoord);
}
`
}
Insert cell
sample.canvas
Insert cell
sample = {
const size = 64
const count = 1024 / 64
const ctx = gpu.context(1024, 1024, '2d')
ctx.fillStyle = '#000'
for (let y = 0; y < count; y++) {
for (let x = -1; x < count; x += 2) {
ctx.fillRect(x * size + y % 2 * size, y * size, size, size)
}
}
return ctx
}
Insert cell
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