shader: Implement indexed textures

3 files changed, 95 insertions, 64 deletions

diff --git a/src/video_core/renderer_vulkan/pipeline_helper.h b/src/video_core/renderer_vulkan/pipeline_helper.h
index aaf9a735e..dd7d2cc0c 100644
--- a/src/video_core/renderer_vulkan/pipeline_helper.h
+++ b/src/video_core/renderer_vulkan/pipeline_helper.h
@@ -85,28 +85,30 @@ public:
     }
 
     void Add(const Shader::Info& info, VkShaderStageFlags stage) {
-        Add(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, stage, info.constant_buffer_descriptors.size());
-        Add(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, stage, info.storage_buffers_descriptors.size());
-        Add(VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, stage, info.texture_buffer_descriptors.size());
-        Add(VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, stage, info.image_buffer_descriptors.size());
-        Add(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, stage, info.texture_descriptors.size());
-        Add(VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, stage, info.image_descriptors.size());
+        Add(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, stage, info.constant_buffer_descriptors);
+        Add(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, stage, info.storage_buffers_descriptors);
+        Add(VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, stage, info.texture_buffer_descriptors);
+        Add(VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, stage, info.image_buffer_descriptors);
+        Add(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, stage, info.texture_descriptors);
+        Add(VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, stage, info.image_descriptors);
     }
 
 private:
