Mesa (main): zink: do compute shader change on bind

[GitLab Mirror]

Module: Mesa
Branch: main
Commit: c39cbd49c16430f55ba1fc8c5e6b0e5ee9f84426
URL:    http://cgit.freedesktop.org/mesa/mesa/commit/?id=c39cbd49c16430f55ba1fc8c5e6b0e5ee9f84426

Author: Mike Blumenkrantz <michael.blumenkrantz at gmail.com>
Date:   Fri May 14 18:47:49 2021 -0400

zink: do compute shader change on bind

we can do this update earlier to optimize the actual compute path

Reviewed-by: Hoe Hao Cheng <haochengho12907 at gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/12428>

---

 src/gallium/drivers/zink/zink_draw.cpp  | 22 ++++++++--------------
 src/gallium/drivers/zink/zink_program.c | 13 ++++++++++---
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/src/gallium/drivers/zink/zink_draw.cpp b/src/gallium/drivers/zink/zink_draw.cpp
index 7ba687d5ed3..15ea1a4527c 100644
--- a/src/gallium/drivers/zink/zink_draw.cpp
+++ b/src/gallium/drivers/zink/zink_draw.cpp
@@ -177,22 +177,14 @@ zink_bind_vertex_buffers(struct zink_batch *batch, struct zink_context *ctx)
 static void
 update_compute_program(struct zink_context *ctx)
 {
-   unsigned bits = 1 << PIPE_SHADER_COMPUTE;
+   const unsigned bits = 1 << PIPE_SHADER_COMPUTE;
    if (ctx->dirty_shader_stages & bits) {
-      struct zink_compute_program *comp = NULL;
-      struct hash_entry *entry = _mesa_hash_table_search(ctx->compute_program_cache,
-                                                         ctx->compute_stage);
-      if (!entry) {
-         comp = zink_create_compute_program(ctx, ctx->compute_stage);
-         entry = _mesa_hash_table_insert(ctx->compute_program_cache, comp->shader, comp);
-      }
-      comp = (struct zink_compute_program*)(entry ? entry->data : NULL);
-      if (comp && comp != ctx->curr_compute) {
-         ctx->compute_pipeline_state.dirty = true;
-         zink_batch_reference_program(&ctx->batch, &comp->base);
-      }
+      struct zink_compute_program *comp = zink_create_compute_program(ctx, ctx->compute_stage);
+      _mesa_hash_table_insert(ctx->compute_program_cache, comp->shader, comp);
+      ctx->compute_pipeline_state.dirty = true;
       ctx->curr_compute = comp;
       ctx->dirty_shader_stages &= bits;
+      zink_batch_reference_program(&ctx->batch, &ctx->curr_compute->base);
    }
 }
 
@@ -786,8 +778,10 @@ zink_launch_grid(struct pipe_context *pctx, const struct pipe_grid_info *info)
    VkPipeline pipeline = zink_get_compute_pipeline(screen, ctx->curr_compute,
                                                &ctx->compute_pipeline_state);
 
-   if (BATCH_CHANGED)
+   if (BATCH_CHANGED) {
       zink_update_descriptor_refs(ctx, true);
+      zink_batch_reference_program(&ctx->batch, &ctx->curr_compute->base);
+   }
 
    if (prev_pipeline != pipeline || BATCH_CHANGED)
       vkCmdBindPipeline(batch->state->cmdbuf, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
diff --git a/src/gallium/drivers/zink/zink_program.c b/src/gallium/drivers/zink/zink_program.c
index c5163e80c5f..592e18f45d4 100644
--- a/src/gallium/drivers/zink/zink_program.c
+++ b/src/gallium/drivers/zink/zink_program.c
@@ -878,16 +878,23 @@ bind_stage(struct zink_context *ctx, enum pipe_shader_type stage,
            struct zink_shader *shader)
 {
    if (stage == PIPE_SHADER_COMPUTE) {
+      if (shader && shader != ctx->compute_stage) {
+         struct hash_entry *entry = _mesa_hash_table_search(ctx->compute_program_cache, shader);
+         if (entry) {
+            ctx->compute_pipeline_state.dirty = true;
+            ctx->curr_compute = entry->data;
+         } else
+            ctx->dirty_shader_stages |= 1 << stage;
+      }
       ctx->compute_stage = shader;
-      if (shader)
-         zink_select_launch_grid(ctx);
+      zink_select_launch_grid(ctx);
    } else {
       ctx->gfx_stages[stage] = shader;
       ctx->gfx_pipeline_state.combined_dirty = true;
       if (!shader)
          ctx->gfx_pipeline_state.modules[stage] = VK_NULL_HANDLE;
+      ctx->dirty_shader_stages |= 1 << stage;
    }
-   ctx->dirty_shader_stages |= 1 << stage;
    if (shader && shader->nir->info.num_inlinable_uniforms)
       ctx->shader_has_inlinable_uniforms_mask |= 1 << stage;
    else