nir: add nir_fixup_is_exported pass

See comment in the pass for motivation. To be used for asahi clc.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Konstantin Seurer <konstantin.seurer@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32226>

src/compiler/nir/meson.build

@@ -93,6 +93,7 @@ files_libnir = files(
   'nir_deref.h',
   'nir_divergence_analysis.c',
   'nir_dominance.c',
+  'nir_fixup_is_exported.c',
   'nir_format_convert.c',
   'nir_format_convert.h',
   'nir_from_ssa.c',

src/compiler/nir/nir.h

@@ -4566,6 +4566,7 @@ nir_shader_get_function_for_name(const nir_shader *shader, const char *name)
  */
 void nir_remove_non_entrypoints(nir_shader *shader);
 void nir_remove_non_exported(nir_shader *shader);
+void nir_fixup_is_exported(nir_shader *shader);
 
 nir_shader *nir_shader_create(void *mem_ctx,
                               gl_shader_stage stage,

src/compiler/nir/nir_fixup_is_exported.c

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2024 Valve Corporation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir.h"
+
+/*
+ * After going through the clang -> LLVM -> SPIR-V translator -> vtn pipes,
+ * OpenCL kernel's end up translated to 2 nir_functions:
+ *
+ * - a "wrapper" function that is_entrypoint but not is_exported
+ * - the "real" function that is_exported
+ *
+ * Confusingly, both functions have the same name.
+ *
+ * Also, workgroup size information is on the wrapper function only, so we can't
+ * just ignore the wrappers. But inlining and removing non-exported would delete
+ * the whole shader and lose that information.
+ *
+ * This pass is a silly solution to the silly problem: it looks for shadowed
+ * function names, which can only come from these wrappers. It then exports the
+ * wrappers and unexports the inner functions. After inlining and removing
+ * non-exported functions, we're left with a single function per kernel with
+ * workgroup size information preserved.
+ *
+ * While we're at it, we unexport _prefixed functions. This is an escape hatch
+ * to allow defining `kernel`s that are not intended for export, to workaround
+ * OpenCL limitations around `static kernel`s and shared local memory outside
+ * `kernel`s.
+ */
+void
+nir_fixup_is_exported(nir_shader *nir)
+{
+   struct set *seen =
+      _mesa_set_create(NULL, _mesa_hash_string, _mesa_key_string_equal);
+   struct set *shadowed =
+      _mesa_set_create(NULL, _mesa_hash_string, _mesa_key_string_equal);
+
+   nir_foreach_function(func, nir) {
+      if (_mesa_set_search(seen, func->name)) {
+         _mesa_set_add(shadowed, func->name);
+      } else {
+         _mesa_set_add(seen, func->name);
+      }
+   }
+
+   nir_foreach_function(func, nir) {
+      if (_mesa_set_search(shadowed, func->name)) {
+         func->is_exported = func->is_entrypoint;
+      }
+
+      if (func->name[0] == '_') {
+         func->is_exported = func->is_entrypoint = false;
+      }
+   }
+
+   _mesa_set_destroy(seen, NULL);
+   _mesa_set_destroy(shadowed, NULL);
+}