radv: Add utilities to serialize and deserialize shader allocation info

Can be used to capture/replay an arbitrary sequence of shader
allocations while preserving VAs.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/23516>

src/amd/vulkan/radv_private.h

@@ -2200,6 +2200,13 @@ struct radv_pipeline_group_handle {
    };
 };
 
+struct radv_serialized_shader_arena_block {
+   uint32_t offset;
+   uint32_t size;
+   uint64_t arena_va;
+   uint32_t arena_size;
+};
+
 enum radv_depth_clamp_mode {
    RADV_DEPTH_CLAMP_MODE_VIEWPORT = 0,    /* Clamp to the viewport min/max depth bounds */
    RADV_DEPTH_CLAMP_MODE_ZERO_TO_ONE = 1, /* Clamp between 0.0f and 1.0f */

src/amd/vulkan/radv_shader.c

@@ -1261,6 +1261,70 @@ radv_free_shader_memory(struct radv_device *device, union radv_shader_arena_bloc
    mtx_unlock(&device->shader_arena_mutex);
 }
 
+struct radv_serialized_shader_arena_block
+radv_serialize_shader_arena_block(union radv_shader_arena_block *block)
+{
+   struct radv_serialized_shader_arena_block serialized_block = {
+      .offset = block->offset,
+      .size = block->size,
+      .arena_va = block->arena->bo->va,
+      .arena_size = block->arena->size,
+   };
+   return serialized_block;
+}
+
+union radv_shader_arena_block *
+radv_replay_shader_arena_block(struct radv_device *device, const struct radv_serialized_shader_arena_block *src,
+                               void *ptr)
+{
+   mtx_lock(&device->shader_arena_mutex);
+   uint64_t va = src->arena_va;
+   void *data = _mesa_hash_table_u64_search(device->capture_replay_arena_vas, va);
+
+   if (!data) {
+      struct radv_shader_arena *arena =
+         radv_create_shader_arena(device, NULL, 0, src->arena_size, false, src->arena_va);
+      if (!arena) {
+         mtx_unlock(&device->shader_arena_mutex);
+         return NULL;
+      }
+
+      _mesa_hash_table_u64_insert(device->capture_replay_arena_vas, src->arena_va, arena);
+      list_addtail(&arena->list, &device->shader_arenas);
+      data = arena;
+   }
+   mtx_unlock(&device->shader_arena_mutex);
+
+   uint32_t block_begin = src->offset;
+   uint32_t block_end = src->offset + src->size;
+
+   struct radv_shader_arena *arena = data;
+   list_for_each_entry (union radv_shader_arena_block, hole, &arena->entries, list) {
+      /* Only consider holes, not allocated shaders */
+      if (!hole->freelist.prev)
+         continue;
+
+      if (hole->offset + hole->size < src->offset)
+         continue;
+
+      uint32_t hole_begin = hole->offset;
+      uint32_t hole_end = hole->offset + hole->size;
+
+      /* If another allocated block overlaps the current replay block, allocation is impossible */
+      if (block_begin > hole_begin || (hole_end < block_end && hole_end >= block_begin))
+         return NULL;
+
+      union radv_shader_arena_block *block = insert_block(device, hole, block_begin - hole_begin, src->size, NULL);
+      if (!block)
+         return NULL;
+
+      block->freelist.prev = NULL;
+      block->freelist.next = ptr;
+      return hole;
+   }
+   return NULL;
+}
+
 void
 radv_init_shader_arenas(struct radv_device *device)
 {

src/amd/vulkan/radv_shader.h

@@ -48,6 +48,7 @@ struct radv_pipeline_key;
 struct radv_shader_args;
 struct radv_vs_input_state;
 struct radv_shader_args;
+struct radv_serialized_shader_arena_block;
 
 enum radv_required_subgroup_size {
    RADV_REQUIRED_NONE = 0,
@@ -646,6 +647,12 @@ bool radv_shader_dma_submit(struct radv_device *device, struct radv_shader_dma_s
 union radv_shader_arena_block *radv_alloc_shader_memory(struct radv_device *device, uint32_t size, bool replayable,
                                                         void *ptr);
 
+union radv_shader_arena_block *radv_replay_shader_arena_block(struct radv_device *device,
+                                                              const struct radv_serialized_shader_arena_block *src,
+                                                              void *ptr);
+
+struct radv_serialized_shader_arena_block radv_serialize_shader_arena_block(union radv_shader_arena_block *block);
+
 void radv_free_shader_memory(struct radv_device *device, union radv_shader_arena_block *alloc);
 
 struct radv_shader *radv_create_trap_handler_shader(struct radv_device *device);