warp-compile-time-optimizer

# Warp cold-start compile time ## Start with a runnable command The probe runs the target in a subprocess, writes only to temporary directories, and needs no network, external tool servers, or Warp checkout. | Script | Purpose | Arguments | | --- | --- | --- | | `scripts/warp_compile_probe.py` | Measure isolated cold/warm compilation and launches. | `measure [OPTIONS] -- COMMAND...`; use `--help`. | Use `run_script("scripts/warp_compile_probe.py", args=[...])` when supported; otherwise use the Python command below. The target must run to completion. ## Compilation model Warp compiles modules, not individual kernels. A module's identity is: ```text (live kernel & function set) x (module options) x (CUDA block_dim) x (generic instances) ``` Each identity requires code generation and native compilation for the full module. Cold-start cost is roughly: ```text number of distinct module identities you touch x size of each module ``` Reduce it in two ways: 1. Stop identity churn. A module that changes after loading compiles again. 2. Stop source duplication. A module used at three block dimensions compiles every kernel three times. Deleting one kernel from a module that still builds saves only part of one build. Removing an unnecessary module identity saves the full build. When neither applies, overlap independent CUDA builds (CS-13). This changes when work happens, not how much is compiled, so judge it on elapsed time. ## When a module's options are fixed Options have two deadlines: 1. **Module creation (Python import).** The module copies `enable_backward`, `max_unroll`, `lineinfo`, `deterministic`, `deterministic_max_records`, and `compile_time_trace` from `warp.config`. Setting a global later is silently ignored by that module. `default_grid_stride` is the exception. 2. **First load.** Before compilation, change the existing module with `wp.set_module_options()` or `wp.get_module(name).options`. Changing an option after load creates a new identity and rebuilds the module (CS-3). | Missed deadline | Symptom | Cost | | --- | --- | --- | | `wp.config.*` set after import | hash unchanged, option silently absent | the entire benefit, invisibly | | module options set after load | a second hash, module builds twice | one extra build, visible in the trace | After changing an option, confirm the hash moved for every target module. An unchanged hash means the option never arrived. ## Preserve behavior Change how Warp compiles the code, not the workload. Do not delete or merge kernels to claim a gain. Apparently redundant stages may preserve ownership, aliasing, retained outputs, numerical boundaries, or API behavior. Fix duplication at the module level. Preserve every launch and its order, dimensions, dtypes, devices, block dimensions, gradients, numerical modes, dynamic/plugin behavior, and public API signatures. Keep kernel names when moving definitions to module scope because logs, cache artifacts, and external tools expose them. ## Instructions ### 1. Find the real cost, and confirm it is compilation Ask what command the user actually waits on, then measure it cold: ```bash python scripts/warp_compile_probe.py measure --samples 3 \ --json baseline.json -- <the user's command> ``` The probe gives each sample private `WARP_CACHE_PATH`, `WARP_CACHE_ROOT`, and `CUDA_CACHE_PATH` directories, enables module timers, and records launches. It creates those directories under the system temporary directory and removes them itself, so isolating a sample never requires writing a cache into the project or deleting anything to re-measure. Isolate a hand-rolled sample the same way. Never clear a live cache with `wp.clear_kernel_cache()` or `wp.clear_lto_cache()`; clearing is not isolated and can disrupt other processes. Read the probe output before source. If compilation is a small part of wall time, report the real bottleneck and stop. For libraries and tests, use the smallest command that compiles the workload's modules. Modules that each compiled once, with no repeated hashes, block-dimension variants, or LTO, have no structural churn. This rules out redundant builds, not oversized builds; still check cache reuse (CS-2), backward codegen (CS-10), unrolling (CS-11), the precompiled header (CS-12), and overlap when several CUDA modules remain (CS-13). Every sample also re-runs the command against the cache it just populated. If warm module work is not near zero, diagnose cache reuse (CS-2) before changing module structure. ### 2. Ask about runtime tradeoffs when needed Apply ordering fixes, lifecycle grouping, and option hoists without asking. Ask before changing `fast_math`, `max_unroll`, or a MathDx/tile implementation: > Some of these knobs cut compile time but can make the compiled kernels > slower or change numerics. Are you optimizing a fast edit-run loop (where > slower kernels are usually fine), or production startup (where they usually > are not)? If the user is unavailable: - Leave numerics-changing options and implementation swaps alone. - Before removing a capability such as backward codegen, test whether it is used; report what is removed, the measured benefit, and how to revert. - Set global `wp.config.*` options at application entry points, not in library code. Record declined options and their measured benefits in the step 6 ledger. #### Match the scope of the change to the scope of the evidence A profile supports a change to the measured application, not every consumer of a shared library. Repository searches also miss out-of-tree and future callers. For example, a forward-only application does not justify disabling gradients inside a solver library that another application differentiates through. Scope the option to the measured process, before importing the library: ```python import warp as wp wp.config.enable_backward = False # must precede the library import import the_library ``` This also reaches every module the application loads. CS-10 covers the silent import-order trap. If only a library change works, send its maintainers the measurement and let them decide the contract. ### 3. Diagnose from the measurement, not from reading the source The probe prints every compiled module identity with its name, hash, device, and block dimension, then names which modules built more than once. Match what you see: | What the probe shows | What it means | Where to look | | --- | --- | --- | | One module name, several **hashes** | Identity churn: its kernel set, options, or generic instances changed after it first loaded | CS-1, CS-3, CS-6 | | One module name, several **block_dim** values | The whole module is recompiled per block dimension (CUDA) | CS-5 | | Many one-kernel modules in one feature | Fixed per-module cost repeated | CS-4 | | A hash-named module per kernel | `module="unique"` used on stable kernels | CS-9 | | Big gap between module time and native compile time, plus `.lto` artifacts | MathDx/LTO setup | CS-7 | | `(compiled)` on a run that should have been warm | Cache is not being reused | CS-2 | | Modules load, then "Failed to find module" | Concurrent CPU JIT first-use race | CS-8 | | Large generated source, no rebuild problem | Unroll budget | CS-11 | | Adjoint code in a module nothing differentiates | Backward codegen | CS-10 | | Compiles slow across the board, or a few small modules on CUDA below toolkit 13 | The precompiled header is turned off, or is not paying for itself | CS-12 | | Several independent modules, each built once, `overlap_factor` near 1.0 | Builds are running one at a time; parallel loading is off by default | CS-13 | | An option you set changed nothing, and that module's hash is unchanged | It was assigned after the module was created, so it never arrived | "When a module's options are fixed" | | No row above fires | Nothing is being built redundantly; the cost is the size of the builds themselves | Step 6 | `references/mechanisms.md` has one section per mechanism: how to confirm it, the fix, its limits, and its failure mode. Read only the sections selected by the measurement. ### 4. Choose module boundaries deliberately Group kernels in one module only when they share: - lifecycle: they are defined, loaded, and invalidated together; - option set: they need the same `fast_math`, `enable_backward`, `max_unroll`, and MathDx settings; - stable block dimension on CUDA. Kernels with the same lifecycle but different stable block dimensions should not share a module because each would compile twice. Separate kernels with independent lifecycles too. Kernels whose block dimension varies at runtime (chosen from input size, say) have no stable mapping, so keep them in their own module rather than dragging a whole shared module into an extra variant. Prefer the least invasive change that removes a build. Ordering fixes and option hoists are cheaper and safer than re-architecting module layout; regroup only when fixed per-module cost or block-dimension duplication dominates. `wp.set_module_options()` targets its calling Python module, not kernels with an explicit `module="pkg.name"`. Either use a real Python module or update the named module before it loads: ```python wp.set_module_options({"enable_backward": False}) # at module scope wp.get_module("pkg.name").options.update({"enable_backward": False}) ``` Do not pass `wp.get_module()` to `wp.set_module_options(module=...)`, or use `@wp.kernel(module_options={...})` without `module="unique"`. Per-kernel `enable_backward=False` has a tile-module exception covered by CS-10. Confirm the module hash after every option change. ### 5. Verify ```bash python scripts/warp_compile_probe.py measure --samples 3 \ --json candidate.json -- <the same command> python scripts/warp_compile_probe.py compare baseline.json candidate.json ``` `compare` rejects changed launch topology and treats a result inside `max(1% of baseline, 2 x baseline MAD)` as inconclusive. For `BUILDS OVERLAPPED`, judge scheduling changes on compile elapsed rather than summed module timers. The required warm pass supplies that clock. See `references/measurement.md`. Then check what the probe cannot see: - Diff numeric output and run the project's tests or entry points. - After changing `enable_backward` or boundaries, verify a gradient path. - After changing `fast_math`, `max_unroll`, MathDx, or an implementation, benchmark steady-state runtime. - Exercise uncovered dynamic kernels, dtypes, profiles, and modes. ### 6. Report results Read "Reporting results" in `references/measurement.md`. Report: Describe every optimization in plain language: name the behavior, the evidence, and the effect. For example, write "moved module options before the first load to avoid a redundant rebuild," not "applied CS-3." Treat `CS-*` labels as internal navigation aids, not user-facing explanations. | Option | Measured | Why not taken | To take it | | --- | ---: | --- | --- | | `enable_backward=False` on `pkg.solver` | −38% cold | a live tape traverses these kernels | set at the entry point, then re-check adjoints | | `max_unroll=4` | −2%, inside noise | changes generated code for no measured gain | — | - before/after medians and sample counts; - the reduction and residual cost, sized against the original complaint; - mechanisms fixed, ruled out, measured and declined, or not reached; - tradeoffs and behavior not verified; - the ledger above for declined or incompletely verified options; - the measured value of relaxing any constraint that blocked a fix. State only what the evidence supports. "No structural churn" does not mean "optimal" or "irreducible." Measure declined levers when practical; label any estimate untested. Before reporting no available fix, check CS-13. For a possible module split, first measure a one-kernel module with the same options to establish the repeated fixed cost. ## Troubleshooting Run the target command directly before debugging the probe. See `references/measurement.md` for cache/noise issues and `references/mechanisms.md` for mechanism-specific failures. ## Limitations Two rules override any gain: - Isolate both Warp and CUDA caches for every cold sample. - Keep `max_workers <= 1` when a load can target CPU, including `device=None` and mixed device lists. Concurrent CPU first loads can lose kernels; retries do not make them safe. CUDA-only loading is unaffected. Measurements are environment-specific: cold times move with CPU, GPU, driver, toolchain, and Warp version. Mechanisms transfer; numbers do not. Read the known unknowns in `references/mechanisms.md` before making broad claims. ## Reference files - `references/mechanisms.md`: the thirteen compile-time mechanisms, each with its confirming signal, fix, applicability limits, and failure mode. Read the sections your measurement points to. - `references/measurement.md`: measurement protocol, what each metric does and does not mean, reporting guidance, log examples, and manual measurement.