-    void Add(VkDescriptorType type, VkShaderStageFlags stage, size_t num) {
+    template <typename Descriptors>
+    void Add(VkDescriptorType type, VkShaderStageFlags stage, const Descriptors& descriptors) {
+        const size_t num{descriptors.size()};
         for (size_t i = 0; i < num; ++i) {
             bindings.push_back({
                 .binding = binding,
                 .descriptorType = type,
-                .descriptorCount = 1,
+                .descriptorCount = descriptors[i].count,
                 .stageFlags = stage,
                 .pImmutableSamplers = nullptr,
             });
             entries.push_back({
                 .dstBinding = binding,
                 .dstArrayElement = 0,
-                .descriptorCount = 1,
+                .descriptorCount = descriptors[i].count,
                 .descriptorType = type,
                 .offset = offset,
                 .stride = sizeof(DescriptorUpdateEntry),
@@ -126,21 +128,29 @@ private:
 inline void PushImageDescriptors(const Shader::Info& info, const VkSampler*& samplers,
                                  const ImageId*& image_view_ids, TextureCache& texture_cache,
                                  VKUpdateDescriptorQueue& update_descriptor_queue) {
-    image_view_ids += info.texture_buffer_descriptors.size();
-    image_view_ids += info.image_buffer_descriptors.size();
+    for (const auto& desc : info.texture_buffer_descriptors) {
+        image_view_ids += desc.count;
+    }
+    for (const auto& desc : info.image_buffer_descriptors) {
+        image_view_ids += desc.count;
+    }
     for (const auto& desc : info.texture_descriptors) {
-        const VkSampler sampler{*(samplers++)};
-        ImageView& image_view{texture_cache.GetImageView(*(image_view_ids++))};
-        const VkImageView vk_image_view{image_view.Handle(desc.type)};
-        update_descriptor_queue.AddSampledImage(vk_image_view, sampler);
+        for (u32 index = 0; index < desc.count; ++index) {
+            const VkSampler sampler{*(samplers++)};
+            ImageView& image_view{texture_cache.GetImageView(*(image_view_ids++))};
+            const VkImageView vk_image_view{image_view.Handle(desc.type)};
+            update_descriptor_queue.AddSampledImage(vk_image_view, sampler);
+        }
     }
     for (const auto& desc : info.image_descriptors) {
-        ImageView& image_view{texture_cache.GetImageView(*(image_view_ids++))};
-        if (desc.is_written) {
-            texture_cache.MarkModification(image_view.image_id);
+        for (u32 index = 0; index < desc.count; ++index) {
+            ImageView& image_view{texture_cache.GetImageView(*(image_view_ids++))};
+            if (desc.is_written) {
+                texture_cache.MarkModification(image_view.image_id);
+            }
+            const VkImageView vk_image_view{image_view.StorageView(desc.type, desc.format)};
+            update_descriptor_queue.AddImage(vk_image_view);
         }
-        const VkImageView vk_image_view{image_view.StorageView(desc.type, desc.format)};
-        update_descriptor_queue.AddImage(vk_image_view);
     }
 }
 
diff --git a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
index 45d837ca4..6e9f66262 100644
--- a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
@@ -91,35 +91,41 @@ void ComputePipeline::Configure(Tegra::Engines::KeplerCompute& kepler_compute,
     const auto& qmd{kepler_compute.launch_description};
     const auto& cbufs{qmd.const_buffer_config};
     const bool via_header_index{qmd.linked_tsc != 0};
-    const auto read_handle{[&](const auto& desc) {
+    const auto read_handle{[&](const auto& desc, u32 index) {
         ASSERT(((qmd.const_buffer_enable_mask >> desc.cbuf_index) & 1) != 0);
+        const u32 index_offset{index << desc.size_shift};
+        const u32 offset{desc.cbuf_offset + index_offset};
         const GPUVAddr addr{cbufs[desc.cbuf_index].Address() + desc.cbuf_offset};
         if constexpr (std::is_same_v<decltype(desc), const Shader::TextureDescriptor&> ||
                       std::is_same_v<decltype(desc), const Shader::TextureBufferDescriptor&>) {
             if (desc.has_secondary) {
                 ASSERT(((qmd.const_buffer_enable_mask >> desc.secondary_cbuf_index) & 1) != 0);
+                const u32 secondary_offset{desc.secondary_cbuf_offset + index_offset};
                 const GPUVAddr separate_addr{cbufs[desc.secondary_cbuf_index].Address() +
-                                             desc.secondary_cbuf_offset};
+                                             secondary_offset};
                 const u32 lhs_raw{gpu_memory.Read<u32>(addr)};
                 const u32 rhs_raw{gpu_memory.Read<u32>(separate_addr)};
-                const u32 raw{lhs_raw | rhs_raw};
-                return TextureHandle{raw, via_header_index};
+                return TextureHandle{lhs_raw | rhs_raw, via_header_index};
             }
         }
         return TextureHandle{gpu_memory.Read<u32>(addr), via_header_index};
     }};
     const auto add_image{[&](const auto& desc) {
-        const TextureHandle handle{read_handle(desc)};
-        image_view_indices.push_back(handle.image);
+        for (u32 index = 0; index < desc.count; ++index) {
+            const TextureHandle handle{read_handle(desc, index)};
+            image_view_indices.push_back(handle.image);
+        }
     }};
     std::ranges::for_each(info.texture_buffer_descriptors, add_image);
     std::ranges::for_each(info.image_buffer_descriptors, add_image);
     for (const auto& desc : info.texture_descriptors) {
-        const TextureHandle handle{read_handle(desc)};
-        image_view_indices.push_back(handle.image);
+        for (u32 index = 0; index < desc.count; ++index) {
+            const TextureHandle handle{read_handle(desc, index)};
+            image_view_indices.push_back(handle.image);
 
-        Sampler* const sampler = texture_cache.GetComputeSampler(handle.sampler);
-        samplers.push_back(sampler->Handle());
+            Sampler* const sampler = texture_cache.GetComputeSampler(handle.sampler);
+            samplers.push_back(sampler->Handle());
+        }
     }
     std::ranges::for_each(info.image_descriptors, add_image);
 
@@ -130,16 +136,18 @@ void ComputePipeline::Configure(Tegra::Engines::KeplerCompute& kepler_compute,
     ImageId* texture_buffer_ids{image_view_ids.data()};
     size_t index{};
     const auto add_buffer{[&](const auto& desc) {
-        ASSERT(desc.count == 1);
-        bool is_written{false};
-        if constexpr (std::is_same_v<decltype(desc), const Shader::ImageBufferDescriptor&>) {
-            is_written = desc.is_written;
+        for (u32 index = 0; index < desc.count; ++index) {
+            bool is_written{false};
+            if constexpr (std::is_same_v<decltype(desc), const Shader::ImageBufferDescriptor&>) {
+                is_written = desc.is_written;
+            }
+            ImageView& image_view = texture_cache.GetImageView(*texture_buffer_ids);
+            buffer_cache.BindComputeTextureBuffer(index, image_view.GpuAddr(),
+                                                  image_view.BufferSize(), image_view.format,
+                                                  is_written);
+            ++texture_buffer_ids;
+            ++index;
         }
-        ImageView& image_view = texture_cache.GetImageView(*texture_buffer_ids);
-        buffer_cache.BindComputeTextureBuffer(index, image_view.GpuAddr(), image_view.BufferSize(),
-                                              image_view.format, is_written);
-        ++texture_buffer_ids;
-        ++index;
     }};
     std::ranges::for_each(info.texture_buffer_descriptors, add_buffer);
     std::ranges::for_each(info.image_buffer_descriptors, add_buffer);
diff --git a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
index 08f00b9ce..b7688aef9 100644
--- a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
@@ -161,23 +161,26 @@ void GraphicsPipeline::Configure(bool is_indexed) {
         const Shader::Info& info{stage_infos[stage]};
         buffer_cache.SetEnabledUniformBuffers(stage, info.constant_buffer_mask);
         buffer_cache.UnbindGraphicsStorageBuffers(stage);
-        size_t index{};
+        size_t ssbo_index{};
         for (const auto& desc : info.storage_buffers_descriptors) {
             ASSERT(desc.count == 1);
-            buffer_cache.BindGraphicsStorageBuffer(stage, index, desc.cbuf_index, desc.cbuf_offset,
-                                                   desc.is_written);
-            ++index;
+            buffer_cache.BindGraphicsStorageBuffer(stage, ssbo_index, desc.cbuf_index,
+                                                   desc.cbuf_offset, desc.is_written);
+            ++ssbo_index;
         }
         const auto& cbufs{maxwell3d.state.shader_stages[stage].const_buffers};
-        const auto read_handle{[&](const auto& desc) {
+        const auto read_handle{[&](const auto& desc, u32 index) {
             ASSERT(cbufs[desc.cbuf_index].enabled);
-            const GPUVAddr addr{cbufs[desc.cbuf_index].address + desc.cbuf_offset};
+            const u32 index_offset{index << desc.size_shift};
+            const u32 offset{desc.cbuf_offset + index_offset};
+            const GPUVAddr addr{cbufs[desc.cbuf_index].address + offset};
             if constexpr (std::is_same_v<decltype(desc), const Shader::TextureDescriptor&> ||
                           std::is_same_v<decltype(desc), const Shader::TextureBufferDescriptor&>) {
                 if (desc.has_secondary) {
                     ASSERT(cbufs[desc.secondary_cbuf_index].enabled);
+                    const u32 second_offset{desc.secondary_cbuf_offset + index_offset};
                     const GPUVAddr separate_addr{cbufs[desc.secondary_cbuf_index].address +
-                                                 desc.secondary_cbuf_offset};
+                                                 second_offset};
                     const u32 lhs_raw{gpu_memory.Read<u32>(addr)};
                     const u32 rhs_raw{gpu_memory.Read<u32>(separate_addr)};
                     const u32 raw{lhs_raw | rhs_raw};
@@ -187,17 +190,21 @@ void GraphicsPipeline::Configure(bool is_indexed) {
             return TextureHandle{gpu_memory.Read<u32>(addr), via_header_index};
         }};
         const auto add_image{[&](const auto& desc) {
-            const TextureHandle handle{read_handle(desc)};
-            image_view_indices.push_back(handle.image);
+            for (u32 index = 0; index < desc.count; ++index) {
+                const TextureHandle handle{read_handle(desc, index)};
+                image_view_indices.push_back(handle.image);
+            }
         }};
         std::ranges::for_each(info.texture_buffer_descriptors, add_image);
         std::ranges::for_each(info.image_buffer_descriptors, add_image);
         for (const auto& desc : info.texture_descriptors) {
-            const TextureHandle handle{read_handle(desc)};
-            image_view_indices.push_back(handle.image);
+            for (u32 index = 0; index < desc.count; ++index) {
+                const TextureHandle handle{read_handle(desc, index)};
+                image_view_indices.push_back(handle.image);
 
-            Sampler* const sampler{texture_cache.GetGraphicsSampler(handle.sampler)};
-            samplers.push_back(sampler->Handle());
+                Sampler* const sampler{texture_cache.GetGraphicsSampler(handle.sampler)};
+                samplers.push_back(sampler->Handle());
+            }
         }
         std::ranges::for_each(info.image_descriptors, add_image);
     }
@@ -208,24 +215,30 @@ void GraphicsPipeline::Configure(bool is_indexed) {
     for (size_t stage = 0; stage < Maxwell::MaxShaderStage; ++stage) {
         size_t index{};
         const auto add_buffer{[&](const auto& desc) {
-            ASSERT(desc.count == 1);
-            bool is_written{false};
-            if constexpr (std::is_same_v<decltype(desc), const Shader::ImageBufferDescriptor&>) {
-                is_written = desc.is_written;
+            for (u32 index = 0; index < desc.count; ++index) {
+                bool is_written{false};
+                if constexpr (std::is_same_v<decltype(desc),
+                                             const Shader::ImageBufferDescriptor&>) {
+                    is_written = desc.is_written;
+                }
+                ImageView& image_view{texture_cache.GetImageView(*texture_buffer_index)};
+                buffer_cache.BindGraphicsTextureBuffer(stage, index, image_view.GpuAddr(),
+                                                       image_view.BufferSize(), image_view.format,
+                                                       is_written);
+                ++index;
+                ++texture_buffer_index;
             }
-            ImageView& image_view{texture_cache.GetImageView(*texture_buffer_index)};
-            buffer_cache.BindGraphicsTextureBuffer(stage, index, image_view.GpuAddr(),
-                                                   image_view.BufferSize(), image_view.format,
-                                                   is_written);
-            ++index;
-            ++texture_buffer_index;
         }};
         const Shader::Info& info{stage_infos[stage]};
         buffer_cache.UnbindGraphicsTextureBuffers(stage);
         std::ranges::for_each(info.texture_buffer_descriptors, add_buffer);
         std::ranges::for_each(info.image_buffer_descriptors, add_buffer);
-        texture_buffer_index += info.texture_descriptors.size();
-        texture_buffer_index += info.image_descriptors.size();
+        for (const auto& desc : info.texture_descriptors) {
+            texture_buffer_index += desc.count;
+        }
+        for (const auto& desc : info.image_descriptors) {
+            texture_buffer_index += desc.count;
+        }
     }
     buffer_cache.UpdateGraphicsBuffers(is_indexed);