nir,mesa: Add helpers for creating uniform state variables.

It's one of the weirder parts of our shader interface's interactions with
the GL API, so let's try to make it a little cleaner.

Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/23111>

src/compiler/nir/nir.c

@@ -323,6 +323,21 @@ nir_local_variable_create(nir_function_impl *impl,
    return var;
 }
 
+nir_variable *
+nir_state_variable_create(nir_shader *shader,
+                          const struct glsl_type *type,
+                          const char *name,
+                          const gl_state_index16 tokens[STATE_LENGTH])
+{
+   nir_variable *var = nir_variable_create(shader, nir_var_uniform, type, name);
+   var->num_state_slots = 1;
+   var->state_slots = rzalloc_array(var, nir_state_slot, 1);
+   memcpy(var->state_slots[0].tokens, tokens,
+          sizeof(var->state_slots[0].tokens));
+   shader->num_uniforms++;
+   return var;
+}
+
 nir_variable *
 nir_create_variable_with_location(nir_shader *shader, nir_variable_mode mode, int location,
                                   const struct glsl_type *type)

src/compiler/nir/nir.h

@@ -4056,6 +4056,13 @@ nir_variable *nir_local_variable_create(nir_function_impl *impl,
                                         const struct glsl_type *type,
                                         const char *name);
 
+/** Creates a uniform builtin state variable. */
+nir_variable *
+nir_state_variable_create(nir_shader *shader,
+                          const struct glsl_type *type,
+                          const char *name,
+                          const gl_state_index16 tokens[STATE_LENGTH]);
+
 /* Gets the variable for the given mode and location, creating it (with the given
  * type) if necessary.
  */

src/compiler/nir/nir_lower_alpha_test.c

@@ -95,15 +95,9 @@ nir_lower_alpha_test(nir_shader *shader, enum compare_func func,
                                       3);
                }
 
-               nir_variable *var = nir_variable_create(shader,
-                                                       nir_var_uniform,
-                                                       glsl_float_type(),
-                                                       "gl_AlphaRefMESA");
-               var->num_state_slots = 1;
-               var->state_slots = rzalloc_array(var, nir_state_slot, 1);
-               memcpy(var->state_slots[0].tokens,
-                      alpha_ref_state_tokens,
-                      sizeof(var->state_slots[0].tokens));
+               nir_variable *var = nir_state_variable_create(shader, glsl_float_type(),
+                                                             "gl_AlphaRefMESA",
+                                                             alpha_ref_state_tokens);
                nir_ssa_def *alpha_ref = nir_load_var(&b, var);
 
                nir_ssa_def *condition =

src/compiler/nir/nir_lower_atomics_to_ssbo.c

@@ -38,13 +38,8 @@ deref_offset_var(nir_builder *b, unsigned binding, unsigned offset_align_state)
    gl_state_index16 tokens[STATE_LENGTH] = {offset_align_state, binding};
    nir_variable *var = nir_find_state_variable(b->shader, tokens);
    if (!var) {
-      var = nir_variable_create(b->shader, nir_var_uniform, glsl_uint_type(), "offset");
-      var->state_slots = rzalloc_array(var, nir_state_slot, 1);
-      var->state_slots[0].tokens[0] = offset_align_state;
-      var->state_slots[0].tokens[1] = binding;
-      var->num_state_slots = 1;
+      var = nir_state_variable_create(b->shader, glsl_uint_type(), "offset", tokens);
       var->data.how_declared = nir_var_hidden;
-      b->shader->num_uniforms++;
    }
    return nir_build_deref_var(b, var);
 }

src/compiler/nir/nir_lower_clip.c

