From be13c3ef9ffe536273d590d32c59363aca5290a2 Mon Sep 17 00:00:00 2001
From: Caio Oliveira <caio.oliveira@intel.com>
Date: Tue, 6 Feb 2024 11:04:25 -0800
Subject: [PATCH] iris: Add stage to iris_compiled_shader

Reviewed-by: Kenneth Graunke <kenneth@whitecape.org>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27646>
---
 src/gallium/drivers/iris/iris_context.h       |  3 +++
 src/gallium/drivers/iris/iris_program.c       | 12 ++++++++----
 src/gallium/drivers/iris/iris_program_cache.c |  8 ++++++--
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/gallium/drivers/iris/iris_context.h b/src/gallium/drivers/iris/iris_context.h
index 65443b596e1..17f8d331404 100644
--- a/src/gallium/drivers/iris/iris_context.h
+++ b/src/gallium/drivers/iris/iris_context.h
@@ -553,6 +553,8 @@ struct iris_compiled_shader {
 
    struct iris_binding_table bt;
 
+   gl_shader_stage stage;
+
    /**
     * Shader packets and other data derived from prog_data.  These must be
     * completely determined from prog_data.
@@ -1141,6 +1143,7 @@ struct iris_compiled_shader *iris_find_cached_shader(struct iris_context *ice,
 
 struct iris_compiled_shader *iris_create_shader_variant(const struct iris_screen *,
                                                         void *mem_ctx,
+                                                        gl_shader_stage stage,
                                                         enum iris_program_cache_id cache_id,
                                                         uint32_t key_size,
                                                         const void *key);
diff --git a/src/gallium/drivers/iris/iris_program.c b/src/gallium/drivers/iris/iris_program.c
index d8bb664456d..5b14a7b8253 100644
--- a/src/gallium/drivers/iris/iris_program.c
+++ b/src/gallium/drivers/iris/iris_program.c
@@ -1225,8 +1225,10 @@ find_or_add_variant(const struct iris_screen *screen,
       }
    }
 
+   gl_shader_stage stage = ish->nir->info.stage;
+
    if (variant == NULL) {
-      variant = iris_create_shader_variant(screen, NULL, cache_id,
+      variant = iris_create_shader_variant(screen, NULL, stage, cache_id,
                                            key_size, key);
 
       /* Append our new variant to the shader's variant list. */
@@ -1240,6 +1242,7 @@ find_or_add_variant(const struct iris_screen *screen,
       util_queue_fence_wait(&variant->ready);
    }
 
+   assert(stage == variant->stage);
    return variant;
 }
 
@@ -1589,6 +1592,7 @@ iris_update_compiled_tcs(struct iris_context *ice)
 
       if (shader == NULL) {
          shader = iris_create_shader_variant(screen, ice->shaders.cache,
+                                             MESA_SHADER_TESS_CTRL,
                                              IRIS_CACHE_TCS, sizeof(key), &key);
          added = true;
       }
@@ -2502,8 +2506,8 @@ iris_create_compute_state(struct pipe_context *ctx,
       struct iris_cs_prog_key key = { KEY_INIT(base) };
 
       struct iris_compiled_shader *shader =
-         iris_create_shader_variant(screen, NULL, IRIS_CACHE_CS,
-                                    sizeof(key), &key);
+         iris_create_shader_variant(screen, NULL, MESA_SHADER_COMPUTE,
+                                    IRIS_CACHE_CS, sizeof(key), &key);
 
       /* Append our new variant to the shader's variant list. */
       list_addtail(&shader->link, &ish->variants);
@@ -2717,7 +2721,7 @@ iris_create_shader_state(struct pipe_context *ctx,
       struct u_upload_mgr *uploader = ice->shaders.uploader_unsync;
 
       struct iris_compiled_shader *shader =
-         iris_create_shader_variant(screen, NULL,
+         iris_create_shader_variant(screen, NULL, info->stage,
                                     (enum iris_program_cache_id) info->stage,
                                     key_size, &key);
 
diff --git a/src/gallium/drivers/iris/iris_program_cache.c b/src/gallium/drivers/iris/iris_program_cache.c
index ab929419ffc..474fbedb935 100644
--- a/src/gallium/drivers/iris/iris_program_cache.c
+++ b/src/gallium/drivers/iris/iris_program_cache.c
@@ -108,6 +108,7 @@ iris_delete_shader_variant(struct iris_compiled_shader *shader)
 struct iris_compiled_shader *
 iris_create_shader_variant(const struct iris_screen *screen,
                            void *mem_ctx,
+                           gl_shader_stage stage,
                            enum iris_program_cache_id cache_id,
                            uint32_t key_size,
                            const void *key)
@@ -141,6 +142,8 @@ iris_create_shader_variant(const struct iris_screen *screen,
       memcpy(&shader->key, key, key_size);
    }
 
+   shader->stage = stage;
+
    return shader;
 }
 
@@ -236,8 +239,8 @@ iris_blorp_upload_shader(struct blorp_batch *blorp_batch, uint32_t stage,
    memset(&bt, 0, sizeof(bt));
 
    struct iris_compiled_shader *shader =
-      iris_create_shader_variant(screen, ice->shaders.cache, IRIS_CACHE_BLORP,
-                                 key_size, key);
+      iris_create_shader_variant(screen, ice->shaders.cache, stage,
+                                 IRIS_CACHE_BLORP, key_size, key);
 
    iris_finalize_program(shader, prog_data, NULL, NULL, 0, 0, 0, &bt);
 
@@ -413,6 +416,7 @@ iris_ensure_indirect_generation_shader(struct iris_batch *batch)
 
    struct iris_compiled_shader *shader =
       iris_create_shader_variant(screen, ice->shaders.cache,
+                                 MESA_SHADER_FRAGMENT,
                                  IRIS_CACHE_BLORP,
                                  sizeof(key), &key);
    iris_finalize_program(shader, &prog_data->base, NULL, NULL, 0, 0, 0, &bt);