[Perf] Cache adstack-sizer metadata per task across SPIR-V + LLVM-GPU; per-snode / DeviceAllocation invalidation (Part 2/2)

[duburcqa]

Cache adstack-sizer metadata per task across SPIR-V + LLVM-GPU; per-snode / DeviceAllocation invalidation

Part 2 of 2 in the adstack-perf series. Builds on the generic launcher infrastructure landed in Part 1 (duburcqa/perf_persistent_context_jit_pointer, #619). This PR adds the per-task metadata cache that turns publish_adstack_metadata's sizer-kernel + DtoH-stride round-trip into a host-only republish on cache hit. The launcher-side persistent scratch from Part 1 is what makes the republish DMA cheap enough to be worth caching.

TL;DR

Every reverse-mode launch with adstacks dispatches a sizer kernel per task to compute adstack_offsets / adstack_max_sizes / strides from the kernel's compile-time SerializedSizeExpr trees plus the current values of FieldLoad / ExternalTensorRead / ExternalTensorShape leaves. The sizer's output is a deterministic function of (per-task SerializedSizeExpr*, snode write generation per leaf snode, ndarray data generation per leaf ndarray). A long sequence of reverse-mode launches over the same kernel pays the sizer dispatch + 3x DtoH-stride round-trips per launch even though the inputs almost never change between launches.

This PR caches the sizer's output keyed by (AdStackSizingInfo *, dependency-generation snapshot). On cache hit the cached (offsets, max_sizes, stride_combined, stride_float, stride_int) is republished into the runtime struct via the same DMA path the host-eval branch already uses; the bytecode-encode + h2d + sizer-kernel + 3x DtoH stalls are skipped. Each launcher (CPU / CUDA / AMDGPU / SPIR-V) bumps snode_write_gen_ / ndarray_data_gen_ for everything its kernel can mutate before publish, so the snapshot catches any input change between launches and forces a fresh sizer dispatch.

// AdStackCache: lives on Program, owned via unique_ptr.
class AdStackCache {
 public:
  // Per-task entries on each backend's cache key.
  bool try_llvm_per_task_ad_stack_cache_hit(const void *task_key, LaunchContextBuilder *ctx,
                                            LlvmPerTaskAdStackCacheEntry &entry);
  void record_llvm_per_task_ad_stack_cache(const void *task_key, LlvmPerTaskAdStackCacheEntry entry);
  // Same shape for SPIR-V (kept separate so the bytecode + sizer-kernel paths can diverge).
  // ...
  // Generation bumps - called by each launcher on every kernel that may write to the source.
  void bump_snode_write_gen(int snode_id);
  void bump_ndarray_data_gen(const DeviceAllocation &devalloc);
};

The metadata cache machinery (struct definitions, methods, generation maps) was extracted from Program into a standalone AdStackCache class living in quadrants/program/adstack_size_expr_eval.{h,cpp}, with Program holding a std::unique_ptr<AdStackCache> and a single getter. Net effect: the surface area changed in Program is small (forward decl + getter + ctor init); all the cache-table ownership lives in adstack-specific files.

Why

Two costs the sizer dispatch was paying every launch:

1. Bytecode encode + h2d + sizer-kernel + 3x DtoH stride read-back. On AMDGPU / CUDA this is 4-5 driver hops per task per launch, plus the CPU stalls waiting on the DtoH. Genesis test_differentiable_rigid runs ~30+ adstack-bearing tasks per kernel; over a 100-substep horizon that's thousands of synchronous round-trips per training step.
2. CPU host-eval also pays a per-launch fixed cost - even though there's no kernel dispatch on CPU, the SizeExpr walk repeatedly reads the same FieldLoad / ExternalTensorRead values across launches whose inputs haven't changed.

The dependency snapshot (per-leaf snode write generation, per-leaf ndarray data generation) catches the actual mutation cases: a kernel that writes to field_a will bump snode_write_gen_[field_a.snode_id] before publish, so any cache entry depending on field_a is forced to re-evaluate. The cache amortises the dispatch + DMA cost to "once per generation bump" rather than "every launch".

Surface API

Zero user-visible change. Same qd.kernel.grad(...) invocation; same gradient values; same per-launch semantics.

Internal API surface (for the launchers):

// In each backend's KernelLauncher::launch_llvm_kernel (or gfx::Runtime::launch_kernel):
bump_writes_for_kernel_llvm(prog, &ctx, offloaded_tasks);   // CUDA / AMDGPU
bump_writes_for_kernel_llvm(prog, &ctx, snode_writes, arr_writes);  // CPU (precomputed at register time)
bump_writes_for_kernel_spirv(prog, &host_ctx, task_attribs, arr_access);  // gfx

Each helper iterates the kernel's compile-time write set (snode writes from codegen + ndarray writes from arg-access metadata) and bumps the matching generation counters on prog->adstack_cache(). The actual cache lookup happens inside publish_adstack_metadata (LLVM-GPU + CPU host-eval) and inside the gfx adstack-sizer-launch helper (SPIR-V).

Mechanism end-to-end

1. AdStackCache class extraction

quadrants/program/adstack_size_expr_eval.{h,cpp} - new class. Holds:

• SizeExpr host-eval cache: (SerializedSizeExpr *) -> {result, observed_reads}. On hit, replays observed FieldLoad / ExternalTensorRead / ExternalTensorShape reads against the live state and confirms a match before reusing the result.
• LlvmPerTaskAdStackCacheEntry: {offsets[], max_sizes[], stride_combined, stride_float, stride_int, snode_gen_snapshot, ndarray_gen_snapshot} keyed by AdStackSizingInfo *.
• SpirvBytecodeCacheEntry + PerTaskAdStackCacheEntry for the gfx backend (kept separate from the LLVM entry so bytecode encode and sizer-launch paths can diverge per backend).
• Generation maps: snode_write_gen_ (int snode_id -> uint64), ndarray_data_gen_ (DeviceAllocation -> uint64).

Program holds std::unique_ptr<AdStackCache> adstack_cache_; with adstack_cache() getter. The forward declaration + getter is the only Program surface change; cache method bodies and table ownership live in adstack-specific files.

2. bump_writes_for_kernel_llvm / bump_writes_for_kernel_spirv

quadrants/program/adstack_size_expr_eval.{h,cpp}. Called by each launcher just before kernel dispatch. Walks the kernel's compile-time write set:

• snode writes: from codegen, attached to each OffloadedTask (LLVM-GPU) or TaskAttributes (gfx).
• ndarray writes: from KernelContextAttributes::arr_access (gfx) or precomputed arg_writes_per_task (CPU launcher).

For each, calls prog->adstack_cache().bump_snode_write_gen(snode_id) / bump_ndarray_data_gen(devalloc). After this, any cache entry whose snapshot diverges from the new generation is invalidated on next lookup.

3. Cache lookup in publish_adstack_metadata

quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp. The GPU branch checks prog->adstack_cache().try_llvm_per_task_ad_stack_cache_hit(...) first; on hit, the cached payload republishes via publish_metadata_pinned_async (the same pinned-host async pipeline the host-eval branch uses) and skips the entire bytecode-encode + h2d + sizer-kernel + DtoH-stride pipeline. On miss, the original sizer dispatch runs and the resulting payload is recorded back into the cache via record_llvm_per_task_ad_stack_cache(...) along with the snapshot.

CPU host-eval has its own cache shape (no bytecode, just SizeExpr walk results); same pattern.

4. SPIR-V (gfx) adstack-sizer-launch helper

quadrants/runtime/gfx/adstack_sizer_launch.cpp - new file holding the SPIR-V sizer-launch path that previously lived inline in gfx/runtime.cpp. Lookup mirrors LLVM-GPU: bytecode cache hit short-circuits the encode + dispatch; per-task metadata cache hit short-circuits the entire sizer-kernel + DtoH pipeline.

5. Per-backend launcher integration

Backend	Where bump fires	Where cache fires
CPU	quadrants/runtime/cpu/kernel_launcher.cpp (host-eval branch of publish_adstack_metadata)	host-eval branch in llvm_adstack_lazy_claim.cpp
CUDA / AMDGPU	quadrants/runtime/{cuda,amdgpu}/kernel_launcher.cpp (just before launch_offloaded_tasks)	GPU branch in llvm_adstack_lazy_claim.cpp
SPIR-V (Metal / Vulkan)	quadrants/runtime/gfx/runtime.cpp (just before cmdlist record)	quadrants/runtime/gfx/adstack_sizer_launch.cpp

Per-backend coverage matrix

Backend	SizeExpr host-eval cache	Per-task metadata cache	Bytecode cache	Generation bump
CPU	✓	(host-eval only)	(n/a, no bytecode)	✓ (bump_writes_for_kernel_llvm)
CUDA	(used by host-eval if reached)	✓	✓ (LLVM-GPU bytecode)	✓
AMDGPU	(used by host-eval if reached)	✓	✓ (LLVM-GPU bytecode)	✓
Metal	(used by host-eval if reached)	✓ (gfx variant)	✓ (SPIR-V bytecode)	✓ (bump_writes_for_kernel_spirv)
Vulkan	(used by host-eval if reached)	✓ (gfx variant)	✓ (SPIR-V bytecode)	✓

Side-effect audit

Concern	Verdict
Snode mutation between launches forces fresh sizer dispatch	Each launcher bumps snode_write_gen_[snode_id] before launch for every snode the kernel writes; cache snapshot diverges, lookup misses.
Ndarray data mutation between launches forces fresh sizer dispatch	Each launcher bumps ndarray_data_gen_[devalloc] before launch; same mechanism.
SizeExprReadObservation replay catches host-eval staleness	On hit the cache replays the recorded FieldLoad / ExternalTensorRead reads and only short-circuits if every value matches; mismatch erases the entry and falls through to a full eval.
AdStackCache lifetime tied to Program	Owned by std::unique_ptr<AdStackCache> member; created in Program ctor, destroyed with the program. Forward decl in program.h keeps the include surface narrow.
Cross-backend cache-entry types do not collide	Distinct entry types per backend (LlvmPerTaskAdStackCacheEntry vs PerTaskAdStackCacheEntry vs SpirvBytecodeCacheEntry); Program does not see any of them.
Cache lookup happens after every snode/ndarray bump in the same launch	Bump calls fire before launch_offloaded_tasks; lookup fires inside publish_adstack_metadata (called from inside the offloaded-task loop). Order is bump-then-lookup so the snapshot read for hit-checking sees the post-bump generation.
Persistent device scratch from Part 1 still safe	Per-handle isolation (Part 1) means recursive snode_reader_* launches via host-eval FieldLoad cannot clobber the parent kernel's arg_buffer / RuntimeContext even when the cache miss path runs the sizer kernel that itself dereferences ctx->arg_buffer.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 42e8ae81fd

ℹ️ 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]

