From 7b33c549b97858129a520483d1eaa254c704d254 Mon Sep 17 00:00:00 2001 From: Alyssa Rosenzweig Date: Fri, 19 Apr 2024 22:00:15 -0400 Subject: [PATCH] agx: handle non-immediate shuffles in divergent CF Signed-off-by: Alyssa Rosenzweig Part-of: --- src/asahi/compiler/agx_compile.c | 7 ++ src/asahi/compiler/agx_compiler.h | 6 ++ .../compiler/agx_lower_divergent_shuffle.c | 71 +++++++++++++++++++ src/asahi/compiler/agx_register_allocate.c | 36 +++++++--- src/asahi/compiler/meson.build | 1 + 5 files changed, 112 insertions(+), 9 deletions(-) create mode 100644 src/asahi/compiler/agx_lower_divergent_shuffle.c diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c index 559dc4fb32a..ee6fe1b3eb9 100644 --- a/src/asahi/compiler/agx_compile.c +++ b/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)) diff --git a/src/asahi/compiler/agx_compiler.h b/src/asahi/compiler/agx_compiler.h index 6ee973c26d0..6fba5cb429d 100644 --- a/src/asahi/compiler/agx_compiler.h +++ b/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); diff --git a/src/asahi/compiler/agx_lower_divergent_shuffle.c b/src/asahi/compiler/agx_lower_divergent_shuffle.c new file mode 100644 index 00000000000..8ba7dca0164 --- /dev/null +++ b/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); + } + } +} diff --git a/src/asahi/compiler/agx_register_allocate.c b/src/asahi/compiler/agx_register_allocate.c index e57c22daf04..2e44cdb6d4f 100644 --- a/src/asahi/compiler/agx_register_allocate.c +++ b/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; diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build index 18c99041fe8..538ea482434 100644 --- a/src/asahi/compiler/meson.build +++ b/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',