Enable metallic, roughness and instance-id options in BMFR

Material and instance-id gbuffers are required only when the options require them, and the code to read them is injected with preprocessor macros. Therefore performance should be unaffected when they are not used.
vga-group · Jan 18, 2025 · 244f6fa · 244f6fa
1 parent ea74316
commit 244f6fa
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 32 deletions.
diff --git a/shader/bmfr_preprocess.comp b/shader/bmfr_preprocess.comp
@@ -56,6 +56,8 @@ layout(binding = 14, set = 0) buffer accept_buf_t
 } 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)
 {

diff --git a/shader/bmfr_weighted_sum.comp b/shader/bmfr_weighted_sum.comp
@@ -28,6 +28,8 @@ layout(binding = 7, set = 0) uniform uniform_buffer_t
 } 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()
 {

diff --git a/src/bmfr_stage.cc b/src/bmfr_stage.cc
@@ -63,15 +63,18 @@ bmfr_stage::bmfr_stage(
 }
 
 static void create_bd_macro(
-        std::vector<std::string> source, 
+        const std::vector<std::string>& source, 
         uint32_t& bd_feature_count, 
         std::string& macro
 ){
-    std::map<std::string, std::string> macro_map =
+    const std::map<std::string, std::string> macro_map =
     {
         {"normal", "curr_normal"},
         {"position", "curr_pos"},
         {"position-2", "curr_pos"},
+        {"roughness", "imageLoad(in_material, p).g"},
+        {"metallic", "imageLoad(in_material, p).r"},
+        {"instance-id", "float(imageLoad(in_instance_id, p).r)"},
 
         {"bounce-count", "bd"},
         {"bounce-contribution", "bd"},
@@ -82,7 +85,7 @@ static void create_bd_macro(
         {"bsdf-nee-contribution", "bd"}
     };
 
-    std::vector<std::string> channel_map = {"x", "y", "z", "w"};
+    const std::vector<std::string> channel_map = {"x", "y", "z", "w"};
 
     macro += "1.f,";
 
@@ -104,6 +107,11 @@ static void create_bd_macro(
             macro += macro_map.at(a) + ".z * " + macro_map.at(a) + ".z,";
             count += 3;
         }
+        else if (a == "roughness" || a == "metallic" || a == "instance-id")
+        {
+            macro += macro_map.at(a) + ",";
+            count += 1;
+        }
         else
         {
             macro += macro_map.at(a) + "." + channel_map.at(indx % 4) + ",";
@@ -130,12 +138,8 @@ shader_source bmfr_stage::load_shader_source(const std::string& path, const opti
         create_bd_macro(opt.bd_vec, _feature_count, macro);
         this->feature_count = _feature_count;
 
-        // features to normalize
-        std::vector<std::string> to_normalize_v =
-        {
-            "position",
-            "position-2"
-        };
+        // Normalization mask is uint32
+        assert(feature_count <= 32); 
 
         uint32_t num_normalized_features = 0;
 
@@ -169,9 +173,6 @@ shader_source bmfr_stage::load_shader_source(const std::string& path, const opti
         defines.insert({ "FEATURES", macro });
         defines.insert({ "FEATURE_COUNT", std::to_string(_feature_count) });
         defines.insert({ "NORMALIZATION_MASK", std::to_string(normalization_mask) });
-
-        std::replace(macro.begin(), macro.end(), ',', '\n');
-        TR_LOG(macro);
     }
 
     if (opt.settings == bmfr_settings::DIFFUSE_ONLY)
@@ -185,11 +186,6 @@ shader_source bmfr_stage::load_shader_source(const std::string& path, const opti
         defines.insert({ "BUFFER_COUNT", "(FEATURE_COUNT + 6)" });
         defines.insert({ "NUM_WEIGHTS_PER_FEATURE", "2" });
     }
-
-    for(auto& a : defines)
-    {
-        TR_LOG(a.first, " : ", a.second);
-    }
 
     return shader_source(path, defines);
 }
@@ -325,7 +321,9 @@ void bmfr_stage::init_resources()
         bmfr_preprocess_desc.set_buffer(dev->id, i, "tmp_buffer", {{tmp_data[i], 0, VK_WHOLE_SIZE}});
         bmfr_preprocess_desc.set_buffer(dev->id, i, "uniform_buffer", {{uniform_buffer[dev->id], 0, VK_WHOLE_SIZE}});
         bmfr_preprocess_desc.set_buffer(dev->id, i, "accept_buffer", {{accepts[i], 0, VK_WHOLE_SIZE}});
-        bmfr_preprocess_desc.set_image(dev->id, i, "in_extra", { {{}, current_features.prob.view, vk::ImageLayout::eGeneral}});
+        bmfr_preprocess_desc.set_image(dev->id, i, "in_extra", { {{}, current_features.prob.view, vk::ImageLayout::eGeneral} });
+        bmfr_preprocess_desc.set_image(dev->id, i, "in_instance_id", { {{}, current_features.instance_id.view, vk::ImageLayout::eGeneral} });
+        bmfr_preprocess_desc.set_image(dev->id, i, "in_material", { {{}, current_features.material.view, vk::ImageLayout::eGeneral}});
 #if 1
         bmfr_fit_desc.set_buffer(dev->id, i, "tmp_buffer", {{tmp_data[i], 0, VK_WHOLE_SIZE}});
         bmfr_fit_desc.set_buffer(dev->id, i, "mins_maxs_buffer", {{min_max_buffer[i], 0, VK_WHOLE_SIZE}});
@@ -345,6 +343,8 @@ void bmfr_stage::init_resources()
         bmfr_weighted_sum_desc.set_image(dev->id, i, "weighted_out", {{{}, weighted_sum[0].view, vk::ImageLayout::eGeneral}, {{}, weighted_sum[1].view, vk::ImageLayout::eGeneral}});
         bmfr_weighted_sum_desc.set_image(dev->id, i, "tmp_noisy", {{{}, tmp_noisy[0].view, vk::ImageLayout::eGeneral}, {{}, tmp_noisy[1].view, vk::ImageLayout::eGeneral}});
         bmfr_weighted_sum_desc.set_image(dev->id, i, "in_extra", { {{}, current_features.prob.view, vk::ImageLayout::eGeneral} });
+        bmfr_weighted_sum_desc.set_image(dev->id, i, "in_instance_id", { {{}, current_features.instance_id.view, vk::ImageLayout::eGeneral} });
+        bmfr_weighted_sum_desc.set_image(dev->id, i, "in_material", { {{}, current_features.material.view, vk::ImageLayout::eGeneral} });
 
         bmfr_accumulate_output_desc.set_image(dev->id, i, "out_color", {{{}, current_features.color.view, vk::ImageLayout::eGeneral}});
         bmfr_accumulate_output_desc.set_image(dev->id, i, "in_screen_motion", {{{}, current_features.screen_motion.view, vk::ImageLayout::eGeneral}});

diff --git a/src/path_tracer_stage.cc b/src/path_tracer_stage.cc
@@ -104,12 +104,9 @@ path_tracer_stage::path_tracer_stage(
     for(auto i = 0; i < (int)opt.bd_vec.size(); i++)
     {
         auto option = opt.bd_vec.at(i);
-        if(option != bd::POSITION && option != bd::POSITION_2 && option != bd::NORMAL)
-        {
-            defines[bd_str_map.at(option)] = std::to_string(id);
-            TR_LOG(bd_str_map.at(option) + " " + std::to_string(id));
-            ++id;
-        }
+        defines[bd_str_map.at(option)] = std::to_string(id);
+        //TR_LOG(bd_str_map.at(option) + " " + std::to_string(id));
+        ++id;
     }
 
 #define TR_GBUFFER_ENTRY(name, ...)\

diff --git a/src/post_processing_renderer.cc b/src/post_processing_renderer.cc
@@ -1,4 +1,5 @@
 #include "post_processing_renderer.hh"
+#include "log.hh"
 
 namespace tr
 {
@@ -58,6 +59,23 @@ void post_processing_renderer::set_gbuffer_spec(gbuffer_spec& spec) const
         spec.albedo_present = true;
         spec.diffuse_present = true;
         spec.prob_present = true;
+
+        bool needs_material = std::find(opt.bmfr->bd_vec.begin(), opt.bmfr->bd_vec.end(), "metallic") != opt.bmfr->bd_vec.end() ||
+                std::find(opt.bmfr->bd_vec.begin(), opt.bmfr->bd_vec.end(), "roughness") != opt.bmfr->bd_vec.end();
+
+        bool needs_instance_id = std::find(opt.bmfr->bd_vec.begin(), opt.bmfr->bd_vec.end(), "instance-id") != opt.bmfr->bd_vec.end();
+
+        if (needs_material)
+        {
+            TR_LOG("BMFR: Enabling material in gbuffer due to specified options\n");
+            spec.material_present = true;
+        }
+
+        if (needs_instance_id)
+        {
+            TR_LOG("BMFR: Enabling instance-id in gbuffer due to specified options\n");
+            spec.instance_id_present = true;
+        }
     }
 
     if(opt.taa.has_value())

diff --git a/src/rt_common.hh b/src/rt_common.hh
@@ -31,9 +31,6 @@ namespace tr
         BSDF_VAR,
         PDF_CONTRIBUTION,
         FULL_PDF_CONTRIBUTION,
-        NORMAL,
-        POSITION,
-        POSITION_2
     };
 
     enum class denoiser_type
@@ -52,9 +49,6 @@ namespace tr
         {"bsdf-variance", bd::BSDF_VAR},
         {"bsdf-contribution", bd::PDF_CONTRIBUTION},
         {"bsdf-nee-contribution", bd::FULL_PDF_CONTRIBUTION},
-        {"normal", bd::NORMAL},
-        {"position", bd::POSITION},
-        {"position-2", bd::POSITION_2}
     };
 
     const std::map<bd, std::string> bd_str_map =

diff --git a/src/tauray.cc b/src/tauray.cc
@@ -491,8 +491,9 @@ renderer* create_renderer(context& ctx, options& opt, scene& s)
                 auto v = parse_bounce_data(opt.noise_data_str);
                 for(auto a : v)
                 {
-                    TR_LOG(a);
-                    rt_opt.bd_vec.push_back(str_bd_map.at(a));
+                    //TR_LOG(a);
+                    if (str_bd_map.count(a) != 0)
+                        rt_opt.bd_vec.push_back(str_bd_map.at(a));
                 }
 
                 rt_opt.dt = (tr::denoiser_type)opt.denoiser;