nir: track variables representing ray queries

v2: Fix missing ray_query variable check (Caio)

Signed-off-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Reviewed-by: Caio Oliveira <caio.oliveira@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/13718>

src/compiler/nir/nir.h

@@ -425,6 +425,11 @@ typedef struct nir_variable {
       unsigned patch:1;
       unsigned invariant:1;
 
+      /**
+       * Is the variable a ray query?
+       */
+      unsigned ray_query:1;
+
      /**
        * Precision qualifier.
        *

src/compiler/nir/nir_gather_info.c

@@ -953,4 +953,22 @@ nir_shader_gather_info(nir_shader *shader, nir_function_impl *entrypoint)
          }
       }
    }
+
+   shader->info.ray_queries = 0;
+   nir_foreach_variable_in_shader(var, shader) {
+      if (!var->data.ray_query)
+         continue;
+
+      shader->info.ray_queries += MAX2(glsl_get_aoa_size(var->type), 1);
+   }
+   nir_foreach_function(func, shader) {
+      if (!func->impl)
+         continue;
+      nir_foreach_function_temp_variable(var, func->impl) {
+         if (!var->data.ray_query)
+            continue;
+
+         shader->info.ray_queries += MAX2(glsl_get_aoa_size(var->type), 1);
+      }
+   }
 }

src/compiler/nir/nir_print.c

@@ -494,8 +494,9 @@ print_var_decl(nir_variable *var, print_state *state)
    const char *const inv = (var->data.invariant) ? "invariant " : "";
    const char *const per_view = (var->data.per_view) ? "per_view " : "";
    const char *const per_primitive = (var->data.per_primitive) ? "per_primitive " : "";
-   fprintf(fp, "%s%s%s%s%s%s%s%s %s ",
-           bindless, cent, samp, patch, inv, per_view, per_primitive,
+   const char *const ray_query = (var->data.ray_query) ? "ray_query " : "";
+   fprintf(fp, "%s%s%s%s%s%s%s%s%s %s ",
+           bindless, cent, samp, patch, inv, per_view, per_primitive, ray_query,
            get_variable_mode_str(var->data.mode, false),
            glsl_interp_mode_name(var->data.interpolation));
 
@@ -1637,6 +1638,7 @@ nir_print_shader_annotated(nir_shader *shader, FILE *fp,
    if (shader->info.num_ubos)
       fprintf(fp, "ubos: %u\n", shader->info.num_ubos);
    fprintf(fp, "shared: %u\n", shader->info.shared_size);
+   fprintf(fp, "ray queries: %u\n", shader->info.ray_queries);
    if (shader->scratch_size)
       fprintf(fp, "scratch: %u\n", shader->scratch_size);
    if (shader->constant_data_size)

src/compiler/nir/nir_split_vars.c

@@ -136,6 +136,7 @@ init_field_for_type(struct field *field, struct field *parent,
       } else {
          field->var = nir_variable_create(state->shader, mode, var_type, name);
       }
+      field->var->data.ray_query = state->base_var->data.ray_query;
    }
 }
 
@@ -525,6 +526,7 @@ create_split_array_vars(struct array_var_info *var_info,
          split->var = nir_variable_create(shader, mode,
                                           var_info->split_var_type, name);
       }
+      split->var->data.ray_query = var_info->base_var->data.ray_query;
    } else {
       assert(var_info->levels[level].split);
       split->num_splits = var_info->levels[level].array_len;

src/compiler/shader_info.h

@@ -213,6 +213,12 @@ typedef struct shader_info {
     */
    unsigned shared_size;
 
+   /**
+    * Number of ray tracing queries in the shader (counts all elements of all
+    * variables).
+    */
+   unsigned ray_queries;
+
    /**
     * Local workgroup size used by compute/task/mesh shaders.
     */