From 03830c9523d89a8afcac896d05ef299864a0d528 Mon Sep 17 00:00:00 2001 From: Alyssa Rosenzweig Date: Sat, 3 Aug 2024 10:03:27 -0400 Subject: [PATCH] agx: optimize popcount(ballot(true)) Signed-off-by: Alyssa Rosenzweig Part-of: --- src/asahi/compiler/agx_nir_lower_subgroups.c | 21 ++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/asahi/compiler/agx_nir_lower_subgroups.c b/src/asahi/compiler/agx_nir_lower_subgroups.c index a59316cadc6..0d6c561dac8 100644 --- a/src/asahi/compiler/agx_nir_lower_subgroups.c +++ b/src/asahi/compiler/agx_nir_lower_subgroups.c @@ -5,6 +5,7 @@ #include "compiler/nir/nir.h" #include "compiler/nir/nir_builder.h" +#include "util/list.h" #include "agx_nir.h" #include "nir_builder_opcodes.h" #include "nir_intrinsics.h" @@ -154,6 +155,26 @@ lower(nir_builder *b, nir_intrinsic_instr *intr, void *data) return true; } + case nir_intrinsic_ballot: { + /* Optimize popcount(ballot(true)) to load_active_subgroup_count_agx() */ + if (!nir_src_is_const(intr->src[0]) || !nir_src_as_bool(intr->src[0]) || + !list_is_singular(&intr->def.uses)) + return false; + + nir_src *use = list_first_entry(&intr->def.uses, nir_src, use_link); + nir_instr *parent = nir_src_parent_instr(use); + if (parent->type != nir_instr_type_alu) + return false; + + nir_alu_instr *alu = nir_instr_as_alu(parent); + if (alu->op != nir_op_bit_count) + return false; + + nir_def_rewrite_uses(&alu->def, + nir_load_active_subgroup_count_agx(b, 32)); + return true; + } + default: return false; }