@@ -227,16 +227,9 @@ get_ucp(nir_builder *b, int plane,
    if (clipplane_state_tokens) {
       char tmp[100];
       snprintf(tmp, ARRAY_SIZE(tmp), "gl_ClipPlane%dMESA", plane);
-      nir_variable *var = nir_variable_create(b->shader,
-                                              nir_var_uniform,
-                                              glsl_vec4_type(),
-                                              tmp);
-
-      var->num_state_slots = 1;
-      var->state_slots = ralloc_array(var, nir_state_slot, 1);
-      memcpy(var->state_slots[0].tokens,
-             clipplane_state_tokens[plane],
-             sizeof(var->state_slots[0].tokens));
+      nir_variable *var = nir_state_variable_create(b->shader,
+                                                    glsl_vec4_type(),
+                                                    tmp, clipplane_state_tokens[plane]);
       return nir_load_var(b, var);
    } else
       return nir_load_user_clip_plane(b, plane);

src/compiler/nir/nir_lower_drawpixels.c

@@ -47,27 +47,12 @@ get_texcoord(nir_builder *b, lower_drawpixels_state *state)
    return nir_load_var(b, state->texcoord);
 }
 
-static nir_variable *
-create_uniform(nir_shader *shader, const char *name,
-               const gl_state_index16 state_tokens[STATE_LENGTH])
-{
-   nir_variable *var = nir_variable_create(shader,
-                                           nir_var_uniform,
-                                           glsl_vec4_type(),
-                                           name);
-   var->num_state_slots = 1;
-   var->state_slots = ralloc_array(var, nir_state_slot, 1);
-   memcpy(var->state_slots[0].tokens, state_tokens,
-          sizeof(var->state_slots[0].tokens));
-   return var;
-}
-
 static nir_ssa_def *
 get_scale(nir_builder *b, lower_drawpixels_state *state)
 {
    if (state->scale == NULL) {
-      state->scale = create_uniform(state->shader, "gl_PTscale",
-                                    state->options->scale_state_tokens);
+      state->scale = nir_state_variable_create(state->shader, glsl_vec4_type(), "gl_PTscale",
+                                               state->options->scale_state_tokens);
    }
    return nir_load_var(b, state->scale);
 }
@@ -76,8 +61,8 @@ static nir_ssa_def *
 get_bias(nir_builder *b, lower_drawpixels_state *state)
 {
    if (state->bias == NULL) {
-      state->bias = create_uniform(state->shader, "gl_PTbias",
-                                   state->options->bias_state_tokens);
+      state->bias = nir_state_variable_create(state->shader, glsl_vec4_type(), "gl_PTbias",
+                                              state->options->bias_state_tokens);
    }
    return nir_load_var(b, state->bias);
 }
@@ -86,9 +71,9 @@ static nir_ssa_def *
 get_texcoord_const(nir_builder *b, lower_drawpixels_state *state)
 {
    if (state->texcoord_const == NULL) {
-      state->texcoord_const = create_uniform(state->shader,
-                                   "gl_MultiTexCoord0",
-                                   state->options->texcoord_state_tokens);
+      state->texcoord_const = nir_state_variable_create(state->shader, glsl_vec4_type(),
+                                                        "gl_MultiTexCoord0",
+                                                        state->options->texcoord_state_tokens);
    }
    return nir_load_var(b, state->texcoord_const);
 }

src/compiler/nir/nir_lower_patch_vertices.c

@@ -30,11 +30,8 @@ make_uniform(nir_shader *nir, const gl_state_index16 *tokens)
     * special handling in uniform setup.
     */
    nir_variable *var =
-      nir_variable_create(nir, nir_var_uniform, glsl_int_type(),
-                          "gl_PatchVerticesIn");
-   var->num_state_slots = 1;
-   var->state_slots = ralloc_array(var, nir_state_slot, var->num_state_slots);
-   memcpy(var->state_slots[0].tokens, tokens, sizeof(*tokens) * STATE_LENGTH);
+      nir_state_variable_create(nir, glsl_int_type(),
+                                "gl_PatchVerticesIn", tokens);
 
    return var;
 }

src/compiler/nir/nir_lower_pntc_ytransform.c

@@ -42,15 +42,11 @@ get_pntc_transform(lower_pntc_ytransform_state *state)
       /* NOTE: name must be prefixed w/ "gl_" to trigger slot based
        * special handling in uniform setup:
        */
