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

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

The first variant is created from nir_shader for uncached shaders.
All other variants are created from serialized NIR.

Reviewed-by: Pierre-Eric Pelloux-Prayer <pierre-eric.pelloux-prayer@amd.com>
Tested-by: Marge Bot <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/2909>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/2909>

src/mesa/state_tracker/st_cb_program.c

@@ -92,6 +92,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->serialized_nir);
+
    /* delete base class */
    _mesa_delete_program( ctx, prog );
 }

src/mesa/state_tracker/st_program.c

@@ -41,6 +41,7 @@
 #include "program/programopt.h"
 
 #include "compiler/nir/nir.h"
+#include "compiler/nir/nir_serialize.h"
 #include "draw/draw_context.h"
 
 #include "pipe/p_context.h"
@@ -502,6 +503,11 @@ st_translate_vertex_program(struct st_context *st,
          if (stp->Base.nir)
             ralloc_free(stp->Base.nir);
 
+         if (stp->serialized_nir) {
+            free(stp->serialized_nir);
+            stp->serialized_nir = NULL;
+         }
+
          stp->state.type = PIPE_SHADER_IR_NIR;
          stp->Base.nir = st_translate_prog_to_nir(st, &stp->Base,
                                                   MESA_SHADER_VERTEX);
@@ -604,6 +610,29 @@ 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->Base.nir) {
+      nir_shader *nir = stp->Base.nir;
+
+      /* The first shader variant takes ownership of NIR, so that there is
+       * no cloning. Additional shader variants are always generated from
+       * serialized NIR to save memory.
+       */
+      stp->Base.nir = NULL;
+      assert(stp->serialized_nir && stp->serialized_nir_size);
+      return 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->serialized_nir, stp->serialized_nir_size);
+   return nir_deserialize(NULL, options, &blob_reader);
+}
+
 static const gl_state_index16 depth_range_state[STATE_LENGTH] =
    { STATE_DEPTH_RANGE };
 
@@ -630,7 +659,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->Base.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;
@@ -827,6 +856,10 @@ st_translate_fragment_program(struct st_context *st,
 
          if (stfp->Base.nir)
             ralloc_free(stfp->Base.nir);
+         if (stfp->serialized_nir) {
+            free(stfp->serialized_nir);
+            stfp->serialized_nir = NULL;
+         }
          stfp->state.type = PIPE_SHADER_IR_NIR;
          stfp->Base.nir = nir;
          return true;
@@ -1184,7 +1217,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->Base.nir);
+      state.ir.nir = get_nir_shader(st, stfp);
 
       if (key->clamp_color) {
          NIR_PASS_V(state.ir.nir, nir_lower_clamp_color_outputs);
@@ -1715,7 +1748,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->Base.nir);
+	    state.ir.nir = get_nir_shader(st, prog);
 
             if (key->clamp_color) {
                NIR_PASS_V(state.ir.nir, nir_lower_clamp_color_outputs);
@@ -1953,6 +1986,20 @@ st_precompile_shader_variant(struct st_context *st,
    }
 }
 
+void
+st_serialize_nir(struct st_program *stp)
+{
+   if (!stp->serialized_nir) {
+      struct blob blob;
+      size_t size;
+
+      blob_init(&blob);
+      nir_serialize(&blob, stp->Base.nir, false);
+      blob_finish_get_buffer(&blob, &stp->serialized_nir, &size);
+      stp->serialized_nir_size = size;
+   }
+}
+
 void
 st_finalize_program(struct st_context *st, struct gl_program *prog)
 {
@@ -1963,9 +2010,16 @@ st_finalize_program(struct st_context *st, struct gl_program *prog)
          st->dirty |= ((struct st_program *)prog)->affected_states;
    }
 
-   if (prog->nir)
+   if (prog->nir) {
       nir_sweep(prog->nir);
 
+      /* This is only needed for ARB_vp/fp programs and when the disk cache
+       * is disabled. If the disk cache is enabled, GLSL programs are
+       * serialized in write_nir_to_cache.
+       */
+      st_serialize_nir(st_program(prog));
+   }
+
    /* Create Gallium shaders now instead of on demand. */
    if (ST_DEBUG & DEBUG_PRECOMPILE ||
        st->shader_has_one_variant[prog->info.stage])

src/mesa/state_tracker/st_program.h

@@ -232,6 +232,9 @@ struct st_program
    struct ati_fragment_shader *ati_fs;
    uint64_t affected_states; /**< ST_NEW_* flags to mark dirty when binding */
 
+   void *serialized_nir;
+   unsigned serialized_nir_size;
+
    /* used when bypassing glsl_to_tgsi: */
    struct gl_shader_program *shader_program;
 
@@ -338,6 +341,9 @@ extern bool
 st_translate_common_program(struct st_context *st,
                             struct st_program *stp);
 
+extern void
+st_serialize_nir(struct st_program *stp);
+
 extern void
 st_finalize_program(struct st_context *st, struct gl_program *prog);
 

src/mesa/state_tracker/st_shader_cache.c

@@ -75,7 +75,13 @@ 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);
+   struct st_program *stp = (struct st_program *)prog;
+
+   st_serialize_nir(stp);
+
+   blob_write_intptr(blob, stp->serialized_nir_size);
+   blob_write_bytes(blob, stp->serialized_nir, stp->serialized_nir_size);
+
    copy_blob_to_driver_cache_blob(blob, prog);
 }
 
@@ -172,8 +178,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);
@@ -203,9 +207,14 @@ 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->serialized_nir == NULL);
+
       stp->state.type = PIPE_SHADER_IR_NIR;
+      stp->serialized_nir_size = blob_read_intptr(&blob_reader);
+      stp->serialized_nir = malloc(stp->serialized_nir_size);
+      blob_copy_bytes(&blob_reader, stp->serialized_nir, stp->serialized_nir_size);
       stp->shader_program = shProg;
-      prog->nir = nir_deserialize(NULL, options, &blob_reader);
    } else {
       read_tgsi_from_cache(&blob_reader, &stp->state.tokens);
    }