glsl/linker: implement find_assignment_visitor for NIR shaders

4 files changed, 88 insertions, 28 deletions

diff --git a/src/compiler/glsl/linker.cpp b/src/compiler/glsl/linker.cpp
index 535a122b3f..d233fdd1a2 100644
--- a/src/compiler/glsl/linker.cpp
+++ b/src/compiler/glsl/linker.cpp
@@ -91,15 +91,6 @@
 
 namespace {
 
-struct find_variable {
-   enum ir_variable_mode mode;
-   unsigned location;
-   bool found;
-
-   find_variable(enum ir_variable_mode mode, unsigned location) :
-      mode(mode), location(location), found(false) {}
-};
-
 /**
  * Visitor that determines whether or not a variable is ever written.
  *
@@ -177,25 +168,33 @@ private:
  * written to.
  */
 static void
-find_assignments(exec_list *ir, find_variable * const *vars)
+find_assignments(gl_linked_shader *sh, find_variable * const *vars)
 {
    unsigned num_variables = 0;
 
    for (find_variable * const *v = vars; *v; ++v)
       num_variables++;
 
-   find_assignment_visitor visitor(num_variables, vars);
-   visitor.run(ir);
+   if (sh->nir) {
+      find_nir_assignments(sh->nir, num_variables, vars);
+   } else {
+      find_assignment_visitor visitor(num_variables, vars);
+      visitor.run(sh->ir);
+   }
 }
 
 /**
  * Determine whether or not the given variable is ever written to.
  */
 static void
-find_assignments(exec_list *ir, find_variable *var)
+find_assignments(gl_linked_shader *sh, find_variable *var)
 {
-   find_assignment_visitor visitor(1, &var);
-   visitor.run(ir);
+   if (sh->nir) {
+      find_nir_assignments(sh->nir, 1, &var);
+   } else {
+      find_assignment_visitor visitor(1, &var);
+      visitor.run(sh->ir);
+   }
 }
 
 /**
@@ -621,7 +620,7 @@ analyze_clip_cull_usage(struct gl_shader_program *prog,
          !prog->IsES ? &gl_ClipVertex : NULL,
          NULL
       };
-      find_assignments(shader->ir, variables);
+      find_assignments(shader, variables);
 
       /* From the ARB_cull_distance spec:
        *
@@ -723,7 +722,7 @@ validate_vertex_shader_executable(struct gl_shader_program *prog,
     */
    if (prog->data->Version < (prog->IsES ? 300 : 140)) {
       find_variable gl_Position(ir_var_shader_out, VARYING_SLOT_POS);
-      find_assignments(shader->ir, &gl_Position);
+      find_assignments(shader, &gl_Position);
       if (!gl_Position.found) {
         if (prog->IsES) {
           linker_warning(prog,
@@ -771,7 +770,7 @@ validate_fragment_shader_executable(struct gl_shader_program *prog,
    find_variable gl_FragColor(ir_var_shader_out, FRAG_RESULT_COLOR);
    find_variable gl_FragData(ir_var_shader_out, FRAG_RESULT_DATA0);
    find_variable * const variables[] = { &gl_FragColor, &gl_FragData, NULL };
-   find_assignments(shader->ir, variables);
+   find_assignments(shader, variables);
 
    if (gl_FragColor.found && gl_FragData.found) {
       linker_error(prog,  "fragment shader writes to both "
diff --git a/src/compiler/glsl/linker.h b/src/compiler/glsl/linker.h
index dd627be5f1..6f3e49bb58 100644
--- a/src/compiler/glsl/linker.h
+++ b/src/compiler/glsl/linker.h
@@ -212,4 +212,13 @@ struct empty_uniform_block {
    unsigned slots;
 };
 
+struct find_variable {
+   enum ir_variable_mode mode;
+   int location;
+   bool found;
+
+   find_variable(enum ir_variable_mode mode, int location) :
+      mode(mode), location(location), found(false) {}
+};
+
 #endif /* GLSL_LINKER_H */
diff --git a/src/compiler/glsl/nir_linker.cpp b/src/compiler/glsl/nir_linker.cpp
index c9f248aa5f..27c3864a31 100644
--- a/src/compiler/glsl/nir_linker.cpp
+++ b/src/compiler/glsl/nir_linker.cpp
@@ -44,6 +44,59 @@ validate_error(gl_shader_program *prog, const char *filename, unsigned line,
                 msg, filename, line, cond);
 }
 
+static ir_variable_mode
+ir_from_nir_variable_mode(nir_variable_mode mode)
+{
+   switch (mode) {
+   case nir_var_shader_in: return ir_var_shader_in;
+   case nir_var_shader_out: return ir_var_shader_out;
+   case nir_var_uniform: return ir_var_uniform;
+   case nir_var_shader_storage: return ir_var_shader_storage;
+   case nir_var_system_value: return ir_var_system_value;
+   case nir_var_shared: return ir_var_shader_shared;
+   default:
+      return ir_var_auto;
+   }
+}
+
+/**
+ * Return a bitmask indicating which variables of the given names are written
+ * to.
+ */
+void
+find_nir_assignments(struct nir_shader *nir, unsigned num_variables,
+                     struct find_variable * const *variables)
+{
+   uint32_t num_found = 0;
+
+   nir_foreach_function(func, nir) {
+      nir_foreach_block(block, func->impl) {
+         nir_foreach_instr(instr, block) {
+            if (instr->type != nir_instr_type_intrinsic)
+               continue;
+
+            nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+            if (intr->intrinsic != nir_intrinsic_store_var)
+               continue;
+
+            nir_variable *var = intr->variables[0]->var;
+
+            for (unsigned i = 0; i < num_variables; ++i) {
+               if (variables[i]->mode == ir_from_nir_variable_mode(var->data.mode) &&
+                   variables[i]->location == var->data.location) {
+                  if (!variables[i]->found) {
+                     variables[i]->found = true;
+                     if (++num_found >= num_variables)
+                        return;
+                  }
+                  break;
+               }
+            }
+         }
+      }
+   }
+}
+
 static ir_depth_layout
 ir_from_nir_depth_layout(nir_depth_layout depth_layout)
 {
@@ -63,16 +116,7 @@ static void
 nir_to_ir_variable(void *mem_ctx, gl_shader_program *prog,
                    gl_linked_shader *linked, nir_variable *nir_var)
 {
-   ir_variable_mode mode;
-
-   switch (nir_var->data.mode) {
-   case nir_var_shader_in: mode = ir_var_shader_in; break;
-   case nir_var_shader_out: mode = ir_var_shader_out; break;
-   case nir_var_uniform: mode = ir_var_uniform; break;
-   default:
-      assert(!"not implemented");
-   }
-
+   ir_variable_mode mode = ir_from_nir_variable_mode(nir_var->data.mode);
    const char *name = nir_var->name;
 
    validate(prog,
diff --git a/src/compiler/glsl/nir_linker.h b/src/compiler/glsl/nir_linker.h
index 6396ffba78..c96d0ddbf6 100644
--- a/src/compiler/glsl/nir_linker.h
+++ b/src/compiler/glsl/nir_linker.h
@@ -24,10 +24,14 @@
 #ifndef NIR_LINKER_H
 #define NIR_LINKER_H
 
+#include <stdint.h>
+
+struct find_variable;
 struct gl_context;
 struct gl_linked_shader;
 struct gl_shader;
 struct gl_shader_program;
+struct nir_shader;
 
 struct gl_linked_shader *
 create_linked_nir_shader(void *mem_ctx,
@@ -35,4 +39,8 @@ create_linked_nir_shader(void *mem_ctx,
                          struct gl_shader_program *prog,
                          struct gl_shader *shader);
 
+void
+find_nir_assignments(struct nir_shader *nir, unsigned num_variables,
+                     struct find_variable * const *variables);
+
 #endif /* NIR_LINKER_H */