-      nir_variable *var = nir_variable_create(state->shader,
-                                              nir_var_uniform,
-                                              glsl_vec4_type(),
-                                              "gl_PntcYTransform");
+      nir_variable *var = nir_state_variable_create(state->shader,
+                                                    glsl_vec4_type(),
+                                                    "gl_PntcYTransform",
+                                                    state->pntc_state_tokens);
 
-      var->num_state_slots = 1;
-      var->state_slots = ralloc_array(var, nir_state_slot, 1);
-      memcpy(var->state_slots[0].tokens, state->pntc_state_tokens,
-             sizeof(var->state_slots[0].tokens));
       var->data.how_declared = nir_var_hidden;
       state->pntc_transform = var;
    }

src/compiler/nir/nir_lower_point_size_mov.c

@@ -42,13 +42,9 @@ lower_impl(nir_function_impl *impl,
 
    nir_builder_init(&b, impl);
 
-   in = nir_variable_create(shader, nir_var_uniform,
-                            glsl_vec4_type(), "gl_PointSizeClampedMESA");
-   in->num_state_slots = 1;
-   in->state_slots = ralloc_array(in, nir_state_slot, 1);
-   memcpy(in->state_slots[0].tokens,
-         pointsize_state_tokens,
-         sizeof(in->state_slots[0].tokens));
+   in = nir_state_variable_create(shader, glsl_vec4_type(),
+                                  "gl_PointSizeClampedMESA",
+                                  pointsize_state_tokens);
 
    /* the existing output can't be removed in order to avoid breaking xfb.
     * drivers must check var->data.explicit_location to find the original output

src/compiler/nir/nir_lower_wpos_ytransform.c

@@ -48,15 +48,11 @@ get_transform(lower_wpos_ytransform_state *state)
       /* NOTE: name must be prefixed w/ "gl_" to trigger slot based
        * special handling in uniform setup:
        */
-      nir_variable *var = nir_variable_create(state->shader,
-                                              nir_var_uniform,
-                                              glsl_vec4_type(),
-                                              "gl_FbWposYTransform");
+      nir_variable *var = nir_state_variable_create(state->shader,
+                                                    glsl_vec4_type(),
+                                                    "gl_FbWposYTransform",
+                                                    state->options->state_tokens);
 
-      var->num_state_slots = 1;
-      var->state_slots = ralloc_array(var, nir_state_slot, 1);
-      memcpy(var->state_slots[0].tokens, state->options->state_tokens,
-             sizeof(var->state_slots[0].tokens));
       var->data.how_declared = nir_var_hidden;
       state->transform = var;
    }

src/gallium/drivers/d3d12/d3d12_lower_point_sprite.c

