zink: split up creating zink_shader objects and VkShaderModule objects

the actual VkShaderModule is only needed when we're creating a program to
draw with, so this can be split off for "uncompiled" and "compiled" shader
objects which will facilitate implementing shader keys

Reviewed-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5970>

6 files changed, 73 insertions, 47 deletions

diff --git a/src/gallium/drivers/zink/zink_compiler.c b/src/gallium/drivers/zink/zink_compiler.c
index db433036f58..55451ad2443 100644
--- a/src/gallium/drivers/zink/zink_compiler.c
+++ b/src/gallium/drivers/zink/zink_compiler.c
@@ -214,9 +214,43 @@ update_so_info(struct zink_shader *sh,
    }
 }
 
+VkShaderModule
+zink_shader_compile(struct zink_screen *screen, struct zink_shader *zs)
+{
+   VkShaderModule mod = NULL;
+   void *streamout = zs->streamout.so_info_slots ? &zs->streamout : NULL;
+   struct spirv_shader *spirv = nir_to_spirv(zs->nir, streamout);
+   assert(spirv);
+
+   if (zink_debug & ZINK_DEBUG_SPIRV) {
+      char buf[256];
+      static int i;
+      snprintf(buf, sizeof(buf), "dump%02d.spv", i++);
+      FILE *fp = fopen(buf, "wb");
+      if (fp) {
+         fwrite(spirv->words, sizeof(uint32_t), spirv->num_words, fp);
+         fclose(fp);
+         fprintf(stderr, "wrote '%s'...\n", buf);
+      }
+   }
+
+   VkShaderModuleCreateInfo smci = {};
+   smci.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
+   smci.codeSize = spirv->num_words * sizeof(uint32_t);
+   smci.pCode = spirv->words;
+
+   if (vkCreateShaderModule(screen->dev, &smci, NULL, &mod) != VK_SUCCESS)
+      mod = NULL;
+
+   /* TODO: determine if there's any reason to cache spirv output? */
+   free(spirv->words);
+   free(spirv);
+   return mod;
+}
+
 struct zink_shader *
-zink_compile_nir(struct zink_screen *screen, struct nir_shader *nir,
-                 const struct pipe_stream_output_info *so_info)
+zink_shader_create(struct zink_screen *screen, struct nir_shader *nir,
+                   const struct pipe_stream_output_info *so_info)
 {
    struct zink_shader *ret = CALLOC_STRUCT(zink_shader);
    bool have_psiz = false;
@@ -276,7 +310,7 @@ zink_compile_nir(struct zink_screen *screen, struct nir_shader *nir,
       }
    }
 
-   ret->info = nir->info;
+   ret->nir = nir;
    if (so_info) {
       memcpy(&ret->streamout.so_info, so_info, sizeof(struct pipe_stream_output_info));
       ret->streamout.so_info_slots = malloc(so_info->num_outputs * sizeof(unsigned int));
@@ -284,32 +318,6 @@ zink_compile_nir(struct zink_screen *screen, struct nir_shader *nir,
       update_so_info(ret, nir->info.outputs_written, have_psiz);
    }
 
-   struct spirv_shader *spirv = nir_to_spirv(nir, so_info ? &ret->streamout : NULL);
-   assert(spirv);
-
-   if (zink_debug & ZINK_DEBUG_SPIRV) {
-      char buf[256];
-      static int i;
-      snprintf(buf, sizeof(buf), "dump%02d.spv", i++);
-      FILE *fp = fopen(buf, "wb");
-      if (fp) {
-         fwrite(spirv->words, sizeof(uint32_t), spirv->num_words, fp);
-         fclose(fp);
-         fprintf(stderr, "wrote '%s'...\n", buf);
-      }
-   }
-
-   VkShaderModuleCreateInfo smci = {};
-   smci.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
-   smci.codeSize = spirv->num_words * sizeof(uint32_t);
-   smci.pCode = spirv->words;
-
-   if (vkCreateShaderModule(screen->dev, &smci, NULL, &ret->shader_module) != VK_SUCCESS)
-      return NULL;
-
-   free(spirv->words);
-   free(spirv);
-
    return ret;
 }
 
