nir: split gathering array sizes from nir_lower_clip_cull_distance_array_vars

nir_lower_clip_cull_distance_array_vars was sneakily updating
shader_info::clip/cull_distance_array_size.

This moves the gathering into a new function
nir_gather_clip_cull_distance_sizes_from_vars.

v2: remove assertions that prevented nir_lower_clip_cull_distance_array_vars
    from being used with non-compact arrays

Reviewed-by: Alyssa Rosenzweig <alyssa.rosenzweig@intel.com> (v1)
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/38465>

src/amd/vulkan/radv_shader.c

@@ -603,6 +603,7 @@ radv_shader_spirv_to_nir(struct radv_device *device, const struct radv_shader_st
 
       NIR_PASS(_, nir, nir_propagate_invariant, pdev->cache_key.invariant_geom);
 
+      nir_gather_clip_cull_distance_sizes_from_vars(nir);
       NIR_PASS(_, nir, nir_lower_clip_cull_distance_array_vars);
 
       if (nir->info.stage == MESA_SHADER_VERTEX || nir->info.stage == MESA_SHADER_TESS_EVAL ||

src/compiler/glsl/gl_nir_linker.c

@@ -1428,15 +1428,15 @@ prelink_lowering(const struct pipe_screen *screen,
       opt_access_options.is_vulkan = false;
       NIR_PASS(_, nir, nir_opt_access, &opt_access_options);
 
+      /* This must be done before calling nir_lower_clip_cull_distance_to_vec4s. */
+      nir_gather_clip_cull_distance_sizes_from_vars(nir);
+
       if (!nir->options->compact_arrays) {
          NIR_PASS(_, nir, nir_lower_clip_cull_distance_to_vec4s);
          NIR_PASS(_, nir, nir_lower_tess_level_array_vars_to_vec);
       }
 
-      /* Combine clip and cull outputs into one array and set:
-       * - shader_info::clip_distance_array_size
-       * - shader_info::cull_distance_array_size
-       */
+      /* Combine clip and cull outputs into one array. */
       NIR_PASS(_, nir, nir_lower_clip_cull_distance_array_vars);
    }
 

src/compiler/nir/nir.h

@@ -5192,6 +5192,7 @@ bool nir_lower_scratch_to_var(nir_shader *nir);
 bool nir_lower_clip_halfz(nir_shader *shader);
 
 void nir_shader_gather_info(nir_shader *shader, nir_function_impl *entrypoint);
+void nir_gather_clip_cull_distance_sizes_from_vars(nir_shader *nir);
 
 void nir_gather_types(nir_function_impl *impl,
                       BITSET_WORD *float_types,

src/compiler/nir/nir_lower_clip_cull_distance_array_vars.c

@@ -26,7 +26,7 @@
 #include "nir_deref.h"
 
 /**
- * This file contains two different lowering passes.
+ * This file contains multiple functions.
  *
  * 1. nir_lower_clip_cull_distance_array_vars()
  *
@@ -52,6 +52,10 @@
  *    is translated into:
  *
  *      gl_ClipDistanceMESA[i>>2][i&3]
+ *
+ * 3. nir_gather_clip_cull_distance_sizes_from_vars
+ *
+ *    It does what it says.
  */
 
 #define GLSL_CLIP_VAR_NAME "gl_ClipDistanceMESA"
@@ -432,9 +436,7 @@ nir_lower_clip_cull_distance_to_vec4s(nir_shader *shader)
 }
 
 static bool
-combine_clip_cull(nir_shader *nir,
-                  nir_variable_mode mode,
-                  bool store_info)
+combine_clip_cull(nir_shader *nir, nir_variable_mode mode)
 {
    nir_variable *cull = NULL;
    nir_variable *clip = NULL;
@@ -447,45 +449,18 @@ combine_clip_cull(nir_shader *nir,
          cull = var;
    }
 
-   if (!cull && !clip) {
-      /* If this is run after optimizations and the variables have been
-       * eliminated, we should update the shader info, because no other
-       * place does that.
-       */
-      if (store_info) {
-         nir->info.clip_distance_array_size = 0;
-         nir->info.cull_distance_array_size = 0;
-      }
+   if (!cull && !clip)
       return false;
-   }
 
-   if (!cull && clip) {
-      /* The GLSL IR lowering pass must have converted these to vectors */
-      if (!clip->data.compact)
-         return false;
+   /* This can't be invoked more than once. */
+   assert(!clip || clip->data.how_declared != nir_var_hidden);
 
-      /* If this pass has already run, don't repeat.  We would think that
-       * the combined clip/cull distance array was clip-only and mess up.
-       */
-      if (clip->data.how_declared == nir_var_hidden)
-         return false;
-   }
-
-   const unsigned clip_array_size = get_unwrapped_array_length(nir, clip);
-   const unsigned cull_array_size = get_unwrapped_array_length(nir, cull);
-
-   if (store_info) {
-      nir->info.clip_distance_array_size = clip_array_size;
-      nir->info.cull_distance_array_size = cull_array_size;
-   }
-
-   if (clip) {
-      assert(clip->data.compact);
+   if (clip)
       clip->data.how_declared = nir_var_hidden;
-   }
 
    if (cull) {
-      assert(cull->data.compact);
+      const unsigned clip_array_size = get_unwrapped_array_length(nir, clip);
+
       cull->data.how_declared = nir_var_hidden;
       cull->data.location = VARYING_SLOT_CLIP_DIST0 + clip_array_size / 4;
       cull->data.location_frac = clip_array_size % 4;
@@ -501,13 +476,11 @@ nir_lower_clip_cull_distance_array_vars(nir_shader *nir)
 
    if (nir->info.stage <= MESA_SHADER_GEOMETRY ||
        nir->info.stage == MESA_SHADER_MESH)
-      progress |= combine_clip_cull(nir, nir_var_shader_out, true);
+      progress |= combine_clip_cull(nir, nir_var_shader_out);
 
    if (nir->info.stage > MESA_SHADER_VERTEX &&
-       nir->info.stage <= MESA_SHADER_FRAGMENT) {
-      progress |= combine_clip_cull(nir, nir_var_shader_in,
-                                    nir->info.stage == MESA_SHADER_FRAGMENT);
-   }
+       nir->info.stage <= MESA_SHADER_FRAGMENT)
+      progress |= combine_clip_cull(nir, nir_var_shader_in);
 
    nir_foreach_function_impl(impl, nir) {
       nir_progress(progress, impl,
@@ -516,3 +489,39 @@ nir_lower_clip_cull_distance_array_vars(nir_shader *nir)
 
    return progress;
 }
+
+/* This initializes shader_info::clip/cull_distance_array_size for the first
+ * time. This should be called before any lowering.
+ */
+void
+nir_gather_clip_cull_distance_sizes_from_vars(nir_shader *nir)
+{
+   nir_variable_mode mode = 0;
+
+   if (nir->info.stage <= MESA_SHADER_GEOMETRY ||
+       nir->info.stage == MESA_SHADER_MESH)
+      mode = nir_var_shader_out;
+   else if (nir->info.stage == MESA_SHADER_FRAGMENT)
+      mode = nir_var_shader_in;
+
+   if (mode) {
+      nir_variable *cull = NULL, *clip = NULL;
+
+      nir_foreach_variable_with_modes(var, nir, mode) {
+         if (var->data.location == VARYING_SLOT_CLIP_DIST0)
+            clip = var;
+         else if (var->data.location == VARYING_SLOT_CULL_DIST0)
+            cull = var;
+      }
+
+      /* The arrays must be "compact". */
+      assert(!clip || clip->data.compact);
+      assert(!cull || cull->data.compact);
+
+      /* nir_lower_clip_cull_distance_array_vars must not have been called. */
+      assert(!clip || clip->data.how_declared != nir_var_hidden);
+
+      nir->info.clip_distance_array_size = get_unwrapped_array_length(nir, clip);
+      nir->info.cull_distance_array_size = get_unwrapped_array_length(nir, cull);
+   }
+}

src/microsoft/spirv_to_dxil/dxil_spirv_nir.c

@@ -1028,6 +1028,7 @@ dxil_spirv_nir_passes(nir_shader *nir,
 
    NIR_PASS(_, nir, dxil_nir_lower_int_cubemaps, false);
 
+   nir_gather_clip_cull_distance_sizes_from_vars(nir);
    NIR_PASS(_, nir, nir_lower_clip_cull_distance_array_vars);
    NIR_PASS(_, nir, nir_lower_io_vars_to_temporaries, nir_shader_get_entrypoint(nir),
             nir_var_shader_out | nir_var_shader_in);

src/vulkan/runtime/vk_nir.c

@@ -195,6 +195,7 @@ vk_spirv_to_nir(struct vk_device *device,
     * insert dead clip/cull vars and we don't want to clip/cull based on
     * uninitialized garbage.
     */
+   nir_gather_clip_cull_distance_sizes_from_vars(nir);
    NIR_PASS(_, nir, nir_lower_clip_cull_distance_array_vars);
 
    if (nir->info.stage == MESA_SHADER_VERTEX ||