anv: avoid looking at the pipeline to flush push descriptors

We do this at the cost of recomputing some values that where available
on the pipeline at vkCmdBindPipeline() time.

We can look at the shaders on graphics/compute which will work nicely
with the runtime.

The runtime doesn't have support for ray tracing pipelines so we keep
using them.

Signed-off-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Reviewed-by: Ivan Briano <ivan.briano@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/36512>

src/intel/vulkan/anv_cmd_buffer.c

@@ -644,6 +644,28 @@ get_pipeline_dirty_stages(struct anv_device *device,
    return bits;
 }
 
+static void
+update_push_descriptor_flags(struct anv_cmd_pipeline_state *state,
+                             struct anv_shader_bin **shaders,
+                             uint32_t shader_count)
+{
+   state->push_buffer_stages = 0;
+   state->push_descriptor_stages = 0;
+
+   for (uint32_t i = 0; i < shader_count; i++) {
+      if (shaders[i] == NULL)
+         continue;
+
+      VkShaderStageFlags stage = mesa_to_vk_shader_stage(shaders[i]->stage);
+
+      if (shaders[i]->push_desc_info.used_descriptors)
+         state->push_descriptor_stages |= stage;
+
+      if (shaders[i]->push_desc_info.push_set_buffer)
+         state->push_buffer_stages |= stage;
+   }
+}
+
 void anv_CmdBindPipeline(
     VkCommandBuffer                             commandBuffer,
     VkPipelineBindPoint                         pipelineBindPoint,
@@ -669,6 +691,8 @@ void anv_CmdBindPipeline(
 
       state = &cmd_buffer->state.compute.base;
       stages = VK_SHADER_STAGE_COMPUTE_BIT;
+
+      update_push_descriptor_flags(state, &compute_pipeline->cs, 1);
       break;
    }
 
@@ -700,6 +724,9 @@ void anv_CmdBindPipeline(
       state = &cmd_buffer->state.gfx.base;
       stages = new_pipeline->base.base.active_stages;
 
+      update_push_descriptor_flags(state,
+                                   new_pipeline->base.shaders,
+                                   ARRAY_SIZE(new_pipeline->base.shaders));
 
       /* When the pipeline is using independent states and dynamic buffers,
        * this will trigger an update of anv_push_constants::dynamic_base_index
@@ -749,6 +776,10 @@ void anv_CmdBindPipeline(
       }
 
       state = &cmd_buffer->state.rt.base;
+
+      state->push_buffer_stages = pipeline->use_push_descriptor_buffer;
+      state->push_descriptor_stages = pipeline->use_push_descriptor_buffer;
+      state->push_descriptor_index = pipeline->layout.push_descriptor_set_index;
       break;
    }
 
@@ -932,6 +963,12 @@ anv_cmd_buffer_bind_descriptor_set(struct anv_cmd_buffer *cmd_buffer,
       }
    }
 
+   /* Update the push descriptor index tracking */
+   if (anv_descriptor_set_is_push(set))
+      pipe_state->push_descriptor_index = set_index;
+   else if (pipe_state->push_descriptor_index == set_index)
+      pipe_state->push_descriptor_index = UINT8_MAX;
+
    if (set->is_push)
       cmd_buffer->state.push_descriptors_dirty |= dirty_stages;
    else

src/intel/vulkan/anv_genX.h

@@ -430,20 +430,18 @@ genX(cmd_buffer_emit_push_descriptor_surfaces)(struct anv_cmd_buffer *cmd_buffer
 
 static inline VkShaderStageFlags
 genX(cmd_buffer_flush_push_descriptors)(struct anv_cmd_buffer *cmd_buffer,
-                                        struct anv_cmd_pipeline_state *state,
-                                        struct anv_pipeline *pipeline)
+                                        struct anv_cmd_pipeline_state *state)
 {
-   if (!pipeline->use_push_descriptor && !pipeline->use_push_descriptor_buffer)
+   if (state->push_buffer_stages == 0 && state->push_descriptor_stages == 0)
       return 0;
 
-   assert(pipeline->layout.push_descriptor_set_index != -1);
+   assert(state->push_descriptor_index != UINT8_MAX);
    struct anv_descriptor_set *set =
-      state->descriptors[pipeline->layout.push_descriptor_set_index];
+      state->descriptors[state->push_descriptor_index];
    assert(set->is_push);
 
    const VkShaderStageFlags push_buffer_dirty =
-      cmd_buffer->state.push_descriptors_dirty &
-      pipeline->use_push_descriptor_buffer;
+      cmd_buffer->state.push_descriptors_dirty & state->push_buffer_stages;
    if (push_buffer_dirty) {
       if (set->desc_surface_state.map == NULL)
          genX(cmd_buffer_emit_push_descriptor_buffer_surface)(cmd_buffer, set);
@@ -453,7 +451,7 @@ genX(cmd_buffer_flush_push_descriptors)(struct anv_cmd_buffer *cmd_buffer,
    }
 
    const VkShaderStageFlags push_descriptor_dirty =
-      cmd_buffer->state.push_descriptors_dirty & pipeline->use_push_descriptor;
+      cmd_buffer->state.push_descriptors_dirty & state->push_descriptor_stages;
    if (push_descriptor_dirty) {
       genX(cmd_buffer_emit_push_descriptor_surfaces)(cmd_buffer, set);
 

src/intel/vulkan/anv_nir.h

@@ -140,10 +140,10 @@ uint32_t anv_nir_compute_used_push_descriptors(nir_shader *shader,
                                                struct anv_descriptor_set_layout * const *set_layouts,
                                                uint32_t set_count);
 
-bool anv_nir_loads_push_desc_buffer(nir_shader *nir,
-                                    struct anv_descriptor_set_layout * const *set_layouts,
-                                    uint32_t set_count,
-                                    const struct anv_pipeline_bind_map *bind_map);
+uint8_t anv_nir_loads_push_desc_buffer(nir_shader *nir,
+                                       struct anv_descriptor_set_layout * const *set_layouts,
+                                       uint32_t set_count,
+                                       const struct anv_pipeline_bind_map *bind_map);
 
 uint32_t anv_nir_push_desc_ubo_fully_promoted(nir_shader *nir,
                                               struct anv_descriptor_set_layout * const *set_layouts,

src/intel/vulkan/anv_nir_push_descriptor_analysis.c

@@ -104,10 +104,11 @@ anv_nir_compute_used_push_descriptors(nir_shader *shader,
    return used_push_bindings;
 }
 
-/* This function checks whether the shader accesses the push descriptor
- * buffer. This function must be called after anv_nir_compute_push_layout().
+/* This function returns a bit field with one bit set to 1 indicating the push
+ * descriptor set used. This function must be called after
+ * anv_nir_compute_push_layout().
  */
-bool
+uint8_t
 anv_nir_loads_push_desc_buffer(nir_shader *nir,
                                struct anv_descriptor_set_layout * const *set_layouts,
                                uint32_t set_count,
@@ -117,7 +118,7 @@ anv_nir_loads_push_desc_buffer(nir_shader *nir,
    const struct anv_descriptor_set_layout *push_set_layout =
       anv_pipeline_layout_get_push_set(set_layouts, set_count, &push_set);
    if (push_set_layout == NULL)
-      return false;
+      return 0;
 
    nir_foreach_function_impl(impl, nir) {
       nir_foreach_block(block, impl) {
@@ -139,13 +140,13 @@ anv_nir_loads_push_desc_buffer(nir_shader *nir,
             if ((binding->set == ANV_DESCRIPTOR_SET_DESCRIPTORS ||
                  binding->set == ANV_DESCRIPTOR_SET_DESCRIPTORS_BUFFER) &&
                 binding->index == push_set) {
-               return true;
+               return BITFIELD_BIT(push_set);
             }
          }
       }
    }
 
-   return false;
+   return 0;
 }
 
 /* This function computes a bitfield of all the UBOs bindings in the push

src/intel/vulkan/anv_pipeline.c

@@ -1165,7 +1165,7 @@ anv_pipeline_lower_nir(struct anv_pipeline *pipeline,
                &stage->prog_data.cs);
    }
 
-   stage->push_desc_info.used_set_buffer =
+   stage->push_desc_info.push_set_buffer =
       anv_nir_loads_push_desc_buffer(nir,
                                      layout->set_layouts,
                                      layout->num_sets,
@@ -1781,7 +1781,7 @@ anv_pipeline_account_shader(struct anv_pipeline *pipeline,
    pipeline->ray_queries = MAX2(pipeline->ray_queries,
                                 shader->prog_data->ray_queries);
 
-   if (shader->push_desc_info.used_set_buffer) {
+   if (shader->push_desc_info.push_set_buffer) {
       pipeline->use_push_descriptor_buffer |=
          mesa_to_vk_shader_stage(shader->stage);
    }

src/intel/vulkan/anv_pipeline_cache.c

@@ -325,7 +325,7 @@ anv_shader_bin_serialize(struct vk_pipeline_cache_object *object,
 
    blob_write_uint32(blob, shader->push_desc_info.used_descriptors);
    blob_write_uint32(blob, shader->push_desc_info.fully_promoted_ubo_descriptors);
-   blob_write_uint8(blob, shader->push_desc_info.used_set_buffer);
+   blob_write_uint8(blob, shader->push_desc_info.push_set_buffer);
 
    blob_write_bytes(blob, shader->bind_map.surface_sha1,
                     sizeof(shader->bind_map.surface_sha1));
@@ -392,7 +392,7 @@ anv_shader_bin_deserialize(struct vk_pipeline_cache *cache,
    struct anv_push_descriptor_info push_desc_info = {};
    push_desc_info.used_descriptors = blob_read_uint32(blob);
    push_desc_info.fully_promoted_ubo_descriptors = blob_read_uint32(blob);
-   push_desc_info.used_set_buffer = blob_read_uint8(blob);
+   push_desc_info.push_set_buffer = blob_read_uint8(blob);
 
    struct anv_pipeline_bind_map bind_map = {};
    blob_copy_bytes(blob, bind_map.surface_sha1, sizeof(bind_map.surface_sha1));

src/intel/vulkan/anv_private.h

@@ -4070,6 +4070,21 @@ struct anv_cmd_pipeline_state {
       uint32_t                                  offsets[MAX_DYNAMIC_BUFFERS];
    }                                            dynamic_offsets[MAX_SETS];
 
+   /**
+    * The current stages using push descriptor buffer.
+    */
+   VkShaderStageFlags push_buffer_stages;
+
+   /**
+    * The current stages using push descriptors.
+    */
+   VkShaderStageFlags push_descriptor_stages;
+
+   /**
+    * Push descriptor index for currently bound shaders (UINT8_MAX if unused).
+    */
+   uint8_t push_descriptor_index;
+
    /**
     * The current bound pipeline.
     */
@@ -4814,8 +4829,8 @@ struct anv_push_descriptor_info {
    /* A bitfield of UBOs bindings fully promoted to push constants. */
    uint32_t fully_promoted_ubo_descriptors;
 
-   /* */
-   uint8_t used_set_buffer;
+   /* A bitfield with one bit set indicating the push descriptor set used. */
+   uint8_t push_set_buffer;
 };
 
 struct anv_shader_upload_params {

src/intel/vulkan/genX_cmd_buffer.c

@@ -2174,7 +2174,7 @@ emit_binding_table(struct anv_cmd_buffer *cmd_buffer,
          /* If the shader doesn't access the set buffer, just put the null
           * surface.
           */
-         if (set->is_push && !shader->push_desc_info.used_set_buffer) {
+         if (set->is_push && shader->push_desc_info.push_set_buffer == 0) {
             bt_map[s] = 0;
             break;
          }

src/intel/vulkan/genX_cmd_compute.c

@@ -170,8 +170,7 @@ cmd_buffer_flush_compute_state(struct anv_cmd_buffer *cmd_buffer)
 
    cmd_buffer->state.descriptors_dirty |=
       genX(cmd_buffer_flush_push_descriptors)(cmd_buffer,
-                                              &cmd_buffer->state.compute.base,
-                                              &pipeline->base);
+                                              &cmd_buffer->state.compute.base);
 
    if ((cmd_buffer->state.descriptors_dirty & VK_SHADER_STAGE_COMPUTE_BIT) ||
        cmd_buffer->state.compute.pipeline_dirty) {
@@ -1191,8 +1190,7 @@ cmd_buffer_trace_rays(struct anv_cmd_buffer *cmd_buffer,
    genX(cmd_buffer_apply_pipe_flushes)(cmd_buffer);
 
    genX(cmd_buffer_flush_push_descriptors)(cmd_buffer,
-                                           &cmd_buffer->state.rt.base,
-                                           &pipeline->base);
+                                           &cmd_buffer->state.rt.base);
 
    /* Add these to the reloc list as they're internal buffers that don't
     * actually have relocs to pick them up manually.

src/intel/vulkan/genX_cmd_draw.c

@@ -814,8 +814,7 @@ cmd_buffer_flush_gfx_state(struct anv_cmd_buffer *cmd_buffer)
 
    descriptors_dirty |=
       genX(cmd_buffer_flush_push_descriptors)(cmd_buffer,
-                                              &cmd_buffer->state.gfx.base,
-                                              &pipeline->base.base);
+                                              &cmd_buffer->state.gfx.base);
 
    if (!cmd_buffer->state.gfx.dirty && !descriptors_dirty &&
        !any_dynamic_state_dirty &&