[techmore]

What

row < m ? x[i] : 0 over a read-only buffer stays a branch on every chip. 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.

A new NIR pass, agx_nir_speculate_guarded_loads, runs in agx_optimize_nir before soft-fault speculation, for compute shaders only. It looks for an if whose arms are single blocks of at most 64 instructions, made of ALU, load_preamble and read-only global loads:

• Each such load reads the zero page when its arm is not taken. AGX_ZERO_PAGE_ADDRESS is always mapped. The base and the unscaled 32-bit index are selected instead of the 64-bit sum, so agx_nir_lower_address still folds base + index into the load.
• Loads the driver already marked speculatable keep their address. These are hk's descriptor reads.
• A read-only, reorderable load_global becomes load_global_constant, the form peephole select accepts.
• Only that if gets nir_selection_control_flatten, so the peephole select that follows turns it into selects.
• Other indirect loads keep their branches, so this does not turn on indirect_load_ok globally. That flag stays off because of the existing XXX about CTS flakes.

ASAHI_NO_GUARDED_LOAD_SELECT=1 turns the pass off for A/B runs.

Why

This change is behind our largest prefill gain on M4. All numbers here are from our receipts, measured on an M4 Mac mini (G16G) with GravityLinux's Apple9 backend on a kernel without soft faults, not on this tree:

• qmm_fma 4-bit matmul went from 174 to 778 GFLOP/s.
• f16 prefill went up 4.5-4.8x (Llama 3.2 1B, 512 tokens: 67 → 310 tok/s), with identical greedy output.

That backend's register allocator paid far more for the joins than this tree's does: 5300 of 6412 instructions were copies. So expect a smaller effect here. The static numbers below are what we can show for this tree without hardware.

Evidence (compile-only, this tree)

drm-shim on an M4 Mac mini host (aarch64 Arch Linux), earlier round, same base. No GPU submits, -Db_ndebug=false (assertions and NIR validation on), shader cache off, robustBufferAccess off. Profiles ASAHI_SHIM_CHIP=g15s / g17p / g13g. Instruction counts come from AGX_MESA_DEBUG=shaders, with the pass on vs. ASAHI_NO_GUARDED_LOAD_SELECT=1:

Kernel / set	Chip	Instructions	Blocks	Moves
qmm_fma.comp	G17P	3687 → 3523	429 → 249	542 → 350
qmm_fma.comp	G15S	3690 → 3620	429 → 249	545 → 447
106 omarchy-mlx shaders	G15S	161180 → 159075 (-1.3%)	19920 → 17886
106 omarchy-mlx shaders	G17P	156231 → 153299 (-1.9%)	19920 → 17886
106 omarchy-mlx shaders	G13G	153260 → 151934 (-0.9%)	18768 → 16734

• Of the 126 omarchy-mlx compute shaders, 106 compile on all three profiles both ways.
  • 8 do not build with our standalone glslang setup.
  • 12 (cooperative-matrix kernels) hit the same existing should've been lowered assertion with and without the pass on these chips.
  • No new failure or validation error.
• 44 shaders change. Most shrink (reduce_general 1965 → 1334 on G17P). Some grow, by up to 7.5% (fast_norm_vjp_dx 828 → 890 on G15S; int8_matmul 2046 → 2103 on G13G), where flattening trades a branch for selects.
• In the guarded loop of a small probe, the G17P loop body becomes one block with no branches.

Tests run

• New this round: GuardedLoadSelect in agx_tests (test/test-guarded-load-select.cpp, 7 tests).
  • The rewrite: the zero-page base and the unscaled index are selected, the result is load_global_constant with ACCESS_CAN_SPECULATE, the if is flattened, and a second run finds nothing left to do.
  • Arms left alone: one with a store, a writable load, a volatile load, dont_flatten, a fragment shader.
  • A descriptor read that is already speculatable keeps its address.
• 2026-11-01, macOS (Apple Silicon host, Apple clang 21):
  • agx_tests on this branch: 509 tests, 508 passed, 1 skipped (base: 502, 501 passed).
  • Merged with PRs 1 and 2: 514, 513 passed, 1 skipped.
  • 17/17 host compiler, packing and layout suites pass under meson test.
• Earlier round (Mini, Linux): the Vulkan driver, drm-shim and agx_tests built (496 passed + 1 skipped in release; 502 + 1 skipped with assertions). The assertion build caught one bug before that version: nir_intrinsic_access on load_preamble. It is fixed.

Remaining test plan

• Linux, drm-shim (mini-run/run-drm-shim.sh): meson test (all suites), drm-smoke.py, and pipelines.py with the pass on and with ASAHI_NO_GUARDED_LOAD_SELECT=1.
• Hardware: A/B with ASAHI_NO_GUARDED_LOAD_SELECT=1 on M3 (no soft faults) and A18 Pro (soft faults):
  • omarchy-mlx qmm_fma GFLOP/s and prefill tok/s.
  • A kernel whose guards are mostly false, to see what the always-issued loads cost in bandwidth.
  • CTS dEQP-VK.compute.* and dEQP-VK.robustness.* with robustBufferAccess off and on.

Risks and limitations

• Not run on a GPU with this tree. Static counts are not run time: flattened arms always issue their loads, which costs bandwidth when guards are mostly false.
• It relies on the zero page being mapped and readable from compute on every chip that runs this compiler. agx_compile.c already assumes that when it replaces addresses with AGX_ZERO_PAGE_ADDRESS.
• It applies to M1/M2 compute too (AGX2 shares agx_optimize_nir). The G13G numbers are mixed: 15 shaders shrink, 28 grow, total -0.9%. Gating it to native AGX3 is a one-line change if preferred.
• It treats read-only, reorderable load_global as load_global_constant. Both lower to the same AGX load, and the reorder flag was already set.
• Only compute shaders. Robust (bounded) loads are not touched.

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; both previous findings are addressed and no new blocking issue was found.

Summary

The PR lets short compute-shader branches issue read-only loads safely before selecting their results.

• Eligible compute shaders now turn guarded read-only loads into selects.

_{Reviews (2) · Last reviewed commit: "agx: keep guarded loads' alignment true ..." · Reviewed by Greptile}

[techmore]

Aurora downstream draft: aurora-silicon#193