[Vulkan] Move adstack-sizer scratch out of Function-scope memory to fix SPIR-V pipeline build failures

[duburcqa]

Move the on-device adstack SizeExpr sizer's per-invocation interpreter state from Function-scope OpVariables into two StorageBuffer scratch SSBOs

Three-commit PR. The headline fix is the SSBO move (commit 3): on Blackwell-class NVIDIA Vulkan drivers (RTX 5090, RTX 6000 Blackwell) the cumulative i64[kMaxNodes=4096] + i64[kMaxVars] + 8x array[kMaxPending] (~34 KiB) of per-thread private memory the sizer shader reserved was over the driver's hidden compute-shader stack-frame budget, and vkCreateComputePipelines failed with VK_ERROR_UNKNOWN. The two earlier commits address a separate but related strict-validation issue the same Blackwell drivers also enforce on Function- / Workgroup-scope OpVariable types under VUID-StandaloneSpirv-None-10684 (Invalid explicit layout decorations on type), surfaced first by the same RTX 5090 reproducer.

TL;DR

On a Blackwell GPU running any reverse-mode kernel whose SizeExpr requires the on-device sizer, publish_adstack_metadata_spirv failed at pipeline creation. The chain of root causes uncovered while debugging:

1. VUID-StandaloneSpirv-None-10684 on the sizer's Function-scope array OpVariables. The SPIR-V validator that ships with Vulkan 1.4 Blackwell drivers refuses any layout decoration on a Function / Private / Workgroup variable type. The sizer's IRBuilder::alloca_variable(get_array_type(i64, 4096)) produced a Function-scope variable whose pointee type carried DecorationArrayStride, tripping the validator. Fix: new IRBuilder::get_function_array_type helper that reuses the existing get_array_type body but drops the ArrayStride decoration, plus a parallel migration in spirv_codegen.cpp:325 for user-kernel tensor allocas in Function / Workgroup storage that hit the same VUID. (commits 1 + 2)

2. VK_ERROR_UNKNOWN from the driver compiler on the sizer pipeline build. With (1) fixed the SPIR-V module validated cleanly but the Blackwell driver's NV-internal SASS lowering still rejected the shader. Bisecting on kMaxNodes showed the cap was the trigger: at 4096 (32 KiB of per-thread i64[] private memory) the build failed; at 256 (2 KiB) it succeeded. The driver caps the per-thread compute-shader stack frame at a budget that varies by SM target, and Blackwell sets it lower than older NVIDIA generations. Fix: rehost every per-invocation array into two StorageBuffer SSBOs (binding 3 = i64 scratch, binding 4 = i32 scratch), allocated once in GfxRuntime and bound on every sizer dispatch. The sizer is 1x1x1 so cross-thread aliasing is impossible. (commit 3)

3. kAdStackSizerMaxNodesPerStack = 4096 was sized to fit per-thread private memory. With the SSBO move that constraint is gone; the cap is now backed by ~512 KiB of device memory per GfxRuntime and exists only to keep the host encoder's hard-error attached to a fixed compile-time ceiling. Raised to 65536 so reverse-mode kernels with deep symbolic trees no longer hit the cap. (commit 4)

Reproduces on:

• RTX 5090 (consumer Blackwell)
• RTX 6000 Blackwell (data-center)

Why

The sizer shader was originally written assuming Function-storage OpVariables were a free-of-charge per-invocation scratch space, and the cumulative ~34 KiB sat well under what older NVIDIA / AMD / Apple Vulkan drivers tolerated. Blackwell drivers ship a stricter SPIR-V validator (catching the layout-decoration-on-Function-variable case) and a tighter per-thread compute-shader stack-frame budget (catching the 32 KiB cumulative). Both have to be addressed; (1) alone leaves the driver still rejecting at SASS lowering with VK_ERROR_UNKNOWN.

Mechanism

IRBuilder::get_function_array_type (commit 1)

New helper in quadrants/codegen/spirv/spirv_ir_builder.{h,cpp}. Emits the same OpTypeArray (or OpTypeRuntimeArray) and id-allocation bookkeeping as get_array_type, but skips the ArrayStride decoration. get_array_type now delegates the emission to get_function_array_type and adds the decoration on top, so the two stay in sync by construction (no duplicated body).

User-kernel tensor allocas (commit 2)

