2 files changed, 24 insertions, 6 deletions

diff --git a/src/mesa/state_tracker/st_glsl_to_nir.cpp b/src/mesa/state_tracker/st_glsl_to_nir.cpp
index 7f4651a3ccd..1912da2404a 100644
--- a/src/mesa/state_tracker/st_glsl_to_nir.cpp
+++ b/src/mesa/state_tracker/st_glsl_to_nir.cpp
@@ -441,34 +441,40 @@ st_nir_get_mesa_program(struct gl_context *ctx,
    struct st_fragment_program *stfp;
    struct st_compute_program *stcp;
 
+   nir_shader *nir = st_glsl_to_nir(st, prog, shader_program, shader->Stage);
+
    switch (shader->Stage) {
    case MESA_SHADER_VERTEX:
       stvp = (struct st_vertex_program *)prog;
       stvp->shader_program = shader_program;
+      stvp->tgsi.type = PIPE_SHADER_IR_NIR;
+      stvp->tgsi.ir.nir = nir;
       break;
    case MESA_SHADER_GEOMETRY:
    case MESA_SHADER_TESS_CTRL:
    case MESA_SHADER_TESS_EVAL:
       stp = (struct st_common_program *)prog;
       stp->shader_program = shader_program;
+      stp->tgsi.type = PIPE_SHADER_IR_NIR;
+      stp->tgsi.ir.nir = nir;
       break;
    case MESA_SHADER_FRAGMENT:
       stfp = (struct st_fragment_program *)prog;
       stfp->shader_program = shader_program;
+      stfp->tgsi.type = PIPE_SHADER_IR_NIR;
+      stfp->tgsi.ir.nir = nir;
       break;
    case MESA_SHADER_COMPUTE:
       stcp = (struct st_compute_program *)prog;
       stcp->shader_program = shader_program;
+      stcp->tgsi.ir_type = PIPE_SHADER_IR_NIR;
+      stcp->tgsi.prog = nir_shader_clone(NULL, nir);
       break;
    default:
       assert(!"should not be reached");
       return NULL;
    }
 
-   struct st_common_program *st_comm_prog = (struct st_common_program *)prog;
-   nir_shader *nir = st_glsl_to_nir(st, prog, shader_program, shader->Stage);
-   st_comm_prog->tgsi.type = PIPE_SHADER_IR_NIR;
-   st_comm_prog->tgsi.ir.nir = nir;
 
    return prog;
 }
diff --git a/src/mesa/state_tracker/st_program.c b/src/mesa/state_tracker/st_program.c
index e3649a8b7cc..70b63a37471 100644
--- a/src/mesa/state_tracker/st_program.c
+++ b/src/mesa/state_tracker/st_program.c
@@ -360,8 +360,20 @@ st_release_cp_variants(struct st_context *st, struct st_compute_program *stcp)
    *variants = NULL;
 
    if (stcp->tgsi.prog) {
-      ureg_free_tokens(stcp->tgsi.prog);
-      stcp->tgsi.prog = NULL;
+      switch (stcp->tgsi.ir_type) {
+      case PIPE_SHADER_IR_TGSI:
+         ureg_free_tokens(stcp->tgsi.prog);
+         stcp->tgsi.prog = NULL;
+         break;
+      case PIPE_SHADER_IR_NIR:
+         /* pipe driver took ownership of prog */
+         break;
+      case PIPE_SHADER_IR_LLVM:
+      case PIPE_SHADER_IR_NATIVE:
+         /* ??? */
+         stcp->tgsi.prog = NULL;
+         break;
+      }
    }
 }