From 3efa723b630dc153252ac7841f60924078684e04 Mon Sep 17 00:00:00 2001 From: Alyssa Rosenzweig Date: Fri, 19 Apr 2024 22:03:37 -0400 Subject: [PATCH] agx: lower nir_intrinsic_load_num_subgroups Signed-off-by: Alyssa Rosenzweig Part-of: --- src/asahi/compiler/agx_nir_lower_subgroups.c | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/asahi/compiler/agx_nir_lower_subgroups.c b/src/asahi/compiler/agx_nir_lower_subgroups.c index c92f796e24c..0e0464bc6f5 100644 --- a/src/asahi/compiler/agx_nir_lower_subgroups.c +++ b/src/asahi/compiler/agx_nir_lower_subgroups.c @@ -9,6 +9,13 @@ #include "nir_builder_opcodes.h" #include "nir_intrinsics.h" +/* XXX: cribbed from nak, move to common */ +static nir_def * +nir_udiv_round_up(nir_builder *b, nir_def *n, nir_def *d) +{ + return nir_udiv(b, nir_iadd(b, n, nir_iadd_imm(b, d, -1)), d); +} + static bool lower(nir_builder *b, nir_intrinsic_instr *intr, void *data) { @@ -67,6 +74,19 @@ lower(nir_builder *b, nir_intrinsic_instr *intr, void *data) return true; } + case nir_intrinsic_load_num_subgroups: { + nir_def *workgroup_size = nir_load_workgroup_size(b); + workgroup_size = nir_imul(b, + nir_imul(b, nir_channel(b, workgroup_size, 0), + nir_channel(b, workgroup_size, 1)), + nir_channel(b, workgroup_size, 2)); + nir_def *subgroup_size = nir_imm_int(b, 32); + nir_def *num_subgroups = + nir_udiv_round_up(b, workgroup_size, subgroup_size); + nir_def_rewrite_uses(&intr->def, num_subgroups); + return true; + } + default: return false; }