@@ -250,15 +250,8 @@ d3d12_lower_point_sprite(nir_shader *shader,
 
    /* Create uniform to retrieve inverse of viewport size and point size:
     * (1/ViewportWidth, 1/ViewportHeight, PointSize, MaxPointSize) */
-   state.uniform = nir_variable_create(shader,
-                                       nir_var_uniform,
-                                       glsl_vec4_type(),
-                                       "d3d12_ViewportSizeRcp");
-   state.uniform->num_state_slots = 1;
-   state.uniform->state_slots = ralloc_array(state.uniform, nir_state_slot, 1);
-   memcpy(state.uniform->state_slots[0].tokens, tokens,
-          sizeof(state.uniform->state_slots[0].tokens));
-   shader->num_uniforms++;
+   state.uniform = nir_state_variable_create(shader, glsl_vec4_type(),
+                                             "d3d12_ViewportSizeRcp", tokens);
 
    /* Create new outputs for point tex coordinates */
    unsigned count = 0;

src/gallium/drivers/d3d12/d3d12_nir_passes.c

@@ -46,17 +46,9 @@ d3d12_get_state_var(nir_builder *b,
 {
    const gl_state_index16 tokens[STATE_LENGTH] = { STATE_INTERNAL_DRIVER, var_enum };
    if (*out_var == NULL) {
-      nir_variable *var = nir_variable_create(b->shader,
-                                              nir_var_uniform,
-                                              var_type,
-                                              var_name);
-
-      var->num_state_slots = 1;
-      var->state_slots = ralloc_array(var, nir_state_slot, 1);
-      memcpy(var->state_slots[0].tokens, tokens,
-             sizeof(var->state_slots[0].tokens));
+      nir_variable *var = nir_state_variable_create(b->shader, var_type,
+                                                    var_name, tokens);
       var->data.how_declared = nir_var_hidden;
-      b->shader->num_uniforms++;
       *out_var = var;
    }
    return nir_load_var(b, *out_var);

src/mesa/main/ffvertex_prog.c

@@ -46,6 +46,7 @@
 #include "util/bitscan.h"
 
 #include "state_tracker/st_program.h"
+#include "state_tracker/st_nir.h"
 
 #include "compiler/nir/nir_builder.h"
 #include "compiler/nir/nir_builtin_builder.h"
@@ -310,20 +311,9 @@ register_state_var(struct tnl_program *p,
    if (var)
       return var;
 
-   int loc = _mesa_add_state_reference(p->state_params, tokens);
+   var = st_nir_state_variable_create(p->b->shader, type, tokens);
+   var->data.driver_location = _mesa_add_state_reference(p->state_params, tokens);
 
-   char *name = _mesa_program_state_string(tokens);
-   var = nir_variable_create(p->b->shader, nir_var_uniform, type,
-                             name);
-   free(name);
-
-   var->num_state_slots = 1;
-   var->state_slots = ralloc_array(var, nir_state_slot, 1);
-   var->data.driver_location = loc;
-   memcpy(var->state_slots[0].tokens, tokens,
-          sizeof(var->state_slots[0].tokens));
-
-   p->b->shader->num_uniforms++;
    return var;
 }
 

src/mesa/state_tracker/st_nir.h

@@ -33,6 +33,7 @@ extern "C" {
 #endif
 
 struct nir_shader;
+struct nir_variable;
 
 void st_nir_lower_builtin(struct nir_shader *shader);
 void st_nir_lower_tex_src_plane(struct nir_shader *shader, unsigned free_slots,
@@ -81,6 +82,11 @@ st_nir_add_point_size(struct nir_shader *nir);
 struct pipe_shader_state *
 st_nir_make_clearcolor_shader(struct st_context *st);
 
+struct nir_variable *
+st_nir_state_variable_create(struct nir_shader *shader,
+                             const struct glsl_type *type,
+                             const gl_state_index16 state[STATE_LENGTH]);
+
 #ifdef __cplusplus
 }
 #endif

src/mesa/state_tracker/st_nir_lower_builtin.c

@@ -64,6 +64,18 @@
 #include "uniforms.h"
 #include "program/prog_instruction.h"
 
+/** Wrapper around nir_state_variable_create to pick the name automatically. */
+nir_variable *
+st_nir_state_variable_create(nir_shader *shader,
+                             const struct glsl_type *type,
+                             const gl_state_index16 tokens[STATE_LENGTH])
+{
+   char *name = _mesa_program_state_string(tokens);
+   nir_variable *var = nir_state_variable_create(shader, type, name, tokens);
+   free(name);
+   return var;
+}
+
 static const struct gl_builtin_uniform_element *
 get_element(const struct gl_builtin_uniform_desc *desc, nir_deref_path *path)
 {
@@ -137,19 +149,8 @@ get_variable(nir_builder *b, nir_deref_path *path,
    if (var)
       return var;
 
-   char *name = _mesa_program_state_string(tokens);
-
    /* variable doesn't exist yet, so create it: */
-   var = nir_variable_create(shader, nir_var_uniform, glsl_vec4_type(), name);
-
-   var->num_state_slots = 1;
-   var->state_slots = rzalloc_array(var, nir_state_slot, 1);
-   memcpy(var->state_slots[0].tokens, tokens,
-          sizeof(var->state_slots[0].tokens));
-
-   free(name);
-
-   return var;
+   return st_nir_state_variable_create(shader, glsl_vec4_type(), tokens);
 }
 
 static bool