TaskCodegen::visit(AllocaStmt) at quadrants/codegen/spirv/spirv_codegen.cpp:325 was building a tensor-type array via get_array_type and feeding the result into either alloca_workgroup_array (Workgroup storage) or alloca_variable (Function storage). Both storage classes are subject to VUID-StandaloneSpirv-None-10684. Switched to get_function_array_type. Same commit also fixes a u1->i32 rewrite regression in get_array_type's body that an earlier draft of get_function_array_type introduced (the rewrite was applied only on the local copy inside the helper, leaving the ArrayStride computation upstream reading the unrewritten u1 and emitting ArrayStride 1 on a 4-byte-element array).

Sizer scratch SSBO migration (commit 3)

The sizer's ShaderState previously held nine Function-storage Values for the interpreter scratch (values_arr, scope_arr, eight pending_*_arrs). Now hosted as compile-time element offsets into two SSBOs:

• Binding 3: int64[] scratch, layout [values_arr | scope_arr | pending_cur_i | pending_end | pending_max_accum]. Total elements = kAdStackSizerScratchI64Elems.
• Binding 4: int32[] scratch, layout [pending_mor_idx | pending_body_start | pending_body_end | pending_var_id | pending_saved_max_k]. Total elements = kAdStackSizerScratchI32Elems.

array_i64_access_ptr / array_i32_access_ptr now take the SSBO + a compile-time base index + a runtime index, route through struct_array_access, and return a StorageBuffer pointer the existing load_variable / store_variable helpers consume unchanged. GfxRuntime allocates both buffers once on first sizer dispatch and reuses them for the lifetime of the runtime; the sizer is 1x1x1 so there is no cross-thread aliasing through the buffers.

kMaxNodes cap raise (commit 4)

Bumped kAdStackSizerMaxNodesPerStack from 4096 to 65536. The cap concept is preserved (the host-side encoder hard-errors above the ceiling so an unexpectedly deep symbolic tree surfaces as an attributable encode-time diagnostic rather than as an out-of-bounds shader access), but at the new value the per-GfxRuntime device-memory allocation is 512 KiB plus the i32 buffer - negligible against any modern GPU's VRAM budget. Reverse-mode kernels with up to 64k symbolic tree nodes per stack now go through the on-device sizer instead of failing.

Per-driver coverage matrix

Driver	Pre-PR sizer pipeline build	Post-PR
NVIDIA Blackwell (RTX 5090, RTX 6000 Blackwell) Vulkan	Fails at vkCreateComputePipelines: VUID-StandaloneSpirv-None-10684 then VK_ERROR_UNKNOWN	Passes
NVIDIA Ada / Ampere / Turing Vulkan	Passes (older drivers tolerated both issues)	Passes
AMD ROCm Vulkan	Passes	Passes
Apple Metal / MoltenVK	Passes	Passes
LLVM CPU / CUDA / AMDGPU paths	Unaffected (no SPIR-V)	Unaffected

User-kernel SPIR-V path	Pre-PR on Blackwell	Post-PR
Tensor / shared-array allocas (spirv_codegen.cpp:325)	Fails VUID-StandaloneSpirv-None-10684 the same way the sizer did	Passes (commit 2 migration)

Tests

• pytest tests/python/test_adstack.py -n 8: 753 passed, 10 xfailed locally on macOS Vulkan / arm64 / Metal (each commit verified independently).
• The Blackwell reproducer test_ad_ndarray_torch.py::test_tape_torch_tensor_grad_none[arch=vulkan] should pass on RTX 5090 / RTX 6000 Blackwell once this lands.

Side-effect audit

Concern	Where checked	Verdict
Storage-buffer / PSB / Uniform interface arrays still get ArrayStride	get_array_type body delegates to get_function_array_type then adds the decoration; every existing call site of get_array_type keeps the decorated path	Identical to pre-PR for those callers
u1->i32 rewrite for layout-decorated array types	get_array_type re-applies the rewrite at the top before computing nbytes, mirroring pre-PR behaviour	Identical to pre-PR
Sizer SSBO scratch layout vs in-shader access pattern	static_asserts pin the i64 / i32 element-count totals against the shader-side kI64Base* / kI32Base* offsets; access helpers route through struct_array_access with the same base + index arithmetic the existing load_buf_* helpers use for the bytecode / args buffers	Layout drift would fail the static_asserts at compile time
Cross-thread aliasing of sizer scratch	The sizer is 1x1x1 (set_work_group_size({1,1,1})); the two scratch SSBOs are bound to that single dispatch and not visible to any other shader	No aliasing concern
Behaviour on non-Blackwell drivers	The pre-PR shader already worked there; the SSBO move is layout-equivalent (both paths read / write the same logical state in the same access order); get_function_array_type removes a decoration that was always invalid for Function-scope variables but tolerated by older validators	Unchanged
Encoder hard-error at the new cap	encode_adstack_size_expr_device_bytecode_for_spirv continues to compare against kAdStackSizerMaxNodesPerStack; the only effect of the bump is that more trees slip below the ceiling	Unchanged contract

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b809440755

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[duburcqa]

