radv: fix and rework shaders upload with GPL

For fast-linking, we really want to upload the binaries directly in
a library to avoid creating and uploading at pipeline creation time.

To achieve that, add a radeon_winsys_bo pointer to radv_shader in
order to indicate that a shader is already uploaded. When a lib is
imported, the pipeline slab BO is also incremented to make sure it's
not freed when the lib is destroyed.

This also allows to free binaries right after they are uploaded.

Signed-off-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Reviewed-by: Tatsuyuki Ishi <ishitatsuyuki@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/18860>

src/amd/vulkan/radv_cmd_buffer.c

@@ -1518,6 +1518,22 @@ radv_emit_graphics_pipeline(struct radv_cmd_buffer *cmd_buffer)
 
    radv_cs_add_buffer(cmd_buffer->device->ws, cmd_buffer->cs, pipeline->base.slab_bo);
 
+   /* With graphics pipeline library, binaries are uploaded from a library and they hold a pointer
+    * to the slab BO.
+    */
+   for (unsigned s = 0; s < MESA_VULKAN_SHADER_STAGES; s++) {
+      struct radv_shader *shader = pipeline->base.shaders[s];
+
+      if (!shader || !shader->bo)
+         continue;
+
+      radv_cs_add_buffer(cmd_buffer->device->ws, cmd_buffer->cs, shader->bo);
+   }
+
+   if (pipeline->base.gs_copy_shader && pipeline->base.gs_copy_shader->bo) {
+      radv_cs_add_buffer(cmd_buffer->device->ws, cmd_buffer->cs, pipeline->base.gs_copy_shader->bo);
+   }
+
    if (unlikely(cmd_buffer->device->trace_bo))
       radv_save_pipeline(cmd_buffer, &pipeline->base);
 

src/amd/vulkan/radv_pipeline.c

@@ -227,11 +227,6 @@ radv_pipeline_destroy(struct radv_device *device, struct radv_pipeline *pipeline
       radv_pipeline_layout_finish(device, &gfx_pipeline_lib->layout);
 
       for (unsigned i = 0; i < MESA_VULKAN_SHADER_STAGES; ++i) {
-         if (pipeline->shaders[i]) {
-            free(pipeline->shaders[i]->binary);
-            pipeline->shaders[i]->binary = NULL;
-         }
-
          ralloc_free(pipeline->retained_shaders[i].nir);
       }
 
@@ -1652,12 +1647,23 @@ radv_graphics_pipeline_import_lib(struct radv_graphics_pipeline *pipeline,
          continue;
 
       pipeline->base.shaders[s] = radv_shader_ref(lib->base.base.shaders[s]);
+
+      /* Hold a pointer to the slab BO to indicate the shader is already uploaded. */
+      pipeline->base.shaders[s]->bo = lib->base.base.slab_bo;
    }
 
    /* Import the GS copy shader if present. */
    if (lib->base.base.gs_copy_shader) {
       assert(!pipeline->base.gs_copy_shader);
       pipeline->base.gs_copy_shader = radv_shader_ref(lib->base.base.gs_copy_shader);
+
+      /* Hold a pointer to the slab BO to indicate the shader is already uploaded. */
+      pipeline->base.gs_copy_shader->bo = lib->base.base.slab_bo;
+   }
+
+   /* Refcount the slab BO to make sure it's not freed when the library is destroyed. */
+   if (lib->base.base.slab) {
+      p_atomic_inc(&lib->base.base.slab->ref_count);
    }
 
    /* Import the PS epilog if present. */
@@ -3222,10 +3228,13 @@ radv_upload_shaders(struct radv_device *device, struct radv_pipeline *pipeline)
       if (!shader)
          continue;
 
+      if (shader->bo)
+         continue;
+
       code_size += align(shader->code_size, RADV_SHADER_ALLOC_ALIGNMENT);
    }
 
-   if (pipeline->gs_copy_shader) {
+   if (pipeline->gs_copy_shader && !pipeline->gs_copy_shader->bo) {
       code_size += align(pipeline->gs_copy_shader->code_size, RADV_SHADER_ALLOC_ALIGNMENT);
    }
 
@@ -3246,6 +3255,9 @@ radv_upload_shaders(struct radv_device *device, struct radv_pipeline *pipeline)
       if (!shader)
          continue;
 
+      if (shader->bo)
+         continue;
+
       shader->va = slab_va + slab_offset;
 
       void *dest_ptr = slab_ptr + slab_offset;
@@ -3255,7 +3267,7 @@ radv_upload_shaders(struct radv_device *device, struct radv_pipeline *pipeline)
       slab_offset += align(shader->code_size, RADV_SHADER_ALLOC_ALIGNMENT);
    }
 
-   if (pipeline->gs_copy_shader) {
+   if (pipeline->gs_copy_shader && !pipeline->gs_copy_shader->bo) {
       pipeline->gs_copy_shader->va = slab_va + slab_offset;
 
       void *dest_ptr = slab_ptr + slab_offset;
@@ -4210,13 +4222,13 @@ radv_create_shaders(struct radv_pipeline *pipeline, struct radv_pipeline_layout
       }
    }
 
-   if (pipeline->gs_copy_shader && !(flags & VK_PIPELINE_CREATE_LIBRARY_BIT_KHR)) {
+   if (pipeline->gs_copy_shader) {
       free(pipeline->gs_copy_shader->binary);
       pipeline->gs_copy_shader->binary = NULL;
    }
 
    for (int i = 0; i < MESA_VULKAN_SHADER_STAGES; ++i) {
-      if (pipeline->shaders[i] && !(flags & VK_PIPELINE_CREATE_LIBRARY_BIT_KHR)) {
+      if (pipeline->shaders[i]) {
          free(pipeline->shaders[i]->binary);
          pipeline->shaders[i]->binary = NULL;
       }

src/amd/vulkan/radv_shader.h

@@ -484,6 +484,7 @@ union radv_shader_arena_block {
 struct radv_shader {
    uint32_t ref_count;
 
+   struct radeon_winsys_bo *bo; /* Not NULL if imported from a lib */
    uint64_t va;
 
    struct ac_shader_config config;