nir: Take a variable remap parameter in nir_inline_function_impl

Acked-by: Daniel Stone <daniels@collabora.com>
Reviewed-by: Jesse Natalie <jenatali@microsoft.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/6411>

src/compiler/nir/nir.h

@@ -4028,7 +4028,8 @@ bool nir_lower_returns(nir_shader *shader);
 
 void nir_inline_function_impl(struct nir_builder *b,
                               const nir_function_impl *impl,
-                              nir_ssa_def **params);
+                              nir_ssa_def **params,
+                              struct hash_table *shader_var_remap);
 bool nir_inline_functions(nir_shader *shader);
 
 bool nir_propagate_invariant(nir_shader *shader);

src/compiler/nir/nir_inline_functions.c

@@ -28,7 +28,8 @@
 
 void nir_inline_function_impl(struct nir_builder *b,
                               const nir_function_impl *impl,
-                              nir_ssa_def **params)
+                              nir_ssa_def **params,
+                              struct hash_table *shader_var_remap)
 {
    nir_function_impl *copy = nir_function_impl_clone(b->shader, impl);
 
@@ -45,6 +46,37 @@ void nir_inline_function_impl(struct nir_builder *b,
    nir_foreach_block(block, copy) {
       nir_foreach_instr_safe(instr, block) {
          switch (instr->type) {
+         case nir_instr_type_deref: {
+            nir_deref_instr *deref = nir_instr_as_deref(instr);
+            if (deref->deref_type != nir_deref_type_var)
+               break;
+
+            /* We don't need to remap function variables.  We already cloned
+             * them as part of nir_function_impl_clone and appended them to
+             * b->impl->locals.
+             */
+            if (deref->var->data.mode == nir_var_function_temp)
+               break;
+
+            /* If no map is provided, we assume that there are either no
+             * shader variables or they already live b->shader (this is the
+             * case for function inlining within a single shader.
+             */
+            if (shader_var_remap == NULL)
+               break;
+
+            struct hash_entry *entry =
+               _mesa_hash_table_search(shader_var_remap, deref->var);
+            if (entry == NULL) {
+               nir_variable *nvar = nir_variable_clone(deref->var, b->shader);
+               nir_shader_add_variable(b->shader, nvar);
+               entry = _mesa_hash_table_insert(shader_var_remap,
+                                               deref->var, nvar);
+            }
+            deref->var = entry->data;
+            break;
+         }
+
          case nir_instr_type_intrinsic: {
             nir_intrinsic_instr *load = nir_instr_as_intrinsic(instr);
             if (load->intrinsic != nir_intrinsic_load_param)
@@ -122,7 +154,7 @@ inline_functions_block(nir_block *block, nir_builder *b,
                                      call->callee->params[i].num_components);
       }
 
-      nir_inline_function_impl(b, call->callee->impl, params);
+      nir_inline_function_impl(b, call->callee->impl, params, NULL);
    }
 
    return progress;

src/compiler/nir/nir_lower_double_ops.c

@@ -600,7 +600,7 @@ lower_doubles_instr_to_soft(nir_builder *b, nir_alu_instr *instr,
       params[i + 1] = nir_mov_alu(b, instr->src[i], 1);
    }
 
-   nir_inline_function_impl(b, func->impl, params);
+   nir_inline_function_impl(b, func->impl, params, NULL);
 
    return nir_load_deref(b, ret_deref);
 }