Lightningbeam/lightningbeam-ui/lightningbeam-editor/src/panes/shaders/warp_apply.wgsl

104 lines
3.3 KiB
WebGPU Shading Language
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

// GPU warp-apply shader.
//
// Two modes selected by grid_cols / grid_rows:
//
// grid_cols == 0 (Liquify / per-pixel mode)
// disp[] is a full canvas-sized array<vec2f>. Each pixel reads its own entry.
//
// grid_cols > 0 (Warp control-point mode)
// disp[] contains only grid_cols * grid_rows vec2f displacements (one per
// control point). The shader bilinearly interpolates them so the CPU never
// needs to build or upload the full per-pixel buffer.
//
// Dispatch: ceil(dst_w / 8) × ceil(dst_h / 8) × 1
struct Params {
src_w: u32,
src_h: u32,
dst_w: u32,
dst_h: u32,
grid_cols: u32, // 0 = per-pixel mode
grid_rows: u32,
_pad0: u32,
_pad1: u32,
}
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var src: texture_2d<f32>;
@group(0) @binding(2) var<storage, read> disp: array<vec2f>;
@group(0) @binding(3) var dst: texture_storage_2d<rgba16float, write>;
// Manual bilinear sample with clamp-to-edge (textureSample forbidden in compute shaders).
fn bilinear_sample(px: f32, py: f32) -> vec4<f32> {
let sw = i32(params.src_w);
let sh = i32(params.src_h);
let ix = i32(floor(px - 0.5));
let iy = i32(floor(py - 0.5));
let fx = fract(px - 0.5);
let fy = fract(py - 0.5);
let x0 = clamp(ix, 0, sw - 1);
let x1 = clamp(ix + 1, 0, sw - 1);
let y0 = clamp(iy, 0, sh - 1);
let y1 = clamp(iy + 1, 0, sh - 1);
let s00 = textureLoad(src, vec2<i32>(x0, y0), 0);
let s10 = textureLoad(src, vec2<i32>(x1, y0), 0);
let s01 = textureLoad(src, vec2<i32>(x0, y1), 0);
let s11 = textureLoad(src, vec2<i32>(x1, y1), 0);
return mix(mix(s00, s10, fx), mix(s01, s11, fx), fy);
}
// Bilinearly interpolate the control-point displacement grid.
fn grid_displacement(px: u32, py: u32) -> vec2f {
let cols = params.grid_cols;
let rows = params.grid_rows;
// Normalised position in grid space [0 .. cols-1] × [0 .. rows-1].
let gx = f32(px) / f32(params.dst_w - 1u) * f32(cols - 1u);
let gy = f32(py) / f32(params.dst_h - 1u) * f32(rows - 1u);
let col0 = u32(floor(gx));
let row0 = u32(floor(gy));
let col1 = min(col0 + 1u, cols - 1u);
let row1 = min(row0 + 1u, rows - 1u);
let fx = gx - floor(gx);
let fy = gy - floor(gy);
let d00 = disp[row0 * cols + col0];
let d10 = disp[row0 * cols + col1];
let d01 = disp[row1 * cols + col0];
let d11 = disp[row1 * cols + col1];
return d00 * (1.0 - fx) * (1.0 - fy)
+ d10 * fx * (1.0 - fy)
+ d01 * (1.0 - fx) * fy
+ d11 * fx * fy;
}
@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
if gid.x >= params.dst_w || gid.y >= params.dst_h { return; }
var d: vec2f;
if params.grid_cols > 0u {
d = grid_displacement(gid.x, gid.y);
} else {
d = disp[gid.y * params.dst_w + gid.x];
}
let sx = f32(gid.x) + d.x;
let sy = f32(gid.y) + d.y;
var color: vec4<f32>;
if sx < 0.0 || sy < 0.0 || sx >= f32(params.src_w) || sy >= f32(params.src_h) {
color = vec4<f32>(0.0);
} else {
color = bilinear_sample(sx + 0.5, sy + 0.5);
}
textureStore(dst, vec2<i32>(i32(gid.x), i32(gid.y)), color);
}