nir/linker: use nir_gather_xfb_info

Instead of a custom ARB_gl_spirv xfb gather info pass.

In fact, this is not only about reusing code, but the current custom
code was not handling properly how many varyings are enumerated from
some complex types. So this change is also about fixing some corner
cases.

v2: Use util_bitcount, simplify current stage check (Kenneth)

Reviewed-by: Timothy Arceri <tarceri@itsqueeze.com>

src/compiler/glsl/gl_nir_link_xfb.c

@@ -22,10 +22,11 @@
  */
 
 #include "nir.h"
+#include "nir_xfb_info.h"
 #include "gl_nir_linker.h"
-#include "ir_uniform.h" /* for gl_uniform_storage */
 #include "linker_util.h"
 #include "main/context.h"
+#include "util/u_math.h"
 
 /*
  * This file does the linking of GLSL transform feedback using NIR.
@@ -34,160 +35,6 @@
  * particularities.
  */
 
-struct active_xfb_buffer {
-   GLuint stride;
-   GLuint num_varyings;
-};
-
-struct active_xfb_varyings {
-   unsigned num_varyings;
-   unsigned num_outputs;
-   unsigned buffer_size;
-   struct nir_variable **varyings;
-   struct active_xfb_buffer buffers[MAX_FEEDBACK_BUFFERS];
-};
-
-static unsigned
-get_num_outputs(nir_variable *var)
-{
-   return glsl_count_attribute_slots(var->type,
-                                     false /* is_vertex_input */);
-}
-
-static void
-add_xfb_varying(struct active_xfb_varyings *active_varyings,
-                nir_variable *var)
-{
-   if (active_varyings->num_varyings >= active_varyings->buffer_size) {
-      if (active_varyings->buffer_size == 0)
-         active_varyings->buffer_size = 1;
-      else
-         active_varyings->buffer_size *= 2;
-
-      active_varyings->varyings = realloc(active_varyings->varyings,
-                                          sizeof(nir_variable*) *
-                                          active_varyings->buffer_size);
-   }
-
-   active_varyings->varyings[active_varyings->num_varyings++] = var;
-
-   active_varyings->num_outputs += get_num_outputs(var);
-}
-
-static int
-cmp_xfb_offset(const void *x_generic, const void *y_generic)
-{
-   const nir_variable *const *x = x_generic;
-   const nir_variable *const *y = y_generic;
-
-   if ((*x)->data.xfb_buffer != (*y)->data.xfb_buffer)
-      return (*x)->data.xfb_buffer - (*y)->data.xfb_buffer;
-   return (*x)->data.offset - (*y)->data.offset;
-}
-
-static void
-get_active_xfb_varyings(struct gl_shader_program *prog,
-                        struct active_xfb_varyings *active_varyings)
-{
-   for (unsigned i = 0; i < MESA_SHADER_STAGES; ++i) {
-      struct gl_linked_shader *sh = prog->_LinkedShaders[i];
-      if (sh == NULL)
-         continue;
-
-      nir_shader *nir = sh->Program->nir;
-
-      nir_foreach_variable(var, &nir->outputs) {
-         if (var->data.explicit_xfb_buffer &&
-             var->data.explicit_xfb_stride) {
-            assert(var->data.xfb_buffer < MAX_FEEDBACK_BUFFERS);
-            active_varyings->buffers[var->data.xfb_buffer].stride =
-               var->data.xfb_stride;
-         }
-
-         if (!var->data.explicit_xfb_buffer ||
-             !var->data.explicit_offset)
-            continue;
-
-         active_varyings->buffers[var->data.xfb_buffer].num_varyings++;
-
-         add_xfb_varying(active_varyings, var);
-      }
-   }
-
-   /* The xfb_offset qualifier does not have to be used in increasing order
-    * however some drivers expect to receive the list of transform feedback
-    * declarations in order so sort it now for convenience.
-    */
-   qsort(active_varyings->varyings,
-         active_varyings->num_varyings,
-         sizeof(*active_varyings->varyings),
-         cmp_xfb_offset);
-}
-
-static unsigned
-add_varying_outputs(nir_variable *var,
-                    const struct glsl_type *type,
-                    unsigned location_offset,
-                    unsigned dest_offset,
-                    struct gl_transform_feedback_output *output)
-{
-   unsigned num_outputs = 0;
-
-   if (glsl_type_is_array(type) || glsl_type_is_matrix(type)) {
-      unsigned length = glsl_get_length(type);
-      const struct glsl_type *child_type = glsl_get_array_element(type);
-      unsigned component_slots = glsl_get_component_slots(child_type);
-
-      for (unsigned i = 0; i < length; i++) {
-         unsigned child_outputs = add_varying_outputs(var,
-                                                      child_type,
-                                                      location_offset,
-                                                      dest_offset,
-                                                      output + num_outputs);
-         num_outputs += child_outputs;
-         location_offset += child_outputs;
-         dest_offset += component_slots;
-      }
-   } else if (glsl_type_is_struct_or_ifc(type)) {
-      unsigned length = glsl_get_length(type);
-      for (unsigned i = 0; i < length; i++) {
-         const struct glsl_type *child_type = glsl_get_struct_field(type, i);
-         unsigned child_outputs = add_varying_outputs(var,
-                                                      child_type,
-                                                      location_offset,
-                                                      dest_offset,
-                                                      output + num_outputs);
-         num_outputs += child_outputs;
-         location_offset += child_outputs;
-         dest_offset += glsl_get_component_slots(child_type);
-      }
-   } else {
-      unsigned location = var->data.location + location_offset;
-      unsigned location_frac = var->data.location_frac;
-      unsigned num_components = glsl_get_component_slots(type);
-
-      while (num_components > 0) {
-         unsigned output_size = MIN2(num_components, 4 - location_frac);
-
-         output->OutputRegister = location;
-         output->OutputBuffer = var->data.xfb_buffer;
-         output->NumComponents = output_size;
-         output->StreamId = var->data.stream;
-         output->DstOffset = var->data.offset / 4 + dest_offset;
-         output->ComponentOffset = location_frac;
-
-         dest_offset += output_size;
-         num_components -= output_size;
-         num_outputs++;
-         output++;
-         location++;
-         location_frac = 0;
-      }
-   }
-
-   return num_outputs;
-}
-
 void
 gl_nir_link_assign_xfb_resources(struct gl_context *ctx,
                                  struct gl_shader_program *prog)