@claude review

[github-actions]

Coverage Report (17a5873bd)

File	Coverage	Missing

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

Full annotated report

[github-actions]

Coverage Report (02dcf40aa)

File	Coverage	Missing
🟢 tests/python/test_adstack.py	100%

Diff coverage: 100% · Overall: 74% · 66 lines, 0 missing

Full annotated report

[github-actions]

Coverage Report (230a98ee8)

File	Coverage	Missing
🟢 tests/python/test_adstack.py	100%

Diff coverage: 100% · Overall: 74% · 66 lines, 0 missing

Full annotated report

[github-actions]

PR change report (f6a8ddd8c)

Code lines (excluding blank lines, comment-only lines, and Python multi-line strings).

File	LoC	Added	Removed
quadrants/program/adstack_size_expr_eval.cpp	1168	+467	-49
quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp	882	+255	-190
quadrants/program/adstack_size_expr_eval.h	155	+131	-1
quadrants/runtime/llvm/llvm_runtime_executor.cpp	684	+86
quadrants/runtime/gfx/adstack_sizer_launch.cpp	328	+74
tests/python/test_adstack.py	2580	+73	-1
quadrants/codegen/llvm/codegen_llvm.cpp	2600	+69	-1
quadrants/program/program.cpp	465	+68	-5
quadrants/codegen/spirv/spirv_codegen.cpp	2479	+57	-22
tests/cpp/program/test_diagnose_adstack_overflow.cpp	48	+48
quadrants/runtime/cpu/kernel_launcher.cpp	138	+36	-24
quadrants/runtime/llvm/runtime_module/runtime.cpp	2010	+25	-7
quadrants/program/program.h	218	+22	-21
quadrants/runtime/gfx/runtime.cpp	907	+20	-6
quadrants/runtime/llvm/runtime_module/internal_functions.h	172	+13	-10
quadrants/runtime/llvm/llvm_runtime_executor.h	161	+8	-1
quadrants/runtime/program_impls/llvm/llvm_program.h	157	+5
quadrants/codegen/llvm/llvm_compiled_data.h	93	+4	-1
quadrants/program/snode_rw_accessors_bank.cpp	88	+4
quadrants/codegen/spirv/kernel_utils.h	211	+3	-1
quadrants/codegen/spirv/detail/spirv_codegen.h	208	+3
quadrants/program/ndarray.cpp	174	+3
quadrants/runtime/amdgpu/kernel_launcher.cpp	257	+3
quadrants/runtime/cpu/kernel_launcher.h	35	+3
quadrants/runtime/cuda/kernel_launcher.cpp	269	+3
quadrants/program/program_impl.h	118	+2

