From 5cef905ea5ce14e8b7d01bd9d99cd3a6f1356907 Mon Sep 17 00:00:00 2001 From: Sean Dolbec Date: Sun, 1 Nov 2026 19:32:58 -0500 Subject: [PATCH 1/2] agx: flatten guarded read-only loads in compute shaders `row < m ? x[i] : 0` over a read-only buffer stays a branch: peephole select keeps indirect loads under their if, because the if may be what keeps the address valid, and soft faults do not change that. omarchy-mlx's FMA prefill matmul (qmm_fma.comp) guards every load this way, so its inner loop is a chain of one-load ifs with a phi at each join. In a short compute-shader arm made only of ALU, preamble reads and read-only global loads, point each load at the zero page when its arm is not taken. The zero page is always mapped and the select formed from the if discards the value, so the load may run either way, on chips with or without soft faults. The base and the unscaled 32-bit index are selected rather than the 64-bit sum, so agx_nir_lower_address still folds base + index into the load. Loads the driver already marks speculatable (descriptor reads) keep their address; read-only, reorderable load_global becomes load_global_constant, the form peephole select accepts; and flattening is requested for that if only. The pass runs before soft-fault speculation marks every load. Other indirect loads keep their branches. ASAHI_NO_GUARDED_LOAD_SELECT=1 turns it off for A/B runs. GuardedLoadSelect in agx_tests covers the rewrite (zero-page base and unscaled index selected, load_global_constant marked speculatable, the if flattened) and the cases left alone: an arm with a store, a writable or volatile load, dont_flatten, a fragment shader, and a descriptor read that is already speculatable. Compile-only, through drm-shim on an M4 Mac mini host (no GPU submits), assertions and NIR validation on, robustBufferAccess off: - qmm_fma.comp, G17P: 3687 -> 3523 instructions, 429 -> 249 blocks, 542 -> 350 moves; G15S: 3690 -> 3620, 429 -> 249, 545 -> 447. - All 106 omarchy-mlx compute shaders that build with a storage-buffer-only layout compile with and without the pass on G15S, G17P and G13G, with no new failure. 44 change; static instructions in total -1.3% (G15S), -1.9% (G17P), -0.9% (G13G), blocks -10%. Some kernels grow, by up to 7.5% (fast_norm_vjp_dx 828 -> 890 on G15S, int8_matmul 2046 -> 2103 on G13G). Not run on a GPU with this tree. On an M4 (G16G) with the GravityLinux Apple9 backend, the same idea took that kernel from 174 to 778 GFLOP/s and f16 prefill 4.5-4.8x, but there the branches cost far more (5300 of 6412 instructions were copies), so expect a smaller effect here. Signed-off-by: Sean Dolbec Co-Authored-By: Claude Opus 5.5 --- src/asahi/compiler/agx_compile.c | 172 ++++++++++++++ src/asahi/compiler/agx_nir.h | 1 + src/asahi/compiler/meson.build | 1 + .../test/test-guarded-load-select.cpp | 218 ++++++++++++++++++ 4 files changed, 392 insertions(+) create mode 100644 src/asahi/compiler/test/test-guarded-load-select.cpp diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c index 5443f4e509f0..8891ad8d3428 100644 --- a/src/asahi/compiler/agx_compile.c +++ b/src/asahi/compiler/agx_compile.c @@ -7933,6 +7933,173 @@ set_speculate(nir_builder *b, nir_instr *instr, UNUSED void *_) return false; } +/* + * `cond ? buffer[i] : 0` over a read-only buffer reaches us as an if around a + * global load. Peephole select leaves such a load under its branch, since the + * branch may be what keeps the address valid, so in a loop every guarded + * value becomes a phi that the register allocator reconciles at each join. + * + * In a short arm of only ALU, preamble reads and read-only global loads, read + * the zero page instead when the arm is not taken. That address is always + * mapped and at least AGX_ZERO_PAGE_SIZE long, and the select formed from + * the if discards the value, so the load is safe to execute either way, with + * or without soft faults. Loads the driver already marked speculatable + * (descriptor reads) keep their address. Read-only, reorderable global loads + * become load_global_constant, the form peephole select knows, and + * flattening is requested for that if alone; other loads keep their + * branches. This runs before soft-fault speculation marks every load. + */ +static bool +agx_guarded_load(nir_intrinsic_instr *intr) +{ + if (intr->intrinsic != nir_intrinsic_load_global_constant && + intr->intrinsic != nir_intrinsic_load_global) + return false; + enum gl_access_qualifier access = nir_intrinsic_access(intr); + if ((access & (ACCESS_VOLATILE | ACCESS_COHERENT)) || + intr->def.num_components * intr->def.bit_size > 128) + return false; + return intr->intrinsic == nir_intrinsic_load_global_constant || + ((access & ACCESS_NON_WRITEABLE) && (access & ACCESS_CAN_REORDER)); +} + +static bool +agx_guarded_load_arm_ok(nir_block *block) +{ + unsigned count = 0; + nir_foreach_instr(instr, block) { + switch (instr->type) { + case nir_instr_type_alu: + case nir_instr_type_load_const: + case nir_instr_type_undef: + break; + case nir_instr_type_intrinsic: { + nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); + if (intr->intrinsic != nir_intrinsic_load_preamble && + !agx_guarded_load(intr)) + return false; + break; + } + default: + return false; + } + if (++count > 64) + return false; + } + return true; +} + +/* The address an arm's load uses when the arm is not taken: the zero page. + * Keep the base + 32-bit index shape agx_nir_lower_address folds into the + * load, selecting the base and the unscaled index rather than the 64-bit + * sum: zero page + 0. + */ +static nir_def * +agx_guarded_address(nir_builder *b, nir_def *addr, nir_def *cond, + bool taken_if_true) +{ + nir_def *zero_page = nir_imm_int64(b, AGX_ZERO_PAGE_ADDRESS); +#define AGX_GUARD(taken, other) \ + (taken_if_true ? nir_bcsel(b, cond, taken, other) \ + : nir_bcsel(b, cond, other, taken)) + + nir_alu_instr *add = nir_def_as_alu_or_null(addr); + if (add && add->op == nir_op_iadd) { + for (unsigned i = 0; i < 2; ++i) { + nir_alu_instr *cvt = nir_def_as_alu_or_null(add->src[1 - i].src.ssa); + if (!cvt || (cvt->op != nir_op_u2u64 && cvt->op != nir_op_i2i64) || + nir_src_bit_size(cvt->src[0].src) != 32) + continue; + + nir_def *base = nir_ssa_for_alu_src(b, add, i); + nir_def *index = nir_ssa_for_alu_src(b, cvt, 0); + nir_def *zero = nir_imm_int(b, 0); + nir_alu_instr *scale = nir_def_as_alu_or_null(index); + nir_def *guarded; + if (scale && (scale->op == nir_op_imul || scale->op == nir_op_amul || + scale->op == nir_op_ishl) && + nir_src_is_const(scale->src[1].src)) { + nir_def *unscaled = nir_ssa_for_alu_src(b, scale, 0); + guarded = nir_build_alu2(b, scale->op, AGX_GUARD(unscaled, zero), + nir_ssa_for_alu_src(b, scale, 1)); + } else { + guarded = AGX_GUARD(index, zero); + } + return nir_iadd(b, AGX_GUARD(base, zero_page), + nir_build_alu1(b, cvt->op, guarded)); + } + } + + return AGX_GUARD(addr, zero_page); +#undef AGX_GUARD +} + +/* Returns whether the arm had a load to redirect. */ +static bool +agx_redirect_arm_loads(nir_block *block, nir_def *cond, bool taken_if_true) +{ + bool loads = false; + nir_foreach_instr_safe(instr, block) { + if (instr->type != nir_instr_type_intrinsic) + continue; + nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); + if (intr->intrinsic == nir_intrinsic_load_preamble) + continue; + enum gl_access_qualifier access = nir_intrinsic_access(intr); + if (intr->intrinsic == nir_intrinsic_load_global_constant && + (access & ACCESS_CAN_SPECULATE)) + continue; + + nir_builder b = nir_builder_at(nir_before_instr(instr)); + nir_def *addr = intr->src[0].ssa; + if (!(access & ACCESS_CAN_SPECULATE)) + addr = agx_guarded_address(&b, addr, cond, taken_if_true); + nir_def *value = nir_load_global_constant( + &b, intr->def.num_components, intr->def.bit_size, addr, + .align_mul = nir_intrinsic_align_mul(intr), + .align_offset = nir_intrinsic_align_offset(intr), + .access = access | ACCESS_CAN_SPECULATE); + nir_def_replace(&intr->def, value); + loads = true; + } + return loads; +} + +bool +agx_nir_speculate_guarded_loads(nir_shader *nir) +{ + if (!mesa_shader_stage_is_compute(nir->info.stage)) + return false; + + bool progress = false; + nir_foreach_function_impl(impl, nir) { + bool impl_progress = false; + nir_foreach_block(block, impl) { + nir_cf_node *next = nir_cf_node_next(&block->cf_node); + if (!next || next->type != nir_cf_node_if) + continue; + nir_if *nif = nir_cf_node_as_if(next); + nir_block *then_block = nir_if_first_then_block(nif); + nir_block *else_block = nir_if_first_else_block(nif); + if (nif->control == nir_selection_control_dont_flatten || + then_block != nir_if_last_then_block(nif) || + else_block != nir_if_last_else_block(nif) || + !agx_guarded_load_arm_ok(then_block) || + !agx_guarded_load_arm_ok(else_block)) + continue; + nir_def *cond = nif->condition.ssa; + bool loads = agx_redirect_arm_loads(then_block, cond, true); + loads |= agx_redirect_arm_loads(else_block, cond, false); + if (loads) { + nif->control = nir_selection_control_flatten; + impl_progress = true; + } + } + progress |= nir_progress(impl_progress, impl, nir_metadata_control_flow); + } + return progress; +} + static bool agx_g17p_discard_word(nir_shader *nir); static bool agx_m5_discard_word(nir_shader *nir); @@ -8136,6 +8303,11 @@ agx_optimize_nir(nir_shader *nir, bool soft_fault, bool uniform_atomics, agx_optimize_loop_nir(nir, m5); + /* ASAHI_NO_GUARDED_LOAD_SELECT=1 keeps guarded loads under their branch + * for A/B comparisons. */ + if (!debug_get_bool_option("ASAHI_NO_GUARDED_LOAD_SELECT", false)) + NIR_PASS(_, nir, agx_nir_speculate_guarded_loads); + /* If soft fault is enabled, we can freely speculate everything. That lets us * peephole select and form preambles more aggressively. */ diff --git a/src/asahi/compiler/agx_nir.h b/src/asahi/compiler/agx_nir.h index d38cdbb132cd..167549be2a6a 100644 --- a/src/asahi/compiler/agx_nir.h +++ b/src/asahi/compiler/agx_nir.h @@ -31,6 +31,7 @@ bool agx_nir_lower_lea(struct nir_shader *shader); bool agx_nir_fuse_selects(struct nir_shader *shader); bool agx_nir_m5_split_or_selects(struct nir_shader *shader); bool agx_nir_opt_select_load(struct nir_shader *shader); +bool agx_nir_speculate_guarded_loads(struct nir_shader *shader); bool agx_nir_opt_sink_loop_sample_uses(struct nir_shader *shader); bool agx_nir_opt_hoist_loop_samples(struct nir_shader *shader); bool agx_nir_fuse_algebraic_late(struct nir_shader *shader); diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build index 2b766a4f14d3..8faac674e4f9 100644 --- a/src/asahi/compiler/meson.build +++ b/src/asahi/compiler/meson.build @@ -231,6 +231,7 @@ if with_tests 'test/test-fminmax-constant-denorm.cpp', 'test/test-shared-byte-offsets.cpp', 'test/test-device-load-batches.cpp', + 'test/test-guarded-load-select.cpp', 'test/test-m5-copy-schedule.cpp', ), c_args : [c_msvc_compat_args, no_override_init_args], diff --git a/src/asahi/compiler/test/test-guarded-load-select.cpp b/src/asahi/compiler/test/test-guarded-load-select.cpp new file mode 100644 index 000000000000..51f292f4dc1a --- /dev/null +++ b/src/asahi/compiler/test/test-guarded-load-select.cpp @@ -0,0 +1,218 @@ +/* Copyright 2026 Sean Dolbec; SPDX-License-Identifier: MIT */ +#include +#include "agx_test.h" +#include "compiler/nir/nir_builder.h" +extern "C" { +#include "agx_nir.h" +#include "asahi/lib/agx_abi.h" +} + +/* + * agx_nir_speculate_guarded_loads: `cond ? buffer[i] : 0` over a read-only + * buffer becomes a load from the selected address (the zero page when the + * arm is not taken) and an if that peephole select may flatten. Anything + * else keeps its branch. + */ + +/* Index-setting helpers: C++ rejects nir_builder's named arguments. */ +static nir_def * +push_constant(nir_builder *b, unsigned bit_size, unsigned offset) +{ + nir_def *def = nir_load_push_constant(b, 1, bit_size, nir_imm_int(b, offset)); + nir_intrinsic_set_range(nir_def_as_intrinsic(def), offset + bit_size / 8); + return def; +} + +static nir_def * +load_global(nir_builder *b, nir_def *addr, gl_access_qualifier access) +{ + nir_def *def = nir_load_global(b, 1, 32, addr); + nir_intrinsic_set_access(nir_def_as_intrinsic(def), access); + nir_intrinsic_set_align(nir_def_as_intrinsic(def), 4, 0); + return def; +} + +static void +store_global(nir_builder *b, nir_def *value, nir_def *addr) +{ + nir_intrinsic_instr *intr = nir_store_global(b, value, addr); + nir_intrinsic_set_align(intr, 4, 0); + nir_intrinsic_set_write_mask(intr, 1); +} + +enum arm_extra { + ARM_LOAD_ONLY, + ARM_STORE, + ARM_WRITABLE_LOAD, + ARM_VOLATILE_LOAD, +}; + +struct guarded { + nir_builder b; + nir_if *nif; + nir_intrinsic_instr *load; + nir_def *base; + nir_def *index; +}; + +static guarded +build(mesa_shader_stage stage, arm_extra extra, bool dont_flatten = false) +{ + guarded g; + g.b = nir_builder_init_simple_shader(stage, &agx_nir_options, + "guarded load"); + nir_builder *b = &g.b; + g.base = push_constant(b, 64, 0); + g.index = push_constant(b, 32, 8); + nir_def *cond = nir_ult_imm(b, g.index, 100); + + g.nif = nir_push_if(b, cond); + if (dont_flatten) + g.nif->control = nir_selection_control_dont_flatten; + nir_def *offset = nir_u2u64(b, nir_imul_imm(b, g.index, 4)); + nir_def *addr = nir_iadd(b, g.base, offset); + gl_access_qualifier access = + (gl_access_qualifier)(ACCESS_NON_WRITEABLE | ACCESS_CAN_REORDER); + if (extra == ARM_WRITABLE_LOAD) + access = ACCESS_CAN_REORDER; + else if (extra == ARM_VOLATILE_LOAD) + access = (gl_access_qualifier)(access | ACCESS_VOLATILE); + nir_def *value = load_global(b, addr, access); + g.load = nir_def_as_intrinsic(value); + if (extra == ARM_STORE) + store_global(b, value, g.base); + nir_push_else(b, NULL); + nir_def *zero = nir_imm_int(b, 0); + nir_pop_if(b, NULL); + + store_global(b, nir_if_phi(b, value, zero), nir_iadd_imm(b, g.base, 64)); + return g; +} + +static nir_intrinsic_instr * +only_global_load(nir_shader *shader) +{ + nir_intrinsic_instr *found = NULL; + nir_foreach_function_impl(impl, shader) { + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type != nir_instr_type_intrinsic) + continue; + nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); + if (intr->intrinsic == nir_intrinsic_load_global || + intr->intrinsic == nir_intrinsic_load_global_constant) { + EXPECT_EQ(found, nullptr); + found = intr; + } + } + } + } + return found; +} + +static bool +is_bcsel_of(nir_def *def, nir_def *taken, uint64_t other) +{ + nir_alu_instr *sel = nir_def_as_alu_or_null(def); + return sel && sel->op == nir_op_bcsel && + sel->src[1].src.ssa == taken && + nir_src_is_const(sel->src[2].src) && + nir_src_as_uint(sel->src[2].src) == other; +} + +TEST(GuardedLoadSelect, ReadOnlyLoadReadsZeroPageWhenNotTaken) +{ + guarded g = build(MESA_SHADER_COMPUTE, ARM_LOAD_ONLY); + ASSERT_TRUE(agx_nir_speculate_guarded_loads(g.b.shader)); + nir_validate_shader(g.b.shader, "guarded load"); + + EXPECT_EQ(g.nif->control, nir_selection_control_flatten); + nir_intrinsic_instr *load = only_global_load(g.b.shader); + ASSERT_NE(load, nullptr); + EXPECT_EQ(load->intrinsic, nir_intrinsic_load_global_constant); + EXPECT_TRUE(nir_intrinsic_access(load) & ACCESS_CAN_SPECULATE); + EXPECT_TRUE(nir_intrinsic_access(load) & ACCESS_NON_WRITEABLE); + + /* base + u2u64(index * 4) keeps that shape for agx_nir_lower_address, + * with the base and the unscaled index selected separately. */ + nir_alu_instr *add = nir_def_as_alu_or_null(load->src[0].ssa); + ASSERT_NE(add, nullptr); + ASSERT_EQ(add->op, nir_op_iadd); + EXPECT_TRUE(is_bcsel_of(add->src[0].src.ssa, g.base, AGX_ZERO_PAGE_ADDRESS)); + nir_alu_instr *cvt = nir_def_as_alu_or_null(add->src[1].src.ssa); + ASSERT_NE(cvt, nullptr); + EXPECT_EQ(cvt->op, nir_op_u2u64); + nir_alu_instr *scale = nir_def_as_alu_or_null(cvt->src[0].src.ssa); + ASSERT_NE(scale, nullptr); + EXPECT_TRUE(is_bcsel_of(scale->src[0].src.ssa, g.index, 0)); + + /* A second run finds nothing more to do. */ + EXPECT_FALSE(agx_nir_speculate_guarded_loads(g.b.shader)); + ralloc_free(g.b.shader); +} + +static void +expect_unchanged(mesa_shader_stage stage, arm_extra extra, + bool dont_flatten = false) +{ + guarded g = build(stage, extra, dont_flatten); + EXPECT_FALSE(agx_nir_speculate_guarded_loads(g.b.shader)); + EXPECT_NE(g.nif->control, nir_selection_control_flatten); + EXPECT_EQ(only_global_load(g.b.shader), g.load); + EXPECT_EQ(g.load->intrinsic, nir_intrinsic_load_global); + EXPECT_FALSE(nir_intrinsic_access(g.load) & ACCESS_CAN_SPECULATE); + ralloc_free(g.b.shader); +} + +TEST(GuardedLoadSelect, ArmWithStoreKeepsBranch) +{ + expect_unchanged(MESA_SHADER_COMPUTE, ARM_STORE); +} + +TEST(GuardedLoadSelect, WritableLoadKeepsBranch) +{ + expect_unchanged(MESA_SHADER_COMPUTE, ARM_WRITABLE_LOAD); +} + +TEST(GuardedLoadSelect, VolatileLoadKeepsBranch) +{ + expect_unchanged(MESA_SHADER_COMPUTE, ARM_VOLATILE_LOAD); +} + +TEST(GuardedLoadSelect, DontFlattenIsRespected) +{ + expect_unchanged(MESA_SHADER_COMPUTE, ARM_LOAD_ONLY, true); +} + +TEST(GuardedLoadSelect, OnlyComputeShaders) +{ + expect_unchanged(MESA_SHADER_FRAGMENT, ARM_LOAD_ONLY); +} + +/* A descriptor read the driver already marked speculatable keeps its + * address; the if is still flattened for the other load. */ +TEST(GuardedLoadSelect, SpeculatableLoadKeepsAddress) +{ + nir_builder b = nir_builder_init_simple_shader( + MESA_SHADER_COMPUTE, &agx_nir_options, "guarded descriptor"); + nir_def *base = push_constant(&b, 64, 0); + nir_def *index = push_constant(&b, 32, 8); + nir_if *nif = nir_push_if(&b, nir_ult_imm(&b, index, 100)); + nir_def *desc = agx_test_load_global_constant( + &b, 1, 64, base, 8, + (gl_access_qualifier)(ACCESS_NON_WRITEABLE | ACCESS_CAN_REORDER | + ACCESS_CAN_SPECULATE)); + nir_def *value = load_global( + &b, nir_iadd(&b, desc, nir_u2u64(&b, index)), + (gl_access_qualifier)(ACCESS_NON_WRITEABLE | ACCESS_CAN_REORDER)); + nir_push_else(&b, NULL); + nir_def *zero = nir_imm_int(&b, 0); + nir_pop_if(&b, NULL); + store_global(&b, nir_if_phi(&b, value, zero), base); + + ASSERT_TRUE(agx_nir_speculate_guarded_loads(b.shader)); + nir_validate_shader(b.shader, "guarded descriptor"); + EXPECT_EQ(nif->control, nir_selection_control_flatten); + EXPECT_EQ(nir_def_as_intrinsic(desc)->src[0].ssa, base); + ralloc_free(b.shader); +} From 3ff99a1e7c83055d7165721ce1b9687634610344 Mon Sep 17 00:00:00 2001 From: Sean Dolbec Date: Sun, 1 Nov 2026 22:30:53 -0500 Subject: [PATCH 2/2] agx: keep guarded loads' alignment true for the zero-page fallback A redirected load kept its align_mul/align_offset, but when its arm is not taken it reads AGX_ZERO_PAGE_ADDRESS, which is 0 modulo any align_mul. For a load with align_mul 16 and align_offset 4 that made the alignment facts false, and the load/store vectorizer that runs after the pass may rely on them. Read zero page + align_offset instead. The zero page is 4 GiB aligned, so that address satisfies the load's facts, and an arm load is at most 16 bytes, so it stays inside the 64 KiB page. If it would not fit, weaken the facts to the largest power of two dividing align_offset, which both addresses satisfy. GuardedLoadSelect gains three tests: - FallbackKeepsAlignmentFacts: a 16/4 load keeps 16/4 and falls back to zero page + 4 (fails without this change). - BothArmsBecomeOneSelect: loads in both arms, run with nir_opt_peephole_select as agx_optimize_nir configures it. The if is gone, the result is bcsel(cond, then load, else load), and each load reads its own address when its arm is taken and the zero page otherwise; the else arm's guard is the reversed select. - PeepholeSelectAloneKeepsBranch: the same shader keeps its branch without this pass. Signed-off-by: Sean Dolbec Co-Authored-By: Claude Opus 5.5 --- src/asahi/compiler/agx_compile.c | 50 ++++-- .../test/test-guarded-load-select.cpp | 167 +++++++++++++++++- 2 files changed, 205 insertions(+), 12 deletions(-) diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c index 8891ad8d3428..7ae13ca760fb 100644 --- a/src/asahi/compiler/agx_compile.c +++ b/src/asahi/compiler/agx_compile.c @@ -7989,16 +7989,41 @@ agx_guarded_load_arm_ok(nir_block *block) return true; } -/* The address an arm's load uses when the arm is not taken: the zero page. +/* The zero-page address a load reads when its arm is not taken. The load + * keeps its alignment facts, which later passes (the load/store vectorizer) + * rely on, so they must hold for this address too. AGX_ZERO_PAGE_ADDRESS is + * 4 GiB aligned, a multiple of any align_mul, so the page base plus + * align_offset satisfies them and the read stays inside the page. If it + * would not, weaken the facts to what the page base satisfies: the largest + * power of two dividing align_offset, offset 0. + */ +static uint64_t +agx_guarded_fallback(nir_intrinsic_instr *intr, unsigned *align_mul, + unsigned *align_offset) +{ + STATIC_ASSERT((AGX_ZERO_PAGE_ADDRESS & ((1ull << 32) - 1)) == 0); + const unsigned bytes = intr->def.num_components * intr->def.bit_size / 8; + *align_mul = nir_intrinsic_align_mul(intr); + *align_offset = nir_intrinsic_align_offset(intr); + + if (*align_offset + bytes <= AGX_ZERO_PAGE_SIZE) + return AGX_ZERO_PAGE_ADDRESS + *align_offset; + + *align_mul = *align_offset & -*align_offset; + *align_offset = 0; + return AGX_ZERO_PAGE_ADDRESS; +} + +/* The address an arm's load uses: `fallback` when the arm is not taken. * Keep the base + 32-bit index shape agx_nir_lower_address folds into the * load, selecting the base and the unscaled index rather than the 64-bit - * sum: zero page + 0. + * sum: fallback + 0. */ static nir_def * agx_guarded_address(nir_builder *b, nir_def *addr, nir_def *cond, - bool taken_if_true) + bool taken_if_true, uint64_t fallback) { - nir_def *zero_page = nir_imm_int64(b, AGX_ZERO_PAGE_ADDRESS); + nir_def *fallback_addr = nir_imm_int64(b, fallback); #define AGX_GUARD(taken, other) \ (taken_if_true ? nir_bcsel(b, cond, taken, other) \ : nir_bcsel(b, cond, other, taken)) @@ -8025,12 +8050,12 @@ agx_guarded_address(nir_builder *b, nir_def *addr, nir_def *cond, } else { guarded = AGX_GUARD(index, zero); } - return nir_iadd(b, AGX_GUARD(base, zero_page), + return nir_iadd(b, AGX_GUARD(base, fallback_addr), nir_build_alu1(b, cvt->op, guarded)); } } - return AGX_GUARD(addr, zero_page); + return AGX_GUARD(addr, fallback_addr); #undef AGX_GUARD } @@ -8052,12 +8077,17 @@ agx_redirect_arm_loads(nir_block *block, nir_def *cond, bool taken_if_true) nir_builder b = nir_builder_at(nir_before_instr(instr)); nir_def *addr = intr->src[0].ssa; - if (!(access & ACCESS_CAN_SPECULATE)) - addr = agx_guarded_address(&b, addr, cond, taken_if_true); + unsigned align_mul = nir_intrinsic_align_mul(intr); + unsigned align_offset = nir_intrinsic_align_offset(intr); + if (!(access & ACCESS_CAN_SPECULATE)) { + const uint64_t fallback = + agx_guarded_fallback(intr, &align_mul, &align_offset); + addr = agx_guarded_address(&b, addr, cond, taken_if_true, fallback); + } nir_def *value = nir_load_global_constant( &b, intr->def.num_components, intr->def.bit_size, addr, - .align_mul = nir_intrinsic_align_mul(intr), - .align_offset = nir_intrinsic_align_offset(intr), + .align_mul = align_mul, + .align_offset = align_offset, .access = access | ACCESS_CAN_SPECULATE); nir_def_replace(&intr->def, value); loads = true; diff --git a/src/asahi/compiler/test/test-guarded-load-select.cpp b/src/asahi/compiler/test/test-guarded-load-select.cpp index 51f292f4dc1a..cc7d8c9ad14b 100644 --- a/src/asahi/compiler/test/test-guarded-load-select.cpp +++ b/src/asahi/compiler/test/test-guarded-load-select.cpp @@ -24,14 +24,18 @@ push_constant(nir_builder *b, unsigned bit_size, unsigned offset) } static nir_def * -load_global(nir_builder *b, nir_def *addr, gl_access_qualifier access) +load_global(nir_builder *b, nir_def *addr, gl_access_qualifier access, + unsigned align_mul = 4, unsigned align_offset = 0) { nir_def *def = nir_load_global(b, 1, 32, addr); nir_intrinsic_set_access(nir_def_as_intrinsic(def), access); - nir_intrinsic_set_align(nir_def_as_intrinsic(def), 4, 0); + nir_intrinsic_set_align(nir_def_as_intrinsic(def), align_mul, align_offset); return def; } +static const gl_access_qualifier read_only = + (gl_access_qualifier)(ACCESS_NON_WRITEABLE | ACCESS_CAN_REORDER); + static void store_global(nir_builder *b, nir_def *value, nir_def *addr) { @@ -216,3 +220,162 @@ TEST(GuardedLoadSelect, SpeculatableLoadKeepsAddress) EXPECT_EQ(nir_def_as_intrinsic(desc)->src[0].ssa, base); ralloc_free(b.shader); } + +/* bcsel(cond, taken, other) for the then arm; the else arm's guard is the + * same select with its operands swapped. */ +static bool +is_guard(nir_def *def, nir_def *cond, nir_def *taken, uint64_t other, + bool then_arm) +{ + nir_alu_instr *sel = nir_def_as_alu_or_null(def); + if (!sel || sel->op != nir_op_bcsel || sel->src[0].src.ssa != cond) + return false; + const nir_src &t = sel->src[then_arm ? 1 : 2].src; + const nir_src &o = sel->src[then_arm ? 2 : 1].src; + return t.ssa == taken && nir_src_is_const(o) && nir_src_as_uint(o) == other; +} + +/* base + u2u64(index * scale), with the base and the unscaled index guarded + * separately: when the arm is not taken the sum is exactly `fallback`. */ +static void +expect_guarded_address(nir_intrinsic_instr *load, nir_def *cond, nir_def *base, + nir_def *index, uint64_t fallback, bool then_arm) +{ + SCOPED_TRACE(then_arm ? "then arm" : "else arm"); + ASSERT_EQ(load->intrinsic, nir_intrinsic_load_global_constant); + EXPECT_TRUE(nir_intrinsic_access(load) & ACCESS_CAN_SPECULATE); + nir_alu_instr *add = nir_def_as_alu_or_null(load->src[0].ssa); + ASSERT_NE(add, nullptr); + ASSERT_EQ(add->op, nir_op_iadd); + EXPECT_TRUE(is_guard(add->src[0].src.ssa, cond, base, fallback, then_arm)); + nir_alu_instr *cvt = nir_def_as_alu_or_null(add->src[1].src.ssa); + ASSERT_NE(cvt, nullptr); + ASSERT_EQ(cvt->op, nir_op_u2u64); + nir_alu_instr *scale = nir_def_as_alu_or_null(cvt->src[0].src.ssa); + ASSERT_NE(scale, nullptr); + EXPECT_TRUE(is_guard(scale->src[0].src.ssa, cond, index, 0, then_arm)); +} + +static unsigned +count_ifs(nir_shader *shader) +{ + unsigned n = 0; + nir_foreach_function_impl(impl, shader) { + nir_foreach_block(block, impl) { + nir_cf_node *next = nir_cf_node_next(&block->cf_node); + n += next && next->type == nir_cf_node_if; + } + } + return n; +} + +static bool +peephole_select(nir_shader *shader) +{ + /* The options agx_optimize_nir uses; indirect loads stay not OK. */ + nir_opt_peephole_select_options options = {}; + options.limit = 64; + options.expensive_alu_ok = true; + return nir_opt_peephole_select(shader, &options); +} + +struct two_arms { + nir_builder b; + nir_def *cond, *base[2], *index[2]; + nir_intrinsic_instr *store; +}; + +/* out = cond ? x[i] : y[j], with x and y read-only. */ +static two_arms +build_two_arms() +{ + two_arms t; + t.b = nir_builder_init_simple_shader(MESA_SHADER_COMPUTE, &agx_nir_options, + "guarded loads in both arms"); + nir_builder *b = &t.b; + t.base[0] = push_constant(b, 64, 0); + t.base[1] = push_constant(b, 64, 8); + t.index[0] = push_constant(b, 32, 16); + t.index[1] = push_constant(b, 32, 20); + t.cond = nir_ult_imm(b, t.index[0], 100); + + nir_push_if(b, t.cond); + nir_def *x = load_global( + b, nir_iadd(b, t.base[0], nir_u2u64(b, nir_imul_imm(b, t.index[0], 4))), + read_only); + nir_push_else(b, NULL); + nir_def *y = load_global( + b, nir_iadd(b, t.base[1], nir_u2u64(b, nir_ishl_imm(b, t.index[1], 3))), + read_only, 8, 0); + nir_pop_if(b, NULL); + + nir_def *out = nir_if_phi(b, x, y); + t.store = nir_store_global(b, out, nir_iadd_imm(b, t.base[0], 64)); + nir_intrinsic_set_align(t.store, 4, 0); + nir_intrinsic_set_write_mask(t.store, 1); + return t; +} + +/* Together with peephole select, the branch disappears: each arm's load + * reads its own address when its arm is taken and the zero page otherwise + * (the else arm's guard is reversed), and the result selects between the + * two loads on the original condition. */ +TEST(GuardedLoadSelect, BothArmsBecomeOneSelect) +{ + two_arms t = build_two_arms(); + ASSERT_TRUE(agx_nir_speculate_guarded_loads(t.b.shader)); + ASSERT_TRUE(peephole_select(t.b.shader)); + nir_validate_shader(t.b.shader, "guarded loads in both arms"); + EXPECT_EQ(count_ifs(t.b.shader), 0u); + + nir_alu_instr *sel = nir_def_as_alu_or_null(t.store->src[0].ssa); + ASSERT_NE(sel, nullptr); + ASSERT_EQ(sel->op, nir_op_bcsel); + EXPECT_EQ(sel->src[0].src.ssa, t.cond); + for (unsigned arm = 0; arm < 2; ++arm) { + nir_def *value = sel->src[1 + arm].src.ssa; + ASSERT_EQ(nir_def_instr(value)->type, nir_instr_type_intrinsic); + expect_guarded_address(nir_def_as_intrinsic(value), t.cond, + t.base[arm], t.index[arm], AGX_ZERO_PAGE_ADDRESS, + arm == 0); + } + ralloc_free(t.b.shader); +} + +/* Without the pass, peephole select keeps both loads under the branch. */ +TEST(GuardedLoadSelect, PeepholeSelectAloneKeepsBranch) +{ + two_arms t = build_two_arms(); + EXPECT_FALSE(peephole_select(t.b.shader)); + EXPECT_EQ(count_ifs(t.b.shader), 1u); + ralloc_free(t.b.shader); +} + +/* The load keeps its alignment facts, so the fallback address must satisfy + * them too: align_mul 16 / align_offset 4 reads zero page + 4. */ +TEST(GuardedLoadSelect, FallbackKeepsAlignmentFacts) +{ + nir_builder b = nir_builder_init_simple_shader( + MESA_SHADER_COMPUTE, &agx_nir_options, "guarded aligned load"); + nir_def *base = push_constant(&b, 64, 0); + nir_def *index = push_constant(&b, 32, 8); + nir_def *cond = nir_ult_imm(&b, index, 100); + nir_push_if(&b, cond); + nir_def *value = load_global( + &b, nir_iadd(&b, base, nir_u2u64(&b, nir_imul_imm(&b, index, 16))), + read_only, 16, 4); + nir_push_else(&b, NULL); + nir_def *zero = nir_imm_int(&b, 0); + nir_pop_if(&b, NULL); + store_global(&b, nir_if_phi(&b, value, zero), base); + + ASSERT_TRUE(agx_nir_speculate_guarded_loads(b.shader)); + nir_validate_shader(b.shader, "guarded aligned load"); + nir_intrinsic_instr *load = only_global_load(b.shader); + ASSERT_NE(load, nullptr); + EXPECT_EQ(nir_intrinsic_align_mul(load), 16u); + EXPECT_EQ(nir_intrinsic_align_offset(load), 4u); + expect_guarded_address(load, cond, base, index, AGX_ZERO_PAGE_ADDRESS + 4, + true); + ralloc_free(b.shader); +}