From 93d0c66b27f440faa7ce3d2ec2f51135137c07ac Mon Sep 17 00:00:00 2001
From: Jason Ekstrand <jason@jlekstrand.net>
Date: Mon, 4 Oct 2021 15:14:29 -0500
Subject: [PATCH] vulkan/pipeline_cache: Add helpers for storing NIR in the
 cache

Reviewed-by: Connor Abbott <cwabbott0@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/13184>
---
 src/vulkan/runtime/meson.build         |  2 +-
 src/vulkan/runtime/vk_pipeline_cache.c | 61 ++++++++++++++++++++++++++++++++--
 src/vulkan/runtime/vk_pipeline_cache.h | 14 ++++++++
 3 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/src/vulkan/runtime/meson.build b/src/vulkan/runtime/meson.build
index 1969974..7b01e2c 100644
--- a/src/vulkan/runtime/meson.build
+++ b/src/vulkan/runtime/meson.build
@@ -79,7 +79,7 @@ vulkan_runtime_files = files(
 vulkan_runtime_deps = [
   vulkan_wsi_deps,
   idep_mesautil,
-  idep_nir_headers,
+  idep_nir,
   idep_vulkan_util,
 ]
 
diff --git a/src/vulkan/runtime/vk_pipeline_cache.c b/src/vulkan/runtime/vk_pipeline_cache.c
index 2214b27..e7068be 100644
--- a/src/vulkan/runtime/vk_pipeline_cache.c
+++ b/src/vulkan/runtime/vk_pipeline_cache.c
@@ -29,6 +29,8 @@
 #include "vk_log.h"
 #include "vk_physical_device.h"
 
+#include "compiler/nir/nir_serialize.h"
+
 #include "util/blob.h"
 #include "util/debug.h"
 #include "util/disk_cache.h"
@@ -327,15 +329,13 @@ vk_pipeline_cache_lookup_object(struct vk_pipeline_cache *cache,
       return NULL;
    }
 
-   if (object->ops == &raw_data_object_ops) {
+   if (object->ops == &raw_data_object_ops && ops != &raw_data_object_ops) {
       /* The object isn't fully formed yet and we need to deserialize it into
        * a real object before it can be used.
        */
       struct raw_data_object *data_obj =
          container_of(object, struct raw_data_object, base);
 
-      assert(ops != &raw_data_object_ops);
-
       struct vk_pipeline_cache_object *real_object =
          vk_pipeline_cache_object_deserialize(cache,
                                               data_obj->base.key_data,
@@ -412,6 +412,61 @@ vk_pipeline_cache_add_object(struct vk_pipeline_cache *cache,
    }
 }
 
+nir_shader *
+vk_pipeline_cache_lookup_nir(struct vk_pipeline_cache *cache,
+                             const void *key_data, size_t key_size,
+                             const struct nir_shader_compiler_options *nir_options,
+                             bool *cache_hit, void *mem_ctx)
+{
+   struct vk_pipeline_cache_object *object =
+      vk_pipeline_cache_lookup_object(cache, key_data, key_size,
+                                      &raw_data_object_ops, cache_hit);
+   if (object == NULL)
+      return NULL;
+
+   struct raw_data_object *data_obj =
+      container_of(object, struct raw_data_object, base);
+
+   struct blob_reader blob;
+   blob_reader_init(&blob, data_obj->data, data_obj->data_size);
+
+   nir_shader *nir = nir_deserialize(mem_ctx, nir_options, &blob);
+   vk_pipeline_cache_object_unref(object);
+
+   if (blob.overrun) {
+      ralloc_free(nir);
+      return NULL;
+   }
+
+   return nir;
+}
+
+void
+vk_pipeline_cache_add_nir(struct vk_pipeline_cache *cache,
+                          const void *key_data, size_t key_size,
+                          const nir_shader *nir)
+{
+   struct blob blob;
+   blob_init(&blob);
+
+   nir_serialize(&blob, nir, false);
+   if (blob.out_of_memory) {
+      vk_logw(VK_LOG_OBJS(cache), "Ran out of memory serializing NIR shader");
+      blob_finish(&blob);
+      return;
+   }
+
+   struct raw_data_object *data_obj =
+      raw_data_object_create(cache->base.device,
+                             key_data, key_size,
+                             blob.data, blob.size);
+   blob_finish(&blob);
+
+   struct vk_pipeline_cache_object *cached =
+      vk_pipeline_cache_add_object(cache, &data_obj->base);
+   vk_pipeline_cache_object_unref(cached);
+}
+
 static int32_t
 find_type_for_ops(const struct vk_physical_device *pdevice,
                   const struct vk_pipeline_cache_object_ops *ops)
diff --git a/src/vulkan/runtime/vk_pipeline_cache.h b/src/vulkan/runtime/vk_pipeline_cache.h
index 2340a2c..2d1eeb1 100644
--- a/src/vulkan/runtime/vk_pipeline_cache.h
+++ b/src/vulkan/runtime/vk_pipeline_cache.h
@@ -39,6 +39,10 @@ struct blob_reader;
 /* #include "util/set.h" */
 struct set;
 
+/* #include "compiler/nir/nir.h" */
+struct nir_shader;
+struct nir_shader_compiler_options;
+
 struct vk_pipeline_cache;
 struct vk_pipeline_cache_object;
 
@@ -236,6 +240,16 @@ struct vk_pipeline_cache_object * MUST_CHECK
 vk_pipeline_cache_add_object(struct vk_pipeline_cache *cache,
                              struct vk_pipeline_cache_object *object);
 
+struct nir_shader *
+vk_pipeline_cache_lookup_nir(struct vk_pipeline_cache *cache,
+                             const void *key_data, size_t key_size,
+                             const struct nir_shader_compiler_options *nir_options,
+                             bool *cache_hit, void *mem_ctx);
+void
+vk_pipeline_cache_add_nir(struct vk_pipeline_cache *cache,
+                          const void *key_data, size_t key_size,
+                          const struct nir_shader *nir);
+
 #ifdef __cplusplus
 }
 #endif
-- 
2.7.4