panfrost: Plumb VS varying_layout in FS

This allows us to use LD_VAR_BUF instead of LD_VAR when the shaders are
linked together.

Signed-off-by: Lorenzo Rossi <lorenzo.rossi@collabora.com>
Reviewed-by: Christoph Pillmayer <christoph.pillmayer@arm.com>
Reviewed-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/40761>

src/gallium/drivers/panfrost/pan_context.h

@@ -346,6 +346,13 @@ struct panfrost_fs_key {
    uint8_t clip_plane_enable;
 
    bool line_smooth;
+
+   /* The VS varying layout determines how the FS is compiled (LD_VAR vs
+    * LD_VAR_BUF and byte offsets). Include the full layout in the key so
+    * the disk cache and in-memory variant cache correctly distinguish FS
+    * binaries compiled against different VS layouts.
+    */
+   struct pan_varying_layout vs_varying_layout;
 };
 
 struct panfrost_vs_key {
@@ -403,6 +410,9 @@ struct panfrost_uncompiled_shader {
    /* Stream output information */
    struct pipe_stream_output_info stream_output;
 
+   /* Varying layout (if known) */
+   struct pan_varying_layout vs_varying_layout;
+
    /** Lock for the variants array */
    simple_mtx_t lock;
 

src/gallium/drivers/panfrost/pan_shader.c

@@ -94,6 +94,7 @@ lower_sample_mask_writes(nir_builder *b, nir_intrinsic_instr *intrin,
 static void
 panfrost_shader_compile(struct panfrost_screen *screen, const nir_shader *ir,
                         struct util_debug_callback *dbg,
+                        const struct pan_varying_layout *varying_layout,
                         struct panfrost_shader_key *key, unsigned req_local_mem,
                         struct panfrost_shader_binary *out)
 {
@@ -133,15 +134,7 @@ panfrost_shader_compile(struct panfrost_screen *screen, const nir_shader *ir,
    /* nir_opt_varyings is replacing all flat highp types with float32, we need
     * to figure out the varying types ourselves */
    inputs.trust_varying_flat_highp_types = false;
-   struct pan_varying_layout varyings_layout;
-   /* TODO: wire up VS layout in FS when linked together */
-   if (s->info.stage == MESA_SHADER_VERTEX) {
-      pan_varying_collect_formats(&varyings_layout, s,
-                                  inputs.gpu_id,
-                                  inputs.trust_varying_flat_highp_types, false);
-      pan_build_varying_layout_compact(&varyings_layout, s, inputs.gpu_id);
-      inputs.varying_layout = &varyings_layout;
-   }
+   inputs.varying_layout = varying_layout;
 
    if (s->info.stage == MESA_SHADER_FRAGMENT) {
       if (key->fs.nr_cbufs_for_fragcolor) {
@@ -262,8 +255,12 @@ panfrost_shader_get(struct pipe_screen *pscreen,
     */
    if (!panfrost_disk_cache_retrieve(screen->disk_cache, uncompiled,
                                      &state->key, &res)) {
-      panfrost_shader_compile(screen, uncompiled->nir, dbg, &state->key,
-                              req_local_mem, &res);
+
+      const struct pan_varying_layout *varying_layout =
+         uncompiled->vs_varying_layout.known != 0
+            ? &uncompiled->vs_varying_layout : NULL;
+      panfrost_shader_compile(screen, uncompiled->nir, dbg, varying_layout,
+                              &state->key, req_local_mem, &res);
 
       panfrost_disk_cache_store(screen->disk_cache, uncompiled, &state->key,
                                 &res);
@@ -332,6 +329,8 @@ panfrost_build_fs_key(struct panfrost_context *ctx,
          key->line_smooth = rast->line_smooth;
    }
 
+   key->vs_varying_layout = uncompiled->vs_varying_layout;
+
    if (dev->arch <= 5) {
       u_foreach_bit(i, (nir->info.outputs_read >> FRAG_RESULT_DATA0)) {
          enum pipe_format fmt = PIPE_FORMAT_R8G8B8A8_UNORM;
@@ -485,6 +484,8 @@ panfrost_default_shader_key(struct panfrost_uncompiled_shader *so)
       */
       if (so->fragcolor_lowered)
          key.fs.nr_cbufs_for_fragcolor = 1;
+
+      key.fs.vs_varying_layout = so->vs_varying_layout;
    }
 
    return key;
@@ -574,6 +575,16 @@ panfrost_create_shader_state(struct pipe_context *pctx,
       so->noperspective_varyings =
          pan_nir_collect_noperspective_varyings_fs(nir);
 
+   if (nir->info.stage == MESA_SHADER_VERTEX) {
+      struct pan_varying_layout *varying_layout = &so->vs_varying_layout;
+      pan_varying_collect_formats(varying_layout, nir,
+                                  panfrost_device_gpu_id(dev),
+                                  false, /* trust_varying_flat_highp_types */
+                                  false /* lower_mediump */);
+      pan_build_varying_layout_compact(varying_layout, nir,
+                                       panfrost_device_gpu_id(dev));
+   }
+
    /* If this shader uses transform feedback, compile the transform
     * feedback program. This is a special shader variant.
     */
@@ -592,6 +603,14 @@ panfrost_create_shader_state(struct pipe_context *pctx,
       nir->info.has_transform_feedback_varyings = false;
    }
 
+   /* If we're not using separate shaders, the FS can use VS varying_layout to
+    * optimize loads (LD_VAR_BUF instead of LD_VAR).  Gallium won't provide us
+    * with the VS directly, so we need to delay the default variant compilation
+    * until link time
+    */
+   if (nir->info.stage == MESA_SHADER_FRAGMENT && !nir->info.separate_shader)
+      return so;
+
    /* Compile the program. We don't use vertex shader keys, so there will
     * be no further vertex shader variants. We do have fragment shader
     * keys, but we can still compile with a default key that will work most
@@ -632,6 +651,39 @@ panfrost_delete_shader_state(struct pipe_context *pctx, void *so)
    ralloc_free(so);
 }
 
+static void
+panfrost_link_shader(struct pipe_context *pctx, void** handles)
+{
+   struct panfrost_context *ctx = pan_context(pctx);
+   struct panfrost_uncompiled_shader *vs = handles[MESA_SHADER_VERTEX];
+   struct panfrost_uncompiled_shader *fs = handles[MESA_SHADER_FRAGMENT];
+
+   if (!fs || fs->nir->info.separate_shader)
+      return;
+
+   /* We only handle VS and FS for now, it's not clear how varying layout will
+    * fit when more shader types are supported.  So assert those are the only
+    * shaders present.
+    */
+   for (unsigned i = 0; i < MESA_SHADER_MESH_STAGES; i++) {
+      if (i != MESA_SHADER_VERTEX && i != MESA_SHADER_FRAGMENT)
+         assert(handles[i] == NULL);
+   }
+
+   /* Only copy the varying layout if we have a VS, sometimes we don't have one
+    * (e.g. fixed-function VS), in those cases we just compile a default FS.
+    */
+   if (vs)
+      fs->vs_varying_layout = vs->vs_varying_layout;
+
+   simple_mtx_lock(&fs->lock);
+
+   struct panfrost_shader_key key = panfrost_default_shader_key(fs);
+   panfrost_new_variant_locked(ctx, fs, &key);
+
+   simple_mtx_unlock(&fs->lock);
+}
+
 /*
  * Create a compute CSO. As compute kernels do not require variants, they are
  * precompiled, creating both the uncompiled and compiled shaders now.
@@ -698,6 +750,8 @@ panfrost_shader_context_init(struct pipe_context *pctx)
    pctx->delete_fs_state = panfrost_delete_shader_state;
    pctx->bind_fs_state = panfrost_bind_fs_state;
 
+   pctx->link_shader = panfrost_link_shader;
+
    pctx->create_compute_state = panfrost_create_compute_state;
    pctx->bind_compute_state = panfrost_bind_compute_state;
    pctx->get_compute_state_info = panfrost_get_compute_state_info;