glsl/nir: Populate nir_shader::xfb_info after linking varyings

Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Reviewed-by: Alyssa Rosenzweig <alyssa@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/16750>

src/compiler/glsl/gl_nir_link_varyings.c

@@ -3190,6 +3190,17 @@ gl_nir_link_varyings(const struct gl_constants *consts,
                prog->last_vert_prog->sh.LinkedTransformFeedback->NumVarying > 0;
          }
       }
+
+      /* Assign NIR XFB info to the last stage before the fragment shader */
+      for (int stage = MESA_SHADER_FRAGMENT - 1; stage >= 0; stage--) {
+         struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
+         if (sh && stage != MESA_SHADER_TESS_CTRL) {
+            sh->Program->nir->xfb_info =
+               gl_to_nir_xfb_info(sh->Program->sh.LinkedTransformFeedback,
+                                  sh->Program->nir);
+            break;
+         }
+      }
    }
 
    ralloc_free(mem_ctx);

src/compiler/glsl/gl_nir_link_xfb.c

@@ -194,3 +194,35 @@ gl_nir_link_assign_xfb_resources(const struct gl_constants *consts,
 
    ralloc_free(xfb_info);
 }
+
+struct nir_xfb_info *
+gl_to_nir_xfb_info(struct gl_transform_feedback_info *info, void *mem_ctx)
+{
+   if (info == NULL || info->NumOutputs == 0)
+      return NULL;
+
+   nir_xfb_info *xfb =
+      ralloc_size(mem_ctx, nir_xfb_info_size(info->NumOutputs));
+
+   xfb->output_count = info->NumOutputs;
+
+   for (unsigned i = 0; i < MAX_FEEDBACK_BUFFERS; i++) {
+      xfb->buffers[i].stride = info->Buffers[i].Stride;
+      xfb->buffers[i].varying_count = info->Buffers[i].NumVaryings;
+      xfb->buffer_to_stream[i] = info->Buffers[i].Stream;
+   }
+
+   for (unsigned i = 0; i < info->NumOutputs; i++) {
+      xfb->outputs[i].buffer = info->Outputs[i].OutputBuffer;
+      xfb->outputs[i].offset = info->Outputs[i].DstOffset * 4;
+      xfb->outputs[i].location = info->Outputs[i].OutputRegister;
+      xfb->outputs[i].component_offset = info->Outputs[i].ComponentOffset;
+      xfb->outputs[i].component_mask =
+         BITFIELD_RANGE(info->Outputs[i].ComponentOffset,
+                        info->Outputs[i].NumComponents);
+      xfb->buffers_written |= BITFIELD_BIT(info->Outputs[i].OutputBuffer);
+      xfb->streams_written |= BITFIELD_BIT(info->Outputs[i].StreamId);
+   }
+
+   return xfb;
+}

src/compiler/glsl/gl_nir_linker.h

@@ -38,7 +38,9 @@ struct gl_constants;
 struct gl_extensions;
 struct gl_linked_shader;
 struct gl_shader_program;
+struct gl_transform_feedback_info;
 struct xfb_decl;
+struct nir_xfb_info;
 
 struct gl_nir_linker_options {
    bool fill_parameters;
@@ -69,6 +71,9 @@ bool gl_nir_link_varyings(const struct gl_constants *consts,
                           const struct gl_extensions *exts,
                           gl_api api, struct gl_shader_program *prog);
 
+struct nir_xfb_info *
+gl_to_nir_xfb_info(struct gl_transform_feedback_info *info, void *mem_ctx);
+
 nir_variable * gl_nir_lower_xfb_varying(nir_shader *shader,
                                         const char *old_var_name,
                                         nir_variable *toplevel_var);

src/mesa/state_tracker/st_glsl_to_nir.cpp

@@ -44,7 +44,6 @@
 
 #include "compiler/nir/nir.h"
 #include "compiler/nir/nir_builder.h"
-#include "nir_xfb_info.h"
 #include "compiler/glsl_types.h"
 #include "compiler/glsl/glsl_to_nir.h"
 #include "compiler/glsl/gl_nir.h"
@@ -1047,41 +1046,6 @@ st_nir_lower_uniforms(struct st_context *st, nir_shader *nir)
                  !st->ctx->Const.NativeIntegers);
 }
 
-static nir_xfb_info *
-st_get_nir_xfb_info(struct gl_program *prog)
-{
-   struct gl_transform_feedback_info *info = prog->sh.LinkedTransformFeedback;
-   if (!info || !info->NumOutputs)
-      return NULL;
-
-   nir_xfb_info *xfb =
-      (nir_xfb_info *)calloc(1, nir_xfb_info_size(info->NumOutputs));
-   if (!xfb)
-      return NULL;
-
-   xfb->output_count = info->NumOutputs;
-
-   for (unsigned i = 0; i < MAX_FEEDBACK_BUFFERS; i++) {
-      xfb->buffers[i].stride = info->Buffers[i].Stride;
-      xfb->buffers[i].varying_count = info->Buffers[i].NumVaryings;
-      xfb->buffer_to_stream[i] = info->Buffers[i].Stream;
-   }
-
-   for (unsigned i = 0; i < info->NumOutputs; i++) {
-      xfb->outputs[i].buffer = info->Outputs[i].OutputBuffer;
-      xfb->outputs[i].offset = info->Outputs[i].DstOffset * 4;
-      xfb->outputs[i].location = info->Outputs[i].OutputRegister;
-      xfb->outputs[i].component_offset = info->Outputs[i].ComponentOffset;
-      xfb->outputs[i].component_mask =
-         BITFIELD_RANGE(info->Outputs[i].ComponentOffset,
-                        info->Outputs[i].NumComponents);
-      xfb->buffers_written |= BITFIELD_BIT(info->Outputs[i].OutputBuffer);
-      xfb->streams_written |= BITFIELD_BIT(info->Outputs[i].StreamId);
-   }
-
-   return xfb;
-}
-
 /* Last third of preparing nir from glsl, which happens after shader
  * variant lowering.
  */
@@ -1111,9 +1075,10 @@ st_finalize_nir(struct st_context *st, struct gl_program *prog,
     * This depends on st_nir_assign_varying_locations.
     */
    if (nir->options->lower_io_variables) {
-      nir_xfb_info *xfb = shader_program ? st_get_nir_xfb_info(prog) : NULL;
+      nir_xfb_info *xfb = shader_program ?
+         gl_to_nir_xfb_info(prog->sh.LinkedTransformFeedback, NULL) : NULL;
       nir_lower_io_passes(nir, xfb);
-      free(xfb);
+      ralloc_free(xfb);
    }
 
    /* Set num_uniforms in number of attribute slots (vec4s) */