zink: more accurately update samplemask for fs shader keys

the fs samplemask needs to be updated on framebuffer rebind and on
fs bind to ensure that the key gets updated in time for the pipeline
change

fixes #5559

cc: mesa-stable

Reviewed-by: Dave Airlie <airlied@redhat.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/13531>
(cherry picked from commit c9ce151ff97b7a8776ed4cb2eff27bfbd836ea3d)

4 files changed, 24 insertions, 10 deletions

diff --git a/.pick_status.json b/.pick_status.json
index 2fb6bcb354d..6417feff873 100644
--- a/.pick_status.json
+++ b/.pick_status.json
@@ -130,7 +130,7 @@
         "description": "zink: more accurately update samplemask for fs shader keys",
         "nominated": true,
         "nomination_type": 0,
-        "resolution": 0,
+        "resolution": 1,
         "main_sha": null,
         "because_sha": null
     },
diff --git a/src/gallium/drivers/zink/zink_context.c b/src/gallium/drivers/zink/zink_context.c
index 476507f6212..5772762fb32 100644
--- a/src/gallium/drivers/zink/zink_context.c
+++ b/src/gallium/drivers/zink/zink_context.c
@@ -2573,15 +2573,8 @@ zink_set_framebuffer_state(struct pipe_context *pctx,
    update_framebuffer_state(ctx, w, h);
 
    uint8_t rast_samples = ctx->fb_state.samples - 1;
-   /* update the shader key if applicable:
-    * if gl_SampleMask[] is written to, we have to ensure that we get a shader with the same sample count:
-    * in GL, rast_samples==1 means ignore gl_SampleMask[]
-    * in VK, gl_SampleMask[] is never ignored
-    */
-   if (rast_samples != ctx->gfx_pipeline_state.rast_samples &&
-       (!ctx->gfx_stages[PIPE_SHADER_FRAGMENT] ||
-        ctx->gfx_stages[PIPE_SHADER_FRAGMENT]->nir->info.outputs_written & (1 << FRAG_RESULT_SAMPLE_MASK)))
-      zink_set_fs_key(ctx)->samples = ctx->fb_state.samples > 0;
+   if (rast_samples != ctx->gfx_pipeline_state.rast_samples)
+      zink_update_fs_key_samples(ctx);
    if (ctx->gfx_pipeline_state.rast_samples != rast_samples) {
       ctx->sample_locations_changed |= ctx->gfx_pipeline_state.sample_locations_enabled;
       ctx->gfx_pipeline_state.dirty = true;
diff --git a/src/gallium/drivers/zink/zink_program.c b/src/gallium/drivers/zink/zink_program.c
index a735f67a6d6..8c23d57851e 100644
--- a/src/gallium/drivers/zink/zink_program.c
+++ b/src/gallium/drivers/zink/zink_program.c
@@ -852,6 +852,23 @@ zink_bind_vs_state(struct pipe_context *pctx,
 
 }
 
+/* if gl_SampleMask[] is written to, we have to ensure that we get a shader with the same sample count:
+ * in GL, samples==1 means ignore gl_SampleMask[]
+ * in VK, gl_SampleMask[] is never ignored
+ */
+void
+zink_update_fs_key_samples(struct zink_context *ctx)
+{
+   if (!ctx->gfx_stages[PIPE_SHADER_FRAGMENT])
+      return;
+   nir_shader *nir = ctx->gfx_stages[PIPE_SHADER_FRAGMENT]->nir;
+   if (nir->info.outputs_written & (1 << FRAG_RESULT_SAMPLE_MASK)) {
+      bool samples = zink_get_fs_key(ctx)->samples;
+      if (samples != (ctx->fb_state.samples > 1))
+         zink_set_fs_key(ctx)->samples = ctx->fb_state.samples > 1;
+   }
+}
+
 static void
 zink_bind_fs_state(struct pipe_context *pctx,
                    void *cso)
@@ -869,6 +886,7 @@ zink_bind_fs_state(struct pipe_context *pctx,
                ctx->fbfetch_outputs |= BITFIELD_BIT(var->data.location - FRAG_RESULT_DATA0);
          }
       }
+      zink_update_fs_key_samples(ctx);
    }
    zink_update_fbfetch(ctx);
 }
diff --git a/src/gallium/drivers/zink/zink_program.h b/src/gallium/drivers/zink/zink_program.h
index 90f9310731c..3c1afe54f51 100644
--- a/src/gallium/drivers/zink/zink_program.h
+++ b/src/gallium/drivers/zink/zink_program.h
@@ -297,6 +297,9 @@ zink_get_fs_key(struct zink_context *ctx)
    return (const struct zink_fs_key *)&ctx->gfx_pipeline_state.shader_keys.key[PIPE_SHADER_FRAGMENT];
 }
 
+void
+zink_update_fs_key_samples(struct zink_context *ctx);
+
 static inline struct zink_vs_key *
 zink_set_vs_key(struct zink_context *ctx)
 {