// The GPU side of one image: its buffers, the four products (two forward, two
// inverse) and the four panels, each with its own small uniform buffer so they
// can all go in one submission.
res = {
const { device, matmul, display, lut } = gpu
const { data, h, w } = image
const hw = h * w
const owned = []
const buffer = (bytes, usage) => {
const b = device.createBuffer({ size: bytes, usage })
owned.push(b)
return b
}
const { STORAGE, UNIFORM, COPY_DST, COPY_SRC } = GPUBufferUsage
const matrix = (n) => buffer(4 * n, STORAGE | COPY_DST | COPY_SRC)
const B = {
Bh: matrix(h * h), BhT: matrix(h * h), Bw: matrix(w * w), BwT: matrix(w * w),
X: matrix(3 * hw), T: matrix(3 * hw), C: matrix(3 * hw), U: matrix(3 * hw), Xh: matrix(3 * hw),
}
device.queue.writeBuffer(B.X, 0, data)
const group = (pipeline, buffers) => device.createBindGroup({
layout: pipeline.getBindGroupLayout(0),
entries: buffers.map((b, binding) => ({ binding, resource: { buffer: b } })),
})
// O = L · R on every channel; `batched` names the side that is the image.
const product = (L, R, O, [M, K, N], batched) => {
const params = buffer(32, UNIFORM | COPY_DST)
const bind = group(matmul, [params, L, R, O])
const set = (t) => {
const bytes = new ArrayBuffer(32)
new Uint32Array(bytes, 0, 6).set([M, N, K, batched === "L" ? hw : 0, batched === "R" ? hw : 0, hw])
new Float32Array(bytes, 24, 1)[0] = t
device.queue.writeBuffer(params, 0, bytes)
}
set(0)
return { set, run: (pass) => { pass.setBindGroup(0, bind); pass.dispatchWorkgroups(N / 64, M / 64, 3) } }
}
const forward = [product(B.Bh, B.X, B.T, [h, h, w], "R"), product(B.T, B.BwT, B.C, [h, w, w], "L")]
const inverse = [product(B.BhT, B.C, B.U, [h, h, w], "R"), product(B.U, B.Bw, B.Xh, [h, w, w], "L")]
const panel = (canvas, A, other = A) => {
const params = buffer(32, UNIFORM | COPY_DST)
const bind = group(display, [params, A, other, lut])
canvas.width = w
canvas.height = h
return {
set: ({ mode = 0, t = 0, span = [0, 1], gain = 1 }) => {
const bytes = new ArrayBuffer(32)
new Uint32Array(bytes, 0, 3).set([w, h, mode])
new Float32Array(bytes, 12, 4).set([t, span[0], span[1], gain])
device.queue.writeBuffer(params, 0, bytes)
},
draw: (encoder) => {
const pass = encoder.beginRenderPass({
colorAttachments: [{ view: canvas.getContext("webgpu").getCurrentTexture().createView(), loadOp: "clear", storeOp: "store" }],
})
pass.setPipeline(display)
pass.setBindGroup(0, bind)
pass.draw(3)
pass.end()
},
}
}
const panels = {
original: panel(canvases.original, B.X),
coefficients: panel(canvases.coefficients, B.C),
sparse: panel(canvases.sparse, B.C),
reconstruction: panel(canvases.reconstruction, B.Xh, B.X),
}
// Redraw one panel with its current settings. A WebGPU canvas reads back
// empty once its frame is on screen, so scripts/export-compression-figures.ts
// calls this and reads the canvas in the same task.
for (const [name, p] of Object.entries(panels)) {
canvases[name].redraw = () => {
const encoder = device.createCommandEncoder()
p.draw(encoder)
device.queue.submit([encoder.finish()])
}
}
// Redraws run one at a time, and while one is waiting a newer one replaces
// it, so a fast drag on a slow GPU renders the latest position instead of
// every position it passed through.
let next = null, busy = false, dead = false
const frame = async (job) => {
next = job
if (busy) return
busy = true
while (next && !dead) { const run = next; next = null; await run() }
busy = false
}
invalidation.then(() => { dead = true; owned.forEach((b) => b.destroy()) })
return { ...B, h, w, forward, inverse, panels, frame }
}