WIP: Merge noise-info features from VGA internal branch

shader/bmfr.glsl

@@ -4,7 +4,15 @@
 #define LOCAL_SIZE 16
 #define BLOCK_EDGE_LENGTH 32
 #define BLOCK_PIXELS (BLOCK_EDGE_LENGTH * BLOCK_EDGE_LENGTH)
+
+#ifndef FEATURE_COUNT
 #define FEATURE_COUNT 10
+#endif
+
+// Bit is set if corresponding feature needs normalization
+#ifndef NORMALIZATION_MASK
+#define NORMALIZATION_MASK 1008u // 0b1111110000
+#endif
 
 #define BLOCK_OFFSETS_COUNT 16u
 ivec2 BLOCK_OFFSETS[BLOCK_OFFSETS_COUNT] = ivec2[](
@@ -41,4 +49,18 @@ float scale(float value, float min, float max)
     return value - min;
 }
 
-#endif
+#ifndef FEATURES
+#define FEATURES \
+        1.f,        \
+        curr_normal.x, \
+        curr_normal.y, \
+        curr_normal.z, \
+        curr_pos.x,    \
+        curr_pos.y,    \
+        curr_pos.z,    \
+        curr_pos.x * curr_pos.x, \
+        curr_pos.y * curr_pos.y, \
+        curr_pos.z * curr_pos.z,
+#endif
+
+#endif

shader/bmfr_fit.comp

@@ -11,9 +11,14 @@
 #include "math.glsl"
 
 #ifdef DIFFUSE_ONLY
-#define R_SIZE ((11) * (11))
+#define R_SIZE ((FEATURE_COUNT + 1) * (FEATURE_COUNT + 1))
 #else
-#define R_SIZE ((12) * (12))
+#define R_SIZE ((FEATURE_COUNT + 2) * (FEATURE_COUNT + 2))
+#endif
+
+#if defined(BD_FEATURE_COUNT)
+    #undef FEATURE_COUNT
+    #define FEATURE_COUNT BD_FEATURE_COUNT
 #endif
 
 #define X_ITERS 2
@@ -52,7 +57,6 @@ shared float block_min_f, block_max_f, vec_length, u_length_squared, v_dot_a;
 shared float16_t u_vec[BLOCK_SIZE * BLOCK_SIZE];
 shared vec3 r_mat[R_SIZE];
 
-
 int get_linear_id()
 {
     return int(gl_LocalInvocationID.y * gl_WorkGroupSize.x + gl_LocalInvocationID.x);
@@ -183,8 +187,11 @@ void main()
     indices += buffer_offset;
 
     // Normalize world_pos and world_pos² 
-    for (int feature_buffer = 4; feature_buffer < FEATURE_COUNT; ++feature_buffer)
+    for (int feature_buffer = 0; feature_buffer < FEATURE_COUNT; ++feature_buffer)
     {
+        bool needs_normalization = ((1 << feature_buffer) & NORMALIZATION_MASK) != 0;
+        if (!needs_normalization) continue;
+
         float tmp_max = -FLT_MAX;
         float tmp_min = FLT_MAX;
 
@@ -205,11 +212,12 @@ void main()
         parallel_reduction_min();
         barrier();
 
-
         if (id == 0)
         {
-            const int index = (get_linear_group_id() * 6 + feature_buffer - 4) * 2;
-            const uint offset = control.workset_size.x * control.workset_size.y * 6 * 2 * p.z;
+            const int normalized_feature_count = bitCount(NORMALIZATION_MASK);
+            int normalize_index =  bitCount(((1 << feature_buffer) - 1) & NORMALIZATION_MASK);
+            const int index = (get_linear_group_id() * normalized_feature_count + normalize_index) * 2;
+            const uint offset = control.workset_size.x * control.workset_size.y * normalized_feature_count * 2 * p.z;
             mins_maxs_buffer.minmax[index + offset + 0] = float16_t(block_min_f);
             mins_maxs_buffer.minmax[index + offset + 1] = float16_t(block_max_f);
         }
@@ -274,7 +282,7 @@ void main()
 
         const uint id_limited = min(id, FEATURE_COUNT);
 #ifdef DIFFUSE_ONLY
-        const uint r_access = 11 * col_limited + id_limited;
+        const uint r_access = (FEATURE_COUNT + 1) * col_limited + id_limited;
 #else
         uint limit = col_limited;
         if (col > 12) limit = 11;
@@ -378,10 +386,10 @@ void main()
         barrier();
     }
 #else
-    vec3 x_s[10];
-    vec3 x_s2[10];
-    const int r_size = 12;
-    const int n = 10;
+    vec3 x_s[FEATURE_COUNT];
+    vec3 x_s2[FEATURE_COUNT];
+    const int r_size = FEATURE_COUNT + 2;
+    const int n = FEATURE_COUNT;
     for (int i = n-1; i >=0; --i)
     {
         x_s[i] = r_mat[(r_size) * i + (r_size - 2)];
@@ -408,7 +416,7 @@ void main()
         weights_buffer.weights[index + offset] = weight;
 #else
         weights_buffer.weights[index + offset] = x_s[id];
-        weights_buffer.weights[index + offset + 10] = x_s2[id];
+        weights_buffer.weights[index + offset + FEATURE_COUNT] = x_s2[id];
 #endif
     }
-}
+}

