shader, renderer/vulkan: Perform gamma correction within shaders when…

… using shader interlock (Vita3K#3298)

vita3k/renderer/include/renderer/vulkan/pipeline_cache.h

@@ -81,7 +81,7 @@ class PipelineCache {
     unordered_map_stable<Sha256Hash, vk::ShaderModule> shaders;
     unordered_map_stable<uint64_t, vk::Pipeline> pipelines;
 
-    vk::PipelineShaderStageCreateInfo retrieve_shader(const SceGxmProgram *program, const Sha256Hash &hash, bool is_vertex, bool maskupdate, MemState &mem, const shader::Hints &hints);
+    vk::PipelineShaderStageCreateInfo retrieve_shader(const SceGxmProgram *program, const Sha256Hash &hash, bool is_vertex, bool maskupdate, MemState &mem, const shader::Hints &hints, bool is_srgb = false);
     vk::PipelineVertexInputStateCreateInfo get_vertex_input_state(const SceGxmVertexProgram &vertex_program, MemState &mem);
 
     // queue containing request sent by the main thread to the compile threads

vita3k/renderer/src/vulkan/pipeline_cache.cpp

@@ -353,12 +353,43 @@ void PipelineCache::save_pipeline_cache() {
     LOG_INFO("Pipeline cache saved");
 }
 
-vk::PipelineShaderStageCreateInfo PipelineCache::retrieve_shader(const SceGxmProgram *program, const Sha256Hash &hash, bool is_vertex, bool maskupdate, MemState &mem, const shader::Hints &hints) {
+// Vulkan structs used to specify a specialization constant
+// Also, booleans in SPIRV are 32bit wide
+static const vk::SpecializationMapEntry srgb_entry = {
+    .constantID = shader::GAMMA_CORRECTION_SPECIALIZATION_ID,
+    .offset = 0,
+    .size = sizeof(uint32_t)
+};
+
+static const uint32_t srgb_entry_true = vk::True;
+static const uint32_t srgb_entry_false = vk::False;
+
+static const vk::SpecializationInfo srgb_info_true = {
+    .mapEntryCount = 1,
+    .pMapEntries = &srgb_entry,
+    .dataSize = sizeof(uint32_t),
+    .pData = &srgb_entry_true
+};
+
+static const vk::SpecializationInfo srgb_info_false = {
+    .mapEntryCount = 1,
+    .pMapEntries = &srgb_entry,
+    .dataSize = sizeof(uint32_t),
+    .pData = &srgb_entry_false
+};
+
+vk::PipelineShaderStageCreateInfo PipelineCache::retrieve_shader(const SceGxmProgram *program, const Sha256Hash &hash, bool is_vertex, bool maskupdate, MemState &mem, const shader::Hints &hints, bool is_srgb) {
     if (maskupdate)
         LOG_CRITICAL("Mask not implemented in the vulkan renderer!");
 
     const vk::ShaderModule shader_compiling = std::bit_cast<vk::ShaderModule>(~0ULL);
 
+    const vk::SpecializationInfo *spec_info = nullptr;
+    if (!is_vertex && state.features.should_use_shader_interlock() && program->is_frag_color_used()) {
+        // if the specialization constant is used in the shader
+        spec_info = is_srgb ? &srgb_info_true : &srgb_info_false;
+    }
+
     vk::ShaderModule *shader_module;
     {
         // look if it is in the cache
@@ -387,7 +418,8 @@ vk::PipelineShaderStageCreateInfo PipelineCache::retrieve_shader(const SceGxmPro
         vk::PipelineShaderStageCreateInfo shader_stage_info{
             .stage = is_vertex ? vk::ShaderStageFlagBits::eVertex : vk::ShaderStageFlagBits::eFragment,
             .module = *shader_module,
-            .pName = is_vertex ? "main_vs" : "main_fs"
+            .pName = is_vertex ? "main_vs" : "main_fs",
+            .pSpecializationInfo = spec_info,
         };
         return shader_stage_info;
     }
@@ -420,7 +452,8 @@ vk::PipelineShaderStageCreateInfo PipelineCache::retrieve_shader(const SceGxmPro
     vk::PipelineShaderStageCreateInfo shader_stage_info{
         .stage = is_vertex ? vk::ShaderStageFlagBits::eVertex : vk::ShaderStageFlagBits::eFragment,
         .module = *shader_module,
-        .pName = is_vertex ? "main_vs" : "main_fs"
+        .pName = is_vertex ? "main_vs" : "main_fs",
+        .pSpecializationInfo = spec_info,
     };
 
     return shader_stage_info;
@@ -707,7 +740,7 @@ vk::Pipeline PipelineCache::compile_pipeline(SceGxmPrimitiveType type, vk::Rende
     const vk::PipelineVertexInputStateCreateInfo vertex_input = get_vertex_input_state(vertex_program_gxm, mem);
 
     const vk::PipelineShaderStageCreateInfo vertex_shader = retrieve_shader(vertex_program_gxm.program.get(mem), vertex_program.hash, true, fragment_program_gxm.is_maskupdate, mem, hints);
-    const vk::PipelineShaderStageCreateInfo fragment_shader = retrieve_shader(gxm_fragment_shader, fragment_program.hash, false, fragment_program_gxm.is_maskupdate, mem, hints);
+    const vk::PipelineShaderStageCreateInfo fragment_shader = retrieve_shader(gxm_fragment_shader, fragment_program.hash, false, fragment_program_gxm.is_maskupdate, mem, hints, record.is_gamma_corrected);
     const vk::PipelineShaderStageCreateInfo shader_stages[] = { vertex_shader, fragment_shader };
     // disable the fragment shader if gxm asks us to
     const bool is_fragment_disabled = record.front_side_fragment_program_mode == SCE_GXM_FRAGMENT_PROGRAM_DISABLED || gxm_fragment_shader->has_no_effect();

vita3k/shader/include/shader/spirv_recompiler.h

@@ -38,6 +38,10 @@ static constexpr int COLOR_ATTACHMENT_TEXTURE_SLOT_IMAGE = 0;
 static constexpr int MASK_TEXTURE_SLOT_IMAGE = 1;
 static constexpr int COLOR_ATTACHMENT_RAW_TEXTURE_SLOT_IMAGE = 3;
 static constexpr uint32_t CURRENT_VERSION = 13;
+// fragment shader using the rendering surface as a storage image (because of shader interlock) have a line
+// layout (constant_id = GAMMA_CORRECTION_SPECIALIZATIO_ID) const bool is_srgb = false;
+// Setting this constant to true performs gamma correction in the shader
+static constexpr uint32_t GAMMA_CORRECTION_SPECIALIZATION_ID = 0;
 
 enum struct Target {
     GLSLOpenGL,

vita3k/shader/include/shader/usse_translator_types.h

@@ -106,6 +106,9 @@ struct SpirvShaderParameters {
     spv::Id thread_buffer;
 
     spv::Id render_info_id;
+
+    // When using shader interlock, specialization constant telling us if the texture is gamma corrected
+    spv::Id is_srgb_constant;
 };
 
 using Coord = std::pair<spv::Id, int>;

vita3k/shader/src/spirv_recompiler.cpp

@@ -771,6 +771,29 @@ static void create_fragment_inputs(spv::Builder &b, SpirvShaderParameters &param
                 source = spv::NoResult;
             } else {
                 source = b.createOp(spv::OpImageRead, v4, { b.createLoad(color_attachment, spv::NoPrecision), current_coord });
+
+                if (translation_state.is_vulkan) {
+                    const spv::Id old_source = source;
+                    // if (is_srgb)
+                    spv::Builder::If cond_builder(parameters.is_srgb_constant, spv::SelectionControlMaskNone, b);
+
+                    // perform gamma correction:
+                    // source.xyz = pow(source.xyz, vec3(2.2));
+                    const spv::Id v3 = b.makeVectorType(f32, 3);
+                    spv::Id rgb = b.createOp(spv::OpVectorShuffle, v3, { { true, source }, { true, source }, { false, 0 }, { false, 1 }, { false, 2 } });
+                    const spv::Id gamma = utils::make_uniform_vector_from_type(b, v3, 2.2f);
+                    rgb = b.createBuiltinCall(v3, utils.std_builtins, GLSLstd450Pow, { rgb, gamma });
+                    source = b.createOp(spv::OpVectorShuffle, v4, { { true, rgb }, { true, source }, { false, 0 }, { false, 1 }, { false, 2 }, { false, 6 } });
+
+                    store_source_result();
+
+                    // else (no shader gamma correction, nothing to do)
+                    cond_builder.makeBeginElse();
+                    source = old_source;
+                    store_source_result();
+                    cond_builder.makeEndIf();
+                    source = 0;
+                }
             }
         } else {
             // Try to initialize outs[0] to some nice value. In case the GPU has garbage data for our shader
@@ -1050,6 +1073,14 @@ static SpirvShaderParameters create_parameters(spv::Builder &b, const SceGxmProg
         b.addDecoration(translation_state.render_info_id, spv::DecorationBinding, translation_state.is_vulkan ? 1 : 3);
         if (translation_state.is_vulkan)
             b.addDecoration(translation_state.render_info_id, spv::DecorationDescriptorSet, 0);
+
+        if (program.is_frag_color_used() && features.should_use_shader_interlock() && translation_state.is_vulkan) {
+            // specialization constant for shader interlock:
+            // layout (constant_id = GAMMA_CORRECTION_SPECIALIZATION_ID) const bool is_srgb = false;
+            spv_params.is_srgb_constant = b.makeBoolConstant(false, true);
+            b.addDecoration(spv_params.is_srgb_constant, spv::DecorationSpecId, (int)GAMMA_CORRECTION_SPECIALIZATION_ID);
+            b.addName(spv_params.is_srgb_constant, "is_srgb");
+        }
     }
 
     spv_params.render_info_id = translation_state.render_info_id;
@@ -1395,7 +1426,31 @@ static spv::Function *make_frag_finalize_function(spv::Builder &b, const SpirvSh
         spv::Id coord_id = b.createLoad(translate_state.frag_coord_id, spv::NoPrecision);
         spv::Id translated_id = b.createUnaryOp(spv::OpConvertFToS, b.makeVectorType(signed_i32, 4), coord_id);
         translated_id = b.createOp(spv::OpVectorShuffle, b.makeVectorType(signed_i32, 2), { { true, translated_id }, { true, translated_id }, { false, 0 }, { false, 1 } });
-        b.createNoResultOp(spv::OpImageWrite, { b.createLoad(translate_state.color_attachment_id, spv::NoPrecision), translated_id, color });
+
+        if (translate_state.is_vulkan) {
+            spv::Id old_color = color;
+            // if (is_srgb)
+            spv::Builder::If cond_builder(parameters.is_srgb_constant, spv::SelectionControlMaskNone, b);
+
+            // perform inverse gamma correction:
+            // color.xyz = pow(color.xyz, vec3(1/2.2));
+            const spv::Id f32 = b.makeFloatType(32);
+            const spv::Id v4 = b.makeVectorType(f32, 4);
+            const spv::Id v3 = b.makeVectorType(f32, 3);
+            spv::Id rgb = b.createOp(spv::OpVectorShuffle, v3, { { true, color }, { true, color }, { false, 0 }, { false, 1 }, { false, 2 } });
+            const spv::Id gamma = utils::make_uniform_vector_from_type(b, v3, 1 / 2.2f);
+            rgb = b.createBuiltinCall(v3, utils.std_builtins, GLSLstd450Pow, { rgb, gamma });
+            color = b.createOp(spv::OpVectorShuffle, v4, { { true, rgb }, { true, color }, { false, 0 }, { false, 1 }, { false, 2 }, { false, 6 } });
+
+            b.createNoResultOp(spv::OpImageWrite, { b.createLoad(translate_state.color_attachment_id, spv::NoPrecision), translated_id, color });
+
+            // else (no shader gamma correction, nothing to do)
+            cond_builder.makeBeginElse();
+            b.createNoResultOp(spv::OpImageWrite, { b.createLoad(translate_state.color_attachment_id, spv::NoPrecision), translated_id, old_color });
+            cond_builder.makeEndIf();
+        } else {
+            b.createNoResultOp(spv::OpImageWrite, { b.createLoad(translate_state.color_attachment_id, spv::NoPrecision), translated_id, color });
+        }
 
         if (features.preserve_f16_nan_as_u16) {
             color_val_operand.type = DataType::UINT16;