vulkan: Fix rescaling push constant usage

8 files changed, 78 insertions, 69 deletions

diff --git a/src/shader_recompiler/backend/spirv/emit_context.cpp b/src/shader_recompiler/backend/spirv/emit_context.cpp
index 8646fe989..723455462 100644
--- a/src/shader_recompiler/backend/spirv/emit_context.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_context.cpp
@@ -1006,47 +1006,47 @@ void EmitContext::DefineRescalingInput(const Info& info) {
         return;
     }
     if (profile.unified_descriptor_binding) {
-        DefineRescalingInputPushConstant(info);
+        DefineRescalingInputPushConstant();
     } else {
         DefineRescalingInputUniformConstant();
     }
 }
 
-void EmitContext::DefineRescalingInputPushConstant(const Info& info) {
-    boost::container::static_vector<Id, 3> members{F32[1]};
+void EmitContext::DefineRescalingInputPushConstant() {
+    boost::container::static_vector<Id, 3> members{};
     u32 member_index{0};
-    if (!info.texture_descriptors.empty()) {
-        rescaling_textures_type = TypeArray(U32[1], Const(4u));
-        Decorate(rescaling_textures_type, spv::Decoration::ArrayStride, 4u);
-        members.push_back(rescaling_textures_type);
-        rescaling_textures_member_index = ++member_index;
-    }
-    if (!info.image_descriptors.empty()) {
-        rescaling_images_type = TypeArray(U32[1], Const(NUM_IMAGE_SCALING_WORDS));
-        if (rescaling_textures_type.value != rescaling_images_type.value) {
-            Decorate(rescaling_images_type, spv::Decoration::ArrayStride, 4u);
-        }
-        members.push_back(rescaling_images_type);
-        rescaling_images_member_index = ++member_index;
+
+    rescaling_textures_type = TypeArray(U32[1], Const(4u));
+    Decorate(rescaling_textures_type, spv::Decoration::ArrayStride, 4u);
+    members.push_back(rescaling_textures_type);
+    rescaling_textures_member_index = member_index++;
+
+    rescaling_images_type = TypeArray(U32[1], Const(NUM_IMAGE_SCALING_WORDS));
+    Decorate(rescaling_images_type, spv::Decoration::ArrayStride, 4u);
+    members.push_back(rescaling_images_type);
+    rescaling_images_member_index = member_index++;
+
+    if (stage != Stage::Compute) {
+        members.push_back(F32[1]);
+        rescaling_downfactor_member_index = member_index++;
     }
     const Id push_constant_struct{TypeStruct(std::span(members.data(), members.size()))};
     Decorate(push_constant_struct, spv::Decoration::Block);
     Name(push_constant_struct, "ResolutionInfo");
 
-    MemberDecorate(push_constant_struct, 0u, spv::Decoration::Offset, 0u);
-    MemberName(push_constant_struct, 0u, "down_factor");
+    MemberDecorate(push_constant_struct, rescaling_textures_member_index, spv::Decoration::Offset,
+                   static_cast<u32>(offsetof(RescalingLayout, rescaling_textures)));
+    MemberName(push_constant_struct, rescaling_textures_member_index, "rescaling_textures");
 
-    const u32 offset_bias = stage == Stage::Compute ? sizeof(u32) : 0;
-    if (!info.texture_descriptors.empty()) {
-        MemberDecorate(
-            push_constant_struct, rescaling_textures_member_index, spv::Decoration::Offset,
-            static_cast<u32>(offsetof(RescalingLayout, rescaling_textures) - offset_bias));
-        MemberName(push_constant_struct, rescaling_textures_member_index, "rescaling_textures");
-    }
-    if (!info.image_descriptors.empty()) {
-        MemberDecorate(push_constant_struct, rescaling_images_member_index, spv::Decoration::Offset,
-                       static_cast<u32>(offsetof(RescalingLayout, rescaling_images) - offset_bias));
-        MemberName(push_constant_struct, rescaling_images_member_index, "rescaling_images");
+    MemberDecorate(push_constant_struct, rescaling_images_member_index, spv::Decoration::Offset,
+                   static_cast<u32>(offsetof(RescalingLayout, rescaling_images)));
+    MemberName(push_constant_struct, rescaling_images_member_index, "rescaling_images");
+
+    if (stage != Stage::Compute) {
+        MemberDecorate(push_constant_struct, rescaling_downfactor_member_index,
+                       spv::Decoration::Offset,
+                       static_cast<u32>(offsetof(RescalingLayout, down_factor)));
+        MemberName(push_constant_struct, rescaling_downfactor_member_index, "down_factor");
     }
     const Id pointer_type{TypePointer(spv::StorageClass::PushConstant, push_constant_struct)};
     rescaling_push_constants = AddGlobalVariable(pointer_type, spv::StorageClass::PushConstant);
diff --git a/src/shader_recompiler/backend/spirv/emit_context.h b/src/shader_recompiler/backend/spirv/emit_context.h
index b67704baa..63f8185d9 100644
--- a/src/shader_recompiler/backend/spirv/emit_context.h
+++ b/src/shader_recompiler/backend/spirv/emit_context.h
@@ -244,6 +244,7 @@ public:
     Id rescaling_images_type{};
     u32 rescaling_textures_member_index{};
     u32 rescaling_images_member_index{};
+    u32 rescaling_downfactor_member_index{};
     u32 texture_rescaling_index{};
     u32 image_rescaling_index{};
 
@@ -324,7 +325,7 @@ private:
     void DefineAttributeMemAccess(const Info& info);
     void DefineGlobalMemoryFunctions(const Info& info);
     void DefineRescalingInput(const Info& info);
-    void DefineRescalingInputPushConstant(const Info& info);
+    void DefineRescalingInputPushConstant();
     void DefineRescalingInputUniformConstant();
 
     void DefineInputs(const IR::Program& program);
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.h b/src/shader_recompiler/backend/spirv/emit_spirv.h
index cf59f2572..4b25534ce 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.h
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.h
@@ -22,11 +22,12 @@ constexpr u32 NUM_TEXTURE_AND_IMAGE_SCALING_WORDS =
     NUM_TEXTURE_SCALING_WORDS + NUM_IMAGE_SCALING_WORDS;
 
 struct RescalingLayout {
-    u32 down_factor;
     alignas(16) std::array<u32, NUM_TEXTURE_SCALING_WORDS> rescaling_textures;
     alignas(16) std::array<u32, NUM_IMAGE_SCALING_WORDS> rescaling_images;
+    alignas(16) u32 down_factor;
 };
-constexpr u32 RESCALING_PUSH_CONSTANT_WORDS_OFFSET = offsetof(RescalingLayout, rescaling_textures);
+constexpr u32 RESCALING_LAYOUT_WORDS_OFFSET = offsetof(RescalingLayout, rescaling_textures);
+constexpr u32 RESCALING_LAYOUT_DOWN_FACTOR_OFFSET = offsetof(RescalingLayout, down_factor);
 
 [[nodiscard]] std::vector<u32> EmitSPIRV(const Profile& profile, const RuntimeInfo& runtime_info,
                                          IR::Program& program, Bindings& bindings);
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp b/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
index c0db7452f..bac683ae1 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
@@ -529,8 +529,8 @@ Id EmitYDirection(EmitContext& ctx) {
 Id EmitResolutionDownFactor(EmitContext& ctx) {
     if (ctx.profile.unified_descriptor_binding) {
         const Id pointer_type{ctx.TypePointer(spv::StorageClass::PushConstant, ctx.F32[1])};
-        const Id pointer{
-            ctx.OpAccessChain(pointer_type, ctx.rescaling_push_constants, ctx.u32_zero_value)};
+        const Id index{ctx.Const(ctx.rescaling_downfactor_member_index)};
+        const Id pointer{ctx.OpAccessChain(pointer_type, ctx.rescaling_push_constants, index)};
         return ctx.OpLoad(ctx.F32[1], pointer);
     } else {
         const Id composite{ctx.OpLoad(ctx.F32[4], ctx.rescaling_uniform_constant)};
diff --git a/src/video_core/renderer_vulkan/pipeline_helper.h b/src/video_core/renderer_vulkan/pipeline_helper.h
index 3612e8a18..ae5e66ef4 100644
--- a/src/video_core/renderer_vulkan/pipeline_helper.h
+++ b/src/video_core/renderer_vulkan/pipeline_helper.h
@@ -22,7 +22,6 @@
 namespace Vulkan {
 
 using Shader::Backend::SPIRV::NUM_TEXTURE_AND_IMAGE_SCALING_WORDS;
-using Shader::Backend::SPIRV::RESCALING_PUSH_CONSTANT_WORDS_OFFSET;
 
 class DescriptorLayoutBuilder {
 public:
@@ -73,12 +72,12 @@ public:
 
     vk::PipelineLayout CreatePipelineLayout(VkDescriptorSetLayout descriptor_set_layout) const {
         using Shader::Backend::SPIRV::RescalingLayout;
-        const u32 push_offset = is_compute ? RESCALING_PUSH_CONSTANT_WORDS_OFFSET : 0;
+        const u32 size_offset = is_compute ? sizeof(RescalingLayout::down_factor) : 0u;
         const VkPushConstantRange range{
             .stageFlags = static_cast<VkShaderStageFlags>(
                 is_compute ? VK_SHADER_STAGE_COMPUTE_BIT : VK_SHADER_STAGE_ALL_GRAPHICS),
             .offset = 0,
-            .size = sizeof(RescalingLayout) - push_offset,
+            .size = sizeof(RescalingLayout) - size_offset,
         };
         return device->GetLogical().CreatePipelineLayout({
             .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
@@ -139,21 +138,21 @@ private:
 
 class RescalingPushConstant {
 public:
-    explicit RescalingPushConstant(u32 num_textures) noexcept {}
+    explicit RescalingPushConstant() noexcept {}
 
     void PushTexture(bool is_rescaled) noexcept {
-        *texture_ptr |= is_rescaled ? texture_bit : 0;
-        texture_bit <<= 1;
-        if (texture_bit == 0) {
+        *texture_ptr |= is_rescaled ? texture_bit : 0u;
+        texture_bit <<= 1u;
+        if (texture_bit == 0u) {
             texture_bit = 1u;
             ++texture_ptr;
         }
     }
 
     void PushImage(bool is_rescaled) noexcept {
-        *image_ptr |= is_rescaled ? image_bit : 0;
-        image_bit <<= 1;
-        if (image_bit == 0) {
+        *image_ptr |= is_rescaled ? image_bit : 0u;
+        image_bit <<= 1u;
+        if (image_bit == 0u) {
             image_bit = 1u;
             ++image_ptr;
         }
@@ -176,8 +175,10 @@ inline void PushImageDescriptors(TextureCache& texture_cache,
                                  const Shader::Info& info, RescalingPushConstant& rescaling,
                                  const VkSampler*& samplers,
                                  const VideoCommon::ImageViewInOut*& views) {
-    views += Shader::NumDescriptors(info.texture_buffer_descriptors);
-    views += Shader::NumDescriptors(info.image_buffer_descriptors);
+    const u32 num_texture_buffers = Shader::NumDescriptors(info.texture_buffer_descriptors);
+    const u32 num_image_buffers = Shader::NumDescriptors(info.image_buffer_descriptors);
+    views += num_texture_buffers;
+    views += num_image_buffers;
     for (const auto& desc : info.texture_descriptors) {
         for (u32 index = 0; index < desc.count; ++index) {
             const VideoCommon::ImageViewId image_view_id{(views++)->id};
diff --git a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
index 6dc52e399..de36bcdb7 100644
--- a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
@@ -22,6 +22,7 @@
 namespace Vulkan {
 
 using Shader::ImageBufferDescriptor;
+using Shader::Backend::SPIRV::RESCALING_LAYOUT_WORDS_OFFSET;
 using Tegra::Texture::TexturePair;
 
 ComputePipeline::ComputePipeline(const Device& device_, DescriptorPool& descriptor_pool,
@@ -185,7 +186,7 @@ void ComputePipeline::Configure(Tegra::Engines::KeplerCompute& kepler_compute,
     buffer_cache.UpdateComputeBuffers();
     buffer_cache.BindHostComputeBuffers();
 
-    RescalingPushConstant rescaling(num_textures);
+    RescalingPushConstant rescaling;
     const VkSampler* samplers_it{samplers.data()};
     const VideoCommon::ImageViewInOut* views_it{views.data()};
     PushImageDescriptors(texture_cache, update_descriptor_queue, info, rescaling, samplers_it,
@@ -199,21 +200,24 @@ void ComputePipeline::Configure(Tegra::Engines::KeplerCompute& kepler_compute,
         });
     }
     const void* const descriptor_data{update_descriptor_queue.UpdateData()};
-    scheduler.Record(
-        [this, descriptor_data, rescaling_data = rescaling.Data()](vk::CommandBuffer cmdbuf) {
-            cmdbuf.BindPipeline(VK_PIPELINE_BIND_POINT_COMPUTE, *pipeline);
-            if (!descriptor_set_layout) {
-                return;
-            }
-            if (num_textures > 0) {
-                cmdbuf.PushConstants(*pipeline_layout, VK_SHADER_STAGE_COMPUTE_BIT, rescaling_data);
-            }
-            const VkDescriptorSet descriptor_set{descriptor_allocator.Commit()};
-            const vk::Device& dev{device.GetLogical()};
-            dev.UpdateDescriptorSet(descriptor_set, *descriptor_update_template, descriptor_data);
-            cmdbuf.BindDescriptorSets(VK_PIPELINE_BIND_POINT_COMPUTE, *pipeline_layout, 0,
-                                      descriptor_set, nullptr);
-        });
+    const bool is_rescaling = !info.texture_descriptors.empty() || !info.image_descriptors.empty();
+    scheduler.Record([this, descriptor_data, is_rescaling,
+                      rescaling_data = rescaling.Data()](vk::CommandBuffer cmdbuf) {
+        cmdbuf.BindPipeline(VK_PIPELINE_BIND_POINT_COMPUTE, *pipeline);
+        if (!descriptor_set_layout) {
+            return;
+        }
+        if (is_rescaling) {
+            cmdbuf.PushConstants(*pipeline_layout, VK_SHADER_STAGE_COMPUTE_BIT,
+                                 RESCALING_LAYOUT_WORDS_OFFSET, sizeof(rescaling_data),
+                                 rescaling_data.data());
+        }
+        const VkDescriptorSet descriptor_set{descriptor_allocator.Commit()};
+        const vk::Device& dev{device.GetLogical()};
+        dev.UpdateDescriptorSet(descriptor_set, *descriptor_update_template, descriptor_data);
+        cmdbuf.BindDescriptorSets(VK_PIPELINE_BIND_POINT_COMPUTE, *pipeline_layout, 0,
+                                  descriptor_set, nullptr);
+    });
 }
 
 } // namespace Vulkan
diff --git a/src/video_core/renderer_vulkan/vk_compute_pipeline.h b/src/video_core/renderer_vulkan/vk_compute_pipeline.h
index e79ce4d7c..8c4b0a301 100644
--- a/src/video_core/renderer_vulkan/vk_compute_pipeline.h
+++ b/src/video_core/renderer_vulkan/vk_compute_pipeline.h
@@ -59,7 +59,6 @@ private:
     vk::PipelineLayout pipeline_layout;
     vk::DescriptorUpdateTemplateKHR descriptor_update_template;
     vk::Pipeline pipeline;
-    u32 num_textures{};
 
     std::condition_variable build_condvar;
     std::mutex build_mutex;
diff --git a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
index f08e9e840..616a7b457 100644
--- a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
@@ -32,6 +32,8 @@ namespace {
 using boost::container::small_vector;
 using boost::container::static_vector;
 using Shader::ImageBufferDescriptor;
+using Shader::Backend::SPIRV::RESCALING_LAYOUT_DOWN_FACTOR_OFFSET;
+using Shader::Backend::SPIRV::RESCALING_LAYOUT_WORDS_OFFSET;
 using Tegra::Texture::TexturePair;
 using VideoCore::Surface::PixelFormat;
 using VideoCore::Surface::PixelFormatFromDepthFormat;
@@ -431,7 +433,7 @@ void GraphicsPipeline::ConfigureImpl(bool is_indexed) {
 
     update_descriptor_queue.Acquire();
 
-    RescalingPushConstant rescaling(num_textures);
+    RescalingPushConstant rescaling;
     const VkSampler* samplers_it{samplers.data()};
     const VideoCommon::ImageViewInOut* views_it{views.data()};
     const auto prepare_stage{[&](size_t stage) LAMBDA_FORCEINLINE {
@@ -477,15 +479,16 @@ void GraphicsPipeline::ConfigureDraw(const RescalingPushConstant& rescaling) {
         if (bind_pipeline) {
             cmdbuf.BindPipeline(VK_PIPELINE_BIND_POINT_GRAPHICS, *pipeline);
         }
+        cmdbuf.PushConstants(*pipeline_layout, VK_SHADER_STAGE_ALL_GRAPHICS,
+                             RESCALING_LAYOUT_WORDS_OFFSET, sizeof(rescaling_data),
+                             rescaling_data.data());
         if (update_rescaling) {
             const f32 config_down_factor{Settings::values.resolution_info.down_factor};
             const f32 scale_down_factor{is_rescaling ? config_down_factor : 1.0f};
-            cmdbuf.PushConstants(*pipeline_layout, VK_SHADER_STAGE_ALL_GRAPHICS, 0,
-                                 sizeof(scale_down_factor), &scale_down_factor);
+            cmdbuf.PushConstants(*pipeline_layout, VK_SHADER_STAGE_ALL_GRAPHICS,
+                                 RESCALING_LAYOUT_DOWN_FACTOR_OFFSET, sizeof(scale_down_factor),
+                                 &scale_down_factor);
         }
-        cmdbuf.PushConstants(*pipeline_layout, VK_SHADER_STAGE_ALL_GRAPHICS,
-                             RESCALING_PUSH_CONSTANT_WORDS_OFFSET, sizeof(rescaling_data),
-                             rescaling_data.data());
         if (!descriptor_set_layout) {
             return;
         }