nir: add a helper for calculating variable slots

this will maybe avoid future bugs, but probably not

Reviewed-by: Connor Abbott <cwabbott0@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24163>

src/compiler/nir/nir.h

@@ -6190,6 +6190,14 @@ nir_variable_is_in_block(const nir_variable *var)
    return nir_variable_is_in_ubo(var) || nir_variable_is_in_ssbo(var);
 }
 
+static inline unsigned
+nir_variable_count_slots(const nir_variable *var, const struct glsl_type *type)
+{
+   return var->data.compact ?
+          DIV_ROUND_UP(var->data.location_frac + glsl_get_length(type), 4) :
+          glsl_count_attribute_slots(type, false);
+}
+
 /* See default_ub_config in nir_range_analysis.c for documentation. */
 typedef struct nir_unsigned_upper_bound_config {
    unsigned min_subgroup_size;

src/compiler/nir/nir_gather_info.c

@@ -210,10 +210,7 @@ mark_whole_variable(nir_shader *shader, nir_variable *var,
       type = glsl_get_array_element(type);
    }
 
-   const unsigned slots =
-      var->data.compact ? DIV_ROUND_UP(var->data.location_frac + glsl_get_length(type), 4)
-                        : glsl_count_attribute_slots(type, false);
-
+   const unsigned slots = nir_variable_count_slots(var, type);
    set_io_mask(shader, var, 0, slots, deref, is_output_read);
 }
 
@@ -287,10 +284,7 @@ try_mask_partial_io(nir_shader *shader, nir_variable *var,
    if (offset == -1)
       return false;
 
-   const unsigned slots =
-      var->data.compact ? DIV_ROUND_UP(var->data.location_frac + glsl_get_length(type), 4)
-                        : glsl_count_attribute_slots(type, false);
-
+   const unsigned slots = nir_variable_count_slots(var, type);
    if (offset >= slots) {
       /* Constant index outside the bounds of the matrix/array.  This could
        * arise as a result of constant folding of a legal GLSL program.

src/freedreno/vulkan/tu_shader.cc

@@ -923,9 +923,7 @@ tu_gather_xfb_info(nir_shader *nir, struct ir3_stream_output_info *info)
    memset(output_map, 0, sizeof(output_map));
 
    nir_foreach_shader_out_variable(var, nir) {
-      unsigned slots =
-         var->data.compact ? DIV_ROUND_UP(var->data.location_frac + glsl_get_length(var->type), 4)
-                           : glsl_count_attribute_slots(var->type, false);
+      unsigned slots = nir_variable_count_slots(var, var->type);
       for (unsigned i = 0; i < slots; i++)
          output_map[var->data.location + i] = var->data.driver_location + i;
    }

src/gallium/auxiliary/nir/nir_to_tgsi_info.c

@@ -335,8 +335,7 @@ void nir_tgsi_scan_shader(const struct nir_shader *nir,
          type = glsl_get_array_element(type);
       }
 
-      unsigned attrib_count = variable->data.compact ? DIV_ROUND_UP(variable->data.location_frac + glsl_get_length(type), 4) :
-         glsl_count_attribute_slots(type, nir->info.stage == MESA_SHADER_VERTEX);
+      unsigned attrib_count = nir_variable_count_slots(variable, type);
 
       i = variable->data.driver_location;
 
@@ -435,8 +434,7 @@ void nir_tgsi_scan_shader(const struct nir_shader *nir,
          type = glsl_get_array_element(type);
       }
 
-      unsigned attrib_count = variable->data.compact ? DIV_ROUND_UP(variable->data.location_frac + glsl_get_length(type), 4) :
-         glsl_count_attribute_slots(type, false);
+      unsigned attrib_count = nir_variable_count_slots(variable, type);
       for (unsigned k = 0; k < attrib_count; k++, i++) {
 
          if (nir->info.stage == MESA_SHADER_FRAGMENT) {

src/gallium/frontends/lavapipe/lvp_pipeline.c

@@ -544,8 +544,7 @@ lvp_shader_xfb_init(struct lvp_shader *shader)
       memset(output_mapping, 0, sizeof(output_mapping));
 
       nir_foreach_shader_out_variable(var, shader->pipeline_nir->nir) {
-         unsigned slots = var->data.compact ? DIV_ROUND_UP(var->data.location_frac + glsl_get_length(var->type), 4)
-                                            : glsl_count_attribute_slots(var->type, false);
+         unsigned slots = nir_variable_count_slots(var, var->type);
          for (unsigned i = 0; i < slots; i++)
             output_mapping[var->data.location + i] = var->data.driver_location + i;
       }

src/intel/compiler/brw_fs_nir.cpp

@@ -71,9 +71,7 @@ fs_visitor::nir_setup_outputs()
     */
    nir_foreach_shader_out_variable(var, nir) {
       const int loc = var->data.driver_location;
-      const unsigned var_vec4s =
-         var->data.compact ? DIV_ROUND_UP(var->data.location_frac + glsl_get_length(var->type), 4)
-                           : type_size_vec4(var->type, true);
+      const unsigned var_vec4s = nir_variable_count_slots(var, var->type);
       vec4s[loc] = MAX2(vec4s[loc], var_vec4s);
    }