shader/bmfr_preprocess.comp

@@ -11,6 +11,21 @@
 #define COS_LIMIT    0.90f
 #define DISCARD_LIMIT 0.50f
 
+#ifndef FEATURES
+#define FEATURES \
+        1.f,        \
+        curr_normal.x, \
+        curr_normal.y, \
+        curr_normal.z, \
+        curr_pos.x,    \
+        curr_pos.y,    \
+        curr_pos.z,    \
+        curr_pos.x * curr_pos.x, \
+        curr_pos.y * curr_pos.y, \
+        curr_pos.z * curr_pos.z,
+
+#endif
+
 layout (local_size_x = 16, local_size_y = 16) in;
 
 layout(binding = 0, set = 0, rgba16f) uniform image2DArray in_color;
@@ -40,6 +55,10 @@ layout(binding = 14, set = 0) buffer accept_buf_t
     uint8_t accepts[];
 } accept_buffer;
 
+layout(binding = 15, set = 0, rgba32f) uniform image2DArray in_extra;
+layout(binding = 16, set = 0, r32i) uniform iimage2DArray in_instance_id;
+layout(binding = 17, set = 0, rgba32f) uniform image2DArray in_material;
+
 vec3 get_specular(ivec3 p)
 {
     return max(vec3(0.0), imageLoad(in_color, p).xyz - (imageLoad(in_albedo, p).xyz * imageLoad(in_diffuse, p).xyz));
@@ -111,6 +130,8 @@ void main()
     vec4 curr_color = imageLoad(in_color, p);
     vec3 curr_normal = unpack_gbuffer_normal(imageLoad(in_normal, p).xy);
     vec3 curr_pos = imageLoad(in_pos, p).xyz;
+    vec4 bd = imageLoad(in_extra, p);
+
     float sum_w = 0.0;
     {
         vec2 motion = vec2(imageLoad(in_screen_motion,p));
@@ -185,26 +206,11 @@ void main()
         imageStore(tmp_noisy[1], p, specular);
     }
 
-
     const float features[BUFFER_COUNT] = {
-        1.f,
-        curr_normal.x,
-        curr_normal.y,
-        curr_normal.z,
-        curr_pos.x,
-        curr_pos.y,
-        curr_pos.z,
-        curr_pos.x * curr_pos.x,
-        curr_pos.y * curr_pos.y,
-        curr_pos.z * curr_pos.z,
+        FEATURES
         diffuse.r,
         diffuse.g,
-        diffuse.b,
-#ifndef DIFFUSE_ONLY
-        specular.r,
-        specular.g,
-        specular.b
-#endif
+        diffuse.b
     };
 
     const int block_edge_length = BLOCK_EDGE_LENGTH;
@@ -232,4 +238,4 @@ void main()
         tmp_buffer.data[location_in_data + buffer_offset] = float16_t(store_value);
     }
 
-}
+}

shader/bmfr_weighted_sum.comp

@@ -20,11 +20,17 @@ layout(binding = 6, set = 0) buffer mins_maxs
 {
     float16_t minmax[];
 } mins_maxs_buffer;
+
+
 layout(binding = 7, set = 0) uniform uniform_buffer_t
 {
     uint frame_counter;
 } uniform_buffer;
 