Total: 26 file(s) changed, +1485 -340 code lines.

Full per-function report

[github-actions]

Coverage Report (f6a8ddd8c)

File	Coverage	Missing
🟢 tests/python/test_adstack.py	100%

Diff coverage: 100% · Overall: 74% · 68 lines, 0 missing

Full annotated report

[github-actions]

PR change report (cf0c297ce)

Code lines (excluding blank lines, comment-only lines, and Python multi-line strings).

File	LoC	Added	Removed
quadrants/program/adstack_size_expr_eval.cpp	1168	+467	-49
quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp	885	+258	-190
quadrants/program/program.cpp	565	+168	-5
quadrants/program/adstack_size_expr_eval.h	155	+131	-1
tests/python/test_adstack.py	2603	+96	-1
quadrants/runtime/llvm/llvm_runtime_executor.cpp	684	+86
quadrants/runtime/gfx/adstack_sizer_launch.cpp	328	+74
quadrants/codegen/llvm/codegen_llvm.cpp	2600	+69	-1
quadrants/codegen/spirv/spirv_codegen.cpp	2479	+57	-22
tests/cpp/program/test_diagnose_adstack_overflow.cpp	48	+48
quadrants/runtime/cpu/kernel_launcher.cpp	138	+36	-24
quadrants/program/program.h	222	+26	-21
quadrants/runtime/llvm/runtime_module/runtime.cpp	2010	+25	-7
quadrants/runtime/gfx/runtime.cpp	902	+21	-12
quadrants/runtime/llvm/runtime_module/internal_functions.h	172	+13	-10
quadrants/runtime/llvm/llvm_runtime_executor.h	161	+8	-1
quadrants/runtime/program_impls/llvm/llvm_program.h	157	+5
quadrants/codegen/llvm/llvm_compiled_data.h	93	+4	-1
quadrants/codegen/spirv/kernel_utils.h	212	+4	-1
quadrants/program/ndarray.cpp	175	+4
quadrants/program/snode_rw_accessors_bank.cpp	88	+4
quadrants/codegen/spirv/detail/spirv_codegen.h	208	+3
quadrants/runtime/amdgpu/kernel_launcher.cpp	257	+3
quadrants/runtime/cpu/kernel_launcher.h	35	+3
quadrants/runtime/cuda/kernel_launcher.cpp	269	+3
quadrants/program/program_impl.h	118	+2

