nir: Split nir_lower_io's input/output/atomic handling into helpers.

The original function was becoming a bit hard to read, with the details
of creating and filling out load/store/atomic atomics all in one
function.

This patch makes helpers for creating each type of intrinsic, and also
combines them with the *_op() helpers, as they're closely coupled and
not too large.

v2: Minor style nits from Jason.

Signed-off-by: Kenneth Graunke <[email protected]>
Reviewed-by: Jason Ekstrand <[email protected]>

1 files changed, 87 insertions, 91 deletions

diff --git a/src/compiler/nir/nir_lower_io.c b/src/compiler/nir/nir_lower_io.c
index 914e0e1e0db..1708e5b77e6 100644
--- a/src/compiler/nir/nir_lower_io.c
+++ b/src/compiler/nir/nir_lower_io.c
@@ -167,18 +167,22 @@ get_io_offset(nir_builder *b, nir_deref_var *deref,
    return offset;
 }
 
-static nir_intrinsic_op
-load_op(nir_variable_mode mode, bool per_vertex)
+static nir_intrinsic_instr *
+lower_load(nir_intrinsic_instr *intrin, struct lower_io_state *state,
+           nir_ssa_def *vertex_index, nir_ssa_def *offset)
 {
+   nir_variable *var = intrin->variables[0]->var;
+   nir_variable_mode mode = var->data.mode;
+
    nir_intrinsic_op op;
    switch (mode) {
    case nir_var_shader_in:
-      op = per_vertex ? nir_intrinsic_load_per_vertex_input :
-                        nir_intrinsic_load_input;
+      op = vertex_index ? nir_intrinsic_load_per_vertex_input :
+                          nir_intrinsic_load_input;
       break;
    case nir_var_shader_out:
-      op = per_vertex ? nir_intrinsic_load_per_vertex_output :
-                        nir_intrinsic_load_output;
+      op = vertex_index ? nir_intrinsic_load_per_vertex_output :
+                          nir_intrinsic_load_output;
       break;
    case nir_var_uniform:
       op = nir_intrinsic_load_uniform;
@@ -189,33 +193,72 @@ load_op(nir_variable_mode mode, bool per_vertex)
    default:
       unreachable("Unknown variable mode");
    }
-   return op;
+
+   nir_intrinsic_instr *load = nir_intrinsic_instr_create(state->mem_ctx, op);
+   load->num_components = intrin->num_components;
+
+   nir_intrinsic_set_base(load, var->data.driver_location);
+   if (mode == nir_var_shader_in || mode == nir_var_shader_out)
+      nir_intrinsic_set_component(load, var->data.location_frac);
+
+   if (load->intrinsic == nir_intrinsic_load_uniform)
+      nir_intrinsic_set_range(load, state->type_size(var->type));
+
+   if (vertex_index)
+      load->src[0] = nir_src_for_ssa(vertex_index);
+
+   load->src[vertex_index ? 1 : 0] = nir_src_for_ssa(offset);
+
+   return load;
 }
 
