nir: Rework function parameters

8 files changed, 67 insertions, 21 deletions

diff --git a/src/compiler/nir/nir.c b/src/compiler/nir/nir.c
index 42a53f6f3db..4ab9fde871a 100644
--- a/src/compiler/nir/nir.c
+++ b/src/compiler/nir/nir.c
@@ -117,6 +117,10 @@ nir_shader_add_variable(nir_shader *shader, nir_variable *var)
       assert(!"nir_shader_add_variable cannot be used for local variables");
       break;
 
+   case nir_var_param:
+      assert(!"nir_shader_add_variable cannot be used for function parameters");
+      break;
+
    case nir_var_global:
       exec_list_push_tail(&shader->globals, &var->node);
       break;
@@ -307,6 +311,20 @@ nir_function_impl_create(nir_function *function)
    impl->params = ralloc_array(function->shader,
                                nir_variable *, impl->num_params);
 
+   for (unsigned i = 0; i < impl->num_params; i++) {
+      impl->params[i] = rzalloc(function->shader, nir_variable);
+      impl->params[i]->type = function->params[i].type;
+      impl->params[i]->data.mode = nir_var_param;
+      impl->params[i]->data.location = i;
+   }
+
+   if (!glsl_type_is_void(function->return_type)) {
+      impl->return_var = rzalloc(function->shader, nir_variable);
+      impl->return_var->type = function->return_type;
+      impl->return_var->data.mode = nir_var_param;
+      impl->return_var->data.location = -1;
+   }
+
    return impl;
 }
 
diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index f130e5e0eb1..7aba195fa69 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -89,7 +89,8 @@ typedef enum {
    nir_var_uniform,
    nir_var_shader_storage,
    nir_var_shared,
-   nir_var_system_value
+   nir_var_system_value,
+   nir_var_param,
 } nir_variable_mode;
 
 /**
@@ -172,7 +173,7 @@ typedef struct nir_variable {
        *
        * \sa nir_variable_mode
        */
