anv: Implement VK_EXT_subgroup_size_control

Reviewed-by: Caio Marcelo de Oliveira Filho <[email protected]>

3 files changed, 46 insertions, 2 deletions

diff --git a/src/intel/vulkan/anv_device.c b/src/intel/vulkan/anv_device.c
index d0ddb3e52bc..217a779576c 100644
--- a/src/intel/vulkan/anv_device.c
+++ b/src/intel/vulkan/anv_device.c
@@ -1590,6 +1590,17 @@ void anv_GetPhysicalDeviceProperties2(
          break;
       }
 
+      case VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_PROPERTIES_EXT: {
+         VkPhysicalDeviceSubgroupSizeControlPropertiesEXT *props =
+            (VkPhysicalDeviceSubgroupSizeControlPropertiesEXT *)ext;
+         STATIC_ASSERT(8 <= BRW_SUBGROUP_SIZE && BRW_SUBGROUP_SIZE <= 32);
+         props->minSubgroupSize = 8;
+         props->maxSubgroupSize = 32;
+         props->maxComputeWorkgroupSubgroups = pdevice->info.max_cs_threads;
+         props->requiredSubgroupSizeStages = VK_SHADER_STAGE_COMPUTE_BIT;
+         break;
+      }
+
       case VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_TEXEL_BUFFER_ALIGNMENT_PROPERTIES_EXT: {
          VkPhysicalDeviceTexelBufferAlignmentPropertiesEXT *props =
             (VkPhysicalDeviceTexelBufferAlignmentPropertiesEXT *)ext;
diff --git a/src/intel/vulkan/anv_extensions.py b/src/intel/vulkan/anv_extensions.py
index 492f8ac4615..9956631eb26 100644
--- a/src/intel/vulkan/anv_extensions.py
+++ b/src/intel/vulkan/anv_extensions.py
@@ -147,6 +147,7 @@ EXTENSIONS = [
     Extension('VK_EXT_shader_demote_to_helper_invocation', 1, True),
     Extension('VK_EXT_shader_stencil_export',             1, 'device->info.gen >= 9'),
     Extension('VK_EXT_shader_viewport_index_layer',       1, True),
+    Extension('VK_EXT_subgroup_size_control',             1, True),
     Extension('VK_EXT_texel_buffer_alignment',            1, True),
     Extension('VK_EXT_transform_feedback',                1, True),
     Extension('VK_EXT_vertex_attribute_divisor',          3, True),
diff --git a/src/intel/vulkan/anv_pipeline.c b/src/intel/vulkan/anv_pipeline.c
index ae3693d76ed..c2a164cc802 100644
--- a/src/intel/vulkan/anv_pipeline.c
+++ b/src/intel/vulkan/anv_pipeline.c
@@ -356,7 +356,10 @@ populate_base_prog_key(const struct gen_device_info *devinfo,
                        VkPipelineShaderStageCreateFlags flags,
                        struct brw_base_prog_key *key)
 {
-   key->subgroup_size_type = BRW_SUBGROUP_SIZE_API_CONSTANT;
+   if (flags & VK_PIPELINE_SHADER_STAGE_CREATE_ALLOW_VARYING_SUBGROUP_SIZE_BIT_EXT)
+      key->subgroup_size_type = BRW_SUBGROUP_SIZE_VARYING;
+   else
+      key->subgroup_size_type = BRW_SUBGROUP_SIZE_API_CONSTANT;
 
    populate_sampler_prog_key(devinfo, &key->tex);
 }
@@ -465,11 +468,36 @@ populate_wm_prog_key(const struct gen_device_info *devinfo,
 static void
 populate_cs_prog_key(const struct gen_device_info *devinfo,
                      VkPipelineShaderStageCreateFlags flags,
+                     const VkPipelineShaderStageRequiredSubgroupSizeCreateInfoEXT *rss_info,
                      struct brw_cs_prog_key *key)
 {
    memset(key, 0, sizeof(*key));
 
    populate_base_prog_key(devinfo, flags, &key->base);
+
+   if (rss_info) {
+      assert(key->base.subgroup_size_type != BRW_SUBGROUP_SIZE_VARYING);
+
+      /* These enum values are expressly chosen to be equal to the subgroup
+       * size that they require.
+       */
+      assert(rss_info->requiredSubgroupSize == 8 ||
+             rss_info->requiredSubgroupSize == 16 ||
+             rss_info->requiredSubgroupSize == 32);
+      key->base.subgroup_size_type = rss_info->requiredSubgroupSize;
+   } else if (flags & VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT_EXT) {
+      /* If the client expressly requests full subgroups and they don't
+       * specify a subgroup size, we need to pick one.  If they're requested
+       * varying subgroup sizes, we set it to UNIFORM and let the back-end
+       * compiler pick.  Otherwise, we specify the API value of 32.
+       * Performance will likely be terrible in this case but there's nothing
+       * we can do about that.  The client should have chosen a size.
+       */
+      if (flags & VK_PIPELINE_SHADER_STAGE_CREATE_ALLOW_VARYING_SUBGROUP_SIZE_BIT_EXT)
+         key->base.subgroup_size_type = BRW_SUBGROUP_SIZE_UNIFORM;
+      else
+         key->base.subgroup_size_type = BRW_SUBGROUP_SIZE_REQUIRE_32;
+   }
 }
 
 struct anv_pipeline_stage {
@@ -1360,8 +1388,12 @@ anv_pipeline_compile_cs(struct anv_pipeline *pipeline,
 
    struct anv_shader_bin *bin = NULL;
 
+   const VkPipelineShaderStageRequiredSubgroupSizeCreateInfoEXT *rss_info =
+      vk_find_struct_const(info->stage.pNext,
+                           PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO_EXT);
+
    populate_cs_prog_key(&pipeline->device->info, info->stage.flags,
-                        &stage.key.cs);
+                        rss_info, &stage.key.cs);
 
    ANV_FROM_HANDLE(anv_pipeline_layout, layout, info->layout);