[Caching] C5 Split the frontend per top-level construct - #880
[Caching] C5 Split the frontend per top-level construct#880hughperkins wants to merge 30 commits into
Conversation
Run the remaining pre-offload + offload frontend (simplify / merge_global_ptrs / offload) PER top-level construct and reassemble, instead of once over the whole kernel, so a later cross-process cache can key and reuse each construct's frontend output independently. Each construct is isolated by its BACKWARD SLICE (operand-def chains recomputed into the construct, cross-construct global-temp hubs dissolved), preserving source order so cross-construct memory ordering holds. Restricted to recompute-safe kernels (autodiff_mode==kNone, non-mesh, recompute-safe); anything else falls back to the whole-kernel path. This PR ships the split with NO reuse tier: every construct is recompiled on every compile (recompiled == total, cache_hit == 0). It defines the construct cache key (get_hashed_per_construct_cache_key) and the observability surface (PerTaskCacheStats construct counts, get_per_task_cache_stats, PerOffloadCacheObservations.frontend_constructs_*, CompileResult per_construct_*), surfaced as kernel._primal.per_offload_cache_observations so the split's structure can be asserted on counts. The cross-process cache PR consumes the key and adds the per-task reuse counts on this surface. Supporting IR helpers: clone_block_subset (clone a chosen top-level subset of a block) and gather_statements' include_containers option.
The launcher consumes a CLONE of the compiled kernel, and LLVMCompiledKernel::clone() rebuilds via the 2-arg constructor, which does not copy per_task_cache_stats -- so the construct-count observability (frontend_constructs_*) was reset to its -1 sentinel before reaching Python. The split itself ran and produced correct IR (an after_offload dump confirms the expected per-construct tasks); only the stats readback was lost on the clone. Copy the field in clone(), alongside the existing per_construct_artifacts copy, so the host sees the real counts. Also sync program.h's per_construct_cache() comments to the reduced "split stats record" role (no reuse tier in this PR).
![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: 335acb9ed8
ℹ️ 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".
Addresses the Codex review on #880. Contact area (AGENTS.md): move the per-construct frontend split -- segmentation, backward slice, recompute-safety gate, and orchestration -- out of the central compile_to_offloads pass into a new transforms/split_frontend_per_construct.cpp that exposes a single irpass::maybe_split_frontend_per_construct(). The central pass now keeps only the one gated call. Recompute-safety gate, two soundness fixes: - Loop-carried local shared across constructs. The old check compared readers against the UNION of writer constructs, so two constructs that both read-modify-write the same local were each "covered" by the union and wrongly accepted -- the second still consumes the value the first produced. Now a loop-produced local is safe only when every access (read or loop-write) stays within a single construct. - Field load snapshotted before a later construct. The backward slice recomputes a serial field/ndarray load into every consuming construct; if a construct between the load and its consumer writes global memory, the recomputed load observes the mutation instead of the source-order snapshot (e.g. base = x[0]; for: x[0] = 2; for i: y[i] = base). Now falls back when a recomputed load is shadowed by an intervening global write. Field identity is not tracked, so this is conservative. Tests: add fallback + correctness regression tests for both miscompiles (test_..._fallback_field_load_shadowed, test_..._fallback_carried_rmw_local).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 74a8a2a89f
ℹ️ 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".
…ructs Second round of Codex review on #880. The backward slice pulls the writers (and their operand chains) of any local a construct reads. When that producer is effectful -- e.g. a local capturing the return of a global AtomicOpStmt -- it already emits its own task in its home segment, so cloning it into a consuming construct's slice runs the effect a second time (e.g. old = atomic_add(counter, 1); for i: out[i] = old would increment counter twice and store the second atomic's return). The previous recompute-safety fixes only covered pure field loads, not effectful producers. split_is_recompute_safe now falls back whenever a construct's backward slice would recompute any side-effecting statement from a different segment. Test: add test_per_construct_frontend_split_fallback_effectful_producer (asserts fallback, plus the atomic runs exactly once: counter == 1, out == 0).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f08c9dcb67
ℹ️ 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".
…ects Third round of Codex review on #880. - RandStmt: qd.random() reports no global side effect, so the effectful-producer check missed it, and each consuming construct's backward slice cloned another PRNG draw (r = random(); for i: x[i] = r; for i: y[i] = r gave the two loops different samples and advanced the PRNG twice). split_is_recompute_safe now also rejects recomputing a RandStmt from another segment. - Sparse ops: the intervening-write check recognized only GlobalStoreStmt and global atomics, so a top-level qd.deactivate/activation between a snapshot load and its consumer was not treated as a mutation and the recomputed load could observe the deactivated (default) value. Switch that check to stmt_is_task_effect, which already covers SNodeOpStmt (activate/deactivate inherit has_global_side_effect == true) plus stores, atomics, and any other real effect. Tests: add test_..._fallback_random_producer and test_..._fallback_sparse_deactivate_shadow (both assert fallback + correctness).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 34e6093eb2
ℹ️ 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".
Move the per-construct frontend-split observability readback out of the LLVM codegen driver and into `KernelCompilationManager::load_or_compile`, reading the counts directly off the program-scoped `PerConstructCache:: last_stats`. This fixes two P2 issues: - In-memory cache hits reused the producer compile's counters, so an equivalent kernel reported nonzero `frontend_constructs_recompiled` even though `cache_hit` was true and no frontend ran; disk hits reported -1. Now stats are read only on a fresh compile (and the entry consumed), so every cache tier reports the no-split sentinel uniformly. - SPIR-V (Vulkan/Metal) never overrode `get_per_task_cache_stats`, so the split's counts were always -1 there even when it ran. Reading from the program-scoped record makes the observability backend-agnostic. Removes the now-unused `PerTaskCacheStats` struct, the virtual `get_per_task_cache_stats` accessor and its LLVM override, the `per_task_cache_stats` member and its clone() copy.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d195db4f2d
ℹ️ 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".
`per_offload_cache_observations` is a single per-Kernel attribute. When `launch_kernel` is handed a ready `compiled_kernel_data` -- a fastcache restore (`_try_load_fastcache`), which bypasses `prog.compile_kernel`, or an already-compiled specialization on the same Kernel object -- it skipped the block that sets the observations, leaving a previous specialization's nonzero split counts visible for a compile that ran no frontend. Reset the observations to the no-split sentinel (-1) on that path so cache-served launches report consistently with the C++ cache-hit path in `KernelCompilationManager::load_or_compile`. Also document the per-construct frontend split and the `per_offload_cache_observations` API in docs/ (optimization_passes.md), per the repo policy to keep user-facing docs in sync with public API. Adds a regression test asserting the counts reset to -1 on a cache-served relaunch.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: db614e11eb
ℹ️ 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".
Two more cases the per-construct split must not handle: - Volatile loads. `qd.volatile_load` lowers to a `GlobalLoadStmt` with `is_volatile`, which the backward slice was free to recompute like any pure load. Volatile semantics require the access to happen exactly once in place, so cloning it into several consuming constructs turns one read of a concurrently-updated cell into several. The recompute-safety gate now rejects recomputing a volatile load from another segment (alongside effects and `RandStmt`). - Concurrently-executed regions. Loops inside `qd.stream_parallel()` / `qd.graph.parallel()` run on separate streams while sharing one kernel context and its single global-temporary buffer. The per-construct split runs `offload` independently per construct, restarting global-temp offset allocation at 0 each time, so two concurrent constructs (e.g. dynamic loop bounds routed through a global temp) would alias the same offset and race. `maybe_split_frontend_per_construct` now falls back when any loop carries a nonzero `stream_parallel_group_id` / `graph_parallel_region_id`. Adds regression tests for both fallbacks.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1bd513da4c
ℹ️ 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".
`get_hashed_per_construct_cache_key` (and its `gather_ir_snode_roots`
helper) has no caller in this PR: the split ships with no reuse tier, so
nothing keys a construct yet. Adding ~110 lines of SNode gathering and
hashing to the shared `offline_cache_util` only to sit unused expands the
contact area of that shared utility for a future feature, against
AGENTS.md ("minimize contact area").
Remove the function, its helper, and the includes they alone pulled in
(`ir/analysis.h`, `ir/offloaded_task_type.h`, `ir/statements.h`, `<map>`).
The cross-process cache PR -- the actual consumer -- reintroduces the key,
ideally in its own feature file. Reword the two comments that named the
removed symbol to describe it generically.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d61d991ebd
ℹ️ 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".
qd.init(print_ir=True) surfaces as `verbose` in compile_to_offloads and prints a whole-kernel IR snapshot before/after every pipeline stage. The per-construct split runs those stages per construct, so print_ir would emit one isolated construct per snapshot instead of the documented whole-kernel console dump. Fall back to the whole-kernel path when print_ir/verbose is set, consistent with the QD_DUMP_CFG/QD_DUMP_IR/QD_KERNEL_COVERAGE diagnostic gates.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 760d7d5157
ℹ️ 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".
The per-construct split's stats flow is already documented at both feature-local sites (split_frontend_per_construct.cpp where they're recorded and kernel_compilation_manager.cpp where they're read back). Remove the redundant explanatory comment so this PR has zero footprint in the shared codegen path, keeping the experimental feature's contact area minimal (AGENTS.md).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 41f58cd221
ℹ️ 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".
QD_DUMP_SIMPLIFY is a whole-kernel diagnostic (simplify.cpp) that snapshots IR after every pass inside each full_simplify loop. Under the split full_simplify runs on a per-construct block, so every snapshot shows only one isolated construct and no dump captures the whole kernel. Add DUMP_SIMPLIFY_ENV to the diagnostic fallback alongside QD_DUMP_CFG/QD_DUMP_IR so the documented whole-kernel dumps are preserved.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e4d55a83ca
ℹ️ 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_to_offloads.cpp uses no Program symbol after the split was extracted, so program.h is a dead feature-only dependency; and <string> was redundant (the file already used std::string via transitive includes before this PR). Remove both so the central pass's only diff is the functional split hook, keeping its contact area minimal (AGENTS.md).
|
@codex review |
This PR had inserted a blank-line-only change into a shared header. Remove it so the header is untouched, keeping the feature's contact area minimal.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b52fc1f55e
ℹ️ 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".
run_construct_frontend ran the frontend stages without the verify_if_debug calls the whole-kernel path makes after each stage, so under debug=True malformed IR from a split-path pass would only surface at codegen. Mirror the whole-kernel verify points after simplify_I, check_out_of_bound, merge_global_ptrs, flag_access, simplify_II, offload and the trailing simplify_III. verify_if_debug is a no-op unless config.debug, so release builds are unaffected.
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep them coming! 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". |
| - `QD_DUMP_IR=1` - writes an IR snapshot at each major pipeline stage (after lowering, before/after each simplify, after offload). | ||
| - `QD_DUMP_SIMPLIFY=1` - writes an IR snapshot after every individual pass on every iteration of the simplify loop. Verbose, but it shows exactly which pass changed what. | ||
| - `QD_DUMP_CFG=1` - writes the control-flow graph itself. (This also forces the CFG pass back onto the whole-kernel path so the complete graph can be dumped.) | ||
| - `QD_DUMP_CFG=1` - writes the control-flow graph itself. (This also forces the CFG pass to run over the whole kernel at once so the complete graph can be dumped.) |
There was a problem hiding this comment.
why not just dump each offload to seapraet file?
There was a problem hiding this comment.
Discssused with agent that QD_DUMP_CFG=1 should not chagne behavior.
There was a problem hiding this comment.
actually, discuss ing
There was a problem hiding this comment.
so apparenlty QD_DUMP_CFG is pre-existing, and what we are changing in this PR is QD_DUMP_IR and similar; so we'll make only QD_DUMP_IR and similar not change behavior, in this PR; and we'll leave QD_DUMP_CFG to another PR.
|
|
||
| This section is for the curious; you never have to think about it to write kernels - the behavior below is transparent and on by default. | ||
|
|
||
| The frontend stages above - the passes that turn your high-level kernel into offloaded tasks - can run either once over the whole kernel or, for eligible kernels, separately for each **top-level construct** (each independent top-level loop or serial run in your kernel). Compiling each construct in isolation is what will let a future cross-process cache reuse the unchanged constructs of a kernel you edited; today it produces the same offloaded tasks and the same results as the whole-kernel path, so the split is transparent. |
There was a problem hiding this comment.
dont talk about the future. just drop discussion of the future.
|
|
||
| The frontend stages above - the passes that turn your high-level kernel into offloaded tasks - can run either once over the whole kernel or, for eligible kernels, separately for each **top-level construct** (each independent top-level loop or serial run in your kernel). Compiling each construct in isolation is what will let a future cross-process cache reuse the unchanged constructs of a kernel you edited; today it produces the same offloaded tasks and the same results as the whole-kernel path, so the split is transparent. | ||
|
|
||
| Quadrants automatically falls back to the whole-kernel path whenever per-construct compilation would not be equivalent: [autodiff](autodiff.md) kernels, certain specialized kernels, and kernels where one construct's value depends on state another construct produced in a way that cannot be recomputed in isolation (for example a local variable that one top-level loop builds up over its iterations and another construct then reads, or a snapshot of a field that a later construct reads after an intervening write). |
There was a problem hiding this comment.
can you give an example of the latter?
…under the split These frontend diagnostics no longer disable the per-construct split; instead the split stays active and their output is emitted per construct, so the dump reflects what was actually compiled rather than perturbing compilation: - QD_DUMP_IR: run_construct_frontend writes <kernel>_construct<i>_<stage>.ll snapshots (construct-indexed so they don't collide), the per-construct counterpart of compile_to_offloads' whole-kernel dumps. - QD_DUMP_SIMPLIFY: simplify.cpp already keys its dump filenames off a global call counter, so each construct's passes land in distinct files with no change needed. - print_ir/verbose: each construct's passes print under a "[per-construct frontend split] <kernel> construct <i>" banner. QD_DUMP_CFG stays gated (it forces the whole-kernel path in cfg_optimization and names dumps by phase not construct, so making it observation-only needs a cfg_optimization rework - deferred to a separate PR). QD_KERNEL_COVERAGE stays gated (correctness, not observability). Docs and tests updated accordingly.
simplify.cpp's QD_DUMP_SIMPLIFY writer does not create debug_dump_path itself (it relies on /tmp/ir already existing in normal use), so the test's rmtree left no directory for it to write into and the glob found no files. Recreate the dump directory after clearing it in both the QD_DUMP_IR and QD_DUMP_SIMPLIFY tests. Cluster-verified: 37 passed, 2 skipped.
|
@codex review |
|
Codex Review: Didn't find any major issues. Chef's kiss. 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". |
1 participant
Summary
Run the remaining pre-offload + offload frontend (
simplify/merge_global_ptrs/offload) per top-level construct and reassemble, instead of once over the whole kernel, so a later cross-process cache can key and reuse each construct's frontend output independently.clone_block_subsetandgather_statements'include_containersopt-in.autodiff_mode == kNone, non-mesh, recompute-safe); anything else falls back to the whole-kernel path.recompiled == total,cache_hit == 0). It defines the construct cache key (get_hashed_per_construct_cache_key) and the construct-level observability surface (PerOffloadCacheObservations.frontend_constructs_*, surfaced askernel._primal.per_offload_cache_observations) so the split's structure can be asserted on counts. The cross-process cache PR consumes the key and adds per-task reuse counts on this surface.LLVMCompiledKernel::clone()(the launcher consumes a clone), so the counts actually reach Python.Test plan
Validated on the cluster (CUDA), incremental build off current
main:test_per_offload_cacheunit tests: 6 passed (3 cases x cpu+cuda) -- split enumerates all constructs (frontend_constructs_total == 4,recompiled == total,hit == 0), a shared serial-def kernel stays correct through backward-slice recompute, and a not-recompute-safe kernel falls back (total == -1).offload offload_cross ndrange ndrange_axes struct_for struct_for_dynamic struct_for_intermediate bls bls_assume_in_range sparse_basics sparse_activate sparse_deactivate sparse_parallel bitmasked mesh matrix dynamic-> 193 passed, 72 skipped, 0 failed.Made with Cursor