-      nir_variable_mode mode:4;
+      nir_variable_mode mode:5;
 
       /**
        * Interpolation mode for shader inputs / outputs
@@ -355,6 +356,12 @@ typedef struct nir_variable {
 #define nir_foreach_variable(var, var_list) \
    foreach_list_typed(nir_variable, var, node, var_list)
 
+static inline bool
+nir_variable_is_global(const nir_variable *var)
+{
+   return var->data.mode != nir_var_local && var->data.mode != nir_var_param;
+}
+
 /**
  * Returns the bits in the inputs_read, outputs_written, or
  * system_values_read bitfield corresponding to this variable.
diff --git a/src/compiler/nir/nir_clone.c b/src/compiler/nir/nir_clone.c
index bc6df56b753..8ece50a69cb 100644
--- a/src/compiler/nir/nir_clone.c
+++ b/src/compiler/nir/nir_clone.c
@@ -106,7 +106,7 @@ remap_reg(clone_state *state, const nir_register *reg)
 static nir_variable *
 remap_var(clone_state *state, const nir_variable *var)
 {
-   return _lookup_ptr(state, var, var->data.mode != nir_var_local);
+   return _lookup_ptr(state, var, nir_variable_is_global(var));
 }
 
 nir_constant *
@@ -591,9 +591,10 @@ clone_function_impl(clone_state *state, const nir_function_impl *fi)
    nfi->num_params = fi->num_params;
    nfi->params = ralloc_array(state->ns, nir_variable *, fi->num_params);
    for (unsigned i = 0; i < fi->num_params; i++) {
-      nfi->params[i] = remap_local(state, fi->params[i]);
+      nfi->params[i] = clone_variable(state, fi->params[i]);
    }
-   nfi->return_var = remap_local(state, fi->return_var);
+   if (fi->return_var)
+      nfi->return_var = clone_variable(state, fi->return_var);
 
    assert(list_empty(&state->phi_srcs));
 
diff --git a/src/compiler/nir/nir_inline_functions.c b/src/compiler/nir/nir_inline_functions.c
index 3cf83279053..b343eb735b1 100644
--- a/src/compiler/nir/nir_inline_functions.c
+++ b/src/compiler/nir/nir_inline_functions.c
@@ -69,12 +69,17 @@ inline_functions_block(nir_block *block, void *void_state)
       /* Add copies of all in parameters */
       assert(call->num_params == callee_copy->num_params);
       for (unsigned i = 0; i < callee_copy->num_params; i++) {
+         nir_variable *param = callee_copy->params[i];
+
+         /* Turn it into a local variable */
+         param->data.mode = nir_var_local;
+         exec_list_push_head(&b->impl->locals, &param->node);
+
          /* Only in or inout parameters */
          if (call->callee->params[i].param_type == nir_parameter_out)
             continue;
 
-         nir_copy_deref_var(b, nir_deref_var_create(b->shader,
-                                                    callee_copy->params[i]),
+         nir_copy_deref_var(b, nir_deref_var_create(b->shader, param),
                                call->params[i]);
       }
 
@@ -97,6 +102,10 @@ inline_functions_block(nir_block *block, void *void_state)
                                                     callee_copy->params[i]));
       }
       if (!glsl_type_is_void(call->callee->return_type)) {
+         /* Turn it into a local variable */
+         callee_copy->return_var->data.mode = nir_var_local;
+         exec_list_push_head(&b->impl->locals, &callee_copy->return_var->node);
+
          nir_copy_deref_var(b, call->return_deref,
                                nir_deref_var_create(b->shader,
                                                     callee_copy->return_var));
diff --git a/src/compiler/nir/nir_print.c b/src/compiler/nir/nir_print.c
index f36b91de6e0..276a948460c 100644
--- a/src/compiler/nir/nir_print.c
+++ b/src/compiler/nir/nir_print.c
@@ -943,6 +943,16 @@ print_function_impl(nir_function_impl *impl, print_state *state)
 
    fprintf(fp, "{\n");
 
+   for (unsigned i = 0; i < impl->num_params; i++) {
+      fprintf(fp, "\t");
+      print_var_decl(impl->params[i], state);
+   }
+
+   if (impl->return_var) {
+      fprintf(fp, "\t");
+      print_var_decl(impl->return_var, state);
+   }
+
    nir_foreach_variable(var, &impl->locals) {
       fprintf(fp, "\t");
       print_var_decl(var, state);
diff --git a/src/compiler/nir/nir_sweep.c b/src/compiler/nir/nir_sweep.c
index 5c62154ec7f..b22f0f56569 100644
--- a/src/compiler/nir/nir_sweep.c
+++ b/src/compiler/nir/nir_sweep.c
@@ -119,6 +119,8 @@ sweep_impl(nir_shader *nir, nir_function_impl *impl)
    ralloc_steal(nir, impl);
 
    ralloc_steal(nir, impl->params);
+   for (unsigned i = 0; i < impl->num_params; i++)
+      ralloc_steal(nir, impl->params[i]);
    ralloc_steal(nir, impl->return_var);
    steal_list(nir, nir_variable, &impl->locals);
    steal_list(nir, nir_register, &impl->registers);
diff --git a/src/compiler/nir/nir_validate.c b/src/compiler/nir/nir_validate.c
index 1a943d76314..00184cabe20 100644
--- a/src/compiler/nir/nir_validate.c
+++ b/src/compiler/nir/nir_validate.c
@@ -864,7 +864,7 @@ postvalidate_reg_decl(nir_register *reg, validate_state *state)
 static void
 validate_var_decl(nir_variable *var, bool is_global, validate_state *state)
 {
-   assert(is_global != (var->data.mode == nir_var_local));
+   assert(is_global == nir_variable_is_global(var));
 
    /*
     * TODO validate some things ir_validate.cpp does (requires more GLSL type
@@ -933,13 +933,19 @@ validate_function_impl(nir_function_impl *impl, validate_state *state)
    assert(impl->cf_node.parent == NULL);
 
    assert(impl->num_params == impl->function->num_params);
-   for (unsigned i = 0; i < impl->num_params; i++)
+   for (unsigned i = 0; i < impl->num_params; i++) {
       assert(impl->params[i]->type == impl->function->params[i].type);
+      assert(impl->params[i]->data.location == i);
+      validate_var_decl(impl->params[i], false, state);
+   }
 
-   if (glsl_type_is_void(impl->function->return_type))
+   if (glsl_type_is_void(impl->function->return_type)) {
       assert(impl->return_var == NULL);
-   else
+   } else {
       assert(impl->return_var->type == impl->function->return_type);
+      assert(impl->return_var->data.location == -1);
+      validate_var_decl(impl->return_var, false, state);
+   }
 
    assert(exec_list_is_empty(&impl->end_block->instr_list));
    assert(impl->end_block->successors[0] == NULL);
diff --git a/src/compiler/nir/spirv/vtn_cfg.c b/src/compiler/nir/spirv/vtn_cfg.c
index 041408b1cfb..144aac315e5 100644
--- a/src/compiler/nir/spirv/vtn_cfg.c
+++ b/src/compiler/nir/spirv/vtn_cfg.c
@@ -73,10 +73,6 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
       func->return_type = glsl_get_function_return_type(func_type);
 
       b->func->impl = nir_function_impl_create(func);
-      if (!glsl_type_is_void(func->return_type)) {
-         b->func->impl->return_var =
-            nir_local_variable_create(b->func->impl, func->return_type, "ret");
-      }
 
       b->func_param_idx = 0;
       break;
@@ -92,13 +88,10 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
          vtn_push_value(b, w[2], vtn_value_type_access_chain);
 
       assert(b->func_param_idx < b->func->impl->num_params);
-      unsigned idx = b->func_param_idx++;
+      nir_variable *param = b->func->impl->params[b->func_param_idx++];
 
-      nir_variable *param =
-         nir_local_variable_create(b->func->impl,
-                                   b->func->impl->function->params[idx].type,
-                                   val->name);
-      b->func->impl->params[idx] = param;
+      /* Name the parameter so it shows up nicely in NIR */
+      param->name = ralloc_strdup(param, val->name);
 
       struct vtn_variable *vtn_var = rzalloc(b, struct vtn_variable);
       vtn_var->mode = vtn_variable_mode_param;