Wave64 compute and mesh programs at one lane per invocation on 32-wide hosts when they never observe the wave

[karanvora2599]

Problem

On a host whose subgroup is 32 wide (NVIDIA), every wave64 compute and mesh program runs two guest lanes per host invocation:

• laneCount is 2 whenever the program is wave64 and the host subgroup is 32 (SpirvEmitter.cpp), whatever the program does.
• The workgroup is halved and flattened to 1D (SpirvModuleSetup.cpp), and every lane-varying value is emitted once per lane (SpirvFlowEmitter.cpp).

#304 already shares the wave-uniform values between the two lanes, but each host invocation still holds both lanes' VGPRs, and NVIDIA caps an invocation at 255 registers. #611 measured PPSA21564 mesh shaders at 255 registers with 624-784 B of local memory per thread, and 448 B left even with every exec select removed. On a 64-wide host the same programs run one lane per invocation. In #357, PPSA21564 runs were pinning one wave64 compute program (0x500597a00) to another layout with APS5_TWO_LANE/APS5_INEXACT_SINGLE_LANE to avoid an NVIDIA spill hang; neither switch is on main now.

Two lanes per invocation are required only when the program looks at the wave as a whole. Many wave64 programs never do: each lane reads its inputs, computes, and writes its own results.

Proposal

Choose the layout per program. A wave64 compute or mesh program on a 32-wide host runs at one guest lane per invocation (the layout APS5_SINGLE_LANE forces today, Recompiler.cpp), with the full workgroup and a required subgroup size of 32, when an IR analysis proves that no result depends on the other half of the wave:

1. No op that reads other lanes or the wave mask: ballots, EXEC or VCC read as a value, v_readlane/v_readfirstlane of a lane-varying value, v_writelane, DPP, v_permlane*, ds_swizzle, ds_permute/ds_bpermute, v_mbcnt, ds_append/ds_consume.
2. Wave-level branches guard only per-lane work. s_cbranch_execz/vccz (or SCC from a lane mask) is taken per host subgroup, so one half can skip a region the other half runs. That is invisible only if the region has no scalar effect seen after it: no SALU result live out of the region, no scalar store, no message. A conservative first version rejects any such region with a scalar definition used later.
3. LDS: no LDS writes, or a single-wave workgroup. Lanes of one wave read each other's LDS writes without s_barrier; across two host subgroups that needs a barrier, which WaveLdsScope can only place when the workgroup is one wave.

Every other wave64 program keeps two lanes per invocation. Nothing changes for wave32 or on 64-wide hosts. The choice depends only on the program and the target, so the cache key needs nothing new.

Measurement plan for the PR: APS5_PIPELINE_STATS=1 register count and local memory per program before and after, the share of wave64 programs that qualify, and execution tests for each rejected case (a ballot, an execz region with a live SALU result, intra-wave LDS exchange in a multi-wave workgroup) next to accepted ones.

Questions

• Is a per-program lane layout acceptable, or does it conflict with anything planned for the recompiler?
• @oneandonlydean, the perf(recompiler): emit wave-uniform values once for both lanes of a two-lane invocation #304 commit message mentions a SingleLane analysis and Astro Bot (PPSA21564) #357 used APS5_INEXACT_SINGLE_LANE. Is a per-program layout already on your branch or planned? I don't want to duplicate it.
• Rule 2: is rejecting any wave-level branch over a scalar definition used later conservative enough, or is there a case where the halves can also diverge through a uniform value that I am missing?
• APS5_SINGLE_LANE=all with APS5_PIPELINE_STATS=1 gives an upper bound on the register and spill savings without any change. Rendering may differ under it, since it ignores rules 1-3, so it is a measurement only. If anyone with a title that has many wave64 compute or mesh programs (PPSA21564, for example) can post the register and local memory numbers from both runs, that would show whether this is worth doing.

[oneandonlydean]

No. A per-program wave64 layout isn't on any branch of mine and I'm not planning one, so it's yours. My performance work has moved off my fork onto a tree based on upstream main. Recompiler changes are reviewed by @adriadam10.

[karanvora2599]

Thanks @oneandonlydean, I'll take it.

Here is where I'd start, with a question about the design for @adriadam10 before I write it.

What APS5_SINGLE_LANE=all does to the tests today. I ran the shader tests on an RTX 4060 with every wave64 program at one lane per invocation: 285 of 308 pass. The BDA tests that hang on driver 591.86 were left out. Of the 22 new failures:

