nir: lower ballot_bit_count_exclusive/inclusive to mbcnt_amd

Reviewed-by: Rhys Perry <pendingchaos02@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/22783>
This commit is contained in:
Georg Lehmann 2023-05-01 19:04:03 +02:00 committed by Marge Bot
parent 5489033fa8
commit b93c92eba3
2 changed files with 15 additions and 4 deletions

View file

@ -5137,6 +5137,7 @@ typedef struct nir_lower_subgroups_options {
bool lower_elect:1;
bool lower_read_invocation_to_cond:1;
bool lower_rotate_to_shuffle:1;
bool lower_ballot_bit_count_to_mbcnt_amd:1;
} nir_lower_subgroups_options;
bool nir_lower_subgroups(nir_shader *shader,

View file

@ -774,6 +774,20 @@ lower_subgroups_instr(nir_builder *b, nir_instr *instr, void *_options)
case nir_intrinsic_ballot_bit_count_exclusive:
case nir_intrinsic_ballot_bit_count_inclusive: {
assert(intrin->src[0].is_ssa);
nir_ssa_def *int_val = ballot_type_to_uint(b, intrin->src[0].ssa,
options);
if (options->lower_ballot_bit_count_to_mbcnt_amd) {
nir_ssa_def *acc;
if (intrin->intrinsic == nir_intrinsic_ballot_bit_count_exclusive) {
acc = nir_imm_int(b, 0);
} else {
acc = nir_iand_imm(b, nir_u2u32(b, int_val), 0x1);
int_val = nir_ushr_imm(b, int_val, 1);
}
return nir_mbcnt_amd(b, int_val, acc);
}
nir_ssa_def *mask;
if (intrin->intrinsic == nir_intrinsic_ballot_bit_count_inclusive) {
mask = nir_inot(b, build_subgroup_gt_mask(b, options));
@ -781,10 +795,6 @@ lower_subgroups_instr(nir_builder *b, nir_instr *instr, void *_options)
mask = nir_inot(b, build_subgroup_ge_mask(b, options));
}
assert(intrin->src[0].is_ssa);
nir_ssa_def *int_val = ballot_type_to_uint(b, intrin->src[0].ssa,
options);
return vec_bit_count(b, nir_iand(b, int_val, mask));
}