anv: reuse embedded samplers across shaders

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

src/intel/vulkan/anv_device.c

@@ -3871,6 +3871,8 @@ VkResult anv_CreateDevice(
 
    anv_device_utrace_init(device);
 
+   anv_device_init_embedded_samplers(device);
+
    BITSET_ONES(device->gfx_dirty_state);
    BITSET_CLEAR(device->gfx_dirty_state, ANV_GFX_STATE_INDEX_BUFFER);
    BITSET_CLEAR(device->gfx_dirty_state, ANV_GFX_STATE_SO_DECL_LIST);
@@ -4026,6 +4028,8 @@ void anv_DestroyDevice(
    vk_pipeline_cache_destroy(device->internal_cache, NULL);
    vk_pipeline_cache_destroy(device->default_pipeline_cache, NULL);
 
+   anv_device_finish_embedded_samplers(device);
+
    anv_device_finish_trtt(device);
 
    if (ANV_SUPPORT_RT && device->info->has_ray_tracing)

src/intel/vulkan/anv_nir_apply_pipeline_layout.c

@@ -2054,17 +2054,17 @@ add_embedded_sampler_entry(struct apply_pipeline_layout_state *state,
       .binding = binding,
    };
 
-   assert(sizeof(sampler->sampler_state) ==
+   assert(sizeof(sampler->key.sampler) ==
           sizeof(bind_layout->immutable_samplers[0]->state_no_bc[0]));
-   memcpy(sampler->sampler_state,
+   memcpy(sampler->key.sampler,
           bind_layout->immutable_samplers[0]->state_no_bc[0],
-          sizeof(sampler->sampler_state));
+          sizeof(sampler->key.sampler));
 
-   assert(sizeof(sampler->border_color) ==
+   assert(sizeof(sampler->key.color) ==
           sizeof(bind_layout->immutable_samplers[0]->vk.border_color_value.uint32));
-   memcpy(sampler->border_color,
+   memcpy(sampler->key.color,
           bind_layout->immutable_samplers[0]->vk.border_color_value.uint32,
-          sizeof(sampler->border_color));
+          sizeof(sampler->key.color));
 }
 
 static bool

src/intel/vulkan/anv_pipeline_cache.c

@@ -33,6 +33,114 @@
 #include "compiler/spirv/nir_spirv.h"
 #include "shaders/float64_spv.h"
 
+/**
+ * Embedded sampler management.
+ */
+
+static unsigned
+embedded_sampler_key_hash(const void *key)
+{
+   return _mesa_hash_data(key, sizeof(struct anv_embedded_sampler_key));
+}
+
+static bool
+embedded_sampler_key_equal(const void *a, const void *b)
+{
+   return memcmp(a, b, sizeof(struct anv_embedded_sampler_key)) == 0;
+}
+
+static void
+anv_embedded_sampler_free(struct anv_device *device,
+                          struct anv_embedded_sampler *sampler)
+{
+   anv_state_pool_free(&device->dynamic_state_db_pool, sampler->sampler_state);
+   anv_state_pool_free(&device->dynamic_state_db_pool, sampler->border_color_state);
+   vk_free(&device->vk.alloc, sampler);
+}
+
+static struct anv_embedded_sampler *
+anv_embedded_sampler_ref(struct anv_embedded_sampler *sampler)
+{
+   sampler->ref_cnt++;
+   return sampler;
+}
+
+static void
+anv_embedded_sampler_unref(struct anv_device *device,
+                           struct anv_embedded_sampler *sampler)
+{
+   simple_mtx_lock(&device->embedded_samplers.mutex);
+   if (--sampler->ref_cnt == 0) {
+      _mesa_hash_table_remove_key(device->embedded_samplers.map,
+                                  &sampler->key);
+      anv_embedded_sampler_free(device, sampler);
+   }
+   simple_mtx_unlock(&device->embedded_samplers.mutex);
+}
+
+void
+anv_device_init_embedded_samplers(struct anv_device *device)
+{
+   simple_mtx_init(&device->embedded_samplers.mutex, mtx_plain);
+   device->embedded_samplers.map =
+      _mesa_hash_table_create(NULL,
+                              embedded_sampler_key_hash,
+                              embedded_sampler_key_equal);
+}
+
+void
+anv_device_finish_embedded_samplers(struct anv_device *device)
+{
+   hash_table_foreach(device->embedded_samplers.map, entry) {
+      anv_embedded_sampler_free(device, entry->data);
+   }
+   ralloc_free(device->embedded_samplers.map);
+   simple_mtx_destroy(&device->embedded_samplers.mutex);
+}
+
+static VkResult
+anv_shader_bin_get_embedded_samplers(struct anv_device *device,
+                                     struct anv_shader_bin *shader,
+                                     const struct anv_pipeline_bind_map *bind_map)
+{
+   VkResult result = VK_SUCCESS;
+
+   simple_mtx_lock(&device->embedded_samplers.mutex);
+
+   for (uint32_t i = 0; i < bind_map->embedded_sampler_count; i++) {
+      struct hash_entry *entry =
+         _mesa_hash_table_search(device->embedded_samplers.map,
+                                 &bind_map->embedded_sampler_to_binding[i].key);
+      if (entry == NULL) {
+         shader->embedded_samplers[i] =
+            vk_zalloc(&device->vk.alloc,
+                      sizeof(struct anv_embedded_sampler), 8,
+                      VK_SYSTEM_ALLOCATION_SCOPE_DEVICE);
+         if (shader->embedded_samplers[i] == NULL) {
+            result = vk_error(device, VK_ERROR_OUT_OF_HOST_MEMORY);
+            goto err;
+         }
+
+         anv_genX(device->info, emit_embedded_sampler)(
+            device, shader->embedded_samplers[i],
+            &bind_map->embedded_sampler_to_binding[i]);
+         _mesa_hash_table_insert(device->embedded_samplers.map,
+                                 &shader->embedded_samplers[i]->key,
+                                 shader->embedded_samplers[i]);
+      } else {
+         shader->embedded_samplers[i] = anv_embedded_sampler_ref(entry->data);
+      }
+   }
+
+ err:
+   simple_mtx_unlock(&device->embedded_samplers.mutex);
+   return result;
+}
+
+/**
+ *
+ */
+
 static bool
 anv_shader_bin_serialize(struct vk_pipeline_cache_object *object,
                          struct blob *blob);
@@ -52,12 +160,8 @@ anv_shader_bin_destroy(struct vk_device *_device,
    struct anv_shader_bin *shader =
       container_of(object, struct anv_shader_bin, base);
 
-   for (uint32_t i = 0; i < shader->bind_map.embedded_sampler_count; i++) {
-      anv_state_pool_free(&device->dynamic_state_db_pool,
-                          shader->embedded_samplers[i].sampler_state);
-      anv_state_pool_free(&device->dynamic_state_db_pool,
-                          shader->embedded_samplers[i].border_color_state);
-   }
+   for (uint32_t i = 0; i < shader->bind_map.embedded_sampler_count; i++)
+      anv_embedded_sampler_unref(device, shader->embedded_samplers[i]);
 
    anv_state_pool_free(&device->instruction_state_pool, shader->kernel);
    vk_pipeline_cache_object_finish(&shader->base);
@@ -87,7 +191,7 @@ anv_shader_bin_rewrite_embedded_samplers(struct anv_device *device,
    for (uint32_t i = 0; i < bind_map->embedded_sampler_count; i++) {
       reloc_values[rv_count++] = (struct brw_shader_reloc_value) {
          .id = BRW_SHADER_RELOC_EMBEDDED_SAMPLER_HANDLE + i,
-         .value = shader->embedded_samplers[i].sampler_state.offset,
+         .value = shader->embedded_samplers[i]->sampler_state.offset,
       };
    }
 
@@ -131,7 +235,7 @@ anv_shader_bin_create(struct anv_device *device,
                       bind_map->embedded_sampler_count);
    VK_MULTIALLOC_DECL(&ma, struct brw_kernel_arg_desc, kernel_args,
                       bind_map->kernel_arg_count);
-   VK_MULTIALLOC_DECL(&ma, struct anv_embedded_sampler, embedded_samplers,
+   VK_MULTIALLOC_DECL(&ma, struct anv_embedded_sampler *, embedded_samplers,
                       bind_map->embedded_sampler_count);
 
    if (!vk_multialloc_alloc(&ma, &device->vk.alloc,
@@ -149,12 +253,14 @@ anv_shader_bin_create(struct anv_device *device,
    memcpy(shader->kernel.map, kernel_data, kernel_size);
    shader->kernel_size = kernel_size;
 
-   for (uint32_t i = 0; i < bind_map->embedded_sampler_count; i++) {
-      anv_genX(device->info, emit_embedded_sampler)(
-         device, &embedded_samplers[i],
-         &bind_map->embedded_sampler_to_binding[i]);
+   if (bind_map->embedded_sampler_count > 0) {
+      shader->embedded_samplers = embedded_samplers;
+      if (anv_shader_bin_get_embedded_samplers(device, shader, bind_map) != VK_SUCCESS) {
+         anv_state_pool_free(&device->instruction_state_pool, shader->kernel);
+         vk_free(&device->vk.alloc, shader);
+         return NULL;
+      }
    }
-   shader->embedded_samplers = embedded_samplers;
 
    uint64_t shader_data_addr =
       device->physical->va.instruction_state_pool.addr +

src/intel/vulkan/anv_private.h

@@ -1362,6 +1362,9 @@ struct anv_pipeline_sets_layout;
 struct anv_push_descriptor_info;
 enum anv_dynamic_push_bits;
 
+void anv_device_init_embedded_samplers(struct anv_device *device);
+void anv_device_finish_embedded_samplers(struct anv_device *device);
+
 extern const struct vk_pipeline_cache_object_ops *const anv_cache_import_ops[2];
 
 struct anv_shader_bin *
@@ -2009,6 +2012,11 @@ struct anv_device {
     struct anv_device_astc_emu                   astc_emu;
 
     struct intel_bind_timeline bind_timeline; /* Xe only */
+
+    struct {
+       simple_mtx_t                              mutex;
+       struct hash_table                        *map;
+    }                                            embedded_samplers;
 };
 
 static inline uint32_t
@@ -2961,6 +2969,22 @@ struct anv_pipeline_binding {
    };
 };
 
+struct anv_embedded_sampler_key {
+   /** No need to track binding elements for embedded samplers as :
+    *
+    *    VUID-VkDescriptorSetLayoutBinding-flags-08006:
+    *
+    *       "If VkDescriptorSetLayoutCreateInfo:flags contains
+    *        VK_DESCRIPTOR_SET_LAYOUT_CREATE_EMBEDDED_IMMUTABLE_SAMPLERS_BIT_EXT,
+    *        descriptorCount must: less than or equal to 1"
+    *
+    * The following struct can be safely hash as it doesn't include in
+    * address/offset.
+    */
+   uint32_t sampler[4];
+   uint32_t color[4];
+};
+
 struct anv_pipeline_embedded_sampler_binding {
    /** The descriptor set this sampler belongs to */
    uint8_t set;
@@ -2968,18 +2992,8 @@ struct anv_pipeline_embedded_sampler_binding {
    /** The binding in the set this sampler belongs to */
    uint32_t binding;
 
-   /* No need to track binding elements for embedded samplers as :
-    *
-    *    VUID-VkDescriptorSetLayoutBinding-flags-08006:
-    *
-    *       "If VkDescriptorSetLayoutCreateInfo:flags contains
-    *        VK_DESCRIPTOR_SET_LAYOUT_CREATE_EMBEDDED_IMMUTABLE_SAMPLERS_BIT_EXT,
-    *        descriptorCount must: less than or equal to 1"
-    */
-
-   uint32_t sampler_state[4];
-
-   uint32_t border_color[4];
+   /** The data configuring the sampler */
+   struct anv_embedded_sampler_key key;
 };
 
 struct anv_push_range {
@@ -4354,6 +4368,10 @@ struct anv_shader_upload_params {
 };
 
 struct anv_embedded_sampler {
+   uint32_t ref_cnt;
+
+   struct anv_embedded_sampler_key key;
+
    struct anv_state sampler_state;
    struct anv_state border_color_state;
 };
@@ -4382,9 +4400,9 @@ struct anv_shader_bin {
 
    /* Not saved in the pipeline cache.
     *
-    * Array of length bind_map.embedded_sampler_count
+    * Array of pointers of length bind_map.embedded_sampler_count
     */
-   struct anv_embedded_sampler *embedded_samplers;
+   struct anv_embedded_sampler **embedded_samplers;
 };
 
 static inline struct anv_shader_bin *

src/intel/vulkan/genX_init_state.c

@@ -1313,12 +1313,15 @@ genX(emit_embedded_sampler)(struct anv_device *device,
                             struct anv_embedded_sampler *sampler,
                             struct anv_pipeline_embedded_sampler_binding *binding)
 {
+   sampler->ref_cnt = 1;
+   memcpy(&sampler->key, &binding->key, sizeof(binding->key));
+
    sampler->border_color_state =
       anv_state_pool_alloc(&device->dynamic_state_db_pool,
                            sizeof(struct gfx8_border_color), 64);
    memcpy(sampler->border_color_state.map,
-          binding->border_color,
-          sizeof(binding->border_color));
+          binding->key.color,
+          sizeof(binding->key.color));
 
    sampler->sampler_state =
       anv_state_pool_alloc(&device->dynamic_state_db_pool,
@@ -1332,7 +1335,7 @@ genX(emit_embedded_sampler)(struct anv_device *device,
 
    for (uint32_t i = 0; i < GENX(SAMPLER_STATE_length); i++) {
       ((uint32_t *)sampler->sampler_state.map)[i] =
-         dwords[i] | binding->sampler_state[i];
+         dwords[i] | binding->key.sampler[i];
    }
 }