nir: Add a more generic helper for gathering constant initializers

The one we had was tied to nir_var_mem_constant but we also need it for
global and, one day, I can imagine us needing it for shared (though
there's currently no spec that requires it).

Reviewed-by: Jesse Natalie <jenatali@microsoft.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/7565>

src/compiler/nir/nir.h

@@ -4424,6 +4424,10 @@ bool
 nir_lower_vars_to_explicit_types(nir_shader *shader,
                                  nir_variable_mode modes,
                                  glsl_type_size_align_func type_info);
+void
+nir_gather_explicit_io_initializers(nir_shader *shader,
+                                    void *dst, size_t dst_size,
+                                    nir_variable_mode mode);
 
 bool nir_lower_mem_constant_vars(nir_shader *shader,
                                  glsl_type_size_align_func type_info);

src/compiler/nir/nir_lower_io.c

@@ -2376,6 +2376,25 @@ write_constant(void *dst, size_t dst_size,
    }
 }
 
+void
+nir_gather_explicit_io_initializers(nir_shader *shader,
+                                    void *dst, size_t dst_size,
+                                    nir_variable_mode mode)
+{
+   /* It doesn't really make sense to gather initializers for more than one
+    * mode at a time.  If this ever becomes well-defined, we can drop the
+    * assert then.
+    */
+   assert(util_bitcount(mode) == 1);
+
+   nir_foreach_variable_with_modes(var, shader, mode) {
+      assert(var->data.driver_location < dst_size);
+      write_constant((char *)dst + var->data.driver_location,
+                     dst_size - var->data.driver_location,
+                     var->constant_initializer, var->type);
+   }
+}
+
 bool
 nir_lower_mem_constant_vars(nir_shader *shader,
                             glsl_type_size_align_func type_info)
@@ -2389,15 +2408,9 @@ nir_lower_mem_constant_vars(nir_shader *shader,
       shader->constant_data = rerzalloc_size(shader, shader->constant_data,
                                              old_constant_data_size,
                                              shader->constant_data_size);
-
+      nir_gather_explicit_io_initializers(shader, shader->constant_data,
-      nir_foreach_variable_with_modes(var, shader, nir_var_mem_constant) {
+                                          shader->constant_data_size,
-         assert(var->data.driver_location < shader->constant_data_size);
+                                          nir_var_mem_constant);
-         write_constant((char *)shader->constant_data +
-                           var->data.driver_location,
-                        shader->constant_data_size -
-                           var->data.driver_location,
-                        var->constant_initializer, var->type);
-      }
       progress = true;
    }