nir/inline_functions: Break inlining into a builder helper

This pulls the guts of function inlining into a builder helper so that
it can be used elsewhere.  The rest of the infrastructure is still
needed for most inlining cases to ensure that everything gets inlined
and only ever once.  However, there are use-cases where you just want to
inline one little thing.  This new helper also has a neat trick where it
can seamlessly inline a function from one nir_shader into another.

Reviewed-by: Matt Turner <[email protected]>
Reviewed-by: Jordan Justen <[email protected]>
Reviewed-by: Kenneth Graunke <[email protected]>

3 files changed, 60 insertions, 40 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 777ba0a4b71..c550f8339ed 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -2807,7 +2807,8 @@ void nir_print_instr(const nir_instr *instr, FILE *fp);
 void nir_print_deref(const nir_deref_instr *deref, FILE *fp);
 
 nir_shader *nir_shader_clone(void *mem_ctx, const nir_shader *s);
-nir_function_impl *nir_function_impl_clone(const nir_function_impl *fi);
+nir_function_impl *nir_function_impl_clone(nir_shader *shader,
+                                           const nir_function_impl *fi);
 nir_constant *nir_constant_clone(const nir_constant *c, nir_variable *var);
 nir_variable *nir_variable_clone(const nir_variable *c, nir_shader *shader);
 
@@ -2940,6 +2941,9 @@ bool nir_split_struct_vars(nir_shader *shader, nir_variable_mode modes);
 bool nir_lower_returns_impl(nir_function_impl *impl);
 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);
 bool nir_inline_functions(nir_shader *shader);
 
 bool nir_propagate_invariant(nir_shader *shader);
diff --git a/src/compiler/nir/nir_clone.c b/src/compiler/nir/nir_clone.c
index 557c7d29f53..b10068928a7 100644
--- a/src/compiler/nir/nir_clone.c
+++ b/src/compiler/nir/nir_clone.c
@@ -656,13 +656,12 @@ clone_function_impl(clone_state *state, const nir_function_impl *fi)
 }
 
 nir_function_impl *
-nir_function_impl_clone(const nir_function_impl *fi)
+nir_function_impl_clone(nir_shader *shader, const nir_function_impl *fi)
 {
    clone_state state;
    init_clone_state(&state, NULL, false, false);
 
-   /* We use the same shader */
-   state.ns = fi->function->shader;
+   state.ns = shader;
 
    nir_function_impl *nfi = clone_function_impl(&state, fi);
 
diff --git a/src/compiler/nir/nir_inline_functions.c b/src/compiler/nir/nir_inline_functions.c
index c57e16d2b44..8226d50b44b 100644
--- a/src/compiler/nir/nir_inline_functions.c
+++ b/src/compiler/nir/nir_inline_functions.c
@@ -26,6 +26,56 @@
 #include "nir_control_flow.h"
 #include "nir_vla.h"
 
+void nir_inline_function_impl(struct nir_builder *b,
+                              const nir_function_impl *impl,
+                              nir_ssa_def **params)
+{
+   nir_function_impl *copy = nir_function_impl_clone(b->shader, impl);
+
+   /* Insert a nop at the cursor so we can keep track of where things are as
+    * we add/remove stuff from the CFG.
+    */
+   nir_intrinsic_instr *nop =
+      nir_intrinsic_instr_create(b->shader, nir_intrinsic_nop);
+   nir_builder_instr_insert(b, &nop->instr);
+
+   exec_list_append(&b->impl->locals, &copy->locals);
+   exec_list_append(&b->impl->registers, &copy->registers);
+
+   nir_foreach_block(block, copy) {
+      nir_foreach_instr_safe(instr, block) {
+         /* Returns have to be lowered for this to work */
+         assert(instr->type != nir_instr_type_jump ||
+                nir_instr_as_jump(instr)->type != nir_jump_return);
+
+         if (instr->type != nir_instr_type_intrinsic)
+            continue;
+
+         nir_intrinsic_instr *load = nir_instr_as_intrinsic(instr);
+         if (load->intrinsic != nir_intrinsic_load_param)
+            continue;
+
+         unsigned param_idx = nir_intrinsic_param_idx(load);
+         assert(param_idx < impl->function->num_params);
+         assert(load->dest.is_ssa);
+         nir_ssa_def_rewrite_uses(&load->dest.ssa,
+                                  nir_src_for_ssa(params[param_idx]));
+
+         /* Remove any left-over load_param intrinsics because they're soon
+          * to be in another function and therefore no longer valid.
+          */
+         nir_instr_remove(&load->instr);
+      }
+   }
+
+   /* Pluck the body out of the function and place it here */
+   nir_cf_list body;
+   nir_cf_list_extract(&body, &copy->body);
+   nir_cf_reinsert(&body, nir_before_instr(&nop->instr));
+
+   b->cursor = nir_instr_remove(&nop->instr);
+}
+
 static bool inline_function_impl(nir_function_impl *impl, struct set *inlined);
 
 static bool
@@ -49,16 +99,10 @@ inline_functions_block(nir_block *block, nir_builder *b,
       nir_call_instr *call = nir_instr_as_call(instr);
       assert(call->callee->impl);
 
+      /* Make sure that the function we're calling is already inlined */
       inline_function_impl(call->callee->impl, inlined);
 
-      nir_function_impl *callee_copy =
-         nir_function_impl_clone(call->callee->impl);
-      callee_copy->function = call->callee;
-
-      exec_list_append(&b->impl->locals, &callee_copy->locals);
-      exec_list_append(&b->impl->registers, &callee_copy->registers);
-
-      b->cursor = nir_before_instr(&call->instr);
+      b->cursor = nir_instr_remove(&call->instr);
 
       /* Rewrite all of the uses of the callee's parameters to use the call
        * instructions sources.  In order to ensure that the "load" happens
@@ -72,34 +116,7 @@ inline_functions_block(nir_block *block, nir_builder *b,
                                      call->callee->params[i].num_components);
       }
 
-      nir_foreach_block(block, callee_copy) {
-         nir_foreach_instr_safe(instr, block) {
-            if (instr->type != nir_instr_type_intrinsic)
-               continue;
-
-            nir_intrinsic_instr *load = nir_instr_as_intrinsic(instr);
-            if (load->intrinsic != nir_intrinsic_load_param)
-               continue;
-
-            unsigned param_idx = nir_intrinsic_param_idx(load);
-            assert(param_idx < num_params);
-            assert(load->dest.is_ssa);
-            nir_ssa_def_rewrite_uses(&load->dest.ssa,
-                                     nir_src_for_ssa(params[param_idx]));
-
-            /* Remove any left-over load_param intrinsics because they're soon
-             * to be in another function and therefore no longer valid.
-             */
-            nir_instr_remove(&load->instr);
-         }
-      }
-
-      /* Pluck the body out of the function and place it here */
-      nir_cf_list body;
-      nir_cf_list_extract(&body, &callee_copy->body);
-      nir_cf_reinsert(&body, b->cursor);
-
-      nir_instr_remove(&call->instr);
+      nir_inline_function_impl(b, call->callee->impl, params);
    }
 
    return progress;