@@ -220,36 +67,51 @@ gl_nir_link_assign_xfb_resources(struct gl_context *ctx,
       free(prog->TransformFeedback.VaryingNames[i]);
    free(prog->TransformFeedback.VaryingNames);
 
-   struct active_xfb_varyings active_varyings = { 0 };
+   nir_xfb_info *xfb_info = NULL;
 
-   get_active_xfb_varyings(prog, &active_varyings);
+   /* Find last stage before fragment shader */
+   for (int stage = MESA_SHADER_FRAGMENT - 1; stage >= 0; stage--) {
+      struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
 
-   for (unsigned buf = 0; buf < MAX_FEEDBACK_BUFFERS; buf++)
-      prog->TransformFeedback.BufferStride[buf] = active_varyings.buffers[buf].stride;
-
-   prog->TransformFeedback.NumVarying = active_varyings.num_varyings;
-   prog->TransformFeedback.VaryingNames =
-      malloc(sizeof(GLchar *) * active_varyings.num_varyings);
+      if (sh && stage != MESA_SHADER_TESS_CTRL) {
+         xfb_info = nir_gather_xfb_info(sh->Program->nir, NULL);
+         break;
+      }
+   }
 
    struct gl_transform_feedback_info *linked_xfb =
       rzalloc(xfb_prog, struct gl_transform_feedback_info);
    xfb_prog->sh.LinkedTransformFeedback = linked_xfb;
 
+   if (!xfb_info) {
+      prog->TransformFeedback.NumVarying = 0;
+      linked_xfb->NumOutputs = 0;
+      linked_xfb->NumVarying = 0;
+      linked_xfb->ActiveBuffers = 0;
+      return;
+   }
+
+   for (unsigned buf = 0; buf < MAX_FEEDBACK_BUFFERS; buf++)
+      prog->TransformFeedback.BufferStride[buf] = xfb_info->buffers[buf].stride;
+
+   prog->TransformFeedback.NumVarying = xfb_info->varying_count;
+   prog->TransformFeedback.VaryingNames =
+      malloc(sizeof(GLchar *) * xfb_info->varying_count);
+
    linked_xfb->Outputs =
       rzalloc_array(xfb_prog,
                     struct gl_transform_feedback_output,
-                    active_varyings.num_outputs);
-   linked_xfb->NumOutputs = active_varyings.num_outputs;
+                    xfb_info->output_count);
+   linked_xfb->NumOutputs = xfb_info->output_count;
 
    linked_xfb->Varyings =
       rzalloc_array(xfb_prog,
                     struct gl_transform_feedback_varying_info,
-                    active_varyings.num_varyings);
-   linked_xfb->NumVarying = active_varyings.num_varyings;
+                    xfb_info->varying_count);
+   linked_xfb->NumVarying = xfb_info->varying_count;
 
