Enable full subgroup together with coop matrices

Author: sergcpp

internal/RendererVK.cpp

@@ -1948,14 +1948,16 @@ bool Ray::Vk::Renderer::InitUNetFilterPipelines(
                            internal_shaders_output_convolution_32_3_img_coop_8x8x16_CF32_comp_spv),
              eShaderType::Comp, false}};
     const int subgroup_size = use_coop_matrix_ ? 32 : -1;
+    const bool require_full_subgroup = use_coop_matrix_;
     parallel_for(0, int(shaders_to_init.size()), [&](const int i) {
         std::get<6>(shaders_to_init[i]) =
             std::get<0>(shaders_to_init[i])
                 .Init(std::get<3>(shaders_to_init[i]), ctx_.get(), Inflate(std::get<4>(shaders_to_init[i])),
                       std::get<5>(shaders_to_init[i]), log);
         std::get<1>(shaders_to_init[i]) =
             Program{std::get<3>(shaders_to_init[i]), ctx_.get(), &std::get<0>(shaders_to_init[i]), log};
-        std::get<2>(shaders_to_init[i]).Init(ctx_.get(), &std::get<1>(shaders_to_init[i]), log, subgroup_size);
+        std::get<2>(shaders_to_init[i])
+            .Init(ctx_.get(), &std::get<1>(shaders_to_init[i]), log, subgroup_size, require_full_subgroup);
     });
 
     bool result = true;

internal/Vk/ContextVK.cpp

@@ -255,6 +255,7 @@ bool Ray::Vk::Context::Init(ILog *log, const VulkanDevice &vk_device, const Vulk
         api_.vkGetPhysicalDeviceFeatures2KHR(physical_device_, &feat2);
 
         subgroup_size_control_supported_ &= (subgroup_size_control_features.subgroupSizeControl == VK_TRUE);
+        subgroup_size_control_supported_ &= (subgroup_size_control_features.computeFullSubgroups == VK_TRUE);
     }
 
     if (!InitCommandBuffers(api_, command_pool_, temp_command_pool_, draw_cmd_bufs_, render_finished_semaphores_,
@@ -893,6 +894,7 @@ bool Ray::Vk::Context::InitVkDevice(const Api &api, VkDevice &device, VkPhysical
     VkPhysicalDeviceSubgroupSizeControlFeaturesEXT subgroup_size_control_features = {
         VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES_EXT};
     subgroup_size_control_features.subgroupSizeControl = VK_TRUE;
+    subgroup_size_control_features.computeFullSubgroups = VK_TRUE;
     if (enable_subgroup_size_control) {
         (*pp_next) = &subgroup_size_control_features;
         pp_next = &subgroup_size_control_features.pNext;

internal/Vk/PipelineVK.cpp

@@ -336,7 +336,7 @@ bool Ray::Vk::Pipeline::Init(Context *ctx, const RastState &rast_state, Program
                 subpass_index, log);
 }
 
-bool Ray::Vk::Pipeline::Init(Context *ctx, Program *prog, ILog *log, const int subgroup_size) {
+bool Ray::Vk::Pipeline::Init(Context *ctx, Program *prog, ILog *log, const int subgroup_size, const bool require_full_subgroup) {
     Destroy();
 
     ePipelineType type = ePipelineType::Undefined;
@@ -413,6 +413,9 @@ bool Ray::Vk::Pipeline::Init(Context *ctx, Program *prog, ILog *log, const int s
         stage_info.module = prog->shader(eShaderType(i))->module();
         stage_info.pName = "main";
         stage_info.pSpecializationInfo = nullptr;
+        if (ctx->subgroup_size_control_supported() && require_full_subgroup) {
+            stage_info.flags |= VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT;
+        }
 
         if (ctx->subgroup_size_control_supported() && subgroup_size != -1) {
             stage_info.pNext = &subgroup_size_info;

internal/Vk/PipelineVK.h

@@ -92,7 +92,7 @@ class Pipeline {
     bool Init(Context *ctx, const RastState &rast_state, Program *prog, const VertexInput *vtx_input,
               Span<const RenderTarget> color_attachments, RenderTarget depth_attachment, uint32_t subpass_index,
               ILog *log);
-    bool Init(Context *ctx, Program *prog, ILog *log, int subgroup_size = -1);
+    bool Init(Context *ctx, Program *prog, ILog *log, int subgroup_size = -1, bool require_full_subgroup = false);
 };
 } // namespace Vk
 } // namespace Ray