nvk: pipeline shader cache

dEQP-VK.pipeline.monolithic.cache.*

Test run totals:
  Passed:        773/773 (100.0%)
  Failed:        0/773 (0.0%)
  Not supported: 0/773 (0.0%)
  Warnings:      0/773 (0.0%)
  Waived:        0/773 (0.0%)

Timing these test:
Before:
real    0m11,304s
user    0m9,442s
sys     0m0,477s

After:
real    0m3,470s
user    0m1,962s
sys     0m0,504s

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

src/nouveau/vulkan/nvk_compute_pipeline.c

@@ -173,32 +173,63 @@ nvk_compute_pipeline_create(struct nvk_device *dev,
    VkPipelineCreateFlags2KHR pipeline_flags =
       vk_compute_pipeline_create_flags(pCreateInfo);
 
+   if (pipeline_flags &
+       VK_PIPELINE_CREATE_CAPTURE_INTERNAL_REPRESENTATIONS_BIT_KHR)
+      cache = NULL;
+
    struct vk_pipeline_robustness_state robustness;
    vk_pipeline_robustness_state_fill(&dev->vk, &robustness,
                                      pCreateInfo->pNext,
                                      pCreateInfo->stage.pNext);
 
-   nir_shader *nir;
-   result = nvk_shader_stage_to_nir(dev, &pCreateInfo->stage, &robustness,
-                                    cache, NULL, &nir);
-   if (result != VK_SUCCESS)
-      goto fail;
-
-   struct nvk_shader *shader = nvk_shader_init(dev);
-   pipeline->base.shaders[MESA_SHADER_COMPUTE] = shader;
-
-   nvk_lower_nir(dev, nir, &robustness, false, pipeline_layout, shader);
-
    unsigned char sha1[SHA1_DIGEST_LENGTH];
    nvk_hash_shader(sha1, &pCreateInfo->stage, &robustness, false,
                    pipeline_layout, NULL);
 
-   result = nvk_compile_nir(dev, nir, pipeline_flags, &robustness, NULL,
-                            shader);
-   ralloc_free(nir);
+   bool cache_hit = false;
+   struct vk_pipeline_cache_object *cache_obj = NULL;
+
+   if (cache) {
+      cache_obj = vk_pipeline_cache_lookup_object(cache, &sha1, sizeof(sha1),
+                                                  &nvk_shader_ops, &cache_hit);
+      pipeline->base.shaders[MESA_SHADER_COMPUTE] =
+         container_of(cache_obj, struct nvk_shader, base);
+      result = VK_SUCCESS;
+   }
+
+   if (!cache_obj) {
+      nir_shader *nir;
+      result = nvk_shader_stage_to_nir(dev, &pCreateInfo->stage, &robustness,
+                                       cache, NULL, &nir);
+      if (result != VK_SUCCESS)
+         goto fail;
+
+      struct nvk_shader *shader = nvk_shader_init(dev, sha1, SHA1_DIGEST_LENGTH);
+      if(shader == NULL)
+         return vk_error(dev, VK_ERROR_OUT_OF_HOST_MEMORY);
+
+      nvk_lower_nir(dev, nir, &robustness, false, pipeline_layout, shader);
+
+      result = nvk_compile_nir(dev, nir, pipeline_flags, &robustness, NULL, cache, shader);
+
+      if (result == VK_SUCCESS) {
+         cache_obj = &shader->base;
+
+         if (cache)
+            cache_obj = vk_pipeline_cache_add_object(cache, cache_obj);
+
+         pipeline->base.shaders[MESA_SHADER_COMPUTE] =
+            container_of(cache_obj, struct nvk_shader, base);
+      }
+
+      ralloc_free(nir);
+   }
+
    if (result != VK_SUCCESS)
       goto fail;
 
+   struct nvk_shader *shader = container_of(cache_obj, struct nvk_shader, base);
+
    result = nvk_shader_upload(dev, shader);
    if (result != VK_SUCCESS)
       goto fail;

src/nouveau/vulkan/nvk_graphics_pipeline.c

@@ -176,6 +176,10 @@ nvk_graphics_pipeline_create(struct nvk_device *dev,
    VkPipelineCreateFlags2KHR pipeline_flags =
       vk_graphics_pipeline_create_flags(pCreateInfo);
 
+   if (pipeline_flags &
+       VK_PIPELINE_CREATE_CAPTURE_INTERNAL_REPRESENTATIONS_BIT_KHR)
+      cache = NULL;
+
    struct vk_graphics_pipeline_all_state all;
    struct vk_graphics_pipeline_state state = {};
    result = vk_graphics_pipeline_state_fill(&dev->vk, &state, pCreateInfo,
@@ -186,6 +190,12 @@ nvk_graphics_pipeline_create(struct nvk_device *dev,
    nir_shader *nir[MESA_SHADER_STAGES] = {};
    struct vk_pipeline_robustness_state robustness[MESA_SHADER_STAGES];
 
+   struct vk_pipeline_cache_object *cache_objs[MESA_SHADER_STAGES] = {};
+
+   struct nak_fs_key fs_key_tmp, *fs_key = NULL;
+   nvk_populate_fs_key(&fs_key_tmp, state.ms, &state);
+   fs_key = &fs_key_tmp;
+
    for (uint32_t i = 0; i < pCreateInfo->stageCount; i++) {
       const VkPipelineShaderStageCreateInfo *sinfo = &pCreateInfo->pStages[i];
       gl_shader_stage stage = vk_to_mesa_shader_stage(sinfo->stage);
@@ -199,6 +209,33 @@ nvk_graphics_pipeline_create(struct nvk_device *dev,
 
       vk_pipeline_robustness_state_fill(&dev->vk, &robustness[stage],
                                         pCreateInfo->pNext, sinfo->pNext);
+   }
+
+   for (gl_shader_stage stage = 0; stage < MESA_SHADER_STAGES; stage++) {
+      const VkPipelineShaderStageCreateInfo *sinfo = infos[stage];
+      if (sinfo == NULL)
+         continue;
+
+      unsigned char sha1[SHA1_DIGEST_LENGTH];
+      nvk_hash_shader(sha1, sinfo, &robustness[stage],
+                      state.rp->view_mask != 0, pipeline_layout,
+                      stage == MESA_SHADER_FRAGMENT ? fs_key : NULL);
+
+      if (cache) {
+         bool cache_hit = false;
+         cache_objs[stage] = vk_pipeline_cache_lookup_object(cache, &sha1, sizeof(sha1),
+                                                             &nvk_shader_ops, &cache_hit);
+         if (cache_objs[stage]) {
+            pipeline->base.shaders[stage] =
+               container_of(cache_objs[stage], struct nvk_shader, base);
+         }
+      }
+   }
+
+   for (gl_shader_stage stage = 0; stage < MESA_SHADER_STAGES; stage++) {
+      const VkPipelineShaderStageCreateInfo *sinfo = infos[stage];
+      if (sinfo == NULL || cache_objs[stage])
+         continue;
 
       result = nvk_shader_stage_to_nir(dev, sinfo, &robustness[stage],
                                        cache, NULL, &nir[stage]);
@@ -216,33 +253,40 @@ nvk_graphics_pipeline_create(struct nvk_device *dev,
       if (sinfo == NULL)
          continue;
 
-      pipeline->base.shaders[stage] = nvk_shader_init(dev);
+      if (!cache_objs[stage]) {
+         unsigned char sha1[SHA1_DIGEST_LENGTH];
+         nvk_hash_shader(sha1, sinfo, &robustness[stage],
+                         state.rp->view_mask != 0, pipeline_layout,
+                         stage == MESA_SHADER_FRAGMENT ? fs_key : NULL);
 
-      nvk_lower_nir(dev, nir[stage], &robustness[stage],
-                    state.rp->view_mask != 0, pipeline_layout,
-                    pipeline->base.shaders[stage]);
-   }
+         struct nvk_shader *shader = nvk_shader_init(dev, sha1, SHA1_DIGEST_LENGTH);
+         if(shader == NULL) {
+            result = vk_error(dev, VK_ERROR_OUT_OF_HOST_MEMORY);
+            goto fail;
+         }
 
-   for (gl_shader_stage stage = 0; stage < MESA_SHADER_STAGES; stage++) {
-      const VkPipelineShaderStageCreateInfo *sinfo = infos[stage];
-      if (sinfo == NULL)
-         continue;
+         nvk_lower_nir(dev, nir[stage], &robustness[stage],
+                       state.rp->view_mask != 0, pipeline_layout, shader);
 
-      struct nak_fs_key fs_key_tmp, *fs_key = NULL;
-      if (stage == MESA_SHADER_FRAGMENT) {
-         nvk_populate_fs_key(&fs_key_tmp, state.ms, &state);
-         fs_key = &fs_key_tmp;
+         result = nvk_compile_nir(dev, nir[stage],
+                                  pipeline_flags, &robustness[stage],
+                                  stage == MESA_SHADER_FRAGMENT ? fs_key : NULL,
+                                  cache, shader);
+
+         if (result == VK_SUCCESS) {
+            cache_objs[stage] = &shader->base;
+
+            if (cache)
+               cache_objs[stage] = vk_pipeline_cache_add_object(cache,
+                                                                cache_objs[stage]);
+
+            pipeline->base.shaders[stage] =
+               container_of(cache_objs[stage], struct nvk_shader, base);
+         }
+
+         ralloc_free(nir[stage]);
       }
 
-      unsigned char sha1[SHA1_DIGEST_LENGTH];
-      nvk_hash_shader(sha1, sinfo, &robustness[stage],
-                      state.rp->view_mask != 0,
-                      pipeline_layout, fs_key);
-
-      result = nvk_compile_nir(dev, nir[stage], pipeline_flags,
-                               &robustness[stage], fs_key,
-                               pipeline->base.shaders[stage]);
-      ralloc_free(nir[stage]);
       if (result != VK_SUCCESS)
          goto fail;
 

src/nouveau/vulkan/nvk_pipeline.c

@@ -32,9 +32,6 @@ nvk_pipeline_free(struct nvk_device *dev,
                   struct nvk_pipeline *pipeline,
                   const VkAllocationCallbacks *pAllocator)
 {
-   for (uint32_t s = 0; s < ARRAY_SIZE(pipeline->shaders); s++)
-      nvk_shader_finish(dev, pipeline->shaders[s]);
-
    vk_object_free(&dev->vk, pAllocator, pipeline);
 }
 

src/nouveau/vulkan/nvk_shader.c

@@ -439,15 +439,21 @@ nvk_compile_nir_with_nak(struct nvk_physical_device *pdev,
 }
 
 struct nvk_shader *
-nvk_shader_init(struct nvk_device *dev)
+nvk_shader_init(struct nvk_device *dev, const void *key_data, size_t key_size)
 {
    VK_MULTIALLOC(ma);
    VK_MULTIALLOC_DECL(&ma, struct nvk_shader, shader, 1);
+   VK_MULTIALLOC_DECL_SIZE(&ma, char, obj_key_data, key_size);
 
    if (!vk_multialloc_zalloc(&ma, &dev->vk.alloc,
                              VK_SYSTEM_ALLOCATION_SCOPE_DEVICE))
       return NULL;
 
+   memcpy(obj_key_data, key_data, key_size);
+
+   vk_pipeline_cache_object_init(&dev->vk, &shader->base,
+                                 &nvk_shader_ops, obj_key_data, key_size);
+
    return shader;
 }
 
@@ -456,6 +462,7 @@ nvk_compile_nir(struct nvk_device *dev, nir_shader *nir,
                 VkPipelineCreateFlagBits2KHR pipeline_flags,
                 const struct vk_pipeline_robustness_state *rs,
                 const struct nak_fs_key *fs_key,
+                struct vk_pipeline_cache *cache,
                 struct nvk_shader *shader)
 {
    struct nvk_physical_device *pdev = nvk_device_physical(dev);
@@ -532,7 +539,7 @@ nvk_shader_finish(struct nvk_device *dev, struct nvk_shader *shader)
    if (shader->nak) {
       nak_shader_bin_destroy(shader->nak);
    } else {
-      /* This came from codegen, just free it */
+      /* This came from codegen or deserialize, just free it */
       free((void *)shader->code_ptr);
    }
 
@@ -574,3 +581,79 @@ nvk_hash_shader(unsigned char *hash,
 
    _mesa_sha1_final(&ctx, hash);
 }
+
+static bool
+nvk_shader_serialize(struct vk_pipeline_cache_object *object,
+                     struct blob *blob);
+
+static struct vk_pipeline_cache_object *
+nvk_shader_deserialize(struct vk_pipeline_cache *cache,
+                       const void *key_data,
+                       size_t key_size,
+                       struct blob_reader *blob);
+
+void
+nvk_shader_destroy(struct vk_device *_dev,
+                   struct vk_pipeline_cache_object *object)
+{
+   struct nvk_device *dev =
+      container_of(_dev, struct nvk_device, vk);
+   struct nvk_shader *shader =
+      container_of(object, struct nvk_shader, base);
+
+   nvk_shader_finish(dev, shader);
+}
+
+const struct vk_pipeline_cache_object_ops nvk_shader_ops = {
+   .serialize = nvk_shader_serialize,
+   .deserialize = nvk_shader_deserialize,
+   .destroy = nvk_shader_destroy,
+};
+
+static bool
+nvk_shader_serialize(struct vk_pipeline_cache_object *object,
+                     struct blob *blob)
+{
+   struct nvk_shader *shader =
+      container_of(object, struct nvk_shader, base);
+
+   blob_write_bytes(blob, &shader->info, sizeof(shader->info));
+   blob_write_bytes(blob, &shader->cbuf_map, sizeof(shader->cbuf_map));
+   blob_write_uint32(blob, shader->code_size);
+   blob_write_bytes(blob, shader->code_ptr, shader->code_size);
+
+   return true;
+}
+
+static struct vk_pipeline_cache_object *
+nvk_shader_deserialize(struct vk_pipeline_cache *cache,
+                       const void *key_data,
+                       size_t key_size,
+                       struct blob_reader *blob)
+{
+   struct nvk_device *dev =
+      container_of(cache->base.device, struct nvk_device, vk);
+   struct nvk_shader *shader =
+      nvk_shader_init(dev, key_data, key_size);
+
+   if (!shader)
+      return NULL;
+
+   blob_copy_bytes(blob, &shader->info, sizeof(shader->info));
+   blob_copy_bytes(blob, &shader->cbuf_map, sizeof(shader->cbuf_map));
+
+   shader->code_size = blob_read_uint32(blob);
+   void *code_ptr = malloc(shader->code_size);
+   if (!code_ptr)
+      goto fail;
+
+   blob_copy_bytes(blob, code_ptr, shader->code_size);
+   shader->code_ptr = code_ptr;
+
+   return &shader->base;
+
+fail:
+   /* nvk_shader_destroy frees both shader and shader->xfb */
+   nvk_shader_destroy(cache->base.device, &shader->base);
+   return NULL;
+}

src/nouveau/vulkan/nvk_shader.h

@@ -8,6 +8,8 @@
 #include "nvk_private.h"
 #include "nvk_device_memory.h"
 
+#include "vk_pipeline_cache.h"
+
 #include "nak.h"
 #include "nir.h"
 #include "nouveau_bo.h"
@@ -46,6 +48,8 @@ struct nvk_cbuf_map {
 };
 
 struct nvk_shader {
+   struct vk_pipeline_cache_object base;
+
    struct nak_shader_info info;
    struct nvk_cbuf_map cbuf_map;
 
@@ -122,13 +126,16 @@ nvk_compile_nir(struct nvk_device *dev, nir_shader *nir,
                 VkPipelineCreateFlagBits2KHR pipeline_flags,
                 const struct vk_pipeline_robustness_state *rstate,
                 const struct nak_fs_key *fs_key,
+                struct vk_pipeline_cache *cache,
                 struct nvk_shader *shader);
 
 VkResult
 nvk_shader_upload(struct nvk_device *dev, struct nvk_shader *shader);
 
 struct nvk_shader *
-nvk_shader_init(struct nvk_device *dev);
+nvk_shader_init(struct nvk_device *dev, const void *key_data, size_t key_size);
+
+extern const struct vk_pipeline_cache_object_ops nvk_shader_ops;
 
 void
 nvk_shader_finish(struct nvk_device *dev, struct nvk_shader *shader);
@@ -141,6 +148,10 @@ nvk_hash_shader(unsigned char *hash,
                 const struct vk_pipeline_layout *layout,
                 const struct nak_fs_key *fs_key);
 
+void
+nvk_shader_destroy(struct vk_device *dev,
+                   struct vk_pipeline_cache_object *object);
+
 /* Codegen wrappers.
  *
  * TODO: Delete these once NAK supports everything.