Mesa (main): radv: Fill task shader info.

[GitLab Mirror]

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

Author: Timur Kristóf <timur.kristof at gmail.com>
Date:   Tue Jan 18 16:37:34 2022 +0100

radv: Fill task shader info.

Signed-off-by: Timur Kristóf <timur.kristof at gmail.com>
Reviewed-by: Rhys Perry <pendingchaos02 at gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/14929>

---

 src/amd/vulkan/radv_pipeline.c    | 10 ++++++++++
 src/amd/vulkan/radv_shader.c      |  2 +-
 src/amd/vulkan/radv_shader.h      |  1 +
 src/amd/vulkan/radv_shader_info.c | 16 ++++++++++++++++
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/amd/vulkan/radv_pipeline.c b/src/amd/vulkan/radv_pipeline.c
index bd23ec721e1..7df2f58717d 100644
--- a/src/amd/vulkan/radv_pipeline.c
+++ b/src/amd/vulkan/radv_pipeline.c
@@ -3494,6 +3494,16 @@ radv_declare_pipeline_args(struct radv_device *device, struct radv_pipeline_stag
       stages[i].info.user_sgprs_locs = stages[i].args.user_sgprs_locs;
       stages[i].info.inline_push_constant_mask = stages[i].args.ac.inline_push_const_mask;
    }
+
+   if (stages[MESA_SHADER_TASK].nir) {
+      /* Task/mesh I/O uses the task ring buffers. */
+      stages[MESA_SHADER_TASK].info.cs.uses_task_rings = true;
+      stages[MESA_SHADER_MESH].info.cs.uses_task_rings = true;
+
+      stages[MESA_SHADER_TASK].info.workgroup_size =
+         ac_compute_cs_workgroup_size(
+            stages[MESA_SHADER_TASK].nir->info.workgroup_size, false, UINT32_MAX);
+   }
 }
 
 static void
diff --git a/src/amd/vulkan/radv_shader.c b/src/amd/vulkan/radv_shader.c
index 460fa0e0b7f..1a3b5d1ef8b 100644
--- a/src/amd/vulkan/radv_shader.c
+++ b/src/amd/vulkan/radv_shader.c
@@ -2327,7 +2327,7 @@ radv_get_max_waves(const struct radv_device *device, struct radv_shader *shader,
       lds_per_wave =
          conf->lds_size * info->lds_encode_granularity + shader->info.ps.num_interp * 48;
       lds_per_wave = align(lds_per_wave, info->lds_alloc_granularity);
-   } else if (stage == MESA_SHADER_COMPUTE) {
+   } else if (stage == MESA_SHADER_COMPUTE || stage == MESA_SHADER_TASK) {
       unsigned max_workgroup_size = shader->info.workgroup_size;
       lds_per_wave =
          align(conf->lds_size * info->lds_encode_granularity, info->lds_alloc_granularity);
diff --git a/src/amd/vulkan/radv_shader.h b/src/amd/vulkan/radv_shader.h
index c158de552ff..f843776f102 100644
--- a/src/amd/vulkan/radv_shader.h
+++ b/src/amd/vulkan/radv_shader.h
@@ -345,6 +345,7 @@ struct radv_shader_info {
 
       bool uses_sbt;
       bool uses_ray_launch_size;
+      bool uses_task_rings;
    } cs;
    struct {
       uint64_t tes_inputs_read;
diff --git a/src/amd/vulkan/radv_shader_info.c b/src/amd/vulkan/radv_shader_info.c
index 0b9ca2ace14..adc30456f3a 100644
--- a/src/amd/vulkan/radv_shader_info.c
+++ b/src/amd/vulkan/radv_shader_info.c
@@ -590,9 +590,25 @@ radv_nir_shader_info_pass(struct radv_device *device, const struct nir_shader *n
                                         BITSET_TEST(nir->info.system_values_read, SYSTEM_VALUE_NUM_SUBGROUPS);
    switch (nir->info.stage) {
    case MESA_SHADER_COMPUTE:
+   case MESA_SHADER_TASK:
       for (int i = 0; i < 3; ++i)
          info->cs.block_size[i] = nir->info.workgroup_size[i];
       info->cs.uses_ray_launch_size = BITSET_TEST(nir->info.system_values_read, SYSTEM_VALUE_RAY_LAUNCH_SIZE);
+
+      /* Task shaders always need these for the I/O lowering even if
+       * the API shader doesn't actually use them.
+       */
+      if (nir->info.stage == MESA_SHADER_TASK) {
+         /* Needed to address the IB to read firstTask. */
+         info->vs.needs_draw_id |=
+            BITSET_TEST(nir->info.system_values_read, SYSTEM_VALUE_WORKGROUP_ID);
+
+         /* Needed to address the task draw/payload rings. */
+         info->cs.uses_block_id[0] = true;
+
+         /* Needed for storing draw ready only on the 1st thread. */
+         info->cs.uses_local_invocation_idx = true;
+      }
       break;
    case MESA_SHADER_FRAGMENT:
       info->ps.can_discard = nir->info.fs.uses_discard;