4 files changed, 45 insertions, 24 deletions

diff --git a/src/video_core/renderer_vulkan/pipeline_helper.h b/src/video_core/renderer_vulkan/pipeline_helper.h
index c6e5e059b..4847db6b6 100644
--- a/src/video_core/renderer_vulkan/pipeline_helper.h
+++ b/src/video_core/renderer_vulkan/pipeline_helper.h
@@ -16,38 +16,50 @@
 #include "video_core/texture_cache/texture_cache.h"
 #include "video_core/texture_cache/types.h"
 #include "video_core/textures/texture.h"
+#include "video_core/vulkan_common/vulkan_device.h"
 
 namespace Vulkan {
 
 class DescriptorLayoutBuilder {
 public:
-    DescriptorLayoutBuilder(const vk::Device& device_) : device{&device_} {}
+    DescriptorLayoutBuilder(const Device& device_) : device{&device_} {}
 
-    vk::DescriptorSetLayout CreateDescriptorSetLayout() const {
+    bool CanUsePushDescriptor() const noexcept {
+        return device->IsKhrPushDescriptorSupported() &&
+               num_descriptors <= device->MaxPushDescriptors();
+    }
+
+    vk::DescriptorSetLayout CreateDescriptorSetLayout(bool use_push_descriptor) const {
         if (bindings.empty()) {
             return nullptr;
         }
-        return device->CreateDescriptorSetLayout({
+        const VkDescriptorSetLayoutCreateFlags flags =
+            use_push_descriptor ? VK_DESCRIPTOR_SET_LAYOUT_CREATE_PUSH_DESCRIPTOR_BIT_KHR : 0;
+        return device->GetLogical().CreateDescriptorSetLayout({
             .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
             .pNext = nullptr,
-            .flags = 0,
+            .flags = flags,
             .bindingCount = static_cast<u32>(bindings.size()),
             .pBindings = bindings.data(),
         });
     }
 
     vk::DescriptorUpdateTemplateKHR CreateTemplate(VkDescriptorSetLayout descriptor_set_layout,
-                                                   VkPipelineLayout pipeline_layout) const {
+                                                   VkPipelineLayout pipeline_layout,
+                                                   bool use_push_descriptor) const {
         if (entries.empty()) {
             return nullptr;
         }
-        return device->CreateDescriptorUpdateTemplateKHR({
+        const VkDescriptorUpdateTemplateType type =
+            use_push_descriptor ? VK_DESCRIPTOR_UPDATE_TEMPLATE_TYPE_PUSH_DESCRIPTORS_KHR
+                                : VK_DESCRIPTOR_UPDATE_TEMPLATE_TYPE_DESCRIPTOR_SET_KHR;
+        return device->GetLogical().CreateDescriptorUpdateTemplateKHR({
             .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_UPDATE_TEMPLATE_CREATE_INFO_KHR,
             .pNext = nullptr,
             .flags = 0,
             .descriptorUpdateEntryCount = static_cast<u32>(entries.size()),
             .pDescriptorUpdateEntries = entries.data(),
-            .templateType = VK_DESCRIPTOR_UPDATE_TEMPLATE_TYPE_DESCRIPTOR_SET_KHR,
+            .templateType = type,
             .descriptorSetLayout = descriptor_set_layout,
             .pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS,
             .pipelineLayout = pipeline_layout,
@@ -56,7 +68,7 @@ public:
     }
 
     vk::PipelineLayout CreatePipelineLayout(VkDescriptorSetLayout descriptor_set_layout) const {
-        return device->CreatePipelineLayout({
+        return device->GetLogical().CreatePipelineLayout({
             .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
             .pNext = nullptr,
             .flags = 0,
@@ -97,14 +109,16 @@ private:
                 .stride = sizeof(DescriptorUpdateEntry),
             });
             ++binding;
+            num_descriptors += descriptors[i].count;
             offset += sizeof(DescriptorUpdateEntry);
         }
     }
 
-    const vk::Device* device{};
+    const Device* device{};
     boost::container::small_vector<VkDescriptorSetLayoutBinding, 32> bindings;
     boost::container::small_vector<VkDescriptorUpdateTemplateEntryKHR, 32> entries;
     u32 binding{};
+    u32 num_descriptors{};
     size_t offset{};
 };
 
diff --git a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
index cc855a62e..70b84c7a6 100644
--- a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
@@ -37,15 +37,14 @@ ComputePipeline::ComputePipeline(const Device& device_, DescriptorPool& descript
                 uniform_buffer_sizes.begin());
 
     auto func{[this, &descriptor_pool, shader_notify] {
-        DescriptorLayoutBuilder builder{device.GetLogical()};
+        DescriptorLayoutBuilder builder{device};
         builder.Add(info, VK_SHADER_STAGE_COMPUTE_BIT);
 
-        descriptor_set_layout = builder.CreateDescriptorSetLayout();
+        descriptor_set_layout = builder.CreateDescriptorSetLayout(false);
         pipeline_layout = builder.CreatePipelineLayout(*descriptor_set_layout);
         descriptor_update_template =
-            builder.CreateTemplate(*descriptor_set_layout, *pipeline_layout);
+            builder.CreateTemplate(*descriptor_set_layout, *pipeline_layout, false);
         descriptor_allocator = descriptor_pool.Allocator(*descriptor_set_layout, info);
-
         const VkPipelineShaderStageRequiredSubgroupSizeCreateInfoEXT subgroup_size_ci{
             .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO_EXT,
             .pNext = nullptr,
@@ -186,7 +185,6 @@ void ComputePipeline::Configure(Tegra::Engines::KeplerCompute& kepler_compute,
     const void* const descriptor_data{update_descriptor_queue.UpdateData()};
     scheduler.Record([this, descriptor_data](vk::CommandBuffer cmdbuf) {
         cmdbuf.BindPipeline(VK_PIPELINE_BIND_POINT_COMPUTE, *pipeline);
-
         if (!descriptor_set_layout) {
             return;
         }
diff --git a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
index e02b1b7ab..2b59a9d88 100644
--- a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
@@ -40,7 +40,7 @@ constexpr size_t NUM_STAGES = Maxwell::MaxShaderStage;
 constexpr size_t MAX_IMAGE_ELEMENTS = 64;
 
 DescriptorLayoutBuilder MakeBuilder(const Device& device, std::span<const Shader::Info> infos) {
-    DescriptorLayoutBuilder builder{device.GetLogical()};
+    DescriptorLayoutBuilder builder{device};
     for (size_t index = 0; index < infos.size(); ++index) {
         static constexpr std::array stages{
             VK_SHADER_STAGE_VERTEX_BIT,
@@ -229,12 +229,15 @@ GraphicsPipeline::GraphicsPipeline(
     }
     auto func{[this, shader_notify, &render_pass_cache, &descriptor_pool] {
         DescriptorLayoutBuilder builder{MakeBuilder(device, stage_infos)};
-        descriptor_set_layout = builder.CreateDescriptorSetLayout();
-        descriptor_allocator = descriptor_pool.Allocator(*descriptor_set_layout, stage_infos);
-
+        uses_push_descriptor = builder.CanUsePushDescriptor();
+        descriptor_set_layout = builder.CreateDescriptorSetLayout(uses_push_descriptor);
+        if (!uses_push_descriptor) {
+            descriptor_allocator = descriptor_pool.Allocator(*descriptor_set_layout, stage_infos);
+        }
         const VkDescriptorSetLayout set_layout{*descriptor_set_layout};
         pipeline_layout = builder.CreatePipelineLayout(set_layout);
-        descriptor_update_template = builder.CreateTemplate(set_layout, *pipeline_layout);
+        descriptor_update_template =
+            builder.CreateTemplate(set_layout, *pipeline_layout, uses_push_descriptor);
 
         const VkRenderPass render_pass{render_pass_cache.Get(MakeRenderPassKey(key.state))};
         Validate();
@@ -462,11 +465,16 @@ void GraphicsPipeline::ConfigureDraw() {
         if (!descriptor_set_layout) {
             return;
         }
-        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_GRAPHICS, *pipeline_layout, 0,
-                                  descriptor_set, nullptr);
+        if (uses_push_descriptor) {
+            cmdbuf.PushDescriptorSetWithTemplateKHR(*descriptor_update_template, *pipeline_layout,
+                                                    0, descriptor_data);
+        } else {
+            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_GRAPHICS, *pipeline_layout, 0,
+                                      descriptor_set, nullptr);
+        }
     });
 }
 
diff --git a/src/video_core/renderer_vulkan/vk_graphics_pipeline.h b/src/video_core/renderer_vulkan/vk_graphics_pipeline.h
index 40d1edabd..622267147 100644
--- a/src/video_core/renderer_vulkan/vk_graphics_pipeline.h
+++ b/src/video_core/renderer_vulkan/vk_graphics_pipeline.h
@@ -148,6 +148,7 @@ private:
     std::condition_variable build_condvar;
     std::mutex build_mutex;
     std::atomic_bool is_built{false};
+    bool uses_push_descriptor{false};
 };
 
 } // namespace Vulkan