agx: lower votes to ballots

not optimal but passes the tests.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/26963>

src/asahi/compiler/agx_compile.c

@@ -2988,6 +2988,7 @@ agx_preprocess_nir(nir_shader *nir, const nir_shader *libagx,
    NIR_PASS_V(nir, nir_shader_intrinsics_pass, agx_lower_front_face,
               nir_metadata_block_index | nir_metadata_dominance, NULL);
    NIR_PASS_V(nir, nir_lower_frag_coord_to_pixel_coord);
+   NIR_PASS_V(nir, agx_nir_lower_subgroups);
 
    /* After lowering, run through the standard suite of NIR optimizations. We
     * will run through the loop later, once we have the shader key, but if we

src/asahi/compiler/agx_nir.h

@@ -17,5 +17,6 @@ bool agx_nir_fence_images(struct nir_shader *shader);
 bool agx_nir_lower_multisampled_image_store(struct nir_shader *s);
 void agx_nir_lower_layer(struct nir_shader *s);
 void agx_nir_lower_cull_distance_vs(struct nir_shader *s);
+void agx_nir_lower_subgroups(struct nir_shader *s);
 
 #endif

src/asahi/compiler/agx_nir_lower_subgroups.c

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2023 Valve Corporation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "compiler/nir/nir.h"
+#include "compiler/nir/nir_builder.h"
+#include "agx_nir.h"
+#include "nir_builder_opcodes.h"
+#include "nir_intrinsics.h"
+
+static bool
+lower(nir_builder *b, nir_intrinsic_instr *intr, void *data)
+{
+   b->cursor = nir_before_instr(&intr->instr);
+
+   switch (intr->intrinsic) {
+   case nir_intrinsic_vote_any: {
+      /* We don't have vote instructions, but we have efficient ballots */
+      nir_def *ballot = nir_ballot(b, 1, 32, intr->src[0].ssa);
+      nir_def_rewrite_uses(&intr->def, nir_ine_imm(b, ballot, 0));
+      return true;
+   }
+
+   case nir_intrinsic_vote_all: {
+      nir_def *ballot = nir_ballot(b, 1, 32, nir_inot(b, intr->src[0].ssa));
+      nir_def_rewrite_uses(&intr->def, nir_ieq_imm(b, ballot, 0));
+      return true;
+   }
+
+   default:
+      return false;
+   }
+}
+
+void
+agx_nir_lower_subgroups(nir_shader *s)
+{
+   /* First, do as much common lowering as we can */
+   nir_lower_subgroups_options opts = {
+      .lower_vote_eq = true,
+      .lower_vote_bool_eq = true,
+      .lower_read_first_invocation = true,
+      .lower_first_invocation_to_ballot = true,
+      .lower_to_scalar = true,
+      .ballot_components = 1,
+      .ballot_bit_size = 32,
+   };
+
+   NIR_PASS_V(s, nir_lower_subgroups, &opts);
+
+   /* Then do AGX-only lowerings on top */
+   nir_shader_intrinsics_pass(
+      s, lower, nir_metadata_block_index | nir_metadata_dominance, NULL);
+}

src/asahi/compiler/meson.build

@@ -18,6 +18,7 @@ libasahi_agx_files = files(
   'agx_nir_lower_layer.c',
   'agx_nir_lower_load_mask.c',
   'agx_nir_lower_shared_bitsize.c',
+  'agx_nir_lower_subgroups.c',
   'agx_nir_opt_preamble.c',
   'agx_lower_64bit.c',
   'agx_lower_parallel_copy.c',