[Cache] Assemble CUDA kernels from per-task modules via a composite JITModule - #875
Conversation
The whole-module path recompiles every task's PTX whenever any part of the kernel changes, because the LLVM text it hashes moves even for tasks that did not. Build one self-contained module per offloaded task instead, assemble each to a relocatable cubin with `ptxas -c`, and device-link them into a single CUmodule with cuLink. The cubins are cached on disk keyed by the hash of each task module's LLVM-IR text, so an unchanged task skips PTX generation and ptxas entirely on a warm edit. Self-containment requires `bls_buffer` and static shared scratch to be internal definitions rather than external declarations: with one cubin per task there is no whole-module definition left to resolve against, and cuLinkComplete fails on the undefined reference. The cubin directory is namespaced by SM version, since `ptxas -arch=` output is not portable across GPUs. The f16 tolerance in test_atomic_float_ops is widened because cuLink changes code layout and so the interleaving of 128 concurrent atomics; f16 cannot represent their exact sum at the old 1e-3 either way. The cubin cache honors offline_cache=False: like PtxCache and the kernel-data cache it then falls back to a mem-only build that touches no disk, so a fresh session with offline_cache=False cannot be served a cubin left by a previous run. init_options.md documents this new disk tier alongside the existing ones. compile_module_to_ptx NUL-terminates the PTX buffer for the driver's cuModuleLoadDataEx, but ptxas (CUDA 13+) treats the embedded NUL as a premature EOF and aborts; write only up to the terminator when handing PTX to ptxas.
1b's assemble path consumes only `module` (the relocatable-cubin cache keys on
the module's LLVM text); the record's other fields -- key, cubin, tasks,
used_tree_ids, struct_for_tls_sizes -- are populated in codegen.cpp but read
nowhere in 1b (the launch path uses the kernel-level LLVMCompiledKernel::tasks).
Per the no-scaffolding rule they are 6's, so keep PerConstructArtifact at just
{ module } and let 6 reintroduce the full record with the artifact cache that is
their sole producer and consumer. The per_construct_artifacts vector stays -- it
carries the modules to the JIT.
Condense the explanatory comments added by the relocatable-cubin path to the essential intent; no code changes.
Drop the per-layer cache breakdown from the user-facing offline_cache option; keep just the cache folder and the cold-start behavior.
![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
Here are some automated review suggestions for this pull request.
Reviewed commit: 4772e5ef6f
ℹ️ 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".
| auto cmd = fmt::format("ptxas -c -arch={} {} -o {} 2>/dev/null", arch, ptx_path.string(), cubin_path.string()); | ||
| int rc = std::system(cmd.c_str()); | ||
| QD_ERROR_IF(rc != 0, "ptxas -c failed (rc={}) arch={}", rc, arch); |
There was a problem hiding this comment.
On a fresh cache, every newly compiled CUDA kernel now reaches this command because the launcher selects the nonempty per-task artifacts, but supported driver/runtime-only installations do not necessarily provide the toolkit's ptxas executable (QD_WITH_CUDA_TOOLKIT is off by default). Such installations now fail with ptxas -c failed instead of using the driver's existing PTX JIT; moreover, the /dev/null redirection and unquoted temporary paths make this command unsuitable for the supported Windows CUDA configuration even when ptxas is installed. Please use the driver APIs or retain the whole-module fallback when this path is unavailable, and keep the installation requirements synchronized if a toolkit is intentionally required.
AGENTS.md reference: AGENTS.md:L15-L22
Useful? React with 👍 / 👎.
| if (!data.per_construct_artifacts.empty()) { | ||
| // Per-task cubin path: cuLink the per-task modules instead of loading the whole-module PTX. | ||
| jit_module = executor->create_jit_module_culink(std::move(data.per_construct_artifacts)); |
There was a problem hiding this comment.
For every freshly compiled CUDA kernel, this branch bypasses JITSessionCUDA::add_module, which is the only path that implements print_kernel_asm and QUADRANTS_LOAD_PTX. Consequently, print_kernel_asm=True no longer writes the documented PTX file, and QUADRANTS_LOAD_PTX=1 silently executes newly generated PTX rather than the user's edited dump, despite the documented workflow requiring both caches to be disabled. Please handle these controls on the cuLink path or select the legacy path when either is active.
AGENTS.md reference: AGENTS.md:L15-L22
Useful? React with 👍 / 👎.
| std::vector<PerConstructArtifact> per_construct_artifacts; | ||
| for (int i = 0; i < (int)data.size(); i++) { | ||
| if (!data[i] || !data[i]->module) | ||
| continue; | ||
| PerConstructArtifact art; | ||
| std::vector<std::unique_ptr<LLVMCompiledTask>> one; | ||
| one.push_back(std::make_unique<LLVMCompiledTask>(data[i]->clone())); | ||
| auto linked_one = tlctx_.link_compiled_tasks(std::move(one)); |
There was a problem hiding this comment.
This common LLVM compilation path now clones, links, and optimizes a self-contained module for every task on CPU and AMDGPU as well as CUDA. The CPU and AMDGPU launchers never consume per_construct_artifacts, so kernels on those backends pay up to one extra runtime-module link and optimization pipeline per task before the original whole-module link still runs. Restrict this CUDA-only feature to the CUDA codegen path to avoid a potentially large compile-time and memory regression on unrelated backends.
AGENTS.md reference: AGENTS.md:L5-L13
Useful? React with 👍 / 👎.
Apply clang-format reflows and reword one comment so it fills to 120 cols rather than breaking early before a long path token.
![graphite-app[bot]](https://avatars.githubusercontent.com/in/158384?s=60&v=4){kind=small}
| if (!dump_ir && fs::exists(cubin_path)) { | ||
| std::ifstream in(cubin_path, std::ios::binary); | ||
| return std::vector<char>((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>()); |
There was a problem hiding this comment.
File handle error not checked: if the file exists but cannot be opened (permissions, race condition, etc.), an empty vector will be returned, causing cuLink to fail later with cryptic errors.
if (!dump_ir && fs::exists(cubin_path)) {
std::ifstream in(cubin_path, std::ios::binary);
if (!in.is_open()) {
QD_ERROR("Failed to open cached cubin: {}", cubin_path);
}
return std::vector<char>((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
}| if (!dump_ir && fs::exists(cubin_path)) { | |
| std::ifstream in(cubin_path, std::ios::binary); | |
| return std::vector<char>((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>()); | |
| if (!dump_ir && fs::exists(cubin_path)) { | |
| std::ifstream in(cubin_path, std::ios::binary); | |
| if (!in.is_open()) { | |
| QD_ERROR("Failed to open cached cubin: {}", cubin_path); | |
| return std::vector<char>(); | |
| } | |
| return std::vector<char>((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>()); | |
Spotted by Graphite
Is this helpful? React 👍 or 👎 to let us know.
Option B: compile each self-contained task module to PTX and load it as its own CUmodule via cuModuleLoadDataEx; JITModuleCUDA is now a composite that resolves task entry points by name across the N modules (memoised). This removes the ptxas subprocess, the relocatable-cubin disk cache, and the cuLink device link, so the per-task path runs on a driver-only machine with no CUDA toolkit and no wheel-size growth. Per-task reuse rides on the existing PtxCache (PTX) plus the NVIDIA driver compute cache (SASS). Rename add_module_culink -> add_module_per_task and create_jit_module_culink -> create_jit_module_per_task, and refresh the now-inaccurate cuLink/relocatable comments on the shared-array linkage fix and the test_atomic tolerance.
hughperkins
changed the title
The composite per-task path dropped the QD_DUMP_IR / QD_LOAD_IR side effects that the old cuLink path emitted, so test_loop_config_name_cuda_ptx_dump found no .ptx under debug_dump_path. Factor add_module's dump/load block into a shared compile_module_to_ptx_with_dump() and call it from add_module_per_task too, so both paths dump <name>_before_ptx.ll / <name>.ptx (and honour QD_LOAD_IR) identically.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 70e9c9a64d
ℹ️ 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".
compile_kernel_to_module built a self-contained per-task module for every task on all LLVM backends, but only the CUDA launcher consumes per_construct_artifacts. That made CPU / AMDGPU pay an extra per-task link + optimize for artifacts they never use. Gate the loop on Arch::cuda (refs 7/8 extend consumption to CPU/AMD).
offline_cache=False disables Quadrants' own on-disk caches, but the NVIDIA driver keeps a separate SASS cache at ~/.nv/ComputeCache that the flag does not disable (the per-process nonce only stops cross-run reuse). Correct the "nothing is read from or written to disk" claim, which misled users profiling disk activity or on read-only filesystems.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 278639b8d2
ℹ️ 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".
| #include "quadrants/runtime/llvm/llvm_fwd.h" | ||
| #include "quadrants/util/lang_util.h" | ||
| #include "quadrants/jit/jit_module.h" | ||
| #include "quadrants/codegen/llvm/llvm_compiled_data.h" // PerConstructArtifact (per-task module path) |
There was a problem hiding this comment.
When configuring a supported QD_WITH_LLVM=OFF build, QuadrantsCore.cmake still compiles quadrants/jit/jit_session.cpp, so this unconditional include now reaches llvm/IR/Module.h even though LLVM discovery and include directories are only configured when LLVM is enabled. Vulkan/Metal-only builds on systems without LLVM headers therefore fail before the existing #ifdef QD_WITH_LLVM in jit_session.cpp; gate or forward-declare the CUDA-only artifact API rather than importing llvm_compiled_data.h into the generic JIT header.
AGENTS.md reference: AGENTS.md:L9-L11
Useful? React with 👍 / 👎.
|
I think I'm going to ignore hte init_opt.md faeilures on this PR. They're being handled seprately, in #873 |
jit_session.cpp is globbed into the core lib unconditionally, so the generic JIT header must not hard-depend on LLVM headers. The per-task include of llvm_compiled_data.h pulled in llvm/IR/Module.h, breaking Vulkan/Metal-only (QD_WITH_LLVM=OFF) builds where LLVM include dirs are not configured. Gate the include and the add_module_per_task declaration behind QD_WITH_LLVM, matching how jit_session.cpp already guards its LLVM guts; the rest of the header keeps using forward-declared llvm types via llvm_fwd.h.
|
@codex review |
|
Codex Review: Didn't find any major issues. What shall we delve into next? Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
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". |
|
running Genesis benchmarks |
|
running Genesis unit tests |
|
benchmarks at this point: 
I think I need to run the before/after script, on a single node. |
|
(running before/after script on single node for convexify and table bussing) |
|
results of before/after, same node, 5 runs of each: |
|
Genesis unit tests ok: 
|
|
I think I will wait for CI to pass-ish, then merge |
1 participant
Summary
Split CUDA kernel compilation per offloaded task, but keep it toolkit-free. Each offloaded task becomes its own self-contained module (
PerConstructArtifact{ module }); at JIT time each is compiled to PTX and loaded as its ownCUmoduleviacuModuleLoadDataEx.JITModuleCUDAbecomes a composite that holds the N modules and resolves task entry points by name across them (memoised). The launcher and CUDA-graph builder are unchanged -- they only ever look tasks up by name.This is "Option B" of the per-task compile design: it deliberately avoids
ptxas/cuLink/a relocatable-cubin cache, so the path needs no CUDA toolkit at runtime and adds nothing to the wheel. (An earlier revision of this branch usedptxas -c+cuLink+ a disk cubin cache; that requires a PTX assembler, which driver-only wheels and the GPU CI runners do not have -- hence the previous push'sptxas -c ... rc=32512failures.)What per-task buys us
Per-task reuse rides on the pre-existing cache tiers, now at task granularity because each task is a separate module:
PtxCache(PTX, keyed on LLVM-IR text) -- an unchanged task skips LLVM->PTX.~/.nv/ComputeCache, SASS, keyed on PTX hash) -- an unchanged task skipsptxasinside the driver.No new on-disk cache is introduced here; the cross-process per-task artifact cache is ref 6. 1b is the enabler: it makes each task an independently-compiled, independently-loadable unit that 6 / 8 / 10 build on.
Changes
runtime/cuda/jit_cuda.{h,cpp}: compositeJITModuleCUDA(vector<CUmodule>+ name->CUfunction memo);add_module_per_taskcompiles each task module to PTX and driver-loads it; deletes theptxas/cuLink/cubin-cache code.jit/jit_session.h,runtime/llvm/llvm_runtime_executor.{h,cpp},runtime/cuda/kernel_launcher.cpp: renameadd_module_culink/create_jit_module_culink->..._per_task.codegen/codegen.cpp,codegen/llvm/{codegen_llvm.cpp,llvm_compiled_data.h}: build one self-contained module per task intoper_construct_artifacts(PerConstructArtifactis{ module }; the full record moves to 6).codegen/cuda/codegen_cuda.cpp: emit static shared scratch andbls_bufferas internal definitions so each per-task module is self-contained.tests/python/test_atomic.py: widen f16 tolerance (per-task codegen changes atomic interleaving order).docs/source/user_guide/init_options.md: trim the offline-cache description.Testing
RTX PRO 6000 (sm_120), CUDA container, editable build:
test_atomic-- 96 passed.reduction graph struct_for offload_cross offline_cache-- 82 passed (exercises BLS / shared-memory reductions, CUDA-graph capture across the composite module, multi-offload struct-for, and offline-cache reuse).pre-commit runon the changed files -- clean.