asahi: wire up shader disk cache support

Note: I (Alyssa) have squashed in some minor changes squashed in pre merge. The
rest is Rose's work :-)

Closes: #8091
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/20835>

src/gallium/drivers/asahi/agx_disk_cache.c

@@ -0,0 +1,168 @@
+/*
+ * Copyright © 2023 Rose Hudson
+ * Copyright (c) 2022 Amazon.com, Inc. or its affiliates.
+ * Copyright © 2018 Intel Corporation
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included
+ * in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+ * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include <assert.h>
+#include <stdio.h>
+
+#include "compiler/shader_enums.h"
+#include "util/blob.h"
+#include "util/build_id.h"
+#include "util/disk_cache.h"
+#include "util/mesa-sha1.h"
+#include "agx_bo.h"
+#include "agx_disk_cache.h"
+#include "agx_state.h"
+
+/**
+ * Compute a disk cache key for the given uncompiled shader and shader key.
+ */
+static void
+agx_disk_cache_compute_key(struct disk_cache *cache,
+                           const struct agx_uncompiled_shader *uncompiled,
+                           const union asahi_shader_key *shader_key,
+                           gl_shader_stage stage, cache_key cache_key)
+{
+   uint8_t data[sizeof(uncompiled->nir_sha1) + sizeof(*shader_key)];
+   int hash_size = sizeof(uncompiled->nir_sha1);
+   int key_size;
+   if (stage == MESA_SHADER_VERTEX)
+      key_size = sizeof(shader_key->vs);
+   else if (stage == MESA_SHADER_FRAGMENT)
+      key_size = sizeof(shader_key->fs);
+   else if (gl_shader_stage_is_compute(stage))
+      key_size = 0;
+   else
+      unreachable("Unsupported shader stage");
+
+   memcpy(data, uncompiled->nir_sha1, hash_size);
+
+   if (key_size)
+      memcpy(data + hash_size, shader_key, key_size);
+
+   disk_cache_compute_key(cache, data, hash_size + key_size, cache_key);
+}
+
+/**
+ * Store the given compiled shader in the disk cache.
+ *
+ * This should only be called on newly compiled shaders.  No checking is
+ * done to prevent repeated stores of the same shader.
+ */
+void
+agx_disk_cache_store(struct disk_cache *cache,
+                     const struct agx_uncompiled_shader *uncompiled,
+                     const union asahi_shader_key *key,
+                     const struct agx_compiled_shader *binary)
+{
+#ifdef ENABLE_SHADER_CACHE
+   if (!cache)
+      return;
+
+   assert(binary->bo->ptr.cpu != NULL && "shaders must be CPU mapped");
+
+   gl_shader_stage stage = uncompiled->nir->info.stage;
+   cache_key cache_key;
+   agx_disk_cache_compute_key(cache, uncompiled, key, stage, cache_key);
+
+   struct blob blob;
+   blob_init(&blob);
+
+   uint32_t shader_size = binary->bo->size;
+   blob_write_uint32(&blob, shader_size);
+   blob_write_bytes(&blob, binary->bo->ptr.cpu, shader_size);
+   blob_write_bytes(&blob, &binary->info, sizeof(binary->info));
+   blob_write_uint32(&blob, binary->push_range_count);
+   blob_write_bytes(&blob, binary->push, sizeof(binary->push));
+
+   disk_cache_put(cache, cache_key, blob.data, blob.size, NULL);
+   blob_finish(&blob);
+#endif
+}
+
+/**
+ * Search for a compiled shader in the disk cache.
+ */
+struct agx_compiled_shader *
+agx_disk_cache_retrieve(struct agx_screen *screen,
+                        const struct agx_uncompiled_shader *uncompiled,
+                        const union asahi_shader_key *key)
+{
+#ifdef ENABLE_SHADER_CACHE
+   struct disk_cache *cache = screen->disk_cache;
+   if (!cache)
+      return NULL;
+
+   gl_shader_stage stage = uncompiled->nir->info.stage;
+   cache_key cache_key;
+   agx_disk_cache_compute_key(cache, uncompiled, key, stage, cache_key);
+
+   size_t size;
+   void *buffer = disk_cache_get(cache, cache_key, &size);
+   if (!buffer)
+      return NULL;
+
+   struct agx_compiled_shader *binary = CALLOC_STRUCT(agx_compiled_shader);
+
+   struct blob_reader blob;
+   blob_reader_init(&blob, buffer, size);
+
+   uint32_t binary_size = blob_read_uint32(&blob);
+   binary->bo = agx_bo_create(&screen->dev, binary_size,
+                              AGX_BO_EXEC | AGX_BO_LOW_VA, "Executable");
+   blob_copy_bytes(&blob, binary->bo->ptr.cpu, binary_size);
+
+   blob_copy_bytes(&blob, &binary->info, sizeof(binary->info));
+   binary->push_range_count = blob_read_uint32(&blob);
+   blob_copy_bytes(&blob, binary->push, sizeof(binary->push));
+
+   free(buffer);
+   return binary;
+#else
+   return NULL;
+#endif
+}
+
+/**
+ * Initialise the on-disk shader cache.
+ */
+void
+agx_disk_cache_init(struct agx_screen *screen)
+{
+#ifdef ENABLE_SHADER_CACHE
+   const char *renderer = screen->pscreen.get_name(&screen->pscreen);
+
+   const struct build_id_note *note =
+      build_id_find_nhdr_for_addr(agx_disk_cache_init);
+   assert(note && build_id_length(note) == 20);
+
+   const uint8_t *id_sha1 = build_id_data(note);
+   assert(id_sha1);
+
+   char timestamp[41];
+   _mesa_sha1_format(timestamp, id_sha1);
+
+   uint64_t driver_flags = screen->dev.debug;
+   screen->disk_cache = disk_cache_create(renderer, timestamp, driver_flags);
+#endif
+}