-static nir_intrinsic_op
-store_op(struct lower_io_state *state,
-         nir_variable_mode mode, bool per_vertex)
+static nir_intrinsic_instr *
+lower_store(nir_intrinsic_instr *intrin, struct lower_io_state *state,
+            nir_ssa_def *vertex_index, nir_ssa_def *offset)
 {
+   nir_variable *var = intrin->variables[0]->var;
+   nir_variable_mode mode = var->data.mode;
+
    nir_intrinsic_op op;
-   switch (mode) {
-   case nir_var_shader_out:
-      op = per_vertex ? nir_intrinsic_store_per_vertex_output :
-                        nir_intrinsic_store_output;
-      break;
-   case nir_var_shared:
+   if (mode == nir_var_shared) {
       op = nir_intrinsic_store_shared;
-      break;
-   default:
-      unreachable("Unknown variable mode");
+   } else {
+      assert(mode == nir_var_shader_out);
+      op = vertex_index ? nir_intrinsic_store_per_vertex_output :
+                          nir_intrinsic_store_output;
    }
-   return op;
+
+   nir_intrinsic_instr *store = nir_intrinsic_instr_create(state->mem_ctx, op);
+   store->num_components = intrin->num_components;
+
+   nir_src_copy(&store->src[0], &intrin->src[0], store);
+
+   nir_intrinsic_set_base(store, var->data.driver_location);
+
+   if (mode == nir_var_shader_out)
+      nir_intrinsic_set_component(store, var->data.location_frac);
+
+   nir_intrinsic_set_write_mask(store, nir_intrinsic_write_mask(intrin));
+
+   if (vertex_index)
+      store->src[1] = nir_src_for_ssa(vertex_index);
+
+   store->src[vertex_index ? 2 : 1] = nir_src_for_ssa(offset);
+
+   return store;
 }
 
-static nir_intrinsic_op
-atomic_op(nir_intrinsic_op opcode)
+static nir_intrinsic_instr *
+lower_atomic(nir_intrinsic_instr *intrin, struct lower_io_state *state,
+             nir_ssa_def *offset)
 {
-   switch (opcode) {
-#define OP(O) case nir_intrinsic_var_##O: return nir_intrinsic_shared_##O;
+   nir_variable *var = intrin->variables[0]->var;
+
+   assert(var->data.mode == nir_var_shared);
+
+   nir_intrinsic_op op;
+   switch (intrin->intrinsic) {
+#define OP(O) case nir_intrinsic_var_##O: op = nir_intrinsic_shared_##O; break;
    OP(atomic_exchange)
    OP(atomic_comp_swap)
    OP(atomic_add)
@@ -230,6 +273,18 @@ atomic_op(nir_intrinsic_op opcode)
    default:
       unreachable("Invalid atomic");
    }
+
+   nir_intrinsic_instr *atomic =
+      nir_intrinsic_instr_create(state->mem_ctx, op);
+
+   atomic->src[0] = nir_src_for_ssa(offset);
+   atomic->const_index[0] = var->data.driver_location;
+
+   for (unsigned i = 0; i < nir_op_infos[intrin->intrinsic].num_inputs; i++) {
+      nir_src_copy(&atomic->src[i+1], &intrin->src[i], atomic);
+   }
+
+   return atomic;
 }
 
 static bool
@@ -282,7 +337,7 @@ nir_lower_io_block(nir_block *block,
          is_per_vertex_input(state, var) || is_per_vertex_output(state, var);
 
       nir_ssa_def *offset;
-      nir_ssa_def *vertex_index;
+      nir_ssa_def *vertex_index = NULL;
 
       offset = get_io_offset(b, intrin->variables[0],
                              per_vertex ? &vertex_index : NULL,
@@ -291,56 +346,13 @@ nir_lower_io_block(nir_block *block,
       nir_intrinsic_instr *replacement;
 
       switch (intrin->intrinsic) {
-      case nir_intrinsic_load_var: {
-         nir_intrinsic_instr *load =
-            nir_intrinsic_instr_create(state->mem_ctx,
-                                       load_op(mode, per_vertex));
-         load->num_components = intrin->num_components;
-
-         nir_intrinsic_set_base(load,
-            var->data.driver_location);
-         if (mode == nir_var_shader_in || mode == nir_var_shader_out) {
-            nir_intrinsic_set_component(load, var->data.location_frac);
-         }
-
-         if (load->intrinsic == nir_intrinsic_load_uniform) {
-            nir_intrinsic_set_range(load, state->type_size(var->type));
-         }
-
-         if (per_vertex)
-            load->src[0] = nir_src_for_ssa(vertex_index);
-
-         load->src[per_vertex ? 1 : 0] = nir_src_for_ssa(offset);
-
-         replacement = load;
+      case nir_intrinsic_load_var:
+         replacement = lower_load(intrin, state, vertex_index, offset);
          break;
-      }
-
-      case nir_intrinsic_store_var: {
-         assert(mode == nir_var_shader_out || mode == nir_var_shared);
-
-         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);
 
-         nir_intrinsic_set_base(store,
-            var->data.driver_location);
-         if (mode == nir_var_shader_out) {
-            nir_intrinsic_set_component(store, var->data.location_frac);
-         }
-         nir_intrinsic_set_write_mask(store, nir_intrinsic_write_mask(intrin));
-
-         if (per_vertex)
-            store->src[1] = nir_src_for_ssa(vertex_index);
-
-         store->src[per_vertex ? 2 : 1] = nir_src_for_ssa(offset);
-
-         replacement = store;
+      case nir_intrinsic_store_var:
+         replacement = lower_store(intrin, state, vertex_index, offset);
          break;
-      }
 
       case nir_intrinsic_var_atomic_add:
       case nir_intrinsic_var_atomic_imin:
@@ -351,26 +363,10 @@ nir_lower_io_block(nir_block *block,
       case nir_intrinsic_var_atomic_or:
       case nir_intrinsic_var_atomic_xor:
       case nir_intrinsic_var_atomic_exchange:
-      case nir_intrinsic_var_atomic_comp_swap: {
-         assert(mode == nir_var_shared);
-
-         nir_intrinsic_instr *atomic =
-            nir_intrinsic_instr_create(state->mem_ctx,
-                                       atomic_op(intrin->intrinsic));
-
-         atomic->src[0] = nir_src_for_ssa(offset);
-
-         atomic->const_index[0] = var->data.driver_location;
-
-         for (unsigned i = 0;
-              i < nir_op_infos[intrin->intrinsic].num_inputs;
-              i++) {
-            nir_src_copy(&atomic->src[i+1], &intrin->src[i], atomic);
-         }
-
-         replacement = atomic;
+      case nir_intrinsic_var_atomic_comp_swap:
+         assert(vertex_index == NULL);
+         replacement = lower_atomic(intrin, state, offset);
          break;
-      }
 
       default:
          break;