diff options
Diffstat (limited to 'cnn_v2/shaders/cnn_v2_static.wgsl')
| -rw-r--r-- | cnn_v2/shaders/cnn_v2_static.wgsl | 75 |
1 files changed, 75 insertions, 0 deletions
diff --git a/cnn_v2/shaders/cnn_v2_static.wgsl b/cnn_v2/shaders/cnn_v2_static.wgsl new file mode 100644 index 0000000..309e832 --- /dev/null +++ b/cnn_v2/shaders/cnn_v2_static.wgsl @@ -0,0 +1,75 @@ +// CNN v2 Static Features Compute Shader +// Generates 8D parametric features: [p0, p1, p2, p3, uv.x, uv.y, sin20_y, bias] +// p0-p3: Parametric features from specified mip level (0=mip0, 1=mip1, 2=mip2, 3=mip3) +// Note: Input image RGBD (mip0) fed separately to Layer 0 +// +// TODO: Binary format should support arbitrary layout and ordering for feature vector (7D). +// Current layout is hardcoded. Future versions should allow runtime-specified +// feature combinations (e.g., [R, G, B, dx, dy, uv_x, bias] or custom encodings). + +struct StaticFeatureParams { + mip_level: u32, + padding0: u32, + padding1: u32, + padding2: u32, +} + +@group(0) @binding(0) var input_tex: texture_2d<f32>; +@group(0) @binding(1) var input_tex_mip1: texture_2d<f32>; +@group(0) @binding(2) var input_tex_mip2: texture_2d<f32>; +@group(0) @binding(3) var depth_tex: texture_2d<f32>; +@group(0) @binding(4) var output_tex: texture_storage_2d<rgba32uint, write>; +@group(0) @binding(5) var<uniform> params: StaticFeatureParams; +@group(0) @binding(6) var linear_sampler: sampler; + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3<u32>) { + let coord = vec2<i32>(id.xy); + let dims = textureDimensions(input_tex); + + if (coord.x >= i32(dims.x) || coord.y >= i32(dims.y)) { + return; + } + + // Parametric features (p0-p3) - bilinear sample from specified mip level + // Use UV coordinates for bilinear interpolation + // Note: Use textureSampleLevel (not textureSample) in compute shaders + let uv = (vec2<f32>(coord) + 0.5) / vec2<f32>(dims); + var rgba: vec4<f32>; + if (params.mip_level == 0u) { + rgba = textureSampleLevel(input_tex, linear_sampler, uv, 0.0); + } else if (params.mip_level == 1u) { + rgba = textureSampleLevel(input_tex_mip1, linear_sampler, uv, 0.0); + } else if (params.mip_level == 2u) { + rgba = textureSampleLevel(input_tex_mip2, linear_sampler, uv, 0.0); + } else { + // Mip 3 or higher: use mip 2 as fallback + rgba = textureSampleLevel(input_tex_mip2, linear_sampler, uv, 0.0); + } + + let p0 = rgba.r; + let p1 = rgba.g; + let p2 = rgba.b; + let p3 = textureLoad(depth_tex, coord, 0).r; + + // UV coordinates (normalized [0,1], top-left origin - matches training) + let uv_x = f32(coord.x) / f32(dims.x); + let uv_y = f32(coord.y) / f32(dims.y); + + // Multi-frequency position encoding + let sin20_y = sin(20.0 * uv_y); + + // Bias dimension (always 1.0) + let bias = 1.0; + + // Pack 8×f16 into 4×u32 (rgba32uint) + // [p0, p1, p2, p3, uv_x, uv_y, sin20_y, bias] + let packed = vec4<u32>( + pack2x16float(vec2<f32>(p0, p1)), + pack2x16float(vec2<f32>(p2, p3)), + pack2x16float(vec2<f32>(uv_x, uv_y)), + pack2x16float(vec2<f32>(sin20_y, bias)) + ); + + textureStore(output_tex, coord, packed); +} |
