st/mesa: keep serialized NIR instead of nir_shader in st_program

This decreases memory usage, because serialized NIR is more compact.

If shader_has_one_variant is true and the shader is uncached, the first
variant is created from nir_shader, otherwise the first variant and
all other variants are created from serialized NIR.

Acked-by: Pierre-Eric Pelloux-Prayer <[email protected]>

4 files changed, 69 insertions, 11 deletions

diff --git a/src/mesa/state_tracker/st_cb_program.c b/src/mesa/state_tracker/st_cb_program.c
index 3a62212cce7..218297390ce 100644
--- a/src/mesa/state_tracker/st_cb_program.c
+++ b/src/mesa/state_tracker/st_cb_program.c
@@ -97,6 +97,8 @@ st_delete_program(struct gl_context *ctx, struct gl_program *prog)
    if (stp->glsl_to_tgsi)
       free_glsl_to_tgsi_visitor(stp->glsl_to_tgsi);
 
+   free(stp->nir_binary);
+
    /* delete base class */
    _mesa_delete_program( ctx, prog );
 }
diff --git a/src/mesa/state_tracker/st_program.c b/src/mesa/state_tracker/st_program.c
index 12061ed91f1..5b0cc4a6a36 100644
--- a/src/mesa/state_tracker/st_program.c
+++ b/src/mesa/state_tracker/st_program.c
@@ -40,7 +40,7 @@
 #include "program/prog_to_nir.h"
 #include "program/programopt.h"
 
-#include "compiler/nir/nir.h"
+#include "compiler/nir/nir_serialize.h"
 
 #include "pipe/p_context.h"
 #include "pipe/p_defines.h"
@@ -664,6 +664,20 @@ st_translate_vertex_program(struct st_context *st,
    return stp->state.tokens != NULL;
 }
 
+static struct nir_shader *
+get_nir_shader(struct st_context *st, struct st_program *stp)
+{
+   if (stp->state.ir.nir)
+      return nir_shader_clone(NULL, stp->state.ir.nir);
+
+   struct blob_reader blob_reader;
+   const struct nir_shader_compiler_options *options =
+      st->ctx->Const.ShaderCompilerOptions[stp->Base.info.stage].NirOptions;
+
+   blob_reader_init(&blob_reader, stp->nir_binary, stp->nir_size);
+   return nir_deserialize(NULL, options, &blob_reader);
+}
+
 static const gl_state_index16 depth_range_state[STATE_LENGTH] =
    { STATE_DEPTH_RANGE };
 
