Module: Mesa Branch: main Commit: 5825e20dbf59f687973303ede2ac0836d6f405a2 URL: http://cgit.freedesktop.org/mesa/mesa/commit/?id=5825e20dbf59f687973303ede2ac0836d6f405a2
Author: Timur Kristóf <[email protected]> Date: Mon Jul 3 17:07:18 2023 +0200 ac/nir: Add new pass to lower intrinsics to shader args. This is beneficial for intrinsics that do an algebraic instruction such as bitfield extract on shader arguments, because it allows NIR to be aware of these instructions and optimize them together with other algebraic instructions in the shader. Currently, just handle subgroup_id and num_subgroups intrinsics. More will be added here in the future. Signed-off-by: Timur Kristóf <[email protected]> Reviewed-by: Marek Olšák <[email protected]> Reviewed-by: Rhys Perry <[email protected]> Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24005> --- src/amd/common/ac_nir.c | 83 +++++++++++++++++++++++++++++++++++++++++++++++++ src/amd/common/ac_nir.h | 10 ++++-- 2 files changed, 90 insertions(+), 3 deletions(-) diff --git a/src/amd/common/ac_nir.c b/src/amd/common/ac_nir.c index 8ed83771d67..0df5d71cb14 100644 --- a/src/amd/common/ac_nir.c +++ b/src/amd/common/ac_nir.c @@ -67,6 +67,89 @@ ac_nir_lower_sin_cos(nir_shader *shader) return nir_shader_lower_instructions(shader, is_sin_cos, lower_sin_cos, NULL); } +typedef struct { + const struct ac_shader_args *const args; + const enum amd_gfx_level gfx_level; + const enum ac_hw_stage hw_stage; +} lower_intrinsics_to_args_state; + +static bool +lower_intrinsic_to_arg(nir_builder *b, nir_instr *instr, void *state) +{ + if (instr->type != nir_instr_type_intrinsic) + return false; + + lower_intrinsics_to_args_state *s = (lower_intrinsics_to_args_state *)state; + nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr); + nir_ssa_def *replacement = NULL; + b->cursor = nir_after_instr(&intrin->instr); + + switch (intrin->intrinsic) { + case nir_intrinsic_load_subgroup_id: { + if (s->hw_stage == AC_HW_COMPUTE_SHADER) { + assert(s->args->tg_size.used); + + if (s->gfx_level >= GFX10_3) { + replacement = ac_nir_unpack_arg(b, s->args, s->args->tg_size, 20, 5); + } else { + /* GFX6-10 don't actually support a wave id, but we can + * use the ordered id because ORDERED_APPEND_* is set to + * zero in the compute dispatch initiatior. + */ + replacement = ac_nir_unpack_arg(b, s->args, s->args->tg_size, 6, 6); + } + } else if (s->hw_stage == AC_HW_HULL_SHADER && s->gfx_level >= GFX11) { + assert(s->args->tcs_wave_id.used); + replacement = ac_nir_unpack_arg(b, s->args, s->args->tcs_wave_id, 0, 3); + } else if (s->hw_stage == AC_HW_LEGACY_GEOMETRY_SHADER || + s->hw_stage == AC_HW_NEXT_GEN_GEOMETRY_SHADER) { + assert(s->args->merged_wave_info.used); + replacement = ac_nir_unpack_arg(b, s->args, s->args->merged_wave_info, 24, 4); + } else { + replacement = nir_imm_int(b, 0); + } + + break; + } + case nir_intrinsic_load_num_subgroups: { + if (s->hw_stage == AC_HW_COMPUTE_SHADER) { + assert(s->args->tg_size.used); + replacement = ac_nir_unpack_arg(b, s->args, s->args->tg_size, 0, 6); + } else if (s->hw_stage == AC_HW_LEGACY_GEOMETRY_SHADER || + s->hw_stage == AC_HW_NEXT_GEN_GEOMETRY_SHADER) { + assert(s->args->merged_wave_info.used); + replacement = ac_nir_unpack_arg(b, s->args, s->args->merged_wave_info, 28, 4); + } else { + replacement = nir_imm_int(b, 1); + } + + break; + } + default: + return false; + } + + assert(replacement); + nir_ssa_def_rewrite_uses(&intrin->dest.ssa, replacement); + nir_instr_remove(&intrin->instr); + return true; +} + +bool +ac_nir_lower_intrinsics_to_args(nir_shader *shader, const enum amd_gfx_level gfx_level, + const enum ac_hw_stage hw_stage, + const struct ac_shader_args *ac_args) +{ + lower_intrinsics_to_args_state state = { + .gfx_level = gfx_level, + .hw_stage = hw_stage, + .args = ac_args, + }; + + return nir_shader_instructions_pass(shader, lower_intrinsic_to_arg, + nir_metadata_block_index | nir_metadata_dominance, &state); +} + void ac_nir_store_var_components(nir_builder *b, nir_variable *var, nir_ssa_def *value, unsigned component, unsigned writemask) diff --git a/src/amd/common/ac_nir.h b/src/amd/common/ac_nir.h index 264355c9dff..3f8585fc0e8 100644 --- a/src/amd/common/ac_nir.h +++ b/src/amd/common/ac_nir.h @@ -8,11 +8,11 @@ #ifndef AC_NIR_H #define AC_NIR_H -#include "nir.h" -#include "nir_builder.h" +#include "ac_hw_stage.h" #include "ac_shader_args.h" #include "ac_shader_util.h" -#include "amd_family.h" +#include "nir.h" +#include "nir_builder.h" #ifdef __cplusplus extern "C" { @@ -66,6 +66,10 @@ ac_nir_unpack_arg(nir_builder *b, const struct ac_shader_args *ac_args, struct a bool ac_nir_lower_sin_cos(nir_shader *shader); +bool ac_nir_lower_intrinsics_to_args(nir_shader *shader, const enum amd_gfx_level gfx_level, + const enum ac_hw_stage hw_stage, + const struct ac_shader_args *ac_args); + void ac_nir_store_var_components(nir_builder *b, nir_variable *var, nir_ssa_def *value, unsigned component, unsigned writemask);