• Programs the rules would reject:
  • DPP (dpp8, dpp16_inactive_source, dpp16_row_share_xmask), ds_swizzle, permlane16, read_first_lane, wave_uniform, look_back_scan (ballot) and ngg_vertex_wave_reduction;
  • gds_append_consume;
  • the prepared_shaders tests (v_readlane_b32 of lane 63; the emitter already refuses it);
  • scalar_program_flow: its execnz_upper_half row is exactly rule 2, a wave-mask branch over a scalar result that the halves decide differently.
• Gaps in the one-lane emitter, which I'd fix rather than reject:
  • lds_addtid, global_store_addtid and global_vcc_base get a lane id of 0–31 per host subgroup instead of 0–63 per wave.
  • sdwa_move_selectors gets lane 36 wrong.
  • mesh renders a triangle in the wrong colour.
    I haven't looked into the last two yet.
• Compile constraint: agc_shader_bindless and agc_shader_memory (CPU tests) stop compiling. A bindless slot that is uniform per wave isn't uniform per host subgroup once the wave spans two of them, and their target has no non-uniform image indexing.

Analysis. I'd do it on the IR rather than the RDNA program, where the wave-level operations are explicit:

1. Reject Ballot, WriteLane, DppMoveU32, Permlane16U32, SwizzleU32, LDS/GDS append/consume, and ReadLane/ReadFirstLane of a value WaveUniformValues doesn't prove uniform.
2. For every block ending in a wave-mask branch (IsWaveMaskBranch), and for loops that exit on one, reject when a wave-uniform value defined in the region it controls is used after the region, or the region has a scalar memory side effect.
3. Reject LDS writes in a workgroup of more than one wave, as WaveLdsScope can't order them there.
4. Reject bindless tables indexed by non-constant slots when the target lacks non-uniform indexing.

Compute only at first. Mesh stays at two lanes until I've understood the mesh failure.

The question: where the decision lives. HostSubgroupSize(request) is read before translation (RequestInputInfo, translation at InstructionTranslator.cpp:192/:368), for result.hostSubgroupSize, and in the cache keys. The analysis needs the IR, so I see two ways:

• (a) Translate as today, analyse, and if the program qualifies, translate again with the 64-wide input info. Carry the choice in the IR program's metadata, so compileVariant and result.hostSubgroupSize use it instead of calling HostSubgroupSize(request) again. Cache keys stay on the target's subgroup size, since the choice is a function of the program and the target.
• (b) Do the analysis on the decoded program and CFG (cached by AcquireFrontend), so that HostSubgroupSize(request) itself returns the choice everywhere. It's cheaper, but rules 1–2 are harder to get right on RDNA than on IR.

I'd go with (a). Is carrying the layout in the program's metadata fine with you, or would you rather keep it a function of the request?

Measurement. In the PR: register count and local memory per program from APS5_PIPELINE_STATS, the share of wave64 programs that qualify, and execution tests for each rejected and accepted case. I can't measure titles here; numbers from PPSA21564 or another wave64-heavy title would still help.

[Raghav314]

@karanvora2599 some title data for this, from Demon's Souls (PPSA01341) on my machine (RTX 4090, Linux). I dumped every shader it registers in its first two minutes (2,388 requests, every one of them wave64: 2,106 compute, 199 vertex, 83 pixel) and recompiled them offline with agc_shader_replay, once as is and once with APS5_SINGLE_LANE=all. The build is #4152's head with two small fixes I posted in #2242, plus #4049 so the game gets through start-up.

These are SPIR-V word counts only, not APS5_PIPELINE_STATS registers or local memory, and APS5_SINGLE_LANE=all ignores your three rules, so it's an upper bound. For the 58 largest compute shaders:

	two lanes (as on main)	one lane
total	36,449,268 words	17,137,939 words (-53%)
per shader		31% to 74% of the two-lane size, mean 47%
over 200k words	56	42
largest	1,372,054	716,421

Why it matters here: the black screen this title shows on a cold start on NVIDIA is the queue waiting on compute pipeline creates. With the driver cache off, one 944k-word shader takes 7.2 s to create on the 4090, and a 785k one 16.8 s. That second one also uses the relative VGPR window, and most of its time goes to that.

