research: cold JAX compile — autotuning is the cost; verdict 2 (#74)

jax_compile/README.md

@@ -145,5 +145,66 @@ one-time-per-machine cost the cache already removes. **Do not restructure.**
    reuse `probe.py` and should track *warm* compile times per cell so cache
    regressions are caught.
 
+## Cold-compile findings (issue #74, 2026-07-17)
+
+The persistent cache (above) solves repeat fits; these findings address the
+**cold** cost it cannot remove.
+
+### 6. The pathological cold compile IS GPU autotuning — one flag removes it
+
+The first run in this stack's history where `--xla_gpu_autotune_level=0`
+actually reached XLA (the pre-#128 wrapper clobbered every earlier attempt —
+including the 2026-07-15 "ruled out" A/B and finding 3's controlled A/B, which
+compared autotune-ON to autotune-ON):
+
+| A100, fresh cache | autotune ON | autotune OFF |
+|---|---|---|
+| FD probe total (kernel-CDF pixelized vag) | 498 s | **29 s** (17×) |
+| full 300×16 adam fit wall | 2081 s | **1253 s** (−40 %) |
+| fixed-input logL | 25536.848940 | 25536.848940 (bit-identical) |
+
+The `input_reduce_fusion` slow-compile alarm was autotune compiling candidate
+kernel configs of that fusion as standalone modules (why it never appeared as
+a dumpable module name). Autotune results cache like everything else, so with
+the cache enabled this is first-fit UX; without the flag a new machine's first
+pathological fit pays ~7m30.
+
+### 7. Steady-state eval does not need autotune (measured cells)
+
+A100 probe matrix, fresh caches per job (jobs 330601/330602):
+
+| cell | steady ON | steady OFF | compile ON | compile OFF |
+|---|---|---|---|---|
+| mge / jit | 0.0042 s | 0.0042 s | 10.1 s | 9.2 s |
+| mge / vag | 0.0096 s | 0.0098 s | 29.4 s | 29.3 s |
+| pix / jit | 0.0571 s | 0.0574 s | 8.5 s | 5.9 s |
+| pix / vag | 0.0871 s | 0.0910 s | 27.3 s | 19.7 s |
+
+Worst case ~4 % on one cell — and the 4800-eval full fit ran *faster* end to
+end with autotune off, so no real eval penalty is observed. (These standard
+cells compile in seconds either way; the 7m30 autotune cost is specific to the
+pathological kernel-CDF no-sparse-operator shape.)
+
+### 8. The tracing floor is jax-internal — no PyAuto lever
+
+cProfile attribution of a 40 s MGE `vag` trace (`trace_profile.py`): 58 % jax
+internals, 34 % stdlib/numpy, 7 % autoarray, ~0 % autofit/autogalaxy/autolens.
+Reducing it means emitting fewer ops (a jax-side concern), not optimizing
+PyAuto Python. Documented and closed as a direction.
+
+## Verdict 2 (cold compile)
+
+**One more setting: default `--xla_gpu_autotune_level=0`.** Evidence: 17×
+pathological cold-probe reduction, −40 % cold full fit, bit-identical fixed-
+input likelihoods, steady-state eval parity across the measured matrix. With
+the cache (#128) plus autotune-off, worst-case first-fit UX drops from ~70 min
+to ~30 s. Recommended as an env-respecting wrapper default (same pattern as
+the cache), so clusters can re-enable autotune where a tuned kernel matters.
+
+Remaining leads deprioritized by these numbers: cache-entry proliferation and
+pre-warming (cold is now ~seconds-to-a-minute); upstream XLA report (the "slow
+fusion" is explained — autotune candidates on a 58 GiB fusion; the HLO dump
+artifact from job 330596 exists if ever needed).
+
 Not indicated: source restructuring, jit boundaries inside likelihoods,
 replacing `lax.map` in MultiStartAdam, autotune flags.

jax_compile/results/local_cpu/trace_profile.json

@@ -0,0 +1,20 @@
+[
+  {
+    "model_type": "mge",
+    "transform": "vag",
+    "ndim": 15,
+    "trace_s": 39.822,
+    "rollup_tottime_s": {
+      "jax": 23.046,
+      "other": 13.66,
+      "autoarray": 2.803,
+      "autogalaxy": 0.123,
+      "autofit": 0.049,
+      "autolens": 0.01,
+      "autoconf": 0.006
+    },
+    "jax_version": "0.10.2",
+    "tag": "first",
+    "timestamp": "2026-07-17T08:26:15"
+  }
+]

jax_compile/results/local_gpu_NVIDIA_A100_80GB_PCIe/mge.json

@@ -0,0 +1,24 @@
+[
+  {
+    "transform": "jit",
+    "trace_s": 4.885,
+    "compile_s": 9.204,
+    "first_s": 1.329,
+    "steady_s": 0.0042,
+    "dataset_class": "imaging",
+    "model_type": "mge",
+    "instrument": "hst",
+    "ndim": 15,
+    "n_batch": 16,
+    "batch_size": 4,
+    "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe",
+    "jax_version": "0.10.2",
+    "cache_dir": "",
+    "mixed_precision": false,
+    "tag": "a100-at0",
+    "hostname": "euclid-ral-gpu-2",
+    "timestamp": "2026-07-17T08:55:59"
+  }
+]
+
+]

jax_compile/results/local_gpu_NVIDIA_A100_80GB_PCIe/pixelization.json

@@ -0,0 +1,82 @@
+[
+  {
+    "transform": "jit",
+    "trace_s": 4.189,
+    "compile_s": 5.875,
+    "first_s": 0.155,
+    "steady_s": 0.0574,
+    "dataset_class": "imaging",
+    "model_type": "pixelization",
+    "instrument": "hst",
+    "ndim": 12,
+    "n_batch": 16,
+    "batch_size": 4,
+    "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe",
+    "jax_version": "0.10.2",
+    "cache_dir": "",
+    "mixed_precision": false,
+    "tag": "a100-at0",
+    "hostname": "euclid-ral-gpu-2",
+    "timestamp": "2026-07-17T08:57:00"
+  },
+  {
+    "transform": "jit",
+    "trace_s": 4.254,
+    "compile_s": 8.497,
+    "first_s": 0.11,
+    "steady_s": 0.0571,
+    "dataset_class": "imaging",
+    "model_type": "pixelization",
+    "instrument": "hst",
+    "ndim": 12,
+    "n_batch": 16,
+    "batch_size": 4,
+    "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe",
+    "jax_version": "0.10.2",
+    "cache_dir": "",
+    "mixed_precision": false,
+    "tag": "a100-at_on",
+    "hostname": "euclid-ral-gpu-2",
+    "timestamp": "2026-07-17T08:57:03"
+  },
+  {
+    "transform": "vag",
+    "trace_s": 7.178,
+    "compile_s": 19.73,
+    "first_s": 0.727,
+    "steady_s": 0.091,
+    "dataset_class": "imaging",
+    "model_type": "pixelization",
+    "instrument": "hst",
+    "ndim": 12,
+    "n_batch": 16,
+    "batch_size": 4,
+    "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe",
+    "jax_version": "0.10.2",
+    "cache_dir": "",
+    "mixed_precision": false,
+    "tag": "a100-at0",
+    "hostname": "euclid-ral-gpu-2",
+    "timestamp": "2026-07-17T08:57:34"
+  },
+  {
+    "transform": "vag",
+    "trace_s": 7.272,
+    "compile_s": 27.256,
+    "first_s": 0.177,
+    "steady_s": 0.0871,
+    "dataset_class": "imaging",
+    "model_type": "pixelization",
+    "instrument": "hst",
+    "ndim": 12,
+    "n_batch": 16,
+    "batch_size": 4,
+    "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe",
+    "jax_version": "0.10.2",
+    "cache_dir": "",
+    "mixed_precision": false,
+    "tag": "a100-at_on",
+    "hostname": "euclid-ral-gpu-2",
+    "timestamp": "2026-07-17T08:57:43"
+  }
+]

jax_compile/trace_profile.py

@@ -0,0 +1,120 @@
+"""
+Where does JAX *tracing* time go? (cold-compile research, issue #74)
+
+Tracing (`jax.jit(fn).lower(x)`) is pure Python, recurs every process, and is
+the one cost the persistent compilation cache cannot remove (#71). This probe
+cProfiles a single `.lower()` call and aggregates cumulative time by library,
+so the uncacheable floor can be attributed (autofit model mapping? autoarray
+grids? jax machinery?) before anyone tries to reduce it.
+
+Usage (from the ``autolens_profiling/`` root)::
+
+    python jax_compile/trace_profile.py --model-type mge --transform vag
+    python jax_compile/trace_profile.py --model-type pixelization --transform jit --top 40
+
+Prints the per-library rollup and the top-N cumulative functions, and appends a
+JSON record under ``jax_compile/results/<hardware>/trace_profile.json``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import cProfile
+import json
+import pstats
+import sys
+import time
+from collections import defaultdict
+from pathlib import Path
+
+_WORKSPACE_ROOT = Path(__file__).resolve().parents[1]
+if str(_WORKSPACE_ROOT) not in sys.path:
+    sys.path.insert(0, str(_WORKSPACE_ROOT))
+
+from jax_compile.probe import build_objective, transformed_fn_and_arg, hardware_label
+
+LIBRARY_KEYS = (
+    "autofit",
+    "autoconf",
+    "autoarray",
+    "autogalaxy",
+    "autolens",
+    "jax",
+)
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description=__doc__.split("\n")[1])
+    p.add_argument("--dataset-class", default="imaging")
+    p.add_argument("--model-type", default="mge")
+    p.add_argument("--instrument", default="hst")
+    p.add_argument("--transform", default="vag")
+    p.add_argument("--n-batch", type=int, default=16)
+    p.add_argument("--batch-size", type=int, default=4)
+    p.add_argument("--top", type=int, default=25)
+    p.add_argument("--mixed-precision", action="store_true")
+    p.add_argument("--tag", default="")
+    return p.parse_args()
+
+
+def library_of(filename: str) -> str:
+    for key in LIBRARY_KEYS:
+        if f"/{key}/" in filename:
+            return key
+    return "other"
+
+
+def main():
+    args = parse_args()
+
+    import jax
+
+    f, x0, ndim = build_objective(args)
+    fn, arg = transformed_fn_and_arg(args.transform, f, x0, args.n_batch, args.batch_size)
+    jitted = jax.jit(fn)
+
+    profiler = cProfile.Profile()
+    t0 = time.perf_counter()
+    profiler.enable()
+    jitted.lower(arg)
+    profiler.disable()
+    trace_s = time.perf_counter() - t0
+
+    stats = pstats.Stats(profiler)
+
+    # Rollup: exclusive (tottime) per library — sums to the trace wall time, so
+    # the attribution is exact rather than double-counted cumulative time.
+    per_library = defaultdict(float)
+    for (filename, _, _), (_, _, tottime, _, _) in stats.stats.items():
+        per_library[library_of(filename)] += tottime
+
+    print(f"[trace_profile] {args.model_type} / {args.transform}: trace {trace_s:.2f}s")
+    print(f"[trace_profile] exclusive-time rollup (sums to ~trace time):")
+    rollup = dict(sorted(per_library.items(), key=lambda kv: -kv[1]))
+    for lib, seconds in rollup.items():
+        print(f"[trace_profile]   {lib:<10} {seconds:7.2f}s  ({100 * seconds / trace_s:4.1f}%)")
+
+    print(f"[trace_profile] top {args.top} functions by cumulative time:")
+    stats.sort_stats("cumulative").print_stats(args.top)
+
+    record = {
+        "model_type": args.model_type,
+        "transform": args.transform,
+        "ndim": ndim,
+        "trace_s": round(trace_s, 3),
+        "rollup_tottime_s": {k: round(v, 3) for k, v in rollup.items()},
+        "jax_version": jax.__version__,
+        "tag": args.tag,
+        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
+    }
+    out_dir = _WORKSPACE_ROOT / "jax_compile" / "results" / hardware_label(jax)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out_path = out_dir / "trace_profile.json"
+    existing = json.loads(out_path.read_text()) if out_path.exists() else []
+    existing.append(record)
+    out_path.write_text(json.dumps(existing, indent=2) + "\n")
+    print(f"[trace_profile] wrote record -> {out_path}")
+
+
+if __name__ == "__main__":
+    main()