-   struct gl_transform_feedback_output *output = linked_xfb->Outputs;
-   for (unsigned i = 0; i < active_varyings.num_varyings; i++) {
-      struct nir_variable *var = active_varyings.varyings[i];
+   for (unsigned i = 0; i < xfb_info->varying_count; i++) {
+      nir_xfb_varying_info *xfb_varying = &xfb_info->varyings[i];
 
       /* From ARB_gl_spirv spec:
        *
@@ -277,23 +139,29 @@ gl_nir_link_assign_xfb_resources(struct gl_context *ctx,
        */
       prog->TransformFeedback.VaryingNames[i] = NULL;
 
-      unsigned varying_outputs = add_varying_outputs(var,
-                                                     var->type,
-                                                     0, /* location_offset */
-                                                     0, /* dest_offset */
-                                                     output);
-      assert(varying_outputs == get_num_outputs(var));
-      output = output + varying_outputs;
-
       struct gl_transform_feedback_varying_info *varying =
          linked_xfb->Varyings + i;
 
       /* ARB_gl_spirv: see above. */
       varying->Name = NULL;
-      varying->Type = glsl_get_gl_type(var->type);
-      varying->BufferIndex = var->data.xfb_buffer;
-      varying->Size = glsl_get_length(var->type);
-      varying->Offset = var->data.offset;
+      varying->Type = glsl_get_gl_type(xfb_varying->type);
+      varying->BufferIndex = xfb_varying->buffer;
+      varying->Size = glsl_get_length(xfb_varying->type);
+      varying->Offset = xfb_varying->offset;
+   }
+
+   for (unsigned i = 0; i < xfb_info->output_count; i++) {
+      nir_xfb_output_info *xfb_output = &xfb_info->outputs[i];
+
+      struct gl_transform_feedback_output *output =
+         linked_xfb->Outputs + i;
+
+      output->OutputRegister = xfb_output->location;
+      output->OutputBuffer = xfb_output->buffer;
+      output->NumComponents = util_bitcount(xfb_output->component_mask);
+      output->StreamId = xfb_info->buffer_to_stream[xfb_output->buffer];
+      output->DstOffset = xfb_output->offset / 4;
+      output->ComponentOffset = xfb_output->component_offset;
    }
 
    /* Make sure MaxTransformFeedbackBuffers is <= 32 so the bitmask for
@@ -303,14 +171,14 @@ gl_nir_link_assign_xfb_resources(struct gl_context *ctx,
    assert(ctx->Const.MaxTransformFeedbackBuffers <= sizeof(buffers) * 8);
 
    for (unsigned buf = 0; buf < MAX_FEEDBACK_BUFFERS; buf++) {
-      if (active_varyings.buffers[buf].stride > 0) {
-         linked_xfb->Buffers[buf].Stride = active_varyings.buffers[buf].stride / 4;
-         linked_xfb->Buffers[buf].NumVaryings = active_varyings.buffers[buf].num_varyings;
+      if (xfb_info->buffers[buf].stride > 0) {
+         linked_xfb->Buffers[buf].Stride = xfb_info->buffers[buf].stride / 4;
+         linked_xfb->Buffers[buf].NumVaryings = xfb_info->buffers[buf].varying_count;
          buffers |= 1 << buf;
       }
    }
 
    linked_xfb->ActiveBuffers = buffers;
 
-   free(active_varyings.varyings);
+   ralloc_free(xfb_info);
 }