Module: Mesa Branch: main Commit: 9a3b902cacb26a46a395c1f0410f7432d59252a0 URL: http://cgit.freedesktop.org/mesa/mesa/commit/?id=9a3b902cacb26a46a395c1f0410f7432d59252a0
Author: Samuel Pitoiset <[email protected]> Date: Thu Sep 21 13:53:59 2023 +0200 ac/nir: add lowering for task shader queries Signed-off-by: Samuel Pitoiset <[email protected]> Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/25331> --- src/amd/common/ac_nir.h | 3 ++- src/amd/common/ac_nir_lower_taskmesh_io_to_mem.c | 26 +++++++++++++++++++++++- src/amd/vulkan/nir/radv_nir_lower_io.c | 4 ++-- 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/src/amd/common/ac_nir.h b/src/amd/common/ac_nir.h index 7c560e3b563..85720a2ba3f 100644 --- a/src/amd/common/ac_nir.h +++ b/src/amd/common/ac_nir.h @@ -201,7 +201,8 @@ ac_nir_lower_ngg_ms(nir_shader *shader, void ac_nir_lower_task_outputs_to_mem(nir_shader *shader, unsigned task_payload_entry_bytes, - unsigned task_num_entries); + unsigned task_num_entries, + bool has_query); void ac_nir_lower_mesh_inputs_to_mem(nir_shader *shader, diff --git a/src/amd/common/ac_nir_lower_taskmesh_io_to_mem.c b/src/amd/common/ac_nir_lower_taskmesh_io_to_mem.c index 334e889a99f..9c81ca8fe30 100644 --- a/src/amd/common/ac_nir_lower_taskmesh_io_to_mem.c +++ b/src/amd/common/ac_nir_lower_taskmesh_io_to_mem.c @@ -20,6 +20,9 @@ typedef struct { unsigned payload_entry_bytes; unsigned draw_entry_bytes; unsigned num_entries; + + /* True if the lowering needs to insert shader query. */ + bool has_query; } lower_tsms_io_state; static nir_def * @@ -139,6 +142,23 @@ filter_task_intrinsics(const nir_instr *instr, intrin->intrinsic == nir_intrinsic_load_task_payload; } +static void +task_invocation_query(nir_builder *b, lower_tsms_io_state *s) +{ + if (!s->has_query) + return; + + const unsigned invocations = b->shader->info.workgroup_size[0] * + b->shader->info.workgroup_size[1] * + b->shader->info.workgroup_size[2]; + + nir_if *if_pipeline_query = nir_push_if(b, nir_load_pipeline_stat_query_enabled_amd(b)); + { + nir_atomic_add_shader_invocation_count_amd(b, nir_imm_int(b, invocations)); + } + nir_pop_if(b, if_pipeline_query); +} + static nir_def * lower_task_launch_mesh_workgroups(nir_builder *b, nir_intrinsic_instr *intrin, @@ -179,6 +199,8 @@ lower_task_launch_mesh_workgroups(nir_builder *b, nir_scoped_memory_barrier(b, SCOPE_INVOCATION, NIR_MEMORY_RELEASE, nir_var_shader_out); /* Ready bit, only write the low 8 bits. */ task_write_draw_ring(b, task_draw_ready_bit(b, s), 12, s); + + task_invocation_query(b, s); } nir_pop_if(b, if_invocation_index_zero); @@ -256,7 +278,8 @@ lower_task_intrinsics(nir_builder *b, void ac_nir_lower_task_outputs_to_mem(nir_shader *shader, unsigned task_payload_entry_bytes, - unsigned task_num_entries) + unsigned task_num_entries, + bool has_query) { assert(util_is_power_of_two_nonzero(task_num_entries)); @@ -269,6 +292,7 @@ ac_nir_lower_task_outputs_to_mem(nir_shader *shader, .draw_entry_bytes = 16, .payload_entry_bytes = task_payload_entry_bytes, .num_entries = task_num_entries, + .has_query = has_query, }; nir_function_impl *impl = nir_shader_get_entrypoint(shader); diff --git a/src/amd/vulkan/nir/radv_nir_lower_io.c b/src/amd/vulkan/nir/radv_nir_lower_io.c index 0b98bcd8d89..4b555a6c819 100644 --- a/src/amd/vulkan/nir/radv_nir_lower_io.c +++ b/src/amd/vulkan/nir/radv_nir_lower_io.c @@ -172,8 +172,8 @@ radv_nir_lower_io_to_mem(struct radv_device *device, struct radv_shader_stage *s NIR_PASS_V(nir, ac_nir_lower_gs_inputs_to_mem, map_input, device->physical_device->rad_info.gfx_level, false); return true; } else if (nir->info.stage == MESA_SHADER_TASK) { - ac_nir_lower_task_outputs_to_mem(nir, AC_TASK_PAYLOAD_ENTRY_BYTES, - device->physical_device->task_info.num_entries); + ac_nir_lower_task_outputs_to_mem(nir, AC_TASK_PAYLOAD_ENTRY_BYTES, device->physical_device->task_info.num_entries, + false); return true; } else if (nir->info.stage == MESA_SHADER_MESH) { ac_nir_lower_mesh_inputs_to_mem(nir, AC_TASK_PAYLOAD_ENTRY_BYTES, device->physical_device->task_info.num_entries);