@@ -690,7 +704,7 @@ st_create_vp_variant(struct st_context *st,
       bool finalize = false;
 
       state.type = PIPE_SHADER_IR_NIR;
-      state.ir.nir = nir_shader_clone(NULL, stvp->state.ir.nir);
+      state.ir.nir = get_nir_shader(st, stvp);
       if (key->clamp_color) {
          NIR_PASS_V(state.ir.nir, nir_lower_clamp_color_outputs);
          finalize = true;
@@ -1246,7 +1260,7 @@ st_create_fp_variant(struct st_context *st,
       bool finalize = false;
 
       state.type = PIPE_SHADER_IR_NIR;
-      state.ir.nir = nir_shader_clone(NULL, stfp->state.ir.nir);
+      state.ir.nir = get_nir_shader(st, stfp);
 
       if (key->clamp_color) {
          NIR_PASS_V(state.ir.nir, nir_lower_clamp_color_outputs);
@@ -1770,7 +1784,7 @@ st_get_common_variant(struct st_context *st,
             bool finalize = false;
 
 	    state.type = PIPE_SHADER_IR_NIR;
-	    state.ir.nir = nir_shader_clone(NULL, prog->state.ir.nir);
+	    state.ir.nir = get_nir_shader(st, prog);
 
             if (key->clamp_color) {
                NIR_PASS_V(state.ir.nir, nir_lower_clamp_color_outputs);
@@ -2057,11 +2071,13 @@ st_precompile_shader_variant(struct st_context *st,
 void
 st_finalize_program(struct st_context *st, struct gl_program *prog)
 {
+   struct st_program *stp = (struct st_program *)prog;
+
    if (st->current_program[prog->info.stage] == prog) {
       if (prog->info.stage == MESA_SHADER_VERTEX)
-         st->dirty |= ST_NEW_VERTEX_PROGRAM(st, (struct st_program *)prog);
+         st->dirty |= ST_NEW_VERTEX_PROGRAM(st, stp);
       else
-         st->dirty |= ((struct st_program *)prog)->affected_states;
+         st->dirty |= stp->affected_states;
    }
 
    if (prog->nir)
@@ -2071,4 +2087,24 @@ st_finalize_program(struct st_context *st, struct gl_program *prog)
    if (ST_DEBUG & DEBUG_PRECOMPILE ||
        st->shader_has_one_variant[prog->info.stage])
       st_precompile_shader_variant(st, prog);
+
+   /* Additional shader variants are always generated from serialized NIR
+    * to save memory.
+    */
+   if (prog->nir) {
+      /* Serialize NIR. */
+      struct blob blob;
+      blob_init(&blob);
+      nir_serialize(&blob, prog->nir, false);
+      stp->nir_binary = malloc(blob.size);
+      memcpy(stp->nir_binary, blob.data, blob.size);
+      stp->nir_size = blob.size;
+      blob_finish(&blob);
+
+      /* Free NIR. */
+      assert(stp->state.ir.nir == prog->nir);
+      ralloc_free(prog->nir);
+      prog->nir = NULL;
+      stp->state.ir.nir = NULL;
+   }
 }
diff --git a/src/mesa/state_tracker/st_program.h b/src/mesa/state_tracker/st_program.h
index 9ad23235bcb..cb085d18bd7 100644
--- a/src/mesa/state_tracker/st_program.h
+++ b/src/mesa/state_tracker/st_program.h
@@ -243,6 +243,9 @@ struct st_program
    struct ati_fragment_shader *ati_fs;
    uint64_t affected_states; /**< ST_NEW_* flags to mark dirty when binding */
 
+   void *nir_binary;
+   unsigned nir_size;
+
   /* used when bypassing glsl_to_tgsi: */
    struct gl_shader_program *shader_program;
 
diff --git a/src/mesa/state_tracker/st_shader_cache.c b/src/mesa/state_tracker/st_shader_cache.c
index b777e8dcbba..474ab4e1a0e 100644
--- a/src/mesa/state_tracker/st_shader_cache.c
+++ b/src/mesa/state_tracker/st_shader_cache.c
@@ -70,7 +70,20 @@ write_tgsi_to_cache(struct blob *blob, const struct tgsi_token *tokens,
 static void
 write_nir_to_cache(struct blob *blob, struct gl_program *prog)
 {
-   nir_serialize(blob, prog->nir, false);
+   if (prog->nir) {
+      /* Reserve intptr_t to store the size. intptr_t is also the alignment
+       * of NIR in the blob, so the NIR size computation will be trivial.
+       */
+      size_t offset = blob_reserve_intptr(blob);
+      nir_serialize(blob, prog->nir, false);
+
+      unsigned nir_size = blob->size - offset - sizeof(intptr_t);
+      *(uintptr_t *)(blob->data + offset) = nir_size;
+   } else {
+      struct st_program *stp = (struct st_program *)prog;
+      blob_write_intptr(blob, stp->nir_size);
+      blob_write_bytes(blob, stp->nir_binary, stp->nir_size);
+   }
    copy_blob_to_driver_cache_blob(blob, prog);
 }
 
@@ -161,8 +174,6 @@ st_deserialise_ir_program(struct gl_context *ctx,
    struct st_context *st = st_context(ctx);
    size_t size = prog->driver_cache_blob_size;
    uint8_t *buffer = (uint8_t *) prog->driver_cache_blob;
-   const struct nir_shader_compiler_options *options =
-      ctx->Const.ShaderCompilerOptions[prog->info.stage].NirOptions;
 
    st_set_prog_affected_state_flags(prog);
    _mesa_associate_uniform_storage(ctx, shProg, prog);
@@ -196,10 +207,16 @@ st_deserialise_ir_program(struct gl_context *ctx,
       read_stream_out_from_cache(&blob_reader, &stp->state);
 
    if (nir) {
+      assert(prog->nir == NULL);
+      assert(stp->state.ir.nir == NULL);
+      assert(stp->nir_binary == NULL);
+
+      /* The remainder of the binary is NIR. */
       stp->state.type = PIPE_SHADER_IR_NIR;
-      stp->state.ir.nir = nir_deserialize(NULL, options, &blob_reader);
+      stp->nir_size = blob_read_intptr(&blob_reader);
+      stp->nir_binary = malloc(stp->nir_size);
+      blob_copy_bytes(&blob_reader, stp->nir_binary, stp->nir_size);
       stp->shader_program = shProg;
-      prog->nir = stp->state.ir.nir;
    } else {
       read_tgsi_from_cache(&blob_reader, &stp->state.tokens);
    }