Public
Edited
May 20, 2024
1 fork
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)

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: vec2<u32> = vec2<u32>(16, 16);

@compute
@workgroup_size(16, 16, 1)
fn cs(@builtin(global_invocation_id) global_id: vec3<u32>) {
let texSize: vec2<u32> = textureDimensions(texInput).xy;

let tileNW: vec2<u32> = (global_id.xy / tileSize) * tileSize;
let localPos: vec2<u32> = global_id.xy % tileSize;
let texPos: vec2<u32> = tileNW + localPos;

if (texPos.x < texSize.x && texPos.y < texSize.y && textureLoad(texInput, texPos, 0).a > 0.0) {
// let color =
// vec4f(f32(localPos.x) / f32(tileSize.x), f32(localPos.y) / f32(tileSize.y), 0.0, 1.0) *
// vec4f(f32(tileNW.x) / f32(texSize.x), f32(tileNW.y) / f32(texSize.y), 0.0, 1.0);
let color = vec4f(f32(localPos.x) / f32(tileSize.x), f32(localPos.y) / f32(tileSize.y), 0.0, 1.0);
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 bindGroup = device.createBindGroup({
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: texInput.createView() },
{ binding: 1, resource: texOutput.createView() }
]
});

const commandEncoder = device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(pipeline);
passEncoder.setBindGroup(0, bindGroup);

passEncoder.dispatchWorkgroups(...DISPATCH_SIZE);
passEncoder.end();

device.queue.submit([commandEncoder.finish()]);

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