Total: 26 file(s) changed, +1618 -346 code lines.

Full per-function report

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 114f817b3e

ℹ️ 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".

[github-actions]

PR change report (114f817b3)

Code lines (excluding blank lines, comment-only lines, and Python multi-line strings).

File	LoC	Added	Removed
quadrants/program/adstack_size_expr_eval.cpp	1168	+467	-49
quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp	893	+265	-189
quadrants/program/program.cpp	559	+162	-5
quadrants/program/adstack_size_expr_eval.h	161	+137	-1
tests/python/test_adstack.py	2634	+127	-1
quadrants/runtime/gfx/adstack_sizer_launch.cpp	370	+116
quadrants/runtime/llvm/llvm_runtime_executor.cpp	684	+86
tests/cpp/program/test_diagnose_adstack_overflow.cpp	73	+73
quadrants/codegen/llvm/codegen_llvm.cpp	2602	+71	-1
quadrants/codegen/spirv/spirv_codegen.cpp	2492	+70	-22
quadrants/runtime/cpu/kernel_launcher.cpp	138	+36	-24
quadrants/runtime/gfx/runtime.cpp	915	+34	-12
quadrants/program/program.h	228	+31	-20
quadrants/runtime/llvm/runtime_module/runtime.cpp	2010	+25	-7
quadrants/runtime/llvm/runtime_module/internal_functions.h	172	+13	-10
quadrants/runtime/llvm/llvm_runtime_executor.h	161	+8	-1
quadrants/codegen/spirv/kernel_utils.h	213	+5	-1
quadrants/runtime/program_impls/llvm/llvm_program.h	157	+5
quadrants/codegen/llvm/llvm_compiled_data.h	93	+4	-1
quadrants/codegen/spirv/detail/spirv_codegen.h	209	+4
quadrants/program/ndarray.cpp	175	+4
quadrants/program/snode_rw_accessors_bank.cpp	88	+4
quadrants/codegen/spirv/kernel_utils.cpp	109	+3
quadrants/runtime/amdgpu/kernel_launcher.cpp	257	+3
quadrants/runtime/cpu/kernel_launcher.h	35	+3
quadrants/runtime/cuda/kernel_launcher.cpp	269	+3
quadrants/program/program_impl.h	118	+2
quadrants/runtime/gfx/runtime.h	151	+2

