anv: track descriptor buffers used & promoted to push constants

Signed-off-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Reviewed-by: Alyssa Rosenzweig <alyssa.rosenzweig@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/39405>

src/intel/vulkan/anv_nir_apply_pipeline_layout.c

@@ -135,13 +135,19 @@ addr_format_for_desc_type(VkDescriptorType desc_type,
 
 static struct anv_binding_apply_layout *
 add_binding(struct apply_pipeline_layout_state *state,
-            uint32_t set, uint32_t binding)
+            uint32_t set, uint32_t binding,
+            bool sampler)
 {
    const struct anv_descriptor_set_layout *set_layout =
       state->set_layouts[set];
    const struct anv_descriptor_set_binding_layout *bind_layout =
       &set_layout->binding[binding];
 
+   if (sampler)
+      state->bind_map->used_sampler_sets |= BITFIELD_BIT(set);
+   else
+      state->bind_map->used_surface_sets |= BITFIELD_BIT(set);
+
    assert(set < state->set_count);
    assert(binding < set_layout->binding_count);
 
@@ -180,7 +186,7 @@ static void
 add_binding_type(struct apply_pipeline_layout_state *state,
                  uint32_t set, uint32_t binding, VkDescriptorType type)
 {
-   add_binding(state, set, binding);
+   add_binding(state, set, binding, false);
 
    const struct anv_descriptor_set_layout *set_layout =
       state->set_layouts[set];
@@ -205,23 +211,27 @@ add_binding_type(struct apply_pipeline_layout_state *state,
 }
 
 static struct anv_binding_apply_layout *
-add_deref_src_binding(struct apply_pipeline_layout_state *state, nir_src src)
+add_deref_src_binding(struct apply_pipeline_layout_state *state, nir_src src,
+                      bool sampler)
 {
    nir_deref_instr *deref = nir_src_as_deref(src);
    nir_variable *var = nir_deref_instr_get_variable(deref);
-   return add_binding(state, var->data.descriptor_set, var->data.binding);
+   return add_binding(state, var->data.descriptor_set, var->data.binding, sampler);
 }
 
 static void
 add_tex_src_binding(struct apply_pipeline_layout_state *state,
-                    nir_tex_instr *tex, nir_tex_src_type deref_src_type)
+                    nir_tex_instr *tex,
+                    nir_tex_src_type deref_src_type)
 {
    int deref_src_idx = nir_tex_instr_src_index(tex, deref_src_type);
    if (deref_src_idx < 0)
       return;
 
    struct anv_binding_apply_layout *layout =
-      add_deref_src_binding(state, tex->src[deref_src_idx].src);
+      add_deref_src_binding(state,
+                            tex->src[deref_src_idx].src,
+                            deref_src_type == nir_tex_src_sampler_deref);
 
    /* Track input attachments use */
    nir_variable *var =
@@ -287,7 +297,7 @@ get_used_bindings(UNUSED nir_builder *_b, nir_instr *instr, void *_state)
       case nir_intrinsic_image_deref_load_raw_intel:
       case nir_intrinsic_image_deref_store_raw_intel:
       case nir_intrinsic_image_deref_sparse_load:
-         add_deref_src_binding(state, intrin->src[0]);
+         add_deref_src_binding(state, intrin->src[0], false);
          break;
 
       case nir_intrinsic_load_constant:

src/intel/vulkan/anv_nir_compute_push_layout.c

@@ -443,6 +443,11 @@ anv_nir_compute_push_layout(nir_shader *nir,
       if (candidate_range->length == 0)
          break;
 
+      if (candidate_range->set == ANV_DESCRIPTOR_SET_DESCRIPTORS) {
+         assert(candidate_range->index < MAX_SETS);
+         map->pushed_sets |= BITFIELD_BIT(candidate_range->index);
+      }
+
       map->push_ranges[n_push_ranges++] = *candidate_range;
       total_push_regs += candidate_range->length;
    }

src/intel/vulkan/anv_private.h

@@ -1174,6 +1174,15 @@ struct anv_pipeline_bind_map {
 
    struct anv_push_range                        push_ranges[4];
 
+   /* Bitfield of sets for which the surfaces are accessed */
+   uint8_t                                      used_surface_sets;
+
+   /* Bitfield of sets for which the samplers are accessed */
+   uint8_t                                      used_sampler_sets;
+
+   /* Bitfield of sets promoted to push constants */
+   uint8_t                                      pushed_sets;
+
    /* Number of dynamic descriptor in each set */
    uint8_t                                      dynamic_descriptors[MAX_SETS];
 };

src/intel/vulkan/anv_shader.c

@@ -86,6 +86,9 @@ anv_shader_deserialize(struct vk_device *vk_device,
    blob_copy_bytes(blob, data.bind_map.input_attachments,
                    sizeof(data.bind_map.input_attachments));
    blob_copy_bytes(blob, data.bind_map.push_ranges, sizeof(data.bind_map.push_ranges));
+   data.bind_map.used_surface_sets = blob_read_uint8(blob);
+   data.bind_map.used_sampler_sets = blob_read_uint8(blob);
+   data.bind_map.pushed_sets = blob_read_uint8(blob);
    blob_copy_bytes(blob, data.bind_map.dynamic_descriptors,
                    sizeof(data.bind_map.dynamic_descriptors));
 
@@ -164,6 +167,9 @@ anv_shader_serialize(struct vk_device *device,
                     sizeof(shader->bind_map.input_attachments));
    blob_write_bytes(blob, shader->bind_map.push_ranges,
                     sizeof(shader->bind_map.push_ranges));
+   blob_write_uint8(blob, shader->bind_map.used_surface_sets);
+   blob_write_uint8(blob, shader->bind_map.used_sampler_sets);
+   blob_write_uint8(blob, shader->bind_map.pushed_sets);
    blob_write_bytes(blob, shader->bind_map.dynamic_descriptors,
                     sizeof(shader->bind_map.dynamic_descriptors));