@@ -317,14 +325,14 @@ void
 zink_shader_free(struct zink_context *ctx, struct zink_shader *shader)
 {
    struct zink_screen *screen = zink_screen(ctx->base.screen);
-   vkDestroyShaderModule(screen->dev, shader->shader_module, NULL);
    set_foreach(shader->programs, entry) {
       struct zink_gfx_program *prog = (void*)entry->key;
-      _mesa_hash_table_remove_key(ctx->program_cache, prog->stages);
-      prog->stages[pipe_shader_type_from_mesa(shader->info.stage)] = NULL;
+      _mesa_hash_table_remove_key(ctx->program_cache, prog->shaders);
+      prog->shaders[pipe_shader_type_from_mesa(shader->nir->info.stage)] = NULL;
       zink_gfx_program_reference(screen, &prog, NULL);
    }
-   _mesa_set_destroy(shader->programs, NULL
+   _mesa_set_destroy(shader->programs, NULL);
    free(shader->streamout.so_info_slots);
+   ralloc_free(shader->nir);
    FREE(shader);
 }
diff --git a/src/gallium/drivers/zink/zink_compiler.h b/src/gallium/drivers/zink/zink_compiler.h
index 50779d7c954..810be0163bd 100644
--- a/src/gallium/drivers/zink/zink_compiler.h
+++ b/src/gallium/drivers/zink/zink_compiler.h
@@ -27,6 +27,7 @@
 #include "pipe/p_defines.h"
 #include "pipe/p_state.h"
 
+#include "compiler/nir/nir.h"
 #include "compiler/shader_info.h"
 
 #include <vulkan/vulkan.h>
@@ -57,9 +58,8 @@ struct nir_shader *
 zink_tgsi_to_nir(struct pipe_screen *screen, const struct tgsi_token *tokens);
 
 struct zink_shader {
-   VkShaderModule shader_module;
+   struct nir_shader *nir;
 
-   shader_info info;
    struct zink_so_info streamout;
 
    struct {
@@ -71,8 +71,11 @@ struct zink_shader {
    struct set *programs;
 };
 
+VkShaderModule
+zink_shader_compile(struct zink_screen *screen, struct zink_shader *zs);
+
 struct zink_shader *
-zink_compile_nir(struct zink_screen *screen, struct nir_shader *nir,
+zink_shader_create(struct zink_screen *screen, struct nir_shader *nir,
                  const struct pipe_stream_output_info *so_info);
 
 void
diff --git a/src/gallium/drivers/zink/zink_draw.c b/src/gallium/drivers/zink/zink_draw.c
index 1a351ee518f..79e61ccb52b 100644
--- a/src/gallium/drivers/zink/zink_draw.c
+++ b/src/gallium/drivers/zink/zink_draw.c
@@ -166,7 +166,7 @@ get_gfx_program(struct zink_context *ctx)
       if (!entry) {
          struct zink_gfx_program *prog;
          prog = zink_create_gfx_program(ctx, ctx->gfx_stages);
-         entry = _mesa_hash_table_insert(ctx->program_cache, prog->stages, prog);
+         entry = _mesa_hash_table_insert(ctx->program_cache, prog->shaders, prog);
          if (!entry)
             return NULL;
       }
diff --git a/src/gallium/drivers/zink/zink_pipeline.c b/src/gallium/drivers/zink/zink_pipeline.c
index 70f859fdd3f..1d45b825830 100644
--- a/src/gallium/drivers/zink/zink_pipeline.c
+++ b/src/gallium/drivers/zink/zink_pipeline.c
@@ -149,7 +149,7 @@ zink_create_gfx_pipeline(struct zink_screen *screen,
       VkPipelineShaderStageCreateInfo stage = {};
       stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
       stage.stage = zink_shader_stage(i);
-      stage.module = prog->stages[i]->shader_module;
+      stage.module = prog->stages[i];
       stage.pName = "main";
       shader_stages[num_stages++] = stage;
    }
diff --git a/src/gallium/drivers/zink/zink_program.c b/src/gallium/drivers/zink/zink_program.c
index c6b0c1cd350..830cf5e9229 100644
--- a/src/gallium/drivers/zink/zink_program.c
+++ b/src/gallium/drivers/zink/zink_program.c
@@ -107,6 +107,17 @@ create_pipeline_layout(VkDevice dev, VkDescriptorSetLayout dsl)
    return layout;
 }
 
+static void
+update_shader_modules(struct zink_context *ctx, struct zink_shader *stages[PIPE_SHADER_TYPES - 1], struct zink_gfx_program *prog)
+{
+   for (int i = 0; i < PIPE_SHADER_TYPES - 1; ++i) {
+      if (stages[i]) {
+         prog->stages[i] = zink_shader_compile(zink_screen(ctx->base.screen), stages[i]);
+         prog->shaders[i] = stages[i];
+      }
+   }
+}
+
 static uint32_t
 hash_gfx_pipeline_state(const void *key)
 {
@@ -130,6 +141,8 @@ zink_create_gfx_program(struct zink_context *ctx,
 
    pipe_reference_init(&prog->reference, 1);
 
+   update_shader_modules(ctx, stages, prog);
+
    for (int i = 0; i < ARRAY_SIZE(prog->pipelines); ++i) {
       prog->pipelines[i] = _mesa_hash_table_create(NULL,
                                                    hash_gfx_pipeline_state,
@@ -139,8 +152,7 @@ zink_create_gfx_program(struct zink_context *ctx,
    }
 
    for (int i = 0; i < PIPE_SHADER_TYPES - 1; ++i) {
-      prog->stages[i] = stages[i];
-      if (stages[i]) {
+      if (prog->stages[i]) {
          _mesa_set_add(stages[i]->programs, prog);
          zink_gfx_program_reference(screen, NULL, prog);
       }
@@ -171,10 +183,10 @@ fail:
 static void
 gfx_program_remove_shader(struct zink_gfx_program *prog, struct zink_shader *shader)
 {
-   enum pipe_shader_type p_stage = pipe_shader_type_from_mesa(shader->info.stage);
+   enum pipe_shader_type p_stage = pipe_shader_type_from_mesa(shader->nir->info.stage);
 
-   assert(prog->stages[p_stage] == shader);
-   prog->stages[p_stage] = NULL;
+   assert(prog->shaders[p_stage] == shader);
+   prog->shaders[p_stage] = NULL;
    _mesa_set_remove_key(shader->programs, prog);
 }
 
@@ -189,8 +201,10 @@ zink_destroy_gfx_program(struct zink_screen *screen,
       vkDestroyDescriptorSetLayout(screen->dev, prog->dsl, NULL);
 
    for (int i = 0; i < PIPE_SHADER_TYPES - 1; ++i) {
+      if (prog->shaders[i])
+         gfx_program_remove_shader(prog, prog->shaders[i]);
       if (prog->stages[i])
-         gfx_program_remove_shader(prog, prog->stages[i]);
+         vkDestroyShaderModule(screen->dev, prog->stages[i], NULL);
    }
 
    /* unref all used render-passes */
@@ -300,7 +314,7 @@ zink_create_vs_state(struct pipe_context *pctx,
    else
       nir = (struct nir_shader *)shader->ir.nir;
 
-   return zink_compile_nir(zink_screen(pctx->screen), nir, &shader->stream_output);
+   return zink_shader_create(zink_screen(pctx->screen), nir, &shader->stream_output);
 }
 
 static void
@@ -336,7 +350,7 @@ zink_create_fs_state(struct pipe_context *pctx,
    else
       nir = (struct nir_shader *)shader->ir.nir;
 
-   return zink_compile_nir(zink_screen(pctx->screen), nir, NULL);
+   return zink_shader_create(zink_screen(pctx->screen), nir, NULL);
 }
 
 static void
diff --git a/src/gallium/drivers/zink/zink_program.h b/src/gallium/drivers/zink/zink_program.h
index 60fc778476f..bd6f7dbaa7a 100644
--- a/src/gallium/drivers/zink/zink_program.h
+++ b/src/gallium/drivers/zink/zink_program.h
@@ -40,7 +40,8 @@ struct set;
 struct zink_gfx_program {
    struct pipe_reference reference;
 
-   struct zink_shader *stages[PIPE_SHADER_TYPES - 1]; // compute stage doesn't belong here
+   VkShaderModule stages[PIPE_SHADER_TYPES - 1]; // compute stage doesn't belong here
+   struct zink_shader *shaders[PIPE_SHADER_TYPES - 1];
    VkDescriptorSetLayout dsl;
    VkPipelineLayout layout;
    unsigned num_descriptors;