nir: Add nir_opt_reuse_constants()

Currently SPIR-V does pull all the NIR constants it creates into the
first block, but we plan to change that behavior to let those constants
be defined as they are used.

This pass was written to provide a fallback to the old behavior, it will
be used for radv to avoid regressions when performing the SPIR-V change.

Reviewed-by: Timur Kristóf <timur.kristof@gmail.com>
Reviewed-by: Faith Ekstrand <faith.ekstrand@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5282>

src/compiler/nir/meson.build

@@ -265,6 +265,7 @@ files_libnir = files(
   'nir_opt_reassociate_bfi.c',
   'nir_opt_rematerialize_compares.c',
   'nir_opt_remove_phis.c',
+  'nir_opt_reuse_constants.c',
   'nir_opt_shrink_stores.c',
   'nir_opt_shrink_vectors.c',
   'nir_opt_sink.c',

src/compiler/nir/nir.h

@@ -6027,6 +6027,8 @@ bool nir_opt_ray_queries(nir_shader *shader);
 
 bool nir_opt_ray_query_ranges(nir_shader *shader);
 
+bool nir_opt_reuse_constants(nir_shader *shader);
+
 void nir_sweep(nir_shader *shader);
 
 void nir_remap_dual_slot_attributes(nir_shader *shader,

src/compiler/nir/nir_opt_reuse_constants.c

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2023 Intel Corporation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir.h"
+#include "nir_instr_set.h"
+
+bool
+nir_opt_reuse_constants(nir_shader *shader)
+{
+   bool progress = false;
+
+   struct set *consts = nir_instr_set_create(NULL);
+   nir_foreach_function_impl(impl, shader) {
+      _mesa_set_clear(consts, NULL);
+
+      nir_block *start_block = nir_start_block(impl);
+      bool func_progress = false;
+
+      nir_foreach_block_safe(block, impl) {
+         const bool in_start_block = start_block == block;
+         nir_foreach_instr_safe(instr, block) {
+            if (instr->type != nir_instr_type_load_const)
+               continue;
+
+            struct set_entry *entry = _mesa_set_search(consts, instr);
+            if (!entry) {
+               if (!in_start_block)
+                  nir_instr_move(nir_after_block_before_jump(start_block), instr);
+               _mesa_set_add(consts, instr);
+            }
+
+            func_progress |= nir_instr_set_add_or_rewrite(consts, instr, nir_instrs_equal);
+         }
+      }
+
+      if (func_progress) {
+         nir_metadata_preserve(impl, nir_metadata_block_index |
+                                        nir_metadata_dominance);
+         progress = true;
+      } else {
+         nir_metadata_preserve(impl, nir_metadata_all);
+      }
+   }
+
+   nir_instr_set_destroy(consts);
+   return progress;
+}