Total: 28 file(s) changed, +1763 -344 code lines.

Full per-function report

[github-actions]

Coverage Report (114f817b3)

File	Coverage	Missing
🟢 tests/python/test_adstack.py	89%	1243-1247,1265,1279,1301-1306

Diff coverage: 89% · Overall: 74% · 122 lines, 13 missing

Full annotated report

[github-actions]

PR change report (9320e5f12)

Code lines (excluding blank lines, comment-only lines, and Python multi-line strings).

File	LoC	Added	Removed
quadrants/program/adstack_size_expr_eval.cpp	1195	+494	-49
quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp	898	+269	-188
quadrants/program/program.cpp	559	+162	-5
tests/python/test_adstack.py	2651	+147	-4
quadrants/program/adstack_size_expr_eval.h	160	+136	-1
quadrants/runtime/gfx/adstack_sizer_launch.cpp	370	+116
quadrants/runtime/llvm/llvm_runtime_executor.cpp	684	+86
quadrants/codegen/llvm/codegen_llvm.cpp	2609	+78	-1
tests/cpp/program/test_diagnose_adstack_overflow.cpp	73	+73
quadrants/codegen/spirv/spirv_codegen.cpp	2492	+70	-22
quadrants/runtime/cpu/kernel_launcher.cpp	142	+40	-24
quadrants/runtime/gfx/runtime.cpp	917	+36	-12
quadrants/program/program.h	228	+31	-20
quadrants/runtime/llvm/runtime_module/runtime.cpp	2011	+26	-7
quadrants/runtime/llvm/runtime_module/internal_functions.h	172	+15	-12
python/quadrants/lang/impl.py	942	+8	-1
quadrants/runtime/llvm/llvm_runtime_executor.h	161	+8	-1
quadrants/codegen/llvm/llvm_compiled_data.h	94	+5	-1
quadrants/codegen/spirv/kernel_utils.h	213	+5	-1
quadrants/runtime/program_impls/llvm/llvm_program.h	157	+5
quadrants/codegen/spirv/detail/spirv_codegen.h	209	+4
quadrants/program/ndarray.cpp	175	+4
quadrants/program/snode_rw_accessors_bank.cpp	88	+4
quadrants/runtime/cpu/kernel_launcher.h	36	+4
quadrants/codegen/spirv/kernel_utils.cpp	109	+3
quadrants/runtime/amdgpu/kernel_launcher.cpp	257	+3
quadrants/runtime/cuda/kernel_launcher.cpp	269	+3
quadrants/program/program_impl.h	118	+2
quadrants/runtime/gfx/runtime.h	151	+2