+layout(binding = 8, set = 0, rgba32f) uniform readonly image2DArray in_extra;
+layout(binding = 9, set = 0, r32i) uniform iimage2DArray in_instance_id;
+layout(binding = 10, set = 0, rgba32f) uniform image2DArray in_material;
+
 void main()
 {
     const ivec3 p = ivec3(gl_GlobalInvocationID).xyz;
@@ -40,51 +46,35 @@ void main()
 
     vec3 curr_pos = imageLoad(in_pos, p).xyz;
     vec3 curr_normal = unpack_gbuffer_normal(imageLoad(in_normal, p).xy);
+    vec4 bd = imageLoad(in_extra, p);
 
-    const float features[FEATURE_COUNT] = {
-        1.f,
-        curr_normal.x,
-        curr_normal.y,
-        curr_normal.z,
-        curr_pos.x,
-        curr_pos.y,
-        curr_pos.z,
-        curr_pos.x * curr_pos.x,
-        curr_pos.y * curr_pos.y,
-        curr_pos.z * curr_pos.z,
+    const float features[FEATURE_COUNT] =
+    {
+        FEATURES
     };
 
     // Weighted sum of the feature buffers
     vec3 color = vec3(0);
-#ifndef DIFFUSE_ONLY
-    vec3 color_spec = vec3(0);
-#endif
     for (int feature_buffer = 0; feature_buffer < FEATURE_COUNT; ++feature_buffer)
     {
         float feature = features[feature_buffer];
-        if (feature_buffer >= 4)
-        {
-            const int min_max_index = (group_index * 6 + feature_buffer - 4) * 2;
-            const uint offset = control.workset_size.x * control.workset_size.y * 6 * 2 * p.z;
+        if (((1 << feature_buffer) & NORMALIZATION_MASK) != 0)
+        { // Normalize feature
+            const int num_normalized_features = bitCount(NORMALIZATION_MASK);
+            const int normalize_index = bitCount(((1 << feature_buffer) - 1) & NORMALIZATION_MASK);
+            const int min_max_index = (group_index * num_normalized_features + normalize_index) * 2;
+            const uint offset = control.workset_size.x * control.workset_size.y * num_normalized_features * 2 * p.z;
             feature = scale(feature, float(mins_maxs_buffer.minmax[offset + min_max_index + 0]), float(mins_maxs_buffer.minmax[offset + min_max_index + 1]));
         }
 
         // Load weight and sum
         const uint offset = control.workset_size.x * control.workset_size.y * FEATURE_COUNT * NUM_WEIGHTS_PER_FEATURE * p.z; 
         vec3 weight = weights_buffer.weights[offset + group_index * FEATURE_COUNT * NUM_WEIGHTS_PER_FEATURE + feature_buffer];
+
         color += weight * feature;
-#ifndef DIFFUSE_ONLY
-        vec3 weight_spec = weights_buffer.weights[offset + group_index * FEATURE_COUNT * NUM_WEIGHTS_PER_FEATURE + feature_buffer + 10];
-        color_spec += weight_spec * feature;
-#endif
     }
 
     color = max(color, vec3(0));
-#ifndef DIFFUSE_ONLY
-    color_spec = max(color_spec, vec3(0));
-#endif
+    //color = imageLoad(in_color, p).rgb;
     imageStore(weighted_out[0], p, vec4(color.rgb, 1.0));
-#ifndef DIFFUSE_ONLY
-    imageStore(weighted_out[1], p, vec4(color_spec.rgb, 1.0));
-#endif
-}
+}

shader/gbuffer.glsl

@@ -575,5 +575,42 @@ vec3 read_gbuffer_linear_depth(ivec3 pos) { return vec3(0.0); }
 
 #endif
 
+//==============================================================================
+// PROB
+//==============================================================================
+
+#ifdef PROB_TARGET_BINDING
+layout(binding = PROB_TARGET_BINDING, set = 0, rgba32f) uniform image2DArray prob_target;
+
+void write_gbuffer_prob(vec4 prob, ivec3 pos)
+{
+    imageStore(prob_target, pos, vec4(prob.rgb, 0));
+}
+
+vec3 read_gbuffer_prob(ivec3 pos)
+{
+    return imageLoad(prob_target, pos).xyz;
+}
+
+#elif defined(PROB_TARGET_LOCATION)
+
+layout(location = PROB_TARGET_LOCATION) out vec3 prob_target;
+
+void write_gbuffer_prob(vec4 prob)
+{
+    prob_target = prob;
+}
+
+#else
+
+void write_gbuffer_prob(vec4 prob, ivec3 pos) {
+    if(pos.xy == ivec2(640, 360))
+        debugPrintfEXT("jei");
+}
+
+void write_gbuffer_prob(vec4 prob) {}
+vec3 read_gbuffer_prob(ivec3 pos) { return vec3(0); }
+
+#endif
 #endif
 

shader/material.glsl

@@ -54,6 +54,19 @@ vec3 modulate_bsdf(sampled_material mat, bsdf_lobes bsdf)
     return mat.albedo.rgb * (bsdf.metallic_reflection + bsdf.transmission + bsdf.diffuse) + bsdf.dielectric_reflection;
 }
 
+vec3 modulate_diffuse(sampled_material mat, float diffuse)
+{
+    return diffuse * mat.albedo.rgb * (1 - mat.metallic);
+}
+
+vec3 modulate_reflection(sampled_material mat, float reflected)
+{
+    float approx_fresnel = 0.02f;
+    return reflected
+        * mix(vec3(approx_fresnel), mat.albedo.rgb, mat.metallic)
+        / mix(approx_fresnel, 1, mat.metallic);
+}
+
 vec3 modulate_color(sampled_material mat, vec3 diffuse, vec3 reflected)
 {
     float approx_fresnel = 0.02f;