[Caching] Per-offload cache - #852
Conversation
The slice-1b cubin cache keyed on a SHA of the sub-module's LLVM-IR *text*, which meant a warm run had to run the whole frontend and per-task CHI->LLVM codegen just to produce the thing to hash -- so it could only ever skip PTX+ptxas. Thread the per-task IR key (get_hashed_per_task_cache_key + "#index", computed from the task IR *before* codegen) from the codegen driver down through LLVMCompiledKernel::per_construct_keys -> kernel_launcher -> create_jit_module_culink -> add_module_culink -> get_or_build_construct_cubin. That key is edit-stable and is the prerequisite for skipping codegen entirely on a warm one-task edit (next increment). Also namespace the cubin dir by SM (culink_cubins_<mcpu>): cubins are built with `ptxas -arch=`, so a cache populated on one GPU must not be loaded on another. per_construct_keys is copied in LLVMCompiledKernel::clone() -- the launcher consumes a clone, and dropping it would silently demote to the text-hash path. Validated on the real qipc SimEngine._step_kernel (freefall, CUDA): run1 writes 164 cubins, run2 reuses all of them, both byte-identical (q_sum 1.7492255125). Default (non-culink) path unchanged; test_per_offload_cache 2 passed/2 skipped.
Part B.2/B.3. A cubin alone cannot be launched -- the launcher and CUDA graph
builder are driven by OffloadedTask metadata (entry name, dims, gdw level,
stream group, checkpoint id, adstack sizing, snode/arg RW sets) and the runtime
needs used_tree_ids. So persist a complete PerTaskArtifact record {tasks,
used_tree_ids, struct_for_tls_sizes, cubin} keyed by the per-task IR key, via
the existing QD_IO_DEF binary serializer (new header-only
codegen/llvm/per_task_artifact_cache.h; records are .qdb, as the serializer
rejects other extensions).
The codegen driver now probes that cache BEFORE compile_task: on a hit the task
skips CHI->LLVM, link, optimize, PTX and ptxas entirely, and carries the cached
cubin to the cuLink assembly. Replaced the parallel module/key vectors with
PerConstructArtifact, which is either a module-to-compile or a prebuilt cubin
plus its metadata, so the JIT (the only holder of the cubin) can write the
complete record.
Consequently a cached task has no LLVM module, so the whole-module link is
skipped when any task is code-only -- which also removes the prototype's
double-build (link+optimize once per task AND once whole-kernel). Guarded the
three whole-module derefs that this exposes: check() (was verifyModule on a null
module -- the segfault found in testing), dump_impl (reports a dump failure so
the manager skips the .qdc; the artifact cache is that kernel's persistence
instead) and debug_dump_to_string.
Gated on QD_PERTASK_ARTIFACT_CACHE=1 + QD_CULINK_PERTASK=1 for A/B.
Measured on the real qipc SimEngine._step_kernel (freefall, CUDA), fresh process
run 2 vs run 1: all 148 tasks served from disk, pertask_compile 1292.8 -> 49.5ms,
per_construct_selfcontained 2622.6 -> 0.1ms, whole_module_link skipped, step0
37.7 -> 24.0s, q_sum byte-identical (1.7492255125). The residual 24s is the
whole-kernel frontend (merge_global_ptrs et al), which is Part A.
Default path unchanged (28.9s, byte-identical); test_per_offload_cache +
test_fields_builder: 25 passed / 2 skipped.
A construct's frontend output cannot be persisted as IR (no round-trippable CHI codec), so persist what it *contributed* instead: the ordered list of per-task artifact keys. On a warm compile a manifest hit lets the split emit one placeholder task per key and skip that construct's entire frontend (merge_global_ptrs / full_simplify / offload); codegen then loads the named artifacts (Part B) and never lowers anything. Mechanics: the split and codegen communicate through a program-scoped ConstructDiskPlan side table (task index -> construct key / artifact key) rather than new OffloadedStmt fields, which would perturb the printed-IR-derived per-task key. Manifests are keyed by a new cross-process construct key (get_hashed_per_construct_disk_key) that restores the full SNode-layout hash and device caps the in-memory key deliberately drops; the layout hash is memoized per snode-tree, since computing it per construct was a >20x blowup before. Because the task index is baked into entry-fn / shared_array / adstack-counter names, a manifest is only reused if the construct lands at the same task offset; the loader verifies this and treats a shift as a miss. Measured on the real qipc SimEngine._step_kernel (freefall, CUDA, fresh process): constructs_manifest_hit 74/130, recompiled 81 -> 7, frontend_ms 3025 -> 21ms; 242/242 tasks served from the artifact cache; compile_kernel 16.89 -> 7.89s; step0 19.3 -> 16.2s. q_sum byte-identical (1.7492255125) throughout.
…ache A kernel whose device code is per-task cubins has no whole-kernel LLVM module, and dump_impl was reporting a dump failure for it -- so it was never written to .qdc, the whole-kernel cache stayed permanently empty, and every run re-paid the per-construct path (measured as a ~50x regression on warm-no-change, 0.17s -> 8.6s). Give such kernels a real .qdc entry instead: serialize per_task_artifact_keys alongside tasks (both now in LLVMCompiledKernel's QD_IO_DEF) and write an empty src_code; on load, rebuild per_construct_artifacts from the PerTaskArtifactCache and leave module null, failing the load if an artifact is missing so the manager recompiles rather than emitting a kernel with a missing task. load_impl sits far from any CompileConfig, so the artifact directory is resolved once in the LlvmProgramImpl constructor rather than threaded through the load path.
The per-construct split cloned the entire top-level block for each construct and then deleted everything outside that construct's backward slice -- O(constructs x block size). On the genesis _step_kernel that is 130 deep clones of a 2685-stmt block, ~5.9 s, and it was paid even when every construct was a cache hit, because the isolated construct is what the cache key is computed from. Compute the slice on the original block instead and clone only the surviving top-level statements, via a new irpass::analysis::clone_block_subset(). It drives the existing IRCloner per statement pair (the source and target blocks no longer line up index-for-index) so operand remapping is unchanged: references between cloned statements are redirected, references outside the subset keep pointing at the original -- same semantics as a whole-block clone. compile_kernel 7.73 s -> 2.13 s on the real _step_kernel, with the split behaviour identical (130 constructs, 49 in-memory / 74 manifest / 7 recompiled) and key_ms 416 -> 176. q_sum byte-identical (1.7492255125); quadrants per_offload/fields_builder/offload tests 31 passed, qipc freefall 4 passed.
get_pos_info builds the source-position banner used in error messages and DebugInfo, and ASTTransformerBase.__call__ calls it for every stmt/expr node visited. It was the single largest cost on the Python side of a recompile: on the genesis _step_kernel, ~9.6s of a 17.5s AST transform, mostly inside textwrap (614k wrap calls, a fresh TextWrapper per node). The same ~63 functions are transformed ~44 times each (once per inlined call site), so the identical source position is formatted over and over. Memoize on what the output depends on -- file, function, indent, line offset and the node's span. The cache is process-global on purpose: the win is reuse ACROSS transforms of a function, which a per-context cache would miss. src is not in the key because for a given file+function+indent+offset it is that function's source, which cannot change under a running interpreter. Measured on the real _step_kernel (step0, time-to-first-step): warm + one offload edit: 10.27s -> 5.80s (machinery on) warm + one offload edit: 27.16s -> 22.06s (baseline, no machinery) cold: 34.28s -> 28.74s (machinery on) It is a pure Python-side fix, so it helps the baseline equally. Warm-unchanged unaffected (0.187s). q_sum byte-identical; 109 tests pass.
The cold regression from the per-task cuLink path was 242 serial 'ptxas -c' shell-outs. Splitting the per-task cubin build showed 9.8s of the 13.7s was ptxas and only 1.6s was LLVM->PTX, with cuLink itself at 6ms and artifact writes at 0ms -- so the whole regression was subprocess launches. Run the per-task builds across hardware_concurrency threads. ptxas is a subprocess on private temp files and holds no shared state, so it parallelises safely; LLVM->PTX stays serialised under a mutex because the per-task modules come off the codegen worker threads and may share an LLVMContext. Each entry writes only its own slot. Also split assemble_and_store_cubin out so the parallelisable half is explicit, and made the diagnostic timers atomic. Cold step0 on the real _step_kernel: 28.74s -> 20.47s (machinery on), i.e. the regression against baseline is gone. Edit case unchanged at 5.66s, warm unchanged at 0.190s, q_sum byte-identical.
hughperkins
changed the title
hughperkins
changed the title
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
https://github.com/Genesis-Embodied-AI/quadrants/blob/4860dbd0936e1329f75d767013b86b4e8da45cfd/codegen/codegen.cpp#L195-L198
For reverse-mode kernels whose offload contains AdStackAllocaStmt, this disk-artifact probe returns before the existing has_adstack guard runs. The in-memory cache below deliberately excludes those tasks because lowering them registers per-task adstack sizing in the Program; a fresh-process artifact hit skips that side effect while still launching the cached cubin. Either run the adstack check before probing/storing artifacts and treat these tasks as misses, or explicitly re-seed the registry from the serialized task metadata.
https://github.com/Genesis-Embodied-AI/quadrants/blob/4860dbd0936e1329f75d767013b86b4e8da45cfd/codegen/codegen.cpp#L127-L128
When QD_PERTASK_ARTIFACT_CACHE=1 and QD_CULINK_PERTASK=1 are set but QD_CULINK_CUBIN is not, this enables loading cached cubins even though add_module_culink() will take its PTX path and call compile_module_to_ptx(art.module) on artifacts whose module is null. The same combination also writes construct manifests while the JIT never stores PerTaskArtifact records. Gate artifact-cache use on cubin mode too, or make the JIT consume prebuilt cubins regardless of QD_CULINK_CUBIN.
https://github.com/Genesis-Embodied-AI/quadrants/blob/4860dbd0936e1329f75d767013b86b4e8da45cfd/transforms/compile_to_offloads.cpp#L767
A manifest hit is accepted solely because the manifest file exists and has task keys, before verifying that the referenced per-task artifacts are still present/readable. If those artifacts were evicted, corrupted, or the manifest was written before the JIT stored them, the split replaces the real construct with placeholders and codegen later hard-errors because it has no IR to fall back to. While cb is still available here, validate all referenced artifacts and treat any missing artifact as a manifest miss.
https://github.com/Genesis-Embodied-AI/quadrants/blob/4860dbd0936e1329f75d767013b86b4e8da45cfd/transforms/compile_to_offloads.cpp#L858-L862
With QD_CONSTRUCT_MANIFEST=1 and QD_CONSTRUCT_CACHE=0, manifest hits above still push placeholder offloads, but this new disk plan is only stored inside the cc != nullptr block. Codegen then sees no artifact_key_by_task, compiles the empty placeholders as ordinary serial tasks, and silently skips the cached construct's real work. Store the ConstructDiskPlan whenever manifests are active, using the program's per_construct_cache() as the side channel even when the in-memory construct cache is disabled.
ℹ️ 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".
|
Superseded. Re-splitting this work by what is actually being changed rather than by the branches that happened to exist, and dropping the env gates and diagnostic scaffolding entirely (everything defaults on). Plan: |
1 participant
Stacked on #851 (which is stacked on #837). This is the top of the stack and carries the headline result.
Result
Editing one offloaded task of the giant qipc nested-graph
_step_kerneland re-running in a fresh process(freefall NB=1, CUDA, caches on, same harness for every cell):
compile_kernelfor the edit case: 18.92 s -> 1.87 s. No regression on cold or warm-unchanged.q_sum = 1.7492255125byte-identical throughout.What this branch adds
OffloadedTaskmetadata — so the record is{cubin, tasks, used_tree_ids, struct_for_tls_sizes}. Probed beforecompile_task, so a hit skips CHI->LLVM, link, optimize, PTX and ptxas entirely.persist what it produced: the ordered per-task artifact keys. A hit skips that construct's whole frontend.
computed after codegen — so a hit could only ever skip ptxas). Directory namespaced by SM.
clone_block_subset— the split cloned the whole 2685-stmt block per construct, O(constructs x block).compile_kernel7.73 s -> 2.13 s.ptxas -cshell-outs were 9.8 s of a 13.7 s cubin build (LLVM->PTXwas only 1.6 s). ptxas is a subprocess on private temp files so it parallelises; LLVM->PTX stays serialised because
those modules may share an
LLVMContext. This removed a cold regression of ~7 s.Reviewer notes
left as follow-up so this can land without changing default behaviour. The numbers above are with the gates on.
(
get_pos_infoformatting a caret hint per AST node). [Lang] Cut per-compilation Python overhead (sys.modules scan + per-node source formatting) #804's approach is better — it avoids the work rather thancaching it — and its blocking review comments are all on its other change (
_inside_class/linecache). Happy todrop the memo here in favour of [Lang] Cut per-compilation Python overhead (sys.modules scan + per-node source formatting) #804's approach, or to have [Lang] Cut per-compilation Python overhead (sys.modules scan + per-node source formatting) #804 split so its second half can land.
perso_hugh/doc/per_offload_cache/: eviction limits for the new cache tiers, task-indexshift handling,
ptxas->libnvptxcompiler_static, and de-gating.Made with Cursor