nir: add driver_functions option to avoid inlining.

This adds a driver control to instruct NIR to not inline
all functions.

It adds a very simple inlining heuristic that works for
what I've played with, but will probably need to grow some
better ideas.

Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24687>

src/compiler/nir/nir.h

@@ -3811,6 +3811,8 @@ typedef struct nir_shader_compiler_options {
     */
    nir_variable_mode force_indirect_unrolling;
 
+   bool driver_functions;
+
    nir_lower_int64_options lower_int64_options;
    nir_lower_doubles_options lower_doubles_options;
    nir_divergence_options divergence_analysis_options;

src/compiler/nir/nir_functions.c

@@ -26,6 +26,26 @@
 #include "nir_control_flow.h"
 #include "nir_vla.h"
 
+/*
+ * TODO: write a proper inliner for GPUs.
+ * This heuristic just inlines small functions,
+ * and tail calls get inlined as well.
+ */
+static bool
+nir_function_can_inline(nir_function *function)
+{
+   bool can_inline = true;
+   if (!function->should_inline) {
+      if (function->impl) {
+         if (function->impl->num_blocks > 2)
+            can_inline = false;
+         if (function->impl->ssa_alloc > 45)
+            can_inline = false;
+      }
+   }
+   return can_inline;
+}
+
 static bool
 function_ends_in_jump(nir_function_impl *impl)
 {
@@ -138,6 +158,14 @@ static bool inline_functions_pass(nir_builder *b,
    nir_call_instr *call = nir_instr_as_call(instr);
    assert(call->callee->impl);
 
+   if (b->shader->options->driver_functions &&
+       b->shader->info.stage == MESA_SHADER_KERNEL) {
+      bool last_instr = (instr == nir_block_last_instr(instr->block));
+      if (!nir_function_can_inline(call->callee) && !last_instr) {
+         return false;
+      }
+   }
+
    /* Make sure that the function we're calling is already inlined */
    inline_function_impl(call->callee->impl, inlined);
 

src/compiler/spirv/spirv_to_nir.c

@@ -6899,6 +6899,7 @@ spirv_to_nir(const uint32_t *words, size_t word_count,
       nir_function *entry_point = b->entry_point->func->nir_func;
       vtn_assert(entry_point);
 
+      entry_point->dont_inline = false;
       /* post process entry_points with input params */
       if (entry_point->num_params && b->shader->info.stage == MESA_SHADER_KERNEL)
          entry_point = vtn_emit_kernel_entry_point_wrapper(b, entry_point);