Fixed! It turns out it was an actual design flaw. I was aware of the limitation of Function Memory but since it worked on all the machines I tested including the AMDGPU box, my mac, and the Nvidia T4 runner, I assumed it was just fine with my 4096 threshold. It turns out it was not. Since lowering the threshold is not acceptable because it would be too limiting, I decided to move to actual storage buffer, which does not suffer such limitation.

[github-actions]

Coverage Report (714a63f35)

File	Coverage	Missing

Diff coverage: 0% · Overall: 73% · 0 lines, 0 missing

Full annotated report

[github-actions]

Coverage Report (4a633b42e)

File	Coverage	Missing

Diff coverage: 0% · Overall: 67% · 0 lines, 0 missing

Full annotated report

[duburcqa]

GS_ENABLE_NDARRAY=1 QD_OFFLINE_CACHE=0 pytest -v -ra --backend gpu --dev --forked ./tests

============ 649 passed, 3 skipped, 2 xfailed in 1115.15s (0:18:35) ============

[duburcqa]

Comparing Quadrants main vs this PR on Genesis main:

env	batch_size	backend	gjk_collision	constraint_solver	runtime_fps_main	runtime_fps_588	runtime_fps_delta_pct
anymal_random	30000	cuda	-	-	9236721	9288483	+0.56
anymal_uniform	30000	cuda	-	-	12290801	12277746	-0.11
anymal_uniform_kinematic	0	cpu	-	-	2054	2000	-2.63
anymal_uniform_kinematic	30000	cuda	-	-	10466141	10382797	-0.80
anymal_zero	0	cpu	-	-	7257	7186	-0.98
anymal_zero	30000	cuda	-	-	18701483	18857814	+0.84
box_pyramid_3	4096	cuda	-	-	969382	978981	+0.99
box_pyramid_4	4096	cuda	-	-	388357	391086	+0.70
box_pyramid_5	4096	cuda	-	-	141016	139273	-1.24
box_pyramid_6	4096	cuda	False	-	59335	58248	-1.83
box_pyramid_6	4096	cuda	True	-	60333	59916	-0.69
dex_hand	4096	cuda	-	-	17048	17364	+1.85
duck_in_box_easy	30000	cuda	False	-	26662843	26593384	-0.26
duck_in_box_easy	30000	cuda	True	-	9549541	9530599	-0.20
duck_in_box_hard	0	cpu	-	-	5185	5144	-0.79
duck_in_box_hard	30000	cuda	False	-	10239009	10162732	-0.74
duck_in_box_hard	30000	cuda	True	-	3405035	3368778	-1.06
franka	30000	cuda	-	-	22058180	21630563	-1.94
franka_accessors	0	cpu	-	-	1227	1203	-1.96
franka_accessors	30000	cuda	-	-	15590143	15764165	+1.12
franka_free	30000	cuda	-	-	31981995	32238005	+0.80
franka_random	0	cpu	-	-	6457	6298	-2.46
franka_random	30000	cuda	-	CG	16737391	16696014	-0.25
franka_random	30000	cuda	-	Newton	16343418	16450966	+0.66
franka_random	30000	cuda	False	-	16534089	16345136	-1.14
franka_random	30000	cuda	True	-	11429393	11410466	-0.17
g1_fall	4096	cuda	-	Newton	921076	922187	+0.12
go2	4096	cuda	False	CG	3678736	3673338	-0.15
go2	4096	cuda	False	Newton	4438961	4464312	+0.57
go2	4096	cuda	True	-	3273783	3243521	-0.92
shadow_hand_cubes	0	cpu	-	-	40	41	+2.50
shadow_hand_cubes_sparse	0	cpu	-	-	66	64	-3.03

speed_test_588.txt

[hughperkins]

checklist:

• no user facing doc changes required
• code changes seem mostly localized to adstack sizer specific files
• genesis unit tests passing
• genesis benchamrks neutral

=> ok to merge

[github-actions]

Coverage Report (a124d46be)

File	Coverage	Missing

Diff coverage: 0% · Overall: 73% · 0 lines, 0 missing

Full annotated report