diff options
Diffstat (limited to 'workspaces/main/shaders/cnn/cnn_layer.wgsl')
| -rw-r--r-- | workspaces/main/shaders/cnn/cnn_layer.wgsl | 46 |
1 files changed, 46 insertions, 0 deletions
diff --git a/workspaces/main/shaders/cnn/cnn_layer.wgsl b/workspaces/main/shaders/cnn/cnn_layer.wgsl new file mode 100644 index 0000000..e026ce8 --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_layer.wgsl @@ -0,0 +1,46 @@ +// CNN layer shader - uses modular convolution snippets +// Supports multi-pass rendering with residual connections + +@group(0) @binding(0) var smplr: sampler; +@group(0) @binding(1) var txt: texture_2d<f32>; + +#include "common_uniforms" +#include "cnn_activation" +#include "cnn_conv3x3" +#include "cnn_weights_generated" + +struct CNNLayerParams { + layer_index: i32, + use_residual: i32, + _pad: vec2<f32>, +}; + +@group(0) @binding(2) var<uniform> uniforms: CommonUniforms; +@group(0) @binding(3) var<uniform> params: CNNLayerParams; + +@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> { + var pos = array<vec2<f32>, 3>( + vec2<f32>(-1.0, -1.0), vec2<f32>(3.0, -1.0), vec2<f32>(-1.0, 3.0) + ); + return vec4<f32>(pos[i], 0.0, 1.0); +} + +@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> { + let uv = p.xy / uniforms.resolution; + var result = vec4<f32>(0.0); + + // Single layer for now (layer 0) + if (params.layer_index == 0) { + result = cnn_conv3x3(txt, smplr, uv, uniforms.resolution, + weights_layer0, bias_layer0); + result = cnn_tanh(result); + } + + // Residual connection + if (params.use_residual != 0) { + let input = textureSample(txt, smplr, uv); + result = input + result * 0.3; + } + + return result; +} |