Total: 29 file(s) changed, +1839 -349 code lines.

Full per-function report

[github-actions]

Coverage Report (9320e5f12)

File	Coverage	Missing
🔴 python/quadrants/lang/impl.py	29%	582-586
🟢 tests/python/test_adstack.py	95%	1244-1248,1291,1305

Diff coverage: 92% · Overall: 72% · 149 lines, 12 missing

Full annotated report

[github-actions]

PR change report (e7de926c7)

Code lines (excluding blank lines, comment-only lines, and Python multi-line strings).

File	LoC	Added	Removed
quadrants/program/adstack_size_expr_eval.cpp	1625	+924	-49
quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp	899	+270	-188
quadrants/program/adstack_size_expr_eval.h	200	+176	-1
tests/python/test_adstack.py	2651	+147	-4
quadrants/runtime/gfx/adstack_sizer_launch.cpp	370	+116
quadrants/runtime/llvm/llvm_runtime_executor.cpp	684	+86
quadrants/codegen/llvm/codegen_llvm.cpp	2610	+79	-1
tests/cpp/program/test_diagnose_adstack_overflow.cpp	79	+79
quadrants/codegen/spirv/spirv_codegen.cpp	2492	+70	-22
quadrants/runtime/cpu/kernel_launcher.cpp	142	+40	-24
quadrants/runtime/gfx/runtime.cpp	917	+36	-12
quadrants/runtime/llvm/runtime_module/runtime.cpp	2011	+26	-7
quadrants/runtime/llvm/runtime_module/internal_functions.h	172	+15	-12
quadrants/program/program.cpp	409	+12	-5
quadrants/program/program.h	205	+9	-21
python/quadrants/lang/impl.py	942	+8	-1
quadrants/runtime/llvm/llvm_runtime_executor.h	161	+8	-1
quadrants/codegen/llvm/llvm_compiled_data.h	94	+5	-1
quadrants/codegen/spirv/kernel_utils.h	213	+5	-1
quadrants/program/launch_context_builder.cpp	389	+5
quadrants/runtime/program_impls/llvm/llvm_program.h	157	+5
quadrants/codegen/spirv/detail/spirv_codegen.h	209	+4
quadrants/program/ndarray.cpp	175	+4
quadrants/program/snode_rw_accessors_bank.cpp	88	+4
quadrants/runtime/cpu/kernel_launcher.h	36	+4
quadrants/codegen/spirv/kernel_utils.cpp	109	+3
quadrants/runtime/amdgpu/kernel_launcher.cpp	257	+3
quadrants/runtime/cuda/kernel_launcher.cpp	269	+3
quadrants/program/program_impl.h	118	+2
quadrants/runtime/gfx/runtime.h	151	+2
quadrants/program/launch_context_builder.h	98	+1