For one shader (0x249dff800, 9,788 code words, 1,252,474 SPIR-V words) I also split the rest: one lane takes it to 655,087, and keeping f32 denormals on top (measurement only, the title's FLOAT_MODE flushes them) to 427,601. So after one lane, the f32 denormal handling is the next biggest part, and with your #2103 already in, I couldn't find much of it that's provably redundant.

I can rerun these requests against your branch when it's up.

LLM assistance used.

[oneandonlydean]

@Raghav314 same result on another title here (RTX 4090, Linux, main d0ca839). 56 wave64 programs dumped with APS5_DUMP_SHADERS, recompiled with agc_shader_replay as is and with APS5_SINGLE_LANE=all. SPIR-V word counts only, and the same upper bound as yours.

	two lanes (as on main)	one lane
49 compute	9,426,728 words	5,281,608 words (56%)
per compute shader		51% to 93% of the two-lane size, mean 58%
7 graphics	2,892,366 words	1,645,715 words (57%)
largest	975,572	509,122

The twelve largest compute shaders are all at 51% to 55%. No pipeline create times or APS5_PIPELINE_STATS from this run.

[karanvora2599]

Thanks @Raghav314, this is the data the proposal was missing. A halved word count lines up with how the two-lane layout works, since every lane-varying instruction is emitted once per half. The pipeline-create times make compile time as strong a reason as register pressure.

The per-program version is up as #4339. It checks the IR after translation, and a wave64 compute program goes to one lane per invocation unless one of these applies:

• it reads other lanes or the wave mask: DPP, v_permlane*, ds_swizzle, ds_[b]permute, ds_append/ds_consume, v_readlane, v_readfirstlane of a lane-varying value, or a ballot used for more than each lane's own bit. Exec masks saved in SGPRs and read back, the usual s_and_saveexec pattern, stay allowed;
• a wave-mask branch (execz/execnz/vccz/vccnz) guards a region that defines a scalar value used after it, or that holds a barrier, or, in a program that writes LDS, any LDS access;
• it writes LDS in a workgroup of more than one wave;
• it indexes a bindless table and the device lacks non-uniform indexing.

Mesh programs stay at two lanes for now. On my RTX 4060 the shader suite passes, and 181 of its wave64 compute compiles go to one lane. The rest are mostly tests that exercise those operations on purpose.

The PR also measures compile time on a synthetic wave64 program of about 65K instructions. NVIDIA pipeline creation drops from 12–283 s to 6–48 s; it's worst for two lanes once a thread hits 255 registers.

APS5_WAVE64_LAYOUT_LOG=1 prints one line per program with its decision and, for a rejection, the opcode that caused it. A rerun of your 58 compute shaders with it on would tell us how many qualify, how much of the 53% they keep, and which rule rejects the rest. That shows which rule to relax first. On the suite the most common blocker is v_readfirstlane of a lane-varying value. If that's true of titles too, a waterfall loop that scalarizes one lane at a time could run per subgroup, and I'd look at that next.

On f32 denormals, there's one redundant part I know of. #2103 removes an input flush when the value can't be a denormal, but it doesn't recognize a flushed v_mul_f32/v_fma_f32 result as denormal-free. So everything that consumes a product flushes it again. That's my next change. Flushes on values loaded from memory, user data or inputs have to stay while NVIDIA has no f32 flush-to-zero mode.

[Celegans12]

I ran #4339's decision log over Ghost of Yōtei's compute shaders and posted the numbers there (#4339 (comment)). It mostly confirms the size pattern. On main, 156 of 362 Yōtei compute programs take one lane, and that's most of the small ones. Yōtei's giant pipelines (over 500k words, 80 to 470 s each in NVIDIA's compiler) almost all keep two lanes, because they use real wave operations: ballot counts, waterfall loops and DPP/v_permlane16 reductions. That puts the saving on them at 0.6% instead of the 45.5% APS5_SINGLE_LANE=all shows. Where a program does flip, a cold pipeline create is 2.5 to 3 times faster (14.9 s to 5.8 s for a 1.08M-word shader).

[karanvora2599]

Thanks @Celegans12, I replied on #4339. In short: the four rule-2-only programs would help, and I'll check whether their phis really depend on the mask. The giants would need ballot counts and reductions passed between the two halves through shared memory and barriers, which waterfall loops can't do, so the one-lane layout won't reach them.