[Mesa-dev] [PATCH 2/2] nir: add shader reference counting

[Rob Clark]

From: Rob Clark <robclark at freedesktop.org>

For gallium, at least, we'll need this to manage shader's lifetimes,
since in some cases both the driver and the state tracker will need
to hold on to a reference for variant managing.

Use nir_shader_mutable() before doing any IR opt/lowering/etc, to
ensure you are not modifying a copy someone else is also holding a
reference to.  In this way, unnecessary nir_shader_clone()s are
avoided whenever possible.

v2: make nir_shader_ref() return itself.. convenient for doing things
like 'nir_shader_mutable(nir_shader_ref(foo))'..

Signed-off-by: Rob Clark <robclark at freedesktop.org>
---
fwiw, patches using refcnt'ing at:
https://github.com/freedreno/mesa/commits/wip-nir-refcnt

 src/gallium/drivers/vc4/vc4_program.c |  2 +-
 src/glsl/nir/nir.c                    |  2 ++
 src/glsl/nir/nir.h                    | 41 +++++++++++++++++++++++++++++++++++
 src/mesa/program/program.c            |  3 ++-
 4 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/gallium/drivers/vc4/vc4_program.c b/src/gallium/drivers/vc4/vc4_program.c
index 52317bd..c9970a0 100644
--- a/src/gallium/drivers/vc4/vc4_program.c
+++ b/src/gallium/drivers/vc4/vc4_program.c
@@ -1741,7 +1741,7 @@ vc4_shader_ntq(struct vc4_context *vc4, enum qstage stage,
                         c->num_uniforms);
         }
 
-        ralloc_free(c->s);
+        nir_shader_unref(c->s);
 
         return c;
 }
diff --git a/src/glsl/nir/nir.c b/src/glsl/nir/nir.c
index 568017a..219f5d9 100644
--- a/src/glsl/nir/nir.c
+++ b/src/glsl/nir/nir.c
@@ -35,6 +35,8 @@ nir_shader_create(gl_shader_stage stage,
 {
    nir_shader *shader = ralloc(NULL, nir_shader);
 
+   p_atomic_set(&shader->refcount, 1);
+
    exec_list_make_empty(&shader->uniforms);
    exec_list_make_empty(&shader->inputs);
    exec_list_make_empty(&shader->outputs);
diff --git a/src/glsl/nir/nir.h b/src/glsl/nir/nir.h
index e378f01..c2a46ea 100644
--- a/src/glsl/nir/nir.h
+++ b/src/glsl/nir/nir.h
@@ -34,6 +34,7 @@
 #include "util/ralloc.h"
 #include "util/set.h"
 #include "util/bitset.h"
+#include "util/u_atomic.h"
 #include "nir_types.h"
 #include "shader_enums.h"
 #include <stdio.h>
@@ -1546,6 +1547,8 @@ typedef struct nir_shader_info {
 } nir_shader_info;
 
 typedef struct nir_shader {
+   int refcount;
+
    /** list of uniforms (nir_variable) */
    struct exec_list uniforms;
 
@@ -1891,6 +1894,44 @@ void nir_print_instr(const nir_instr *instr, FILE *fp);
 
 nir_shader * nir_shader_clone(const nir_shader *s);
 
+static inline nir_shader *
+nir_shader_ref(nir_shader *shader)
+{
+   p_atomic_inc(&shader->refcount);
+   return shader;
+}
+
+static inline void
+nir_shader_unref(nir_shader *shader)
+{
+   if (p_atomic_dec_zero(&shader->refcount)) {
+      ralloc_free(shader);
+   }
+}
+
+/* A shader with only a single reference is mutable: */
+static inline bool
+nir_shader_is_mutable(nir_shader *shader)
+{
+   return p_atomic_read(&shader->refcount) == 1;
+}
+
+/* Convert a shader reference into a mutable shader reference.  Ie. if
+ * there is only a single reference to the shader, then return that,
+ * otherwise clone and drop reference to existing shader.
+ */
+static inline nir_shader *
+nir_shader_mutable(nir_shader *shader)
+{
+   if (nir_shader_is_mutable(shader)) {
+      return shader;
+   } else {
+      nir_shader *ns = nir_shader_clone(shader);
+      nir_shader_unref(shader);
+      return ns;
+   }
+}
+
 #ifdef DEBUG
 void nir_validate_shader(nir_shader *shader);
 void nir_metadata_set_validation_flag(nir_shader *shader);
diff --git a/src/mesa/program/program.c b/src/mesa/program/program.c
index 0e78e6a..c2da66e 100644
--- a/src/mesa/program/program.c
+++ b/src/mesa/program/program.c
@@ -38,6 +38,7 @@
 #include "prog_parameter.h"
 #include "prog_instruction.h"
 #include "util/ralloc.h"
+#include "nir.h"
 
 
 /**
@@ -273,7 +274,7 @@ _mesa_delete_program(struct gl_context *ctx, struct gl_program *prog)
    }
 
    if (prog->nir) {
-      ralloc_free(prog->nir);
+      nir_shader_unref(prog->nir);
    }
 
    mtx_destroy(&prog->Mutex);
-- 
2.5.0