render-zig

A 3D rendering engine written in Zig
git clone git://git.christianermann.dev/render-zig
Log | Files | Refs

mesh_clusters.wgsl (5731B)


      1 struct VertexInput {
      2     @location(0) position: vec4<f32>,
      3     @location(1) normal: vec3<f32>,
      4     @location(2) tex_coord: vec3<f32>,
      5     @location(3) tangent: vec3<f32>,
      6     @location(4) bitangent: vec3<f32>,
      7 };
      8 
      9 struct VertexOutput {
     10     @builtin(position) clip_position: vec4<f32>,
     11     @location(0) tex_coord: vec3<f32>,
     12     @location(1) pt_fragment: vec3<f32>,
     13     @location(2) pt_light: vec3<f32>,
     14     @location(3) pt_camera: vec3<f32>,
     15     @location(4) @interpolate(flat) instance_id: u32,
     16 };
     17 
     18 @group(0) @binding(0)
     19 var<storage, read> instances: array<Instance>;
     20 struct Instance {
     21     model_matrix: mat4x4<f32>,
     22     material: Material,
     23 };
     24 struct Material {
     25     albedo: u32,
     26     normal: u32,
     27     roughness: u32,
     28     metalness: u32,
     29 };
     30 
     31 @group(0) @binding(4)
     32 var<storage, read> clusters: array<Cluster>;
     33 struct Cluster {
     34     model_id: u32,
     35 };
     36 
     37 @group(0) @binding(1)
     38 var<uniform> uniform_data: UniformData;
     39 struct UniformData {
     40     pw_camera: vec3<f32>,
     41     view_matrix: mat4x4<f32>,
     42     proj_matrix: mat4x4<f32>,
     43 };
     44 
     45 @vertex fn vertex(
     46     in: VertexInput,
     47     @builtin(instance_index) idx: u32,
     48 ) -> VertexOutput {
     49     let model_id = clusters[idx].model_id;
     50     let model_matrix = instances[model_id].model_matrix;
     51     let camera_matrix = uniform_data.proj_matrix * uniform_data.view_matrix;
     52 
     53     let t = normalize((model_matrix * vec4(in.tangent, 0.0)).xyz);
     54     let b = normalize((model_matrix * vec4(in.bitangent, 0.0)).xyz);
     55     let n = normalize((model_matrix * vec4(in.normal, 0.0)).xyz);
     56     let tangent_matrix = mat3x3<f32>(t, b, n);
     57 
     58     let world_position = model_matrix * in.position;
     59     let tangent_position = tangent_matrix * world_position.xyz;
     60 
     61     let light_position = vec3<f32>(20, 80, 20);
     62     let tangent_light_position = tangent_matrix * light_position;
     63 
     64     let pt_camera = tangent_matrix * uniform_data.pw_camera;
     65 
     66 
     67     return VertexOutput(
     68         camera_matrix * world_position,
     69         in.tex_coord,
     70         tangent_position,
     71         tangent_light_position,
     72         pt_camera,
     73         idx,
     74     );
     75 }
     76 
     77 
     78 @group(0) @binding(2) var material_sampler: sampler;
     79 @group(0) @binding(3) var material_texture: texture_2d_array<f32>;
     80 
     81 fn diffuse_lambert(albedo: vec3<f32>) -> vec3<f32> {
     82     return albedo * 0.31830988618;
     83 }
     84 
     85 fn distribution_trowbridge_reitz_ggx(
     86     n_dot_h: f32,
     87     alpha: f32,
     88 ) -> f32 {
     89     let alpha_squared = alpha * alpha;
     90     let n_dot_h_squared = n_dot_h * n_dot_h;
     91     let denom = n_dot_h_squared * (alpha_squared - 1) + 1;
     92     let denom_squared = denom * denom;
     93     return alpha_squared / (3.14159265359 * denom_squared);
     94 }
     95 
     96 fn geometry_schlick_ggx(
     97     n_dot_x: f32,
     98     alpha: f32,
     99 ) -> f32 {
    100     let k = 0.5 * alpha;
    101     let denom = n_dot_x * (1.0 - k) + k;
    102     return n_dot_x / denom;
    103 }
    104 
    105 fn geometry_smith(
    106     n_dot_l: f32,
    107     n_dot_v: f32,
    108     alpha: f32,
    109 ) -> f32 {
    110     let g1 = geometry_schlick_ggx(n_dot_l, alpha);
    111     let g2 = geometry_schlick_ggx(n_dot_v, alpha);
    112     return g1 * g2;
    113 }
    114 
    115 fn fresnel_schlick(
    116     f0: vec3<f32>,
    117     v_dot_h: f32,
    118 ) -> vec3<f32> {
    119     return f0 + (1.0 - f0) * pow(1.0 - v_dot_h, 5.0);
    120 }
    121 
    122 fn cook_torrance(
    123     n_dot_v: f32,
    124     n_dot_l: f32,
    125     n_dot_h: f32,
    126     roughness: f32,
    127 ) -> f32 {
    128     let alpha = roughness * roughness;
    129     let distribution = distribution_trowbridge_reitz_ggx(
    130         n_dot_h,
    131         alpha,
    132     );
    133     let geometry = geometry_smith(
    134         n_dot_l,
    135         n_dot_v,
    136         alpha,
    137     );
    138     let denom = 4.0 * n_dot_v * n_dot_l + 0.000001;
    139     return (distribution * geometry) / denom;
    140 }
    141 
    142 fn lighting(
    143     normal: vec3<f32>,
    144     p_fragment: vec3<f32>,
    145     p_camera: vec3<f32>,
    146     p_light: vec3<f32>,
    147     c_light: vec3<f32>,
    148     albedo: vec3<f32>,
    149     roughness: f32,
    150     metalness: f32,
    151 ) -> vec3<f32> {
    152     let v = normalize(p_camera - p_fragment);
    153     let l = normalize(p_light - p_fragment);
    154     let h = normalize(v + l);
    155 
    156     let n_dot_v = max(dot(normal, v), 0.0);
    157     let n_dot_l = max(dot(normal, l), 0.0);
    158     let n_dot_h = max(dot(normal, h), 0.0);
    159     let v_dot_h = max(dot(v, h), 0.0);
    160 
    161     let f0 = mix(vec3<f32>(0.04), albedo, metalness);
    162     let k_s = fresnel_schlick(f0, v_dot_h);
    163     let k_d = 1.0 - k_s;
    164 
    165     let diffuse = k_d * diffuse_lambert(albedo);
    166     let specular = k_s * cook_torrance(n_dot_v, n_dot_l, n_dot_h, roughness);
    167     let brdf = (diffuse + specular) * n_dot_l;
    168     return brdf * c_light;
    169 }
    170 
    171 @fragment fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {
    172     let material = instances[in.instance_id].material;
    173     let uv = vec2<f32>(in.tex_coord.x, 1 - in.tex_coord.y);
    174     var albedo = textureSample(material_texture, material_sampler, uv, material.albedo).rgb;
    175     albedo = pow(albedo, vec3<f32>(2.2));
    176     var normal = textureSample(material_texture, material_sampler, uv, material.normal).rgb;
    177     normal = normalize(normal * 2.0 - 1.0);
    178     var roughness = textureSample(material_texture, material_sampler, uv, material.roughness).r;
    179     var metalness = textureSample(material_texture, material_sampler, uv, material.metalness).r;
    180 
    181     let light_color = normalize(vec3<f32>(23.47, 21.31, 20.79));
    182     //let light_color = normalize(vec3<f32>(0.9, 0.9, 0.9));
    183     var c = lighting(
    184         normal,
    185         in.pt_fragment,
    186         in.pt_camera,
    187         in.pt_light,
    188         light_color,
    189         albedo.rgb,
    190         roughness,
    191         metalness,
    192     );
    193 
    194     let ambient = vec3<f32>(0.03) * albedo;
    195     c += ambient;
    196 
    197     c = c / (c + 1.0);
    198     c = pow(c, vec3<f32>(1.0/2.2));
    199 
    200     let colors = array<vec3<f32>, 6>(
    201         vec3<f32>(1, 1, 0),
    202         vec3<f32>(0, 1, 1),
    203         vec3<f32>(1, 0, 1),
    204         vec3<f32>(1, 0, 0),
    205         vec3<f32>(0, 1, 0),
    206         vec3<f32>(0, 0, 1),
    207     );
    208     let color = colors[in.instance_id % 6];
    209 
    210     return vec4<f32>(color, 1.0);
    211 }