Total: 31 file(s) changed, +2151 -350 code lines.

Full per-function report

[hughperkins]

Is it guaranteed to be the next? I feel it is more likely to be the next launch or so on that happens after the error has occurred, (and not the next launch after the launch that led to a kernel erroring)?

[hughperkins]

eg, we launhc:

• A
• B
• C
• at this point, A finishes running, and it had an error
• D
  • receives the exception/error

?

[hughperkins]

(I suppose it depends on our definition of 'next')

[duburcqa]

Only on CPU. Updated to be more accurate.

[hughperkins]

nice explanation I feel

[hughperkins]

"File a bug with QD_DUMP_IR=1 set". Whilst I can undersatnd what you mean, this seems to be a confusing way to state this? Maybe omsething like:

"File a bug report as follows, as an issue against the Quadrants repo:

• rm /tmp/ir/*
• run your job with the env var QD_DUMP_IR=1 defined
• zip the contents of /tmp/ir
• create an issue against the Quadrants repo
• attach the zip of /tmp/ir to the issue"

?

[duburcqa]

updated

[hughperkins]

👀

[duburcqa]

updated

[hughperkins]

• left some comments on the doc
• waiting also for the agent-based checks to run
• this touches a ton of different files, so lets have a full Genesis unit test run, and benchmarks, pelase

[hughperkins]

Description of how to file the issue seems clear, complete and concise. Great! 🙌

[hughperkins]

doc looks good to me 🙌

I'm waiting now on:

• agent-based checks to run
• results of full Genesis unit test run (at least on field)
• results of Genesis benchmarks run (at least on field)

[github-actions]

PR change report (11577f99e)

Code lines (excluding blank lines, comment-only lines, and Python multi-line strings).

File	LoC	Added	Removed
quadrants/program/adstack_size_expr_eval.cpp	1637	+936	-49
quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp	895	+266	-188
quadrants/program/adstack_size_expr_eval.h	205	+181	-1
tests/python/test_adstack.py	2680	+147	-4
quadrants/runtime/gfx/adstack_sizer_launch.cpp	373	+119
quadrants/runtime/llvm/llvm_runtime_executor.cpp	684	+86
quadrants/codegen/llvm/codegen_llvm.cpp	2610	+79	-1
tests/cpp/program/test_diagnose_adstack_overflow.cpp	79	+79
quadrants/codegen/spirv/spirv_codegen.cpp	2492	+70	-22
quadrants/runtime/cpu/kernel_launcher.cpp	142	+40	-24
quadrants/runtime/gfx/runtime.cpp	917	+36	-12
quadrants/program/program.cpp	424	+27	-5
quadrants/runtime/llvm/runtime_module/runtime.cpp	2011	+26	-7
quadrants/runtime/llvm/runtime_module/internal_functions.h	172	+15	-12
quadrants/program/program.h	205	+9	-21
python/quadrants/lang/impl.py	942	+8	-1
quadrants/runtime/llvm/llvm_runtime_executor.h	161	+8	-1
quadrants/codegen/llvm/llvm_compiled_data.h	94	+5	-1
quadrants/codegen/spirv/kernel_utils.h	213	+5	-1
quadrants/program/launch_context_builder.cpp	389	+5
quadrants/runtime/program_impls/llvm/llvm_program.h	157	+5
quadrants/codegen/spirv/detail/spirv_codegen.h	209	+4
quadrants/program/ndarray.cpp	175	+4
quadrants/program/snode_rw_accessors_bank.cpp	88	+4
quadrants/runtime/cpu/kernel_launcher.h	36	+4
quadrants/codegen/spirv/kernel_utils.cpp	109	+3
quadrants/runtime/amdgpu/kernel_launcher.cpp	257	+3
quadrants/runtime/cuda/kernel_launcher.cpp	269	+3
quadrants/program/program_impl.h	118	+2
quadrants/runtime/gfx/runtime.h	151	+2
quadrants/program/launch_context_builder.h	98	+1

Total: 31 file(s) changed, +2182 -350 code lines.

Full per-function report

[duburcqa]

env	batch_size	backend	gjk_collision	constraint_solver	runtime_fps_main	runtime_fps_620	runtime_fps_delta_pct
anymal_random	30000	cuda	-	-	9443191	9390292	-0.56
anymal_uniform	30000	cuda	-	-	12346008	12387632	+0.34
anymal_uniform_kinematic	0	cpu	-	-	1982	1948	-1.72
anymal_uniform_kinematic	30000	cuda	-	-	10631243	10352282	-2.62
anymal_zero	0	cpu	-	-	7267	7139	-1.76
anymal_zero	30000	cuda	-	-	19146598	19301149	+0.81
box_pyramid_3	4096	cuda	-	-	979160	988222	+0.93
box_pyramid_4	4096	cuda	-	-	384826	387764	+0.76
box_pyramid_5	4096	cuda	-	-	141026	137764	-2.31
box_pyramid_6	4096	cuda	False	-	58596	58888	+0.50
box_pyramid_6	4096	cuda	True	-	60147	59231	-1.52
dex_hand	4096	cuda	-	-	17240	17173	-0.39
duck_in_box_easy	30000	cuda	False	-	27101190	26767881	-1.23
duck_in_box_easy	30000	cuda	True	-	9526647	9501811	-0.26
duck_in_box_hard	0	cpu	-	-	5079	5121	+0.83
duck_in_box_hard	30000	cuda	False	-	10245189	10305593	+0.59
duck_in_box_hard	30000	cuda	True	-	3386673	3416781	+0.89
franka	30000	cuda	-	-	22146368	21472201	-3.04
franka_accessors	0	cpu	-	-	1136	1193	+5.02
franka_accessors	30000	cuda	-	-	15354421	15706391	+2.29
franka_free	30000	cuda	-	-	32056108	32810552	+2.35
franka_random	0	cpu	-	-	6249	6271	+0.35
franka_random	30000	cuda	-	CG	16548461	16839003	+1.76
franka_random	30000	cuda	-	Newton	16720811	16688250	-0.19
franka_random	30000	cuda	False	-	16509323	16401148	-0.66
franka_random	30000	cuda	True	-	11570179	11542843	-0.24
g1_fall	4096	cuda	-	Newton	907003	922819	+1.74
go2	4096	cuda	False	CG	3646422	3725625	+2.17
go2	4096	cuda	False	Newton	4624845	4583004	-0.90
go2	4096	cuda	True	-	3330016	3318628	-0.34
shadow_hand_cubes	0	cpu	-	-	41	41	+0.00
shadow_hand_cubes_sparse	0	cpu	-	-	65	66	+1.54

speed_test_620.txt
speed_test_main.txt

[duburcqa]

All the unit tests are passing against Genesis-Embodied-AI/genesis-world#2742.

However, there is a bug in load_store elimination for static for-loops in eliminate_recomputable_ad_stack_pushes which breaks the unit test test_differentiable_push on genesis/main. It is unrelated to this PR. It was introduced by #621. I was not able to spot it because I ran the unit test against 2742 instead of main at that time... I suggest to fix this bug in a follow-up since it is unrelated?

[github-actions]

Diff coverage: 92% · 149 lines, 12 missing

[duburcqa]

After applying #633, all the unit tests are passing on Genesis main
======= 642 passed, 3 skipped, 2 xfailed in 1086.07s (0:18:06) =======

[hughperkins]

checklist:

• docs look good to me
• agent based checks passed
• genesis unit tests passing
• genesis benchmarks ~neutral

=> ok to merge