agx: handle non-immediate shuffles in divergent CF

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

src/asahi/compiler/agx_compile.c

@@ -1488,6 +1488,9 @@ agx_emit_intrinsic(agx_builder *b, nir_intrinsic_instr *instr)
    }
 
    case nir_intrinsic_read_invocation: {
+      /* TODO: Check if we're actually inside divergent control flow */
+      b->shader->any_quad_divergent_shuffle |= b->shader->any_cf;
+
       /* Lane ID guaranteed to be uniform */
       return agx_shuffle_to(b, dst, agx_src_index(&instr->src[0]),
                             agx_src_index(&instr->src[1]));
@@ -3096,6 +3099,10 @@ agx_compile_function_nir(nir_shader *nir, nir_function_impl *impl,
    agx_opt_empty_else(ctx);
    agx_opt_break_if(ctx);
    agx_opt_jmp_none(ctx);
+
+   if (ctx->any_quad_divergent_shuffle)
+      agx_lower_divergent_shuffle(ctx);
+
    agx_lower_pseudo(ctx);
 
    if (agx_should_dump(nir, AGX_DBG_SHADERS))

src/asahi/compiler/agx_compiler.h

@@ -455,6 +455,11 @@ typedef struct {
    /* Has r0l been zeroed yet due to control flow? */
    bool any_cf;
 
+   /* Do we need r0h zero throughout the program to handle quad-divergent
+    * shuffle?
+    */
+   bool any_quad_divergent_shuffle;
+
    /* Number of nested control flow structures within the innermost loop. Since
     * NIR is just loop and if-else, this is the number of nested if-else
     * statements in the loop */
@@ -946,6 +951,7 @@ void agx_print_instr(const agx_instr *I, FILE *fp);
 void agx_print_block(const agx_block *block, FILE *fp);
 void agx_print_shader(const agx_context *ctx, FILE *fp);
 void agx_optimizer(agx_context *ctx);
+void agx_lower_divergent_shuffle(agx_context *ctx);
 void agx_lower_pseudo(agx_context *ctx);
 void agx_lower_spill(agx_context *ctx);
 void agx_lower_uniform_sources(agx_context *ctx);

src/asahi/compiler/agx_lower_divergent_shuffle.c

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2024 Alyssa Rosenzweig
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "agx_builder.h"
+#include "agx_compiler.h"
+#include "agx_opcodes.h"
+
+static bool
+is_shuffle(enum agx_opcode op)
+{
+   switch (op) {
+   case AGX_OPCODE_SHUFFLE:
+   case AGX_OPCODE_SHUFFLE_UP:
+   case AGX_OPCODE_SHUFFLE_DOWN:
+   case AGX_OPCODE_SHUFFLE_XOR:
+   case AGX_OPCODE_QUAD_SHUFFLE:
+   case AGX_OPCODE_QUAD_SHUFFLE_UP:
+   case AGX_OPCODE_QUAD_SHUFFLE_DOWN:
+   case AGX_OPCODE_QUAD_SHUFFLE_XOR:
+      return true;
+   default:
+      return false;
+   }
+}
+
+/*
+ * AGX shuffle instructions read indices to shuffle with from the entire quad
+ * and accumulate them. That means that an inactive thread anywhere in the quad
+ * can make the whole shuffle undefined! To workaround, we reserve a scratch
+ * register (r0h) which we keep zero throughout the program... except for when
+ * actually shuffling, when we copy the shuffle index into r0h for the
+ * operation. This ensures that inactive threads read 0 for their index and
+ * hence do not contribute to the accumulated index.
+ */
+void
+agx_lower_divergent_shuffle(agx_context *ctx)
+{
+   agx_builder b = agx_init_builder(ctx, agx_before_function(ctx));
+   agx_index scratch = agx_register(1, AGX_SIZE_16);
+
+   assert(ctx->any_quad_divergent_shuffle);
+   agx_mov_imm_to(&b, scratch, 0);
+
+   agx_foreach_block(ctx, block) {
+      bool needs_zero = false;
+
+      agx_foreach_instr_in_block_safe(block, I) {
+         if (is_shuffle(I->op) && I->src[1].type == AGX_INDEX_REGISTER) {
+            assert(I->dest[0].value != scratch.value);
+            assert(I->src[0].value != scratch.value);
+            assert(I->src[1].value != scratch.value);
+
+            /* Use scratch register for our input, then zero it at the end of
+             * the block so all inactive threads read zero.
+             */
+            b.cursor = agx_before_instr(I);
+            agx_mov_to(&b, scratch, I->src[1]);
+            needs_zero = true;
+
+            I->src[1] = scratch;
+         }
+      }
+
+      if (needs_zero) {
+         b.cursor = agx_after_block_logical(block);
+         agx_mov_imm_to(&b, scratch, 0);
+      }
+   }
+}

src/asahi/compiler/agx_register_allocate.c

@@ -177,6 +177,9 @@ agx_calc_register_demand(agx_context *ctx)
       if (ctx->any_cf)
          demand++;
 
+      if (ctx->any_quad_divergent_shuffle)
+         demand++;
+
       /* Everything live-in */
       {
          int i;
@@ -299,16 +302,23 @@ find_best_region_to_evict(struct ra_ctx *rctx, enum ra_class cls, unsigned size,
    unsigned best_base = ~0;
    unsigned best_moves = ~0;
 
+   /* r0h unevictable only when r0l unevictable */
+   assert(!rctx->shader->any_quad_divergent_shuffle || rctx->shader->any_cf);
+
    for (unsigned base = 0; base + size <= rctx->bound[cls]; base += size) {
-      /* r0l is unevictable, skip it. By itself, this does not pose a problem.
-       * We are allocating n registers, but the region containing r0l has at
-       * most n-1 free. Since there are at least n free registers total, there
-       * is at least 1 free register outside this region. Thus the region
-       * containing that free register contains at most n-1 occupied registers.
-       * In the worst case, those n-1 occupied registers are moved to the region
-       * with r0l and then the n free registers are used for the destination.
-       * Thus, we do not need extra registers to handle "single point"
-       * unevictability.
+      /* The first k registers are preallocated and unevictable, so must be
+       * skipped. By itself, this does not pose a problem. We are allocating n
+       * registers, but this region has at most n-k free.  Since there are at
+       * least n free registers total, there is at least k free registers
+       * outside this region. Choose any such free register. The region
+       * containing it has at most n-1 occupied registers. In the worst case,
+       * n-k of those registers are are moved to the beginning region and the
+       * remaining (n-1)-(n-k) = k-1 registers are moved to the k-1 free
+       * registers in other regions, given there are k free registers total.
+       * These recursive shuffles work out because everything is power-of-two
+       * sized and naturally aligned, so the sizes shuffled are strictly
+       * descending. So, we do not need extra registers to handle "single
+       * region" unevictability.
        */
       if (base == 0 && rctx->shader->any_cf)
          continue;
@@ -509,6 +519,8 @@ insert_copies_for_clobbered_killed(struct ra_ctx *rctx, unsigned reg,
     * have to move it.  find_best_region_to_evict knows better than to try.
     */
    assert(!(reg == 0 && rctx->shader->any_cf) && "r0l is never moved");
+   assert(!(reg == 1 && rctx->shader->any_quad_divergent_shuffle) &&
+          "r0h is never moved");
 
    /* Consider the destination clobbered for the purpose of source collection.
     * This way, killed sources already in the destination will be preserved
@@ -1069,6 +1081,12 @@ agx_ra_assign_local(struct ra_ctx *rctx)
    if (rctx->shader->any_cf)
       BITSET_SET(used_regs_gpr, 0);
 
+   /* Force the zero r0h live throughout shaders using divergent shuffles. */
+   if (rctx->shader->any_quad_divergent_shuffle) {
+      assert(rctx->shader->any_cf);
+      BITSET_SET(used_regs_gpr, 1);
+   }
+
    agx_foreach_instr_in_block(block, I) {
       rctx->instr = I;
 

src/asahi/compiler/meson.build

@@ -17,6 +17,7 @@ libasahi_agx_files = files(
   'agx_nir_lower_subgroups.c',
   'agx_nir_opt_preamble.c',
   'agx_lower_64bit.c',
+  'agx_lower_divergent_shuffle.c',
   'agx_lower_parallel_copy.c',
   'agx_lower_pseudo.c',
   'agx_lower_spill.c',