freedreno/ir3: Create UBO variables for driver-UBOs

Some nir passes, like lower_amul, expect to have varibles declared for
things that are accessed via load_ubo().

Fixes: 76e417ca59 ("turnip,ir3/a750: Implement consts loading via preamble")
Signed-off-by: Rob Clark <robdclark@chromium.org>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/31611>
(cherry picked from commit e548f90edb)

.pick_status.json

@@ -724,7 +724,7 @@
         "description": "freedreno/ir3: Create UBO variables for driver-UBOs",
         "nominated": true,
         "nomination_type": 1,
-        "resolution": 0,
+        "resolution": 1,
         "main_sha": null,
         "because_sha": "76e417ca593866080731da59c479a99542e3a529",
         "notes": null

src/freedreno/ir3/ir3_nir.c

@@ -51,6 +51,37 @@ ir3_get_driver_ubo(nir_builder *b, struct ir3_driver_ubo *ubo)
    return nir_imm_int(b, ubo->idx);
 }
 
+static const struct glsl_type *
+get_driver_ubo_type(const struct ir3_driver_ubo *ubo)
+{
+   return glsl_array_type(glsl_uint_type(), ubo->size, 0);
+}
+
+/* Create or update the size of a driver-ubo: */
+void
+ir3_update_driver_ubo(nir_shader *nir, const struct ir3_driver_ubo *ubo, const char *name)
+{
+   if (ubo->idx < 0)
+      return;
+
+
+   nir_foreach_variable_in_shader(var, nir) {
+      if (var->data.mode != nir_var_mem_ubo)
+         continue;
+      if (var->data.binding != ubo->idx)
+         continue;
+
+      /* UBO already exists, make sure it is big enough: */
+      if (glsl_array_size(var->type) < ubo->size)
+         var->type = get_driver_ubo_type(ubo);
+   }
+
+   /* UBO variable does not exist yet, so create it: */
+   nir_variable *var =
+      nir_variable_create(nir, nir_var_mem_ubo, get_driver_ubo_type(ubo), name);
+   var->data.driver_location = ubo->idx;
+}
+
 nir_def *
 ir3_load_driver_ubo(nir_builder *b, unsigned components,
                     struct ir3_driver_ubo *ubo,

src/freedreno/ir3/ir3_nir.h

@@ -93,6 +93,7 @@ nir_def *ir3_nir_try_propagate_bit_shift(nir_builder *b,
 bool ir3_nir_opt_subgroups(nir_shader *nir, struct ir3_shader_variant *v);
 
 nir_def *ir3_get_driver_ubo(nir_builder *b, struct ir3_driver_ubo *ubo);
+void ir3_update_driver_ubo(nir_shader *nir, const struct ir3_driver_ubo *ubo, const char *name);
 nir_def *ir3_load_driver_ubo(nir_builder *b, unsigned components,
                              struct ir3_driver_ubo *ubo,
                              unsigned offset);

src/freedreno/ir3/ir3_nir_analyze_ubo_ranges.c

@@ -872,6 +872,9 @@ ir3_nir_lower_load_constant(nir_shader *nir, struct ir3_shader_variant *v)
                compiler->const_upload_unit * 4 * sizeof(uint32_t));
       v->constant_data = rzalloc_size(v, v->constant_data_size);
       memcpy(v->constant_data, nir->constant_data, nir->constant_data_size);
+
+      const struct ir3_const_state *const_state = ir3_const_state(v);
+      ir3_update_driver_ubo(nir, &const_state->consts_ubo, "$consts");
    }
 
    return progress;

src/freedreno/ir3/ir3_nir_lower_driver_params_to_ubo.c

@@ -82,5 +82,13 @@ ir3_nir_lower_driver_params_to_ubo(nir_shader *nir,
       nir, lower_driver_param_to_ubo,
       nir_metadata_control_flow, ir3_const_state(v));
 
+   if (result) {
+      const struct ir3_const_state *const_state = ir3_const_state(v);
+
+      ir3_update_driver_ubo(nir, &const_state->primitive_map_ubo, "$primitive_map");
+      ir3_update_driver_ubo(nir, &const_state->primitive_param_ubo, "$primitive_param");
+      ir3_update_driver_ubo(nir, &const_state->driver_params_ubo, "$driver_params");
+   }
+
    return result;
 }