nir: Add a nir_xfb_info to nir_shader

We want to be able to carry this along with the shader instead of always
having to re-generate it from scratch.  A new nir_gather_xfb_info()
helper is also added which, instead of returning it, adds it to the
shader.

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/nir/nir.h

@@ -3681,6 +3681,8 @@ typedef struct nir_shader {
    /** Size of the constant data associated with the shader, in bytes */
    unsigned constant_data_size;
 
+   struct nir_xfb_info *xfb_info;
+
    unsigned printf_info_count;
    nir_printf_info *printf_info;
 } nir_shader;

src/compiler/nir/nir_clone.c

@@ -23,6 +23,7 @@
 
 #include "nir.h"
 #include "nir_control_flow.h"
+#include "nir_xfb_info.h"
 
 /* Secret Decoder Ring:
  *   clone_foo():
@@ -785,6 +786,12 @@ nir_shader_clone(void *mem_ctx, const nir_shader *s)
       memcpy(ns->constant_data, s->constant_data, s->constant_data_size);
    }
 
+   if (s->xfb_info) {
+      size_t size = nir_xfb_info_size(s->xfb_info->output_count);
+      ns->xfb_info = ralloc_size(ns, size);
+      memcpy(ns->xfb_info, s->xfb_info, size);
+   }
+
    free_clone_state(&state);
 
    return ns;

src/compiler/nir/nir_gather_xfb_info.c

@@ -184,6 +184,13 @@ nir_shader_get_xfb_info(const nir_shader *shader, void *mem_ctx)
    return nir_gather_xfb_info_with_varyings(shader, mem_ctx, NULL);
 }
 
+void
+nir_shader_gather_xfb_info(nir_shader *shader)
+{
+   ralloc_free(shader->xfb_info);
+   shader->xfb_info = nir_gather_xfb_info_with_varyings(shader, shader, NULL);
+}
+
 nir_xfb_info *
 nir_gather_xfb_info_with_varyings(const nir_shader *shader,
                                   void *mem_ctx,

src/compiler/nir/nir_serialize.c

@@ -23,6 +23,7 @@
 
 #include "nir_serialize.h"
 #include "nir_control_flow.h"
+#include "nir_xfb_info.h"
 #include "util/u_dynarray.h"
 #include "util/u_math.h"
 
@@ -2050,6 +2051,32 @@ read_function(read_ctx *ctx)
       fxn->impl = NIR_SERIALIZE_FUNC_HAS_IMPL;
 }
 
+static void
+write_xfb_info(write_ctx *ctx, const nir_xfb_info *xfb)
+{
+   if (xfb == NULL) {
+      blob_write_uint32(ctx->blob, 0);
+   } else {
+      size_t size = nir_xfb_info_size(xfb->output_count);
+      assert(size <= UINT32_MAX);
+      blob_write_uint32(ctx->blob, size);
+      blob_write_bytes(ctx->blob, xfb, size);
+   }
+}
+
+static nir_xfb_info *
+read_xfb_info(read_ctx *ctx)
+{
+   uint32_t size = blob_read_uint32(ctx->blob);
+   if (size == 0)
+      return NULL;
+
+   struct nir_xfb_info *xfb = ralloc_size(ctx->nir, size);
+   blob_copy_bytes(ctx->blob, (void *)xfb, size);
+
+   return xfb;
+}
+
 /**
  * Serialize NIR into a binary blob.
  *
@@ -2104,6 +2131,8 @@ nir_serialize(struct blob *blob, const nir_shader *nir, bool strip)
    if (nir->constant_data_size > 0)
       blob_write_bytes(blob, nir->constant_data, nir->constant_data_size);
 
+   write_xfb_info(&ctx, nir->xfb_info);
+
    blob_overwrite_uint32(blob, idx_size_offset, ctx.next_idx);
 
    _mesa_hash_table_destroy(ctx.remap_table, NULL);
@@ -2159,6 +2188,8 @@ nir_deserialize(void *mem_ctx,
                       ctx.nir->constant_data_size);
    }
 
+   ctx.nir->xfb_info = read_xfb_info(&ctx);
+
    free(ctx.idx_table);
 
    nir_validate_shader(ctx.nir, "after deserialize");

src/compiler/nir/nir_sweep.c

@@ -163,6 +163,7 @@ nir_sweep(nir_shader *nir)
    assert(list_is_empty(&instr_gc_list));
 
    ralloc_steal(nir, nir->constant_data);
+   ralloc_steal(nir, nir->xfb_info);
    ralloc_steal(nir, nir->printf_info);
    for (int i = 0; i < nir->printf_info_count; i++) {
       ralloc_steal(nir, nir->printf_info[i].arg_sizes);

src/compiler/nir/nir_validate.c

@@ -26,6 +26,7 @@
  */
 
 #include "nir.h"
+#include "nir_xfb_info.h"
 #include "c11/threads.h"
 #include <assert.h>
 
@@ -1798,6 +1799,18 @@ nir_validate_shader(nir_shader *shader, const char *when)
       validate_function(func, &state);
    }
 
+   if (shader->xfb_info != NULL) {
+      /* At least validate that, if nir_shader::xfb_info exists, the shader
+       * has real transform feedback going on.
+       */
+      validate_assert(&state, shader->info.stage == MESA_SHADER_VERTEX ||
+                              shader->info.stage == MESA_SHADER_TESS_EVAL ||
+                              shader->info.stage == MESA_SHADER_GEOMETRY);
+      validate_assert(&state, shader->xfb_info->buffers_written != 0);
+      validate_assert(&state, shader->xfb_info->streams_written != 0);
+      validate_assert(&state, shader->xfb_info->output_count > 0);
+   }
+
    if (_mesa_hash_table_num_entries(state.errors) > 0)
       dump_errors(&state, when);
 

src/compiler/nir/nir_xfb_info.h

@@ -77,6 +77,8 @@ nir_xfb_info_size(uint16_t output_count)
 nir_xfb_info *
 nir_shader_get_xfb_info(const nir_shader *shader, void *mem_ctx);
 
+void nir_shader_gather_xfb_info(nir_shader *shader);
+
 nir_xfb_info *
 nir_gather_xfb_info_with_varyings(const nir_shader *shader,
                                   void *mem_ctx,