nir: Add support for lowering load/stores of shared variables

Signed-off-by: Jordan Justen <[email protected]>
Reviewed-by: Jason Ekstrand <[email protected]>

5 files changed, 32 insertions, 8 deletions

diff --git a/src/compiler/nir/nir.c b/src/compiler/nir/nir.c
index 386cdafe161..655dc884382 100644
--- a/src/compiler/nir/nir.c
+++ b/src/compiler/nir/nir.c
@@ -53,6 +53,7 @@ nir_shader_create(void *mem_ctx,
    shader->num_inputs = 0;
    shader->num_outputs = 0;
    shader->num_uniforms = 0;
+   shader->num_shared = 0;
 
    shader->stage = stage;
 
diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 7d2bd03c914..00f107d4243 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -1692,7 +1692,7 @@ typedef struct nir_shader {
     * the highest index a load_input_*, load_uniform_*, etc. intrinsic can
     * access plus one
     */
-   unsigned num_inputs, num_uniforms, num_outputs;
+   unsigned num_inputs, num_uniforms, num_outputs, num_shared;
 
    /** The shader stage, such as MESA_SHADER_VERTEX. */
    gl_shader_stage stage;
diff --git a/src/compiler/nir/nir_clone.c b/src/compiler/nir/nir_clone.c
index d17d384ed61..0b426e940b4 100644
--- a/src/compiler/nir/nir_clone.c
+++ b/src/compiler/nir/nir_clone.c
@@ -705,6 +705,7 @@ nir_shader_clone(void *mem_ctx, const nir_shader *s)
    ns->num_inputs = s->num_inputs;
    ns->num_uniforms = s->num_uniforms;
    ns->num_outputs = s->num_outputs;
+   ns->num_shared = s->num_shared;
 
    free_clone_state(&state);
 
diff --git a/src/compiler/nir/nir_lower_io.c b/src/compiler/nir/nir_lower_io.c
index 11fb973a237..f844947549d 100644
--- a/src/compiler/nir/nir_lower_io.c
+++ b/src/compiler/nir/nir_lower_io.c
@@ -160,6 +160,29 @@ load_op(struct lower_io_state *state,
    case nir_var_uniform:
       op = nir_intrinsic_load_uniform;
       break;
+   case nir_var_shared:
+      op = nir_intrinsic_load_shared;
+      break;
+   default:
+      unreachable("Unknown variable mode");
+   }
+   return op;
+}
+
+static nir_intrinsic_op
+store_op(struct lower_io_state *state,
+         nir_variable_mode mode, bool per_vertex)
+{
+   nir_intrinsic_op op;
+   switch (mode) {
+   case nir_var_shader_in:
+   case nir_var_shader_out:
+      op = per_vertex ? nir_intrinsic_store_per_vertex_output :
+                        nir_intrinsic_store_output;
+      break;
+   case nir_var_shared:
+      op = nir_intrinsic_store_shared;
+      break;
    default:
       unreachable("Unknown variable mode");
    }
@@ -190,6 +213,7 @@ nir_lower_io_block(nir_block *block, void *void_state)
 
       if (mode != nir_var_shader_in &&
           mode != nir_var_shader_out &&
+          mode != nir_var_shared &&
           mode != nir_var_uniform)
          continue;
 
@@ -236,7 +260,7 @@ nir_lower_io_block(nir_block *block, void *void_state)
       }
 
       case nir_intrinsic_store_var: {
-         assert(mode == nir_var_shader_out);
+         assert(mode == nir_var_shader_out || mode == nir_var_shared);
 
          nir_ssa_def *offset;
          nir_ssa_def *vertex_index;
@@ -248,12 +272,9 @@ nir_lower_io_block(nir_block *block, void *void_state)
                                 per_vertex ? &vertex_index : NULL,
                                 state->type_size);
 
-         nir_intrinsic_op store_op =
-            per_vertex ? nir_intrinsic_store_per_vertex_output :
-                         nir_intrinsic_store_output;
-
-         nir_intrinsic_instr *store = nir_intrinsic_instr_create(state->mem_ctx,
-                                                                 store_op);
+         nir_intrinsic_instr *store =
+            nir_intrinsic_instr_create(state->mem_ctx,
+                                       store_op(state, mode, per_vertex));
          store->num_components = intrin->num_components;
 
          nir_src_copy(&store->src[0], &intrin->src[0], store);
diff --git a/src/compiler/nir/nir_print.c b/src/compiler/nir/nir_print.c
index 644a21463b1..b8943b83f46 100644
--- a/src/compiler/nir/nir_print.c
+++ b/src/compiler/nir/nir_print.c
@@ -1057,6 +1057,7 @@ nir_print_shader(nir_shader *shader, FILE *fp)
    fprintf(fp, "inputs: %u\n", shader->num_inputs);
    fprintf(fp, "outputs: %u\n", shader->num_outputs);
    fprintf(fp, "uniforms: %u\n", shader->num_uniforms);
+   fprintf(fp, "shared: %u\n", shader->num_shared);
 
    nir_foreach_variable(var, &shader->uniforms) {
       print_var_decl(var, &state);