src/gallium/drivers/asahi/agx_disk_cache.h

@@ -0,0 +1,44 @@
+/*
+ * © Copyright 2023 Rose Hudson
+ * © Copyright 2018 Alyssa Rosenzweig
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice (including the next
+ * paragraph) shall be included in all copies or substantial portions of the
+ * Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ */
+
+#include "util/disk_cache.h"
+#include "agx_state.h"
+
+#ifndef AGX_DISK_CACHE_H
+#define AGX_DISK_CACHE_H
+
+void agx_disk_cache_store(struct disk_cache *cache,
+                          const struct agx_uncompiled_shader *uncompiled,
+                          const union asahi_shader_key *key,
+                          const struct agx_compiled_shader *binary);
+
+struct agx_compiled_shader *
+agx_disk_cache_retrieve(struct agx_screen *screen,
+                        const struct agx_uncompiled_shader *uncompiled,
+                        const union asahi_shader_key *key);
+
+void agx_disk_cache_init(struct agx_screen *screen);
+
+#endif

src/gallium/drivers/asahi/agx_pipe.c

@@ -49,6 +49,7 @@
 #include "util/u_screen.h"
 #include "util/u_upload_mgr.h"
 #include "agx_device.h"
+#include "agx_disk_cache.h"
 #include "agx_public.h"
 #include "agx_state.h"
 #include "magic.h"
@@ -1600,10 +1601,13 @@ agx_is_dmabuf_modifier_supported(struct pipe_screen *screen, uint64_t modifier,
 }
 
 static void
-agx_destroy_screen(struct pipe_screen *screen)
+agx_destroy_screen(struct pipe_screen *pscreen)
 {
-   u_transfer_helper_destroy(screen->transfer_helper);
-   agx_close_device(agx_device(screen));
+   struct agx_screen *screen = agx_screen(pscreen);
+
+   u_transfer_helper_destroy(pscreen->transfer_helper);
+   agx_close_device(&screen->dev);
+   disk_cache_destroy(screen->disk_cache);
    ralloc_free(screen);
 }
 
@@ -1646,6 +1650,12 @@ agx_resource_get_internal_format(struct pipe_resource *prsrc)
    return agx_resource(prsrc)->layout.format;
 }
 
+static struct disk_cache *
+agx_get_disk_shader_cache(struct pipe_screen *pscreen)
+{
+   return agx_screen(pscreen)->disk_cache;
+}
+
 static const struct u_transfer_vtbl transfer_vtbl = {
    .resource_create = agx_resource_create,
    .resource_destroy = agx_resource_destroy,
@@ -1717,6 +1727,7 @@ agx_screen_create(int fd, struct renderonly *ro, struct sw_winsys *winsys)
    screen->fence_reference = agx_fence_reference;
    screen->fence_finish = agx_fence_finish;
    screen->get_compiler_options = agx_get_compiler_options;
+   screen->get_disk_shader_cache = agx_get_disk_shader_cache;
 
    screen->resource_create = u_transfer_helper_resource_create;
    screen->resource_destroy = u_transfer_helper_resource_destroy;
@@ -1725,5 +1736,7 @@ agx_screen_create(int fd, struct renderonly *ro, struct sw_winsys *winsys)
       U_TRANSFER_HELPER_SEPARATE_Z32S8 | U_TRANSFER_HELPER_SEPARATE_STENCIL |
          U_TRANSFER_HELPER_MSAA_MAP | U_TRANSFER_HELPER_Z24_IN_Z32F);
 
