diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c index 5443f4e509f0..7ae13ca760fb 100644 --- a/src/asahi/compiler/agx_compile.c +++ b/src/asahi/compiler/agx_compile.c @@ -7933,6 +7933,203 @@ 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 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: fallback + 0. + */ +static nir_def * +agx_guarded_address(nir_builder *b, nir_def *addr, nir_def *cond, + bool taken_if_true, uint64_t fallback) +{ + 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)) + + 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, fallback_addr), + nir_build_alu1(b, cvt->op, guarded)); + } + } + + return AGX_GUARD(addr, fallback_addr); +#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; + 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 = align_mul, + .align_offset = align_offset, + .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 +8333,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..cc7d8c9ad14b --- /dev/null +++ b/src/asahi/compiler/test/test-guarded-load-select.cpp @@ -0,0 +1,381 @@ +/* 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, + 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), 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) +{ + 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); +} + +/* 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); +}