+   agx_disk_cache_init(agx_screen);
+
    return screen;
 }

src/gallium/drivers/asahi/agx_state.c

@@ -31,6 +31,7 @@
 #include "asahi/lib/agx_ppp.h"
 #include "asahi/lib/agx_usc.h"
 #include "compiler/nir/nir.h"
+#include "compiler/nir/nir_serialize.h"
 #include "gallium/auxiliary/nir/tgsi_to_nir.h"
 #include "gallium/auxiliary/tgsi/tgsi_from_mesa.h"
 #include "gallium/auxiliary/util/u_blend.h"
@@ -47,6 +48,7 @@
 #include "util/u_prim.h"
 #include "util/u_transfer.h"
 #include "agx_state.h"
+#include "agx_disk_cache.h"
 
 static struct pipe_stream_output_target *
 agx_create_stream_output_target(struct pipe_context *pctx,
@@ -1307,16 +1309,32 @@ agx_compile_variant(struct agx_device *dev, struct agx_uncompiled_shader *so,
    ralloc_free(nir);
    util_dynarray_fini(&binary);
 
+   return compiled;
+}
+
+static struct agx_compiled_shader *
+agx_get_shader_variant(struct agx_screen *screen,
+                       struct agx_uncompiled_shader *so,
+                       struct util_debug_callback *debug,
+                       union asahi_shader_key *key)
+{
+   struct agx_compiled_shader *compiled =
+      agx_disk_cache_retrieve(screen, so, key);
+
+   if (!compiled) {
+      compiled = agx_compile_variant(&screen->dev, so, debug, key);
+      agx_disk_cache_store(screen->disk_cache, so, key, compiled);
+   }
+
    /* key may be destroyed after we return, so clone it before using it as a
     * hash table key. The clone is logically owned by the hash table.
     */
    union asahi_shader_key *cloned_key =
       ralloc(so->variants, union asahi_shader_key);
-   memcpy(cloned_key, key_, sizeof(union asahi_shader_key));
+   memcpy(cloned_key, key, sizeof(union asahi_shader_key));
+   _mesa_hash_table_insert(so->variants, cloned_key, compiled);
 
-   struct hash_entry *he =
-      _mesa_hash_table_insert(so->variants, cloned_key, compiled);
-   return he->data;
+   return compiled;
 }
 
 static void *
@@ -1346,6 +1364,12 @@ agx_create_shader_state(struct pipe_context *pctx,
                                              asahi_fs_shader_key_equal);
    }
 
+   struct blob blob;
+   blob_init(&blob);
+   nir_serialize(&blob, so->nir, true);
+   _mesa_sha1_compute(blob.data, blob.size, so->nir_sha1);
+   blob_finish(&blob);
+
    /* For shader-db, precompile a shader with a default key. This could be
     * improved but hopefully this is acceptable for now.
     */
@@ -1397,8 +1421,8 @@ agx_update_shader(struct agx_context *ctx, struct agx_compiled_shader **out,
       return true;
    }
 
-   struct agx_device *dev = agx_device(ctx->base.screen);
-   *out = agx_compile_variant(dev, so, &ctx->base.debug, key);
+   struct agx_screen *screen = agx_screen(ctx->base.screen);
+   *out = agx_get_shader_variant(screen, so, &ctx->base.debug, key);
    return true;
 }
 

src/gallium/drivers/asahi/agx_state.h

@@ -39,6 +39,7 @@
 #include "gallium/include/pipe/p_screen.h"
 #include "gallium/include/pipe/p_state.h"
 #include "util/bitset.h"
+#include "util/disk_cache.h"
 #include "util/hash_table.h"
 #include "agx_meta.h"
 
@@ -116,6 +117,7 @@ struct agx_compiled_shader {
 struct agx_uncompiled_shader {
    struct pipe_shader_state base;
    struct nir_shader *nir;
+   uint8_t nir_sha1[20];
    struct hash_table *variants;
 
    /* Set on VS, passed to FS for linkage */
@@ -355,6 +357,7 @@ struct agx_screen {
    struct pipe_screen pscreen;
    struct agx_device dev;
    struct sw_winsys *winsys;
+   struct disk_cache *disk_cache;
 };
 
 static inline struct agx_screen *

src/gallium/drivers/asahi/meson.build

@@ -21,6 +21,7 @@
 files_asahi = files(
   'agx_batch.c',
   'agx_blit.c',
+  'agx_disk_cache.c',
   'agx_pipe.c',
   'agx_nir_lower_sysvals.c',
   'agx_query.c',